diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 926b6dac98f8..6146111738cb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1467,7 +1467,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|^tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/|^tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt b/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt index 2c4a27545e16..e5bc72b46f7d 100644 --- a/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/moe/CMakeLists.txt @@ -20,6 +20,7 @@ # into kernels_src, exactly as it did before the move. add_subdirectory(cuteDsl) # cute_dsl_src add_subdirectory(trtllmGen) # trtllm_gen_fp8_block_scale_moe{,_routing} +add_subdirectory(loadBalance/dynamicEplb) # moe_dynamic_eplb_src # cutlass/: moe_gemm_src and its per-arch OBJECT libraries are defined in # ../cutlass_kernels/CMakeLists.txt, next to the CUTLASS instantiation generator # they depend on. @@ -28,7 +29,9 @@ add_subdirectory(trtllmGen) # trtllm_gen_fp8_block_scale_moe{,_routing} # kernels_src, so the backend directories are named here, next to the # add_subdirectory() calls they correspond to, and nowhere else. file(GLOB_RECURSE MOE_SRC *.cpp *.cu) -list(FILTER MOE_SRC EXCLUDE REGEX "/kernels/moe/(cutlass|trtllmGen|cuteDsl)/.*") +list(FILTER MOE_SRC EXCLUDE REGEX + "/kernels/moe/(cutlass|trtllmGen|cuteDsl|loadBalance/dynamicEplb)/.*") +list(APPEND MOE_SRC $) set(MOE_KERNELS_SRC ${MOE_SRC} PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt new file mode 100644 index 000000000000..d2ada7394c9c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/CMakeLists.txt @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +set(DYNAMIC_EPLB_ARCHS) +foreach(ARCH 100 103 107) + if("${ARCH}" IN_LIST CMAKE_CUDA_ARCHITECTURES_ORIG) + list(APPEND DYNAMIC_EPLB_ARCHS ${ARCH}) + endif() +endforeach() + +set(DYNAMIC_EPLB_SOURCES) +set(DYNAMIC_EPLB_DEFINITIONS) + +if(DYNAMIC_EPLB_ARCHS) + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalanceHaloQ.cu) +else() + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalancePlaceholder.cpp) + list(APPEND DYNAMIC_EPLB_DEFINITIONS TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_HALO_Q) + message( + STATUS + "Per-layer dynamic EPLB HALO-Q disabled: none of SM100, SM103, or SM107 was requested" + ) +endif() + +if(DYNAMIC_EPLB_ARCHS AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL + 13.1) + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalanceTma.cu) +else() + list(APPEND DYNAMIC_EPLB_SOURCES moeRebalancePlaceholder.cpp) + list(APPEND DYNAMIC_EPLB_DEFINITIONS TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_TMA) + if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.1) + message( + STATUS + "Per-layer dynamic EPLB TMA disabled: CUDA 13.1 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION})" + ) + else() + message( + STATUS + "Per-layer dynamic EPLB TMA disabled: none of SM100, SM103, or SM107 was requested" + ) + endif() +endif() + +list(REMOVE_DUPLICATES DYNAMIC_EPLB_SOURCES) +add_library(moe_dynamic_eplb_src OBJECT ${DYNAMIC_EPLB_SOURCES}) + +if(DYNAMIC_EPLB_ARCHS) + # These kernels use SM100-family PTX. The family target avoids inheriting the + # generic kernels_src SM80/SM90 compilation passes. + set_cuda_architectures(moe_dynamic_eplb_src 100f) +endif() + +if(DYNAMIC_EPLB_DEFINITIONS) + target_compile_definitions(moe_dynamic_eplb_src + PRIVATE ${DYNAMIC_EPLB_DEFINITIONS}) +endif() + +set_property(TARGET moe_dynamic_eplb_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET moe_dynamic_eplb_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu new file mode 100644 index 000000000000..59d440ba29f3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.cu @@ -0,0 +1,46 @@ +/* + * 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/assert.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params) +{ + TLLM_CHECK_WITH_INFO((params.routes != nullptr || params.validRouteCount == 0) && params.outSlots != nullptr + && params.outIds != nullptr && params.outLevels != nullptr && params.outOwners != nullptr + && params.peerBases != nullptr && params.status != nullptr && params.partial != nullptr + && params.routeAux != nullptr && params.gridSync != nullptr && params.planWorkspace != nullptr + && params.routePrefix != nullptr && params.ep >= 2 && params.ep <= kMaxEp && params.ep % 2 == 0 + && params.experts > 0 && params.experts <= kMaxExperts && params.experts % params.ep == 0 + && params.helpers > 0 && params.helpers <= INT_MAX / params.ep - params.experts / params.ep + && params.localRank >= 0 && params.localRank < params.ep && params.validRouteCount >= 0 + && params.validRouteCount <= params.routeCapacity && params.routeCapacity > 0 + && params.routeCapacity <= INT_MAX / params.ep && params.routeCapacity <= INT_MAX - params.ctas + && params.ctas > 0 && params.ctas <= 128 && params.threads == 512, + "Invalid HALO-Q scheduler launch parameters"); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h new file mode 100644 index 000000000000..01901208d624 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h @@ -0,0 +1,63 @@ +/* + * 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. + */ + +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H +#define TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H + +#include "tensorrt_llm/common/config.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +// The HALO-Q scheduler phase of the fused scheduler + in-switch copy kernel +// (megamoe_tma_copy_submit_fused). +struct MoeRebalanceHaloQParams +{ + int const* routes; + int* outSlots; + int* outIds; + int* outLevels; + int* outOwners; + std::uint64_t const* peerBases; + int* status; + int* partial; + int* routeAux; + int* gridSync; + int* planWorkspace; + int* routePrefix; + int ep; + int experts; + int helpers; + int localRank; + int routeCapacity; + int ctas; + int threads; + std::uint64_t spinCycles; + int validRouteCount; +}; + +// Throws unless params describe a valid scheduler launch. +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const& params); + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_H diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh new file mode 100644 index 000000000000..3657a67671c6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh @@ -0,0 +1,1914 @@ +/* + * 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. + */ + +/* HALO-Q scheduler device code, the first phase of the fused scheduler + + * in-switch copy kernel. */ +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH +#define TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ +constexpr int kMaxEp = 32; +constexpr int kMaxExperts = 384; +constexpr int kMaxBroadcasts = kMaxExperts; +constexpr int kHelperMaskWords = (kMaxBroadcasts + 31) / 32; +constexpr int kMinGroup = 4; +/* Rank loops run in fully unrolled chunks of this many predicated ranks, so + * every load of a chunk is independent and in flight together. */ +constexpr int kRankChunk = 8; +constexpr int kSymPayloadOffset = 64; +using Epoch = std::uint32_t; + +enum PlanField +{ + kPlanExpert = 0, + kPlanLevel = 1, + kPlanGroupBegin = 2, + kPlanGroupSize = 3, + kPlanOwner = 4, + kPlanHelper = 5, +}; + +__device__ __forceinline__ int imin(int a, int b) +{ + return a < b ? a : b; +} + +__device__ __forceinline__ int imax(int a, int b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ Epoch fetch_add_acq_rel_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.fetch_add(value, cuda::memory_order_acq_rel); +} + +__device__ __forceinline__ Epoch load_acquire_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_device(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_relaxed_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_device(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ Epoch load_acquire_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_acquire); +} + +__device__ __forceinline__ Epoch load_relaxed_system(int* address) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + return atom.load(cuda::memory_order_relaxed); +} + +__device__ __forceinline__ void store_release_system(int* address, Epoch value) +{ + auto* valueAddress = reinterpret_cast(address); + cuda::atomic_ref atom(*valueAddress); + atom.store(value, cuda::memory_order_release); +} + +__device__ __forceinline__ bool epoch_before(Epoch observed, Epoch expected) +{ + // RFC-1982-style serial comparison: all arithmetic remains unsigned and + // is unambiguous while peers differ by fewer than half the epoch space. + return observed != expected && expected - observed < (Epoch{1} << 31); +} + +__device__ __forceinline__ void record_status(int* address, Epoch epoch) +{ + // Zero is reserved for "never failed". Preserve a nonzero sticky error at + // the one modular epoch whose bit pattern is zero. + Epoch const value = epoch == 0 ? ~Epoch{0} : epoch; + atomicExch(reinterpret_cast(address), value); +} + +__device__ __forceinline__ int* plan_field(int* plan, int field) +{ + return plan + 2 + field * kMaxBroadcasts; +} + +__device__ __forceinline__ bool group_contains(int rank, int begin, int size) +{ + return rank >= begin && rank < begin + size; +} + +__host__ __device__ __forceinline__ unsigned rank_group_mask(int begin, int size) +{ + return size == 32 ? 0xffffffffu : ((1u << size) - 1u) << begin; +} + +__host__ __device__ __forceinline__ int hierarchy_depth(int ep) +{ + int depth = 1; + for (int level = 1; (1 << level) <= ep; ++level) + { + int groups = 1 << level; + if (ep % groups || ep / groups < kMinGroup) + break; + ++depth; + } + return depth; +} + +/* Distinct experts and total receive capacity independently bound broadcasts. + * Actual S is retained in the public outputs; only scratch uses this bound. */ +__host__ __device__ __forceinline__ int broadcast_capacity(int ep, int experts, int helpers) +{ + int const bounded_helpers = helpers < experts ? helpers : experts; + /* Every hierarchy level divides ep exactly. This cancels two integer + * divisions without changing capacity, including non-power-of-two EP. */ + int const groups = 1 << (hierarchy_depth(ep) - 1); + int const receive_bound = groups * bounded_helpers; + return receive_bound < experts ? receive_bound : experts; +} + +__device__ bool grid_rendezvous( + int* sync, int blocks, unsigned long long spin_cycles, int* status, Epoch* completedEpoch) +{ + __shared__ Epoch generation; + __shared__ int success; + /* The block barrier orders every writer's stores before thread zero's + * device-scope release advertises this CTA. */ + __syncthreads(); + if (threadIdx.x == 0) + { + success = 1; + Epoch const current = load_relaxed_device(sync + 1); + Epoch const target = current + 1; + generation = target; + Epoch const old = fetch_add_acq_rel_device(sync, 1); + if (old == static_cast(blocks - 1)) + { + // The last CTA resets arrivals before publishing the next modular + // epoch. Stream ordering prevents a new launch from entering until + // every CTA in this launch has observed the release. + store_relaxed_device(sync, 0); + store_release_device(sync + 1, target); + } + else if (old >= static_cast(blocks)) + { + record_status(status + 1, target); + success = 0; + } + else + { + unsigned long long const start = clock64(); + Epoch observed = load_relaxed_device(sync + 1); + while (epoch_before(observed, target)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 1, target); + success = 0; + break; + } + observed = load_relaxed_device(sync + 1); + } + if (!epoch_before(observed, target)) + (void) load_acquire_device(sync + 1); + } + } + __syncthreads(); + *completedEpoch = generation; + return success != 0; +} + +__device__ void local_histogram_pass(int const* routes, int* partial, int route_count, int expert_count, int ctas) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shards = warps > 8 ? 8 : warps; + int const shared_count = shards * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const cta_span = (route_count + ctas - 1) / ctas; + int const cta_begin = blockIdx.x * cta_span; + int const cta_end = imin(route_count, cta_begin + cta_span); + int const warp_span = ((cta_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = cta_begin + warp * warp_span; + int const end = imin(cta_end, begin + warp_span); + /* Issue a small load group before consuming it to expose memory-level + * parallelism. Batching does not change which bin a route lands in, and + * atomicAdd order does not matter because integer addition commutes. */ + constexpr int kBatch = 4; + for (int base_index = begin; base_index < end; base_index += 32 * kBatch) + { + int expert[kBatch]; +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + expert[step] = index < end ? routes[index] : -1; + } +#pragma unroll + for (int step = 0; step < kBatch; ++step) + { + int const index = base_index + step * 32 + lane; + if (index < end) + { + int const value = expert[step]; + int const bin = (value >= 0 && value < expert_count) ? value : expert_count; + atomicAdd(shared + (warp & 7) * bins + bin, 1); + } + } + } + __syncthreads(); + + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int shard = 0; shard < shards; ++shard) + count += shared[shard * bins + expert]; + partial[blockIdx.x * expert_count + expert] = count; + } +} + +/* Re-rank all local routes on the C-1 worker CTAs while CTA0 performs the + * cross-rank exchange and planner. route_aux[0:2] is a reusable worker + * rendezvous; the following NWORK*(E+1) entries hold per-worker totals. + * Per-warp cursors remain in this CTA's shared memory until materialization. */ +__device__ bool stable_worker_ordinal_pass(int const* routes, int* ordinal, int* route_aux, int* status, + int route_count, int expert_count, int worker, int workers, unsigned long long spin_cycles) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = expert_count + 1; + int const shared_count = warps * bins; + for (int index = tid; index < shared_count; index += blockDim.x) + shared[index] = 0; + __syncthreads(); + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int base_index = begin; base_index < end; base_index += 32) + { + int const index = base_index + lane; + bool const active = index < end; + int const expert = active ? routes[index] : -1; + // Padding has its own key so it cannot contribute to real/invalid bins. + int const bin = !active ? bins : ((expert >= 0 && expert < expert_count) ? expert : expert_count); + unsigned const peers = __match_any_sync(0xffffffffu, bin); + int const leader = __ffs(peers) - 1; + int const rank = __popc(peers & ((1u << lane) - 1u)); + int base = 0; + if (active && lane == leader) + base = atomicAdd(shared + warp * bins + bin, __popc(peers)); + // One collective for the entire warp; each lane selects its own leader. + base = __shfl_sync(0xffffffffu, base, leader); + if (active) + ordinal[index] = base + rank; + } + __syncthreads(); + + int* worker_totals = route_aux + 2; + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int total = 0; + for (int local_warp = 0; local_warp < warps; ++local_warp) + total += shared[local_warp * bins + bin]; + worker_totals[worker * bins + bin] = total; + } + + /* The first all-CTA rendezvous guarantees every worker is resident, so a + * worker-only global barrier cannot wait on an unscheduled CTA. */ + Epoch workerEpoch; + if (!grid_rendezvous(route_aux, workers, spin_cycles, status, &workerEpoch)) + return false; + + for (int bin = tid; bin < bins; bin += blockDim.x) + { + int cursor = 0; + for (int previous = 0; previous < worker; ++previous) + cursor += worker_totals[previous * bins + bin]; + for (int local_warp = 0; local_warp < warps; ++local_warp) + { + int index = local_warp * bins + bin; + int count = shared[index]; + shared[index] = cursor; + cursor += count; + } + } + __syncthreads(); + return true; +} + +__device__ bool exchange_counts(unsigned long long const* peer_bases, int const* partial, int* shared, int ep, + int expert_count, int local_rank, int ctas, Epoch epoch, unsigned long long spin_cycles, int* status) +{ + int const tid = threadIdx.x; + __shared__ int wait_failed; + if (tid == 0) + wait_failed = 0; + /* Keep the simple accumulation loop; the compiler provides the required + * pipelining without explicit batching. */ + for (int expert = tid; expert < expert_count; expert += blockDim.x) + { + int count = 0; + for (int cta = 0; cta < ctas; ++cta) + count += partial[cta * expert_count + expert]; + shared[expert] = count; + } + __syncthreads(); + + int const parity = epoch & 1; + /* One warp owns each destination row. Its warp barrier orders every + * lane's stores before lane zero's system-scope release publishes the + * complete row to the peer acquire. */ + int const warp = tid >> 5; + int const lane = tid & 31; + for (int destination = warp; destination < ep; destination += blockDim.x >> 5) + { + int* peer = reinterpret_cast(peer_bases[destination]); + int* row = peer + kSymPayloadOffset + (parity * ep + local_rank) * expert_count; + /* Keep scalar remote stores for the publication path; the local mirror + * copy below uses vector transfers when alignment permits. */ + for (int expert = lane; expert < expert_count; expert += 32) + row[expert] = shared[expert]; + /* Warp synchronization orders every lane's row stores before lane + * zero resumes; its system-scope release then publishes that complete + * transitive happens-before set to the peer acquire. */ + __syncwarp(); + if (lane == 0) + store_release_system(peer + local_rank, epoch); + } + __syncthreads(); + + int* local = reinterpret_cast(peer_bases[local_rank]); + if (tid < ep) + { + unsigned long long start = clock64(); + Epoch observed = load_relaxed_system(local + tid); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 0, epoch); + atomicExch(&wait_failed, 1); + break; + } + observed = load_relaxed_system(local + tid); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_system(local + tid); + } + __syncthreads(); + if (wait_failed) + return false; + int const* payload = local + kSymPayloadOffset + parity * ep * expert_count; + int const total = ep * expert_count; + /* Vectorize the local mirror copy when shape and alignment permit. */ + if ((total & 3) == 0 && (reinterpret_cast(payload) & 15) == 0) + { + int4 const* source = reinterpret_cast(payload); + int4* destination = reinterpret_cast(shared); + int const vectors = total >> 2; + for (int vector = tid; vector < vectors; vector += blockDim.x) + destination[vector] = source[vector]; + } + else + { + for (int index = tid; index < total; index += blockDim.x) + shared[index] = payload[index]; + } + __syncthreads(); + return true; +} + +/* Rank and home bounds of the register planner path. Every hierarchy level + * of EP <= 8 fits one 8-lane warp segment, and a home of at most 64 experts + * fits two 32-lane rounds. */ +constexpr int kRegisterEp = 8; +constexpr int kRegisterRounds = 2; + +__device__ __forceinline__ bool uses_register_prepare(int ep, int experts) +{ + return ep <= kRegisterEp && experts / ep <= 32 * kRegisterRounds; +} + +/* Pick the owner's hot list from the lane-resident home masses: the heaviest + * remaining expert per slot, the lowest home index on a tie. Each lane holds + * home experts lane + 32 * round; entries outside the home carry -1. */ +template +__device__ __forceinline__ void select_owner_candidates(int (&local_mass)[Rounds], int* histogram, int* candidates, + int owner, int home, int candidate_count, bool write_mass) +{ + int const lane = threadIdx.x & 31; + for (int slot = 0; slot < candidate_count; ++slot) + { + int best_mass = -1; + int best_local = INT_MAX; +#pragma unroll + for (int round = 0; round < Rounds; ++round) + { + int const mass = local_mass[round]; + int const local = lane + 32 * round; + if (mass > best_mass || (mass == best_mass && local < best_local)) + { + best_mass = mass; + best_local = local; + } + } + int const warp_mass = __reduce_max_sync(0xffffffffu, best_mass); + int const winner = __reduce_min_sync(0xffffffffu, best_mass == warp_mass ? best_local : INT_MAX); + if (lane == 0) + { + int const expert = owner * home + winner; + candidates[owner * candidate_count + slot] = expert; + if (write_mass) + histogram[expert] = warp_mass; + } +#pragma unroll + for (int round = 0; round < Rounds; ++round) + if (lane + 32 * round == winner) + local_mass[round] = -1; + } +} + +__device__ __forceinline__ void reset_owner_state( + int owner, int rank_load, int* load, int* recv_used, int* send_used, int* next_index, int* frozen) +{ + load[owner] = rank_load; + recv_used[owner] = 0; + send_used[owner] = 0; + next_index[owner] = 0; + frozen[owner] = 0; +} + +/* Rank loads and per-owner hot lists for uses_register_prepare(): one warp + * per owner sums its home masses over the source ranks straight into + * registers. Every load is unconditional (out-of-range sources and experts + * read a clamped valid address and contribute nothing), so all of them are in + * flight together. Only the candidates' masses reach histogram[]; the + * planners read no other entry. The owner warps, warp 0 among them, meet at + * a named barrier; the other warps leave at once. */ +__device__ void prepare_halo_m_register(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) +{ + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + if (warp >= ep) + return; + int const owner = warp; + int const home = experts / ep; + int local_mass[kRegisterRounds]; + int rank_load = 0; +#pragma unroll + for (int round = 0; round < kRegisterRounds; ++round) + { + int const local = lane + 32 * round; + bool const in_home = local < home; + int const expert = owner * home + imin(local, home - 1); + int mass = 0; +#pragma unroll + for (int source = 0; source < kRegisterEp; ++source) + { + int const count = counts[imin(source, ep - 1) * experts + expert]; + mass += source < ep ? count : 0; + } + rank_load += in_home ? mass : 0; + local_mass[round] = in_home ? mass : -1; + } + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + reset_owner_state(owner, rank_load, load, recv_used, send_used, next_index, frozen); + select_owner_candidates(local_mass, histogram, candidates, owner, home, imin(helpers, home), true); + asm volatile("bar.sync 4, %0;" ::"r"(ep * 32) : "memory"); +} + +/* Rounds of 32 home experts per lane; EP2 has the largest home, 192 experts. */ +constexpr int kMaxHomeRounds = (kMaxExperts / 2 + 31) / 32; + +/* Dense histogram, rank loads and per-owner hot lists for every other + * geometry, with one warp per owner. */ +__device__ void prepare_halo_m(int const* counts, int* histogram, int* load, int* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int ep, int experts, int helpers) +{ + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int const home = experts / ep; + for (int expert = tid; expert < experts; expert += blockDim.x) + { + int mass = 0; +#pragma unroll + for (int rank = 0; rank < ep; ++rank) + mass += counts[rank * experts + expert]; + histogram[expert] = mass; + } + __syncthreads(); + + for (int owner = warp; owner < ep; owner += blockDim.x >> 5) + { + int rank_load = 0; + for (int local = lane; local < home; local += 32) + rank_load += histogram[owner * home + local]; + rank_load = __reduce_add_sync(0xffffffffu, rank_load); + if (lane == 0) + reset_owner_state(owner, rank_load, load, recv_used, send_used, next_index, frozen); + int local_mass[kMaxHomeRounds]; +#pragma unroll + for (int round = 0; round < kMaxHomeRounds; ++round) + { + int const local = lane + 32 * round; + local_mass[round] = local < home ? histogram[owner * home + local] : -1; + } + select_owner_candidates(local_mass, histogram, candidates, owner, home, imin(helpers, home), false); + } + __syncthreads(); +} + +/* EP<=8 fast path: every participating lane owns one rank's complete + * mutable HALO-M state. Group and destination decisions use only shuffles, + * avoiding repeated shared-memory rank-state traffic and warp barriers. */ +template +__device__ __forceinline__ int plan_halo_m_register8_impl( + int const* histogram, int lane_load, int total, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const lane = threadIdx.x & 31; + constexpr unsigned mask = 0xffu; + int lane_recv = 0; + int lane_send = 0; + int lane_next = 0; + int lane_frozen = 0; + + int const target = total / ep + (total % ep != 0); + int const depth = hierarchy_depth(ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + int trips = 0; + for (int guard = 0; guard < ep * (helpers + 2); ++guard) + { + ++trips; + bool const eligible = lane < ep && !lane_frozen && lane_load > target; + int owner, owner_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_max_sync( + mask, eligible ? ((static_cast(lane_load) << 4) | (15u - static_cast(lane))) : 0u); + if (!key) + break; + owner = 15 - static_cast(key & 15u); + owner_load = static_cast(key >> 4); + } + else + { + owner_load = __reduce_max_sync(mask, eligible ? lane_load : -1); + if (owner_load < 0) + break; + owner = __reduce_min_sync(mask, eligible && lane_load == owner_load ? lane : INT_MAX); + } + + unsigned const packed + = __shfl_sync(mask, (static_cast(lane_send) << 4) | static_cast(lane_next), owner, 8); + int const owner_send = static_cast(packed >> 4); + int const owner_next = static_cast(packed & 15u); + if (owner_send >= helpers || owner_next >= helpers) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + /* owner and owner_next are warp-uniform here, so every lane can read + * candidates[] and histogram[] directly. Both are written by + * the prepare step before its barrier and never inside this + * loop, so every lane loads the same values and no broadcast from the + * owner is needed. Only the owner advances lane_next, and it does so + * before the mass test. */ + int const expert = candidates[owner * helpers + owner_next]; + if (lane == owner) + ++lane_next; + int const mass = histogram[expert]; + if (mass == 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int give = imin(owner_load - target, mass); + if (give <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + /* EP <= 8 has at most two hierarchy levels: level 0 is every rank + * and, at EP8, level 1 is the two halves of four ranks. Two + * independent full-mask sums give the two half sums. Integer addition + * is exact, so zeroing the other half's addends yields each half sum + * exactly; the two reduces do not depend on each other, and neither + * redux.sync uses a per-lane membermask. The lower half wins a tie. */ + int const half = ep >> 1; + unsigned const rank_mask = (1u << ep) - 1u; + unsigned const low_mask = (1u << half) - 1u; + unsigned const high_mask = rank_mask & ~low_mask; + int const slack = lane < ep && lane != owner ? imax(0, target - lane_load) : 0; + int const room_low = __reduce_add_sync(mask, lane < half ? slack : 0); + int const room_high = __reduce_add_sync(mask, lane >= half ? slack : 0); + unsigned const receive_mask = __ballot_sync(mask, lane < ep && lane_recv < helpers); + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + if (depth > 1 && mass <= half * target) + { + bool const valid_low = (receive_mask & low_mask) == low_mask && room_low > 0; + bool const valid_high = (receive_mask & high_mask) == high_mask && room_high > 0; + if (valid_low || valid_high) + { + chosen_level = 1; + chosen_begin = valid_low && (!valid_high || room_low >= room_high) ? 0 : half; + chosen_size = half; + } + } + if (chosen_level < 0 && (receive_mask & rank_mask) == rank_mask && room_low + room_high > 0) + { + chosen_level = 0; + chosen_begin = 0; + chosen_size = ep; + } + if (chosen_level < 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int lane_quota = lane == owner ? mass : 0; + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int destination, destination_load; + if constexpr (NarrowKeys) + { + unsigned const key = __reduce_min_sync(mask, + can_receive ? ((static_cast(lane_load) << 4) | static_cast(lane)) + : 0xffffffffu); + if (key == 0xffffffffu) + break; + destination = static_cast(key & 15u); + destination_load = static_cast(key >> 4); + } + else + { + destination_load = __reduce_min_sync(mask, can_receive ? lane_load : INT_MAX); + destination = __reduce_min_sync(mask, can_receive && lane_load == destination_load ? lane : INT_MAX); + if (destination == INT_MAX) + break; + } + visited |= 1u << destination; + int room = target - destination_load; + if (room <= 0) + break; + int amount = imin(room, remaining); + if (lane == owner) + lane_quota -= amount; + if (lane == destination) + lane_quota += amount; + remaining -= amount; + } + + int moved = give - remaining; + if (moved <= 0) + { + if (lane == owner) + lane_frozen = 1; + continue; + } + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < ep) + quota[b * ep + lane] = lane_quota; + + bool in_group = lane >= chosen_begin && lane < chosen_begin + chosen_size; + if (in_group) + ++lane_recv; + if (lane == owner) + { + ++lane_send; + lane_load -= moved; + } + else if (in_group) + { + lane_load += lane_quota; + } + } + return broadcasts; +} + +/* Select the key representation once, from the actual complete rank load. + * This guard keeps the owner and destination load keys representable. + * Keep the wide path for large valid int32 inputs without putting a runtime + * representation branch on each iteration of either planning loop. */ +__device__ __forceinline__ int plan_halo_m_register8( + int const* histogram, int const* load, int* quota, int const* candidates, int* plan, int ep, int helpers) +{ + int const lane = threadIdx.x & 31; + int const lane_load = lane < ep ? load[lane] : 0; + int const total = __reduce_add_sync(0xffu, lane_load); + if (total <= 0x0fffffff) + return plan_halo_m_register8_impl(histogram, lane_load, total, quota, candidates, plan, ep, helpers); + return plan_halo_m_register8_impl(histogram, lane_load, total, quota, candidates, plan, ep, helpers); +} + +/* Exact HALO-M state machine executed by warp 0. Each lane owns one rank; + * rank/group/destination scans are warp reductions rather than serial loops. */ +__device__ int plan_halo_m_warp(int const* histogram, int* load, int* quota, int const* candidates, int* recv_used, + int* send_used, int* next_index, int* frozen, int* plan, int ep, int experts, int helpers) +{ + int const candidate_count = imin(helpers, experts / ep); + int const lane = threadIdx.x & 31; + int const reduce_width = ep <= 8 ? 8 : 32; + int const reduce_start = reduce_width >> 1; + unsigned const reduce_mask = ep <= 8 ? 0xffu : 0xffffffffu; + int total = lane < ep ? load[lane] : 0; + for (int offset = reduce_start; offset; offset >>= 1) + { + int const other = __shfl_down_sync(reduce_mask, total, offset, reduce_width); + if (lane + offset < reduce_width) + total += other; + } + total = __shfl_sync(reduce_mask, total, 0, reduce_width); + int const target = total / ep + (total % ep != 0); + int const depth = hierarchy_depth(ep); + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int broadcasts = 0; + + for (int guard = 0; guard < ep * (candidate_count + 2); ++guard) + { + bool eligible = lane < ep && !frozen[lane] && load[lane] > target; + int best_load = eligible ? load[lane] : -1; + int best_rank = eligible ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, best_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, best_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load > best_load || (other_load == best_load && other_rank < best_rank))) + { + best_load = other_load; + best_rank = other_rank; + } + } + int owner = __shfl_sync(reduce_mask, best_rank, 0, reduce_width); + if (owner == INT_MAX) + break; + /* One lane owns the capacity check and shared cursor transition. + * Publish either the chosen expert or -1 for exhaustion explicitly; + * no reader can observe a partially advanced cursor. */ + int expert = -1; + if (lane == 0) + { + if (send_used[owner] >= helpers || next_index[owner] >= candidate_count) + { + frozen[owner] = 1; + } + else + { + int const candidate_index = next_index[owner]++; + expert = candidates[owner * candidate_count + candidate_index]; + } + } + expert = __shfl_sync(reduce_mask, expert, 0, reduce_width); + __syncwarp(reduce_mask); + if (expert < 0) + continue; + int mass = histogram[expert]; + if (mass == 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + int give = imin(load[owner] - target, mass); + if (give <= 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + + int need = imax(1, mass / imax(target, 1) + (mass % imax(target, 1) != 0)); + int wanted = depth - 1; + for (int level = depth - 1; level >= 0; --level) + { + if (ep / (1 << level) >= need) + { + wanted = level; + break; + } + } + int chosen_level = -1; + int chosen_begin = -1; + int chosen_size = 0; + for (int level = wanted; level >= 0 && chosen_level < 0; --level) + { + int size = ep / (1 << level); + int begin = lane * size; + bool active = begin < ep; + bool affordable = active; + int room = 0; + if (active) + { + for (int rank = begin; rank < begin + size; ++rank) + { + affordable &= recv_used[rank] < helpers; + if (rank != owner) + room += imax(0, target - load[rank]); + } + } + int candidate_room = affordable ? room : -1; + int candidate_begin = affordable ? begin : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_room = __shfl_down_sync(reduce_mask, candidate_room, offset, reduce_width); + int other_begin = __shfl_down_sync(reduce_mask, candidate_begin, offset, reduce_width); + if (lane + offset < reduce_width + && (other_room > candidate_room || (other_room == candidate_room && other_begin < candidate_begin))) + { + candidate_room = other_room; + candidate_begin = other_begin; + } + } + int best_room = __shfl_sync(reduce_mask, candidate_room, 0, reduce_width); + int best_begin = __shfl_sync(reduce_mask, candidate_begin, 0, reduce_width); + if (best_begin != INT_MAX && best_room > 0) + { + chosen_level = level; + chosen_begin = best_begin; + chosen_size = size; + } + } + if (chosen_level < 0) + { + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (broadcasts >= kMaxBroadcasts) + break; + + int b = broadcasts++; + if (lane == 0) + { + p_expert[b] = expert; + p_level[b] = chosen_level; + p_begin[b] = chosen_begin; + p_size[b] = chosen_size; + p_owner[b] = owner; + p_helper[b] = -1; + } + if (lane < ep) + quota[b * ep + lane] = 0; + __syncwarp(reduce_mask); + if (lane == 0) + quota[b * ep + owner] = mass; + __syncwarp(reduce_mask); + + int remaining = give; + unsigned visited = 0; + for (int step = 0; step < chosen_size && remaining > 0; ++step) + { + bool can_receive = lane >= chosen_begin && lane < chosen_begin + chosen_size && lane != owner + && !(visited & (1u << lane)); + int cold_load = can_receive ? load[lane] : INT_MAX; + int cold_rank = can_receive ? lane : INT_MAX; + for (int offset = reduce_start; offset; offset >>= 1) + { + int other_load = __shfl_down_sync(reduce_mask, cold_load, offset, reduce_width); + int other_rank = __shfl_down_sync(reduce_mask, cold_rank, offset, reduce_width); + if (lane + offset < reduce_width + && (other_load < cold_load || (other_load == cold_load && other_rank < cold_rank))) + { + cold_load = other_load; + cold_rank = other_rank; + } + } + int destination = __shfl_sync(reduce_mask, cold_rank, 0, reduce_width); + if (destination == INT_MAX) + break; + visited |= 1u << destination; + int room = target - load[destination]; + if (room <= 0) + continue; + int amount = imin(room, remaining); + if (lane == 0) + { + quota[b * ep + owner] -= amount; + quota[b * ep + destination] += amount; + } + remaining -= amount; + __syncwarp(reduce_mask); + } + int moved = give - remaining; + if (moved <= 0) + { + --broadcasts; + if (lane == 0) + frozen[owner] = 1; + __syncwarp(reduce_mask); + continue; + } + if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + ++recv_used[lane]; + if (lane == owner) + { + ++send_used[owner]; + load[owner] -= moved; + } + else if (lane >= chosen_begin && lane < chosen_begin + chosen_size) + { + load[lane] += quota[b * ep + lane]; + } + __syncwarp(reduce_mask); + } + return broadcasts; +} + +__device__ __forceinline__ int warp_sum(int value) +{ + for (int offset = 16; offset; offset >>= 1) + { + int const other = __shfl_down_sync(0xffffffffu, value, offset); + if ((threadIdx.x & 31) + offset < 32) + value += other; + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__device__ __forceinline__ int warp_exclusive_scan(int value) +{ + int sum = value; + int const lane = threadIdx.x & 31; + for (int offset = 1; offset < 32; offset <<= 1) + { + int other = __shfl_up_sync(0xffffffffu, sum, offset); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ void halo_q_repair(int receiver, int donor, int broadcasts, int const* order, int const* host_mask, + int const* counts, int const* p_expert, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + int local_pull = 0, local_give = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + local_pull += imin(imax(0, cr - qr), imax(0, qd - cd)); + local_give += imax(0, qr - cr); + } + } + int const delta = imin(warp_sum(local_pull), warp_sum(local_give)); + int pull_before_tiles = 0, give_before_tiles = 0; + for (int base = 0; base < broadcasts; base += 32) + { + int const b = base + lane < broadcasts ? order[base + lane] : -1; + int pull = 0, give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int const expert = p_expert[b]; + int const cr = counts[receiver * experts + expert]; + int const cd = counts[donor * experts + expert]; + int const qr = quota[b * ep + receiver]; + int const qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int const pull_before = pull_before_tiles + warp_exclusive_scan(pull); + int const give_before = give_before_tiles + warp_exclusive_scan(give); + int const take = imin(pull, imax(0, delta - pull_before)); + int const release = imin(give, imax(0, delta - give_before)); + pull_before_tiles += warp_sum(pull); + give_before_tiles += warp_sum(give); + if (b >= 0) + { + quota[b * ep + receiver] += take - release; + quota[b * ep + donor] += release - take; + } + __syncwarp(); + } + return; + } + int b = lane < broadcasts ? order[lane] : -1; + int pull = 0; + int give = 0; + if (b >= 0 && (host_mask[b] & (1u << receiver)) && (host_mask[b] & (1u << donor))) + { + int expert = p_expert[b]; + int cr = counts[receiver * experts + expert]; + int cd = counts[donor * experts + expert]; + int qr = quota[b * ep + receiver]; + int qd = quota[b * ep + donor]; + pull = imin(imax(0, cr - qr), imax(0, qd - cd)); + give = imax(0, qr - cr); + } + int delta = imin(warp_sum(pull), warp_sum(give)); + int pull_before = warp_exclusive_scan(pull); + int give_before = warp_exclusive_scan(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + if (b >= 0) + { + quota[b * ep + receiver] += pull_take - give_take; + quota[b * ep + donor] += give_take - pull_take; + } + __syncwarp(); +} + +__device__ __forceinline__ int warp_sum8(int value) +{ + return __reduce_add_sync(0xffu, value); +} + +__device__ __forceinline__ int warp_exclusive_scan8(int value) +{ + int sum = value; + int const lane = threadIdx.x & 7; + for (int offset = 1; offset < 8; offset <<= 1) + { + int other = __shfl_up_sync(0xffu, sum, offset, 8); + if (lane >= offset) + sum += other; + } + return sum - value; +} + +__device__ __forceinline__ void halo_q_direction8( + bool common, int receiver_count, int donor_count, int& receiver_quota, int& donor_quota) +{ + int pull = common ? imin(imax(0, receiver_count - receiver_quota), imax(0, donor_quota - donor_count)) : 0; + int give = common ? imax(0, receiver_quota - receiver_count) : 0; + int pull_total = warp_sum8(pull); + int give_total = warp_sum8(give); + int delta = imin(pull_total, give_total); + if (!delta) + return; + int pull_before = warp_exclusive_scan8(pull); + int give_before = warp_exclusive_scan8(give); + int pull_take = imin(pull, imax(0, delta - pull_before)); + int give_take = imin(give, imax(0, delta - give_before)); + receiver_quota += pull_take - give_take; + donor_quota += give_take - pull_take; +} + +/* Fuse both ordered directions of one rank pair. Each broadcast lane loads + * its two count/quota entries once, executes left<-right then right<-left in + * registers, and writes the final pair once. + * + * `b`, `hosts` and `expert` are the lane's own broadcast identity and do not + * change between phases or rounds, so the caller hoists them. Reading them + * here would put a four-deep dependent load chain (order[lane] -> b -> host_mask[b] + * and the GLOBAL p_expert[b] -> expert -> counts[...]) at the head of every + * phase. */ +__device__ __forceinline__ void halo_q_pair8( + int left, int right, int b, int hosts, int expert, int const* counts, int* quota, int ep, int experts) +{ + int const lane = threadIdx.x & 31; + if (lane >= 8) + return; + bool common = false; + int count_left = 0; + int count_right = 0; + int quota_left = 0; + int quota_right = 0; + if (b >= 0) + { + common = (hosts & (1u << left)) && (hosts & (1u << right)); + if (common) + { + count_left = counts[left * experts + expert]; + count_right = counts[right * experts + expert]; + quota_left = quota[b * ep + left]; + quota_right = quota[b * ep + right]; + } + } + + halo_q_direction8(common, count_left, count_right, quota_left, quota_right); + halo_q_direction8(common, count_right, count_left, quota_right, quota_left); + + if (common) + { + quota[b * ep + left] = quota_left; + quota[b * ep + right] = quota_right; + } +} + +__device__ __forceinline__ void halo_q_active_barrier(int threads) +{ + asm volatile("bar.sync 1, %0;" ::"r"(threads) : "memory"); +} + +__device__ __forceinline__ void scheduler_tail_barrier() +{ + asm volatile("bar.sync 3, 64;" ::: "memory"); +} + +/* Build order[] and host_mask[] with one lane per broadcast. + * + * order[]: lane b's position is the number of broadcasts that sort before it, + * ordered by expert. The `j < lane` tiebreak orders equal experts by + * broadcast index, so the order is stable. + * host_mask: the destination group is a contiguous rank range, so the mask is + * closed form; no loop over group members is needed. + * rank_group_mask handles the full EP32 word without a shift by32. + */ +__device__ __forceinline__ void halo_q_build_order_and_masks(int const* p_expert, int const* p_begin, int const* p_size, + int const* p_owner, int* order, int* host_mask, int broadcasts) +{ + int const lane = threadIdx.x & 31; + if (broadcasts > 32) + { + for (int b = lane; b < broadcasts; b += 32) + { + int const mine = p_expert[b]; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + position += p_expert[j] < mine || (p_expert[j] == mine && j < b); + order[position] = b; + host_mask[b] = static_cast(rank_group_mask(p_begin[b], p_size[b]) | (1u << p_owner[b])); + } + return; + } + /* Read p_expert cooperatively and distribute values with shuffles. All + * lanes must participate before inactive lanes return; inactive lanes + * carry INT_MAX and never contribute to a real comparison. */ + bool const active = lane < broadcasts; + int const mine = active ? p_expert[lane] : INT_MAX; + int position = 0; + for (int j = 0; j < broadcasts; ++j) + { + int const other = __shfl_sync(0xffffffffu, mine, j); + position += (other < mine) || (other == mine && j < lane); + } + if (!active) + return; + order[position] = lane; + host_mask[lane] = static_cast(rank_group_mask(p_begin[lane], p_size[lane]) | (1u << p_owner[lane])); +} + +__device__ void run_halo_q(int const* counts, int* quota, int* plan, int broadcasts, int ep, int experts, int* ring, + int* order, int* host_mask) +{ + int const tid = threadIdx.x; + int const warp = tid >> 5; + int const lane = tid & 31; + int* p_expert = plan_field(plan, kPlanExpert); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + if (ep <= kRegisterEp && broadcasts <= 8) + { + if (warp < ep / 2) + { + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + halo_q_active_barrier((ep / 2) * 32); + /* Loop-invariant broadcast identity, read once after the barrier + * that publishes order[] and host_mask[]. All ep / 2 participating + * warps read the same lane-indexed entries. */ + int const my_b = lane < broadcasts ? order[lane] : -1; + int const my_hosts = my_b >= 0 ? host_mask[my_b] : 0; + int const my_expert = my_b >= 0 ? p_expert[my_b] : -1; + if (broadcasts >= 2) + { + int const ring_size = ep - 1; + /* Run two repair rounds unconditionally. A phase with no required + * repair is a no-op, so the plan equals running only the required + * repairs. */ + int const repair_rounds = 2; + for (int round = 0; round < repair_rounds; ++round) + { + /* At most kRegisterEp - 1 phases, unrolled; the guard is + * uniform across the block, so every pair warp meets + * every barrier it executes. */ +#pragma unroll + for (int phase = 0; phase < kRegisterEp - 1; ++phase) + { + if (phase >= ep - 1) + break; + int right_index = warp == 0 ? 0 : ep - 1 - warp; + int right_position = right_index - phase; + if (right_position < 0) + right_position += ring_size; + int right = 1 + right_position; + int left = 0; + if (warp != 0) + { + int left_position = warp - phase; + if (left_position < 0) + left_position += ring_size; + left = 1 + left_position; + } + halo_q_pair8(left, right, my_b, my_hosts, my_expert, counts, quota, ep, experts); + /* Must stay unconditional. Omitting it after a + * round's last phase races twice: warp 0 enters the + * next round's phase 0 while the last pair warp still + * writes the same quota[], and after the final phase + * the other pair warps leave run_halo_q while warp 0 + * -- which only meets warp 8 at scheduler_tail_barrier + * -- walks into build_route_prefix and reads every + * quota[]. */ + halo_q_active_barrier((ep / 2) * 32); + } + } + } + } + return; + } + if (warp == 0) + halo_q_build_order_and_masks(p_expert, p_begin, p_size, p_owner, order, host_mask, broadcasts); + if (tid < ep - 1) + ring[tid] = tid + 1; + __syncthreads(); + if (broadcasts < 2) + return; + + for (int round = 0; round < 2; ++round) + { + for (int phase = 0; phase < ep - 1; ++phase) + { + if (warp < ep / 2) + { + int left, right; + if (warp == 0) + { + left = 0; + right = ring[0]; + } + else + { + left = ring[warp]; + right = ring[ep - 1 - warp]; + } + halo_q_repair(left, right, broadcasts, order, host_mask, counts, p_expert, quota, ep, experts); + halo_q_repair(right, left, broadcasts, order, host_mask, counts, p_expert, quota, ep, experts); + } + __syncthreads(); + if (tid == 0) + { + int last = ring[ep - 2]; + for (int index = ep - 2; index > 0; --index) + ring[index] = ring[index - 1]; + ring[0] = last; + } + __syncthreads(); + } + } + (void) lane; +} + +__device__ void color_and_publish_local(int* plan, int broadcasts, int ep, int helpers, int local_rank, int* used, + int* out_ids, int* out_levels, int* out_owners, int* status, Epoch epoch) +{ + int* p_expert = plan_field(plan, kPlanExpert); + int* p_level = plan_field(plan, kPlanLevel); + int* p_begin = plan_field(plan, kPlanGroupBegin); + int* p_size = plan_field(plan, kPlanGroupSize); + int* p_owner = plan_field(plan, kPlanOwner); + int* p_helper = plan_field(plan, kPlanHelper); + int const active_helpers = imin(helpers, kMaxBroadcasts); + int const words = (active_helpers + 31) / 32; + for (int rank = 0; rank < ep; ++rank) + for (int word = 0; word < words; ++word) + used[rank * kHelperMaskWords + word] = 0; + for (int level = 0; level < hierarchy_depth(ep); ++level) + { + for (int b = 0; b < broadcasts; ++b) + { + if (p_level[b] != level) + continue; + /* The lowest slot free on every group rank, one mask word at a + * time, so the scan keeps no per-word array. */ + int helper = -1; + for (int word = 0; word < words && helper < 0; ++word) + { + unsigned unavailable = 0; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + unavailable |= static_cast(used[rank * kHelperMaskWords + word]); + int const slots = imin(32, active_helpers - word * 32); + unsigned const valid = slots == 32 ? 0xffffffffu : (1u << slots) - 1u; + unsigned const open_slots = ~unavailable & valid; + if (open_slots) + helper = word * 32 + __ffs(open_slots) - 1; + } + if (helper < 0) + { + record_status(status + 4, epoch); + helper = 0; + } + p_helper[b] = helper; + for (int rank = p_begin[b]; rank < p_begin[b] + p_size[b]; ++rank) + used[rank * kHelperMaskWords + helper / 32] |= static_cast(1u << (helper % 32)); + } + } + for (int slot = 0; slot < helpers; ++slot) + { + out_ids[slot] = -1; + out_levels[slot] = -1; + out_owners[slot] = -1; + } + for (int b = 0; b < broadcasts; ++b) + { + if (!group_contains(local_rank, p_begin[b], p_size[b])) + continue; + int helper = p_helper[b]; + if (out_ids[helper] != -1) + record_status(status + 4, epoch); + out_ids[helper] = p_expert[b]; + out_levels[helper] = p_level[b]; + out_owners[helper] = p_owner[b]; + } +} + +__device__ void build_route_prefix(int const* counts, int const* quota, int* prefix, int const* plan, int broadcasts, + int ep, int experts, int local_rank, int participating_threads) +{ + int* p_expert = plan_field(const_cast(plan), kPlanExpert); + for (int b = threadIdx.x; b < broadcasts; b += participating_threads) + { + int expert = p_expert[b]; + int source_before = 0; + /* Ranks past the bound load a clamped valid entry and add nothing, + * so every load of a chunk is unconditional. */ + for (int base = 0; base < local_rank; base += kRankChunk) + { +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) + { + int const rank = imin(base + step, local_rank - 1); + int const count = counts[rank * experts + expert]; + int const diagonal = imin(count, quota[b * ep + rank]); + source_before += base + step < local_rank ? count - diagonal : 0; + } + } + int source_count = counts[local_rank * experts + expert]; + int source_diagonal = imin(source_count, quota[b * ep + local_rank]); + int residual_supply = source_count - source_diagonal; + /* route_destination loads prefix[destination + step] only under + * `candidate < ep`, so entry ep is never read: the loop stops before + * destination == ep and skips the accumulation after the last + * destination, which would only feed that entry. Entry 0 is written + * although route_destination never reads it: it costs one store and + * leaves no stale entry below ep. */ + int destination_before = 0; + for (int base = 0; base < ep; base += kRankChunk) + { + /* Load the whole chunk before the first prefix store, which the + * compiler cannot order after loads it might alias. */ + int destination_supply[kRankChunk]; +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) + { + int const destination = imin(base + step, ep - 1); + int const destination_quota = quota[b * ep + destination]; + destination_supply[step] + = destination_quota - imin(counts[destination * experts + expert], destination_quota); + } +#pragma unroll + for (int step = 0; step < kRankChunk; ++step) + { + int const destination = base + step; + int const overlap = imax(0, imin(residual_supply, destination_before - source_before)); + int const diagonal_prefix = local_rank < destination ? source_diagonal : 0; + if (destination < ep) + prefix[b * (ep + 1) + destination] = diagonal_prefix + overlap; + destination_before += destination + 1 < ep ? destination_supply[step] : 0; + } + } + } +} + +__device__ __forceinline__ int* worker_expert_lut(int experts) +{ + extern __shared__ int shared[]; + return shared + (blockDim.x >> 5) * (experts + 1); +} + +__device__ void prepare_worker_default_slots(int* expert_lut, int ep, int experts, int helpers) +{ + int const home = experts / ep; + int const local_slots = home + helpers; + for (int expert = threadIdx.x; expert < experts; expert += blockDim.x) + { + int owner = expert / home; + int local = expert - owner * home; + /* Negative values distinguish the precomputed cold slot from a + * nonnegative broadcast index. Slot zero therefore maps to -1. */ + expert_lut[expert] = -(owner * local_slots + local) - 1; + } + __syncthreads(); +} + +__device__ void prepare_worker_selected_experts(int* expert_lut, int const* plan) +{ + int const broadcasts = plan[1]; + int const* p_expert = plan + 2 + kPlanExpert * kMaxBroadcasts; + for (int b = threadIdx.x; b < broadcasts; b += blockDim.x) + expert_lut[p_expert[b]] = b; + __syncthreads(); +} + +__device__ bool wait_for_plan_epoch(int* marker, Epoch epoch, unsigned long long spin_cycles, int* status, int* ready) +{ + if (threadIdx.x == 0) + { + *ready = 1; + unsigned long long start = clock64(); + Epoch observed = load_relaxed_device(marker); + while (epoch_before(observed, epoch)) + { + if (clock64() - start > spin_cycles) + { + record_status(status + 2, epoch); + *ready = 0; + break; + } + observed = load_relaxed_device(marker); + } + if (!epoch_before(observed, epoch)) + (void) load_acquire_device(marker); + } + __syncthreads(); + return *ready != 0; +} + +/* Prefix entries are monotonic, so binary lifting returns the number of + * thresholds <= ordinal, counting repeated entries. */ +__device__ __forceinline__ int route_destination(int const* prefix, int ordinal, int ep) +{ + int destination = 0; +#pragma unroll + for (int step = 16; step; step >>= 1) + { + int candidate = destination + step; + if (candidate < ep && ordinal >= prefix[candidate]) + destination = candidate; + } + return destination; +} + +/* Once HALO-M publishes its selected experts, workers can finish every cold + * route while CTA0 runs HALO-Q. Selected route indices are compacted into + * disjoint slices of route_aux; their ordinal values remain in output. */ +__device__ void pre_materialize_and_compact_worker_routes(int const* routes, int* output, int* route_aux, int* status, + int route_count, int experts, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = 0; + for (int base = begin; base < end; base += 32) + { + int index = base + lane; + bool active = index < end; + int expert = active ? routes[index] : -1; + int slot = -1; + bool hot = false; + if (active && expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + hot = entry >= 0; + if (!hot) + slot = -entry - 1; + } + else if (active && expert != -1) + { + record_status(status + 3, epoch); + } + unsigned hot_mask = __ballot_sync(0xffffffffu, hot); + if (hot) + { + int offset = __popc(hot_mask & ((1u << lane) - 1u)); + hot_indices[begin + hot_count + offset] = index; + } + else if (active) + { + output[index] = slot; + } + hot_count += __popc(hot_mask); + } + if (lane == 0) + shared[warp * bins + experts] = hot_count; +} + +/* The final epoch only has to resolve selected routes. Stable ordinal + * cursors remain in the per-warp shared histogram from the initial pass. */ +__device__ void materialize_worker_hot_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* route_aux, int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, + Epoch epoch) +{ + extern __shared__ int shared[]; + int const lane = threadIdx.x & 31; + int const warp = threadIdx.x >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int* hot_indices = route_aux + 2 + workers * bins; + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + int hot_count = shared[warp * bins + experts]; + hot_count = imin(hot_count, imax(0, end - begin)); + /* Each lane resolves kHotBatch routes at once: every dependent load of a + * stage is issued for the whole batch before the next stage consumes it. + * Batch members are distinct routes, so the order of their stores does + * not matter. */ + constexpr int kHotBatch = 4; + for (int first = lane; first < hot_count; first += 32 * kHotBatch) + { + int index[kHotBatch]; + int expert[kHotBatch]; + int ordinal[kHotBatch]; +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) + { + int const offset = first + 32 * step; + index[step] = offset < hot_count ? hot_indices[begin + offset] : -1; + } +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) + { + expert[step] = index[step] >= 0 ? routes[index[step]] : 0; + ordinal[step] = index[step] >= 0 ? output[index[step]] : 0; + } +#pragma unroll + for (int step = 0; step < kHotBatch; ++step) + { + if (index[step] < 0) + continue; + int b = expert_lut[expert[step]]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal[step] + shared[warp * bins + expert[step]], ep); + int slot = -1; + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert[step] - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + output[index[step]] = slot; + } + } +} + +__device__ void materialize_worker_routes(int const* routes, int* output, int const* route_prefix, int const* plan, + int* status, int route_count, int ep, int experts, int helpers, int worker, int workers, Epoch epoch) +{ + extern __shared__ int shared[]; + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const warps = blockDim.x >> 5; + int const bins = experts + 1; + int* expert_lut = worker_expert_lut(experts); + int const home = experts / ep; + int const local_slots = home + helpers; + int const* p_begin = plan + 2 + kPlanGroupBegin * kMaxBroadcasts; + int const* p_size = plan + 2 + kPlanGroupSize * kMaxBroadcasts; + int const* p_owner = plan + 2 + kPlanOwner * kMaxBroadcasts; + int const* p_helper = plan + 2 + kPlanHelper * kMaxBroadcasts; + + int const worker_span = (route_count + workers - 1) / workers; + int const worker_begin = worker * worker_span; + int const worker_end = imin(route_count, worker_begin + worker_span); + int const warp_span = ((worker_span + warps * 32 - 1) / (warps * 32)) * 32; + int const begin = worker_begin + warp * warp_span; + int const end = imin(worker_end, begin + warp_span); + for (int index = begin + lane; index < end; index += 32) + { + int expert = routes[index]; + int slot = -1; + if (expert >= 0 && expert < experts) + { + int entry = expert_lut[expert]; + if (entry < 0) + { + slot = -entry - 1; + } + else + { + int b = entry; + int ordinal = output[index] + shared[warp * bins + expert]; + int const* prefix = route_prefix + b * (ep + 1); + int destination = route_destination(prefix, ordinal, ep); + if (group_contains(destination, p_begin[b], p_size[b])) + { + slot = destination * local_slots + home + p_helper[b]; + } + else if (destination == p_owner[b]) + { + slot = destination * local_slots + expert - p_owner[b] * home; + } + else + { + record_status(status + 5, epoch); + } + } + } + else if (expert != -1) + { + record_status(status + 3, epoch); + } + output[index] = slot; + } +} + +/* One scheduler generation on a grid of ctas resident CTAs of 512 threads: + * CTA 0 plans and publishes the helper placement, the other CTAs write routes. + * The fused scheduler + in-switch copy kernel runs it as its first phase. */ +__device__ __forceinline__ void halo_q_scheduler_body(int const* routes, int* out_slots, int* out_ids, int* out_levels, + int* out_owners, unsigned long long const* peer_bases, int* status, int* partial, int* route_aux, int* grid_sync, + int* plan, int* route_prefix, int ep, int experts, int helpers, int local_rank, int route_capacity, int ctas, + unsigned long long spin_cycles, int capacity, int route_count) +{ + // Only the caller's valid prefix is readable. Keep the public output at + // its configured capacity, clearing stale routes inside this same kernel. + // Empty ranks still execute every collective; no route load occurs at N=0. + for (int index = route_count + blockIdx.x * blockDim.x + threadIdx.x; index < route_capacity; + index += gridDim.x * blockDim.x) + out_slots[index] = -1; + local_histogram_pass(routes, partial, route_count, experts, ctas); + Epoch epoch; + if (!grid_rendezvous(grid_sync, ctas, spin_cycles, status, &epoch)) + return; + + if (blockIdx.x == 0) + { + extern __shared__ int shared[]; + int offset = ep * experts; + int* counts = shared; + int* histogram = shared + offset; + offset += experts; + int* load = shared + offset; + offset += ep; + // Reuse the exact capacity already validated/computed by launch(). + int* quota = shared + offset; + offset += capacity * ep; + int* selected_flags = shared + offset; + offset += experts; + int* recv_used = shared + offset; + offset += ep * kHelperMaskWords; + int* send_used = shared + offset; + offset += ep; + int* next_index = shared + offset; + offset += ep; + int* frozen = shared + offset; + offset += ep; + int* ring = shared + offset; + offset += ep; + int* order = shared + offset; + offset += capacity; + int* host_mask = shared + offset; + + if (!exchange_counts(peer_bases, partial, shared, ep, experts, local_rank, ctas, epoch, spin_cycles, status)) + return; + if (uses_register_prepare(ep, experts)) + prepare_halo_m_register(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, + ep, experts, helpers); + else + prepare_halo_m(counts, histogram, load, selected_flags, recv_used, send_used, next_index, frozen, ep, + experts, helpers); + + int broadcasts = 0; + if (threadIdx.x < (ep <= 8 ? 8 : 32)) + { + // These short-circuit bounds keep helpers*ep <= 120. + if (ep <= 8 && helpers <= 15 && helpers * ep <= experts) + { + broadcasts = plan_halo_m_register8(histogram, load, quota, selected_flags, plan, ep, helpers); + } + else + { + broadcasts = plan_halo_m_warp(histogram, load, quota, selected_flags, recv_used, send_used, next_index, + frozen, plan, ep, experts, helpers); + } + if (threadIdx.x == 0) + plan[1] = broadcasts; + } + __syncthreads(); + broadcasts = plan[1]; + /* p_expert and plan[1] are produced by thread zero. Publish this + * placement-only epoch before HALO-Q so workers can construct their + * selected-expert LUT while quota repair is still running. */ + if (ctas > 1 && threadIdx.x == 0) + { + store_release_device(plan + 0, epoch); + } + /* Small-EP repair excludes warp8 and can overlap coloring. Generic + * repair first meets the coloring warp at its full-block barrier. */ + if (threadIdx.x == 8 * 32) + color_and_publish_local( + plan, broadcasts, ep, helpers, local_rank, recv_used, out_ids, out_levels, out_owners, status, epoch); + run_halo_q(counts, quota, plan, broadcasts, ep, experts, ring, order, host_mask); + if (ctas > 1) + { + int const warp = threadIdx.x >> 5; + /* build_route_prefix is independent of warp 8's outputs. Defer their + * rendezvous until both paths reach their first shared consumer so + * prefix construction and coloring can overlap safely. */ + if (warp == 0) + { + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, 32); + __syncwarp(); + } + if (warp == 0 || warp == 8) + scheduler_tail_barrier(); + /* The worker CTAs gate on grid_sync+2 and consume route_prefix, + * which is complete above. */ + if (threadIdx.x == 0) + store_release_device(grid_sync + 2, epoch); + return; + } + if (threadIdx.x == 0) + store_release_device(plan + 0, epoch); + __syncthreads(); + build_route_prefix(counts, quota, route_prefix, plan, broadcasts, ep, experts, local_rank, blockDim.x); + /* The block barrier orders route-prefix/helper writes before release. */ + __syncthreads(); + } + else + { + int const worker = blockIdx.x - 1; + int const workers = ctas - 1; + if (!stable_worker_ordinal_pass( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + __shared__ int plan_ready; + if (!wait_for_plan_epoch(plan + 0, epoch, spin_cycles, status, &plan_ready)) + return; + prepare_worker_selected_experts(expert_lut, plan); + pre_materialize_and_compact_worker_routes( + routes, out_slots, route_aux, status, route_count, experts, worker, workers, epoch); + if (!wait_for_plan_epoch(grid_sync + 2, epoch, spin_cycles, status, &plan_ready)) + return; + materialize_worker_hot_routes(routes, out_slots, route_prefix, plan, route_aux, status, route_count, ep, + experts, helpers, worker, workers, epoch); + return; + } + + /* CTA-count one remains a correct serial fallback. */ + if (!stable_worker_ordinal_pass(routes, out_slots, route_aux, status, route_count, experts, 0, 1, spin_cycles)) + return; + int* expert_lut = worker_expert_lut(experts); + prepare_worker_default_slots(expert_lut, ep, experts, helpers); + prepare_worker_selected_experts(expert_lut, plan); + materialize_worker_routes( + routes, out_slots, route_prefix, plan, status, route_count, ep, experts, helpers, 0, 1, epoch); +} + +/* Dynamic shared-memory ints that the planner CTA and a worker CTA need. */ +__host__ __device__ __forceinline__ int halo_q_shared_ints(int ep, int experts, int capacity, int threads) +{ + int const workerInts = threads / 32 * (experts + 1) + experts; + int const plannerInts + = ep * experts + experts + ep + capacity * ep + experts + (kHelperMaskWords + 4) * ep + 2 * capacity; + return workerInts > plannerInts ? workerInts : plannerInts; +} + +} // namespace + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_HALO_Q_KERNEL_CUH diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp new file mode 100644 index 000000000000..929a18cce151 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlaceholder.cpp @@ -0,0 +1,101 @@ +/* + * 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/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" + +#include "tensorrt_llm/common/assert.h" + +#include + +#ifdef TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_HALO_Q +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +void checkMoeRebalanceHaloQParams(MoeRebalanceHaloQParams const&) +{ + TLLM_THROW("Per-layer dynamic EPLB HALO-Q requires a requested SM100, SM103, or SM107 target"); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END +#endif + +#ifdef TRTLLM_DYNAMIC_EPLB_PLACEHOLDER_TMA +namespace +{ + +int unsupported() +{ + return static_cast(cudaErrorNotSupported); +} + +} // namespace + +extern "C" int megamoe_tma_copy_create(uint64_t, int, int, MegamoeTmaCopyState** out) +{ + if (out != nullptr) + { + *out = nullptr; + } + return unsupported(); +} + +extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState*, MegamoeTmaGpuPlanConfig const*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_bind_gpu_direct( + MegamoeTmaCopyState*, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, int, int) +{ + return unsupported(); +} + +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState*, tensorrt_llm::kernels::MoeRebalanceHaloQParams const&, + uint64_t*, uint64_t, uint64_t, uint64_t, void*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState*, MegamoeTmaGpuPlanResult*) +{ + return unsupported(); +} + +extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state) +{ + return state == nullptr || *state == nullptr ? static_cast(cudaSuccess) : unsupported(); +} + +extern "C" int megamoe_tma_copy_config_info(MegamoeTmaCopyState const*, MegamoeTmaCopyConfig*) +{ + return unsupported(); +} + +extern "C" char const* megamoe_tma_copy_error_string(int error) +{ + if (error == static_cast(cudaErrorNotSupported)) + { + return "Per-layer dynamic EPLB TMA requires CUDA 13.1+ and a requested SM100, SM103, or SM107 target"; + } + return cudaGetErrorString(static_cast(error)); +} +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h new file mode 100644 index 000000000000..ccb2f820b867 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h @@ -0,0 +1,380 @@ +/* + * 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. + */ + +#ifndef TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H +#define TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +// Copy-owned GPU policy adapter. The scheduler supplies only placement +// geometry. +// Fusing this adapter into its publication tail avoids another kernel launch or +// any CPU policy work. Without a bound copy channel, the adapter does nothing. +namespace sami_copy +{ +constexpr int kMaxEp = 32; +constexpr int kMinGroup = 4; + +struct PlacementView +{ + int count, capacity; + int const *experts, *levels, *begins, *sizes, *owners, *helpers; +}; + +__device__ __forceinline__ int imin(int a, int b) +{ + return a < b ? a : b; +} + +__device__ __forceinline__ int imax(int a, int b) +{ + return a > b ? a : b; +} + +__device__ __forceinline__ bool group_contains(int rank, int begin, int size) +{ + return rank >= begin && rank < begin + size; +} + +__host__ __device__ __forceinline__ int hierarchy_depth(int ep) +{ + int depth = 1; + for (int level = 1; (1 << level) <= ep; ++level) + { + int groups = 1 << level; + if (ep % groups || ep / groups < kMinGroup) + break; + ++depth; + } + return depth; +} + +__device__ __forceinline__ void store_release_system(int* address, int value) +{ + cuda::atomic_ref atom(*address); + atom.store(value, cuda::memory_order_release); +} + +/* ABI6 publishes explicit per-slot modes and bounded owner jobs. All scans + * cover the complete global plan, including broadcasts beyond one warp. Only + * the releasing lane stores mapped words, preserving the SYS release contract. + */ +template +__device__ void publish_plan_channel(int* channel, int const* ids, int const* levels, int const* owners, + PlacementView const& plan, int helpers, int ep, int local_rank, int expert_count, int plan_abi_version, + int route_features) +{ + if (!channel || threadIdx.x != 0) + return; + int const stride = (helpers + 3) & ~3; + int const broadcasts = plan.count; + int const* p_expert = plan.experts; + int const* p_level = plan.levels; + int const* p_begin = plan.begins; + int const* p_size = plan.sizes; + int const* p_owner = plan.owners; + int const* p_helper = plan.helpers; + int owner_count[kMaxEp] = {0}; + int status = plan_abi_version != 6 || (route_features & ~1) || broadcasts < 0 || broadcasts > plan.capacity; + int global_max_send = 0; + int common_size = broadcasts > 0 ? p_size[0] : 0; + int foreign_begin = -1, target_active = 0; + bool common_proper_size = common_size > 0 && common_size < ep; + bool one_foreign_target = true; + int const home = expert_count / ep; + for (int b = 0; b < broadcasts && !status; ++b) + { + int const level = p_level[b]; + int const size = level >= 0 && level < hierarchy_depth(ep) ? ep / (1 << level) : 0; + if (p_expert[b] < 0 || p_expert[b] >= expert_count || p_owner[b] != p_expert[b] / home || !size + || p_size[b] != size || p_begin[b] < 0 || p_begin[b] > ep - size || p_begin[b] % size || p_helper[b] < 0 + || p_helper[b] >= helpers) + { + status = 1; + break; + } + for (int j = 0; j < b; ++j) + { + bool const intersects = p_begin[j] < p_begin[b] + size && p_begin[b] < p_begin[j] + p_size[j]; + if (p_expert[j] == p_expert[b] || (intersects && p_helper[j] == p_helper[b])) + status = 1; + } + if (group_contains(local_rank, p_begin[b], size) + && (ids[p_helper[b]] != p_expert[b] || levels[p_helper[b]] != level || owners[p_helper[b]] != p_owner[b])) + status = 1; + int const count = ++owner_count[p_owner[b]]; + if (count > imin(helpers, home)) + status = 1; + global_max_send = imax(global_max_send, count); + common_proper_size &= size == common_size; + if (!group_contains(p_owner[b], p_begin[b], size)) + { + if (foreign_begin < 0) + foreign_begin = p_begin[b]; + else + one_foreign_target &= foreign_begin == p_begin[b]; + } + } + if (common_proper_size && foreign_begin >= 0) + for (int b = 0; b < broadcasts; ++b) + target_active += p_begin[b] == foreign_begin && p_size[b] == common_size; + bool const external = !status && (route_features & 1) && common_proper_size && one_foreign_target + && foreign_begin >= 0 && (global_max_send <= 1 || (global_max_send <= 2 && target_active >= 3)); + + for (int slot = 0; slot < stride; ++slot) + { + bool const inactive = status || slot >= helpers; + channel[4 + slot] = inactive ? -1 : ids[slot]; + channel[4 + stride + slot] = inactive ? -1 : levels[slot]; + channel[4 + 2 * stride + slot] = inactive ? -1 : owners[slot]; + channel[4 + 3 * stride + slot] = 0; + channel[4 + 4 * stride + slot] = -1; + channel[4 + 5 * stride + slot] = -1; + channel[4 + 6 * stride + slot] = -1; + } + int outgoing = 0; + for (int b = 0; b < broadcasts && !status; ++b) + { + int const owner = p_owner[b], begin = p_begin[b], size = p_size[b]; + bool const local_owner = group_contains(owner, begin, size); + int mode = 0; + if (local_owner) + { + bool pure = true, best = true; + for (int j = 0; j < broadcasts; ++j) + { + bool const intersects = p_begin[j] < begin + size && begin < p_begin[j] + p_size[j]; + if (intersects + && (p_begin[j] != begin || p_size[j] != size || !group_contains(p_owner[j], begin, size))) + pure = false; + if (p_owner[j] == owner && group_contains(owner, p_begin[j], p_size[j]) + && (p_size[j] > size || (p_size[j] == size && j < b))) + best = false; + } + if (pure || best) + mode = 1; + } + else if (external) + { + mode = 2; + if (owner == local_rank) + { + int target_id = begin / size; + for (int level = 1; level < p_level[b]; ++level) + target_id += 1 << level; + channel[4 + 4 * stride + outgoing] = p_expert[b]; + channel[4 + 5 * stride + outgoing] = target_id; + channel[4 + 6 * stride + outgoing] = p_helper[b]; + ++outgoing; + } + } + if (group_contains(local_rank, begin, size)) + channel[4 + 3 * stride + p_helper[b]] = mode; + } + channel[1] = global_max_send; + channel[2] = status; + channel[3] = outgoing; + if constexpr (PublishHost) + store_release_system(channel, 0); + else + channel[0] = 0; +} + +/* Small ABI6 plans fit in one warp. Load each record once and compare its + * register values through shuffles. Two mode bits per helper fit in a register + * for S<=16. Larger or malformed plans use the scalar publisher. + * Only lane 0 writes mapped words; its SYS release orders every field. */ +template +__device__ void publish_plan_channel_warp(int* channel, int const* ids, int const* levels, int const* owners, + PlacementView const& plan, int helpers, int ep, int local_rank, int expert_count, int plan_abi_version, + int route_features) +{ + if (!channel || (threadIdx.x >> 5) != 0) + return; + // One shared cold fallback keeps its scalar O(B^2) implementation out + // of the hot sections instead of inlining a copy at every failure gate. + do + { + int const broadcasts = plan.count; + if (broadcasts < 0 || broadcasts > plan.capacity || broadcasts > 32 || helpers <= 0 || helpers > 16 + || plan_abi_version != 6 || (route_features & ~1)) + { + break; + } + constexpr unsigned mask = 0xffffffffu; + int const lane = threadIdx.x & 31; + bool const active = lane < broadcasts; + int const expert = active ? plan.experts[lane] : -1; + int const level = active ? plan.levels[lane] : 0; + int const begin = active ? plan.begins[lane] : 0; + int const size = active ? plan.sizes[lane] : 0; + int const owner = active ? plan.owners[lane] : -1; + int const helper = active ? plan.helpers[lane] : -1; + // Coalesced helper reads avoid three dependent, short-circuited global + // loads for every local record. Invalid helpers select a harmless lane + // and are rejected below before any publication. + int const slot_id = lane < helpers ? ids[lane] : -1; + int const slot_level = lane < helpers ? levels[lane] : -1; + int const slot_owner = lane < helpers ? owners[lane] : -1; + int const check_id = __shfl_sync(mask, slot_id, helper & 31); + int const check_level = __shfl_sync(mask, slot_level, helper & 31); + int const check_owner = __shfl_sync(mask, slot_owner, helper & 31); + int const home = expert_count / ep; + bool const level_ok = static_cast(level) < static_cast(hierarchy_depth(ep)); + int const expected_size = level_ok ? ep >> level : 0; + bool const aligned + = (ep & (ep - 1)) == 0 ? (begin & (expected_size - 1)) == 0 : begin % imax(expected_size, 1) == 0; + // Unsigned interval checks stay defined even for malformed int32 fields. + // The native launch contract guarantees positive home and even EP<=32. + bool const owner_ok = (static_cast(owner) < static_cast(ep)) + & (static_cast(expert) - static_cast(owner) * static_cast(home) + < static_cast(home)); + bool const local + = static_cast(local_rank) - static_cast(begin) < static_cast(expected_size); + bool const record_valid = (expert >= 0) & (expert < expert_count) & owner_ok & (expected_size > 0) + & (size == expected_size) & (begin >= 0) & (begin <= ep - expected_size) & aligned & (helper >= 0) + & (helper < helpers); + bool const local_valid = (check_id == expert) & (check_level == level) & (check_owner == owner); + bool const valid = !active | (record_valid & (!local | local_valid)); + // Reject invalid addresses before the pairwise geometry calculations. + if (__any_sync(mask, !valid)) + { + break; + } + unsigned const same_owner = __match_any_sync(mask, owner); + int const owned = active ? __popc(same_owner) : 0; + int const global_max_send = __reduce_max_sync(mask, owned); + unsigned const same_expert = __match_any_sync(mask, expert); + int bad = active & ((__popc(same_expert) != 1) | (owned > imin(helpers, home))); + int pure = 1, best = 1; + // Validated EP<=32/S<=16 geometry needs only 19 bits. Exchange one + // packed register per pair rather than four independent shuffles. + unsigned const pair_record = static_cast(begin) | (static_cast(size - 1) << 5) + | (static_cast(owner) << 10) | (static_cast(helper) << 15); +// Eager bitwise predicates avoid divergent short-circuit branches and +// repeated reconvergence inside this register-only comparison loop. +// All lanes execute every shuffle, including inactive plan lanes. +#pragma unroll 1 + for (int j = 0; j < broadcasts; ++j) + { + unsigned const other = __shfl_sync(mask, pair_record, j); + int const other_begin = other & 31u; + int const other_size = ((other >> 5) & 31u) + 1; + int const other_owner = (other >> 10) & 31u; + int const other_helper = (other >> 15) & 15u; + int const intersects = (other_begin < begin + size) & (begin < other_begin + other_size); + bad |= active & (j < lane) & intersects & (other_helper == helper); + pure &= !intersects + | ((other_begin == begin) & (other_size == size) & (other_owner >= begin) + & (other_owner < begin + size)); + best &= !((other_owner == owner) & (owner >= other_begin) & (owner < other_begin + other_size) + & ((other_size > size) | ((other_size == size) & (j < lane)))); + } + if (__any_sync(mask, bad)) + { + // Preserve even the scalar error header/partial owner-count semantics. + break; + } + int const common_size = __shfl_sync(mask, size, 0); + bool const common_proper_size + = broadcasts > 0 && common_size < ep && !__any_sync(mask, active && size != common_size); + bool const local_owner = active && group_contains(owner, begin, size); + bool const foreign = active && !local_owner; + unsigned const foreign_mask = __ballot_sync(mask, foreign); + int const first_foreign = foreign_mask ? __ffs(foreign_mask) - 1 : 0; + int const foreign_begin = __shfl_sync(mask, begin, first_foreign); + bool const one_foreign_target = !__any_sync(mask, foreign && begin != foreign_begin); + int const target_active = __popc(__ballot_sync(mask, active && begin == foreign_begin && size == common_size)); + bool const external = (route_features & 1) && common_proper_size && one_foreign_target && foreign_mask + && (global_max_send <= 1 || (global_max_send <= 2 && target_active >= 3)); + int const mode = local_owner ? (pure | best) : (foreign && external ? 2 : 0); + unsigned const local_mode + = active && group_contains(local_rank, begin, size) ? static_cast(mode) << (2 * helper) : 0; + unsigned const mode_bits = __reduce_or_sync(mask, local_mode); + unsigned const outgoing_mask = __ballot_sync(mask, mode == 2 && owner == local_rank); + int const stride = (helpers + 3) & ~3; + // Reuse the coalesced helper registers from validation. All lanes + // participate in every shuffle; lane 0 remains the only mapped writer. + for (int slot = 0; slot < stride; slot += 4) + { + const int4 ids4 = make_int4(__shfl_sync(mask, slot_id, slot), __shfl_sync(mask, slot_id, slot + 1), + __shfl_sync(mask, slot_id, slot + 2), __shfl_sync(mask, slot_id, slot + 3)); + const int4 levels4 = make_int4(__shfl_sync(mask, slot_level, slot), __shfl_sync(mask, slot_level, slot + 1), + __shfl_sync(mask, slot_level, slot + 2), __shfl_sync(mask, slot_level, slot + 3)); + const int4 owners4 = make_int4(__shfl_sync(mask, slot_owner, slot), __shfl_sync(mask, slot_owner, slot + 1), + __shfl_sync(mask, slot_owner, slot + 2), __shfl_sync(mask, slot_owner, slot + 3)); + if (lane == 0) + { + // Native channel/stride alignment permits 16-byte stores. + const int4 empty = make_int4(-1, -1, -1, -1); + *reinterpret_cast(channel + 4 + slot) = ids4; + *reinterpret_cast(channel + 4 + stride + slot) = levels4; + *reinterpret_cast(channel + 4 + 2 * stride + slot) = owners4; + unsigned const packed = mode_bits >> (2 * slot); + *reinterpret_cast(channel + 4 + 3 * stride + slot) + = make_int4(packed & 3u, (packed >> 2) & 3u, (packed >> 4) & 3u, (packed >> 6) & 3u); + *reinterpret_cast(channel + 4 + 4 * stride + slot) = empty; + *reinterpret_cast(channel + 4 + 5 * stride + slot) = empty; + *reinterpret_cast(channel + 4 + 6 * stride + slot) = empty; + } + } + // Canonical target IDs enumerate proper hierarchy levels, widest first. + int const target = mode == 2 ? begin / size + (1 << level) - 2 : -1; + int outgoing = 0; + for (unsigned remaining = outgoing_mask; remaining; remaining &= remaining - 1) + { + int const source_lane = __ffs(remaining) - 1; + int const out_expert = __shfl_sync(mask, expert, source_lane); + int const out_target = __shfl_sync(mask, target, source_lane); + int const out_helper = __shfl_sync(mask, helper, source_lane); + if (lane == 0) + { + channel[4 + 4 * stride + outgoing] = out_expert; + channel[4 + 5 * stride + outgoing] = out_target; + channel[4 + 6 * stride + outgoing] = out_helper; + } + ++outgoing; + } + if (lane == 0) + { + channel[1] = global_max_send; + channel[2] = 0; + channel[3] = outgoing; + if constexpr (PublishHost) + store_release_system(channel, 0); + else + channel[0] = 0; + } + return; + } while (false); + publish_plan_channel( + channel, ids, levels, owners, plan, helpers, ep, local_rank, expert_count, plan_abi_version, route_features); +} + +} // namespace sami_copy + +} // namespace kernels + +TRTLLM_NAMESPACE_END + +#endif // TRTLLM_KERNELS_MOE_REBALANCE_PLAN_POLICY_H diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu new file mode 100644 index 000000000000..cbd487168663 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.cu @@ -0,0 +1,1316 @@ +/* + * 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. + */ +/* SM-driven unicast -> multicast copy, PTX 9.1 / CUDA 13.1 or newer. + * One aligned TMA load handles each eligible 8 KiB logical slice. Physical + * scratch slots include 16 bytes of padding for aligned head/tail accesses. + * + * Each communication warp owns a near-equal share of all whole 8 KiB slots, split into two + * banks. A bank is one bulk group containing all of its 8 KiB slice slots. + * A producer warp issues UC loads; its consumer warp issues MC stores. + * Per-bank alternating tickets protect slot metadata and shared-memory reuse. + * This permits every allocated slot to participate even when one warp owns + * many slots, without exceeding the ISA's outstanding bulk-group limit. + * + * The copy runs as the second phase of one kernel whose first phase is the + * HALO-Q scheduler that produces its plan. + */ +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQKernel.cuh" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalancePlanPolicy.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#if CUDART_VERSION < 13010 +#error "TMA multicast copy requires CUDA >= 13.1 (PTX >= 9.1)" +#endif +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 900 +#error "TMA multicast copy requires SM90 or newer" +#endif + +namespace +{ +constexpr unsigned kSlice = MEGAMOE_TMA_COPY_SLICE_BYTES; +constexpr unsigned kSlotControlBytes = MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES; +constexpr unsigned kSlotDataBytes = MEGAMOE_TMA_COPY_SLOT_DATA_BYTES; +constexpr unsigned kCoalesced = 4U; +constexpr unsigned kFallbackBulk = 8U; + +struct alignas(16) SlotControl +{ + unsigned long long barrier; + unsigned phase; + unsigned reserved; // head=1, tail=2, coalesced=4, fallback-has-bulk=8 + uint64_t fast_dst; + unsigned fast_bytes; + unsigned fast_padding; // first slot in each bank: even=free, odd=ready ticket +}; + +static_assert(kSlotControlBytes == 32, "cached consume requires 32-byte control"); +static_assert(offsetof(SlotControl, fast_dst) == 16, "cached destination alignment"); +static_assert(offsetof(SlotControl, fast_bytes) == 24, "cached byte-count offset"); +static_assert(kSlotDataBytes == kSlice + 16, "one aligned edge of slot padding"); +static_assert(sizeof(SlotControl) == kSlotControlBytes, "shared slot control ABI"); +static_assert(kSlotDataBytes % 16 == 0, "TMA slot alignment"); +static_assert(sizeof(MegamoeTmaCopySegment) == 4 * sizeof(uint64_t), "segment ABI"); +static_assert(sizeof(MegamoeTmaCopyRange) == 9 * sizeof(uint64_t), "range ABI"); +static_assert(sizeof(MegamoeTmaCopyConfig) == 88, "configuration ABI"); + +__device__ __forceinline__ unsigned shared_address(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ void initialize_barrier(SlotControl* slot) +{ + unsigned bar = shared_address(&slot->barrier); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + slot->phase = 0; + slot->reserved = 0; + slot->fast_padding = 0; +} + +__device__ __forceinline__ void expect_slice(SlotControl* slot, unsigned bytes) +{ + unsigned bar = shared_address(&slot->barrier); + unsigned long long arrival; + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 %0, [%1], %2;" + : "=l"(arrival) + : "r"(bar), "r"(bytes) + : "memory"); + asm volatile("" ::"l"(arrival)); +} + +__device__ __forceinline__ void load_fragment(unsigned char* buffer, SlotControl* slot, uint64_t src, unsigned bytes) +{ + unsigned bar = shared_address(&slot->barrier); + unsigned smem = shared_address(buffer); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];" ::"r"(smem), + "l"(src), "r"(bytes), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void await_slice(SlotControl* slot) +{ + unsigned ready; + unsigned bar = shared_address(&slot->barrier); + unsigned phase = slot->phase; + do + { + asm volatile( + "{ .reg .pred p; " + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], %2; " + "selp.u32 %0, 1, 0, p; }" + : "=r"(ready) + : "r"(bar), "r"(phase) + : "memory"); + if (!ready) + __nanosleep(32); + } while (!ready); + // Only a slot that actually issued and completed TMA advances its phase. + slot->phase = phase ^ 1U; +} + +__device__ __forceinline__ void multicast_slice(uint64_t dst, unsigned char const* buffer, unsigned bytes) +{ + unsigned smem = shared_address(buffer); + asm volatile( + "multimem.cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst), "r"(smem), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ void multicast_edge(uint64_t dst, unsigned char const* buffer, uint16_t mask) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + unsigned smem = shared_address(buffer); + asm volatile( + "multimem.cp.async.bulk.global.shared::cta.bulk_group.cp_mask " + "[%0], [%1], 16, %2;" ::"l"(dst), + "r"(smem), "h"(mask) + : "memory"); +#else + // coalesced_copy_flags() is zero on these targets; never use .cp_mask on SM90. + asm volatile("trap;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void wait_older_reads() +{ + asm volatile("cp.async.bulk.wait_group.read 1;" ::: "memory"); +} + +__device__ __forceinline__ void wait_all_reads() +{ + // Buffer reuse only requires completion of the SMEM reads. Destination + // writes are drained separately before publishing READY. + asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory"); +} + +__device__ __forceinline__ void wait_all_writes() +{ + // Deliberately no .read: the multicast destination writes must be complete. + asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); +} + +__device__ __forceinline__ void system_publish_fence() +{ + asm volatile("fence.proxy.alias;" ::: "memory"); + asm volatile("fence.release.sys;" ::: "memory"); +} + +__device__ __forceinline__ void multicast_generation(uint64_t flag_mc, uint64_t generation) +{ + // The notification uses the same MC backing/UC acquire protocol as the + // former copy-engine terminal. Publish only after all issuers' writes. + system_publish_fence(); + asm volatile("multimem.st.release.sys.global.u64 [%0], %1;" ::"l"(flag_mc), "l"(generation) : "memory"); + asm volatile("fence.proxy.alias;" ::: "memory"); +} + +/* Helper residency of one layer: the expert that each rank's helper slot holds, + * int32 [world][helper_count], -1 when unknown. Every rank keeps an identical + * table because it changes only from the replicated HALO-Q placement. */ +struct HelperResidency +{ + int* table; + int const* placement; // HALO-Q workspace: epoch, count, six capacity-strided fields. + int capacity; + int helpers; + int world; +}; + +// After a copy, every member of each broadcast's team holds its expert in the +// broadcast's helper slot. +__device__ __forceinline__ void update_helper_residency(HelperResidency const& residency) +{ + int const count = __ldcg(residency.placement + 1); + int const used = count < 0 ? 0 : (count > residency.capacity ? residency.capacity : count); + int const* fields = residency.placement + 2; + for (int broadcast = 0; broadcast < used; ++broadcast) + { + int const expert = __ldcg(fields + broadcast); + int const begin = __ldcg(fields + 2 * residency.capacity + broadcast); + int const size = __ldcg(fields + 3 * residency.capacity + broadcast); + int const helper = __ldcg(fields + 5 * residency.capacity + broadcast); + if (helper < 0 || helper >= residency.helpers || begin < 0 || size <= 0 || begin + size > residency.world) + continue; + for (int rank = begin; rank < begin + size; ++rank) + residency.table[rank * residency.helpers + helper] = expert; + } +} + +__device__ __forceinline__ void finish_cta_and_notify( + unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation, HelperResidency const* residency = nullptr) +{ + unsigned previous; + // A release/acquire RMW chain combines the completed payload of all CTAs. + // No CTA waits for another: oversubscribed grids cannot deadlock here. + asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], 1;" : "=r"(previous) : "l"(completed_ctas) : "memory"); + if (previous == gridDim.x - 1) + { + // All RMWs are done. The next launch is ordered after this kernel, so + // the last CTA can rearm the counter without a host memset each time. + asm volatile("st.relaxed.gpu.global.u32 [%0], 0;" ::"l"(completed_ctas) : "memory"); + multicast_generation(flag_mc, generation); + // Every CTA read the table before joining the chain; the next copy of + // this layer is a later kernel on the same stream. + if (residency != nullptr) + update_helper_residency(*residency); + } +} + +struct Slice +{ + uint64_t begin; + uint64_t end; + uint64_t segment_begin; + uint64_t segment_end; +}; + +struct Fragment +{ + uint64_t src; + uint64_t dst; + unsigned bytes; + unsigned scalar_prefix; + unsigned bulk_bytes; + unsigned smem_offset; +}; + +__device__ __forceinline__ Slice locate_slice( + MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, uint64_t range_count, uint64_t ordinal) +{ + uint64_t lo = 0, hi = range_count; + while (lo < hi) + { + uint64_t mid = lo + (hi - lo) / 2; + if (ordinal < ranges[mid].prefix_end) + hi = mid; + else + lo = mid + 1; + } + const MegamoeTmaCopyRange range = ranges[lo]; + const uint64_t local = ordinal - (range.prefix_end - range.slice_count); + uint64_t slice = range.first_slice + local * range.slice_stride; + if (range.slice_divisor != 1) + { + // floor(n/d) from a cold-computed reciprocal; the low estimate is at + // most one short. No runtime integer division in the TMA issue loop. + const uint64_t estimate = __umul64hi(slice, range.slice_reciprocal); + slice = estimate + (slice - estimate * range.slice_divisor >= range.slice_divisor); + } + const uint64_t begin = slice * kSlice; + const uint64_t remaining = range.bytes - begin; + const uint64_t end = begin + (remaining < kSlice ? remaining : kSlice); + const uint64_t segment_end = range.segment_begin + range.segment_count; + lo = range.segment_begin; + hi = segment_end; + while (lo < hi) + { + const uint64_t mid = lo + (hi - lo) / 2; + if (segments[mid].virtual_begin + segments[mid].bytes <= begin) + lo = mid + 1; + else + hi = mid; + } + return Slice{begin, end, lo, segment_end}; +} + +__device__ __forceinline__ Fragment locate_fragment(Slice const& slice, MegamoeTmaCopySegment const& segment) +{ + const uint64_t begin = slice.begin > segment.virtual_begin ? slice.begin : segment.virtual_begin; + const uint64_t segment_end = segment.virtual_begin + segment.bytes; + const uint64_t end = slice.end < segment_end ? slice.end : segment_end; + unsigned const bytes = static_cast(end - begin); + const uint64_t offset = begin - segment.virtual_begin; + const uint64_t src = segment.src + offset, dst = segment.dst + offset; + // UC -> SMEM alignment depends only on the source. A differently aligned + // MC destination can consume the staged data with exact 32-bit stores. + unsigned prefix = static_cast((16U - (src & 15U)) & 15U); + if (prefix > bytes) + prefix = bytes; + unsigned const bulk = (bytes - prefix) & ~15U; + // The virtual concat is NOT a physical address space. Align each bulk + // fragment's scratch area independently; a preceding 4/8/12-byte plane + // must not force the next large weight onto the scalar path. + // end_of_bulk <= align_up(fragment_virtual_end,16) <= 8192, and the + // next fragment's aligned start is no smaller. Scratch regions are disjoint. + unsigned const smem_offset = static_cast((begin - slice.begin + prefix + 15U) & ~uint64_t(15U)); + return Fragment{src, dst, bytes, prefix, bulk, smem_offset}; +} + +// Only a virtual slice entirely contained in one physical segment can use +// one padded, aligned TMA load. Other slices retain exact scalar edge copies. +// The NOMINAL aligned 16-byte accesses, not just their masks, must lie within +// the original source and destination intervals. No small-plane overread. +__device__ __forceinline__ unsigned coalesced_copy_flags( + Slice const& slice, MegamoeTmaCopySegment const& segment, Fragment const& f) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + if (slice.begin < segment.virtual_begin || slice.end > segment.virtual_begin + segment.bytes + || ((f.src ^ f.dst) & 15U) || segment.bytes < 16) + return 0; + unsigned const tail = f.bytes - f.scalar_prefix - f.bulk_bytes; + unsigned const flags = (f.scalar_prefix ? 1U : 0U) | (tail ? 2U : 0U); + const uint64_t first_src = f.src & ~uint64_t(15); + const uint64_t first_dst = f.dst & ~uint64_t(15); + const uint64_t last_src = (f.src + f.bytes - 1) & ~uint64_t(15); + const uint64_t last_dst = (f.dst + f.bytes - 1) & ~uint64_t(15); + // Subtraction form also avoids overflow of a one-past/rounded-up pointer. + if (first_src < segment.src || first_dst < segment.dst || last_src < segment.src || last_dst < segment.dst + || last_src - segment.src > segment.bytes - 16 || last_dst - segment.dst > segment.bytes - 16) + return 0; + return flags | kCoalesced; +#else + return 0; +#endif +} + +__device__ __forceinline__ uint16_t edge_byte_mask(unsigned begin, unsigned bytes) +{ + // begin+bytes <= 16; the shift is evaluated in 32 bits. + return static_cast(((1U << bytes) - 1U) << begin); +} + +__device__ __forceinline__ void prefetch_bank(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t ordinal_base, uint64_t ordinal_stride, uint64_t sequence, unsigned char* warp_buffer, + SlotControl* warp_control, int slot_start, int count) +{ + for (int i = 0; i < count; ++i) + { + const Slice slice = locate_slice(segments, ranges, range_count, ordinal_base + ordinal_stride * (sequence + i)); + SlotControl* slot = warp_control + slot_start + i; + unsigned char* slot_buffer = warp_buffer + static_cast(slot_start + i) * kSlotDataBytes; + const MegamoeTmaCopySegment first_segment = segments[slice.segment_begin]; + const Fragment first = locate_fragment(slice, first_segment); + unsigned const fast = coalesced_copy_flags(slice, first_segment, first); + if (fast) + { + // ONE aligned TMA load supplies head, body and tail. The only extra + // storage is the final 16B of this physical slot; logical RR remains 8KiB. + slot->reserved = fast; + slot->fast_dst = first.dst; + slot->fast_bytes = first.bytes; + unsigned const load_bytes = (static_cast(first.src & 15U) + first.bytes + 15U) & ~15U; + expect_slice(slot, load_bytes); + load_fragment(slot_buffer, slot, first.src & ~uint64_t(15), load_bytes); + continue; + } + unsigned bulk_bytes = first.bulk_bytes; + for (uint64_t index = slice.segment_begin + 1; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + bulk_bytes += locate_fragment(slice, segments[index]).bulk_bytes; + slot->reserved = bulk_bytes ? kFallbackBulk : 0U; + if (!bulk_bytes) + continue; + expect_slice(slot, bulk_bytes); + if (first.bulk_bytes) + load_fragment(slot_buffer + first.smem_offset, slot, first.src + first.scalar_prefix, first.bulk_bytes); + for (uint64_t index = slice.segment_begin + 1; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + { + const Fragment f = locate_fragment(slice, segments[index]); + if (f.bulk_bytes) + load_fragment(slot_buffer + f.smem_offset, slot, f.src + f.scalar_prefix, f.bulk_bytes); + } + } +} + +__device__ __forceinline__ void scalar_words(Fragment const& f, unsigned begin, unsigned end) +{ + for (unsigned off = begin; off < end; off += 4) + { + unsigned value; + // UC means unicast. .cg avoids a stale read-only (.nc) cache. + asm volatile("ld.global.cg.u32 %0, [%1];" : "=r"(value) : "l"(f.src + off) : "memory"); + asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" ::"l"(f.dst + off), "r"(value) : "memory"); + } +} + +__device__ __forceinline__ void multicast_shared_words(Fragment const& f, unsigned char const* buffer) +{ + // The slot mbarrier has completed: TMA data is visible to generic shared + // loads. Only the source-aligned body is staged; no rounded overread. + for (unsigned off = 0; off < f.bulk_bytes; off += 4) + { + unsigned value; + asm volatile("ld.shared.u32 %0, [%1];" : "=r"(value) : "r"(shared_address(buffer + off)) : "memory"); + asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" ::"l"(f.dst + f.scalar_prefix + off), "r"(value) + : "memory"); + } +} + +__device__ __forceinline__ void consume_bank(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t ordinal_base, uint64_t ordinal_stride, uint64_t sequence, unsigned char* warp_buffer, + SlotControl* warp_control, int slot_start, int count, int bank, bool* pending, int& outstanding) +{ + bool issued = false; + for (int i = 0; i < count; ++i) + { + SlotControl* slot = warp_control + slot_start + i; + unsigned const flags = slot->reserved; + unsigned char const* slot_buffer = warp_buffer + static_cast(slot_start + i) * kSlotDataBytes; + if (flags & kCoalesced) + { + // Prefetch cached these values before issuing the TMA load. + // No range search, segment search, or descriptor reload on consume. + const uint64_t dst = slot->fast_dst; + unsigned const bytes = slot->fast_bytes; + unsigned const shift = static_cast(dst & 15U); + unsigned prefix = (16U - shift) & 15U; + if (prefix > bytes) + prefix = bytes; + unsigned const bulk = (bytes - prefix) & ~15U; + unsigned const tail_offset = prefix + bulk; + const uint16_t head_mask = edge_byte_mask(shift, prefix); + const uint16_t tail_mask = edge_byte_mask(0, bytes - tail_offset); + await_slice(slot); + if (bulk) + multicast_slice(dst + prefix, slot_buffer + shift + prefix, bulk); + if (flags & 1U) + multicast_edge(dst & ~uint64_t(15), slot_buffer, head_mask); + if (flags & 2U) + multicast_edge(dst + tail_offset, slot_buffer + shift + tail_offset, tail_mask); + issued = true; + continue; + } + const Slice slice = locate_slice(segments, ranges, range_count, ordinal_base + ordinal_stride * (sequence + i)); + // Prefetch already determined whether this slice has any async data: + // no duplicate fragment scan or second range lookup for scalar edges. + if (flags & kFallbackBulk) + await_slice(slot); + for (uint64_t index = slice.segment_begin; + index < slice.segment_end && segments[index].virtual_begin < slice.end; ++index) + { + const Fragment f = locate_fragment(slice, segments[index]); + if (f.bulk_bytes) + { + if (((f.dst + f.scalar_prefix) & 15U) == 0) + { + multicast_slice(f.dst + f.scalar_prefix, slot_buffer + f.smem_offset, f.bulk_bytes); + issued = true; + } + else + { + // Different UC/MC alignment cannot use a bulk MC store. + // Reuse the prefetched bytes instead of serial remote loads. + multicast_shared_words(f, slot_buffer + f.smem_offset); + } + } + // These exact scalar intervals are disjoint from ALL TMA writes. + // No intermediate group drain: the final full-write wait and fence cover them. + scalar_words(f, 0, f.scalar_prefix); + scalar_words(f, f.scalar_prefix + f.bulk_bytes, f.bytes); + } + } + if (issued) + { + asm volatile("cp.async.bulk.commit_group;" ::: "memory"); + pending[bank] = true; + ++outstanding; + } +} + +// The two peers are the only writers of a bank ticket. Unsigned wrap is +// safe: a peer cannot publish its next ticket until the other responds. +__device__ __forceinline__ void wait_bank_ticket(unsigned const* address, unsigned expected) +{ + unsigned observed; + do + { + asm volatile("ld.acquire.cta.shared::cta.u32 %0, [%1];" + : "=r"(observed) + : "r"(shared_address(address)) + : "memory"); + if (observed != expected) + __nanosleep(32); + } while (observed != expected); +} + +__device__ __forceinline__ void publish_bank_ticket(unsigned* address, unsigned value) +{ + asm volatile("st.release.cta.shared::cta.u32 [%0], %1;" ::"r"(shared_address(address)), "r"(value) : "memory"); +} + +__device__ __forceinline__ void tma_copy_body(MegamoeTmaCopySegment const* segments, MegamoeTmaCopyRange const* ranges, + uint64_t range_count, uint64_t total_slices, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, + uint64_t generation, HelperResidency const* residency) +{ + if (!total_slices) + { + // Empty source ranks still publish READY. One tiny CTA, same kernel; + // no descriptor upload or additional host-side notification launch. + if (threadIdx.x == 0 && flag_mc) + multicast_generation(flag_mc, generation); + return; + } + extern __shared__ __align__(128) unsigned char smem[]; + int const physical_warp = threadIdx.x / 32, lane = threadIdx.x % 32; + int const warp = physical_warp / 2; // logical RR worker, shared by its two warps + bool const producer = (physical_warp & 1) == 0; + // Warps beyond the 2 * warps copy workers only join the CTA barriers. + bool const worker = warp < warps; + int const slots_per_warp = megamoe_tma_warp_slot_count(total_slots, warps, warp); + int const first_slot = megamoe_tma_warp_slot_begin(total_slots, warps, warp); + SlotControl* control = reinterpret_cast(smem + static_cast(total_slots) * kSlotDataBytes); + SlotControl* warp_control = control + first_slot; + unsigned char* warp_buffer = smem + static_cast(first_slot) * kSlotDataBytes; + if (worker && producer && lane == 0) + { + for (int slot = 0; slot < slots_per_warp; ++slot) + initialize_barrier(warp_control + slot); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + } + __syncthreads(); + // Elect once for the entire pipeline so ptxas can prove one TMA issuer per warp. + if (worker && cuda::ptx::elect_sync(0xffffffffu)) + { + // CTA/SM is the fastest-varying dimension, then warp, then sequence. + const uint64_t ordinal_base = megamoe_tma_worker_first_slice(blockIdx.x, gridDim.x, warp); + const uint64_t ordinal_stride = megamoe_tma_worker_slice_stride(gridDim.x, warps); + const uint64_t tasks = ordinal_base < total_slices ? 1 + (total_slices - 1 - ordinal_base) / ordinal_stride : 0; + int const capacity[2] = {(slots_per_warp + 1) / 2, slots_per_warp / 2}; + int const start[2] = {0, capacity[0]}; + unsigned ticket[2] = {producer ? 0U : 1U, producer ? 0U : 1U}; + uint64_t next_sequence = 0; + if (producer) + { + // Each bank is independent: bank 1 can already be loading while + // the consumer drains bank 0's shared-memory reads and releases it. + for (int bank = 0; next_sequence < tasks; bank ^= 1) + { + const uint64_t left = tasks - next_sequence; + int const count + = static_cast(left < static_cast(capacity[bank]) ? left : capacity[bank]); + unsigned* handshake = &warp_control[start[bank]].fast_padding; + wait_bank_ticket(handshake, ticket[bank]); + prefetch_bank(segments, ranges, range_count, ordinal_base, ordinal_stride, next_sequence, warp_buffer, + warp_control, start[bank], count); + // Publish generic slot metadata after all TMA loads have been + // issued. The consumer still waits on each load's mbarrier. + publish_bank_ticket(handshake, ticket[bank] + 1U); + ticket[bank] += 2U; + next_sequence += count; + } + } + else + { + bool pending[2] = {false, false}; + int outstanding = 0; + int held_bank = -1; + for (int bank = 0; next_sequence < tasks; bank ^= 1) + { + const uint64_t left = tasks - next_sequence; + int const count + = static_cast(left < static_cast(capacity[bank]) ? left : capacity[bank]); + unsigned* handshake = &warp_control[start[bank]].fast_padding; + wait_bank_ticket(handshake, ticket[bank]); + consume_bank(segments, ranges, range_count, ordinal_base, ordinal_stride, next_sequence, warp_buffer, + warp_control, start[bank], count, bank, pending, outstanding); + if (held_bank >= 0) + { + // Consume both banks before releasing the older one. The + // producer refills it while this bank's MC reads remain + // outstanding; at most two read-pending groups exist. + if (pending[held_bank]) + { + if (outstanding == 2) + wait_older_reads(); + else + wait_all_reads(); // this bank was scalar-only + pending[held_bank] = false; + --outstanding; + } + publish_bank_ticket(&warp_control[start[held_bank]].fast_padding, ticket[held_bank] + 1U); + ticket[held_bank] += 2U; + } + held_bank = bank; + next_sequence += count; + } + // The final bank has no following consume to release it. Flush + // remaining reads and publish its final free ticket explicitly. + if (held_bank >= 0) + { + if (pending[held_bank]) + wait_all_reads(); + publish_bank_ticket(&warp_control[start[held_bank]].fast_padding, ticket[held_bank] + 1U); + } + // READY still covers destination completion, not merely SMEM reads. + wait_all_writes(); + // Every issuer joins through the CTA barrier and the acq_rel + // counter chain before the final CTA publishes READY. + } + } + __syncthreads(); + if (threadIdx.x == 0 && flag_mc) + finish_cta_and_notify(completed_ctas, flag_mc, generation, residency); +} + +// Per-CTA scratch is allocated once. No plan tensor is copied between kernels. +struct GpuDirectStorage +{ + MegamoeTmaGpuPlanConfig config; + int const *ids, *levels, *owners, *workspace, *status, *grid_sync; + int capacity; + uint64_t max_segments; + int* scratch; + MegamoeTmaCopySegment* segments; + MegamoeTmaCopyRange* ranges; + MegamoeTmaGpuPlanResult* results; + // Helper residency table, or null when helper slots are shared across layers. + int* resident; +}; + +__device__ __forceinline__ void tma_copy_gpu_direct_body(GpuDirectStorage const& direct, int warps, int total_slots, + unsigned* completed_ctas, uint64_t flag_mc, uint64_t generation) +{ + auto const& c = direct.config; + const uint64_t index = blockIdx.x; + int* plan = direct.scratch + index * (c.plan_words + 4ULL * c.helper_count); + auto* segments = direct.segments + index * direct.max_segments; + auto* ranges = direct.ranges + index * direct.max_segments; + auto* result = direct.results + index; + __shared__ int schedulerPlanReady; + if (threadIdx.x == 0) + { + bool statusOk = true; + for (int field = 0; field < 6; ++field) + statusOk &= direct.status[field] == 0; + auto const planEpoch = static_cast(direct.workspace[0]); + auto const schedulerEpoch = static_cast(direct.grid_sync[1]); + schedulerPlanReady = statusOk && planEpoch == schedulerEpoch; + if (!schedulerPlanReady && blockIdx.x == 0) + printf("GPU-direct in-switch copy: stale or failed scheduler plan; no READY published\n"); + } + __syncthreads(); + if (!schedulerPlanReady) + { + asm volatile("trap;" ::: "memory"); + return; + } + // HALO-Q's existing workspace: epoch, count, six capacity-strided fields. + int const* fields = direct.workspace + 2; + int const cap = direct.capacity; + tensorrt_llm::kernels::sami_copy::PlacementView placement{direct.workspace[1], cap, fields, fields + cap, + fields + 2 * cap, fields + 3 * cap, fields + 4 * cap, fields + 5 * cap}; + // Private routing scratch, not a CPU channel: no SYS publication, extra + // CUDA operation or scheduler-side copy of the placement outputs. + tensorrt_llm::kernels::sami_copy::publish_plan_channel_warp(plan, direct.ids, direct.levels, direct.owners, + placement, c.helper_count, c.world, c.rank, c.global_experts, c.plan_abi_version, c.route_features); + __syncthreads(); + if (threadIdx.x == 0) + { + megamoe_tma_build_gpu_plan(&c, plan, plan + c.plan_words, segments, direct.max_segments, ranges, + direct.max_segments, result, direct.resident); + // Malformed plans must never publish READY. The CUDA error surfaces at + // the caller's normal completion check, without a host wait in submit. + if (result->error) + { + printf("GPU-direct in-switch copy: invalid plan, error=%d; no READY published\n", result->error); + asm volatile("trap;" ::: "memory"); + } + } + __syncthreads(); + HelperResidency const residency{direct.resident, direct.workspace, cap, c.helper_count, c.world}; + HelperResidency const* const update = direct.resident != nullptr ? &residency : nullptr; + if (!result->total_slices) + { + // Work size is device-derived, so the launch always has config.sms CTAs. + // Empty ranks must still join the last-CTA chain before one READY store. + if (threadIdx.x == 0) + finish_cta_and_notify(completed_ctas, flag_mc, generation, update); + return; + } + tma_copy_body(segments, ranges, result->range_count, result->total_slices, warps, total_slots, completed_ctas, + flag_mc, generation, update); +} + +// Scheduler arguments of the fused launch, as halo_q_scheduler_body takes them. +struct FusedSchedulerArgs +{ + int const* routes; + int* out_slots; + int* out_ids; + int* out_levels; + int* out_owners; + unsigned long long const* peer_bases; + int* status; + int* partial; + int* route_aux; + int* grid_sync; + int* plan; + int* route_prefix; + unsigned long long spin_cycles; + int ep; + int experts; + int helpers; + int local_rank; + int route_capacity; + int capacity; + int route_count; +}; + +/* HALO-Q and the GPU-direct copy of its plan in one grid of config.sms CTAs. + * Every CTA first takes its scheduler role (CTA 0 plans, the others write + * routes) and then adds one to route_epoch. Route consumers on other streams + * wait until route_epoch reaches route_target with a stream memory wait, which + * holds no SM. The kernel runs in the remote memory-sync domain while its + * consumers run in the default one, so the release is system-scope. Each CTA + * also waits for route_target before it reads the plan, which CTA 0 finishes + * last, with the scheduler's spin bound. */ +__global__ void __launch_bounds__(512, 1) halo_q_tma_fused_kernel(FusedSchedulerArgs const scheduler, + GpuDirectStorage const direct, int warps, int total_slots, unsigned* completed_ctas, uint64_t flag_mc, + uint64_t generation, unsigned long long* route_epoch, unsigned long long route_target) +{ + tensorrt_llm::kernels::halo_q_scheduler_body(scheduler.routes, scheduler.out_slots, scheduler.out_ids, + scheduler.out_levels, scheduler.out_owners, scheduler.peer_bases, scheduler.status, scheduler.partial, + scheduler.route_aux, scheduler.grid_sync, scheduler.plan, scheduler.route_prefix, scheduler.ep, + scheduler.experts, scheduler.helpers, scheduler.local_rank, scheduler.route_capacity, + static_cast(gridDim.x), scheduler.spin_cycles, scheduler.capacity, scheduler.route_count); + // Orders every thread's scheduler writes before thread 0's release. + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned long long observed; + asm volatile("atom.release.sys.global.add.u64 %0, [%1], 1;" : "=l"(observed) : "l"(route_epoch) : "memory"); + long long const start = clock64(); + for (++observed; static_cast(observed - route_target) < 0;) + { + // A CTA that never became resident would hang the copy; fail like a stale plan. + if (clock64() - start > static_cast(scheduler.spin_cycles)) + { + printf("GPU-direct in-switch copy: route epoch timeout; no READY published\n"); + asm volatile("trap;" ::: "memory"); + } + __nanosleep(32); + asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(observed) : "l"(route_epoch) : "memory"); + } + } + __syncthreads(); + tma_copy_gpu_direct_body(direct, warps, total_slots, completed_ctas, flag_mc, generation); +} + +} // namespace + +struct MegamoeTmaCopyState +{ + MegamoeTmaCopyConfig config{}; + unsigned* completed_ctas = nullptr; + cudaStream_t last_stream = nullptr; + bool work_queued = false; + cudaError_t poison = cudaSuccess; + GpuDirectStorage direct{}; + uint64_t* direct_tables = nullptr; + bool direct_configured = false; + bool direct_bound = false; + // The fused kernel, prepared by megamoe_tma_copy_bind_gpu_direct. + void const* fused_kernel = nullptr; + int fused_total_slots = 0; + int fused_dynamic_bytes = 0; + cudaLaunchAttribute fused_attributes[2]{}; + unsigned fused_attribute_count = 0; + std::mutex submit_mutex; +}; + +namespace +{ + +// Every fused launch has one HALO-Q CTA of this size per SM. +constexpr int kFusedThreads = 512; + +/* Loads the fused kernel and reserves its local memory at cold setup: doing + * either at the first launch waits for the device, which deadlocks against + * kernels that spin on work enqueued after that launch. Also resolves the slots + * that fit beside both phases' static shared memory; HALO-Q reuses the same + * dynamic allocation before the copy initializes its slots. */ +cudaError_t prepare_fused_launch(MegamoeTmaCopyState& state) +{ + auto const& c = state.config; + auto const& plan = state.direct.config; + void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel); + cudaFuncAttributes attributes{}; + cudaError_t error = cudaFuncGetAttributes(&attributes, kernel); + if (error != cudaSuccess) + return error; + size_t stack = 0; + error = cudaDeviceGetLimit(&stack, cudaLimitStackSize); + if (error == cudaSuccess && stack < attributes.localSizeBytes) + error = cudaDeviceSetLimit(cudaLimitStackSize, attributes.localSizeBytes); + if (error != cudaSuccess) + return error; + int const static_bytes = static_cast(attributes.sharedSizeBytes); + int const slots = megamoe_tma_copy_total_slots(c.device_optin_shared_bytes - static_bytes); + int const capacity = tensorrt_llm::kernels::broadcast_capacity(plan.world, plan.global_experts, plan.helper_count); + int const scheduler_bytes = static_cast(sizeof(int)) + * tensorrt_llm::kernels::halo_q_shared_ints(plan.world, plan.global_experts, capacity, kFusedThreads); + int const slot_bytes = static_cast((slots * (kSlotDataBytes + kSlotControlBytes) + 127) & ~127U); + int const bytes = std::max(slot_bytes, scheduler_bytes); + if (attributes.maxThreadsPerBlock < kFusedThreads || slots < 2 * c.warps + || bytes + static_bytes > c.device_optin_shared_bytes) + return cudaErrorInvalidConfiguration; + error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + if (error == cudaSuccess) + error = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + int active = 0; + if (error == cudaSuccess) + error = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, kernel, kFusedThreads, bytes); + if (error != cudaSuccess) + return error; + if (active != 1) + return cudaErrorNotSupported; + state.fused_kernel = kernel; + state.fused_total_slots = slots; + state.fused_dynamic_bytes = bytes; + // The copy's NVLink traffic runs in the remote memory-sync domain, so fences + // and completions of kernels in the default domain do not wait for it. + int domains = 0; + if (cudaDeviceGetAttribute(&domains, cudaDevAttrMemSyncDomainCount, c.device) == cudaSuccess && domains >= 2) + { + state.fused_attributes[0].id = cudaLaunchAttributeMemSyncDomain; + state.fused_attributes[0].val.memSyncDomain = cudaLaunchMemSyncDomainRemote; + state.fused_attributes[1].id = cudaLaunchAttributeMemSyncDomainMap; + state.fused_attributes[1].val.memSyncDomainMap.default_ = 0; + state.fused_attributes[1].val.memSyncDomainMap.remote = 1; + state.fused_attribute_count = 2; + } + return cudaSuccess; +} + +} // namespace + +extern "C" int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state_ptr) +{ + if (!state_ptr || !*state_ptr) + return static_cast(cudaSuccess); + MegamoeTmaCopyState* state = *state_ptr; + cudaError_t first = cudaSuccess; + auto record = [&first](cudaError_t error) + { + if (first == cudaSuccess) + first = error; + }; + int previous_device = -1; + cudaError_t error = cudaGetDevice(&previous_device); + record(error); + if (error == cudaSuccess && previous_device != state->config.device) + { + error = cudaSetDevice(state->config.device); + record(error); + } + if (error != cudaSuccess) + return static_cast(first); // Preserve state for a later cleanup attempt. + // GPU-direct submissions use one stream; synchronize it before releasing + // the bound plan storage and completion counter. + if (state->work_queued) + record(cudaStreamSynchronize(state->last_stream)); + if (state->direct_tables) + record(cudaFree(state->direct_tables)); + if (state->direct.scratch) + record(cudaFree(state->direct.scratch)); + if (state->direct.segments) + record(cudaFree(state->direct.segments)); + if (state->direct.ranges) + record(cudaFree(state->direct.ranges)); + if (state->direct.results) + record(cudaFree(state->direct.results)); + if (state->direct.resident) + record(cudaFree(state->direct.resident)); + if (state->completed_ctas) + record(cudaFree(state->completed_ctas)); + if (previous_device != state->config.device) + record(cudaSetDevice(previous_device)); + delete state; + *state_ptr = nullptr; + return static_cast(first); +} + +extern "C" int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps, MegamoeTmaCopyState** out) +{ + if (!out) + return static_cast(cudaErrorInvalidValue); + *out = nullptr; + if (!max_segments || max_segments > SIZE_MAX / (sizeof(MegamoeTmaCopySegment) + sizeof(MegamoeTmaCopyRange)) + || sms < 0 || sms > MEGAMOE_TMA_COPY_MAX_SMS || warps < 0 || warps > MEGAMOE_TMA_COPY_MAX_WARPS) + return static_cast(cudaErrorInvalidValue); + auto* state = new (std::nothrow) MegamoeTmaCopyState; + if (!state) + return static_cast(cudaErrorMemoryAllocation); + auto fail = [&state](cudaError_t error) + { + megamoe_tma_copy_destroy(&state); + return static_cast(error); + }; + cudaError_t error = cudaGetDevice(&state->config.device); + if (error != cudaSuccess) + { + delete state; + return static_cast(error); + } + cudaDeviceProp prop{}; + error = cudaGetDeviceProperties(&prop, state->config.device); + if (error != cudaSuccess) + return fail(error); + int runtime_version = 0, multicast_supported = 0; + error = cudaRuntimeGetVersion(&runtime_version); + if (error != cudaSuccess) + return fail(error); + if (runtime_version < 13010 || prop.major < 9) + return fail(cudaErrorNotSupported); + // Multicast support is exposed by the Driver API, not cudaDeviceAttr. + CUdevice driver_device; + CUresult driver_error = cuDeviceGet(&driver_device, state->config.device); + if (driver_error != CUDA_SUCCESS) + return fail(cudaErrorUnknown); + driver_error = cuDeviceGetAttribute(&multicast_supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, driver_device); + if (driver_error != CUDA_SUCCESS) + return fail(driver_error == CUDA_ERROR_NOT_SUPPORTED ? cudaErrorNotSupported : cudaErrorUnknown); + if (!multicast_supported) + return fail(cudaErrorNotSupported); + auto& config = state->config; + config.abi_version = MEGAMOE_TMA_COPY_ABI_VERSION; + config.device_sm_count = prop.multiProcessorCount; + config.sms = megamoe_tma_copy_sm_count(sms, prop.multiProcessorCount); + config.warps = warps ? warps : 7; + config.threads_per_cta = kFusedThreads; + config.slice_bytes = kSlice; + config.max_segments = max_segments; + config.compute_major = prop.major; + config.compute_minor = prop.minor; + if (!config.sms) + return fail(cudaErrorInvalidValue); + error = cudaDeviceGetAttribute( + &config.device_optin_shared_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, config.device); + if (error != cudaSuccess) + return fail(error); + error = cudaDeviceGetAttribute( + &config.device_shared_bytes_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, config.device); + if (error != cudaSuccess) + return fail(error); + // The fused kernel's static shared memory bounds the copy slots; binding + // prepares its launch for the bound plan geometry. + void const* kernel = reinterpret_cast(&halo_q_tma_fused_kernel); + cudaFuncAttributes attributes{}; + error = cudaFuncGetAttributes(&attributes, kernel); + if (error != cudaSuccess) + return fail(error); + const size_t available = config.device_optin_shared_bytes > attributes.sharedSizeBytes + ? config.device_optin_shared_bytes - attributes.sharedSizeBytes + : 0; + config.total_slots = megamoe_tma_copy_total_slots(static_cast(available)); + config.max_warps = megamoe_tma_copy_max_warps( + config.total_slots, std::min(prop.maxThreadsPerBlock, attributes.maxThreadsPerBlock)); + if (config.warps > config.max_warps) + return fail(cudaErrorInvalidConfiguration); + const size_t raw = static_cast(config.total_slots) * (kSlotDataBytes + kSlotControlBytes); + const size_t dynamic_bytes = (raw + 127) & ~size_t(127); + if (dynamic_bytes + attributes.sharedSizeBytes <= static_cast(config.device_shared_bytes_per_sm) / 2) + return fail(cudaErrorNotSupported); + config.slots_per_warp = config.total_slots / config.warps; + config.extra_slot_warps = config.total_slots % config.warps; + config.max_slots_per_warp = config.slots_per_warp + (config.extra_slot_warps != 0); + config.bank0_slots_per_warp = (config.slots_per_warp + 1) / 2; + config.bank1_slots_per_warp = config.slots_per_warp / 2; + config.dynamic_shared_bytes = static_cast(dynamic_bytes); + error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, config.dynamic_shared_bytes); + if (error != cudaSuccess) + return fail(error); + error + = cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (error != cudaSuccess) + return fail(error); + error = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &config.max_active_ctas_per_sm, kernel, config.threads_per_cta, config.dynamic_shared_bytes); + if (error != cudaSuccess) + return fail(error); + if (config.max_active_ctas_per_sm != 1) + return fail(cudaErrorNotSupported); + error = cudaMalloc(reinterpret_cast(&state->completed_ctas), sizeof(unsigned)); + if (error != cudaSuccess) + return fail(error); + error = cudaMemset(state->completed_ctas, 0, sizeof(unsigned)); + if (error != cudaSuccess) + return fail(error); + // Cold initialization must be visible even when the first launch uses a + // non-default stream. Subsequent launches rearm entirely on the GPU. + error = cudaStreamSynchronize(nullptr); + if (error != cudaSuccess) + return fail(error); + *out = state; + return static_cast(cudaSuccess); +} + +extern "C" int megamoe_tma_copy_config_info(MegamoeTmaCopyState const* state, MegamoeTmaCopyConfig* out) +{ + if (!state || !out) + return static_cast(cudaErrorInvalidValue); + *out = state->config; + return static_cast(cudaSuccess); +} + +extern "C" char const* megamoe_tma_copy_error_string(int error) +{ + if (error == static_cast(cudaErrorInvalidConfiguration)) + return "invalid TMA launch geometry: each communication warp requires at least " + "two 8 KiB slices; shared-memory capacity and kernel/device thread limits apply"; + return cudaGetErrorString(static_cast(error)); +} + +extern "C" int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* input) +{ + if (!state || !input) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (state->work_queued || state->direct_configured || state->direct_tables) + return cudaErrorInvalidValue; + auto& d = state->direct; + auto const& c = *input; + if (c.world < 2 || c.world > 32 || (c.world & 1) || c.rank < 0 || c.rank >= c.world || c.global_experts % c.world + || c.home_count != c.global_experts / c.world || c.plan_abi_version != 6 + || (c.route_features != 0 && c.route_features != 1) || c.tma_route < 0 || c.tma_route > 2 + || c.tma_source_load_percent < 1 || c.tma_source_load_percent > 199 || c.helper_count <= 0 + || c.helper_count > (INT32_MAX - 25) / 7 || c.global_experts <= 0 || c.global_experts > 384 || c.planes <= 0 + || c.planes > 32 || c.level_count <= 0 || c.level_count > 4 || c.target_count < 0 || c.target_count > 32 + || !c.src_table || !c.dst_table || !c.plane_bytes || (c.route_features && !c.foreign_dst_table)) + return cudaErrorInvalidValue; + if (c.owner_stride != ((c.helper_count + 3) & ~3) || c.plan_words != 4 + 7 * c.owner_stride) + return cudaErrorInvalidValue; + int expected_targets = 0; + for (int level = 0; level < c.level_count; ++level) + { + int group = c.group_sizes[level]; + if (group < (level ? 4 : 2) || group > c.world || c.world % group || (!level && group != c.world) + || (level && (group >= c.group_sizes[level - 1] || c.group_sizes[level - 1] % group))) + return cudaErrorInvalidValue; + if (level) + expected_targets += c.world / group; + } + if (c.target_count != expected_targets || (c.route_features && !expected_targets) + || (!c.route_features && c.foreign_dst_table)) + return cudaErrorInvalidValue; + uint64_t total_bytes = 0; + for (int plane = 0; plane < c.planes; ++plane) + { + const uint64_t bytes = c.plane_bytes[plane]; + if (!bytes || (bytes & 3) || bytes > UINT64_MAX / c.world - total_bytes) + return cudaErrorInvalidValue; + total_bytes += bytes; + } + auto valid_address = [&](uint64_t address, int plane) + { return address && !(address & 3) && c.plane_bytes[plane] <= UINT64_MAX - address; }; + for (uint64_t i = 0; i < uint64_t(c.global_experts) * c.planes; ++i) + if (!valid_address(c.src_table[i], i % c.planes)) + return cudaErrorInvalidValue; + for (uint64_t i = 0; i < uint64_t(c.level_count) * c.helper_count * c.planes; ++i) + if (!valid_address(c.dst_table[i], i % c.planes)) + return cudaErrorInvalidValue; + if (c.route_features) + { + uint64_t index = 0; + for (int level = 1; level < c.level_count; ++level) + { + int group = c.group_sizes[level]; + for (int begin = 0; begin < c.world; begin += group) + for (int slot = 0; slot < c.helper_count; ++slot) + for (int plane = 0; plane < c.planes; ++plane, ++index) + { + const uint64_t address = c.foreign_dst_table[index]; + bool const member = c.rank >= begin && c.rank < begin + group; + if (member ? address != 0 : !valid_address(address, plane)) + return cudaErrorInvalidValue; + } + } + } + int device = -1; + cudaError_t error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + auto fail = [state](cudaError_t e) + { + state->poison = e; + return static_cast(e); + }; + d.config = c; + const uint64_t sources = uint64_t(c.global_experts) * c.planes; + const uint64_t destinations = uint64_t(c.level_count) * c.helper_count * c.planes; + const uint64_t foreign = c.route_features ? uint64_t(c.target_count) * c.helper_count * c.planes : 0; + const uint64_t words = sources + destinations + foreign + c.planes; + if (words > SIZE_MAX / sizeof(uint64_t)) + return cudaErrorInvalidValue; + error = cudaMalloc(reinterpret_cast(&state->direct_tables), words * sizeof(uint64_t)); + if (error != cudaSuccess) + return fail(error); + uint64_t* cursor = state->direct_tables; + auto upload = [&](uint64_t const* host, uint64_t count, uint64_t const** target) + { + *target = count ? cursor : nullptr; + auto e = count ? cudaMemcpy(cursor, host, count * sizeof(uint64_t), cudaMemcpyHostToDevice) : cudaSuccess; + cursor += count; + return e; + }; + if ((error = upload(c.src_table, sources, &d.config.src_table)) != cudaSuccess + || (error = upload(c.dst_table, destinations, &d.config.dst_table)) != cudaSuccess + || (error = upload(c.foreign_dst_table, foreign, &d.config.foreign_dst_table)) != cudaSuccess + || (error = upload(c.plane_bytes, c.planes, &d.config.plane_bytes)) != cudaSuccess) + return fail(error); + auto const& config = state->config; + const uint64_t n = config.sms; + d.max_segments = config.max_segments; + if (d.max_segments > SIZE_MAX / n / sizeof(MegamoeTmaCopyRange)) + return fail(cudaErrorInvalidValue); + error = cudaMalloc(reinterpret_cast(&d.scratch), n * (c.plan_words + 4ULL * c.helper_count) * sizeof(int)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.segments), n * d.max_segments * sizeof(MegamoeTmaCopySegment)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.ranges), n * d.max_segments * sizeof(MegamoeTmaCopyRange)); + if (error != cudaSuccess) + return fail(error); + error = cudaMalloc(reinterpret_cast(&d.results), n * sizeof(MegamoeTmaGpuPlanResult)); + if (error != cudaSuccess) + return fail(error); + error = cudaMemset(d.results, 0, n * sizeof(MegamoeTmaGpuPlanResult)); + if (error != cudaSuccess) + return fail(error); + error = cudaStreamSynchronize(nullptr); // Cold initialization only. + if (error != cudaSuccess) + return fail(error); + state->direct_configured = true; + return cudaSuccess; +} + +extern "C" int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, + uint64_t owners, uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity, int helper_residency) +{ + if (!state || !ids || !levels || !owners || !workspace || !status || !grid_sync + || (ids | levels | owners | workspace | status | grid_sync) % alignof(int) || capacity < 1 || capacity > 384 + || (helper_residency != 0 && helper_residency != 1)) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (!state->direct_configured || state->direct_bound || state->work_queued) + return cudaErrorInvalidValue; + int device = -1; + auto error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + for (auto ptr : {ids, levels, owners, workspace, status, grid_sync}) + { + cudaPointerAttributes attr{}; + error = cudaPointerGetAttributes(&attr, reinterpret_cast(ptr)); + if (error != cudaSuccess) + return error; + if (attr.type != cudaMemoryTypeDevice || attr.device != device) + return cudaErrorInvalidValue; + } + auto& d = state->direct; + d.ids = reinterpret_cast(ids); + d.levels = reinterpret_cast(levels); + d.owners = reinterpret_cast(owners); + d.workspace = reinterpret_cast(workspace); + d.status = reinterpret_cast(status); + d.grid_sync = reinterpret_cast(grid_sync); + d.capacity = capacity; + if (helper_residency) + { + size_t const bytes + = sizeof(int) * static_cast(d.config.world) * static_cast(d.config.helper_count); + error = cudaMalloc(reinterpret_cast(&d.resident), bytes); + if (error == cudaSuccess) + error = cudaMemset(d.resident, 0xFF, bytes); // -1: no expert is known to be resident. + if (error == cudaSuccess) + error = cudaStreamSynchronize(nullptr); // Cold initialization only. + if (error != cudaSuccess) + return error; + } + error = prepare_fused_launch(*state); + if (error != cudaSuccess) + return error; + state->direct_bound = true; + return cudaSuccess; +} + +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, + tensorrt_llm::kernels::MoeRebalanceHaloQParams const& scheduler, uint64_t* route_epoch, uint64_t route_target, + uint64_t flag_mc, uint64_t generation, void* cuda_stream) +{ + if (!state || !route_epoch || (reinterpret_cast(route_epoch) & 7) || !flag_mc || (flag_mc & 7) + || !generation) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (state->poison != cudaSuccess) + return state->poison; + auto const& c = state->config; + auto const& d = state->direct; + // The copy reads the bound scheduler outputs, so the scheduler must write exactly those. + if (!state->direct_bound || scheduler.ctas != c.sms || scheduler.threads != kFusedThreads + || scheduler.outIds != d.ids || scheduler.outLevels != d.levels || scheduler.outOwners != d.owners + || scheduler.planWorkspace != d.workspace || scheduler.status != d.status || scheduler.gridSync != d.grid_sync + || scheduler.ep != d.config.world || scheduler.experts != d.config.global_experts + || scheduler.helpers != d.config.helper_count) + return cudaErrorInvalidValue; + int device = -1; + cudaError_t error = cudaGetDevice(&device); + if (error != cudaSuccess) + return error; + if (device != c.device) + return cudaErrorInvalidDevice; + auto stream = reinterpret_cast(cuda_stream); + cudaStreamCaptureStatus capture; + error = cudaStreamIsCapturing(stream, &capture); + if (error != cudaSuccess) + return error; + if (capture != cudaStreamCaptureStatusNone) + return cudaErrorStreamCaptureUnsupported; + // One stream orders the scheduler workspace, the plan scratch and the residency table. + if (state->work_queued && stream != state->last_stream) + return cudaErrorInvalidValue; + FusedSchedulerArgs scheduler_args{scheduler.routes, scheduler.outSlots, scheduler.outIds, scheduler.outLevels, + scheduler.outOwners, reinterpret_cast(scheduler.peerBases), scheduler.status, + scheduler.partial, scheduler.routeAux, scheduler.gridSync, scheduler.planWorkspace, scheduler.routePrefix, + static_cast(scheduler.spinCycles), scheduler.ep, scheduler.experts, scheduler.helpers, + scheduler.localRank, scheduler.routeCapacity, + tensorrt_llm::kernels::broadcast_capacity(scheduler.ep, scheduler.experts, scheduler.helpers), + scheduler.validRouteCount}; + GpuDirectStorage direct = d; + int warps = c.warps; + int total_slots = state->fused_total_slots; + unsigned* completed_ctas = state->completed_ctas; + auto* epoch = reinterpret_cast(route_epoch); + unsigned long long target = route_target; + void* arguments[] + = {&scheduler_args, &direct, &warps, &total_slots, &completed_ctas, &flag_mc, &generation, &epoch, &target}; + cudaLaunchConfig_t launch{}; + launch.gridDim = dim3(c.sms, 1, 1); + launch.blockDim = dim3(kFusedThreads, 1, 1); + launch.dynamicSmemBytes = state->fused_dynamic_bytes; + launch.stream = stream; + launch.attrs = state->fused_attributes; + launch.numAttrs = state->fused_attribute_count; + state->last_stream = stream; + state->work_queued = true; + error = cudaLaunchKernelExC(&launch, state->fused_kernel, arguments); + if (error != cudaSuccess) + { + state->poison = error; + return error; + } + return cudaSuccess; +} + +extern "C" int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result) +{ + if (!state || !result) + return cudaErrorInvalidValue; + std::lock_guard lock(state->submit_mutex); + if (!state->direct_configured) + return cudaErrorInvalidValue; + if (state->poison != cudaSuccess) + return state->poison; + int device = -1; + auto device_error = cudaGetDevice(&device); + if (device_error != cudaSuccess) + return device_error; + if (device != state->config.device) + return cudaErrorInvalidDevice; + if (state->work_queued) + { + auto error = cudaStreamSynchronize(state->last_stream); + if (error != cudaSuccess) + { + state->poison = error; + return error; + } + } + return cudaMemcpy(result, state->direct.results, sizeof(*result), cudaMemcpyDeviceToHost); +} diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h new file mode 100644 index 000000000000..39668c395f18 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h @@ -0,0 +1,220 @@ +/* + * 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. + */ +/* CUDA TMA multicast-copy backend. C ABI; no CUDA or Python headers required. */ +#ifndef MEGAMOE_TMA_COPY_H +#define MEGAMOE_TMA_COPY_H + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h" +#include + +#ifdef __cplusplus +extern "C" +{ +#endif + +#define MEGAMOE_TMA_COPY_ABI_VERSION 6 +#define MEGAMOE_TMA_COPY_SLICE_BYTES 8192 +#define MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES 32 +#define MEGAMOE_TMA_COPY_SLOT_DATA_BYTES 8208 +#define MEGAMOE_TMA_COPY_MAX_SMS 8 +#define MEGAMOE_TMA_COPY_MAX_WARPS 32 +/* Experimental candidate: warps counts logical producer/consumer worker pairs. */ +#define MEGAMOE_TMA_COPY_THREADS_PER_WORKER 64 + + /* Zero requests the default; zero return denotes invalid geometry. The entire + * payload/notification kernel stays within this per-GPU copy SM budget. */ + static inline int megamoe_tma_copy_sm_count(int requested_sms, int device_sms) + { + if (device_sms <= 0 || requested_sms < 0 || requested_sms > MEGAMOE_TMA_COPY_MAX_SMS + || requested_sms > device_sms) + return 0; + return requested_sms ? requested_sms + : (device_sms < MEGAMOE_TMA_COPY_MAX_SMS ? device_sms : MEGAMOE_TMA_COPY_MAX_SMS); + } + + /* Each logical slice remains 8192B; physical slots add 16B for one aligned + * load that straddles the logical slice boundary. + * Maximum whole slots after control records and 128-byte launch alignment. + * Capacity is independent of warp count; remainder slots are not discarded. */ + static inline int megamoe_tma_copy_total_slots(int available_bytes) + { + if (available_bytes <= 0) + return 0; + int const slot_bytes = MEGAMOE_TMA_COPY_SLOT_DATA_BYTES + MEGAMOE_TMA_COPY_SLOT_CONTROL_BYTES; + int slots = available_bytes / slot_bytes; + while (slots > 0 && (((uint64_t) slots * slot_bytes + 127) & ~(uint64_t) 127) > (uint64_t) available_bytes) + --slots; + return slots; + } + + static inline int megamoe_tma_copy_max_warps(int total_slots, int max_threads) + { + if (total_slots < 2 || max_threads < MEGAMOE_TMA_COPY_THREADS_PER_WORKER) + return 0; + int warps = total_slots / 2; + if (warps > MEGAMOE_TMA_COPY_MAX_WARPS) + warps = MEGAMOE_TMA_COPY_MAX_WARPS; + if (warps > max_threads / MEGAMOE_TMA_COPY_THREADS_PER_WORKER) + warps = max_threads / MEGAMOE_TMA_COPY_THREADS_PER_WORKER; + return warps; + } + +#if defined(__CUDACC__) +#define MEGAMOE_TMA_HD __host__ __device__ +#else +#define MEGAMOE_TMA_HD +#endif + /* Validated geometry has warps>=1 and total_slots>=2*warps. The first + * total_slots%warps warps own one extra slot; every slot has exactly one owner. */ + static inline MEGAMOE_TMA_HD int megamoe_tma_warp_slot_count(int total_slots, int warps, int warp) + { + return total_slots / warps + (warp < total_slots % warps); + } + + static inline MEGAMOE_TMA_HD int megamoe_tma_warp_slot_begin(int total_slots, int warps, int warp) + { + int extra = total_slots % warps; + return warp * (total_slots / warps) + (warp < extra ? warp : extra); + } + + /* Shared with CPU coverage checks: CTA/SM first, warp second, sequence last. */ + static inline MEGAMOE_TMA_HD uint64_t megamoe_tma_worker_first_slice(int cta, int sms, int warp) + { + return (uint64_t) cta + (uint64_t) sms * (uint64_t) warp; + } + + static inline MEGAMOE_TMA_HD uint64_t megamoe_tma_worker_slice_stride(int sms, int warps) + { + return (uint64_t) sms * (uint64_t) warps; + } + +#undef MEGAMOE_TMA_HD + + typedef struct MegamoeTmaCopyState MegamoeTmaCopyState; + + /* A physical interval within one logically concatenated byte stream. Entries + * in a range are contiguous in virtual space (no per-weight/plane rounding). + * Addresses/lengths are multiples of four; source is unicast, dst is multicast. + */ + typedef struct MegamoeTmaCopySegment + { + uint64_t src; + uint64_t dst; + uint64_t bytes; + uint64_t virtual_begin; + } MegamoeTmaCopySegment; + + /* A whole routing domain: source-owned slots, or one scatter destination team. + * Slice j spans [j*8192,min((j+1)*8192,bytes)) in its virtual byte stream and can + * cross any number of segments. It keeps ONE GPU/CTA/warp owner throughout. + * This rank owns floor((first_slice+j*slice_stride)/slice_divisor). Uniform + * routing uses divisor=1, reciprocal=0 and skips the division. + * Otherwise reciprocal=UINT64_MAX/divisor enables an exact mulhi+correction. + * prefix_end compacts rank-local slices across ranges for SM/worker round-robin. + * No reset at weight boundaries. + */ + typedef struct MegamoeTmaCopyRange + { + uint64_t segment_begin; + uint64_t segment_count; + uint64_t bytes; + uint64_t first_slice; + uint64_t slice_stride; + uint64_t slice_count; + uint64_t prefix_end; + uint64_t slice_divisor; + uint64_t slice_reciprocal; + } MegamoeTmaCopyRange; + + typedef struct MegamoeTmaCopyConfig + { + int abi_version; + int device; + int device_sm_count; + int sms; + int warps; + int threads_per_cta; + int slots_per_warp; /* Minimum over all warps, floor(total_slots/warps). */ + int bank0_slots_per_warp; /* Minimum bank 0 capacity, ceil(slots_per_warp/2). */ + int bank1_slots_per_warp; /* Minimum bank 1 capacity, floor(slots_per_warp/2). */ + int slice_bytes; + int dynamic_shared_bytes; + int device_optin_shared_bytes; + int device_shared_bytes_per_sm; + int max_active_ctas_per_sm; + int compute_major; + int compute_minor; + int total_slots; + int max_slots_per_warp; + int extra_slot_warps; /* First this many warps receive max_slots_per_warp. */ + int max_warps; /* Shared-memory and kernel/device thread capacity. */ + uint64_t max_segments; + } MegamoeTmaCopyConfig; + + /* Every function returns a CUDA runtime error integer (0 == success). + * sms=0 selects min(device_sm_count,8); otherwise 1..min(device_sm_count,8). + * Explicit counts above eight are rejected; READY stays in this same kernel. + * warps=0 selects seven communication warps; explicit requests are 1..32. + * Runtime geometry requires at least two 8 KiB slots per warp and respects + * the active kernel and device limits. + * All whole shared-memory slots are divided near-equally among warps, reserving + * an mbarrier/phase record for every slot. Two banks use ALL allocated slots. + * Excessive runtime warp counts return cudaErrorInvalidConfiguration. + * Creation verifies that this shared-memory footprint permits one CTA/SM. + */ + int megamoe_tma_copy_create(uint64_t max_segments, int sms, int warps, MegamoeTmaCopyState** out); + + /* GPU-direct cold setup: input config contains HOST pointer tables. They are + * uploaded once; current HALO-Q outputs are bound by device address, never copied. + * Submit is stream-ordered and does not read or wait for plan data on the CPU. + * All bound tensors must remain alive and only be reused on the same stream + * after the copy kernel has read them. The Python endpoint enforces this lease. + * Binding also loads the fused scheduler + copy kernel. helper_residency=1 + * keeps a per-layer table of the expert each helper slot holds and skips a + * broadcast whose whole team already holds its expert; it requires helper + * rows that only this copy writes. + */ + int megamoe_tma_copy_configure_gpu_plan(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanConfig const* config); + int megamoe_tma_copy_bind_gpu_direct(MegamoeTmaCopyState* state, uint64_t ids, uint64_t levels, uint64_t owners, + uint64_t workspace, uint64_t status, uint64_t grid_sync, int capacity, int helper_residency); + /* Explicit diagnostic only: synchronizes last work and copies one small result. */ + int megamoe_tma_copy_gpu_plan_result(MegamoeTmaCopyState* state, MegamoeTmaGpuPlanResult* result); + /* Nulls *state once ownership is consumed. If device activation fails before + * cleanup begins, the pointer is preserved so the caller can retry. */ + int megamoe_tma_copy_destroy(MegamoeTmaCopyState** state); + int megamoe_tma_copy_config_info(MegamoeTmaCopyState const* state, MegamoeTmaCopyConfig* out); + char const* megamoe_tma_copy_error_string(int error); + +#ifdef __cplusplus +} + +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" + +/* Launches HALO-Q and the GPU-direct copy of its plan as one kernel of config.sms + * CTAs. scheduler must pass checkMoeRebalanceHaloQParams, use config.sms CTAs and + * write the outputs bound by megamoe_tma_copy_bind_gpu_direct. + * route_epoch is a caller-owned uint64 that only fused launches change: each CTA + * adds one with a system-scope release once its routes are written, and + * route_target is its value after this launch. A cuStreamWaitValue64 (GEQ + * route_target) on another stream orders that stream after the routes. READY is + * published as by the copy alone. + */ +int megamoe_tma_copy_submit_fused(MegamoeTmaCopyState* state, + tensorrt_llm::kernels::MoeRebalanceHaloQParams const& scheduler, uint64_t* route_epoch, uint64_t route_target, + uint64_t flag_mc, uint64_t generation, void* cuda_stream); +#endif +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h new file mode 100644 index 000000000000..9ff8abd2eaf3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h @@ -0,0 +1,53 @@ +/* + * 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. + */ +/* Device-resident ABI6 plan input for the fused TMA copy kernel. */ +#ifndef MEGAMOE_TMA_GPU_PLAN_H +#define MEGAMOE_TMA_GPU_PLAN_H +#include + +#define MEGAMOE_TMA_GPU_PLAN_MAX_WORLD 32 +#define MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS 4 + +/* Immutable geometry and device addresses. Upload these tables once at bind. + * Source entries: [global_experts][planes]. Destinations: [level][helper][plane]. + * External destinations: [target][helper][plane]. Bytes: [plane]. No entry here + * depends on the current HALO-Q result. All pointer values are device pointers. + */ +typedef struct MegamoeTmaGpuPlanConfig +{ + int world, rank, helper_count, planes, global_experts, home_count; + int owner_stride, level_count, target_count; + int plan_abi_version, plan_words, route_features; + int tma_route; /* 0: published modes; 1: source; 2: scatter. */ + int tma_source_load_percent; /* Source quota = slices * percent / (100 * team). */ + int group_sizes[MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS]; + uint64_t const *src_table, *dst_table, *foreign_dst_table, *plane_bytes; +} MegamoeTmaGpuPlanConfig; + +/* Per-CTA output; only successful construction makes counts consumable. + * Error 1: malformed plan/config; 2: forced source outside its destination team; + * 3: invalid address, byte count, capacity, or arithmetic overflow. + */ +typedef struct MegamoeTmaGpuPlanResult +{ + int error; + int direct_slots, scatter_slots; + int weighted_domains, uniform_domains, mixed_source_domains, foreign_source_domains; + uint64_t segment_count, range_count, total_slices; + uint64_t weighted_total_slices, weighted_source_quota_slices, weighted_local_slices; +} MegamoeTmaGpuPlanResult; +#endif diff --git a/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh new file mode 100644 index 000000000000..6441e46a10fe --- /dev/null +++ b/cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlanBuilder.cuh @@ -0,0 +1,515 @@ +/* + * 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. + */ +/* Construct the existing TMA descriptor ABI inside the payload kernel. + * Each CTA has private global scratch: 4*helper_count int32 slot words, + * max_segments segments, max_ranges ranges, and one result. CTA thread zero + * calls the builder and all threads synchronize before consuming its output. + * No CPU plan read, descriptor DMA, allocation, atomics or device-wide barrier. + * Scheduling already orders the producer before this kernel on its stream. + */ +#ifndef MEGAMOE_TMA_GPU_PLAN_BUILDER_CUH +#define MEGAMOE_TMA_GPU_PLAN_BUILDER_CUH +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/kernels/moe/loadBalance/dynamicEplb/moeRebalanceTmaGpuPlan.h" +#include +#if defined(__CUDACC__) +#define MEGAMOE_GP_HD __host__ __device__ +#else +#define MEGAMOE_GP_HD +#endif +namespace megamoe_gpu_plan_detail +{ +/* A helper slot whose whole team already holds its expert needs no copy. */ +static constexpr int kSlotModeResident = 3; + +struct Builder : MegamoeTmaGpuPlanConfig +{ + int32_t const* plan; + int32_t const* resident; /* [world][helper_count] expert held per helper slot, or null. */ + int32_t* slot_scratch; + MegamoeTmaCopySegment* tma_segments; + MegamoeTmaCopyRange* tma_ranges; + size_t max_commands, tma_max_ranges; + size_t tma_segment_count, tma_range_count, last_commands; + int last_direct_slots, last_scatter_slots; + int last_tma_weighted_domains, last_tma_uniform_domains; + int last_tma_mixed_source_domains, last_tma_foreign_source_domains; + uint64_t last_tma_weighted_total_slices, last_tma_weighted_source_quota_slices; + uint64_t last_tma_weighted_local_slices, last_tma_local_slices; +}; + +static MEGAMOE_GP_HD inline uint64_t tma_weighted_source_quota(uint64_t slices, unsigned members, unsigned percent) +{ + return slices * percent / (100U * members); +} + +static MEGAMOE_GP_HD inline int target_geometry(Builder const* c, int target, int* level_out, int* begin_out) +{ + if (target < 0 || target >= c->target_count) + return -1; + for (int level = 1; level < c->level_count; ++level) + { + int count = c->world / c->group_sizes[level]; + if (target < count) + { + *level_out = level; + *begin_out = target * c->group_sizes[level]; + return 0; + } + target -= count; + } + return -1; +} + +static MEGAMOE_GP_HD inline int append_tma_segment( + Builder* c, size_t* segments, uint64_t* cursor, uint64_t src, uint64_t dst, size_t bytes) +{ + if (*segments >= c->max_commands || !src || !dst || !bytes || (src & 3) || (dst & 3) || (bytes & 3) + || bytes > UINT64_MAX - *cursor || bytes > UINT64_MAX - src || bytes > UINT64_MAX - dst) + return -1; + MegamoeTmaCopySegment* s = &c->tma_segments[(*segments)++]; + s->src = src; + s->dst = dst; + s->bytes = bytes; + s->virtual_begin = *cursor; + *cursor += bytes; + return 0; +} + +/* The validated arena binds one UC object per (owner, plane) and one MC + * object per (team, plane). The caller resets identity at plane/domain + * boundaries and gives normal and external teams separate namespaces. + * Numerical pointer adjacency alone is not enough to merge MC destinations. */ +static MEGAMOE_GP_HD inline int append_tma_coalesced_segment(Builder* c, size_t* segments, uint64_t* cursor, + uint64_t src, uint64_t dst, size_t bytes, uint64_t identity, uint64_t* previous_identity) +{ + if (!src || !dst || !bytes || (src & 3) || (dst & 3) || (bytes & 3) || bytes > UINT64_MAX - *cursor + || bytes > UINT64_MAX - src || bytes > UINT64_MAX - dst) + return -1; + if (*segments && *previous_identity == identity) + { + MegamoeTmaCopySegment* previous = &c->tma_segments[*segments - 1]; + if (previous->bytes <= UINT64_MAX - previous->src && previous->bytes <= UINT64_MAX - previous->dst + && previous->bytes <= UINT64_MAX - previous->virtual_begin && previous->src + previous->bytes == src + && previous->dst + previous->bytes == dst && previous->virtual_begin + previous->bytes == *cursor + && bytes <= UINT64_MAX - previous->bytes) + { + previous->bytes += bytes; + *cursor += bytes; + return 0; + } + } + if (append_tma_segment(c, segments, cursor, src, dst, bytes)) + return -1; + *previous_identity = identity; + return 0; +} + +static MEGAMOE_GP_HD inline int append_tma_range(Builder* c, size_t* ranges, uint64_t* prefix, size_t segment_begin, + size_t segment_count, uint64_t bytes, uint64_t first, uint64_t step, uint64_t slices, uint64_t divisor, + uint64_t reciprocal) +{ + /* These constant bounds also avoid runtime division for overflow checks. + * Valid 8KiB streams are smaller than this conservative slice bound. */ + if (!step || !divisor || divisor > 100U * MEGAMOE_TMA_GPU_PLAN_MAX_WORLD + || slices > UINT64_MAX / (100U * MEGAMOE_TMA_GPU_PLAN_MAX_WORLD)) + return -1; + uint64_t scaled_limit = slices * divisor; + if (first >= scaled_limit) + return 0; + uint64_t count = 1 + (scaled_limit - 1 - first) / step; + if (*ranges >= c->tma_max_ranges || count > UINT64_MAX - *prefix) + return -1; + MegamoeTmaCopyRange* r = &c->tma_ranges[(*ranges)++]; + r->segment_begin = segment_begin; + r->segment_count = segment_count; + r->bytes = bytes; + r->first_slice = first; + r->slice_stride = step; + r->slice_count = count; + *prefix += count; + r->prefix_end = *prefix; + r->slice_divisor = divisor; + r->slice_reciprocal = reciprocal; + return 0; +} + +static MEGAMOE_GP_HD inline bool team_holds_expert(Builder const* c, int slot, int expert, int level) +{ + int group = c->group_sizes[level]; + int begin = c->rank / group * group; + for (int rank = begin; rank < begin + group; ++rank) + if (c->resident[rank * c->helper_count + slot] != expert) + return false; + return true; +} + +static MEGAMOE_GP_HD inline size_t build_tma_descriptors(Builder* c, int outgoing_count, int* plan_error) +{ + *plan_error = 3; + int slots = c->helper_count, stride = c->owner_stride; + int *experts = c->slot_scratch, *levels = experts + slots; + int *owners = levels + slots, *modes = owners + slots; + size_t segments = 0, ranges = 0; + uint64_t prefix = 0; + int direct = 0, scatter = 0; + c->last_tma_weighted_domains = c->last_tma_uniform_domains = 0; + c->last_tma_mixed_source_domains = c->last_tma_foreign_source_domains = 0; + c->last_tma_weighted_total_slices = c->last_tma_weighted_source_quota_slices = 0; + c->last_tma_weighted_local_slices = c->last_tma_local_slices = 0; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] == 2) + continue; + if (c->resident && team_holds_expert(c, slot, experts[slot], levels[slot])) + { + modes[slot] = kSlotModeResident; + continue; + } + if (c->tma_route == 1) + { + int group = c->group_sizes[levels[slot]]; + int begin = c->rank / group * group; + if (owners[slot] < begin || owners[slot] >= begin + group) + { + *plan_error = 2; + return 0; + } + modes[slot] = 1; + } + else if (c->tma_route == 2) + { + modes[slot] = 0; + } + direct += modes[slot] == 1; + scatter += modes[slot] == 0; + } + /* Domain zero is source-owned work (including outgoing direct). A scatter + * domain must have one common member set, so it is separated per level; + * all ranks in that aligned team see the same plane/slot byte order. */ + for (int domain = 0; domain <= c->level_count; ++domain) + { + size_t begin_segment = segments; + uint64_t cursor = 0; + /* Keep one uninterrupted virtual stream per routing domain. Plane-major + * order groups the large aligned weights before the scalar planes; the + * physical arena and source ownership are unchanged. Within each plane, + * normal source-owned slots precede existing external outgoing records. */ + for (int plane = 0; plane < c->planes; ++plane) + { + uint64_t previous_identity = UINT64_MAX; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] == 2) + continue; + if (domain == 0 ? (modes[slot] != 1 || owners[slot] != c->rank) + : (modes[slot] != 0 || levels[slot] != domain - 1)) + continue; + int level = levels[slot]; + size_t si = (size_t) experts[slot] * c->planes + plane; + size_t di = ((size_t) level * slots + slot) * c->planes + plane; + uint64_t identity = ((uint64_t) owners[slot] << 32) | (uint32_t) level; + if (append_tma_coalesced_segment(c, &segments, &cursor, c->src_table[si], c->dst_table[di], + c->plane_bytes[plane], identity, &previous_identity)) + return 0; + } + if (domain == 0) + { + for (int i = 0; i < outgoing_count; ++i) + { + int expert = c->plan[4 + 4 * stride + i]; + int target = c->plan[4 + 5 * stride + i]; + int helper = c->plan[4 + 6 * stride + i]; + size_t si = (size_t) expert * c->planes + plane; + size_t di = ((size_t) target * slots + helper) * c->planes + plane; + uint64_t identity = ((uint64_t) c->rank << 32) | (uint32_t) (c->level_count + target); + if (append_tma_coalesced_segment(c, &segments, &cursor, c->src_table[si], c->foreign_dst_table[di], + c->plane_bytes[plane], identity, &previous_identity)) + return 0; + } + } + } + if (!cursor) + continue; + uint64_t step = domain ? (uint64_t) c->group_sizes[domain - 1] : 1; + uint64_t first = domain ? (uint64_t) c->rank % step : 0; + uint64_t slices = cursor / MEGAMOE_TMA_COPY_SLICE_BYTES + (cursor % MEGAMOE_TMA_COPY_SLICE_BYTES != 0); + size_t before_ranges = ranges; + int common_source = -1; + int weighted = 0; + if (domain && c->tma_source_load_percent != 100) + { + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0 || modes[slot] != 0 || levels[slot] != domain - 1) + continue; + if (common_source == -1) + common_source = owners[slot]; + else if (common_source != owners[slot]) + { + common_source = -2; + break; + } + } + int team_begin = c->rank / (int) step * (int) step; + if (common_source == -2) + ++c->last_tma_mixed_source_domains; + else if (common_source < team_begin || common_source >= team_begin + (int) step) + ++c->last_tma_foreign_source_domains; + else + weighted = 1; + } + if (weighted) + { + unsigned members = (unsigned) step; + unsigned source_local = (unsigned) common_source % members; + unsigned percent = (unsigned) c->tma_source_load_percent; + uint64_t denominator = 100U * members; + uint64_t before_prefix = prefix; + uint64_t source_quota = tma_weighted_source_quota(slices, members, percent); + if (slices > UINT64_MAX - c->last_tma_weighted_total_slices + || source_quota > UINT64_MAX - c->last_tma_weighted_source_quota_slices) + return 0; + ++c->last_tma_weighted_domains; + c->last_tma_weighted_total_slices += slices; + c->last_tma_weighted_source_quota_slices += source_quota; + uint64_t divisor, reciprocal; + if (c->rank == common_source) + { + first = denominator - 1; + step = denominator; + divisor = percent; + reciprocal = divisor == 1 ? 0 : UINT64_MAX / divisor; + } + else + { + unsigned peer_order = ((unsigned) c->rank % members + members - source_local) % members - 1; + first = peer_order * denominator; + step = (members - 1) * denominator; + divisor = denominator - percent; + reciprocal = divisor == 1 ? 0 : UINT64_MAX / divisor; + } + /* floor((first+j*step)/divisor) is the inverse of the source + * quota events and peer round-robin. The whole existing virtual + * stream stays intact, with increasing slices and a continuous + * rank-local prefix across domains for SM/worker round-robin. */ + if (append_tma_range(c, &ranges, &prefix, begin_segment, segments - begin_segment, cursor, first, step, + slices, divisor, reciprocal)) + return 0; + c->last_tma_weighted_local_slices += prefix - before_prefix; + } + else + { + if (domain) + ++c->last_tma_uniform_domains; + if (append_tma_range( + c, &ranges, &prefix, begin_segment, segments - begin_segment, cursor, first, step, slices, 1, 0)) + return 0; + } + if (ranges == before_ranges) + segments = begin_segment; /* This rank owns none of this stream's slices. */ + } + c->last_tma_local_slices = prefix; + c->tma_segment_count = segments; + c->tma_range_count = ranges; + c->last_commands = segments; + c->last_direct_slots = direct; + c->last_scatter_slots = scatter; + *plan_error = 0; + return segments; +} + +static MEGAMOE_GP_HD inline size_t build_descriptors(Builder* c, int* plan_error) +{ + const volatile int32_t* plan = c->plan; + int const slots = c->helper_count; + int const stride = c->owner_stride; + int const quota = slots < c->home_count ? slots : c->home_count; + int* experts = c->slot_scratch; + int* levels = experts + slots; + int* owners = levels + slots; + int* modes = owners + slots; + int owner_counts[MEGAMOE_TMA_GPU_PLAN_MAX_WORLD] = {0}; + unsigned char seen_experts[384] = {0}; + int incoming_slots = 0; + int active = 0, common_level = -1; + int global_max_send = plan[1]; + int outgoing_count = plan[3]; + *plan_error = 1; + if (c->plan_abi_version != 6 || c->plan_words != 4 + 7 * stride || plan[2] != 0 || global_max_send < 0 + || global_max_send > quota || outgoing_count < 0 || outgoing_count > quota + || (!c->route_features && outgoing_count)) + return 0; + + /* Published modes are global decisions. Validate the complete local row + * before constructing any descriptor; inactive capacity never becomes work. */ + for (int slot = 0; slot < slots; ++slot) + { + int expert = experts[slot] = plan[4 + slot]; + int level = levels[slot] = plan[4 + stride + slot]; + int owner = owners[slot] = plan[4 + 2 * stride + slot]; + int mode = modes[slot] = plan[4 + 3 * stride + slot]; + if (expert == -1) + { + if (level != -1 || owner != -1 || mode != 0) + return 0; + continue; + } + if (expert < 0 || expert >= c->global_experts || level < 0 || level >= c->level_count || owner < 0 + || owner >= c->world || owner != expert / c->home_count || mode < 0 || mode > 2 || seen_experts[expert]) + return 0; + seen_experts[expert] = 1; + if (++owner_counts[owner] > global_max_send) + return 0; + int group = c->group_sizes[level]; + int begin = c->rank / group * group; + int local = owner >= begin && owner < begin + group; + if ((mode == 1 && !local) + || (mode == 2 && (!c->route_features || !c->foreign_dst_table || local || level == 0))) + return 0; + if (!active) + common_level = level; + else if (common_level != level) + common_level = -2; + ++active; + incoming_slots += mode == 2; + } + /* Source-direct is valid whenever its owner belongs to the target team. + * The scheduler may explicitly choose it for multiple weights or levels; + * auto-policy heuristics are not descriptor safety constraints. */ + if (incoming_slots) + { + if (common_level <= 0 || global_max_send == 0 || global_max_send > 2 || (global_max_send == 2 && active < 3)) + return 0; + int group = c->group_sizes[common_level]; + int begin = c->rank / group * group; + for (int slot = 0; slot < slots; ++slot) + { + if (experts[slot] < 0) + continue; + int external = owners[slot] < begin || owners[slot] >= begin + group; + if (external != (modes[slot] == 2)) + return 0; + } + } + if (outgoing_count && (global_max_send == 0 || global_max_send > 2 || (active && common_level <= 0))) + return 0; + if (owner_counts[c->rank] + outgoing_count > global_max_send) + return 0; + for (int i = 0; i < outgoing_count; ++i) + { + int expert = plan[4 + 4 * stride + i]; + int target = plan[4 + 5 * stride + i]; + int helper = plan[4 + 6 * stride + i]; + int level, begin; + if (expert < c->rank * c->home_count || expert >= (c->rank + 1) * c->home_count || seen_experts[expert] + || helper < 0 || helper >= slots || target_geometry(c, target, &level, &begin)) + return 0; + int group = c->group_sizes[level]; + if (c->rank >= begin && c->rank < begin + group) + return 0; + seen_experts[expert] = 1; + for (int j = 0; j < i; ++j) + if (plan[4 + 5 * stride + j] == target && plan[4 + 6 * stride + j] == helper) + return 0; + } + for (int i = outgoing_count; i < stride; ++i) + if (plan[4 + 4 * stride + i] != -1 || plan[4 + 5 * stride + i] != -1 || plan[4 + 6 * stride + i] != -1) + return 0; + for (int slot = slots; slot < stride; ++slot) + if (plan[4 + slot] != -1 || plan[4 + stride + slot] != -1 || plan[4 + 2 * stride + slot] != -1 + || plan[4 + 3 * stride + slot] != 0) + return 0; + + return build_tma_descriptors(c, outgoing_count, plan_error); +} + +} // namespace megamoe_gpu_plan_detail + +/* Host/device parity permits differential tests against the CPU submitter. + * Counts remain zero on failure: callers must not publish successful READY. + * Caller supplies valid memory for the full immutable tables and scratch. + * A non-null resident table removes the slots whose team already holds the + * planned expert from every routing domain. + */ +static MEGAMOE_GP_HD inline int megamoe_tma_build_gpu_plan(MegamoeTmaGpuPlanConfig const* config, + int32_t const* device_plan, int32_t* slot_scratch, MegamoeTmaCopySegment* segments, uint64_t max_segments, + MegamoeTmaCopyRange* ranges, uint64_t max_ranges, MegamoeTmaGpuPlanResult* result, + int32_t const* resident = nullptr) +{ + if (!result) + return 1; + *result = MegamoeTmaGpuPlanResult{}; + result->error = 1; + if (!config || !device_plan || !slot_scratch || !segments || !ranges || !max_segments || !max_ranges + || max_segments > SIZE_MAX || max_ranges > SIZE_MAX) + return 1; + MegamoeTmaGpuPlanConfig const& g = *config; + if (g.world < 2 || g.world > MEGAMOE_TMA_GPU_PLAN_MAX_WORLD || (g.world & 1) || g.rank < 0 || g.rank >= g.world + || g.helper_count <= 0 || g.planes <= 0 || g.global_experts <= 0 || g.global_experts > 384 + || g.global_experts % g.world || g.home_count != g.global_experts / g.world || g.level_count <= 0 + || g.level_count > MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS + || g.owner_stride != ((int64_t(g.helper_count) + 3) & ~INT64_C(3)) + || int64_t(g.plan_words) != 4 + 7 * int64_t(g.owner_stride) || g.plan_abi_version != 6 + || (g.route_features != 0 && g.route_features != 1) || g.tma_route < 0 || g.tma_route > 2 + || g.tma_source_load_percent < 1 || g.tma_source_load_percent > 199 || !g.src_table || !g.dst_table + || !g.plane_bytes || (g.route_features && !g.foreign_dst_table) || (!g.route_features && g.foreign_dst_table) + || device_plan[0] != 0) + return 1; + int targets = 0; + for (int level = 0; level < g.level_count; ++level) + { + int group = g.group_sizes[level]; + if (group < (level ? 4 : 2) || group > g.world || g.world % group || (!level && group != g.world) + || (level && (group >= g.group_sizes[level - 1] || g.group_sizes[level - 1] % group))) + return 1; + if (level) + targets += g.world / group; + } + if (g.target_count != targets || (g.route_features && !targets)) + return 1; + megamoe_gpu_plan_detail::Builder c{}; + static_cast(c) = g; + c.plan = device_plan; + c.resident = resident; + c.slot_scratch = slot_scratch; + c.tma_segments = segments; + c.tma_ranges = ranges; + c.max_commands = (size_t) max_segments; + c.tma_max_ranges = (size_t) max_ranges; + int error = 1; + megamoe_gpu_plan_detail::build_descriptors(&c, &error); + result->error = error; + if (error) + return error; + result->direct_slots = c.last_direct_slots; + result->scatter_slots = c.last_scatter_slots; + result->weighted_domains = c.last_tma_weighted_domains; + result->uniform_domains = c.last_tma_uniform_domains; + result->mixed_source_domains = c.last_tma_mixed_source_domains; + result->foreign_source_domains = c.last_tma_foreign_source_domains; + result->segment_count = c.tma_segment_count; + result->range_count = c.tma_range_count; + result->total_slices = c.last_tma_local_slices; + result->weighted_total_slices = c.last_tma_weighted_total_slices; + result->weighted_source_quota_slices = c.last_tma_weighted_source_quota_slices; + result->weighted_local_slices = c.last_tma_weighted_local_slices; + return 0; +} + +#undef MEGAMOE_GP_HD +#endif diff --git a/cpp/tensorrt_llm/thop/fp4Quantize.cpp b/cpp/tensorrt_llm/thop/fp4Quantize.cpp index 859f1489b2b6..266a825329bf 100644 --- a/cpp/tensorrt_llm/thop/fp4Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp4Quantize.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -42,8 +42,9 @@ namespace torch_ext // returns self_fp4, self_block_scale_factors // self_fp4: [M, K / 2], FLOAT4_E2M1X2 // self_block_scale_factors: ceil(M / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0) -std::tuple fp4_quantize(at::Tensor const& self, std::optional const& globalScale, - int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout) +std::tuple fp4_quantize_sm_budget(at::Tensor const& self, + std::optional const& globalScale, int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout, + int64_t reservedSms) { CHECK_TH_CUDA(self); CHECK_CONTIGUOUS(self); @@ -88,6 +89,9 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt = at::detail::empty_cuda({SFSize}, SF_DTYPE, self.device(), /* stride */ std::nullopt); // 1D tensor const thread_local int mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + TORCH_CHECK(reservedSms >= 0 && reservedSms < mMultiProcessorCount, + "reserved_sms must be nonnegative and smaller than the device SM count"); + int const availableSms = mMultiProcessorCount - static_cast(reservedSms); auto const layout = isSfSwizzledLayout ? tensorrt_llm::QuantizationSFLayout::SWIZZLED : tensorrt_llm::QuantizationSFLayout::LINEAR; @@ -95,7 +99,7 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt #define LAUNCH_FP4_QUANTIZE_KERNEL(T, SF_VEC_SIZE) \ tensorrt_llm::kernels::invokeFP4Quantization(1, m, k, reinterpret_cast(self.data_ptr()), \ globalScalePtr, reinterpret_cast(valueE2M1.data_ptr()), \ - reinterpret_cast(scaleFP8SF.data_ptr()), sfUseUE8M0, layout, mMultiProcessorCount, \ + reinterpret_cast(scaleFP8SF.data_ptr()), sfUseUE8M0, layout, availableSms, \ at::cuda::getCurrentCUDAStream(self.get_device())); if (sfUseUE8M0) @@ -158,6 +162,12 @@ std::tuple fp4_quantize(at::Tensor const& self, std::opt return {valueE2M1, scaleFP8SF}; } +std::tuple fp4_quantize(at::Tensor const& self, std::optional const& globalScale, + int64_t sfVecSize, bool sfUseUE8M0, bool isSfSwizzledLayout) +{ + return fp4_quantize_sm_budget(self, globalScale, sfVecSize, sfUseUE8M0, isSfSwizzledLayout, 0); +} + at::Tensor calculate_nvfp4_global_scale(at::Tensor const& input, std::optional const& tokensPerBatch) { CHECK_TH_CUDA(input); @@ -377,6 +387,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "fp4_quantize(Tensor input, Tensor? globalScale, int sfVecSize, bool sfUseUE8M0=False, bool " "isSfSwizzledLayout=True) -> (Tensor, Tensor)"); + m.def( + "fp4_quantize.sm_budget(Tensor input, Tensor? globalScale, int sfVecSize, bool sfUseUE8M0=False, bool " + "isSfSwizzledLayout=True, int reserved_sms=0) -> (Tensor, Tensor)"); m.def("calculate_nvfp4_global_scale(Tensor input, Tensor? tokensPerBatch) -> Tensor"); m.def( "fp4_quantize_with_reorder_residual(Tensor X, Tensor input_scale, Tensor reorder_index, int KE, bool is_act) " @@ -387,6 +400,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fp4_quantize", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize)); + m.impl("fp4_quantize.sm_budget", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_sm_budget)); m.impl("calculate_nvfp4_global_scale", TORCH_FN(tensorrt_llm::torch_ext::calculate_nvfp4_global_scale)); m.impl("fp4_quantize_with_reorder_residual", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_with_reorder_residual)); m.impl("fp4_quantize_with_residual", TORCH_FN(tensorrt_llm::torch_ext::fp4_quantize_with_residual)); diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index e80184ff215c..86388ab8ea10 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -151,7 +151,8 @@ std::tuple fp8_batched_quantize_1x128_permute102(at::Ten // so each UE8M0 scale is replicated into four adjacent K32 slots. The legacy // deep_gemm MN-major packed layout remains available while that consumer still // requires it. -std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self, bool useR128c4Layout) +std::tuple fp8_quantize_1x128_packed_ue8m0_sm_budget( + at::Tensor const& self, bool useR128c4Layout, int64_t reservedSms) { CHECK_TH_CUDA(self); CHECK_CONTIGUOUS(self); @@ -168,6 +169,8 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co TORCH_CHECK(n <= std::numeric_limits::max(), "N must be within int32"); auto const num_n_blocks = (n + 127) / 128; auto const num_packed_sf_k = (num_n_blocks + 3) / 4; + TORCH_CHECK(reservedSms >= 0, "reserved_sms must be nonnegative"); + TORCH_CHECK(reservedSms == 0 || useR128c4Layout, "reserved_sms requires the R128c4 scale layout"); auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); if (useR128c4Layout) @@ -188,11 +191,13 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co { #ifdef ENABLE_BF16 const thread_local int multiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); + TORCH_CHECK(reservedSms < multiProcessorCount, "reserved_sms must be smaller than the device SM count"); + int const availableSms = multiProcessorCount - static_cast(reservedSms); tensorrt_llm::kernels::invokeMxFP8Quantization<__nv_bfloat16, kQuantSfVecSize, kOutputSfVecSize>(1, static_cast(m), static_cast(n), static_cast(n), reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), reinterpret_cast(valueE4M3.data_ptr()), reinterpret_cast(scaleFP8SF.data_ptr()), - tensorrt_llm::QuantizationSFLayout::SWIZZLED, multiProcessorCount, stream); + tensorrt_llm::QuantizationSFLayout::SWIZZLED, availableSms, stream); #else C10_THROW_ERROR(NotImplementedError, "BFloat16 must be enabled to quantize a BF16 tensor to MXFP8."); #endif @@ -232,6 +237,11 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co return {valueE4M3.slice(0, 0, m), packedScale}; } +std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self, bool useR128c4Layout) +{ + return fp8_quantize_1x128_packed_ue8m0_sm_budget(self, useR128c4Layout, 0); +} + std::tuple fp8_quantize_1x128_cutedsl_ue8m0(at::Tensor const& self) { CHECK_TH_CUDA(self); @@ -336,6 +346,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor)"); m.def("fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input, bool use_r128c4_layout=True) -> (Tensor, Tensor)"); + m.def( + "fp8_quantize_1x128_packed_ue8m0.sm_budget(Tensor input, bool use_r128c4_layout=True, " + "int reserved_sms=0) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_cutedsl_ue8m0(Tensor input) -> (Tensor, Tensor)"); m.def( "silu_and_mul_fp8_quantize_1x128_packed_ue8m0(Tensor input, float? swiglu_limit=None, bool " @@ -347,6 +360,8 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("fp8_quantize_1x128", &tensorrt_llm::torch_ext::fp8_quantize_1x128); m.impl("fp8_batched_quantize_1x128_permute102", &tensorrt_llm::torch_ext::fp8_batched_quantize_1x128_permute102); m.impl("fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0); + m.impl("fp8_quantize_1x128_packed_ue8m0.sm_budget", + &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0_sm_budget); m.impl("fp8_quantize_1x128_cutedsl_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_cutedsl_ue8m0); m.impl("silu_and_mul_fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::silu_and_mul_fp8_quantize_1x128_packed_ue8m0); diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp new file mode 100644 index 000000000000..334e72cf9332 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.cpp @@ -0,0 +1,126 @@ +/* + * 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/thop/moe/loadBalance/moeRebalanceHaloQOp.h" + +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void checkCudaTensor(torch::Tensor const& tensor, c10::Device const& device, c10::ScalarType dtype, + std::int64_t minElements, char const* name) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.device() == device, name, " must be on the routes CUDA device"); + TORCH_CHECK(tensor.scalar_type() == dtype, name, " has an invalid dtype"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.numel() >= minElements, name, " is too small: expected at least ", minElements, + " elements, got ", tensor.numel()); +} + +int checkedInt(std::int64_t value, char const* name) +{ + TORCH_CHECK(value >= INT_MIN && value <= INT_MAX, name, " does not fit in int32"); + return static_cast(value); +} + +} // namespace + +kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, + torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, + torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t epValue, std::int64_t expertsValue, + std::int64_t helpersValue, std::int64_t localRankValue, std::int64_t routeCapacityValue, std::int64_t ctasValue, + std::int64_t threadsValue, std::int64_t spinCycles, std::int64_t validRouteCountValue) +{ + int const ep = checkedInt(epValue, "ep"); + int const experts = checkedInt(expertsValue, "experts"); + int const helpers = checkedInt(helpersValue, "helpers"); + int const localRank = checkedInt(localRankValue, "local_rank"); + int const routeCapacity = checkedInt(routeCapacityValue, "route_capacity"); + int const ctas = checkedInt(ctasValue, "ctas"); + int const threads = checkedInt(threadsValue, "threads"); + int const validRouteCount = checkedInt(validRouteCountValue, "valid_route_count"); + TORCH_CHECK(ep >= 2 && ep <= 32 && ep % 2 == 0, "ep must be even and in [2, 32]"); + TORCH_CHECK(experts > 0 && experts <= 384 && experts % ep == 0, + "experts must be positive, no greater than 384, and divisible by ep"); + TORCH_CHECK(helpers > 0 && helpers <= INT_MAX / ep - experts / ep, "invalid helpers"); + TORCH_CHECK(localRank >= 0 && localRank < ep, "local_rank must be in [0, ep)"); + TORCH_CHECK(routeCapacity > 0 && routeCapacity <= INT_MAX / ep, "invalid route_capacity"); + TORCH_CHECK( + validRouteCount >= 0 && validRouteCount <= routeCapacity, "valid_route_count must be in [0, route_capacity]"); + TORCH_CHECK(ctas > 0 && ctas <= 128 && routeCapacity <= INT_MAX - ctas, "invalid ctas"); + TORCH_CHECK(threads == 512, "threads must be 512"); + TORCH_CHECK(spinCycles >= 0, "spin_cycles cannot be negative"); + + c10::Device const device = routes.device(); + checkCudaTensor(routes, device, torch::kInt32, validRouteCount, "routes"); + checkCudaTensor(outSlots, device, torch::kInt32, routeCapacity, "out_slots"); + checkCudaTensor(outIds, device, torch::kInt32, helpers, "out_ids"); + checkCudaTensor(outLevels, device, torch::kInt32, helpers, "out_levels"); + checkCudaTensor(outOwners, device, torch::kInt32, helpers, "out_owners"); + checkCudaTensor(peerBases, device, torch::kInt64, ep, "peer_bases"); + checkCudaTensor(status, device, torch::kInt32, 6, "status"); + checkCudaTensor(partial, device, torch::kInt32, static_cast(ctas) * experts, "partial"); + std::int64_t const bins = static_cast(experts) + 1; + std::int64_t routeAuxElements = static_cast(ctas) * (threads / 32) * bins; + if (ctas > 1) + { + routeAuxElements = std::max(routeAuxElements, 2 + static_cast(ctas - 1) * bins + routeCapacity); + } + checkCudaTensor(routeAux, device, torch::kInt32, routeAuxElements, "route_aux"); + checkCudaTensor(gridSync, device, torch::kInt32, 4, "grid_sync"); + checkCudaTensor(planWorkspace, device, torch::kInt32, 2 + 6 * 384, "plan_workspace"); + checkCudaTensor(routePrefix, device, torch::kInt32, static_cast(384) * (ep + 1), "route_prefix"); + + return kernels::MoeRebalanceHaloQParams{ + routes.data_ptr(), + outSlots.data_ptr(), + outIds.data_ptr(), + outLevels.data_ptr(), + outOwners.data_ptr(), + reinterpret_cast(peerBases.data_ptr()), + status.data_ptr(), + partial.data_ptr(), + routeAux.data_ptr(), + gridSync.data_ptr(), + planWorkspace.data_ptr(), + routePrefix.data_ptr(), + ep, + experts, + helpers, + localRank, + routeCapacity, + ctas, + threads, + static_cast(spinCycles), + validRouteCount, + }; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h new file mode 100644 index 000000000000..ef9e075da699 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h @@ -0,0 +1,40 @@ +/* + * 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/kernels/moe/loadBalance/dynamicEplb/moeRebalanceHaloQ.h" + +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Validates the HALO-Q scheduler arguments that moe_rebalance_tma_bind_fused binds and returns its launch parameters. +kernels::MoeRebalanceHaloQParams makeMoeRebalanceHaloQParams(torch::Tensor const& routes, torch::Tensor& outSlots, + torch::Tensor& outIds, torch::Tensor& outLevels, torch::Tensor& outOwners, torch::Tensor const& peerBases, + torch::Tensor& status, torch::Tensor& partial, torch::Tensor& routeAux, torch::Tensor& gridSync, + torch::Tensor& planWorkspace, torch::Tensor& routePrefix, std::int64_t ep, std::int64_t experts, + std::int64_t helpers, std::int64_t localRank, std::int64_t routeCapacity, std::int64_t ctas, std::int64_t threads, + std::int64_t spinCycles, std::int64_t validRouteCount); + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp new file mode 100644 index 000000000000..4a43b222ec98 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp @@ -0,0 +1,631 @@ +/* + * 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/moe/loadBalance/dynamicEplb/moeRebalanceTma.h" +#include "tensorrt_llm/thop/moe/loadBalance/moeRebalanceHaloQOp.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace +{ + +void checkTmaResult(int error, char const* operation) +{ + TORCH_CHECK(error == 0, operation, " failed: ", megamoe_tma_copy_error_string(error), " (", error, ")"); +} + +void checkCudaResult(cudaError_t error, char const* operation) +{ + TORCH_CHECK(error == cudaSuccess, operation, " failed: ", cudaGetErrorString(error)); +} + +int checkedInt(int64_t value, char const* name) +{ + TORCH_CHECK(value >= std::numeric_limits::min() && value <= std::numeric_limits::max(), name, + " is outside the int32 range"); + return static_cast(value); +} + +uint64_t checkedUnsigned(int64_t value, char const* name) +{ + TORCH_CHECK(value >= 0, name, " must be nonnegative"); + return static_cast(value); +} + +std::vector checkedU64Vector(std::vector const& values, uint64_t expectedSize, char const* name) +{ + TORCH_CHECK(expectedSize <= std::numeric_limits::max(), name, " size overflows size_t"); + TORCH_CHECK(values.size() == static_cast(expectedSize), name, " has ", values.size(), " entries; expected ", + expectedSize); + std::vector result; + result.reserve(values.size()); + for (int64_t value : values) + { + result.push_back(checkedUnsigned(value, name)); + } + return result; +} + +void checkBoundTensor(torch::Tensor const& tensor, int device, int64_t minimumElements, char const* name) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.scalar_type() == at::ScalarType::Int, name, " must have dtype torch.int32"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.get_device() == device, name, " must be on CUDA device ", device); + TORCH_CHECK(tensor.numel() >= minimumElements, name, " has ", tensor.numel(), " elements; expected at least ", + minimumElements); +} + +} // namespace + +class MoeRebalanceTmaState : public torch::CustomClassHolder +{ +public: + MoeRebalanceTmaState(int64_t maxSegments, int64_t sms, int64_t warps) + { + TORCH_CHECK(maxSegments > 0, "max_segments must be positive"); + auto const maxSegmentsUnsigned = checkedUnsigned(maxSegments, "max_segments"); + int const error + = megamoe_tma_copy_create(maxSegmentsUnsigned, checkedInt(sms, "sms"), checkedInt(warps, "warps"), &mState); + checkTmaResult(error, "moe rebalance TMA create"); + + MegamoeTmaCopyConfig config{}; + int const configError = megamoe_tma_copy_config_info(mState, &config); + if (configError != 0) + { + megamoe_tma_copy_destroy(&mState); + checkTmaResult(configError, "moe rebalance TMA config query"); + } + mDevice = config.device; + } + + MoeRebalanceTmaState(MoeRebalanceTmaState const&) = delete; + MoeRebalanceTmaState& operator=(MoeRebalanceTmaState const&) = delete; + + ~MoeRebalanceTmaState() noexcept override + { + std::lock_guard lock(mMutex); + if (mState != nullptr) + { + // Destructors cannot report CUDA cleanup failures. Explicit destroy() is + // available when the caller needs a synchronous error result. + static_cast(megamoe_tma_copy_destroy(&mState)); + } + releaseFusedEvents(); + } + + void configureGpuPlan(int64_t world, int64_t rank, int64_t helperCount, int64_t planes, int64_t globalExperts, + int64_t homeCount, int64_t ownerStride, int64_t levelCount, int64_t targetCount, int64_t planAbiVersion, + int64_t planWords, int64_t routeFeatures, int64_t tmaRoute, int64_t tmaSourceLoadPercent, + std::vector const& groupSizes, std::vector const& sourceTable, + std::vector const& destinationTable, std::vector const& foreignDestinationTable, + std::vector const& planeBytes) + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(!mConfigured, "moe rebalance TMA state is already configured"); + + int const worldValue = checkedInt(world, "world"); + int const helperCountValue = checkedInt(helperCount, "helper_count"); + int const planesValue = checkedInt(planes, "planes"); + int const globalExpertsValue = checkedInt(globalExperts, "global_experts"); + int const levelCountValue = checkedInt(levelCount, "level_count"); + int const targetCountValue = checkedInt(targetCount, "target_count"); + int const routeFeaturesValue = checkedInt(routeFeatures, "route_features"); + TORCH_CHECK(worldValue > 0 && helperCountValue > 0 && planesValue > 0 && globalExpertsValue > 0, + "world, helper_count, planes, and global_experts must be positive"); + TORCH_CHECK(worldValue <= MEGAMOE_TMA_GPU_PLAN_MAX_WORLD && planesValue <= 32 && globalExpertsValue <= 384, + "world, planes, or global_experts exceeds the GPU plan ABI"); + TORCH_CHECK( + helperCountValue <= (std::numeric_limits::max() - 25) / 7, "helper_count exceeds the GPU plan ABI"); + TORCH_CHECK(levelCountValue > 0 && targetCountValue >= 0, + "level_count must be positive and target_count must be nonnegative"); + TORCH_CHECK(targetCountValue <= 32 && (routeFeaturesValue == 0 || routeFeaturesValue == 1), + "target_count or route_features exceeds the GPU plan ABI"); + TORCH_CHECK(levelCountValue <= MEGAMOE_TMA_GPU_PLAN_MAX_LEVELS, "level_count exceeds the GPU plan ABI"); + TORCH_CHECK( + groupSizes.size() == static_cast(levelCountValue), "group_sizes must contain level_count entries"); + + std::array groups{}; + for (int index = 0; index < levelCountValue; ++index) + { + groups[index] = checkedInt(groupSizes[index], "group_sizes entry"); + } + + uint64_t const sourceCount = static_cast(globalExpertsValue) * planesValue; + uint64_t const destinationCount = static_cast(levelCountValue) * helperCountValue * planesValue; + uint64_t const foreignCount + = routeFeaturesValue ? static_cast(targetCountValue) * helperCountValue * planesValue : 0; + auto sources = checkedU64Vector(sourceTable, sourceCount, "source_table"); + auto destinations = checkedU64Vector(destinationTable, destinationCount, "destination_table"); + auto foreign = checkedU64Vector(foreignDestinationTable, foreignCount, "foreign_destination_table"); + auto bytes = checkedU64Vector(planeBytes, planesValue, "plane_bytes"); + + MegamoeTmaGpuPlanConfig config{}; + config.world = worldValue; + config.rank = checkedInt(rank, "rank"); + config.helper_count = helperCountValue; + config.planes = planesValue; + config.global_experts = globalExpertsValue; + config.home_count = checkedInt(homeCount, "home_count"); + config.owner_stride = checkedInt(ownerStride, "owner_stride"); + config.level_count = levelCountValue; + config.target_count = targetCountValue; + config.plan_abi_version = checkedInt(planAbiVersion, "plan_abi_version"); + config.plan_words = checkedInt(planWords, "plan_words"); + config.route_features = routeFeaturesValue; + config.tma_route = checkedInt(tmaRoute, "tma_route"); + config.tma_source_load_percent = checkedInt(tmaSourceLoadPercent, "tma_source_load_percent"); + for (size_t index = 0; index < groups.size(); ++index) + { + config.group_sizes[index] = groups[index]; + } + config.src_table = sources.data(); + config.dst_table = destinations.data(); + config.foreign_dst_table = foreign.empty() ? nullptr : foreign.data(); + config.plane_bytes = bytes.data(); + + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + checkTmaResult(megamoe_tma_copy_configure_gpu_plan(mState, &config), "moe rebalance TMA GPU plan configure"); + mConfigured = true; + mHelperCount = helperCountValue; + } + + void bindGpuDirect(torch::Tensor ids, torch::Tensor levels, torch::Tensor owners, torch::Tensor workspace, + torch::Tensor status, torch::Tensor gridSync, int64_t capacity, bool helperResidency) + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(mConfigured, "moe rebalance TMA state must be configured before bind"); + TORCH_CHECK(!mBound, "moe rebalance TMA state is already bound"); + int const capacityValue = checkedInt(capacity, "capacity"); + TORCH_CHECK(capacityValue > 0, "capacity must be positive"); + checkBoundTensor(ids, mDevice, mHelperCount, "ids"); + checkBoundTensor(levels, mDevice, mHelperCount, "levels"); + checkBoundTensor(owners, mDevice, mHelperCount, "owners"); + checkBoundTensor(workspace, mDevice, 2 + 6LL * capacityValue, "workspace"); + checkBoundTensor(status, mDevice, 6, "status"); + checkBoundTensor(gridSync, mDevice, 2, "grid_sync"); + + c10::cuda::CUDAGuard const deviceGuard(ids.device()); + checkTmaResult( + megamoe_tma_copy_bind_gpu_direct(mState, reinterpret_cast(ids.data_ptr()), + reinterpret_cast(levels.data_ptr()), reinterpret_cast(owners.data_ptr()), + reinterpret_cast(workspace.data_ptr()), reinterpret_cast(status.data_ptr()), + reinterpret_cast(gridSync.data_ptr()), capacityValue, helperResidency ? 1 : 0), + "moe rebalance TMA GPU-direct bind"); + + // The native state stores raw device addresses. Retaining the tensors in + // the owner prevents allocator reuse until destroy() has drained work. + mIds = std::move(ids); + mLevels = std::move(levels); + mOwners = std::move(owners); + mWorkspace = std::move(workspace); + mStatus = std::move(status); + mGridSync = std::move(gridSync); + mBound = true; + } + + // Binds every argument that the fused launches share: the HALO-Q scheduler workspace and outputs, the route + // epoch, and the static copy stream. A launch then passes only its routes. + void bindFused(int64_t flagMc, torch::Tensor routeEpoch, torch::Tensor outSlots, torch::Tensor outIds, + torch::Tensor outLevels, torch::Tensor outOwners, torch::Tensor peerBases, torch::Tensor status, + torch::Tensor partial, torch::Tensor routeAux, torch::Tensor gridSync, torch::Tensor planWorkspace, + torch::Tensor routePrefix, int64_t ep, int64_t experts, int64_t helpers, int64_t localRank, int64_t topk, + int64_t routeCapacity, int64_t ctas, int64_t threads, int64_t copyStream) + { + std::lock_guard lock(mMutex); + requireBound(); + TORCH_CHECK(!mFusedBound, "moe rebalance fused launch is already bound"); + TORCH_CHECK( + routeEpoch.is_cuda() && routeEpoch.get_device() == mDevice, "route_epoch must be on CUDA device ", mDevice); + TORCH_CHECK( + routeEpoch.scalar_type() == at::ScalarType::Long && routeEpoch.is_contiguous() && routeEpoch.numel() == 1, + "route_epoch must be a contiguous one-element torch.int64 tensor"); + uint64_t const flag = checkedUnsigned(flagMc, "flag_mc"); + TORCH_CHECK(flag != 0, "flag_mc must be nonzero"); + TORCH_CHECK(topk > 0 && routeCapacity % topk == 0, "route_capacity must be a positive multiple of topk"); + TORCH_CHECK(copyStream != 0, "copy_stream must be a non-default CUDA stream"); + TORCH_CHECK( + outSlots.is_cuda() && outSlots.get_device() == mDevice, "out_slots must be on CUDA device ", mDevice); + // Validate the shared arguments once, with no routes; a launch only adds its routes. + auto params = makeMoeRebalanceHaloQParams(torch::empty({0, topk}, outSlots.options()), outSlots, outIds, + outLevels, outOwners, peerBases, status, partial, routeAux, gridSync, planWorkspace, routePrefix, ep, + experts, helpers, localRank, routeCapacity, ctas, threads, 0, 0); + kernels::checkMoeRebalanceHaloQParams(params); + + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + cudaEvent_t inputReady{nullptr}; + checkCudaResult( + cudaEventCreateWithFlags(&inputReady, cudaEventDisableTiming), "cudaEventCreate(rebalance input ready)"); + cudaEvent_t consumerDone{nullptr}; + cudaError_t const error = cudaEventCreateWithFlags(&consumerDone, cudaEventDisableTiming); + if (error != cudaSuccess) + { + static_cast(cudaEventDestroy(inputReady)); + checkCudaResult(error, "cudaEventCreate(rebalance consumer done)"); + } + + mRouteEpochPtr = reinterpret_cast(routeEpoch.data_ptr()); + // The fused launch reads and writes these by address; retaining them prevents allocator reuse. + mFusedTensors = {std::move(routeEpoch), std::move(outSlots), std::move(outIds), std::move(outLevels), + std::move(outOwners), std::move(peerBases), std::move(status), std::move(partial), std::move(routeAux), + std::move(gridSync), std::move(planWorkspace), std::move(routePrefix)}; + mFusedParams = params; + mFlagMc = flag; + mTopk = topk; + mCopyStream = reinterpret_cast(copyStream); + mInputReady = inputReady; + mConsumerDone = consumerDone; + mFusedBound = true; + } + + // Launches the bound HALO-Q scheduler and this copy as one kernel on the copy stream, ordered after the work + // already queued on the current stream; see megamoe_tma_copy_submit_fused. + void submitFused(torch::Tensor const& routes, int64_t routeTarget, int64_t spinCycles) + { + std::lock_guard lock(mMutex); + requireFusedBound(); + TORCH_CHECK_VALUE(routes.is_cuda() && routes.get_device() == mDevice + && routes.scalar_type() == at::ScalarType::Int && routes.is_contiguous() && routes.dim() == 2 + && routes.size(1) == mTopk && routes.numel() <= mFusedParams.routeCapacity, + "routes must be a contiguous CUDA int32 [tokens, ", mTopk, "] tensor on device ", mDevice, " with at most ", + mFusedParams.routeCapacity / mTopk, " tokens"); + TORCH_CHECK_VALUE(spinCycles > 0, "spin_cycles must be positive"); + auto scheduler = mFusedParams; + scheduler.routes = routes.data_ptr(); + scheduler.validRouteCount = static_cast(routes.numel()); + scheduler.spinCycles = static_cast(spinCycles); + + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + cudaStream_t const execution = currentExecutionStream(); + checkCudaResult(cudaEventRecord(mInputReady, execution), "cudaEventRecord(rebalance input ready)"); + checkCudaResult(cudaStreamWaitEvent(mCopyStream, mInputReady, 0), "cudaStreamWaitEvent(rebalance input ready)"); + checkTmaResult( + megamoe_tma_copy_submit_fused(mState, scheduler, mRouteEpochPtr, static_cast(routeTarget), + mFlagMc, mGeneration, reinterpret_cast(mCopyStream)), + "moe rebalance fused HALO-Q + TMA submit"); + ++mGeneration; + TORCH_CHECK(mGeneration != 0, "moe rebalance TMA generation overflow"); + } + + // Orders the copy stream after the work already queued on the current stream, such as the consumer of the + // live bank that the next fused launch overwrites. + void releaseAfterCurrentStream() + { + std::lock_guard lock(mMutex); + requireFusedBound(); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + cudaStream_t const execution = currentExecutionStream(); + checkCudaResult(cudaEventRecord(mConsumerDone, execution), "cudaEventRecord(rebalance consumer done)"); + checkCudaResult( + cudaStreamWaitEvent(mCopyStream, mConsumerDone, 0), "cudaStreamWaitEvent(rebalance consumer done)"); + } + + int64_t currentGeneration() const + { + std::lock_guard lock(mMutex); + requireLive(); + return checkedDiagnostic(mGeneration); + } + + // Result order: error, direct/scatter slots, four domain counts, segment/range/slice counts, + // weighted total/quota/local slices. This diagnostic synchronizes the last submitted stream. + + std::vector result() + { + std::lock_guard lock(mMutex); + requireLive(); + TORCH_CHECK(mConfigured, "moe rebalance TMA state is not configured"); + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + MegamoeTmaGpuPlanResult value{}; + checkTmaResult(megamoe_tma_copy_gpu_plan_result(mState, &value), "moe rebalance TMA diagnostic result"); + return {value.error, value.direct_slots, value.scatter_slots, value.weighted_domains, value.uniform_domains, + value.mixed_source_domains, value.foreign_source_domains, checkedDiagnostic(value.segment_count), + checkedDiagnostic(value.range_count), checkedDiagnostic(value.total_slices), + checkedDiagnostic(value.weighted_total_slices), checkedDiagnostic(value.weighted_source_quota_slices), + checkedDiagnostic(value.weighted_local_slices)}; + } + + // Config order follows MegamoeTmaCopyConfig through max_segments. + + std::vector config() const + { + std::lock_guard lock(mMutex); + requireLive(); + MegamoeTmaCopyConfig value{}; + checkTmaResult(megamoe_tma_copy_config_info(mState, &value), "moe rebalance TMA config query"); + return {value.abi_version, value.device, value.device_sm_count, value.sms, value.warps, value.threads_per_cta, + value.slots_per_warp, value.bank0_slots_per_warp, value.bank1_slots_per_warp, value.slice_bytes, + value.dynamic_shared_bytes, value.device_optin_shared_bytes, value.device_shared_bytes_per_sm, + value.max_active_ctas_per_sm, value.compute_major, value.compute_minor, value.total_slots, + value.max_slots_per_warp, value.extra_slot_warps, value.max_warps, checkedDiagnostic(value.max_segments)}; + } + + void destroy() + { + std::lock_guard lock(mMutex); + if (mState == nullptr) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard(c10::Device(c10::DeviceType::CUDA, mDevice)); + int const error = megamoe_tma_copy_destroy(&mState); + if (mState == nullptr) + { + mIds = torch::Tensor(); + mLevels = torch::Tensor(); + mOwners = torch::Tensor(); + mWorkspace = torch::Tensor(); + mStatus = torch::Tensor(); + mGridSync = torch::Tensor(); + mConfigured = false; + mBound = false; + releaseFusedEvents(); + mFusedTensors.clear(); + mRouteEpochPtr = nullptr; + mCopyStream = nullptr; + mFusedBound = false; + } + checkTmaResult(error, "moe rebalance TMA destroy"); + } + +private: + void requireLive() const + { + TORCH_CHECK(mState != nullptr, "moe rebalance TMA state has been destroyed"); + } + + void requireBound() const + { + requireLive(); + TORCH_CHECK(mBound, "moe rebalance TMA state must be bound before submit"); + } + + void requireFusedBound() const + { + requireBound(); + TORCH_CHECK(mFusedBound, "moe rebalance fused launch must be bound before submit"); + } + + // The stream whose queued work a launch or release follows; it must differ from the copy stream. + cudaStream_t currentExecutionStream() const + { + cudaStream_t const stream = at::cuda::getCurrentCUDAStream(mDevice).stream(); + TORCH_CHECK(stream != mCopyStream, "the current stream must differ from the bound copy stream"); + return stream; + } + + void releaseFusedEvents() noexcept + { + // Pending records complete normally after an event is destroyed. + if (mInputReady != nullptr) + { + static_cast(cudaEventDestroy(mInputReady)); + mInputReady = nullptr; + } + if (mConsumerDone != nullptr) + { + static_cast(cudaEventDestroy(mConsumerDone)); + mConsumerDone = nullptr; + } + } + + static int64_t checkedDiagnostic(uint64_t value) + { + TORCH_CHECK(value <= static_cast(std::numeric_limits::max()), + "moe rebalance TMA diagnostic counter exceeds int64"); + return static_cast(value); + } + + mutable std::mutex mMutex; + MegamoeTmaCopyState* mState{nullptr}; + int mDevice{-1}; + int mHelperCount{0}; + uint64_t mGeneration{1}; + bool mConfigured{false}; + bool mBound{false}; + torch::Tensor mIds; + torch::Tensor mLevels; + torch::Tensor mOwners; + torch::Tensor mWorkspace; + torch::Tensor mStatus; + torch::Tensor mGridSync; + // Arguments bound by bindFused() and shared by every fused launch. + kernels::MoeRebalanceHaloQParams mFusedParams{}; + std::vector mFusedTensors; + uint64_t* mRouteEpochPtr{nullptr}; + uint64_t mFlagMc{0}; + int64_t mTopk{0}; + cudaStream_t mCopyStream{nullptr}; + cudaEvent_t mInputReady{nullptr}; + cudaEvent_t mConsumerDone{nullptr}; + bool mFusedBound{false}; +}; + +using MoeRebalanceTmaStatePtr = c10::intrusive_ptr; + +MoeRebalanceTmaStatePtr moeRebalanceTmaCreate(int64_t maxSegments, int64_t sms, int64_t warps) +{ + return c10::make_intrusive(maxSegments, sms, warps); +} + +void moeRebalanceTmaConfigureGpuPlan(MoeRebalanceTmaStatePtr const& state, int64_t world, int64_t rank, + int64_t helperCount, int64_t planes, int64_t globalExperts, int64_t homeCount, int64_t ownerStride, + int64_t levelCount, int64_t targetCount, int64_t planAbiVersion, int64_t planWords, int64_t routeFeatures, + int64_t tmaRoute, int64_t tmaSourceLoadPercent, std::vector const& groupSizes, + std::vector const& sourceTable, std::vector const& destinationTable, + std::vector const& foreignDestinationTable, std::vector const& planeBytes) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->configureGpuPlan(world, rank, helperCount, planes, globalExperts, homeCount, ownerStride, levelCount, + targetCount, planAbiVersion, planWords, routeFeatures, tmaRoute, tmaSourceLoadPercent, groupSizes, sourceTable, + destinationTable, foreignDestinationTable, planeBytes); +} + +void moeRebalanceTmaBindGpuDirect(MoeRebalanceTmaStatePtr const& state, torch::Tensor ids, torch::Tensor levels, + torch::Tensor owners, torch::Tensor workspace, torch::Tensor status, torch::Tensor gridSync, int64_t capacity, + bool helperResidency) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->bindGpuDirect(std::move(ids), std::move(levels), std::move(owners), std::move(workspace), std::move(status), + std::move(gridSync), capacity, helperResidency); +} + +void moeRebalanceTmaBindFused(MoeRebalanceTmaStatePtr const& state, int64_t flagMc, torch::Tensor routeEpoch, + torch::Tensor outSlots, torch::Tensor outIds, torch::Tensor outLevels, torch::Tensor outOwners, + torch::Tensor peerBases, torch::Tensor status, torch::Tensor partial, torch::Tensor routeAux, + torch::Tensor gridSync, torch::Tensor planWorkspace, torch::Tensor routePrefix, int64_t ep, int64_t experts, + int64_t helpers, int64_t localRank, int64_t topk, int64_t routeCapacity, int64_t ctas, int64_t threads, + int64_t copyStream) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->bindFused(flagMc, std::move(routeEpoch), std::move(outSlots), std::move(outIds), std::move(outLevels), + std::move(outOwners), std::move(peerBases), std::move(status), std::move(partial), std::move(routeAux), + std::move(gridSync), std::move(planWorkspace), std::move(routePrefix), ep, experts, helpers, localRank, topk, + routeCapacity, ctas, threads, copyStream); +} + +// The per-layer launch and release ops take this handle instead of the custom class, whose conversion costs several +// microseconds of host time per operator call. The caller keeps the state alive while it uses the handle. +int64_t moeRebalanceTmaHandle(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return reinterpret_cast(state.get()); +} + +namespace +{ + +MoeRebalanceTmaState& stateFromHandle(int64_t handle) +{ + TORCH_CHECK(handle != 0, "moe rebalance TMA handle is null"); + return *reinterpret_cast(handle); +} + +} // namespace + +void moeRebalanceTmaSubmitFused(int64_t handle, torch::Tensor const& routes, int64_t routeTarget, int64_t spinCycles) +{ + stateFromHandle(handle).submitFused(routes, routeTarget, spinCycles); +} + +void moeRebalanceTmaRelease(int64_t handle) +{ + stateFromHandle(handle).releaseAfterCurrentStream(); +} + +int64_t moeRebalanceTmaCurrentGeneration(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->currentGeneration(); +} + +std::vector moeRebalanceTmaResult(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->result(); +} + +std::vector moeRebalanceTmaConfig(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + return state->config(); +} + +void moeRebalanceTmaDestroy(MoeRebalanceTmaStatePtr const& state) +{ + TORCH_CHECK(state, "moe rebalance TMA state is null"); + state->destroy(); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.class_("MoeRebalanceTmaState"); + m.def( + "moe_rebalance_tma_create(int max_segments, int sms=0, int warps=0) -> " + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState"); + m.def( + "moe_rebalance_tma_configure_gpu_plan(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int world, int rank, int helper_count, " + "int planes, int global_experts, int home_count, int owner_stride, int level_count, int target_count, " + "int plan_abi_version, int plan_words, int route_features, int tma_route, int tma_source_load_percent, " + "int[] group_sizes, int[] source_table, int[] destination_table, int[] foreign_destination_table, " + "int[] plane_bytes) -> ()"); + m.def( + "moe_rebalance_tma_bind_gpu_direct(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, Tensor ids, Tensor levels, Tensor owners, " + "Tensor workspace, Tensor status, Tensor grid_sync, int capacity, bool helper_residency) -> ()"); + m.def( + "moe_rebalance_tma_bind_fused(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state, int flag_mc, Tensor route_epoch, " + "Tensor out_slots, Tensor out_ids, Tensor out_levels, Tensor out_owners, Tensor peer_bases, Tensor status, " + "Tensor partial, Tensor route_aux, Tensor grid_sync, Tensor plan_workspace, Tensor route_prefix, int ep, " + "int experts, int helpers, int local_rank, int topk, int route_capacity, int ctas, int threads, " + "int copy_stream) -> ()"); + m.def( + "moe_rebalance_tma_handle(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); + m.def("moe_rebalance_tma_submit_fused(int handle, Tensor routes, int route_target, int spin_cycles) -> ()"); + m.def("moe_rebalance_tma_release(int handle) -> ()"); + m.def( + "moe_rebalance_tma_current_gen(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int"); + m.def( + "moe_rebalance_tma_result(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int[]"); + m.def( + "moe_rebalance_tma_config(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> int[]"); + m.def( + "moe_rebalance_tma_destroy(" + "__torch__.torch.classes.trtllm.MoeRebalanceTmaState state) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) +{ + m.impl("moe_rebalance_tma_create", &tensorrt_llm::torch_ext::moeRebalanceTmaCreate); + m.impl("moe_rebalance_tma_configure_gpu_plan", &tensorrt_llm::torch_ext::moeRebalanceTmaConfigureGpuPlan); + m.impl("moe_rebalance_tma_bind_gpu_direct", &tensorrt_llm::torch_ext::moeRebalanceTmaBindGpuDirect); + m.impl("moe_rebalance_tma_bind_fused", &tensorrt_llm::torch_ext::moeRebalanceTmaBindFused); + m.impl("moe_rebalance_tma_handle", &tensorrt_llm::torch_ext::moeRebalanceTmaHandle); + m.impl("moe_rebalance_tma_submit_fused", &tensorrt_llm::torch_ext::moeRebalanceTmaSubmitFused); + m.impl("moe_rebalance_tma_release", &tensorrt_llm::torch_ext::moeRebalanceTmaRelease); + m.impl("moe_rebalance_tma_current_gen", &tensorrt_llm::torch_ext::moeRebalanceTmaCurrentGeneration); + m.impl("moe_rebalance_tma_result", &tensorrt_llm::torch_ext::moeRebalanceTmaResult); + m.impl("moe_rebalance_tma_config", &tensorrt_llm::torch_ext::moeRebalanceTmaConfig); + m.impl("moe_rebalance_tma_destroy", &tensorrt_llm::torch_ext::moeRebalanceTmaDestroy); +} diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md index c0dcebe7b94d..37a66edc725d 100644 --- a/docs/source/developer-guide/perf-analysis.md +++ b/docs/source/developer-guide/perf-analysis.md @@ -237,6 +237,41 @@ ENABLE_PERFECT_ROUTER=1 trtllm-serve ... | Similar performance with/without perfect router | Router load balancing is not a bottleneck; focus optimization efforts elsewhere | | Significant improvement with perfect router | The learned router is causing load imbalance; consider router optimization or load balancing strategies | +### Replacing Expert IDs Instead of Logits + +`ENABLE_PERFECT_ROUTER=expert_ids` keeps the learned router and its routing +method, then replaces the routed expert IDs with a perfectly balanced +assignment: every EP rank and every expert receives an equal share of the +routes (within one) for any token count, and the receiver order is rotated by +the sender rank. The routing scales are left as computed. + +Use it where logits cannot steer the routing, for example the DeepSeek-V4 hash +layers, which pick experts from the token IDs, or where the token count changes +every iteration, as in context-only benchmarks: the IDs come from one cached +table per expert geometry, sized to `max_num_tokens`, so no logits are +synthesized per token count. + +It applies where the MoE scheduler computes the routed expert IDs outside the +MoE kernel: always with `MEGAMOE_CUTEDSL`, `CUTLASS`, `DEEPGEMM` and `CUTEDSL`, +and with other backends under attention DP or with DeepSeek-V4 routing. A +backend that routes inside its kernel, such as `TRTLLM` without attention DP, +has no routed expert IDs to replace: the mode logs a warning and has no effect +there, so use `ENABLE_PERFECT_ROUTER=1` instead. + +With the load balancer off, the gap to a plain run bounds what any MoE load +balancer can gain: + +```bash +# Learned routing, load balancer off +trtllm-bench ... +# Perfectly balanced routing, load balancer off (upper bound) +ENABLE_PERFECT_ROUTER=expert_ids trtllm-bench ... +``` + +The layer-wise benchmarks (`examples/layer_wise_benchmarks`) accept the same +environment variable with `--balance-method NotModified` and without +`--replay-file-path`, whose recorded routing would replace the balanced IDs. + ### Supported Models ```{note} diff --git a/docs/source/developer-guide/perf-benchmarking.md b/docs/source/developer-guide/perf-benchmarking.md index 0f1db502d914..9738e4c0ceb0 100644 --- a/docs/source/developer-guide/perf-benchmarking.md +++ b/docs/source/developer-guide/perf-benchmarking.md @@ -37,6 +37,7 @@ The following guidance will mostly focus on benchmarks using `trtllm-bench` CLI. - [Supported Quantization Modes](#supported-quantization-modes) - [Preparing a Dataset](#preparing-a-dataset) - [Running with the PyTorch Workflow](#running-with-the-pytorch-workflow) + - [Keeping Host Stalls out of the Measured Run](#keeping-host-stalls-out-of-the-measured-run) - [Benchmarking with LoRA Adapters in PyTorch workflow](#benchmarking-with-lora-adapters-in-pytorch-workflow) - [Running multi-modal models in the PyTorch Workflow](#running-multi-modal-models-in-the-pytorch-workflow) - [Quantization in the PyTorch Flow](#quantization-in-the-pytorch-flow) @@ -251,6 +252,26 @@ trtllm-bench --model meta-llama/Llama-3.1-8B \ --backend pytorch ``` +#### Keeping Host Stalls out of the Measured Run + +With many short requests, for example context-only runs at high concurrency, +rank 0 can be bound by host work, and a pause there makes every other rank wait +at the next collective. With `TRTLLM_BENCH_DISABLE_GC=1`, the `throughput` and +`latency` commands collect once after warmup, freeze the surviving objects, and +disable Python's cyclic garbage collector for the measured run, so a full +collection over the in-flight requests cannot pause the benchmark process. +`TRTLLM_WORKER_DISABLE_GC=1` does the same in the worker processes. Memory held +in reference cycles is then not freed until the run ends, so set the variables +identically in every run you compare. + +With `TRTLLM_PROXY_BATCH_REQUESTS=1`, the executor proxy pickles each request on +the submitting thread and a background thread sends each burst to the rank-0 +worker as one message, so the worker receives and authenticates one message per +burst instead of one per request. The sender thread must take the GIL first, so +while the submitting thread keeps running Python, as a busy `trtllm-serve` event +loop does, a lone request can wait up to the interpreter's switch interval (5 ms +by default) before it is sent. + #### Benchmarking with LoRA Adapters in PyTorch workflow The PyTorch workflow supports benchmarking with LoRA (Low-Rank Adaptation) adapters. This requires preparing a dataset with LoRA metadata and configuring the LoRA settings. diff --git a/docs/source/features/parallel-strategy.md b/docs/source/features/parallel-strategy.md index 58732079936c..96ebb9a2198c 100644 --- a/docs/source/features/parallel-strategy.md +++ b/docs/source/features/parallel-strategy.md @@ -141,7 +141,7 @@ Large-scale MoE models like DeepSeek-V3/R1, LLaMA4, and Qwen3 use fine-grained e Wide-EP introduces the concept of **expert slots** that are decoupled from specific experts. This allows: - Multiple replicas of hot experts across different GPUs - Dynamic expert placement based on workload patterns -- Both offline and online load balancing strategies +- Offline, online, and per-layer dynamic load balancing strategies #### 2. Custom EP Communication Kernels - Optimized for NVIDIA GB200 Multi-Node NVLink (MNNVL) @@ -151,6 +151,7 @@ Wide-EP introduces the concept of **expert slots** that are decoupled from speci #### 3. Expert Parallelism Load Balancer (EPLB) - **Offline EPLB**: Pre-computed expert placement based on historical workload statistics - **Online EPLB**: Dynamic expert placement that adapts to real-time traffic patterns +- **Per-layer dynamic EPLB**: Expert placement re-planned for every MoE layer on every forward pass, with hot-expert weights copied into helper slots on the GPU (`MEGAMOE_CUTEDSL` backend only) - Layer-wise weight redistribution to minimize inference disruption ### Architecture Overview @@ -166,9 +167,10 @@ The system maintains a routing table that maps Expert IDs to Slot IDs, which can 1. **Start with offline EPLB** for production deployments with known workload patterns 2. **Use online EPLB** for dynamic workloads or when traffic patterns change frequently -3. **Monitor expert statistics** to understand workload distribution -4. **Tune max_num_tokens** based on your memory constraints and EP size -5. **Test with representative datasets** to validate load balancing effectiveness +3. **Use per-layer dynamic EPLB** for context-heavy (prefill) serving on the `MEGAMOE_CUTEDSL` backend when the hot experts change from one batch to the next. Online EPLB refreshes a few layers per iteration from decayed statistics, so it follows the recent average load; per-layer dynamic EPLB re-plans each layer from the routing of the current forward pass. It runs without CUDA graphs, so prefer online EPLB for generation-heavy deployments that rely on them. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for its requirements. +4. **Monitor expert statistics** to understand workload distribution +5. **Tune max_num_tokens** based on your memory constraints and EP size +6. **Test with representative datasets** to validate load balancing effectiveness ### References diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index 9fecec99a595..63f798341355 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -122,6 +122,29 @@ NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6, # Run Kimi K3 at the serving parallelism (DEP16, see examples/kimi_k3/eval_extra_llm_options.yaml) NP=16 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --seq-len-q 1 --moe-max-num-tokens 33024 --use-low-precision-moe-combine +# Run DeepSeek-V4 with MegaMoE, $DSV4 being a local checkpoint directory, as three arms: +# load balancer off, per-layer dynamic EPLB (num_slots = experts + helper slots per rank * EP, +# here 384 + 3 * 8), and perfectly balanced routing with the load balancer off (the upper +# bound for any MoE load balancer). MEGAMOE_CUTEDSL is not in the balance-method whitelist, +# so the arms use --balance-method NotModified. DeepSeek-V4 needs --tokens-per-block 128 or +# 256, and --max-seq-len above 256 (an assert in its sparse-attention top-k); 512 is tested. +# MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 selects the FP4 combine. A slice may start mid-model +# (here 7,8). +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified --moe-load-balancer eplb_per_layer_dynamic.yaml +NP=8 ./mpi_launch.sh -x MEGAMOE_COMBINE_FORMAT=16e2m1xbf16 -x ENABLE_PERFECT_ROUTER=expert_ids ./run.sh config_ctx.yaml --model $DSV4 --layer-indices 7,8 --tokens-per-block 128 --batch-size 64 --seq-len-q 128 --max-seq-len 512 --moe-backend MEGAMOE_CUTEDSL --balance-method NotModified +# with eplb_per_layer_dynamic.yaml holding: +# mode: per_layer_dynamic +# num_slots: 408 +# Per-layer dynamic EPLB runs without CUDA graphs only, and online standard EPLB is not +# supported on a layer slice; the runner rejects both. +# To replay the routing of an end-to-end run, add --replay-file-path (see "Performance +# alignment" below) to the first two arms only, each with a recording made with the same +# num_slots: a recording with the load balancer off (384 slots) cannot be replayed into the +# per-layer dynamic arm (408 slots), nor the reverse. Replay overrides the expert IDs that +# ENABLE_PERFECT_ROUTER=expert_ids substitutes, so run the third arm without it; run.py +# rejects that combination. Replay with MEGAMOE_CUTEDSL and the CTX phase is untested. + # Run with DeepEP A2A NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_gen.yaml --moe-backend CUTEDSL diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index c88f47bc5df6..f848043462cc 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -76,6 +76,14 @@ def comma_separated_floats(s): # whose Literal is the authoritative list. A second list here can only drift. parser.add_argument("--moe-backend-for-prefill", type=str) parser.add_argument("--moe-max-num-tokens", type=int) +parser.add_argument( + "--moe-load-balancer", + type=str, + help="MoeConfig.load_balancer: a YAML file path, or a mapping when set in the config" + " file, e.g. {mode: per_layer_dynamic, num_slots: 408}. Supports per_layer_dynamic" + " without CUDA graphs, and static standard EPLB (layer_updates_per_iter 0)." + " Default: no load balancer.", +) parser.add_argument( "--spec-max-draft-len", type=int, @@ -177,6 +185,11 @@ def comma_separated_floats(s): args.use_cuda_graph = False if (args.replay_start_iter is None) != (args.replay_stop_iter is None): parser.error("Both --replay-start-iter and --replay-stop-iter must be provided or none") +if args.replay_file_path and os.environ.get("ENABLE_PERFECT_ROUTER") == "expert_ids": + parser.error( + "--replay-file-path replays the recorded expert IDs over the ones that" + " ENABLE_PERFECT_ROUTER=expert_ids substitutes; run the perfect-router arm without it" + ) if args.replay_verify_metadata is None: args.replay_verify_metadata = True if args.spec_max_draft_len is not None: @@ -250,6 +263,7 @@ def comma_separated_floats(s): use_cuda_graph=args.use_cuda_graph, spec_config=spec_config, vision_config=args.vision_config, + moe_load_balancer=args.moe_load_balancer, ) logger.info("Layer-wise benchmarks: Create runner ... Done") @@ -418,10 +432,15 @@ def comma_separated_floats(s): with runner.replace_routing_method_ctx( balance_method=BalanceMethod[args.balance_method], balance_ratio=balance_ratio ): - capture_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(capture_stream): + if args.use_cuda_graph: + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + run_pack(check=True) + torch.cuda.current_stream().wait_stream(capture_stream) + else: + # Eager runs warm up on the stream they are timed on: per-layer + # dynamic EPLB binds its submission stream on first use. run_pack(check=True) - torch.cuda.current_stream().wait_stream(capture_stream) torch.cuda.synchronize() logger.info("Layer-wise benchmarks: Warmup ... Done") @@ -507,3 +526,4 @@ def comma_separated_floats(s): f" (us)" ) torch.cuda.cudart().cudaProfilerStop() +runner.shutdown_moe_load_balancer() diff --git a/examples/wide_ep/README.md b/examples/wide_ep/README.md index 4cac840c86d7..2973ce4db0ad 100644 --- a/examples/wide_ep/README.md +++ b/examples/wide_ep/README.md @@ -67,9 +67,9 @@ If set, at most `max_num_tokens` tokens will be sent to `torch.ops.trtllm.fused_ #### `load_balancer` -Configuration for MoE load balancing, users can directly set `num_slots` and `layer_updates_per_iter` as online EPLB settings, while set path to a YAML file that also includes `initial_global_assignments` for offline EPLB. +Configuration for MoE load balancing, users can directly set `num_slots` and `layer_updates_per_iter` as online EPLB settings, while set path to a YAML file that also includes `initial_global_assignments` for offline EPLB. The [`mode`](#mode) field selects between this standard EPLB and per-layer dynamic EPLB. -EPLB is only set up when all of the following hold. Otherwise the `load_balancer` section is silently ignored: +EPLB is only set up when all of the following hold. Otherwise a `standard` `load_balancer` section is silently ignored, while `mode: per_layer_dynamic` raises an error: * The model architecture is listed in `moe_model_arch_list` in [`moe_load_balancer.py`](../../tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py) — DeepSeek-V3/V3.2/V4, GLM MoE DSA, GPT-OSS, Kimi K2.5, Mixtral, Llama 4, Nemotron-H, Qwen2/Qwen3/Qwen3.5 MoE, and the Qwen3-VL / Qwen3.5 MoE conditional-generation wrappers. * Expert parallelism is enabled (`moe_expert_parallel_size` greater than 1). @@ -77,6 +77,13 @@ EPLB is only set up when all of the following hold. Otherwise the `load_balancer Attention DP (`enable_attention_dp: true`) is not part of this gate, so EPLB is also set up without it. It is nonetheless the configuration wide-EP deployments normally use, and it is enabled throughout the [EP Load Balancer examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer). +#### `mode` + +Load-balancing mode, `standard` or `per_layer_dynamic`. Configs without `mode` are treated as `standard`. + +* `standard`: offline or online EPLB, selected by [`layer_updates_per_iter`](#layer_updates_per_iter). +* `per_layer_dynamic`: per-layer dynamic EPLB, which re-plans the expert replicas of every MoE layer on every forward pass from that pass's routing. It requires the `MEGAMOE_CUTEDSL` backend (`AUTO` resolves to it) and accepts only `num_slots`, `auxiliary_sms` and `helper_slot_mode` besides `mode`. See [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) for an example and its constraints. + #### `num_slots` Total number of expert slots. It must be ≥ total experts, and divisible by the EP size so that every rank holds the same number of slots. Two typical settings: @@ -85,12 +92,16 @@ Total number of expert slots. It must be ≥ total experts, and divisible by the 2. Set to number of total experts + EP size, such as 288 for DeepSeek R1, 32-way EP. * This means there is 1 extra expert on each EP rank, so that there is more room for the per-rank token distribution to be more balanced. +With `mode: per_layer_dynamic`, `num_slots` is required and must be greater than the total expert count, which must itself be divisible by the EP size. Each rank keeps its `num_experts / ep_size` experts and gets `(num_slots - num_experts) / ep_size` extra slots. Unlike `standard` mode, the extra slots are helper slots rather than static replicas: each layer refills them with copies of its hot experts on every forward pass. By default (`helper_slot_mode: per_layer_resident`) each layer has its own helper slots and skips experts they already hold; `shared_banks` has all MoE layers alternate between two shared helper banks. For example, 384 experts with `num_slots: 416` on 8 EP ranks give each rank 48 experts and 4 helper slots. + To disable MoE load balancing, omit the whole `load_balancer` section. EPLB is enabled by the presence of `load_balancer`, so setting `num_slots` to 0 does not disable it. #### `layer_updates_per_iter` Number of layers updated per iteration, defaults to `0`. `0` means offline, while `>0` means online EPLB. +`layer_updates_per_iter` and `initial_global_assignments` apply only to `standard` mode; a `per_layer_dynamic` configuration that sets either of them is rejected. + This is the number of layers refreshed per iteration, not how often a given layer is refreshed. Layers take turns over `ceil(num_moe_layers / layer_updates_per_iter)` iterations before the plan repeats, so with `layer_updates_per_iter: 1` a DeepSeek-R1 layer (58 MoE layers) is refreshed once every 59 iterations. See [Update Strategy](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#update-strategy) for the statistics, the replication and placement rules, and a worked example. ### Offline Load Balancer Configuration @@ -99,6 +110,10 @@ Refer to the [Offline EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/t *Note: Online EP Load Balancer is more suitable for production deployments that need to react timely to online traffic changes.* +### Per-Layer Dynamic Load Balancer Configuration + +Refer to the [Per-Layer Dynamic EP Load Balancer](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer#per-layer-dynamic-ep-load-balancer) documentation. + ### Execute Wide-EP on SLURM Clusters Refer to the [slurm_scripts](./slurm_scripts/) directory, which reuses [disaggregated slurm scripts](../disaggregated/slurm/) for submitting jobs to SLURM clusters. diff --git a/examples/wide_ep/ep_load_balancer/README.md b/examples/wide_ep/ep_load_balancer/README.md index e7365d41739e..0f1883ea4b62 100644 --- a/examples/wide_ep/ep_load_balancer/README.md +++ b/examples/wide_ep/ep_load_balancer/README.md @@ -270,3 +270,47 @@ R3 = 26.7 + 26.7 + 10 = 63.4 ``` The hottest rank is now at 70.0 against a mean of 65, so the imbalance drops from 177% to 7.7%. The three replicas of expert 0 land on three different ranks and share its traffic. This example traces the algorithm on the given input rather than reporting measured performance; for measured results, see the [tech blog](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog04_Scaling_Expert_Parallelism_in_TensorRT-LLM.md#online-ep-load-balancer-1). + + +## Per-Layer Dynamic EP Load Balancer + +Offline EPLB keeps one expert placement for the whole run, and online EPLB refreshes `layer_updates_per_iter` layers per iteration from decayed statistics. Per-layer dynamic EPLB instead re-plans every MoE layer on every forward pass, from the routing of that pass: + +1. Each rank keeps its own `num_experts / ep_size` experts and has `(num_slots - num_experts) / ep_size` helper slots. By default (`helper_slot_mode: per_layer_resident`) every MoE layer has its own helper slots; with `shared_banks`, all MoE layers alternate between two shared helper banks. +2. Before a layer's experts run, a GPU scheduler (HALO-Q) exchanges routed-token counts across the EP ranks, picks the hot experts to copy into helper slots, and maps every routed token to a resident or helper slot. +3. The same kernel then moves the picked expert weights over NVLink on a high-priority auxiliary stream, while the main stream runs other work such as input quantization and, for DeepSeek-V4, the shared experts. By default, an expert that a helper slot already holds from that layer's previous forward pass is not copied again. +4. The MoE kernel reads a helper slot only after its weights are marked ready. + +Expert weights are not staged in host shared memory, so the GDRCopy and shared-memory setup that online EPLB needs does not apply. + +An example configuration for a model with 384 routed experts on 8 GPUs: + +```bash +cat > ./config_per_layer_dynamic.yaml < dict: from tensorrt_llm.llmapi.llm_args import ( SERVE_CONFIG_ALIASES, MoeLoadBalancerConfig, + MoePerLayerDynamicLoadBalancerConfig, TorchLlmArgs, ) @@ -141,10 +142,17 @@ def generate_serve_schema() -> dict: ) # Serving treats null as unset before constructing the non-optional telemetry model. properties["telemetry_config"] = _nullable(properties["telemetry_config"]) - # Keep accepted forms aligned with TorchLlmArgs.validate_load_balancer(). - schema["$defs"]["MoeConfig"]["properties"]["load_balancer"] = { - "anyOf": [_schema_for(MoeLoadBalancerConfig, schema), {"type": ["string", "null"]}] - } + # Keep accepted forms aligned with TorchLlmArgs.validate_load_balancer(): a mapping + # without mode is standard EPLB, so only the per-layer dynamic form requires it. + standard, per_layer_dynamic = ( + {"$ref": f"#/$defs/{model.__name__}"} + for model in (MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig) + ) + schema["$defs"]["MoeConfig"]["properties"]["load_balancer"]["anyOf"] = [ + standard, + {**per_layer_dynamic, "required": ["mode"]}, + {"type": ["string", "null"]}, + ] # These object/Any annotations are Python API hooks, not arbitrary YAML mappings. # Update these restrictions when hooks are added or gain YAML construction support. for name in ("batched_logits_processor", "checkpoint_loader", "_mpi_session"): diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 1a1afaf02929..71ee8c6ffd78 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -574,7 +574,7 @@ def _(activation: torch.Tensor): activation.new_empty(scale_shape, dtype=activation.dtype)) @torch.library.register_fake("trtllm::fp4_quantize") - def _( + def _fp4_quantize_fake( input: torch.Tensor, global_scale: torch.Tensor, sf_vec_size: int, @@ -587,6 +587,18 @@ def _( return (input.new_empty(output_shape, dtype=torch.uint8), global_scale.new_empty(scale_shape, dtype=torch.uint8)) + if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): + + @torch.library.register_fake("trtllm::fp4_quantize.sm_budget") + def _(input: torch.Tensor, + global_scale: torch.Tensor, + sf_vec_size: int, + sf_use_ue8m0: bool = False, + swizzled_layout: bool = True, + reserved_sms: int = 0): + return _fp4_quantize_fake(input, global_scale, sf_vec_size, + sf_use_ue8m0, swizzled_layout) + @torch.library.register_fake("trtllm::fp4_quantize_with_reorder_residual") def _( X: torch.Tensor, @@ -1051,8 +1063,10 @@ def _(pe: torch.Tensor, nope: torch.Tensor): return packed, scale @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") - def _(input: torch.Tensor, - use_r128c4_layout: bool = True) -> tuple[torch.Tensor, torch.Tensor]: + def _fp8_quantize_packed_fake( + input: torch.Tensor, + use_r128c4_layout: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: m, k = input.shape[0], input.shape[1] num_n_blocks = (k + 127) // 128 num_packed_sf_k = (num_n_blocks + 3) // 4 @@ -1067,6 +1081,15 @@ def _(input: torch.Tensor, dtype=torch.int32) return torch.empty_like(input, dtype=torch.float8_e4m3fn), scale + if hasattr(torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0, "sm_budget"): + + @torch.library.register_fake( + "trtllm::fp8_quantize_1x128_packed_ue8m0.sm_budget") + def _(input: torch.Tensor, + use_r128c4_layout: bool = True, + reserved_sms: int = 0): + return _fp8_quantize_packed_fake(input, use_r128c4_layout) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_cutedsl_ue8m0") def _(input: torch.Tensor): m, k = input.shape diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 25ac615f0879..c35a272f6963 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -39,10 +39,13 @@ torch_dtype_to_binding) from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy +# Callers import MoeLoadBalancerConfig from this module. +from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig # noqa: F401 from tensorrt_llm.llmapi.llm_args import (CSA2SparseAttentionConfig, DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, - KvCacheConfig, MoeLoadBalancerConfig, + KvCacheConfig, + MoeLoadBalancerConfigType, MultimodalConfig) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -286,7 +289,7 @@ class ModelConfig(Generic[TConfig]): _moe_max_num_tokens_is_default: Optional[bool] = field(default=None, repr=False, compare=False) - moe_load_balancer: Optional[MoeLoadBalancerConfig] = None + moe_load_balancer: Optional[MoeLoadBalancerConfigType] = None attn_backend: str = 'TRTLLM' # Effective threshold requested for trtllm-gen MLA skip-correction. Zero diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 4892807e0c6e..9d559ac32b3f 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -47,7 +47,7 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._ipc_utils import can_access_peer from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._utils import get_sm_version +from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -72,7 +72,7 @@ from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig from ..modules.mhc.hyper_connection import HCHead, HCState, mHC -from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.multi_stream_utils import do_multi_stream, maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import ( DEFAULT_MOE_ACTIVATION, @@ -1697,25 +1697,75 @@ def __init__( self.mapping = model_config.mapping - # FIXME: incompatible with mixed quantization mode (including excluding modules from quantization) - block_size = 1 - if model_config.quant_config and model_config.quant_config.group_size is not None: - block_size = model_config.quant_config.group_size + gate_up_quant_config, down_quant_config = self._get_shared_expert_projection_quant_configs( + model_config, layer_idx + ) + shared_model_config = model_config + if ( + gate_up_quant_config == down_quant_config + and gate_up_quant_config is not model_config.quant_config + ): + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = gate_up_quant_config + quantized_block_sizes = [ + quant_config.group_size + for quant_config in (gate_up_quant_config, down_quant_config) + if ( + quant_config is not None + and quant_config.quant_algo is not None + and quant_config.group_size is not None + ) + ] + block_size = max(quantized_block_sizes, default=1) shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( shared_expert_intermediate_size, block_size ) - self.shared_experts = GatedMLP( + load_balancer_config = getattr(model_config, "moe_load_balancer", None) + self._per_layer_dynamic_eplb_enabled = ( + getattr(load_balancer_config, "mode", None) == "per_layer_dynamic" + ) + self._per_layer_dynamic_eplb_auxiliary_sms = ( + int(load_balancer_config.auxiliary_sms) if self._per_layer_dynamic_eplb_enabled else 0 + ) + rebalance_backend = getattr(self.experts, "backend", self.experts) + # The fused shared FC1/activation/FC2 kernel comes with the routed MegaMoE + # backend; ask that backend for the capability instead of naming a concrete + # implementation here. Per-layer dynamic EPLB is not a condition, but the + # kernel runs only in eager mode, which that load balancer requires. + use_fused_fc12 = ( + self._shared_fc12_backend_supported(rebalance_backend) + and is_sm_100f() + and self._shared_fc12_quantization_supported(gate_up_quant_config, down_quant_config) + and not model_config.use_cuda_graph + ) + + shared_mlp_cls = GatedMLP + if use_fused_fc12: + from ..modules.megamoe_shared_mlp import MegaMoESharedMLP + + shared_mlp_cls = MegaMoESharedMLP + + # Only fused FC12 accepts an SM budget; other shared MLP backends stay unrestricted. + shared_reserved_sms = ( + self._per_layer_dynamic_eplb_auxiliary_sms + if use_fused_fc12 and self._per_layer_dynamic_eplb_enabled + else 0 + ) + shared_mlp_kwargs = {"fc12_reserved_sms": shared_reserved_sms} if use_fused_fc12 else {} + + self.shared_experts = shared_mlp_cls( hidden_size=hidden_size, intermediate_size=shared_expert_intermediate_size, bias=False, dtype=dtype, - config=model_config, + config=shared_model_config, overridden_tp_size=shared_tp_size, reduce_output=False, use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, swiglu_limit=swiglu_limit, + **shared_mlp_kwargs, ) self.allreduce = None @@ -1788,6 +1838,62 @@ def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: return model_config.quant_config return experts_quant_config + @staticmethod + def _get_shared_expert_projection_quant_configs( + model_config, layer_idx: int + ) -> tuple[Optional[QuantConfig], Optional[QuantConfig]]: + base_name = f"model.layers.{layer_idx}.mlp.shared_experts" + quant_config_dict = getattr(model_config, "quant_config_dict", None) or {} + global_quant_config = model_config.quant_config + + gate_up_name = f"{base_name}.gate_up_proj" + gate_name = f"{base_name}.gate_proj" + up_name = f"{base_name}.up_proj" + down_name = f"{base_name}.down_proj" + + # Match apply_layerwise_quant_config exactly for the actual fused + # gate_up Linear: the first gate_proj/gate_up_proj entry wins, while an + # up-only entry does not independently reconfigure a fused weight. + gate_up_quant_config = global_quant_config + for name, candidate in quant_config_dict.items(): + if gate_name in name or gate_up_name in name: + gate_up_quant_config = candidate + break + down_quant_config = quant_config_dict.get(down_name, global_quant_config) + + # apply_quant_config_exclude_modules runs after layerwise overrides and + # excludes the fused gate_up Linear when any fused constituent matches. + if global_quant_config is not None: + unquantized = QuantConfig( + quant_algo=None, + kv_cache_quant_algo=global_quant_config.kv_cache_quant_algo, + ) + if any( + global_quant_config.is_module_excluded_from_quantization(name) + for name in (gate_up_name, gate_name, up_name) + ): + gate_up_quant_config = unquantized + if global_quant_config.is_module_excluded_from_quantization(down_name): + down_quant_config = unquantized + return gate_up_quant_config, down_quant_config + + @staticmethod + def _shared_fc12_backend_supported(backend) -> bool: + capabilities = getattr(backend, "capabilities", None) + return bool(getattr(capabilities, "supports_per_layer_dynamic_eplb", False)) + + @staticmethod + def _shared_fc12_quantization_supported( + gate_up_quant_config: Optional[QuantConfig], + down_quant_config: Optional[QuantConfig], + ) -> bool: + return ( + gate_up_quant_config is not None + and gate_up_quant_config == down_quant_config + and gate_up_quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + and gate_up_quant_config.group_size == 128 + ) + def compute_routed_output( self, hidden_states, @@ -1874,13 +1980,36 @@ def _compute_routed_output(): # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames - routed_output, shared_output = maybe_execute_in_parallel( - _compute_routed_output, - _compute_shared_output, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - self.aux_stream, - ) + # With per-layer dynamic EPLB, the route runs the shared experts after + # MAIN has enqueued both HALO-Q and the TMA copy on the high-priority + # copy stream, so they overlap planning and the helper-weight copy. + # Every rank installs the same function. + if ( + self._per_layer_dynamic_eplb_enabled + and self.shared_experts is not None + and not do_multi_stream() + ): + _shared_box = {} + + def _overlap_shared(): + # Shared-expert scaling must run at most once per forward. + if "out" not in _shared_box: + _shared_box["out"] = _compute_shared_output() + + with self.experts.load_balancer_overlap(_overlap_shared): + routed_output = _compute_routed_output() + # Compute shared experts here if the routed path did not run them. + if "out" not in _shared_box: + _shared_box["out"] = _compute_shared_output() + shared_output = _shared_box["out"] + else: + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) if not do_finalize: return [shared_output, *routed_output] diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 922a02dc1515..bc5fd8622680 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1592,7 +1592,10 @@ def transform_weights(self, module: Linear) -> None: or module.disable_deep_gemm)) or get_sm_version() == 120 use_indexer_q_cutedsl_layout = (use_deep_gemm_layout and getattr( module, "use_indexer_q_cutedsl_fusion", False)) - if use_deep_gemm_layout or use_indexer_q_cutedsl_layout: + # Consumers that read UE8M0 scales from the FP32 K128 layout (e.g. the + # fused shared FC12 kernel) ask for the rounding alone. + round_to_ue8m0 = getattr(module, "round_block_scales_to_ue8m0", False) + if use_deep_gemm_layout or use_indexer_q_cutedsl_layout or round_to_ue8m0: weight, weight_scale = resmooth_to_fp8_e8m0(module.weight, module.weight_scale) @@ -1638,6 +1641,15 @@ def transform_weights(self, module: Linear) -> None: if get_sm_version() != 107: module.fp8_low_m_gemm = Fp8BlockScaleLowMGemmRange.for_weight( *weight.shape) + elif round_to_ue8m0: + # Same FP32 K128 layout; the rounded tensors may live on another device. + replace_parameter_and_save_metadata( + module, "weight", nn.Parameter(weight, requires_grad=False), + module.rebuild_tensor_metadata) + replace_parameter_and_save_metadata( + module, "weight_scale", + nn.Parameter(weight_scale, requires_grad=False), + module.rebuild_tensor_metadata) class NVFP4LinearMethod(LinearMethodBase): diff --git a/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py new file mode 100644 index 000000000000..798306e7415a --- /dev/null +++ b/tensorrt_llm/_torch/modules/megamoe_shared_mlp.py @@ -0,0 +1,259 @@ +# 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. + +"""Shared FP8 MLP using a single MegaMoE FC1/FC2 compute kernel.""" + +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from tensorrt_llm.quantization.utils.fp8_utils import transform_k128_scales_to_cutedsl_mxfp8_layout + +from ..distributed import AllReduceParams +from ..model_config import ModelConfig +from .gated_mlp import GatedMLP + + +def _unpack_r128c4(scales: torch.Tensor, rows: int, columns: int) -> torch.Tensor: + padded_rows = (rows + 127) // 128 * 128 + padded_columns = (columns + 3) // 4 * 4 + if scales.dtype != torch.uint8 or scales.ndim != 1: + raise ValueError("Shared FC12 requires loaded uint8 R128c4 weight scales") + if not scales.is_contiguous() or scales.numel() != padded_rows * padded_columns: + raise ValueError("Shared FC12 weight scale storage has an incompatible shape") + return ( + scales.reshape(padded_rows // 128, padded_columns // 4, 32, 4, 4) + .permute(0, 3, 2, 1, 4) + .reshape(padded_rows, padded_columns)[:rows, :columns] + .contiguous() + ) + + +def _pack_r128c4(scales: torch.Tensor) -> torch.Tensor: + rows, columns = scales.shape + padded = F.pad(scales, (0, -columns % 4, 0, -rows % 128)) + return ( + padded.reshape((rows + 127) // 128, 4, 32, (columns + 3) // 4, 4) + .permute(0, 3, 2, 1, 4) + .contiguous() + .reshape(1, -1) + ) + + +def _interleave_gate_up(tensor: torch.Tensor) -> torch.Tensor: + intermediate = tensor.shape[0] // 2 + # Shared Linear stores gate || up; the FC12 epilogue consumes alternating 16-row blocks. + return ( + tensor.reshape(2, intermediate // 16, 16, tensor.shape[1]) + .permute(1, 0, 2, 3) + .contiguous() + .reshape(tensor.shape) + ) + + +class MegaMoESharedMLP(GatedMLP): + """Preserve GatedMLP loading and TP shards while fusing both projections. + + Linear owns checkpoint loading and its existing FP8 resmoothing. This adapter + only rearranges the loaded FP8/UE8M0 bytes. The caller owns shared-output + scaling and reduction; this module always returns the local TP contribution. + """ + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + bias: bool, + activation: Callable[[torch.Tensor], torch.Tensor] = F.silu, + dtype: torch.dtype | None = None, + config: ModelConfig | None = None, + overridden_tp_size: int | None = None, + reduce_output: bool = False, + layer_idx: int | None = None, + use_cute_dsl_blockscaling_mm: bool = False, + disable_deep_gemm: bool = False, + use_custom_cublas_mm: bool = False, + is_shared_expert: bool = False, + swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, + fc12_reserved_sms: int = 0, + ): + if bias or activation is not F.silu: + raise ValueError("Shared FC12 requires bias=False and SwiGLU activation") + if swiglu_alpha not in (None, 1.0) or swiglu_beta not in (None, 0.0): + raise ValueError("Shared FC12 requires SwiGLU alpha=1 and beta=0") + if swiglu_limit is not None and not float(swiglu_limit) > 0: + raise ValueError("Shared FC12 clamp must be positive or None") + if dtype not in (None, torch.bfloat16): + raise ValueError("Shared FC12 requires BF16 activations") + if reduce_output or use_custom_cublas_mm: + raise ValueError("Shared FC12 requires local output and the CuTe DSL backend") + if type(fc12_reserved_sms) is not int or fc12_reserved_sms < 0: + raise ValueError("Shared FC12 reserved-SM budget must be a non-negative integer") + config = config or ModelConfig() + if config.use_cuda_graph: + raise ValueError("Shared FC12 currently requires eager execution") + super().__init__( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=bias, + activation=activation, + dtype=dtype, + config=config, + overridden_tp_size=overridden_tp_size, + reduce_output=False, + layer_idx=layer_idx, + # On SM107 this selects Linear's post-load CuTe DSL MXFP8 conversion. On + # SM100/103 Linear then keeps FP32 K128 scales instead of DeepGEMM's layout. + use_cute_dsl_blockscaling_mm=True, + disable_deep_gemm=disable_deep_gemm, + is_shared_expert=is_shared_expert, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + self._validate_fp8_quantization() + # The fused kernel reads UE8M0 scales; Linear rounds FP32 block scales at load. + for projection in (self.gate_up_proj, self.down_proj): + projection.round_block_scales_to_ue8m0 = True + self._fc12_reserved_sms = fc12_reserved_sms + self._fc12_max_tokens = int(config.max_num_tokens) + if self._fc12_max_tokens <= 0: + raise ValueError("Shared FC12 max_num_tokens must be positive") + self._fc12_source_state: tuple[tuple[int, int, int | None], ...] | None = None + self._fc12_sm_count: int | None = None + for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): + self.register_buffer(f"_fc12_{name}", None, persistent=False) + + backend_name = "cutedsl_megamoe_shared_fc12" + kernel_name = "BlockScaledSwapAbFc12Kernel" + quantization_kind = "mxfp8_e4m3" + + def _validate_fp8_quantization(self) -> None: + for projection in (self.gate_up_proj, self.down_proj): + quant_config = projection.quant_config + if quant_config is None or not quant_config.layer_quant_mode.has_fp8_block_scales(): + raise ValueError("Shared FC12 requires FP8_BLOCK_SCALES for both projections") + + @property + def sm_count(self) -> int: + if not self.gate_up_proj._weights_created: + raise RuntimeError("Shared FC12 SM count requires materialized CUDA weights") + device = self.gate_up_proj.weight.device + if device.type != "cuda": + raise RuntimeError("Shared FC12 SM count requires CUDA weights") + total_sms = torch.cuda.get_device_properties(device).multi_processor_count + reserved = self._fc12_reserved_sms + if not 0 <= reserved < total_sms: + raise ValueError("Shared FC12 reserved-SM budget must leave at least one SM") + return total_sms - reserved + + def cache_derived_state(self) -> None: + """Drop derived views after loading, including staged in-place restoration.""" + self._fc12_source_state = None + for name in ("fc1_weight", "fc1_weight_sf", "fc2_weight", "fc2_weight_sf"): + setattr(self, f"_fc12_{name}", None) + + def post_load_weights(self) -> None: + self.cache_derived_state() + + @torch.no_grad() + def prepare_fc12_weights(self) -> None: + """Cache FC12 layouts after both Linear post-load transformations finish.""" + gate_up, down = self.gate_up_proj, self.down_proj + self._validate_fp8_quantization() + if not gate_up._weights_transformed or not down._weights_transformed: + raise RuntimeError("Shared FC12 requires both Linear post_load_weights calls first") + sources = (gate_up.weight, gate_up.weight_scale, down.weight, down.weight_scale) + source_state = tuple( + (id(t), t.data_ptr(), None if t.is_inference() else t._version) for t in sources + ) + if source_state == self._fc12_source_state: + return + fc1, fc1_sf, fc2, fc2_sf = sources + if fc1.dtype != torch.float8_e4m3fn or fc2.dtype != torch.float8_e4m3fn: + raise ValueError("Shared FC12 requires loaded E4M3 weights") + if fc1.ndim != 2 or fc2.ndim != 2 or not fc1.is_contiguous() or not fc2.is_contiguous(): + raise ValueError("Shared FC12 weights must be contiguous matrices") + intermediate = fc2.shape[1] + if fc1.shape != (2 * intermediate, self.hidden_size) or fc2.shape[0] != self.hidden_size: + raise ValueError("Shared FC12 projection shapes do not describe the same TP shard") + if ( + intermediate <= 0 + or self.hidden_size <= 0 + or intermediate % 128 + or self.hidden_size % 128 + ): + raise ValueError( + "Shared FC12 hidden and local intermediate sizes must be divisible by 128" + ) + if any(t.device != fc1.device for t in sources): + raise ValueError("Shared FC12 weights and scales must be on the same device") + if fc1_sf.dtype == torch.float32 and fc2_sf.dtype == torch.float32: + # Linear kept UE8M0-valued FP32 K128 scales (SM100/103). + fc1_sf = transform_k128_scales_to_cutedsl_mxfp8_layout( + fc1_sf, mn=fc1.shape[0], k=fc1.shape[1] + ) + fc2_sf = transform_k128_scales_to_cutedsl_mxfp8_layout( + fc2_sf, mn=fc2.shape[0], k=fc2.shape[1] + ) + logical_fc1_sf = _unpack_r128c4(fc1_sf, 2 * intermediate, self.hidden_size // 32) + _unpack_r128c4(fc2_sf, self.hidden_size, intermediate // 32) + # Byte views keep packing independent of PyTorch FP8 indexing support. + packed_fc1 = _interleave_gate_up(fc1.view(torch.uint8)).view(fc1.dtype).unsqueeze(0) + packed_fc1_sf = _pack_r128c4(_interleave_gate_up(logical_fc1_sf)) + self._fc12_fc1_weight = packed_fc1 + self._fc12_fc1_weight_sf = packed_fc1_sf + self._fc12_fc2_weight = fc2.unsqueeze(0) + self._fc12_fc2_weight_sf = fc2_sf.unsqueeze(0) + self._fc12_sm_count = self.sm_count + self._fc12_source_state = source_state + + def forward( + self, + x: torch.Tensor, + all_rank_num_tokens: list[int] | None = None, + final_all_reduce_params: AllReduceParams | None = None, + lora_params: dict | None = None, + **kwargs, + ) -> torch.Tensor: + if lora_params or kwargs: + raise ValueError("Shared FC12 does not support LoRA or extra forward options") + if final_all_reduce_params is not None: + raise ValueError("Shared FC12 reduction belongs to the enclosing MoE module") + if not isinstance(x, torch.Tensor) or x.dtype != torch.bfloat16 or not x.is_cuda: + raise ValueError("Shared FC12 input must be a CUDA BF16 tensor") + if x.ndim != 2 or x.shape[1] != self.hidden_size or not x.is_contiguous(): + raise ValueError("Shared FC12 input must be contiguous [tokens, hidden_size]") + if x.shape[0] > self._fc12_max_tokens: + raise ValueError("Shared FC12 input exceeds configured max_num_tokens") + self.prepare_fc12_weights() + if x.device != self._fc12_fc1_weight.device: + raise ValueError("Shared FC12 input and weights must be on the same device") + from ..moe.kernels.megamoe_shared_fc12 import run_shared_fc12 + + return run_shared_fc12( + x, + self._fc12_fc1_weight, + self._fc12_fc1_weight_sf, + self._fc12_fc2_weight, + self._fc12_fc2_weight_sf, + swiglu_limit=self.swiglu_limit, + sm_count=self._fc12_sm_count, + max_tokens=self._fc12_max_tokens, + ) diff --git a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py index 475619e3582b..faac62e78056 100644 --- a/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py @@ -35,6 +35,7 @@ import dataclasses import functools +import math import weakref from typing import Any, List, Optional, Tuple @@ -58,6 +59,9 @@ from tensorrt_llm.logger import logger from tensorrt_llm.math_utils import ceil_div, pad_up +# Default tactics validated once when the runtime becomes available. +_PREVALIDATED_TACTICS: List[Tuple] = [] + _NVFP4_BLOCK_SIZE = 16 _SF_PADDING_BLOCK = 128 _SUPPORTED_MMA_TILE_M = (128, 256) @@ -71,6 +75,7 @@ "enumerate_megamoe_candidate_tactics", "megamoe_activation_sf_bytes_per_row", "validate_megamoe_tactic", + "synthesize_profiling_topk", ] # Set to ``True`` if every symbol the op registration needs imports @@ -134,7 +139,7 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # --------------------------------------------------------------------------- -# Tactic representation (v4: 10-tuple perf knobs) +# Tactic representation (v5: 11-tuple perf knobs) # --------------------------------------------------------------------------- # # A tactic is a tuple of JSON-friendly primitives (lists / ints / bools / @@ -151,7 +156,8 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # use_bulk_fc2_store, # bool (epi -> peer UBLK; non-epi -> local TMA) # fc2_tma_stages, # optional staged-bulk depth # flag_batch, # int >= 1 (standalone_warps requires == 1) -# epi_flag_batch) # (int, int) fc1/fc2 done-counter publish batch +# epi_flag_batch, # (int, int) fc1/fc2 done-counter publish batch +# token_back_ready_granularity) # "expert" | "token_tile" # # Derived / out-of-tuple: # use_2cta_instrs = (mma_tiler_mnk[0] == 256), derived in _build_kernel. @@ -160,8 +166,11 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: # Tuple wrapping makes the tactic hashable, which AutoTuner needs for the # tactics cache. -_TACTIC_LEN = 10 +_TACTIC_LEN = 11 +_TACTIC_LEN_V4 = 10 _LEGACY_TACTIC_LEN = 8 +_DEFAULT_TOKEN_BACK_READY_GRANULARITY = "expert" # nosec B105 -- scheduler mode name, not a credential +_TOKEN_BACK_READY_GRANULARITIES = ("expert", "token_tile") # Kernel-side ceiling: ``flag_batch`` is hard-checked ``[1, 32]`` and # ``epi_flag_batch`` entries are SILENTLY clamped to ``[1, 32]``; reject > 32 @@ -169,14 +178,109 @@ def _megamoe_fence_key_live(fkey: Tuple[int, int, int, str]) -> bool: _FLAG_BATCH_MAX = 32 +_AUTOTUNE_PL_SEED = 0xA107 + + +def _autotune_pl_alpha(helper_expert_count: int = 0) -> float: + """Use the balanced-rank power-law workload only for load-balanced tuning.""" + return 0.8 if helper_expert_count > 0 else 0.0 + + +def synthesize_profiling_topk( + *, + num_tokens: int, + num_topk: int, + num_experts_per_rank: int, + world_size: int, + alpha: float, + device: torch.device, + source_rank: int = 0, +) -> torch.Tensor: + """Build physical routes with equal rank totals and ascending local load. + + Every source rank has T tokens and T*K routes. Rotating the destination + round-robin by source_rank*T*K makes the aggregate receive count exactly + T*K on every destination, including small buckets where EP does not + divide T*K. A single source's destination totals differ by at most one. + + Within each destination, sample with weights (E, ..., 1)**(-alpha). + Sampling is without replacement only within a token and destination; + equal local expert IDs on different ranks are distinct physical experts. + Relabel by each destination's realized counts so its expert load is + non-decreasing, including after summing all source ranks. Integer counts + fluctuate around the power-law weights; rank totals are exact. + + alpha=0 retains the OFF uniform round-robin byte for byte. This helper + runs only during profiling setup, outside candidate timing. + """ + tokens, topk = int(num_tokens), int(num_topk) + experts, ranks = int(num_experts_per_rank), int(world_size) + if tokens < 0 or experts <= 0 or ranks <= 0 or not 0 < topk <= experts * ranks: + raise ValueError("Invalid MegaMoE profiling route dimensions") + if not 0 <= source_rank < ranks or not math.isfinite(alpha) or alpha < 0: + raise ValueError("Invalid MegaMoE profiling source rank or exponent") + flat = torch.arange(tokens * topk, dtype=torch.long, device=device) + if alpha == 0: + return (flat % (experts * ranks)).view(tokens, topk) + destination = ((flat + source_rank * tokens * topk) % ranks).view(tokens, topk) + local_draws = (topk + ranks - 1) // ranks + log_weights = -alpha * torch.arange(experts, 0, -1, dtype=torch.float32, device=device).log() + generator = torch.Generator(device=device) + generator.manual_seed(_AUTOTUNE_PL_SEED) + noise = torch.rand((tokens, ranks, experts), generator=generator, device=device).clamp_( + min=torch.finfo(torch.float32).tiny + ) + gumbel = -torch.log(-torch.log(noise)) + candidates = torch.topk(log_weights + gumbel, local_draws, dim=-1).indices + local = candidates[ + torch.arange(tokens, device=device).unsqueeze(1), + destination, + torch.arange(topk, device=device) // ranks, + ] + counts = torch.bincount( + (destination * experts + local).reshape(-1), minlength=ranks * experts + ).view(ranks, experts) + order = torch.argsort(counts, dim=1, stable=True) + relabel = torch.empty_like(order) + relabel.scatter_(1, order, torch.arange(experts, device=device).expand(ranks, experts)) + return destination * experts + relabel[destination, local] + + +def _token_back_ready_granularity_choices(*, load_balance: bool = False) -> tuple[str, ...]: + """Compare both readiness modes only for load-balanced tuning.""" + return _TOKEN_BACK_READY_GRANULARITIES if load_balance else ("expert",) + + +@functools.cache +def _megamoe_tuning_op_name(helper_expert_count: int) -> str: + """Separate timing caches without changing the registered op or kernel ABI. + + Every launch that may tune asks for it, so each name is built once. + """ + base = "trtllm::cute_dsl_megamoe_nvfp4_blackwell" + load_balance = helper_expert_count > 0 + alpha = _autotune_pl_alpha(helper_expert_count) + granularities = _token_back_ready_granularity_choices(load_balance=load_balance) + if not load_balance and alpha == 0 and granularities == ("expert",): + return base + return ( + f"{base}:profile_v2:load_balance={int(load_balance)}:alpha={alpha!r}" + f":ready={','.join(granularities)}" + ) + + def _unpack_tactic(tactic: Tuple) -> Tuple: - """Return the canonical 10 fields, accepting the legacy 8-tuple. + """Return the canonical 11 fields, accepting the v4 10- and v3 8-tuples. Legacy tactics map ``static`` to grid-stride work IDs and preserve their grouped scheduling hint. Their single-stage bulk store is represented explicitly so old autotuner cache entries keep identical code generation. + Pre-v5 tactics take ``token_back_ready_granularity="expert"`` -- the only + value this op compiled before the field entered the tuple, so a cache + written by an older build keeps its exact meaning. Validation remains centralized in :func:`validate_megamoe_tactic`. """ + token_back_ready_granularity = _DEFAULT_TOKEN_BACK_READY_GRANULARITY if len(tactic) == _LEGACY_TACTIC_LEN: ( mma_tiler, @@ -200,6 +304,19 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: f"'atomic_counter', got {load_balance_mode!r}." ) fc2_tma_stages = 1 if use_bulk_fc2_store else None + elif len(tactic) == _TACTIC_LEN_V4: + ( + mma_tiler, + cluster_shape, + fallback_cluster_shape, + schedule_policy, + work_id_mode, + token_back_mode, + use_bulk_fc2_store, + fc2_tma_stages, + flag_batch, + epi_flag_batch, + ) = tactic else: ( mma_tiler, @@ -212,6 +329,7 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = tactic return ( @@ -225,14 +343,60 @@ def _unpack_tactic(tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) +# Reuse default tactic objects so validation can be memoized by identity. +# These shared constants must not be mutated. Value-only memoization could +# conflate bool with int and bypass tactic type checks. +_DEFAULT_TACTIC_DECODE: Tuple = ( + [256, 128, 256], + [2, 1, 1], + 512, + "static", + "epi_warps", + True, + 1, + (1, 1), +) +_DEFAULT_TACTIC_MID: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "static", + "epi_warps", + True, + 4, + (1, 1), +) +_DEFAULT_TACTIC_PREFILL_STATIC: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "static", + "reuse_dispatch_warps", + False, + 8, + (2, 4), +) +_DEFAULT_TACTIC_PREFILL_ATOMIC: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "atomic_counter", + "reuse_dispatch_warps", + False, + 8, + (2, 4), +) + + def default_megamoe_tactic(num_tokens: int) -> Tuple: """Deterministic token-bucket fallback when no shape-tuned default exists; never profiled by the autotuner.""" if num_tokens <= 1024: - # decode winner: N128 only helps epi_warps + bulk. + # Deterministic fallback for the smallest token bucket. return ( [256, 128, 256], [2, 1, 1], @@ -244,6 +408,7 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: 1, 1, (1, 1), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) if num_tokens <= 8192: return ( @@ -257,8 +422,9 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: 1, 4, (1, 1), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) - # prefill: atomic_counter only helps the very large tail (>=16384). + # Use the atomic work distributor for the upper token bucket. return ( [256, 256, 256], [2, 1, 1], @@ -270,13 +436,13 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: None, 8, (2, 4), + _DEFAULT_TOKEN_BACK_READY_GRANULARITY, ) -# Static defaults for the SM107 DSv4 Pro DEP4 and DEP8 shapes. -# Routing distribution is not a -# rank-identical runtime input, so each bucket uses one static tactic. The full -# runner key prevents applying the table beyond the measured shape and codegen modes. +# Static defaults for the explicitly keyed SM107 configurations. Routing is not +# rank-identical, so each bucket uses one deterministic tactic. The complete +# runner key prevents applying an entry to a different shape or codegen mode. _SM107_DSV4_PRO_DEFAULT_TACTICS: dict[Tuple, Tuple] = { ( 107, @@ -296,6 +462,7 @@ def default_megamoe_tactic(num_tokens: int) -> Tuple: "bf16", None, None, + 0, ): ( list(mma_tiler), [4, 1, 1], @@ -524,6 +691,7 @@ def _megamoe_problem_key( combine_format: str, swiglu_alpha: Optional[float] = None, swiglu_beta: Optional[float] = None, + helper_expert_count: int = 0, ) -> Tuple: return ( int(sm_version), @@ -543,6 +711,8 @@ def _megamoe_problem_key( str(combine_format), None if swiglu_alpha is None else float(swiglu_alpha), None if swiglu_beta is None else float(swiglu_beta), + # The helper count changes generated code and must be part of its cache key. + int(helper_expert_count), ) @@ -557,19 +727,19 @@ def _default_megamoe_tactic_for_problem( shape_tuned_tactic: Optional[Tuple] = None, ) -> Tuple: if combine_format != "bf16": - return _MEGAMOE_NONBULK_STANDALONE_TACTIC + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) bucket = _sm107_tactic_bucket(max_tokens_per_rank) if sm_version == 107 else None if bucket is not None: expected_in_kernel_fc2_reduce = _sm107_use_in_kernel_fc2_reduce( max_tokens_per_rank, apply_topk_in_fc1 ) if in_kernel_fc2_reduce != expected_in_kernel_fc2_reduce: - return _MEGAMOE_NONBULK_STANDALONE_TACTIC - return _SM107_GENPHASE_TACTICS[bucket] + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + return _unpack_tactic(_SM107_GENPHASE_TACTICS[bucket]) if in_kernel_fc2_reduce: - return _MEGAMOE_NONBULK_STANDALONE_TACTIC + return _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) if shape_tuned_tactic is not None: - return shape_tuned_tactic + return _unpack_tactic(shape_tuned_tactic) return default_megamoe_tactic(num_tokens) @@ -587,14 +757,16 @@ def _default_megamoe_tactic_for_key(problem_key: Tuple, *, num_tokens: int) -> T def _megamoe_default_tactic_candidates(tactic: Tuple) -> Tuple[Tuple, ...]: """Return the heuristic tactic followed by shape-safe fallbacks.""" + tactic = _unpack_tactic(tactic) candidates = [tactic] - schedule_policy = _unpack_tactic(tactic)[3] + schedule_policy = tactic[3] if schedule_policy[0] == "phase_interleave" and schedule_policy[1] is not None: auto_hint_tactic = list(tactic) auto_hint_tactic[3] = ("phase_interleave", None) candidates.append(tuple(auto_hint_tactic)) - if _MEGAMOE_NONBULK_STANDALONE_TACTIC not in candidates: - candidates.append(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + nonbulk_tactic = _unpack_tactic(_MEGAMOE_NONBULK_STANDALONE_TACTIC) + if nonbulk_tactic not in candidates: + candidates.append(nonbulk_tactic) return tuple(candidates) @@ -713,13 +885,13 @@ def _construct_megamoe_kernel( """Construct the active architecture kernel and report a narrow gen rejection.""" from cutlass.cute.nvgpu import OperandMajorMode - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe import ( + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import CombineFormat, QuantKind + from tensorrt_llm._torch.moe.kernels.megamoe_kernels import ( BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceMegaMoE, ) - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.api import ImplDesc, ProblemDesc - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.quant_def import CombineFormat, QuantKind local_experts, intermediate_gateup_size, hidden_size = common["static_expert_shape"] problem_desc = ProblemDesc( @@ -754,12 +926,15 @@ def _construct_megamoe_kernel( "work_id_mode": str(common["work_id_mode"]), "fc2_use_bulk": not bool(common["non_ubulk_fc2_store"]), "epi_flag_batches": tuple(common["epi_flag_batch"]), + # Kernel construction and READY arguments must use the same helper count. + "helper_expert_count": int(common.get("helper_expert_count", 0)), "launch_cluster_count": int(common["launch_cluster_count"]), "fallback_cluster_shape_mn": common["fallback_cluster_shape_mn"], "preferred_cluster_count": common["preferred_cluster_count"], "fallback_cluster_count": common["fallback_cluster_count"], "token_in_flag_batch": int(common["flag_batch"]), "token_back_mode": str(common["token_back_mode"]), + "token_back_ready_granularity": str(common.get("token_back_ready_granularity", "expert")), "reduce_topk_in_kernel": bool(common["in_kernel_fc2_reduce"]), } if sm_version == 107: @@ -775,15 +950,25 @@ def _construct_megamoe_kernel( impl_values["fc2_tma_stages"] = int(common["fc2_tma_stages"]) impl_desc = ImplDesc(impl_values) + def construct(kernel_cls): + # ProblemDesc.validate ignores undeclared keys, so a kernel without the SwiGLU + # constants would silently run plain SwiGLU. + if ( + common["swiglu_alpha"] is not None + and "swiglu_alpha" not in kernel_cls.problem_desc_require() + ): + raise NotImplementedError(f"{kernel_cls.__name__} does not support swiglu_alpha.") + return kernel_cls(problem_desc, impl_desc) + bucket = _sm107_tactic_bucket(int(common["max_tokens_per_rank"])) if sm_version == 107 else None rejection = None if bucket is not None and bucket not in _SM107_GENERIC_KERNEL_BUCKETS: try: - return RubinInferenceGenphaseMegaMoE(problem_desc, impl_desc), None + return construct(RubinInferenceGenphaseMegaMoE), None except (ValueError, TypeError, NotImplementedError) as exc: rejection = exc kernel_cls = RubinInferenceMegaMoE if sm_version == 107 else BlackwellInferenceMegaMoE - return kernel_cls(problem_desc, impl_desc), rejection + return construct(kernel_cls), rejection def _megamoe_barrier_reset_slices(kernel) -> Tuple[int, int, int, int, str]: @@ -822,11 +1007,12 @@ def validate_megamoe_tactic(tactic: Tuple, sm_version: int = 100) -> None: """ if (not isinstance(tactic, tuple)) or len(tactic) not in { _TACTIC_LEN, + _TACTIC_LEN_V4, _LEGACY_TACTIC_LEN, }: raise ValueError( f"MegaMoE tactic must be a {_TACTIC_LEN}-tuple (or legacy " - f"{_LEGACY_TACTIC_LEN}-tuple), got " + f"{_TACTIC_LEN_V4}- / {_LEGACY_TACTIC_LEN}-tuple), got " f"{type(tactic).__name__} len={len(tactic) if isinstance(tactic, tuple) else 'NA'}={tactic!r}" ) is_legacy_tactic = len(tactic) == _LEGACY_TACTIC_LEN @@ -841,6 +1027,7 @@ def validate_megamoe_tactic(tactic: Tuple, sm_version: int = 100) -> None: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) if (not isinstance(mma_tiler, (list, tuple))) or len(mma_tiler) != 3: @@ -1001,10 +1188,89 @@ def validate_cluster_shape(shape, field_name: str) -> None: f"(epilogue clamp range), got {epi_flag_batch!r}." ) + if token_back_ready_granularity not in _TOKEN_BACK_READY_GRANULARITIES: + raise ValueError( + f"token_back_ready_granularity must be one of " + f"{_TOKEN_BACK_READY_GRANULARITIES}, got {token_back_ready_granularity!r}." + ) + if token_back_ready_granularity == "token_tile": # nosec B105 -- scheduler mode name, not a credential + # Mirrors the kernel's own guard (kernel_src/*/inference/mega/ + # block_scaled_swap_ab_mega_moe_kernel.py, _validate_geometry). Checking + # it HERE is the point of putting the granularity in the tuple: the + # kernel raises at build time, which is a crash, while enumeration wraps + # every candidate in this function and simply drops what fails. The + # kernel-side copy stays as the backstop for hand-supplied tactics. + # + # Two of the kernel's nine conditions are not tactic-local and are not + # re-checked here: + # quant_kind == nvfp4 -- this op's ProblemDesc hardcodes nvfp4. + # not reduce_topk_in_kernel -- implied transitively: token_tile + # requires use_bulk_fc2_store, and get_valid_tactics already drops + # every bulk candidate when in_kernel_fc2_reduce is set. + unmet = [] + if work_id_mode != "atomic_counter": + unmet.append(f"work_id_mode={work_id_mode!r} (need 'atomic_counter')") + if token_back_mode not in ("standalone_warps", "reuse_dispatch_warps"): + unmet.append( + f"token_back_mode={token_back_mode!r} " + f"(need 'standalone_warps' or 'reuse_dispatch_warps')" + ) + if mma_tiler[0] != 256: + unmet.append(f"mma_tiler_mnk[0]={mma_tiler[0]} (need 256, i.e. use_2cta_instrs)") + if not use_bulk_fc2_store: + unmet.append("use_bulk_fc2_store=False (need True, FC2 TMA stores)") + if unmet: + raise ValueError( + "token_back_ready_granularity='token_tile' is unmet by: " + "; ".join(unmet) + ) + + +def _expand_megamoe_ready_tactics( + candidates: list[Tuple], sm_version: int, *, load_balance: bool +) -> list[Tuple]: + """Return canonical tactics for every legal token-back READY mode.""" + choices = _token_back_ready_granularity_choices(load_balance=load_balance) + if sm_version == 107 and choices != ("expert",): + # At every token bucket include a legal local-TMA/atomic-counter + # base; the curated SM107 candidates otherwise exclude token_tile. + for cluster, fallback, policy, mode, flag, epi in ( + ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (1, 1)), + ([2, 1, 1], None, ("grouped", 512), "reuse_dispatch_warps", 4, (2, 4)), + ([2, 1, 1], None, ("grouped", 512), "standalone_warps", 1, (2, 4)), + ([4, 1, 1], [2, 1, 1], ("phase_interleave", None), "reuse_dispatch_warps", 4, (2, 4)), + ): + candidates.append( + ( + [256, 256, 256], + cluster, + fallback, + policy, + "atomic_counter", + mode, + True, + 2, + flag, + epi, + ) + ) + expanded = [] + for base in candidates: + for granularity in choices: + tactic = (*_unpack_tactic(base)[:_TACTIC_LEN_V4], granularity) + try: + validate_megamoe_tactic(tactic, sm_version=sm_version) + except ValueError: + continue + if tactic not in expanded: + expanded.append(tactic) + return expanded + def enumerate_megamoe_candidate_tactics( num_tokens: int, sm_version: int = 100, + *, + load_balance: bool = False, ) -> List[Tuple]: """Return the curated candidate tactic list for the current token bucket. @@ -1093,7 +1359,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(mixed, sm_version=sm_version) candidates.append(mixed) - return candidates + return _expand_megamoe_ready_tactics(candidates, sm_version, load_balance=load_balance) if sm_version not in (100, 103): raise ValueError( f"MegaMoE tactic enumeration supports SM100, SM103, or SM107; got SM{sm_version}." @@ -1125,11 +1391,8 @@ def enumerate_megamoe_candidate_tactics( continue candidates.append(tactic) - # These additions are intentionally curated instead of crossed with every - # legacy axis: staged FC2 UBLK and mixed-CGA/phase scheduling are expensive - # to compile and autotuning is opt-in. These entries expose the upstream - # performance mechanisms while preserving the existing deterministic - # production defaults. + # Add a bounded set of valid staged and mixed-scheduling candidates rather + # than crossing every independent tactic axis. if num_tokens >= 2048: for fc2_tma_stages in (2, 4): tactic = ( @@ -1146,8 +1409,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(tactic, sm_version=sm_version) candidates.append(tactic) - # Local TMA stage 2 was the only consistently useful depth in B200 - # A/B tests: stage 1 tied it while stage 4 regressed from extra SMEM. + # Include the resource-balanced local TMA configuration. tactic = ( [256, 256, 256], [2, 1, 1], @@ -1177,7 +1439,7 @@ def enumerate_megamoe_candidate_tactics( ) validate_megamoe_tactic(tactic, sm_version=sm_version) candidates.append(tactic) - return candidates + return _expand_megamoe_ready_tactics(candidates, sm_version, load_balance=load_balance) if IS_CUTLASS_DSL_AVAILABLE: @@ -1200,15 +1462,15 @@ def enumerate_megamoe_candidate_tactics( except ImportError: from cuda import cuda - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe import ( # noqa: F401 + from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.communication.nvlink_domain.symmetric_buffer import ( + SymmetricBufferHost, # noqa: F401 + ) + from tensorrt_llm._torch.moe.kernels.megamoe_kernels import ( # noqa: F401 BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceLocalMegaMoE, RubinInferenceMegaMoE, ) - from tensorrt_llm._torch.moe.kernels.cutedsl_megamoe.communication.nvlink_domain.symmetric_buffer import ( - SymmetricBufferHost, # noqa: F401 - ) IS_MEGAMOE_OP_AVAILABLE = True except Exception as _megamoe_import_err: # pragma: no cover - env-specific @@ -1227,6 +1489,27 @@ def enumerate_megamoe_candidate_tactics( if IS_MEGAMOE_OP_AVAILABLE: + # Validate default tactics before enabling identity-based validation bypass. + # If validation fails, retain normal per-call checks rather than failing import. + for _prevalidate_t in ( + _DEFAULT_TACTIC_DECODE, + _DEFAULT_TACTIC_MID, + _DEFAULT_TACTIC_PREFILL_STATIC, + _DEFAULT_TACTIC_PREFILL_ATOMIC, + _MEGAMOE_NONBULK_STANDALONE_TACTIC, + ): + try: + validate_megamoe_tactic(_prevalidate_t) + except Exception as _prevalidate_err: # pragma: no cover - partial installation + logger.debug( + "[MegaMoE] tactic prevalidation skipped (%s); falling back to " + "per-launch validation.", + _prevalidate_err, + ) + _PREVALIDATED_TACTICS.clear() + break + _PREVALIDATED_TACTICS.append(_prevalidate_t) + # ----- Local workspace cache -------------------------------------------- # # ``local_workspace`` is per-rank CUDA-only and sized by @@ -1243,6 +1526,11 @@ def enumerate_megamoe_candidate_tactics( # freed when the next one allocates (only the winner is reused). _MEGAMOE_TUNING_WORKSPACE_KEYS: set = set() + # Cache runners by constructor inputs, except while the AutoTuner profiles + # an autotune opt-in. Reset profiling scratch on reuse and overwrite READY + # inputs for every call. + _MEGAMOE_RUNNER_CACHE: dict = {} + @functools.lru_cache(maxsize=1) def _cute_launch_helpers(): """Import the cutlass-dsl launch surface once per process and build the @@ -1268,6 +1556,60 @@ def to_cute_ptr(t, assumed_align=16): return to_cute, to_cute_ptr, SymBufferHost, cute + def _cute_view_key(t: torch.Tensor, assumed_align: int) -> Tuple: + """Everything a cute view of ``t`` encodes.""" + return (t.data_ptr(), tuple(t.shape), t.stride(), t.dtype, t.device.index, assumed_align) + + def _to_cute_cached(cache: dict, t: torch.Tensor, assumed_align: int = 16): + """Return the cute view of ``t`` from ``cache``, converting it on a miss. + + The key holds everything the view encodes, so a hit equals a fresh + conversion even if the address was reused. Only for tensors whose + layout is fixed (module weights, launch buffers): one entry per tensor. + """ + key = _cute_view_key(t, assumed_align) + cute_tensor = cache.get(key) + if cute_tensor is None: + cute_tensor = cache[key] = _cute_launch_helpers()[0](t, assumed_align) + return cute_tensor + + def _to_cute_slot(cache: dict, slot: str, t: torch.Tensor, assumed_align: int = 16): + """Like ``_to_cute_cached`` but keep only the latest view under ``slot``. + + For a buffer whose live row count changes between iterations but not + between the layers of one iteration. + """ + key = _cute_view_key(t, assumed_align) + entry = cache.get(slot) + if entry is None or entry[0] != key: + entry = cache[slot] = (key, _cute_launch_helpers()[0](t, assumed_align)) + return entry[1] + + # cute pointers by (address, alignment): a pointer encodes nothing else, so + # a hit equals a fresh ``make_ptr``. Bounded because tuning visits many + # workspaces; a pointer holds no tensor reference. + _MEGAMOE_CUTE_PTR_CACHE: dict = {} + _MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES = 256 + + def _to_cute_ptr_cached(t: torch.Tensor, assumed_align: int = 16): + key = (t.data_ptr(), assumed_align) + ptr = _MEGAMOE_CUTE_PTR_CACHE.get(key) + if ptr is None: + if len(_MEGAMOE_CUTE_PTR_CACHE) >= _MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES: + _MEGAMOE_CUTE_PTR_CACHE.clear() + ptr = _MEGAMOE_CUTE_PTR_CACHE[key] = _cute_launch_helpers()[1](t, assumed_align) + return ptr + + # cuda-python stream wrappers by raw handle. + _MEGAMOE_CU_STREAMS: dict = {} + + def _current_cu_stream(): + handle = torch._C._cuda_getCurrentRawStream(torch._C._cuda_getDevice()) + stream = _MEGAMOE_CU_STREAMS.get(handle) + if stream is None: + stream = _MEGAMOE_CU_STREAMS[handle] = cuda.CUstream(handle) + return stream + def _evict_latched_tuning_workspaces() -> None: """Pop every tuning-latched local workspace and prune its fence keys. @@ -1358,7 +1700,7 @@ class MegaMoeSymmRegions: # (max_T, 1, hidden): both reduction forms collapse the top-k axis in-op. combine_output: torch.Tensor shared_workspace: torch.Tensor # (shared_ws_bytes,) uint8 - peer_offsets: List[int] # symmetric peer-pointer deltas + peer_offsets: Tuple[int, ...] # symmetric peer-pointer deltas rank: int world_size: int @@ -1479,6 +1821,13 @@ def __init__( peer_ptr = int(self._handle.buffer_ptrs[r]) self.peer_offsets.append(peer_ptr - local_base) + # Built on first use per inference-mode state: a view made under + # inference mode cannot be updated in place outside it. + self._regions: dict = {} + # Cute views of these regions, shared by every layer launching on + # this provider (set as the op's buffer cache by the backend). + self.cute_views: dict = {} + logger.debug( "[MegaMoeSymmMemProvider] group=%s rank=%d/%d total_bytes=%d " "(activation=%d sf=%d topk_weights=%d combine=%d shared=%d)", @@ -1505,6 +1854,13 @@ def _region_view( return byte_view.view(dtype).view(shape) def get_regions(self) -> MegaMoeSymmRegions: + inference_mode = torch.is_inference_mode_enabled() + regions = self._regions.get(inference_mode) + if regions is None: + regions = self._regions[inference_mode] = self._build_regions() + return regions + + def _build_regions(self) -> MegaMoeSymmRegions: hidden = self.hidden_size max_t = self.max_tokens_per_rank top_k = self.num_topk @@ -1529,7 +1885,8 @@ def get_regions(self) -> MegaMoeSymmRegions: shared_workspace=self._region_view( "shared_workspace", (self._region_sizes["shared_workspace"],), torch.uint8 ), - peer_offsets=list(self.peer_offsets), + # A tuple: the regions object is shared across launches. + peer_offsets=tuple(self.peer_offsets), rank=self.rank, world_size=self.world_size, ) @@ -1647,6 +2004,27 @@ def get_megamoe_profiling_scratch( _MEGAMOE_PROFILING_SCRATCH_CACHE[cache_key] = provider return provider + # Cute views of the calling module's weights, set by the backend around + # the op call like the profiling scratch. A cute view keeps its tensor + # alive, so the module owns the dict and the views never outlive it. + _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE: Optional[dict] = None + + def set_active_megamoe_cute_weight_cache(cache: Optional[dict]) -> None: + """Set (``None`` clears) the calling module's cute weight-view cache.""" + global _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE + _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE = cache + + # Cute views of the symmetric launch buffers, set by the backend only + # when activation, activation_sf, topk_weights and the combine output are + # its provider's regions and topk_idx is its own staging buffer. The + # provider owns the dict; topk_idx views go to the module's weight cache. + _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE: Optional[dict] = None + + def set_active_megamoe_cute_buffer_cache(cache: Optional[dict]) -> None: + """Set (``None`` clears) the active provider's cute launch-buffer cache.""" + global _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE + _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE = cache + def set_active_megamoe_profiling_scratch(regions) -> None: """Set (``None`` clears) the :class:`MegaMoeSymmRegions` used for AutoTuner profiling on the current call (set by the backend around @@ -1752,6 +2130,7 @@ def build(candidate: Tuple): fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(candidate) mma_tiler = tuple(mma_tiler) ( @@ -1796,6 +2175,7 @@ def build(candidate: Tuple): fc2_output_dtype=cutlass.BFloat16, in_kernel_fc2_reduce=bool(in_kernel_fc2_reduce), token_back_mode=str(token_back_mode), + token_back_ready_granularity=str(token_back_ready_granularity), non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), flag_batch=int(flag_batch), epi_flag_batch=tuple(epi_flag_batch), @@ -1877,6 +2257,7 @@ def __init__( in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, + helper_expert_count: int = 0, ) -> None: super().__init__() if (sm_version := get_sm_version()) not in (100, 103, 107): @@ -1929,12 +2310,28 @@ def __init__( # generated kernel, so deliberately EXCLUDED from unique_id / # _tactic_cache_key: opted-in and opted-out runs share caches. self.tactic_autotune = bool(tactic_autotune) + # ``S``. CODEGEN-CHANGING: a helper-bearing build carries the READY + # gate and a static ``H`` boundary in the sched warp, so it MUST NOT + # share a compile-cache entry with the ungated baseline. Hence it + # is in unique_id() (which _tactic_cache_key wraps). + self.helper_expert_count = int(helper_expert_count) + # Per-launch READY runtime bundle (device pointer + generation). + # Set by the op right after construction; never a codegen input. + self._helper_ready_flags = None + self._helper_ready_generation = 0 def unique_id(self): # local_rank is intentionally excluded: every EP rank must run the # SAME tactic and MERGE merges timings by cache key across ranks, # so the key MUST be rank-identical. world_size stays so - # single-rank and multi-rank never share entries. + # single-rank and multi-rank never share entries. The key depends + # only on constructor inputs, so each runner builds it once. + key = self.__dict__.get("_unique_id") + if key is None: + key = self._unique_id = self._build_unique_id() + return key + + def _build_unique_id(self): return _megamoe_problem_key( sm_version=self.sm_version, world_size=self.world_size, @@ -1955,6 +2352,7 @@ def unique_id(self): situ_linear_beta=self.situ_linear_beta, in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, combine_format=self.combine_format, + helper_expert_count=self.helper_expert_count, ) def get_valid_tactics( @@ -1965,7 +2363,11 @@ def get_valid_tactics( ) -> List[Tuple]: del profile, kwargs num_tokens = int(inputs[0].shape[0]) - candidates = enumerate_megamoe_candidate_tactics(num_tokens, sm_version=self.sm_version) + candidates = enumerate_megamoe_candidate_tactics( + num_tokens, + sm_version=self.sm_version, + load_balance=self.helper_expert_count > 0, + ) # Form-B remains non-bulk. FP4 only rejects peer UBLK; local TMA # supports its packed payload. Fields 5/6 are token-back mode and # use_bulk_fc2_store respectively. @@ -1988,17 +2390,22 @@ def get_valid_tactics( if default_tactic not in candidates: candidates.append(default_tactic) default_tactic = _SM107_DSV4_PRO_DEFAULT_TACTICS.get(self.unique_id()) - if default_tactic is not None and default_tactic not in candidates: - candidates.append(default_tactic) - return candidates + if default_tactic is not None: + default_tactic = _unpack_tactic(default_tactic) + if default_tactic not in candidates: + candidates.append(default_tactic) + choices = _token_back_ready_granularity_choices( + load_balance=self.helper_expert_count > 0 + ) + return [t for t in candidates if _unpack_tactic(t)[10] in choices] def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: """Sanitize ONLY the autotuner-regenerated fake inputs. AutoTuner rebuilds inputs 0-3 and 11; the static inputs (4-10) are the caller's REAL weights/scales, passed by reference -- filling - those would clobber them. topk_idx becomes a valid round-robin - (random ints index OOB); SF / weights are filled with 1.0, NOT 0 + those would clobber them. ON topk_idx uses rank-balanced power-law + physical routes; OFF keeps round-robin. SF / weights use 1.0, NOT 0 (SF==0 degenerates the GEMMs and skews tactic timing). """ # Runs ONLY on a real profiling MISS (a tuning-mode cache HIT @@ -2026,15 +2433,19 @@ def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.T topk_idx = inputs[2] if isinstance(topk_idx, torch.Tensor) and topk_idx.dim() == 2: T, K = topk_idx.shape - valid = ( - torch.arange( - T * K, - dtype=topk_idx.dtype, - device=topk_idx.device, - ) - % total_experts - ).view(T, K) - topk_idx.copy_(valid) + alpha = _autotune_pl_alpha(self.helper_expert_count) + valid = synthesize_profiling_topk( + num_tokens=T, + num_topk=K, + num_experts_per_rank=self.num_experts_per_rank, + world_size=self.world_size, + source_rank=self.local_rank, + alpha=alpha, + device=topk_idx.device, + ) + topk_idx.copy_(valid.to(topk_idx.dtype)) + if self.local_rank == 0 and alpha > 0: + logger.debug("[MegaMoE] helper-bearing autotune workload prepared") scratch = self._profiling_scratch if scratch is not None: @@ -2082,13 +2493,15 @@ def get_tuning_config(self) -> TuningConfig: autotuner does not double-enumerate tile sizes for independent token axes. - Rebuilt on every call, NEVER cached across runners: + Built once per runner and NEVER shared across runners: ``inputs_pre_hook`` is bound to THIS runner's ``_profiling_scratch``; - a cached hook bound to a dead runner would mix non-symmetric - profiling tensors with the live runner's peer offsets (IMA + hang). + a hook bound to another runner would mix non-symmetric profiling + tensors with the live runner's peer offsets (IMA + hang). """ - - config = TuningConfig( + config = self.__dict__.get("_tuning_config") + if config is not None: + return config + config = self._tuning_config = TuningConfig( dynamic_tensor_specs=( DynamicTensorSpec( 0, 0, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2 @@ -2143,6 +2556,7 @@ def _build_kernel_once(self, tactic: Tuple): fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) ( launch_cluster_count, @@ -2186,6 +2600,7 @@ def _build_kernel_once(self, tactic: Tuple): fc2_output_dtype=cutlass.BFloat16, in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, token_back_mode=str(token_back_mode), + token_back_ready_granularity=str(token_back_ready_granularity), non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), flag_batch=int(flag_batch), epi_flag_batch=tuple(epi_flag_batch), @@ -2195,6 +2610,7 @@ def _build_kernel_once(self, tactic: Tuple): gate_up_clamp=self.gate_up_clamp, situ_beta=self.situ_beta, situ_linear_beta=self.situ_linear_beta, + helper_expert_count=self.helper_expert_count, **_LOCKED_KERNEL_KWARGS, ) kernel, rejection = _construct_megamoe_kernel( @@ -2228,6 +2644,7 @@ def _tactic_cache_key(self, tactic: Tuple) -> Tuple: fc2_tma_stages, flag_batch, epi_flag_batch, + token_back_ready_granularity, ) = _unpack_tactic(tactic) launch_config = _launch_cluster_configuration( cluster_shape, @@ -2246,6 +2663,7 @@ def _tactic_cache_key(self, tactic: Tuple) -> Tuple: fc2_tma_stages, int(flag_batch), tuple(epi_flag_batch), + str(token_back_ready_granularity), launch_config, ) @@ -2327,10 +2745,14 @@ def forward( # -- and a full zero wastes decode time). form-A overwrites. # Cached per FULL tactic: local workspace SIZE is tactic-dependent - # (see _tactic_cache_key). + # (see _tactic_cache_key). The key is computed HERE, not at the + # launch site further down: the local-workspace cache needs it + # first, and `tactic_t` is final from `_resolve_megamoe_kernel_tactic` + # above. + cache_key = self._tactic_cache_key(tactic_t) local_workspace = _get_or_alloc_local_workspace( kernel, - cache_key=self._tactic_cache_key(tactic_t), + cache_key=cache_key, device=activation.device, latch_in_tuning_mode=self.tactic_autotune, ) @@ -2464,30 +2886,53 @@ def forward( # uint8 workspaces MUST be cute.Pointer, NOT cute.Tensor: the # 32-bit memref shape field overflows once shared_workspace passes # 2 GiB (the kernel addresses by raw base + Int64 byte offset). - _to_cute, _to_cute_ptr, SymBufferHost, cute = _cute_launch_helpers() + _to_cute, _, SymBufferHost, cute = _cute_launch_helpers() + # Weights reuse the calling module's cute views and the symmetric + # launch buffers the active provider's; tuning launches convert afresh. + tuning = AutoTuner.get().is_tuning_mode + weight_cache = None if tuning else _ACTIVE_MEGAMOE_CUTE_WEIGHT_CACHE + buffer_cache = None if tuning else _ACTIVE_MEGAMOE_CUTE_BUFFER_CACHE + if weight_cache is None: + _to_cute_w = _to_cute + else: + _to_cute_w = functools.partial(_to_cute_cached, weight_cache) # combine_output (max_T, 1, hidden) reshapes freely to 2D. Weights # present K stride-1 via a transpose VIEW (DLPack carries the # strides); do NOT ``.contiguous()``. output_activation = combine_output.reshape(combine_output.shape[0], self.hidden_size) + if buffer_cache is None: + cute_activation = _to_cute(activation) + cute_activation_sf = _to_cute(activation_sf) + cute_topk_weights = _to_cute(topk_weights) + cute_output = _to_cute(output_activation) + cute_topk_idx = _to_cute(topk_idx) + else: + cute_activation = _to_cute_cached(buffer_cache, activation) + cute_activation_sf = _to_cute_cached(buffer_cache, activation_sf) + cute_topk_weights = _to_cute_cached(buffer_cache, topk_weights) + # Its live rows change per iteration, not between the layers of one. + cute_output = _to_cute_slot(buffer_cache, "output_activation", output_activation) + # topk_idx is the calling module's own staging buffer. + cute_topk_idx = _to_cute_w(topk_idx) runtime_kwargs = dict( - activation=_to_cute(activation), - activation_sf=_to_cute(activation_sf), - fc1_weight=_to_cute(fc1_weight.transpose(1, 2)), - fc1_weight_sf=_to_cute(fc1_weight_sf), - fc2_weight=_to_cute(fc2_weight.transpose(1, 2)), - fc2_weight_sf=_to_cute(fc2_weight_sf), - fc1_alpha=_to_cute(fc1_alpha, assumed_align=4), - fc2_alpha=_to_cute(fc2_alpha, assumed_align=4), - fc1_norm_const=_to_cute(fc1_norm_const, assumed_align=4), - output_activation=_to_cute(output_activation), - local_workspace=_to_cute_ptr(local_workspace), - shared_workspace=_to_cute_ptr(shared_workspace), - stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), + activation=cute_activation, + activation_sf=cute_activation_sf, + fc1_weight=_to_cute_w(fc1_weight.transpose(1, 2)), + fc1_weight_sf=_to_cute_w(fc1_weight_sf), + fc2_weight=_to_cute_w(fc2_weight.transpose(1, 2)), + fc2_weight_sf=_to_cute_w(fc2_weight_sf), + fc1_alpha=_to_cute_w(fc1_alpha, assumed_align=4), + fc2_alpha=_to_cute_w(fc2_alpha, assumed_align=4), + fc1_norm_const=_to_cute_w(fc1_norm_const, assumed_align=4), + output_activation=cute_output, + local_workspace=_to_cute_ptr_cached(local_workspace), + shared_workspace=_to_cute_ptr_cached(shared_workspace), + stream=_current_cu_stream(), ) runtime_kwargs.update( - topk_indices=_to_cute(topk_idx), - topk_scores=_to_cute(topk_weights), + topk_indices=cute_topk_idx, + topk_scores=cute_topk_weights, peer_rank_ptr_mapper_host=SymBufferHost( base_address=int(activation.untyped_storage().data_ptr()), offsets=tuple(int(off) for off in peer_offsets), @@ -2495,7 +2940,35 @@ def forward( max_ranks=int(self.world_size), ), ) - cache_key = self._tactic_cache_key(tactic_t) + # READY ABI tail: attached only when S > 0, i.e. rebalance is on. + # Contract: with S > 0 both values must be supplied. Fail loudly + # here rather than hand the kernel an unarmed gate. + if self.helper_expert_count > 0: + if self._helper_ready_flags is None: + raise RuntimeError( + "MegaMoE-CuteDSL: helper_expert_count > 0 requires " + "hot_expert_weight_ready_flags (uint64[EP] terminal " + "flags); got None." + ) + if not (1 <= int(self._helper_ready_generation) < (1 << 63)): + raise RuntimeError( + "MegaMoE-CuteDSL: hot_expert_weight_ready_generation " + f"must be in [1, 2**63); got " + f"{self._helper_ready_generation}." + ) + if self._helper_ready_flags.numel() < self.world_size: + raise RuntimeError( + "MegaMoE-CuteDSL: hot_expert_weight_ready_flags must " + f"hold at least world_size={self.world_size} uint64 " + f"cells; got numel={self._helper_ready_flags.numel()}." + ) + runtime_kwargs["hot_expert_weight_ready_flags"] = _to_cute_ptr_cached( + self._helper_ready_flags + ) + runtime_kwargs["hot_expert_weight_ready_generation"] = cutlass.Uint64( + int(self._helper_ready_generation) + ) + # unique_id() includes helper_expert_count, separating cached ABIs. cached = self.__class__.kernel_cache.get(cache_key) if cached is None: compiled = cute.compile(kernel, **runtime_kwargs) @@ -2505,14 +2978,94 @@ def forward( compiled(**runtime_kwargs) return combine_output + def _megamoe_get_runner( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + output_dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + situ_beta, + situ_linear_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, + ): + """Reuse a runner by its complete constructor key. + + Reset profiling scratch on reuse; each op call supplies fresh READY inputs. + """ + args = ( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + output_dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + situ_beta, + situ_linear_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, + ) + + def _make(): + return Sm100MegaMoENvfp4Runner( + world_size=world_size, + local_rank=local_rank, + num_topk=num_topk, + num_experts_per_rank=num_experts_per_rank, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + expand_intermediate_size_per_partition=expand_intermediate_size_per_partition, + max_tokens_per_rank=max_tokens_per_rank, + output_dtype=output_dtype, + apply_topk_in_fc1=apply_topk_in_fc1, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + gate_up_clamp=gate_up_clamp, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + in_kernel_fc2_reduce=in_kernel_fc2_reduce, + combine_format=combine_format, + tactic_autotune=tactic_autotune, + helper_expert_count=helper_expert_count, + ) + + # A runner that may profile owns profiling scratch, so tuning-mode calls + # with the autotune opt-in get a fresh one. Outside tuning mode + # choose_one only looks up the tuned tactic, and the runner is reused. + if tactic_autotune and AutoTuner.get().is_tuning_mode: + return _make() + runner = _MEGAMOE_RUNNER_CACHE.get(args) + if runner is None: + runner = _make() + _MEGAMOE_RUNNER_CACHE[args] = runner + else: + # Reset profiling scratch on reuse; the caller supplies fresh READY fields. + runner._profiling_scratch = None + runner._profiling_scratch_factory = None + return runner + # ----- torch op --------------------------------------------------------- - @torch.library.custom_op( - "trtllm::cute_dsl_megamoe_nvfp4_blackwell", - mutates_args=("combine_output", "shared_workspace"), - device_types="cuda", - ) - def cute_dsl_megamoe_nvfp4_blackwell( + def run_cute_dsl_megamoe_nvfp4_blackwell( activation: torch.Tensor, activation_sf: torch.Tensor, topk_idx: torch.Tensor, @@ -2545,6 +3098,14 @@ def cute_dsl_megamoe_nvfp4_blackwell( combine_format: str = "bf16", tactic_autotune: bool = False, num_tokens: int = -1, + # READY ABI tail. ``helper_expert_count`` (S) is codegen-time; + # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal + # view published by the weight transport, ``..._generation`` the + # expected generation g in [1, 2**63). All three default to the + # S == 0 kernel without helper slots. + helper_expert_count: int = 0, + hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, + hot_expert_weight_ready_generation: int = 0, ) -> None: """Run the fused MegaMoE CuteDSL NVFP4 kernel. @@ -2586,26 +3147,31 @@ def cute_dsl_megamoe_nvfp4_blackwell( if num_tokens > 0: combine_output = combine_output[: min(num_tokens, combine_output.shape[0])] - runner = Sm100MegaMoENvfp4Runner( - world_size=world_size, - local_rank=local_rank, - num_topk=num_topk, - num_experts_per_rank=num_experts_per_rank, - hidden_size=hidden_size, - intermediate_size_per_partition=intermediate_size_per_partition, - expand_intermediate_size_per_partition=expand_intermediate_size_per_partition, - max_tokens_per_rank=max_tokens_per_rank, - output_dtype=combine_output.dtype, - apply_topk_in_fc1=apply_topk_in_fc1, - swiglu_alpha=swiglu_alpha, - swiglu_beta=swiglu_beta, - gate_up_clamp=gate_up_clamp, - situ_beta=act_alpha, - situ_linear_beta=act_beta, - in_kernel_fc2_reduce=in_kernel_fc2_reduce, - combine_format=combine_format, - tactic_autotune=tactic_autotune, + runner = _megamoe_get_runner( + world_size, + local_rank, + num_topk, + num_experts_per_rank, + hidden_size, + intermediate_size_per_partition, + expand_intermediate_size_per_partition, + max_tokens_per_rank, + combine_output.dtype, + apply_topk_in_fc1, + swiglu_alpha, + swiglu_beta, + gate_up_clamp, + act_alpha, + act_beta, + in_kernel_fc2_reduce, + combine_format, + tactic_autotune, + helper_expert_count, ) + # READY inputs belong to this launch. Never regenerate them or retain stale + # values across calls, even when the runner itself is cached. + runner._helper_ready_flags = hot_expert_weight_ready_flags + runner._helper_ready_generation = int(hot_expert_weight_ready_generation) inputs = [ activation, activation_sf, @@ -2626,12 +3192,17 @@ def cute_dsl_megamoe_nvfp4_blackwell( # materializes the multi-GiB scratch and MERGE-sweeps every # candidates. tactic=-1 guarantees deterministic default selection # with no tuning collectives, even inside a global autotune(). - runner( - inputs, - tactic=-1, - peer_offsets=peer_offsets, - shared_workspace=shared_workspace, - ) + try: + runner( + inputs, + tactic=-1, + peer_offsets=peer_offsets, + shared_workspace=shared_workspace, + ) + finally: + # Release the cached runner reference to this generation's terminal tensor. + # The next call supplies fresh READY inputs. + runner._helper_ready_flags = None return tuner = AutoTuner.get() # Opt-IN: in tuning mode the MERGE lockstep sweep runs (made safe by @@ -2670,7 +3241,7 @@ def cute_dsl_megamoe_nvfp4_blackwell( prof_shared_workspace = shared_workspace try: _, best_tactic = tuner.choose_one( - "trtllm::cute_dsl_megamoe_nvfp4_blackwell", + _megamoe_tuning_op_name(runner.helper_expert_count), [runner], runner.get_tuning_config(), inputs, @@ -2681,12 +3252,23 @@ def cute_dsl_megamoe_nvfp4_blackwell( runner._profiling_scratch = None runner._profiling_scratch_factory = None # Real run always uses the caller's staging buffer + peer_offsets. - runner( - inputs, - tactic=best_tactic, - peer_offsets=peer_offsets, - shared_workspace=shared_workspace, - ) + try: + runner( + inputs, + tactic=best_tactic, + peer_offsets=peer_offsets, + shared_workspace=shared_workspace, + ) + finally: + runner._helper_ready_flags = None + + # Eager callers may call ``run_cute_dsl_megamoe_nvfp4_blackwell`` directly + # and skip the custom-op dispatch; traced code must use the registered op. + cute_dsl_megamoe_nvfp4_blackwell = torch.library.custom_op( + "trtllm::cute_dsl_megamoe_nvfp4_blackwell", + mutates_args=("combine_output", "shared_workspace"), + device_types="cuda", + )(run_cute_dsl_megamoe_nvfp4_blackwell) @torch.library.register_fake("trtllm::cute_dsl_megamoe_nvfp4_blackwell") def _( @@ -2722,5 +3304,13 @@ def _( combine_format: str = "bf16", tactic_autotune: bool = False, num_tokens: int = -1, + # READY ABI tail. ``helper_expert_count`` (S) is codegen-time; + # ``hot_expert_weight_ready_flags`` is the local uint64[EP] terminal + # view published by the weight transport, ``..._generation`` the + # expected generation g in [1, 2**63). All three default to the + # S == 0 kernel without helper slots. + helper_expert_count: int = 0, + hot_expert_weight_ready_flags: Optional[torch.Tensor] = None, + hot_expert_weight_ready_generation: int = 0, ) -> None: return None diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 57063ebb6368..e657247910ae 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -216,6 +216,7 @@ Still on old path (standalone, with embedded communication): | `routing.py` | Routing methods (`TopKRouting`, etc.) | | `moe_comm_timeout_guard.py` | Device-side timeouts of communication kernels, relaxed during warmup — budget resolution and `MoECommTimeoutProxy` registration (see [Communication timeouts](#communication-timeouts)) | | `moe_load_balancer.py` | EPLB implementation | +| `per_layer_dynamic_eplb/` | Per-layer dynamic EPLB, created by `maybe_create_moe_load_balancer`: model-level `manager`, per-layer handle `layer`, live-weight `arena`, and HALO-Q/TMA `planner` | | `moe_op_backend.py` | Op backend registry for TRTLLMGen (flashinfer/trtllm ops) | ### Backends (`fused_moe/`) @@ -433,8 +434,10 @@ literal also selects, descend from the CUTLASS family and stay in | `KERNEL_INTERNALS.html` | Reference for the underlying DeepGEMM kernel layout | The exported CuteDSL kernel sources for `TrtllmCutedslMegaMoeNvfp4Impl` live under -`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, preserving the upstream -package structure and `KernelClass` entry points. The runner constructs +`tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/`, a copy of the upstream +cutedsl_megamoe sources maintained in-tree; update it by copying newer upstream +code in. `moe/kernels/megamoe_kernels.py` exposes the `KernelClass` entry +points. The runner constructs `ProblemDesc` and `ImplDesc` and dispatches internally; the public backend remains `MEGAMOE_CUTEDSL`. Rubin's smaller token buckets use the genphase kernel when supported, otherwise the generic kernel. During warmup, the @@ -851,6 +854,75 @@ set wrong has a test that catches it: - `FusedCommMoEScheduler` calls `backend.quantize_input(...)` for every chunk including zero-token chunks (so peer ranks can cross the in-kernel NVLink barrier). Each fused-comm backend therefore MUST make `quantize_input` tolerate `x.shape[0] == 0` and return its own empty tensor layout; the scheduler does NOT synthesize backend-specific empty tensors. - Dynamic EPLB requires backend and quantization-method support; backends gate wrapper-level constraints in `validate_configurable_moe`. `MegaMoEDeepGemm` routes to slot IDs and migrates the transformed DG weights its quantization method registers (`num_slots % ep_size == 0`). `TrtllmCutedslMegaMoeNvfp4Impl` declares `supports_eplb=True`; its quantization method registers the MegaMoE-format derived params and `fc1_norm_const` with the load balancer alongside the raw NVFP4 family, so per-slot migration stays byte-consistent. - `FUSED_COMM` backends use `ignore_allreduce=False` for EPLB statistic update because the fused kernel AllReduces routing stats internally. +- Per-layer dynamic EPLB uses `MoeConfig.load_balancer.mode=per_layer_dynamic` with + `num_slots` greater than the model expert count. Backend resolution selects + `MEGAMOE_CUTEDSL` when no concrete backend is requested and rejects an + incompatible explicit backend before model construction. +- Per-layer dynamic EPLB requires `CUDA_SCALE_LAUNCH_QUEUES=4x` before CUDA context + initialization. TensorRT-LLM applies the setting before `TorchLlmArgs` + validation and CUDA initialization/device probes, and propagates it to MPI, + direct, and Ray workers. Applications that initialize CUDA before TensorRT-LLM + must set it at process startup. +- HALO-Q and the in-switch TMA copy run as one kernel on the higher-priority + copy stream, one CTA per auxiliary SM. Each CTA raises a device route epoch + once its routes are written. The main stream may enqueue shared-expert work + and input quantization before its first routed-expert use, where a stream + memory wait on that epoch orders it after the routes. Each layer binds one + main stream, the executor's execution stream, on which warmup and serving + submit from their own threads; the path does not require a CUDA green + context or a steady-state host wait. +- HALO-Q waits for every EP rank inside the kernel and traps once a peer is + later than its spin bound. Each launch passes the bound, which + `SchedulerSpinBound` takes from the MoE communication timeout guard: the + warmup timeout during warmup and AutoTuner profiling, the native bound + otherwise. +- `helper_slot_mode=per_layer_resident` gives each layer its own helper-slot + set, and the copy keeps a per-layer table of the expert each helper slot + holds so that it skips experts a whole team already holds. Any other writer + of those helper rows would invalidate the table. +- HALO-Q and TMA follow the native-operator path under + `cpp/tensorrt_llm/kernels/moe/loadBalance/dynamicEplb` and + `cpp/tensorrt_llm/thop/moe/loadBalance`. The scheduler accepts the original + contiguous CUDA `int32` route tensor; do not restore route staging, tail + fills, or a host-generated copy plan. +- Rebalance resources are model-scoped. Worker shutdown and construction + rollback must release the scheduler, TMA bindings, shared-slot mappings, + and fabric allocations collectively. Forward execution must not retain + activation tensors after their consumer has been enqueued. +- Per-layer dynamic EPLB lives in `per_layer_dynamic_eplb/` (`manager`, + `layer`, `arena`, `planner`) and reuses the EPLB seams. + `maybe_create_moe_load_balancer` builds `PerLayerDynamicMoeLoadBalancer`, + which shares the `BaseMoeLoadBalancer` lifecycle with `MoeLoadBalancer`, and + `add_layer` returns a `PerLayerDynamicMoeLayerLoadBalancer` with + `produces_replica_plan = True`. In every chunk `_load_balancer_route` plans, + the backend reads `MoERunContext.replica_plan`, and the post-`run_moe` hook + releases the plan. The planner consumes logical routes, so the expert + statistics and the layer-wise calibrator record (or replay) them before the + route; standard EPLB records the slots after it. When the forward raises, + `MoeLoadBalancerIterContext` calls `abort_iter`, which releases the pending + plans. Plans run during warmup too, so there is no warmup switch; only the + HALO-Q spin bound changes, as described above. Model code passes independent + main-stream work such as shared experts through + `MoE.load_balancer_overlap(fn)`; the route runs it once, after submitting the + copy stream. +- The quant method owns the seven live kernel planes. Its `create_weights` + allocates them from its plane spec through the layer handle's + `allocate_arena` and binds them; its `cache_derived_state` checks them and + calls the handle's `bind_backend`, which builds the HALO-Q/TMA producer. The + backend only declares `supports_per_layer_dynamic_eplb`, uses + `compute_slot_count` (H + S, with S from the layer handle) for the planes, + symmetric workspace, and launch geometry while `expert_size_per_partition` + stays H, and consumes the replica plan, including the SMs its input + quantization leaves to the copy kernel. `bind_backend` reads the backend's + `ep_process_group` and `max_tokens_per_rank_capacity`. +- DeepSeek-V4 selects the fused shared FC12 kernel when the routed backend + declares per-layer dynamic support, the final projection quantization and + hardware (SM100 family) support that kernel, and the model runs eagerly + (`cuda_graph_config: null`). Per-layer dynamic EPLB is not a condition, but it + requires eager mode, so an EPLB-off baseline runs the same shared kernel only + with `cuda_graph_config: null`. Other runs keep the model's existing shared-MLP + backend. With per-layer dynamic EPLB, the active configuration supplies the + fused kernel an SM budget that leaves the copy kernel its auxiliary SMs. ## Canonical Examples @@ -972,7 +1044,7 @@ their real class name. - **Do NOT set a communication kernel's timeout directly or through environment variables at runtime** — A backend whose kernels wait for peers under a timeout registers a `MoECommTimeoutProxy` with the timeout guard when it is constructed and unregisters it before releasing what the proxy controls; only the guard decides the value (see [Communication timeouts](#communication-timeouts)). Writing environment variables at runtime races with NCCL threads that read them - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. - **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. -- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the 10-field tactic pattern (with legacy 8-field compatibility) (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. +- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py` for the canonical 11-field tactic pattern (with legacy 10- and 8-field compatibility), including token-back READY granularity; `_unpack_tactic` is the single source of truth for field order and compatibility defaults. The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. - **Use `distributed_tuning_strategy=DistributedTuningStrategy.MERGE` on a multi-rank `FUSED_COMM` backend's `TuningConfig`** — Every EP rank must converge on the same compiled tactic for every chunk, otherwise the in-kernel NVLink dispatch barrier deadlocks. `PARALLEL` can profile different tactics on different ranks and is unsafe for fused collectives. Reference: `Sm100MegaMoENvfp4Runner.get_tuning_config`. ### CUTLASS eager scratch reclamation diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 37601f07b41b..074e8a8f2a31 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -115,9 +115,16 @@ def create_moe_backend( # backend that cannot run under a load balancer already declines in # ``can_implement``. eplb_enabled = get_moe_load_balancer() is not None + load_balancer_config = getattr(model_config, "moe_load_balancer", None) + eplb_mode = (getattr(load_balancer_config, "mode", "standard") + if eplb_enabled else "disabled") if eplb_enabled and not moe_cls.capabilities.supports_eplb: raise ValueError( f"{moe_cls.__name__} does not support the MoE load balancer.") + if (eplb_mode == "per_layer_dynamic" + and not moe_cls.capabilities.supports_per_layer_dynamic_eplb): + raise ValueError( + f"{moe_cls.__name__} does not support per-layer dynamic EPLB.") if bias and not moe_cls.capabilities.supports_expert_bias: raise ValueError(f"bias not supported in {moe_cls.__name__}.") diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 4f8dd3ac5d9e..2b6aba68f06b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -15,9 +15,9 @@ """Contracts for declaring, selecting, and executing MoE implementations.""" import hashlib -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple import torch @@ -47,6 +47,9 @@ class MoEStaticCapability: # ``_supports_load_balancer()``, which TrtllmGenFusedMoEBase overrides to mean # "separated routing is used". supports_eplb: bool = False + # Whether the backend consumes a per-forward helper-slot plan whose live + # weights may be published concurrently with shared-expert computation. + supports_per_layer_dynamic_eplb: bool = False # Legacy gate: the ``assert moe_cls in [...]`` bias allow-list in # ``create_moe_backend``. Per-expert FC bias from the checkpoint, added # before the activation functor runs -- not an activation constant. @@ -280,6 +283,8 @@ class MoEDeployment: cluster_size: int = 1 # True only when an EPLB load balancer is registered. eplb_enabled: bool = False + # Unified EPLB mode: disabled, standard, or per_layer_dynamic. + eplb_mode: str = "disabled" # True only for routed-expert LoRA targets. moe_lora_enabled: bool = False # False when ``moe_disable_finalize_fusion`` is set or any LoRA is @@ -603,6 +608,7 @@ def to_dict(self) -> Dict[str, object]: "use_dp": self.deployment.use_dp, "num_slots": self.deployment.num_slots, "eplb_enabled": self.deployment.eplb_enabled, + "eplb_mode": self.deployment.eplb_mode, "moe_lora_enabled": self.deployment.moe_lora_enabled, "force_dynamic_quantization": self.deployment.force_dynamic_quantization, "oss_cutlass_moe": self.deployment.env.oss_cutlass_moe, @@ -676,6 +682,48 @@ class MoECommPlan: use_deep_ep_direct_metadata: bool = False +@dataclass(frozen=True) +class MoEReplicaPlan: + """Per-forward physical replica layout produced by the EPLB layer. + + The backend consumes these facts but does not own the planner, copy stream, + weight arena, or generation lifecycle. The private wait callback preserves + the late dependency: independent input preparation remains overlapped until + the fused kernel is about to read the physical routes. + """ + + resident_experts_per_rank: int + compute_slots_per_rank: int + ready_flags: torch.Tensor + ready_generation: int + reserved_sms: int + _wait_for_routes: Callable[[], None] = field(repr=False, compare=False) + + def __post_init__(self) -> None: + resident = int(self.resident_experts_per_rank) + compute = int(self.compute_slots_per_rank) + generation = int(self.ready_generation) + if resident <= 0 or compute <= resident: + raise ValueError( + "A replica plan requires compute slots M greater than resident " + f"experts H; got H={resident}, M={compute}." + ) + if not 1 <= generation < (1 << 63): + raise ValueError( + f"A replica plan READY generation must be in [1, 2**63); got {generation}." + ) + if int(self.reserved_sms) <= 0: + raise ValueError("A replica plan must reserve at least one auxiliary SM") + + @property + def helper_slots_per_rank(self) -> int: + return int(self.compute_slots_per_rank) - int(self.resident_experts_per_rank) + + def wait_for_routes(self) -> None: + """Order the first route consumer after the plan producer on device.""" + self._wait_for_routes() + + @dataclass(frozen=True) class MoERunContext: """Everything ``run_moe`` needs that the CALLER genuinely produces. @@ -706,6 +754,8 @@ class MoERunContext: all_rank_num_tokens: Optional[List[int]] = None # produced by the comm strategy, one object per forward comm_plan: Optional[MoECommPlan] = None + # produced by per-layer dynamic EPLB; None for standard EPLB and OFF. + replica_plan: Optional[MoEReplicaPlan] = None def require_comm_plan(impl: object, ctx: MoERunContext) -> MoECommPlan: diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 3c14d8a174b6..6d7754a74586 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -13,11 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import os import weakref from abc import abstractmethod from enum import Enum -from typing import Dict, List, Optional, Tuple, Union, final +from typing import Callable, Dict, Iterator, List, Optional, Tuple, Union, final import torch from torch import nn @@ -52,6 +53,7 @@ def _reject(reason: MoERejectReason, detail: str) -> MoEEligibility: get_model_extra_attrs, is_gated_activation, is_torch_compiling) from .routing import (BaseMoeRoutingMethod, RoutingMethodType, + get_cached_perfect_router_expert_ids, get_cached_perfect_router_logits, precompute_common_perfect_router_logits) @@ -232,6 +234,8 @@ class MoE(MoEExecutionContractMixin, MoEWeightOwnerMixin, # Set by callers that compose a ``FUSED_COMM`` backend's globally combined # routed output themselves; read through ``returns_global_routed_output``. accepts_global_routed_output: bool = False + # Set only inside ``load_balancer_overlap``; consumed by the next route. + _load_balancer_overlap_fn: Optional[Callable[[], None]] = None @property def returns_global_routed_output(self) -> bool: @@ -388,7 +392,7 @@ def __init__( # un-overridden default would be ``expert_size_per_partition = # num_experts`` (each rank allocating storage for every expert). self._init_dwdp_expert_layout() - self._init_perfect_router() + self._init_perfect_router(model_config) def _init_dwdp_expert_layout(self): """Override expert layout when DWDP is enabled. @@ -462,9 +466,17 @@ def _get_perfect_router_dtype(self) -> torch.dtype: return torch.float32 return self.dtype if self.dtype is not None else torch.float32 - def _init_perfect_router(self): - self._enable_perfect_router = os.environ.get("ENABLE_PERFECT_ROUTER", - "0") == "1" + def _init_perfect_router(self, model_config: ModelConfig): + # ENABLE_PERFECT_ROUTER=1 replaces the router logits before routing; + # ENABLE_PERFECT_ROUTER=expert_ids replaces the routed expert ids after + # routing. Both are for performance analysis only. + perfect_router = os.environ.get("ENABLE_PERFECT_ROUTER", "0") + self._enable_perfect_router = perfect_router == "1" + self._enable_perfect_router_expert_ids = perfect_router == "expert_ids" + # No MoE call routes more local tokens than this. It sizes the + # expert-id table, which is built on the first call, not here: the + # module may still be on the meta device. + self._perfect_router_max_num_tokens = model_config.max_num_tokens if not self._enable_perfect_router: return @@ -493,6 +505,28 @@ def _maybe_get_perfect_router_logits( dtype=router_logits.dtype, routing_method=self.routing_method) + def _maybe_get_perfect_router_expert_ids( + self, token_selected_experts: torch.Tensor) -> torch.Tensor: + """Return balanced expert ids under ENABLE_PERFECT_ROUTER=expert_ids. + + The router still runs, so its cost is kept; only its expert ids are + replaced, by a view of a cached table (see + ``get_cached_perfect_router_expert_ids``). The routing scales are left + as computed. Model outputs are wrong by design. + """ + if not self._enable_perfect_router_expert_ids: + return token_selected_experts + + num_tokens, experts_per_token = token_selected_experts.shape + return get_cached_perfect_router_expert_ids( + num_tokens=num_tokens, + num_experts=self.num_experts, + experts_per_token=experts_per_token, + moe_ep_size=self.ep_size, + ep_rank=self.ep_rank, + max_num_tokens=self._perfect_router_max_num_tokens, + device=token_selected_experts.device) + def _init_load_balancer( self, model_config: ModelConfig, @@ -514,8 +548,12 @@ def _init_load_balancer( moe_load_balancer = get_moe_load_balancer() moe_load_balancer_config = model_config.moe_load_balancer - # Calculate initial expert assignments - if moe_load_balancer_config: + # Standard EPLB checkpoints carry one expert id per physical slot. + # Per-layer dynamic EPLB keeps checkpoint/source storage on the resident + # axis H=E/EP; only the live kernel planes are widened to M=P/EP. + per_layer_dynamic_eplb = (getattr(moe_load_balancer_config, "mode", + None) == "per_layer_dynamic") + if moe_load_balancer_config and not per_layer_dynamic_eplb: init_expert_size_per_partition = moe_load_balancer_config.num_local_slots self.initial_global_assignments = [ (ep_rank * self.num_experts // self.ep_size + local_slot_id) % @@ -523,7 +561,7 @@ def _init_load_balancer( for local_slot_id in range(init_expert_size_per_partition) ] else: - # Sequential mapping: expert i → slot i; covers all experts regardless of divisibility + # Sequential resident mapping covers every logical expert exactly once. self.initial_global_assignments = list(range(self.num_experts)) # Setup load balancer if available @@ -543,23 +581,44 @@ def _init_load_balancer( f"by ep_size ({self.ep_size}) so each rank holds the same " f"number of slots.") top_k = self.routing_method.experts_per_token - self.expert_size_per_partition = moe_load_balancer_config.num_local_slots + if per_layer_dynamic_eplb: + validate_capacity = getattr(moe_load_balancer_config, + "validate_expert_capacity", None) + if callable(validate_capacity): + validate_capacity(self.num_experts) + if self.num_experts % self.ep_size != 0: + raise ValueError( + f"{type(self).__name__}: per-layer dynamic EPLB requires " + f"num_experts ({self.num_experts}) divisible by ep_size " + f"({self.ep_size}) so resident H is rank-identical.") + self.expert_size_per_partition = self.num_experts // self.ep_size + else: + self.expert_size_per_partition = moe_load_balancer_config.num_local_slots - # Add this layer to the load balancer + # Add this layer to the load balancer. slot_count_per_rank is the + # compute capacity M for per-layer dynamic mode and the ordinary local + # slot count for standard mode. aux_stream = getattr(self, '_get_load_balancer_aux_stream', lambda: None)() + slot_count_per_rank = (moe_load_balancer_config.num_local_slots + if per_layer_dynamic_eplb else + self.expert_size_per_partition) self.layer_load_balancer = moe_load_balancer.add_layer( self.num_experts, top_k, - self.expert_size_per_partition, + slot_count_per_rank, aux_stream=aux_stream) self.repeat_count = self.layer_load_balancer.get_repeat_count() # Handle initial global assignments - loaded_initial_global_assignments = ( - moe_load_balancer_config.get_layer_initial_global_assignments( - self.layer_idx)) + get_initial_assignments = getattr( + moe_load_balancer_config, + "get_layer_initial_global_assignments", + None, + ) + loaded_initial_global_assignments = (get_initial_assignments( + self.layer_idx) if callable(get_initial_assignments) else None) self.num_slots = moe_load_balancer_config.num_slots if loaded_initial_global_assignments is not None: @@ -581,7 +640,8 @@ def _init_load_balancer( f"initial_global_assignments (layer {self.layer_idx}) = {self.initial_global_assignments}" ) - # Slot boundaries for EPLB (uniform: all ranks hold same num_local_slots) + # Source/checkpoint boundaries are resident H in per-layer dynamic + # mode. The layer manager separately owns compute slots M. self.slot_start = self.ep_rank * self.expert_size_per_partition self.slot_end = self.slot_start + self.expert_size_per_partition self.initial_local_expert_ids = self.initial_global_assignments[ @@ -673,18 +733,52 @@ def _load_balancer_update_statistic(self, is_last_stage=is_last_call, allreduce=self.allreduce) + @contextlib.contextmanager + def load_balancer_overlap(self, + overlap_fn: Callable[[], None]) -> Iterator[None]: + """Run ``overlap_fn`` inside the next load-balancer route call. + + A per-layer dynamic EPLB layer calls it after the copy stream has been + submitted and before the routes are returned, so independent + main-stream work such as shared experts overlaps planning and the + helper-weight copy. Only the first route call consumes it, and it is + cleared on every exit; if no route call consumes it, the caller still + owns that work. Every EP rank must install the same function. + """ + if not callable(overlap_fn): + raise TypeError("overlap_fn must be callable") + self._load_balancer_overlap_fn = overlap_fn + try: + yield + finally: + self._load_balancer_overlap_fn = None + def _load_balancer_route(self, token_selected_experts: torch.Tensor, use_dp: bool) -> torch.Tensor: - """Route tokens using load balancer.""" - if self.layer_load_balancer: + """Route tokens and enqueue a per-layer dynamic replica plan when enabled.""" + if not self.layer_load_balancer: + return token_selected_experts + + overlap_fn = self._load_balancer_overlap_fn + if overlap_fn is None: return self.layer_load_balancer.route(token_selected_experts, use_dp) - else: - return token_selected_experts + self._load_balancer_overlap_fn = None + return self.layer_load_balancer.route(token_selected_experts, + use_dp, + overlap_fn=overlap_fn) + + def _load_balancer_uses_replica_plan(self) -> bool: + """Whether routing is produced by the per-layer dynamic device planner.""" + return bool(self.layer_load_balancer + and self.layer_load_balancer.produces_replica_plan) def _load_balancer_start_set_cpu_stage(self, is_last_call: bool): """Start CPU stage in load balancer.""" - if self._using_dynamic_load_balancer() and is_last_call: + # A per-layer dynamic EPLB layer plans every chunk and releases that + # chunk's replica plan here, after the chunk's MoE kernel. + if self._load_balancer_uses_replica_plan() or ( + self._using_dynamic_load_balancer() and is_last_call): self.layer_load_balancer.start_set_cpu_stage() def _load_balancer_done_set_cpu_stage(self, is_last_call: bool): diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index 638892f0f871..3ca2c6889890 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -68,12 +68,12 @@ from __future__ import annotations +import functools import os import socket -import time import weakref from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.distributed as dist @@ -240,7 +240,7 @@ def is_megamoe_cute_dsl_runtime_available() -> Tuple[bool, Optional[str]]: return False, reason try: - from ....moe.kernels.cutedsl_megamoe import ( # noqa: F401 + from ....moe.kernels.megamoe_kernels import ( # noqa: F401 BlackwellInferenceMegaMoE, RubinInferenceGenphaseMegaMoE, RubinInferenceLocalMegaMoE, @@ -286,6 +286,17 @@ def _as_fp8_sf(t: torch.Tensor) -> torch.Tensor: return t if t.dtype == torch.float8_e4m3fn else t.view(torch.float8_e4m3fn) +@functools.cache +def _fp4_quantize_has_sm_budget() -> bool: + if hasattr(torch.ops.trtllm.fp4_quantize, "sm_budget"): + return True + logger.warning( + "fp4_quantize.sm_budget is unavailable; per-layer dynamic EPLB " + "quantization will use the selected backend without SM control" + ) + return False + + # --------------------------------------------------------------------------- # Weight view passed to ``run_moe`` # --------------------------------------------------------------------------- @@ -356,7 +367,7 @@ class _MegaMoeBuffers: topk_weights: torch.Tensor # (max_T, top_k) float32 combine_output: torch.Tensor # (max_T, top_k, hidden) bf16 shared_workspace: torch.Tensor # (shared_ws_bytes,) uint8 - peer_offsets: List[int] # length == world_size; [0] for single-rank + peer_offsets: Tuple[int, ...] # length == world_size; (0,) for single-rank topk_idx_local: torch.Tensor # (max_T, top_k), int32 on Rubin, otherwise int64 @@ -404,7 +415,10 @@ class TrtllmCutedslMegaMoeNvfp4Impl(MoEImplBase): # Static and dynamic EPLB both work: see ``_supports_load_balancer`` # below for why the MegaMoE-format derived parameters migrate # atomically. - capabilities=MoEStaticCapability(supports_eplb=True), + capabilities=MoEStaticCapability( + supports_eplb=True, + supports_per_layer_dynamic_eplb=True, + ), doc="MegaMoE CuteDSL fused NVFP4 kernels: NVFP4 weights and activations, SM100/103/107.", ) @@ -598,10 +612,7 @@ def __init__( ) self._topk_idx_dtype = torch.int32 if get_sm_version() == 107 else torch.int64 - # Cross-rank combine wire format is selected with - # MEGAMOE_COMBINE_FORMAT (default bf16). It MUST be rank-identical: - # it sizes the symmetric workspace and picks the compiled kernel, so - # divergence desyncs the rendezvous / NVLink barrier. + # The wire format must match on all EP ranks: it sizes the symmetric workspace. combine_format = os.environ.get("MEGAMOE_COMBINE_FORMAT", "bf16") if combine_format not in self._SUPPORTED_COMBINE_FORMATS: raise ValueError( @@ -639,6 +650,8 @@ def __init__( # Per-bucket symmetric providers (bucket -> MegaMoeSymmMemProvider); # filled in ``create_weights`` on the multi-rank EP path. self._symm_providers = {} + # Launch buffers per (bucket, top_k, hidden, device, dtype, inference mode). + self._launch_buffers: Dict[Tuple, _MegaMoeBuffers] = {} self._genphase_direct_input_buckets: set[int] = set() # Lockstep per-chunk cross-rank max for the NEXT run_moe (set via # ``set_adaptive_launch_tokens``); ``None`` -> full bucket. @@ -678,6 +691,8 @@ def __init__( # deferred factory in ``_launch_megamoe_kernel``. self._profiling_scratch_provider_ref = None self._profiling_scratch_kwargs = None + # Cute views of this module's weights, reused across launches. + self._cute_weight_cache: Dict[tuple, Any] = {} self._weights_loaded = False self._weights_created = False self.quant_method = None @@ -692,6 +707,37 @@ def __init__( if not model_config.skip_create_weights_in_init: self.create_weights() + # ------------------------------------------------------------------ + # Compute axis + # ------------------------------------------------------------------ + @property + def compute_slot_count(self) -> int: + """Return the kernel-only slot axis ``M = H + S``. + + ``expert_size_per_partition`` remains the resident/checkpoint axis H. + The per-layer dynamic EPLB layer handle supplies the S helper slots; + without one, M is H. Only the seven live kernel planes, symmetric + workspace, and launch geometry use this wider axis. + """ + resident_experts = int(self.expert_size_per_partition) + layer_load_balancer = self.layer_load_balancer + if layer_load_balancer is None or not layer_load_balancer.produces_replica_plan: + return resident_experts + return resident_experts + int(layer_load_balancer.helper_slots_per_rank) + + # ------------------------------------------------------------------ + # Backend facts the per-layer dynamic EPLB layer handle reads + # ------------------------------------------------------------------ + @property + def ep_process_group(self): + """The EP ProcessGroup resolved at construction, or ``None`` on one rank.""" + return self._ep_pg + + @property + def max_tokens_per_rank_capacity(self) -> int: + """The largest per-rank token count of any launch bucket.""" + return int(self._maxt_buckets[-1]) + # ------------------------------------------------------------------ # Topology # ------------------------------------------------------------------ @@ -1011,6 +1057,7 @@ def _alloc_symm_provider(self): top_k = self.routing_method.experts_per_token full_T = self._maxt_buckets[-1] self._profiling_scratch_provider_ref = None + self._launch_buffers.clear() for max_T in self._maxt_buckets: ( shared_workspace_bytes, @@ -1018,7 +1065,7 @@ def _alloc_symm_provider(self): ) = query_megamoe_kernel_properties( world_size=self.ep_size, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=self.compute_slot_count, hidden_size=self.hidden_size, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1156,6 +1203,7 @@ def quantize_input( x: Union[torch.Tensor, "Fp4QuantizedTensor"], *, post_quant_comm: bool = False, + reserved_sms: int = 0, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: """BF16 -> NVFP4 packed activation + plain K-major FP8 SF. @@ -1173,6 +1221,10 @@ def quantize_input( ``fc31_input_scale`` was calibrated on the scaled activation. The kernel itself only ever sees the quantized input. + ``reserved_sms`` is the number of SMs left to the concurrent + HALO-Q/TMA kernel of a replica plan; zero uses the default + quantization overload. + Empty input (``x.shape[0] == 0``) short-circuits to empty NVFP4 + empty SF without launching a quantization kernel, so the ``FusedCommMoEScheduler`` can call ``quantize_input`` uniformly @@ -1200,13 +1252,21 @@ def quantize_input( # admits bf16 activations only, so ``fc31_act_scale``, created at # ``module.dtype``, is bf16 as well. x_bf16 = x_bf16 * fc31_act_scale - x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( - x_bf16, - self.fc31_input_scale, - 16, # scaling_vector_size == Nvfp4BlockSize - False, # sf_use_ue8m0 - False, # is_sf_swizzled - MegaMoE expects plain K-major - ) + # Leave a replica plan's SMs to its concurrent HALO-Q/TMA kernel. The + # default overload remains unchanged without a plan. + if reserved_sms > 0 and _fp4_quantize_has_sm_budget(): + # Positional: a keyword argument costs the op dispatch about a microsecond. + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize.sm_budget( + x_bf16, self.fc31_input_scale, 16, False, False, reserved_sms + ) + else: + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize( + x_bf16, + self.fc31_input_scale, + 16, # scaling_vector_size == Nvfp4BlockSize + False, # sf_use_ue8m0 + False, # is_sf_swizzled - MegaMoE expects plain K-major + ) # ``fp4_quantize(is_sf_swizzled=False)`` returns LINEAR layout # ``(rows, ceil(hidden/16))`` with no column pad. The kernel TMA # load needs ``pad_up(ceil_div(hidden, 16), 4)`` bytes per row @@ -1221,6 +1281,42 @@ def quantize_input( padded_sf[:, :raw_cols] = x_sf_raw return x_fp4, padded_sf + def _replica_plan_launch_args(self, replica_plan): + """Validate and unpack the caller-owned per-forward replica plan.""" + if replica_plan is None: + return None + resident_experts = int(replica_plan.resident_experts_per_rank) + compute_slots = int(replica_plan.compute_slots_per_rank) + helper_slots = compute_slots - resident_experts + expected_resident = int(self.expert_size_per_partition) + expected_compute = self.compute_slot_count + if resident_experts != expected_resident or compute_slots != expected_compute: + raise RuntimeError( + "MegaMoECuteDsl replica-plan geometry does not match its allocated " + f"weights: plan H/M={resident_experts}/{compute_slots}, " + f"backend H/M={expected_resident}/{expected_compute}." + ) + generation = int(replica_plan.ready_generation) + reserved_sms = int(replica_plan.reserved_sms) + if helper_slots <= 0 or not 1 <= generation < (1 << 63) or reserved_sms <= 0: + raise RuntimeError( + "MegaMoECuteDsl received an invalid replica plan: " + f"S={helper_slots}, generation={generation}, " + f"reserved_sms={reserved_sms}." + ) + return ( + compute_slots, + helper_slots, + replica_plan.ready_flags, + generation, + reserved_sms, + ) + + def _wait_rebalance_routes(self, replica_plan=None) -> None: + """Order route consumers after HALO-Q while leaving setup overlapped.""" + if replica_plan is not None: + replica_plan.wait_for_routes() + def run_moe( self, ctx: MoERunContext, @@ -1239,6 +1335,7 @@ def run_moe( token_selected_experts = ctx.token_selected_experts token_final_scales = ctx.token_final_scales x_sf = ctx.x_sf + replica_plan = getattr(ctx, "replica_plan", None) output_dtype = ctx.output_dtype if output_dtype is None: output_dtype = self.dtype or torch.bfloat16 @@ -1271,13 +1368,16 @@ def run_moe( top_k = int(token_selected_experts.shape[-1]) device = x.device - topk_idx = token_selected_experts.to(self._topk_idx_dtype).contiguous() - topk_weights_f32 = token_final_scales.to(torch.float32).contiguous() + topk_weights_f32 = token_final_scales + if topk_weights_f32.dtype != torch.float32 or not topk_weights_f32.is_contiguous(): + topk_weights_f32 = topk_weights_f32.to(torch.float32).contiguous() + # Routes stay as produced: the staging copy into the kernel's index + # buffer widens int32 to its dtype, and it is the first route reader. return self._run_moe( x=x, x_sf=x_sf, - topk_idx=topk_idx, + topk_idx=token_selected_experts, topk_weights=topk_weights_f32, weight_view=weight_view, num_tokens=num_tokens, @@ -1285,6 +1385,7 @@ def run_moe( hidden=hidden, device=device, output_dtype=output_dtype, + replica_plan=replica_plan, ) def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype, max_T: int): @@ -1355,7 +1456,7 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype shared_bytes, supports_direct_inputs = query_megamoe_kernel_properties( world_size=1, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=self.compute_slot_count, hidden_size=hidden, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1401,7 +1502,32 @@ def _acquire_buffers( in plain CUDA memory. ``launch_max_T`` selects both the local staging set and (multi-rank) the symmetric provider, which MUST share the same ``max_T`` as the kernel's compile-time ``max_tokens_per_rank``. + Each bundle is built once; views made under inference mode cannot be + updated in place outside it, so the mode is part of the key. """ + key = ( + launch_max_T, + top_k, + hidden, + device, + output_dtype, + torch.is_inference_mode_enabled(), + ) + bufs = self._launch_buffers.get(key) + if bufs is None: + bufs = self._build_buffers( + top_k=top_k, + hidden=hidden, + device=device, + output_dtype=output_dtype, + launch_max_T=launch_max_T, + ) + self._launch_buffers[key] = bufs + return bufs + + def _build_buffers( + self, *, top_k: int, hidden: int, device, output_dtype, launch_max_T: int + ) -> _MegaMoeBuffers: staging = self._ensure_local_staging( top_k=top_k, hidden=hidden, device=device, output_dtype=output_dtype, max_T=launch_max_T ) @@ -1412,7 +1538,7 @@ def _acquire_buffers( topk_weights=staging["topk_weights"], combine_output=staging["combine_output"], shared_workspace=staging["shared_workspace"], - peer_offsets=[0], + peer_offsets=(0,), topk_idx_local=staging["topk_idx"], ) # Multi-rank: the provider must have been rendezvous'd at @@ -1461,6 +1587,7 @@ def _stage_inputs( num_tokens: int, top_k: int, stage_activation: bool, + replica_plan=None, ) -> None: """Copy live rows of the user-domain inputs into the kernel's pre-allocated buffers and refresh the padded tail. @@ -1490,13 +1617,15 @@ def _stage_inputs( if last_T is not None and last_T > num_tokens: bufs.topk_idx_local[num_tokens:last_T].fill_(-1) if num_tokens > 0: - bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) if stage_activation: bufs.activation[:num_tokens].copy_(x.view(torch.uint8), non_blocking=True) bufs.activation_sf[:num_tokens].copy_(x_sf.view(torch.uint8), non_blocking=True) bufs.topk_weights[:num_tokens, :top_k].copy_(topk_weights, non_blocking=True) if num_tokens < max_T: bufs.topk_weights[num_tokens:max_T, :top_k].zero_() + if num_tokens > 0: + self._wait_rebalance_routes(replica_plan) + bufs.topk_idx_local[:num_tokens].copy_(topk_idx, non_blocking=True) self._last_staged_T[max_T] = num_tokens def _launch_megamoe_kernel( @@ -1517,6 +1646,8 @@ def _launch_megamoe_kernel( num_tokens: int, output_dtype: torch.dtype, launch_max_T: int, + replica_plan=None, + staged_inputs: bool = False, ) -> torch.Tensor: """Launch the fused MegaMoE CuteDSL kernel and return its reduced output. @@ -1525,13 +1656,18 @@ def _launch_megamoe_kernel( upstream by :meth:`_acquire_buffers`). NVFP4 / FP8-SF dtype views happen through the module-level :func:`_as_nvfp4` / :func:`_as_fp8_sf` helpers (the kernel rejects raw uint8 byte - tensors). + tensors). ``staged_inputs`` states that activation, activation_sf, + topk_idx and topk_weights are this bucket's launch buffers, whose cute + views the op may then reuse. """ # Hand the symmetric profiling scratch (or a deferred factory) to the # op for tuning-mode profiling launches; cleared in the ``finally`` so # a later same-process call for a different shape never sees a stale # scratch (ownership: see ``_profiling_scratch_provider_ref``). from tensorrt_llm._torch.moe.custom_ops.cute_dsl_megamoe_custom_op import ( + run_cute_dsl_megamoe_nvfp4_blackwell, + set_active_megamoe_cute_buffer_cache, + set_active_megamoe_cute_weight_cache, set_active_megamoe_profiling_scratch, set_active_megamoe_profiling_scratch_factory, ) @@ -1574,8 +1710,39 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): else None ) set_active_megamoe_profiling_scratch_factory(scratch_factory) + # READY generation comes only from the caller-owned replica plan. + # Helper-free launches (OFF and internal bucket priming) keep the + # M-shaped allocation but compile out the READY gate with S=0. + helper_ready_flags = None + helper_ready_generation = 0 + compute_slots = self.compute_slot_count + plan_args = self._replica_plan_launch_args(replica_plan) + if plan_args is None: + _S_active = 0 + else: + ( + compute_slots, + _S_active, + helper_ready_flags, + helper_ready_generation, + _reserved_sms, + ) = plan_args + # Quantization, independent staging, output clears, and host setup + # are already enqueued. The router inside this op first reads routes. + self._wait_rebalance_routes(replica_plan) + # Eager launches call the op's implementation directly: the custom-op + # dispatch costs tens of microseconds of host time per layer. + megamoe_op = ( + torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell + if torch.compiler.is_compiling() + else run_cute_dsl_megamoe_nvfp4_blackwell + ) + set_active_megamoe_cute_weight_cache(self._cute_weight_cache) + # The provider owns its symmetric regions, so it owns their cute views. + provider = self._symm_providers.get(launch_max_T) if staged_inputs else None + set_active_megamoe_cute_buffer_cache(provider.cute_views if provider is not None else None) try: - torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( + megamoe_op( activation=_as_nvfp4(activation), activation_sf=_as_fp8_sf(activation_sf), topk_idx=topk_idx, @@ -1592,7 +1759,7 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): world_size=world_size, local_rank=local_rank, num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), + num_experts_per_rank=compute_slots, hidden_size=hidden, intermediate_size_per_partition=int(self.intermediate_size_per_partition), expand_intermediate_size_per_partition=int( @@ -1615,12 +1782,18 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): combine_format=self.combine_format, tactic_autotune=bool(self.tactic_autotune), num_tokens=num_tokens, + # Helper-free launches omit the READY gate. + helper_expert_count=_S_active, + hot_expert_weight_ready_flags=helper_ready_flags, + hot_expert_weight_ready_generation=int(helper_ready_generation), ) finally: # Always clear: a raising op call must not leave the factory # global alive (its closure strongly references this module). set_active_megamoe_profiling_scratch(None) set_active_megamoe_profiling_scratch_factory(None) + set_active_megamoe_cute_weight_cache(None) + set_active_megamoe_cute_buffer_cache(None) if num_tokens == 0: return torch.empty((0, hidden), dtype=output_dtype, device=combine_output.device) # The kernel already collapsed the top-k axis into the (T, 1, hidden) @@ -1651,7 +1824,7 @@ def _prime_ladder( self.ep_size, hidden, int(self.intermediate_size_per_partition), - int(self.expert_size_per_partition), + self.compute_slot_count, top_k, self.combine_format, self.act_clamp, @@ -1663,8 +1836,35 @@ def _prime_ladder( ) if key in _MEGAMOE_PRIMED_LADDERS: return - t0 = time.perf_counter() primed = [] + self._prime_ladder_buckets( + live_T=live_T, + top_k=top_k, + hidden=hidden, + device=device, + output_dtype=output_dtype, + weight_view=weight_view, + primed=primed, + ) + _MEGAMOE_PRIMED_LADDERS.add(key) + if self.ep_rank == 0: + logger.debug( + "[MegaMoECuteDsl] adaptive bucket priming complete: count=%s", + len(primed), + ) + + def _prime_ladder_buckets( + self, + *, + live_T, + top_k, + hidden, + device, + output_dtype, + weight_view, + primed, + ): + """Prime each max-token bucket.""" for max_T in self._maxt_buckets: if max_T == live_T: continue @@ -1685,6 +1885,7 @@ def _prime_ladder( top_k=top_k, stage_activation=False, ) + # Bucket priming has no replica plan and therefore compiles the S=0 path. self._launch_megamoe_kernel( activation=bufs.activation, activation_sf=bufs.activation_sf, @@ -1703,12 +1904,6 @@ def _prime_ladder( launch_max_T=max_T, ) primed.append(max_T) - _MEGAMOE_PRIMED_LADDERS.add(key) - if self.ep_rank == 0: - logger.info( - f"[MegaMoECuteDsl] primed adaptive buckets {primed} (live bucket " - f"{live_T}) in {time.perf_counter() - t0:.1f}s" - ) def _run_moe( self, @@ -1723,6 +1918,7 @@ def _run_moe( hidden: int, device, output_dtype: torch.dtype, + replica_plan=None, ) -> torch.Tensor: """Unified MegaMoE CuteDSL forward: acquire -> stage -> launch. @@ -1812,6 +2008,7 @@ def _run_moe( num_tokens=num_tokens, top_k=top_k, stage_activation=not use_direct_inputs, + replica_plan=replica_plan, ) return self._launch_megamoe_kernel( activation=x if use_direct_inputs else bufs.activation, @@ -1831,6 +2028,9 @@ def _run_moe( num_tokens=num_tokens, output_dtype=output_dtype, launch_max_T=launch_max_T, + replica_plan=replica_plan, + # Direct routing inputs imply direct activation inputs. + staged_inputs=not use_direct_inputs, ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index 37f50886ee56..224153392bc9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -188,28 +188,34 @@ def release_symm_buffer_cache() -> None: """ if not _MEGA_MOE_SYMM_BUFFER_CACHE: return - # Detach before freeing so a raising free cannot leave a half-destroyed - # buffer in the cache for a later lookup to trip over. - entries = list(_MEGA_MOE_SYMM_BUFFER_CACHE.values()) + # Detach before freeing so successful entries cannot be reused while + # teardown is in progress. Failed entries are restored below so the worker + # can retry them before destroying their EP process group. + entries = list(_MEGA_MOE_SYMM_BUFFER_CACHE.items()) _MEGA_MOE_SYMM_BUFFER_CACHE.clear() - # Free one at a time and keep going on failure: the entries are already - # out of the cache, so any buffer skipped by a raising predecessor would - # be unreclaimable for the rest of the process lifetime. released = 0 total_bytes = 0 - for buffered, _ in entries: + first_error = None + for key, (buffered, process_group) in entries: try: total_bytes += _free_symm_buffer(buffered) released += 1 - except Exception as e: + except Exception as error: + _MEGA_MOE_SYMM_BUFFER_CACHE[key] = (buffered, process_group) + if first_error is None: + first_error = error logger.error( f"[MegaMoE] failed to release a DG SymmBuffer during executor " - f"teardown, continuing with the remaining buffers: {e}" + f"teardown, retaining it for retry: {error}" ) logger.info( f"[MegaMoE] released {released}/{len(entries)} DG SymmBuffer(s): " f"{total_bytes / 2**30:.2f} GiB" ) + if first_error is not None: + raise RuntimeError( + f"failed to release {len(_MEGA_MOE_SYMM_BUFFER_CACHE)} MegaMoE DG SymmBuffer(s)" + ) from first_error # ---- Fused MXFP8 per-token quant backends -------------------------------- diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py index 21943f2f06a4..4b3ac2c3b189 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py @@ -4,6 +4,7 @@ import gc import os import threading +from abc import ABC, abstractmethod from contextlib import AbstractContextManager, nullcontext from multiprocessing import resource_tracker, shared_memory from typing import Callable, Dict, List, Optional, Tuple @@ -303,6 +304,9 @@ class SingleLayerMoeLoadBalancer: This class wraps the C++ implementation of SingleLayerMoeLoadBalancer. """ + # Standard EPLB routes to slot IDs and has no per-invocation replica plan. + produces_replica_plan = False + def __init__( self, single_layer_load_balancer_impl: _tbr.SingleLayerMoeLoadBalancer, @@ -761,11 +765,117 @@ def py_post_shutdown_cleanup(self): self.host_tensor_sharer.post_shutdown_cleanup() -# Global variable to store the current active MoeLoadBalancer instance +# Global variable to store the current active MoE load balancer instance _current_moe_load_balancer = threading.local() -class MoeLoadBalancer: +class BaseMoeLoadBalancer(ABC): + """Model-level lifecycle shared by every MoE load-balancer mode. + + Model construction enters the balancer as a context manager so that MoE + layers register through ``get_moe_load_balancer``. The engine drives each + forward through ``MoeLoadBalancerIterContext`` (``set_iter_info``, + ``start_iter``, and ``end_iter`` or, when the forward raises, + ``abort_iter``), and workers release resources with a collective + ``shutdown``. Subclasses initialize ``_previous_balancer``, + ``enable_statistic`` and ``enable_update_weights``. + """ + + @abstractmethod + def is_static_routing(self) -> bool: + """Whether the placement stays fixed, so no iteration hooks run.""" + + def is_dynamic_routing(self) -> bool: + return not self.is_static_routing() + + @abstractmethod + def set_repeated_for_next_layer(self, repeated_count: int) -> None: + """Set how many times the next added layer is invoked per forward.""" + + @abstractmethod + def add_layer(self, + expert_count: int, + top_k: int, + slot_count_per_rank: int, + aux_stream: Optional[torch.cuda.Stream] = None): + """Register one MoE layer and return its layer load balancer.""" + + @abstractmethod + def register_weight_slots_after_to_cuda(self) -> None: + """Register weight slots after the model has been moved to CUDA.""" + + @abstractmethod + def finalize_model(self) -> None: + """Finish model-level setup after every layer has been added.""" + + @abstractmethod + def set_warm_up_iter_count(self, iter_count: int) -> None: + """Set the number of warm-up iterations.""" + + @abstractmethod + def start_iter(self) -> None: + """Start one forward iteration.""" + + @abstractmethod + def end_iter(self) -> None: + """End the current forward iteration.""" + + def abort_iter(self, error: BaseException) -> None: + """End an iteration whose forward raised ``error``; defaults to ``end_iter``.""" + self.end_iter() + + @abstractmethod + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively release resources; return False if any rank is unsafe.""" + + def set_iter_info(self, enable_statistic: Optional[bool], + enable_update_weights: Optional[bool]): + if enable_statistic is not None: + self.enable_statistic = enable_statistic + if enable_update_weights is not None: + self.enable_update_weights = enable_update_weights + + def __enter__(self): + """ + Enter the context manager. + + Returns: + The load balancer instance + """ + # Save the previous load balancer (if any) + try: + self._previous_balancer = getattr(_current_moe_load_balancer, + 'instance', None) + except AttributeError: + self._previous_balancer = None + + # Set the current instance as the active load balancer + _current_moe_load_balancer.instance = self + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Exit the context manager. + + Args: + exc_type: The exception type + exc_val: The exception value + exc_tb: The exception traceback + + Returns: + False to not suppress exceptions + """ + # Restore the previous load balancer + if self._previous_balancer is not None: + _current_moe_load_balancer.instance = self._previous_balancer + else: + delattr(_current_moe_load_balancer, 'instance') + + # Do not suppress exceptions + return False + + +class MoeLoadBalancer(BaseMoeLoadBalancer): """ A class representing a Mixture of Experts (MoE) load balancer. This class can be used as a context manager to manage the lifecycle of a MoeLoadBalancer. @@ -806,17 +916,18 @@ def __init__(self, self.next_layer_repeated_count = None + # Without per-layer dynamic EPLB, executor teardown is rank-local and this + # balancer is released by its own finalizer, which unlinks its shared memory. + shutdown_on_finalize = True + def __del__(self): - if not self.is_shutdown: + if not getattr(self, "is_shutdown", True): self.shutdown() def is_static_routing(self): # if we don't update, then it is statistic routing. return self.layer_updates_per_iter == 0 - def is_dynamic_routing(self): - return not self.is_static_routing() - def _setup_mpi_comm(self): global_mpi_comm = tensorrt_llm.mpi_comm() shared_mpi_comm = global_mpi_comm.Split_type( @@ -908,13 +1019,6 @@ def set_warm_up_iter_count(self, iter_count: int): """ self.load_balancer_impl.set_warm_up_iter_count(iter_count) - def set_iter_info(self, enable_statistic: Optional[bool], - enable_update_weights: Optional[bool]): - if enable_statistic is not None: - self.enable_statistic = enable_statistic - if enable_update_weights is not None: - self.enable_update_weights = enable_update_weights - def reconfigure_mask_only(self, dead_ranks: list[int]) -> None: """ Reconfigure EPLB routing so slots on dead EP ranks are unreachable. @@ -949,10 +1053,13 @@ def end_iter(self): self.in_iter = False self.iter_id += 1 - def shutdown(self): - """ - Shutdown the load balancer and release resources. - """ + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively release resources after an all-rank safety preflight.""" + if self.is_shutdown: + return True + local_safe = bool(local_safe and not self.in_iter) + if not all(self.shared_mpi_comm.allgather(local_safe)): + return False for single_layer_load_balancer in self.single_layer_load_balancers: single_layer_load_balancer.py_pre_shutdown_cleanup() self.load_balancer_impl.shutdown() @@ -962,6 +1069,7 @@ def shutdown(self): single_layer_load_balancer.py_post_shutdown_cleanup() self.shared_mpi_comm.barrier() self.is_shutdown = True + return True def __repr__(self): """ @@ -972,45 +1080,6 @@ def __repr__(self): """ return f"MoeLoadBalancer(ep_rank={self.ep_rank}, ep_size={self.ep_size}, layer_updates_per_iter={self.layer_updates_per_iter})" - def __enter__(self): - """ - Enter the context manager. - - Returns: - The MoeLoadBalancer instance - """ - # Save the previous load balancer (if any) - try: - self._previous_balancer = getattr(_current_moe_load_balancer, - 'instance', None) - except AttributeError: - self._previous_balancer = None - - # Set the current instance as the active load balancer - _current_moe_load_balancer.instance = self - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """ - Exit the context manager. - - Args: - exc_type: The exception type - exc_val: The exception value - exc_tb: The exception traceback - - Returns: - False to not suppress exceptions - """ - # Restore the previous load balancer - if self._previous_balancer is not None: - _current_moe_load_balancer.instance = self._previous_balancer - else: - delattr(_current_moe_load_balancer, 'instance') - - # Do not suppress exceptions - return False - moe_model_arch_list = [ 'DeepseekV3ForCausalLM', @@ -1041,38 +1110,72 @@ def __exit__(self, exc_type, exc_val, exc_tb): def maybe_create_moe_load_balancer( model_config, mapping: Optional[Mapping] -) -> AbstractContextManager[MoeLoadBalancer | None]: +) -> AbstractContextManager[BaseMoeLoadBalancer | None]: ep_rank = model_config.mapping.moe_ep_rank ep_size = model_config.mapping.moe_ep_size model_arch = model_config.pretrained_config.architectures[0] using_ep = mapping and mapping.moe_ep_size > 1 in_supported_model_arch = model_arch in moe_model_arch_list using_smart_router = mapping and mapping.moe_cluster_size > 1 - moe_load_balancer: AbstractContextManager[MoeLoadBalancer + moe_load_balancer: AbstractContextManager[BaseMoeLoadBalancer | None] = nullcontext() - if in_supported_model_arch and using_ep and not using_smart_router and model_config.moe_load_balancer is not None: - model_config.moe_load_balancer.setup(ep_rank=ep_rank, ep_size=ep_size) - if model_config.moe_load_balancer.layer_updates_per_iter > 0: - # TODO: remove this when supported. - # cpu_arch = platform.machine().lower() - # assert cpu_arch == 'aarch64', "online load balancer only support aarch64, e.g. GB200 now, x86 coming soon." - pass - - moe_load_balancer = MoeLoadBalancer( - ep_rank=ep_rank, - ep_size=ep_size, - layer_updates_per_iter=model_config.moe_load_balancer. - layer_updates_per_iter) - logger.info( - f"Created MoE LoadBalancer, layer_updates_per_iter={model_config.moe_load_balancer.layer_updates_per_iter}..." - ) + config = model_config.moe_load_balancer + mode = getattr(config, "mode", "standard") + if config is not None and mode == "per_layer_dynamic": + if not in_supported_model_arch: + raise ValueError( + "per-layer dynamic EPLB does not support model architecture " + f"{model_arch!r}") + if not using_ep: + raise ValueError("per-layer dynamic EPLB requires expert " + "parallelism (moe_ep_size > 1)") + if using_smart_router: + raise ValueError( + "per-layer dynamic EPLB is incompatible with smart routing " + "(moe_cluster_size > 1)") + + if (in_supported_model_arch and using_ep and not using_smart_router + and config is not None): + config.setup(ep_rank=ep_rank, ep_size=ep_size) + if mode == "per_layer_dynamic": + from .per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer + + moe_load_balancer = PerLayerDynamicMoeLoadBalancer( + ep_rank=ep_rank, + ep_size=ep_size, + config=config, + mapping=mapping, + ) + logger.info( + "Created per-layer dynamic MoE LoadBalancer, num_slots=%s, " + "auxiliary_sms=%s...", + config.num_slots, + getattr(config, "auxiliary_sms", 4), + ) + elif mode == "standard": + if config.layer_updates_per_iter > 0: + # TODO: remove this when supported. + # cpu_arch = platform.machine().lower() + # assert cpu_arch == 'aarch64', "online load balancer only support aarch64, e.g. GB200 now, x86 coming soon." + pass + + moe_load_balancer = MoeLoadBalancer( + ep_rank=ep_rank, + ep_size=ep_size, + layer_updates_per_iter=config.layer_updates_per_iter) + logger.info( + "Created standard MoE LoadBalancer, layer_updates_per_iter=%s...", + config.layer_updates_per_iter, + ) + else: + raise ValueError(f"Unknown MoE load-balancer mode: {mode!r}") return moe_load_balancer class MoeLoadBalancerIterContext: def __init__(self, - moe_load_balancer: Optional[MoeLoadBalancer], + moe_load_balancer: Optional[BaseMoeLoadBalancer], enable_statistic: Optional[bool] = None, enable_updates: Optional[bool] = None): self.moe_load_balancer = moe_load_balancer @@ -1107,16 +1210,19 @@ def __exit__(self, exc_type, exc_val, exc_tb): """ if self.moe_load_balancer is not None and not self.moe_load_balancer.is_static_routing( ): - self.moe_load_balancer.end_iter() + if exc_type is None: + self.moe_load_balancer.end_iter() + else: + self.moe_load_balancer.abort_iter(exc_val) return False -def get_moe_load_balancer() -> Optional[MoeLoadBalancer]: +def get_moe_load_balancer() -> Optional[BaseMoeLoadBalancer]: """ - Get the current active MoeLoadBalancer instance. + Get the current active MoE load balancer instance. Returns: - The current active MoeLoadBalancer instance, or None if not in a MoeLoadBalancer context + The current active load balancer instance, or None if not in a load-balancer context """ try: return getattr(_current_moe_load_balancer, 'instance', None) diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index f8bd6ef77942..c345860b3911 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -374,6 +374,7 @@ def build_moe_deployment( # num_experts, so the deployment must say the same. eplb_enabled = get_moe_load_balancer() is not None balancer_config = getattr(model_config, "moe_load_balancer", None) + eplb_mode = getattr(balancer_config, "mode", "standard") if eplb_enabled else "disabled" num_slots = getattr(balancer_config, "num_slots", None) if eplb_enabled else None if num_slots is None: num_slots = num_experts if num_experts is not None else 0 @@ -390,6 +391,7 @@ def build_moe_deployment( env=environment if environment is not None else collect_moe_environment(), # Registered balancer only; config alone does not enable EPLB. eplb_enabled=eplb_enabled, + eplb_mode=eplb_mode, # Routed-expert LoRA only; attention-only LoRA stays False. moe_lora_enabled=has_moe_lora_targets(lora_config), force_dynamic_quantization=model_config.force_dynamic_quantization, @@ -572,6 +574,27 @@ def resolve_moe_impl( rejected = [] eligible: List[MoEImplClass] = [] for candidate in candidates: + if ( + deployment.eplb_mode == "per_layer_dynamic" + and not candidate.capabilities.supports_per_layer_dynamic_eplb + ): + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.EPLB_UNSUPPORTED, + f"{candidate.__name__} does not consume per-layer dynamic EPLB plans", + ) + ) + continue + if deployment.eplb_enabled and not candidate.capabilities.supports_eplb: + rejected.append( + MoERejection( + _legacy_backend_name(candidate), + MoERejectReason.EPLB_UNSUPPORTED, + f"{candidate.__name__} does not support EPLB slot layouts", + ) + ) + continue if deployment.moe_lora_enabled and not candidate.capabilities.supports_moe_lora: rejected.append( MoERejection( diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py index 468c9d3d4f4a..b465348ae3a8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_scheduler.py @@ -51,6 +51,7 @@ from tensorrt_llm._torch.moe.expert_statistic import ExpertStatistic from tensorrt_llm._torch.route_capture import get_active_route_capture # R3 from tensorrt_llm._torch.utils import EventType, Fp4QuantizedTensor +from tensorrt_llm.logger import logger from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided @@ -70,6 +71,31 @@ from .configurable_moe import ConfigurableMoE +def _leading_rows(t: torch.Tensor, rows: int) -> torch.Tensor: + """Return ``t[:rows]``, or ``t`` itself when it already has ``rows`` rows.""" + return t if t.shape[0] == rows else t[:rows] + + +def _split_rows(t: torch.Tensor, sizes: List[int]) -> List[torch.Tensor]: + """Return ``list(t.split(sizes))``; one chunk covering ``t`` is ``t`` itself.""" + if len(sizes) == 1 and sizes[0] == t.shape[0]: + return [t] + return list(t.split(sizes)) + + +def _record_routes( + moe: "ConfigurableMoE", routes: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + """Record ``routes`` for expert statistics and the layer-wise calibrator. + + Returns ``routes``, or the routes the calibrator replays in their place. + """ + if routes is not None: + ExpertStatistic.set_layer(moe.layer_idx) + ExpertStatistic.maybe_add_info(moe.num_slots, routes) + return get_calibrator().maybe_collect_or_replay_slots(moe.num_slots, routes) + + def _wait_input_ready(event: Optional[torch.cuda.Event]) -> None: """Join a pending ``input_ready_event`` before the first read of ``x``. @@ -491,7 +517,9 @@ def _forward_chunk_impl( token_selected_experts, token_final_scales, x, x_sf = fused_result used_fused_route_quant = True - token_selected_experts = token_selected_experts.to(torch.int32) + token_selected_experts = moe._maybe_get_perfect_router_expert_ids( + token_selected_experts.to(torch.int32) + ) if row_offset is not None: route_capture = get_active_route_capture() if route_capture is not None: # R3: in-graph device-buffer capture @@ -541,6 +569,16 @@ def _forward_chunk_impl( f"EPLB is enabled but backend {moe.backend.__class__.__name__} " f"has fused routing (does not support routing separation)" ) + if moe._enable_perfect_router_expert_ids: + logger.warning_once( + "ENABLE_PERFECT_ROUTER=expert_ids has no effect on " + f"{moe.backend.__class__.__name__} here: it routes inside its " + "kernel, so there are no routed expert ids to replace. Use " + "ENABLE_PERFECT_ROUTER=1 instead; for TRTLLM, setting " + "TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1 in every compared run " + "also works.", + key="perfect_router_expert_ids_fused_routing", + ) token_selected_experts = None token_final_scales = None @@ -1198,12 +1236,12 @@ def forward( if input_ids_real is not None: chunk_size_list = all_rank_chunk_size_list[ep_rank] input_ids_chunks = ( - list(input_ids_real.split(chunk_size_list)) if input_ids_real.numel() > 0 else [] + _split_rows(input_ids_real, chunk_size_list) if input_ids_real.numel() > 0 else [] ) routing_aux_chunks = [] if routing_aux is not None and x_real.shape[0] > 0: - routing_aux_chunks = list( - routing_aux[: x_real.shape[0]].split(all_rank_chunk_size_list[ep_rank]) + routing_aux_chunks = _split_rows( + _leading_rows(routing_aux, x_real.shape[0]), all_rank_chunk_size_list[ep_rank] ) outputs = self._run_chunks( x_chunks, @@ -1290,9 +1328,9 @@ def _strip_adp_padding( ) # ADP padding stripped before split, else trailing rows silently # drift into chunk-0 or torch.split shape-errors. - x_real = x[:real_local] - rl_real = router_logits[:real_local] - input_ids_real = input_ids[:real_local] if input_ids is not None else None + x_real = _leading_rows(x, real_local) + rl_real = _leading_rows(router_logits, real_local) + input_ids_real = _leading_rows(input_ids, real_local) if input_ids is not None else None return x_real, rl_real, input_ids_real, real_all_rank_num_tokens, ep_rank, had_meta def _compute_chunk_layout( @@ -1332,8 +1370,8 @@ def _compute_chunk_layout( # ``else []`` shortcut for real_local == 0: equivalent to # x_real.split([0]*num_chunks) but skips the no-op torch call. The # zero-token fallback in ``_run_chunks`` then fires for every chunk. - x_chunks = list(x_real.split(chunk_size_list)) if real_local > 0 else [] - rl_chunks = list(rl_real.split(chunk_size_list)) if real_local > 0 else [] + x_chunks = _split_rows(x_real, chunk_size_list) if real_local > 0 else [] + rl_chunks = _split_rows(rl_real, chunk_size_list) if real_local > 0 else [] return num_chunks, x_chunks, rl_chunks, all_rank_chunk_size_list def _run_chunks( @@ -1431,7 +1469,10 @@ def _forward_chunk( quantize -> ``run_moe`` -> ``start_set_cpu_stage`` -> ``done_set_cpu_stage``. ``start/done_set_cpu_stage`` are placed AFTER ``run_moe``; otherwise dynamic-EPLB weight migration would - race with the fused kernel using those weights. + race with the fused kernel using those weights. A per-layer dynamic + EPLB layer has no CPU stage and no host statistics: it only gets + ``route`` and the post-``run_moe`` ``start_set_cpu_stage``, which + releases its replica plan. """ moe = self.moe assert not moe.apply_router_weight_on_input, ( @@ -1454,14 +1495,21 @@ def _forward_chunk( num_tokens = x.shape[0] assert num_tokens <= x.shape[0], f"num_tokens ({num_tokens}) > x.shape[0] ({x.shape[0]})" - x_chunk_real = x[:num_tokens] - router_logits_chunk_real = router_logits[:num_tokens] - input_ids_chunk_real = input_ids[:num_tokens] if input_ids is not None else None + x_chunk_real = _leading_rows(x, num_tokens) + router_logits_chunk_real = _leading_rows(router_logits, num_tokens) + input_ids_chunk_real = ( + _leading_rows(input_ids, num_tokens) if input_ids is not None else None + ) + + # Per-layer dynamic EPLB plans on the GPU; the CPU-stage and + # statistics hooks below are no-ops for it, so it skips them. + uses_replica_plan = moe._load_balancer_uses_replica_plan() # ----- EPLB: drain previous CPU rebalance ----- # Static EPLB early-returns inside the helper; only the dynamic # balancer actually waits. - moe._load_balancer_start_wait_gpu_stage(is_first_call) + if not uses_replica_plan: + moe._load_balancer_start_wait_gpu_stage(is_first_call) # ----- routing ----- # int32 matches the EPLB stats kernel contract used by the external-comm @@ -1473,13 +1521,20 @@ def _forward_chunk( ) else: token_selected_experts, token_final_scales = moe.routing_method.apply_with_aux( - router_logits_chunk_real, input_ids_chunk_real, routing_aux[:num_tokens] + router_logits_chunk_real, + input_ids_chunk_real, + _leading_rows(routing_aux, num_tokens), ) - token_selected_experts = token_selected_experts.to(torch.int32) + if token_selected_experts.dtype != torch.int32: + token_selected_experts = token_selected_experts.to(torch.int32) + token_selected_experts = moe._maybe_get_perfect_router_expert_ids( + token_selected_experts + ) route_capture = get_active_route_capture() if route_capture is not None: # R3: in-graph device-buffer capture route_capture.capture(moe.layer_idx, token_selected_experts, row_offset) - token_final_scales = token_final_scales.to(torch.float32) + if token_final_scales.dtype != torch.float32: + token_final_scales = token_final_scales.to(torch.float32) else: device = x.device token_selected_experts = torch.empty( @@ -1494,27 +1549,30 @@ def _forward_chunk( ) # ----- EPLB: update stats + remap expert ids -> slot ids ----- - if moe.layer_load_balancer: - moe._load_balancer_done_wait_gpu_stage(is_first_call) - # ignore_allreduce=False: the fused kernel has no side channel - # for an external stats gather. The base helper runs its own - # EP-wide AllReduce, gated to is_last_call=True. - moe._load_balancer_update_statistic( - token_selected_experts, - is_first_call, - is_last_call, - ignore_allreduce=False, - ) + # Per-layer dynamic EPLB plans from the routes that the layer-wise + # calibrator records or replays, so they are recorded before the route; + # otherwise the recorded routes are the slots routed to. + if uses_replica_plan: + token_selected_experts = _record_routes(moe, token_selected_experts) token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + replica_plan = moe.layer_load_balancer.require_replica_plan() else: - token_selected_slots = token_selected_experts - - if token_selected_slots is not None: - ExpertStatistic.set_layer(moe.layer_idx) - ExpertStatistic.maybe_add_info(moe.num_slots, token_selected_slots) - token_selected_slots = get_calibrator().maybe_collect_or_replay_slots( - moe.num_slots, token_selected_slots - ) + if moe.layer_load_balancer: + moe._load_balancer_done_wait_gpu_stage(is_first_call) + # ignore_allreduce=False: the fused kernel has no side channel + # for an external stats gather. The base helper runs its own + # EP-wide AllReduce, gated to is_last_call=True. + moe._load_balancer_update_statistic( + token_selected_experts, + is_first_call, + is_last_call, + ignore_allreduce=False, + ) + token_selected_slots = moe._load_balancer_route(token_selected_experts, moe.use_dp) + else: + token_selected_slots = token_selected_experts + token_selected_slots = _record_routes(moe, token_selected_slots) + replica_plan = None # ----- quantize / prepare ----- if getattr(moe.backend, "supports_fused_prepare", lambda: False)(): @@ -1523,12 +1581,17 @@ def _forward_chunk( # prepare its workspace. moe_input = x_chunk_real x_sf = None - else: + elif replica_plan is None: # Delegate to ``backend.quantize_input`` so each fused-comm backend # owns its own empty-tensor layout. Both MegaMoEDeepGemm and # MegaMoECuteDsl short-circuit ``x.shape[0] == 0`` inside their # quantize_input contracts. moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) + else: + # Leave the plan's SMs to its concurrent HALO-Q/TMA kernel. + moe_input, x_sf = moe.backend.quantize_input( + x_chunk_real, reserved_sms=replica_plan.reserved_sms + ) # CuteDSL needs the scheduler's rank-identical chunk maximum to select # one adaptive bucket on every EP rank; using a local token count could @@ -1550,6 +1613,7 @@ def _forward_chunk( x=moe_input, x_sf=x_sf, output_dtype=output_dtype, + replica_plan=replica_plan, ) ) @@ -1560,7 +1624,8 @@ def _forward_chunk( # start_set before run_moe would let dynamic-EPLB migration race the # kernel. moe._load_balancer_start_set_cpu_stage(is_last_call) - moe._load_balancer_done_set_cpu_stage(is_last_call) + if not uses_replica_plan: + moe._load_balancer_done_set_cpu_stage(is_last_call) return out diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py new file mode 100644 index 000000000000..d37f18c5d3d6 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model-owned lifecycle for per-layer dynamic expert load balancing. + +The standard EPLB implementation couples route updates to host weight +migration. Per-layer dynamic EPLB instead produces one device-side replica plan +for every MoE invocation and copies helper weights on a CUDA stream. It keeps +the existing load-balancer seams: ``manager`` owns the layers and their GPU +resources for the whole model, ``layer`` is the per-layer handle that plans each +route and releases its plan, ``arena`` allocates the live-weight storage, and +``planner`` submits HALO-Q and the TMA copy of its plan. +""" + +from .layer import PerLayerDynamicMoeLayerLoadBalancer +from .manager import PerLayerDynamicMoeLoadBalancer + +__all__ = ["PerLayerDynamicMoeLayerLoadBalancer", "PerLayerDynamicMoeLoadBalancer"] diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/arena.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/arena.py new file mode 100644 index 000000000000..81e8655615c8 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/arena.py @@ -0,0 +1,1066 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""EP-scoped live-weight storage for HALO-Q and hierarchical SAMI copies. + +Every layer owns its home-expert storage and maps a helper-slot set directly +after it in virtual address space: either its own set, which lets helper experts +stay resident across forwards, or one of two banks that all layers alternate +between. Per-layer READY terminals and small planes remain independent. +Allocation and binding are collective across the EP group; vendor imports stay +lazy until the rebalance path is enabled. +""" + +import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import torch +from torch.utils._python_dispatch import _disable_current_modes + +from tensorrt_llm.logger import logger + +__all__ = ["SharedSlotArenaProvider", "allocate_rebalance_arena"] + +BUFFER_ALIGNMENT = 2 * 1024 * 1024 + +#: The ``uint64[EP]`` terminal array is addressed by the SAMI copy kernel and +#: wants its own page. +TERMINAL_ALIGNMENT = 4096 + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _plane_dtypes() -> Dict[str, torch.dtype]: + """Return the canonical kernel-plane dtypes. + + Resolve optional Torch dtypes lazily so unsupported wheels fail only when + constructing an enabled rebalance arena. + """ + fp4 = getattr(torch, "float4_e2m1fn_x2", None) + fp8 = getattr(torch, "float8_e4m3fn", None) + if fp4 is None or fp8 is None: + raise RuntimeError( + "MoE rebalance live arena needs torch.float4_e2m1fn_x2 and " + "torch.float8_e4m3fn to type the NVFP4 weight planes; this torch " + f"({torch.__version__}) exposes float4_e2m1fn_x2=" + f"{fp4 is not None}, float8_e4m3fn={fp8 is not None}." + ) + return { + "uint8": torch.uint8, + "float4_e2m1fn_x2": fp4, + "float8_e4m3fn": fp8, + "float32": torch.float32, + } + + +def _plane_view( + flat_u8: torch.Tensor, + plane: Any, + slots: int, + *, + kernel: bool, +) -> torch.Tensor: + """Create a kernel or framework view from the canonical plane descriptor.""" + metadata = plane.kernel_metadata(slots) if kernel else plane.storage_metadata(slots) + shape, stride, dtype_name, element_size, pointer_alignment = metadata + dtype = _plane_dtypes().get(dtype_name) + if dtype is None: + raise RuntimeError(f"unsupported live-plane dtype {dtype_name!r}") + typed = flat_u8.view(dtype) + if typed.element_size() != element_size: + raise RuntimeError( + f"{plane.name} descriptor element size {element_size} does not " + f"match torch dtype {dtype_name} ({typed.element_size()})" + ) + if int(typed.data_ptr()) % pointer_alignment: + raise RuntimeError(f"{plane.name} pointer is not {pointer_alignment}-byte aligned") + return typed.as_strided(shape, stride) + + +@dataclass +class RebalanceLiveArena: + """Own a bound live arena and both views of its seven weight planes. + + Kernel views preserve the allocator's bound identities. Framework aliases + back nn.Parameters without changing the allocation or plane offsets. + """ + + #: The allocator owns these bytes; the arena only retains and exposes it. + provider: Any + arena: Any + bound: Any + bundle: Any + plane_names: Tuple[str, ...] + local_plane_views: Tuple[torch.Tensor, ...] + tekit_alias_views: Tuple[torch.Tensor, ...] + home_experts: int + helper_slots: int + #: READY generation ``g`` currently published into this arena's helper + #: slots. 0 means "nothing published yet" -- the kernel ABI rejects it + #: (the gate requires g in [1, 2**63)), which is the intended fail-closed + #: state before the weight transport has run for this generation. + ready_generation_value: int = 0 + + @property + def slot_count(self) -> int: + return self.home_experts + self.helper_slots + + def terminal_flags_tensor(self) -> torch.Tensor: + """This rank's ``uint64[EP]`` terminal view -- the contract's + ``hot_expert_weight_ready_flags``. Sources publish into the multicast + alias of the SAME bytes, so a publish by any EP rank becomes visible + here without any host-side copy.""" + # Hierarchical arenas expose terminals through their base arena. + arena = getattr(self.arena, "base", self.arena) + return arena.terminals.local_view + + def ready_generation(self) -> int: + return int(self.ready_generation_value) + + def publish_ready_generation(self, generation: int) -> None: + """Record the generation the kernel must wait for. Called by the + weight transport AFTER it has issued the payload copies for ``g``; + the terminal cells themselves are written by the transport.""" + if type(generation) is not int or not (1 <= generation < (1 << 63)): + raise ValueError( + f"READY generation must be an exact int in [1, 2**63); got {generation!r}" + ) + self.ready_generation_value = generation + + def assert_identity(self) -> None: + """Run the optional allocator identity guard. + + The caller must also check framework parameter pointers: allocator checks + cannot detect replacement of a Parameter's storage. + """ + # The allocator guard is optional; framework pointer checks remain mandatory. + guard = getattr(self.bound, "assert_identity", None) + if guard is not None: + guard() + + +class _EpComm: + """Adapt the resolved torch EP ProcessGroup to SAMI's collectives. + + Using the default distributed group could include ranks outside this arena + and deadlock its collective allocation or binding. + """ + + def __init__(self, process_group: Any) -> None: + if process_group is None: + raise ValueError("MoE rebalance requires an EP ProcessGroup") + self.process_group = process_group + self.rank = int(torch.distributed.get_rank(group=process_group)) + self.world = int(torch.distributed.get_world_size(group=process_group)) + + def validate_geometry(self, *, expected_rank: int, expected_world: int) -> None: + expected = (int(expected_rank), int(expected_world)) + actual = (self.rank, self.world) + if actual != expected: + raise ValueError( + "EP ProcessGroup rank/size differs from the model mapping: " + f"process_group={actual}, mapping={expected}." + ) + + def barrier(self) -> None: + with _disable_current_modes(): + torch.distributed.barrier(group=self.process_group) + + def bcast(self, payload: object, root: int) -> object: + if type(root) is not int or not 0 <= root < self.world: + raise ValueError(f"broadcast root must be in [0, {self.world}); got {root!r}") + values = [payload] + source_rank = torch.distributed.get_global_rank(self.process_group, root) + with _disable_current_modes(): + torch.distributed.broadcast_object_list( + values, src=source_rank, group=self.process_group + ) + return values[0] + + def allgather(self, value: object) -> list: + values = [None] * self.world + with _disable_current_modes(): + torch.distributed.all_gather_object(values, value, group=self.process_group) + return values + + +class _RawPointer: + """Expose a CUDA address through __cuda_array_interface__. + + Tensor views do not own this allocation; the provider retains its owners. + """ + + def __init__(self, pointer: int, nbytes: int) -> None: + self.__cuda_array_interface__ = { + "data": (int(pointer), False), + "shape": (int(nbytes),), + "typestr": "|u1", + "strides": None, + "version": 3, + } + + +#: Planes whose whole [H + S] span is at most this many bytes (the fp32 scale +#: planes) live in a per-layer record of a small shared page instead of in +#: their own granularity-sized allocations. +_RECORD_PLANE_BYTES = 64 * 1024 +_RECORD_PLANE_ALIGNMENT = 256 + + +_HELPER_BANK_COUNT = 2 + + +class _ViewSpec: + """Address and layout of a plane view SAMI only binds, never dereferences. + + Peer and multicast views feed SAMI's bind-time pointer tables. A multicast + view's row 0 lies in a reserved but unmapped prefix, which a torch tensor + cannot wrap; only its helper rows are mapped and written by the copy. + """ + + __slots__ = ("_pointer", "shape", "_stride", "dtype", "_element_size", "device", "is_cuda") + + def __init__(self, pointer: int, like: torch.Tensor) -> None: + self._pointer = int(pointer) + self.shape = tuple(like.shape) + self._stride = tuple(like.stride()) + self.dtype = like.dtype + self._element_size = int(like.element_size()) + self.device = like.device + self.is_cuda = True + + def data_ptr(self) -> int: + return self._pointer + + def stride(self) -> Tuple[int, ...]: + return self._stride + + def element_size(self) -> int: + return self._element_size + + +class _Vmm: + """CUDA VMM calls on FABRIC-exportable memory with explicit ownership.""" + + def __init__(self, device: int, group_sizes: Tuple[int, ...]) -> None: + from cuda.bindings import driver as cuda + + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami._util import check_cuda + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami.fabric import ( + FABRIC, + _allocation_properties, + _read_write_access, + ) + + self.cuda = cuda + self.check = check_cuda + self.fabric = FABRIC + self.device = int(device) + self.properties = _allocation_properties(self.device) + self._access = [_read_write_access(self.device)] + # Every chunk is multicast-bindable at every level, so align to the + # coarsest MINIMUM granularity rather than the 512 MiB recommendation. + values = [ + BUFFER_ALIGNMENT, + int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + self.properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ), + "cuMemGetAllocationGranularity(shared slots)", + ) + ), + ] + for members in sorted(set(group_sizes)): + values.append( + int( + check_cuda( + cuda.cuMulticastGetGranularity( + self.multicast_properties(members, BUFFER_ALIGNMENT), + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_MINIMUM, + ), + "cuMulticastGetGranularity(shared slots)", + ) + ) + ) + self.granularity = math.lcm(*values) + self.local_allocation_handles: List[Any] = [] + self.imported_allocation_handles: List[Any] = [] + self.created_multicast_handles: List[Any] = [] + self.imported_multicast_handles: List[Any] = [] + self.mappings: List[Tuple[int, int]] = [] + self.reservations: List[Tuple[int, int]] = [] + self.multicast_binds: List[Tuple[Any, int, int]] = [] + self.allocated_bytes = 0 + + def multicast_properties(self, members: int, nbytes: int) -> Any: + properties = self.cuda.CUmulticastObjectProp() + properties.numDevices = int(members) + properties.size = int(nbytes) + properties.handleTypes = int(self.fabric) + properties.flags = 0 + return properties + + def create(self, nbytes: int) -> Tuple[Any, bytes]: + handle = self.check( + self.cuda.cuMemCreate(nbytes, self.properties, 0), "cuMemCreate(shared slots)" + ) + self.local_allocation_handles.append(handle) + self.allocated_bytes += nbytes + exported = self.check( + self.cuda.cuMemExportToShareableHandle(handle, self.fabric, 0), + "cuMemExportToShareableHandle(shared slots)", + ) + shareable = bytes(exported.data) + if len(shareable) != 64: + raise RuntimeError(f"fabric handle has {len(shareable)} bytes; expected 64") + return handle, shareable + + def import_(self, shareable: bytes) -> Any: + if not isinstance(shareable, bytes) or len(shareable) != 64: + raise RuntimeError("peer fabric handle must contain exactly 64 bytes") + handle = self.check( + self.cuda.cuMemImportFromShareableHandle(bytearray(shareable), self.fabric), + "cuMemImportFromShareableHandle(shared slots)", + ) + self.imported_allocation_handles.append(handle) + return handle + + def register_created_multicast(self, handle: Any) -> None: + self.created_multicast_handles.append(handle) + + def import_multicast(self, shareable: bytes) -> Any: + if not isinstance(shareable, bytes) or len(shareable) != 64: + raise RuntimeError("multicast fabric handle must contain exactly 64 bytes") + handle = self.check( + self.cuda.cuMemImportFromShareableHandle(bytearray(shareable), self.fabric), + "cuMemImportFromShareableHandle(shared multicast)", + ) + self.imported_multicast_handles.append(handle) + return handle + + def reserve(self, nbytes: int) -> int: + address = int( + self.check( + self.cuda.cuMemAddressReserve(nbytes, self.granularity, 0, 0), + "cuMemAddressReserve(shared slots)", + ) + ) + self.reservations.append((address, int(nbytes))) + return address + + def map(self, address: int, nbytes: int, handle: Any) -> None: + self.check(self.cuda.cuMemMap(address, nbytes, 0, handle, 0), "cuMemMap(shared slots)") + self.mappings.append((int(address), int(nbytes))) + self.check( + self.cuda.cuMemSetAccess(address, nbytes, self._access, 1), + "cuMemSetAccess(shared slots)", + ) + + def record_multicast_bind(self, handle: Any, offset: int, nbytes: int) -> None: + self.multicast_binds.append((handle, int(offset), int(nbytes))) + + def unmap_all(self) -> None: + while self.mappings: + address, nbytes = self.mappings[-1] + self.check(self.cuda.cuMemUnmap(address, nbytes), "cuMemUnmap(shared slots)") + self.mappings.pop() + + def unbind_all(self) -> None: + while self.multicast_binds: + handle, offset, nbytes = self.multicast_binds[-1] + self.check( + self.cuda.cuMulticastUnbind(handle, self.device, offset, nbytes), + "cuMulticastUnbind(shared slots)", + ) + self.multicast_binds.pop() + + def free_reservations(self) -> None: + while self.reservations: + address, nbytes = self.reservations[-1] + self.check( + self.cuda.cuMemAddressFree(address, nbytes), + "cuMemAddressFree(shared slots)", + ) + self.reservations.pop() + + def _release_handles(self, handles: List[Any], label: str) -> None: + while handles: + self.check(self.cuda.cuMemRelease(handles[-1]), label) + handles.pop() + + def release_imported_handles(self) -> None: + self._release_handles( + self.imported_multicast_handles, + "cuMemRelease(imported shared multicast)", + ) + self._release_handles( + self.imported_allocation_handles, + "cuMemRelease(imported shared allocation)", + ) + + def release_created_handles(self) -> None: + self._release_handles( + self.created_multicast_handles, + "cuMemRelease(created shared multicast)", + ) + self._release_handles( + self.local_allocation_handles, + "cuMemRelease(local shared allocation)", + ) + + def zero(self, address: int, nbytes: int) -> None: + self.check( + self.cuda.cuMemsetD8(self.cuda.CUdeviceptr(address), 0, nbytes), + "cuMemsetD8(shared slots)", + ) + + +class _MulticastSpec: + """One aligned group's multicast pair: member binds and local mappings.""" + + def __init__( + self, + members: Tuple[int, ...], + nbytes: int, + binds: List[Tuple[int, Any, int]], + mappings: List[int], + ) -> None: + self.members = members + self.nbytes = nbytes + #: (multicast offset, member allocation handle, bytes) bound on every member. + self.binds = binds + #: Reserved addresses where this rank maps its selected object. + self.mappings = mappings + self.handles: List[Any] = [] + + +def _create_multicast(vmm: _Vmm, comm: Any, rank: int, specs: List[_MulticastSpec]) -> None: + """Collectively create, bind and map dual-creator multicast objects. + + Follows MulticastTeamPair: the first two members each create one object, + every member binds its chunks to both, and the first member maps the second + member's object while every other member maps the first member's. Ranks + outside a group join only the collectives; specs are rank-ordered alike. + """ + cuda, check = vmm.cuda, vmm.check + created: Dict[Tuple[int, int], Any] = {} + exported: Dict[Tuple[int, int], bytes] = {} + for index, spec in enumerate(specs): + for alias, creator in enumerate(spec.members[:2]): + if rank == creator: + handle = check( + cuda.cuMulticastCreate( + vmm.multicast_properties(len(spec.members), spec.nbytes) + ), + "cuMulticastCreate(shared slots)", + ) + vmm.register_created_multicast(handle) + created[(index, alias)] = handle + exported[(index, alias)] = bytes( + check( + cuda.cuMemExportToShareableHandle(handle, vmm.fabric, 0), + "cuMemExportToShareableHandle(shared multicast)", + ).data + ) + gathered = comm.allgather(exported) + joined = [(index, spec) for index, spec in enumerate(specs) if rank in spec.members] + for index, spec in joined: + for alias, creator in enumerate(spec.members[:2]): + handle = created.get((index, alias)) + if handle is None: + handle = vmm.import_multicast(gathered[creator][(index, alias)]) + spec.handles.append(handle) + check( + cuda.cuMulticastAddDevice(handle, vmm.device), "cuMulticastAddDevice(shared slots)" + ) + comm.barrier() + for _, spec in joined: + for handle in spec.handles: + for offset, memory, nbytes in spec.binds: + check( + cuda.cuMulticastBindMem_v2(handle, vmm.device, offset, memory, 0, nbytes, 0), + "cuMulticastBindMem_v2(shared slots)", + ) + vmm.record_multicast_bind(handle, offset, nbytes) + comm.barrier() + for _, spec in joined: + selected = spec.handles[1 if len(spec.handles) > 1 and rank == spec.members[0] else 0] + for address in spec.mappings: + vmm.map(address, spec.nbytes, selected) + comm.barrier() + + +@dataclass +class _HelperSet: + """One physical helper-slot set: a chunk per large plane, one multicast + object per hierarchy level, and that object's view base per plane.""" + + chunks: Dict[str, Any] + level_bases: Tuple[Dict[str, int], ...] + level_specs: Tuple[_MulticastSpec, ...] + zeroed: bool = False + + +@dataclass +class _RecordPage: + """Per-layer records of READY terminals and small planes, mapped on every + rank and bound to one multicast object per hierarchy level.""" + + uc_ptrs: Tuple[int, ...] + mc_ptrs: Tuple[int, ...] + level_specs: Tuple[_MulticastSpec, ...] + + +class _SharedLayer: + """Anti-GC owner of one layer's home chunks and raw pointer holders.""" + + def __init__(self, pool: "_SharedSlotPool", helper_set: int) -> None: + self.pool = pool + self.helper_set = helper_set + self.home_handles: List[Any] = [] + self.holders: List[_RawPointer] = [] + + +class _SharedSlotPool: + """Helper-slot sets of every MoE layer of one geometry. + + Layer ``i`` maps helper set ``i % sets`` directly after its home rows, so + each plane is still one contiguous ``[H + S]`` kernel view. With + ``per_layer_sets`` the pool adds one set per layer as the layers are built, + so only a layer's own copies write its helper rows; otherwise every layer + shares two sets. No extra synchronization guards reuse of a shared set: + layer ``i``'s copy follows HALO-Q's all-rank exchange, which every rank + enters only after MAIN finished the previous layers' MegaMoE. + + READY terminals stay per layer because one forward shares its generation + across layers. They and the fp32 planes live in per-layer records. + """ + + def __init__( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: Any, + device: int, + comm: Any, + owner_registry: List["_SharedSlotPool"], + per_layer_sets: bool, + ) -> None: + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import hierarchy_group_sizes + + torch.cuda.set_device(device) + self.world = world + self.rank = rank + self.device = device + self.comm = comm + self._closed = False + self._closing = False + self._producers: Dict[int, Any] = {} + self.home_count = home_count + self.helper_count = helper_count + self.slots = home_count + helper_count + self.bundle = bundle + self.group_sizes = hierarchy_group_sizes(world) + self.groups = tuple( + tuple(range(rank // size * size, rank // size * size + size)) + for size in self.group_sizes + ) + self.vmm = _Vmm(device, self.group_sizes) + granularity = self.vmm.granularity + + self.large = tuple(p for p in bundle.planes if p.nbytes * self.slots > _RECORD_PLANE_BYTES) + self.home_bytes = { + p.name: _align_up(home_count * p.nbytes, granularity) for p in self.large + } + self.helper_bytes = { + p.name: _align_up(helper_count * p.nbytes, granularity) for p in self.large + } + # One multicast object per set and level holds every large plane's chunk. + self.mc_offsets: Dict[str, int] = {} + cursor = 0 + for plane in self.large: + self.mc_offsets[plane.name] = cursor + cursor += self.helper_bytes[plane.name] + self.set_bytes = cursor + + # A record keeps its READY terminals on their own page. + cursor = _align_up(world * 8, TERMINAL_ALIGNMENT) + self.record_offsets: Dict[str, int] = {} + for plane in bundle.planes: + if plane.name not in self.home_bytes: + cursor = _align_up(cursor, _RECORD_PLANE_ALIGNMENT) + self.record_offsets[plane.name] = cursor + cursor += self.slots * plane.nbytes + self.record_stride = _align_up(cursor, TERMINAL_ALIGNMENT) + self.page_bytes = _align_up(self.record_stride, granularity) + self.records_per_page = self.page_bytes // self.record_stride + self.pages: List[_RecordPage] = [] + self.layer_count = 0 + self.sets: Tuple[_HelperSet, ...] = () + self.per_layer_sets = per_layer_sets + self._construction_complete = False + self._owner_registry = owner_registry + self._owner_registry.append(self) + if not per_layer_sets: + self.sets = self._create_sets(_HELPER_BANK_COUNT, comm) + self._construction_complete = True + + def register_producer(self, producer: Any) -> None: + if self._closed: + raise RuntimeError("cannot register a producer on a closed shared-slot pool") + self._producers.setdefault(id(producer), producer) + + def _cleanup_manifest(self) -> tuple: + return ( + 1, + self._construction_complete, + self.world, + self.home_count, + self.helper_count, + self.bundle.hidden, + self.bundle.intermediate, + self.group_sizes, + self.layer_count, + len(self.pages), + len(self.sets), + tuple((plane.name, int(plane.nbytes)) for plane in self.bundle.planes), + tuple( + ( + type(producer).__qualname__, + getattr(producer, "layer_idx", getattr(producer, "_layer_idx", None)), + ) + for producer in self._producers.values() + ), + ) + + def _collective_phase(self, label: str, operation: Any) -> None: + error = None + try: + operation() + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + errors = self.comm.allgather(error) + if len(errors) != self.world: + raise RuntimeError( + f"shared-slot cleanup {label} returned {len(errors)} ranks; expected {self.world}" + ) + failures = [f"rank {rank}: {item}" for rank, item in enumerate(errors) if item] + if failures: + raise RuntimeError(f"shared-slot cleanup {label} failed: {'; '.join(failures)}") + + def _close_producers(self) -> None: + for producer in self._producers.values(): + close = getattr(producer, "close", None) + if close is None: + raise TypeError(f"registered producer {type(producer).__qualname__} has no close") + if close() is False: + raise RuntimeError( + f"registered producer {type(producer).__qualname__} did not close" + ) + + def _on_device(self, operation: Any) -> None: + with torch.cuda.device(self.device): + operation() + + def close_collectively(self, local_safe: bool = True) -> bool: + if self._closed: + return True + if self._closing: + raise RuntimeError("shared-slot pool cleanup is already in progress") + preflight = self.comm.allgather( + { + "local_safe": bool(local_safe), + "manifest": self._cleanup_manifest(), + }, + ) + if len(preflight) != self.world: + raise RuntimeError( + f"shared-slot cleanup preflight returned {len(preflight)} ranks; " + f"expected {self.world}" + ) + if not all(item["local_safe"] for item in preflight): + return False + manifests = [item["manifest"] for item in preflight] + expected = manifests[0] + if any(manifest != expected for manifest in manifests[1:]): + raise RuntimeError("shared-slot cleanup manifests differ across the EP communicator") + self._closing = True + try: + self._collective_phase( + "device synchronize", + lambda: self._on_device(lambda: torch.cuda.synchronize(self.device)), + ) + self._collective_phase("producers", lambda: self._on_device(self._close_producers)) + self._producers.clear() + self._collective_phase("unmap", lambda: self._on_device(self.vmm.unmap_all)) + self._collective_phase("multicast unbind", lambda: self._on_device(self.vmm.unbind_all)) + self._collective_phase( + "address free", lambda: self._on_device(self.vmm.free_reservations) + ) + self._collective_phase( + "release imported handles", + lambda: self._on_device(self.vmm.release_imported_handles), + ) + self._collective_phase( + "release created handles", + lambda: self._on_device(self.vmm.release_created_handles), + ) + self._closed = True + if self in self._owner_registry: + self._owner_registry.remove(self) + self.sets = () + self.pages.clear() + return True + finally: + self._closing = False + + def _create_sets(self, count: int, comm: Any) -> Tuple[_HelperSet, ...]: + if not self.large: + return tuple(_HelperSet({}, tuple({} for _ in self.groups), ()) for _ in range(count)) + sets, specs = [], [] + for _ in range(count): + chunks = {p.name: self.vmm.create(self.helper_bytes[p.name])[0] for p in self.large} + level_bases, level_specs = [], [] + for members in self.groups: + bases: Dict[str, int] = {} + mappings: List[int] = [] + for plane in self.large: + # Row 0 of the view sits H rows before the plane's chunk. + prefix = self.home_bytes[plane.name] + address = self.vmm.reserve(prefix + self.set_bytes) + mappings.append(address + prefix) + bases[plane.name] = ( + address + + prefix + + self.mc_offsets[plane.name] + - self.home_count * plane.nbytes + ) + spec = _MulticastSpec( + members, + self.set_bytes, + [ + (self.mc_offsets[p.name], chunks[p.name], self.helper_bytes[p.name]) + for p in self.large + ], + mappings, + ) + specs.append(spec) + level_specs.append(spec) + level_bases.append(bases) + sets.append(_HelperSet(chunks, tuple(level_bases), tuple(level_specs))) + _create_multicast(self.vmm, comm, self.rank, specs) + return tuple(sets) + + def _create_page(self, comm: Any) -> _RecordPage: + handle, shareable = self.vmm.create(self.page_bytes) + local = self.vmm.reserve(self.page_bytes) + self.vmm.map(local, self.page_bytes, handle) + self.vmm.zero(local, self.page_bytes) + torch.cuda.synchronize(self.device) + shareables = comm.allgather(shareable) + if len(shareables) != self.world: + raise RuntimeError("shared-slot record exchange returned an incomplete EP") + uc_ptrs = [] + for peer in range(self.world): + if peer == self.rank: + uc_ptrs.append(local) + continue + address = self.vmm.reserve(self.page_bytes) + self.vmm.map(address, self.page_bytes, self.vmm.import_(shareables[peer])) + uc_ptrs.append(address) + mc_ptrs = [self.vmm.reserve(self.page_bytes) for _ in self.groups] + specs = [ + _MulticastSpec(members, self.page_bytes, [(0, handle, self.page_bytes)], [address]) + for members, address in zip(self.groups, mc_ptrs) + ] + _create_multicast(self.vmm, comm, self.rank, specs) + return _RecordPage(tuple(uc_ptrs), tuple(mc_ptrs), tuple(specs)) + + def _flat(self, owner: _SharedLayer, pointer: int, nbytes: int) -> torch.Tensor: + holder = _RawPointer(pointer, nbytes) + owner.holders.append(holder) + return torch.as_tensor(holder, device=f"cuda:{self.device}") + + def build_layer(self, comm: Any) -> Tuple[Any, Dict[str, torch.Tensor], _SharedLayer, int]: + """Collectively map one layer; returns its arena, framework views, + owner and own allocation bytes.""" + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import ( + HierarchicalLiveWeightArena, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + ) + + world, rank, slots = self.world, self.rank, self.slots + index = self.layer_count + if index % self.records_per_page == 0: + self.pages.append(self._create_page(comm)) + page = self.pages[-1] + record = (index % self.records_per_page) * self.record_stride + if self.per_layer_sets: + self.sets += self._create_sets(1, comm) + helper_set = self.sets[index % len(self.sets)] + layer = _SharedLayer(self, index % len(self.sets)) + self.layer_count += 1 + + homes = {p.name: self.vmm.create(self.home_bytes[p.name]) for p in self.large} + layer.home_handles = [homes[p.name][0] for p in self.large] + peer_handles = comm.allgather(tuple(homes[p.name][1] for p in self.large)) + if len(peer_handles) != world: + raise RuntimeError("shared-slot home exchange returned an incomplete EP") + + planes, group_views, group_owners = [], [], [] + tekit_views: Dict[str, torch.Tensor] = {} + for plane in self.bundle.planes: + span = slots * plane.nbytes + if plane.name in self.home_bytes: + large_index = [p.name for p in self.large].index(plane.name) + home_bytes = self.home_bytes[plane.name] + helper_bytes = self.helper_bytes[plane.name] + # Home rows end where the helper chunk begins. + pad = home_bytes - self.home_count * plane.nbytes + address = self.vmm.reserve(home_bytes + helper_bytes) + self.vmm.map(address, home_bytes, homes[plane.name][0]) + self.vmm.map(address + home_bytes, helper_bytes, helper_set.chunks[plane.name]) + self.vmm.zero(address, home_bytes) + if not helper_set.zeroed: + self.vmm.zero(address + home_bytes, helper_bytes) + uc_ptrs = [] + for peer in range(world): + if peer == rank: + uc_ptrs.append(address + pad) + continue + # Peers expose only home rows; the reservation keeps the + # view's helper tail clear of every other mapping. + peer_address = self.vmm.reserve(home_bytes + helper_bytes) + self.vmm.map( + peer_address, + home_bytes, + self.vmm.import_(peer_handles[peer][large_index]), + ) + uc_ptrs.append(peer_address + pad) + level_ptrs = tuple(bases[plane.name] for bases in helper_set.level_bases) + owners = helper_set.level_specs + else: + offset = record + self.record_offsets[plane.name] + uc_ptrs = [pointer + offset for pointer in page.uc_ptrs] + level_ptrs = tuple(pointer + offset for pointer in page.mc_ptrs) + owners = page.level_specs + flat = self._flat(layer, uc_ptrs[rank], span) + kernel_view = _plane_view(flat, plane, slots, kernel=True) + uc_views = tuple( + kernel_view if peer == rank else _ViewSpec(pointer, kernel_view) + for peer, pointer in enumerate(uc_ptrs) + ) + level_views = tuple(_ViewSpec(pointer, kernel_view) for pointer in level_ptrs) + tekit_views[plane.name] = _plane_view(flat, plane, slots, kernel=False) + planes.append( + LivePlaneView( + name=plane.name, + uc_views=uc_views, + mc_view=level_views[0], + aliases_same_backing=True, + backing_owner=layer, + ) + ) + group_views.append(level_views) + group_owners.append(tuple(owners)) + helper_set.zeroed = True + + terminals = LiveTerminalView( + local_view=self._flat(layer, page.uc_ptrs[rank] + record, world * 8).view(torch.uint64), + mc_view=self._flat(layer, page.mc_ptrs[0] + record, world * 8).view(torch.uint64), + aliases_same_backing=True, + backing_owner=layer, + ) + arena = HierarchicalLiveWeightArena( + base=LiveWeightArena(planes=tuple(planes), terminals=terminals), + group_sizes=self.group_sizes, + group_mc_views=tuple(group_views), + group_mc_owners=tuple(group_owners), + aliases_same_backing=True, + ) + return arena, tekit_views, layer, sum(self.home_bytes.values()) + + +class SharedSlotArenaProvider: + """Provide one layer's arena over its helper-slot set. + + ``per_layer_helper_sets`` gives the layer its own helper set; otherwise it + maps one of the two banks shared by all layers. The broadcaster requests + the arena again after Parameters are bound, so a provider memoizes its + first build and requires identical geometry. + """ + + def __init__(self, owner: Any, *, per_layer_helper_sets: bool = False) -> None: + # The model-level per-layer dynamic manager owns both the geometry cache + # and creation order. No Mapping attribute or process-global registry + # participates in allocation lifetime. + self._owner = owner + pools = getattr(owner, "_rebalance_shared_slot_pools", None) + order = getattr(owner, "_rebalance_shared_slot_pool_order", None) + if not isinstance(pools, dict) or not isinstance(order, list): + raise TypeError("SharedSlotArenaProvider owner must expose model-local pool state") + self._pools = pools + self._pool_order = order + self._per_layer_helper_sets = bool(per_layer_helper_sets) + self.buffer_bytes = 0 + self.group_sizes: Tuple[int, ...] = () + self.pool: Optional[_SharedSlotPool] = None + self.layer: Optional[_SharedLayer] = None + self._built: Optional[Any] = None + self._signature: Optional[tuple] = None + self._tekit_views: Dict[str, torch.Tensor] = {} + + def build_hierarchical_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: Any, + device: int, + comm: Any, + ) -> Any: + signature = ( + int(world), + int(rank), + int(home_count), + int(helper_count), + int(device), + bundle, + self._per_layer_helper_sets, + ) + if self._built is not None: + if self._signature != signature: + raise RuntimeError( + "shared-slot arena provider is per-layer and single-geometry; " + f"rebuilt with {signature} after {self._signature}" + ) + if self.pool is not None and not self.pool._closed: + return self._built + self._built = None + self._tekit_views = {} + self.layer = None + self.pool = None + self.buffer_bytes = 0 + self.group_sizes = () + self._signature = signature + key = signature + pool = self._pools.get(key) + if pool is not None and pool._closed: + self._pools.pop(key) + pool = None + if pool is None: + pool = _SharedSlotPool( + world=int(world), + rank=int(rank), + home_count=int(home_count), + helper_count=int(helper_count), + bundle=bundle, + device=int(device), + comm=comm, + owner_registry=self._pool_order, + per_layer_sets=self._per_layer_helper_sets, + ) + self._pools[key] = pool + self._built, self._tekit_views, self.layer, self.buffer_bytes = pool.build_layer(comm) + self.pool = pool + self.group_sizes = pool.group_sizes + return self._built + + def register_producer(self, producer: Any) -> None: + if self.pool is None: + raise RuntimeError("build the shared-slot arena before registering its producer") + self.pool.register_producer(producer) + + def tekit_alias(self, name: str) -> torch.Tensor: + """Return this rank's framework storage view for a weight plane.""" + return self._tekit_views[name] + + +def allocate_rebalance_arena( + *, + home_experts: int, + helper_slots: int, + bundle: Any, + ep_comm: Any, + mapping: Any, + provider: SharedSlotArenaProvider, + device: int, + layer_idx: Optional[int], +) -> RebalanceLiveArena: + """Allocate and bind one layer's arena collectively across its EP group. + + Every rank must call this with identical geometry and collective order. + """ + world = int(mapping.moe_ep_size) + rank = int(mapping.moe_ep_rank) + if world < 2: + raise RuntimeError(f"MoE rebalance needs at least two EP ranks; got moe_ep_size={world}.") + + if bundle is None or not getattr(bundle, "planes", None): + raise TypeError("rebalance arena requires a quant-method live-weight plane spec") + if not isinstance(ep_comm, _EpComm): + raise TypeError("rebalance arena requires its model-owned EP communicator") + ep_comm.validate_geometry(expected_rank=rank, expected_world=world) + if not isinstance(provider, SharedSlotArenaProvider): + raise TypeError("rebalance arena provider must be owned by the model load balancer") + arena = provider.build_hierarchical_live_arena( + world=world, + rank=rank, + home_count=int(home_experts), + helper_count=int(helper_slots), + bundle=bundle, + device=int(device), + comm=ep_comm, + ) + # Initialize READY terminals explicitly before peers can observe them. + # This startup-only reset does not depend on allocator initialization details. + arena.base.terminals.local_view.zero_() + torch.cuda.synchronize() + ep_comm.barrier() + bound = arena.bind( + world=world, + rank=rank, + home_count=int(home_experts), + helper_count=int(helper_slots), + bundle=bundle, + device=int(device), + ) + plane_names = tuple(plane.name for plane in bundle.planes) + # Retain the hierarchical arena and its bound base views. + live = RebalanceLiveArena( + provider=provider, + arena=arena, + bound=bound, + bundle=bundle, + plane_names=plane_names, + # The bound hierarchy exposes local plane views through its base arena. + local_plane_views=tuple(bound.base.local_plane_views), + tekit_alias_views=tuple(provider.tekit_alias(name) for name in plane_names), + home_experts=int(home_experts), + helper_slots=int(helper_slots), + ) + pool = provider.pool + layer = provider.layer + assert pool is not None and layer is not None + logger.debug( + f"[MegaMoECuteDsl] layer={layer_idx} MoE rebalance SHARED-SLOT live arena: " + f"{provider.buffer_bytes} B of home rows for M={live.slot_count} slots " + f"(H={home_experts} + S={helper_slots}) on helper set " + f"{layer.helper_set}/{len(pool.sets)}; pool sets " + f"{len(pool.sets)} x {pool.set_bytes} B, record pages {len(pool.pages)} x " + f"{pool.page_bytes} B, granularity {pool.vmm.granularity} B, VMM total " + f"{pool.vmm.allocated_bytes} B after {pool.layer_count} layers; ep_size={world} " + f"ep_rank={rank} device={device}, group_sizes={provider.group_sizes}." + ) + return live diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py new file mode 100644 index 000000000000..a81315445f82 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/layer.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-layer handle of per-layer dynamic EPLB: it plans each route and releases the plan.""" + +from __future__ import annotations + +import weakref +from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence + +import torch + +from ..impl_contract import MoEReplicaPlan + +if TYPE_CHECKING: + from .manager import PerLayerDynamicMoeLoadBalancer + + +def _append_note(error: BaseException, note: str) -> None: + """Attach ``note`` to ``error`` like ``BaseException.add_note``, which Python 3.10 lacks.""" + notes = getattr(error, "__notes__", None) + if type(notes) is not list: + notes = [] + error.__notes__ = notes + notes.append(note) + + +class PerLayerDynamicMoeLayerLoadBalancer: + """Adapter between the common MoE routing seam and one replica planner.""" + + # ``route`` publishes a ``MoEReplicaPlan`` that the caller must finish. + produces_replica_plan = True + + def __init__( + self, + manager: "PerLayerDynamicMoeLoadBalancer", + *, + layer_idx: int, + expert_count: int, + top_k: int, + compute_slots_per_rank: int, + repeated_count: int = 1, + ) -> None: + if expert_count <= 0 or expert_count % manager.ep_size != 0: + raise ValueError( + "Per-layer dynamic EPLB requires the expert count to be positive " + f"and divisible by EP size; got E={expert_count}, EP={manager.ep_size}." + ) + resident = expert_count // manager.ep_size + compute = int(compute_slots_per_rank) + if compute <= resident: + raise ValueError( + "Per-layer dynamic EPLB requires at least one helper slot per rank; " + f"got H={resident}, M={compute}." + ) + if top_k <= 0 or repeated_count <= 0: + raise ValueError("top_k and repeated_count must be positive") + + self.manager = manager + self.layer_idx = int(layer_idx) + self.expert_count = int(expert_count) + self.top_k = int(top_k) + self.resident_experts_per_rank = resident + self.compute_slots_per_rank = compute + self.helper_slots_per_rank = compute - resident + self.repeated_count = int(repeated_count) + + self._arena: Optional[Any] = None + # The arena's READY terminal view, resolved once the runtime is bound. + self._ready_flags: Optional[torch.Tensor] = None + self._scheduler_group: Optional[Any] = None + self._bound_backend_ref: Optional[weakref.ReferenceType[Any]] = None + self._pending_replica_plan: Optional[Any] = None + + def get_layer_idx(self) -> int: + return self.layer_idx + + def get_repeat_count(self) -> int: + return self.repeated_count + + def get_load_expert_ids(self) -> list[int]: + start = self.manager.ep_rank * self.resident_experts_per_rank + return list(range(start, start + self.resident_experts_per_rank)) + + def is_static_routing(self) -> bool: + return False + + def is_dynamic_routing(self) -> bool: + """Keep the legacy host-migration path disabled. + + The common route seam still calls :meth:`route` whenever a layer + adapter exists. The legacy meaning of ``is_dynamic_routing`` is + broader: it enables CPU statistics and host-weight migration, neither + of which belongs to the per-layer dynamic CUDA planner. + """ + return False + + def need_load_shared_weights(self) -> bool: + return False + + def set_initial_weight_assignments( + self, initial_weight_assignments: Optional[Sequence[int]] + ) -> None: + if initial_weight_assignments is None: + return + assignments = tuple(int(expert) for expert in initial_weight_assignments) + canonical = tuple(range(self.expert_count)) + if assignments != canonical: + raise ValueError( + "Per-layer dynamic EPLB requires the canonical initial assignment " + "range(num_experts); custom initial_global_assignments are only " + "valid in standard mode." + ) + + def start_set_cpu_stage(self) -> None: + """Release the plan of the chunk whose MoE kernel was just enqueued.""" + self.finish_replica_plan() + + # The other hooks of the standard EPLB layer interface are intentionally + # empty: HALO-Q produces the routing and statistics on the GPU, and no CPU + # stage migrates weights. + def start_wait_gpu_stage(self) -> None: + pass + + def done_wait_gpu_stage(self) -> None: + pass + + def done_set_cpu_stage(self) -> None: + pass + + def update_local_statistic(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_local_ids(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_global_ids(self, *args, **kwargs) -> None: + pass + + def update_statistic_with_gathered_statistic(self, *args, **kwargs) -> None: + pass + + def get_local_statistic_tensor(self) -> None: + return None + + def bind_runtime(self, *, arena: Any, scheduler_group: Any) -> None: + """Attach resources built after final weight storage is available.""" + if arena is None or scheduler_group is None: + raise ValueError("arena and scheduler_group are both required") + if self._scheduler_group is not None: + if arena is self._arena and scheduler_group is self._scheduler_group: + return + raise RuntimeError(f"per-layer dynamic EPLB layer {self.layer_idx} is already bound") + if self._arena is not None and arena is not self._arena: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} cannot bind an arena " + "it did not allocate" + ) + + arena_home = getattr(arena, "home_experts", self.resident_experts_per_rank) + arena_helpers = getattr(arena, "helper_slots", self.helper_slots_per_rank) + group_home = getattr(scheduler_group, "home_experts", self.resident_experts_per_rank) + group_helpers = getattr(scheduler_group, "helper_slots", self.helper_slots_per_rank) + if (int(arena_home), int(arena_helpers)) != ( + self.resident_experts_per_rank, + self.helper_slots_per_rank, + ): + raise ValueError("arena H/S geometry does not match the EPLB layer") + if (int(group_home), int(group_helpers)) != ( + self.resident_experts_per_rank, + self.helper_slots_per_rank, + ): + raise ValueError("scheduler H/S geometry does not match the EPLB layer") + + self._arena = arena + self._ready_flags = arena.terminal_flags_tensor() + self._scheduler_group = scheduler_group + self.manager._register_runtime(self) + + def allocate_arena( + self, + *, + bundle: Any, + ep_process_group: Any, + mapping: Any, + device: int, + layer_idx: Optional[int], + ) -> Any: + """Allocate this layer's H + S live-weight arena from model-owned resources. + + Collective across the EP group: the quant method calls it from its + ``create_weights`` on every rank, before it registers the kernel-plane + Parameters. ``bundle`` is the quant method's live-weight plane spec. + A later call returns the same arena. + """ + if self._arena is not None: + if bundle != self._arena.bundle: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} already allocated " + "its arena for a different live-weight plane spec" + ) + return self._arena + from .arena import allocate_rebalance_arena + + self._arena = allocate_rebalance_arena( + home_experts=self.resident_experts_per_rank, + helper_slots=self.helper_slots_per_rank, + bundle=bundle, + ep_comm=self.manager.get_ep_comm(ep_process_group), + mapping=mapping, + provider=self.manager.create_arena_provider(), + device=int(device), + layer_idx=layer_idx, + ) + return self._arena + + @property + def live_weight_arena(self) -> Optional[Any]: + """The H + S arena that the live kernel planes alias, or ``None``.""" + return self._arena + + def bind_backend(self, backend: Any) -> None: + """Build and own the producer after the backend's live aliases exist.""" + arena = self._arena + if arena is None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no live-weight arena" + ) + if self._scheduler_group is not None: + bound_backend = ( + self._bound_backend_ref() if self._bound_backend_ref is not None else None + ) + if bound_backend is not backend: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} is already " + "bound to a different backend" + ) + return + + backend_ref = weakref.ref(backend) + from .planner import build_rebalance_slot_scheduler_group + + if bool(getattr(backend, "tactic_autotune", False)): + # Synthetic profiling may visit every helper row before a real + # plan selects it. Seed those rows from resident weights once; + # serving generations still replace active rows and publish READY. + home = int(arena.home_experts) + with torch.no_grad(): + for plane in arena.tekit_alias_views: + for helper in range(int(arena.helper_slots)): + plane[home + helper].copy_(plane[helper % home], non_blocking=True) + + # ``layer_idx`` is the backend's model layer index: shared-slot pools + # compare producers across ranks by it. + scheduler_group = build_rebalance_slot_scheduler_group( + arena=arena, + manager=self.manager, + mapping=backend.mapping, + ep_process_group=backend.ep_process_group, + home_experts=self.resident_experts_per_rank, + helper_slots=self.helper_slots_per_rank, + topk=self.top_k, + max_tokens_per_rank=backend.max_tokens_per_rank_capacity, + layer_idx=backend.layer_idx, + ) + self.bind_runtime(arena=arena, scheduler_group=scheduler_group) + self._bound_backend_ref = backend_ref + + @property + def replica_plan(self) -> Optional[Any]: + return self._pending_replica_plan + + def route( + self, + token_selected_experts: torch.Tensor, + offset_by_ep_rank: bool = False, + *, + overlap_fn: Optional[Callable[[], None]] = None, + ) -> torch.Tensor: + """Submit HALO-Q/TMA and return its physical routes without a host wait. + + ``overlap_fn`` runs after the copy stream is submitted but before the + physical routes are lent to the caller, so independent main-stream work + such as shared experts overlaps planning and the helper-weight copy. + """ + del offset_by_ep_rank # HALO-Q consumes global logical expert IDs. + if overlap_fn is not None and not callable(overlap_fn): + raise TypeError("overlap_fn must be callable") + if self._scheduler_group is None or self._arena is None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no bound runtime" + ) + if self._pending_replica_plan is not None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has an unfinished plan" + ) + + if overlap_fn is None: + physical, generation = self._scheduler_group.plan( + token_selected_experts, defer_wait=True + ) + else: + part = self._scheduler_group.plan_schedule(token_selected_experts) + try: + overlap_fn() + except BaseException as error: # noqa: BLE001 + try: + self._scheduler_group.discard_plan(part) + cleanup_note = "the unused rebalance generation was safely released." + except BaseException as cleanup_error: # noqa: BLE001 + cleanup_note = ( + "rebalance generation cleanup failed and the group " + f"requires teardown: {cleanup_error!r}" + ) + _append_note(error, "overlap_fn failed after HALO-Q/TMA enqueue; " + cleanup_note) + raise + physical, generation = self._scheduler_group.plan_finish(part, defer_wait=True) + generation = int(generation) + self._arena.publish_ready_generation(generation) + + plan = MoEReplicaPlan( + resident_experts_per_rank=self.resident_experts_per_rank, + compute_slots_per_rank=self.compute_slots_per_rank, + ready_flags=self._ready_flags, + ready_generation=generation, + reserved_sms=self.manager.auxiliary_sms, + _wait_for_routes=self._scheduler_group.wait_for_routes, + ) + self._pending_replica_plan = plan + return physical + + def require_replica_plan(self) -> Any: + if self._pending_replica_plan is None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has not produced a plan" + ) + return self._pending_replica_plan + + def finish_replica_plan( + self, + plan: Optional[Any] = None, + *, + completion_event: Optional[torch.cuda.Event] = None, + ) -> None: + pending = self._pending_replica_plan + if pending is None: + raise RuntimeError( + f"per-layer dynamic EPLB layer {self.layer_idx} has no plan to finish" + ) + if plan is not None and plan is not pending: + raise RuntimeError("attempted to finish a different replica plan") + assert self._scheduler_group is not None + # Also makes cleanup safe when input preparation failed before the + # backend reached its first route consumer. This is the plan's route + # wait, called on the group directly. + self._scheduler_group.wait_for_routes() + self._scheduler_group.finish(completion_event=completion_event) + self._pending_replica_plan = None + + def _clear_runtime(self) -> None: + self._pending_replica_plan = None + self._scheduler_group = None + self._arena = None + self._ready_flags = None + self._bound_backend_ref = None diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py new file mode 100644 index 000000000000..61d1fb288c09 --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/manager.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Model-level owner of the per-layer dynamic EPLB layers and their GPU resources.""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from ..moe_load_balancer import BaseMoeLoadBalancer +from .layer import PerLayerDynamicMoeLayerLoadBalancer, _append_note + + +class PerLayerDynamicMoeLoadBalancer(BaseMoeLoadBalancer): + """Model-level owner for per-layer dynamic EPLB layers and GPU resources.""" + + def __init__( + self, + *, + ep_rank: int, + ep_size: int, + config: Any, + mapping: Optional[Any] = None, + ) -> None: + # Per-layer dynamic EPLB has no C++ standard-EPLB object and creates no + # MPI split; its EP communicator is resolved by the first MoE layer. + if ep_size <= 1 or not 0 <= ep_rank < ep_size: + raise ValueError(f"invalid EPLB EP rank/size: rank={ep_rank}, size={ep_size}") + num_slots = int(config.num_slots) + if num_slots <= 0 or num_slots % ep_size != 0: + raise ValueError( + f"num_slots must be positive and divisible by EP size; got {num_slots}" + ) + + self.is_shutdown = False + self.ep_rank = int(ep_rank) + self.ep_size = int(ep_size) + self.config = config + self.mapping = mapping + self.num_slots = num_slots + self.num_local_slots = num_slots // ep_size + self.auxiliary_sms = int(getattr(config, "auxiliary_sms", 4)) + if self.auxiliary_sms <= 0: + raise ValueError("auxiliary_sms must be positive") + helper_slot_mode = getattr(config, "helper_slot_mode", "per_layer_resident") + if helper_slot_mode not in ("per_layer_resident", "shared_banks"): + raise ValueError(f"unknown helper_slot_mode {helper_slot_mode!r}") + # Resident helper experts need helper rows that only their layer writes. + self.helper_residency = helper_slot_mode == "per_layer_resident" + + # State read by the shared lifecycle and by the engine. + self._previous_balancer = None + self.single_layer_load_balancers: list[PerLayerDynamicMoeLayerLoadBalancer] = [] + self.next_layer_repeated_count: Optional[int] = None + self.iter_id = 0 + self.in_iter = False + self.enable_statistic = False + self.enable_update_weights = False + self._copy_stream: Optional[torch.cuda.Stream] = None + self._copy_stream_device: Optional[int] = None + self._ep_comm: Optional[Any] = None + self._scheduler_spin_bound: Optional[Any] = None + + # Providers use this manager as their scope owner. Resource references + # stay reachable here, never through Mapping or a process-global registry. + self._arena_providers: list[Any] = [] + self._rebalance_shared_slot_pools: dict[tuple, Any] = {} + self._rebalance_shared_slot_pool_order: list[Any] = [] + + def __del__(self) -> None: + # Teardown is collective and must be driven explicitly by the worker; + # a Python finalizer cannot safely enter EP collectives. + pass + + def __repr__(self) -> str: + return ( + f"PerLayerDynamicMoeLoadBalancer(ep_rank={self.ep_rank}, " + f"ep_size={self.ep_size}, num_slots={self.num_slots}, " + f"auxiliary_sms={self.auxiliary_sms})" + ) + + def is_static_routing(self) -> bool: + return False + + def set_repeated_for_next_layer(self, repeated_count: int) -> None: + if repeated_count <= 0: + raise ValueError("repeat count must be positive") + self.next_layer_repeated_count = int(repeated_count) + + def add_layer( + self, + expert_count: int, + top_k: int, + slot_count_per_rank: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ) -> PerLayerDynamicMoeLayerLoadBalancer: + del aux_stream # The scheduler owns its high-priority copy stream. + self.config.validate_expert_capacity(expert_count) + if int(slot_count_per_rank) != self.num_local_slots: + raise ValueError( + "per-layer dynamic EPLB layer capacity must match the configured " + f"local slot count; got {slot_count_per_rank}, expected " + f"{self.num_local_slots}" + ) + repeated_count = self.next_layer_repeated_count or 1 + self.next_layer_repeated_count = None + layer = PerLayerDynamicMoeLayerLoadBalancer( + self, + layer_idx=len(self.single_layer_load_balancers), + expert_count=expert_count, + top_k=top_k, + compute_slots_per_rank=slot_count_per_rank, + repeated_count=repeated_count, + ) + self.single_layer_load_balancers.append(layer) + return layer + + def get_copy_stream(self, device: int) -> torch.cuda.Stream: + """Return this model's single highest-priority auxiliary stream.""" + device = int(device) + if self._copy_stream is None: + total_sms = int(torch.cuda.get_device_properties(device).multi_processor_count) + if self.auxiliary_sms >= total_sms: + raise ValueError( + "auxiliary_sms must leave at least one SM for the compute stream; " + f"got {self.auxiliary_sms} of {total_sms}" + ) + _, priority = torch.cuda.Stream.priority_range() + self._copy_stream = torch.cuda.Stream(device=device, priority=priority) + self._copy_stream_device = device + elif self._copy_stream_device != device: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span CUDA devices") + return self._copy_stream + + def get_ep_comm(self, process_group: Any) -> Any: + """Return this model's adapter for its resolved EP ProcessGroup.""" + from .arena import _EpComm + + if self._ep_comm is None: + self._ep_comm = _EpComm(process_group) + elif self._ep_comm.process_group is not process_group: + raise RuntimeError("one per-layer dynamic EPLB manager cannot span EP ProcessGroups") + self._ep_comm.validate_geometry(expected_rank=self.ep_rank, expected_world=self.ep_size) + return self._ep_comm + + def get_scheduler_spin_bound(self) -> Any: + """Return the HALO-Q spin bound shared by this model's layers. + + It is registered with the MoE communication timeout guard, which + lengthens it while the engine warms up. + """ + if self._scheduler_spin_bound is None: + from ....autotuner import AutoTuner + from ..moe_comm_timeout_guard import ( + get_moe_comm_timeout_budgets, + register_moe_comm_timeout_proxy, + ) + from .planner import SchedulerSpinBound + + spin_bound = SchedulerSpinBound( + warmup_seconds=get_moe_comm_timeout_budgets().warmup_seconds, + is_tuning=lambda: AutoTuner.get().is_tuning_mode, + ) + register_moe_comm_timeout_proxy(spin_bound) + self._scheduler_spin_bound = spin_bound + return self._scheduler_spin_bound + + def create_arena_provider(self) -> Any: + """Create a per-layer provider scoped to this model manager.""" + from .arena import SharedSlotArenaProvider + + provider = SharedSlotArenaProvider(self, per_layer_helper_sets=self.helper_residency) + self._arena_providers.append(provider) + return provider + + def adopt_arena_provider(self, provider: Any) -> None: + """Retain a transitional provider already created by the allocator.""" + if provider not in self._arena_providers: + self._arena_providers.append(provider) + + def _register_runtime(self, layer: PerLayerDynamicMoeLayerLoadBalancer) -> None: + if layer not in self.single_layer_load_balancers: + raise ValueError("cannot bind a layer owned by another EPLB manager") + assert layer._arena is not None + self.adopt_arena_provider(layer._arena.provider) + + def start_iter(self) -> None: + if self.in_iter: + raise RuntimeError("already in a MoE load-balancer iteration") + self.in_iter = True + + def end_iter(self) -> None: + if not self.in_iter: + raise RuntimeError("not in a MoE load-balancer iteration") + unfinished = [ + layer.layer_idx + for layer in self.single_layer_load_balancers + if layer.replica_plan is not None + ] + if unfinished: + raise RuntimeError(f"unfinished replica plans at iteration end: {unfinished}") + self.in_iter = False + self.iter_id += 1 + + def abort_iter(self, error: BaseException) -> None: + """End a failed iteration after releasing the plans it left unfinished. + + A forward that raises between a layer's route and its post-``run_moe`` + hook leaves that layer's plan pending. + """ + try: + for layer in self.single_layer_load_balancers: + if layer.replica_plan is None: + continue + try: + layer.finish_replica_plan() + except BaseException as cleanup_error: # noqa: BLE001 + _append_note( + error, + "per-layer dynamic EPLB could not release the plan of layer " + f"{layer.layer_idx}: {cleanup_error!r}", + ) + finally: + self.in_iter = False + self.iter_id += 1 + + def register_weight_slots_after_to_cuda(self) -> None: + # Live aliases are registered by the per-layer dynamic arena binder. + pass + + def finalize_model(self) -> None: + # Scheduler groups are built when each layer's final aliases exist. + pass + + def set_warm_up_iter_count(self, iter_count: int) -> None: + del iter_count + + def shutdown(self, local_safe: bool = True) -> bool: + """Collectively close this model's pools in deterministic layer order.""" + if self.is_shutdown: + return True + + # Agree on local readiness before any rank starts destructive teardown; + # a rank is unsafe while any plan is pending. Every MoE layer resolves + # the EP communicator in create_weights, and a rank whose model load + # failed never owns an engine that reaches this method, so the + # communicator exists on every EP rank or on none. + local_safe = bool( + local_safe + and not self.in_iter + and all(layer.replica_plan is None for layer in self.single_layer_load_balancers) + ) + comm = self._ep_comm + if comm is not None: + states = comm.allgather( + { + "local_safe": local_safe, + "pool_count": len(self._rebalance_shared_slot_pool_order), + } + ) + pool_counts = {state["pool_count"] for state in states} + if len(pool_counts) != 1: + raise RuntimeError( + "per-layer dynamic EPLB pool counts differ across the EP communicator" + ) + local_safe = all(state["local_safe"] for state in states) + if not local_safe: + return False + + # Snapshot creation order because each successful close removes itself. + for pool in list(self._rebalance_shared_slot_pool_order): + if not pool.close_collectively(local_safe=True): + return False + + # Pool teardown closes its registered producers. Calling close again is + # deliberately allowed and also covers a transitional group that has + # not yet been registered with a pool. + for layer in self.single_layer_load_balancers: + group = layer._scheduler_group + if group is not None: + group.close() + layer._clear_runtime() + self._arena_providers.clear() + self._ep_comm = None + if self._scheduler_spin_bound is not None: + from ..moe_comm_timeout_guard import unregister_moe_comm_timeout_proxy + + unregister_moe_comm_timeout_proxy(self._scheduler_spin_bound) + self._scheduler_spin_bound = None + self._rebalance_shared_slot_pools.clear() + self._rebalance_shared_slot_pool_order.clear() + self._copy_stream = None + self._copy_stream_device = None + self.is_shutdown = True + return True diff --git a/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py new file mode 100644 index 000000000000..00c429880bbf --- /dev/null +++ b/tensorrt_llm/_torch/moe/fused_moe/per_layer_dynamic_eplb/planner.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HALO-Q and in-switch TMA weight-copy integration on one MAIN stream. + +One highest-priority stream is shared across layers. plan_schedule() enqueues +one kernel that runs HALO-Q and then the TMA copy of its plan on the auxiliary +SMs, before returning to the shared-expert hook. Each CTA of that kernel raises +a device route epoch once its routes are written. plan_finish() can lend routes +before they exist; the first route consumer calls wait_for_routes(), a stream +memory wait on that epoch that holds no SM. MegaMoE's system-acquire READY gate +protects helper reads while the copy finishes concurrently. + +For each rank and generation g: + MAIN: MegaMoE(g) -> ... -> routes(g+1) -> input_ready(g+1) + COPY: wait input_ready(g+1) -> [HALO-Q(g+1) EP exchange -> TMA(g+1)] +Every launch orders COPY after the MAIN work queued before it, which includes +the previous consumer MegaMoE(g), so finish() queues no device work of its own. +All ranks enter the bounded EP exchange after their own previous consumer, so no +next-generation overwrite can precede any previous-generation consumer. No host +collective or auxiliary copy submitter is needed in the serving path. +""" + +from __future__ import annotations + +from typing import Any, Callable, Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +__all__ = [ + "RebalanceSlotSchedulerGroup", + "SchedulerSpinBound", + "build_rebalance_slot_scheduler_group", +] + + +class SchedulerSpinBound: + """Bound of the HALO-Q scheduler's cross-rank waits, chosen per launch. + + HALO-Q exchanges route counts with every EP rank inside the kernel and + traps once a peer is later than this bound. During warmup, one-time JIT + compilation and autotuning can delay a healthy rank by minutes. As a MoE + communication timeout proxy, the bound follows every warmup transition: + warmup gets the warmup timeout, and serving the serving timeout or the + scheduler's native ``SPIN_CYCLES``. Launches while ``is_tuning`` returns + True, such as AutoTuner profiling outside warmup, also get the warmup + timeout. One bound serves every layer of a model. + + Args: + warmup_seconds: The MoE communication warmup timeout. + is_tuning: Returns whether kernels are being profiled. + """ + + name = "per-layer dynamic EPLB HALO-Q" + + def __init__(self, *, warmup_seconds: int, is_tuning: Callable[[], bool]) -> None: + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime import ( + spin_cycles_for_timeout, + ) + + self._spin_cycles_for_timeout = spin_cycles_for_timeout + self._is_tuning = is_tuning + self.warmup_cycles = spin_cycles_for_timeout(warmup_seconds) + self._timeout_cycles = spin_cycles_for_timeout(None) + + def set_timeout_seconds(self, seconds: Optional[int]) -> None: + """Apply a timeout to the launches issued after this call; ``None`` restores the default.""" + self._timeout_cycles = self._spin_cycles_for_timeout(seconds) + + def cycles(self) -> int: + """Return the bound of the next launch.""" + if self._is_tuning(): + return self.warmup_cycles + return self._timeout_cycles + + +class _LiveBankLeaseProvider: + """Authorize ordered reuse, not a claim of completed GPU work on the host. + + release_generation_after() registers the local consumer; the next fused + HALO-Q + TMA launch orders COPY after it on MAIN. prove() is legal right + before that launch. The launch's bounded all-rank count exchange precedes its plan, + and its copy follows the plan in the same CTAs. Therefore a delayed peer + consumer blocks every rank's overwrite without a host-side barrier. + """ + + collective_safe = True + producer_reuse_guarded = True + live_bank_count = 1 + + def __init__(self, group: RebalanceSlotSchedulerGroup) -> None: + self.bound_copy_module = group.broadcaster + self._group = group + self.proved_generation = 0 + + def prove(self, generation: int) -> None: + group = self._group + if ( + generation != group._finished_generation + or group._pending_release is None + or group._pending_release[0] != generation + or group._scheduled_generation != generation + 1 + or group.scheduler_stream is not group.copy_stream + ): + raise RuntimeError("Live-bank reuse lacks the consumer/HALO-Q ordering chain") + self.bound_copy_module.mark_collective_reuse_safe(self, generation) + self.proved_generation = generation + + +class RebalanceSlotSchedulerGroup: + """One layer's scheduler, copy endpoint, and paired consumer generations.""" + + def __init__( + self, + *, + arena: Any, + mapping: Any, + ep_comm: Any, + device: int, + home_experts: int, + helper_slots: int, + topk: int, + max_tokens_per_rank: int, + copy_stream: torch.cuda.Stream, + auxiliary_sms: int, + helper_residency: bool, + spin_bound: SchedulerSpinBound, + layer_idx: Optional[int], + ) -> None: + from cuda.bindings import driver + + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + ) + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.multirank_dist import ( + FabricSymmetricBuffer, + ) + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime import ( + symmetric_buffer_ints, + ) + from tensorrt_llm._torch.moe.kernels.megamoe_scheduler.sami import ( + HierarchicalSamiWeightBroadcast, + ) + + ep_size = int(mapping.moe_ep_size) + ep_rank = int(mapping.moe_ep_rank) + self.layer_idx = layer_idx + self.device = int(device) + self.ep_size = ep_size + self.ep_rank = ep_rank + self.home_experts = int(home_experts) + self.helper_slots = int(helper_slots) + self.topk = int(topk) + self.max_tokens_per_rank = int(max_tokens_per_rank) + self.auxiliary_sms = int(auxiliary_sms) + self._spin_bound = spin_bound + if copy_stream.device.index != self.device: + raise ValueError("copy stream device does not match scheduler device") + self.arena = arena + self.plan_calls = 0 + self._plan_part: Optional[Tuple[Any, int]] = None + self._generation = 0 + self._pending_release: Optional[Tuple[int, Any]] = None + self._finished_generation = 0 + self._scheduled_generation = 0 + self._route_wait_generation = 0 + # Routes of the latest launch, held until MAIN waits for its route epoch. + self._routes_in_flight: Optional[torch.Tensor] = None + # The lent physical routes view and its row count: every layer of a + # forward plans the same number of tokens, so one view serves them. + self._lent_rows = -1 + self._lent_routes: Optional[torch.Tensor] = None + self._execution_stream_handle: Optional[int] = None + self._closed = False + self.scheduler = None + self.symmetric = None + self.broadcaster = None + self.lease = None + + # Keep a partially initialized producer reachable by the arena's + # collective shutdown path. Constructor failures must not attempt a + # rank-local teardown while peers may still be initializing. + ep_comm.validate_geometry(expected_rank=ep_rank, expected_world=ep_size) + self._ep_comm = ep_comm + pool = getattr(arena.provider, "pool", None) + if pool is None or pool.comm is not ep_comm: + raise RuntimeError("rebalance arena and scheduler must use the same EP ProcessGroup") + arena.provider.register_producer(self) + + logical_expert_count = self.home_experts * ep_size + + # Use the EP rank, not the CUDA ordinal. HALO-Q and the TMA copy run in + # the same CTAs, one per auxiliary SM. + cfg = CudaSchedulerConfig( + ep_size=ep_size, + logical_expert_count=logical_expert_count, + extra_slots_per_rank=self.helper_slots, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + local_rank=ep_rank, + ctas=self.auxiliary_sms, + ) + cfg.validate() + self.cfg = cfg + self.scheduler = CudaPhysicalSlotScheduler(cfg, device=f"cuda:{self.device}") + + # All layers share the model manager's highest-priority copy stream. + self.copy_stream = copy_stream + self.scheduler_stream = self.copy_stream + self._stream_handle = int(self.copy_stream.cuda_stream) + self._driver = driver + # The route wait runs once per layer; resolve its driver constants once. + self._wait_value_geq = driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ + self._cuda_success = driver.CUresult.CUDA_SUCCESS + error, cu_device = driver.cuDeviceGet(self.device) + supported = 0 + if error == driver.CUresult.CUDA_SUCCESS: + error, supported = driver.cuDeviceGetAttribute( + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_CAN_USE_64_BIT_STREAM_MEM_OPS, + cu_device, + ) + if error != driver.CUresult.CUDA_SUCCESS or not supported: + raise RuntimeError("MoE rebalance requires 64-bit CUDA stream memory operations") + # CTAs of the fused launches whose routes are written, and its value + # after the latest launch. Only those launches change it. + self._route_epoch = torch.zeros(1, dtype=torch.int64, device=f"cuda:{self.device}") + self._route_epoch_ptr = int(self._route_epoch.data_ptr()) + self._route_target = 0 + + # Every peer address must be connected before the first fused launch; + # HALO-Q performs an all-rank exchange on the COPY stream. + self.symmetric = FabricSymmetricBuffer( + 4 * symmetric_buffer_ints(ep_size, logical_expert_count), self.device + ) + handles = list(ep_comm.allgather(self.symmetric.shareable)) + if len(handles) != ep_size or len(set(handles)) != ep_size: + raise RuntimeError( + f"MoE rebalance requires {ep_size} unique EP fabric handles; " + f"received {len(handles)} handles with {len(set(handles))} unique values." + ) + self.scheduler.connect_peer_bases( + [ + self.symmetric.ptr if r == ep_rank else self.symmetric.import_peer(handles[r]) + for r in range(ep_size) + ] + ) + + # connect_peer_bases() synchronized this device, so its zeroed exchange + # buffer and peer mappings are complete. The barrier extends that to every + # rank before any rank binds its copy and can launch the fused kernel, + # whose HALO-Q exchange writes into its peers' buffers. + ep_comm.barrier() + + # The provider must return the existing arena with identical bundle geometry. + # A second allocation would separate the loader aliases from the copy target. + bundle = arena.bundle + self.broadcaster = HierarchicalSamiWeightBroadcast( + comm=ep_comm, + device=self.device, + bundle=bundle, + helper_count=self.helper_slots, + global_expert_count=logical_expert_count, + arena_provider=arena.provider, + tma_sm_count=self.auxiliary_sms, + tma_warps=7, + tma_route="plan", + helper_residency=bool(helper_residency), + ) + self.broadcaster.bind_scheduler(self.scheduler, self.scheduler_stream, self._route_epoch) + + # The previous consumer precedes the next HALO-Q EP rendezvous. + self.lease = _LiveBankLeaseProvider(self) + self.broadcaster.bind_generation_reuse_authority(self.lease) + + logger.debug("[MegaMoECuteDsl] layer=%s rebalance producer initialized", layer_idx) + + def close(self) -> None: + """Release this layer after the owning pool has drained device work.""" + if self._closed: + return + with torch.cuda.device(self.device): + if self.broadcaster is not None: + self.broadcaster.close() + if self.symmetric is not None: + self.symmetric.close() + self._closed = True + self._execution_stream_handle = None + self._plan_part = None + self._pending_release = None + self._routes_in_flight = None + self._lent_rows = -1 + self._lent_routes = None + self.lease = None + self.broadcaster = None + self.scheduler = None + self.symmetric = None + self.arena = None + + def _check_owner(self) -> None: + if self._closed: + raise RuntimeError("Rebalance scheduler group is closed") + # Runs several times per layer: read the raw device and stream instead + # of building a torch.cuda.Stream, and return early for the bound MAIN. + if torch._C._cuda_getDevice() != self.device: + raise RuntimeError("Rebalance submission requires its bound CUDA device") + handle = torch._C._cuda_getCurrentRawStream(self.device) + if handle == self._execution_stream_handle: + return + if self._execution_stream_handle is not None: + raise RuntimeError( + "Rebalance MAIN stream cannot change; every executor of a model " + "engine must submit on its one execution stream" + ) + stream = torch.cuda.current_stream(self.device) + if handle == self._stream_handle or stream.priority <= self.copy_stream.priority: + raise RuntimeError("MAIN must be distinct from the higher-priority COPY stream") + self._execution_stream_handle = handle + + def _wait_route_epoch(self, stream_handle: int) -> None: + (error,) = self._driver.cuStreamWaitValue64( + stream_handle, self._route_epoch_ptr, self._route_target, self._wait_value_geq + ) + if error != self._cuda_success: + raise RuntimeError(f"cuStreamWaitValue64(rebalance) failed: {error!r}") + + def plan_schedule(self, logical_expert_ids: torch.Tensor) -> Tuple[Any, int]: + """Submit contiguous CUDA int32 routes [tokens, topk] without staging. + + The route producer must allocate and write this tensor on MAIN and must + not overwrite it until HALO-Q finishes. The group holds the tensor until + wait_for_routes() orders MAIN after the route epoch. Every HALO-Q CTA + reads its routes before it raises that epoch, so the caching allocator + can only reuse the block for MAIN work queued after that wait. This needs + no COPY-stream allocator event, which every later allocation would poll + while COPY runs behind MAIN. Empty ranks still submit and join the EP + exchange. + """ + self._check_owner() + if self._plan_part is not None or self._generation != self._finished_generation: + raise RuntimeError("Previous rebalance plan has no paired finish") + if self._routes_in_flight is not None: + # Only a failed submission leaves routes that MAIN never waited for. + raise RuntimeError("Routes of an earlier HALO-Q launch were never ordered on MAIN") + if ( + not isinstance(logical_expert_ids, torch.Tensor) + or logical_expert_ids.ndim != 2 + or logical_expert_ids.shape[1] != self.topk + or logical_expert_ids.dtype != torch.int32 + or not logical_expert_ids.is_contiguous() + # -1 for a CPU tensor; cheaper than building and comparing torch.device objects. + or logical_expert_ids.get_device() != self.device + ): + raise ValueError("Routes must be contiguous CUDA int32 [tokens, topk] on this device") + num_tokens = int(logical_expert_ids.shape[0]) + if num_tokens > self.max_tokens_per_rank: + raise RuntimeError(f"Route chunk {num_tokens} exceeds {self.max_tokens_per_rank}") + + outputs = self.scheduler.outputs + self._scheduled_generation = self._generation + 1 + if self._pending_release is not None: + self.lease.prove(self._pending_release[0]) + self._pending_release = None + route_target = self._route_target + self.auxiliary_sms + self._routes_in_flight = logical_expert_ids + # The launch follows the MAIN work queued so far, which wrote the + # routes; no host synchronization is needed. + generation = self.broadcaster.submit_fused( + outputs, + logical_expert_ids, + route_target, + spin_cycles=self._spin_bound.cycles(), + ) + if generation != self._scheduled_generation: + raise RuntimeError("TMA and HALO-Q generation mismatch") + self._route_target = route_target + self._generation = generation + part = (outputs, num_tokens) + self._plan_part = part + return part + + def plan_finish( + self, part: Tuple[Any, int], *, defer_wait: bool = False + ) -> Tuple[torch.Tensor, int]: + """Lend routes until finish(); optionally defer their device dependency. + + It queues no device work itself: plan_schedule() checked the owner of + this pending plan, and the route wait checks it again. + """ + if part is not self._plan_part or part is None: + raise RuntimeError("plan_finish requires this group's pending plan") + outputs, num_tokens = part + if not defer_wait: + self.wait_for_routes() + physical = self._lent_routes + if num_tokens != self._lent_rows: + # A view of the scheduler's fixed output buffer; slicing it is + # the cost, so reuse it while the token count repeats. + physical = self._lent_routes = outputs.physical_slot_ids[:num_tokens] + self._lent_rows = num_tokens + self.plan_calls += 1 + self._plan_part = None + return physical, self._generation + + def wait_for_routes(self) -> None: + """Order the first MAIN route read after HALO-Q, without a host wait. + + The first call of a generation checks the owner and queues the stream + wait. Later calls of that generation return at once: their consumers + are already ordered after it on MAIN. + """ + generation = self._generation + if generation > 0 and generation == self._route_wait_generation and not self._closed: + return + self._check_owner() + if generation <= 0: + raise RuntimeError("No scheduler generation is available") + self._wait_route_epoch(self._execution_stream_handle) + self._route_wait_generation = generation + # MAIN work queued from here on runs after every HALO-Q read of the routes. + self._routes_in_flight = None + + def plan( + self, logical_expert_ids: torch.Tensor, *, defer_wait: bool = False + ) -> Tuple[torch.Tensor, int]: + return self.plan_finish(self.plan_schedule(logical_expert_ids), defer_wait=defer_wait) + + def discard_plan(self, part: Tuple[Any, int]) -> None: + """Release an enqueued generation that has no MegaMoE consumer.""" + self.plan_finish(part, defer_wait=False) + self.finish() + + def finish(self, completion_event: torch.cuda.Event | None = None) -> None: + """Release this generation to the next-generation producer. + + The consumer is ``completion_event``, or without it the MAIN work + queued so far: the MegaMoE launch that read this generation, if any. + The next fused launch, submitted on MAIN, orders COPY after it. + """ + self._check_owner() + if ( + self._plan_part is not None + or self._generation <= 0 + or self._generation != self.plan_calls + or self._route_wait_generation != self._generation + or self._generation != self._finished_generation + 1 + ): + raise RuntimeError("finish must pair exactly once with a completed plan") + self.broadcaster.release_generation_after(self._generation, completion_event) + self._pending_release = (self._generation, completion_event) + self._finished_generation = self._generation + + +def build_rebalance_slot_scheduler_group( + *, + arena: Any, + manager: Any, + mapping: Any, + ep_process_group: Any, + home_experts: int, + helper_slots: int, + topk: int, + max_tokens_per_rank: int, + layer_idx: Optional[int], +) -> RebalanceSlotSchedulerGroup: + """Build one layer's scheduler/copy group from its manager's resources. + + The manager supplies the model-owned EP communicator, copy stream, + auxiliary SM budget, helper-slot mode and scheduler spin bound; the + communicator is resolved before the stream. + """ + device = int(mapping.local_rank) + return RebalanceSlotSchedulerGroup( + arena=arena, + mapping=mapping, + ep_comm=manager.get_ep_comm(ep_process_group), + home_experts=int(home_experts), + helper_slots=int(helper_slots), + max_tokens_per_rank=int(max_tokens_per_rank), + topk=int(topk), + device=device, + copy_stream=manager.get_copy_stream(device), + auxiliary_sms=int(manager.auxiliary_sms), + helper_residency=bool(manager.helper_residency), + spin_bound=manager.get_scheduler_spin_bound(), + layer_idx=layer_idx, + ) diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 764d45775cfb..cc190601c03e 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -4279,38 +4279,45 @@ def _get_fc2_alpha_input_scale( return entry['w2'][...].reshape([]).to(dtype=torch.float32).reciprocal() # ----------------------------------------------------------------- - # Shape helpers (kernel-side authoritative; the SF flat sizes match - # kernel_fc12.py as cited per-method below). + # Shape helpers delegate to the seven-plane ABI used by the arena. # ----------------------------------------------------------------- @staticmethod - def _ceil_div_int(a: int, b: int) -> int: - return (a + b - 1) // b + def _live_weight_bundle_layout(hidden: int, + intermediate: int, + expand_intermediate: Optional[int] = None): + from ..kernels.megamoe_scheduler.sami.geometry import BundleLayout + + return BundleLayout.create( + hidden=int(hidden), + intermediate=int(intermediate), + expand_intermediate=(None if expand_intermediate is None else + int(expand_intermediate)), + ) @staticmethod - def _round_up_int(a: int, b: int) -> int: - return ((a + b - 1) // b) * b + def _live_weight_plane(bundle, name: str): + try: + return next(plane for plane in bundle.planes if plane.name == name) + except StopIteration as error: + raise RuntimeError( + f"canonical live-weight bundle is missing plane {name!r}" + ) from error @classmethod def fc1_sf_flat_size(cls, intermediate: int, hidden: int) -> int: - """``round_up(expand_intermediate, SfPaddingBlock=128) * - round_up(ceil(hidden / 16), 4)`` -- matches the FC1 weight-SF view - in ``kernel_fc12.py`` (``intermediate_gateup_padded`` / - ``expected_fc1_weight_sf_cols``). - ``expand_intermediate = 2 * intermediate``. - """ - expand_intermediate = intermediate * 2 - return (cls._round_up_int(expand_intermediate, 128) * - cls._round_up_int(cls._ceil_div_int(hidden, 16), 4)) + """Return the canonical FC1 scale-factor elements per slot.""" + bundle = cls._live_weight_bundle_layout(hidden, intermediate) + plane = cls._live_weight_plane(bundle, "mega_fc1_weight_sf") + assert plane.storage is not None + return plane.storage.elements_per_slot @classmethod def fc2_sf_flat_size(cls, hidden: int, intermediate: int) -> int: - """``round_up(hidden, SfPaddingBlock=128) * - round_up(ceil(intermediate / 16), 4)`` -- matches the FC2 weight-SF - view in ``kernel_fc12.py`` (``hidden_padded_fc2`` / - ``expected_fc2_weight_sf_cols``). - """ - return (cls._round_up_int(hidden, 128) * - cls._round_up_int(cls._ceil_div_int(intermediate, 16), 4)) + """Return the canonical FC2 scale-factor elements per slot.""" + bundle = cls._live_weight_bundle_layout(hidden, intermediate) + plane = cls._live_weight_plane(bundle, "mega_fc2_weight_sf") + assert plane.storage is not None + return plane.storage.elements_per_slot # Source-checkpoint tensors kept as 0-element placeholders outside the # load window so the full source set and the mega buffers never coexist @@ -4324,34 +4331,137 @@ def fc2_sf_flat_size(cls, hidden: int, intermediate: int) -> int: # "peak = steady set + ONE layer" bound. Class-level on purpose. _streamed_transient_lock = threading.Lock() + def live_weight_plane_spec(self, module: torch.nn.Module): + """Return the canonical seven-plane ABI for this quantized module.""" + return self._live_weight_bundle_layout( + hidden=int(module.hidden_size), + intermediate=int(module.intermediate_size_per_partition), + expand_intermediate=int( + module.expand_intermediate_size_per_partition), + ) + + def assert_live_weight_aliases(self, module: torch.nn.Module, + arena) -> None: + """Fail if loading or transformation replaced an arena-backed plane.""" + if self.live_weight_plane_spec(module) != arena.bundle: + raise RuntimeError( + "MoE rebalance arena does not match the quant-method plane spec" + ) + if not (len(arena.bundle.planes) == len(arena.local_plane_views) == len( + arena.tekit_alias_views)): + raise RuntimeError( + "MoE rebalance live-plane collections differ in length") + for plane, arena_view, alias in zip( + arena.bundle.planes, + arena.local_plane_views, + arena.tekit_alias_views, + ): + current = getattr(module, plane.name, None) + if current is None or bool(getattr(current, "is_meta", False)): + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} is not materialized" + ) + if (tuple(current.shape) != tuple(alias.shape) + or tuple(current.stride()) != tuple(alias.stride()) + or current.dtype != alias.dtype + or current.device != alias.device): + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} changed layout after binding" + ) + pointers = ( + int(current.data_ptr()), + int(alias.data_ptr()), + int(arena_view.data_ptr()), + ) + if pointers[0] != pointers[1] or pointers[1] != pointers[2]: + raise RuntimeError( + f"MoE rebalance live weight plane {plane.name!r} changed storage after binding" + ) + + def _bind_live_weight_planes(self, module: torch.nn.Module, arena) -> None: + """Alias MegaMoE's seven live parameters onto its rebalance arena.""" + expected_bundle = self.live_weight_plane_spec(module) + if expected_bundle != arena.bundle: + raise RuntimeError( + "MoE rebalance arena does not match the quant-method plane spec" + ) + expected_names = tuple(plane.name for plane in expected_bundle.planes) + if tuple(arena.plane_names) != expected_names: + raise RuntimeError( + "MoE rebalance arena plane order is not canonical") + rebound = False + for name, arena_view, alias in zip(arena.plane_names, + arena.local_plane_views, + arena.tekit_alias_views): + old = getattr(module, name, None) + if old is None: + raise RuntimeError( + f"MoE rebalance live weight plane {name!r} was not registered" + ) + if (tuple(alias.shape) != tuple(old.shape) + or tuple(alias.stride()) != tuple(old.stride()) + or alias.dtype != old.dtype): + raise RuntimeError( + f"MoE rebalance arena alias for {name!r} has an incompatible layout" + ) + if (not alias.is_cuda + or int(alias.data_ptr()) != int(arena_view.data_ptr())): + raise RuntimeError( + f"MoE rebalance arena alias for {name!r} is not its allocator view" + ) + if old.is_meta: + setattr( + module, + name, + torch.nn.Parameter(alias, requires_grad=old.requires_grad), + ) + rebound = True + else: + try: + old.data = alias + except RuntimeError as error: + raise RuntimeError( + f"MoE rebalance failed to bind live plane {name!r}: {error}" + ) from error + if rebound: + self.setup_quant_scales(module) + # Arena pages start zeroed. Restore the create-time value of the + # per-expert scalars on every row, so an expert whose checkpoint leaves + # them unset computes as it does without the load balancer. + for name in ("fc31_alpha", "fc2_alpha", "fc1_norm_const"): + plane = getattr(module, name, None) + if plane is not None: + plane.data.fill_(1.0) + self.assert_live_weight_aliases(module, arena) + # ----------------------------------------------------------------- # create_weights: register MegaMoE-format parameters in addition to # the grandparent's standard NVFP4 parameters. # ----------------------------------------------------------------- def create_weights(self, module: torch.nn.Module): - # The MegaMoE NVFP4 weight + SF pipeline hard-codes the gated - # 2x expansion (``expand_intermediate == 2 * intermediate``): - # ``fc1_sf_flat_size`` computes ``round_up(2 * intermediate, 128)`` - # and ``_build_mega_format_buffers`` slices ``w3_w1_weight`` at - # ``[:intermediate, :]`` / ``[intermediate:, :]`` before the - # 16-atom gate/up interleave. A non-2x configuration would - # silently mis-size the registered ``mega_fc1_weight*`` buffers - # and the loader. Fail fast at create time instead of breaking - # inside ``_build_mega_format_buffers``. - if (module.expand_intermediate_size_per_partition - != 2 * module.intermediate_size_per_partition): - raise NotImplementedError( - "NVFP4MegaMoECuteDslMethod currently requires the gated " - "2x expansion (expand_intermediate == 2 * intermediate); " - f"got expand_intermediate=" - f"{module.expand_intermediate_size_per_partition}, " - f"intermediate={module.intermediate_size_per_partition}.") + # Resolve the seven-plane descriptor before the parent materializes any + # source or derived Parameters; a per-layer dynamic EPLB arena is + # allocated from it first. + bundle = self.live_weight_plane_spec(module) # Fail fast: the derived-only EPLB migration path registers no bias # staging. if module.bias and self.need_load_shared_weights(module): raise NotImplementedError( "NVFP4MegaMoECuteDslMethod does not support expert bias " "together with dynamic-EPLB shared weight loading.") + layer_load_balancer = getattr(module, "layer_load_balancer", None) + arena = None + if (layer_load_balancer is not None + and layer_load_balancer.produces_replica_plan): + # Collective over EP: every rank allocates its H + S live-weight + # arena here. + arena = layer_load_balancer.allocate_arena( + bundle=bundle, + ep_process_group=module.ep_process_group, + mapping=module.mapping, + device=int(module.mapping.local_rank), + layer_idx=module.layer_idx, + ) weight_vec_size = torch.iinfo(self.weight_dtype).bits // 4 self.block_scales_vec_size = torch.iinfo( @@ -4382,69 +4492,91 @@ def create_weights(self, module: torch.nn.Module): # init tensors are actually released. self.setup_quant_scales(module) - num_local_slots = module.expert_size_per_partition - hidden = module.hidden_size - intermediate = module.intermediate_size_per_partition - expand_intermediate = module.expand_intermediate_size_per_partition - # NVFP4 packs 2 elements per byte along K (= hidden for fc1, = - # intermediate for fc2), so the natural HF ``(slots, N, K_bytes)`` - # storage already has K as the stride-1 (innermost) dim. The MegaMoE - # CuteDSL kernel reads the weight K-major with K innermost; the - # backend hands it a ``.transpose(1, 2)`` VIEW (NOT ``.contiguous()``) - # of this storage so the logical shape becomes ``(slots, K_bytes, N)`` - # while K stays stride-1 (see ``mega_moe_cute_dsl`` kernel-input prep). - # ``mega_fc1_weight`` keeps the 16-atom gate/up interleave along - # expand_intermediate (the N axis). ``mega_fc2_weight`` is byte- - # equivalent to ``w2_weight``. - mega_fc1_weight = nn.Parameter( - torch.empty(num_local_slots, - expand_intermediate, - hidden // 2, - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc1_weight", mega_fc1_weight) - - mega_fc2_weight = nn.Parameter( - torch.empty(num_local_slots, - hidden, - intermediate // 2, - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc2_weight", mega_fc2_weight) + # The generic NVFP4 source/checkpoint Parameters created above use + # the resident axis H (``expert_size_per_partition``). Only the seven + # tensors consumed directly by the MegaMoE kernel use M = H + S. + num_resident_experts = int(module.expert_size_per_partition) + num_compute_slots = int(module.compute_slot_count) + if num_compute_slots < num_resident_experts: + raise ValueError( + "MegaMoE compute_slot_count must cover every resident expert: " + f"M={num_compute_slots}, H={num_resident_experts}.") + + planes = {plane.name: plane for plane in bundle.planes} + + def storage_metadata(name: str, slots: int): + plane = planes[name] + return plane.storage_metadata(slots) + + def storage_parameter(name: str, factory=torch.empty) -> nn.Parameter: + shape, stride, dtype_name, _, _ = storage_metadata( + name, num_compute_slots) + tensor = factory(shape, dtype=getattr(torch, dtype_name)) + if tuple(tensor.stride()) != stride: + raise RuntimeError( + f"{name} canonical storage descriptor is not contiguous") + return nn.Parameter(tensor, requires_grad=False) + + # The NVFP4 parent owns alpha construction. Derive the resident and + # compute shapes from its canonical planes, preserve the resident + # values, and refresh setup_quant_scales after the final identities are + # installed. + for name in ("fc31_alpha", "fc2_alpha"): + resident = getattr(module, name) + resident_shape, _, dtype_name, _, _ = storage_metadata( + name, num_resident_experts) + target_shape, _, _, _, _ = storage_metadata(name, num_compute_slots) + if (tuple(resident.shape) != resident_shape + or resident.dtype != getattr(torch, dtype_name)): + raise RuntimeError( + f"{name} source layout must be {resident_shape} " + f"{dtype_name}; got {tuple(resident.shape)} " + f"{resident.dtype}.") + if target_shape != resident_shape: + expanded = resident.detach().new_ones(target_shape) + expanded[:num_resident_experts].copy_(resident.detach()) + module.register_parameter( + name, nn.Parameter(expanded, requires_grad=False)) - mega_fc1_weight_sf = nn.Parameter( - torch.empty(num_local_slots, - self.fc1_sf_flat_size(intermediate, hidden), - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc1_weight_sf", mega_fc1_weight_sf) + # Only the four derived weight/SF buffers are newly registered here. + # Their storage descriptors match the loader; the launch path exposes + # the two weight buffers as zero-copy transposed kernel views. + for name in ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + ): + module.register_parameter(name, storage_parameter(name)) - mega_fc2_weight_sf = nn.Parameter( - torch.empty(num_local_slots, - self.fc2_sf_flat_size(hidden, intermediate), - dtype=torch.uint8), - requires_grad=False, - ) - module.register_parameter("mega_fc2_weight_sf", mega_fc2_weight_sf) - - # Per-expert FC1-output (= FC2-input) NVFP4 quantization norm_const. - # The MegaMoE CuteDSL kernel ABI is per-expert ``(num_local_slots,)``. - # This buffer is filled in ``process_weights_after_loading`` from each - # local expert's raw ``w2.input_scale`` as ``1 / w2.input_scale`` and is - # a stable, contiguous, device-local tensor (NOT a stride-0 expand view) - # so the runner's - # ``from_dlpack(...).mark_layout_dynamic(...)`` and the compile cache - # see a normal 1-D fp32 layout. Because the value is genuinely - # per-expert, EPLB shared-load paths also register CPU staging for this - # parameter. - fc1_norm_const = nn.Parameter( - torch.ones(num_local_slots, dtype=torch.float32), - requires_grad=False, - ) - module.register_parameter("fc1_norm_const", fc1_norm_const) + # Per-expert FC1-output (= FC2-input) NVFP4 quantization norm_const is + # initialized to one and populated after loading from w2.input_scale. + # The ones come from the factory: MetaInitMode keeps torch.ones tensors + # real but rejects any later write into a meta tensor. + module.register_parameter( + "fc1_norm_const", storage_parameter("fc1_norm_const", torch.ones)) + + # Alpha Parameters may have been resized or rebound after the parent + # built its view; refresh it against the final seven-plane identities. + self.setup_quant_scales(module) + if arena is not None: + self._bind_live_weight_planes(module, arena) + + def cache_derived_state(self, module: torch.nn.Module) -> None: + super().cache_derived_state(module) + layer_load_balancer = getattr(module, "layer_load_balancer", None) + if (layer_load_balancer is None + or not layer_load_balancer.produces_replica_plan): + return + # Loading and transformation must leave the seven planes on the arena; + # the layer handle then builds its HALO-Q/TMA producer against them. + arena = layer_load_balancer.live_weight_arena + if arena is None: + raise RuntimeError( + "per-layer dynamic EPLB layer has no live-weight arena") + arena.assert_identity() + self.assert_live_weight_aliases(module, arena) + layer_load_balancer.bind_backend(module) def _materialize_source_params(self, module: torch.nn.Module): """Rematerialize this module's streamed source params (full shape) @@ -4698,7 +4830,7 @@ def _rows_covered(set_name: str, param) -> int: def _check_initial_aux_scale_coverage(self, module: torch.nn.Module) -> None: """Reject partially populated NVFP4 auxiliary-scale families.""" - n_slots = module.expert_size_per_partition + n_slots = int(module.expert_size_per_partition) # A whole-checkpoint load is handed every expert's input_scale, because # the weights dict holds the entire checkpoint; a streaming EP load # only ever reads its own rank's experts, so its complete answer is @@ -4768,7 +4900,7 @@ def process_weights_after_loading(self, module: torch.nn.Module): if module.w3_w1_weight.data.numel() == 0: return - n_slots = module.expert_size_per_partition + n_slots = int(module.expert_size_per_partition) coverage = self._streamed_coverage(module) incomplete = {k: v for k, v in coverage.items() if v < n_slots} if incomplete: @@ -4915,7 +5047,12 @@ def _build_fc1_norm_const(self, module: torch.nn.Module) -> None: raw_input_scales, module.initial_local_expert_ids, device=module.fc1_norm_const.device) - module.fc1_norm_const.data.copy_(routed_norm_const) + n_home = int(module.expert_size_per_partition) + if routed_norm_const.numel() != n_home: + raise RuntimeError( + "MegaMoE-CuteDSL fc1_norm_const expected one entry per " + f"resident expert ({n_home}), got {routed_norm_const.numel()}.") + module.fc1_norm_const.data[:n_home].copy_(routed_norm_const) if self.need_load_shared_weights(module): local_shared_load_expert_ids = module.layer_load_balancer.get_load_expert_ids( @@ -5011,11 +5148,24 @@ def _build_mega_format_buffers( passes ``module.w3_w1_weight.data`` etc., the EPLB staging path passes ``module.local_shared_w3_w1_tensors`` etc. """ + bundle = self._live_weight_bundle_layout(hidden, intermediate, + expand_intermediate) + planes = {plane.name: plane for plane in bundle.planes} + + def storage(name: str): + view = planes[name].storage + if view is None: + raise RuntimeError( + f"canonical live-weight plane {name!r} has no storage view") + return view + if intermediate % 16 != 0: raise ValueError( f"MegaMoE NVFP4 FC1 transform requires intermediate % 16 == 0" f" (Fc1GateUpInterleave); got intermediate={intermediate}.") - h_bytes = hidden // 2 + expand_intermediate = bundle.expand_intermediate + assert expand_intermediate is not None + h_bytes = storage("mega_fc1_weight").shape[1] n_pairs = intermediate // 16 # The parent NVFP4 layout stores raw weight tensors as int64 @@ -5043,13 +5193,18 @@ def _build_mega_format_buffers( # 2)`` view ``(slots, hidden//2, expand_intermediate)`` so the kernel # sees K-major with K stride-1 -- WITHOUT materializing a contiguous # copy (which would move K off the innermost axis and corrupt the GEMM). - mega_fc1 = interleaved.view(num_slots, expand_intermediate, - h_bytes).contiguous() + fc1_storage_shape = (num_slots, *storage("mega_fc1_weight").shape) + mega_fc1 = interleaved.view(fc1_storage_shape).contiguous() # ----- FC2 weight: byte-equivalent clone ----- # ``raw_w2`` is ``(slots, hidden, intermediate//2)`` (N, K_bytes) with # K = intermediate//2 innermost; the kernel-input prep transposes the # last two dims to expose K-major as a view (see fc1 note above). + fc2_storage_shape = (num_slots, *storage("mega_fc2_weight").shape) + if tuple(raw_w2.shape) != fc2_storage_shape: + raise ValueError( + "raw FC2 weight does not match canonical storage shape: " + f"got {tuple(raw_w2.shape)}, expected {fc2_storage_shape}") mega_fc2 = raw_w2.detach().clone().contiguous() # ----- FC1 weight SF: same 16-atom interleave + to_blocked ----- @@ -5059,7 +5214,7 @@ def _build_mega_format_buffers( gate_up_interleave_intermediate=intermediate, n_pairs=n_pairs, expand_intermediate=expand_intermediate, - flat_size=self.fc1_sf_flat_size(intermediate, hidden), + flat_size=storage("mega_fc1_weight_sf").elements_per_slot, ) # ----- FC2 weight SF: per-slot to_blocked only (no gate/up) ----- @@ -5069,7 +5224,7 @@ def _build_mega_format_buffers( gate_up_interleave_intermediate=None, n_pairs=None, expand_intermediate=expand_intermediate, - flat_size=self.fc2_sf_flat_size(hidden, intermediate), + flat_size=storage("mega_fc2_weight_sf").elements_per_slot, ) return mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf diff --git a/tensorrt_llm/_torch/moe/fused_moe/routing.py b/tensorrt_llm/_torch/moe/fused_moe/routing.py index 432f93c4557a..7e9823b57122 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/routing.py +++ b/tensorrt_llm/_torch/moe/fused_moe/routing.py @@ -15,6 +15,9 @@ # Global cache for perfect router logits shared across MoE modules _PERFECT_ROUTER_LOGITS_CACHE = {} +# Global cache for perfect router expert-id tables shared across MoE modules, +# keyed by expert geometry, sender rank, token bound and device +_PERFECT_ROUTER_EXPERT_IDS_CACHE = {} def _get_perfect_router_cache_signature(routing_method): @@ -234,6 +237,73 @@ def get_cached_perfect_router_logits(num_tokens: int, num_experts: int, return logits +def get_cached_perfect_router_expert_ids(num_tokens: int, num_experts: int, + experts_per_token: int, + moe_ep_size: int, ep_rank: int, + max_num_tokens: int, + device: torch.device) -> torch.Tensor: + """ + Return perfectly balanced expert ids for ``num_tokens`` local tokens. + + This is the ``ENABLE_PERFECT_ROUTER=expert_ids`` counterpart of + ``get_cached_perfect_router_logits``: instead of synthesizing logits for + the routing method to consume, it replaces the routed expert ids after + routing. It therefore also covers routers whose ids do not come from the + logits (DeepSeek-V4 hash layers) and needs no per-token-count planning. + + The ids follow ``FlatTopKPerfectRouterPlanner._get_balanced_expert_indices``: + slot ``t * experts_per_token + j``, rotated by ``ep_rank``, goes to EP rank + ``slot % moe_ep_size`` and to expert ``(slot // moe_ep_size) % + (num_experts // moe_ep_size)`` on it. A slot's expert depends only on its + own index, so the ids for ``n`` tokens are the leading ``n`` rows of a + longer table. One table of ``max_num_tokens`` rows per geometry, bound and + device is built on the first call and never replaced, so CUDA graphs + captured later keep reading valid memory. Every call returns a view of its + leading rows, with no kernel launch. + + Args: + num_tokens: Number of local tokens on the current EP rank. + num_experts: Total number of experts across all EP ranks. + experts_per_token: Number of experts selected for each token. + moe_ep_size: Number of EP ranks. + ep_rank: Sender EP rank used to rotate the assignment stream. + max_num_tokens: Upper bound on ``num_tokens`` in every call; it sizes + the table. + device: Device of the routed expert ids. + + Returns: + torch.Tensor: Contiguous int32 expert ids of shape + ``[num_tokens, experts_per_token]``. + + Raises: + ValueError: If ``num_tokens`` exceeds ``max_num_tokens``. + RuntimeError: If the table would be built while a CUDA graph is being + captured. + """ + if num_tokens > max_num_tokens: + raise ValueError( + f"ENABLE_PERFECT_ROUTER=expert_ids got {num_tokens} tokens in one " + f"MoE call, more than max_num_tokens ({max_num_tokens}), which " + "sizes its expert-id table") + cache_key = (num_experts, experts_per_token, moe_ep_size, ep_rank, + max_num_tokens, device) + table = _PERFECT_ROUTER_EXPERT_IDS_CACHE.get(cache_key) + if table is None: + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "ENABLE_PERFECT_ROUTER=expert_ids builds its expert-id table on " + "the first MoE call, which must run before CUDA graph capture") + table = FlatTopKPerfectRouterPlanner._get_balanced_expert_indices( + max_num_tokens, num_experts, experts_per_token, moe_ep_size, + ep_rank, device).to(torch.int32) + if table.is_cuda: + # Other streams, such as the MoE chunking aux stream, read the + # table without waiting for the stream that built it. + torch.cuda.current_stream(table.device).synchronize() + _PERFECT_ROUTER_EXPERT_IDS_CACHE[cache_key] = table + return table[:num_tokens] + + # The type of method in top-K routing, for use in torch custom op # Please keep this in sync with the counterpart defined in cpp/tensorrt_llm/kernels/moe/trtllmGen/runner.h class RoutingMethodType(IntEnum): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py index 786802035913..8feb6089feb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/__init__.py @@ -1,24 +1 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Exported MegaMoE KernelClass entry points.""" - -from .kernel_src.blackwell.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( - BlockScaledSwapAbMegaMoeKernel as BlackwellInferenceMegaMoE, -) -from .kernel_src.rubin.inference.local_mega.block_scaled_swap_ab_local_mega_moe_kernel import ( - BlockScaledSwapAbLocalMegaMoeKernel as RubinInferenceLocalMegaMoE, -) -from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( - BlockScaledSwapAbMegaMoeKernel as RubinInferenceMegaMoE, -) -from .kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel_gen_specialized import ( - BlockScaledSwapAbGenphaseMoeKernel as RubinInferenceGenphaseMegaMoE, -) - -__all__ = [ - "BlackwellInferenceMegaMoE", - "RubinInferenceGenphaseMegaMoE", - "RubinInferenceLocalMegaMoE", - "RubinInferenceMegaMoE", -] +"""Integration-ready CuTeDSL MegaMoE sources.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py index 06bc3a8d0426..c3fd0af463b0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/api.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Stable construction API for composable kernel implementations.""" import types @@ -11,6 +8,7 @@ from cutlass.cute.typing import SymInt + RuntimeIntegerType = SymInt StaticIntegerType = int StaticOrRuntimeIntegerType = Union[int, SymInt] @@ -29,15 +27,16 @@ def __init__(self, expected_type) -> None: def _required_type(requirement: Requirement): - return ( - requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement - ) + return requirement.expected_type if isinstance(requirement, OptionalRequirement) else requirement def _matches_type(value, expected_type) -> bool: origin = get_origin(expected_type) if origin in (Union, types.UnionType): - return any(_matches_type(value, candidate) for candidate in get_args(expected_type)) + return any( + _matches_type(value, candidate) + for candidate in get_args(expected_type) + ) return isinstance(value, expected_type) @@ -66,7 +65,9 @@ def validate( if name not in self._values: if isinstance(requirement, OptionalRequirement): continue - raise KeyError(f"{component_name} requires descriptor field {name!r}.") + raise KeyError( + f"{component_name} requires descriptor field {name!r}." + ) expected_type = _required_type(requirement) value = self._values[name] if not _matches_type(value, expected_type): @@ -96,11 +97,13 @@ class KernelComponent(ABC, metaclass=_KernelComponentMeta): @classmethod @abstractmethod - def problem_desc_require(cls) -> Dict[str, Requirement]: ... + def problem_desc_require(cls) -> Dict[str, Requirement]: + ... @classmethod @abstractmethod - def impl_desc_require(cls) -> Dict[str, Requirement]: ... + def impl_desc_require(cls) -> Dict[str, Requirement]: + ... def _validate_desc_inputs( self, @@ -108,10 +111,14 @@ def _validate_desc_inputs( impl_desc: ImplDesc, ) -> None: component_name = type(self).__name__ - overlap = self.problem_desc_require().keys() & self.impl_desc_require().keys() + overlap = ( + self.problem_desc_require().keys() + & self.impl_desc_require().keys() + ) if overlap: raise ValueError( - f"{component_name} requires fields from both descriptors: {sorted(overlap)}." + f"{component_name} requires fields from both descriptors: " + f"{sorted(overlap)}." ) problem_desc.validate( self.problem_desc_require(), @@ -132,26 +139,33 @@ def _validate_component_init(self) -> None: if not hasattr(self, name): if isinstance(requirement, OptionalRequirement): continue - raise RuntimeError(f"{component_name} did not bind required field {name!r}.") + raise RuntimeError( + f"{component_name} did not bind required field {name!r}." + ) expected_type = _required_type(requirement) value = getattr(self, name) if not _matches_type(value, expected_type): raise TypeError( - f"{component_name}.{name} must have type {expected_type}, got {type(value)}." + f"{component_name}.{name} must have type " + f"{expected_type}, got {type(value)}." ) for name, value in vars(self).items(): if isinstance(value, Desc): - raise RuntimeError(f"{component_name}.{name} retains a descriptor.") + raise RuntimeError( + f"{component_name}.{name} retains a descriptor." + ) class KernelClass(KernelComponent): """Top-level host wrapper for one composable kernel implementation.""" @abstractmethod - def name(self) -> str: ... + def name(self) -> str: + ... @abstractmethod - def aot_compile(self): ... + def aot_compile(self): + ... __all__ = [ diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py index 52a7a9daf028..cc1ef51580d3 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Cross-rank communication protocols and implementations.""" + +from ..quant_def import CombineFormat, QuantKind +from .nvlink_domain.token_comm import ( + TokenBackScheduleMode, + TokenBackMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .nvlink_domain.token_comm_deterministic import TokenCommDeterministic +from .token_protocol import TokenSrcMetadata + +__all__ = [ + "CombineFormat", + "QuantKind", + "TokenBackScheduleMode", + "TokenBackMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", + "TokenSrcMetadata", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py index 52a7a9daf028..c72d7918ac70 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/__init__.py @@ -1,2 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""NVLink-domain pointer mapping and token communication.""" + +from ...helpers.software_sync import NvlinkBarrier, SoftwareGridSync +from ...quant_def import QuantKind +from .symmetric_buffer import SymmetricBufferDevice, SymmetricBufferHost +from .token_comm import ( + TokenBackMode, + TokenBackScheduleMode, + TokenCommArgs, + TokenCommNonDeterministic, +) +from .token_comm_deterministic import TokenCommDeterministic + +__all__ = [ + "NvlinkBarrier", + "QuantKind", + "SoftwareGridSync", + "SymmetricBufferDevice", + "SymmetricBufferHost", + "TokenBackMode", + "TokenBackScheduleMode", + "TokenCommArgs", + "TokenCommDeterministic", + "TokenCommNonDeterministic", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py index 4ac013c3b5f0..42c8a7b6c20e 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/symmetric_buffer.py @@ -1,11 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Symmetric-heap peer pointer mapping carried in kernel arguments.""" from dataclasses import dataclass from typing import Any, Optional +from packaging.version import Version + import cutlass import cutlass.cute as cute from cutlass._mlir import ir @@ -15,7 +14,7 @@ from cutlass.base_dsl.typing import get_c_pointers from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import Int32, Int64, dsl_user_op -from packaging.version import Version + try: from cutlass.base_dsl.typing import MLIR_DYNAMIC_INDEX @@ -64,9 +63,7 @@ def __extract_mlir_attributes__(self) -> list: return [ir.DictAttr.get({})] @cute.jit - def map( - self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0) - ) -> Int64: + def map(self, local_address: Int64, destination_rank: Int32, byte_offset: Int64 = Int64(0)) -> Int64: if cutlass.const_expr(self.max_ranks <= _byval_rank_limit): i64_type = ir.Type.parse("i64") offset_pointer = llvm.getelementptr( @@ -89,10 +86,7 @@ def map_pointer(self, pointer, destination_rank: Int32, byte_alignment: Optional if cutlass.const_expr(byte_alignment is None): byte_alignment = pointer.max_alignment return cute.make_ptr( - pointer.dtype, - self.map(pointer.toint(), destination_rank), - pointer.memspace, - assumed_align=byte_alignment, + pointer.dtype, self.map(pointer.toint(), destination_rank), pointer.memspace, assumed_align=byte_alignment ) @@ -119,27 +113,11 @@ def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: pointer_type = ir.Type.parse("!llvm.ptr") struct_type = _byval_struct_type(self.max_ranks) i64_type = ir.Type.parse("i64") - one = arith.constant( - value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip - ) - buffer = llvm.alloca( - res=pointer_type, - elem_type=struct_type, - array_size=one, - alignment=64, - loc=loc, - ip=ip, - ) + one = arith.constant(value=ir.IntegerAttr.get(i64_type, 1), result=i64_type, loc=loc, ip=ip) + buffer = llvm.alloca(res=pointer_type, elem_type=struct_type, array_size=one, alignment=64, loc=loc, ip=ip) for index, offset in enumerate(offsets): slot = llvm.getelementptr( - pointer_type, - buffer, - [], - [index], - i64_type, - no_wrap_flags="None", - loc=loc, - ip=ip, + pointer_type, buffer, [], [index], i64_type, no_wrap_flags="None", loc=loc, ip=ip ) llvm.store(self._as_int64(offset).ir_value(), slot, loc=loc, ip=ip) return SymmetricBufferDevice(buffer, self.max_ranks) @@ -148,12 +126,8 @@ def make_device_object(self, *, loc=None, ip=None) -> SymmetricBufferDevice: vector_type = ir.Type.parse(f"vector<{self.max_ranks}xi64>") vector = llvm.mlir_zero(vector_type, loc=loc, ip=ip) for index, offset in enumerate(offsets): - element_index = arith.constant( - value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip - ) - vector = llvm.insertelement( - vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip - ) + element_index = arith.constant(value=ir.IntegerAttr.get(i32_type, index), result=i32_type, loc=loc, ip=ip) + vector = llvm.insertelement(vector, self._as_int64(offset).ir_value(), element_index, loc=loc, ip=ip) return SymmetricBufferDevice(vector, self.max_ranks) @@ -163,14 +137,8 @@ def __init__(self, argument: SymmetricBufferHost) -> None: self._argument = argument offsets = tuple(argument.offsets) if len(offsets) != int(argument.max_ranks): - raise ValueError( - f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}." - ) - self._fields = ( - Int64(argument.base_address), - *(Int64(offset) for offset in offsets), - Int32(argument.rank), - ) + raise ValueError(f"Expected {int(argument.max_ranks)} peer offsets, got {len(offsets)}.") + self._fields = (Int64(argument.base_address), *(Int64(offset) for offset in offsets), Int32(argument.rank)) def __c_pointers__(self) -> list[Any]: pointers: list[Any] = [] @@ -195,9 +163,7 @@ def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymmetricBufferHos rebuilt = [] for field in self._fields: field_value_count = len(get_mlir_types(field)) - rebuilt.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError(f"Consumed {value_index} MLIR values, got {len(values)}.") diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py index 79fae23df062..2fbb50436b53 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Metadata-push routing and fused non-deterministic token communication.""" import dataclasses @@ -8,6 +5,7 @@ from typing import Callable, ClassVar, Literal, Optional, Tuple, Union import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -19,8 +17,9 @@ from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc from ...helpers.device_workspace import DeviceWorkspace from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix -from ...helpers.flag_batch import make_flag_batch_tracker from ...helpers.iket_compat import iket +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.software_sync import NvlinkBarrier from ...helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, @@ -33,13 +32,13 @@ tma_load_1d, ) from ...helpers.smem_workspace import SmemWorkspace -from ...helpers.software_sync import NvlinkBarrier from ...helpers.utils import ceil_div, round_up from ...quant_def import CombineFormat, QuantKind from ..token_protocol import TokenSrcMetadata from .symmetric_buffer import SymmetricBufferDevice -TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps"] + +TokenBackMode = Literal["epi_warps", "standalone_warps", "reuse_dispatch_warps", "hybrid_warps"] TokenBackScheduleMode = Literal["static", "atomic_counter"] @@ -51,21 +50,37 @@ class TokenCommArgs: under in-kernel top-k reduction where it degenerates to a view of the caller's 2D output. That REDG accumulates, so in that mode the incoming content is the caller's accumulation base -- zero, or a shared-expert result. + + ``fc1_activation_pool`` is the caller-owned ``(pool_capacity, hidden)`` dispatch pool that replaces + TokenComm's own ``fc1_activation_region`` under ``external_dispatch_pool``; ``None`` otherwise. """ activation: cute.Tensor - activation_sf: cute.Tensor + activation_sf: Optional[cute.Tensor] pre_reduced_activation: cute.Tensor pre_reduced_activation_sf: Optional[cute.Tensor] peer_rank_ptr_mapper: SymmetricBufferDevice + fc1_activation_pool: Optional[cute.Tensor] = None - def __extract_mlir_values__(self) -> list: - return [] - def __new_from_mlir_values__(self, values: list) -> "TokenCommArgs": - if values: - raise ValueError(f"TokenCommArgs expected no MLIR values, got {len(values)}.") - return self +@cute.jit +def _worker_coordinates( + warp_start: cutlass.Constexpr[int], + warp_count: cutlass.Constexpr[int], + warp_idx: Int32, + fused_warp_idx: Int32, + linear_cta_idx: Int32, + cta_ordinal: Optional[Int32], +): + """Return the local warp slot and dense worker index for a transfer operation.""" + if cutlass.const_expr(cta_ordinal is not None): + role_warp_idx = warp_idx - Int32(warp_start) + worker_cta = cta_ordinal + else: + # Legacy kernels select four transfer warps before calling the operation. + role_warp_idx = fused_warp_idx + worker_cta = linear_cta_idx + return role_warp_idx, worker_cta * Int32(warp_count) + role_warp_idx @dataclasses.dataclass(frozen=True) @@ -87,16 +102,12 @@ def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": return cls(flat_topk_index=Int32(packed), topk_score=None) scratch = cute.make_rmem_tensor((2,), cutlass.Int32) cute.recast_tensor(scratch, cutlass.Int64)[0] = packed - return cls( - flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1] - ) + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) @cute.jit def _copy_atom(dtype, num_bits_per_copy: int): - return cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy - ) + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) class _MetadataPushRouter(KernelComponent): @@ -114,6 +125,8 @@ class _MetadataPushRouter(KernelComponent): pool_expert_base_region = "nvlink.token_comm.pool_expert_base" token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + # Pool-ordered fp32 per-token FC1 activation scale, pushed next to the metadata like the scores. + fc1_activation_scales_region = "nvlink.token_comm.fc1_activation_per_token_scales" source_expert_base_region = "nvlink.token_comm.source_expert_base" push_destination_base_region = "nvlink.token_comm.push_destination_base" sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" @@ -141,6 +154,8 @@ def problem_desc_require(cls) -> dict[str, type]: "topk": int, "max_tokens_per_rank": int, "apply_topk_at_fc1": bool, + # The FC1 activation was quantized from x / s_t; the router carries s_t to the pool. + "fc1_activation_per_token_scale": OptionalRequirement(bool), } @classmethod @@ -159,6 +174,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.topk = problem_desc["topk"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.fc1_activation_per_token_scale = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.token_padding_block = impl_desc["token_padding_block"] self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] @@ -167,16 +183,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_router_configuration() self.expert_count_padded = round_up(self.expert_count, 4) self.expert_count_with_trash = self.expert_count_padded + 1 - self.router_elements_per_lane, self.router_data_cta_count = ( - self._router_launch_configuration() - ) + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) if self.router_grid_cta_count > self.promised_launchable_sm_count: raise ValueError( - "Router grid exceeds promised_launchable_sm_count; " - "all metadata-push CTAs must be concurrently resident." + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." ) self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) self._router_smem_workspace = self._build_router_smem_workspace() @@ -204,6 +217,11 @@ def _validate_router_configuration(self) -> None: value = getattr(self, field_name) if value <= 0: raise ValueError(f"{field_name} must be positive, got {value}.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}, got {self.max_tokens_per_rank}." + ) if self.expert_count % self.world_size != 0: raise ValueError( f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." @@ -211,9 +229,7 @@ def _validate_router_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("TokenComm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") @property def experts_per_rank(self) -> int: @@ -224,9 +240,7 @@ def worst_case_padded_tokens(self, block: int) -> int: routes_per_source_token = min(self.topk, self.experts_per_rank) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.experts_per_rank, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -240,13 +254,8 @@ def next_power_of_two(value: int) -> int: maximum_data_cta_count = 128 maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count if routed_token_capacity > maximum_supported_tokens: - raise NotImplementedError( - f"The router supports at most {maximum_supported_tokens} routed tokens per rank." - ) - cta_capacity = min( - maximum_cta_capacity, - next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)), - ) + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) data_cta_count = ceil_div(routed_token_capacity, cta_capacity) return elements_per_lane, data_cta_count @@ -256,18 +265,11 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) data_lifetime.register_tensor( - self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,) - ) - data_lifetime.register_tensor( - self.router_data_prefix_region, - cutlass.Int32, - (self.expert_count_with_trash,), - byte_alignment=16, - ) - data_lifetime.register_tensor( - self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,) + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) data_lifetime.register_tensor( self.router_data_sorted_region, (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), @@ -276,10 +278,7 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: ) if self.router_data_cta_count > 1: data_lifetime.register_tensor( - self.router_data_base_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime = overlay.add_lifetime("helper_cta") @@ -291,22 +290,13 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: byte_alignment=16, ) helper_lifetime.register_tensor( - self.router_helper_totals_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_prefix_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_warp_totals_region, - cutlass.Int32, - (self.router_warps_per_cta,), - byte_alignment=16, + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 ) workspace.finalize(max_bytes=self.router_smem_limit_bytes) return workspace @@ -329,34 +319,14 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: stride=(self.expert_count_padded, 1), ) workspace.register( - self.sizes_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) - workspace.register( - self.metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sorted_metadata_region, - cutlass.Int64, - (maximum_routed_tokens,), - buffer_space="local", + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") if self.apply_topk_at_fc1: workspace.register( - self.sorted_scores_region, - cutlass.Float32, - (maximum_routed_tokens,), - buffer_space="local", + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" ) workspace.register( self.token_src_metadata_region, @@ -367,43 +337,26 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.apply_topk_at_fc1: workspace.register( - self.fc1_topk_scores_region, + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + if self.fc1_activation_per_token_scale: + workspace.register( + self.fc1_activation_scales_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared", ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") workspace.register( - self.pool_expert_base_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - ) - workspace.register( - self.source_expert_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", - ) - workspace.register( - self.push_destination_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.sorted_metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.push_table_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") if self.router_data_cta_count > 1: workspace.register( self.router_size_counter_region, @@ -413,18 +366,10 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: reset="tail_reset", ) workspace.register( - self.router_histogram_done_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) workspace.register( - self.source_base_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) def __extract_mlir_values__(self) -> list: @@ -446,14 +391,18 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace: DeviceWorkspace, stream: cuda.CUstream, + activation_scales: Optional[cute.Tensor] = None, # (max_tokens_per_rank,) fp32, local ) -> None: """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(self.fc1_activation_per_token_scale != (activation_scales is not None)): + raise ValueError("The router carries activation_scales iff fc1_activation_per_token_scale is enabled.") peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() self._router_kernel( topk_indices, topk_scores, + activation_scales, local_rank, local_workspace, shared_workspace, @@ -471,6 +420,7 @@ def _router_kernel( self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], + activation_scales: Optional[cute.Tensor], local_rank: Int32, local_workspace: cute.Pointer, shared_workspace: cute.Pointer, @@ -499,13 +449,14 @@ def _router_kernel( self._router_grid_thread_idx = grid_thread_idx self._router_warp_idx = warp_idx self._router_lane_idx = lane_idx + self._router_activation_scales = activation_scales if cutlass.const_expr(self.router_data_cta_count == 1): self._router_single_cta(topk_indices, topk_scores, smem_base) else: self._router_multiple_ctas(topk_indices, topk_scores, smem_base) if linear_cta_idx < Int32(self.router_push_cta_count): - self._router_push_metadata() + self._router_push_metadata(smem_base) device_workspace.remove_device_members() self._device_workspace = None @@ -516,6 +467,7 @@ def _router_kernel( self._router_grid_thread_idx = None self._router_warp_idx = None self._router_lane_idx = None + self._router_activation_scales = None @cute.jit def _router_single_cta( @@ -524,16 +476,10 @@ def _router_single_cta( if self._router_linear_cta_idx < Int32(self.router_data_cta_count): block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) for zero_round in cutlass.range_constexpr(zero_round_count): @@ -574,12 +520,7 @@ def _router_single_cta( iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -588,10 +529,7 @@ def _router_single_cta( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -605,16 +543,10 @@ def _router_multiple_ctas( block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) @@ -635,17 +567,12 @@ def _router_multiple_ctas( expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if expert < Int32(self.expert_count_padded): dump_base[expert] = Int32( - cute.arch.atomic_add( - size_counter + expert, histogram[expert], sem="relaxed", scope="gpu" - ) + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") ) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.router_histogram_done_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" ) total_valid_routes = smem_exclusive_prefix( @@ -663,12 +590,7 @@ def _router_multiple_ctas( iket.range_pop() iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -676,9 +598,7 @@ def _router_multiple_ctas( iket.range_push("router.wait_source_base") source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) if self._router_thread_idx == Int32(0): - while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32( - 1 - ): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): nanosleep(150) cute.arch.sync_threads() iket.range_pop() @@ -694,10 +614,7 @@ def _router_multiple_ctas( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -714,9 +631,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) size_counter = self._device_workspace.tensor(self.router_size_counter_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) @@ -755,10 +670,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.source_base_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" ) publish_sizes() iket.range_pop() @@ -768,7 +680,26 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: iket.range_pop() @cute.jit - def _router_push_metadata(self) -> None: + def _expert_of_route(self, source_expert_base: cute.Tensor, route: Int32) -> Int32: + """Expert whose segment of this rank's sorted route array contains ``route``. + + ``source_expert_base`` is the exclusive prefix of this rank's per-expert route counts (non-decreasing), so the + answer is the largest expert with ``source_expert_base[expert] <= route`` (upper_bound - 1). A zero-size expert + shares its base with its successor and is therefore never returned. Binary search over the padded expert + count: 256 experts -> 8 steps, each one smem load. + """ + search_steps = (self.expert_count_padded - 1).bit_length() + expert = Int32(0) + for step_idx in cutlass.range_constexpr(search_steps): + step = Int32(1 << (search_steps - 1 - step_idx)) + probe = expert + step + if probe < Int32(self.expert_count_padded): + if source_expert_base[probe] <= route: + expert = probe + return expert + + @cute.jit + def _router_push_metadata(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) @@ -781,45 +712,73 @@ def _router_push_metadata(self) -> None: nanosleep(150) cute.arch.sync_threads() - global_expert = ( - self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx - ) - if global_expert < Int32(self.expert_count): - sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) - source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) - push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) - route_count = sizes_by_rank[self._router_local_rank, global_expert] - source_begin = source_expert_base[global_expert] - destination_begin = push_destination_base[global_expert] - destination_rank = global_expert // Int32(self.experts_per_rank) - peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) - source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) - destination_metadata_address = ( - self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset - ) - if cutlass.const_expr(self.apply_topk_at_fc1): - source_scores = self._device_workspace.ptr(self.sorted_scores_region) - destination_scores_address = ( - self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + # Partition the push by route index, not by expert: one warp per expert serialises that expert's whole route + # run (32 routes per latency-bound remote-store round), so under skewed routing the hottest expert's warp runs + # for as long as the whole kernel while the other warps idle. Every push warp owns an equal, 32-aligned slice + # of this rank's flat sorted route array and finds the owning expert with a binary search over + # source_expert_base (exclusive prefix, non-decreasing; the largest expert with base <= route wins, which skips + # zero-size experts). Completion signalling is unchanged: per CTA (below), not per expert. + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base_gmem = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base_gmem = self._device_workspace.tensor(self.push_destination_base_region) + # Stage both tables in smem for the per-lane lookups below (the DATA-phase prefix/histogram regions are dead + # once sorted_metadata_ready == data_cta_count, which this CTA already waited for); the binary search would + # otherwise issue 8 dependent global loads per route. + source_expert_base = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + push_destination_base = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + table_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for table_round in cutlass.range_constexpr(table_rounds): + table_expert = Int32(table_round * block_thread_count) + self._router_thread_idx + if table_expert < Int32(self.expert_count_padded): + source_expert_base[table_expert] = source_expert_base_gmem[table_expert] + push_destination_base[table_expert] = push_destination_base_gmem[table_expert] + cute.arch.sync_threads() + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_base = self._device_workspace.ptr(self.token_src_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_base = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + if cutlass.const_expr(self.fc1_activation_per_token_scale): + destination_activation_scales_base = self._device_workspace.ptr(self.fc1_activation_scales_region).toint() + # This rank's routes are one flat array of total_routes records, sorted by expert; every push warp takes an + # equal, 32-aligned slice of it (the last warps may get an empty slice) so the work per warp is total_routes / + # push_warp_count regardless of how skewed the per-expert counts are. + last_expert = Int32(self.expert_count - 1) + total_routes = source_expert_base[last_expert] + sizes_by_rank[self._router_local_rank, last_expert] + push_warp_count = self.router_push_cta_count * self.router_warps_per_cta + push_warp_idx = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + # ceil(total_routes / push_warp_count), rounded up to a multiple of 32 (one record per lane per round). + routes_per_warp = ((total_routes + Int32(push_warp_count - 1)) // Int32(push_warp_count) + Int32(31)) // Int32(32) * Int32(32) + slice_begin = push_warp_idx * routes_per_warp + slice_end = cutlass.min(slice_begin + routes_per_warp, total_routes) + route_round_count = Int32(0) + if slice_begin < total_routes: + route_round_count = (slice_end - slice_begin + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = slice_begin + Int32(route_round) * Int32(32) + self._router_lane_idx + if route < slice_end: + # The record does not carry its expert; the slice may span several experts, so each lane looks up + # the expert whose segment contains its route and maps the in-segment offset to the destination. + expert = self._expert_of_route(source_expert_base, route) + destination_position = push_destination_base[expert] + (route - source_expert_base[expert]) + destination_rank = expert // Int32(self.experts_per_rank) + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + metadata = cute.arch.load(source_metadata + route, cutlass.Int64) + stg_b64( + destination_metadata_base + peer_offset + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, ) - route_round_count = (route_count + Int32(31)) // Int32(32) - for route_round in cutlass.range(route_round_count, unroll=1): - route = Int32(route_round) * Int32(32) + self._router_lane_idx - if route < route_count: - source_position = source_begin + route - destination_position = destination_begin + route - metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) - stg_b64( - destination_metadata_address - + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), - metadata, + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + route, cutlass.Float32) + stg_f32(destination_scores_base + peer_offset + Int64(destination_position) * Int64(4), score) + if cutlass.const_expr(self.fc1_activation_per_token_scale): + # Gather the source token's scale into the same pool slot as its metadata. + source_token = TokenSrcMetadata.unpack(metadata).src_token + activation_scale = self._router_activation_scales[source_token] + stg_f32( + destination_activation_scales_base + peer_offset + Int64(destination_position) * Int64(4), + activation_scale, ) - if cutlass.const_expr(self.apply_topk_at_fc1): - score = cute.arch.load(source_scores + source_position, cutlass.Float32) - stg_f32( - destination_scores_address + Int64(destination_position) * Int64(4), - score, - ) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): @@ -832,10 +791,7 @@ def _router_push_metadata(self) -> None: destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx if destination_rank < Int32(self.world_size): red_add_relaxed_sys_s32( - self._peer_rank_ptr_mapper.map( - metadata_ready_address, destination_rank, Int64(0) - ), - Int32(1), + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) ) @cute.jit @@ -852,8 +808,7 @@ def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable destination_row_address = ( matrix_address + peer_offset - + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) - * Int64(4) + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) ) destination_row = cute.make_ptr( cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 @@ -863,9 +818,7 @@ def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable ) with cute.arch.elect_one(): cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) - cp_reduce_async_bulk_add_u32_s2g( - destination_total, smem_expert_counts.iterator, row_bytes - ) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) cute.arch.cp_async_bulk_commit_group() def finalize() -> None: @@ -880,8 +833,7 @@ def finalize() -> None: destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx if destination_rank < Int32(self.world_size): red_add_relaxed_sys_s32( - self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), - Int32(1), + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) ) return finalize @@ -892,19 +844,11 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) matrix_bytes = self.world_size * self.expert_count_padded * 4 - size_matrix = self._router_smem_workspace.tensor( - self.router_helper_size_matrix_region, smem_base - ) - padded_totals = self._router_smem_workspace.tensor( - self.router_helper_totals_region, smem_base - ) + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) - load_mbarrier = self._router_smem_workspace.ptr( - self.router_helper_load_mbarrier_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) sizes = self._device_workspace.tensor(self.sizes_region) if self._router_thread_idx == Int32(0): @@ -913,9 +857,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) iket.range_push("router.wait_sizes_ready") if self._router_thread_idx == Int32(0): - while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32( - self.world_size - ): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): nanosleep(150) cute.arch.mbarrier_init_fence() cute.arch.sync_threads() @@ -937,8 +879,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: if expert < Int32(self.expert_count_padded): expert_size = sizes[expert] padded_totals[expert] = ( - (expert_size + Int32(self.token_padding_block - 1)) - // Int32(self.token_padding_block) + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) ) * Int32(self.token_padding_block) cute.arch.sync_threads() smem_exclusive_prefix( @@ -955,9 +896,7 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: for expert_round in cutlass.range_constexpr(local_expert_rounds): local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if local_expert < Int32(self.experts_per_rank): - pool_expert_base[local_expert] = ( - prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] - ) + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] cute.arch.mbarrier_wait(load_mbarrier, 0) iket.range_pop() @@ -978,17 +917,12 @@ def _compute_push_tables(self, smem_base: cute.Pointer) -> None: for ring_position in cutlass.range_constexpr(self.world_size): source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) if Int32(ring_position) < local_ring_position: - source_ring_offset = ( - source_ring_offset + size_matrix[source_rank, global_expert] - ) + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] push_destination_base[global_expert] = destination_pool_base + source_ring_offset cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.push_table_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() @@ -1006,62 +940,46 @@ def _load_router_inputs( elements_per_thread = load_round_count * elements_per_vector topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) - topk_vectors = cute.logical_divide( - cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None) - ) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) load_atom = _copy_atom(topk_indices.dtype, 128) expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) if cutlass.const_expr(topk_indices.dtype.width == 64): raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), raw_vectors[None, load_round], ) else: expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), expert_vectors[None, load_round], ) score_registers = None if cutlass.const_expr(self.apply_topk_at_fc1): score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) - scores_flat = cute.make_tensor( - topk_scores.iterator, cute.make_layout((maximum_elements,)) - ) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) score_vectors = cute.logical_divide( cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) ) score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( score_atom, - mark_alignment( - score_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), score_register_vectors[None, load_round], ) @@ -1072,26 +990,21 @@ def _load_router_inputs( if cutlass.const_expr(topk_indices.dtype.width == 64): expert_registers[register_idx] = raw_indices_i32[2 * register_idx] token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) - is_invalid = ( - expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count) - ) | (token_idx >= actual_token_count) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) if is_invalid: expert_registers[register_idx] = Int32(self.expert_count_padded) return expert_registers, score_registers @cute.jit - def _build_histogram( - self, expert_registers: cute.Tensor, histogram: cute.Tensor - ) -> cute.Tensor: + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: register_count = cute.size(expert_registers) within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) for register_idx in cutlass.range_constexpr(register_count): within_expert_indices[register_idx] = Int32( cute.arch.atomic_add( - histogram.iterator + expert_registers[register_idx], - Int32(1), - sem="relaxed", - scope="cta", + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" ) ) cute.arch.sync_threads() @@ -1111,41 +1024,28 @@ def _sort_router_elements( for register_idx in cutlass.range_constexpr(register_count): token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) flat_topk_index = token_idx * Int32(self.topk) + topk_slot - destination = ( - expert_run_starts[expert_registers[register_idx]] - + within_expert_indices[register_idx] - ) + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] if cutlass.const_expr(self.apply_topk_at_fc1): - sorted_elements[destination] = _SortedElement( - flat_topk_index, score_registers[register_idx] - ).pack() + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() else: sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() @cute.jit - def _router_value_coordinate( - self, register_idx: int, topk_index_type: type - ) -> Tuple[Int32, Int32]: + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: elements_per_vector = 128 // topk_index_type.width - tile_span = ( - elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 - ) + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 flat_index = Int32( register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector ) + self._router_grid_thread_idx * Int32(elements_per_vector) return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) @cute.jit - def _dump_contiguous_router_output( - self, sorted_elements: cute.Tensor, total_valid_routes: Int32 - ) -> None: + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: block_thread_count = self.router_warps_per_cta * 32 metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() - dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32( - block_thread_count - ) + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) for dump_round in cutlass.range(dump_round_count, unroll=4): position = Int32(dump_round * block_thread_count) + self._router_thread_idx predicate = Int32(position < total_valid_routes) @@ -1155,11 +1055,7 @@ def _dump_contiguous_router_output( src_token=(element.flat_topk_index // Int32(self.topk)), src_topk=(element.flat_topk_index % Int32(self.topk)), ) - stg_b64( - metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), - metadata.pack(), - predicate, - ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) if cutlass.const_expr(self.apply_topk_at_fc1): stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) @@ -1185,9 +1081,7 @@ def _dump_router_output_by_expert( for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx predicate = Int32(route < run_length) - element = _SortedElement.from_packed( - sorted_elements[predicate * (run_begin + route)] - ) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) output_position = dump_begin + route metadata = TokenSrcMetadata( src_rank=Int32(self._router_local_rank), @@ -1200,22 +1094,14 @@ def _dump_router_output_by_expert( predicate, ) if cutlass.const_expr(self.apply_topk_at_fc1): - stg_f32( - score_address + Int64(output_position) * Int64(4), - element.topk_score, - predicate, - ) + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) @cute.jit - def local_expert_sizes( - self, device_workspace: DeviceWorkspace, local_rank: Int32 - ) -> cute.Tensor: + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: """Return this rank's contiguous expert-size view.""" sizes = device_workspace.tensor(self.sizes_region) expert_begin = local_rank * Int32(self.experts_per_rank) - return cute.make_tensor( - sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,)) - ) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) @property def metadata_ready_target(self) -> int: @@ -1234,23 +1120,17 @@ def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute. return device_workspace.ptr(self.token_src_metadata_region) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): sizes_ready = device_workspace.ptr(self.sizes_ready_region) - while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32( - self.world_size - ): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): nanosleep(sleep_cycles) cute.arch.sync_warp() @cute.jit - def wait_for_metadata_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1271,6 +1151,13 @@ def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[ return None return device_workspace.tensor(self.fc1_topk_scores_region) + @cute.jit + def fc1_activation_per_token_scales_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Pool-ordered (worst_case_token_count,) fp32 view, or None when the feature is off.""" + if cutlass.const_expr(not self.fc1_activation_per_token_scale): + return None + return device_workspace.tensor(self.fc1_activation_scales_region) + class TokenCommNonDeterministic(KernelComponent): """Public fused-kernel token communication component.""" @@ -1301,6 +1188,7 @@ class TokenCommNonDeterministic(KernelComponent): expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 + token_in_sf_group_region = "nvlink.token_comm.main_smem.token_in_sf_group" # nosec B105 token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 @@ -1328,8 +1216,16 @@ def impl_desc_require(cls) -> dict[str, type]: "token_in_flag_batch": int, "token_back_mode": str, "token_back_schedule_mode": str, + "token_back_ready_granularity": OptionalRequirement(str), "reduce_topk_in_kernel": bool, "router_smem_limit_bytes": OptionalRequirement(int), + "comm_cta_count": OptionalRequirement(int), + "comm_warps": OptionalRequirement(tuple), + "token_back_on_compute_ctas": OptionalRequirement(bool), + "comm_stage_depths": OptionalRequirement(tuple), + "token_back_chunk_bytes": OptionalRequirement(int), + "dispatch_sf_group": OptionalRequirement(int), + "token_in_stage_split": OptionalRequirement(bool), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -1340,7 +1236,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.topk = problem_desc["topk"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.hidden_size = problem_desc["hidden_size"] - self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.quant_kind = "bf16" if problem_desc["quant_kind"] == "bf16" else QuantKind(problem_desc["quant_kind"]) + self.activation_quant_kind = None if self.quant_kind == "bf16" else self.quant_kind self.combine_format = problem_desc["combine_format"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] @@ -1352,20 +1249,82 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + comm_cta_count = impl_desc.get("comm_cta_count", 0) + if ( + isinstance(comm_cta_count, bool) + or not isinstance(comm_cta_count, int) + or not 0 <= comm_cta_count < self.promised_launchable_sm_count + ): + raise ValueError("comm_cta_count must be a nonnegative integer leaving at least one compute CTA.") + self.cta_specialized = comm_cta_count > 0 + if self.cta_specialized: + comm_warps = impl_desc.get("comm_warps") + if ( + not isinstance(comm_warps, tuple) + or len(comm_warps) != 2 + or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in comm_warps) + ): + raise ValueError("Dedicated communication CTAs require comm_warps as two positive integers.") + self.dispatch_warp_count, self.combine_warp_count = comm_warps + back_on_compute = impl_desc.get("token_back_on_compute_ctas", False) + self.dispatch_warp_start = 0 + self.combine_warp_start = 8 if back_on_compute else self.dispatch_warp_count + self.dispatch_cta_count = comm_cta_count + self.combine_cta_count = ( + self.promised_launchable_sm_count - comm_cta_count if back_on_compute else comm_cta_count + ) + else: + # FC12 occupies warps 0..7. Legacy primitive callers retain their + # modulo-four numbering by omitting the explicit CTA ordinal. + self.dispatch_warp_start = 8 + self.combine_warp_start = 12 if self.token_back_mode == "standalone_warps" else 8 # nosec B105 + self.dispatch_warp_count = self.transfer_warp_count + self.combine_warp_count = self.transfer_warp_count + self.dispatch_cta_count = self.promised_launchable_sm_count + self.combine_cta_count = self.promised_launchable_sm_count + self.comm_stage_depths = tuple(impl_desc.get("comm_stage_depths", (1, 1))) + self.token_back_chunk_bytes = impl_desc.get("token_back_chunk_bytes", self.standalone_chunk_bytes) + # Tokens a token_in warp batches before writing their scale rows; see _token_in_multi_stage. + self.dispatch_sf_group = impl_desc.get("dispatch_sf_group", 1) + self.token_in_stage_split = impl_desc.get("token_in_stage_split", False) + self._token_in_stage_chunk_count = 2 if self.token_in_stage_split else 1 self._validate_configuration() self._router = _MetadataPushRouter(problem_desc, impl_desc) - self._nvlink_barrier = NvlinkBarrier( - world_size=self.world_size, barrier_id=self.grid_sync_barrier_id - ) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) self._device_workspace = None self._token_comm_args = None self._local_rank = None self._linear_cta_idx = None self._transfer_warp_idx = None self._lane_idx = None + self._dispatch_warp_idx = None + self._combine_warp_idx = None + self._dispatch_worker_idx = None + self._combine_worker_idx = None + + @property + def dispatch_stage_depth(self) -> int: + """Tokens one token_in warp keeps in flight; ``1`` is the single-buffered loop.""" + return self.comm_stage_depths[0] + + @property + def combine_stage_depth(self) -> int: + """Output chunks one token_back warp keeps in flight.""" + return self.comm_stage_depths[1] + + @property + def combine_output_token_bytes(self) -> int: + """Bytes one combine output row occupies on the wire.""" + return self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + + @property + def sf_k_group_count(self) -> int: + """Scale vectors of one token, in groups of four -- the width the destination interleaves at.""" + return self.activation_sf_hidden_padded // 4 def _validate_configuration(self) -> None: positive_fields = ( @@ -1391,50 +1350,74 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("TokenComm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") if self.token_back_schedule_mode not in ("static", "atomic_counter"): raise ValueError( f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." ) + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") + if self.token_back_ready_granularity == "token_tile" and ( # nosec B105 + self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps") + or self.token_back_schedule_mode != "atomic_counter" # nosec B105 + ): + raise ValueError("token_tile return-ready requires standalone_warps or reuse_dispatch_warps with atomic_counter.") if not 1 <= self.token_in_flag_batch <= 32: - raise ValueError( - f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}." - ) + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: - raise ValueError( - "fc2_done_signals_per_token_tile must be positive when token-back is enabled." - ) + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: raise ValueError("In-kernel top-k reduction requires BF16 combine data.") - element_block = self.activation_sf_vector_size * 4 - if self.hidden_size % element_block != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {element_block}." - ) + if self.activation_is_quantized: + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + else: + # The 1D bulk copies move whole token rows, so each row must be whole 16-byte lines. + row_element_block = 128 // int(self.activation_dtype.width) + if self.hidden_size % row_element_block != 0: + raise ValueError(f"The BF16 wire requires hidden_size divisible by {row_element_block}.") + if self.dispatch_stage_depth != 1: + raise NotImplementedError("The BF16 wire requires comm_stage_depths[0]=1.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") + if self.token_in_stage_split: + if self.activation_is_quantized or self.token_back_enabled: + raise NotImplementedError( + "token_in_stage_split requires the BF16 wire with token_back disabled (epi_warps)." + ) + if self.bytes_per_token % (self._token_in_stage_chunk_count * 16) != 0: + raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") + + if self.token_back_enabled and self.combine_stage_depth > 1: + if self.token_back_schedule_mode != "static": # nosec B105 + raise ValueError("Multi-stage token-back requires static scheduling.") + if self.combine_format.is_quantized: + raise ValueError("Multi-stage token-back supports BF16 combine only.") @property def experts_per_rank(self) -> int: return self.expert_count // self.world_size + @property + def activation_is_quantized(self) -> bool: + return self.activation_quant_kind is not None + @property def activation_dtype(self) -> type: - return self.quant_kind.activation_dtype + return cutlass.BFloat16 if self.activation_quant_kind is None else self.activation_quant_kind.activation_dtype @property - def activation_sf_dtype(self) -> type: - return self.quant_kind.sf_dtype + def activation_sf_dtype(self) -> Optional[type]: + return None if self.activation_quant_kind is None else self.activation_quant_kind.sf_dtype @property def activation_sf_vector_size(self) -> int: - return self.quant_kind.sf_vec_size + return 0 if self.activation_quant_kind is None else self.activation_quant_kind.sf_vec_size @property def bytes_per_token(self) -> int: @@ -1442,6 +1425,8 @@ def bytes_per_token(self) -> int: @property def activation_sf_hidden_padded(self) -> int: + if not self.activation_is_quantized: + return 0 valid_hidden = self.hidden_size // self.activation_sf_vector_size elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) return int(round_up(valid_hidden, elements_per_16_bytes)) @@ -1476,10 +1461,7 @@ def worst_case_sf_token_count(self) -> int: @property def max_fc1_ready_slot_count(self) -> int: - return ( - self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) - // self.tokens_per_fc1_ready_slot - ) + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot @property def router_smem_workspace(self) -> SmemWorkspace: @@ -1537,6 +1519,7 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace: DeviceWorkspace, stream: cuda.CUstream, + activation_scales: Optional[cute.Tensor] = None, ) -> None: self._router.launch_router( topk_indices, @@ -1547,18 +1530,15 @@ def launch_router( peer_rank_ptr_mapper_host, device_workspace, stream, + activation_scales, ) @cute.jit - def local_expert_sizes( - self, device_workspace: DeviceWorkspace, local_rank: Int32 - ) -> cute.Tensor: + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: return self._router.local_expert_sizes(device_workspace, local_rank) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) @cute.jit @@ -1569,6 +1549,10 @@ def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.T def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: return self._router.fc1_topk_scores_tensor(device_workspace) + @cute.jit + def fc1_activation_per_token_scales_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_activation_per_token_scales_tensor(device_workspace) + @cute.jit def assign_device_members( self, @@ -1577,18 +1561,40 @@ def assign_device_members( token_comm_args: TokenCommArgs, local_rank: Int32, linear_cta_idx: Int32, + dispatch_cta_ordinal: Optional[Int32] = None, + combine_cta_ordinal: Optional[Int32] = None, ) -> None: + if cutlass.const_expr( + self.cta_specialized and (dispatch_cta_ordinal is None or combine_cta_ordinal is None) + ): + raise ValueError("Specialized token communication requires explicit dispatch and combine CTA ordinals.") self._device_workspace = device_workspace self._token_comm_args = token_comm_args self._local_rank = local_rank self._linear_cta_idx = linear_cta_idx thread_idx, _, _ = cute.arch.thread_idx() - transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) - self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) - self._lane_idx = transfer_thread_idx % Int32(32) - self._nvlink_barrier.assign_device_members( - device_workspace, token_comm_args.peer_rank_ptr_mapper - ) + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + self._transfer_warp_idx = cute.arch.make_warp_uniform( + (thread_idx % Int32(self.transfer_thread_count)) // Int32(32) + ) + self._lane_idx = thread_idx % Int32(32) + self._dispatch_warp_idx, self._dispatch_worker_idx = _worker_coordinates( + self.dispatch_warp_start, + self.dispatch_warp_count, + warp_idx, + self._transfer_warp_idx, + linear_cta_idx, + dispatch_cta_ordinal, + ) + self._combine_warp_idx, self._combine_worker_idx = _worker_coordinates( + self.combine_warp_start, + self.combine_warp_count, + warp_idx, + self._transfer_warp_idx, + linear_cta_idx, + combine_cta_ordinal, + ) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) def remove_device_members(self) -> None: self._nvlink_barrier.remove_device_members() @@ -1598,6 +1604,10 @@ def remove_device_members(self) -> None: self._linear_cta_idx = None self._transfer_warp_idx = None self._lane_idx = None + self._dispatch_warp_idx = None + self._combine_warp_idx = None + self._dispatch_worker_idx = None + self._combine_worker_idx = None def __extract_mlir_values__(self) -> list: return [] @@ -1643,24 +1653,24 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", byte_alignment=128, ) - activation_sf_element_count = ( - self.worst_case_sf_token_count * self.activation_sf_hidden_padded - ) - workspace.register( - self.fc1_activation_sf_region, - self.activation_sf_dtype, - (activation_sf_element_count,), - buffer_space="local", - byte_alignment=128, - ) + if self.activation_is_quantized: + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) if self.token_back_enabled: + fc2_done_slots = ( + self.max_fc1_ready_slot_count + if self.token_back_ready_granularity == "token_tile" # nosec B105 + else self.experts_per_rank + ) workspace.register( - self.fc2_done_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.fc2_done_region, cutlass.Int32, (fc2_done_slots,), buffer_space="local", reset="tail_reset" ) if self.token_back_push_data: fc2_element_count = self.worst_case_token_count * self.hidden_size @@ -1682,11 +1692,7 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 workspace.register( - self.token_back_schedule_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) if not self.reduce_topk_in_kernel: @@ -1709,9 +1715,16 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: ) def register_smem_regions(self, workspace: SmemWorkspace) -> None: - workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + """Legacy colocated placement, preserving its declaration order and reuse overlay.""" + workspace.register_mbarrier( + self.token_in_mbarrier_region, + self.dispatch_warp_count * self.dispatch_stage_depth, + ) if self.token_back_enabled: - workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_mbarrier( + self.token_back_mbarrier_region, + self.combine_warp_count * self.combine_stage_depth, + ) workspace.register_tensor( self.expert_sizes_smem_region, cutlass.Int32, @@ -1719,76 +1732,103 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: byte_alignment=16, ) transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") - token_in_lifetime = transfer_overlay.add_lifetime("token_in") - token_in_lifetime.register_tensor( + self._register_token_in_staging(transfer_overlay.add_lifetime("token_in")) + if not self.token_back_enabled: + return + if self.token_back_mode == "standalone_warps": # nosec B105 + back_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back") + else: + back_overlay = transfer_overlay + self._register_token_back_staging(back_overlay.add_lifetime("token_back")) + + def _register_token_in_staging(self, staging_lifetime) -> None: + """Declare input tensors in the caller-selected lifetime.""" + dispatch_rings = self.dispatch_warp_count + dispatch_depth = self.dispatch_stage_depth + sf_padded = self.activation_sf_hidden_padded + if dispatch_depth == 1: + activation_shape = (dispatch_rings, self.hidden_size // self._token_in_stage_chunk_count) + sf_shape = (dispatch_rings, (self.activation_sf_vector_size, sf_padded)) + sf_stride = (sf_padded, (0, 1)) + else: + activation_shape = (dispatch_rings, dispatch_depth, self.hidden_size) + sf_shape = ( + dispatch_rings, + dispatch_depth, + (self.activation_sf_vector_size, sf_padded), + ) + sf_stride = (dispatch_depth * sf_padded, sf_padded, (0, 1)) + staging_lifetime.register_tensor( self.token_in_activation_smem_region, self.activation_dtype, - (self.transfer_warp_count, self.hidden_size), - byte_alignment=16, - ) - token_in_lifetime.register_tensor( - self.token_in_sf_smem_region, - self.activation_sf_dtype, - ( - self.transfer_warp_count, - (self.activation_sf_vector_size, self.activation_sf_hidden_padded), - ), - stride=(self.activation_sf_hidden_padded, (0, 1)), + activation_shape, byte_alignment=16, ) + if self.activation_is_quantized: + staging_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + sf_shape, + stride=sf_stride, + byte_alignment=16, + ) + if self.activation_is_quantized and self.dispatch_stage_depth > 1: + # One k-group of a token group, laid out the way the destination wants it. Shares the token_in + # lifetime because it is only live between a group's first arrival and its scale store. The + # condition is the depth rather than the group size, because it is the depth that selects the + # grouped loop -- a group of one still stages its scales here. + staging_lifetime.register_tensor( + self.token_in_sf_group_region, + self.activation_sf_dtype, + (dispatch_rings, self.sf_k_group_count, 4 * self.dispatch_sf_group), + byte_alignment=16, + ) + + def _register_token_back_staging(self, staging_lifetime) -> None: + """Declare output tensors in the caller-selected lifetime.""" if not self.token_back_enabled: return - - if self.token_back_mode == "standalone_warps": # nosec B105 - token_back_lifetime = workspace.create_overlay( - "nvlink.token_comm.main_smem.standalone_token_back" - ).add_lifetime("token_back") - else: - token_back_lifetime = transfer_overlay.add_lifetime("token_back") - + combine_rings = self.combine_warp_count + combine_depth = self.combine_stage_depth + sf_padded = self.activation_sf_hidden_padded if self.token_back_mode == "standalone_warps": # nosec B105 - available_bytes_per_warp = self.standalone_chunk_bytes + available_bytes_per_warp = self.token_back_chunk_bytes else: - activation_bytes = self.bytes_per_token - sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 - available_bytes_per_warp = activation_bytes + sf_bytes + sf_bytes = sf_padded * int(self.activation_sf_dtype.width) // 8 if self.activation_is_quantized else 0 + available_bytes_per_warp = self.bytes_per_token + sf_bytes if self.token_back_push_data: - bytes_per_output_token = ( - self.hidden_size * int(self.combine_format.act_dtype.width) // 8 - ) + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 if self.token_back_mode == "standalone_warps": # nosec B105 - chunk_bytes = self.standalone_chunk_bytes + chunk_bytes = self.token_back_chunk_bytes elif available_bytes_per_warp < bytes_per_output_token: chunk_bytes = self.bytes_per_token else: chunk_bytes = bytes_per_output_token - if ( - self.token_back_mode != "standalone_warps" # nosec B105 - and bytes_per_output_token % chunk_bytes != 0 - ): + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 raise ValueError("Token-back data chunk bytes must divide one row.") chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) - token_back_lifetime.register_tensor( + shape = ( + (combine_rings, chunk_elements) + if combine_depth == 1 + else (combine_rings, combine_depth, chunk_elements) + ) + staging_lifetime.register_tensor( self.token_back_activation_smem_region, self.combine_format.act_dtype, - (self.transfer_warp_count, chunk_elements), + shape, byte_alignment=16, ) if self.token_back_push_sf: - sf_row_bytes = ( - self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 - ) + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 if sf_row_bytes > available_bytes_per_warp: raise ValueError("Token-back scale row exceeds its per-warp stage.") - token_back_lifetime.register_tensor( + combine_sf = self.combine_sf_hidden_padded + staging_lifetime.register_tensor( self.token_back_sf_smem_region, self.combine_format.scale_dtype, - ( - self.transfer_warp_count, - (self.combine_format.scale_block, self.combine_sf_hidden_padded), - ), - stride=(self.combine_sf_hidden_padded, (0, 1)), + (combine_rings, (self.combine_format.scale_block, combine_sf)), + stride=(combine_sf, (0, 1)), byte_alignment=16, ) @@ -1800,19 +1840,17 @@ def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.P def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: return cute.make_tensor( device_workspace.ptr(self.fc1_activation_region), - cute.make_layout( - (self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1) - ), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), ) @cute.jit - def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.activation_is_quantized): + return None layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) - return cute.make_tensor( - device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1]) - ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) @cute.jit def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: @@ -1827,8 +1865,7 @@ def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[c return cute.make_tensor( device_workspace.ptr(self.fc2_activation_region), cute.make_layout( - (self.worst_case_token_count, 1, self.hidden_size), - stride=(self.hidden_size, self.hidden_size, 1), + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) ), ) @@ -1842,73 +1879,406 @@ def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optiona ( self.worst_case_token_count, 1, - ( - self.combine_format.scale_block, - self.hidden_size // self.combine_format.scale_block, - ), + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), ), stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), ), ) @cute.jit - def pre_reduced_activation_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" if cutlass.const_expr(self.reduce_topk_in_kernel): return None return device_workspace.tensor(self.pre_reduced_activation_region) @cute.jit - def pre_reduced_activation_sf_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The scale plane parallel to ``pre_reduced_activation_tensor``; only quantized wire formats carry one.""" if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): return None return device_workspace.tensor(self.pre_reduced_activation_sf_region) @cute.jit - def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: - """Wait for pushed metadata, then pull activation payloads into local pools.""" - transfer_warp_idx = self._transfer_warp_idx - lane_idx = self._lane_idx - global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx - global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + def _token_in_multi_stage( + self, + smem_workspace, + smem_base, + owned_sizes, + pool_expert_bases, + transfer_warp_idx, + lane_idx, + global_warp_idx, + global_warp_count, + ) -> None: + """Pull tokens through a depth-D ring, writing their scale rows a group at a time. + + Where a scale byte lands: ``r * 512 + (t % 32) * 16 + (t // 32) * 4 + k``. One 128-token atom block + spreads over 112 such k-groups, and inside each one the 512 bytes are laid out by token, not in + token order:: + + k-group r: r*512 r*512 + 511 + |<------------------ 512 B ------------------------>| + +---------+---------+---------+ +---------+ + | j = 0 | j = 1 | j = 2 | ... | j = 31 | 16 B per j + +---------+---------+---------+ +---------+ + | t | t | t | t | + 0 32 64 96 4 B per token + + So token ``t`` owns four bytes here, four more 512 B later, 112 times over: written per token that + is 112 scattered stores against the ONE bulk copy its 3584-byte activation takes. The cost is the + store count, not the bytes. + + The four tokens sharing a ``j`` are exactly one contiguous sixteen. Staging ``{j, j+32, j+64, j+96}`` + together and writing that run turns the 112 four-byte stores into 112 sixteen-byte ones, which is why + the worker space strides over ``(block, j)`` pairs rather than tokens. + + Those four stages stay pinned until the run is written, so the ring is consumed a group at a time and + the depth must be a multiple of the group. + """ + group_size = self.dispatch_sf_group + depth = self.dispatch_stage_depth + group_stage_sets = depth // group_size + k_groups = self.sf_k_group_count + # One 128-token atom block is 112 atoms of 512 bytes, and that is also the stride between blocks. + sf_block_bytes = k_groups * 512 - sizes = self._router.sizes_tensor(self._device_workspace) - pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) - token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + sf_group = smem_workspace.tensor(self.token_in_sf_group_region, smem_base) + warp_mbarrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + transfer_warp_idx * Int32(depth) + if lane_idx == Int32(0): + for init_stage in cutlass.range_constexpr(depth): + cute.arch.mbarrier_init(warp_mbarrier + Int32(init_stage), 1) + cute.arch.sync_warp() - iket.range_push("token_in.wait_sizes_ready") - self._router.wait_for_sizes_ready(self._device_workspace) - iket.range_pop() - iket.range_push("token_in.stage_sizes") - owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) - owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) - source_sizes = cute.make_tensor( - sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,)) - ) - copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 - source_size_vectors = cute.zipped_divide( - ( - mark_alignment(source_sizes, 16) - if cutlass.const_expr(copy_elements == 4) - else source_sizes - ), - (copy_elements,), + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf_pointer = self._device_workspace.ptr(self.fc1_activation_sf_region) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = self.bytes_per_token + scatter_atom = _copy_atom( + self.activation_sf_dtype, + 4 * int(self.activation_sf_dtype.width), ) - destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) - size_vector_count = cute.size(destination_size_vectors, mode=[1]) - size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) - size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) - size_copy_registers = cute.make_rmem_tensor( - (copy_elements, size_copy_rounds), cutlass.Int32 + store_atom = _copy_atom( + self.activation_sf_dtype, + 4 * group_size * int(self.activation_sf_dtype.width), ) - transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx - for size_copy_round in cutlass.range_constexpr(size_copy_rounds): - vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + + first_group_in_this_expert = global_warp_idx + expert_group_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase_mask = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + block_count = (expert_token_count + Int32(127)) // Int32(128) + expert_group_end = expert_group_begin + block_count * Int32(32) + group_count = Int32(0) + if first_group_in_this_expert < expert_group_end: + group_count = ( + expert_group_end - first_group_in_this_expert + global_warp_count - Int32(1) + ) // global_warp_count + pool_expert_base = pool_expert_bases[local_expert] + + # One group in flight before the first is consumed. The stage cursor restarts at 0 every expert + # because every group issued here is consumed here, so the ring holds nothing unconsumed at the + # boundary; the per-stage mbarrier phases do NOT restart, since a stage's parity is the count of + # waits it has served. + if group_count > Int32(0): + self._issue_token_group( + smem_workspace, + smem_base, + Int32(0), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + Int32(0), + transfer_warp_idx, + lane_idx, + ) + + stage_base = Int32(0) + for group_round in cutlass.range(group_count, unroll=1): + next_stage_base = stage_base + Int32(group_size) + # Both arms assign: a value carried out of a staged conditional has to be written on + # every path, so the else is required rather than redundant. + if next_stage_base >= Int32(depth): + next_stage_base = next_stage_base - Int32(depth) + else: + next_stage_base = next_stage_base + if cutlass.const_expr(group_stage_sets > 1): + # Spare stages exist, so the next group's reads run under this group's drain. + if group_round + Int32(1) < group_count: + self._issue_token_group( + smem_workspace, + smem_base, + group_round + Int32(1), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + next_stage_base, + transfer_warp_idx, + lane_idx, + ) + + group_in_this_expert = first_group_in_this_expert + group_round * global_warp_count - expert_group_begin + block_idx = group_in_this_expert // Int32(32) + column_in_block = group_in_this_expert - block_idx * Int32(32) + + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + stage_idx = stage_base + Int32(token_in_group) + cute.arch.mbarrier_wait( + warp_mbarrier + stage_idx, + (pull_phase_mask >> stage_idx) & Int32(1), + ) + pull_phase_mask = pull_phase_mask ^ (Int32(1) << stage_idx) + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_expert_base + token_in_expert) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + if lane_idx == Int32(0): + cp_async_bulk_s2g( + destination_activation, + token_in_activation[transfer_warp_idx, stage_idx, None].iterator, + Int32(activation_bytes), + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + # Transpose this token's contiguous scale row into the group buffer's destination + # order. These are SMEM stores; the global ones happen once per group below. + stage_sf = token_in_sf[transfer_warp_idx, stage_idx, (None, None)] + stage_sf_vectors = cute.zipped_divide(cute.slice_(stage_sf, (0, None)), (4,)) + for scatter_round in cutlass.range_constexpr(ceil_div(k_groups, 32)): + k_group_idx = Int32(scatter_round * 32) + lane_idx + if k_group_idx < Int32(k_groups): + cute.copy( + scatter_atom, + stage_sf_vectors[None, k_group_idx], + cute.zipped_divide( + sf_group[transfer_warp_idx, k_group_idx, None], + (4,), + )[None, Int32(token_in_group)], + ) + else: + pull_phase_mask = pull_phase_mask + cute.arch.sync_warp() + + # The activation rows have to be visible before the flags, and the flags carry release + # semantics that cover the scale stores, so one drain covers the whole group. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_warp() + + sf_block_byte_base = Int64(expert_sf_begin + block_idx * Int32(128)) // Int64(128) * Int64( + sf_block_bytes + ) + Int64(column_in_block) * Int64(4 * group_size) + for store_round in cutlass.range_constexpr(ceil_div(k_groups, 32)): + k_group_idx = Int32(store_round * 32) + lane_idx + if k_group_idx < Int32(k_groups): + destination_sf = cute.make_tensor( + cute.make_ptr( + self.activation_sf_dtype, + fc1_activation_sf_pointer.toint() + + sf_block_byte_base + + Int64(k_group_idx) * Int64(512), + AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout(4 * group_size), + ) + cute.copy( + store_atom, + sf_group[transfer_warp_idx, k_group_idx, None], + destination_sf, + ) + cute.arch.sync_warp() + + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( + self.tokens_per_fc1_ready_slot + ) + flag_tracker = flag_tracker.accumulate( + Int32(0), + self.token_in_flag_batch, + (fc1_ready_counter + ready_slot_idx).toint(), + ) + else: + flag_tracker = flag_tracker + cute.arch.sync_warp() + + if cutlass.const_expr(group_stage_sets == 1): + # The ring is exactly one group wide, so the stages are free only now. + if group_round + Int32(1) < group_count: + self._issue_token_group( + smem_workspace, + smem_base, + group_round + Int32(1), + first_group_in_this_expert, + expert_group_begin, + global_warp_count, + expert_token_count, + pool_expert_base, + next_stage_base, + transfer_warp_idx, + lane_idx, + ) + stage_base = next_stage_base + + first_group_in_this_expert = first_group_in_this_expert + group_count * global_warp_count + expert_group_begin = expert_group_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + + @cute.jit + def _issue_token_group( + self, + smem_workspace, + smem_base, + group_round, + first_group_in_this_expert, + expert_group_begin, + worker_count, + expert_token_count, + pool_expert_base, + stage_base, + warp_slot, + lane_idx, + ) -> None: + """Start the remote reads for one token group into stages ``stage_base .. stage_base + G - 1``. + + A group is ``{j, j + 32, j + 64, j + 96}`` of one 128-token atom block, and the unit the worker space + strides over is the ``(block, j)`` pair rather than the token. Only the local staging address is + needed here, and that is pure arithmetic on the static schedule, so the lookahead costs no extra + dependent load; the outbound destination needs ``TokenSrcMetadata`` and stays on the drain side. + + A method rather than a closure because CuTe-DSL will not let a closure defined in one traced region + capture staged values, and that includes ``self``. + """ + group_size = self.dispatch_sf_group + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + warp_slot * Int32( + self.dispatch_stage_depth + ) + metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + activation_bytes = self.bytes_per_token + activation_sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + + group_in_this_expert = first_group_in_this_expert + group_round * worker_count - expert_group_begin + block_idx = group_in_this_expert // Int32(32) + column_in_block = group_in_this_expert - block_idx * Int32(32) + for token_in_group in cutlass.range_constexpr(group_size): + token_in_expert = block_idx * Int32(128) + column_in_block + Int32(32 * token_in_group) + if token_in_expert < expert_token_count: + pool_token_idx = pool_expert_base + token_in_expert + metadata = TokenSrcMetadata.load( + metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + stage_idx = stage_base + Int32(token_in_group) + stage_mbarrier = warp_mbarrier + stage_idx + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx( + stage_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + token_in_activation[warp_slot, stage_idx, None].iterator, + Int64(remote_activation_address), + stage_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + token_in_sf[warp_slot, stage_idx, (None, None)].iterator, + remote_sf_row.iterator, + stage_mbarrier, + Int32(activation_sf_bytes), + ) + cute.arch.sync_warp() + + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._dispatch_warp_idx + dispatch_worker_idx = self._dispatch_worker_idx + lane_idx = self._lane_idx + global_warp_idx = dispatch_worker_idx + global_warp_count = Int32(self.dispatch_cta_count * self.dispatch_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor( + sizes.iterator + owner_expert_begin, + cute.make_layout((self.experts_per_rank,)), + ) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) + # Stride by the role's own thread count, not the class constant: with a split dispatch owns + # twelve warps rather than the fused world's four, and the barrier below counts the same set. + dispatch_thread_count = self.dispatch_warp_count * 32 + size_copy_rounds = ceil_div(size_vector_count, dispatch_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * dispatch_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( size_copy_atom, @@ -1922,7 +2292,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No iket.range_pop() for size_copy_round in cutlass.range_constexpr(size_copy_rounds): - vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + vector_idx = Int32(size_copy_round * dispatch_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( size_copy_atom, @@ -1931,45 +2301,58 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No ) iket.range_push("token_in.size_barrier") token_in_size_barrier = pipeline.NamedBarrier( - barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + barrier_id=self.token_in_size_barrier_id, + num_threads=dispatch_thread_count, ) token_in_size_barrier.arrive_and_wait() iket.range_pop() - if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, ) sizes_ready_barrier.arrive() - iket.range_push("token_in.pull_payload") + if cutlass.const_expr(self.dispatch_stage_depth > 1): + self._token_in_multi_stage( + smem_workspace, + smem_base, + owned_sizes, + pool_expert_bases, + transfer_warp_idx, + lane_idx, + global_warp_idx, + global_warp_count, + ) + return + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) - token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) warp_mbarrier = token_in_mbarriers + transfer_warp_idx warp_activation_stage = token_in_activation[transfer_warp_idx, None] - warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if cutlass.const_expr(self.activation_is_quantized): + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] if lane_idx == Int32(0): cute.arch.mbarrier_init(warp_mbarrier, 1) cute.arch.sync_warp() fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) - fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) - activation_bytes = ( - cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 - ) - activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 - sf_copy_elements = 4 - source_sf_values = cute.slice_(warp_sf_stage, (0, None)) - source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) - sf_copy_atom = _copy_atom( - self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width) - ) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = 0 + if cutlass.const_expr(self.activation_is_quantized): + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) next_dense_token = global_warp_idx expert_valid_begin = Int32(0) - expert_sf_begin = Int32(0) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = Int32(0) expert_ready_slot_begin = Int32(0) pull_phase = Int32(0) flag_tracker = make_flag_batch_tracker( @@ -1986,108 +2369,127 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No expert_valid_end = expert_valid_begin + expert_token_count pull_count = Int32(0) if next_dense_token < expert_valid_end: - pull_count = ( - expert_valid_end - next_dense_token + global_warp_count - Int32(1) - ) // global_warp_count + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count for pull_round in cutlass.range(pull_count, unroll=1): dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_bases[local_expert] + token_in_expert - sf_token_idx = expert_sf_begin + token_in_expert - ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( - self.tokens_per_fc1_ready_slot - ) + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = TokenSrcMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) - ) - peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map( - Int64(0), metadata.src_rank, Int64(0) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) remote_activation_address = ( self._token_comm_args.activation.iterator.toint() + peer_offset + Int64(metadata.src_token) * Int64(self.bytes_per_token) ) - remote_sf = cute.make_tensor( - cute.make_ptr( - self._token_comm_args.activation_sf.dtype, - self._token_comm_args.activation_sf.iterator.toint() + peer_offset, - AddressSpace.gmem, - assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), - ), - self._token_comm_args.activation_sf.layout, - ) - remote_sf_row = remote_sf[Int64(metadata.src_token), None] - - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) - ) - tma_load_1d( - warp_activation_stage.iterator, - Int64(remote_activation_address), - warp_mbarrier, - Int32(activation_bytes), - ) - tma_load_1d( - warp_sf_stage.iterator, - remote_sf_row.iterator, - warp_mbarrier, - Int32(activation_sf_bytes), + if cutlass.const_expr(self.activation_is_quantized): + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, ) - cute.arch.sync_warp() - cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] - destination_activation = cute.make_ptr( - self.activation_dtype, - fc1_activation_pointer.toint() - + Int64(pool_token_idx) * Int64(self.bytes_per_token), - AddressSpace.gmem, - assumed_align=16, - ) - with cute.arch.elect_one(): - cp_async_bulk_s2g( - destination_activation, - warp_activation_stage.iterator, - Int32(activation_bytes), + if cutlass.const_expr(self.token_in_stage_split): + # Chunked pull: the halved stage bounces each token row in two pieces, + # trading one extra NVLink round trip per token for SMEM. + for chunk_idx in cutlass.range_constexpr(self._token_in_stage_chunk_count): + chunk_offset = Int64(chunk_idx) * Int64(activation_bytes) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_bytes)) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address) + chunk_offset, + warp_mbarrier, + Int32(activation_bytes), + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + pull_phase = pull_phase ^ Int32(1) + + chunk_destination = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.bytes_per_token) + + chunk_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(chunk_destination, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + # Drain before the stage is reused by the next chunk (and before the flag). + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + else: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + if cutlass.const_expr(self.activation_is_quantized): + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, ) - cute.arch.sync_warp() - cute.arch.cp_async_bulk_commit_group() + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() - destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] - destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) - destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide( - destination_sf_values, (sf_copy_elements,) - ) - sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) - for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): - sf_vector_idx = Int32(sf_round * 32) + lane_idx - if sf_vector_idx < Int32(sf_vector_count): - cute.copy( - sf_copy_atom, - source_sf_vectors[None, sf_vector_idx], - destination_sf_vectors[None, sf_vector_idx], - ) + if cutlass.const_expr(self.activation_is_quantized): + sf_token_idx = expert_sf_begin + token_in_expert + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) - cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.sync_warp() + if cutlass.const_expr(not self.token_in_stage_split): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() - flag_tracker = flag_tracker.accumulate( - Int32(0), self.token_in_flag_batch, ready_address - ) + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() - pull_phase = pull_phase ^ Int32(1) + if cutlass.const_expr(not self.token_in_stage_split): + pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end - expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) - // Int32(self.sf_padding_block) - ) * Int32(self.sf_padding_block) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32(self.tokens_per_fc1_ready_slot) @@ -2096,13 +2498,11 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No flag_tracker.fire() cute.arch.sync_warp() - iket.range_pop() - if cutlass.const_expr( - self.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 - ): + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 iket.range_push("token_in.transfer_barrier") transfer_lifetime_barrier = pipeline.NamedBarrier( - barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + barrier_id=self.transfer_lifetime_barrier_id, + num_threads=dispatch_thread_count, ) transfer_lifetime_barrier.arrive_and_wait() iket.range_pop() @@ -2114,31 +2514,31 @@ def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> Non sleep_cycles = reference_window - current_window elif current_window > reference_window: sleep_cycles = cutlass.min( - current_window - reference_window, Int32(self.standalone_max_backoff_cycles) + current_window - reference_window, + Int32(self.standalone_max_backoff_cycles), ) if sleep_cycles > Int32(0): nanosleep(sleep_cycles) @cute.jit def _adaptive_pace( - self, average_window: Int32, current_window: Int32, low_window: int, high_window: int + self, + average_window: Int32, + current_window: Int32, + low_window: int, + high_window: int, ) -> Int32: sleep_cycles = Int32(0) if current_window > average_window: sleep_cycles = current_window - average_window - average_window = average_window + ( - (current_window - average_window + Int32(3)) // Int32(4) - ) + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) if sleep_cycles > Int32(high_window): sleep_cycles = Int32(high_window) else: sleep_cycles = average_window - current_window - average_window = average_window - ( - (average_window - current_window + Int32(3)) // Int32(4) - ) + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) # with cute.arch.elect_one(): - # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, - # sleep_cycles) + # cute.printf("avg_window: {}, current_window: {}, sleep_cycles: {}", average_window, current_window, sleep_cycles) if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): nanosleep(sleep_cycles) if average_window > Int32(high_window): @@ -2147,29 +2547,269 @@ def _adaptive_pace( average_window = Int32(low_window) return average_window + @cute.jit + def _issue_combine_read( + self, + smem_workspace, + smem_base, + chunk_round, + stage_index, + first_dense_token, + expert_valid_begin, + pool_expert_base, + worker_count, + warp_slot, + lane_idx, + ) -> None: + """Read chunk ``chunk_round`` of this worker's stream from the local FC2 staging into ``stage_index``. + + Needs no metadata: the local address is arithmetic on the static schedule's ``first + (r // C) * W``, + so running the read ahead costs one integer divide and no extra dependent load. It is the OUTBOUND + destination that needs ``TokenSrcMetadata``, and that stays on the drain side where the wait for this + read has already been paid. + """ + chunk_bytes = self.token_back_chunk_bytes + chunk_count = ceil_div(self.combine_output_token_bytes, chunk_bytes) + stage = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + stage_mbarrier = ( + smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + + warp_slot * Int32(self.combine_stage_depth) + + stage_index + ) + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + + token_round = chunk_round // Int32(chunk_count) + chunk_idx = chunk_round - token_round * Int32(chunk_count) + dense_token_idx = first_dense_token + token_round * worker_count + pool_token_idx = pool_expert_base + (dense_token_idx - expert_valid_begin) + current_chunk_bytes = Int32(chunk_bytes) + remaining_bytes = Int32(self.combine_output_token_bytes) - chunk_idx * Int32(chunk_bytes) + if remaining_bytes < current_chunk_bytes: + current_chunk_bytes = remaining_bytes + else: + current_chunk_bytes = current_chunk_bytes + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(stage_mbarrier, current_chunk_bytes) + tma_load_1d( + stage[warp_slot, stage_index, None].iterator, + fc2_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.combine_output_token_bytes) + + Int64(chunk_idx) * Int64(chunk_bytes), + stage_mbarrier, + current_chunk_bytes, + ) + + @cute.jit + def _token_back_multi_stage( + self, + smem_workspace, + smem_base, + owned_sizes, + transfer_warp_idx, + lane_idx, + global_worker_idx, + global_worker_count, + ) -> None: + """Push output chunks with the local read running the whole ring ahead of the outbound store.""" + depth = self.combine_stage_depth + chunk_bytes = self.token_back_chunk_bytes + chunk_count = ceil_div(self.combine_output_token_bytes, chunk_bytes) + + stage = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_mbarrier = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + transfer_warp_idx * Int32( + depth + ) + if lane_idx == Int32(0): + for init_stage in cutlass.range_constexpr(depth): + cute.arch.mbarrier_init(worker_mbarrier + Int32(init_stage), 1) + cute.arch.sync_warp() + + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, + num_threads=2 * self.transfer_thread_count, + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + + next_dense_token = global_worker_idx + push_phase_mask = Int32(0) + expert_valid_begin = Int32(0) + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + token_count = Int32(0) + if next_dense_token < expert_valid_end: + token_count = ( + expert_valid_end - next_dense_token + global_worker_count - Int32(1) + ) // global_worker_count + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + chunk_total = token_count * Int32(chunk_count) + pool_expert_base = pool_expert_bases[local_expert] + + # The prologue issues with no recycle in front of it, so its stages must be free of the previous + # expert's stores as well as its own. One full read drain buys that, once per expert against + # thousands of chunks. + if chunk_total > Int32(0): + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + for prologue_chunk in cutlass.range_constexpr(depth): + if Int32(prologue_chunk) < chunk_total: + self._issue_combine_read( + smem_workspace, + smem_base, + Int32(prologue_chunk), + Int32(prologue_chunk), + next_dense_token, + expert_valid_begin, + pool_expert_base, + global_worker_count, + transfer_warp_idx, + lane_idx, + ) + cute.arch.sync_warp() + + stage_idx = Int32(0) + for chunk_round in cutlass.range(chunk_total, unroll=1): + token_round = chunk_round // Int32(chunk_count) + chunk_idx = chunk_round - token_round * Int32(chunk_count) + dense_token_idx = next_dense_token + token_round * global_worker_count + pool_token_idx = pool_expert_base + (dense_token_idx - expert_valid_begin) + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + current_chunk_bytes = Int32(chunk_bytes) + remaining_bytes = Int32(self.combine_output_token_bytes) - chunk_idx * Int32(chunk_bytes) + if remaining_bytes < current_chunk_bytes: + current_chunk_bytes = remaining_bytes + else: + current_chunk_bytes = current_chunk_bytes + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + Int64(chunk_idx) * Int64(chunk_bytes), + AddressSpace.gmem, + assumed_align=16, + ) + cute.arch.mbarrier_wait( + worker_mbarrier + stage_idx, + (push_phase_mask >> stage_idx) & Int32(1), + ) + push_phase_mask = push_phase_mask ^ (Int32(1) << stage_idx) + if lane_idx == Int32(0): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, + stage[transfer_warp_idx, stage_idx, None].iterator, + current_chunk_bytes, + ) + else: + cp_async_bulk_s2g( + destination_chunk, + stage[transfer_warp_idx, stage_idx, None].iterator, + current_chunk_bytes, + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + + refill_round = chunk_round + Int32(depth) + if refill_round < chunk_total: + # Every stage is given to the reads, so the stage being refilled is the one whose store + # was issued `depth` chunks ago and nothing newer has to be complete. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + self._issue_combine_read( + smem_workspace, + smem_base, + refill_round, + stage_idx, + next_dense_token, + expert_valid_begin, + pool_expert_base, + global_worker_count, + transfer_warp_idx, + lane_idx, + ) + cute.arch.sync_warp() + stage_idx = stage_idx + Int32(1) + if stage_idx == Int32(depth): + stage_idx = Int32(0) + + next_dense_token = next_dense_token + token_count * global_worker_count + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + + # token_back publishes no flag, but the tail NvlinkBarrier's fence does not drain async-proxy bulk + # copies, so without this a peer could read a half-written row after the rendezvous. + if lane_idx == Int32(0): + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_warp() + @cute.jit def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: """Push completed FC2 data and scale rows to source ranks.""" - transfer_warp_idx = self._transfer_warp_idx + transfer_warp_idx = self._combine_warp_idx + combine_worker_idx = self._combine_worker_idx lane_idx = self._lane_idx if cutlass.const_expr(not self.token_back_enabled): return if cutlass.const_expr( - self.combine_format.is_quantized - and self._token_comm_args.pre_reduced_activation_sf is None + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None ): raise ValueError("Quantized token-back requires a scale destination.") - global_worker_idx = ( - self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx - ) - global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + if cutlass.const_expr(self.cta_specialized): + # A dedicated token-back CTA need not have token-in's SMEM scratch. + # pool_expert_base is built after sizes_ready; acquire metadata_ready too. + self._router.wait_for_sizes_ready(self._device_workspace) + owned_sizes = self._router.local_expert_sizes(self._device_workspace, self._local_rank) + self._router.wait_for_metadata_ready(self._device_workspace) + else: + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + + global_worker_idx = combine_worker_idx + global_worker_count = Int32(self.combine_cta_count * self.combine_warp_count) + if cutlass.const_expr(self.combine_stage_depth > 1): + self._token_back_multi_stage( + smem_workspace, + smem_base, + owned_sizes, + transfer_warp_idx, + lane_idx, + global_worker_idx, + global_worker_count, + ) + return token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) worker_mbarrier = token_back_mbarriers + transfer_warp_idx if cutlass.const_expr(self.token_back_push_data): - token_back_activation = smem_workspace.tensor( - self.token_back_activation_smem_region, smem_base - ) + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) worker_activation_stage = token_back_activation[transfer_warp_idx, None] activation_chunk_bytes = ( cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 @@ -2177,15 +2817,12 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if cutlass.const_expr(self.token_back_push_sf): token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] - sf_chunk_bytes = ( - cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 - ) + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 if lane_idx == Int32(0): cute.arch.mbarrier_init(worker_mbarrier, 1) cute.arch.sync_warp() - owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) - if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + if cutlass.const_expr(not self.cta_specialized and self.token_back_mode == "standalone_warps"): # nosec B105 sizes_ready_barrier = pipeline.NamedBarrier( barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count, @@ -2201,8 +2838,7 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) data_window_unit = activation_chunk_bytes * 2 reuse_data_pacing_enabled = ( - self.token_back_mode == "reuse_dispatch_warps" # nosec B105 - and data_window_unit > self.minimum_pacing_window_cycles + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 ) # Preserve the empirical low:initial:high ratio of 1:2.5:5. data_average_window = Int32(data_window_unit) @@ -2210,14 +2846,11 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> data_high_window = data_window_unit * 2 if cutlass.const_expr(self.token_back_push_sf): fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) - output_sf_bytes = ( - self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 - ) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) reuse_sf_pacing_enabled = ( - self.token_back_mode == "reuse_dispatch_warps" # nosec B105 - and sf_window_unit > self.minimum_pacing_window_cycles + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 ) sf_average_window = Int32(sf_window_unit) sf_low_window = sf_window_unit * 2 // 5 @@ -2229,60 +2862,63 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> next_dense_token = self.next_token(next_dense_token) expert_valid_begin = Int32(0) transfer_phase = Int32(0) + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + expert_ready_slot_begin = Int32(0) + last_ready_slot = Int32(-1) - iket.range_push("token_back.work") local_expert = Int32(0) while local_expert < Int32(self.experts_per_rank): expert_token_count = owned_sizes[local_expert] expert_valid_end = expert_valid_begin + expert_token_count - if next_dense_token < expert_valid_end: - token_tile_count = ( - expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) - ) // Int32(self.tokens_per_fc1_ready_slot) - completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) - iket.range_push("token_back.wait_fc2") - while ( - cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") - < completion_target - ): - nanosleep(500) - iket.range_pop() + if cutlass.const_expr(self.token_back_ready_granularity == "expert"): # nosec B105 + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() while next_dense_token < expert_valid_end: token_in_expert = next_dense_token - expert_valid_begin + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Same rank-local expert/tile prefix as FC1 ready; separate counter storage. + ready_slot = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + if ready_slot != last_ready_slot: + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + ready_slot, Int32, sem="acquire", scope="gpu") < Int32( + self.fc2_done_signals_per_token_tile + ): + nanosleep(500) + iket.range_pop() + last_ready_slot = ready_slot pool_token_idx = pool_expert_bases[local_expert] + token_in_expert metadata = TokenSrcMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) ) destination_topk = metadata.src_topk if cutlass.const_expr(self.reduce_topk_in_kernel): destination_topk = Int32(0) - peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map( - Int64(0), metadata.src_rank, Int64(0) - ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) is_remote_token = metadata.src_rank != self._local_rank if cutlass.const_expr(self.token_back_push_data): iket.range_push("token_back.push_data") - local_activation_address = fc2_activation_pointer.toint() + Int64( - pool_token_idx - ) * Int64(output_token_bytes) + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) remote_activation = cute.make_tensor( cute.make_ptr( self._token_comm_args.pre_reduced_activation.dtype, - self._token_comm_args.pre_reduced_activation.iterator.toint() - + peer_offset, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, AddressSpace.gmem, - assumed_align=( - self._token_comm_args.pre_reduced_activation.iterator.max_alignment - ), + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), ), self._token_comm_args.pre_reduced_activation.layout, ) - destination_row = remote_activation[ - Int64(metadata.src_token), destination_topk, None - ] + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] for chunk_idx in cutlass.range_constexpr(activation_chunk_count): chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) chunk_bytes_this_round = min( @@ -2297,17 +2933,13 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) round_start_clock = Int64(0) cute.arch.sync_warp() - if cutlass.const_expr( - reuse_data_pacing_enabled or stateless_data_pacing_enabled - ): + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): if is_remote_token: round_start_clock = read_clock64() else: round_start_clock = round_start_clock with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - worker_mbarrier, current_chunk_bytes - ) + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) tma_load_1d( worker_activation_stage.iterator, local_activation_address + chunk_byte_offset, @@ -2355,29 +2987,20 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if cutlass.const_expr(self.token_back_push_sf): iket.range_push("token_back.push_sf") - local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64( - output_sf_bytes - ) + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) remote_sf = cute.make_tensor( cute.make_ptr( self._token_comm_args.pre_reduced_activation_sf.dtype, - self._token_comm_args.pre_reduced_activation_sf.iterator.toint() - + peer_offset, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, AddressSpace.gmem, - assumed_align=( - self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment - ), + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), ), self._token_comm_args.pre_reduced_activation_sf.layout, ) - destination_sf_row = remote_sf[ - Int64(metadata.src_token), destination_topk, None - ] + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] for chunk_idx in cutlass.range_constexpr(sf_chunk_count): chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) - chunk_bytes_this_round = min( - sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes - ) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) current_chunk_bytes = Int32(chunk_bytes_this_round) current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) stateless_sf_pacing_enabled = ( @@ -2386,17 +3009,13 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) round_start_clock = Int64(0) cute.arch.sync_warp() - if cutlass.const_expr( - reuse_sf_pacing_enabled or stateless_sf_pacing_enabled - ): + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): if is_remote_token: round_start_clock = read_clock64() else: round_start_clock = round_start_clock with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - worker_mbarrier, current_chunk_bytes - ) + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) tma_load_1d( worker_sf_stage.iterator, local_sf_address + chunk_byte_offset, @@ -2412,7 +3031,9 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> ) with cute.arch.elect_one(): cp_async_bulk_s2g( - destination_chunk, worker_sf_stage.iterator, current_chunk_bytes + destination_chunk, + worker_sf_stage.iterator, + current_chunk_bytes, ) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0) @@ -2422,7 +3043,10 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> if is_remote_token: current_window = Int32(read_clock64() - round_start_clock) sf_average_window = self._adaptive_pace( - sf_average_window, current_window, sf_low_window, sf_high_window + sf_average_window, + current_window, + sf_low_window, + sf_high_window, ) elif cutlass.const_expr(stateless_sf_pacing_enabled): if is_remote_token: @@ -2433,24 +3057,28 @@ def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> cute.arch.sync_warp() next_dense_token = self.next_token(next_dense_token) + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # Advance even when this worker skipped the expert (empty experts contribute zero). + expert_ready_slot_begin = expert_ready_slot_begin + ( + expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1) + ) // Int32(self.tokens_per_fc1_ready_slot) expert_valid_begin = expert_valid_end local_expert = local_expert + Int32(1) - iket.range_pop() - # with cute.arch.elect_one(): - # cute.printf(" final data_average_window: {} ", data_average_window) @cute.jit def next_token(self, current_token: Int32) -> Int32: - global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + # The static stride is the number of warps that actually walk the pool. With a split that is the + # communication CTAs' combine warps, not every persistent CTA's transfer warps. + global_worker_count = self.combine_cta_count * self.combine_warp_count schedule_counter = None if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 claimed_token = Int32(0) - if self._lane_idx == Int32(0): - claimed_token = cute.arch.atomic_add( - schedule_counter, Int32(1), sem="relaxed", scope="gpu" - ) + if lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) return current_token + global_worker_count @@ -2471,9 +3099,7 @@ def reset_tail(self) -> None: ) iket.range_pop() total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count - global_reset_thread = ( - self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx - ) + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx iket.range_push("tail.reset_workspace") self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) @@ -2500,4 +3126,4 @@ def reset_tail(self) -> None: iket.range_pop() -__all__ = ["TokenBackMode", "TokenBackScheduleMode", "TokenCommArgs", "TokenCommNonDeterministic"] +__all__ = ["TokenCommNonDeterministic", "TokenCommArgs"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py new file mode 100644 index 000000000000..ecc713fcbf9d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_cta_mainloop.py @@ -0,0 +1,93 @@ +"""Warp execution and scratch storage for dedicated token communication CTAs.""" + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int32 + +from ...helpers.smem_workspace import SmemWorkspace +from .token_comm import TokenCommNonDeterministic + + +class TokenCommCtaMainloop: + """Run token-in and optional concurrent token-back on a dedicated CTA. + + Device buffers and counters belong to the shared TokenComm protocol. This + mainloop owns the communication CTA's warp branches and SMEM placement; + each transfer operation initializes and drains its own asynchronous copies. + The enclosing kernel performs the all-CTA reset after both sides finish. + """ + + def __init__( + self, + token_comm: TokenCommNonDeterministic, + *, + token_back_on_compute_ctas: bool, + ) -> None: + self.token_comm = token_comm + self.token_back_on_compute_ctas = token_back_on_compute_ctas + self.dispatch_warp_start = token_comm.dispatch_warp_start + self.dispatch_warp_count = token_comm.dispatch_warp_count + self.combine_warp_start = token_comm.combine_warp_start + self.combine_warp_count = token_comm.combine_warp_count + if token_comm.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("Dedicated communication CTAs require standalone token-back warps.") + if not token_back_on_compute_ctas: + dispatch_end = self.dispatch_warp_start + self.dispatch_warp_count + combine_end = self.combine_warp_start + self.combine_warp_count + if self.dispatch_warp_start < combine_end and self.combine_warp_start < dispatch_end: + raise ValueError("Concurrent token-in and token-back warp intervals must not overlap.") + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + """Declare communication CTA scratch independently of fused placement.""" + self.register_token_in_smem(workspace) + if not self.token_back_on_compute_ctas: + self.register_token_back_smem(workspace) + + def register_token_in_smem(self, workspace: SmemWorkspace) -> None: + """Declare input scratch for the dedicated dispatch warp group.""" + comm = self.token_comm + workspace.register_mbarrier( + comm.token_in_mbarrier_region, + self.dispatch_warp_count * comm.dispatch_stage_depth, + ) + workspace.register_tensor( + comm.expert_sizes_smem_region, + cutlass.Int32, + (comm.experts_per_rank,), + byte_alignment=16, + ) + staging = workspace.create_overlay("nvlink.token_comm.main_smem.transfer").add_lifetime("token_in") + comm._register_token_in_staging(staging) + + def register_token_back_smem(self, workspace: SmemWorkspace) -> None: + """Declare output scratch on whichever CTA executes standalone token-back.""" + comm = self.token_comm + if not comm.token_back_enabled: + return + workspace.register_mbarrier( + comm.token_back_mbarrier_region, + self.combine_warp_count * comm.combine_stage_depth, + ) + staging = workspace.create_overlay("nvlink.token_comm.main_smem.standalone_token_back").add_lifetime( + "token_back" + ) + comm._register_token_back_staging(staging) + + @cute.jit + def run( + self, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + ) -> None: + """Run independent warp groups without joining them before their drains.""" + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + if (warp_idx >= Int32(self.dispatch_warp_start)) & ( + warp_idx < Int32(self.dispatch_warp_start + self.dispatch_warp_count) + ): + TokenCommNonDeterministic.token_in(self.token_comm, smem_workspace, smem_base) + elif cutlass.const_expr(not self.token_back_on_compute_ctas): + if (warp_idx >= Int32(self.combine_warp_start)) & ( + warp_idx < Int32(self.combine_warp_start + self.combine_warp_count) + ): + TokenCommNonDeterministic.token_back(self.token_comm, smem_workspace, smem_base) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py new file mode 100644 index 000000000000..8efed0023e67 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_dedup.py @@ -0,0 +1,2457 @@ +"""Combine-deduplicating token communication: dispatch as in token_comm, combine reduces a token's expert copies +on the destination rank and pushes one row per (src_rank, src_token).""" + +import dataclasses +import math +import os +from typing import Callable, ClassVar, Literal, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import mark_alignment, smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + red_add_relaxed_sys_s32, + cvt_rn_bf16x2_f32x2, + f32x2_pack, + fma_rn_f32x2_bf16x2, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice + + +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs # noqa: E402 + + + +@dataclasses.dataclass(frozen=True) +class _CopyInfo: + """Per-route order and count-minus-one, packed into two eight-bit fields.""" + + dst_order_id: Int32 + dst_token_cnt: Int32 + + FIELD_BITS: ClassVar[int] = 8 + FIELD_MASK: ClassVar[int] = (1 << FIELD_BITS) - 1 + + def pack(self) -> Int32: + return self.dst_order_id | ((self.dst_token_cnt - Int32(1)) << Int32(self.FIELD_BITS)) + + @classmethod + def unpack(cls, packed: Int32) -> "_CopyInfo": + return cls( + dst_order_id=packed & Int32(cls.FIELD_MASK), + dst_token_cnt=((packed >> Int32(cls.FIELD_BITS)) & Int32(cls.FIELD_MASK)) + Int32(1), + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + """Sort-buffer element: Int128 ``[flat_topk_index | copy_info | topk_score | pad]`` with scores, else Int64 + ``[flat_topk_index | copy_info]``; ``flat_topk_index = token * topk + slot``.""" + + flat_topk_index: Int32 + copy_info: Int32 + topk_score: Optional[cutlass.Float32] + + @staticmethod + def storage_type(with_score: bool): + return cutlass.Int128 if with_score else cutlass.Int64 + + def pack(self) -> Union[cutlass.Int128, Int64]: + if cutlass.const_expr(self.topk_score is None): + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + scratch[1] = self.copy_info + return cute.recast_tensor(scratch, cutlass.Int64)[0] + scratch = cute.make_rmem_tensor((4,), cutlass.Int32) + scratch[0] = self.flat_topk_index + scratch[1] = self.copy_info + cute.recast_tensor(scratch, cutlass.Float32)[2] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int128)[0] + + @classmethod + def from_packed(cls, packed: Union[cutlass.Int128, Int64]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 64): + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], copy_info=scratch[1], topk_score=None) + scratch = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int128)[0] = packed + return cls( + flat_topk_index=scratch[0], + copy_info=scratch[1], + topk_score=cute.recast_tensor(scratch, cutlass.Float32)[2], + ) + + +def _odd_even_merge_sort_pairs(length: int) -> Tuple[Tuple[int, int], ...]: + """Batcher odd-even merge sort as (low, high) compare-exchange pairs; generated for the next power of two and + pruned to pairs inside ``length`` (a pruned partner would be +inf, a no-op).""" + width = 1 << (max(length, 1) - 1).bit_length() + pairs = [] + + def merge(lo, n, r): + step = r * 2 + if step < n: + merge(lo, n, step) + merge(lo + r, n, step) + pairs.extend((i, i + r) for i in range(lo + r, lo + n - r, step)) + else: + pairs.append((lo, lo + r)) + + def sort(lo, n): + if n > 1: + half = n // 2 + sort(lo, half) + sort(lo + half, half) + merge(lo, n, 1) + + sort(0, width) + return tuple((i, j) for i, j in pairs if j < length) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +class _DedupMetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + # Supported top-k: one token's routes are sorted in registers by a per-thread network (191 compare-exchanges at 32). + router_max_topk: ClassVar[int] = 32 + # Two Int32 registers per route (expert; within-expert index | copy info); the largest footprint without spills. + router_max_elements_per_thread: ClassVar[int] = 24 + # SMEM reserved for the expert tables / mbarriers next to the per-route buffers. + router_table_smem_bytes: ClassVar[int] = 24 * 1024 + + sizes_by_rank_region = "nvlink.token_comm_dedup.sizes_by_rank" + sizes_region = "nvlink.token_comm_dedup.sizes" + sizes_ready_region = "nvlink.token_comm_dedup.sizes_ready" + metadata_ready_region = "nvlink.token_comm_dedup.metadata_ready" + sorted_metadata_region = "nvlink.token_comm_dedup.sorted_metadata" + sorted_scores_region = "nvlink.token_comm_dedup.sorted_scores" + pool_expert_base_region = "nvlink.token_comm_dedup.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm_dedup.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm_dedup.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm_dedup.source_expert_base" + push_destination_base_region = "nvlink.token_comm_dedup.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm_dedup.sorted_metadata_ready" + push_table_ready_region = "nvlink.token_comm_dedup.push_table_ready" + router_size_counter_region = "nvlink.token_comm_dedup.router_size_counter" + router_histogram_done_region = "nvlink.token_comm_dedup.router_histogram_done" + source_base_ready_region = "nvlink.token_comm_dedup.source_base_ready" + valid_rank_mask_region = "nvlink.token_comm_dedup.valid_rank_mask" + + router_data_histogram_region = "nvlink.token_comm_dedup.router_smem.data_histogram" + router_data_prefix_region = "nvlink.token_comm_dedup.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm_dedup.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm_dedup.router_smem.data_sorted" + router_data_scores_region = "nvlink.token_comm_dedup.router_smem.data_scores" + router_data_base_region = "nvlink.token_comm_dedup.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm_dedup.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm_dedup.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm_dedup.router_smem.helper_prefix" + router_push_expert_base_region = "nvlink.token_comm_dedup.router_smem.push_expert_base" + router_push_destination_base_region = "nvlink.token_comm_dedup.router_smem.push_destination_base" + router_push_warp_start_counts_region = "nvlink.token_comm_dedup.router_smem.push_warp_start_counts" + router_helper_warp_totals_region = "nvlink.token_comm_dedup.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm_dedup.router_smem.helper_load_mbarrier" + router_data_load_mbarrier_region = "nvlink.token_comm_dedup.router_smem.data_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "promised_launchable_sm_count": int, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.carry_topk_scores = True + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + + self._validate_router_configuration() + # Large topk: drop warps so ``threads * topk`` routes still fit the per-CTA sort capacity. + while self.router_warps_per_cta > 1 and self.router_warps_per_cta * 32 * self.topk > self.router_max_routes_per_cta: + self.router_warps_per_cta //= 2 + if 32 * self.topk > self.router_max_routes_per_cta: + raise ValueError( + f"router_smem_limit_bytes={self.router_smem_limit_bytes} leaves room for {self.router_max_routes_per_cta} " + f"routes per CTA, below the 32 x topk={32 * self.topk} that one warp holding one token per thread needs." + ) + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + # Historical name: this field counts routes (tokens_per_cta * topk) rather than source tokens. + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + # Same grid footprint as the non-dedup router (DATA by route capacity, PUSH by experts / warps). + self.router_push_cta_count = min(ceil_div(self.expert_count, self.router_warps_per_cta), + self.promised_launchable_sm_count) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self.worst_case_token_count = self.worst_case_padded_tokens(self.token_padding_block) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._router_smem_base = None + + def _validate_router_configuration(self) -> None: + if self.topk > self.router_max_topk: + raise ValueError(f"Combine dedup sorts a token's routes in registers; topk must be <= {self.router_max_topk}.") + if self.world_size > 32: + raise ValueError("Combine dedup keeps the destination-rank mask in an Int32; world_size must be <= 32.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}." + ) + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def worst_case_padded_tokens(self, block: int) -> int: + source_token_capacity = self.world_size * self.max_tokens_per_rank + routes_per_source_token = min(self.topk, self.experts_per_rank) + route_capacity = source_token_capacity * routes_per_source_token + active_expert_capacity = min(self.experts_per_rank, route_capacity) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block + expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) + return min(route_budget_blocks, expert_bound_blocks) * block + + @property + def router_max_routes_per_cta(self) -> int: + # Per-route SMEM: sort record (16 B with scores, 8 B without) + fp32 score slab; capped like the base router. + record_bytes = int(_SortedElement.storage_type(self.carry_topk_scores).width) // 8 + per_route = max(record_bytes, 8) + (4 if self.carry_topk_scores else 0) + return min(16384, (self.router_smem_limit_bytes - self.router_table_smem_bytes) // per_route) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + # DATA grid = routes / capacity; each thread owns whole tokens so a token's copies never straddle threads. + routed_token_capacity = self.max_tokens_per_rank * self.topk + threads_per_cta = self.router_warps_per_cta * 32 + minimum_cta_capacity = 2048 + cta_capacity = min( + self.router_max_routes_per_cta, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)) + ) + # Whole tokens per thread, bounded by SMEM capacity and register footprint; at least one token. + tokens_per_thread = max( + 1, min(cta_capacity // (threads_per_cta * self.topk), self.router_max_elements_per_thread // self.topk) + ) + self.router_tokens_per_thread = tokens_per_thread + if tokens_per_thread * threads_per_cta * self.topk > 1 << 16: + raise ValueError("The within-expert index is packed into 16 bits; routes per DATA CTA must be <= 65536.") + elements_per_lane = tokens_per_thread * self.topk + data_cta_count = ceil_div(self.max_tokens_per_rank, tokens_per_thread * threads_per_cta) + # Slab-read bank swizzle (ROUTER_SMEM_BANK_CONFLICT_PLAN.md section 4): lane l reads vector chunk + # ``c XOR f(l)``, f(l) = (l // (32 / gcd(K, 32))) % (gcd(K, 32) / V), which makes every wavefront cover the + # 32 banks once; ``_router_value_coordinate`` inverts it. + k = self.topk + self.router_load_vector_words = 4 if k % 4 == 0 else (2 if k % 2 == 0 else 1) + gcd_k = math.gcd(k, 32) + self.router_swizzle_period_lanes = 32 // gcd_k + self.router_swizzle_groups = max(1, gcd_k // self.router_load_vector_words) + assert (k // self.router_load_vector_words) % self.router_swizzle_groups == 0 + return elements_per_lane, data_cta_count + + @cute.jit + def _router_slot_swizzle(self) -> Int32: + """f(lane) of the slab-read swizzle: the XOR applied to a thread's vector-chunk index (0 when identity).""" + if cutlass.const_expr(self.router_swizzle_groups == 1): + return Int32(0) + lane = self._router_thread_idx % Int32(32) + return (lane // Int32(self.router_swizzle_period_lanes)) % Int32(self.router_swizzle_groups) + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + workspace.register_mbarrier(self.router_data_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm_dedup.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + # The sort buffer doubles as the landing slab of the bulk-copied topk indices; scores get their own slab. + sorted_element_bytes = int(_SortedElement.storage_type(self.carry_topk_scores).width) // 8 + slab_bytes = self.router_tokens_per_cta * 8 + sorted_count = max(self.router_tokens_per_cta, ceil_div(slab_bytes, sorted_element_bytes)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + _SortedElement.storage_type(self.carry_topk_scores), + (sorted_count,), + byte_alignment=16, + ) + if self.carry_topk_scores: + data_lifetime.register_tensor( + self.router_data_scores_region, cutlass.Float32, (self.router_tokens_per_cta,), byte_alignment=16 + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + # Dense PUSH tables: per-expert source / destination start positions plus one ballot slot per warp. + push_lifetime = overlay.add_lifetime("push_cta") + push_lifetime.register_tensor( + self.router_push_expert_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + push_lifetime.register_tensor( + self.router_push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + push_lifetime.register_tensor( + self.router_push_warp_start_counts_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.carry_topk_scores: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.carry_topk_scores: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + # Per source token: bitmask of destination ranks holding >= 1 copy (rewritten every launch). + workspace.register(self.valid_rank_mask_region, cutlass.Int32, (self.max_tokens_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_size_counter_region, + cutlass.Int32, + (self.expert_count_with_trash,), + buffer_space="local", + reset="tail_reset", + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_DedupMetadataPushRouter": + if values: + raise ValueError("_DedupMetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.carry_topk_scores and topk_scores is None): + raise ValueError("Combine dedup routing requires router topk_scores.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._router_smem_base = smem_base + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata(smem_base) + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._router_smem_base = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + # Copy fields are computed while the helper exchanges sizes and builds push tables. + iket.range_push("router.sort") + self._sort_token_routes_and_scatter( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + expert = Int32(zero_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + histogram[expert] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + within_expert_indices = self._build_histogram(expert_registers, histogram) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + size_counter = self._device_workspace.ptr(self.router_size_counter_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = Int32( + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") + ) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(histogram.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + # Copy fields are computed while the helper exchanges sizes and builds push tables. + iket.range_push("router.sort") + self._sort_token_routes_and_scatter( + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = dump_base[expert] + source_expert_base[expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(histogram, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + size_counter = self._device_workspace.tensor(self.router_size_counter_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + totals[expert] = size_counter[expert] + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self, smem_base: cute.Pointer) -> None: + """Dense PUSH of the sorted records: each CTA sweeps a contiguous interval block-stride and derives a + record's expert from its position by a forward scan over the SMEM prefix table.""" + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + global_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + source_sizes = self._device_workspace.tensor(self.sizes_by_rank_region) + last_expert = Int32(self.expert_count - 1) + # Invalid routes sit in the trash bucket after the last expert and are never pushed. + total_routes = source_expert_base[last_expert] + source_sizes[Int32(self._router_local_rank), last_expert] + # CTA c owns records [c * routes_per_cta, (c + 1) * routes_per_cta), a multiple of the block (coalesced). + push_threads = Int32(self.router_push_cta_count * block_thread_count) + routes_per_cta = (total_routes + push_threads - Int32(1)) // push_threads * Int32(block_thread_count) + cta_begin = self._router_linear_cta_idx * routes_per_cta + cta_end = cutlass.min(cta_begin + routes_per_cta, total_routes) + + # Stage the per-expert tables in SMEM and ballot the expert of this CTA's first record. + expert_base = self._router_smem_workspace.tensor(self.router_push_expert_base_region, smem_base) + push_destination_base = self._router_smem_workspace.tensor(self.router_push_destination_base_region, smem_base) + warp_start_counts = self._router_smem_workspace.tensor(self.router_push_warp_start_counts_region, smem_base) + warp_start_count = Int32(0) + for expert_round in cutlass.range_constexpr(ceil_div(self.expert_count_padded, block_thread_count)): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + starts_before = Int32(0) + if expert < Int32(self.expert_count_padded): + expert_start = source_expert_base[expert] + expert_base[expert] = expert_start + push_destination_base[expert] = global_destination_base[expert] + if expert <= last_expert: + if expert_start <= cta_begin: + starts_before = Int32(1) + warp_start_count = warp_start_count + Int32( + cute.arch.popc(cute.arch.vote_ballot_sync(starts_before != Int32(0))) + ) + if self._router_lane_idx == Int32(0): + warp_start_counts[self._router_warp_idx] = warp_start_count + cute.arch.sync_threads() + start_count = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + start_count = start_count + warp_start_counts[warp] + expert = cutlass.max(start_count - Int32(1), Int32(0)) + + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_base = self._device_workspace.ptr(self.token_src_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_base = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + source_position = cta_begin + self._router_thread_idx + # Per-expert state lives in registers and is refreshed only when the cursor advances. + next_expert_start = Int32(0x7FFFFFFF) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + while (expert < last_expert) and (next_expert_start <= source_position): + expert = expert + Int32(1) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + else: + next_expert_start = Int32(0x7FFFFFFF) + destination_shift = push_destination_base[expert] - expert_base[expert] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), expert // Int32(self.experts_per_rank), Int64(0)) + while source_position < cta_end: + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + if next_expert_start <= source_position: + # Advance the cursor to the last expert whose first position is <= source_position. + while (expert < last_expert) and (next_expert_start <= source_position): + expert = expert + Int32(1) + if expert < last_expert: + next_expert_start = expert_base[expert + Int32(1)] + else: + next_expert_start = Int32(0x7FFFFFFF) + destination_shift = push_destination_base[expert] - expert_base[expert] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), expert // Int32(self.experts_per_rank), Int64(0)) + destination_position = source_position + destination_shift + stg_b64(destination_metadata_base + peer_offset + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), metadata) + if cutlass.const_expr(self.carry_topk_scores): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32(destination_scores_base + peer_offset + Int64(destination_position) * Int64(4), score) + source_position = source_position + Int32(block_thread_count) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + source_ring_offset = Int32(0) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + """Bulk-copy this CTA's token slab into SMEM; thread ``tid`` then owns tokens ``tid, tid + threads, ...`` + (``topk`` contiguous routes each) in registers.""" + topk = self.topk + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_thread = self.router_tokens_per_thread + tokens_per_cta = tokens_per_thread * threads_per_cta + routes_per_cta = tokens_per_cta * topk + elements_per_thread = tokens_per_thread * topk + experts_per_rank = self.experts_per_rank + sentinel_expert = Int32(self.expert_count_padded) + actual_token_count = Int32(self.max_tokens_per_rank) + index_bytes = int(topk_indices.dtype.width) // 8 + maximum_elements = self.max_tokens_per_rank * topk + + cta_token_begin = self._router_linear_cta_idx * Int32(tokens_per_cta) + cta_route_begin = cta_token_begin * Int32(topk) + valid_routes = Int32(maximum_elements) - cta_route_begin + if valid_routes > Int32(routes_per_cta): + valid_routes = Int32(routes_per_cta) + if valid_routes < Int32(0): + valid_routes = Int32(0) + + # Index slab aliases the sorted buffer; the score slab is its own region (it must survive the sort). + smem_base_address = self._router_smem_workspace.ptr(self.router_data_sorted_region, self._router_smem_base).toint() + slab_indices = cute.make_tensor( + cute.make_ptr(topk_indices.dtype, smem_base_address, AddressSpace.smem, assumed_align=16), + cute.make_layout((routes_per_cta,)), + ) + slab_scores = None + if cutlass.const_expr(self.carry_topk_scores): + slab_scores = self._router_smem_workspace.tensor(self.router_data_scores_region, self._router_smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_data_load_mbarrier_region, self._router_smem_base) + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + scores_flat = None + if cutlass.const_expr(self.carry_topk_scores): + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + # Bulk part: the largest 16 B multiple of the valid bytes; tail: the < 16 B remainder, one element per thread. + index_bytes_valid = valid_routes * Int32(index_bytes) + index_bytes_bulk = index_bytes_valid & Int32(~15) + score_bytes_valid = valid_routes * Int32(4) + score_bytes_bulk = score_bytes_valid & Int32(~15) + expected_bytes = index_bytes_bulk + if cutlass.const_expr(self.carry_topk_scores): + expected_bytes = expected_bytes + score_bytes_bulk + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, expected_bytes) + if index_bytes_bulk > Int32(0): + tma_load_1d( + slab_indices.iterator, + cute.make_ptr( + topk_indices.dtype, + topk_flat.iterator.toint() + Int64(cta_route_begin) * Int64(index_bytes), + AddressSpace.gmem, + assumed_align=16, + ), + load_mbarrier, + index_bytes_bulk, + ) + if cutlass.const_expr(self.carry_topk_scores): + if score_bytes_bulk > Int32(0): + tma_load_1d( + slab_scores.iterator, + cute.make_ptr( + cutlass.Float32, + scores_flat.iterator.toint() + Int64(cta_route_begin) * Int64(4), + AddressSpace.gmem, + assumed_align=16, + ), + load_mbarrier, + score_bytes_bulk, + ) + index_tail_begin = index_bytes_bulk // Int32(index_bytes) + if self._router_thread_idx < valid_routes - index_tail_begin: + route = index_tail_begin + self._router_thread_idx + slab_indices[route] = topk_flat[cta_route_begin + route] + if cutlass.const_expr(self.carry_topk_scores): + score_tail_begin = score_bytes_bulk // Int32(4) + if self._router_thread_idx < valid_routes - score_tail_begin: + route = score_tail_begin + self._router_thread_idx + slab_scores[route] = scores_flat[cta_route_begin + route] + cute.arch.mbarrier_wait(load_mbarrier, 0) + cute.arch.sync_threads() + + # Token t of this thread = routes [(t * threads + tid) * topk, +topk); only expert ids go to registers. + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + slab_indices_i32 = cute.make_tensor( + cute.make_ptr(cutlass.Int32, smem_base_address, AddressSpace.smem, assumed_align=16), + cute.make_layout((routes_per_cta * (index_bytes // 4),)), + ) + # Register s holds physical slot ``(chunk(s) XOR swizzle) * V + s % V``; _router_value_coordinate inverts it. + vector_words: cutlass.Constexpr[int] = self.router_load_vector_words + swizzle = self._router_slot_swizzle() + for token_slot in cutlass.range_constexpr(tokens_per_thread): + local_token = Int32(token_slot * threads_per_cta) + self._router_thread_idx + token_idx = cta_token_begin + local_token + token_valid = token_idx < actual_token_count + slab_route_begin = local_token * Int32(topk) + for chunk in cutlass.range_constexpr(topk // vector_words): + chunk_route_begin = slab_route_begin + (Int32(chunk) ^ swizzle) * Int32(vector_words) + for word in cutlass.range_constexpr(vector_words): + register_idx = token_slot * topk + chunk * vector_words + word + expert = slab_indices_i32[(chunk_route_begin + Int32(word)) * Int32(index_bytes // 4)] + is_invalid = (cutlass.Uint32(expert) >= cutlass.Uint32(self.expert_count)) | (not token_valid) + if is_invalid: + expert = sentinel_expert + expert_registers[register_idx] = expert + + return expert_registers, slab_scores + + @cute.jit + def _sort_token_routes_and_scatter( + self, + expert_registers: cute.Tensor, + within_expert_indices: cute.Tensor, + score_slab: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + topk_index_type: type, + ) -> None: + """Sort each token's routes by expert in registers (invalid routes sink to the end); a rank's routes form one + run, so ``dst_order_id`` = position in the run and ``dst_token_cnt`` = run length. Scatter from the sorted + registers (payload = within-expert index | physical topk slot << 16).""" + topk = self.topk + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_thread = self.router_tokens_per_thread + experts_per_rank = self.experts_per_rank + expert_count = Int32(self.expert_count) + actual_token_count = Int32(self.max_tokens_per_rank) + cta_token_begin = self._router_linear_cta_idx * Int32(tokens_per_thread * threads_per_cta) + valid_rank_mask = self._device_workspace.tensor(self.valid_rank_mask_region) + pairs = _odd_even_merge_sort_pairs(topk) + for token_slot in cutlass.range_constexpr(tokens_per_thread): + local_token = Int32(token_slot * threads_per_cta) + self._router_thread_idx + token_idx = cta_token_begin + local_token + keys = cute.make_rmem_tensor((topk,), cutlass.Int32) + payload = cute.make_rmem_tensor((topk,), cutlass.Int32) + for slot in cutlass.range_constexpr(topk): + register_idx = token_slot * topk + slot + _, physical_slot = self._router_value_coordinate(register_idx, topk_index_type) + keys[slot] = expert_registers[register_idx] + payload[slot] = within_expert_indices[register_idx] | (physical_slot << Int32(16)) + for low, high in pairs: + key_low = keys[low] + key_high = keys[high] + swap_mask = Int32(0) - Int32(key_low > key_high) + keys[low] = cutlass.min(key_low, key_high) + keys[high] = cutlass.max(key_low, key_high) + payload_low = payload[low] + payload_high = payload[high] + payload_delta = (payload_low ^ payload_high) & swap_mask + payload[low] = payload_low ^ payload_delta + payload[high] = payload_high ^ payload_delta + # Forward pass: run start (first position of this rank) and the rank bit of every valid route. + order = cute.make_rmem_tensor((topk,), cutlass.Int32) + run_start = Int32(0) + run_end_expert = Int32(0) + block_mask = Int32(0) + for position in cutlass.range_constexpr(topk): + expert = keys[position] + rank = expert // Int32(experts_per_rank) + new_run = Int32(expert >= run_end_expert) + run_start = run_start + (Int32(position) - run_start) * new_run + run_end_expert = run_end_expert + ((rank + Int32(1)) * Int32(experts_per_rank) - run_end_expert) * new_run + order[position] = Int32(position) - run_start + valid = Int32(expert < expert_count) + block_mask = block_mask | ((Int32(1) << rank) & (Int32(0) - valid)) + if token_idx < actual_token_count: + valid_rank_mask[token_idx] = block_mask + # Backward pass: run length, then the record scatter in sorted order. + run_end = Int32(topk) + for position in cutlass.range_constexpr(topk - 1, -1, -1): + if cutlass.const_expr(position + 1 < topk): + starts_next = Int32(order[position + 1] == Int32(0)) + run_end = run_end + (Int32(position + 1) - run_end) * starts_next + expert = keys[position] + valid = Int32(expert < expert_count) + order_id = order[position] * valid + copy_count = Int32(1) + (run_end - (Int32(position) - order[position]) - Int32(1)) * valid + within_expert_index = payload[position] & Int32(0xFFFF) + physical_slot = payload[position] >> Int32(16) + flat_topk_index = token_idx * Int32(topk) + physical_slot + copy_info = _CopyInfo(dst_order_id=order_id, dst_token_cnt=copy_count).pack() + destination = expert_run_starts[expert] + within_expert_index + score = None + if cutlass.const_expr(self.carry_topk_scores): + score = score_slab[local_token * Int32(topk) + physical_slot] + sorted_elements[destination] = _SortedElement(flat_topk_index, copy_info, score).pack() + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) + for register_idx in cutlass.range_constexpr(register_count): + within_expert_indices[register_idx] = Int32( + cute.arch.atomic_add( + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" + ) + ) + cute.arch.sync_threads() + return within_expert_indices + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + """(token index, topk slot) of this thread's register ``register_idx`` under the whole-token ownership.""" + threads_per_cta = self.router_warps_per_cta * 32 + tokens_per_cta = self.router_tokens_per_thread * threads_per_cta + token_slot = register_idx // self.topk + token_idx = ( + self._router_linear_cta_idx * Int32(tokens_per_cta) + + Int32(token_slot * threads_per_cta) + + self._router_thread_idx + ) + # Undo the slab-read swizzle: register slot s came from physical slot (chunk(s) XOR f(lane)) * V + s % V. + slot = register_idx % self.topk + vector_words = self.router_load_vector_words + physical_slot = (Int32(slot // vector_words) ^ self._router_slot_swizzle()) * Int32(vector_words) + Int32( + slot % vector_words + ) + return (token_idx, physical_slot) + + @cute.jit + def _route_metadata(self, flat_topk_index: Int32, copy_info: Int32) -> TokenSrcMetadata: + """Unfold a sorted element (flat route index + packed _CopyInfo) into the wire metadata record.""" + info = _CopyInfo.unpack(copy_info) + return TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=flat_topk_index // Int32(self.topk), + src_topk=flat_topk_index % Int32(self.topk), + dst_order_id=info.dst_order_id, + dst_token_cnt=info.dst_token_cnt, + ) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = self._route_metadata(element.flat_topk_index, element.copy_info) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.carry_topk_scores): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + histogram: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.carry_topk_scores): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = histogram[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = self._route_metadata(element.flat_topk_index, element.copy_info) + packed_metadata = metadata.pack() + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + packed_metadata, + predicate, + ) + if cutlass.const_expr(self.carry_topk_scores): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.world_size): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def valid_rank_mask_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.valid_rank_mask_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.carry_topk_scores): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDedup(KernelComponent): + """Public fused-kernel token communication component.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + # Order-buffer hidden chunk (= stage slot); 1 KB chunks measured 0.1 ms slower at DSV4. + standalone_chunk_bytes: ClassVar[int] = 2048 + reduction_round_group: ClassVar[int] = 4 + # Last-copy reduce: two copies per 1-D TMA through two buffers, so per-warp SMEM is 2 x 2 x chunk. + token_back_copies_per_fill: ClassVar[int] = 2 + token_back_fill_buffers: ClassVar[int] = 2 + # Combine-warp fc2_done prefix-wait spin cadence (ns). Diagnostic override: MEGA_DEDUP_WAIT_SLEEP_NS. + fc2_done_wait_sleep_ns: ClassVar[int] = int(os.environ.get("MEGA_DEDUP_WAIT_SLEEP_NS", "500")) + # Backoff cap (ns) of that spin; polling every 0.5 us slows the MMA/epilogue warps (DSV4 T8192: -12 us). + fc2_done_wait_sleep_max_ns: ClassVar[int] = int(os.environ.get("MEGA_DEDUP_WAIT_SLEEP_MAX_NS", "8000")) + # SMEM cap per combine warp for the fill buffers (reuse mode shares the token_in overlay). + stage_ring_budget_bytes: ClassVar[int] = 12288 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm_dedup.fc1_ready" + fc1_activation_region = "nvlink.token_comm_dedup.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm_dedup.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm_dedup.fc2_done" + # Order buffer the FC2 epilogue scatters into at [src_token, src_rank, dst_order_id]. + fc2_order_output_region = "nvlink.token_comm_dedup.fc2_order_output" + pre_reduced_activation_region = "nvlink.token_comm_dedup.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm_dedup.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm_dedup.token_back_schedule" # nosec B105 + # Rank-major copies pulled by a token's first same-rank occurrence (dispatch dedup). + + token_in_mbarrier_region = "nvlink.token_comm_dedup.main_smem.token_in_mbarriers" # nosec B105 + token_back_mbarrier_region = "nvlink.token_comm_dedup.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm_dedup.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_in_activation" # nosec B105 + token_in_sf_smem_region = "nvlink.token_comm_dedup.main_smem.token_in_sf" # nosec B105 + token_back_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_activation" # nosec B105 + # hybrid_warps: second set of rings for the dispatch warps, overlaid on their token_in stage. + token_back_reuse_activation_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_reuse_activation" # nosec B105 + token_back_sf_smem_region = "nvlink.token_comm_dedup.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + "token_in_flag_batch": int, + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "router_smem_limit_bytes": OptionalRequirement(int), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.carry_topk_scores = True + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.router_smem_limit_bytes = impl_desc.get("router_smem_limit_bytes", 227 * 1024) + self._validate_configuration() + self._router = _DedupMetadataPushRouter(problem_desc, impl_desc) + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + self._transfer_group = None + + def _validate_configuration(self) -> None: + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "hidden_size", + "token_padding_block", + "sf_padding_block", + "tokens_per_fc1_ready_slot", + "promised_launchable_sm_count", + "router_smem_limit_bytes", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + if self.token_back_mode not in ("standalone_warps", "reuse_dispatch_warps", "hybrid_warps"): + raise ValueError( + f"Combine dedup pushes the reduced rows from the combine warps; token_back_mode {self.token_back_mode!r} " + "is not supported (use standalone_warps, reuse_dispatch_warps or hybrid_warps)." + ) + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("Combine dedup supports the BF16 combine format only.") + if self.hidden_size % 128 != 0: + raise ValueError("Combine dedup requires hidden_size divisible by the 128-wide FC2 epilogue tile.") + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_dtype(self) -> type: + return self.quant_kind.activation_dtype + + @property + def activation_sf_dtype(self) -> type: + return self.quant_kind.sf_dtype + + @property + def activation_sf_vector_size(self) -> int: + return self.quant_kind.sf_vec_size + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + # Dedup always pushes from the combine warps (reuse_dispatch_warps / standalone_warps). + return True + + @property + def output_token_bytes(self) -> int: + return self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + + # The FC2 epilogue writes 128-hidden CTA tiles; a chunk must hold whole tiles so one tile never straddles two. + order_chunk_granule_bytes: ClassVar[int] = 256 + + @property + def token_back_stage_slots(self) -> int: + """Chunk slots of SMEM per combine warp: fill buffers x copies per fill.""" + return self.token_back_fill_buffers * self.token_back_copies_per_fill + + def token_back_chunk_bytes(self, available_bytes_per_warp: int) -> int: + """Bytes per hidden chunk of the order buffer / last-copy reduce: the per-warp budget divided by the + ``token_back_stage_slots`` chunk slots, capped at the standalone chunk and the row, rounded down to the 256 B + granule; the row's last chunk may be partial (padded in the buffer).""" + row_bytes = self.output_token_bytes + granule = self.order_chunk_granule_bytes + chunk = min(self.standalone_chunk_bytes, available_bytes_per_warp // self.token_back_stage_slots) + chunk = min(chunk, ceil_div(row_bytes, granule) * granule) + chunk = chunk // granule * granule + if chunk < granule: + raise ValueError( + f"Token-back fill buffers too small: {available_bytes_per_warp} B per warp for " + f"{self.token_back_stage_slots} chunk slots gives a {chunk} B chunk (need >= {granule} B)." + ) + if chunk > 32 * 16 * self.reduction_round_group: + raise ValueError(f"Token-back chunk {chunk} B exceeds one reduce round; its accumulators must fit registers.") + return chunk + + @property + def hybrid_token_back(self) -> bool: + """hybrid_warps: the standalone combine warps AND the dispatch warps (after token_in) both run token_back. + Twice the combine workers for the same chunk size; wins where the combine tail dominates (Qwen T>=8k: + -5..-7 %), loses where mainloop interference dominates (DSV4 8k: +6 %).""" + return self.token_back_mode == "hybrid_warps" # nosec B105 + + @property + def token_back_worker_groups(self) -> int: + return 2 if self.hybrid_token_back else 1 + + def _token_back_available_bytes_per_warp(self) -> int: + if self.token_back_mode in ("standalone_warps", "hybrid_warps"): + # hybrid: one chunk size for both groups (the order-buffer layout is shared). + return self.stage_ring_budget_bytes + activation_bytes = self.bytes_per_token + sf_bytes = self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + return max(activation_bytes + sf_bytes, self.stage_ring_budget_bytes) + + @property + def order_chunk_bytes(self) -> int: + """Hidden chunk of the order buffer layout [T, R, chunks_per_row, K, chunk].""" + return self.token_back_chunk_bytes(self._token_back_available_bytes_per_warp()) + + @property + def order_chunks_per_row(self) -> int: + return ceil_div(self.output_token_bytes, self.order_chunk_bytes) + + @property + def order_block_bytes(self) -> int: + """Bytes of one (src_token, src_rank) block: chunks_per_row x K x chunk.""" + return self.order_chunks_per_row * self.topk * self.order_chunk_bytes + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + return self._router.worst_case_padded_tokens(self.sf_padding_block) + + @property + def max_fc1_ready_slot_count(self) -> int: + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + # Transfer-warp group: 0 = dispatch warps (8-11), 1 = standalone combine warps (12-15). + self._transfer_group = cute.arch.make_warp_uniform( + Int32(thread_idx // Int32(self.transfer_thread_count) == Int32(3)) + ) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + self._transfer_group = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDedup": + if values: + raise ValueError("TokenCommDedup carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + """Register the fused kernel's GMEM regions: FC1 pools, the token-back machinery (``fc2_done`` + order + buffer) and, unless ``reduce_topk_in_kernel``, the symmetric (T, R, H) combine plane (never cleared: every + exposed cell is rewritten each launch).""" + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + # Order buffer [T, R, chunks_per_row, K, chunk]; never cleared (every byte read was written this launch). + order_bytes = self.max_tokens_per_rank * self.world_size * self.order_block_bytes + workspace.register( + self.fc2_order_output_region, + self.combine_format.act_dtype, + (order_bytes * 8 // int(self.combine_format.act_dtype.width),), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + if not self.reduce_topk_in_kernel: + # One slot per pushing rank (not per topk): the receiver sums the slots named by valid_rank_mask. + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.world_size, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier( + self.token_back_mbarrier_region, + self.transfer_warp_count * self.token_back_worker_groups * self.token_back_fill_buffers, + ) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm_dedup.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size), + byte_alignment=16, + ) + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode in ("standalone_warps", "hybrid_warps"): + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm_dedup.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_push_data: + # Per warp: two fill buffers of two chunk slots; reuse mode shares the token_in overlay. + chunk_bytes = self.order_chunk_bytes + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, self.token_back_stage_slots, chunk_elements), + byte_alignment=16, + ) + if self.hybrid_token_back: + # Dispatch warps' buffers live in their own token_in overlay (free once token_in is done). + transfer_overlay.add_lifetime("token_back").register_tensor( + self.token_back_reuse_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, self.token_back_stage_slots, chunk_elements), + byte_alignment=16, + ) + if False: # bf16-only combine: no scale stage + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout((self.worst_case_token_count, self.hidden_size), stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Base of the order buffer as a (T*R*K, 1, hidden) tensor. Only the base pointer and the hidden extent + (shape[2]) are meaningful: the epilogue addresses rows through the chunked layout + [T, R, chunks_per_row, K, chunk] (``dedup_chunk_bytes`` / ``dedup_chunks_per_row`` in its impl_desc).""" + rows = self.max_tokens_per_rank * self.world_size * self.topk + return cute.make_tensor( + device_workspace.ptr(self.fc2_order_output_region), + cute.make_layout((rows, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1)), + ) + + @cute.jit + def valid_rank_mask_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.valid_rank_mask_tensor(device_workspace) + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return None + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """The (tokens, world_size, hidden) rank-slot staging plane, or None under in-kernel top-k reduction.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return None + @cute.jit + def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools (same as the non-dedup dispatch: + every copy pulls its row from the source rank; dedup happens on the combine side only).""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = _copy_atom(cutlass.Int32, copy_elements * 32) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + size_copy_registers = cute.make_rmem_tensor((copy_elements, size_copy_rounds), cutlass.Int32) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + source_size_vectors[None, vector_idx], + size_copy_registers[None, size_copy_round], + ) + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, + size_copy_registers[None, size_copy_round], + destination_size_vectors[None, vector_idx], + ) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_flag_batch == 1, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + sf_token_idx = expert_sf_begin + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr( + self.token_back_enabled and self.token_back_mode in ("reuse_dispatch_warps", "hybrid_warps") + ): + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Combine warps: walk the pool metadata in expert order, wait for the ``fc2_done`` prefix up to the current + expert, and for every last copy (``order == cnt - 1``) reduce the block's copies in fp32 (order 0..cnt-1), + round once to bf16 and push it to the source rank. Single-copy rows are pushed by the FC2 epilogue.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + + workers_per_cta = self.transfer_warp_count * self.token_back_worker_groups + worker_in_cta = transfer_warp_idx + if cutlass.const_expr(self.hybrid_token_back): + worker_in_cta = self._transfer_group * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_idx = self._linear_cta_idx * Int32(workers_per_cta) + worker_in_cta + global_worker_count = Int32(self.promised_launchable_sm_count * workers_per_cta) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarriers = token_back_mbarriers + worker_in_cta * Int32(self.token_back_fill_buffers) + token_back_stages = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_stages = token_back_stages[transfer_warp_idx, None, None] # (stage_slots, chunk_elements) + if cutlass.const_expr(self.hybrid_token_back): + # Dispatch warps (group 0) reduce in the rings overlaid on their finished token_in stage. + reuse_stages = smem_workspace.tensor(self.token_back_reuse_activation_smem_region, smem_base) + if self._transfer_group == Int32(0): + worker_stages = reuse_stages[transfer_warp_idx, None, None] + chunk_bytes = cute.size(worker_stages, mode=[1]) * int(self.combine_format.act_dtype.width) // 8 + if lane_idx == Int32(0): + for buffer_idx in cutlass.range_constexpr(self.token_back_fill_buffers): + cute.arch.mbarrier_init(worker_mbarriers + buffer_idx, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): + # The dispatch warps already met this barrier inside token_in (hybrid); only the standalone group waits. + if cutlass.const_expr(self.token_back_mode == "standalone_warps") or self._transfer_group == Int32(1): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + order_pointer = self._device_workspace.ptr(self.fc2_order_output_region) + output_token_bytes = self.output_token_bytes + # Fills issued by this worker so far: fill s uses buffer s % 2 and phase (s // 2) & 1 of its mbarrier. + fill_sequence = Int32(0) + + iket.range_push("token_back.work") + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + completed_expert_prefix = Int32(0) + expert_valid_begin = Int32(0) + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + # Single-copy rows were pushed straight from the FC2 epilogue; only multi-copy blocks are reduced here. + if (metadata.is_last_copy != Int32(0)) and (metadata.dst_token_cnt > Int32(1)): + iket.range_push("token_back.wait_fc2_warp_prefix") + wait_sleep_ns = Int32(self.fc2_done_wait_sleep_ns) + while completed_expert_prefix <= local_expert: + prefix_base = (completed_expert_prefix // Int32(32)) * Int32(32) + checked_expert = prefix_base + Int32(31) - lane_idx + lane_ready = checked_expert >= Int32(self.experts_per_rank) + if checked_expert < Int32(self.experts_per_rank): + checked_count = owned_sizes[checked_expert] + checked_target = ((checked_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot)) * Int32(self.fc2_done_signals_per_token_tile) + lane_ready = cute.arch.load(fc2_done + checked_expert, Int32, + sem="acquire", scope="gpu") >= checked_target + ready_mask = cute.arch.vote_ballot_sync(lane_ready) + completed_expert_prefix = cutlass.min( + prefix_base + Int32(cute.arch.clz(~ready_mask)), Int32(self.experts_per_rank)) + if completed_expert_prefix <= local_expert: + nanosleep(wait_sleep_ns) + wait_sleep_ns = cutlass.min(wait_sleep_ns * Int32(2), Int32(self.fc2_done_wait_sleep_max_ns)) + # Order every lane's fc2_done acquire before the TMA that reads the completed copies. + cute.arch.sync_warp() + iket.range_pop() + block_index = Int64(metadata.src_token) * Int64(self.world_size) + Int64(metadata.src_rank) + block_address = order_pointer.toint() + block_index * Int64(self.order_block_bytes) + fill_sequence = self._reduce_and_push( + block_address, + metadata.dst_token_cnt, + metadata.src_rank, + metadata.src_token, + worker_mbarriers, + worker_stages, + fill_sequence, + chunk_bytes, + ) + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def _issue_fill( + self, + sequence: Int32, + fill_idx: Int32, + fills_per_chunk: Int32, + copy_count: Int32, + worker_mbarriers: cute.Pointer, + stage_base, + block_address: Int64, + chunk_bytes: int, + ) -> None: + """Elected lane: one 1-D TMA of the two copies of fill ``fill_idx`` (chunk-major, copies ascending) into + buffer ``sequence % 2``, arming that buffer's mbarrier with the byte count.""" + chunk_idx = fill_idx // fills_per_chunk + first_copy = (fill_idx - chunk_idx * fills_per_chunk) * Int32(self.token_back_copies_per_fill) + copies = cutlass.min(Int32(self.token_back_copies_per_fill), copy_count - first_copy) + buffer_idx = sequence % Int32(self.token_back_fill_buffers) + mbarrier = worker_mbarriers + buffer_idx + destination = cute.make_ptr( + cutlass.Uint8, + stage_base + buffer_idx * Int32(self.token_back_copies_per_fill * chunk_bytes), + AddressSpace.smem, + assumed_align=16, + ) + source = block_address + (Int64(chunk_idx) * Int64(self.topk) + Int64(first_copy)) * Int64(chunk_bytes) + cute.arch.mbarrier_arrive_and_expect_tx(mbarrier, copies * Int32(chunk_bytes)) + tma_load_1d(destination, source, mbarrier, copies * Int32(chunk_bytes)) + + @cute.jit + def _reduce_and_push( + self, + block_address: Int64, + copy_count: Int32, + src_rank: Int32, + src_token: Int32, + worker_mbarriers: cute.Pointer, + worker_stages: cute.Tensor, + fill_sequence: Int32, + chunk_bytes: int, + ) -> Int32: + """Reduce the ``copy_count`` copies of the order block at ``block_address`` chunk by chunk (two copies per + TMA through two buffers, fp32 accumulation in order 0..cnt-1, one bf16 rounding) and push each chunk. + Returns the advanced fill sequence.""" + lane_idx = self._lane_idx + copies_per_fill = self.token_back_copies_per_fill + output_token_bytes = self.output_token_bytes + chunk_count = ceil_div(output_token_bytes, chunk_bytes) + round_group = min(self.reduction_round_group, ceil_div(chunk_bytes // 16, 32)) + destination_slot = Int32(self._local_rank) + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_slot = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), src_rank, Int64(0)) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(src_token), destination_slot, None] + word_atom = _copy_atom(cutlass.Int32, 128) + vector_layout = cute.make_layout((4,)) + stage_base = worker_stages.iterator.toint() + one2 = f32x2_pack(cutlass.Float32(1.0), cutlass.Float32(1.0)) # {1.0, 1.0} weight for the FFMA2 adds + fills_per_chunk = (copy_count + Int32(copies_per_fill - 1)) // Int32(copies_per_fill) + fill_count = fills_per_chunk * Int32(chunk_count) + + iket.range_push("token_back.push_data") + with cute.arch.elect_one(): + self._issue_fill(fill_sequence, Int32(0), fills_per_chunk, copy_count, worker_mbarriers, stage_base, block_address, chunk_bytes) + # Whole-chunk fp32 accumulators (round_group x 32 lanes x 16 B covers the chunk, checked at construction). + words = cute.make_rmem_tensor((4,), cutlass.Int32) + acc = cute.make_rmem_tensor((4, round_group), cutlass.Int64) + # Rolled loop over the fills of the token: the body exists once in SASS (icache footprint). + for fill_idx in cutlass.range(fill_count, unroll=1): + chunk_idx = fill_idx // fills_per_chunk + fill_in_chunk = fill_idx - chunk_idx * fills_per_chunk + first_copy = fill_in_chunk * Int32(copies_per_fill) + sequence = fill_sequence + fill_idx + buffer_base = stage_base + (sequence % Int32(self.token_back_fill_buffers)) * Int32(copies_per_fill * chunk_bytes) + current_vectors = cutlass.min(Int32(chunk_bytes), Int32(output_token_bytes) - chunk_idx * Int32(chunk_bytes)) // Int32(16) + # Prefetch fill i+1 into the other buffer. If that buffer just fed a bulk copy, wait for its READ only. + if fill_idx + Int32(1) < fill_count: + cute.arch.cp_async_bulk_wait_group(0, read=True) + with cute.arch.elect_one(): + self._issue_fill(sequence + Int32(1), fill_idx + Int32(1), fills_per_chunk, copy_count, + worker_mbarriers, stage_base, block_address, chunk_bytes) + if fill_in_chunk == Int32(0): + for group in cutlass.range_constexpr(round_group): + for word in cutlass.range_constexpr(4): + acc[word, group] = Int64(0) + cute.arch.mbarrier_wait( + worker_mbarriers + sequence % Int32(self.token_back_fill_buffers), + (sequence // Int32(self.token_back_fill_buffers)) & Int32(1), + ) + iket.range_push("token_back.reduce") + # Copies accumulate in 0..cnt-1 order; groups are 512 B apart so adjacent lanes keep contiguous 16 B. + for copy_in_fill in cutlass.range_constexpr(copies_per_fill): + if first_copy + Int32(copy_in_fill) < copy_count: + for group in cutlass.range_constexpr(round_group): + vector_idx = Int32(group * 32) + lane_idx + if vector_idx < current_vectors: + slot_vector = cute.make_tensor( + cute.make_ptr(cutlass.Int32, buffer_base + Int32(copy_in_fill * chunk_bytes) + vector_idx * Int32(16), + AddressSpace.smem, assumed_align=16), + vector_layout, + ) + cute.copy(word_atom, slot_vector, words) + for word in cutlass.range_constexpr(4): + acc[word, group] = fma_rn_f32x2_bf16x2(words[word], one2, acc[word, group]) + iket.range_pop() + if fill_in_chunk == fills_per_chunk - Int32(1): + # Round once to bf16 into slot 0 of this buffer and push the chunk from there. + for group in cutlass.range_constexpr(round_group): + vector_idx = Int32(group * 32) + lane_idx + if vector_idx < current_vectors: + slot0_vector = cute.make_tensor( + cute.make_ptr(cutlass.Int32, buffer_base + vector_idx * Int32(16), AddressSpace.smem, assumed_align=16), + vector_layout, + ) + for word in cutlass.range_constexpr(4): + words[word] = cvt_rn_bf16x2_f32x2(acc[word, group]) + cute.copy(word_atom, words, slot0_vector) + # Generic-proxy SMEM writes must be visible to the bulk copy (async proxy) that reads slot 0. + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_warp() + chunk_source = cute.make_ptr(cutlass.Uint8, buffer_base, AddressSpace.smem, assumed_align=16) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + Int64(chunk_idx) * Int64(chunk_bytes), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g(destination_chunk, chunk_source, current_vectors * Int32(16)) + else: + cp_async_bulk_s2g(destination_chunk, chunk_source, current_vectors * Int32(16)) + cute.arch.cp_async_bulk_commit_group() + cute.arch.sync_warp() + # The pushes of this token must have landed before the worker moves on (matches the previous per-chunk wait). + cute.arch.cp_async_bulk_wait_group(0) + iket.range_pop() + cute.arch.sync_warp() + return fill_sequence + fill_count + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + # Static stride over every token_back worker in the grid; must match ``global_worker_count`` in ``token_back``. + workers_per_cta = self.transfer_warp_count * self.token_back_worker_groups + global_worker_count = self.promised_launchable_sm_count * workers_per_cta + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with the four token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDedup"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py new file mode 100644 index 000000000000..98b651a96143 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/nvlink_domain/token_comm_deterministic.py @@ -0,0 +1,2706 @@ +"""Metadata-push routing with a fixed token sequence and fused communication.""" + +import dataclasses +import os +from typing import Callable, ClassVar, Optional, Tuple, Union + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 +from cutlass.utils.blockscaled_layout import tile_atom_to_shape_SF + +from ...api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.dsl_helpers import smem_exclusive_prefix +from ...helpers.flag_batch import make_flag_batch_tracker +from ...helpers.iket_compat import iket +from ...helpers.software_sync import NvlinkBarrier +from ...helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, + cp_reduce_async_bulk_add_u32_s2g, + nanosleep, + read_clock64, + red_add_relaxed_sys_s32, + stg_b64, + stg_f32, + tma_load_1d, +) +from ...helpers.smem_workspace import SmemWorkspace +from ...helpers.utils import ceil_div, round_up +from ...quant_def import CombineFormat, QuantKind +from ..token_protocol import TokenSrcMetadata +from .symmetric_buffer import SymmetricBufferDevice +from .token_comm import TokenBackMode, TokenBackScheduleMode, TokenCommArgs + + +_quant_spec = { + "bf16": (cutlass.BFloat16, None, 0), + "nvfp4": (cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), + "mxfp4": (cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e4m3": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + "mxfp8_e5m2": (cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), +} + + +@dataclasses.dataclass(frozen=True) +class _ReceiveCapacity: + raw_route_count: int + logical_route_count: int + padded_route_count: int + + +def _compute_receive_capacity( + *, + world_size: int, + max_tokens_per_rank: int, + topk: int, + experts_per_rank: int, + max_recv_size_per_rank: int, + padding_block: int, +) -> _ReceiveCapacity: + raw_route_count = world_size * max_tokens_per_rank * topk + logical_route_count = min(max_recv_size_per_rank, raw_route_count) + active_expert_count = min(experts_per_rank, logical_route_count) + padded_block_count = active_expert_count + (logical_route_count - active_expert_count) // padding_block + return _ReceiveCapacity( + raw_route_count=raw_route_count, + logical_route_count=logical_route_count, + padded_route_count=padded_block_count * padding_block, + ) + + +@dataclasses.dataclass(frozen=True) +class _SortedElement: + flat_topk_index: Int32 + topk_score: Optional[cutlass.Float32] + + def pack(self) -> Union[Int64, Int32]: + if cutlass.const_expr(self.topk_score is None): + return self.flat_topk_index + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + scratch[0] = self.flat_topk_index + cute.recast_tensor(scratch, cutlass.Float32)[1] = self.topk_score + return cute.recast_tensor(scratch, cutlass.Int64)[0] + + @classmethod + def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": + if cutlass.const_expr(type(packed).width == 32): + return cls(flat_topk_index=Int32(packed), topk_score=None) + scratch = cute.make_rmem_tensor((2,), cutlass.Int32) + cute.recast_tensor(scratch, cutlass.Int64)[0] = packed + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) + + +@cute.jit +def _copy_atom(dtype, num_bits_per_copy: int): + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) + + +@cute.jit +def _mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: + pointer = tensor.iterator + return cute.make_tensor( + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout + ) + + +@cute.jit +def _device_trap() -> None: + """Terminate the current kernel after its failure state is published.""" + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class _MetadataPushRouter(KernelComponent): + """Sort and push routing metadata into each destination rank's final pool.""" + + router_smem_limit_bytes: ClassVar[int] = 227 * 1024 + router_warps_per_cta: ClassVar[int] = 16 + + sizes_by_rank_region = "nvlink.token_comm.sizes_by_rank" + sizes_region = "nvlink.token_comm.sizes" + sizes_ready_region = "nvlink.token_comm.sizes_ready" + metadata_ready_region = "nvlink.token_comm.metadata_ready" + sorted_metadata_region = "nvlink.token_comm.sorted_metadata" + sorted_scores_region = "nvlink.token_comm.sorted_scores" + pool_expert_base_region = "nvlink.token_comm.pool_expert_base" + token_src_metadata_region = "nvlink.token_comm.token_src_metadata" # nosec B105 + fc1_topk_scores_region = "nvlink.token_comm.fc1_topk_scores" + source_expert_base_region = "nvlink.token_comm.source_expert_base" + push_destination_base_region = "nvlink.token_comm.push_destination_base" + sorted_metadata_ready_region = "nvlink.token_comm.sorted_metadata_ready" + # Private continuation: zero means this (source token, topk slot) is absent. + input_dedup_routes_region = "nvlink.token_comm.input_dedup_routes" + push_table_ready_region = "nvlink.token_comm.push_table_ready" + router_histogram_done_region = "nvlink.token_comm.router_histogram_done" + router_cta_histograms_region = "nvlink.token_comm.router_cta_histograms" + source_base_ready_region = "nvlink.token_comm.source_base_ready" + + router_data_histogram_region = "nvlink.token_comm.router_smem.data_histogram" + router_data_totals_region = "nvlink.token_comm.router_smem.data_totals" + router_data_prefix_region = "nvlink.token_comm.router_smem.data_prefix" + router_data_warp_totals_region = "nvlink.token_comm.router_smem.data_warp_totals" + router_data_sorted_region = "nvlink.token_comm.router_smem.data_sorted" + router_data_base_region = "nvlink.token_comm.router_smem.data_base" + router_helper_size_matrix_region = "nvlink.token_comm.router_smem.helper_size_matrix" + router_helper_totals_region = "nvlink.token_comm.router_smem.helper_totals" + router_helper_prefix_region = "nvlink.token_comm.router_smem.helper_prefix" + router_helper_warp_totals_region = "nvlink.token_comm.router_smem.helper_warp_totals" + router_helper_load_mbarrier_region = "nvlink.token_comm.router_smem.helper_load_mbarrier" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return {"token_padding_block": int, "promised_launchable_sm_count": int, "drop_on_overflow": bool, + "input_dedup_mode": OptionalRequirement(str), + "interleave_source_ranks": OptionalRequirement(bool), + "sf_padding_block": OptionalRequirement(int), + "tokens_per_fc1_ready_slot": OptionalRequirement(int)} + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + self.interleave_source_ranks = impl_desc.get("interleave_source_ranks", False) + + self.input_dedup_mode = impl_desc.get("input_dedup_mode", "off") + if self.input_dedup_mode not in ("off", "rank"): + raise ValueError("input_dedup_mode must be off/rank") + # The retained specialization uses expert-major ownership and publishes + # readiness only after every local fanout store has completed. + self.input_dedup_store_mode = "batched" if self.input_dedup_mode != "off" else "serial" + self.input_dedup_order = "expert" if self.input_dedup_mode != "off" else "token" + self._validate_router_configuration() + token_capacity = self.receive_capacity(self.token_padding_block) + if self.input_dedup_mode != "off": + # Route encoding uses three 20-bit row/ready fields. The smallest + # destination row selects the owner without encoding an expert ID. + if not self.apply_topk_at_fc1: + raise ValueError("Input dedup requires backward routing") + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + if self.sf_padding_block <= 0 or self.tokens_per_fc1_ready_slot <= 0: + raise ValueError("Input dedup requires positive SF padding and ready-slot granularity") + if self.max_recv_size_per_rank != self.world_size * self.max_tokens_per_rank * self.topk: + raise ValueError("Input dedup requires uncapped receive capacity") + # The low field stores pool_row + 1; zero denotes an unused route. + if token_capacity.padded_route_count >= (1 << 20): + raise ValueError("Input dedup requires padded receive capacity < 2**20 rows") + # SF rows and ready slots have separate zero-based 20-bit fields. + sf_capacity = self.receive_capacity(self.sf_padding_block).padded_route_count + ready_slots = ( + self.receive_capacity(self.tokens_per_fc1_ready_slot).padded_route_count + // self.tokens_per_fc1_ready_slot + ) + if sf_capacity > (1 << 20): + raise ValueError("Input dedup requires padded SF capacity <= 2**20 rows") + if ready_slots > (1 << 20): + raise ValueError("Input dedup requires at most 2**20 ready slots") + self.raw_route_count = token_capacity.raw_route_count + self.logical_route_capacity = token_capacity.logical_route_count + self.worst_case_token_count = token_capacity.padded_route_count + self.expert_count_padded = round_up(self.expert_count, 4) + self.expert_count_with_trash = self.expert_count_padded + 1 + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() + self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 + self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) + self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) + if self.router_grid_cta_count > self.promised_launchable_sm_count: + raise ValueError( + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." + ) + self._router_smem_workspace = self._build_router_smem_workspace() + + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + def _validate_router_configuration(self) -> None: + if type(self.max_recv_size_per_rank) is not int: + raise TypeError(f"max_recv_size_per_rank must be an int, got {type(self.max_recv_size_per_rank).__name__}.") + if type(self.drop_on_overflow) is not bool: + raise TypeError(f"drop_on_overflow must be a bool, got {type(self.drop_on_overflow).__name__}.") + positive_fields = ( + "world_size", + "expert_count", + "topk", + "max_tokens_per_rank", + "max_recv_size_per_rank", + "token_padding_block", + "promised_launchable_sm_count", + ) + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}, got {self.max_tokens_per_rank}." + ) + if self.expert_count % self.world_size != 0: + raise ValueError( + f"expert_count must be divisible by world_size, got {self.expert_count} and {self.world_size}." + ) + if self.expert_count > 16384: + raise NotImplementedError("TokenComm supports at most 16384 global experts.") + if self.topk > self.expert_count: + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + def receive_capacity(self, padding_block: int) -> _ReceiveCapacity: + return _compute_receive_capacity( + world_size=self.world_size, + max_tokens_per_rank=self.max_tokens_per_rank, + topk=self.topk, + experts_per_rank=self.experts_per_rank, + max_recv_size_per_rank=self.max_recv_size_per_rank, + padding_block=padding_block, + ) + + def _router_launch_configuration(self) -> Tuple[int, int]: + def next_power_of_two(value: int) -> int: + return 1 << (max(value, 1) - 1).bit_length() + + routed_token_capacity = self.max_tokens_per_rank * self.topk + minimum_cta_capacity = 2048 + maximum_cta_capacity = 16384 + maximum_data_cta_count = 128 + maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count + if routed_token_capacity > maximum_supported_tokens: + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens per rank.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) + elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) + data_cta_count = ceil_div(routed_token_capacity, cta_capacity) + return elements_per_lane, data_cta_count + + def _build_router_smem_workspace(self) -> SmemWorkspace: + workspace = SmemWorkspace() + workspace.register_mbarrier(self.router_helper_load_mbarrier_region, 1) + overlay = workspace.create_overlay("nvlink.token_comm.router_smem.role") + data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor( + self.router_data_histogram_region, + cutlass.Int32, + (self.expert_count_with_trash, self.router_warps_per_cta + 1), + stride=(self.router_warps_per_cta + 1, 1), + ) + data_lifetime.register_tensor( + self.router_data_totals_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor( + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 + ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) + data_lifetime.register_tensor( + self.router_data_sorted_region, + (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), + (self.router_tokens_per_cta,), + byte_alignment=16, + ) + if self.router_data_cta_count > 1: + data_lifetime.register_tensor( + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + + helper_lifetime = overlay.add_lifetime("helper_cta") + helper_lifetime.register_tensor( + self.router_helper_size_matrix_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + stride=(self.expert_count_padded, 1), + byte_alignment=16, + ) + helper_lifetime.register_tensor( + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 + ) + helper_lifetime.register_tensor( + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 + ) + workspace.finalize(max_bytes=self.router_smem_limit_bytes) + return workspace + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router_smem_workspace + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + """Register router-private state and Router-to-Main outputs.""" + self._register_router_workspace(workspace) + + def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: + maximum_routed_tokens = self.max_tokens_per_rank * self.topk + if self.input_dedup_mode != "off": + workspace.register( + self.input_dedup_routes_region, cutlass.Int64, + (self.world_size * maximum_routed_tokens,), + buffer_space="shared", byte_alignment=128, reset="tail_reset", + ) + workspace.register( + self.sizes_by_rank_region, + cutlass.Int32, + (self.world_size, self.expert_count_padded), + buffer_space="shared", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" + ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") + if self.apply_topk_at_fc1: + workspace.register( + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" + ) + workspace.register( + self.token_src_metadata_region, + cutlass.Int64, + (self.worst_case_token_count,), + buffer_space="shared", + byte_alignment=16, + ) + if self.apply_topk_at_fc1: + workspace.register( + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" + ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local") + workspace.register( + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.push_destination_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" + ) + workspace.register( + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") + if self.router_data_cta_count > 1: + workspace.register( + self.router_cta_histograms_region, + cutlass.Int32, + (self.router_data_cta_count, self.expert_count_padded), + buffer_space="local", + stride=(self.expert_count_padded, 1), + ) + workspace.register( + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + workspace.register( + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "_MetadataPushRouter": + if values: + raise ValueError("_MetadataPushRouter carries no MLIR values.") + return self + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + """Launch counting-sort DATA, size-exchange HELPER, and metadata PUSH roles.""" + if cutlass.const_expr(self.apply_topk_at_fc1 and topk_scores is None): + raise ValueError("apply_topk_at_fc1 requires router topk_scores.") + if cutlass.const_expr(overflow_flag.iterator.dtype is not cutlass.Int32): + raise TypeError("overflow_flag must use Int32 elements.") + if cutlass.const_expr(cute.size(overflow_flag) != 1): + raise ValueError("overflow_flag must contain exactly one element.") + peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() + self._router_kernel( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper, + device_workspace, + ).launch( + grid=[self.router_grid_cta_count, 1, 1], + block=[self.router_warps_per_cta * 32, 1, 1], + min_blocks_per_mp=1, + stream=stream, + ) + + @cute.kernel + def _router_kernel( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper: SymmetricBufferDevice, + device_workspace: DeviceWorkspace, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + linear_cta_idx, _, _ = cute.arch.block_idx() + cute.arch.griddepcontrol_launch_dependents() + block_thread_count = self.router_warps_per_cta * 32 + grid_thread_idx = thread_idx + linear_cta_idx * block_thread_count + warp_idx = cute.arch.make_warp_uniform(thread_idx // Int32(32)) + lane_idx = thread_idx % Int32(32) + + storage_type = self._router_smem_workspace.storage_class() + smem_allocator = cutlass.utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + + device_workspace.assign_device_members(local_workspace, shared_workspace) + self._device_workspace = device_workspace + self._peer_rank_ptr_mapper = peer_rank_ptr_mapper + self._router_local_rank = local_rank + self._router_thread_idx = thread_idx + self._router_linear_cta_idx = linear_cta_idx + self._router_grid_thread_idx = grid_thread_idx + self._router_warp_idx = warp_idx + self._router_lane_idx = lane_idx + self._overflow_flag = overflow_flag + + if cutlass.const_expr(self.router_data_cta_count == 1): + self._router_single_cta(topk_indices, topk_scores, smem_base) + else: + self._router_multiple_ctas(topk_indices, topk_scores, smem_base) + if linear_cta_idx < Int32(self.router_push_cta_count): + self._router_push_metadata() + + device_workspace.remove_device_members() + self._device_workspace = None + self._peer_rank_ptr_mapper = None + self._router_local_rank = None + self._router_thread_idx = None + self._router_linear_cta_idx = None + self._router_grid_thread_idx = None + self._router_warp_idx = None + self._router_lane_idx = None + self._overflow_flag = None + + @cute.jit + def _router_single_cta( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.prefix_and_publish") + publish_sizes = self._broadcast_sizes_to_peers( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))) + ) + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + publish_sizes() + iket.range_pop() + + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.write_out") + self._dump_contiguous_router_output(sorted_elements, total_valid_routes) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_single_cta(smem_base) + + @cute.jit + def _router_multiple_ctas( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor], smem_base: cute.Pointer + ) -> None: + if self._router_linear_cta_idx < Int32(self.router_data_cta_count): + block_thread_count = self.router_warps_per_cta * 32 + trash_bucket = self.expert_count_padded + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) + totals = self._router_smem_workspace.tensor(self.router_data_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) + dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) + + histogram_element_count = self.expert_count_with_trash * (self.router_warps_per_cta + 1) + histogram_flat = cute.make_tensor(histogram.iterator, cute.make_layout((histogram_element_count,))) + zero_round_count = ceil_div(histogram_element_count, block_thread_count) + for zero_round in cutlass.range_constexpr(zero_round_count): + slot = Int32(zero_round * block_thread_count) + self._router_thread_idx + if slot < Int32(histogram_element_count): + histogram_flat[slot] = Int32(0) + + iket.range_push("router.histogram") + expert_registers, score_registers = self._load_router_inputs(topk_indices, topk_scores) + cute.arch.sync_threads() + match_masks = self._build_histogram(expert_registers, histogram) + self._prefix_warp_histogram(histogram, totals) + iket.range_pop() + + iket.range_push("router.reserve_and_prefix") + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + cta_histograms[self._router_linear_cta_idx, expert] = totals[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" + ) + + total_valid_routes = smem_exclusive_prefix( + cute.make_tensor(totals.iterator, cute.make_layout((self.expert_count_padded,))), + cute.make_tensor(prefix.iterator, cute.make_layout((self.expert_count_padded,))), + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + if self._router_thread_idx == Int32(0): + prefix[trash_bucket] = total_valid_routes + cute.arch.sync_threads() + iket.range_pop() + iket.range_push("router.sort") + self._sort_router_elements( + expert_registers, match_masks, score_registers, sorted_elements, prefix, histogram, topk_indices.dtype + ) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.wait_source_base") + source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.write_out") + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + dump_base[expert] = source_expert_base[expert] + cta_histograms[self._router_linear_cta_idx, expert] + cute.arch.sync_threads() + self._dump_router_output_by_expert(totals, prefix, dump_base, sorted_elements) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" + ) + iket.range_pop() + elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): + self._router_helper_multiple_ctas(smem_base) + + @cute.jit + def _router_helper_single_cta(self, smem_base: cute.Pointer) -> None: + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + cta_histograms = self._device_workspace.tensor(self.router_cta_histograms_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) + + histogram_done = self._device_workspace.ptr(self.router_histogram_done_region) + iket.range_push("router.wait_histogram") + if self._router_thread_idx == Int32(0): + while cute.arch.load(histogram_done, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.broadcast_sizes") + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + running = Int32(0) + for cta_idx in cutlass.range_constexpr(self.router_data_cta_count): + count = cta_histograms[cta_idx, expert] + cta_histograms[cta_idx, expert] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + publish_sizes = self._broadcast_sizes_to_peers(totals) + smem_exclusive_prefix( + totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + source_expert_base[expert] = prefix[expert] + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.atomic_add( + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" + ) + publish_sizes() + iket.range_pop() + + iket.range_push("router.compute_push_tables") + self._compute_push_tables(smem_base) + iket.range_pop() + + @cute.jit + def _router_push_metadata(self) -> None: + block_thread_count = self.router_warps_per_cta * 32 + sorted_metadata_ready = self._device_workspace.ptr(self.sorted_metadata_ready_region) + push_table_ready = self._device_workspace.ptr(self.push_table_ready_region) + if self._router_thread_idx == Int32(0): + while cute.arch.load(sorted_metadata_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.router_data_cta_count + ): + nanosleep(150) + while cute.arch.load(push_table_ready, Int32, sem="acquire", scope="gpu") != Int32(1): + nanosleep(150) + cute.arch.sync_threads() + + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx + if global_expert < Int32(self.expert_count): + sizes_by_rank = self._device_workspace.tensor(self.sizes_by_rank_region) + source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + route_count = sizes_by_rank[self._router_local_rank, global_expert] + source_begin = source_expert_base[global_expert] + destination_begin = push_destination_base[global_expert] + destination_rank = global_expert // Int32(self.experts_per_rank) + if cutlass.const_expr(self.interleave_source_ranks): + source_counts = [sizes_by_rank[r, global_expert] for r in range(self.world_size)] + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) + destination_metadata_address = ( + self._device_workspace.ptr(self.token_src_metadata_region).toint() + peer_offset + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + source_scores = self._device_workspace.ptr(self.sorted_scores_region) + destination_scores_address = ( + self._device_workspace.ptr(self.fc1_topk_scores_region).toint() + peer_offset + ) + if cutlass.const_expr(self.input_dedup_mode != "off"): + destination_routes_address = self._device_workspace.ptr(self.input_dedup_routes_region).toint() + peer_offset + local_expert = global_expert % Int32(self.experts_per_rank) + global_sizes = self._device_workspace.tensor(self.sizes_region) + pool_begin = Int32(0) + sf_begin = Int32(0) + ready_begin = Int32(0) + for earlier in cutlass.range_constexpr(self.experts_per_rank): + if Int32(earlier) < local_expert: + count = global_sizes[destination_rank * Int32(self.experts_per_rank) + Int32(earlier)] + pool_begin = pool_begin + ( + (count + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + sf_begin = sf_begin + ( + (count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + ready_begin = ready_begin + ( + (count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + route_round_count = (route_count + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + if route < route_count: + source_position = source_begin + route + destination_position = destination_begin + route + if cutlass.const_expr(self.interleave_source_ranks): + # Visit source ranks in order, skipping exhausted ranks. + # This is also the canonical row order of fc1_preact: + # FWD and BWD must select the same ordering. + round_robin_offset = Int32(0) + for source_rank in cutlass.range_constexpr(self.world_size): + count = source_counts[source_rank] + round_robin_offset = round_robin_offset + cutlass.min(route, count) + if Int32(source_rank) < Int32(self._router_local_rank) and route < count: + round_robin_offset = round_robin_offset + Int32(1) + destination_position = destination_begin + round_robin_offset + physical_store = Int32(destination_position < Int32(self.worst_case_token_count)) + metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) + stg_b64( + destination_metadata_address + Int64(destination_position) * Int64(TokenSrcMetadata.nbytes), + metadata, + physical_store, + ) + if cutlass.const_expr(self.input_dedup_mode != "off"): + # Decode through the protocol class: the record layout changed (token:u16 | topk:u8 | + # rank:u8 | order:u8 | cnt:u8) and a hand-rolled mask here silently indexed the wrong route. + decoded = TokenSrcMetadata.unpack(Int64(metadata)) + source_token = decoded.src_token + source_slot = decoded.src_topk + route_index = (self._router_local_rank * Int32(self.max_tokens_per_rank) + source_token) * Int32(self.topk) + source_slot + token_in_expert = destination_position - pool_begin + sf_position = sf_begin + token_in_expert + ready_slot = ready_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + # Zero denotes an absent route; the destination row + # also gives each source token/receiver a unique owner. + packed_route = ( + (Int64(ready_slot) << Int64(40)) | (Int64(sf_position) << Int64(20)) + | Int64(destination_position + Int32(1)) + ) + stg_b64(destination_routes_address + Int64(route_index) * Int64(8), packed_route, physical_store) + if cutlass.const_expr(self.apply_topk_at_fc1): + score = cute.arch.load(source_scores + source_position, cutlass.Float32) + stg_f32( + destination_scores_address + Int64(destination_position) * Int64(4), score, physical_store + ) + + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + # Keep notifier threads behind the leader's system fence. + cute.arch.sync_threads() + metadata_ready_address = self._device_workspace.ptr(self.metadata_ready_region).toint() + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(metadata_ready_address, destination_rank, Int64(0)), Int32(1) + ) + + @cute.jit + def _broadcast_sizes_to_peers(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: + block_thread_count = self.router_warps_per_cta * 32 + row_bytes = Int32(self.expert_count_padded * 4) + matrix_address = self._device_workspace.ptr(self.sizes_by_rank_region).toint() + total_address = self._device_workspace.ptr(self.sizes_region).toint() + rank_round_count = ceil_div(self.world_size, self.router_warps_per_cta) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = self._router_warp_idx + Int32(rank_round * self.router_warps_per_cta) + if destination_rank < Int32(self.world_size): + peer_offset = self._peer_rank_ptr_mapper.map(Int64(0), destination_rank, Int64(0)) + destination_row_address = ( + matrix_address + + peer_offset + + Int64(Int32(self._router_local_rank) * Int32(self.expert_count_padded)) * Int64(4) + ) + destination_row = cute.make_ptr( + cutlass.Int32, destination_row_address, AddressSpace.gmem, assumed_align=16 + ) + destination_total = cute.make_ptr( + cutlass.Int32, total_address + peer_offset, AddressSpace.gmem, assumed_align=16 + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_row, smem_expert_counts.iterator, row_bytes) + cp_reduce_async_bulk_add_u32_s2g(destination_total, smem_expert_counts.iterator, row_bytes) + cute.arch.cp_async_bulk_commit_group() + + def finalize() -> None: + cute.arch.cp_async_bulk_wait_group(0) + cute.arch.sync_threads() + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + cute.arch.sync_threads() + ready_address = self._device_workspace.ptr(self.sizes_ready_region).toint() + ready_round_count = ceil_div(self.world_size, block_thread_count) + for ready_round in cutlass.range_constexpr(ready_round_count): + destination_rank = Int32(ready_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + red_add_relaxed_sys_s32( + self._peer_rank_ptr_mapper.map(ready_address, destination_rank, Int64(0)), Int32(1) + ) + + return finalize + + @cute.jit + def _compute_push_tables(self, smem_base: cute.Pointer) -> None: + block_thread_count = self.router_warps_per_cta * 32 + owner_expert_begin = Int32(self._router_local_rank) * Int32(self.experts_per_rank) + matrix_bytes = self.world_size * self.expert_count_padded * 4 + + size_matrix = self._router_smem_workspace.tensor(self.router_helper_size_matrix_region, smem_base) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) + prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) + load_mbarrier = self._router_smem_workspace.ptr(self.router_helper_load_mbarrier_region, smem_base) + sizes = self._device_workspace.tensor(self.sizes_region) + + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_init(load_mbarrier, 1) + + sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) + iket.range_push("router.wait_sizes_ready") + if self._router_thread_idx == Int32(0): + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="sys") != Int32(self.raw_sizes_ready_target): + nanosleep(150) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.load_sizes_and_prefix") + if self._router_thread_idx == Int32(0): + cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(matrix_bytes)) + tma_load_1d( + size_matrix.iterator, + self._device_workspace.ptr(self.sizes_by_rank_region), + load_mbarrier, + Int32(matrix_bytes), + ) + + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_padded): + expert_size = sizes[expert] + padded_totals[expert] = ( + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) + ) * Int32(self.token_padding_block) + cute.arch.sync_threads() + smem_exclusive_prefix( + padded_totals, + prefix, + warp_totals, + block_thread_count, + self._router_thread_idx, + self._router_lane_idx, + self._router_warp_idx, + ) + pool_expert_base = self._device_workspace.tensor(self.pool_expert_base_region) + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + pool_expert_base[local_expert] = prefix[owner_expert_begin + local_expert] - prefix[owner_expert_begin] + + cute.arch.mbarrier_wait(load_mbarrier, 0) + iket.range_pop() + + iket.range_push("router.build_push_destinations") + push_destination_base = self._device_workspace.tensor(self.push_destination_base_region) + padded_expert_rounds = ceil_div(self.expert_count_padded, block_thread_count) + for expert_round in cutlass.range_constexpr(padded_expert_rounds): + global_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if global_expert < Int32(self.expert_count): + destination_rank = global_expert // Int32(self.experts_per_rank) + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + destination_pool_base = prefix[global_expert] - prefix[destination_expert_begin] + source_ring_offset = Int32(0) + if cutlass.const_expr(not self.interleave_source_ranks): + local_ring_position = ( + Int32(self._router_local_rank) - destination_rank + Int32(self.world_size) + ) % Int32(self.world_size) + for ring_position in cutlass.range_constexpr(self.world_size): + source_rank = (destination_rank + Int32(ring_position)) % Int32(self.world_size) + if Int32(ring_position) < local_ring_position: + source_ring_offset = source_ring_offset + size_matrix[source_rank, global_expert] + push_destination_base[global_expert] = destination_pool_base + source_ring_offset + cute.arch.sync_threads() + iket.range_pop() + + iket.range_push("router.apply_receive_limit") + if self._router_thread_idx == Int32(0): + warp_totals[0] = Int32(0) + warp_totals[1] = Int32(0) + cute.arch.sync_threads() + + # sizes[] contains the same pre-truncation, all-source expert totals on + # every rank once raw_sizes_ready_target has been observed. Compute the + # group-wide policy bit locally so no post-kernel collective is needed. + # Slot 0 counts overflowing destinations; slot 1 records this rank's + # destination bit for local-only truncation. + rank_round_count = ceil_div(self.world_size, block_thread_count) + for rank_round in cutlass.range_constexpr(rank_round_count): + destination_rank = Int32(rank_round * block_thread_count) + self._router_thread_idx + if destination_rank < Int32(self.world_size): + destination_expert_begin = destination_rank * Int32(self.experts_per_rank) + raw_destination_total = Int32(0) + for local_expert in cutlass.range_constexpr(self.experts_per_rank): + raw_destination_total = ( + raw_destination_total + + sizes[destination_expert_begin + Int32(local_expert)] + ) + destination_overflow = Int32( + raw_destination_total > Int32(self.max_recv_size_per_rank) + ) + if destination_overflow: + cute.arch.atomic_add( + warp_totals.iterator, + Int32(1), + scope="cta", + ) + if destination_rank == Int32(self._router_local_rank): + warp_totals[1] = destination_overflow + cute.arch.sync_threads() + global_did_overflow = Int32(warp_totals[0] > Int32(0)) + local_did_overflow = warp_totals[1] + + if cutlass.const_expr(self.drop_on_overflow): + local_expert_rounds = ceil_div(self.experts_per_rank, block_thread_count) + for expert_round in cutlass.range_constexpr(local_expert_rounds): + local_expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if local_expert < Int32(self.experts_per_rank): + # Keep every raw row that still lands inside the statically allocated padded pool. + raw_size = sizes[owner_expert_begin + local_expert] + retained_size = Int32(0) + remaining_capacity = Int32(self.worst_case_token_count) - pool_expert_base[local_expert] + if remaining_capacity > Int32(0): + retained_size = cutlass.min(raw_size, remaining_capacity) + if local_did_overflow: + sizes[owner_expert_begin + local_expert] = retained_size + + if self._router_thread_idx == Int32(0): + cute.arch.store( + self._overflow_flag.iterator, + global_did_overflow, + sem="relaxed", + scope="sys", + ) + cute.arch.sync_threads() + + if cutlass.const_expr(self.drop_on_overflow): + self._publish_push_table_and_sizes_ready(sizes_ready) + else: + if global_did_overflow: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_sys() + _device_trap() + else: + self._publish_push_table_and_sizes_ready(sizes_ready) + iket.range_pop() + + @cute.jit + def _publish_push_table_and_sizes_ready(self, sizes_ready: cute.Pointer) -> None: + if self._router_thread_idx == Int32(0): + cute.arch.fence_acq_rel_gpu() + cute.arch.atomic_add( + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="relaxed", scope="gpu" + ) + cute.arch.atomic_add(sizes_ready, Int32(1), sem="relaxed", scope="gpu") + + @cute.jit + def _load_router_inputs( + self, topk_indices: cute.Tensor, topk_scores: Optional[cute.Tensor] + ) -> Tuple[cute.Tensor, Optional[cute.Tensor]]: + elements_per_vector = 128 // topk_indices.dtype.width + grid_thread_count = self.router_data_cta_count * self.router_warps_per_cta * 32 + tile_span = elements_per_vector * grid_thread_count + maximum_elements = self.max_tokens_per_rank * self.topk + actual_token_count = Int32(self.max_tokens_per_rank) + actual_elements = Int32(maximum_elements) + load_round_count = ceil_div(maximum_elements, tile_span) + elements_per_thread = load_round_count * elements_per_vector + + topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) + load_atom = _copy_atom(topk_indices.dtype, 128) + expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) + raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + raw_vectors[None, load_round], + ) + else: + expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + load_atom, + _mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), + expert_vectors[None, load_round], + ) + + score_registers = None + if cutlass.const_expr(self.apply_topk_at_fc1): + score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) + score_vectors = cute.logical_divide( + cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) + ) + score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) + score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) + for load_round in cutlass.range_constexpr(load_round_count): + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) + if tile_begin < actual_elements: + cute.copy( + score_atom, + _mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), + score_register_vectors[None, load_round], + ) + + expert_registers_u32 = cute.recast_tensor(expert_registers, cutlass.Uint32) + if cutlass.const_expr(topk_indices.dtype.width == 64): + raw_indices_i32 = cute.recast_tensor(raw_indices, cutlass.Int32) + for register_idx in cutlass.range_constexpr(elements_per_thread): + if cutlass.const_expr(topk_indices.dtype.width == 64): + expert_registers[register_idx] = raw_indices_i32[2 * register_idx] + token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) + if is_invalid: + expert_registers[register_idx] = Int32(self.expert_count_padded) + return expert_registers, score_registers + + @cute.jit + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: + register_count = cute.size(expert_registers) + match_masks = cute.make_rmem_tensor((register_count,), cutlass.Int32) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = Int32(cute.arch.match_sync(0xFFFFFFFF, expert, kind="any")) + match_masks[register_idx] = match_mask + histogram_slot = expert * Int32(self.router_warps_per_cta + 1) + self._router_warp_idx + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + if rank_in_group == Int32(0): + cute.arch.atomic_add( + histogram.iterator + histogram_slot, Int32(cute.arch.popc(match_mask)), sem="relaxed", scope="cta" + ) + cute.arch.sync_threads() + return match_masks + + @cute.jit + def _prefix_warp_histogram(self, histogram: cute.Tensor, totals: cute.Tensor) -> None: + """Turn per-warp counts into fixed warp bases and expert totals.""" + block_thread_count = self.router_warps_per_cta * 32 + expert_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = Int32(expert_round * block_thread_count) + self._router_thread_idx + if expert < Int32(self.expert_count_with_trash): + running = Int32(0) + for warp in cutlass.range_constexpr(self.router_warps_per_cta): + count = histogram[expert, warp] + histogram[expert, warp] = running + running = running + count + totals[expert] = running + cute.arch.sync_threads() + + @cute.jit + def _sort_router_elements( + self, + expert_registers: cute.Tensor, + match_masks: cute.Tensor, + score_registers: Optional[cute.Tensor], + sorted_elements: cute.Tensor, + expert_run_starts: cute.Tensor, + warp_histogram: cute.Tensor, + topk_index_type: type, + ) -> None: + """Stable CTA-local counting-sort scatter. + + The cursor for an expert is private to one warp. ``match.any`` gives + every equal-expert lane a fixed rank, and only that group's first lane + advances the cursor. Therefore the output order is a pure function of + warp id, register round, and lane id rather than atomic arrival order. + """ + register_count = cute.size(expert_registers) + lane_mask_less_than = (Int32(1) << self._router_lane_idx) - Int32(1) + for register_idx in cutlass.range_constexpr(register_count): + expert = expert_registers[register_idx] + match_mask = match_masks[register_idx] + group_size = Int32(cute.arch.popc(match_mask)) + rank_in_group = Int32(cute.arch.popc(match_mask & lane_mask_less_than)) + warp_base = warp_histogram[expert, self._router_warp_idx] + token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) + flat_topk_index = token_idx * Int32(self.topk) + topk_slot + destination = expert_run_starts[expert] + warp_base + rank_in_group + + # Every lane in an equal-expert group reads the old cursor before + # its leader advances it for the next register round. + cute.arch.sync_warp() + if rank_in_group == Int32(0): + warp_histogram[expert, self._router_warp_idx] = warp_base + group_size + if cutlass.const_expr(self.apply_topk_at_fc1): + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() + else: + sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() + cute.arch.sync_warp() + + @cute.jit + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: + elements_per_vector = 128 // topk_index_type.width + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 + flat_index = Int32( + register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector + ) + self._router_grid_thread_idx * Int32(elements_per_vector) + return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) + + @cute.jit + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: + block_thread_count = self.router_warps_per_cta * 32 + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) + for dump_round in cutlass.range(dump_round_count, unroll=4): + position = Int32(dump_round * block_thread_count) + self._router_thread_idx + predicate = Int32(position < total_valid_routes) + element = _SortedElement.from_packed(sorted_elements[position]) + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64(metadata_address + Int64(position) * Int64(TokenSrcMetadata.nbytes), metadata.pack(), predicate) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def _dump_router_output_by_expert( + self, + expert_totals: cute.Tensor, + expert_run_starts: cute.Tensor, + expert_dump_bases: cute.Tensor, + sorted_elements: cute.Tensor, + ) -> None: + metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() + if cutlass.const_expr(self.apply_topk_at_fc1): + score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() + expert_round_count = ceil_div(self.expert_count_padded, self.router_warps_per_cta) + for expert_round in cutlass.range_constexpr(expert_round_count): + expert = self._router_warp_idx + Int32(expert_round * self.router_warps_per_cta) + if expert < Int32(self.expert_count_padded): + run_begin = expert_run_starts[expert] + run_length = expert_totals[expert] + dump_begin = expert_dump_bases[expert] + route_round_count = (run_length + Int32(31)) // Int32(32) + for route_round in cutlass.range(route_round_count, unroll=1): + route = Int32(route_round) * Int32(32) + self._router_lane_idx + predicate = Int32(route < run_length) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) + output_position = dump_begin + route + metadata = TokenSrcMetadata( + src_rank=Int32(self._router_local_rank), + src_token=(element.flat_topk_index // Int32(self.topk)), + src_topk=(element.flat_topk_index % Int32(self.topk)), + ) + stg_b64( + metadata_address + Int64(output_position) * Int64(TokenSrcMetadata.nbytes), + metadata.pack(), + predicate, + ) + if cutlass.const_expr(self.apply_topk_at_fc1): + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + """Return this rank's contiguous expert-size view.""" + sizes = device_workspace.tensor(self.sizes_region) + expert_begin = local_rank * Int32(self.experts_per_rank) + return cute.make_tensor(sizes.iterator + expert_begin, cute.make_layout((self.experts_per_rank,))) + + @property + def metadata_ready_target(self) -> int: + return self.router_push_cta_count * self.world_size + + @property + def raw_sizes_ready_target(self) -> int: + return self.world_size + + @property + def published_sizes_ready_target(self) -> int: + return self.world_size + 1 + + @cute.jit + def sizes_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.sizes_region) + + @cute.jit + def pool_expert_base_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.pool_expert_base_region) + + @cute.jit + def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.token_src_metadata_region) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + sizes_ready = device_workspace.ptr(self.sizes_ready_region) + while cute.arch.load(sizes_ready, Int32, sem="acquire", scope="gpu") != Int32( + self.published_sizes_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx == Int32(0): + metadata_ready = device_workspace.ptr(self.metadata_ready_region) + while cute.arch.load(metadata_ready, Int32, sem="acquire", scope="sys") != Int32( + self.metadata_ready_target + ): + nanosleep(sleep_cycles) + cute.arch.sync_warp() + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return device_workspace.tensor(self.token_src_metadata_region) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.apply_topk_at_fc1): + return None + return device_workspace.tensor(self.fc1_topk_scores_region) + + +class TokenCommDeterministic(KernelComponent): + """Fused communication with a fixed logical token-position sequence.""" + + transfer_warp_count: ClassVar[int] = 4 + transfer_thread_count: ClassVar[int] = transfer_warp_count * 32 + standalone_chunk_bytes: ClassVar[int] = 2048 + minimum_pacing_window_cycles: ClassVar[int] = 512 + standalone_max_backoff_cycles: ClassVar[int] = 500 + adaptive_minimum_sleep_cycles: ClassVar[int] = 50 + transfer_lifetime_barrier_id: ClassVar[int] = 9 + grid_sync_barrier_id: ClassVar[int] = 10 + standalone_size_barrier_id: ClassVar[int] = 11 + token_in_size_barrier_id: ClassVar[int] = 12 + + fc1_ready_region = "nvlink.token_comm.fc1_ready" + fc1_activation_region = "nvlink.token_comm.fc1_activation" + fc1_activation_sf_region = "nvlink.token_comm.fc1_activation_sf" + fc2_done_region = "nvlink.token_comm.fc2_done" + fc2_activation_region = "nvlink.token_comm.fc2_activation" + fc2_activation_sf_region = "nvlink.token_comm.fc2_activation_sf" + pre_reduced_activation_region = "nvlink.token_comm.pre_reduced_activation" + pre_reduced_activation_sf_region = "nvlink.token_comm.pre_reduced_activation_sf" + token_back_schedule_region = "nvlink.token_comm.token_back_schedule" # nosec B105 + + token_in_mbarrier_region = "nvlink.token_comm.main_smem.token_in_mbarriers" # nosec B105 + token_back_mbarrier_region = "nvlink.token_comm.main_smem.token_back_mbarriers" # nosec B105 + expert_sizes_smem_region = "nvlink.token_comm.main_smem.expert_sizes" + token_in_activation_smem_region = "nvlink.token_comm.main_smem.token_in_activation" # nosec B105 + token_in_sf_smem_region = "nvlink.token_comm.main_smem.token_in_sf" # nosec B105 + token_back_activation_smem_region = "nvlink.token_comm.main_smem.token_back_activation" # nosec B105 + token_back_sf_smem_region = "nvlink.token_comm.main_smem.token_back_sf" # nosec B105 + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return { + "world_size": int, + "expert_count": int, + "topk": int, + "max_tokens_per_rank": int, + "max_recv_size_per_rank": int, + "hidden_size": int, + "quant_kind": str, + "combine_format": CombineFormat, + "apply_topk_at_fc1": bool, + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "token_padding_block": int, + "sf_padding_block": int, + "tokens_per_fc1_ready_slot": int, + "fc2_done_signals_per_token_tile": int, + "promised_launchable_sm_count": int, + # Optional so existing component users keep the four-warp ABI. + # Embedding kernels must still launch exactly this many transfer + # warps and include them in their block-wide synchronization. + "token_in_transfer_warp_count": OptionalRequirement(int), + "token_in_flag_batch": int, + "token_in_async_flags": OptionalRequirement(bool), + "input_dedup_mode": OptionalRequirement(str), + "interleave_source_ranks": OptionalRequirement(bool), + "token_back_mode": str, + "token_back_schedule_mode": str, + "reduce_topk_in_kernel": bool, + "drop_on_overflow": bool, + "token_in_stage_split": OptionalRequirement(bool), + "external_dispatch_pool": OptionalRequirement(bool), + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.world_size = problem_desc["world_size"] + self.expert_count = problem_desc["expert_count"] + self.topk = problem_desc["topk"] + self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] + self.max_recv_size_per_rank = min( + problem_desc["max_recv_size_per_rank"], self.world_size * self.max_tokens_per_rank * self.topk + ) + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = problem_desc["quant_kind"] + self.combine_format = problem_desc["combine_format"] + self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.tokens_per_fc1_ready_slot = impl_desc["tokens_per_fc1_ready_slot"] + self.fc2_done_signals_per_token_tile = impl_desc["fc2_done_signals_per_token_tile"] + self.promised_launchable_sm_count = impl_desc["promised_launchable_sm_count"] + self.transfer_warp_count = impl_desc.get( + "token_in_transfer_warp_count", type(self).transfer_warp_count + ) + self.transfer_thread_count = self.transfer_warp_count * 32 + self.token_in_flag_batch = impl_desc["token_in_flag_batch"] + self.token_in_async_flags = impl_desc.get("token_in_async_flags", False) + self.token_back_mode: TokenBackMode = impl_desc["token_back_mode"] + self.token_back_schedule_mode: TokenBackScheduleMode = impl_desc["token_back_schedule_mode"] + self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.drop_on_overflow = impl_desc["drop_on_overflow"] + self.token_in_stage_split = impl_desc.get("token_in_stage_split", False) + self._token_in_stage_chunk_count = 2 if self.token_in_stage_split else 1 + self.external_dispatch_pool = impl_desc.get("external_dispatch_pool", False) + + self._validate_configuration() + self._router = _MetadataPushRouter(problem_desc, impl_desc) + self.interleave_source_ranks = self._router.interleave_source_ranks + self.input_dedup_mode = self._router.input_dedup_mode + self.input_dedup_store_mode = self._router.input_dedup_store_mode + self.input_dedup_order = self._router.input_dedup_order + if self.input_dedup_mode != "off" and self.quant_kind not in ("mxfp8_e4m3", "mxfp8_e5m2"): + raise ValueError("Input dedup requires MXFP8") + self._nvlink_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.grid_sync_barrier_id) + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def _validate_configuration(self) -> None: + positive_fields = ("hidden_size", "token_padding_block", "sf_padding_block", "tokens_per_fc1_ready_slot") + for field_name in positive_fields: + value = getattr(self, field_name) + if value <= 0: + raise ValueError(f"{field_name} must be positive, got {value}.") + if self.quant_kind not in _quant_spec: + raise ValueError(f"Unsupported quant_kind {self.quant_kind!r}.") + if self.token_back_mode not in ("epi_warps", "standalone_warps", "reuse_dispatch_warps"): + raise ValueError(f"Unsupported token_back_mode {self.token_back_mode!r}.") + if self.token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + f"token_back_schedule_mode must be static or atomic_counter, got {self.token_back_schedule_mode!r}." + ) + if self.token_in_async_flags and self.token_in_flag_batch != 1: + raise ValueError("Asynchronous token-in flags require token_in_flag_batch=1.") + if not 1 <= self.token_in_flag_batch <= 32: + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") + if ( + isinstance(self.transfer_warp_count, bool) + or not isinstance(self.transfer_warp_count, int) + or not 1 <= self.transfer_warp_count <= 32 + ): + raise ValueError( + "token_in_transfer_warp_count must be in [1, 32], got " + f"{self.transfer_warp_count}." + ) + if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: + raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") + _pad_lo = min(self.sf_padding_block, self.token_padding_block) + _pad_hi = max(self.sf_padding_block, self.token_padding_block) + if _pad_hi % _pad_lo != 0: + raise ValueError("sf_padding_block and token_padding_block must be power-of-two multiples of each other.") + if self.token_back_enabled and self.fc2_done_signals_per_token_tile <= 0: + raise ValueError("fc2_done_signals_per_token_tile must be positive when token-back is enabled.") + if self.reduce_topk_in_kernel and self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("In-kernel top-k reduction requires BF16 combine data.") + if self.activation_is_quantized: + element_block = self.activation_sf_vector_size * 4 + if self.hidden_size % element_block != 0: + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") + else: + # The 1D bulk copies move whole token rows, so each row must be whole 16-byte lines. + row_element_block = 128 // int(self.activation_dtype.width) + if self.hidden_size % row_element_block != 0: + raise ValueError(f"The BF16 wire requires hidden_size divisible by {row_element_block}.") + if self.sf_padding_block % 128 != 0: + raise ValueError("sf_padding_block must be a multiple of 128.") + if self.token_in_stage_split: + # Scale-free wire only: the SF pull and the token-back stage sizing both assume one + # whole-row stage per transfer warp. + if self.activation_is_quantized or self.token_back_enabled: + raise NotImplementedError( + "token_in_stage_split requires the BF16 wire with token_back disabled (epi_warps)." + ) + if self.bytes_per_token % (self._token_in_stage_chunk_count * 16) != 0: + raise ValueError("token_in_stage_split must cut a token row into whole 16-byte chunks.") + + @property + def experts_per_rank(self) -> int: + return self.expert_count // self.world_size + + @property + def activation_is_quantized(self) -> bool: + return _quant_spec[self.quant_kind][1] is not None + + @property + def activation_dtype(self) -> type: + return _quant_spec[self.quant_kind][0] + + @property + def activation_sf_dtype(self) -> type: + return _quant_spec[self.quant_kind][1] + + @property + def activation_sf_vector_size(self) -> int: + return _quant_spec[self.quant_kind][2] + + @property + def bytes_per_token(self) -> int: + return self.hidden_size * int(self.activation_dtype.width) // 8 + + @property + def activation_sf_hidden_padded(self) -> int: + if not self.activation_is_quantized: + return 0 + valid_hidden = self.hidden_size // self.activation_sf_vector_size + elements_per_16_bytes = 128 // int(self.activation_sf_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def combine_sf_hidden_padded(self) -> int: + if not self.combine_format.is_quantized: + return 0 + valid_hidden = self.hidden_size // self.combine_format.scale_block + elements_per_16_bytes = 128 // int(self.combine_format.scale_dtype.width) + return int(round_up(valid_hidden, elements_per_16_bytes)) + + @property + def token_back_push_data(self) -> bool: + return self.token_back_mode != "epi_warps" # nosec B105 + + @property + def token_back_push_sf(self) -> bool: + return self.combine_format.is_quantized + + @property + def token_back_enabled(self) -> bool: + return self.token_back_push_data or self.token_back_push_sf + + @property + def worst_case_token_count(self) -> int: + return self._router.worst_case_token_count + + @property + def worst_case_sf_token_count(self) -> int: + # The SF pool has to be at least as tall as the data pool: once + # token_padding_block may exceed sf_padding_block, an expert's data + # cursor can walk past the SF-capacity-derived row count. + sf_capacity = self._router.receive_capacity(self.sf_padding_block).padded_route_count + return max(sf_capacity, self.worst_case_token_count) + + @property + def max_fc1_ready_slot_count(self) -> int: + return ( + self._router.receive_capacity(self.tokens_per_fc1_ready_slot).padded_route_count + // self.tokens_per_fc1_ready_slot + ) + + @property + def router_smem_workspace(self) -> SmemWorkspace: + return self._router.router_smem_workspace + + @property + def expert_count_padded(self) -> int: + return self._router.expert_count_padded + + @property + def expert_count_with_trash(self) -> int: + return self._router.expert_count_with_trash + + @property + def router_elements_per_lane(self) -> int: + return self._router.router_elements_per_lane + + @property + def router_warps_per_cta(self) -> int: + return self._router.router_warps_per_cta + + @property + def router_data_cta_count(self) -> int: + return self._router.router_data_cta_count + + @property + def router_tokens_per_cta(self) -> int: + return self._router.router_tokens_per_cta + + @property + def router_push_cta_count(self) -> int: + return self._router.router_push_cta_count + + @property + def router_grid_cta_count(self) -> int: + return self._router.router_grid_cta_count + + @property + def metadata_ready_target(self) -> int: + return self._router.metadata_ready_target + + def register_device_workspace(self, workspace: DeviceWorkspace) -> None: + self._router.register_device_workspace(workspace) + self._register_main_workspace(workspace) + self._nvlink_barrier.register_device_workspace(workspace) + + @cute.jit + def launch_router( + self, + topk_indices: cute.Tensor, + topk_scores: Optional[cute.Tensor], + overflow_flag: cute.Tensor, + local_rank: Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + peer_rank_ptr_mapper_host, + device_workspace: DeviceWorkspace, + stream: cuda.CUstream, + ) -> None: + self._router.launch_router( + topk_indices, + topk_scores, + overflow_flag, + local_rank, + local_workspace, + shared_workspace, + peer_rank_ptr_mapper_host, + device_workspace, + stream, + ) + + @cute.jit + def local_expert_sizes(self, device_workspace: DeviceWorkspace, local_rank: Int32) -> cute.Tensor: + return self._router.local_expert_sizes(device_workspace, local_rank) + + @cute.jit + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: + self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) + + @cute.jit + def token_src_metadata_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + return self._router.token_src_metadata_tensor(device_workspace) + + @cute.jit + def fc1_topk_scores_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + return self._router.fc1_topk_scores_tensor(device_workspace) + + @cute.jit + def assign_device_members( + self, + *, + device_workspace: DeviceWorkspace, + token_comm_args: TokenCommArgs, + local_rank: Int32, + linear_cta_idx: Int32, + ) -> None: + if cutlass.const_expr(self.external_dispatch_pool and token_comm_args.fc1_activation_pool is None): + raise ValueError("external_dispatch_pool requires TokenCommArgs.fc1_activation_pool.") + if cutlass.const_expr(not self.external_dispatch_pool and token_comm_args.fc1_activation_pool is not None): + raise ValueError("TokenCommArgs.fc1_activation_pool is only consumed under external_dispatch_pool.") + self._device_workspace = device_workspace + self._token_comm_args = token_comm_args + self._local_rank = local_rank + self._linear_cta_idx = linear_cta_idx + thread_idx, _, _ = cute.arch.thread_idx() + transfer_thread_idx = thread_idx % Int32(self.transfer_thread_count) + self._transfer_warp_idx = cute.arch.make_warp_uniform(transfer_thread_idx // Int32(32)) + self._lane_idx = transfer_thread_idx % Int32(32) + self._nvlink_barrier.assign_device_members(device_workspace, token_comm_args.peer_rank_ptr_mapper) + + def remove_device_members(self) -> None: + self._nvlink_barrier.remove_device_members() + self._device_workspace = None + self._token_comm_args = None + self._local_rank = None + self._linear_cta_idx = None + self._transfer_warp_idx = None + self._lane_idx = None + + def __extract_mlir_values__(self) -> list: + return [] + + def __new_from_mlir_values__(self, values: list) -> "TokenCommDeterministic": + if values: + raise ValueError("TokenCommDeterministic carries no MLIR values.") + return self + + def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: + workspace.register( + self.fc1_ready_region, + cutlass.Int32, + (self.max_fc1_ready_slot_count,), + buffer_space="local", + reset="tail_reset", + ) + if not self.external_dispatch_pool: + activation_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc1_activation_region, + self.activation_dtype, + (activation_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.activation_is_quantized: + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded + workspace.register( + self.fc1_activation_sf_region, + self.activation_sf_dtype, + (activation_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled: + workspace.register( + self.fc2_done_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" + ) + if self.token_back_push_data: + fc2_element_count = self.worst_case_token_count * self.hidden_size + workspace.register( + self.fc2_activation_region, + self.combine_format.act_dtype, + (fc2_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_push_sf: + fc2_sf_element_count = self.worst_case_token_count * self.combine_sf_hidden_padded + workspace.register( + self.fc2_activation_sf_region, + self.combine_format.scale_dtype, + (fc2_sf_element_count,), + buffer_space="local", + byte_alignment=128, + ) + if self.token_back_enabled and self.token_back_schedule_mode == "atomic_counter": # nosec B105 + workspace.register( + self.token_back_schedule_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" + ) + if not self.reduce_topk_in_kernel: + workspace.register( + self.pre_reduced_activation_region, + self.combine_format.act_dtype, + (self.max_tokens_per_rank, self.topk, self.hidden_size), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + if self.combine_format.is_quantized: + workspace.register( + self.pre_reduced_activation_sf_region, + self.combine_format.scale_dtype, + (self.max_tokens_per_rank, self.topk, self.combine_sf_hidden_padded), + buffer_space="shared", + mem_order=(2, 1, 0), + byte_alignment=128, + ) + + def register_smem_regions(self, workspace: SmemWorkspace) -> None: + workspace.register_mbarrier(self.token_in_mbarrier_region, self.transfer_warp_count) + if self.token_back_enabled: + workspace.register_mbarrier(self.token_back_mbarrier_region, self.transfer_warp_count) + workspace.register_tensor( + self.expert_sizes_smem_region, cutlass.Int32, (self.experts_per_rank,), byte_alignment=16 + ) + transfer_overlay = workspace.create_overlay("nvlink.token_comm.main_smem.transfer") + token_in_lifetime = transfer_overlay.add_lifetime("token_in") + token_in_lifetime.register_tensor( + self.token_in_activation_smem_region, + self.activation_dtype, + (self.transfer_warp_count, self.hidden_size // self._token_in_stage_chunk_count), + byte_alignment=16, + ) + if self.activation_is_quantized: + token_in_lifetime.register_tensor( + self.token_in_sf_smem_region, + self.activation_sf_dtype, + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), + stride=(self.activation_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + if not self.token_back_enabled: + return + + if self.token_back_mode == "standalone_warps": # nosec B105 + token_back_lifetime = workspace.create_overlay( + "nvlink.token_comm.main_smem.standalone_token_back" + ).add_lifetime("token_back") + else: + token_back_lifetime = transfer_overlay.add_lifetime("token_back") + + if self.token_back_mode == "standalone_warps": # nosec B105 + available_bytes_per_warp = self.standalone_chunk_bytes + else: + activation_bytes = self.bytes_per_token + sf_bytes = ( + self.activation_sf_hidden_padded * int(self.activation_sf_dtype.width) // 8 + if self.activation_is_quantized + else 0 + ) + available_bytes_per_warp = activation_bytes + sf_bytes + + if self.token_back_push_data: + bytes_per_output_token = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + if self.token_back_mode == "standalone_warps": # nosec B105 + chunk_bytes = self.standalone_chunk_bytes + elif available_bytes_per_warp < bytes_per_output_token: + chunk_bytes = self.bytes_per_token + else: + chunk_bytes = bytes_per_output_token + if self.token_back_mode != "standalone_warps" and bytes_per_output_token % chunk_bytes != 0: # nosec B105 + raise ValueError("Token-back data chunk bytes must divide one row.") + chunk_elements = chunk_bytes * 8 // int(self.combine_format.act_dtype.width) + token_back_lifetime.register_tensor( + self.token_back_activation_smem_region, + self.combine_format.act_dtype, + (self.transfer_warp_count, chunk_elements), + byte_alignment=16, + ) + if self.token_back_push_sf: + sf_row_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + if sf_row_bytes > available_bytes_per_warp: + raise ValueError("Token-back scale row exceeds its per-warp stage.") + token_back_lifetime.register_tensor( + self.token_back_sf_smem_region, + self.combine_format.scale_dtype, + (self.transfer_warp_count, (self.combine_format.scale_block, self.combine_sf_hidden_padded)), + stride=(self.combine_sf_hidden_padded, (0, 1)), + byte_alignment=16, + ) + + @cute.jit + def fc1_ready_counter_pointer(self, device_workspace: DeviceWorkspace) -> cute.Pointer: + return device_workspace.ptr(self.fc1_ready_region) + + @property + def dispatch_pool_shape(self) -> Tuple[int, int]: + """Row-major ``(pool_capacity, hidden)`` geometry of the dispatch pool, workspace-owned or external.""" + return (self.worst_case_token_count, self.hidden_size) + + @cute.jit + def fc1_activation_tensor(self, device_workspace: DeviceWorkspace) -> cute.Tensor: + if cutlass.const_expr(self.external_dispatch_pool): + raise ValueError("The dispatch pool is caller-owned under external_dispatch_pool; use the passed tensor.") + return cute.make_tensor( + device_workspace.ptr(self.fc1_activation_region), + cute.make_layout(self.dispatch_pool_shape, stride=(self.hidden_size, 1)), + ) + + @cute.jit + def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.activation_is_quantized): + return None + layout = tile_atom_to_shape_SF( + (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size + ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) + + @cute.jit + def fc2_done_counter_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_enabled): + return None + return device_workspace.tensor(self.fc2_done_region) + + @cute.jit + def fc2_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_data): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_region), + cute.make_layout( + (self.worst_case_token_count, 1, self.hidden_size), stride=(self.hidden_size, self.hidden_size, 1) + ), + ) + + @cute.jit + def fc2_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + if cutlass.const_expr(not self.token_back_push_sf): + return None + return cute.make_tensor( + device_workspace.ptr(self.fc2_activation_sf_region), + cute.make_layout( + ( + self.worst_case_token_count, + 1, + (self.combine_format.scale_block, self.hidden_size // self.combine_format.scale_block), + ), + stride=(self.combine_sf_hidden_padded, self.combine_sf_hidden_padded, (0, 1)), + ), + ) + + @cute.jit + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine plane when top-k reduction is separate.""" + if cutlass.const_expr(self.reduce_topk_in_kernel): + return None + return device_workspace.tensor(self.pre_reduced_activation_region) + + @cute.jit + def pre_reduced_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: + """Return the source-domain combine scale plane when one is required.""" + if cutlass.const_expr(self.reduce_topk_in_kernel or not self.combine_format.is_quantized): + return None + return device_workspace.tensor(self.pre_reduced_activation_sf_region) + + @cute.jit + def _load_input_dedup_route_batch( + self, routes: cute.Pointer, source_global_token: Int32, route_batch: Int32, + ) -> Int64: + """Keep one route per lane; predicate the final partial batch.""" + route_slot = route_batch * Int32(32) + self._lane_idx + record = Int64(0) + if route_slot < Int32(self.topk): + record = cute.arch.load( + routes + source_global_token * Int32(self.topk) + route_slot, Int64 + ) + return record + + @cute.jit + def _token_in_deduplicated(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """One remote TMA payload per source token/receiver group, original row outputs. + + Router writes a direct slot index before its original metadata publication. + Each slot has one source-router writer. Full tail drain precedes index reset. + The original resident communication warps guarantee progress; no new kernel, + cross-CTA waiting, mathematical operation, or reduced output is introduced. + """ + iket.range_push("token_in.deduplicated_payload") + lane_idx = self._lane_idx + transfer_warp_idx = self._transfer_warp_idx + route_batch_count = ceil_div(self.topk, 32) + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + routes = self._device_workspace.ptr(self._router.input_dedup_routes_region) + stage = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base)[transfer_warp_idx, None] + sf_stage = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base)[transfer_warp_idx, (None, None)] + barrier = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + transfer_warp_idx + if lane_idx == Int32(0): + cute.arch.mbarrier_init(barrier, 1) + cute.arch.sync_warp() + activation_bytes = cute.cosize(stage) * int(self.activation_dtype.width) // 8 + sf_bytes = cute.cosize(sf_stage) * int(self.activation_sf_dtype.width) // 8 + source_sf_values = cute.slice_(sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (4,)) + sf_atom = _copy_atom(self.activation_sf_dtype, 4 * int(self.activation_sf_dtype.width)) + destination_sf = self.fc1_activation_sf_tensor(self._device_workspace) + if cutlass.const_expr(self.external_dispatch_pool): + destination_data = self._token_comm_args.fc1_activation_pool.iterator + else: + destination_data = self._device_workspace.ptr(self.fc1_activation_region) + ready = self._device_workspace.ptr(self.fc1_ready_region) + flags = make_flag_batch_tracker(use_async=self.token_in_async_flags, flag_address=Int64(0), + accumulated_flags=Int32(0), phase=Int32(0), thread_idx=lane_idx) + pull_phase = Int32(0) + # Traverse the same dense, expert-major valid rows as original + # token_in; physical pool padding never becomes a work item. + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + pool_bases = self._router.pool_expert_base_tensor(self._device_workspace) + metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + work_limit = Int32(0) + for expert_sum in cutlass.range_constexpr(self.experts_per_rank): + work_limit = work_limit + owned_sizes[Int32(expert_sum)] + current_expert = Int32(0) + expert_valid_begin = Int32(0) + expert_valid_end = owned_sizes[Int32(0)] + work = global_warp_idx + while work < work_limit: + # work < sum(sizes) proves a nonempty expert remains, including + # leading/interior zero-size experts and the final partial wave. + while work >= expert_valid_end: + expert_valid_begin = expert_valid_end + current_expert = current_expert + Int32(1) + expert_valid_end = expert_valid_end + owned_sizes[current_expert] + pool_row = pool_bases[current_expert] + work - expert_valid_begin + source_rank_from_row = Int32(0) + source_token_from_row = Int32(0) + if lane_idx == Int32(0): + metadata = TokenSrcMetadata.load( + metadata_pointer.toint() + Int64(pool_row) * Int64(TokenSrcMetadata.nbytes) + ) + source_rank_from_row = metadata.src_rank + source_token_from_row = metadata.src_token + source_rank_from_row = Int32(cute.arch.shuffle_sync(source_rank_from_row, Int32(0))) + source_token_from_row = Int32(cute.arch.shuffle_sync(source_token_from_row, Int32(0))) + source_global_token = source_rank_from_row * Int32(self.max_tokens_per_rank) + source_token_from_row + # Every route owns a distinct destination row. Select the smallest + # row across ALL batches, preserving expert-major ownership without + # packing a local expert ID. Repeated experts still have distinct rows. + record = self._load_input_dedup_route_batch(routes, source_global_token, Int32(0)) + route_row = Int32(record & Int64(0xFFFFF)) + earlier_owner = (route_row != Int32(0)) & (route_row < pool_row + Int32(1)) + has_earlier_owner = cute.arch.vote_ballot_sync(earlier_owner) + if cutlass.const_expr(self.topk > 32): + for route_batch in cutlass.range(1, route_batch_count, unroll=1): + other = self._load_input_dedup_route_batch(routes, source_global_token, route_batch) + other_row = Int32(other & Int64(0xFFFFF)) + earlier_owner = (other_row != Int32(0)) & (other_row < pool_row + Int32(1)) + has_earlier_owner = has_earlier_owner | cute.arch.vote_ballot_sync(earlier_owner) + # The current valid row guarantees at least one local route, even + # when this receiver's first matching route is in a later batch. + if has_earlier_owner == Int32(0): + source_rank = source_global_token // Int32(self.max_tokens_per_rank) + source_token = source_global_token % Int32(self.max_tokens_per_rank) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), source_rank, Int64(0)) + source_data_address = self._token_comm_args.activation.iterator.toint() + peer_offset + Int64(source_token) * Int64(self.bytes_per_token) + remote_sf = cute.make_tensor(cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, assumed_align=self._token_comm_args.activation_sf.iterator.max_alignment), + self._token_comm_args.activation_sf.layout) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(barrier, Int32(activation_bytes + sf_bytes)) + tma_load_1d(stage.iterator, source_data_address, barrier, Int32(activation_bytes)) + tma_load_1d(sf_stage.iterator, remote_sf[Int64(source_token), None].iterator, barrier, Int32(sf_bytes)) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(barrier, pull_phase) + for route_batch in cutlass.range(route_batch_count, unroll=1): + batch_record = record + if cutlass.const_expr(self.topk > 32): + if route_batch != Int32(0): + batch_record = self._load_input_dedup_route_batch( + routes, source_global_token, route_batch + ) + for slot in cutlass.range_constexpr(min(self.topk, 32)): + lo = cute.arch.shuffle_sync(Int32(batch_record & Int64(0xFFFFFFFF)), Int32(slot)) + hi = cute.arch.shuffle_sync(Int32(batch_record >> Int64(32)), Int32(slot)) + route = (Int64(hi) << Int64(32)) | (Int64(lo) & Int64(0xFFFFFFFF)) + if route != Int64(0): + destination_row = Int32(route & Int64(0xFFFFF)) - Int32(1) + sf_row = Int32((route >> Int64(20)) & Int64(0xFFFFF)) + dst = cute.make_ptr(self.activation_dtype, + destination_data.toint() + Int64(destination_row) * Int64(self.bytes_per_token), + AddressSpace.gmem, assumed_align=16) + with cute.arch.elect_one(): + cp_async_bulk_s2g(dst, stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + sf_row_view = destination_sf[Int64(sf_row), ((None, None), None)] + sf_values = cute.group_modes(cute.slice_(sf_row_view, (0, None, None)), 0, 2) + sf_vectors = cute.zipped_divide(sf_values, (4,)) + for sf_round in cutlass.range_constexpr(ceil_div(cute.size(sf_vectors, mode=[1]), 32)): + vector = Int32(sf_round * 32) + lane_idx + if vector < Int32(cute.size(sf_vectors, mode=[1])): + cute.copy(sf_atom, source_sf_vectors[None, vector], sf_vectors[None, vector]) + # Every batch has issued its stores before any ready is + # published. All destinations are distinct pool rows. Their S2G + # copies may share one bulk async-group without write-order + # dependencies. Keep both source stages immutable until the + # full wait below; .read alone is insufficient for ready. + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + # Join every bulk issuer and every generic SF writer before + # the elected lane publishes any destination's release flag. + cute.arch.sync_warp() + for route_batch in cutlass.range(route_batch_count, unroll=1): + batch_record = record + if cutlass.const_expr(self.topk > 32): + if route_batch != Int32(0): + batch_record = self._load_input_dedup_route_batch( + routes, source_global_token, route_batch + ) + for slot in cutlass.range_constexpr(min(self.topk, 32)): + lo = cute.arch.shuffle_sync(Int32(batch_record & Int64(0xFFFFFFFF)), Int32(slot)) + hi = cute.arch.shuffle_sync(Int32(batch_record >> Int64(32)), Int32(slot)) + route = (Int64(hi) << Int64(32)) | (Int64(lo) & Int64(0xFFFFFFFF)) + if route != Int64(0): + ready_slot = Int32((route >> Int64(40)) & Int64(0xFFFFF)) + flags = flags.accumulate( + Int32(0), self.token_in_flag_batch, (ready + ready_slot).toint() + ) + cute.arch.sync_warp() + pull_phase = pull_phase ^ Int32(1) + work = work + global_warp_count + flags.fire() + cute.arch.sync_warp() + iket.range_pop() + + @cute.jit + def token_in( + self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, + ) -> None: + """Wait for pushed metadata, then pull activation payloads into local pools.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + global_warp_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_warp_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + + sizes = self._router.sizes_tensor(self._device_workspace) + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + + iket.range_push("token_in.wait_sizes_ready") + self._router.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("token_in.stage_sizes") + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + owner_expert_begin = self._local_rank * Int32(self.experts_per_rank) + source_sizes = cute.make_tensor(sizes.iterator + owner_expert_begin, cute.make_layout((self.experts_per_rank,))) + copy_elements = 4 if self.experts_per_rank % 4 == 0 else 1 + source_size_vectors = cute.zipped_divide( + (_mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), + (copy_elements,), + ) + destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) + size_vector_count = cute.size(destination_size_vectors, mode=[1]) + size_copy_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS + ), + cutlass.Int32, + num_bits_per_copy=copy_elements * 32, + ) + size_copy_rounds = ceil_div(size_vector_count, self.transfer_thread_count) + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + for size_copy_round in cutlass.range_constexpr(size_copy_rounds): + vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx + if vector_idx < Int32(size_vector_count): + cute.copy( + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] + ) + cute.arch.cp_async_commit_group() + iket.range_pop() + + iket.range_push("token_in.wait_metadata_ready") + self._router.wait_for_metadata_ready(self._device_workspace) + iket.range_pop() + + cute.arch.cp_async_wait_group(0) + iket.range_push("token_in.size_barrier") + token_in_size_barrier = pipeline.NamedBarrier( + barrier_id=self.token_in_size_barrier_id, num_threads=self.transfer_thread_count + ) + token_in_size_barrier.arrive_and_wait() + iket.range_pop() + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive() + + if cutlass.const_expr(self.input_dedup_mode != "off"): + self._token_in_deduplicated(smem_workspace, smem_base) + if cutlass.const_expr(self.input_dedup_mode == "off"): + iket.range_push("token_in.pull_payload") + token_in_mbarriers = smem_workspace.ptr(self.token_in_mbarrier_region, smem_base) + token_in_activation = smem_workspace.tensor(self.token_in_activation_smem_region, smem_base) + warp_mbarrier = token_in_mbarriers + transfer_warp_idx + warp_activation_stage = token_in_activation[transfer_warp_idx, None] + if cutlass.const_expr(self.activation_is_quantized): + token_in_sf = smem_workspace.tensor(self.token_in_sf_smem_region, smem_base) + warp_sf_stage = token_in_sf[transfer_warp_idx, (None, None)] + if lane_idx == Int32(0): + cute.arch.mbarrier_init(warp_mbarrier, 1) + cute.arch.sync_warp() + + if cutlass.const_expr(self.external_dispatch_pool): + fc1_activation_pointer = self._token_comm_args.fc1_activation_pool.iterator + else: + fc1_activation_pointer = self._device_workspace.ptr(self.fc1_activation_region) + fc1_ready_counter = self._device_workspace.ptr(self.fc1_ready_region) + activation_bytes = cute.cosize(warp_activation_stage) * int(self.activation_dtype.width) // 8 + activation_sf_bytes = 0 + if cutlass.const_expr(self.activation_is_quantized): + fc1_activation_sf = self.fc1_activation_sf_tensor(self._device_workspace) + activation_sf_bytes = cute.cosize(warp_sf_stage) * int(self.activation_sf_dtype.width) // 8 + sf_copy_elements = 4 + source_sf_values = cute.slice_(warp_sf_stage, (0, None)) + source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) + + next_dense_token = global_warp_idx + expert_valid_begin = Int32(0) + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = Int32(0) + expert_ready_slot_begin = Int32(0) + pull_phase = Int32(0) + flag_tracker = make_flag_batch_tracker( + use_async=self.token_in_async_flags, + flag_address=Int64(0), + accumulated_flags=Int32(0), + phase=Int32(0), + thread_idx=lane_idx, + ) + + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + pool_expert_base = pool_expert_bases[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + pull_count = Int32(0) + if next_dense_token < expert_valid_end: + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count + + for pull_round in cutlass.range(pull_count, unroll=1): + dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count + token_in_expert = dense_token_idx - expert_valid_begin + pool_token_idx = pool_expert_base + token_in_expert + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) + + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + remote_activation_address = ( + self._token_comm_args.activation.iterator.toint() + + peer_offset + + Int64(metadata.src_token) * Int64(self.bytes_per_token) + ) + if cutlass.const_expr(self.activation_is_quantized): + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.activation_sf.dtype, + self._token_comm_args.activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.activation_sf.iterator.max_alignment), + ), + self._token_comm_args.activation_sf.layout, + ) + remote_sf_row = remote_sf[Int64(metadata.src_token), None] + + if cutlass.const_expr(self.token_in_stage_split): + # Chunked pull: the halved stage bounces each token row in two pieces, + # trading one extra NVLink round trip per token for SMEM. + for chunk_idx in cutlass.range_constexpr(self._token_in_stage_chunk_count): + chunk_offset = Int64(chunk_idx) * Int64(activation_bytes) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_bytes)) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address) + chunk_offset, + warp_mbarrier, + Int32(activation_bytes), + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + pull_phase = pull_phase ^ Int32(1) + + chunk_destination = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + + Int64(pool_token_idx) * Int64(self.bytes_per_token) + + chunk_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(chunk_destination, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + # Drain before the stage is reused by the next chunk (and before the flag). + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + else: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + warp_mbarrier, Int32(activation_bytes + activation_sf_bytes) + ) + tma_load_1d( + warp_activation_stage.iterator, + Int64(remote_activation_address), + warp_mbarrier, + Int32(activation_bytes), + ) + if cutlass.const_expr(self.activation_is_quantized): + tma_load_1d( + warp_sf_stage.iterator, remote_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) + ) + cute.arch.sync_warp() + cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) + + destination_activation = cute.make_ptr( + self.activation_dtype, + fc1_activation_pointer.toint() + Int64(pool_token_idx) * Int64(self.bytes_per_token), + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_activation, warp_activation_stage.iterator, Int32(activation_bytes)) + cute.arch.sync_warp() + cute.arch.cp_async_bulk_commit_group() + + if cutlass.const_expr(self.activation_is_quantized): + sf_token_idx = expert_sf_begin + token_in_expert + destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] + destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) + destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) + sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) + for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): + sf_vector_idx = Int32(sf_round * 32) + lane_idx + if sf_vector_idx < Int32(sf_vector_count): + cute.copy( + sf_copy_atom, + source_sf_vectors[None, sf_vector_idx], + destination_sf_vectors[None, sf_vector_idx], + ) + + if cutlass.const_expr(self.token_in_flag_batch == 1): + cute.arch.cp_async_bulk_wait_group(0) + elif cutlass.const_expr(not self.token_in_stage_split): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.sync_warp() + ready_address = (fc1_ready_counter + ready_slot_idx).toint() + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) + cute.arch.sync_warp() + if cutlass.const_expr(not self.token_in_stage_split): + pull_phase = pull_phase ^ Int32(1) + + next_dense_token = next_dense_token + pull_count * global_warp_count + expert_valid_begin = expert_valid_end + if cutlass.const_expr(self.activation_is_quantized): + expert_sf_begin = expert_sf_begin + ( + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) + ) * Int32(self.sf_padding_block) + expert_ready_slot_begin = expert_ready_slot_begin + ( + (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) + // Int32(self.tokens_per_fc1_ready_slot) + ) + local_expert = local_expert + Int32(1) + + flag_tracker.fire() + cute.arch.sync_warp() + iket.range_pop() + if cutlass.const_expr(self.token_back_enabled and self.token_back_mode != "standalone_warps"): # nosec B105 + iket.range_push("token_in.transfer_barrier") + transfer_lifetime_barrier = pipeline.NamedBarrier( + barrier_id=self.transfer_lifetime_barrier_id, num_threads=self.transfer_thread_count + ) + transfer_lifetime_barrier.arrive_and_wait() + iket.range_pop() + + @cute.jit + def _stateless_pace(self, reference_window: Int32, current_window: Int32) -> None: + sleep_cycles = Int32(0) + if current_window < reference_window: + sleep_cycles = reference_window - current_window + elif current_window > reference_window: + sleep_cycles = cutlass.min(current_window - reference_window, Int32(self.standalone_max_backoff_cycles)) + if sleep_cycles > Int32(0): + nanosleep(sleep_cycles) + + @cute.jit + def _adaptive_pace(self, average_window: Int32, current_window: Int32, low_window: int, high_window: int) -> Int32: + sleep_cycles = Int32(0) + if current_window > average_window: + average_window = average_window + ((current_window - average_window + Int32(3)) // Int32(4)) + sleep_cycles = current_window - average_window + if sleep_cycles > Int32(high_window): + sleep_cycles = Int32(high_window) + else: + average_window = average_window - ((average_window - current_window + Int32(3)) // Int32(4)) + sleep_cycles = average_window - current_window + if sleep_cycles > Int32(self.adaptive_minimum_sleep_cycles): + nanosleep(sleep_cycles) + if average_window > Int32(high_window): + average_window = Int32(high_window) + if average_window < Int32(low_window): + average_window = Int32(low_window) + return average_window + + @cute.jit + def token_back(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: + """Push completed FC2 data and scale rows to source ranks.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + if cutlass.const_expr(not self.token_back_enabled): + return + if cutlass.const_expr( + self.combine_format.is_quantized and self._token_comm_args.pre_reduced_activation_sf is None + ): + raise ValueError("Quantized token-back requires a scale destination.") + + global_worker_idx = self._linear_cta_idx * Int32(self.transfer_warp_count) + transfer_warp_idx + global_worker_count = Int32(self.promised_launchable_sm_count * self.transfer_warp_count) + token_back_mbarriers = smem_workspace.ptr(self.token_back_mbarrier_region, smem_base) + worker_mbarrier = token_back_mbarriers + transfer_warp_idx + if cutlass.const_expr(self.token_back_push_data): + token_back_activation = smem_workspace.tensor(self.token_back_activation_smem_region, smem_base) + worker_activation_stage = token_back_activation[transfer_warp_idx, None] + activation_chunk_bytes = ( + cute.cosize(worker_activation_stage) * int(self.combine_format.act_dtype.width) // 8 + ) + if cutlass.const_expr(self.token_back_push_sf): + token_back_sf = smem_workspace.tensor(self.token_back_sf_smem_region, smem_base) + worker_sf_stage = token_back_sf[transfer_warp_idx, (None, None)] + sf_chunk_bytes = cute.cosize(worker_sf_stage) * int(self.combine_format.scale_dtype.width) // 8 + if lane_idx == Int32(0): + cute.arch.mbarrier_init(worker_mbarrier, 1) + cute.arch.sync_warp() + + owned_sizes = smem_workspace.tensor(self.expert_sizes_smem_region, smem_base) + if cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + sizes_ready_barrier = pipeline.NamedBarrier( + barrier_id=self.standalone_size_barrier_id, num_threads=2 * self.transfer_thread_count + ) + sizes_ready_barrier.arrive_and_wait() + + pool_expert_bases = self._router.pool_expert_base_tensor(self._device_workspace) + token_metadata_pointer = self._router.token_src_metadata_pointer(self._device_workspace) + fc2_done = self._device_workspace.ptr(self.fc2_done_region) + if cutlass.const_expr(self.token_back_push_data): + fc2_activation_pointer = self._device_workspace.ptr(self.fc2_activation_region) + output_token_bytes = self.hidden_size * int(self.combine_format.act_dtype.width) // 8 + activation_chunk_count = ceil_div(output_token_bytes, activation_chunk_bytes) + data_window_unit = ceil_div(activation_chunk_bytes * 2, 3) + reuse_data_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and data_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + # Preserve the empirical low:initial:high ratio of 1:2.5:5. + data_average_window = Int32(data_window_unit) + data_low_window = data_window_unit * 2 // 5 + data_high_window = data_window_unit * 2 + print(f"[{data_window_unit}, {data_low_window}, {data_high_window}]") + if cutlass.const_expr(self.token_back_push_sf): + fc2_sf_pointer = self._device_workspace.ptr(self.fc2_activation_sf_region) + output_sf_bytes = self.combine_sf_hidden_padded * int(self.combine_format.scale_dtype.width) // 8 + sf_chunk_count = ceil_div(output_sf_bytes, sf_chunk_bytes) + sf_window_unit = ceil_div(sf_chunk_bytes * 2, 3) + reuse_sf_pacing_enabled = ( + self.token_back_mode == "reuse_dispatch_warps" and sf_window_unit > self.minimum_pacing_window_cycles # nosec B105 + ) + sf_average_window = Int32(sf_window_unit) + sf_low_window = sf_window_unit * 2 // 5 + sf_high_window = sf_window_unit * 2 + + next_dense_token = global_worker_idx - global_worker_count + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + next_dense_token = Int32(0) + next_dense_token = self.next_token(next_dense_token) + expert_valid_begin = Int32(0) + transfer_phase = Int32(0) + + iket.range_push("token_back.work") + local_expert = Int32(0) + while local_expert < Int32(self.experts_per_rank): + expert_token_count = owned_sizes[local_expert] + expert_valid_end = expert_valid_begin + expert_token_count + if next_dense_token < expert_valid_end: + token_tile_count = (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) // Int32( + self.tokens_per_fc1_ready_slot + ) + completion_target = token_tile_count * Int32(self.fc2_done_signals_per_token_tile) + iket.range_push("token_back.wait_fc2") + while cute.arch.load(fc2_done + local_expert, Int32, sem="acquire", scope="gpu") < completion_target: + nanosleep(500) + iket.range_pop() + + while next_dense_token < expert_valid_end: + token_in_expert = next_dense_token - expert_valid_begin + pool_token_idx = pool_expert_bases[local_expert] + token_in_expert + metadata = TokenSrcMetadata.load( + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(TokenSrcMetadata.nbytes) + ) + destination_topk = metadata.src_topk + if cutlass.const_expr(self.reduce_topk_in_kernel): + destination_topk = Int32(0) + peer_offset = self._token_comm_args.peer_rank_ptr_mapper.map(Int64(0), metadata.src_rank, Int64(0)) + is_remote_token = metadata.src_rank != self._local_rank + + if cutlass.const_expr(self.token_back_push_data): + iket.range_push("token_back.push_data") + local_activation_address = fc2_activation_pointer.toint() + Int64(pool_token_idx) * Int64( + output_token_bytes + ) + remote_activation = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation.dtype, + self._token_comm_args.pre_reduced_activation.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation.layout, + ) + destination_row = remote_activation[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(activation_chunk_count): + chunk_byte_offset = Int64(chunk_idx * activation_chunk_bytes) + chunk_bytes_this_round = min( + activation_chunk_bytes, output_token_bytes - chunk_idx * activation_chunk_bytes + ) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_data_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_data_pacing_enabled or stateless_data_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_activation_stage.iterator, + local_activation_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + if cutlass.const_expr(self.reduce_topk_in_kernel): + cp_reduce_async_bulk_add_bf16_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + else: + cp_async_bulk_s2g( + destination_chunk, worker_activation_stage.iterator, current_chunk_bytes + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + data_average_window = self._adaptive_pace( + data_average_window, current_window, data_low_window, data_high_window + ) + elif cutlass.const_expr(stateless_data_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + if cutlass.const_expr(self.token_back_push_sf): + iket.range_push("token_back.push_sf") + local_sf_address = fc2_sf_pointer.toint() + Int64(pool_token_idx) * Int64(output_sf_bytes) + remote_sf = cute.make_tensor( + cute.make_ptr( + self._token_comm_args.pre_reduced_activation_sf.dtype, + self._token_comm_args.pre_reduced_activation_sf.iterator.toint() + peer_offset, + AddressSpace.gmem, + assumed_align=(self._token_comm_args.pre_reduced_activation_sf.iterator.max_alignment), + ), + self._token_comm_args.pre_reduced_activation_sf.layout, + ) + destination_sf_row = remote_sf[Int64(metadata.src_token), destination_topk, None] + for chunk_idx in cutlass.range_constexpr(sf_chunk_count): + chunk_byte_offset = Int64(chunk_idx * sf_chunk_bytes) + chunk_bytes_this_round = min(sf_chunk_bytes, output_sf_bytes - chunk_idx * sf_chunk_bytes) + current_chunk_bytes = Int32(chunk_bytes_this_round) + current_window_unit = ceil_div(chunk_bytes_this_round * 2, 3) + stateless_sf_pacing_enabled = ( + self.token_back_mode != "reuse_dispatch_warps" # nosec B105 + and current_window_unit > self.minimum_pacing_window_cycles + ) + round_start_clock = Int64(0) + if cutlass.const_expr(reuse_sf_pacing_enabled or stateless_sf_pacing_enabled): + if is_remote_token: + round_start_clock = read_clock64() + else: + round_start_clock = round_start_clock + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(worker_mbarrier, current_chunk_bytes) + tma_load_1d( + worker_sf_stage.iterator, + local_sf_address + chunk_byte_offset, + worker_mbarrier, + current_chunk_bytes, + ) + cute.arch.mbarrier_wait(worker_mbarrier, transfer_phase) + destination_chunk = cute.make_ptr( + cutlass.Uint8, + destination_sf_row.iterator.toint() + chunk_byte_offset, + AddressSpace.gmem, + assumed_align=16, + ) + with cute.arch.elect_one(): + cp_async_bulk_s2g(destination_chunk, worker_sf_stage.iterator, current_chunk_bytes) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + transfer_phase = transfer_phase ^ Int32(1) + if cutlass.const_expr(reuse_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + sf_average_window = self._adaptive_pace( + sf_average_window, current_window, sf_low_window, sf_high_window + ) + elif cutlass.const_expr(stateless_sf_pacing_enabled): + if is_remote_token: + current_window = Int32(read_clock64() - round_start_clock) + self._stateless_pace(Int32(current_window_unit), current_window) + iket.range_pop() + + cute.arch.sync_warp() + next_dense_token = self.next_token(next_dense_token) + + expert_valid_begin = expert_valid_end + local_expert = local_expert + Int32(1) + iket.range_pop() + + @cute.jit + def next_token(self, current_token: Int32) -> Int32: + global_worker_count = self.promised_launchable_sm_count * self.transfer_warp_count + schedule_counter = None + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + schedule_counter = self._device_workspace.ptr(self.token_back_schedule_region) + if cutlass.const_expr(self.token_back_schedule_mode == "atomic_counter"): # nosec B105 + claimed_token = Int32(0) + if self._lane_idx == Int32(0): + claimed_token = cute.arch.atomic_add(schedule_counter, Int32(1), sem="relaxed", scope="gpu") + return Int32(cute.arch.shuffle_sync(claimed_token, Int32(0))) + return current_token + global_worker_count + + @cute.jit + def reset_tail(self) -> None: + """Reset communication state with all configured token-in transfer warps.""" + transfer_warp_idx = self._transfer_warp_idx + lane_idx = self._lane_idx + transfer_thread_idx = transfer_warp_idx * Int32(32) + lane_idx + iket.range_push("tail.nvlink_drain") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx + iket.range_push("tail.reset_workspace") + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + iket.range_push("tail.nvlink_publish") + self._nvlink_barrier.arrive_and_wait( + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + prologue_grid_sync=False, + epilogue_grid_sync=True, + ) + iket.range_pop() + if cutlass.const_expr(os.environ.get("MEGA_USE_NCU", "0") == "1"): + iket.range_push("tail.ncu_finalize") + self._nvlink_barrier.finalize( + 2, + self.transfer_thread_count, + Int32(self.promised_launchable_sm_count), + self._linear_cta_idx, + transfer_thread_idx, + ) + iket.range_pop() + + +__all__ = ["TokenCommDeterministic"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py index 8ed3b33dd75d..e7f5e1252bfe 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/communication/token_protocol.py @@ -9,24 +9,45 @@ from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import Int32, Int64 - @dataclasses.dataclass(frozen=True) class TokenSrcMetadata: - """One i64 routing record: rank:u16, topk:u16, token:u32.""" + """One i64 routing record: ``token:u16 | topk:u8 | rank:u8 | dst_order_id:u8 | (dst_token_cnt-1):u8 | reserved:u16``. + . ``dst_token_cnt`` = the token's routes landing on the destination rank, ``dst_order_id`` = this + route's position among them in ascending expert order (``order == cnt - 1`` is the last copy). Producers that do + not track copies use order 0, count 1.""" src_rank: Int32 src_token: Int32 src_topk: Int32 + dst_order_id: Int32 = Int32(0) + dst_token_cnt: Int32 = Int32(1) nbytes: ClassVar[int] = 8 + FIELD_BITS: ClassVar[int] = 8 + TOKEN_BITS: ClassVar[int] = 16 + MAX_RANKS: ClassVar[int] = 1 << 8 + MAX_TOPK: ClassVar[int] = 1 << 8 + MAX_TOKENS: ClassVar[int] = 1 << 16 def pack(self) -> Int64: - high = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) - return (high << Int64(32)) | (Int64(self.src_token) & Int64(0xFFFFFFFF)) + high = Int64(self.dst_order_id) | ((Int64(self.dst_token_cnt) - Int64(1)) << Int64(8)) + low = ( + (Int64(self.src_rank) << Int64(24)) + | (Int64(self.src_topk) << Int64(16)) + | (Int64(self.src_token) & Int64(0xFFFF)) + ) + return (high << Int64(32)) | low + + @property + def is_last_copy(self) -> Int32: + """1 when this row is the destination rank's last copy of the source token.""" + return Int32(self.dst_order_id == self.dst_token_cnt - Int32(1)) @staticmethod def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: - raw_address = address if isinstance(address, Int64) else address.toint() + raw_address = ( + address if isinstance(address, Int64) else address.toint() + ) return cute.make_ptr( Int64, raw_address, @@ -37,18 +58,24 @@ def _pointer(address: Union[cute.Pointer, Int64]) -> cute.Pointer: def store(self, address: Union[cute.Pointer, Int64]) -> None: cute.arch.store(self._pointer(address), self.pack(), scope="gpu") + @classmethod + def unpack(cls, packed: Int64) -> "TokenSrcMetadata": + """Inverse of ``pack`` for a record already in a register; every reader must decode through here.""" + high = packed >> Int64(32) + return cls( + src_rank=Int32((packed >> Int64(24)) & Int64(0xFF)), + src_token=Int32(packed & Int64(0xFFFF)), + src_topk=Int32((packed >> Int64(16)) & Int64(0xFF)), + dst_order_id=Int32(high & Int64(0xFF)), + dst_token_cnt=Int32((high >> Int64(8)) & Int64(0xFF)) + Int32(1), + ) + @classmethod def load( cls, address: Union[cute.Pointer, Int64], ) -> "TokenSrcMetadata": - packed = Int64(cute.arch.load(cls._pointer(address), Int64, scope="gpu")) - high = packed >> Int64(32) - return cls( - src_rank=Int32((high >> Int64(16)) & Int64(0xFFFF)), - src_token=Int32(packed & Int64(0xFFFFFFFF)), - src_topk=Int32(high & Int64(0xFFFF)), - ) + return cls.unpack(Int64(cute.arch.load(cls._pointer(address), Int64, scope="gpu"))) __all__ = ["TokenSrcMetadata"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py index 52a7a9daf028..0e7d0e5264dc 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/__init__.py @@ -1,2 +1,47 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Low-level workspace, synchronization, and PTX helpers.""" + +from .device_workspace import DeviceWorkspace +from .dsl_helpers import spin_peek, spin_wait +from .flag_batch import GpuAsyncReleaseFlagBatchTracker, GpuReleaseFlagBatchTracker, make_flag_batch_tracker +from .iket_compat import iket +from .ptx_helpers import ( + cvt_f32_to_fp8_to_f32, + cvt_f32x4_to_f8x4_pack_i32, + stg_e8m0_from_f32, + stg_e8m0x8_from_f32, +) +from .smem_workspace import SmemWorkspace +from .software_sync import NvlinkBarrier, SoftwareGridSync +from .utils import ( + IntegerType, + ceil_div, + cosize_from_shape_stride_tuples, + product, + row_major_stride, + round_up, + validate_static_integer_tuple, +) + +__all__ = [ + "DeviceWorkspace", + "GpuAsyncReleaseFlagBatchTracker", + "GpuReleaseFlagBatchTracker", + "IntegerType", + "NvlinkBarrier", + "SmemWorkspace", + "SoftwareGridSync", + "ceil_div", + "cosize_from_shape_stride_tuples", + "cvt_f32_to_fp8_to_f32", + "cvt_f32x4_to_f8x4_pack_i32", + "make_flag_batch_tracker", + "product", + "row_major_stride", + "round_up", + "iket", + "spin_peek", + "spin_wait", + "stg_e8m0_from_f32", + "stg_e8m0x8_from_f32", + "validate_static_integer_tuple", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py index ce2c01076f54..7870293a0bf7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/constants.py @@ -1,8 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Specification-defined numeric constants shared across kernel components.""" + Log2E = 1.4426950408889634 Fp32Max = 3.40282346638528859812e38 diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py index 625075c71335..8cf623735b84 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/cute_py_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - from dataclasses import dataclass from math import gcd from typing import Optional, Tuple, Union @@ -11,6 +8,7 @@ from .smem_workspace import SmemRegion + # Every block-scaled tcgen05 MMA kind accumulates in F32 and no other accumulator is legal, so # this belongs to the instruction family rather than being a caller choice. tcgen05_block_scaled_acc_dtype = cutlass.Float32 @@ -54,9 +52,7 @@ def make_tcgen05_tmem_plan( accumulators; other supported cases maximize disjoint accumulator stages. """ if mma_instruction.participates not in (1, 2): - raise ValueError( - f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}." - ) + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_instruction.participates}.") if len(mma_instruction.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") if any(dimension <= 0 for dimension in (*mma_instruction.instruction_mnk, *mma_tiler_mnk)): @@ -64,10 +60,7 @@ def make_tcgen05_tmem_plan( instruction_m, instruction_n, instruction_k = mma_instruction.instruction_mnk tile_m, tile_n, tile_k = mma_tiler_mnk - if ( - instruction_m % mma_instruction.participates != 0 - or instruction_n % mma_instruction.participates != 0 - ): + if instruction_m % mma_instruction.participates != 0 or instruction_n % mma_instruction.participates != 0: raise ValueError("MMA instruction M and N must be divisible by participates.") if tile_m % instruction_m != 0 or tile_n % instruction_n != 0 or tile_k % instruction_k != 0: raise ValueError("MMA instruction dimensions must divide mma_tiler_mnk.") @@ -97,9 +90,7 @@ def make_tcgen05_tmem_plan( tmem_column_capacity = cute.arch.get_max_tmem_alloc_cols(arch) accumulator_stage_columns = tile_n scale_factor_columns = sfa_columns + sfb_columns - accumulator_stage_count = ( - tmem_column_capacity - scale_factor_columns - ) // accumulator_stage_columns + accumulator_stage_count = (tmem_column_capacity - scale_factor_columns) // accumulator_stage_columns if accumulator_stage_count < 1: raise ValueError("Scale-factor TMEM leaves no accumulator stage.") @@ -161,9 +152,7 @@ def _round_tmem_allocation_columns(used_columns: int, capacity_columns: int) -> else: allocation_columns = _round_up(used_columns, 32) if allocation_columns > capacity_columns: - raise ValueError( - f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}." - ) + raise ValueError(f"TMEM plan needs {allocation_columns} columns, exceeding {capacity_columns}.") return allocation_columns @@ -192,9 +181,7 @@ def make_smem_layouts( arch: str, ) -> Union[Tuple[SmemRegion, SmemRegion], Tuple[SmemRegion, SmemRegion, SmemRegion, SmemRegion]]: """Derive workspace-ready TCGen05 operand regions without an MLIR context.""" - has_scale_factors = any( - value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) - ) + has_scale_factors = any(value is not None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size)) if has_scale_factors and any( value is None for value in (mma_inst.sfa_type, mma_inst.sfb_type, mma_inst.sf_vec_size) ): @@ -211,9 +198,7 @@ def make_smem_layouts( raise ValueError(f"All stage counts must be positive, got {operand_stages}.") if mma_inst.participates not in (1, 2): - raise ValueError( - f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}." - ) + raise ValueError(f"TCGen05 MMA participates must be one or two, got {mma_inst.participates}.") if len(mma_inst.instruction_mnk) != 3 or len(mma_tiler_mnk) != 3: raise ValueError("instruction_mnk and mma_tiler_mnk must each contain three dimensions.") if any(dimension <= 0 for dimension in (*mma_inst.instruction_mnk, *mma_tiler_mnk)): @@ -401,12 +386,8 @@ def _tiled_divide_2d(shape, stride, tile_shape): ): raise ValueError("A 2-D tiled divide requires rank-2 shape, stride, and tile.") - mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile( - shape[0], stride[0], tile_shape[0] - ) - k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile( - shape[1], stride[1], tile_shape[1] - ) + mn_tile_stride, mn_rest_shape, mn_rest_stride = _divide_mode_by_tile(shape[0], stride[0], tile_shape[0]) + k_tile_stride, k_rest_shape, k_rest_stride = _divide_mode_by_tile(shape[1], stride[1], tile_shape[1]) return ( ((tile_shape[0], tile_shape[1]), mn_rest_shape, k_rest_shape), ((mn_tile_stride, k_tile_stride), mn_rest_stride, k_rest_stride), @@ -471,9 +452,7 @@ def _make_ab_region( mma_mn_shape = _prefix_tile_profile(mn_shape, value_shape[0]) mma_k_shape = _prefix_tile_profile(k_shape, value_shape[1]) - mma_shape, mma_stride = _tiled_divide_2d( - (mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape) - ) + mma_shape, mma_stride = _tiled_divide_2d((mn_shape, k_shape), (mn_stride, k_stride), (mma_mn_shape, mma_k_shape)) return SmemRegion( name="", kind="tensor", @@ -512,10 +491,7 @@ def _make_sf_region( full_k_stride = ((0, 1), 512 if basic_chunk_repetitions > 1 else 0) mma_mn_shape = ((32, 4), instruction_mn_blocks) - mma_k_shape = ( - sf_vec_size, - _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size), - ) + mma_k_shape = (sf_vec_size, _prefix_tile_profile((4, basic_chunk_repetitions), instruction_k // sf_vec_size)) mma_shape, mma_stride = _tiled_divide_2d( (full_mn_shape, full_k_shape), (full_mn_stride, full_k_stride), (mma_mn_shape, mma_k_shape) ) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py index ec982a384ac9..df72be3e4037 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/device_workspace.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Declarative GMEM workspace shared by host layout and device access.""" import dataclasses @@ -18,11 +15,12 @@ is_nested_shape, is_power_of_two, ordered_stride, - round_up, row_major_stride, + round_up, validate_static_integer_tuple, ) + BufferResetAttr = Literal["data", "zero_on_first_allocate", "tail_reset"] BufferSpace = Literal["local", "shared"] @@ -144,18 +142,14 @@ def finalize(self) -> None: if self._finalized: raise RuntimeError("DeviceWorkspace.finalize() may only be called once.") for buffer_space in ("local", "shared"): - ordered = sorted( - self._registered[buffer_space], key=lambda region: _reset_order[region.reset] - ) + ordered = sorted(self._registered[buffer_space], key=lambda region: _reset_order[region.reset]) cursor = 0 for region_index, region in enumerate(ordered): stride = region.stride if stride is None: raise RuntimeError(f"Region {region.name!r} stride was not resolved.") cosize = cosize_from_shape_stride_tuples(region.shape, stride) - nbytes = ( - _footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7 - ) // 8 + nbytes = (_footprint_from_shape_stride(region.shape, stride) * int(region.dtype.width) + 7) // 8 if region.reset == "tail_reset" and ( region_index == 0 or ordered[region_index - 1].reset != "tail_reset" ): @@ -170,8 +164,7 @@ def finalize(self) -> None: self._byte_alignment[region.name] = byte_alignment cursor += nbytes is_last_tail_reset_region = region.reset == "tail_reset" and ( - region_index + 1 == len(ordered) - or ordered[region_index + 1].reset != "tail_reset" + region_index + 1 == len(ordered) or ordered[region_index + 1].reset != "tail_reset" ) if is_last_tail_reset_region: cursor = round_up(cursor, 16) @@ -246,16 +239,12 @@ def ptr(self, name: str) -> cute.Pointer: region = self._region_by_name[name] base = self._base[region.buffer_space] address = base.toint() + Int64(self._offset[name]) - return cute.make_ptr( - region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name] - ) + return cute.make_ptr(region.dtype, address, AddressSpace.gmem, assumed_align=self._byte_alignment[name]) @cute.jit def tensor(self, name: str) -> cute.Tensor: region = self._region_by_name[name] - return cute.make_tensor( - self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name]) - ) + return cute.make_tensor(self.ptr(name), cute.make_layout(region.shape, stride=self._stride[name])) @cute.jit def reset_tail(self, tid: Int32, total_threads: int) -> None: @@ -267,17 +256,13 @@ def reset_tail_space(self, buffer_space: BufferSpace, tid: Int32, total_threads: num_vectors = self._tail_leading[buffer_space] // 16 if cutlass.const_expr(num_vectors > 0): vectors = cute.make_tensor( - cute.make_ptr( - Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16 - ), + cute.make_ptr(Int32, self._base[buffer_space].toint(), AddressSpace.gmem, assumed_align=16), cute.make_layout((num_vectors, 4), stride=(4, 1)), ) zero = cute.make_rmem_tensor((4,), Int32) for element in cutlass.range_constexpr(4): zero[element] = Int32(0) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Int32, num_bits_per_copy=128) reset_round_count = ceil_div(num_vectors, total_threads) for reset_round in cutlass.range_constexpr(reset_round_count): vector_index = Int32(reset_round * total_threads) + tid diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py new file mode 100644 index 000000000000..596362cf6bb9 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dlb_expert_count.py @@ -0,0 +1,30 @@ +"""Resolve logical experts and helper capacity at the host construction boundary.""" + +from ..api import ImplDesc, ProblemDesc + + +def resolve_dlb_expert_counts( + problem_desc: ProblemDesc, impl_desc: ImplDesc, *, world_size: int +) -> tuple[ProblemDesc, ImplDesc]: + """Normalize the new logical-count API while retaining legacy physical counts.""" + if "total_helper_slots" not in impl_desc: + return problem_desc, impl_desc + if "helper_expert_count" in impl_desc: + raise ValueError("Use total_helper_slots or legacy helper_expert_count, not both.") + + logical_experts = problem_desc["expert_count"] + total_helpers = impl_desc["total_helper_slots"] + if type(world_size) is not int or world_size <= 0: + raise ValueError("world_size must be a positive exact int.") + if type(logical_experts) is not int or logical_experts <= 0: + raise ValueError("Logical expert_count must be a positive exact int.") + if type(total_helpers) is not int or total_helpers < 0: + raise ValueError("total_helper_slots must be a non-negative exact int.") + if logical_experts % world_size or total_helpers % world_size: + raise ValueError("Logical expert_count and total_helper_slots must be divisible by world_size.") + + physical_problem = ProblemDesc({**problem_desc, "expert_count": logical_experts + total_helpers}) + physical_impl = dict(impl_desc) + del physical_impl["total_helper_slots"] + physical_impl["helper_expert_count"] = total_helpers // world_size + return physical_problem, ImplDesc(physical_impl) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py index bdb83e4b0641..807d85c27b59 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/dsl_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """General-purpose CuTe DSL helpers.""" from typing import Callable, Literal, Optional @@ -28,12 +25,8 @@ def smem_exclusive_prefix( num_elements = cute.size(input_tensor) scan_rows = num_elements // 4 num_warps = block_thread_count // 32 - input_vectors = cute.make_tensor( - input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1)) - ) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) values = cute.make_rmem_tensor((4,), cutlass.Int32) carry = Int32(0) @@ -51,12 +44,7 @@ def smem_exclusive_prefix( for element in cutlass.range_constexpr(4): values[element] = Int32(0) - local_prefix = ( - Int32(0), - values[0], - values[0] + values[1], - values[0] + values[1] + values[2], - ) + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) lane_total = local_prefix[3] + values[3] inclusive = lane_total for step_log in cutlass.range_constexpr(5): @@ -95,17 +83,12 @@ def smem_exclusive_prefix( def mark_alignment(tensor: cute.Tensor, byte_alignment: int) -> cute.Tensor: pointer = tensor.iterator return cute.make_tensor( - cute.make_ptr( - pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment - ), - tensor.layout, + cute.make_ptr(pointer.dtype, pointer.toint(), pointer.memspace, assumed_align=byte_alignment), tensor.layout ) @cute.jit -def spin_peek( - pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu" -) -> Boolean: +def spin_peek(pointer: Pointer, condition: Callable[[Int32], Boolean], scope: str = "gpu") -> Boolean: """Perform one acquire load and test its value.""" value = cute.arch.load(pointer, pointer.dtype, sem="acquire", scope=scope) return Boolean(condition(value)) @@ -138,10 +121,7 @@ def _tma_multicast_pattern( if tensor_role in ("a", "sfa"): return sum(1 << (cluster_n_index * cluster_m) for cluster_n_index in range(cluster_n)) if tensor_role == "b": - return sum( - 1 << (cluster_m_index * mma_cta_count) - for cluster_m_index in range(cluster_m // mma_cta_count) - ) + return sum(1 << (cluster_m_index * mma_cta_count) for cluster_m_index in range(cluster_m // mma_cta_count)) if tensor_role == "sfb": return (1 << cluster_m) - 1 raise ValueError(f"Unsupported TMA tensor role {tensor_role!r}.") @@ -176,11 +156,7 @@ def tma_multicast_mask( ) else: cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 offset;\n\t" @@ -210,11 +186,7 @@ def tma_multicast_mask( ) else: cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 offset;\n\t" @@ -227,9 +199,7 @@ def tma_multicast_mask( return Int16(result) if cutlass.const_expr(is_preferred is None): - raise ValueError( - "is_preferred is required when the preferred and fallback multicast patterns differ." - ) + raise ValueError("is_preferred is required when the preferred and fallback multicast patterns differ.") if cutlass.const_expr(tensor_role in ("a", "sfa")): result = cute.arch.inline_ptx( @@ -243,15 +213,9 @@ def tma_multicast_mask( ) else: if cutlass.const_expr(fallback_pattern & preferred_pattern != fallback_pattern): - raise ValueError( - "Fallback B/SFB multicast pattern must be a subset of the preferred pattern." - ) + raise ValueError("Fallback B/SFB multicast pattern must be a subset of the preferred pattern.") cta_n = Int32(cta_coord_in_cluster[1]) - mma_cta_index = ( - Int32(0) - if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) - else cta_m % 2 - ) + mma_cta_index = Int32(0) if cutlass.const_expr(tensor_role == "sfb" or mma_cta_count == 1) else cta_m % 2 result = cute.arch.inline_ptx( "{\n\t" ".reg .u32 cluster_m, offset;\n\t" @@ -268,10 +232,4 @@ def tma_multicast_mask( return Int16(result) -__all__ = [ - "mark_alignment", - "smem_exclusive_prefix", - "spin_peek", - "spin_wait", - "tma_multicast_mask", -] +__all__ = ["mark_alignment", "smem_exclusive_prefix", "spin_peek", "spin_wait", "tma_multicast_mask"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py new file mode 100644 index 000000000000..0ec5abde4407 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/fc2_ready_geometry.py @@ -0,0 +1,36 @@ +"""Host geometry for FC2 completion counters in the swap-AB kernel.""" + + +def derive_fc2_ready_geometry( + *, + hidden_size: int, + mma_tiler_mnk: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, + fallback_cluster_shape_mn: tuple[int, int] | None = None, +) -> tuple[int, int]: + """Return (tokens per ready slot, CTA completions per ready slot). + + Work IDs and token slots use the preferred *logical* cluster. Fixed groups + of fallback clusters partition that cluster's CTA coordinates; they do not + create extra token slots or smaller completion targets. Every logical CTA + publishes once, including padded hidden tiles in the last cluster tile. + """ + mma_cta_count = 2 if use_2cta_instrs else 1 + if cluster_shape_mn[1] != 1: + raise ValueError("FC2 token-ready geometry requires hidden-only cluster splitting (N=1).") + if mma_tiler_mnk[0] % mma_cta_count or cluster_shape_mn[0] % mma_cta_count: + raise ValueError("FC2 token-ready geometry must preserve complete MMA CTA groups.") + if fallback_cluster_shape_mn is not None: + fallback_m, fallback_n = fallback_cluster_shape_mn + if ( + fallback_n != 1 + or fallback_m <= 0 + or cluster_shape_mn[0] % fallback_m + or fallback_m % mma_cta_count + ): + raise ValueError("FC2 fallback clusters must partition the preferred hidden axis into complete MMA groups.") + cta_hidden = mma_tiler_mnk[0] // mma_cta_count + hidden_per_logical_cluster = cta_hidden * cluster_shape_mn[0] + hidden_cluster_tiles = (hidden_size + hidden_per_logical_cluster - 1) // hidden_per_logical_cluster + return mma_tiler_mnk[1], hidden_cluster_tiles * cluster_shape_mn[0] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py index b38730b01cfc..28088ec8f0ca 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/flag_batch.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Lane- and warp-distributed GPU release-counter batching.""" import dataclasses @@ -24,22 +21,15 @@ class GpuReleaseFlagBatchTracker: thread_idx: Int32 @cute.jit - def _make( - self, flag_address: Int64, accumulated_flags: Int32, phase: Int32 - ) -> "GpuReleaseFlagBatchTracker": + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuReleaseFlagBatchTracker": return type(self)( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - thread_idx=self.thread_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=self.thread_idx ) @cute.jit def fire(self) -> None: if self.flag_address != Int64(0): - pointer = cute.make_ptr( - cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4 - ) + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) red_add_release_gpu_s32(pointer, Int32(1)) @cute.jit @@ -51,12 +41,8 @@ def accumulate( lane_address = Int64(0) if self.thread_idx == Int32(0): lane_address = flag_address - self._make( - flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase - ).fire() - return self._make( - flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase) - ) + self._make(flag_address=lane_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) current_address = self.flag_address accumulated_flags = self.accumulated_flags @@ -66,19 +52,11 @@ def accumulate( if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: if cutlass.const_expr(not no_fire): - self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=self.phase, - ).fire() + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() accumulated_flags = Int32(0) current_address = Int64(0) - return self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=Int32(next_phase), - ) + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) @dataclasses.dataclass(frozen=True) @@ -91,22 +69,15 @@ class GpuAsyncReleaseFlagBatchTracker: warp_idx: Int32 @cute.jit - def _make( - self, flag_address: Int64, accumulated_flags: Int32, phase: Int32 - ) -> "GpuAsyncReleaseFlagBatchTracker": + def _make(self, flag_address: Int64, accumulated_flags: Int32, phase: Int32) -> "GpuAsyncReleaseFlagBatchTracker": return type(self)( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - warp_idx=self.warp_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, warp_idx=self.warp_idx ) @cute.jit def fire(self) -> None: if self.flag_address != Int64(0): - pointer = cute.make_ptr( - cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4 - ) + pointer = cute.make_ptr(cutlass.Int32, self.flag_address, AddressSpace.gmem, assumed_align=4) with cute.arch.elect_one(): red_async_add_release_gpu_s32(pointer, Int32(1)) @@ -119,12 +90,8 @@ def accumulate( warp_address = Int64(0) if self.warp_idx == Int32(0): warp_address = flag_address - self._make( - flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase - ).fire() - return self._make( - flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase) - ) + self._make(flag_address=warp_address, accumulated_flags=Int32(1), phase=self.phase).fire() + return self._make(flag_address=Int64(0), accumulated_flags=Int32(0), phase=Int32(next_phase)) current_address = self.flag_address accumulated_flags = self.accumulated_flags @@ -134,29 +101,16 @@ def accumulate( if accumulated_flags == Int32(flush_threshold) or next_phase != self.phase: if cutlass.const_expr(not no_fire): - self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=self.phase, - ).fire() + self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=self.phase).fire() accumulated_flags = Int32(0) current_address = Int64(0) - return self._make( - flag_address=current_address, - accumulated_flags=accumulated_flags, - phase=Int32(next_phase), - ) + return self._make(flag_address=current_address, accumulated_flags=accumulated_flags, phase=Int32(next_phase)) @cute.jit def make_flag_batch_tracker( - use_async: bool, - *, - flag_address: Int64, - accumulated_flags: Int32, - phase: Int32, - thread_idx: Int32, + use_async: bool, *, flag_address: Int64, accumulated_flags: Int32, phase: Int32, thread_idx: Int32 ) -> Union[GpuReleaseFlagBatchTracker, GpuAsyncReleaseFlagBatchTracker]: """Construct a lane-batched synchronous or warp-batched asynchronous tracker. @@ -172,15 +126,8 @@ def make_flag_batch_tracker( warp_idx=cute.arch.make_warp_uniform(thread_idx // Int32(32)), ) return GpuReleaseFlagBatchTracker( - flag_address=flag_address, - accumulated_flags=accumulated_flags, - phase=phase, - thread_idx=thread_idx, + flag_address=flag_address, accumulated_flags=accumulated_flags, phase=phase, thread_idx=thread_idx ) -__all__ = [ - "GpuAsyncReleaseFlagBatchTracker", - "GpuReleaseFlagBatchTracker", - "make_flag_batch_tracker", -] +__all__ = ["GpuAsyncReleaseFlagBatchTracker", "GpuReleaseFlagBatchTracker", "make_flag_batch_tracker"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py new file mode 100644 index 000000000000..25122eb351e2 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/helper_weight_ready_gate.py @@ -0,0 +1,318 @@ +"""Shared Blackwell/Rubin kernel gate for dynamically installed helper weights. + +The gate is deliberately optional at code generation time. A specialization +with ``S == 0`` constructs no gate and therefore carries neither of the two +runtime values. A helper-bearing specialization derives ``H = M - S`` and +waits once per scheduler CTA, immediately before its first physical helper +work tile is published. + +``terminal_flags`` is the local ``uint64[EP]`` view published by the weight +transport. Bit 63 is reserved. A source is terminal-ready when that bit is +clear and its low-63-bit generation is greater than or equal to the expected +generation. The expected generation is a CPU-captured launch-time scalar, not +device READY storage. After READY publication, callers keep the weights +immutable until the MegaMoE launch completes. +""" + +import dataclasses +from enum import IntEnum +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace, Pointer +from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values + +from .ptx_helpers import nanosleep, read_globaltimer_ns + + +HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT = 1 << 63 +HELPER_WEIGHT_READY_GENERATION_MASK = HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT - 1 +HELPER_WEIGHT_READY_POLL_SLEEP_NS = 100 +# Elapsed time after the first unsuccessful READY sweep, independent of SM clocks. +# This is deliberately compile-time policy, not another production ABI knob. +HELPER_WEIGHT_READY_TIMEOUT_NS = 120_000_000_000 + + +class HelperWeightReadyStatus(IntEnum): + """Warp-uniform result of one terminal-flags sweep.""" + + WAITING = 0 + READY = 1 + TERMINAL_ERROR = 2 + INVALID_EXPECTED_GENERATION = 3 + TIMEOUT = 4 + + +@dataclasses.dataclass(frozen=True) +class HelperWeightReadyTopology: + """Compile-time ``[H home slots][S helper slots]`` topology.""" + + memory_slot_count: int + helper_count: int + home_expert_count: int + gate_required: bool + +def derive_helper_weight_ready_topology( + memory_slot_count: int, + helper_count: int, +) -> HelperWeightReadyTopology: + """Validate static ``M``/``S`` and derive the only helper boundary ``H``.""" + + if type(memory_slot_count) is not int or memory_slot_count <= 0: + raise ValueError("memory_slot_count (M) must be a positive exact int") + if type(helper_count) is not int or helper_count < 0: + raise ValueError("helper_count (S) must be a non-negative exact int") + if helper_count >= memory_slot_count: + raise ValueError("helper_count (S) must leave at least one home slot") + + home_expert_count = memory_slot_count - helper_count + return HelperWeightReadyTopology( + memory_slot_count=memory_slot_count, + helper_count=helper_count, + home_expert_count=home_expert_count, + gate_required=helper_count > 0, + ) + + +def _validate_helper_weight_ready_binding( + *, + topology: HelperWeightReadyTopology, + source_count: Optional[int], + terminal_flags: Optional[Pointer], + expected_generation: Optional[object], +) -> None: + """Enforce exact runtime-bundle presence from the compile-time topology.""" + + if not topology.gate_required: + if terminal_flags is not None or expected_generation is not None: + raise ValueError("S=0 must not carry READY runtime ABI values") + return + + if type(source_count) is not int or source_count <= 0: + raise ValueError("source_count (EP) must be a positive exact int when S>0") + if terminal_flags is None or expected_generation is None: + raise ValueError("S>0 requires terminal_flags and expected_generation") + if isinstance(expected_generation, int): + if type(expected_generation) is not int or not ( + 1 <= expected_generation <= HELPER_WEIGHT_READY_GENERATION_MASK + ): + raise ValueError("expected_generation must be in [1, 2**63)") + + +@cute.jit +def _helper_weight_ready_fail_closed_trap() -> None: + """Prevent publication after a terminal, invalid, or timed-out wait.""" + + llvm.inline_asm( + res=None, + operands_=[], + asm_string="trap;", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class HelperWeightReadyGenerationGate: + """One scheduler CTA's wait-all gate with an explicit caller-owned READY latch.""" + + def __init__( + self, + *, + topology: HelperWeightReadyTopology, + source_count: int, + terminal_flags: Pointer, + expected_generation, + ) -> None: + _validate_helper_weight_ready_binding( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + if not topology.gate_required: + raise ValueError("construct no HelperWeightReadyGenerationGate when S=0") + + self.topology = topology + self.source_count = source_count + self.terminal_flags = terminal_flags + self.expected_generation = expected_generation + + def __extract_mlir_values__(self) -> list: + values = [] + for field in ( + self.terminal_flags, + self.expected_generation, + ): + values.extend(extract_mlir_values(field)) + return values + + def __new_from_mlir_values__(self, values: list) -> "HelperWeightReadyGenerationGate": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values( + field, + values[value_index : value_index + field_value_count], + ) + value_index += field_value_count + return result + + result = type(self)( + topology=self.topology, + source_count=self.source_count, + terminal_flags=rebuild(self.terminal_flags), + expected_generation=rebuild(self.expected_generation), + ) + if value_index != len(values): + raise ValueError( + "HelperWeightReadyGenerationGate MLIR value count mismatch: " + f"consumed {value_index}, got {len(values)}." + ) + return result + + @cute.jit + def _typed_terminal_flags(self): + return cute.make_ptr( + cutlass.Uint64, + self.terminal_flags.toint(), + AddressSpace.gmem, + assumed_align=8, + ) + + @cute.jit + def _sweep_terminal_flags(self): + """Read every source cell; acquire all releases before admitting READY.""" + + # The launch generation is a uniform scalar. This terminal pointer is + # the only global-memory READY authority loaded by the spin loop. + flags = self._typed_terminal_flags() + all_sources_geq = Int32(1) + terminal_error_seen = Int32(0) + generation_mask = cutlass.Uint64(HELPER_WEIGHT_READY_GENERATION_MASK) + + for source_rank in cutlass.range_constexpr(self.source_count): + observed = cute.arch.load( + flags + Int32(source_rank), + cutlass.Uint64, + sem="relaxed", + scope="sys", + ) + observed_generation = observed & generation_mask + if observed > generation_mask: + terminal_error_seen = Int32(1) + if observed_generation < self.expected_generation: + all_sources_geq = Int32(0) + + status = Int32(HelperWeightReadyStatus.WAITING) + if terminal_error_seen != Int32(0): + status = Int32(HelperWeightReadyStatus.TERMINAL_ERROR) + elif all_sources_geq != Int32(0): + # Each system-scope strong read plus this acquire fence forms an + # acquire pattern for its source's release (PTX ISA, section 8.8). + llvm.inline_asm( + res=None, + operands_=[], + asm_string="fence.acquire.sys;", + constraints="~{memory}", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + status = Int32(HelperWeightReadyStatus.READY) + return status + + @cute.jit + def _wait_all(self): + """Lane 0 polls; fence the READY acquire before returning success.""" + + status = Int32(HelperWeightReadyStatus.WAITING) + if cute.arch.lane_idx() == Int32(0): + generation_mask = cutlass.Uint64(HELPER_WEIGHT_READY_GENERATION_MASK) + if (self.expected_generation == cutlass.Uint64(0)) | ( + self.expected_generation > generation_mask + ): + status = Int32(HelperWeightReadyStatus.INVALID_EXPECTED_GENERATION) + else: + status = self._sweep_terminal_flags() + if status == Int32(HelperWeightReadyStatus.WAITING): + # Already-ready weights take no timer reads. Unsigned elapsed + # time also handles wraparound of the 64-bit nanosecond timer. + wait_started_ns = read_globaltimer_ns() + while status == Int32(HelperWeightReadyStatus.WAITING): + elapsed_ns = read_globaltimer_ns() - wait_started_ns + if elapsed_ns >= cutlass.Uint64(HELPER_WEIGHT_READY_TIMEOUT_NS): + status = Int32(HelperWeightReadyStatus.TIMEOUT) + else: + nanosleep(HELPER_WEIGHT_READY_POLL_SLEEP_NS) + status = self._sweep_terminal_flags() + + status = Int32(cute.arch.shuffle_sync(status, offset=0)) + if status == Int32(HelperWeightReadyStatus.READY): + if cute.arch.lane_idx() == Int32(0): + cute.arch.fence_proxy("async.global") + return status + + @cute.jit + def wait_before_first_helper_publish(self, expert_idx, ready_latched: Int32) -> Int32: + """Return the next CTA latch after admitting physical ``m in [H,M)``. + + The scheduler must carry this scalar through its work loop. Mutating a + gate nested in a frozen extension does not produce a CuTe loop result. + """ + + needs_admission = ( + (expert_idx >= Int32(self.topology.home_expert_count)) + & (expert_idx < Int32(self.topology.memory_slot_count)) + & (ready_latched == Int32(0)) + ) + if needs_admission: + status = self._wait_all() + if status != Int32(HelperWeightReadyStatus.READY): + _helper_weight_ready_fail_closed_trap() + ready_latched = Int32(1) + return ready_latched + + +def make_helper_weight_ready_generation_gate( + *, + memory_slot_count: int, + helper_count: int, + source_count: Optional[int] = None, + terminal_flags: Optional[Pointer] = None, + expected_generation: Optional[object] = None, +) -> Optional[HelperWeightReadyGenerationGate]: + """Create the mandatory ``S>0`` gate or the ABI-free ``S==0`` sentinel.""" + + topology = derive_helper_weight_ready_topology(memory_slot_count, helper_count) + _validate_helper_weight_ready_binding( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + if not topology.gate_required: + return None + return HelperWeightReadyGenerationGate( + topology=topology, + source_count=source_count, + terminal_flags=terminal_flags, + expected_generation=expected_generation, + ) + + +__all__ = [ + "HELPER_WEIGHT_READY_GENERATION_MASK", + "HELPER_WEIGHT_READY_TIMEOUT_NS", + "HELPER_WEIGHT_READY_TERMINAL_ERROR_BIT", + "HelperWeightReadyGenerationGate", + "HelperWeightReadyStatus", + "HelperWeightReadyTopology", + "derive_helper_weight_ready_topology", + "make_helper_weight_ready_generation_gate", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py index cfcce57b687a..d4cfd4c972b9 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/iket_compat.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Compatibility wrapper for optional in-kernel event tracing support.""" try: @@ -9,7 +6,6 @@ try: from cutlass.cute import iket # type: ignore except (ImportError, NotImplementedError): - class _IketShim: """No-op IKET interface for toolchains without the dialect.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py new file mode 100644 index 000000000000..51534b776661 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/megamoe_aot.py @@ -0,0 +1,18 @@ +"""Restore argument names and READY defaults for loaded MegaMoE functions.""" + +from inspect import signature + + +def make_megamoe_aot_callable(function, kernel): + """Use TVM-FFI's native adapter without attaching scheduler metadata.""" + from tvm_ffi.utils.kwargs_wrapper import make_kwargs_wrapper + + argument_names = list(signature(kernel.__call__).parameters) + # Home-only entries can omit the unused READY pair, as with direct JIT + # calls. Helper-enabled entries require both READY operands explicitly. + # Kernels that take the trailing optional fc1_activation_per_token_scale + # may omit it too (defaults are right-aligned, hence its position last). + trailing_optional = 1 if argument_names[-1] == "fc1_activation_per_token_scale" else 0 + ready_optional = 2 if kernel.helper_expert_count == 0 else 0 + defaults = (None,) * (ready_optional + trailing_optional) + return make_kwargs_wrapper(function, argument_names, arg_defaults=defaults) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py index 055d71d7c48d..84b7cb973a34 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/ptx_helpers.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Minimal inline-PTX primitives required by the greenfield NVFP4 kernels.""" from typing import Optional @@ -11,6 +8,7 @@ from cutlass._mlir.dialects import arith, llvm, vector from cutlass.cutlass_dsl import Float32, Int32, Int64, T, dsl_user_op + TmaCacheHintEvictFirst = 0x12F0000000000000 @@ -21,10 +19,8 @@ def _address_value(pointer_or_address, *, loc=None, ip=None): @dsl_user_op -def nanosleep( - sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None -) -> None: - """Suspend the calling thread for up to the requested clock cycles.""" +def nanosleep(sleep_cycles: int, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> None: + """Suspend the calling thread for approximately the requested nanoseconds.""" if cutlass.const_expr(hasattr(cute.arch, "nanosleep")): cute.arch.nanosleep(sleep_time=sleep_cycles, loc=loc, ip=ip) return @@ -58,20 +54,21 @@ def exit(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = @dsl_user_op -def read_clock64( - *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None -) -> Int64: +def read_clock64(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: """Read the per-SM 64-bit cycle counter.""" return Int64( llvm.inline_asm( - T.i64(), - [], - "mov.u64 $0, %clock64;", - "=l", - has_side_effects=True, - asm_dialect=0, - loc=loc, - ip=ip, + T.i64(), [], "mov.u64 $0, %clock64;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip + ) + ) + + +@dsl_user_op +def read_globaltimer_ns(*, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> cutlass.Uint64: + """Read the 64-bit global nanosecond timer, independent of SM clock rate.""" + return cutlass.Uint64( + llvm.inline_asm( + T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, asm_dialect=0, loc=loc, ip=ip ) ) @@ -80,9 +77,7 @@ def read_clock64( def movmatrix_b16(input_regs: cute.Tensor) -> cute.Tensor: """Transpose every packed m8n8 b16 register fragment across the warp.""" if cutlass.const_expr(input_regs.element_type.width != 32): - raise TypeError( - f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}." - ) + raise TypeError(f"movmatrix_b16 expects packed 32-bit registers, got {input_regs.element_type}.") input_words = cute.coalesce(cute.flatten(cute.recast_tensor(input_regs, Int32))) output_words = cute.make_rmem_tensor((cute.size(input_words),), Int32) @@ -145,9 +140,7 @@ def cvt_f32_to_fp8_to_f32( @dsl_user_op -def tma_load_1d( - destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None -) -> None: +def tma_load_1d(destination_smem, source_gmem, mbarrier_smem, num_bytes, *, loc=None, ip=None) -> None: """Issue a cache-hinted 1D GMEM-to-SMEM bulk copy.""" llvm.inline_asm( None, @@ -330,9 +323,7 @@ def lds128_v4_b32(smem_pointer, *, loc=None, ip=None): @dsl_user_op -def stg_f32( - address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None -) -> None: +def stg_f32(address: Int64, value: Float32, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: if predicate is None: llvm.inline_asm( None, @@ -358,9 +349,7 @@ def stg_f32( @dsl_user_op -def stg_b64( - address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None -) -> None: +def stg_b64(address: Int64, value: Int64, predicate: Optional[Int32] = None, *, loc=None, ip=None) -> None: if predicate is None: llvm.inline_asm( None, @@ -401,11 +390,7 @@ def red_add_relaxed_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) @dsl_user_op def red_add_relaxed_sys_f32( - address, - value: Float32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + address, value: Float32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -436,11 +421,7 @@ def red_add_release_sys_s32(address: Int64, value: Int32, *, loc=None, ip=None) @dsl_user_op def red_add_release_gpu_s32( - counter_pointer, - value: Int32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -457,11 +438,7 @@ def red_add_release_gpu_s32( @dsl_user_op def red_async_add_release_gpu_s32( - counter_pointer, - value: Int32, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + counter_pointer, value: Int32, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -503,12 +480,7 @@ def red_async_add_release_sys_u32(address: Int64, value: Int32, *, loc=None, ip= @dsl_user_op def red_add_relaxed_sys_v2_bf16x2( - address, - value0, - value1, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, + address, value0, value1, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None ) -> None: llvm.inline_asm( None, @@ -523,6 +495,64 @@ def red_add_relaxed_sys_v2_bf16x2( ) +@dsl_user_op +def f32x2_pack(lo, hi, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int64: + """Pack two f32 into one b64 f32x2 operand (low word ``lo``, high word ``hi``).""" + packed = llvm.inline_asm( + T.i64(), + [Float32(lo).ir_value(loc=loc, ip=ip), Float32(hi).ir_value(loc=loc, ip=ip)], + "mov.b64 $0, {$1, $2};", + "=l,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int64(packed) + + +@dsl_user_op +def fma_rn_f32x2_bf16x2( + word, scale2, acc2, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None +) -> Int64: + """``acc2 + word * scale2`` on a packed bf16x2 / f32x2 pair (``fma.rn.f32x2.bf16x2.f32x2.f32x2``, one rounding + per lane; scale 1.0 gives an exact ``acc + bf16``).""" + result = llvm.inline_asm( + T.i64(), + [ + Int32(word).ir_value(loc=loc, ip=ip), + Int64(scale2).ir_value(loc=loc, ip=ip), + Int64(acc2).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f32x2.bf16x2.f32x2.f32x2 $0, $1, $2, $3;", + "=l,r,l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int64(result) + + +@dsl_user_op +def cvt_rn_bf16x2_f32x2(acc2, *, loc: Optional[ir.Location] = None, ip: Optional[ir.InsertionPoint] = None) -> Int32: + """Round a packed f32x2 (low word -> low bf16 half) to one bf16x2 word with RN rounding (one F2FP).""" + packed = llvm.inline_asm( + T.i32(), + [Int64(acc2).ir_value(loc=loc, ip=ip)], + "{\n .reg .f32 lo, hi;\n mov.b64 {lo, hi}, $1;\n cvt.rn.bf16x2.f32 $0, hi, lo;\n}", + "=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed) + + @cute.jit def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=None) -> Int32: """Round four f32 lanes to the selected FP8 format and pack them into one i32.""" @@ -561,6 +591,136 @@ def cvt_f32x4_to_f8x4_pack_i32(fp32x4: cute.Tensor, fp8_type, *, loc=None, ip=No return Int32(packed_i32) +@dsl_user_op +def cvt_f32_to_ue8m0_raw_i32( + value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> Int32: + """Round one f32 up to a raw UE8M0 byte in the low bits of an i32.""" + packed_scale = llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .b16 sf2;\n" + " cvt.rp.satfinite.ue8m0x2.f32 sf2, 0f00000000, $1;\n" + " cvt.u32.u16 $0, sf2;\n" + "}", + "=r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_scale) + + +@cute.jit +def cvt_f32x4_to_f8x4_scaled_pack_i32( + fp32x4: cute.Tensor, + fp8_type, + scale_raw_i32, + *, + loc=None, + ip=None, +) -> Int32: + """Pack four f32 values to FP8 while dividing by one UE8M0 scale.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + srcs = [ + Float32(vector.extract(src_vec4, [], [index])).ir_value(loc=loc, ip=ip) + for index in range(4) + ] + scale_i32 = Int32(scale_raw_i32).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e5m2x2.f32" + else: + raise ValueError(f"Scaled FP8 conversion does not support {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [*srcs, scale_i32], + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + " .reg .b8 s;\n" + " cvt.u8.u32 s, $5;\n" + f" {cvt_instruction} lo, $2, $1, s;\n" + f" {cvt_instruction} hi, $4, $3, s;\n" + " mov.b32 $0, {lo, hi};\n" + "}", + "=r,f,f,f,f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_i32) + + +@cute.jit +def cvt_f32x4_to_f8x4_scaled4_pack_i32( + fp32x4: cute.Tensor, + fp8_type, + scales_packed_i32, + *, + loc=None, + ip=None, +) -> Int32: + """Pack four f32 values to FP8 using four independent UE8M0 scales.""" + fp32x4 = fp32x4.load() + src_vec4 = fp32x4.ir_value(loc=loc, ip=ip) if hasattr(fp32x4, "ir_value") else fp32x4 + srcs = [ + Float32(vector.extract(src_vec4, [], [index])).ir_value(loc=loc, ip=ip) + for index in range(4) + ] + scales_i32 = Int32(scales_packed_i32).ir_value(loc=loc, ip=ip) + + if cutlass.const_expr(fp8_type is cutlass.Float8E4M3FN): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e4m3x2.f32" + elif cutlass.const_expr(fp8_type is cutlass.Float8E5M2): + cvt_instruction = "cvt.rn.satfinite.scaled::n1::ue8m0.e5m2x2.f32" + else: + raise ValueError(f"Scaled FP8 conversion does not support {fp8_type}.") + + packed_i32 = llvm.inline_asm( + T.i32(), + [*srcs, scales_i32], + "{\n" + " .reg .b8 s0, s1, s2, s3;\n" + " .reg .b16 t0, t1, t2, t3;\n" + " .reg .b32 e1, e2, e3, p01, p23;\n" + " shr.u32 e1, $5, 8;\n" + " shr.u32 e2, $5, 16;\n" + " shr.u32 e3, $5, 24;\n" + " cvt.u8.u32 s0, $5;\n" + " cvt.u8.u32 s1, e1;\n" + " cvt.u8.u32 s2, e2;\n" + " cvt.u8.u32 s3, e3;\n" + f" {cvt_instruction} t0, $1, $1, s0;\n" + f" {cvt_instruction} t1, $2, $2, s1;\n" + f" {cvt_instruction} t2, $3, $3, s2;\n" + f" {cvt_instruction} t3, $4, $4, s3;\n" + " mov.b32 p01, {t0, t1};\n" + " mov.b32 p23, {t2, t3};\n" + " prmt.b32 $0, p01, p23, 0x6420;\n" + "}", + "=r,f,f,f,f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(packed_i32) + + @dsl_user_op def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> None: """Convert ``fp32_val`` to E8M0 via PTX and store the 1-byte result to global memory. @@ -590,17 +750,9 @@ def stg_e8m0_from_f32(addr: Int64, fp32_val: Float32, *, loc=None, ip=None) -> N @dsl_user_op def stg_e8m0x8_from_f32( addr: Int64, - v0: Float32, - v1: Float32, - v2: Float32, - v3: Float32, - v4: Float32, - v5: Float32, - v6: Float32, - v7: Float32, - *, - loc=None, - ip=None, + v0: Float32, v1: Float32, v2: Float32, v3: Float32, + v4: Float32, v5: Float32, v6: Float32, v7: Float32, + *, loc=None, ip=None, ) -> None: """Convert 8 fp32 values to E8M0 and store them as 8 contiguous bytes in one shot. @@ -613,14 +765,8 @@ def stg_e8m0x8_from_f32( None, [ addr.ir_value(), - v0.ir_value(), - v1.ir_value(), - v2.ir_value(), - v3.ir_value(), - v4.ir_value(), - v5.ir_value(), - v6.ir_value(), - v7.ir_value(), + v0.ir_value(), v1.ir_value(), v2.ir_value(), v3.ir_value(), + v4.ir_value(), v5.ir_value(), v6.ir_value(), v7.ir_value(), ], "{\n" " .reg .b16 p0, p1, p2, p3;\n" @@ -647,13 +793,17 @@ def stg_e8m0x8_from_f32( "cp_reduce_async_bulk_add_bf16_s2g", "cp_reduce_async_bulk_add_u32_s2g", "cvt_f32_to_fp8_to_f32", + "cvt_f32_to_ue8m0_raw_i32", "cvt_f32x4_to_f8x4_pack_i32", + "cvt_f32x4_to_f8x4_scaled4_pack_i32", + "cvt_f32x4_to_f8x4_scaled_pack_i32", "exit", "lds128_v4_b32", "mbarrier_arrive_expect_tx_on_peer", "movmatrix_b16", "nanosleep", "read_clock64", + "read_globaltimer_ns", "red_add_relaxed_sys_f32", "red_add_relaxed_sys_s32", "red_add_relaxed_sys_v2_bf16x2", @@ -667,4 +817,8 @@ def stg_e8m0x8_from_f32( "stg_e8m0x8_from_f32", "stg_f32", "tma_load_1d", + "fence_acq_rel_gpu", + "f32x2_pack", + "fma_rn_f32x2_bf16x2", + "cvt_rn_bf16x2_f32x2", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py index eb6c51d75cb3..7e5d4af1b466 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/smem_workspace.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Pure-Python SMEM declarations with lifetime-aware overlay placement.""" import dataclasses @@ -12,11 +9,12 @@ from .utils import ( cosize_from_shape_stride_tuples, is_power_of_two, - round_up, row_major_stride, + round_up, validate_static_integer_tuple, ) + SmemRegionKind = Literal["mbarrier", "tensor"] SwizzleSpec = Tuple[int, int, int] @@ -110,7 +108,9 @@ def add_lifetime(self, name: str) -> SmemLifetime: if not name: raise ValueError("An SMEM lifetime needs a non-empty name.") if name in self._lifetime_names: - raise ValueError(f"Duplicate lifetime {name!r} in overlay {self.name!r}.") + raise ValueError( + f"Duplicate lifetime {name!r} in overlay {self.name!r}." + ) lifetime = SmemLifetime(self._workspace, self, name) self._lifetimes.append(lifetime) self._lifetime_names.add(name) @@ -121,7 +121,13 @@ def add_lifetime(self, name: str) -> SmemLifetime: class SmemWorkspace: - """Collect static SMEM declarations and finalize one physical placement.""" + """Collect static SMEM declarations and finalize one physical placement. + + A kernel whose CTAs take different roles builds one workspace per role and lays every resulting offset + table over the same arena. The tables overlap by design: a CTA only reads the table of the role it took, + so a region name a role never registers raises at trace time rather than silently reading another role's + bytes. + """ def __init__( self, @@ -235,17 +241,29 @@ def _make_region( stride = row_major_stride(shape) validate_static_integer_tuple(stride, field_name=f"{name}.stride") if len(shape) != len(stride): - raise ValueError(f"SMEM region {name!r} shape and stride ranks differ.") + raise ValueError( + f"SMEM region {name!r} shape and stride ranks differ." + ) if swizzle is not None: - if len(swizzle) != 3 or not all(isinstance(parameter, int) for parameter in swizzle): - raise TypeError(f"SMEM region {name!r} swizzle must be three Python ints.") + if len(swizzle) != 3 or not all( + isinstance(parameter, int) for parameter in swizzle + ): + raise TypeError( + f"SMEM region {name!r} swizzle must be three Python ints." + ) if swizzle[0] < 0 or swizzle[1] < 0: - raise ValueError(f"SMEM region {name!r} swizzle bits/base must be non-negative.") + raise ValueError( + f"SMEM region {name!r} swizzle bits/base must be non-negative." + ) natural_alignment = max(1, (int(dtype.width) + 7) // 8) - explicit_alignment = natural_alignment if byte_alignment is None else byte_alignment + explicit_alignment = ( + natural_alignment if byte_alignment is None else byte_alignment + ) if not is_power_of_two(explicit_alignment): - raise ValueError(f"SMEM region {name!r} alignment must be a positive power of two.") + raise ValueError( + f"SMEM region {name!r} alignment must be a positive power of two." + ) swizzle_alignment = _swizzle_alignment(swizzle) effective_alignment = max( explicit_alignment, @@ -268,7 +286,10 @@ def _make_region( ) def _claim_region_name(self, region: SmemRegion) -> None: - if region.name in self._region_by_name or region.name in self._overlay_names: + if ( + region.name in self._region_by_name + or region.name in self._overlay_names + ): raise ValueError(f"Duplicate SMEM region {region.name!r}.") self._region_by_name[region.name] = region @@ -287,7 +308,9 @@ def finalize(self, *, max_bytes: Optional[int] = None) -> None: raise RuntimeError("SmemWorkspace.finalize() may only be called once.") offsets, total_bytes = self._build_placement() if max_bytes is not None and total_bytes > max_bytes: - raise ValueError(f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes.") + raise ValueError( + f"SMEM plan needs {total_bytes} bytes, exceeding {max_bytes} bytes." + ) self._offset = offsets self._total_bytes = total_bytes self._finalized = True @@ -305,7 +328,9 @@ def _build_placement(self) -> Tuple[Dict[str, int], int]: offsets[declaration.name] = cursor cursor += declaration.nbytes continue - relative_offsets, overlay_alignment, overlay_bytes = self._layout_overlay(declaration) + relative_offsets, overlay_alignment, overlay_bytes = ( + self._layout_overlay(declaration) + ) cursor = round_up(cursor, overlay_alignment) for region_name, relative_offset in relative_offsets.items(): offsets[region_name] = cursor + relative_offset @@ -323,7 +348,9 @@ def _layout_overlay( overlay_bytes = 0 for lifetime in overlay.lifetimes: if not lifetime.regions: - raise ValueError(f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions.") + raise ValueError( + f"SMEM lifetime {overlay.name}.{lifetime.name} has no regions." + ) lifetime_cursor = 0 for region in lifetime.regions: overlay_alignment = max( @@ -369,10 +396,24 @@ class SmemStorage: return SmemStorage + def assert_fits_arena(self, arena_bytes: int) -> None: + """Check this role's placement against the arena every CTA of the launch allocates.""" + self._require_finalized() + if self._total_bytes > arena_bytes: + raise ValueError( + f"SMEM workspace needs {self._total_bytes} bytes but the shared arena is only " + f"{arena_bytes} bytes. The arena is what every CTA allocates, so it cannot be widened for " + "one role without taking the difference away from every other role." + ) + @cute.jit def ptr(self, name: str, smem_base: cute.Pointer) -> cute.Pointer: region = self._region_by_name[name] - swizzle = None if region.swizzle is None else cute.make_swizzle(*region.swizzle) + swizzle = ( + None + if region.swizzle is None + else cute.make_swizzle(*region.swizzle) + ) return cute.make_ptr( region.dtype, smem_base.toint() + self._offset[name], @@ -400,3 +441,4 @@ def _require_finalized(self) -> None: "SmemWorkspace", "SwizzleSpec", ] + diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py index 135612ecea50..9b1ef4010cb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/software_sync.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Software grid and NVLink synchronization for persistent kernels.""" import cutlass @@ -172,7 +169,10 @@ def arrive_and_wait( signal_base_address = self._signal.toint() signal_byte_offset = Int64(signal_phase * Int32(4)) for rank_round in cutlass.range_constexpr(rank_round_count): - destination_rank = Int32(rank_round * participating_threads) + thread_idx_in_group + destination_rank = ( + Int32(rank_round * participating_threads) + + thread_idx_in_group + ) if destination_rank < Int32(self.world_size): destination_address = self._peer_rank_ptr_mapper.map( signal_base_address, @@ -193,15 +193,12 @@ def arrive_and_wait( scope="gpu", ) local_signal = self._signal + signal_phase - while ( - cute.arch.load( - local_signal, - Int32, - sem="acquire", - scope="sys", - ) - != signal_target - ): + while cute.arch.load( + local_signal, + Int32, + sem="acquire", + scope="sys", + ) != signal_target: pass if cutlass.const_expr(epilogue_grid_sync): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py similarity index 82% rename from tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py rename to tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py index 5ee767bde5b7..ce4cb6465a0f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/topk_reduce.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/topk_reduce.py @@ -23,14 +23,16 @@ from typing import ClassVar, Dict, Optional import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute -from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import Float32, Int32, T +from cutlass._mlir.dialects import llvm + +from ..quant_def import CombineFormat +from .constants import Nvfp4E2M1RcpLimit +from .dsl_helpers import mark_alignment -from .....helpers.constants import Nvfp4E2M1RcpLimit -from .....helpers.dsl_helpers import mark_alignment -from .....quant_def import CombineFormat # --------------------------------------------------------------------------- # fp4 (e2m1) -> fp32 register decode. @@ -198,9 +200,7 @@ def __call__( combine_quant = cute.make_tensor( combine_quant.iterator, - cute.make_layout( - (combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride - ), + cute.make_layout((combine_quant.shape[0], self.num_topk, self.hidden), stride=combine_quant.stride), ) reduced_output = cute.make_tensor( reduced_output.iterator, @@ -208,14 +208,11 @@ def __call__( ) if cutlass.const_expr(topk_score is not None): topk_score = cute.make_tensor( - topk_score.iterator, - cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride), + topk_score.iterator, cute.make_layout((topk_score.shape[0], self.num_topk), stride=topk_score.stride) ) if cutlass.const_expr(not self.combine_format.is_quantized): - self._reduce_bf16(combine_quant, topk_score, reduced_output).launch( - grid=grid, block=block, stream=stream - ) + self._reduce_bf16(combine_quant, topk_score, reduced_output).launch(grid=grid, block=block, stream=stream) return # The mega kernel hands sf in already as the depth-2 broadcast layout; a @@ -226,8 +223,7 @@ def __call__( sf = cute.make_tensor( combine_sf.iterator, cute.make_layout( - (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), - stride=combine_sf.stride, + (combine_sf.shape[0], self.num_topk, (sf_vec, self.hidden // sf_vec)), stride=combine_sf.stride ), ) else: @@ -251,12 +247,7 @@ def __call__( # -- kernels -------------------------------------------------------------- @cute.kernel - def _reduce_bf16( - self, - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - reduced_output: cute.Tensor, - ): + def _reduce_bf16(self, combine_output: cute.Tensor, topk_score: Optional[cute.Tensor], reduced_output: cute.Tensor): threads = self._threads hidden_per_thread = self.hidden_per_thread hidden_tiles = self.hidden_tiles @@ -269,20 +260,16 @@ def _reduce_bf16( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - terms = cute.zipped_divide( - combine_output[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + terms = cute.zipped_divide(combine_output[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -291,17 +278,13 @@ def _reduce_bf16( for k in cutlass.range_constexpr(num_topk): score_reg[k] = score_dtype(1) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=128) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) cute.copy( - load_atom, - mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), - term, + load_atom, mark_alignment(terms[k, None], hidden_per_thread * cutlass.BFloat16.width // 8), term ) if cutlass.const_expr(topk_score is not None and not prefetch): score_reg[k] = topk_score[token_idx, Int32(k)] @@ -310,15 +293,12 @@ def _reduce_bf16( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): value_pair = (Float32(term[i]), Float32(term[i + 1])) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - value_pair, score_pair, (acc[i], acc[i + 1]) + value_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - else: - if cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(value_pair, score_pair) - else: - acc[i] = value_pair[0] - acc[i + 1] = value_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) @@ -348,25 +328,21 @@ def _reduce_mxfp8( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.Float8E8M0FNU) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - codes = cute.zipped_divide( - combine_quant[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] # (token, topk, hidden) -> (topk, hidden_per_thread) - sf = cute.zipped_divide( - combine_sf[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -375,23 +351,15 @@ def _reduce_mxfp8( for k in cutlass.range_constexpr(num_topk): score_reg[k] = score_dtype(1) if cutlass.const_expr(prefetch): - cute.autovec_copy( - sf[None, 0], scale_reg - ) # one scale per topk slot (stride-0 broadcast) + cute.autovec_copy(sf[None, 0], scale_reg) # one scale per topk slot (stride-0 broadcast) fp8_dtype = self.combine_format.act_dtype - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), fp8_dtype, num_bits_per_copy=128) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), fp8_dtype) - cute.copy( - load_atom, - mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), - term, - ) + cute.copy(load_atom, mark_alignment(codes[k, None], hidden_per_thread * fp8_dtype.width // 8), term) value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) value.store(term.load().to(cutlass.Float32)) @@ -407,17 +375,12 @@ def _reduce_mxfp8( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - dequant_pair, score_pair, (acc[i], acc[i + 1]) + dequant_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - else: - if cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2( - dequant_pair, score_pair - ) - else: - acc[i] = dequant_pair[0] - acc[i + 1] = dequant_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) @@ -447,25 +410,21 @@ def _reduce_fp4( token_idx = worker_idx // hidden_tiles hidden_tile_idx = worker_idx % hidden_tiles - score_dtype = ( - topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 - ) + score_dtype = topk_score.dtype if cutlass.const_expr(topk_score is not None) else cutlass.Float32 score_reg = cute.make_rmem_tensor((num_topk,), score_dtype) scale_reg = cute.make_rmem_tensor((num_topk,), cutlass.BFloat16) if (not needs_guard) or token_idx < reduced_output.shape[0]: # (token, topk, hidden) -> (topk, hidden_per_thread) - codes = cute.zipped_divide( - combine_quant[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] + codes = cute.zipped_divide(combine_quant[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] # (token, topk, hidden) -> (topk, hidden_per_thread) - sf = cute.zipped_divide( - combine_sf[token_idx, None, None], (num_topk, hidden_per_thread) - )[(None, None), (0, hidden_tile_idx)] - # (token, hidden) -> (hidden_per_thread) - dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[ - (None,), (hidden_tile_idx,) + sf = cute.zipped_divide(combine_sf[token_idx, None, None], (num_topk, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) ] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] if cutlass.const_expr(topk_score is not None): if cutlass.const_expr(prefetch): @@ -476,19 +435,13 @@ def _reduce_fp4( if cutlass.const_expr(prefetch): cute.autovec_copy(sf[None, 0], scale_reg) - load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64 - ) + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float4E2M1FN, num_bits_per_copy=64) acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) for k in cutlass.range_constexpr(0, num_topk, 1): term = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float4E2M1FN) cute.copy( - load_atom, - mark_alignment( - codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8 - ), - term, + load_atom, mark_alignment(codes[k, None], hidden_per_thread * cutlass.Float4E2M1FN.width // 8), term ) value = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) # Dev-only knob (MEGA_F4CVT_USE_MANUAL): manual LUT/PRMT decode vs @@ -513,14 +466,12 @@ def _reduce_fp4( for i in cutlass.range_constexpr(0, hidden_per_thread, 2): dequant_pair = cute.arch.mul_packed_f32x2((value[i], value[i + 1]), scale_pair) if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + # Include the +0 seed: direct assignment/multiply preserves -0. acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( - dequant_pair, score_pair, (acc[i], acc[i + 1]) + dequant_pair, score_pair, (Float32(0.0), Float32(0.0)) ) - elif cutlass.const_expr(topk_score is not None): - acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2(dequant_pair, score_pair) - else: - acc[i] = dequant_pair[0] - acc[i + 1] = dequant_pair[1] out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) out.store(acc.load().to(out_dtype)) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py index 60f768c3f43d..0c3738be2418 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/helpers/utils.py @@ -1,12 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Small integer and layout helpers shared by workspace implementations.""" from typing import Iterable, List, Tuple, Union import cutlass + IntegerType = Union[int, cutlass.Int32, cutlass.Int64, cutlass.Uint32, cutlass.Uint64] @@ -69,9 +67,7 @@ def strides_equal_ignoring_singletons(shape, lhs_stride, rhs_stride) -> bool: return shape == 1 or lhs_stride == rhs_stride -def ordered_stride( - shape: Tuple[int, ...], mem_order: Tuple[int, ...] -) -> Tuple[Tuple[int, ...], int]: +def ordered_stride(shape: Tuple[int, ...], mem_order: Tuple[int, ...]) -> Tuple[Tuple[int, ...], int]: stride = [0] * len(shape) cosize = 1 for mode in sorted(range(len(shape)), key=lambda index: mem_order[index]): diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py index 52a7a9daf028..337547020a79 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/__init__.py @@ -1,2 +1,26 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Kernel source modules independent of repository-only runners.""" + +from .schedulers import ( + BlackwellFusedFc12Scheduler, + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SchedulerBase, + SchedulerConsumer, + SchedulerWorkTileBase, + SwapAbFc12WorkTileInfo, + WorkIdAcquisitionMode, +) + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py index 52a7a9daf028..a84af22de952 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell kernel implementations.""" + +from .inference import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Kernel, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py index 162ec81ab183..32f380cdada0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/custom_mix_cga_helpers.py @@ -12,16 +12,17 @@ from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.nvgpu.cpasync.copy import ( - TMA_DESC_PTR_FIELD_NAME, - TMA_MBAR_PTR_FIELD_NAME, - TMA_MCAST_MASK_FIELD_NAME, CopyBulkTensorTileG2SMulticastNonExecTrait, CopyBulkTensorTileG2SMulticastTrait, CopyBulkTensorTileG2SNonExecTrait, CopyBulkTensorTileG2STrait, + TMA_DESC_PTR_FIELD_NAME, + TMA_MBAR_PTR_FIELD_NAME, + TMA_MCAST_MASK_FIELD_NAME, ) from cutlass.cutlass_dsl import Boolean, Int16, Int32 + TmaAtomOrPair = Union[cute.CopyAtom, Tuple[cute.CopyAtom, cute.CopyAtom]] @@ -42,17 +43,12 @@ def make_executable_tma_atom(atom: cute.CopyAtom) -> cute.CopyAtom: def make_executable_tma_atom_or_pair(atom_or_pair: TmaAtomOrPair) -> TmaAtomOrPair: """Materialize one executable atom or a preferred/fallback executable pair.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): - return ( - make_executable_tma_atom(atom_or_pair[0]), - make_executable_tma_atom(atom_or_pair[1]), - ) + return (make_executable_tma_atom(atom_or_pair[0]), make_executable_tma_atom(atom_or_pair[1])) return make_executable_tma_atom(atom_or_pair) @cute.jit -def select_tma_atom_for_cluster( - atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean -) -> cute.CopyAtom: +def select_tma_atom_for_cluster(atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean) -> cute.CopyAtom: """Select one preferred/fallback atom after both candidates are materialized.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): atom = atom_or_pair[0] @@ -64,34 +60,18 @@ def select_tma_atom_for_cluster( @cute.jit def bind_executable_tma_load_fields( - atom: cute.CopyAtom, - *, - tma_bar_ptr: cute.Pointer, - tma_desc_ptr: Optional[cute.Pointer], - mcast_mask: Optional[Int32], + atom: cute.CopyAtom, *, tma_bar_ptr: cute.Pointer, tma_desc_ptr: Optional[cute.Pointer], mcast_mask: Optional[Int32] ) -> cute.CopyAtom: """Bind per-issue fields to an already executable TMA load atom.""" executable_value = atom._trait.value if cutlass.const_expr(mcast_mask is not None): - mcast_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MCAST_MASK_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, mcast_attr, Int16(mcast_mask).ir_value() - ) + mcast_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MCAST_MASK_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, mcast_attr, Int16(mcast_mask).ir_value()) if cutlass.const_expr(tma_desc_ptr is not None): - descriptor_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, descriptor_attr, tma_desc_ptr.value - ) - barrier_attr = ir.Attribute.parse( - f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MBAR_PTR_FIELD_NAME}>" - ) - executable_value = _cute_nvgpu_ir.atom_set_value( - executable_value, barrier_attr, tma_bar_ptr.value - ) + descriptor_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, descriptor_attr, tma_desc_ptr.value) + barrier_attr = ir.Attribute.parse(f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_MBAR_PTR_FIELD_NAME}>") + executable_value = _cute_nvgpu_ir.atom_set_value(executable_value, barrier_attr, tma_bar_ptr.value) return cute.CopyAtom(atom.op, atom._trait.__class__(executable_value)) @@ -102,9 +82,7 @@ def _static_vmnk_shape(cta_layout_vmnk: cute.Layout) -> Tuple[int, int, int, int return shape -def _multicast_group_count( - shape_vmnk: Tuple[int, int, int, int], mcast_mode_mn: Tuple[int, int] -) -> int: +def _multicast_group_count(shape_vmnk: Tuple[int, int, int, int], mcast_mode_mn: Tuple[int, int]) -> int: mcast_a_count = shape_vmnk[2] if mcast_mode_mn[0] else 0 mcast_b_count = shape_vmnk[1] if mcast_mode_mn[1] else 0 overlap = 1 if mcast_mode_mn[0] and mcast_mode_mn[1] else 0 @@ -116,10 +94,7 @@ def _multicast_group_count( @cute.jit def _reader_mask( - cta_layout_vmnk: cute.Layout, - cta_coord_vmnk: cute.Coord, - mcast_mode_mn: Tuple[int, int], - mma_cta_count: int, + cta_layout_vmnk: cute.Layout, cta_coord_vmnk: cute.Coord, mcast_mode_mn: Tuple[int, int], mma_cta_count: int ) -> Int32: """Build the general reader union for non-single-sided cluster layouts. @@ -127,23 +102,15 @@ def _reader_mask( """ mask = Int32(0) if cutlass.const_expr(mcast_mode_mn[0]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=2) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=2)) if cutlass.const_expr(mcast_mode_mn[1]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=1) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, cta_coord_vmnk, mcast_mode=1)) if cutlass.const_expr(mma_cta_count == 2): peer_coord_vmnk = (cta_coord_vmnk[0] ^ Int32(1), *cta_coord_vmnk[1:]) if cutlass.const_expr(mcast_mode_mn[0]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=2) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=2)) if cutlass.const_expr(mcast_mode_mn[1]): - mask = mask | Int32( - cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=1) - ) + mask = mask | Int32(cpasync.create_tma_multicast_mask(cta_layout_vmnk, peer_coord_vmnk, mcast_mode=1)) return mask @@ -156,9 +123,7 @@ def _active_cluster_uses_cluster_sync( preferred_uses_cluster = preferred_cluster_shape_mn[0] * preferred_cluster_shape_mn[1] > 1 fallback_uses_cluster = fallback_cluster_shape_mn[0] * fallback_cluster_shape_mn[1] > 1 cluster_sync_behavior_differs = preferred_uses_cluster != fallback_uses_cluster - return Boolean(preferred_uses_cluster) ^ ( - is_fallback_cluster & Boolean(cluster_sync_behavior_differs) - ) + return Boolean(preferred_uses_cluster) ^ (is_fallback_cluster & Boolean(cluster_sync_behavior_differs)) @cute.jit @@ -233,17 +198,11 @@ def create( or num_mma_consumer_warps <= 0 ): raise ValueError("num_mma_consumer_warps must be a positive Python int.") - if cutlass.const_expr( - len(mcast_mode_mn) != 2 or any(mode not in (0, 1) for mode in mcast_mode_mn) - ): + if cutlass.const_expr(len(mcast_mode_mn) != 2 or any(mode not in (0, 1) for mode in mcast_mode_mn)): raise ValueError("mcast_mode_mn must contain two zero-or-one values.") preferred_shape_vmnk = _static_vmnk_shape(preferred_cta_layout_vmnk) - fallback_layout = ( - preferred_cta_layout_vmnk - if fallback_cta_layout_vmnk is None - else fallback_cta_layout_vmnk - ) + fallback_layout = preferred_cta_layout_vmnk if fallback_cta_layout_vmnk is None else fallback_cta_layout_vmnk fallback_shape_vmnk = _static_vmnk_shape(fallback_layout) if cutlass.const_expr(preferred_shape_vmnk[0] != fallback_shape_vmnk[0]): raise ValueError("Preferred and fallback layouts must use the same MMA V extent.") @@ -256,21 +215,15 @@ def create( cta_group = tcgen05.CtaGroup.TWO if mma_cta_count == 2 else tcgen05.CtaGroup.ONE is_true_mixed = preferred_shape_vmnk != fallback_shape_vmnk if cutlass.const_expr(is_true_mixed and is_fallback_cluster is None): - raise ValueError( - "is_fallback_cluster is required when preferred and fallback layouts differ." - ) + raise ValueError("is_fallback_cluster is required when preferred and fallback layouts differ.") preferred_group_count = _multicast_group_count(preferred_shape_vmnk, mcast_mode_mn) fallback_group_count = _multicast_group_count(fallback_shape_vmnk, mcast_mode_mn) preferred_arrival_count = preferred_group_count * num_mma_consumer_warps fallback_arrival_count = fallback_group_count * num_mma_consumer_warps - preferred_cluster_size = ( - preferred_shape_vmnk[0] * preferred_shape_vmnk[1] * preferred_shape_vmnk[2] - ) - fallback_cluster_size = ( - fallback_shape_vmnk[0] * fallback_shape_vmnk[1] * fallback_shape_vmnk[2] - ) + preferred_cluster_size = preferred_shape_vmnk[0] * preferred_shape_vmnk[1] * preferred_shape_vmnk[2] + fallback_cluster_size = fallback_shape_vmnk[0] * fallback_shape_vmnk[1] * fallback_shape_vmnk[2] tracks_both_operand_families = mcast_mode_mn == (1, 1) is_single_sided = tracks_both_operand_families and ( (preferred_shape_vmnk[1] == fallback_shape_vmnk[1] == 1) @@ -285,9 +238,7 @@ def create( if cutlass.const_expr(is_single_sided): consumer_mask = Int32((1 << fallback_cluster_size) - 1) else: - consumer_mask = _reader_mask( - fallback_layout, cta_coord_vmnk, mcast_mode_mn, mma_cta_count - ) + consumer_mask = _reader_mask(fallback_layout, cta_coord_vmnk, mcast_mode_mn, mma_cta_count) else: active_arrival_count = Int32(preferred_arrival_count) if cutlass.const_expr(is_single_sided): @@ -297,9 +248,7 @@ def create( preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count ) elif cutlass.const_expr(not is_single_sided): - consumer_mask = _reader_mask( - preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count - ) + consumer_mask = _reader_mask(preferred_cta_layout_vmnk, cta_coord_vmnk, mcast_mode_mn, mma_cta_count) full_sync_object = pipeline.MbarrierArray( barrier_storage=barrier_storage.align(min_align=8), diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py index 52a7a9daf028..b0d6b3eb2ace 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/__init__.py @@ -1,2 +1,22 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell inference kernels.""" + +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Kernel, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + SwapABGatedActEpilogue, + TensorRole, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py index 52a7a9daf028..52aa7758812a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/__init__.py @@ -1,2 +1,30 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Blackwell MegaMoE inference kernels and issue helpers.""" + +from .block_scaled_swap_ab_fc12_extension import ( + BlockScaledSwapAbFc12Extension, + TensorRole, +) +from .block_scaled_swap_ab_fc12_epilogue import ( + GatedActEpilogueArgs, + SwapABGatedActEpilogue, +) +from .block_scaled_swap_ab_fc12_kernel import ( + BlockScaledSwapAbFc12Kernel, +) +from .block_scaled_swap_ab_fc12_mainloop import ( + BlockScaledSwapAbFc12Mainloop, +) +from .block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Kernel", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index 33b754f7b9cd..a363fcdd831a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -10,32 +10,29 @@ """ import dataclasses +import math from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union import cutlass import cutlass.cute as cute -import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace import cutlass.utils as utils +import cutlass.pipeline as pipeline import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Int64, T -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice from .....communication.token_protocol import TokenSrcMetadata -from .....helpers.constants import Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max, Nvfp4E2M1RcpLimit +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment -from .....helpers.flag_batch import GpuReleaseFlagBatchTracker +from .....helpers.flag_batch import make_flag_batch_tracker from .....helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, @@ -44,10 +41,8 @@ red_add_relaxed_sys_v2_bf16x2, ) from .....helpers.smem_workspace import SmemRegion, SmemWorkspace -from .....quant_def import CombineFormat, QuantKind from ....function_mapping import CoordinateSpace, FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -82,43 +77,27 @@ class QuantImpl: warp_idx: Optional[Any] = None pair_exchange_barrier: Optional[Any] = None - _directions: ClassVar[Tuple[str, ...]] = ( - "regs_in_thread", - "threads_with_the_same_reg", - "regs_in_pair_threads", - ) + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") # -- config / validation -------------------------------------------------- def __post_init__(self): if isinstance(self.quant_kind, CombineFormat): if not self.quant_kind.is_quantized: - raise ValueError( - f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}." - ) + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") elif not isinstance(self.quant_kind, QuantKind): - raise ValueError( - f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}." - ) + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") if self.sf_vec_direction not in self._directions: - raise ValueError( - f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}." - ) + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") if self.sf_vec_direction == "regs_in_pair_threads": if isinstance(self.quant_kind, CombineFormat): - raise ValueError( - "The paired direction is orthodox-only; combine blocks never straddle warps." - ) + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") if self.sf_vec_size % 2 != 0: - raise ValueError( - f"The paired direction needs an even sf_vec, got {self.sf_vec_size}." - ) + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: - raise ValueError( - "The paired direction needs lane_idx, warp_idx and pair_exchange_barrier." - ) + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: # A wider orthodox block cannot be reduced inside one thread: see the paired direction. raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") @@ -159,20 +138,14 @@ def _data_rcp_limit(self) -> float: @property def _block_share_per_thread(self) -> int: """How many of a block's elements this thread holds.""" - return ( - self.sf_vec_size // 2 - if self.sf_vec_direction == "regs_in_pair_threads" - else self.sf_vec_size - ) + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size # -- dispatch ------------------------------------------------------------- @cute.jit def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): - raise ValueError( - "prequant_reg size must be divisible by this thread's share of a block." - ) + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 if cutlass.const_expr(prequant_reg.element_type is not expected_in): @@ -181,9 +154,7 @@ def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermedi ) needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): - raise ValueError( - "smem_intermediate must be supplied for the paired direction and only for it." - ) + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") if cutlass.const_expr(not self._is_combine): if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) @@ -255,14 +226,10 @@ def mx_quant_regs_in_pair_threads_impl( if cutlass.const_expr(smem_intermediate.stride[0] != 1): # Thread-major would still be correct but would silently split the exchange into # per-block scalar accesses. - raise ValueError( - "The paired amax exchange must be block-major to stay one vector access." - ) + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") exchange_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Float32, - num_bits_per_copy=n_blocks * cutlass.Float32.width, + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width ) # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of # intermediate outputs, so a block that is twice that run wide joins adjacent warps. @@ -273,20 +240,12 @@ def mx_quant_regs_in_pair_threads_impl( half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) for vec_block_idx in cutlass.range_constexpr(n_blocks): half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) - cute.copy( - exchange_atom, - cute.coalesce(half_amax), - cute.coalesce(smem_intermediate[None, own_thread]), - ) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) self.pair_exchange_barrier.arrive_and_wait() partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) - cute.copy( - exchange_atom, - cute.coalesce(smem_intermediate[None, partner_thread]), - cute.coalesce(partner_amax), - ) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) data_blocks = [] sf_values = [] @@ -294,17 +253,13 @@ def mx_quant_regs_in_pair_threads_impl( block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) scale_e8m0, scale_f32 = self._e8m0(block_amax) sf_values.append(scale_e8m0) - data_blocks.append( - self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32)) - ) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) self._store_packed_blocks(data, data_blocks) sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) return data, sf @cute.jit - def nvfp4_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -361,9 +316,7 @@ def nvfp4_combine_quant_threads_with_the_same_reg_impl( return data, selected_sf @cute.jit - def mxfp8_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -439,8 +392,7 @@ def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.Tenso @cute.jit def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: vec = vector.from_elements( - T.vector(len(values), dtype.mlir_type), - [values[i].ir_value() for i in range(len(values))], + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] ) return cute.TensorSSA(vec, (len(values),), dtype) @@ -484,9 +436,7 @@ def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: acc = block[0] for elem_idx in cutlass.range_constexpr(1, cute.size(block)): acc = max_abs(acc, block[elem_idx]) - mag_bits = cutlass.Int32( - llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value()) - ) & cutlass.Int32(0x7FFFFFFF) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) @cute.jit @@ -592,15 +542,12 @@ class _TmemTranspose16x32Core: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @staticmethod def load_subtile_raw_acc( @@ -646,31 +593,19 @@ def load_subtile_raw_acc( # offsets here would wrap them as SSA values that cute treats as # alignment-unknown, tripping the atom's verifier. first_top_view = cute.make_tensor(ptr, _TmemTranspose16x32Core._tmem_layout(16, 32)) - first_bot_view = cute.make_tensor( - ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32) - ) + first_bot_view = cute.make_tensor(ptr + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) second_top_view = cute.make_tensor(ptr + 32, _TmemTranspose16x32Core._tmem_layout(16, 32)) - second_bot_view = cute.make_tensor( - ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32) - ) + second_bot_view = cute.make_tensor(ptr + 32 + half_lane_off, _TmemTranspose16x32Core._tmem_layout(16, 32)) first_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) first_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) second_top = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) second_bot = cute.make_rmem_tensor((16,), _TmemTranspose16x32Core._io_dtype) - cute.copy( - atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16) - ) - cute.copy( - atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16) - ) - cute.copy( - atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16) - ) - cute.copy( - atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16) - ) + cute.copy(atom_ld16x64, first_top_view, _TmemTranspose16x32Core._rmem_copy_view(first_top, 16)) + cute.copy(atom_ld16x64, first_bot_view, _TmemTranspose16x32Core._rmem_copy_view(first_bot, 16)) + cute.copy(atom_ld16x64, second_top_view, _TmemTranspose16x32Core._rmem_copy_view(second_top, 16)) + cute.copy(atom_ld16x64, second_bot_view, _TmemTranspose16x32Core._rmem_copy_view(second_bot, 16)) return (first_top, first_bot, second_top, second_bot) @@ -703,21 +638,11 @@ def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = No self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) - self._atom_ld16x64 = cute.make_copy_atom( - tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_st16x128 = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype - ) - self._atom_st32x32 = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_ld16x256 = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype - ) - self._atom_ld16x128 = cute.make_copy_atom( - tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype - ) + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) @@ -772,18 +697,10 @@ def r2_store(self) -> None: # -- R3 ------------------------------------------------------------------ def r3_load_top(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r3_load_bot(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r3_perm(self) -> None: for r in range(16): @@ -795,18 +712,10 @@ def r3_store(self) -> None: # -- R4 ------------------------------------------------------------------ def r4_load_top(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r4_load_bot(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r4_perm(self) -> None: for r in range(16): @@ -866,10 +775,7 @@ class TmemTranspose32x32Inplace: """ def __init__( - self, - tmem_ptr, - reg_tensor_top: Optional[cute.Tensor] = None, - reg_tensor_bot: Optional[cute.Tensor] = None, + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None ) -> None: if (reg_tensor_top is None) != (reg_tensor_bot is None): raise ValueError( @@ -962,26 +868,19 @@ class TmemTranspose32x64B16Movm: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @cute.jit def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): - raise TypeError( - f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}." - ) + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") if cutlass.const_expr(cute.size(reg_tensor) != 64): - raise ValueError( - f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}." - ) + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) @@ -992,20 +891,14 @@ def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) - self._store_atom = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32 - ) - self._load_atom = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32 - ) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) @cute.jit def __call__(self) -> cute.Tensor: movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) - cute.copy( - self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom - ) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) output_words = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) @@ -1020,9 +913,13 @@ class GatedActEpilogueArgs: fc2_alpha: Optional[cute.Tensor] fc1_norm_const: Optional[cute.Tensor] # ----------------------------------- - # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the - # score mul. + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. topk_scores: Optional[cute.Tensor] + # ----------------------------------- + # Pool-ordered (token,) fp32. The FC1 activation was quantized from ``x / s_t``, so the FC1 + # epilogue multiplies ``s_t`` back onto the raw accumulator before clamp / activation. It has + # no other home: the e4m3 block scale cannot carry an fp32 factor and alpha is per-expert. + fc1_activation_per_token_scale: Optional[cute.Tensor] = None class SwapABGatedActEpilogue(KernelComponent): @@ -1052,6 +949,10 @@ class SwapABGatedActEpilogue(KernelComponent): # warp 32 accumulator rows, which the gate/up interleave halves. _EpilogueFc1IntermediateDownPerWarp = 16 smem_scratch_overlay: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.scratch" + # (subtile, warp, 32) fp32 de-interleave buffer for the per-token FC1 activation scale (see + # SwapABFc1Epilogue.run_subtile): 512 B per subtile, 2 KB at cta_tile_n=256. Registered only + # when the descriptor enables the scale (``fc1_activation_per_token_scale``). + fc1_act_scale_scratch_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_act_scale_scratch" fc1_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc1_staging" fc2_staging_region: ClassVar[str] = "blackwell.swap_ab_gated_act_epilogue.fc2_staging" # A 144 B FP8 stride rotates each row by four SMEM banks. @@ -1073,8 +974,17 @@ def problem_desc_require(cls) -> dict[str, type]: # existing SwiGLU descriptors stay valid unchanged. See _resolve_situ_betas. "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu": up * gate * sigmoid(1.702 * gate). Same core as + # SwiGLU with the sigmoid argument scaled, so it costs nothing extra; exclusive with SiTU. + "gate_activation": OptionalRequirement(str), + # Whether run() will receive the per-token FC1 activation scale. Absent means False: + # the SMEM scratch it needs is reserved only for descriptors that enable it. + "fc1_activation_per_token_scale": OptionalRequirement(bool), } + _supported_gate_activations: ClassVar[Tuple[str, ...]] = ("swiglu", "quick_geglu") + _quick_geglu_sigmoid_scale: ClassVar[float] = 1.702 + @classmethod def impl_desc_require(cls) -> dict[str, object]: return { @@ -1087,6 +997,7 @@ def impl_desc_require(cls) -> dict[str, object]: "fc2_epi_flag_batch": int, "fc2_tma_stages": OptionalRequirement(int), "reduce_topk_in_kernel": OptionalRequirement(bool), + "epi_async_release": OptionalRequirement(bool), "token_back_push_data": OptionalRequirement(bool), } @@ -1102,23 +1013,38 @@ def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Opt beta = problem_desc.get("situ_beta") linear_beta = problem_desc.get("situ_linear_beta") if (beta is None) != (linear_beta is None): - raise ValueError( - f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}." - ) + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") if beta is None: return None, None if beta <= 0 or linear_beta <= 0: - raise ValueError( - f"SiTU beta parameters must be positive, got {beta} and {linear_beta}." - ) + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") if problem_desc["gate_up_clamp"] is not None: - raise ValueError( - "SiTU does not support gate_up_clamp; the two activation variants are exclusive." - ) + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") if problem_desc.get("swiglu_alpha") is not None: raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) + @classmethod + def _resolve_gate_activation(cls, problem_desc: ProblemDesc) -> str: + """Resolve the gated-activation core name; absent means SwiGLU. + + quick_geglu only rescales the sigmoid argument of the SwiGLU core (1.702), so it keeps the + clamp semantics; SiTU is a different core and cannot be combined with it. Baked in at + codegen time, so the enclosing KernelClass folds it into ``name()``. + """ + gate_activation = problem_desc.get("gate_activation") + if gate_activation is None: + return "swiglu" + if gate_activation not in cls._supported_gate_activations: + raise ValueError( + f"gate_activation must be one of {cls._supported_gate_activations}, got {gate_activation!r}." + ) + if gate_activation != "swiglu" and problem_desc.get("situ_beta") is not None: + raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SiTU betas.") + if gate_activation != "swiglu" and problem_desc.get("swiglu_alpha") is not None: + raise ValueError(f"gate_activation={gate_activation!r} cannot be combined with SwiGLUBias constants.") + return gate_activation + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) @@ -1136,10 +1062,16 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + self.gate_activation = self._resolve_gate_activation(problem_desc) + # The scratch costs mainloop SMEM budget, so only a descriptor that enables the scale + # reserves it; passing the tensor without the flag is rejected in run(). + self.act_scale_scratch_enabled = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + # fc1/fc2 done flags published with red.async.release (default) instead of the blocking red.release (see flag_batch). + self.epi_async_release = impl_desc.get("epi_async_release", True) self.communication_enabled = impl_desc["communication_enabled"] self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] @@ -1147,17 +1079,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.communication_enabled: for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): if field_name not in impl_desc: - raise KeyError( - f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}." - ) + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) self.token_back_push_data = impl_desc.get("token_back_push_data", False) - if not self.communication_enabled and ( - self.reduce_topk_in_kernel or self.token_back_push_data - ): - raise ValueError( - "A communication-disabled Epilogue cannot enable communication policies." - ) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. self.fc1_output_dtype = self.quant_kind.activation_dtype @@ -1177,9 +1103,10 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel and not self.token_back_push_data if not 1 <= self.fc1_epi_flag_batch <= 32 or not 1 <= self.fc2_epi_flag_batch <= 32: raise ValueError("Epilogue flag batch sizes must be in [1, 32].") - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + if self.epi_async_release and max(self.fc1_epi_flag_batch, self.fc2_epi_flag_batch) > self._EpilogueWarpCnt: + # The async tracker batches per warp (one elected publisher per warp), the blocking one per lane. + raise ValueError(f"epi_async_release=True needs epilogue flag batch sizes <= {self._EpilogueWarpCnt}.") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -1199,17 +1126,10 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns self.overlapping_accum = self.accumulator_overlap_columns > 0 - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) - if ( - isinstance(self.hidden_size, int) - and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 - ): + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: self.fc2_hidden_needs_predicate: bool = False else: self.fc2_hidden_needs_predicate: bool = True @@ -1223,10 +1143,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's @@ -1235,32 +1152,21 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # bytes -- see the lifetime argument in fc1_quant. self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") # Additional stages trade mainloop SMEM for store overlap. - self.fc2_tma_stages = ( - requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 - ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + self.fc2_tma_stages = requested_fc2_tma_stages if requested_fc2_tma_stages is not None else 1 + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None @@ -1274,9 +1180,7 @@ def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to different work tiles. """ - return pipeline.NamedBarrier( - barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt - ) + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. @@ -1284,6 +1188,15 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: FC1's staging tile and FC2's store tile belong to different work tiles, separated by the tile-boundary TMA drain and rendezvous in ``run()``, so they never coexist. """ + if self.act_scale_scratch_enabled: + # (subtile, warp, 32): every subtile of the work tile keeps its own slot so the whole + # tile's scales can be staged once, ahead of the accumulator wait. + smem_workspace.register_tensor( + self.fc1_act_scale_scratch_region, + cutlass.Float32, + (self.subtile_cnt * self._EpilogueWarpCnt * 32,), + byte_alignment=64, + ) overlay = smem_workspace.create_overlay(self.smem_scratch_overlay) overlay.add_lifetime("fc1_staging").register_tensor( self.fc1_staging_region, cutlass.Int8, (self.fc1_staging_bytes,), byte_alignment=128 @@ -1331,13 +1244,9 @@ def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[i raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") return shape, stride, swizzle - def fc2_tma_staged_smem_layout( - self, n_stages: int, without_stage_mode: bool = False - ) -> cute.ComposedLayout: + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) - layout = cute.make_composed_layout( - cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride) - ) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) if without_stage_mode: return cute.select(layout, mode=[0, 1]) return layout @@ -1377,10 +1286,7 @@ def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: kind="tensor", dtype=single_stage_region.dtype, shape=(*single_stage_region.shape, stage_count), - stride=( - *single_stage_region.stride, - single_stage_region.cosize if stage_count > 1 else 0, - ), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), swizzle=single_stage_region.swizzle, byte_alignment=single_stage_region.byte_alignment, ) @@ -1450,21 +1356,25 @@ def run( fc1_alpha=None, fc2_alpha=None, fc1_norm_const=None, topk_scores=None ) fc1_staging_pointer = smem_workspace.ptr(self.fc1_staging_region, smem_base) + if cutlass.const_expr(self.act_scale_scratch_enabled): + fc1_act_scale_scratch_pointer = smem_workspace.ptr(self.fc1_act_scale_scratch_region, smem_base) + else: + if cutlass.const_expr(optional_epi_args.fc1_activation_per_token_scale is not None): + raise ValueError( + "fc1_activation_per_token_scale was passed, but the descriptor does not enable it " + "(problem_desc['fc1_activation_per_token_scale'] must be True), so no SMEM scratch was reserved." + ) + fc1_act_scale_scratch_pointer = None if cutlass.const_expr(self.fc2_tma_stages > 0): fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) else: fc2_smem_tensor = None - if cutlass.const_expr( - self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None) - ): + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") if cutlass.const_expr( - self.communication_enabled - and (token_src_metadata is None or peer_rank_ptr_mapper is None) + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) ): - raise ValueError( - "Communication-enabled Epilogue requires token metadata and a peer pointer mapper." - ) + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): raise ValueError("Token-back requires an FC2 done counter.") if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): @@ -1478,6 +1388,7 @@ def run( self, tidx, fc1_staging_pointer, + fc1_act_scale_scratch_pointer, kernel_extension, tma_atom_fc1_output, fc1_output, @@ -1505,8 +1416,20 @@ def run( wait_only_named_barrier = self.epilogue_sync_barrier() is_odd_turn = cutlass.Int32(1) work_tile_info = sched_consumer.consume_work() - - flag_tracker = GpuReleaseFlagBatchTracker( + # With per-token scales the epilogue runs one work tile ahead on the scheduler ring: the + # next tile's scale loads are issued before the current tile is processed, so their + # latency hides behind a whole tile instead of sitting in front of the activation. A + # prefetch may only be issued once at least one acc wait lies behind it: in the EP kernel + # that wait is what orders the pool reads after every rank's router pushes. So the CTA's + # first FC1 tile loads its own scales after its acc wait (the one exposed load per CTA) and + # prefetches the next tile from there; every other tile is prefetched at the loop top. + prefetch_act_scales = cutlass.const_expr(optional_epi_args.fc1_activation_per_token_scale is not None) + # Whether the CURRENT tile's first-half scales sit in the prefetch registers. + act_scales_prefetched = cutlass.Int32(0) + loop_iteration = cutlass.Int32(0) + + flag_tracker = make_flag_batch_tracker( + self.epi_async_release, flag_address=Int64(0), accumulated_flags=cutlass.Int32(0), phase=cutlass.Int32(work_tile_info.phase), @@ -1514,6 +1437,23 @@ def run( ) while work_tile_info.is_valid_tile: + if cutlass.const_expr(prefetch_act_scales): + # The current tile is valid, so the scheduler publishes at least a terminal tile + # after it: consuming here cannot block forever. Stage the current tile's scales + # (from the registers prefetched last iteration) before those registers are + # reloaded for the next tile. + next_work_tile_info = sched_consumer.consume_work() + next_act_scales_prefetched = cutlass.Int32(0) + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + if act_scales_prefetched == cutlass.Int32(1): + fc1_epi.stage_act_scales(work_tile_info) + fc1_epi.prefetch_act_scales(next_work_tile_info) + # else: fc1_epi loads + stages this tile after its acc wait, then prefetches + # the next one itself (the registers are busy until then). + next_act_scales_prefetched = cutlass.Int32(1) + elif loop_iteration > cutlass.Int32(0): + fc1_epi.prefetch_act_scales(next_work_tile_info) + next_act_scales_prefetched = cutlass.Int32(1) if cutlass.const_expr(self.overlapping_accum): tmem_stage_idx = acc_consumer_state.phase else: @@ -1527,6 +1467,8 @@ def run( acc_pipeline=acc_pipeline, acc_consumer_state=acc_consumer_state, is_odd_turn=is_odd_turn, + act_scales_staged=act_scales_prefetched, + next_work_tile_info=next_work_tile_info if cutlass.const_expr(prefetch_act_scales) else None, ) else: # The __call__ args should only take the while loop args, leave all loop irrevalent args to the init. @@ -1537,6 +1479,9 @@ def run( acc_consumer_state=acc_consumer_state, is_odd_turn=is_odd_turn, ) + if cutlass.const_expr(prefetch_act_scales): + act_scales_prefetched = next_act_scales_prefetched + loop_iteration = loop_iteration + cutlass.Int32(1) prev_work_tile_info = work_tile_info cur_was_linear1 = prev_work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1) @@ -1544,7 +1489,10 @@ def run( if cutlass.const_expr(self.overlapping_accum): is_odd_turn = cutlass.Int32(1) - is_odd_turn - work_tile_info = sched_consumer.consume_work() + if cutlass.const_expr(prefetch_act_scales): + work_tile_info = next_work_tile_info + else: + work_tile_info = sched_consumer.consume_work() # Every asynchronous store commits at issue; drain before completion or scratch reuse. cute.arch.cp_async_bulk_wait_group(0) @@ -1553,13 +1501,9 @@ def run( # Publish completion for the work tile snapshotted above. if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: - flag_tracker = fc2_epi.signal_fc2_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) # Tail flush flag_tracker.fire() @@ -1569,9 +1513,7 @@ class _ImmutableAfterInit: def __setattr__(self, name, value): if self.__dict__.get("_frozen_", False): - raise AttributeError( - f"{type(self).__name__} is immutable after __init__ (cannot set {name!r})." - ) + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") object.__setattr__(self, name, value) def _freeze(self) -> None: @@ -1585,6 +1527,7 @@ def __init__( base: SwapABGatedActEpilogue, tidx: cutlass.Int32, staging_pointer: cute.Pointer, + act_scale_scratch_pointer: cute.Pointer, kernel_extension: BlockScaledSwapAbFc12Extension, tma_atom_fc1_output: cute.CopyAtom, fc1_output: cute.Tensor, # fake (m,n,l) domain @@ -1600,12 +1543,27 @@ def __init__( # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) self.smem_tensor = cute.make_tensor( - cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), - staged_layout.outer, + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer ) # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one # staging stage without going through the staging tensor's swizzle. self.staging_pointer = staging_pointer + # (subtile_cnt * warp_cnt * 32,) fp32: per (subtile, warp) slot holding the first-half + # token scales de-interleaved by token parity (see stage_act_scales / run_subtile). + # Prefetch registers hold the NEXT FC1 tile's first-half scales (one per subtile, this + # lane's token) and its expert alpha: loaded while the current tile is processed so the + # global-load latency never sits in front of the activation. + if optional_epi_args.fc1_activation_per_token_scale is not None: + self.act_scale_scratch = cute.make_tensor( + cute.recast_ptr(act_scale_scratch_pointer, dtype=cutlass.Float32), + cute.make_layout((base.subtile_cnt * base._EpilogueWarpCnt * 32,)), + ) + self.prefetched_act_scales = cute.make_rmem_tensor((base.subtile_cnt,), cutlass.Float32) + self.prefetched_act_alpha = cute.make_rmem_tensor((1,), cutlass.Float32) + else: + self.act_scale_scratch = None + self.prefetched_act_scales = None + self.prefetched_act_alpha = None self.kernel_extension = kernel_extension self.fc1_tma_atom = tma_atom_fc1_output self.fc1_output = fc1_output @@ -1629,6 +1587,57 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "SwapABFc1Epilogue assert len(values) == 0 return self + @cute.jit + def prefetch_act_scales(self, work_tile_info: SwapAbFc12WorkTileInfo) -> None: + """Issue the global loads for ``work_tile_info``'s first-half token scales (and its alpha). + + Normally called one tile ahead of the tile's processing, so the loads land while the + current tile runs; a CTA's first FC1 tile calls it for itself after its acc wait. Only FC1 + tiles load, and only lanes whose token is valid (the pool view ends at the expert's valid + rows). The values stay in this lane's registers until stage_act_scales moves them to SMEM. + """ + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is None): + return + if work_tile_info.is_valid_tile: + if work_tile_info.phase == cutlass.Int32(BlockPhase.Linear1): + act_scale_tensor, _ = self.kernel_extension.get_gmem_tensor( + "fc1_act_scale", self.optional_epi_args.fc1_activation_per_token_scale, work_tile_info + ) + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + self.prefetched_act_alpha[0] = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] + tile_token_base = work_tile_info.tile_n_idx * self.cta_tile_n + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + for subtile in cutlass.range_constexpr(self.subtile_cnt): + token_in_tile = cutlass.Int32(subtile * self._EpilogueTokenTileSize) + self.lane_idx + if token_in_tile < valid_tokens: + self.prefetched_act_scales[subtile] = act_scale_tensor[tile_token_base + token_in_tile] + + @cute.jit + def stage_act_scales(self, work_tile_info: SwapAbFc12WorkTileInfo) -> None: + """Move the prefetched first-half scales of ``work_tile_info`` into this warp's SMEM slots. + + Per subtile, lane L owns token L of the first 32 tokens; it stores the alpha-folded scale + at slot (L%2)*16 + L//2 of the (subtile, warp) scratch row, so each token parity's 16 + scales are contiguous for run_subtile's four 16-byte loads. Lanes past the valid tokens + hold no scale (nothing was loaded for them) and stage 1.0, like the second half's default, + so the pad rows the subtile store still writes stay deterministic. + """ + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is None): + return + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + own_slot_offset = (self.lane_idx % cutlass.Int32(2)) * cutlass.Int32(16) + self.lane_idx // 2 + for subtile in cutlass.range_constexpr(self.subtile_cnt): + staged_scale = cutlass.Float32(1.0) + if cutlass.Int32(subtile * self._EpilogueTokenTileSize) + self.lane_idx < valid_tokens: + folded_scale = self.prefetched_act_scales[subtile] + if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): + folded_scale = folded_scale * self.prefetched_act_alpha[0] + staged_scale = folded_scale + self.act_scale_scratch[ + cutlass.Int32((subtile * self._EpilogueWarpCnt) * 32) + self.warp_idx * cutlass.Int32(32) + own_slot_offset + ] = staged_scale + cute.arch.sync_warp() + @cute.jit def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): # Only in-bound intermediate_downproj tiles signal; OOB -> null slot. @@ -1637,19 +1646,14 @@ def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 ) if cutlass.const_expr(needs_intermediate_guard): - in_bound = ( - work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize - < self.fc1_output.shape[1] - ) + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] else: in_bound = True slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx flag_address = Int64(0) if in_bound: flag_address = (self.fc1_done_counter.iterator + slot).toint() - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) @cute.jit def __call__( @@ -1659,15 +1663,16 @@ def __call__( acc_pipeline, acc_consumer_state, is_odd_turn: cutlass.Int32, + # 0 for a CTA's first FC1 tile: nothing was prefetched for it, so after the acc wait it + # loads and stages its own scales and then prefetches ``next_work_tile_info``'s (see + # SwapABGatedActEpilogue.run). + act_scales_staged: cutlass.Int32, + next_work_tile_info: Optional[SwapAbFc12WorkTileInfo] = None, ): # (tokens_this_expert, intermediate_down, 1) - real_fc1_output, _ = self.kernel_extension.get_gmem_tensor( - "c", self.fc1_output, work_tile_info - ) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) # (tokens_this_expert, intermediate_down, 1) - real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor( - "sfc", self.fc1_output_sf, work_tile_info - ) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) # subtile-irrevalent hoist out here. if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] @@ -1677,26 +1682,37 @@ def __call__( norm_const = self.optional_epi_args.fc1_norm_const[work_tile_info.expert_idx] else: norm_const = None + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if cutlass.const_expr(self.optional_epi_args.fc1_activation_per_token_scale is not None): + # (tokens_this_expert,), same pool-ordered view as topk scores. The first-half scales + # were staged into SMEM by stage_act_scales before this call. + act_scale_tensor, _ = self.kernel_extension.get_gmem_tensor( + "fc1_act_scale", self.optional_epi_args.fc1_activation_per_token_scale, work_tile_info + ) + else: + act_scale_tensor = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, - (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize), + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) )[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) - valid_tokens = work_tile_info.valid_tokens_in_cta_tile + if cutlass.const_expr(act_scale_tensor is not None): + # The acc wait orders these pool reads after the router's writes (dispatch pulled these + # tokens only after metadata_ready). Exposed once per CTA: the next tile's prefetch is + # issued right here, so every later FC1 tile finds its scales already in registers. + if act_scales_staged == cutlass.Int32(0): + self.prefetch_act_scales(work_tile_info) + self.stage_act_scales(work_tile_info) + self.prefetch_act_scales(next_work_tile_info) # Overlap path preloads two subtiles before releasing acc TMEM. unroll_tile_cnt = 2 if cutlass.const_expr(self.overlapping_accum) else 0 remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt if cutlass.const_expr(unroll_tile_cnt > 0): - subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32( - self.subtile_cnt - ) - subtile_idx_second = ( - cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's @@ -1735,22 +1751,19 @@ def __call__( self.run_subtile( work_tile_info=work_tile_info, subtile_idx=subtile_idx_pair[i], - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], preload_acc=preload_pair[i], fc1_output=real_fc1_output, fc1_output_sf=real_fc1_output_sf, alpha_val=alpha_val, norm_const=norm_const, + act_scale_tensor=act_scale_tensor, ) for i in cutlass.range(remain_subtile_cnt, unroll=1): real_i = i + unroll_tile_cnt if cutlass.const_expr(self.overlapping_accum): - subtile_idx = ( - cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) else: subtile_idx = cutlass.Int32(real_i) @@ -1764,6 +1777,7 @@ def __call__( fc1_output_sf=real_fc1_output_sf, alpha_val=alpha_val, norm_const=norm_const, + act_scale_tensor=act_scale_tensor, ) # Non-overlap-path release: at the natural task-tile boundary. @@ -1785,6 +1799,8 @@ def run_subtile( fc1_output_sf: cute.Tensor, alpha_val: Optional[cutlass.Float32], norm_const: Optional[cutlass.Float32], + # (tokens_this_expert,) fp32 per-token activation scale view, or None. + act_scale_tensor: Optional[cute.Tensor] = None, ): if cutlass.const_expr(self.optional_epi_args.topk_scores is not None): # This means we need to perform DeepGEMM computation graph, topk_score at fc1 pre-quant @@ -1802,9 +1818,7 @@ def run_subtile( # Step -1: preload topk scores. current_two_token_idices = ( - work_tile_info.tile_n_idx * self.cta_tile_n - + subtile_idx * self._EpilogueTokenTileSize - + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx @@ -1817,6 +1831,53 @@ def run_subtile( ) else: topk_scores = None + # Per-token activation scale, with the expert alpha folded in once per lane so the + # activation's dequant multiply stays a single packed FMUL per register pair. + # second half (token 32..63): activated after its transpose, one token per lane -> the + # lane's own scalar. + # first half (token 0..31): activated before its transpose, where register i of lane L + # belongs to token 2i + (L//2)%2 -> the lane needs the 16 scales of one parity. Lane L + # writes token L's scale to slot (L%2)*16 + L//2 of this warp's scratch (even tokens + # first, then odd), so each parity's 16 scales are contiguous and come back with four + # 16-byte LDS instead of 16 shuffles. + # Both halves end up as a complete dequant multiplier (alpha x s_t), so alpha_swiglu_clamp + # receives it INSTEAD of alpha_val and multiplies exactly once per register. + if cutlass.const_expr(act_scale_tensor is not None): + # Second half: one token per lane (token lane+32). Loaded here, alpha folded here; it is + # consumed only after the 32x32 transpose, which hides the load latency. Lanes past the + # valid tokens skip the load (the pool view ends at the expert's valid rows). + dequant_scale_token_32_64 = cutlass.Float32(1.0) + second_half_token_in_tile = subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) + self.lane_idx + 32 + if second_half_token_in_tile < work_tile_info.valid_tokens_in_cta_tile: + dequant_scale_token_32_64 = act_scale_tensor[current_two_token_idices[1]] + if cutlass.const_expr(alpha_val is not None): + dequant_scale_token_32_64 = dequant_scale_token_32_64 * alpha_val + # First half: 16 tokens per lane. The 16 values were loaded (prefetch_act_scales), + # multiplied by alpha and written de-interleaved to SMEM by stage_act_scales before this + # tile started; here they are only read back. Alpha is therefore already included. + lane_parity = (self.lane_idx // cutlass.Int32(2)) % cutlass.Int32(2) + scale_vec_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128) + # ((subtile*warps + warp)*32 + parity*16) * 4 B is a multiple of 64 B; the dynamic + # offset hides that from the copy verifier. + parity_scales_smem = mark_alignment( + cute.make_tensor( + self.act_scale_scratch.iterator + + ( + (subtile_idx * cutlass.Int32(self._EpilogueWarpCnt) + self.warp_idx) * cutlass.Int32(32) + + lane_parity * cutlass.Int32(16) + ), + cute.make_layout((16,)), + ), + 64, + ) + dequant_scales_token_0_32 = cute.make_rmem_tensor((16,), cutlass.Float32) + cute.copy(scale_vec_atom, parity_scales_smem, dequant_scales_token_0_32) + # With a per-token scale the dequant multiplier carries alpha; do not apply it twice. + activation_alpha = None + else: + dequant_scales_token_0_32 = None + dequant_scale_token_32_64 = None + activation_alpha = alpha_val # Step 0: load tmem if cutlass.const_expr(preload_acc is not None): @@ -1828,11 +1889,8 @@ def run_subtile( up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) - # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, - # self.warp_idx), None] - tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, 0), None - ] + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) tmem_subtile_tensor_in_first_load_view = cute.logical_divide( cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) @@ -1861,15 +1919,13 @@ def run_subtile( # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( - gate_token_0_32, up_token_0_32, alpha_val + gate_token_0_32, up_token_0_32, activation_alpha, dequant_scales=dequant_scales_token_0_32 ) # gate_token_32_64 / up_token_32_64 are already in the transpose input # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). token_32_64_tmem_trans = TmemTranspose32x32Inplace( - tmem_subtile_tensor.iterator, - reg_tensor_top=gate_token_32_64, - reg_tensor_bot=up_token_32_64, + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 ) # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp @@ -1879,7 +1935,10 @@ def run_subtile( ) token_32_64_pre_quant = self.alpha_swiglu_clamp( - gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act, alpha_val + gate_token_32_64_trans_pre_act, + up_token_32_64_trans_pre_act, + activation_alpha, + dequant_scale=dequant_scale_token_32_64, ) token_0_32_tmem_trans = TmemTranspose16x32( @@ -1907,8 +1966,7 @@ def run_subtile( ] # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) fc1_gmem_subtile_view = cute.flat_divide( - fc1_gmem_cta_view, - (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) )[None, None, subtile_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( self.fc1_tma_atom, @@ -1938,27 +1996,28 @@ def alpha_swiglu_clamp( gate_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem up_rmem: cute.Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem alpha_val: Optional[cutlass.Float32], + *, + # Full dequant multiplier with alpha ALREADY folded in (the caller multiplies once per + # lane), replacing alpha_val when given. Exactly one form applies: + # dequant_scale : one token per lane (post-transpose distribution), one scalar. + # dequant_scales: (n,) rmem, one value per register (pre-transpose distribution). + dequant_scale: Optional[cutlass.Float32] = None, + dequant_scales: Optional[cute.Tensor] = None, ) -> cute.Tensor: # ── Input contract checks (compile-time): fp32, 1D, even-count, rmem ── # Wrapped in const_expr so the DSL evaluates them at trace time and the # raise fires during compilation rather than emitting a runtime branch. for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): if cutlass.const_expr(_t.element_type is not cutlass.Float32): - raise TypeError( - f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}" - ) + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") if cutlass.const_expr(_t.memspace != AddressSpace.rmem): raise ValueError( f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" ) if cutlass.const_expr(cute.rank(_t) != 1): - raise ValueError( - f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") if cutlass.const_expr(cute.size(_t) % 2 != 0): - raise ValueError( - f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): raise ValueError( "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " @@ -1968,8 +2027,9 @@ def alpha_swiglu_clamp( # gate_rmem / up_rmem are the RAW fc1 fp32 accumulator (pre-dequant). # Order follows the NVFP4 -> fp32 -> SwiGLU contract and MUST be: # - # 1. dequant: gate = alpha * gate_raw ; up = alpha * up_raw - # (alpha = expert-wise global scale on the acc; None => alpha == 1.) + # 1. dequant: gate = alpha * s_t * gate_raw ; up = alpha * s_t * up_raw + # (alpha = expert-wise global scale on the acc; s_t = per-token activation scale + # the input was divided by before quantization; None => 1 for either.) # 2. clamp the DEQUANTED (real) values, gpt-oss ``_apply_gate`` style: # gate = min(gate, +limit) (upper bound only) # up = clamp(up, -limit, +limit) (symmetric) @@ -1987,17 +2047,28 @@ def alpha_swiglu_clamp( out = cute.make_rmem_tensor((n,), cutlass.Float32) log2_e = 1.4426950408889634 - activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha - neg_activation_alpha_log2e = -activation_alpha * log2_e - neg_log2e_pair = ( - cutlass.Float32(neg_activation_alpha_log2e), - cutlass.Float32(neg_activation_alpha_log2e), - ) + # quick_geglu: out = up * gate * sigmoid(1.702 * gate). The 1.702 folds into the exp2 constant + # (one Python-side double product, rounded once to fp32), so the SwiGLU instruction sequence + # is reused unchanged. SiTU never reaches here with quick_geglu (_resolve_gate_activation). + sigmoid_scale = self._quick_geglu_sigmoid_scale if self.gate_activation == "quick_geglu" else 1.0 + if cutlass.const_expr(self.swiglu_alpha is not None): + sigmoid_scale = self.swiglu_alpha + neg_log2e_pair = (cutlass.Float32(-sigmoid_scale * log2_e), cutlass.Float32(-sigmoid_scale * log2_e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) if cutlass.const_expr(self.swiglu_beta is not None): beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) if cutlass.const_expr(self.gate_up_clamp is not None): limit = cutlass.Float32(self.gate_up_clamp) + if cutlass.const_expr(dequant_scale is not None and dequant_scales is not None): + raise ValueError("alpha_swiglu_clamp: pass at most one of dequant_scale and dequant_scales.") + if cutlass.const_expr(alpha_val is not None and (dequant_scale is not None or dequant_scales is not None)): + raise ValueError( + "alpha_swiglu_clamp: a dequant multiplier already carries alpha; pass alpha_val=None with it." + ) + if cutlass.const_expr(dequant_scales is not None and cute.size(dequant_scales) != n): + raise ValueError( + f"alpha_swiglu_clamp: dequant_scales must hold {n} values, got {cute.size(dequant_scales)}" + ) for i in cutlass.range_constexpr(0, n, 2): g0 = gate_rmem[i] @@ -2005,11 +2076,18 @@ def alpha_swiglu_clamp( u0 = up_rmem[i] u1 = up_rmem[i + 1] - # 1) dequant raw acc to real values (skip entirely when alpha is None). - if cutlass.const_expr(alpha_val is not None): - alpha_pair = (alpha_val, alpha_val) - g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), alpha_pair) - u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), alpha_pair) + # 1) dequant raw acc to real values (skip entirely when alpha and scale are None). + if cutlass.const_expr(dequant_scales is not None): + scale_pair = (dequant_scales[i], dequant_scales[i + 1]) + elif cutlass.const_expr(dequant_scale is not None): + scale_pair = (dequant_scale, dequant_scale) + elif cutlass.const_expr(alpha_val is not None): + scale_pair = (alpha_val, alpha_val) + else: + scale_pair = None + if cutlass.const_expr(scale_pair is not None): + g0, g1 = cute.arch.mul_packed_f32x2((g0, g1), scale_pair) + u0, u1 = cute.arch.mul_packed_f32x2((u0, u1), scale_pair) # 2) clamp the real values (skip when no clamp configured). if cutlass.const_expr(self.gate_up_clamp is not None): @@ -2071,6 +2149,7 @@ def _sigmoid(p0, p1): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) # The SwiGLUBias offset follows the up clamp. if cutlass.const_expr(self.swiglu_beta is not None): u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) @@ -2116,9 +2195,7 @@ def _sigmoid(p0, p1): def fc1_quant( self, work_tile_info: SwapAbFc12WorkTileInfo, - two_token: Tuple[ - cute.Tensor, cute.Tensor - ], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], norm_const: Optional[cutlass.Float32], intermediate_output_size: cutlass.Int32, @@ -2152,8 +2229,7 @@ def fc1_quant( cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), cute.make_layout( - (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), - stride=(1, self.fc1_amax_token_chunks), + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) ), ), 16, @@ -2180,13 +2256,8 @@ def fc1_quant( stores_scale_factor = self.warp_idx % 2 == 0 else: stores_scale_factor = True - subtile_token_start = ( - work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize - ) - token_idx_pair = ( - subtile_token_start + self.lane_idx, - subtile_token_start + self.lane_idx + 32, - ) + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous @@ -2243,9 +2314,7 @@ def fc1_quant( # 4) STS the quantized values into this subtile's shared output stage. # ((1, 16), (token_tile_size, warp_cnt)) -> (16) smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] - cute.copy( - store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row) - ) + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) @dataclasses.dataclass(frozen=True) @@ -2343,20 +2412,15 @@ def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): if cutlass.const_expr(publish): flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() no_fire: cutlass.Constexpr = not publish - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -2365,9 +2429,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out metadata = None peer_rank_ptr_mapper = None data_token_base = task_tile_data_row_start - if cutlass.const_expr( - self.token_src_metadata is not None and not self.token_back_push_data - ): + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) peer_rank_ptr_mapper = self.peer_rank_ptr_mapper data_token_base = None @@ -2375,10 +2437,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out if cutlass.const_expr(self.combine_format.is_quantized): base_outputs = (self.fc2_output, self.fc2_output_sf) token_bases = (data_token_base, task_tile_data_row_start) - output_mappings = ( - self.process_pipeline.store_out_mapping, - self.process_pipeline.sf_store_out_mapping, - ) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) else: base_outputs = self.fc2_output token_bases = data_token_base @@ -2427,21 +2486,13 @@ def __call__( # Overlap path preloads two subtiles before releasing acc TMEM. unroll_tile_cnt = ( - 2 - if cutlass.const_expr( - self.overlapping_accum and self.process_pipeline.require_tmem_trans - ) - else 0 + 2 if cutlass.const_expr(self.overlapping_accum and self.process_pipeline.require_tmem_trans) else 0 ) remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt if cutlass.const_expr(unroll_tile_cnt > 0): - subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32( - self.subtile_cnt - ) - subtile_idx_second = ( - cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) # Each warp preloads its local 32-hidden x 64-token accumulator as two # (32,) FP32 tensors: @@ -2453,10 +2504,7 @@ def __call__( # consecutive register pair holds adjacent token columns at one # hidden coordinate. After F2FP, that pair is one BF16x2 MOVM input. preload_subtile_first: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_first - ], - epi=self, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_first], epi=self ) # Release acc to next MMA unconditionally. @@ -2466,10 +2514,7 @@ def __call__( # Same two-half mapping for the other token subtile. Both preloaded # subtiles can now use this second subtile's TMEM as transpose workspace. preload_subtile_second: Tuple[cute.Tensor, ...] = self.process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], - epi=self, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], epi=self ) # Both unrolled subtiles borrow tmem_subtile_second as workspace. @@ -2481,9 +2526,7 @@ def __call__( work_tile_info=work_tile_info, epilogue_iter_idx=cutlass.Int32(i), subtile_idx=subtile_idx_pair[i], - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second - ], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx_second], preload_acc=preload_pair[i], fc2_output_router=fc2_output_router, alpha_val=alpha_val, @@ -2500,9 +2543,7 @@ def __call__( # for i in cutlass.range_constexpr(remain_subtile_cnt): real_i = i + unroll_tile_cnt if cutlass.const_expr(self.overlapping_accum): - subtile_idx = ( - cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn - ) % cutlass.Int32(self.subtile_cnt) + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn) % cutlass.Int32(self.subtile_cnt) else: subtile_idx = cutlass.Int32(real_i) @@ -2543,9 +2584,7 @@ def run_subtile( ): process_pipeline = self.process_pipeline if cutlass.const_expr(preload_acc is None): - loaded = process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_subtile_tensor, epi=self - ) + loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) if release_after_ldtm: cute.arch.fence_view_async_tmem_load() acc_pipeline.consumer_release(acc_consumer_state) @@ -2555,9 +2594,7 @@ def run_subtile( casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) # reorder returns a bare RMEM fragment in the store's expected pre-store # distribution; reorder + store are paired 1:1 inside the pipeline. - pre_store = process_pipeline.post_f2fp_reorder( - casted=casted, tmem_subtile_view=tmem_subtile_tensor - ) + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) process_pipeline.store_function( epi=self, subtile=pre_store, @@ -2586,9 +2623,7 @@ class Fc2OutputRouter: epi_tid: cutlass.Int32 combine_format: CombineFormat # After metadata prefetch - dst_ptrs: Optional[cute.Tensor] = ( - None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. - ) + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) @property @@ -2610,11 +2645,7 @@ def sf_token_base(self) -> Optional[cutlass.Int32]: @property def data_mapping(self) -> FunctionMapping: - return ( - self.output_mappings[0] - if isinstance(self.output_mappings, tuple) - else self.output_mappings - ) + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings @property def sf_mapping(self) -> Optional[FunctionMapping]: @@ -2622,9 +2653,7 @@ def sf_mapping(self) -> Optional[FunctionMapping]: def __post_init__(self) -> None: if (self.metadata is None) == (self.data_token_base is None): - raise ValueError( - "Fc2OutputRouter requires exactly one of metadata or a (data) token base." - ) + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") if self.reduce_topk_in_epilogue and self.metadata is None: @@ -2643,8 +2672,7 @@ def prefetch(self) -> "Fc2OutputRouter": valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) - # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + - # strength_reduce) + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) # We should check the SASS to ensure this happens. for iter_idx in cutlass.range_constexpr(copy_iters): coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) @@ -2663,14 +2691,11 @@ def prefetch(self) -> "Fc2OutputRouter": dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile # Int64 token coord: dst_tokens*K*H overflows int32 once # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). - dst_ptrs[iter_idx] = self.data_output[ - Int64(dst_tokens), None, dst_hidden - ].iterator.toint() + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() else: md = TokenSrcMetadata.load( - self.metadata.iterator.toint() - + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) ) dst_rank = md.src_rank dst_token = md.src_token @@ -2682,9 +2707,7 @@ def prefetch(self) -> "Fc2OutputRouter": # Int64 token coord: domain_offset on (token, topk, hidden) # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( - cute.domain_offset( - (Int64(dst_token), dst_topk, dst_hidden), self.data_output - ).iterator, + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, dst_rank, byte_alignment=32, ).toint() @@ -2692,9 +2715,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: """Per-issue DATA destination: gmem pointer + validity predicate. The router owns ``data_output`` so the caller never re-assembles a @@ -2726,9 +2747,7 @@ def get_data_dst( # comm: read the pointer / validity prefetched by prefetch(). addr = self.dst_ptrs[iter_idx] pred = self.valid[iter_idx] - ptr = cute.make_ptr( - self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32 - ) + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) return ptr, pred @cute.jit @@ -2755,9 +2774,7 @@ def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() # Per-block scale offsets are element-granular; claim the scale dtype's # natural element alignment (e8m0 1 B / bf16 2 B). - sf_ptr = cute.make_ptr( - self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4 - ) + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) return sf_ptr, pred @@ -2781,16 +2798,13 @@ def make_fc2_stg_cta_store_out_mapping( }, ) store_out_mapping = FunctionMapping( - domain=CoordinateSpace( - ("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32) - ), + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), codomain=CoordinateSpace( ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, - "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg - + epi_tid // 32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, }, ) sf_store_out_mapping = None @@ -2804,8 +2818,7 @@ def stg_sf_mapping(epi_tid, iter_idx): sf_store_out_mapping = FunctionMapping( domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), codomain=CoordinateSpace( - ("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size), + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=stg_sf_mapping, ) @@ -2835,10 +2848,7 @@ def make_fc2_redg_cta_store_out_mapping( + (epi_tid % 32) // 4 ), "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 - + (epi_tid % 4) * 4 - + (((elem_idx // 4) % 4) // 2) * 16 - + elem_idx % 4 + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 ), }, ) @@ -2858,17 +2868,13 @@ def make_fc2_redg_cta_store_out_mapping( + ((iter_idx % 4) % 2) * 8 + (epi_tid % 32) // 4 ), - "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16 - ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), }, ) return store_out_mapping, None, fundamental_mapping -def make_fc2_ublk_store_out_mapping( - combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int -): +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): assert cta_hidden_tile_size == 128 assert cta_token_tile_size % 64 == 0 # UBLK pushes whole hidden rows by byte count, so the token/hidden mapping @@ -2897,9 +2903,7 @@ def make_fc2_ublk_store_out_mapping( ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { - "token_in_cta_tile": ( - iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16 - ), + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), "hidden_in_cta_tile": 0, }, ) @@ -2933,9 +2937,7 @@ def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute elems_processed = 0 for t in tensors: current_tensor_size = cute.size(t) - dst = cute.make_tensor( - converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,)) - ) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) if cutlass.const_expr(alpha_val is None): dst.store(t.load().to(reorder_dtype)) else: @@ -2965,9 +2967,7 @@ def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): ptr = tmem_subtile_tensor.iterator half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) - bottom_view = cute.make_tensor( - ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64) - ) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) top = cute.make_rmem_tensor((32,), cutlass.Float32) bottom = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) @@ -2980,13 +2980,9 @@ def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes # the CTA-level 128-hidden x 64-token subtile view, so select this epi # warp's hidden block before issuing LDTM.x64. - tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, epi.warp_idx), None - ] + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) - atom_ld32x32_x64 = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32 - ) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) cute.copy( atom_ld32x32_x64, wrap_into_copy_standard_layout(tmem_subtile_per_warp), @@ -3011,20 +3007,14 @@ def fc2_stg_post_f2fp_reorder( def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) - core_matrix_reorder_sttm_atom = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32 - ) - core_matrix_reorder_ldtm_atom = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32 - ) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) # ((16, 2), token_32_group) token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) out_as_i32 = cute.recast_tensor(out, cutlass.Float32) # (32, 64) - tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[ - None, None, 0, 0 - ] + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): @@ -3064,9 +3054,7 @@ def fc2_stg_store_function( sf_regs = None stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width - copy_atom_vec = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits - ) + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): raise ValueError( @@ -3078,9 +3066,7 @@ def fc2_stg_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3097,11 +3083,7 @@ def fc2_stg_store_function( if pred != cutlass.Int32(0): src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) - cute.copy( - copy_atom_vec, - cute.recast_tensor(src_i, cutlass.Int32), - cute.recast_tensor(dst_i, cutlass.Int32), - ) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) @cute.jit @@ -3115,12 +3097,8 @@ def fc2_tma_store_function( fc2_output_router: Fc2OutputRouter, **_, ): - if cutlass.const_expr( - epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None - ): - raise ValueError( - "FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor." - ) + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") if cutlass.const_expr(epi.combine_format.is_quantized): data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) @@ -3132,9 +3110,7 @@ def fc2_tma_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3148,15 +3124,11 @@ def fc2_tma_store_function( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 TMA store requires a contiguous two-token by 32-hidden register fragment." - ) + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) smem_stage = epi.smem_tensor[None, None, stage_idx] - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3170,9 +3142,7 @@ def fc2_tma_store_function( + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) ) // cutlass.Int32(epi._EpilogueTokenTileSize) - tiled_output = cute.flat_divide( - epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize) - ) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( epi.tma_atom_fc2_output, @@ -3212,9 +3182,7 @@ def fc2_ublk_store_function_impl( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -3241,12 +3209,8 @@ def fc2_ublk_store_function_impl( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment." - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3256,9 +3220,7 @@ def fc2_ublk_store_function_impl( cute.copy(tiled_store, register_partition, smem_partition) else: if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): - raise ValueError( - "BF16 UBLK staging requires one register per token in the 64-token subtile." - ) + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") warp_hidden_base = cutlass.Int32(warp_idx * 32) for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] @@ -3312,18 +3274,14 @@ def fc2_redg_store_function( if pred != cutlass.Int32(0): bf16x4 = subtile_by_redg_issue[None, local_iter] packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) - red_add_relaxed_sys_v2_bf16x2( - dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1]) - ) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) def make_fc2_stg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3340,10 +3298,8 @@ def make_fc2_stg_process_pipeline( def make_fc2_tma_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3360,10 +3316,8 @@ def make_fc2_tma_process_pipeline( def make_fc2_redg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_redg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py index b9318a28b60f..bcead8bcb700 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Block-scaled SwapAb adapter between FC12 scheduling and kernel tensor views.""" import dataclasses @@ -16,7 +13,8 @@ from .....helpers.dsl_helpers import spin_peek, spin_wait from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit -TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] + +TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk", "fc1_act_scale"] @cute.jit @@ -54,18 +52,14 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAb def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) fc2_spin_threshold = rebuild(self.fc2_spin_threshold) fc1_ready_counter_pointer = ( - rebuild(self.fc1_ready_counter_pointer) - if self.fc1_ready_counter_pointer is not None - else None + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None ) if value_index != len(values): raise ValueError( @@ -91,9 +85,7 @@ def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12Work if is_fc1: counter_pointer = self.fc1_ready_counter_pointer + counter_slot peek_flag = Int32(0) - if spin_peek( - counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile - ): + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): peek_flag = Int32(peek_ready_bit) phase_and_flags = work_tile.phase_and_flags | peek_flag @@ -151,25 +143,20 @@ def get_gmem_tensor( result = cute.domain_offset((0, 0, expert_idx), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name in ("sfb", "sfc")): result = cute.domain_offset((sf_token_offset, 0, 0), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name == "c"): result = cute.domain_offset((data_token_offset, 0, 0), tensor) return (_rewrite_tensor_shape(result, (shape[0], shape[1], singleton)), None) - if cutlass.const_expr(tensor_name == "topk"): + if cutlass.const_expr(tensor_name in ("topk", "fc1_act_scale")): + # Both follow the data-token pool (not the independently padded SF pool). return (cute.domain_offset((data_token_offset,), tensor), None) raise ValueError(f"Unknown tensor_name: {tensor_name!r}.") diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py new file mode 100644 index 000000000000..c35a024f48a7 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_kernel.py @@ -0,0 +1,712 @@ +"""Lean block-scaled swap-AB FC12 kernel composition for Blackwell.""" + +from typing import ClassVar, Optional, Tuple + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cutlass_dsl import Boolean +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype +from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.utils import ceil_div, round_up +from .....quant_def import CombineFormat, QuantKind +from ....schedulers.base import WorkIdAcquisitionMode +from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler +from ...custom_mix_cga_helpers import TmaAtomOrPair +from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop + + +class BlockScaledSwapAbFc12Kernel(KernelClass): + """Lean FC12-only top-level with one explicit eight-warp topology.""" + + fc1_output_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_output" + fc1_output_sf_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_output_sf" + fc1_done_counter_region: ClassVar[str] = "blackwell.swap_ab_fc12.fc1_done_counter" + epilogue_warp_ids: ClassVar[Tuple[int, int, int, int]] = (0, 1, 2, 3) + mma_warp_id: ClassVar[int] = 4 + tma_a_warp_id: ClassVar[int] = 5 + tma_b_warp_id: ClassVar[int] = 6 + scheduler_warp_id: ClassVar[int] = 7 + threads_per_cta: ClassVar[int] = 8 * 32 + + @classmethod + def problem_desc_require(cls) -> dict[str, object]: + return { + "expert_count": StaticOrRuntimeIntegerType, + "intermediate_gateup_size": StaticOrRuntimeIntegerType, + "hidden_size": StaticOrRuntimeIntegerType, + "quant_kind": str, + "a_major_mode": OperandMajorMode, + "b_major_mode": OperandMajorMode, + "combine_format": CombineFormat, + "gate_up_clamp": Optional[float], + "situ_beta": OptionalRequirement(Optional[float]), + "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu"; validated by the epilogue. + "gate_activation": OptionalRequirement(str), + # Whether __call__ receives fc1_activation_per_token_scale; reserves the epilogue's + # SMEM scratch for it. Absent means False. + "fc1_activation_per_token_scale": OptionalRequirement(bool), + } + + @classmethod + def impl_desc_require(cls) -> dict[str, type]: + return { + "mma_tiler_mnk": tuple, + "cluster_shape_mn": tuple, + "use_2cta_instrs": bool, + "group_hint": int, + "token_padding_block": int, + "sf_padding_block": int, + "work_id_mode": str, + "max_tokens": int, + "launch_cluster_count": int, + "fc2_use_bulk": bool, + "fc1_epi_flag_batch": int, + "fc2_epi_flag_batch": int, + } + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + self._validate_desc_inputs(problem_desc, impl_desc) + + self.expert_count = problem_desc["expert_count"] + self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] + self.hidden_size = problem_desc["hidden_size"] + self.quant_kind = QuantKind(problem_desc["quant_kind"]) + self.a_dtype = self.quant_kind.weight_dtype + self.b_dtype = self.quant_kind.activation_dtype + self.sf_dtype = self.quant_kind.sf_dtype + self.sf_vec_size = self.quant_kind.sf_vec_size + self.acc_dtype = tcgen05_block_scaled_acc_dtype + self.a_major_mode = problem_desc["a_major_mode"] + self.b_major_mode = problem_desc["b_major_mode"] + self.combine_format = problem_desc["combine_format"] + self.gate_up_clamp = problem_desc["gate_up_clamp"] + self.situ_beta = problem_desc.get("situ_beta") + self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.gate_activation = problem_desc.get("gate_activation") or "swiglu" + + self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] + self.cluster_shape_mn = impl_desc["cluster_shape_mn"] + self.use_2cta_instrs = impl_desc["use_2cta_instrs"] + self.group_hint = impl_desc["group_hint"] + self.token_padding_block = impl_desc["token_padding_block"] + self.sf_padding_block = impl_desc["sf_padding_block"] + self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] + self.max_tokens = impl_desc["max_tokens"] + self.launch_cluster_count = impl_desc["launch_cluster_count"] + self.fc2_use_bulk = impl_desc["fc2_use_bulk"] + self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] + self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + + self.occupancy = 1 + self.architecture = "sm_100" + + self._validate_geometry() + mma_cta_count = 2 if self.use_2cta_instrs else 1 + mma_instruction = Tcgen05MmaInstruction( + a_type=self.a_dtype, + b_type=self.b_dtype, + acc_type=self.acc_dtype, + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), + participates=mma_cta_count, + sfa_type=self.sf_dtype, + sfb_type=self.sf_dtype, + sf_vec_size=self.sf_vec_size, + ) + tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) + resolved_impl_desc = ImplDesc( + { + **impl_desc, + "hint": self.group_hint, + "tmem_plan": tmem_plan, + "is_swap_ab": True, + "num_scheduler_consumer_threads": self.threads_per_cta - 32, + "num_accumulator_consumer_warps_per_cta": len(self.epilogue_warp_ids), + "communication_enabled": False, + } + ) + self.scheduler = BlackwellFusedFc12Scheduler(problem_desc, resolved_impl_desc) + self.epilogue = SwapABGatedActEpilogue(problem_desc, resolved_impl_desc) + if self.epilogue.fc1_output_dtype is not self.b_dtype: + raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") + if self.epilogue.fc1_output_sf_dtype is not self.sf_dtype: + raise ValueError("Epilogue FC1 output scale dtype must match the Mainloop scale dtype.") + if self.epilogue.sf_vec_size != self.sf_vec_size: + raise ValueError("Epilogue and Mainloop scale vector sizes must match.") + self._device_workspace = self._build_device_workspace() + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(problem_desc, resolved_impl_desc) + + def _validate_geometry(self) -> None: + static_expert_dimensions = ( + isinstance(self.expert_count, int), + isinstance(self.intermediate_gateup_size, int), + isinstance(self.hidden_size, int), + ) + if any(static_expert_dimensions) and not all(static_expert_dimensions): + raise ValueError("FC12 expert dimensions must be either all static or all runtime.") + if not all(static_expert_dimensions): + raise NotImplementedError("The Lean FC12 Kernel currently requires static expert dimensions.") + if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: + raise ValueError("FC12 expert dimensions must be positive.") + if self.launch_cluster_count <= 0: + raise ValueError("launch_cluster_count must be positive.") + if self.intermediate_gateup_size % 2 != 0: + raise ValueError("The SwiGLU intermediate dimension must be even.") + if self.max_tokens <= 0: + raise ValueError("max_tokens must be positive.") + if self.token_padding_block <= 0: + raise ValueError("token_padding_block must be positive.") + if self.token_padding_block % 64 != 0: + raise ValueError("token_padding_block must be a multiple of 64.") + if self.sf_vec_size <= 0: + raise ValueError("sf_vec_size must be positive.") + if self.quant_kind.needs_unpack_tma(self.architecture): + for name, extent in ( + ("hidden_size", self.hidden_size), + ("intermediate_downproj", self.intermediate_gateup_size // 2), + ): + if extent % 128 != 0: + raise ValueError( + f"{self.quant_kind} loads its fp4 weight through the unpacking TMA, which " + f"requires {name} to be a multiple of 128, got {extent}." + ) + mma_m, mma_n, mma_k = self.mma_tiler_mnk + cluster_m, cluster_n = self.cluster_shape_mn + expected_mma_m = 256 if self.use_2cta_instrs else 128 + if mma_m != expected_mma_m: + raise ValueError(f"mma_tiler M must be {expected_mma_m}, got {mma_m}.") + if mma_n not in (64, 128, 256): + raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") + if mma_k % (self.sf_vec_size * 4) != 0: + raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") + if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") + if cluster_n != 1: + raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") + if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): + raise ValueError("cluster M must be a power of two no greater than 16.") + if self.use_2cta_instrs and cluster_m % 2 != 0: + raise ValueError("Two-CTA MMA requires an even cluster M.") + + def _build_device_workspace(self) -> DeviceWorkspace: + intermediate_downproj = self.intermediate_gateup_size // 2 + sf_column_count = round_up(ceil_div(intermediate_downproj, self.sf_vec_size), 4) + max_sf_rows = self.max_tokens + self.expert_count * self.sf_padding_block + counter_slot_count = ceil_div(self.max_tokens, self.mma_tiler_mnk[1]) + self.expert_count + + device_workspace = DeviceWorkspace() + device_workspace.register( + self.fc1_output_region, + self.epilogue.fc1_output_dtype, + (self.max_tokens, intermediate_downproj), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_output_sf_region, + self.epilogue.fc1_output_sf_dtype, + (max_sf_rows, sf_column_count), + buffer_space="local", + mem_order=(1, 0), + byte_alignment=128, + ) + device_workspace.register( + self.fc1_done_counter_region, + cutlass.Int32, + (counter_slot_count,), + buffer_space="local", + byte_alignment=16, + reset="tail_reset", + ) + self.scheduler.register_device_workspace(device_workspace) + device_workspace.finalize() + return device_workspace + + def get_workspace_sizes(self) -> Tuple[int, int]: + """Return required local and shared workspace bytes.""" + return self._device_workspace.local_and_shared_bytes + + @property + def require_zero_workspace_leading_bytes(self) -> Tuple[int, int]: + return self._device_workspace.require_zero_workspace_leading_bytes + + def name(self) -> str: + return "" + + def aot_compile(self): + return None + + def _build_mainloop_and_smem( + self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc + ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: + smem_limit = utils.get_smem_capacity_in_bytes() // self.occupancy + smem_workspace = SmemWorkspace() + self.scheduler.register_smem_regions(smem_workspace) + self.epilogue.register_smem_regions(smem_workspace) + alignment_shift_bytes = max(region.byte_alignment for region in smem_workspace.regions()) + mainloop_smem_budget_bytes = smem_limit - smem_workspace.estimate_total_bytes() - alignment_shift_bytes + if mainloop_smem_budget_bytes <= 0: + raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") + + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) + mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) + + mainloop.register_smem_regions(smem_workspace) + smem_workspace.finalize(max_bytes=smem_limit) + return mainloop, smem_workspace + + @cute.jit + def __call__( + self, + activation: cute.Tensor, # (tokens, hidden) + fc1_weight: cute.Tensor, # (experts, hidden, intermediate_gateup) + activation_sf: cute.Tensor, # (padded_tokens, padded_hidden // sf_vec_size) + fc1_weight_sf: cute.Tensor, # (experts, packed_fc1_sf) + fc2_weight: cute.Tensor, # (experts, intermediate_downproj, hidden) + fc2_weight_sf: cute.Tensor, # (experts, packed_fc2_sf) + fc2_output: cute.Tensor, # (tokens, topk, hidden) + local_workspace: cute.Pointer, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + expert_token_sizes: Optional[cute.Tensor] = None, # (experts,) + expert_token_prefix_sum: Optional[cute.Tensor] = None, # (experts,) + topk_scores: Optional[cute.Tensor] = None, # (tokens,) + fc1_alpha: Optional[cute.Tensor] = None, + fc2_alpha: Optional[cute.Tensor] = None, + fc1_norm_const: Optional[cute.Tensor] = None, + # (tokens,) fp32, same grouped layout as topk_scores. The activation was quantized from + # x / scale; the FC1 epilogue multiplies it back before the gated activation. Requires + # problem_desc["fc1_activation_per_token_scale"] = True (the epilogue rejects it otherwise). + fc1_activation_per_token_scale: Optional[cute.Tensor] = None, + ) -> None: + """Prepare and launch the lean fused FC1+FC2 kernel. + + The caller must zero the workspace prefix reported by + require_zero_workspace_leading_bytes before its first launch. + """ + if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") + if cutlass.const_expr( + fc1_activation_per_token_scale is not None + and fc1_activation_per_token_scale.element_type is not cutlass.Float32 + ): + raise TypeError( + f"fc1_activation_per_token_scale must be Float32, got {fc1_activation_per_token_scale.element_type}." + ) + # The quant kind fixes every operand format, so the tensors are where a caller mismatch can + # still show up. See BlockScaledSwapAbMegaMoeKernel.__call__. + for operand_name, operand, expected_dtype in ( + ("activation", activation, self.b_dtype), + ("activation_sf", activation_sf, self.sf_dtype), + ("fc1_weight", fc1_weight, self.a_dtype), + ("fc1_weight_sf", fc1_weight_sf, self.sf_dtype), + ("fc2_weight", fc2_weight, self.a_dtype), + ("fc2_weight_sf", fc2_weight_sf, self.sf_dtype), + ): + if cutlass.const_expr(operand.element_type is not expected_dtype): + raise TypeError( + f"{self.quant_kind} requires {operand_name} to be {expected_dtype}, got {operand.element_type}." + ) + if cutlass.const_expr(not self.quant_kind.uses_global_scale): + for scalar_name, scalar in ( + ("fc1_alpha", fc1_alpha), + ("fc2_alpha", fc2_alpha), + ("fc1_norm_const", fc1_norm_const), + ): + if cutlass.const_expr(scalar is not None): + raise ValueError( + f"{self.quant_kind} folds its rescale into the e8m0 scale factors, " + f"so {scalar_name} must be None." + ) + + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: + return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) + + self._device_workspace.assign_device_members(local_workspace) + fc1_output = self._device_workspace.tensor(self.fc1_output_region) + fc1_output_sf = self._device_workspace.tensor(self.fc1_output_sf_region) + + if cutlass.const_expr(isinstance(self.expert_count, int)): + experts = self.expert_count + else: + experts = fc1_weight.shape[0] + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = fc1_weight.shape[1] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = fc1_weight.shape[2] + intermediate_downproj = intermediate_gateup // 2 + fc1_weight = rewrite_tensor_shape(fc1_weight, (experts, hidden, intermediate_gateup)) + fc2_weight = rewrite_tensor_shape(fc2_weight, (experts, intermediate_downproj, hidden)) + activation = rewrite_tensor_shape(activation, (activation.shape[0], hidden)) + if cutlass.const_expr(isinstance(activation.shape[0], int)): + if cutlass.const_expr(activation.shape[0] > self.max_tokens): + raise ValueError("activation rows exceed the configured max_tokens.") + fc2_output = rewrite_tensor_shape(fc2_output, (fc2_output.shape[0], fc2_output.shape[1], hidden)) + + singleton = cutlass.Int32(1) + token_rows = activation.shape[0] + + fc1_a = cute.make_tensor( + fc1_weight.iterator, + cute.make_layout( + (cutlass.Int32(intermediate_gateup), cutlass.Int32(hidden), cutlass.Int32(experts)), + stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), + ), + ) + fc1_b = cute.make_tensor( + activation.iterator, + cute.make_layout( + (token_rows, cutlass.Int32(hidden), singleton), stride=(activation.stride[0], activation.stride[1], 0) + ), + ) + fc1_output_gemm = cute.make_tensor( + fc1_output.iterator, + cute.make_layout( + (token_rows, cutlass.Int32(intermediate_downproj), singleton), + stride=(fc1_output.stride[0], fc1_output.stride[1], 0), + ), + ) + + sf_vec_size = self.sf_vec_size + padded_token_rows = activation_sf.shape[0] + padded_hidden = activation_sf.shape[1] * sf_vec_size + fc1_sfb = cute.make_tensor( + activation_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (padded_token_rows, cutlass.Int32(padded_hidden), singleton), sf_vec_size + ), + ) + padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) + expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size + if cutlass.const_expr( + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) + ): + if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): + raise ValueError("fc1_weight_sf has an incompatible column count.") + fc1_sfa = cute.make_tensor( + fc1_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), + sf_vec_size, + ), + ) + + experts_fc2, intermediate_fc2, hidden_fc2 = fc2_weight.shape + fc2_a = cute.make_tensor( + fc2_weight.iterator, + cute.make_layout( + (cutlass.Int32(hidden_fc2), cutlass.Int32(intermediate_fc2), cutlass.Int32(experts_fc2)), + stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), + ), + ) + fc2_b = fc1_output_gemm + padded_fc2_hidden = cute.round_up(hidden_fc2, 128) + padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) + expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size + if cutlass.const_expr( + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) + ): + if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): + raise ValueError("fc2_weight_sf has an incompatible column count.") + fc2_sfa = cute.make_tensor( + fc2_weight_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), + sf_vec_size, + ), + ) + fc2_sfb = cute.make_tensor( + fc1_output_sf.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size + ), + ) + + self._mainloop.materialize_codegen_members() + ( + fc1_tma_a_tensor, + fc1_tma_a_atom, + fc1_tma_sfa_tensor, + fc1_tma_sfa_atom, + fc2_tma_a_tensor, + fc2_tma_a_atom, + fc2_tma_sfa_tensor, + fc2_tma_sfa_atom, + fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair, + ) = self._mainloop.prepare_tma_load_params( + fc1_a=fc1_a, + fc1_b=fc1_b, + fc1_sfa=fc1_sfa, + fc1_sfb=fc1_sfb, + fc2_a=fc2_a, + fc2_b=fc2_b, + fc2_sfa=fc2_sfa, + fc2_sfb=fc2_sfb, + ) + + (fc1_output_tma_atom, fc1_output_tma_tensor, fc2_output_tma_atom, fc2_output_tma_tensor) = ( + self.epilogue.prepare_tma_store_params(fc1_output_gemm, fc2_output) + ) + + grid = self.scheduler.get_grid_shape(max_active_clusters=max_active_clusters) + self._kernel( + fc1_tma_a_tensor, + fc1_tma_a_atom, + fc1_tma_sfa_tensor, + fc1_tma_sfa_atom, + fc2_tma_a_tensor, + fc2_tma_a_atom, + fc2_tma_sfa_tensor, + fc2_tma_sfa_atom, + fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + expert_token_sizes, + expert_token_prefix_sum, + (experts, intermediate_gateup, hidden), + local_workspace, + fc1_alpha, + fc2_alpha, + fc1_norm_const, + topk_scores, + fc1_activation_per_token_scale, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + @cute.kernel + def _kernel( + self, + fc1_tma_a_tensor: cute.Tensor, + fc1_tma_a_atom: cute.CopyAtom, + fc1_tma_sfa_tensor: cute.Tensor, + fc1_tma_sfa_atom: cute.CopyAtom, + fc2_tma_a_tensor: cute.Tensor, + fc2_tma_a_atom: cute.CopyAtom, + fc2_tma_sfa_tensor: cute.Tensor, + fc2_tma_sfa_atom: cute.CopyAtom, + fc1_tma_b_tensor: cute.Tensor, + fc1_tma_b_atom_or_pair: TmaAtomOrPair, + fc1_tma_sfb_tensor: cute.Tensor, + fc1_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc2_tma_b_tensor: cute.Tensor, + fc2_tma_b_atom_or_pair: TmaAtomOrPair, + fc2_tma_sfb_tensor: cute.Tensor, + fc2_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc1_output_tma_atom: cute.CopyAtom, + fc1_output_tma_tensor: cute.Tensor, + fc2_output_tma_atom: Optional[cute.CopyAtom], + fc2_output_tma_tensor: Optional[cute.Tensor], + fc2_output: cute.Tensor, + expert_token_sizes: Optional[cute.Tensor], + expert_token_prefix_sum: Optional[cute.Tensor], + actual_expert_shape: Tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + local_workspace: cute.Pointer, + fc1_alpha: Optional[cute.Tensor], + fc2_alpha: Optional[cute.Tensor], + fc1_norm_const: Optional[cute.Tensor], + topk_scores: Optional[cute.Tensor], + fc1_activation_per_token_scale: Optional[cute.Tensor], + ): + """Compose Scheduler, Mainloop, and Epilogue.""" + self._mainloop.materialize_codegen_members() + storage_type = self._smem_workspace.storage_class() + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + self._device_workspace.assign_device_members(local_workspace) + fc1_done_counter = self._device_workspace.tensor(self.fc1_done_counter_region) + fc1_output_sf_storage = self._device_workspace.tensor(self.fc1_output_sf_region) + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = actual_expert_shape[2] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = actual_expert_shape[1] + fc1_output_sf = cute.make_tensor( + fc1_output_sf_storage.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size + ), + ) + + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(thread_idx // 32) + block_idx = cute.arch.block_idx() + + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_coord_in_cluster = ( + cta_rank_in_cluster % self.cluster_shape_mn[0], + cta_rank_in_cluster // self.cluster_shape_mn[0], + cutlass.Int32(0), + ) + + # Standalone FC12 binds the preferred layout before pipeline creation. + self._mainloop.assign_device_members( + self._smem_workspace, smem_base, cta_coord_in_cluster, Boolean(False), hidden, intermediate_gateup + ) + ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) + tma_a_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + tma_b_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + mma_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._mainloop.num_ab_pipeline_stages + ) + acc_pipeline = self._mainloop.create_acc_pipeline(self._smem_workspace, smem_base) + tmem_allocator = self._mainloop.create_tmem_allocator( + self._smem_workspace, smem_base, allocator_warp_id=self.epilogue_warp_ids[0] + ) + self.scheduler.assign_device_members( + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + actual_expert_shape=actual_expert_shape, + block_idx=block_idx, + smem_workspace=self._smem_workspace, + smem_base=smem_base, + device_workspace=self._device_workspace, + ) + scheduler = self.scheduler + + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + kernel_extension = BlockScaledSwapAbFc12Extension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=None, + ) + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=topk_scores, + fc1_activation_per_token_scale=fc1_activation_per_token_scale, + ) + + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + if warp_idx == self.scheduler_warp_id: + work_tile = scheduler.gen_next_work() + while work_tile.is_valid_tile: + scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + work_tile = scheduler.gen_next_work() + scheduler.publish_work(work_tile) + scheduler.produce_tail() + + if warp_idx == self.tma_a_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_a( + fc1_tma_a_tensor=fc1_tma_a_tensor, + fc1_tma_a_atom=fc1_tma_a_atom, + fc1_tma_sfa_tensor=fc1_tma_sfa_tensor, + fc1_tma_sfa_atom=fc1_tma_sfa_atom, + fc2_tma_a_tensor=fc2_tma_a_tensor, + fc2_tma_a_atom=fc2_tma_a_atom, + fc2_tma_sfa_tensor=fc2_tma_sfa_tensor, + fc2_tma_sfa_atom=fc2_tma_sfa_atom, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_a_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + ) + + if warp_idx == self.tma_b_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_b( + fc1_tma_b_tensor=fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair=fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor=fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair=fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor=fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair=fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor=fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair=fc2_tma_sfb_atom_or_pair, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_b_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + ) + + if warp_idx == self.mma_warp_id: + sched_consumer = scheduler.make_consumer() + self._mainloop.run_mma( + tmem_allocator=tmem_allocator, + ab_pipeline=ab_pipeline, + ab_pipeline_state=mma_pipeline_state, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + ) + + if warp_idx < len(self.epilogue_warp_ids): + sched_consumer = scheduler.make_consumer() + tmem_allocator.allocate(self._mainloop.num_tmem_alloc_cols) + tmem_allocator.wait_for_alloc() + tmem_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) + self.epilogue.run( + self._smem_workspace, + smem_base, + tmem_pointer, + acc_pipeline, + sched_consumer, + kernel_extension, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc1_output_sf, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + fc1_done_counter, + thread_idx, + None, + None, + None, + None, + optional_epilogue_args, + ) + tmem_allocator.relinquish_alloc_permit() + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py index c04f39fc3686..93d61ff8f69c 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Optional, Tuple @@ -14,13 +11,7 @@ from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....helpers.cute_py_helpers import ( Tcgen05MmaInstruction, make_smem_layouts, @@ -45,9 +36,7 @@ @cute.jit -def _make_executable_tma_atom_for_cluster( - atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean -) -> cute.CopyAtom: +def _make_executable_tma_atom_for_cluster(atom_or_pair: TmaAtomOrPair, is_fallback_cluster: Boolean) -> cute.CopyAtom: """CuTeDSL < 4.7 WAR: materialize exec atoms inside the selecting scf.if branches.""" if cutlass.const_expr(isinstance(atom_or_pair, tuple)): executable_atom = make_executable_tma_atom(atom_or_pair[0]) @@ -62,22 +51,14 @@ def _make_executable_tma_atom_for_cluster( class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) - tmem_holding_buffer_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.tmem_holding_buffer" - ) - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "blackwell.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" + tmem_holding_buffer_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.tmem_holding_buffer" + tmem_deallocation_mbarrier_region: ClassVar[str] = "blackwell.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 @classmethod @@ -120,9 +101,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self._validate_configuration() self.architecture = "sm_100" @@ -141,11 +120,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: tmem_plan.accumulator_stage_columns - tmem_plan.accumulator_stage_stride_columns ) self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.overlapping_accum = self.accumulator_overlap_columns > 0 @@ -153,11 +128,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: a_type=self.a_dtype, b_type=self.b_dtype, acc_type=self.acc_dtype, - instruction_mnk=( - self.mma_tiler_mnk[0], - self.mma_tiler_mnk[1], - self.quant_kind.instruction_k("1x"), - ), + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), participates=self.mma_cta_count, sfa_type=self.sf_dtype, sfb_type=self.sf_dtype, @@ -165,12 +136,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) # A mixed-width pair reaches SMEM through the unpacking TMA, so the narrow operand's SMEM # image is byte-per-element even though the MMA still sees its logical type. - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.mma_tiler_sfb = (self.cta_tile_m, round_up(self.cta_tile_n, 128), self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -180,9 +147,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.cluster_layout_sfb_shape_vmnk = ((1,), *self.cluster_shape_mn, 1) self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) self.fallback_cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -190,11 +155,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.resolved_fallback_cluster_shape_mn[1], 1, ) - self.fallback_cluster_layout_sfb_shape_vmnk = ( - (1,), - *self.resolved_fallback_cluster_shape_mn, - 1, - ) + self.fallback_cluster_layout_sfb_shape_vmnk = ((1,), *self.resolved_fallback_cluster_shape_mn, 1) self.is_mixed_cga = self.resolved_fallback_cluster_shape_mn != self.cluster_shape_mn self.num_mcast_ctas_a = self.cluster_layout_shape_vmnk[2] self.num_mcast_ctas_b = self.cluster_layout_shape_vmnk[1] @@ -214,9 +175,7 @@ def _validate_configuration(self) -> None: if any(dimension <= 0 for dimension in self.cluster_shape_mn): raise ValueError("cluster_shape_mn dimensions must be positive.") if self.mma_tiler_mnk[1] not in (64, 128, 256): - raise ValueError( - f"The component supports mma_tiler N in (64, 128, 256); got {self.mma_tiler_mnk[1]}." - ) + raise ValueError(f"The component supports mma_tiler N in (64, 128, 256); got {self.mma_tiler_mnk[1]}.") mma_cta_count = 2 if self.use_2cta_instrs else 1 if self.mma_tiler_mnk[0] % mma_cta_count != 0: raise ValueError("mma_tiler M must be divisible by the MMA CTA count.") @@ -240,9 +199,7 @@ def _validate_configuration(self) -> None: self.cluster_shape_mn, self.fallback_cluster_shape_mn ) ): - raise ValueError( - "Preferred cluster dimensions must be divisible by fallback dimensions." - ) + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") if self.mainloop_smem_budget_bytes <= 0: raise ValueError("mainloop_smem_budget_bytes must be positive.") if self.num_accumulator_consumer_warps_per_cta <= 0: @@ -257,9 +214,7 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) smem_workspace.register_tensor( self.a_smem_tensor_region, self.a_smem_region.dtype, @@ -292,21 +247,13 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=self.sfb_smem_region.swizzle, byte_alignment=self.sfb_smem_region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) a_region, b_region, sfa_region, sfb_region = stage_one_regions @@ -323,32 +270,23 @@ def packed_bytes(region, dtype) -> int: mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( "One AB pipeline stage needs " f"{ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def materialize_codegen_members(self) -> None: @@ -376,12 +314,10 @@ def materialize_codegen_members(self) -> None: cute.make_layout((*self.cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) self.fallback_cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma.thr_id.shape,) ) self.fallback_cluster_layout_sfb_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma_sfb.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) if self.cluster_layout_vmnk.shape != self.cluster_layout_shape_vmnk: raise ValueError( @@ -398,10 +334,7 @@ def materialize_codegen_members(self) -> None: f"Fallback cluster layout mismatch: Python {self.fallback_cluster_layout_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_vmnk.shape}." ) - if ( - self.fallback_cluster_layout_sfb_vmnk.shape - != self.fallback_cluster_layout_sfb_shape_vmnk - ): + if self.fallback_cluster_layout_sfb_vmnk.shape != self.fallback_cluster_layout_sfb_shape_vmnk: raise ValueError( f"Fallback SFB cluster layout mismatch: Python {self.fallback_cluster_layout_sfb_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_sfb_vmnk.shape}." @@ -429,11 +362,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -459,27 +389,15 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) # An operand whose SMEM container is wider than its element type must be loaded by the # unpacking TMA variant; naming the container as internal_type is what selects it. - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None def make_tma_maybe_pair( make_atom, @@ -648,17 +566,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 2) return PipelineTmaUmmaMixedCga.create( @@ -679,9 +590,7 @@ def create_ab_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -693,9 +602,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -706,9 +613,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -745,17 +650,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -765,16 +663,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -829,16 +721,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -927,15 +813,11 @@ def run_tma_b( fc2_tma_b_partition_atom = fc2_tma_b_atom_or_pair[0] fc2_tma_sfb_partition_atom = fc2_tma_sfb_atom_or_pair[0] - fc1_tma_b_exec_atom = _make_executable_tma_atom_for_cluster( - fc1_tma_b_atom_or_pair, self.is_fallback_cluster - ) + fc1_tma_b_exec_atom = _make_executable_tma_atom_for_cluster(fc1_tma_b_atom_or_pair, self.is_fallback_cluster) fc1_tma_sfb_exec_atom = _make_executable_tma_atom_for_cluster( fc1_tma_sfb_atom_or_pair, self.is_fallback_cluster ) - fc2_tma_b_exec_atom = _make_executable_tma_atom_for_cluster( - fc2_tma_b_atom_or_pair, self.is_fallback_cluster - ) + fc2_tma_b_exec_atom = _make_executable_tma_atom_for_cluster(fc2_tma_b_atom_or_pair, self.is_fallback_cluster) fc2_tma_sfb_exec_atom = _make_executable_tma_atom_for_cluster( fc2_tma_sfb_atom_or_pair, self.is_fallback_cluster ) @@ -960,20 +842,14 @@ def run_tma_b( "sfb", ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) # Rescale fallback CTA coordinates into the preferred tile partition. b_partition_coord = self.main_vmnk_coord[1] sfb_partition_coord = self.sfb_vmnk_coord[1] if cutlass.const_expr(self.is_mixed_cga): if self.is_fallback_cluster: - split_factor = ( - self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] - ) + split_factor = self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] b_partition_coord = b_partition_coord * Int32(split_factor) sfb_partition_coord = sfb_partition_coord * Int32(split_factor) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -988,24 +864,17 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc1_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc1_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1049,11 +918,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc1_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1061,9 +926,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) else: iket.range_push("tma_token_fc2") @@ -1079,24 +942,17 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1140,11 +996,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc2_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1152,9 +1004,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) iket.range_pop() @@ -1178,14 +1028,10 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( self.tiled_mma, self.mma_tiler_mnk, @@ -1195,8 +1041,7 @@ def run_mma( tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( self.tiled_mma, @@ -1237,10 +1082,7 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) @@ -1256,12 +1098,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa - ) - cute.copy( - tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) stage_coord = (None, None, None, ab_pipeline_state.index) dynamic_mainloop.issue_dynamic_block_scaled_mma_tile( quant_kind=self.quant_kind, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index e5482f1c4491..2e0497d3aa72 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Fused pull-dispatch MegaMoE kernel composition for Blackwell.""" +import hashlib + from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -14,23 +14,20 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.token_comm import TokenCommArgs, TokenCommNonDeterministic -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, +) from .....quant_def import CombineFormat, QuantKind from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import ( @@ -39,15 +36,11 @@ minimum_phase_interleave_hint, ) from ....schedulers.non_clc_mixed_cga import NonClcMixedCgaConfig -from ...custom_mix_cga_helpers import ( - TmaAtomOrPair, - pipeline_init_arrive_mixed_cga, - pipeline_init_wait_mixed_cga, -) +from ...custom_mix_cga_helpers import TmaAtomOrPair, pipeline_init_arrive_mixed_cga, pipeline_init_wait_mixed_cga from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop -from .topk_reduce import TopkReduce + _aot_symbol_prefix = "blackwell_mega_moe_aot" @@ -84,11 +77,17 @@ def problem_desc_require(cls) -> dict[str, object]: "swiglu_beta": OptionalRequirement(Optional[float]), "situ_beta": OptionalRequirement(Optional[float]), "situ_linear_beta": OptionalRequirement(Optional[float]), + # "swiglu" (default) or "quick_geglu"; validated by the epilogue, keyed in name(). + "gate_activation": OptionalRequirement(str), "world_size": int, "topk": int, "topk_index_dtype": type, "max_tokens_per_rank": int, "apply_topk_at_fc1": bool, + # The FC1 activation was quantized from x / s_t with a per-token fp32 s_t; the kernel + # multiplies s_t back before the gated activation. Optional so existing descriptors + # stay valid; it keys the compiled kernel (workspace region + epilogue path). + "fc1_activation_per_token_scale": OptionalRequirement(bool), } @classmethod @@ -102,6 +101,7 @@ def impl_desc_require(cls) -> dict[str, type]: "sf_padding_block": int, "work_id_mode": str, "fc2_use_bulk": bool, + "epi_async_release": OptionalRequirement(bool), "epi_flag_batches": tuple, "launch_cluster_count": int, "fallback_cluster_shape_mn": OptionalRequirement(Optional[tuple]), @@ -110,10 +110,18 @@ def impl_desc_require(cls) -> dict[str, type]: "token_in_flag_batch": int, "token_back_mode": str, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) + # The Blackwell epilogue publishes expert counters, not token-tile counters. + if impl_desc.get("token_back_ready_granularity", "expert") != "expert": + raise ValueError("Blackwell requires token_back_ready_granularity='expert'.") self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -132,11 +140,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_beta = problem_desc.get("swiglu_beta") self.situ_beta = problem_desc.get("situ_beta") self.situ_linear_beta = problem_desc.get("situ_linear_beta") + self.gate_activation = problem_desc.get("gate_activation") or "swiglu" self.world_size = problem_desc["world_size"] self.topk = problem_desc["topk"] self.topk_index_dtype = problem_desc["topk_index_dtype"] self.max_tokens_per_rank = problem_desc["max_tokens_per_rank"] self.apply_topk_at_fc1 = problem_desc["apply_topk_at_fc1"] + self.fc1_activation_per_token_scale = bool(problem_desc.get("fc1_activation_per_token_scale", False)) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] @@ -154,10 +164,23 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode = impl_desc["token_back_mode"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] - + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") self.occupancy = 1 self.architecture = "sm_100" self.local_expert_count = self.expert_count // self.world_size + if self.helper_expert_count > 0: + helper_topology = derive_helper_weight_ready_topology( + self.local_expert_count, + self.helper_expert_count, + ) + self.home_expert_count = helper_topology.home_expert_count + else: + # Preserve the baseline's runtime-expert-count specialization. A + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 @@ -172,9 +195,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.preferred_cluster_count = self.mixed_cga_config.preferred_cluster_count self.fallback_cluster_count = self.mixed_cga_config.fallback_cluster_count self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) if self.launch_cluster_count != self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred: raise ValueError( @@ -189,11 +210,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: a_type=self.a_dtype, b_type=self.b_dtype, acc_type=self.acc_dtype, - instruction_mnk=( - self.mma_tiler_mnk[0], - self.mma_tiler_mnk[1], - self.quant_kind.instruction_k("1x"), - ), + instruction_mnk=(self.mma_tiler_mnk[0], self.mma_tiler_mnk[1], self.quant_kind.instruction_k("1x")), participates=mma_cta_count, sfa_type=self.sf_dtype, sfb_type=self.sf_dtype, @@ -201,9 +218,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] - hidden_per_fc2_cluster_tile = ( - self.mma_tiler_mnk[0] // mma_cta_count * self.cluster_shape_mn[0] - ) + hidden_per_fc2_cluster_tile = self.mma_tiler_mnk[0] // mma_cta_count * self.cluster_shape_mn[0] fc2_done_signals_per_token_tile = ( ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) * self.cluster_shape_mn[0] @@ -218,9 +233,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, "promised_launchable_sm_count": self.mixed_cga_config.total_cta_cnt, - "token_back_schedule_mode": ( - "atomic_counter" if self.work_id_mode == "atomic_counter" else "static" - ), + "token_back_schedule_mode": ("atomic_counter" if self.work_id_mode == "atomic_counter" else "static"), } ) self.token_comm = TokenCommNonDeterministic(problem_desc, token_comm_impl_desc) @@ -242,13 +255,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "communication_enabled": True, "token_back_push_data": self.token_comm.token_back_push_data, "fc1_epi_flag_batch": self.epi_flag_batches[0], + "epi_async_release": impl_desc.get("epi_async_release", True), "fc2_epi_flag_batch": self.epi_flag_batches[1], } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -259,13 +271,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) def _validate_geometry(self) -> None: @@ -277,9 +285,7 @@ def _validate_geometry(self) -> None: if any(static_expert_dimensions) and not all(static_expert_dimensions): raise ValueError("MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The MegaMoE Kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The MegaMoE Kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -287,9 +293,7 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") if self.intermediate_gateup_size % 2 != 0: @@ -306,12 +310,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: @@ -345,9 +345,7 @@ def _validate_geometry(self) -> None: if mma_k % (self.sf_vec_size * 4) != 0: raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): @@ -357,14 +355,8 @@ def _validate_geometry(self) -> None: if self.fallback_cluster_shape_mn is not None: fallback_cluster_m, fallback_cluster_n = self.fallback_cluster_shape_mn if fallback_cluster_n != 1: - raise ValueError( - f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}." - ) - if ( - fallback_cluster_m <= 0 - or fallback_cluster_m > 16 - or fallback_cluster_m & (fallback_cluster_m - 1) - ): + raise ValueError(f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}.") + if fallback_cluster_m <= 0 or fallback_cluster_m > 16 or fallback_cluster_m & (fallback_cluster_m - 1): raise ValueError("fallback cluster M must be a power of two no greater than 16.") if self.use_2cta_instrs and fallback_cluster_m % 2 != 0: raise ValueError("Two-CTA MMA requires an even fallback cluster M.") @@ -462,6 +454,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"tile{tile}_cluster{cluster}_{'2cta' if self.use_2cta_instrs else '1cta'}_" @@ -470,7 +463,7 @@ def dtype_name(dtype: type) -> str: f"sched{self.schedule_mode}_hint{self.hint}_pad{self.token_padding_block}x{self.sf_padding_block}_" f"work{self.work_id_mode}_" f"fc2store{'tma' if self.epilogue.fc2_use_tma else 'ublk' if self.epilogue.fc2_use_ublk else 'stg'}_" - f"fc2tmastages{self.epilogue.fc2_tma_stages}_" + f"fc2tmastages{self.epilogue.fc2_tma_stages}_asyncrel{int(self.epilogue.epi_async_release)}_" f"epiflag{epi_flags}_clusters{self.launch_cluster_count}_ctas{self.mixed_cga_config.total_cta_cnt}_" f"tokeninflag{self.token_in_flag_batch}_tokenback{self.token_back_mode}_" f"combine{self.combine_format.name}_clamp{self.gate_up_clamp}_" @@ -479,6 +472,10 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + "_dlb_v1" + # Off / SwiGLU keep the historical name so existing compiled-kernel caches stay valid. + + ("_fc1actscale" if self.fc1_activation_per_token_scale else "") + + (f"_act{self.gate_activation}" if self.gate_activation != "swiglu" else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -500,14 +497,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -516,20 +509,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so the token wire dtype is TokenComm's, not a_dtype. - activation=fake_tensor( - self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.token_comm.activation_sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), @@ -539,18 +526,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -565,6 +544,16 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, + 0, + AddressSpace.gmem, + assumed_align=8, + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: # nvfp4 carries per-expert dequant scalars. Under an e8m0 scale they do not exist at # all, and declaring them here would put three tensors in the ABI that the caller has @@ -574,20 +563,31 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fc2_alpha=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), fc1_norm_const=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), ) + if self.fc1_activation_per_token_scale: + # Local (tokens,) fp32 read by the router on the source rank; only in the ABI when enabled. + fake_arguments.update( + fc1_activation_per_token_scale=fake_tensor(cutlass.Float32, (tokens,), (0,), {0}, 4), + ) compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: return compiled compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True + out_path, + function_name=self._aot_symbol_name(), + export_only_tvm_ffi_symbols=True, ) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc @@ -603,9 +603,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -631,16 +629,62 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, + # (tokens,) fp32, this rank's own tokens, plain local memory (the router reads it here and + # pushes it to the pool). Required iff the descriptor enables fc1_activation_per_token_scale. + # Last on purpose: the AOT kwargs wrapper right-aligns its None defaults, so appending keeps + # every existing positional caller valid. + fc1_activation_per_token_scale: Optional[cute.Tensor] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) - if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): + if cutlass.const_expr(self.fc1_activation_per_token_scale != (fc1_activation_per_token_scale is not None)): + raise ValueError( + "fc1_activation_per_token_scale must be passed iff the descriptor enables it " + f"(enabled={self.fc1_activation_per_token_scale})." + ) + if cutlass.const_expr( + fc1_activation_per_token_scale is not None + and fc1_activation_per_token_scale.element_type is not cutlass.Float32 + ): raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." + f"fc1_activation_per_token_scale must be Float32, got {fc1_activation_per_token_scale.element_type}." + ) + if cutlass.const_expr(fc1_activation_per_token_scale is not None): + # The router gathers scale[token] for every routed token, so the scale needs one entry + # per topk_indices row. Only checkable when both extents are static (AOT leaves them + # dynamic); a short tensor would otherwise read past its end silently. + scale_rows = fc1_activation_per_token_scale.shape[0] + token_rows = topk_indices.shape[0] + if cutlass.const_expr(isinstance(scale_rows, int) and isinstance(token_rows, int) and scale_rows < token_rows): + raise ValueError( + f"fc1_activation_per_token_scale has {scale_rows} entries but topk_indices has {token_rows} rows." + ) + + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + ( + self.helper_expert_count > 0 + and not all(value is not None for value in helper_ready_bundle) + ) + or ( + self.helper_expert_count == 0 + and any(value is not None for value in helper_ready_bundle) ) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" + ) + + if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -676,6 +720,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: peer_rank_ptr_mapper_host, self._device_workspace, stream, + fc1_activation_per_token_scale, ) peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() @@ -729,19 +774,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -759,31 +799,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -791,28 +821,19 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) pre_reduced_activation_sf = None else: - pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor( - self._device_workspace - ) - pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor( - self._device_workspace - ) + pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor(self._device_workspace) + pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor(self._device_workspace) if cutlass.const_expr(self.token_comm.token_back_push_data): fc2_output = self.token_comm.fc2_activation_tensor(self._device_workspace) else: fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) fc2_output_sf = self.token_comm.fc2_activation_sf_tensor(self._device_workspace) @@ -886,6 +907,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ) if cutlass.const_expr(self.mixed_cga_config.is_mixed): kernel.launch( @@ -909,11 +932,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores self._topk_reduce( - pre_reduced_activation, - pre_reduced_activation_sf, - output_activation, - reduce_scores, - stream, + pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream ) @cute.kernel @@ -953,6 +972,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose TokenComm, Scheduler, Mainloop, and Epilogue.""" self._mainloop.materialize_codegen_members() @@ -974,8 +995,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -986,9 +1006,7 @@ def _kernel( is_fallback_cluster = Boolean(False) if cutlass.const_expr(self.mixed_cga_config.is_mixed): active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() - is_fallback_cluster = active_cluster_m == Int32( - self.resolved_fallback_cluster_shape_mn[0] - ) + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) cta_rank_in_cluster = cute.arch.block_idx_in_cluster() cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) @@ -1009,15 +1027,11 @@ def _kernel( token_src_metadata = self.token_comm.token_src_metadata_tensor(self._device_workspace) fc2_done_counter = self.token_comm.fc2_done_counter_tensor(self._device_workspace) pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) + pool_activation_scales = self.token_comm.fc1_activation_per_token_scales_tensor(self._device_workspace) # Bind the active cluster before deriving the AB pipeline arrival count. self._mainloop.assign_device_members( - self._smem_workspace, - smem_base, - cta_coord_in_cluster, - is_fallback_cluster, - hidden, - intermediate_gateup, + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup ) ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) tma_a_pipeline_state = pipeline.make_pipeline_state( @@ -1045,39 +1059,38 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer( - self._device_workspace - ), + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores, + fc1_activation_per_token_scale=pool_activation_scales, ) pipeline_init_arrive_mixed_cga( - self.cluster_shape_mn, - self.resolved_fallback_cluster_shape_mn, - is_fallback_cluster, - is_relaxed=True, + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True ) pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) - finalize_barrier_thread_count = ( - 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - ) - finalize_barrier = pipeline.NamedBarrier( - barrier_id=13, num_threads=finalize_barrier_thread_count - ) + finalize_barrier_thread_count = 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) @@ -1091,9 +1104,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1188,9 +1207,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() finalize_barrier.arrive_unaligned() @@ -1201,8 +1218,7 @@ def _kernel( self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() if cutlass.const_expr( - self.token_comm.token_back_enabled - and self.token_back_mode != "standalone_warps" # nosec B105 + self.token_comm.token_back_enabled and self.token_back_mode != "standalone_warps" # nosec B105 ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py index b7fedaa598a1..6216b50fb00d 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/blackwell/inference/mega/dynamic_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - # Dynamic UMMA N for fused fc1+fc2 swap-AB MegaMoE kernel. # # Emits the block-scaled UMMA instruction as literal PTX text (via @@ -34,13 +31,16 @@ from typing import Optional +import cutlass import cutlass.cute as cute +from cutlass.cutlass_dsl import dsl_user_op, Int32, Boolean from cutlass._mlir import ir -from cutlass._mlir.dialects import builtin, llvm -from cutlass.cutlass_dsl import Boolean, Int32, dsl_user_op +from cutlass._mlir.dialects import builtin +from cutlass._mlir.dialects import llvm from .....quant_def import QuantKind + # ============================================================================= # Alignment policy (single source of truth) # ============================================================================= @@ -223,20 +223,10 @@ def _tcgen05_mma_block_scaled( """ assert cta_group in (1, 2), f"cta_group must be 1 or 2, got {cta_group}" assert umma_kind in ("mxf4nvf4", "mxf4", "mxf8f6f4"), f"Unsupported UMMA kind {umma_kind!r}" - assert scale_vec_suffix in ("", ".block16", ".block32"), ( - f"Unsupported suffix {scale_vec_suffix!r}" - ) + assert scale_vec_suffix in ("", ".block16", ".block32"), f"Unsupported suffix {scale_vec_suffix!r}" llvm.inline_asm( None, - [ - d_tmem_i32, - a_desc_i64, - b_desc_i64, - idesc_i32, - enable_input_d_i32, - sfa_tmem_i32, - sfb_tmem_i32, - ], + [d_tmem_i32, a_desc_i64, b_desc_i64, idesc_i32, enable_input_d_i32, sfa_tmem_i32, sfb_tmem_i32], "{\n\t" ".reg .pred p;\n\t" "setp.ne.b32 p, $4, 0;\n\t" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py index c222690227be..83cacc303661 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/function_mapping.py @@ -23,13 +23,17 @@ def __post_init__(self) -> None: names = tuple(self.names) sizes = tuple(self.sizes) if not names or len(names) != len(sizes): - raise FunctionMappingError(f"Coordinate-space rank mismatch: {names!r}, {sizes!r}.") + raise FunctionMappingError( + f"Coordinate-space rank mismatch: {names!r}, {sizes!r}." + ) if any(not isinstance(name, str) or not name for name in names): raise FunctionMappingError("Axis names must be non-empty strings.") if len(set(names)) != len(names): raise FunctionMappingError(f"Axis names must be unique: {names!r}.") if any(not isinstance(size, int) or size <= 0 for size in sizes): - raise FunctionMappingError(f"Axis sizes must be positive Python ints: {sizes!r}.") + raise FunctionMappingError( + f"Axis sizes must be positive Python ints: {sizes!r}." + ) object.__setattr__(self, "names", names) object.__setattr__(self, "sizes", sizes) @@ -45,7 +49,9 @@ def axis_size(self, name: str) -> int: def delinearize(self, linear_index: int) -> tuple[int, ...]: if linear_index < 0 or linear_index >= self.size: - raise FunctionMappingError(f"Linear index {linear_index} is outside [0, {self.size}).") + raise FunctionMappingError( + f"Linear index {linear_index} is outside [0, {self.size})." + ) remaining = linear_index coordinate = [] for size in self.sizes: @@ -92,7 +98,9 @@ def _validate_signature(self) -> None: or parameter.default is not inspect.Parameter.empty ] if unsupported: - raise FunctionMappingError(f"Unsupported mapping parameters {unsupported!r}.") + raise FunctionMappingError( + f"Unsupported mapping parameters {unsupported!r}." + ) if set(parameters) != set(self.domain.names): raise FunctionMappingError( f"Mapping parameters {tuple(parameters)!r} must match " @@ -111,14 +119,17 @@ def _normalize_result( f"Mapping result keys {tuple(result)!r} must match " f"codomain axes {self.codomain.names!r}." ) - coordinate = {name: result[name] for name in self.codomain.names} + coordinate = { + name: result[name] for name in self.codomain.names + } elif isinstance(result, Sequence) and not isinstance( result, (str, bytes), ): if len(result) != len(self.codomain.names): raise FunctionMappingError( - f"Mapping result rank {len(result)} must equal {len(self.codomain.names)}." + f"Mapping result rank {len(result)} must equal " + f"{len(self.codomain.names)}." ) coordinate = dict(zip(self.codomain.names, result)) elif len(self.codomain.names) == 1: diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py index 52a7a9daf028..b84f4d01b814 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/__init__.py @@ -1,2 +1 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin kernel implementations.""" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py new file mode 100644 index 000000000000..dececcded527 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/custom_mix_cga_helpers.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Rubin source-copy shim for shared mixed-CGA helpers.""" + +# <<>> +from ..blackwell.custom_mix_cga_helpers import ( + PipelineTmaUmmaMixedCga, + TmaAtomOrPair, + bind_executable_tma_load_fields, + make_executable_tma_atom_or_pair, + pipeline_init_arrive_mixed_cga, + pipeline_init_wait_mixed_cga, + select_tma_atom_for_cluster, +) + + +__all__ = [ + "PipelineTmaUmmaMixedCga", + "TmaAtomOrPair", + "bind_executable_tma_load_fields", + "make_executable_tma_atom_or_pair", + "pipeline_init_arrive_mixed_cga", + "pipeline_init_wait_mixed_cga", + "select_tma_atom_for_cluster", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py index 52a7a9daf028..9f3a0e63c832 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/__init__.py @@ -1,2 +1,30 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin inference kernels.""" + +from .mega import ( + BlockScaledSwapAbFc12Extension, + BlockScaledSwapAbFc12Mainloop, + BlockScaledSwapAbFc12MainloopGenSpecialized, + BlockScaledSwapAbGenphaseMoeKernel, + BlockScaledSwapAbMegaMoeKernel, + GatedActEpilogueArgs, + GenphaseTokenComm, + GenphaseTokenCommArgs, + SwapABGatedActEpilogue, + TensorRole, + dynamic_mainloop, +) + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbFc12MainloopGenSpecialized", + "BlockScaledSwapAbGenphaseMoeKernel", + "BlockScaledSwapAbMegaMoeKernel", + "GenphaseTokenComm", + "GenphaseTokenCommArgs", + "GatedActEpilogueArgs", + "SwapABGatedActEpilogue", + "TensorRole", + "dynamic_mainloop", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py index 52a7a9daf028..97f3397bcf5a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/__init__.py @@ -1,2 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin Local MegaMoE inference kernel.""" + +from .block_scaled_swap_ab_local_mega_moe_kernel import BlockScaledSwapAbLocalMegaMoeKernel + + +__all__ = ["BlockScaledSwapAbLocalMegaMoeKernel"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py index 1315d5179679..911ced925fb0 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_epilogue.py @@ -19,20 +19,17 @@ from .....helpers.smem_workspace import SmemRegion, SmemWorkspace from .....quant_def import QuantKind from ....function_mapping import FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from ..mega.block_scaled_swap_ab_fc12_epilogue import ( GatedActEpilogueArgs, SwapABFc1Epilogue, + SwapABFc2Epilogue as _SharedSwapABFc2Epilogue, + SwapABGatedActEpilogue as _SharedSwapABGatedActEpilogue, make_fc2_redg_process_pipeline, make_fc2_stg_process_pipeline, make_fc2_ublk_process_pipeline, ) -from ..mega.block_scaled_swap_ab_fc12_epilogue import SwapABFc2Epilogue as _SharedSwapABFc2Epilogue -from ..mega.block_scaled_swap_ab_fc12_epilogue import ( - SwapABGatedActEpilogue as _SharedSwapABGatedActEpilogue, -) -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .local_routing import RoutingMetadata @@ -74,6 +71,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + # quick_geglu is Blackwell-only so far; refuse rather than silently computing SwiGLU. + if problem_desc.get("gate_activation") not in (None, "swiglu"): + raise NotImplementedError( + f"Rubin local epilogue does not support gate_activation={problem_desc.get('gate_activation')!r} yet." + ) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] @@ -92,12 +94,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fc2_use_ublk = self.fc2_use_bulk self.reduce_topk_in_epilogue = self.reduce_topk_in_kernel if not 1 <= self.fc1_epi_flag_batch <= self._EpilogueWarpCnt: - raise ValueError( - f"Asynchronous FC1 flag batch size must be in [1, {self._EpilogueWarpCnt}]." - ) - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + raise ValueError(f"Asynchronous FC1 flag batch size must be in [1, {self._EpilogueWarpCnt}].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -115,15 +113,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin Local MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.fc2_hidden_needs_predicate = not ( @@ -131,36 +123,24 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 ) self.intermediate_downproj = ( - self.intermediate_gateup_size // 2 - if isinstance(self.intermediate_gateup_size, int) - else None + self.intermediate_gateup_size // 2 if isinstance(self.intermediate_gateup_size, int) else None ) self.subtile_cnt = self.cta_tile_n // self._EpilogueTokenTileSize self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one FC1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if single_stage_region.nbytes % 16 != 0: @@ -170,27 +150,16 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: else: bf16_baseline_stages = min(2, self.subtile_cnt) bf16_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc2HiddenTileSize - * cutlass.BFloat16.width - // 8 - ) - available_staging_bytes = max( - self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes - ) - self.fc2_tma_stages = min( - self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes + self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + available_staging_bytes = max(self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None - def prepare_tma_store_params( - self, fc1_output_template: cute.Tensor - ) -> Tuple[cute.CopyAtom, cute.Tensor]: + def prepare_tma_store_params(self, fc1_output_template: cute.Tensor) -> Tuple[cute.CopyAtom, cute.Tensor]: fc1_operation = cpasync.CopyBulkTensorTileS2GOp() fc1_smem_layout = self.fc1_staged_smem_layout(1, without_stage_mode=True) fc1_tile = (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) @@ -295,9 +264,7 @@ def run( cute.arch.cp_async_bulk_wait_group(0) wait_only_named_barrier.arrive_and_wait() if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: flag_tracker = fc2_epi.advance_fc1_flag_phase(work_tile_info, flag_tracker) flag_tracker.fire() @@ -364,13 +331,10 @@ def advance_fc1_flag_phase(self, next_work_tile_info, flag_tracker): @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -425,14 +389,9 @@ def prefetch(self) -> "Fc2OutputRouter": if token_valid and hidden_valid: valid[iter_idx] = cutlass.Int32(1) metadata = RoutingMetadata.load( - self.routing_metadata.iterator.toint() - + Int64(token_in_tile) * Int64(RoutingMetadata.nbytes) - ) - dst_topk = ( - cutlass.Int32(0) - if cutlass.const_expr(self.reduce_topk_in_epilogue) - else metadata.topk_slot + self.routing_metadata.iterator.toint() + Int64(token_in_tile) * Int64(RoutingMetadata.nbytes) ) + dst_topk = cutlass.Int32(0) if cutlass.const_expr(self.reduce_topk_in_epilogue) else metadata.topk_slot dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile dst_ptrs[iter_idx] = cute.domain_offset( (Int64(metadata.token), dst_topk, dst_hidden), self.output @@ -441,9 +400,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: ptr = cute.make_ptr( self.output.element_type, self.dst_ptrs[iter_idx], diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py new file mode 100644 index 000000000000..0bc99ae344f1 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_extension.py @@ -0,0 +1,7 @@ +"""Local MegaMoE shim for the compatible Rubin MegaMoE FC12 extension.""" + +# <<>> +from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole + + +__all__ = ["BlockScaledSwapAbFc12Extension", "TensorRole"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py index 243336bbe19a..345070033ba4 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the Rubin block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Tuple @@ -29,7 +26,7 @@ from .....quant_def import QuantKind from ....schedulers.base import SchedulerConsumer from ....schedulers.fc12_mapping import BlockPhase -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .local_routing import RoutingMetadata from .tma_gather import sm107_tma_gather4_load @@ -37,20 +34,14 @@ class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 fc1_gather_sync_barrier_id: ClassVar[int] = 3 @@ -95,9 +86,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -118,15 +107,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin Local MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Local MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -139,12 +122,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -160,11 +139,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -187,17 +162,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ): raise ValueError("Local MegaMoE AB stage transaction byte accounting is inconsistent.") if self.num_tma_load_bytes % 2 != 0: - raise ValueError( - "Local MegaMoE total AB transaction bytes must split evenly across two producer arrivals." - ) + raise ValueError("Local MegaMoE total AB transaction bytes must split evenly across two producer arrivals.") self.ab_pipeline_tx_count_per_producer = self.num_tma_load_bytes // 2 self.num_gather_b_warps = 4 self.gather_b_rows_per_cta = self.cta_tile_n // self.mma_cta_count self.gather_b_rows_per_warp = self.gather_b_rows_per_cta // self.num_gather_b_warps if self.gather_b_rows_per_warp % 4 != 0: - raise ValueError( - "Local MegaMoE gather-B rows per warp must be divisible by gather4 width." - ) + raise ValueError("Local MegaMoE gather-B rows per warp must be divisible by gather4 width.") self.gather_b_calls_per_warp = self.gather_b_rows_per_warp // 4 self.gather_b_stage_bytes = self.cta_tile_n * self.mma_tile_k * int(self.b_dtype.width) // 8 if self.gather_b_stage_bytes != self.b_stage_tma_bytes: @@ -211,9 +182,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: fc1_b_side_stage_bytes = self.gather_b_stage_bytes + self.fc1_sfb_stage_bytes fc2_b_side_stage_bytes = self.b_stage_tma_bytes + self.sfb_stage_tma_bytes if fc1_b_side_stage_bytes != fc2_b_side_stage_bytes: - raise ValueError( - "Local MegaMoE FC1 and FC2 B-side transaction bytes must be identical." - ) + raise ValueError("Local MegaMoE FC1 and FC2 B-side transaction bytes must be identical.") if self.fc1_sfb_stage_bytes <= 0: raise ValueError("Local MegaMoE FC1 SFB transaction byte count must be positive.") self.fc1_gather_sync_barrier = pipeline.NamedBarrier( @@ -223,9 +192,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: def _validate_configuration(self) -> None: if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin Local MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin Local MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}.") for field_name, dimensions, expected_rank in ( ("mma_instruction_mnk", self.mma_instruction_mnk, 3), ("mma_tiler_mnk", self.mma_tiler_mnk, 3), @@ -233,10 +200,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -253,9 +217,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse.") if self.mma_tiler_mnk[1] not in (64, 128, 256): raise ValueError("Rubin Local MegaMoE supports tile N in (64, 128, 256).") if self.mma_tiler_mnk[2] % self.instruction_k != 0: @@ -276,9 +238,7 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected Rubin mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -293,37 +253,23 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -347,31 +293,22 @@ def packed_bytes(region, dtype) -> int: mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( f"One AB stage needs {ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -439,11 +376,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -469,24 +403,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None fc1_a_atom, fc1_a_tensor = cute.nvgpu.make_tiled_tma_atom_A( a_operation, @@ -502,9 +424,7 @@ def prepare_tma_load_params( # through TMA TILE_GATHER4 while landing in the native K_SW128 B layout. fc1_b_2d = fc1_b[(None, None, 0)] gather_b_base = cute.make_layout((1, self.mma_tile_k), stride=(self.mma_tile_k, 1)) - gather_b_smem_layout = cute.make_composed_layout( - cute.make_swizzle(3, 4, 3), 0, gather_b_base - ) + gather_b_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, gather_b_base) fc1_b_atom, fc1_b_tensor = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), fc1_b_2d, @@ -615,17 +535,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline. Each producer arrival contributes the same accounting share. The two @@ -662,9 +575,7 @@ def advance_pipeline_state( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -676,9 +587,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -689,9 +598,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -705,18 +612,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -745,17 +646,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -765,16 +659,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -829,16 +717,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -920,28 +802,14 @@ def run_tma_b( sfb_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) sfb_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "sfb", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "sfb" ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) thread_mma_sfb = self.tiled_mma_sfb.get_slice(self.mma_cta_index) @@ -954,9 +822,7 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) @@ -1013,16 +879,10 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1118,7 +978,9 @@ def run_tma_gather_b1( self.use_2cta_instrs, "b", ) - is_gather_b_mcast_leader = (not self.is_b_mcast) or (self.main_vmnk_coord[1] == Int32(0)) + is_gather_b_mcast_leader = (not self.is_b_mcast) or ( + self.main_vmnk_coord[1] == Int32(0) + ) work_tile = sched_consumer.consume_work() while work_tile.is_valid_tile: @@ -1127,8 +989,7 @@ def run_tma_gather_b1( iket.range_push("tma_gather_b1_fc1") kernel_extension.wait_for_input(work_tile) tile_pool_base = ( - work_tile.cumulative_data_physical_row - + work_tile.tile_n_idx * Int32(self.cta_tile_n) + work_tile.cumulative_data_physical_row + work_tile.tile_n_idx * Int32(self.cta_tile_n) ) cta_row_base = self.mma_cta_index * Int32(self.gather_b_rows_per_cta) warp_row_base = gather_warp_idx * Int32(rows_per_warp) @@ -1151,8 +1012,8 @@ def run_tma_gather_b1( stage_base_elements = ( stage_index * self.gather_b_rows_per_cta * self.mma_tile_k ) - warp_base_elements = gather_warp_idx * Int32( - rows_per_warp * self.mma_tile_k + warp_base_elements = ( + gather_warp_idx * Int32(rows_per_warp * self.mma_tile_k) ) for gather_idx in cutlass.range_constexpr(self.gather_b_calls_per_warp): row_start = gather_idx * 4 @@ -1204,33 +1065,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1263,10 +1115,7 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) @@ -1281,12 +1130,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa - ) - cute.copy( - tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py index 616bc2f576b4..3e86b1ad34e2 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/block_scaled_swap_ab_local_mega_moe_kernel.py @@ -1,11 +1,11 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Local MegaMoE kernel composition for Rubin.""" +import hashlib + from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -15,31 +15,28 @@ from cutlass.cutlass_dsl import Int32 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, +) from .....quant_def import CombineFormat, QuantKind -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler -from ..mega.block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop from .local_routing import LocalRouting, LocalRoutingArgs + _aot_symbol_prefix = "rubin_local_mega_moe_aot" @@ -97,10 +94,16 @@ def impl_desc_require(cls) -> dict[str, type]: "launch_cluster_count": int, "routing_flag_batch": int, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), + "weight_ready_source_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=1 + ) self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -141,6 +144,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.launch_cluster_count = impl_desc["launch_cluster_count"] self.routing_flag_batch = impl_desc["routing_flag_batch"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + self.weight_ready_source_count = impl_desc.get("weight_ready_source_count", 1) + self._validate_helper_configuration() self.occupancy = 1 self.architecture = "sm_107" @@ -170,9 +176,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "promised_launchable_sm_count": ( self.launch_cluster_count * self.cluster_shape_mn[0] * self.cluster_shape_mn[1] ), - "router_smem_limit_bytes": cutlass.memory.get_smem_capacity_in_bytes( - self.architecture - ), + "router_smem_limit_bytes": utils.get_smem_capacity_in_bytes(self.architecture), } ) resolved_problem_desc = ProblemDesc( @@ -185,9 +189,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.local_routing = LocalRouting(resolved_problem_desc, local_routing_impl_desc) self.max_routed_tokens = self.local_routing.worst_case_token_count - fc12_problem_desc = ProblemDesc( - {**resolved_problem_desc, "expert_count": self.expert_count} - ) + fc12_problem_desc = ProblemDesc({**resolved_problem_desc, "expert_count": self.expert_count}) resolved_impl_desc = ImplDesc( { **impl_desc, @@ -201,9 +203,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -214,15 +214,19 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) + def _validate_helper_configuration(self) -> None: + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") + if type(self.weight_ready_source_count) is not int or self.weight_ready_source_count <= 0: + raise ValueError("weight_ready_source_count must be a positive exact int") + if self.helper_expert_count > 0: + derive_helper_weight_ready_topology(self.expert_count, self.helper_expert_count) + def _validate_geometry(self) -> None: static_expert_dimensions = ( isinstance(self.expert_count, int), @@ -230,21 +234,15 @@ def _validate_geometry(self) -> None: isinstance(self.hidden_size, int), ) if any(static_expert_dimensions) and not all(static_expert_dimensions): - raise ValueError( - "Local MegaMoE expert dimensions must be either all static or all runtime." - ) + raise ValueError("Local MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The Local MegaMoE kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The Local MegaMoE kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("Local MegaMoE expert dimensions must be positive.") if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") if self.intermediate_gateup_size % 2 != 0: @@ -261,12 +259,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if not 1 <= self.fc1_epi_flag_batch <= len(self.epilogue_warp_ids): @@ -282,17 +276,12 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin Local MegaMoE requires mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin Local MegaMoE requires mma_k_mode='2x', got {self.mma_k_mode!r}.") mma_m, mma_n, mma_k = self.mma_tiler_mnk instruction_m, instruction_n, instruction_k = self.mma_instruction_mnk @@ -309,9 +298,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin Local MegaMoE does not implement M/N instruction repetition or B-reuse.") if mma_n not in (64, 128, 256): raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") if mma_k % instruction_k != 0: @@ -323,9 +310,7 @@ def _validate_geometry(self) -> None: "Rubin Local MegaMoE excludes NVFP4 N256 K512 because its multi-window path has no measured gain." ) if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 4 or cluster_m & (cluster_m - 1): @@ -342,9 +327,7 @@ def _resolve_schedule_hint(self) -> None: blocks_fc1 = ceil_div(self.intermediate_gateup_size, cluster_feature_tile) blocks_fc2 = ceil_div(self.hidden_size, cluster_feature_tile) # Cover the worst token-block alignment of one full persistent-cluster FC2 claim wave. - max_dependent_token_blocks = ceil_div( - self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2 - ) + max_dependent_token_blocks = ceil_div(self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2) required_fc1_work = max_dependent_token_blocks * blocks_fc1 raw_minimum_hint = max(1, ceil_div(required_fc1_work, self.launch_cluster_count)) minimum_safe_hint = round_up(raw_minimum_hint, self.fc1_epi_flag_batch) @@ -419,6 +402,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_weightsources{self.weight_ready_source_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -435,6 +419,7 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + + ("_dlb_v1" if self.helper_expert_count > 0 else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -453,14 +438,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens hidden = self.hidden_size @@ -469,20 +450,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so activation uses the routing input dtype. - activation=fake_tensor( - self.local_routing.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.local_routing.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.local_routing.activation_sf_dtype, (tokens, self.local_routing.activation_sf_hidden_padded), @@ -492,18 +467,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -517,24 +484,34 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): fc1_norm_const=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, 0, AddressSpace.gmem, assumed_align=8 + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) + compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: return compiled - compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True - ) + compiled.export_to_c(out_path, function_name=self._aot_symbol_name(), export_only_tvm_ffi_symbols=True) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() self.local_routing.register_smem_regions(smem_workspace) self.scheduler.register_smem_regions(smem_workspace) @@ -545,9 +522,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -572,16 +547,28 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[weight_ready_source_count] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + (self.helper_expert_count > 0 and not all(value is not None for value in helper_ready_bundle)) + or (self.helper_expert_count == 0 and any(value is not None for value in helper_ready_bundle)) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" + ) + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -654,19 +641,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -684,31 +666,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -716,21 +688,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) else: - pre_reduced_activation = self.local_routing.pre_reduced_activation_tensor( - self._device_workspace - ) + pre_reduced_activation = self.local_routing.pre_reduced_activation_tensor(self._device_workspace) fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) self._mainloop.materialize_codegen_members() @@ -762,9 +727,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc2_sfb=fc2_sfb, ) - fc1_output_tma_atom, fc1_output_tma_tensor = self.epilogue.prepare_tma_store_params( - fc1_output_gemm - ) + fc1_output_tma_atom, fc1_output_tma_tensor = self.epilogue.prepare_tma_store_params(fc1_output_gemm) grid = self.scheduler.get_grid_shape(max_active_clusters=self.launch_cluster_count) self._device_workspace.remove_device_members() @@ -797,6 +760,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -807,9 +772,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: ) if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, None, output_activation, reduce_scores, stream - ) + self._topk_reduce(pre_reduced_activation, None, output_activation, reduce_scores, stream) @cute.kernel def _kernel( @@ -842,6 +805,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose local routing, scheduler, mainloop, and epilogue.""" self._mainloop.materialize_codegen_members() @@ -863,8 +828,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -918,22 +882,24 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.expert_count, + helper_count=self.helper_expert_count, + source_count=self.weight_ready_source_count, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.local_routing.fc1_ready_counter_pointer( - self._device_workspace - ), + fc1_ready_counter_pointer=self.local_routing.fc1_ready_counter_pointer(self._device_workspace), ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) @@ -951,9 +917,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1043,9 +1015,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() if warp_idx >= Int32(self.transfer_warp_idx_start): @@ -1074,9 +1044,7 @@ def _kernel( iket.range_push("local_mega.kernel_tail_wait") cute.arch.sync_threads() iket.range_pop() - if (warp_idx >= Int32(self.transfer_warp_idx_start)) & ( - warp_idx < Int32(self.token_in_end_warp_idx) - ): + if (warp_idx >= Int32(self.transfer_warp_idx_start)) & (warp_idx < Int32(self.token_in_end_warp_idx)): iket.range_push("local_mega.tail_reset") self.local_routing.reset_tail() iket.range_pop() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py index 3ed254dcfc0c..fd3ceb087e89 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/local_routing.py @@ -1,12 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Routing support for the Rubin Local MegaMoE kernel.""" import dataclasses from typing import Callable, ClassVar, Optional, Tuple, Union import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -20,7 +18,13 @@ from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import make_flag_batch_tracker from .....helpers.iket_compat import iket -from .....helpers.ptx_helpers import cp_async_bulk_s2g, nanosleep, stg_b64, stg_f32, tma_load_1d +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + nanosleep, + stg_b64, + stg_f32, + tma_load_1d, +) from .....helpers.smem_workspace import SmemWorkspace from .....helpers.software_sync import SoftwareGridSync from .....helpers.utils import ceil_div, round_up @@ -87,16 +91,12 @@ def from_packed(cls, packed: Union[Int64, Int32]) -> "_SortedElement": return cls(flat_topk_index=Int32(packed), topk_score=None) scratch = cute.make_rmem_tensor((2,), cutlass.Int32) cute.recast_tensor(scratch, cutlass.Int64)[0] = packed - return cls( - flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1] - ) + return cls(flat_topk_index=scratch[0], topk_score=cute.recast_tensor(scratch, cutlass.Float32)[1]) @cute.jit def _copy_atom(dtype, num_bits_per_copy: int): - return cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy - ) + return cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype, num_bits_per_copy=num_bits_per_copy) @cute.jit @@ -113,9 +113,7 @@ def _smem_exclusive_prefix( num_elements = cute.size(input_tensor) scan_rows = num_elements // 4 num_warps = block_thread_count // 32 - input_vectors = cute.make_tensor( - input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1)) - ) + input_vectors = cute.make_tensor(input_tensor.iterator, cute.make_layout((scan_rows, 4), stride=(4, 1))) load_atom = _copy_atom(cutlass.Int32, 128) values = cute.make_rmem_tensor((4,), cutlass.Int32) @@ -129,12 +127,7 @@ def _smem_exclusive_prefix( for element in cutlass.range_constexpr(4): values[element] = Int32(0) - local_prefix = ( - Int32(0), - values[0], - values[0] + values[1], - values[0] + values[1] + values[2], - ) + local_prefix = (Int32(0), values[0], values[0] + values[1], values[0] + values[1] + values[2]) lane_total = local_prefix[3] + values[3] inclusive = lane_total for step_log in cutlass.range_constexpr(5): @@ -231,16 +224,13 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_router_configuration() self.expert_count_padded = round_up(self.expert_count, 4) self.expert_count_with_trash = self.expert_count_padded + 1 - self.router_elements_per_lane, self.router_data_cta_count = ( - self._router_launch_configuration() - ) + self.router_elements_per_lane, self.router_data_cta_count = self._router_launch_configuration() self.router_tokens_per_cta = self.router_elements_per_lane * self.router_warps_per_cta * 32 self.router_push_cta_count = ceil_div(self.expert_count, self.router_warps_per_cta) self.router_grid_cta_count = max(self.router_data_cta_count + 1, self.router_push_cta_count) if self.router_grid_cta_count > self.promised_launchable_sm_count: raise ValueError( - "Router grid exceeds promised_launchable_sm_count; " - "all metadata-push CTAs must be concurrently resident." + "Router grid exceeds promised_launchable_sm_count; all metadata-push CTAs must be concurrently resident." ) self._router_smem_workspace = self._build_router_smem_workspace() @@ -267,18 +257,14 @@ def _validate_router_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Local routing supports at most 16384 experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") def worst_case_padded_tokens(self, block: int) -> int: source_token_capacity = self.max_tokens routes_per_source_token = min(self.topk, self.expert_count) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.expert_count, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -296,13 +282,8 @@ def next_power_of_two(value: int) -> int: maximum_data_cta_count = 128 maximum_supported_tokens = maximum_cta_capacity * maximum_data_cta_count if routed_token_capacity > maximum_supported_tokens: - raise NotImplementedError( - f"The router supports at most {maximum_supported_tokens} routed tokens." - ) - cta_capacity = min( - maximum_cta_capacity, - next_power_of_two(max(routed_token_capacity, minimum_cta_capacity)), - ) + raise NotImplementedError(f"The router supports at most {maximum_supported_tokens} routed tokens.") + cta_capacity = min(maximum_cta_capacity, next_power_of_two(max(routed_token_capacity, minimum_cta_capacity))) elements_per_lane = cta_capacity // (self.router_warps_per_cta * 32) data_cta_count = ceil_div(routed_token_capacity, cta_capacity) return elements_per_lane, data_cta_count @@ -311,18 +292,11 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: workspace = SmemWorkspace() overlay = workspace.create_overlay("local_mega.routing.router_smem.role") data_lifetime = overlay.add_lifetime("data_cta") + data_lifetime.register_tensor(self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,)) data_lifetime.register_tensor( - self.router_data_histogram_region, cutlass.Int32, (self.expert_count_with_trash,) - ) - data_lifetime.register_tensor( - self.router_data_prefix_region, - cutlass.Int32, - (self.expert_count_with_trash,), - byte_alignment=16, - ) - data_lifetime.register_tensor( - self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,) + self.router_data_prefix_region, cutlass.Int32, (self.expert_count_with_trash,), byte_alignment=16 ) + data_lifetime.register_tensor(self.router_data_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,)) data_lifetime.register_tensor( self.router_data_sorted_region, (cutlass.Int64 if self.apply_topk_at_fc1 else cutlass.Int32), @@ -331,30 +305,18 @@ def _build_router_smem_workspace(self) -> SmemWorkspace: ) if self.router_data_cta_count > 1: data_lifetime.register_tensor( - self.router_data_base_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_data_base_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime = overlay.add_lifetime("helper_cta") helper_lifetime.register_tensor( - self.router_helper_totals_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_totals_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_prefix_region, - cutlass.Int32, - (self.expert_count_padded,), - byte_alignment=16, + self.router_helper_prefix_region, cutlass.Int32, (self.expert_count_padded,), byte_alignment=16 ) helper_lifetime.register_tensor( - self.router_helper_warp_totals_region, - cutlass.Int32, - (self.router_warps_per_cta,), - byte_alignment=16, + self.router_helper_warp_totals_region, cutlass.Int32, (self.router_warps_per_cta,), byte_alignment=16 ) workspace.finalize(max_bytes=self.router_smem_limit_bytes) return workspace @@ -370,34 +332,14 @@ def register_device_workspace(self, workspace: DeviceWorkspace) -> None: def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: maximum_routed_tokens = self.max_tokens * self.topk workspace.register( - self.sizes_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) - workspace.register( - self.metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="shared", - reset="tail_reset", - ) - workspace.register( - self.sorted_metadata_region, - cutlass.Int64, - (maximum_routed_tokens,), - buffer_space="local", + self.sizes_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="shared", reset="tail_reset" ) + workspace.register(self.sizes_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.metadata_ready_region, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") + workspace.register(self.sorted_metadata_region, cutlass.Int64, (maximum_routed_tokens,), buffer_space="local") if self.apply_topk_at_fc1: workspace.register( - self.sorted_scores_region, - cutlass.Float32, - (maximum_routed_tokens,), - buffer_space="local", + self.sorted_scores_region, cutlass.Float32, (maximum_routed_tokens,), buffer_space="local" ) workspace.register( self.token_src_metadata_region, @@ -408,34 +350,16 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: ) if self.apply_topk_at_fc1: workspace.register( - self.fc1_topk_scores_region, - cutlass.Float32, - (self.worst_case_token_count,), - buffer_space="shared", + self.fc1_topk_scores_region, cutlass.Float32, (self.worst_case_token_count,), buffer_space="shared" ) + workspace.register(self.pool_expert_base_region, cutlass.Int32, (self.expert_count,), buffer_space="local") workspace.register( - self.pool_expert_base_region, cutlass.Int32, (self.expert_count,), buffer_space="local" + self.source_expert_base_region, cutlass.Int32, (self.expert_count_padded,), buffer_space="local" ) workspace.register( - self.source_expert_base_region, - cutlass.Int32, - (self.expert_count_padded,), - buffer_space="local", - ) - workspace.register( - self.sorted_metadata_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", - ) - workspace.register( - self.push_table_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.sorted_metadata_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) + workspace.register(self.push_table_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") if self.router_data_cta_count > 1: workspace.register( self.router_size_counter_region, @@ -445,18 +369,10 @@ def _register_router_workspace(self, workspace: DeviceWorkspace) -> None: reset="tail_reset", ) workspace.register( - self.router_histogram_done_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.router_histogram_done_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) workspace.register( - self.source_base_ready_region, - cutlass.Int32, - (1,), - buffer_space="local", - reset="tail_reset", + self.source_base_ready_region, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" ) def __extract_mlir_values__(self) -> list: @@ -545,16 +461,10 @@ def _router_single_cta( if self._router_linear_cta_idx < Int32(self.router_data_cta_count): block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) for zero_round in cutlass.range_constexpr(zero_round_count): @@ -595,12 +505,7 @@ def _router_single_cta( iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -609,10 +514,7 @@ def _router_single_cta( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -626,16 +528,10 @@ def _router_multiple_ctas( block_thread_count = self.router_warps_per_cta * 32 trash_bucket = self.expert_count_padded expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) - histogram = self._router_smem_workspace.tensor( - self.router_data_histogram_region, smem_base - ) + histogram = self._router_smem_workspace.tensor(self.router_data_histogram_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_data_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_data_warp_totals_region, smem_base - ) - sorted_elements = self._router_smem_workspace.tensor( - self.router_data_sorted_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_data_warp_totals_region, smem_base) + sorted_elements = self._router_smem_workspace.tensor(self.router_data_sorted_region, smem_base) dump_base = self._router_smem_workspace.tensor(self.router_data_base_region, smem_base) zero_round_count = ceil_div(self.expert_count_with_trash, block_thread_count) @@ -656,17 +552,12 @@ def _router_multiple_ctas( expert = Int32(expert_round * block_thread_count) + self._router_thread_idx if expert < Int32(self.expert_count_padded): dump_base[expert] = Int32( - cute.arch.atomic_add( - size_counter + expert, histogram[expert], sem="relaxed", scope="gpu" - ) + cute.arch.atomic_add(size_counter + expert, histogram[expert], sem="relaxed", scope="gpu") ) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.router_histogram_done_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.router_histogram_done_region), Int32(1), sem="release", scope="gpu" ) total_valid_routes = _smem_exclusive_prefix( @@ -684,12 +575,7 @@ def _router_multiple_ctas( iket.range_pop() iket.range_push("router.sort") self._sort_router_elements( - expert_registers, - within_expert_indices, - score_registers, - sorted_elements, - prefix, - topk_indices.dtype, + expert_registers, within_expert_indices, score_registers, sorted_elements, prefix, topk_indices.dtype ) cute.arch.sync_threads() iket.range_pop() @@ -697,9 +583,7 @@ def _router_multiple_ctas( iket.range_push("router.wait_source_base") source_base_ready = self._device_workspace.ptr(self.source_base_ready_region) if self._router_thread_idx == Int32(0): - while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32( - 1 - ): + while cute.arch.load(source_base_ready, Int32, sem="acquire", scope="gpu") != Int32(1): nanosleep(150) cute.arch.sync_threads() iket.range_pop() @@ -715,10 +599,7 @@ def _router_multiple_ctas( cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sorted_metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sorted_metadata_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() elif self._router_linear_cta_idx == Int32(self.router_data_cta_count): @@ -735,9 +616,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) size_counter = self._device_workspace.tensor(self.router_size_counter_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) expert_round_count = ceil_div(self.expert_count_padded, block_thread_count) @@ -776,10 +655,7 @@ def _router_helper_multiple_ctas(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.source_base_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.source_base_ready_region), Int32(1), sem="release", scope="gpu" ) publish_sizes() iket.range_pop() @@ -801,9 +677,7 @@ def _router_push_metadata(self) -> None: nanosleep(150) cute.arch.sync_threads() - global_expert = ( - self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx - ) + global_expert = self._router_linear_cta_idx * Int32(self.router_warps_per_cta) + self._router_warp_idx if global_expert < Int32(self.expert_count): sizes = self._device_workspace.tensor(self.sizes_region) source_expert_base = self._device_workspace.tensor(self.source_expert_base_region) @@ -812,14 +686,10 @@ def _router_push_metadata(self) -> None: source_begin = source_expert_base[global_expert] destination_begin = pool_expert_base[global_expert] source_metadata = self._device_workspace.ptr(self.sorted_metadata_region) - destination_metadata_address = self._device_workspace.ptr( - self.token_src_metadata_region - ).toint() + destination_metadata_address = self._device_workspace.ptr(self.token_src_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): source_scores = self._device_workspace.ptr(self.sorted_scores_region) - destination_scores_address = self._device_workspace.ptr( - self.fc1_topk_scores_region - ).toint() + destination_scores_address = self._device_workspace.ptr(self.fc1_topk_scores_region).toint() route_round_count = (route_count + Int32(31)) // Int32(32) for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx @@ -828,16 +698,12 @@ def _router_push_metadata(self) -> None: destination_position = destination_begin + route metadata = cute.arch.load(source_metadata + source_position, cutlass.Int64) stg_b64( - destination_metadata_address - + Int64(destination_position) * Int64(RoutingMetadata.nbytes), + destination_metadata_address + Int64(destination_position) * Int64(RoutingMetadata.nbytes), metadata, ) if cutlass.const_expr(self.apply_topk_at_fc1): score = cute.arch.load(source_scores + source_position, cutlass.Float32) - stg_f32( - destination_scores_address + Int64(destination_position) * Int64(4), - score, - ) + stg_f32(destination_scores_address + Int64(destination_position) * Int64(4), score) cute.arch.sync_threads() if self._router_thread_idx == Int32(0): @@ -846,10 +712,7 @@ def _router_push_metadata(self) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.metadata_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.metadata_ready_region), Int32(1), sem="release", scope="gpu" ) @cute.jit @@ -858,9 +721,7 @@ def _publish_sizes(self, smem_expert_counts: cute.Tensor) -> Callable[[], None]: if self._router_warp_idx == Int32(0): with cute.arch.elect_one(): cp_async_bulk_s2g( - self._device_workspace.ptr(self.sizes_region), - smem_expert_counts.iterator, - row_bytes, + self._device_workspace.ptr(self.sizes_region), smem_expert_counts.iterator, row_bytes ) cute.arch.cp_async_bulk_commit_group() @@ -872,10 +733,7 @@ def finalize() -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.sizes_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.sizes_ready_region), Int32(1), sem="release", scope="gpu" ) return finalize @@ -884,13 +742,9 @@ def finalize() -> None: def _compute_local_tables(self, smem_base: cute.Pointer) -> None: block_thread_count = self.router_warps_per_cta * 32 - padded_totals = self._router_smem_workspace.tensor( - self.router_helper_totals_region, smem_base - ) + padded_totals = self._router_smem_workspace.tensor(self.router_helper_totals_region, smem_base) prefix = self._router_smem_workspace.tensor(self.router_helper_prefix_region, smem_base) - warp_totals = self._router_smem_workspace.tensor( - self.router_helper_warp_totals_region, smem_base - ) + warp_totals = self._router_smem_workspace.tensor(self.router_helper_warp_totals_region, smem_base) sizes = self._device_workspace.tensor(self.sizes_region) sizes_ready = self._device_workspace.ptr(self.sizes_ready_region) @@ -908,8 +762,7 @@ def _compute_local_tables(self, smem_base: cute.Pointer) -> None: if expert < Int32(self.expert_count_padded): expert_size = sizes[expert] padded_totals[expert] = ( - (expert_size + Int32(self.token_padding_block - 1)) - // Int32(self.token_padding_block) + (expert_size + Int32(self.token_padding_block - 1)) // Int32(self.token_padding_block) ) * Int32(self.token_padding_block) cute.arch.sync_threads() _smem_exclusive_prefix( @@ -929,10 +782,7 @@ def _compute_local_tables(self, smem_base: cute.Pointer) -> None: cute.arch.sync_threads() if self._router_thread_idx == Int32(0): cute.arch.atomic_add( - self._device_workspace.ptr(self.push_table_ready_region), - Int32(1), - sem="release", - scope="gpu", + self._device_workspace.ptr(self.push_table_ready_region), Int32(1), sem="release", scope="gpu" ) iket.range_pop() @@ -950,62 +800,46 @@ def _load_router_inputs( elements_per_thread = load_round_count * elements_per_vector topk_flat = cute.make_tensor(topk_indices.iterator, cute.make_layout((maximum_elements,))) - topk_vectors = cute.logical_divide( - cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None) - ) + topk_vectors = cute.logical_divide(cute.zipped_divide(topk_flat, (tile_span,)), (elements_per_vector, None)) load_atom = _copy_atom(topk_indices.dtype, 128) expert_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Int32) if cutlass.const_expr(topk_indices.dtype.width == 64): raw_indices = cute.make_rmem_tensor((elements_per_thread,), topk_indices.dtype) raw_vectors = cute.zipped_divide(raw_indices, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), raw_vectors[None, load_round], ) else: expert_vectors = cute.zipped_divide(expert_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( load_atom, - mark_alignment( - topk_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(topk_vectors[(None, self._router_grid_thread_idx), load_round], 16), expert_vectors[None, load_round], ) score_registers = None if cutlass.const_expr(self.apply_topk_at_fc1): score_registers = cute.make_rmem_tensor((elements_per_thread,), cutlass.Float32) - scores_flat = cute.make_tensor( - topk_scores.iterator, cute.make_layout((maximum_elements,)) - ) + scores_flat = cute.make_tensor(topk_scores.iterator, cute.make_layout((maximum_elements,))) score_vectors = cute.logical_divide( cute.zipped_divide(scores_flat, (tile_span,)), (elements_per_vector, None) ) score_atom = _copy_atom(cutlass.Float32, elements_per_vector * 32) score_register_vectors = cute.zipped_divide(score_registers, (elements_per_vector,)) for load_round in cutlass.range_constexpr(load_round_count): - tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32( - elements_per_vector - ) + tile_begin = Int32(load_round * tile_span) + self._router_grid_thread_idx * Int32(elements_per_vector) if tile_begin < actual_elements: cute.copy( score_atom, - mark_alignment( - score_vectors[(None, self._router_grid_thread_idx), load_round], 16 - ), + mark_alignment(score_vectors[(None, self._router_grid_thread_idx), load_round], 16), score_register_vectors[None, load_round], ) @@ -1016,26 +850,21 @@ def _load_router_inputs( if cutlass.const_expr(topk_indices.dtype.width == 64): expert_registers[register_idx] = raw_indices_i32[2 * register_idx] token_idx, _ = self._router_value_coordinate(register_idx, topk_indices.dtype) - is_invalid = ( - expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count) - ) | (token_idx >= actual_token_count) + is_invalid = (expert_registers_u32[register_idx] >= cutlass.Uint32(self.expert_count)) | ( + token_idx >= actual_token_count + ) if is_invalid: expert_registers[register_idx] = Int32(self.expert_count_padded) return expert_registers, score_registers @cute.jit - def _build_histogram( - self, expert_registers: cute.Tensor, histogram: cute.Tensor - ) -> cute.Tensor: + def _build_histogram(self, expert_registers: cute.Tensor, histogram: cute.Tensor) -> cute.Tensor: register_count = cute.size(expert_registers) within_expert_indices = cute.make_rmem_tensor((register_count,), cutlass.Int32) for register_idx in cutlass.range_constexpr(register_count): within_expert_indices[register_idx] = Int32( cute.arch.atomic_add( - histogram.iterator + expert_registers[register_idx], - Int32(1), - sem="relaxed", - scope="cta", + histogram.iterator + expert_registers[register_idx], Int32(1), sem="relaxed", scope="cta" ) ) cute.arch.sync_threads() @@ -1055,41 +884,28 @@ def _sort_router_elements( for register_idx in cutlass.range_constexpr(register_count): token_idx, topk_slot = self._router_value_coordinate(register_idx, topk_index_type) flat_topk_index = token_idx * Int32(self.topk) + topk_slot - destination = ( - expert_run_starts[expert_registers[register_idx]] - + within_expert_indices[register_idx] - ) + destination = expert_run_starts[expert_registers[register_idx]] + within_expert_indices[register_idx] if cutlass.const_expr(self.apply_topk_at_fc1): - sorted_elements[destination] = _SortedElement( - flat_topk_index, score_registers[register_idx] - ).pack() + sorted_elements[destination] = _SortedElement(flat_topk_index, score_registers[register_idx]).pack() else: sorted_elements[destination] = _SortedElement(flat_topk_index, None).pack() @cute.jit - def _router_value_coordinate( - self, register_idx: int, topk_index_type: type - ) -> Tuple[Int32, Int32]: + def _router_value_coordinate(self, register_idx: int, topk_index_type: type) -> Tuple[Int32, Int32]: elements_per_vector = 128 // topk_index_type.width - tile_span = ( - elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 - ) + tile_span = elements_per_vector * self.router_data_cta_count * self.router_warps_per_cta * 32 flat_index = Int32( register_idx // elements_per_vector * tile_span + register_idx % elements_per_vector ) + self._router_grid_thread_idx * Int32(elements_per_vector) return (flat_index // Int32(self.topk), flat_index % Int32(self.topk)) @cute.jit - def _dump_contiguous_router_output( - self, sorted_elements: cute.Tensor, total_valid_routes: Int32 - ) -> None: + def _dump_contiguous_router_output(self, sorted_elements: cute.Tensor, total_valid_routes: Int32) -> None: block_thread_count = self.router_warps_per_cta * 32 metadata_address = self._device_workspace.ptr(self.sorted_metadata_region).toint() if cutlass.const_expr(self.apply_topk_at_fc1): score_address = self._device_workspace.ptr(self.sorted_scores_region).toint() - dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32( - block_thread_count - ) + dump_round_count = (total_valid_routes + Int32(block_thread_count - 1)) // Int32(block_thread_count) for dump_round in cutlass.range(dump_round_count, unroll=4): position = Int32(dump_round * block_thread_count) + self._router_thread_idx predicate = Int32(position < total_valid_routes) @@ -1098,11 +914,7 @@ def _dump_contiguous_router_output( token=(element.flat_topk_index // Int32(self.topk)), topk_slot=(element.flat_topk_index % Int32(self.topk)), ) - stg_b64( - metadata_address + Int64(position) * Int64(RoutingMetadata.nbytes), - metadata.pack(), - predicate, - ) + stg_b64(metadata_address + Int64(position) * Int64(RoutingMetadata.nbytes), metadata.pack(), predicate) if cutlass.const_expr(self.apply_topk_at_fc1): stg_f32(score_address + Int64(position) * Int64(4), element.topk_score, predicate) @@ -1128,9 +940,7 @@ def _dump_router_output_by_expert( for route_round in cutlass.range(route_round_count, unroll=1): route = Int32(route_round) * Int32(32) + self._router_lane_idx predicate = Int32(route < run_length) - element = _SortedElement.from_packed( - sorted_elements[predicate * (run_begin + route)] - ) + element = _SortedElement.from_packed(sorted_elements[predicate * (run_begin + route)]) output_position = dump_begin + route metadata = RoutingMetadata( token=(element.flat_topk_index // Int32(self.topk)), @@ -1142,11 +952,7 @@ def _dump_router_output_by_expert( predicate, ) if cutlass.const_expr(self.apply_topk_at_fc1): - stg_f32( - score_address + Int64(output_position) * Int64(4), - element.topk_score, - predicate, - ) + stg_f32(score_address + Int64(output_position) * Int64(4), element.topk_score, predicate) @cute.jit def expert_sizes(self, device_workspace: DeviceWorkspace) -> cute.Tensor: @@ -1171,9 +977,7 @@ def token_src_metadata_pointer(self, device_workspace: DeviceWorkspace) -> cute. return device_workspace.ptr(self.token_src_metadata_region) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1183,9 +987,7 @@ def wait_for_sizes_ready( cute.arch.sync_warp() @cute.jit - def wait_for_metadata_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_metadata_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) if lane_idx == Int32(0): @@ -1291,20 +1093,14 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Local routing supports at most 16384 experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if not 1 <= self.token_in_flag_batch <= 32: - raise ValueError( - f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}." - ) + raise ValueError(f"token_in_flag_batch must be in [1, 32], got {self.token_in_flag_batch}.") if self.tokens_per_fc1_ready_slot % self.token_padding_block != 0: raise ValueError("tokens_per_fc1_ready_slot must be divisible by token_padding_block.") element_block = self.activation_sf_vector_size * 4 if self.hidden_size % element_block != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {element_block}." - ) + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {element_block}.") if self.sf_padding_block % 128 != 0: raise ValueError("sf_padding_block must be a multiple of 128.") @@ -1336,10 +1132,7 @@ def worst_case_sf_token_count(self) -> int: @property def max_fc1_ready_slot_count(self) -> int: - return ( - self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) - // self.tokens_per_fc1_ready_slot - ) + return self._router.worst_case_padded_tokens(self.tokens_per_fc1_ready_slot) // self.tokens_per_fc1_ready_slot @property def token_src_metadata_region(self) -> str: @@ -1374,9 +1167,7 @@ def expert_sizes(self, device_workspace: DeviceWorkspace) -> cute.Tensor: return self._router.expert_sizes(device_workspace) @cute.jit - def wait_for_sizes_ready( - self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000 - ) -> None: + def wait_for_sizes_ready(self, device_workspace: DeviceWorkspace, sleep_cycles: int = 1000) -> None: self._router.wait_for_sizes_ready(device_workspace, sleep_cycles) @cute.jit @@ -1439,9 +1230,7 @@ def _register_main_workspace(self, workspace: DeviceWorkspace) -> None: buffer_space="local", reset="tail_reset", ) - activation_sf_element_count = ( - self.worst_case_sf_token_count * self.activation_sf_hidden_padded - ) + activation_sf_element_count = self.worst_case_sf_token_count * self.activation_sf_hidden_padded workspace.register( self.fc1_activation_sf_region, self.activation_sf_dtype, @@ -1470,10 +1259,7 @@ def register_smem_regions(self, workspace: SmemWorkspace) -> None: token_in_lifetime.register_tensor( self.token_in_sf_smem_region, self.activation_sf_dtype, - ( - self.transfer_warp_count, - (self.activation_sf_vector_size, self.activation_sf_hidden_padded), - ), + (self.transfer_warp_count, (self.activation_sf_vector_size, self.activation_sf_hidden_padded)), stride=(self.activation_sf_hidden_padded, (0, 1)), byte_alignment=16, ) @@ -1487,14 +1273,10 @@ def fc1_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> cute.Te layout = tile_atom_to_shape_SF( (self.worst_case_sf_token_count, self.hidden_size, 1), self.activation_sf_vector_size ) - return cute.make_tensor( - device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1]) - ) + return cute.make_tensor(device_workspace.ptr(self.fc1_activation_sf_region), cute.select(layout, mode=[0, 1])) @cute.jit - def pre_reduced_activation_tensor( - self, device_workspace: DeviceWorkspace - ) -> Optional[cute.Tensor]: + def pre_reduced_activation_tensor(self, device_workspace: DeviceWorkspace) -> Optional[cute.Tensor]: """The (tokens, topk, hidden) combine staging plane, or None under in-kernel top-k reduction.""" if cutlass.const_expr(self.reduce_topk_in_kernel): return None @@ -1524,20 +1306,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No source_sizes = cute.make_tensor(sizes.iterator, cute.make_layout((self.expert_count,))) copy_elements = 4 if self.expert_count % 4 == 0 else 1 source_size_vectors = cute.zipped_divide( - ( - mark_alignment(source_sizes, 16) - if cutlass.const_expr(copy_elements == 4) - else source_sizes - ), + (mark_alignment(source_sizes, 16) if cutlass.const_expr(copy_elements == 4) else source_sizes), (copy_elements,), ) destination_size_vectors = cute.zipped_divide(owned_sizes, (copy_elements,)) size_vector_count = cute.size(destination_size_vectors, mode=[1]) size_copy_atom = cute.make_copy_atom( cute.nvgpu.cpasync.CopyG2SOp( - cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL - if copy_elements == 4 - else cute.nvgpu.LoadCacheMode.ALWAYS + cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL if copy_elements == 4 else cute.nvgpu.LoadCacheMode.ALWAYS ), cutlass.Int32, num_bits_per_copy=copy_elements * 32, @@ -1548,9 +1324,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No vector_idx = Int32(size_copy_round * self.transfer_thread_count) + transfer_thread_idx if vector_idx < Int32(size_vector_count): cute.copy( - size_copy_atom, - source_size_vectors[None, vector_idx], - destination_size_vectors[None, vector_idx], + size_copy_atom, source_size_vectors[None, vector_idx], destination_size_vectors[None, vector_idx] ) cute.arch.cp_async_commit_group() iket.range_pop() @@ -1582,9 +1356,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No sf_copy_elements = 4 source_sf_values = cute.slice_(warp_sf_stage, (0, None)) source_sf_vectors = cute.zipped_divide(source_sf_values, (sf_copy_elements,)) - sf_copy_atom = _copy_atom( - self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width) - ) + sf_copy_atom = _copy_atom(self.activation_sf_dtype, sf_copy_elements * int(self.activation_sf_dtype.width)) next_dense_token = global_warp_idx expert_valid_begin = Int32(0) @@ -1605,34 +1377,24 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No expert_valid_end = expert_valid_begin + expert_token_count pull_count = Int32(0) if next_dense_token < expert_valid_end: - pull_count = ( - expert_valid_end - next_dense_token + global_warp_count - Int32(1) - ) // global_warp_count + pull_count = (expert_valid_end - next_dense_token + global_warp_count - Int32(1)) // global_warp_count for pull_round in cutlass.range(pull_count, unroll=1): dense_token_idx = next_dense_token + Int32(pull_round) * global_warp_count token_in_expert = dense_token_idx - expert_valid_begin pool_token_idx = pool_expert_bases[local_expert] + token_in_expert sf_token_idx = expert_sf_begin + token_in_expert - ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32( - self.tokens_per_fc1_ready_slot - ) + ready_slot_idx = expert_ready_slot_begin + token_in_expert // Int32(self.tokens_per_fc1_ready_slot) metadata = RoutingMetadata.load( - token_metadata_pointer.toint() - + Int64(pool_token_idx) * Int64(RoutingMetadata.nbytes) + token_metadata_pointer.toint() + Int64(pool_token_idx) * Int64(RoutingMetadata.nbytes) ) source_sf_row = self._local_routing_args.activation_sf[Int64(metadata.token), None] with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - warp_mbarrier, Int32(activation_sf_bytes) - ) + cute.arch.mbarrier_arrive_and_expect_tx(warp_mbarrier, Int32(activation_sf_bytes)) tma_load_1d( - warp_sf_stage.iterator, - source_sf_row.iterator, - warp_mbarrier, - Int32(activation_sf_bytes), + warp_sf_stage.iterator, source_sf_row.iterator, warp_mbarrier, Int32(activation_sf_bytes) ) cute.arch.sync_warp() cute.arch.mbarrier_wait(warp_mbarrier, pull_phase) @@ -1640,9 +1402,7 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No destination_sf_row = fc1_activation_sf[Int64(sf_token_idx), ((None, None), None)] destination_sf_values = cute.slice_(destination_sf_row, (0, None, None)) destination_sf_values = cute.group_modes(destination_sf_values, 0, 2) - destination_sf_vectors = cute.zipped_divide( - destination_sf_values, (sf_copy_elements,) - ) + destination_sf_vectors = cute.zipped_divide(destination_sf_values, (sf_copy_elements,)) sf_vector_count = cute.size(destination_sf_vectors, mode=[1]) for sf_round in cutlass.range_constexpr(ceil_div(sf_vector_count, 32)): sf_vector_idx = Int32(sf_round * 32) + lane_idx @@ -1655,17 +1415,14 @@ def token_in(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> No cute.arch.sync_warp() ready_address = (fc1_ready_counter + ready_slot_idx).toint() - flag_tracker = flag_tracker.accumulate( - Int32(0), self.token_in_flag_batch, ready_address - ) + flag_tracker = flag_tracker.accumulate(Int32(0), self.token_in_flag_batch, ready_address) cute.arch.sync_warp() pull_phase = pull_phase ^ Int32(1) next_dense_token = next_dense_token + pull_count * global_warp_count expert_valid_begin = expert_valid_end expert_sf_begin = expert_sf_begin + ( - (expert_token_count + Int32(self.sf_padding_block - 1)) - // Int32(self.sf_padding_block) + (expert_token_count + Int32(self.sf_padding_block - 1)) // Int32(self.sf_padding_block) ) * Int32(self.sf_padding_block) expert_ready_slot_begin = expert_ready_slot_begin + ( (expert_token_count + Int32(self.tokens_per_fc1_ready_slot - 1)) @@ -1692,9 +1449,7 @@ def reset_tail(self) -> None: ) iket.range_pop() total_reset_threads = self.promised_launchable_sm_count * self.transfer_thread_count - global_reset_thread = ( - self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx - ) + global_reset_thread = self._linear_cta_idx * Int32(self.transfer_thread_count) + transfer_thread_idx iket.range_push("tail.reset_workspace") self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py index 4e26bd3d722c..b96869775c12 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/local_mega/tma_gather.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Rubin TMA Gather4 load primitives used by Local MegaMoE.""" import cutlass @@ -10,6 +7,7 @@ from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import dsl_user_op + _CTA_GROUP_2_MBAR_MASK = 0xFEFFFFFF diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py index 52a7a9daf028..51f74b497db7 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/__init__.py @@ -1,2 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Rubin MegaMoE inference kernel components.""" + +from . import dynamic_mainloop +from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension, TensorRole +from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop +from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized +from .block_scaled_swap_ab_mega_moe_kernel import BlockScaledSwapAbMegaMoeKernel +from .block_scaled_swap_ab_mega_moe_kernel_gen_specialized import BlockScaledSwapAbGenphaseMoeKernel +from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs + + +__all__ = [ + "BlockScaledSwapAbFc12Extension", + "BlockScaledSwapAbFc12Mainloop", + "BlockScaledSwapAbFc12MainloopGenSpecialized", + "BlockScaledSwapAbGenphaseMoeKernel", + "BlockScaledSwapAbMegaMoeKernel", + "GatedActEpilogueArgs", + "GenphaseTokenComm", + "GenphaseTokenCommArgs", + "SwapABGatedActEpilogue", + "TensorRole", + "dynamic_mainloop", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py index db92666ff682..4d76ea7ffe3a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue.py @@ -10,32 +10,30 @@ """ import dataclasses +import math from typing import Any, Callable, ClassVar, List, Literal, Optional, Tuple, Type, Union import cutlass import cutlass.cute as cute -import cutlass.pipeline as pipeline +from cutlass.cutlass_dsl import Int64, T +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import AddressSpace import cutlass.utils as utils +import cutlass.pipeline as pipeline import cutlass.utils.blackwell_helpers as sm100_utils + from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Int64, T -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice from .....communication.token_protocol import TokenSrcMetadata -from .....helpers.constants import Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max, Nvfp4E2M1RcpLimit +from .....quant_def import CombineFormat, QuantKind +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType +from .....helpers.constants import Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, Fp8E5M2RcpLimit, Fp32Max from .....helpers.cute_py_helpers import tcgen05_block_scaled_acc_dtype from .....helpers.dsl_helpers import mark_alignment from .....helpers.flag_batch import make_flag_batch_tracker +from .....helpers.iket_compat import iket from .....helpers.ptx_helpers import ( cp_async_bulk_s2g, cp_reduce_async_bulk_add_bf16_s2g, @@ -44,10 +42,8 @@ red_add_relaxed_sys_v2_bf16x2, ) from .....helpers.smem_workspace import SmemRegion, SmemWorkspace -from .....quant_def import CombineFormat, QuantKind from ....function_mapping import CoordinateSpace, FunctionMapping -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo +from ....schedulers import BlockPhase, SchedulerConsumer, SwapAbFc12WorkTileInfo from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -82,43 +78,27 @@ class QuantImpl: warp_idx: Optional[Any] = None pair_exchange_barrier: Optional[Any] = None - _directions: ClassVar[Tuple[str, ...]] = ( - "regs_in_thread", - "threads_with_the_same_reg", - "regs_in_pair_threads", - ) + _directions: ClassVar[Tuple[str, ...]] = ("regs_in_thread", "threads_with_the_same_reg", "regs_in_pair_threads") # -- config / validation -------------------------------------------------- def __post_init__(self): if isinstance(self.quant_kind, CombineFormat): if not self.quant_kind.is_quantized: - raise ValueError( - f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}." - ) + raise ValueError(f"QuantImpl combine path needs a quantized CombineFormat, got {self.quant_kind}.") elif not isinstance(self.quant_kind, QuantKind): - raise ValueError( - f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}." - ) + raise ValueError(f"quant_kind must be a QuantKind or a CombineFormat, got {self.quant_kind!r}.") if self.sf_vec_direction not in self._directions: - raise ValueError( - f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}." - ) + raise ValueError(f"sf_vec_direction must be one of {self._directions}, got {self.sf_vec_direction!r}.") if self.sf_vec_direction == "threads_with_the_same_reg" and self.lane_idx is None: raise ValueError("across-lane quant needs lane_idx for the CREDUX half-warp predicate.") if self.sf_vec_direction == "regs_in_pair_threads": if isinstance(self.quant_kind, CombineFormat): - raise ValueError( - "The paired direction is orthodox-only; combine blocks never straddle warps." - ) + raise ValueError("The paired direction is orthodox-only; combine blocks never straddle warps.") if self.sf_vec_size % 2 != 0: - raise ValueError( - f"The paired direction needs an even sf_vec, got {self.sf_vec_size}." - ) + raise ValueError(f"The paired direction needs an even sf_vec, got {self.sf_vec_size}.") if self.lane_idx is None or self.warp_idx is None or self.pair_exchange_barrier is None: - raise ValueError( - "The paired direction needs lane_idx, warp_idx and pair_exchange_barrier." - ) + raise ValueError("The paired direction needs lane_idx, warp_idx and pair_exchange_barrier.") elif not isinstance(self.quant_kind, CombineFormat) and self.sf_vec_size != 16: # A wider orthodox block cannot be reduced inside one thread: see the paired direction. raise ValueError(f"Orthodox sf_vec {self.sf_vec_size} needs the paired direction.") @@ -159,20 +139,14 @@ def _data_rcp_limit(self) -> float: @property def _block_share_per_thread(self) -> int: """How many of a block's elements this thread holds.""" - return ( - self.sf_vec_size // 2 - if self.sf_vec_direction == "regs_in_pair_threads" - else self.sf_vec_size - ) + return self.sf_vec_size // 2 if self.sf_vec_direction == "regs_in_pair_threads" else self.sf_vec_size # -- dispatch ------------------------------------------------------------- @cute.jit def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermediate=None): if cutlass.const_expr(cute.size(prequant_reg) % self._block_share_per_thread != 0): - raise ValueError( - "prequant_reg size must be divisible by this thread's share of a block." - ) + raise ValueError("prequant_reg size must be divisible by this thread's share of a block.") # Combine quantizes fc2's bf16 reorder regs; orthodox the fp32 swiglu. expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 if cutlass.const_expr(prequant_reg.element_type is not expected_in): @@ -181,9 +155,7 @@ def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None, smem_intermedi ) needs_smem = cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads") if cutlass.const_expr(needs_smem != (smem_intermediate is not None)): - raise ValueError( - "smem_intermediate must be supplied for the paired direction and only for it." - ) + raise ValueError("smem_intermediate must be supplied for the paired direction and only for it.") if cutlass.const_expr(not self._is_combine): if cutlass.const_expr(self.sf_vec_direction == "regs_in_pair_threads"): return self.mx_quant_regs_in_pair_threads_impl(prequant_reg, smem_intermediate) @@ -255,14 +227,10 @@ def mx_quant_regs_in_pair_threads_impl( if cutlass.const_expr(smem_intermediate.stride[0] != 1): # Thread-major would still be correct but would silently split the exchange into # per-block scalar accesses. - raise ValueError( - "The paired amax exchange must be block-major to stay one vector access." - ) + raise ValueError("The paired amax exchange must be block-major to stay one vector access.") exchange_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Float32, - num_bits_per_copy=n_blocks * cutlass.Float32.width, + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=n_blocks * cutlass.Float32.width ) # Which warps pair up is epilogue knowledge: warp w owns one contiguous run of # intermediate outputs, so a block that is twice that run wide joins adjacent warps. @@ -273,20 +241,12 @@ def mx_quant_regs_in_pair_threads_impl( half_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) for vec_block_idx in cutlass.range_constexpr(n_blocks): half_amax[vec_block_idx] = self._amax_thread_fp32(in_blocks[None, vec_block_idx]) - cute.copy( - exchange_atom, - cute.coalesce(half_amax), - cute.coalesce(smem_intermediate[None, own_thread]), - ) + cute.copy(exchange_atom, cute.coalesce(half_amax), cute.coalesce(smem_intermediate[None, own_thread])) self.pair_exchange_barrier.arrive_and_wait() partner_amax = cute.make_rmem_tensor((n_blocks,), cutlass.Float32) - cute.copy( - exchange_atom, - cute.coalesce(smem_intermediate[None, partner_thread]), - cute.coalesce(partner_amax), - ) + cute.copy(exchange_atom, cute.coalesce(smem_intermediate[None, partner_thread]), cute.coalesce(partner_amax)) data_blocks = [] sf_values = [] @@ -294,17 +254,13 @@ def mx_quant_regs_in_pair_threads_impl( block_amax = cute.arch.fmax(half_amax[vec_block_idx], partner_amax[vec_block_idx]) scale_e8m0, scale_f32 = self._e8m0(block_amax) sf_values.append(scale_e8m0) - data_blocks.append( - self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32)) - ) + data_blocks.append(self._scale_to_data_ssa(in_blocks[None, vec_block_idx], self._enc_mxfp8(scale_f32))) self._store_packed_blocks(data, data_blocks) sf.store(self._values_to_ssa(sf_values, self.scale_dtype)) return data, sf @cute.jit - def nvfp4_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def nvfp4_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -361,9 +317,7 @@ def nvfp4_combine_quant_threads_with_the_same_reg_impl( return data, selected_sf @cute.jit - def mxfp8_combine_quant_regs_in_thread_impl( - self, prequant_reg: cute.Tensor - ) -> Tuple[cute.Tensor, cute.Tensor]: + def mxfp8_combine_quant_regs_in_thread_impl(self, prequant_reg: cute.Tensor) -> Tuple[cute.Tensor, cute.Tensor]: # bf16 in -> e4m3/e5m2 data + per-32 e8m0. amax found on bf16 (packed). vec = self.sf_vec_size n_blocks = cute.size(prequant_reg) // vec @@ -439,8 +393,7 @@ def _concat_blocks_ssa(self, blocks, dtype: Type[cutlass.Numeric]) -> cute.Tenso @cute.jit def _values_to_ssa(self, values, dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: vec = vector.from_elements( - T.vector(len(values), dtype.mlir_type), - [values[i].ir_value() for i in range(len(values))], + T.vector(len(values), dtype.mlir_type), [values[i].ir_value() for i in range(len(values))] ) return cute.TensorSSA(vec, (len(values),), dtype) @@ -484,9 +437,7 @@ def max_abs(lhs: cutlass.Float32, rhs: cutlass.Float32) -> cutlass.Float32: acc = block[0] for elem_idx in cutlass.range_constexpr(1, cute.size(block)): acc = max_abs(acc, block[elem_idx]) - mag_bits = cutlass.Int32( - llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value()) - ) & cutlass.Int32(0x7FFFFFFF) + mag_bits = cutlass.Int32(llvm.bitcast(T.i32(), cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) @cute.jit @@ -592,15 +543,12 @@ class _TmemTranspose16x32Core: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((_TmemTranspose16x32Core._TmemRowStride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = None) -> None: # The whole transpose is built from 32-bit element atoms; _io_dtype @@ -631,21 +579,11 @@ def __init__(self, tmem_ptr, region: int, reg_tensor: Optional[cute.Tensor] = No self._tmem_dst_top = cute.make_tensor(dst_ptr, self._tmem_layout(16, 16)) self._tmem_dst_bot = cute.make_tensor(dst_ptr + half_lane_off, self._tmem_layout(16, 16)) - self._atom_ld16x64 = cute.make_copy_atom( - tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_st16x128 = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype - ) - self._atom_st32x32 = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype - ) - self._atom_ld16x256 = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype - ) - self._atom_ld16x128 = cute.make_copy_atom( - tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype - ) + self._atom_ld16x64 = cute.make_copy_atom(tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_st16x128 = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), self._io_dtype) + self._atom_st32x32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), self._io_dtype) + self._atom_ld16x256 = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), self._io_dtype) + self._atom_ld16x128 = cute.make_copy_atom(tcgen05.Ld16x128bOp(tcgen05.Repetition.x4), self._io_dtype) self._src_regs = cute.make_rmem_tensor((16,), self._io_dtype) # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) @@ -700,18 +638,10 @@ def r2_store(self) -> None: # -- R3 ------------------------------------------------------------------ def r3_load_top(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r3_load_bot(self) -> None: - cute.copy( - self._atom_ld16x256, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x256, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r3_perm(self) -> None: for r in range(16): @@ -723,18 +653,10 @@ def r3_store(self) -> None: # -- R4 ------------------------------------------------------------------ def r4_load_top(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_top, - self._rmem_copy_view(self._src_regs, 8, offset=0), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_top, self._rmem_copy_view(self._src_regs, 8, offset=0)) def r4_load_bot(self) -> None: - cute.copy( - self._atom_ld16x128, - self._tmem_dst_bot, - self._rmem_copy_view(self._src_regs, 8, offset=8), - ) + cute.copy(self._atom_ld16x128, self._tmem_dst_bot, self._rmem_copy_view(self._src_regs, 8, offset=8)) def r4_perm(self) -> None: for r in range(16): @@ -794,10 +716,7 @@ class TmemTranspose32x32Inplace: """ def __init__( - self, - tmem_ptr, - reg_tensor_top: Optional[cute.Tensor] = None, - reg_tensor_bot: Optional[cute.Tensor] = None, + self, tmem_ptr, reg_tensor_top: Optional[cute.Tensor] = None, reg_tensor_bot: Optional[cute.Tensor] = None ) -> None: if (reg_tensor_top is None) != (reg_tensor_bot is None): raise ValueError( @@ -890,26 +809,19 @@ class TmemTranspose32x64B16Movm: @staticmethod def _tmem_layout(num_lanes: int, num_cols: int) -> cute.Layout: return cute.make_layout( - (((num_lanes, num_cols), 1),), - stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),), + (((num_lanes, num_cols), 1),), stride=(((TmemTranspose32x64B16Movm._tmem_row_stride, 1), 0),) ) @staticmethod def _rmem_copy_view(rmem: cute.Tensor, num_regs: int, offset: int = 0) -> cute.Tensor: - return cute.make_tensor( - rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),)) - ) + return cute.make_tensor(rmem.iterator + offset, cute.make_layout((((num_regs,), 1),), stride=(((1,), 0),))) @cute.jit def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: if cutlass.const_expr(reg_tensor.element_type is not cutlass.BFloat16): - raise TypeError( - f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}." - ) + raise TypeError(f"{type(self).__name__} expects BF16 input after f2fp, got {reg_tensor.element_type}.") if cutlass.const_expr(cute.size(reg_tensor) != 64): - raise ValueError( - f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}." - ) + raise ValueError(f"{type(self).__name__} expects 64 BF16 elements, got {cute.size(reg_tensor)}.") movm_words = movmatrix_b16(cute.recast_tensor(reg_tensor, cutlass.Int32)) self._store_words = cute.make_rmem_tensor(movm_words.layout, movm_words.element_type) @@ -920,20 +832,14 @@ def __init__(self, tmem_ptr, reg_tensor: cute.Tensor) -> None: self._tmem_top = cute.make_tensor(tmem_ptr, self._tmem_layout(16, 32)) self._tmem_bottom = cute.make_tensor(tmem_ptr + half_lane_offset, self._tmem_layout(16, 32)) self._tmem_full = cute.make_tensor(tmem_ptr, self._tmem_layout(32, 32)) - self._store_atom = cute.make_copy_atom( - tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32 - ) - self._load_atom = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32 - ) + self._store_atom = cute.make_copy_atom(tcgen05.St16x128bOp(tcgen05.Repetition.x8), cutlass.Float32) + self._load_atom = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) @cute.jit def __call__(self) -> cute.Tensor: movm_words_f32 = cute.recast_tensor(self._store_words, cutlass.Float32) cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16), self._tmem_top) - cute.copy( - self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom - ) + cute.copy(self._store_atom, self._rmem_copy_view(movm_words_f32, 16, offset=16), self._tmem_bottom) output_words = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(self._load_atom, self._tmem_full, self._rmem_copy_view(output_words, 32)) @@ -948,8 +854,7 @@ class GatedActEpilogueArgs: fc2_alpha: Optional[cute.Tensor] fc1_norm_const: Optional[cute.Tensor] # ----------------------------------- - # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the - # score mul. + # MoE domain (token, topk), deepgemm graph only? for transformer graph, we want reduce kernel to perform the score mul. topk_scores: Optional[cute.Tensor] @@ -1013,8 +918,13 @@ def impl_desc_require(cls) -> dict[str, object]: "fc2_tma_stages": OptionalRequirement(int), "reduce_topk_in_kernel": OptionalRequirement(bool), "token_back_push_data": OptionalRequirement(bool), + "token_back_ready_granularity": OptionalRequirement(str), } + def _make_fc2_epilogue(self, *args) -> "SwapABFc2Epilogue": + """Device-side FC2 epilogue factory; variants substitute their own ``SwapABFc2Epilogue`` subclass.""" + return SwapABFc2Epilogue(self, *args) + @staticmethod def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Optional[float]]: """Resolve the SiTU (Kimi K3) betas from the ProblemDesc; ``(None, None)`` means SwiGLU. @@ -1027,19 +937,13 @@ def _resolve_situ_betas(problem_desc: ProblemDesc) -> Tuple[Optional[float], Opt beta = problem_desc.get("situ_beta") linear_beta = problem_desc.get("situ_linear_beta") if (beta is None) != (linear_beta is None): - raise ValueError( - f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}." - ) + raise ValueError(f"situ_beta and situ_linear_beta must be set together, got {beta} and {linear_beta}.") if beta is None: return None, None if beta <= 0 or linear_beta <= 0: - raise ValueError( - f"SiTU beta parameters must be positive, got {beta} and {linear_beta}." - ) + raise ValueError(f"SiTU beta parameters must be positive, got {beta} and {linear_beta}.") if problem_desc["gate_up_clamp"] is not None: - raise ValueError( - "SiTU does not support gate_up_clamp; the two activation variants are exclusive." - ) + raise ValueError("SiTU does not support gate_up_clamp; the two activation variants are exclusive.") if problem_desc.get("swiglu_alpha") is not None: raise ValueError("SiTU does not support SwiGLUBias constants.") return float(beta), float(linear_beta) @@ -1061,6 +965,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.swiglu_alpha = float(self.swiglu_alpha) self.swiglu_beta = float(self.swiglu_beta) self.situ_beta, self.situ_linear_beta = self._resolve_situ_betas(problem_desc) + # The Blackwell epilogue accepts gate_activation="quick_geglu"; this one has not been + # ported yet. Refuse rather than silently computing SwiGLU for a quick_geglu descriptor. + if problem_desc.get("gate_activation") not in (None, "swiglu"): + raise NotImplementedError( + f"Rubin epilogue does not support gate_activation={problem_desc.get('gate_activation')!r} yet." + ) self.mma_tiler_mnk = impl_desc["mma_tiler_mnk"] self.cluster_shape_mn = impl_desc["cluster_shape_mn"] self.use_2cta_instrs = impl_desc["use_2cta_instrs"] @@ -1068,21 +978,18 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.communication_enabled = impl_desc["communication_enabled"] self.fc1_epi_flag_batch = impl_desc["fc1_epi_flag_batch"] self.fc2_epi_flag_batch = impl_desc["fc2_epi_flag_batch"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") if self.communication_enabled: for field_name in ("reduce_topk_in_kernel", "token_back_push_data"): if field_name not in impl_desc: - raise KeyError( - f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}." - ) + raise KeyError(f"Communication-enabled Epilogue requires ImplDesc field {field_name!r}.") self.reduce_topk_in_kernel = impl_desc.get("reduce_topk_in_kernel", False) self.token_back_push_data = impl_desc.get("token_back_push_data", False) - if not self.communication_enabled and ( - self.reduce_topk_in_kernel or self.token_back_push_data - ): - raise ValueError( - "A communication-disabled Epilogue cannot enable communication policies." - ) + if not self.communication_enabled and (self.reduce_topk_in_kernel or self.token_back_push_data): + raise ValueError("A communication-disabled Epilogue cannot enable communication policies.") # FC1 emits what FC2 consumes as its activation, in the kind's own scale format. self.fc1_output_dtype = self.quant_kind.activation_dtype @@ -1104,12 +1011,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: not 1 <= self.fc1_epi_flag_batch <= self._EpilogueWarpCnt or not 1 <= self.fc2_epi_flag_batch <= self._EpilogueWarpCnt ): - raise ValueError( - f"Asynchronous epilogue flag batch sizes must be in [1, {self._EpilogueWarpCnt}]." - ) - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] - ) + raise ValueError(f"Asynchronous epilogue flag batch sizes must be in [1, {self._EpilogueWarpCnt}].") + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * self.cluster_shape_mn[0] atom_thr_size = 2 if self.use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize @@ -1127,21 +1030,12 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) - if ( - isinstance(self.hidden_size, int) - and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0 - ): + if isinstance(self.hidden_size, int) and self.hidden_size % (self.cta_tile_m * self.cluster_shape_mn[0]) == 0: self.fc2_hidden_needs_predicate: bool = False else: self.fc2_hidden_needs_predicate: bool = True @@ -1155,10 +1049,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # One staging stage per token subtile, each an (epi_tile_n, epi_tile_m // 2) quantized tile. self.fc1_staging_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc1IntermediateDownTileSize - * self.fc1_output_dtype.width - // 8 + self._EpilogueTokenTileSize * self._EpilogueFc1IntermediateDownTileSize * self.fc1_output_dtype.width // 8 ) self.fc1_staging_bytes = self.subtile_cnt * self.fc1_staging_stage_bytes # The amax exchange plane is (token chunk, epilogue thread), block-major so each thread's @@ -1167,21 +1058,14 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: # bytes -- see the lifetime argument in fc1_quant. self.fc1_amax_token_chunks = self._EpilogueTokenTileSize // 32 self.fc1_amax_slot_count = ( - self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 - if self.needs_pair_amax_exchange - else 0 + self.fc1_amax_token_chunks * self._EpilogueWarpCnt * 32 if self.needs_pair_amax_exchange else 0 ) if self.fc1_amax_slot_count * 4 > self.fc1_staging_stage_bytes: raise ValueError("The paired amax exchange does not fit in one fc1 staging stage.") requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") - if ( - requested_fc2_tma_stages is not None - and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt - ): - raise ValueError( - f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}." - ) + if requested_fc2_tma_stages is not None and not 1 <= requested_fc2_tma_stages <= self.subtile_cnt: + raise ValueError(f"fc2_tma_stages must be in [1, {self.subtile_cnt}], got {requested_fc2_tma_stages}.") if self.fc2_use_bulk: single_stage_region = self._make_fc2_single_stage_region() if self.fc2_use_ublk and single_stage_region.nbytes % 16 != 0: @@ -1191,20 +1075,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: else: bf16_baseline_stages = min(2, self.subtile_cnt) bf16_stage_bytes = ( - self._EpilogueTokenTileSize - * self._EpilogueFc2HiddenTileSize - * cutlass.BFloat16.width - // 8 - ) - available_staging_bytes = max( - self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes + self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 ) - self.fc2_tma_stages = min( - self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes - ) - self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region( - self.fc2_tma_stages - ) + available_staging_bytes = max(self.fc1_staging_bytes, bf16_baseline_stages * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec: Optional[SmemRegion] = self._make_fc2_staging_region(self.fc2_tma_stages) else: self.fc2_tma_stages = 0 self.fc2_staging_spec = None @@ -1218,9 +1093,7 @@ def epilogue_sync_barrier(cls) -> pipeline.NamedBarrier: outside the subtile loop, FC1's amax exchange and FC2's bulk-store handshake belong to different work tiles. """ - return pipeline.NamedBarrier( - barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt - ) + return pipeline.NamedBarrier(barrier_id=cls._EpilogueSyncWaitBarId, num_threads=32 * cls._EpilogueWarpCnt) def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Declare the epilogue scratch as one allocation shared by two exclusive lifetimes. @@ -1275,13 +1148,9 @@ def fc2_tma_staged_smem_spec(self, n_stages: int) -> Tuple[Tuple, Tuple, Tuple[i raise ValueError(f"Unsupported FC2 TMA staging dtype {wire_dtype}.") return shape, stride, swizzle - def fc2_tma_staged_smem_layout( - self, n_stages: int, without_stage_mode: bool = False - ) -> cute.ComposedLayout: + def fc2_tma_staged_smem_layout(self, n_stages: int, without_stage_mode: bool = False) -> cute.ComposedLayout: shape, stride, swizzle = self.fc2_tma_staged_smem_spec(n_stages) - layout = cute.make_composed_layout( - cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride) - ) + layout = cute.make_composed_layout(cute.make_swizzle(*swizzle), 0, cute.make_layout(shape, stride=stride)) if without_stage_mode: return cute.select(layout, mode=[0, 1]) return layout @@ -1322,10 +1191,7 @@ def _make_fc2_staging_region(self, stage_count: int) -> SmemRegion: kind="tensor", dtype=single_stage_region.dtype, shape=(*single_stage_region.shape, stage_count), - stride=( - *single_stage_region.stride, - single_stage_region.cosize if stage_count > 1 else 0, - ), + stride=(*single_stage_region.stride, single_stage_region.cosize if stage_count > 1 else 0), swizzle=single_stage_region.swizzle, byte_alignment=single_stage_region.byte_alignment, ) @@ -1399,17 +1265,12 @@ def run( fc2_smem_tensor = smem_workspace.tensor(self.fc2_staging_region, smem_base) else: fc2_smem_tensor = None - if cutlass.const_expr( - self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None) - ): + if cutlass.const_expr(self.fc2_use_tma and (tma_atom_fc2_output is None or fc2_tma_output is None)): raise ValueError("FC2 TMA store requires a TMA atom and token-major output tensor.") if cutlass.const_expr( - self.communication_enabled - and (token_src_metadata is None or peer_rank_ptr_mapper is None) + self.communication_enabled and (token_src_metadata is None or peer_rank_ptr_mapper is None) ): - raise ValueError( - "Communication-enabled Epilogue requires token metadata and a peer pointer mapper." - ) + raise ValueError("Communication-enabled Epilogue requires token metadata and a peer pointer mapper.") if cutlass.const_expr(self.token_back_enabled and fc2_done_counter is None): raise ValueError("Token-back requires an FC2 done counter.") if cutlass.const_expr(self.token_back_push_sf and fc2_output_sf is None): @@ -1430,8 +1291,7 @@ def run( fc1_done_counter, optional_epi_args, ) - fc2_epi = SwapABFc2Epilogue( - self, + fc2_epi = self._make_fc2_epilogue( tidx, fc2_smem_tensor, tma_atom_fc2_output, @@ -1488,13 +1348,9 @@ def run( # Publish completion for the work tile snapshotted above. if cur_was_linear1: - flag_tracker = fc1_epi.signal_fc1_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, work_tile_info, flag_tracker) else: - flag_tracker = fc2_epi.signal_fc2_done( - prev_work_tile_info, work_tile_info, flag_tracker - ) + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, work_tile_info, flag_tracker) # Tail flush flag_tracker.fire() @@ -1504,9 +1360,7 @@ class _ImmutableAfterInit: def __setattr__(self, name, value): if self.__dict__.get("_frozen_", False): - raise AttributeError( - f"{type(self).__name__} is immutable after __init__ (cannot set {name!r})." - ) + raise AttributeError(f"{type(self).__name__} is immutable after __init__ (cannot set {name!r}).") object.__setattr__(self, name, value) def _freeze(self) -> None: @@ -1535,8 +1389,7 @@ def __init__( # spelled out here, so an 8-bit fc1 output picks up its own atom without a code change. staged_layout = base.fc1_staged_smem_layout(base.subtile_cnt) self.smem_tensor = cute.make_tensor( - cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), - staged_layout.outer, + cute.recast_ptr(staging_pointer, staged_layout.inner, dtype=base.fc1_output_dtype), staged_layout.outer ) # Kept unswizzled and untyped so fc1_quant can carve an fp32 amax-exchange view out of one # staging stage without going through the staging tensor's swizzle. @@ -1572,19 +1425,14 @@ def signal_fc1_done(self, work_tile_info, next_work_tile_info, flag_tracker): or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0 ) if cutlass.const_expr(needs_intermediate_guard): - in_bound = ( - work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize - < self.fc1_output.shape[1] - ) + in_bound = work_tile_info.tile_m_idx * self._EpilogueFc1IntermediateDownTileSize < self.fc1_output.shape[1] else: in_bound = True slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx flag_address = Int64(0) if in_bound: flag_address = (self.fc1_done_counter.iterator + slot).toint() - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc1_epi_flag_batch, flag_address) @cute.jit def __call__( @@ -1595,13 +1443,9 @@ def __call__( acc_consumer_state, ): # (tokens_this_expert, intermediate_down, 1) - real_fc1_output, _ = self.kernel_extension.get_gmem_tensor( - "c", self.fc1_output, work_tile_info - ) + real_fc1_output, _ = self.kernel_extension.get_gmem_tensor("c", self.fc1_output, work_tile_info) # (tokens_this_expert, intermediate_down, 1) - real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor( - "sfc", self.fc1_output_sf, work_tile_info - ) + real_fc1_output_sf, _ = self.kernel_extension.get_gmem_tensor("sfc", self.fc1_output_sf, work_tile_info) # subtile-irrevalent hoist out here. if cutlass.const_expr(self.optional_epi_args.fc1_alpha is not None): alpha_val = self.optional_epi_args.fc1_alpha[work_tile_info.expert_idx] @@ -1613,8 +1457,7 @@ def __call__( norm_const = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, - (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize), + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, self._EpilogueTokenTileSize) )[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) @@ -1664,9 +1507,7 @@ def run_subtile( # Step -1: preload topk scores. current_two_token_idices = ( - work_tile_info.tile_n_idx * self.cta_tile_n - + subtile_idx * self._EpilogueTokenTileSize - + self.lane_idx, + work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx @@ -1687,11 +1528,8 @@ def run_subtile( up_token_32_64 = cute.make_rmem_tensor((16,), cutlass.Float32) # Although hardcode is not right, but since the whole tmem transpose is too tricky, I have to hardcode... # (epi_tile_m, epi_tile_n) -> (warp_local_epi_tile_m, epi_tile_n) - # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, - # self.warp_idx), None] - tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, 0), None - ] + # tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, self.warp_idx), None] + tmem_subtile_tensor_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, 0), None] # (warp_local_epi_tile_m, epi_tile_n) -> (((16, 32), 1), (2, 2)) tmem_subtile_tensor_in_first_load_view = cute.logical_divide( cute.zipped_divide(tmem_subtile_tensor_per_warp, (16, 32)), ((16, 32), 1) @@ -1719,16 +1557,12 @@ def run_subtile( ) # Step 1: perform swiglu on the first part, interleave with the second's 32x32 tmem transpose. - token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( - gate_token_0_32, up_token_0_32, alpha_val - ) + token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp(gate_token_0_32, up_token_0_32, alpha_val) # gate_token_32_64 / up_token_32_64 are already in the transpose input # distribution documented by TmemTranspose16x32. token_32_64_tmem_trans = TmemTranspose32x32Inplace( - tmem_subtile_tensor.iterator, - reg_tensor_top=gate_token_32_64, - reg_tensor_bot=up_token_32_64, + tmem_subtile_tensor.iterator, reg_tensor_top=gate_token_32_64, reg_tensor_bot=up_token_32_64 ) # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp @@ -1766,8 +1600,7 @@ def run_subtile( ] # (cta_token, cta_intermediate_down) -> (token_64, intermediate_64) fc1_gmem_subtile_view = cute.flat_divide( - fc1_gmem_cta_view, - (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize), + fc1_gmem_cta_view, (self._EpilogueTokenTileSize, self._EpilogueFc1IntermediateDownTileSize) )[None, None, subtile_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( self.fc1_tma_atom, @@ -1803,21 +1636,15 @@ def alpha_swiglu_clamp( # raise fires during compilation rather than emitting a runtime branch. for _name, _t in (("gate_rmem", gate_rmem), ("up_rmem", up_rmem)): if cutlass.const_expr(_t.element_type is not cutlass.Float32): - raise TypeError( - f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}" - ) + raise TypeError(f"alpha_swiglu_clamp: {_name} must be Float32, got {_t.element_type}") if cutlass.const_expr(_t.memspace != AddressSpace.rmem): raise ValueError( f"alpha_swiglu_clamp: {_name} must be a register (rmem) tensor, got address space {_t.memspace}" ) if cutlass.const_expr(cute.rank(_t) != 1): - raise ValueError( - f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} must be 1D, got rank {cute.rank(_t)}") if cutlass.const_expr(cute.size(_t) % 2 != 0): - raise ValueError( - f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}" - ) + raise ValueError(f"alpha_swiglu_clamp: {_name} element count must be even, got {cute.size(_t)}") if cutlass.const_expr(cute.size(gate_rmem) != cute.size(up_rmem)): raise ValueError( "alpha_swiglu_clamp: gate_rmem and up_rmem must have equal size, got " @@ -1848,10 +1675,7 @@ def alpha_swiglu_clamp( activation_alpha = 1.0 if self.swiglu_alpha is None else self.swiglu_alpha neg_activation_alpha_log2e = -activation_alpha * log2_e - neg_log2e_pair = ( - cutlass.Float32(neg_activation_alpha_log2e), - cutlass.Float32(neg_activation_alpha_log2e), - ) + neg_log2e_pair = (cutlass.Float32(neg_activation_alpha_log2e), cutlass.Float32(neg_activation_alpha_log2e)) one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) if cutlass.const_expr(self.swiglu_beta is not None): beta_pair = (cutlass.Float32(self.swiglu_beta), cutlass.Float32(self.swiglu_beta)) @@ -1930,6 +1754,7 @@ def _sigmoid(p0, p1): sigmoid_pair = _sigmoid(g0, g1) if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = (up + swiglu_beta) * gate * sigmoid(swiglu_alpha * gate) # The SwiGLUBias offset follows the up clamp. if cutlass.const_expr(self.swiglu_beta is not None): u0, u1 = cute.arch.add_packed_f32x2((u0, u1), beta_pair) @@ -1975,9 +1800,7 @@ def _sigmoid(p0, p1): def fc1_quant( self, work_tile_info: SwapAbFc12WorkTileInfo, - two_token: Tuple[ - cute.Tensor, cute.Tensor - ], # two rmem tensor, each fp32 @ (token_1, intermediate_16) + two_token: Tuple[cute.Tensor, cute.Tensor], # two rmem tensor, each fp32 @ (token_1, intermediate_16) topk_scores: Optional[Tuple[cutlass.Float32, cutlass.Float32]], norm_const: Optional[cutlass.Float32], intermediate_output_size: cutlass.Int32, @@ -2011,8 +1834,7 @@ def fc1_quant( cute.recast_ptr(self.staging_pointer, dtype=cutlass.Float32) + subtile_idx * cutlass.Int32(self.fc1_staging_stage_bytes // 4), cute.make_layout( - (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), - stride=(1, self.fc1_amax_token_chunks), + (self.fc1_amax_token_chunks, self._EpilogueWarpCnt * 32), stride=(1, self.fc1_amax_token_chunks) ), ), 16, @@ -2039,13 +1861,8 @@ def fc1_quant( stores_scale_factor = self.warp_idx % 2 == 0 else: stores_scale_factor = True - subtile_token_start = ( - work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize - ) - token_idx_pair = ( - subtile_token_start + self.lane_idx, - subtile_token_start + self.lane_idx + 32, - ) + subtile_token_start = work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + token_idx_pair = (subtile_token_start + self.lane_idx, subtile_token_start + self.lane_idx + 32) # This subtile's (token, intermediate) shared output stage, tiled into (1, 16) blocks so # each thread's cells slice out directly (zipped_divide + slice; avoids the ambiguous @@ -2102,9 +1919,7 @@ def fc1_quant( # 4) STS the quantized values into this subtile's shared output stage. # ((1, 16), (token_tile_size, warp_cnt)) -> (16) smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] - cute.copy( - store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row) - ) + cute.copy(store_atom, cute.coalesce(data_by_token[None, half]), cute.coalesce(smem_thread_row)) @dataclasses.dataclass(frozen=True) @@ -2199,22 +2014,24 @@ def signal_fc2_done(self, work_tile_info, next_work_tile_info, flag_tracker): publish: cutlass.Constexpr = self.token_back_enabled flag_address = Int64(0) if cutlass.const_expr(publish): - flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() + if cutlass.const_expr(self.token_back_ready_granularity == "token_tile"): # nosec B105 + # The scheduler maps every fallback CTA into the preferred + # logical cluster. N=1 keeps this slot common to all hidden + # shards; padded hidden CTAs must also publish their completion. + slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + flag_address = (self.fc2_done_counter.iterator + slot).toint() + else: + flag_address = (self.fc2_done_counter.iterator + work_tile_info.expert_idx).toint() no_fire: cutlass.Constexpr = not publish - return flag_tracker.accumulate( - next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire - ) + return flag_tracker.accumulate(next_work_tile_info.phase, self.fc2_epi_flag_batch, flag_address, no_fire) @cute.jit def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2OutputRouter": task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row - + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n) ) hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32(self.cta_tile_m) - valid_hidden_this_cta_tile = ( - cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile - ) + valid_hidden_this_cta_tile = cutlass.Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile if valid_hidden_this_cta_tile < 0: valid_hidden_this_cta_tile = 0 if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: @@ -2223,9 +2040,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out metadata = None peer_rank_ptr_mapper = None data_token_base = task_tile_data_row_start - if cutlass.const_expr( - self.token_src_metadata is not None and not self.token_back_push_data - ): + if cutlass.const_expr(self.token_src_metadata is not None and not self.token_back_push_data): metadata = cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata) peer_rank_ptr_mapper = self.peer_rank_ptr_mapper data_token_base = None @@ -2233,10 +2048,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out if cutlass.const_expr(self.combine_format.is_quantized): base_outputs = (self.fc2_output, self.fc2_output_sf) token_bases = (data_token_base, task_tile_data_row_start) - output_mappings = ( - self.process_pipeline.store_out_mapping, - self.process_pipeline.sf_store_out_mapping, - ) + output_mappings = (self.process_pipeline.store_out_mapping, self.process_pipeline.sf_store_out_mapping) else: base_outputs = self.fc2_output token_bases = data_token_base @@ -2258,11 +2070,7 @@ def _make_output_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> "Fc2Out @cute.jit def __call__( - self, - work_tile_info: SwapAbFc12WorkTileInfo, - tmem_acc_tensor: cute.Tensor, - acc_pipeline, - acc_consumer_state, + self, work_tile_info: SwapAbFc12WorkTileInfo, tmem_acc_tensor: cute.Tensor, acc_pipeline, acc_consumer_state ): # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): @@ -2308,14 +2116,17 @@ def run_subtile( alpha_val: Optional[cutlass.Float32], ): process_pipeline = self.process_pipeline + iket.range_push("fc2.tmem_load") loaded = process_pipeline.tmem_acc_load(tmem_subtile_tensor=tmem_subtile_tensor, epi=self) + iket.range_pop() + iket.range_push("fc2.f2fp_reorder") casted = process_pipeline.f2fp(*loaded, alpha_val=alpha_val) # reorder returns a bare RMEM fragment in the store's expected pre-store # distribution; reorder + store are paired 1:1 inside the pipeline. - pre_store = process_pipeline.post_f2fp_reorder( - casted=casted, tmem_subtile_view=tmem_subtile_tensor - ) + pre_store = process_pipeline.post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_tensor) + iket.range_pop() + iket.range_push("fc2.store") process_pipeline.store_function( epi=self, subtile=pre_store, @@ -2324,6 +2135,7 @@ def run_subtile( subtile_idx=subtile_idx, fc2_output_router=fc2_output_router, ) + iket.range_pop() @dataclasses.dataclass(frozen=True) @@ -2344,9 +2156,7 @@ class Fc2OutputRouter: epi_tid: cutlass.Int32 combine_format: CombineFormat # After metadata prefetch - dst_ptrs: Optional[cute.Tensor] = ( - None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. - ) + dst_ptrs: Optional[cute.Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) @property @@ -2368,11 +2178,7 @@ def sf_token_base(self) -> Optional[cutlass.Int32]: @property def data_mapping(self) -> FunctionMapping: - return ( - self.output_mappings[0] - if isinstance(self.output_mappings, tuple) - else self.output_mappings - ) + return self.output_mappings[0] if isinstance(self.output_mappings, tuple) else self.output_mappings @property def sf_mapping(self) -> Optional[FunctionMapping]: @@ -2380,9 +2186,7 @@ def sf_mapping(self) -> Optional[FunctionMapping]: def __post_init__(self) -> None: if (self.metadata is None) == (self.data_token_base is None): - raise ValueError( - "Fc2OutputRouter requires exactly one of metadata or a (data) token base." - ) + raise ValueError("Fc2OutputRouter requires exactly one of metadata or a (data) token base.") if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): raise ValueError("Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set.") if self.reduce_topk_in_epilogue and self.metadata is None: @@ -2401,8 +2205,7 @@ def prefetch(self) -> "Fc2OutputRouter": valid = cute.make_rmem_tensor((copy_iters,), cutlass.Int32) dst_ptrs = cute.make_rmem_tensor((copy_iters,), cutlass.Int64) - # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + - # strength_reduce) + # Compiler should be able to optimize the same token_copy_group's offset add. (Fundamental cse + strength_reduce) # We should check the SASS to ensure this happens. for iter_idx in cutlass.range_constexpr(copy_iters): coord = self.data_mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx) @@ -2421,14 +2224,11 @@ def prefetch(self) -> "Fc2OutputRouter": dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile # Int64 token coord: dst_tokens*K*H overflows int32 once # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). - dst_ptrs[iter_idx] = self.data_output[ - Int64(dst_tokens), None, dst_hidden - ].iterator.toint() + dst_ptrs[iter_idx] = self.data_output[Int64(dst_tokens), None, dst_hidden].iterator.toint() else: md = TokenSrcMetadata.load( - self.metadata.iterator.toint() - + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) ) dst_rank = md.src_rank dst_token = md.src_token @@ -2440,9 +2240,7 @@ def prefetch(self) -> "Fc2OutputRouter": # Int64 token coord: domain_offset on (token, topk, hidden) # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map_pointer( - cute.domain_offset( - (Int64(dst_token), dst_topk, dst_hidden), self.data_output - ).iterator, + cute.domain_offset((Int64(dst_token), dst_topk, dst_hidden), self.data_output).iterator, dst_rank, byte_alignment=32, ).toint() @@ -2450,9 +2248,7 @@ def prefetch(self) -> "Fc2OutputRouter": return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid) @cute.jit - def get_data_dst( - self, iter_idx: Union[int, cutlass.Int32] - ) -> Tuple[cute.Pointer, cutlass.Int32]: + def get_data_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, cutlass.Int32]: """Per-issue DATA destination: gmem pointer + validity predicate. The router owns ``data_output`` so the caller never re-assembles a @@ -2484,9 +2280,7 @@ def get_data_dst( # comm: read the pointer / validity prefetched by prefetch(). addr = self.dst_ptrs[iter_idx] pred = self.valid[iter_idx] - ptr = cute.make_ptr( - self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32 - ) + ptr = cute.make_ptr(self.data_output.element_type, addr, AddressSpace.gmem, assumed_align=32) return ptr, pred @cute.jit @@ -2513,9 +2307,7 @@ def get_sf_dst(self, iter_idx: Union[int, cutlass.Int32]) -> Tuple[cute.Pointer, addr = self.sf_output[Int64(sf_row), None, sf_hidden].iterator.toint() # Per-block scale offsets are element-granular; claim the scale dtype's # natural element alignment (e8m0 1 B / bf16 2 B). - sf_ptr = cute.make_ptr( - self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4 - ) + sf_ptr = cute.make_ptr(self.sf_output.element_type, addr, AddressSpace.gmem, assumed_align=4) return sf_ptr, pred @@ -2539,16 +2331,13 @@ def make_fc2_stg_cta_store_out_mapping( }, ) store_out_mapping = FunctionMapping( - domain=CoordinateSpace( - ("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32) - ), + domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, stgs_per_hidden32 * cta_token_tile_size // 32)), codomain=CoordinateSpace( ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { "token_in_cta_tile": epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, - "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg - + epi_tid // 32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * elems_per_stg + epi_tid // 32 * 32, }, ) sf_store_out_mapping = None @@ -2562,8 +2351,7 @@ def stg_sf_mapping(epi_tid, iter_idx): sf_store_out_mapping = FunctionMapping( domain=CoordinateSpace(("epi_tid", "iter_idx"), (128, cta_token_tile_size // 32)), codomain=CoordinateSpace( - ("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size), + ("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size) ), function=stg_sf_mapping, ) @@ -2593,10 +2381,7 @@ def make_fc2_redg_cta_store_out_mapping( + (epi_tid % 32) // 4 ), "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 - + (epi_tid % 4) * 4 - + (((elem_idx // 4) % 4) // 2) * 16 - + elem_idx % 4 + (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (((elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4 ), }, ) @@ -2616,17 +2401,13 @@ def make_fc2_redg_cta_store_out_mapping( + ((iter_idx % 4) % 2) * 8 + (epi_tid % 32) // 4 ), - "hidden_in_cta_tile": ( - (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16 - ), + "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), }, ) return store_out_mapping, None, fundamental_mapping -def make_fc2_ublk_store_out_mapping( - combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int -): +def make_fc2_ublk_store_out_mapping(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): assert cta_hidden_tile_size == 128 assert cta_token_tile_size % 64 == 0 assert combine_format.act_dtype.width in (8, 16), "fc2 UBLK wire dtype must be fp8 or bf16." @@ -2652,9 +2433,7 @@ def make_fc2_ublk_store_out_mapping( ("token_in_cta_tile", "hidden_in_cta_tile"), (max_token_cta_tile, cta_hidden_tile_size) ), function=lambda epi_tid, iter_idx: { - "token_in_cta_tile": ( - iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16 - ), + "token_in_cta_tile": (iter_idx * 128 + ((epi_tid % 32) // 16) * 64 + (epi_tid // 32) * 16 + epi_tid % 16), "hidden_in_cta_tile": 0, }, ) @@ -2687,9 +2466,7 @@ def fc2_f2fp(*tensors, alpha_val: Optional[cutlass.Float32] = None, **_) -> cute elems_processed = 0 for t in tensors: current_tensor_size = cute.size(t) - dst = cute.make_tensor( - converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,)) - ) + dst = cute.make_tensor(converted_acc.iterator + elems_processed, cute.make_layout((current_tensor_size,))) if cutlass.const_expr(alpha_val is None): dst.store(t.load().to(reorder_dtype)) else: @@ -2719,9 +2496,7 @@ def fc2_stg_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, **_): ptr = tmem_subtile_tensor.iterator half_lane_offset = 16 * TmemTranspose32x64B16Movm._tmem_row_stride top_view = cute.make_tensor(ptr, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) - bottom_view = cute.make_tensor( - ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64) - ) + bottom_view = cute.make_tensor(ptr + half_lane_offset, TmemTranspose32x64B16Movm._tmem_layout(16, 64)) top = cute.make_rmem_tensor((32,), cutlass.Float32) bottom = cute.make_rmem_tensor((32,), cutlass.Float32) cute.copy(atom_ld16x256, top_view, TmemTranspose32x64B16Movm._rmem_copy_view(top, 32)) @@ -2734,13 +2509,9 @@ def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): # UBLK consumes a warp-local 32-hidden x 64-token slice. The caller passes # the CTA-level 128-hidden x 64-token subtile view, so select this epi # warp's hidden block before issuing LDTM.x64. - tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[ - (None, epi.warp_idx), None - ] + tmem_subtile_per_warp = cute.logical_divide(tmem_subtile_tensor, (32, None))[(None, epi.warp_idx), None] raw_regs = cute.make_rmem_tensor((64,), cutlass.Float32) - atom_ld32x32_x64 = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32 - ) + atom_ld32x32_x64 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x64), cutlass.Float32) cute.copy( atom_ld32x32_x64, wrap_into_copy_standard_layout(tmem_subtile_per_warp), @@ -2765,20 +2536,14 @@ def fc2_stg_post_f2fp_reorder( def fc2_redg_post_f2fp_reorder(*, casted: cute.Tensor, tmem_subtile_view: cute.Tensor, **_): # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 natural = fc2_stg_post_f2fp_reorder(casted=casted, tmem_subtile_view=tmem_subtile_view) - core_matrix_reorder_sttm_atom = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32 - ) - core_matrix_reorder_ldtm_atom = cute.make_copy_atom( - tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32 - ) + core_matrix_reorder_sttm_atom = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + core_matrix_reorder_ldtm_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition.x2), cutlass.Float32) # ((16, 2), token_32_group) token_groups = cute.logical_divide(cute.zipped_divide(natural, (32,)), (16, None)) out = cute.make_rmem_tensor(token_groups.shape, casted.dtype) out_as_i32 = cute.recast_tensor(out, cutlass.Float32) # (32, 64) - tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[ - None, None, 0, 0 - ] + tmem_warp = cute.flat_divide(tmem_subtile_view, (32, cute.size(tmem_subtile_view, 1)))[None, None, 0, 0] # (16, 16, 16dp_group, token_32_groups). Note, this tmem can provide 2x cols since the original is bf16. tmem_groups = cute.flat_divide(tmem_warp, (16, 16)) for group_idx in cutlass.range_constexpr(cute.size(token_groups, 1)): @@ -2818,9 +2583,7 @@ def fc2_stg_store_function( sf_regs = None stg_width_elems: cutlass.Constexpr[int] = min(32, 256 // data_subtile.element_type.width) stg_bits: cutlass.Constexpr[int] = stg_width_elems * data_subtile.element_type.width - copy_atom_vec = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits - ) + copy_atom_vec = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=stg_bits) elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): raise ValueError( @@ -2832,9 +2595,7 @@ def fc2_stg_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2851,11 +2612,7 @@ def fc2_stg_store_function( if pred != cutlass.Int32(0): src_i = cute.make_tensor(copy_src[None, local_iter].iterator, single_copy_layout) dst_i = cute.make_tensor(dst_ptr, single_copy_layout) - cute.copy( - copy_atom_vec, - cute.recast_tensor(src_i, cutlass.Int32), - cute.recast_tensor(dst_i, cutlass.Int32), - ) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), cute.recast_tensor(dst_i, cutlass.Int32)) @cute.jit @@ -2869,12 +2626,8 @@ def fc2_tma_store_function( fc2_output_router: Fc2OutputRouter, **_, ): - if cutlass.const_expr( - epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None - ): - raise ValueError( - "FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor." - ) + if cutlass.const_expr(epi.smem_tensor is None or epi.tma_atom_fc2_output is None or epi.fc2_tma_output is None): + raise ValueError("FC2 TMA store requires staged SMEM, a TMA atom, and a token-major output tensor.") if cutlass.const_expr(epi.combine_format.is_quantized): data_subtile, sf_regs = QuantImpl(epi.combine_format, "regs_in_thread")(subtile) @@ -2886,9 +2639,7 @@ def fc2_tma_store_function( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2902,15 +2653,11 @@ def fc2_tma_store_function( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 TMA store requires a contiguous two-token by 32-hidden register fragment." - ) + raise ValueError("FC2 TMA store requires a contiguous two-token by 32-hidden register fragment.") stage_idx = epilogue_iter_idx % cutlass.Int32(epi.fc2_tma_stages) smem_stage = epi.smem_tensor[None, None, stage_idx] - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -2924,9 +2671,7 @@ def fc2_tma_store_function( + work_tile_info.tile_n_idx * cutlass.Int32(epi.cta_tile_n) + subtile_idx * cutlass.Int32(epi._EpilogueTokenTileSize) ) // cutlass.Int32(epi._EpilogueTokenTileSize) - tiled_output = cute.flat_divide( - epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize) - ) + tiled_output = cute.flat_divide(epi.fc2_tma_output, (epi._EpilogueTokenTileSize, epi._EpilogueFc2HiddenTileSize)) gmem_subtile = tiled_output[None, None, token_tile_idx, work_tile_info.tile_m_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( epi.tma_atom_fc2_output, @@ -2965,9 +2710,7 @@ def fc2_ublk_store_function_impl( sf_scales_per_stg: cutlass.Constexpr[int] = 32 // epi.combine_format.scale_block token_groups_per_subtile: cutlass.Constexpr[int] = epi._EpilogueTokenTileSize // 32 for token_group in cutlass.range_constexpr(token_groups_per_subtile): - sf_iter = ( - cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group - ) + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32(token_groups_per_subtile) + token_group sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) if sf_pred != cutlass.Int32(0): sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((sf_scales_per_stg,))) @@ -2990,12 +2733,8 @@ def fc2_ublk_store_function_impl( or data_subtile.stride[0] != 1 or 32 % vector_elements != 0 ): - raise ValueError( - "FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment." - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128 - ) + raise ValueError("FC2 UBLK store requires a contiguous two-token by 32-hidden register fragment.") + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), data_subtile.element_type, num_bits_per_copy=128) tiler_mn = (64, 128) layout_copy_tv = cute.make_layout(((32, 4), (32, 2)), stride=((1, 2048), (64, 32))) tiled_store = cute.make_tiled_copy(store_atom, layout_copy_tv, tiler_mn) @@ -3005,9 +2744,7 @@ def fc2_ublk_store_function_impl( cute.copy(tiled_store, register_partition, smem_partition) else: if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): - raise ValueError( - "BF16 UBLK staging requires one register per token in the 64-token subtile." - ) + raise ValueError("BF16 UBLK staging requires one register per token in the 64-token subtile.") warp_hidden_base = cutlass.Int32(warp_idx * 32) for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] @@ -3063,18 +2800,14 @@ def fc2_redg_store_function( if pred != cutlass.Int32(0): bf16x4 = subtile_by_redg_issue[None, local_iter] packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) - red_add_relaxed_sys_v2_bf16x2( - dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1]) - ) + red_add_relaxed_sys_v2_bf16x2(dst_ptr, cutlass.Float32(packed_bf16x2[0]), cutlass.Float32(packed_bf16x2[1])) def make_fc2_stg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3090,10 +2823,8 @@ def make_fc2_stg_process_pipeline( def make_fc2_tma_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_stg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, @@ -3109,10 +2840,8 @@ def make_fc2_tma_process_pipeline( def make_fc2_redg_process_pipeline( *, combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int ) -> Fc2ProcessPipeline: - store_out_mapping, sf_store_out_mapping, fundamental_mapping = ( - make_fc2_redg_cta_store_out_mapping( - combine_format, cta_token_tile_size, cta_hidden_tile_size - ) + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_mapping( + combine_format, cta_token_tile_size, cta_hidden_tile_size ) return Fc2ProcessPipeline( tmem_acc_load=fc2_stg_tmem_acc_load, diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py new file mode 100644 index 000000000000..5364aa737f04 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_epilogue_dedup.py @@ -0,0 +1,349 @@ +"""Dedup FC2 epilogue: SMEM-staged subtiles pushed with one ``cp.async.bulk`` per token. ``dst_token_cnt == 1`` +rows go straight to the source rank's ``pre_reduced[src_token][my_rank]``, others into the local order buffer +``[T, R, chunks_per_row, K, chunk]`` for the combine warps. Top-k scores are applied at FC1 only.""" + +from __future__ import annotations + +import dataclasses +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int32, Int64 + +from .....api import ImplDesc, ProblemDesc +from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice +from .....communication.token_protocol import TokenSrcMetadata +from .....helpers.iket_compat import iket +from .....helpers.ptx_helpers import ( + cp_async_bulk_s2g, + cp_reduce_async_bulk_add_bf16_s2g, +) +from ....function_mapping import FunctionMapping +from ....schedulers import SwapAbFc12WorkTileInfo +from .block_scaled_swap_ab_fc12_epilogue import ( + GatedActEpilogueArgs, + SwapABFc2Epilogue, + SwapABGatedActEpilogue, + make_fc2_ublk_process_pipeline, +) + + +@dataclasses.dataclass(frozen=True) +class GatedActEpilogueArgsDedup(GatedActEpilogueArgs): + """Base epilogue args plus the pre-reduced (T, R, H) buffer and this rank's index (the kernel's host argument, + which selects the ``pre_reduced[src_token][local_rank]`` slot of single-copy rows).""" + + fc2_remote_output: Optional[cute.Tensor] = None + local_rank: Optional[Int32] = None + + +@dataclasses.dataclass(frozen=True) +class Fc2DedupBulkRouter: + """Per work tile: the bulk-copy destination of every token this thread owns. Order-buffer slot = + ``(((src_token * R + src_rank) * chunks_per_row + hidden_chunk) * K + dst_order_id) * chunk``; single-copy + tokens map to ``pre_reduced[src_token][slot][hidden_base]`` (slot = my rank, 0 under in-kernel REDG).""" + + metadata: cute.Tensor # domain-offset to the tile's first pool row + order_base: Int64 + remote_output: cute.Tensor # local pre-reduced buffer (T, R, H); (T, 1, H) view of the output under REDG + peer_rank_ptr_mapper: SymmetricBufferDevice + local_rank: Int32 + hidden_base: Int32 # first hidden element of this CTA tile + valid_tokens_this_cta_tile: Int32 + valid_hidden_this_cta_tile: Int32 + world_size: int + topk: int + chunk_bytes: int + chunks_per_row: int + element_bytes: int + reduce_topk_in_kernel: bool + mapping: FunctionMapping # per-issue (epi_tid, iter_idx) -> token_in_cta_tile + epi_tid: Int32 + dst_ptrs: Optional[cute.Tensor] = None # i64 x copy_iters + valid: Optional[cute.Tensor] = None # i32 x copy_iters: 0 skip, 1 bulk copy, 2 bulk reduce-add + copy_iters: int = 0 + + @cute.jit + def prefetch(self) -> "Fc2DedupBulkRouter": + copy_iters: cutlass.Constexpr[int] = self.mapping.domain.axis_size("iter_idx") + valid = cute.make_rmem_tensor((copy_iters,), Int32) + dst_ptrs = cute.make_rmem_tensor((copy_iters,), Int64) + hidden_base_bytes = Int64(self.hidden_base) * Int64(self.element_bytes) + chunk_index = hidden_base_bytes // Int64(self.chunk_bytes) + byte_in_chunk = hidden_base_bytes - chunk_index * Int64(self.chunk_bytes) + remote_slot = Int32(0) if cutlass.const_expr(self.reduce_topk_in_kernel) else self.local_rank + for iter_idx in cutlass.range_constexpr(copy_iters): + token_in_tile = Int32(self.mapping.evaluate(epi_tid=self.epi_tid, iter_idx=iter_idx)["token_in_cta_tile"]) + valid[iter_idx] = Int32(0) + dst_ptrs[iter_idx] = Int64(0) + if (token_in_tile < self.valid_tokens_this_cta_tile) and (self.valid_hidden_this_cta_tile > Int32(0)): + md = TokenSrcMetadata.load( + self.metadata.iterator.toint() + Int64(token_in_tile) * Int64(TokenSrcMetadata.nbytes) + ) + if md.dst_token_cnt == Int32(1): + # Single copy on this rank: push straight to the source rank's pre-reduced row. + local_row = cute.domain_offset( + (Int64(md.src_token), remote_slot, self.hidden_base), self.remote_output + ).iterator.toint() + dst_ptrs[iter_idx] = self.peer_rank_ptr_mapper.map(local_row, md.src_rank) + valid[iter_idx] = Int32(2) if cutlass.const_expr(self.reduce_topk_in_kernel) else Int32(1) + else: + block_index = Int64(md.src_token) * Int64(self.world_size) + Int64(md.src_rank) + slot_index = (block_index * Int64(self.chunks_per_row) + chunk_index) * Int64(self.topk) + Int64( + md.dst_order_id + ) + dst_ptrs[iter_idx] = self.order_base + slot_index * Int64(self.chunk_bytes) + byte_in_chunk + valid[iter_idx] = Int32(1) + return dataclasses.replace(self, dst_ptrs=dst_ptrs, valid=valid, copy_iters=copy_iters) + + @cute.jit + def get_data_dst(self, iter_idx: int) -> Tuple[cute.Pointer, Int32]: + ptr = cute.make_ptr(cutlass.BFloat16, self.dst_ptrs[iter_idx], AddressSpace.gmem, assumed_align=16) + return ptr, self.valid[iter_idx] + + +@cute.jit +def fc2_dedup_bulk_store_function( + *, + epi, + subtile: cute.Tensor, # (64,) bf16: one element per token of the subtile (hidden = warp*32 + lane) + epilogue_iter_idx: Int32, + subtile_idx: Int32, + fc2_output_router: Fc2DedupBulkRouter, + **_, +): + smem_tensor = epi.smem_tensor + if cutlass.const_expr(smem_tensor is None): + raise ValueError("The dedup bulk store requires epi.smem_tensor.") + data_subtile = subtile + if cutlass.const_expr(data_subtile.element_type is not cutlass.BFloat16): + raise ValueError("The dedup FC2 store expects a BF16 payload.") + if cutlass.const_expr(cute.size(data_subtile) != epi._EpilogueTokenTileSize): + raise ValueError("BF16 bulk staging requires one register per token in the 64-token subtile.") + warp_idx = epi.warp_idx + lane_idx = epi.lane_idx + smem_read_write_bar = SwapABGatedActEpilogue.epilogue_sync_barrier() + stage_idx = epilogue_iter_idx % Int32(epi.fc2_tma_stages) + smem_stage = smem_tensor[None, None, stage_idx] + + warp_hidden_base = Int32(warp_idx * 32) + for token_idx in cutlass.range_constexpr(epi._EpilogueTokenTileSize): + smem_stage[token_idx, warp_hidden_base + lane_idx] = data_subtile[token_idx] + + cute.arch.fence_proxy("async.shared", space="cta") + smem_read_write_bar.arrive_and_wait() + + copy_elems = Int32(epi._EpilogueFc2HiddenTileSize) + if cutlass.const_expr(epi.fc2_hidden_needs_predicate): + copy_elems = fc2_output_router.valid_hidden_this_cta_tile + copy_bytes = copy_elems * Int32(epi.combine_format.act_dtype.width // 8) + scratch_row = warp_idx * Int32(16) + lane_idx % Int32(16) + copy_iters: cutlass.Constexpr[int] = fc2_output_router.mapping.domain.axis_size("iter_idx") + for ublk_iter_idx in cutlass.range_constexpr(copy_iters): + owned_subtile = Int32(ublk_iter_idx * 2) + lane_idx // Int32(16) + if owned_subtile == subtile_idx: + dst_ptr, kind = fc2_output_router.get_data_dst(ublk_iter_idx) + if kind != Int32(0): + src_row = cute.slice_(smem_stage, (scratch_row, None)) + if cutlass.const_expr(fc2_output_router.reduce_topk_in_kernel): + if kind == Int32(2): + cp_reduce_async_bulk_add_bf16_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + else: + cp_async_bulk_s2g(dst_ptr, src_row.iterator, copy_bytes) + + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(epi.fc2_tma_stages - 1, read=True) + smem_read_write_bar.arrive_and_wait() + + +class SwapABFc2EpilogueDedup(SwapABFc2Epilogue): + """FC2 epilogue with SMEM-staged, metadata-routed bulk stores (order buffer or direct peer push).""" + + def __init__( + self, + base: "SwapABGatedActEpilogueDedup", + tidx: Int32, + smem_tensor: Optional[cute.Tensor], + tma_atom_fc2_output: Optional[cute.CopyAtom], + fc2_tma_output: Optional[cute.Tensor], + fc2_output: cute.Tensor, # order buffer base as a (T*R*K, 1, hidden) tensor + token_src_metadata: Optional[cute.Tensor], + fc2_done_counter: Optional[cute.Tensor], + fc2_output_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper, + optional_epi_args: GatedActEpilogueArgs, + ): + # Not super().__init__(): the base picks its store pipeline and freezes; dedup installs its own. + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + self.tma_atom_fc2_output = None + self.fc2_tma_output = None + self.fc2_output = fc2_output + self.token_src_metadata = token_src_metadata + self.fc2_done_counter = fc2_done_counter + self.fc2_output_sf = None + self.peer_rank_ptr_mapper = peer_rank_ptr_mapper + self.optional_epi_args = optional_epi_args + if smem_tensor is None: + raise ValueError("The dedup FC2 epilogue needs the FC2 staging SMEM (fc2_use_bulk).") + self.smem_tensor = smem_tensor + self.process_pipeline = dataclasses.replace( + make_fc2_ublk_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ), + store_function=fc2_dedup_bulk_store_function, + ) + self._freeze() + + @cute.jit + def _make_bulk_router(self, work_tile_info: SwapAbFc12WorkTileInfo) -> Fc2DedupBulkRouter: + task_tile_data_row_start = ( + work_tile_info.cumulative_data_physical_row + work_tile_info.tile_n_idx * Int32(self.cta_tile_n) + ) + hidden_base_this_cta_tile = work_tile_info.tile_m_idx * Int32(self.cta_tile_m) + valid_hidden = Int32(self.fc2_output.shape[2]) - hidden_base_this_cta_tile + if valid_hidden < 0: + valid_hidden = 0 + if valid_hidden > self._EpilogueFc2HiddenTileSize: + valid_hidden = self._EpilogueFc2HiddenTileSize + remote_output = self.optional_epi_args.fc2_remote_output + if cutlass.const_expr(remote_output is None): + raise ValueError("The dedup FC2 epilogue needs fc2_remote_output (the pre-reduced (T, R, H) buffer).") + # The pre_reduced slot of single-copy rows is the kernel's local_rank argument (not inferred from the mapper). + local_rank = self.optional_epi_args.local_rank + if cutlass.const_expr(local_rank is None): + raise ValueError("The dedup FC2 epilogue needs local_rank (the kernel's host-provided rank argument).") + return Fc2DedupBulkRouter( + metadata=cute.domain_offset((task_tile_data_row_start,), self.token_src_metadata), + order_base=self.fc2_output.iterator.toint(), + remote_output=remote_output, + peer_rank_ptr_mapper=self.peer_rank_ptr_mapper, + local_rank=Int32(local_rank), + hidden_base=hidden_base_this_cta_tile, + valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_cta_tile, + valid_hidden_this_cta_tile=valid_hidden, + world_size=self.dedup_world_size, + topk=self.dedup_topk, + chunk_bytes=self.dedup_chunk_bytes, + chunks_per_row=self.dedup_chunks_per_row, + element_bytes=self.combine_format.act_dtype.width // 8, + reduce_topk_in_kernel=self.reduce_topk_in_kernel, + mapping=self.process_pipeline.store_out_mapping, + epi_tid=self.tidx, + ) + + @cute.jit + def __call__( + self, work_tile_info: SwapAbFc12WorkTileInfo, tmem_acc_tensor: cute.Tensor, acc_pipeline, acc_consumer_state + ): + if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): + alpha_val = self.optional_epi_args.fc2_alpha[work_tile_info.expert_idx] + else: + alpha_val = None + router = self._make_bulk_router(work_tile_info) + acc_ready = False + if not work_tile_info.peek_ready: + acc_ready = True + acc_pipeline.consumer_wait(acc_consumer_state) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, self._EpilogueTokenTileSize) + )[None, None, 0, None] + acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) + # Read the metadata only after the accumulator wait: under PDL the router may still be pushing records. + router = router.prefetch() + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + for subtile_idx in cutlass.range(self.subtile_cnt, unroll=1): + if subtile_idx * Int32(self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + work_tile_info=work_tile_info, + epilogue_iter_idx=subtile_idx, + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], + fc2_output_router=router, + alpha_val=alpha_val, + ) + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + +class SwapABGatedActEpilogueDedup(SwapABGatedActEpilogue): + """``SwapABGatedActEpilogue`` whose FC2 half is ``SwapABFc2EpilogueDedup`` (bulk stores routed by metadata).""" + + @classmethod + def problem_desc_require(cls) -> dict[str, type]: + return {**super().problem_desc_require(), "apply_topk_at_fc1": bool} + + @classmethod + def impl_desc_require(cls) -> dict[str, object]: + return { + **super().impl_desc_require(), + "dedup_world_size": int, + "dedup_topk": int, + # Order buffer layout [T, R, chunks_per_row, K, chunk_bytes] (from TokenCommDedup). + "dedup_chunk_bytes": int, + "dedup_chunks_per_row": int, + } + + def _make_fc2_epilogue(self, *args) -> SwapABFc2EpilogueDedup: + return SwapABFc2EpilogueDedup(self, *args) + + def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: + super().__init__(problem_desc, impl_desc) + self.dedup_world_size = int(impl_desc["dedup_world_size"]) + self.dedup_topk = int(impl_desc["dedup_topk"]) + self.dedup_chunk_bytes = int(impl_desc["dedup_chunk_bytes"]) + self.dedup_chunks_per_row = int(impl_desc["dedup_chunks_per_row"]) + self.apply_topk_at_fc1 = bool(problem_desc["apply_topk_at_fc1"]) + if not self.apply_topk_at_fc1: + raise ValueError("Combine dedup applies top-k scores at FC1 only (pre_fc1_quant); post_fc2 is not supported.") + if not self.communication_enabled: + raise ValueError("Combine dedup requires token communication.") + if not self.token_back_push_data: + raise ValueError( + "Combine dedup reduces and pushes the order-buffer blocks from the combine warps; epi_warps is not supported." + ) + if not self.fc2_use_bulk: + raise ValueError("Combine dedup stages FC2 subtiles in SMEM for per-token bulk stores; fc2_use_bulk must be on.") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("Combine dedup supports the BF16 combine format only.") + if self.reduce_topk_in_epilogue: + raise ValueError("Combine dedup never reduces in the epilogue.") + if self.dedup_world_size <= 0 or self.dedup_topk <= 0 or self.dedup_chunks_per_row <= 0: + raise ValueError("Combine dedup needs positive dedup_world_size, dedup_topk and dedup_chunks_per_row.") + tile_bytes = self._EpilogueFc2HiddenTileSize * self.combine_format.act_dtype.width // 8 + if self.dedup_chunk_bytes <= 0 or self.dedup_chunk_bytes % tile_bytes != 0: + raise ValueError( + f"Combine dedup needs dedup_chunk_bytes as a multiple of the {tile_bytes} B hidden tile, got " + f"{self.dedup_chunk_bytes}." + ) + # The base resolved fc2_use_bulk + local output as the TMA tile store; dedup stores per token with UBLK. + self.fc2_use_tma = False + self.fc2_use_ublk = True + single_stage_region = self._make_fc2_single_stage_region() + if single_stage_region.nbytes % 16 != 0: + raise ValueError("Each FC2 UBLK staging stage must occupy a multiple of 16 bytes.") + requested_fc2_tma_stages = impl_desc.get("fc2_tma_stages") + if requested_fc2_tma_stages is not None: + self.fc2_tma_stages = int(requested_fc2_tma_stages) + else: + bf16_stage_bytes = self._EpilogueTokenTileSize * self._EpilogueFc2HiddenTileSize * cutlass.BFloat16.width // 8 + available_staging_bytes = max(self.fc1_staging_bytes, min(2, self.subtile_cnt) * bf16_stage_bytes) + self.fc2_tma_stages = min(self.subtile_cnt, available_staging_bytes // single_stage_region.nbytes) + self.fc2_staging_spec = self._make_fc2_staging_region(self.fc2_tma_stages) + + +__all__ = [ + "Fc2DedupBulkRouter", + "GatedActEpilogueArgsDedup", + "SwapABFc2EpilogueDedup", + "SwapABGatedActEpilogueDedup", + "fc2_dedup_bulk_store_function", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py index c17849379ec3..a8900c010327 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_extension.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Rubin block-scaled SwapAB adapter between FC12 scheduling and kernel tensor views.""" import dataclasses @@ -16,6 +13,7 @@ from .....helpers.dsl_helpers import spin_peek, spin_wait from ....schedulers.fc12_mapping import BlockPhase, SwapAbFc12WorkTileInfo, peek_ready_bit + TensorRole = Literal["a", "b", "sfa", "sfb", "c", "sfc", "topk"] @@ -59,18 +57,14 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "BlockScaledSwapAb def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result fc1_done_counter_pointer = rebuild(self.fc1_done_counter_pointer) fc2_spin_threshold = rebuild(self.fc2_spin_threshold) fc1_ready_counter_pointer = ( - rebuild(self.fc1_ready_counter_pointer) - if self.fc1_ready_counter_pointer is not None - else None + rebuild(self.fc1_ready_counter_pointer) if self.fc1_ready_counter_pointer is not None else None ) expert_remap = rebuild(self.expert_remap) if self.expert_remap is not None else None if value_index != len(values): @@ -110,9 +104,7 @@ def prepare_work_tile(self, work_tile: SwapAbFc12WorkTileInfo) -> SwapAbFc12Work if is_fc1: counter_pointer = self.fc1_ready_counter_pointer + counter_slot peek_flag = Int32(0) - if spin_peek( - counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile - ): + if spin_peek(counter_pointer, lambda value: value >= work_tile.valid_tokens_in_cta_tile): peek_flag = Int32(peek_ready_bit) phase_and_flags = work_tile.phase_and_flags | peek_flag @@ -170,19 +162,13 @@ def get_gmem_tensor( result = cute.domain_offset((0, 0, expert_idx), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name in ("sfb", "sfc")): result = cute.domain_offset((sf_token_offset, 0, 0), tensor) per_expert_shape = (shape[0], shape[1], singleton) sf_layout = tile_atom_to_shape_SF(per_expert_shape, self.sf_vec_size) - return ( - cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), - None, - ) + return (cute.make_tensor(result.iterator, cute.make_layout(sf_layout.shape, stride=stride)), None) if cutlass.const_expr(tensor_name == "c"): result = cute.domain_offset((data_token_offset, 0, 0), tensor) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py index b26743df47e8..b8fd0871d71f 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Mainloop component for the Rubin block-scaled swap-AB FC12 kernel.""" from typing import ClassVar, Optional, Tuple @@ -15,13 +12,7 @@ from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelComponent, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelComponent, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....helpers.cute_py_helpers import ( Tcgen05MmaInstruction, make_smem_layouts, @@ -33,15 +24,15 @@ from .....helpers.smem_workspace import SmemWorkspace from .....helpers.utils import ceil_div, round_up, strides_equal_ignoring_singletons from .....quant_def import QuantKind -from ....blackwell.custom_mix_cga_helpers import ( +from ....schedulers.base import SchedulerConsumer +from ....schedulers.fc12_mapping import BlockPhase +from ...custom_mix_cga_helpers import ( PipelineTmaUmmaMixedCga, TmaAtomOrPair, bind_executable_tma_load_fields, make_executable_tma_atom_or_pair, select_tma_atom_for_cluster, ) -from ....schedulers.base import SchedulerConsumer -from ....schedulers.fc12_mapping import BlockPhase from . import dynamic_mainloop from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension @@ -49,20 +40,14 @@ class BlockScaledSwapAbFc12Mainloop(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.ab_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.swap_ab_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 @classmethod @@ -108,9 +93,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -131,15 +114,9 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: raise ValueError("Rubin MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -152,12 +129,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -173,11 +146,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -186,9 +155,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self.cluster_layout_sfb_shape_vmnk = ((1,), *self.cluster_shape_mn, 1) self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) self.fallback_cluster_layout_shape_vmnk = ( (self.mma_cta_count,), @@ -196,11 +163,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.resolved_fallback_cluster_shape_mn[1], 1, ) - self.fallback_cluster_layout_sfb_shape_vmnk = ( - (1,), - *self.resolved_fallback_cluster_shape_mn, - 1, - ) + self.fallback_cluster_layout_sfb_shape_vmnk = ((1,), *self.resolved_fallback_cluster_shape_mn, 1) self.is_mixed_cga = self.resolved_fallback_cluster_shape_mn != self.cluster_shape_mn self.num_mcast_ctas_a = self.cluster_layout_shape_vmnk[2] self.num_mcast_ctas_b = self.cluster_layout_shape_vmnk[1] @@ -212,9 +175,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: def _validate_configuration(self) -> None: if self.mma_k_mode != "2x": - raise ValueError( - f"Rubin MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}." - ) + raise ValueError(f"Rubin MegaMoE only supports mma_k_mode='2x', got {self.mma_k_mode!r}.") for field_name, dimensions, expected_rank in ( ("mma_instruction_mnk", self.mma_instruction_mnk, 3), ("mma_tiler_mnk", self.mma_tiler_mnk, 3), @@ -222,10 +183,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -242,9 +200,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin MegaMoE does not implement M/N instruction repetition or B-reuse.") if self.mma_tiler_mnk[1] not in (64, 128, 256): raise ValueError("Rubin MegaMoE supports tile N in (64, 128, 256).") if self.mma_tiler_mnk[2] % self.instruction_k != 0: @@ -269,9 +225,7 @@ def _validate_configuration(self) -> None: self.cluster_shape_mn, self.fallback_cluster_shape_mn ) ): - raise ValueError( - "Preferred cluster dimensions must be divisible by fallback dimensions." - ) + raise ValueError("Preferred cluster dimensions must be divisible by fallback dimensions.") if self.mainloop_smem_budget_bytes <= 0: raise ValueError("mainloop_smem_budget_bytes must be positive.") if self.num_accumulator_consumer_warps_per_cta <= 0: @@ -286,9 +240,9 @@ def _validate_configuration(self) -> None: def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the selected Rubin mainloop SMEM plan.""" - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_ab_pipeline_stages * 2) + # Register the 4 B holding word before the 128 B-aligned stage tensors so it lands in the alignment gap. + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -303,37 +257,22 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -349,32 +288,25 @@ def packed_bytes(region, dtype) -> int: ) mbarrier_bytes = int(cutlass.Int64.width) // 8 ab_stage_cost_bytes = self.ab_stage_payload_bytes + 2 * mbarrier_bytes - plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes - self.num_ab_pipeline_stages = ( - self.mainloop_smem_budget_bytes - plan_tail_bytes - ) // ab_stage_cost_bytes + # Accumulator + TMEM-deallocation mbarriers plus the 8 B-aligned TMEM holding word. + tmem_holding_buffer_bytes = 8 + plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes + tmem_holding_buffer_bytes + self.num_ab_pipeline_stages = (self.mainloop_smem_budget_bytes - plan_tail_bytes) // ab_stage_cost_bytes if self.num_ab_pipeline_stages < 1: raise ValueError( f"One AB stage needs {ab_stage_cost_bytes + plan_tail_bytes} bytes, exceeding the " f"{self.mainloop_smem_budget_bytes}-byte mainloop budget." ) - (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = ( - make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - self.num_ab_pipeline_stages, - (self.a_major_mode, self.b_major_mode), - self.architecture, - ) + (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) = make_smem_layouts( + self.mma_instruction, + self.mma_tiler_mnk, + self.num_ab_pipeline_stages, + (self.a_major_mode, self.b_major_mode), + self.architecture, ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -412,12 +344,10 @@ def materialize_codegen_members(self) -> None: cute.make_layout((*self.cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) self.fallback_cluster_layout_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma.thr_id.shape,) ) self.fallback_cluster_layout_sfb_vmnk = cute.tiled_divide( - cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), - (self.tiled_mma_sfb.thr_id.shape,), + cute.make_layout((*self.resolved_fallback_cluster_shape_mn, 1)), (self.tiled_mma_sfb.thr_id.shape,) ) if self.cluster_layout_vmnk.shape != self.cluster_layout_shape_vmnk: raise ValueError( @@ -434,10 +364,7 @@ def materialize_codegen_members(self) -> None: f"Fallback cluster layout mismatch: Python {self.fallback_cluster_layout_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_vmnk.shape}." ) - if ( - self.fallback_cluster_layout_sfb_vmnk.shape - != self.fallback_cluster_layout_sfb_shape_vmnk - ): + if self.fallback_cluster_layout_sfb_vmnk.shape != self.fallback_cluster_layout_sfb_shape_vmnk: raise ValueError( f"Fallback SFB cluster layout mismatch: Python {self.fallback_cluster_layout_sfb_shape_vmnk}, " f"CuTe {self.fallback_cluster_layout_sfb_vmnk.shape}." @@ -463,11 +390,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -493,24 +417,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None def make_tma_maybe_pair( make_atom, @@ -678,17 +590,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 2) return PipelineTmaUmmaMixedCga.create( @@ -709,9 +614,7 @@ def create_ab_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -723,9 +626,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -736,9 +637,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -752,18 +651,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -792,17 +685,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -812,16 +698,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -876,16 +756,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1006,19 +880,13 @@ def run_tma_b( "sfb", ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) b_partition_coord = self.main_vmnk_coord[1] sfb_partition_coord = self.sfb_vmnk_coord[1] if cutlass.const_expr(self.is_mixed_cga): if self.is_fallback_cluster: - split_factor = ( - self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] - ) + split_factor = self.cluster_shape_mn[0] // self.resolved_fallback_cluster_shape_mn[0] b_partition_coord = b_partition_coord * Int32(split_factor) sfb_partition_coord = sfb_partition_coord * Int32(split_factor) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -1033,24 +901,17 @@ def run_tma_b( kernel_extension.wait_for_input(work_tile) iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc1_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc1_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1094,11 +955,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc1_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1106,9 +963,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) else: iket.range_push("tma_token_fc2") @@ -1124,24 +979,17 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1185,11 +1033,7 @@ def run_tma_b( tma_desc_ptr=b_descriptor_ptr, mcast_mask=b_multicast_mask, ) - cute.copy( - b_issue_atom, - global_b_slice[(None, k_tile_idx)], - partitioned_smem_b[(None, stage_index)], - ) + cute.copy(b_issue_atom, global_b_slice[(None, k_tile_idx)], partitioned_smem_b[(None, stage_index)]) sfb_issue_atom = bind_executable_tma_load_fields( fc2_tma_sfb_exec_atom, tma_bar_ptr=tma_barrier, @@ -1197,9 +1041,7 @@ def run_tma_b( mcast_mask=sfb_multicast_mask, ) cute.copy( - sfb_issue_atom, - global_sfb_slice[(None, k_tile_idx)], - partitioned_smem_sfb[(None, stage_index)], + sfb_issue_atom, global_sfb_slice[(None, k_tile_idx)], partitioned_smem_sfb[(None, stage_index)] ) iket.range_pop() @@ -1224,33 +1066,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1283,20 +1116,13 @@ def run_mma( if cutlass.const_expr(self.cta_tile_n == 64): sfb_shift = (work_tile.tile_n_idx % Int32(2)) * Int32(2) shifted_sfb_pointer = cute.recast_ptr( - accumulator_pointer - + self.num_accumulator_tmem_cols - + self.num_sfa_tmem_cols - + sfb_shift, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + sfb_shift, dtype=self.sf_dtype, ) mma_tmem_sfb = cute.make_tensor(shifted_sfb_pointer, sfb_layout) use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - aligned_instruction_n = ( - work_tile.valid_tokens_in_cta_tile + Int32(15) - ) & Int32(-16) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + aligned_instruction_n = (work_tile.valid_tokens_in_cta_tile + Int32(15)) & Int32(-16) use_n128_sfb_mapping = aligned_instruction_n <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) @@ -1311,35 +1137,17 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. - lower_n_lower_k_source_coord = ( - None, - 0, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_lower_k_source_coord = (None, 0, None, instruction_index, ab_pipeline_state.index) lower_n_lower_k_destination_coord = (None, 0, None, instruction_index) cute.copy( tiled_copy_sfb, partitioned_smem_sfb[lower_n_lower_k_source_coord], partitioned_tmem_sfb[lower_n_lower_k_destination_coord], ) - lower_n_upper_k_source_coord = ( - None, - 2, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_upper_k_source_coord = (None, 2, None, instruction_index, ab_pipeline_state.index) compact_upper_k_destination_coord = (None, 1, None, instruction_index) cute.copy( tiled_copy_sfb, @@ -1378,16 +1186,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[s2t_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py index 9533ff4ff895..029cd8a08ee9 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_fc12_mainloop_gen_specialized.py @@ -1,7 +1,5 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause """Mainloop component for the Rubin generation-phase block-scaled swap-AB FC12 kernel. Three properties separate this from the two mainloops it descends from. @@ -54,31 +52,23 @@ from .....quant_def import QuantKind from ....schedulers.base import SchedulerConsumer from ....schedulers.fc12_mapping import BlockPhase -from ..local_mega.tma_gather import sm107_tma_gather4_load from . import dynamic_mainloop from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension +from .tma_gather import sm107_tma_gather4_load class BlockScaledSwapAbFc12MainloopGenSpecialized(KernelComponent): """Own all Rubin load, MMA, pipeline, SMEM, and TMEM mainloop state.""" - ab_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.ab_pipeline_mbarriers" - ) - token_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.token_pipeline_mbarriers" - ) + ab_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.ab_pipeline_mbarriers" + token_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.token_pipeline_mbarriers" a_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.a_smem_tensor" b_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.b_smem_tensor" sfa_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.sfa_smem_tensor" sfb_smem_tensor_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.sfb_smem_tensor" - acc_pipeline_mbarriers_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.acc_pipeline_mbarriers" - ) + acc_pipeline_mbarriers_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.acc_pipeline_mbarriers" tmem_holding_buffer_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.tmem_holding_buffer" - tmem_deallocation_mbarrier_region: ClassVar[str] = ( - "rubin.genphase_fc12.mainloop.tmem_deallocation_mbarrier" - ) + tmem_deallocation_mbarrier_region: ClassVar[str] = "rubin.genphase_fc12.mainloop.tmem_deallocation_mbarrier" tmem_allocation_barrier_id: ClassVar[int] = 2 fc1_gather_sync_barrier_id: ClassVar[int] = 3 # Four warps rather than one because the gather4 row operands live in uniform @@ -113,7 +103,8 @@ class BlockScaledSwapAbFc12MainloopGenSpecialized(KernelComponent): # profiling put it -- the token operand is read out of L2, so a handful of # stages covers its latency, but not as few as the L2 latency alone suggests. asymmetric_token_stages: ClassVar[int] = 6 - # The token tile this specialization applies to; other tiles use the shared plan. + # The token tile this specialization applies to. At 128 the shared plan + # already reaches 93% of DRAM peak and is left alone. asymmetric_token_tile: ClassVar[int] = 256 @classmethod @@ -153,9 +144,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.use_2cta_instrs = impl_desc["use_2cta_instrs"] self.tmem_plan = impl_desc["tmem_plan"] self.mainloop_smem_budget_bytes = impl_desc["mainloop_smem_budget_bytes"] - self.num_accumulator_consumer_warps_per_cta = impl_desc[ - "num_accumulator_consumer_warps_per_cta" - ] + self.num_accumulator_consumer_warps_per_cta = impl_desc["num_accumulator_consumer_warps_per_cta"] self.architecture = "sm_107" self.mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -174,19 +163,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.num_accumulator_stages = tmem_plan.accumulator_stage_count self.num_accumulator_pipeline_stages = tmem_plan.accumulator_pipeline_stages if tmem_plan.accumulator_stage_stride_columns != tmem_plan.accumulator_stage_columns: - raise ValueError( - "Rubin Genphase MegaMoE does not support overlapping accumulator stages." - ) + raise ValueError("Rubin Genphase MegaMoE does not support overlapping accumulator stages.") if tmem_plan.accumulator_pipeline_stages != tmem_plan.accumulator_stage_count: - raise ValueError( - "Rubin Genphase MegaMoE requires one pipeline stage per disjoint accumulator stage." - ) + raise ValueError("Rubin Genphase MegaMoE requires one pipeline stage per disjoint accumulator stage.") self.num_accumulator_tmem_cols = tmem_plan.accumulator_columns - self.accumulator_shape = ( - self.cta_tile_m, - self.cta_tile_n, - tmem_plan.accumulator_stage_count, - ) + self.accumulator_shape = (self.cta_tile_m, self.cta_tile_n, tmem_plan.accumulator_stage_count) self.accumulator_stride = (1 << 16, 1, tmem_plan.accumulator_stage_stride_columns) self.mma_instruction = Tcgen05MmaInstruction( @@ -199,12 +180,8 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.a_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.a_dtype, self.b_dtype, self.architecture - ) - self.b_smem_alloc_dtype = tcgen05_smem_alloc_type( - self.b_dtype, self.a_dtype, self.architecture - ) + self.a_smem_alloc_dtype = tcgen05_smem_alloc_type(self.a_dtype, self.b_dtype, self.architecture) + self.b_smem_alloc_dtype = tcgen05_smem_alloc_type(self.b_dtype, self.a_dtype, self.architecture) self.sfb_instruction_shape_mnk = ( self.mma_instruction_mnk[0] // self.mma_cta_count, round_up(self.mma_instruction_mnk[1], 128), @@ -220,11 +197,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sfb_type=self.sf_dtype, sf_vec_size=self.sf_vec_size, ) - self.mma_tiler_sfb = ( - self.sfb_instruction_shape_mnk[0], - self.sfb_instruction_shape_mnk[1], - self.mma_tile_k, - ) + self.mma_tiler_sfb = (self.sfb_instruction_shape_mnk[0], self.sfb_instruction_shape_mnk[1], self.mma_tile_k) self.cluster_layout_shape_vmnk = ( (self.mma_cta_count,), self.cluster_shape_mn[0] // self.mma_cta_count, @@ -237,19 +210,15 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.is_a_mcast = self.num_mcast_ctas_a > 1 self.is_b_mcast = self.num_mcast_ctas_b > 1 - # Only the selected token tile uses asymmetric operand stages. + # Only the 256-wide token tile is specialized: at 128 the shared plan + # already reaches 93% of DRAM peak, so it stays byte-for-byte as it was. self.uses_asymmetric_ab_stages = self.cta_tile_n == self.asymmetric_token_tile self._select_ab_stage_plan() self.num_tma_load_bytes = self.ab_stage_tma_bytes * self.mma_cta_count if self.num_tma_load_bytes != ( - self.a_stage_tma_bytes - + self.b_stage_tma_bytes - + self.sfa_stage_tma_bytes - + self.sfb_stage_tma_bytes + self.a_stage_tma_bytes + self.b_stage_tma_bytes + self.sfa_stage_tma_bytes + self.sfb_stage_tma_bytes ): - raise ValueError( - "Genphase MegaMoE AB stage transaction byte accounting is inconsistent." - ) + raise ValueError("Genphase MegaMoE AB stage transaction byte accounting is inconsistent.") # Under the split plan each pipeline has a single producer whose expectation # is exactly the bytes it issues. The shared plan instead has two producers # splitting one combined expectation, so neither share matches the bytes @@ -258,14 +227,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_pipeline_tx_count = self.b_stage_tma_bytes + self.sfb_stage_tma_bytes if not self.uses_asymmetric_ab_stages: if self.num_tma_load_bytes % 2 != 0: - raise ValueError( - "Genphase MegaMoE total AB transaction bytes must split evenly across two arrivals." - ) + raise ValueError("Genphase MegaMoE total AB transaction bytes must split evenly across two arrivals.") self.ab_pipeline_tx_count_per_producer = self.num_tma_load_bytes // 2 self._plan_gather_geometry() self.fc1_gather_sync_barrier = pipeline.NamedBarrier( - barrier_id=self.fc1_gather_sync_barrier_id, - num_threads=32 * (1 + self.num_gather_b_warps), + barrier_id=self.fc1_gather_sync_barrier_id, num_threads=32 * (1 + self.num_gather_b_warps) ) def _plan_gather_geometry(self) -> None: @@ -291,9 +257,7 @@ def _plan_gather_geometry(self) -> None: if self.gather_b_rows_per_cta % run_span_rows != 0: # The MMA CTA count cancels out of both sides, so the requirement is # purely on tile N and cluster M. - required_multiple = ( - self.gather_index_load_rows * self.num_gather_b_warps * self.cluster_shape_mn[0] - ) + required_multiple = self.gather_index_load_rows * self.num_gather_b_warps * self.cluster_shape_mn[0] raise ValueError( f"Genphase MegaMoE gather runs must divide evenly across {self.gather_issuing_units} issuing " f"units: tile N {self.cta_tile_n} must be a multiple of {required_multiple}." @@ -335,10 +299,7 @@ def _validate_configuration(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -355,9 +316,7 @@ def _validate_configuration(self) -> None: f"got {self.mma_instruction_mnk[0]}." ) if self.mma_instruction_mnk[:2] != self.mma_tiler_mnk[:2]: - raise NotImplementedError( - "Rubin Genphase MegaMoE does not implement M/N instruction repetition." - ) + raise NotImplementedError("Rubin Genphase MegaMoE does not implement M/N instruction repetition.") # Tile N 64 is excluded: it forces the halved SFB tile index and the SFB # TMEM shift, two special paths whose interaction with a dynamic # instruction N is not worth carrying for a token extent this small. @@ -382,13 +341,9 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: # Under the shared plan this one array serves all four operands. Under the # split plan it becomes the weight side's and the token side gets its own, # because the two depths no longer index the same stage. - smem_workspace.register_mbarrier( - self.ab_pipeline_mbarriers_region, self.num_weight_ab_stages * 2 - ) + smem_workspace.register_mbarrier(self.ab_pipeline_mbarriers_region, self.num_weight_ab_stages * 2) if self.uses_asymmetric_ab_stages: - smem_workspace.register_mbarrier( - self.token_pipeline_mbarriers_region, self.num_token_ab_stages * 2 - ) + smem_workspace.register_mbarrier(self.token_pipeline_mbarriers_region, self.num_token_ab_stages * 2) for region_name, region in ( (self.a_smem_tensor_region, self.a_smem_region), (self.b_smem_tensor_region, self.b_smem_region), @@ -403,37 +358,23 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: swizzle=region.swizzle, byte_alignment=region.byte_alignment, ) - smem_workspace.register_mbarrier( - self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2 - ) + smem_workspace.register_mbarrier(self.acc_pipeline_mbarriers_region, self.num_accumulator_pipeline_stages * 2) smem_workspace.register_mbarrier(self.tmem_deallocation_mbarrier_region, 1) - smem_workspace.register_tensor( - self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8 - ) + smem_workspace.register_tensor(self.tmem_holding_buffer_region, cutlass.Int32, (1,), byte_alignment=8) def _select_ab_stage_plan(self) -> None: stage_one_regions = make_smem_layouts( - self.mma_instruction, - self.mma_tiler_mnk, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.mma_instruction, self.mma_tiler_mnk, 1, (self.a_major_mode, self.b_major_mode), self.architecture ) independent_sfb_region = make_smem_layouts( - self.sfb_mma_instruction, - self.mma_tiler_sfb, - 1, - (self.a_major_mode, self.b_major_mode), - self.architecture, + self.sfb_mma_instruction, self.mma_tiler_sfb, 1, (self.a_major_mode, self.b_major_mode), self.architecture )[3] if ( independent_sfb_region.shape != stage_one_regions[3].shape or independent_sfb_region.stride != stage_one_regions[3].stride or independent_sfb_region.nbytes != stage_one_regions[3].nbytes ): - raise ValueError( - "The Rubin independent SFB plan does not match the primary MMA SFB operand." - ) + raise ValueError("The Rubin independent SFB plan does not match the primary MMA SFB operand.") def packed_bytes(region, dtype) -> int: return (region.cosize * int(dtype.width) + 7) // 8 @@ -448,24 +389,17 @@ def packed_bytes(region, dtype) -> int: stage_one_regions[3].nbytes, ) self.ab_stage_tma_bytes = sum(stage_operand_tma_bytes) - ( - self.a_stage_tma_bytes, - self.b_stage_tma_bytes, - self.sfa_stage_tma_bytes, - self.sfb_stage_tma_bytes, - ) = tuple(stage_bytes * self.mma_cta_count for stage_bytes in stage_operand_tma_bytes) + (self.a_stage_tma_bytes, self.b_stage_tma_bytes, self.sfa_stage_tma_bytes, self.sfb_stage_tma_bytes) = tuple( + stage_bytes * self.mma_cta_count for stage_bytes in stage_operand_tma_bytes + ) mbarrier_bytes = int(cutlass.Int64.width) // 8 plan_tail_bytes = (2 * self.num_accumulator_pipeline_stages + 2) * mbarrier_bytes stage_budget_bytes = self.mainloop_smem_budget_bytes - plan_tail_bytes # Each pipeline spends two mbarriers per stage, full and empty. The shared # plan pays that once for all four operands; the split plan pays it twice, # which is the only overhead the split introduces. - weight_stage_cost_bytes = ( - stage_one_regions[0].nbytes + stage_one_regions[2].nbytes + 2 * mbarrier_bytes - ) - token_stage_cost_bytes = ( - stage_one_regions[1].nbytes + stage_one_regions[3].nbytes + 2 * mbarrier_bytes - ) + weight_stage_cost_bytes = stage_one_regions[0].nbytes + stage_one_regions[2].nbytes + 2 * mbarrier_bytes + token_stage_cost_bytes = stage_one_regions[1].nbytes + stage_one_regions[3].nbytes + 2 * mbarrier_bytes if self.uses_asymmetric_ab_stages: self.num_token_ab_stages = self.asymmetric_token_stages @@ -516,12 +450,7 @@ def packed_bytes(region, dtype) -> int: ) self.selected_smem_bytes = sum( region.nbytes - for region in ( - self.a_smem_region, - self.b_smem_region, - self.sfa_smem_region, - self.sfb_smem_region, - ) + for region in (self.a_smem_region, self.b_smem_region, self.sfa_smem_region, self.sfb_smem_region) ) def make_tiled_mma(self) -> cute.TiledMma: @@ -591,11 +520,8 @@ def materialize_codegen_members(self) -> None: ("SFA", self.sfa_smem_layout_staged, self.sfa_smem_region), ("SFB", self.sfb_smem_layout_staged, self.sfb_smem_region), ): - if ( - materialized_layout.shape != planned_region.shape - or not strides_equal_ignoring_singletons( - planned_region.shape, materialized_layout.stride, planned_region.stride - ) + if materialized_layout.shape != planned_region.shape or not strides_equal_ignoring_singletons( + planned_region.shape, materialized_layout.stride, planned_region.stride ): raise ValueError( f"{tensor_name} SMEM plan mismatch: Python " @@ -621,24 +547,12 @@ def prepare_tma_load_params( sfa_stage_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) sfb_stage_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) - a_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - b_operation = sm100_utils.cluster_shape_to_tma_atom_B( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB( - self.cluster_shape_mn, self.tiled_mma.thr_id - ) - a_internal_type = ( - self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None - ) - b_internal_type = ( - self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None - ) + a_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + b_operation = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfa_operation = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, self.tiled_mma.thr_id) + sfb_operation = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, self.tiled_mma.thr_id) + a_internal_type = self.a_smem_alloc_dtype if self.a_smem_alloc_dtype is not self.a_dtype else None + b_internal_type = self.b_smem_alloc_dtype if self.b_smem_alloc_dtype is not self.b_dtype else None fc1_a_atom, fc1_a_tensor = cute.nvgpu.make_tiled_tma_atom_A( a_operation, @@ -656,9 +570,7 @@ def prepare_tma_load_params( # gather supplies coordinates directly and never partitions a tensor view. fc1_b_2d = fc1_b[(None, None, 0)] gather_b_base = cute.make_layout((1, self.mma_tile_k), stride=(self.mma_tile_k, 1)) - gather_b_smem_layout = cute.make_composed_layout( - cute.make_swizzle(3, 4, 3), 0, gather_b_base - ) + gather_b_smem_layout = cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, gather_b_base) fc1_b_atom, _ = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), fc1_b_2d, gather_b_smem_layout, (1, self.mma_tile_k) ) @@ -765,17 +677,10 @@ def assign_device_members( ) self.fc1_k_tile_count = ceil_div(bound_hidden_size, self.mma_tile_k) self.fc2_k_tile_count = ceil_div(bound_intermediate_gateup_size // 2, self.mma_tile_k) - self.sfb_vmnk_coord = ( - 0, - cta_coord_in_cluster[0], - cta_coord_in_cluster[1], - cta_coord_in_cluster[2], - ) + self.sfb_vmnk_coord = (0, cta_coord_in_cluster[0], cta_coord_in_cluster[1], cta_coord_in_cluster[2]) @cute.jit - def create_ab_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_ab_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the shared two-producer TMA-to-UMMA pipeline. Each producer arrival contributes the same accounting share. The two @@ -798,9 +703,7 @@ def create_ab_pipeline( ) @cute.jit - def create_weight_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_weight_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the weight half of the split TMA-to-UMMA pipeline. One producer, TMA-A, and its expectation is exactly the A and SFA bytes it @@ -821,9 +724,7 @@ def create_weight_pipeline( ) @cute.jit - def create_token_pipeline( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> pipeline.PipelineAsync: + def create_token_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> pipeline.PipelineAsync: """Create the token half of the split TMA-to-UMMA pipeline. One producer, TMA-B, registering the B and SFB bytes. Gather-B still never @@ -848,9 +749,7 @@ def create_token_pipeline( def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer): """Create the UMMA-to-epilogue accumulator pipeline.""" producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - consumer_thread_count = ( - self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count - ) + consumer_thread_count = self.num_accumulator_consumer_warps_per_cta * 32 * self.mma_cta_count consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, consumer_thread_count) return pipeline.PipelineUmmaAsync.create( barrier_storage=smem_workspace.ptr(self.acc_pipeline_mbarriers_region, smem_base), @@ -862,9 +761,7 @@ def create_acc_pipeline(self, smem_workspace: SmemWorkspace, smem_base: cute.Poi ) @cute.jit - def create_tmem_allocator( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int - ): + def create_tmem_allocator(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer, *, allocator_warp_id: int): """Bind the SM107 TMEM allocator to finalized SMEM regions.""" allocation_barrier = pipeline.NamedBarrier( barrier_id=self.tmem_allocation_barrier_id, @@ -875,9 +772,7 @@ def create_tmem_allocator( barrier_for_retrieve=allocation_barrier, allocator_warp_id=allocator_warp_id, is_two_cta=self.use_2cta_instrs, - two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr( - self.tmem_deallocation_mbarrier_region, smem_base - ), + two_cta_tmem_dealloc_mbar_ptr=smem_workspace.ptr(self.tmem_deallocation_mbarrier_region, smem_base), arch=self.architecture, ) @@ -891,18 +786,12 @@ def _s2t_copy_and_partition(self, smem_tensor: cute.Tensor, tmem_tensor: cute.Te mn_mode = cute.get(compact_smem_tensor.layout, mode=[0, 0]) mn_mode = cute.append(mn_mode, cute.make_layout((4,), stride=(0,))) - broadcast_layout = cute.append( - cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1]) - ) + broadcast_layout = cute.append(cute.group_modes(mn_mode, 0), cute.get(compact_smem_tensor.layout, mode=[0, 1])) broadcast_layout = cute.append( cute.group_modes(broadcast_layout, 0), cute.get(compact_smem_tensor.layout, mode=[1]) ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2]) - ) - broadcast_layout = cute.append( - broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3]) - ) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[2])) + broadcast_layout = cute.append(broadcast_layout, cute.get(compact_smem_tensor.layout, mode=[3])) broadcast_smem_tensor = cute.make_tensor(compact_smem_tensor.iterator, broadcast_layout) partitioned_smem = thread_copy.partition_S(broadcast_smem_tensor) @@ -931,17 +820,10 @@ def run_tma_a( multicast_mask = None if cutlass.const_expr(self.is_a_mcast or self.use_2cta_instrs): multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "a", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "a" ) - a_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape - ) + a_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, 0, None, 0)).shape) sfa_cta_layout = a_cta_layout thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) @@ -954,16 +836,10 @@ def run_tma_a( if is_fc1: iket.range_push("tma_weight_fc1") k_tile_count = self.fc1_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc1_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc1_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc1_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc1_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1018,16 +894,10 @@ def run_tma_a( else: iket.range_push("tma_weight_fc2") k_tile_count = self.fc2_k_tile_count - real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor( - "a", fc2_tma_a_tensor, work_tile - ) - real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfa", fc2_tma_sfa_tensor, work_tile - ) + real_a, a_descriptor_ptr = kernel_extension.get_gmem_tensor("a", fc2_tma_a_tensor, work_tile) + real_sfa, sfa_descriptor_ptr = kernel_extension.get_gmem_tensor("sfa", fc2_tma_sfa_tensor, work_tile) - global_a = cute.local_tile( - real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) - ) + global_a = cute.local_tile(real_a, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None)) global_sfa = cute.local_tile( real_sfa, cute.slice_(self.mma_tiler_mnk, (None, 0, None)), (None, None, None) ) @@ -1120,28 +990,14 @@ def run_tma_b( sfb_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) sfb_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "sfb", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "sfb" ) - b_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape - ) - sfb_cta_layout = cute.make_layout( - cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape - ) + b_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_vmnk, (0, None, 0, 0)).shape) + sfb_cta_layout = cute.make_layout(cute.slice_(self.cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape) thread_mma = self.tiled_mma.get_slice(self.mma_cta_index) thread_mma_sfb = self.tiled_mma_sfb.get_slice(self.mma_cta_index) @@ -1153,16 +1009,10 @@ def run_tma_b( iket.range_push("tma_b.wait_expert_sf_ready") expert_idx = work_tile.expert_idx expert_sf_blocks = ceil_div(expert_sizes[expert_idx], Int32(self.sf_padding_block)) - spin_wait( - fc1_sf_ready_pointer + expert_idx, - lambda value: value == expert_sf_blocks, - scope="gpu", - ) + spin_wait(fc1_sf_ready_pointer + expert_idx, lambda value: value == expert_sf_blocks, scope="gpu") iket.range_pop() k_tile_count = self.fc1_k_tile_count - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc1_tma_sfb_tensor, work_tile - ) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc1_tma_sfb_tensor, work_tile) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) @@ -1216,12 +1066,8 @@ def run_tma_b( iket.range_pop() k_tile_count = self.fc2_k_tile_count - real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor( - "b", fc2_tma_b_tensor, work_tile - ) - real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor( - "sfb", fc2_tma_sfb_tensor, work_tile - ) + real_b, b_descriptor_ptr = kernel_extension.get_gmem_tensor("b", fc2_tma_b_tensor, work_tile) + real_sfb, sfb_descriptor_ptr = kernel_extension.get_gmem_tensor("sfb", fc2_tma_sfb_tensor, work_tile) # Under two-CTA instructions the dynamic N window is split across # the pair, so the non-leader's rows start at the runtime midpoint @@ -1229,14 +1075,11 @@ def run_tma_b( if cutlass.const_expr(self.use_2cta_instrs): if not self.is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( - valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, - mma_tiler_n=self.mma_tiler_mnk[1], + valid_tokens_in_tile=work_tile.valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler_mnk[1] ) real_b = cute.domain_offset((load_shift, 0, 0), real_b) - global_b = cute.local_tile( - real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None) - ) + global_b = cute.local_tile(real_b, cute.slice_(self.mma_tiler_mnk, (0, None, None)), (None, None, None)) global_sfb = cute.local_tile( real_sfb, cute.slice_(self.mma_tiler_sfb, (0, None, None)), (None, None, None) ) @@ -1315,9 +1158,7 @@ def run_tma_gather_b1( producer. FC2 is fully produced by TMA-A and TMA-B, so on those tiles this warp only keeps its register-held pipeline state aligned. """ - row_ids = cute.make_rmem_tensor( - cute.make_layout((self.gather_b_rows_per_unit,)), cutlass.Int32 - ) + row_ids = cute.make_rmem_tensor(cute.make_layout((self.gather_b_rows_per_unit,)), cutlass.Int32) row_id_runs = cute.zipped_divide(row_ids, (self.gather_index_load_rows,)) index_load_atom = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), @@ -1328,12 +1169,7 @@ def run_tma_gather_b1( gather_b_multicast_mask = None if cutlass.const_expr(self.is_b_mcast or self.use_2cta_instrs): gather_b_multicast_mask = tma_multicast_mask( - self.cluster_shape_mn, - None, - self.cta_coord_in_cluster, - None, - self.use_2cta_instrs, - "b", + self.cluster_shape_mn, None, self.cta_coord_in_cluster, None, self.use_2cta_instrs, "b" ) issuing_unit = self.main_vmnk_coord[1] * Int32(self.num_gather_b_warps) + gather_warp_idx @@ -1342,10 +1178,7 @@ def run_tma_gather_b1( is_fc1 = work_tile.phase == Int32(BlockPhase.Linear1) if is_fc1: iket.range_push("tma_gather_b1_fc1") - tile_pool_base = ( - work_tile.cumulative_data_physical_row - + work_tile.tile_n_idx * Int32(self.cta_tile_n) - ) + tile_pool_base = work_tile.cumulative_data_physical_row + work_tile.tile_n_idx * Int32(self.cta_tile_n) # The MMA issues a dynamic instruction N, so under two-CTA # instructions the pair splits that runtime extent rather than the # static tile: the non-leader's local row zero is the midpoint of @@ -1353,12 +1186,9 @@ def run_tma_gather_b1( # `compute_non_leader_cta_load_shift` does for the tiled TMA legs. # Rounding to sixteen keeps that midpoint a multiple of eight, so # every run stays aligned to its 256-bit index load. - cta_row_base = self.mma_cta_index * ( - round_up(work_tile.valid_tokens_in_cta_tile, 16) // Int32(2) - ) + cta_row_base = self.mma_cta_index * (round_up(work_tile.valid_tokens_in_cta_tile, 16) // Int32(2)) index_runs = cute.zipped_divide( - cute.domain_offset((tile_pool_base + cta_row_base,), gather_index), - (self.gather_index_load_rows,), + cute.domain_offset((tile_pool_base + cta_row_base,), gather_index), (self.gather_index_load_rows,) ) for local_run in cutlass.range_constexpr(self.gather_b_runs_per_unit): run_index = Int32(local_run * self.gather_issuing_units) + issuing_unit @@ -1376,9 +1206,7 @@ def run_tma_gather_b1( tma_barrier = ab_pipeline.producer_get_barrier(ab_pipeline_state) stage_index = ab_pipeline_state.index col_k = k_tile_idx * Int32(self.mma_tile_k) - stage_base_elements = stage_index * Int32( - self.gather_b_rows_per_cta * self.mma_tile_k - ) + stage_base_elements = stage_index * Int32(self.gather_b_rows_per_cta * self.mma_tile_k) for local_run in cutlass.range_constexpr(self.gather_b_runs_per_unit): run_index = Int32(local_run * self.gather_issuing_units) + issuing_unit @@ -1389,10 +1217,7 @@ def run_tma_gather_b1( destination_offset = run_base_elements + Int32( group_in_run * self.gather4_width * self.mma_tile_k ) - row_start = ( - local_run * self.gather_index_load_rows - + group_in_run * self.gather4_width - ) + row_start = local_run * self.gather_index_load_rows + group_in_run * self.gather4_width with cute.arch.elect_one(): sm107_tma_gather4_load( fc1_tma_b_atom, @@ -1437,33 +1262,24 @@ def run_mma( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1496,12 +1312,8 @@ def run_mma( # window compact, so its two K halves come from a different pair of # SMEM rest coordinates than the full-width case. use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - use_n128_sfb_mapping = round_up( - work_tile.valid_tokens_in_cta_tile, 16 - ) <= Int32(128) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + use_n128_sfb_mapping = round_up(work_tile.valid_tokens_in_cta_tile, 16) <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) if use_n128_sfb_mapping: @@ -1515,35 +1327,17 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. - lower_n_lower_k_source_coord = ( - None, - 0, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_lower_k_source_coord = (None, 0, None, instruction_index, ab_pipeline_state.index) lower_n_lower_k_destination_coord = (None, 0, None, instruction_index) cute.copy( tiled_copy_sfb, partitioned_smem_sfb[lower_n_lower_k_source_coord], partitioned_tmem_sfb[lower_n_lower_k_destination_coord], ) - lower_n_upper_k_source_coord = ( - None, - 2, - None, - instruction_index, - ab_pipeline_state.index, - ) + lower_n_upper_k_source_coord = (None, 2, None, instruction_index, ab_pipeline_state.index) compact_upper_k_destination_coord = (None, 1, None, instruction_index) cute.copy( tiled_copy_sfb, @@ -1582,16 +1376,8 @@ def run_mma( next_full_status = ab_pipeline.consumer_try_wait(next_ab_pipeline_state) s2t_stage_coord = (None, None, None, None, ab_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[s2t_stage_coord], - partitioned_tmem_sfa, - ) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[s2t_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[s2t_stage_coord], partitioned_tmem_sfa) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[s2t_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, ab_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, ab_pipeline_state.index)] @@ -1651,33 +1437,24 @@ def run_mma_asymmetric( tmem_allocator.wait_for_alloc() accumulator_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) - accumulator_layout = cute.make_layout( - self.accumulator_shape, stride=self.accumulator_stride - ) + accumulator_layout = cute.make_layout(self.accumulator_shape, stride=self.accumulator_stride) accumulator_base = cute.make_tensor(accumulator_pointer, accumulator_layout) sfa_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfa( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfa(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) sfb_stage_smem_layout = cute.slice_( - blockscaled_utils.make_smem_layout_sfb( - tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1 - ), + blockscaled_utils.make_smem_layout_sfb(tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, 1), (None, None, None, 0), ) - sfa_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype - ) + sfa_pointer = cute.recast_ptr(accumulator_pointer + self.num_accumulator_tmem_cols, dtype=self.sf_dtype) sfa_layout = blockscaled_utils.make_tmem_layout_sfa( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfa_stage_smem_layout ) tmem_sfa = cute.make_tensor(sfa_pointer, sfa_layout) sfb_pointer = cute.recast_ptr( - accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, - dtype=self.sf_dtype, + accumulator_pointer + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, dtype=self.sf_dtype ) sfb_layout = blockscaled_utils.make_tmem_layout_sfb( tiled_mma, self.mma_tiler_mnk, self.sf_vec_size, sfb_stage_smem_layout @@ -1713,12 +1490,8 @@ def run_mma_asymmetric( # state machine out of the branch: the predicate is warp-uniform # and this loop has no state machine to duplicate. use_n128_sfb_mapping = False - if cutlass.const_expr( - self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256 - ): - use_n128_sfb_mapping = round_up( - work_tile.valid_tokens_in_cta_tile, 16 - ) <= Int32(128) + if cutlass.const_expr(self.quant_kind == QuantKind.nvfp4 and self.cta_tile_n == 256): + use_n128_sfb_mapping = round_up(work_tile.valid_tokens_in_cta_tile, 16) <= Int32(128) acc_pipeline.producer_acquire(acc_producer_state) for k_tile_idx in cutlass.range(k_tile_count, unroll=1): @@ -1726,15 +1499,9 @@ def run_mma_asymmetric( token_pipeline.consumer_wait(token_pipeline_state) weight_stage_coord = (None, None, None, None, weight_pipeline_state.index) - cute.copy( - tiled_copy_sfa, - partitioned_smem_sfa[weight_stage_coord], - partitioned_tmem_sfa, - ) + cute.copy(tiled_copy_sfa, partitioned_smem_sfa[weight_stage_coord], partitioned_tmem_sfa) if use_n128_sfb_mapping: - for instruction_index in cutlass.range_constexpr( - self.num_mma_instructions_per_ab_stage - ): + for instruction_index in cutlass.range_constexpr(self.num_mma_instructions_per_ab_stage): # Rest_Tiler order: N0K0, N128K0, N0K64, N128K64. lower_n_lower_k_source_coord = ( None, @@ -1762,11 +1529,7 @@ def run_mma_asymmetric( ) else: token_stage_coord = (None, None, None, None, token_pipeline_state.index) - cute.copy( - tiled_copy_sfb, - partitioned_smem_sfb[token_stage_coord], - partitioned_tmem_sfb, - ) + cute.copy(tiled_copy_sfb, partitioned_smem_sfb[token_stage_coord], partitioned_tmem_sfb) a_stage_frag = fragment_a[(None, None, None, weight_pipeline_state.index)] b_stage_frag = fragment_b[(None, None, None, token_pipeline_state.index)] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py index b894e963af64..b8d4910a0cdc 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel.py @@ -1,11 +1,13 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Fused pull-dispatch MegaMoE kernel composition for Rubin.""" +import hashlib + from typing import ClassVar, Optional, Tuple +import os + import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -14,30 +16,25 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cutlass_dsl import Boolean, Int32 -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.token_comm import TokenCommArgs, TokenCommNonDeterministic -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, +from .....communication.nvlink_domain.token_comm_cta_mainloop import TokenCommCtaMainloop +from .....helpers.helper_weight_ready_gate import ( + derive_helper_weight_ready_topology, + make_helper_weight_ready_generation_gate, ) +from .....communication.nvlink_domain.token_comm_dedup import TokenCommDedup +from .....communication.token_protocol import TokenSrcMetadata +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts +from .....helpers.fc2_ready_geometry import derive_fc2_ready_geometry from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up from .....quant_def import CombineFormat, QuantKind -from ....blackwell.custom_mix_cga_helpers import ( - TmaAtomOrPair, - pipeline_init_arrive_mixed_cga, - pipeline_init_wait_mixed_cga, -) -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import ( BlackwellFusedFc12Scheduler, @@ -45,9 +42,13 @@ minimum_phase_interleave_hint, ) from ....schedulers.non_clc_mixed_cga import NonClcMixedCgaConfig +from ...custom_mix_cga_helpers import TmaAtomOrPair, pipeline_init_arrive_mixed_cga, pipeline_init_wait_mixed_cga from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue +from .block_scaled_swap_ab_fc12_epilogue_dedup import GatedActEpilogueArgsDedup, SwapABGatedActEpilogueDedup from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension from .block_scaled_swap_ab_fc12_mainloop import BlockScaledSwapAbFc12Mainloop +from .topk_reduce_dedup import TopkReduceDedup + _aot_symbol_prefix = "rubin_mega_moe_aot" @@ -111,11 +112,25 @@ def impl_desc_require(cls) -> dict[str, type]: "fallback_cluster_count": OptionalRequirement(Optional[int]), "token_in_flag_batch": int, "token_back_mode": str, + "comm_clusters": OptionalRequirement(int), + "token_back_on_compute_ctas": OptionalRequirement(bool), + "comm_warps": OptionalRequirement(tuple), + "comm_stage_depths": OptionalRequirement(tuple), + "token_back_chunk_bytes": OptionalRequirement(int), + "dispatch_sf_group": OptionalRequirement(int), + "token_back_schedule_mode": OptionalRequirement(str), + "token_back_ready_granularity": OptionalRequirement(str), "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), + "combine_dedup": OptionalRequirement(bool), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) self.expert_count = problem_desc["expert_count"] self.intermediate_gateup_size = problem_desc["intermediate_gateup_size"] @@ -157,13 +172,54 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.launch_cluster_count = impl_desc["launch_cluster_count"] self.token_in_flag_batch = impl_desc["token_in_flag_batch"] self.token_back_mode = impl_desc["token_back_mode"] + self.token_back_ready_granularity = impl_desc.get("token_back_ready_granularity", "expert") self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] + self.comm_clusters = impl_desc.get("comm_clusters", 0) + self.token_back_on_compute_ctas = impl_desc.get("token_back_on_compute_ctas", False) + self.compute_cluster_count = self.launch_cluster_count - self.comm_clusters + # Left unbound rather than None when absent: the framework skips an optional field that was never + # bound but type-checks one that was, and only a split reads this, so there is no fused tuple to + # fall back to. + if impl_desc.get("comm_warps") is not None: + self.comm_warps = tuple(impl_desc["comm_warps"]) + self.comm_stage_depths = tuple(impl_desc.get("comm_stage_depths", (1, 1))) + # The fused default follows token_comm's own rather than restating it, so the two cannot drift. + self.token_back_chunk_bytes = impl_desc.get( + "token_back_chunk_bytes", TokenCommNonDeterministic.standalone_chunk_bytes + ) + self.dispatch_sf_group = impl_desc.get("dispatch_sf_group", 1) + # Dedicated communication CTAs use the static staged loop by default. + # Fused kernels retain the original work-ID-dependent token-back schedule. + default_schedule = ( + "static" if self.comm_clusters > 0 or self.work_id_mode != "atomic_counter" else "atomic_counter" + ) + self.token_back_schedule_mode = impl_desc.get("token_back_schedule_mode", default_schedule) + + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") + self.combine_dedup = bool(impl_desc.get("combine_dedup", False)) + # Diagnostic only (MEGA_DIAG_SKIP_TOKEN_BACK=1): combine warps do nothing, so the timed kernel is + # MMA + epilogue + dispatch alone. Output is wrong; run with MEGA_TIMING_ONLY=1. + self._diag_skip_token_back = os.environ.get("MEGA_DIAG_SKIP_TOKEN_BACK", "0") == "1" self.occupancy = 1 self.architecture = "sm_107" self.local_expert_count = self.expert_count // self.world_size - self.threads_per_cta = 16 * 32 if self.token_back_mode == "standalone_warps" else 12 * 32 # nosec B105 - self.other_warp_register_count = 64 if self.token_back_mode == "standalone_warps" else 72 # nosec B105 + if self.helper_expert_count > 0: + helper_topology = derive_helper_weight_ready_topology( + self.local_expert_count, + self.helper_expert_count, + ) + self.home_expert_count = helper_topology.home_expert_count + else: + # Preserve the baseline's runtime-expert-count specialization. A + # static M is required only when an S>0 helper boundary must be + # compiled into the scheduler gate. + self.home_expert_count = self.local_expert_count + n_warps = 16 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 12 + self.threads_per_cta = n_warps * 32 + self.other_warp_register_count = 64 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 72 self.mixed_cga_config = NonClcMixedCgaConfig( preferred_cluster_shape=self.cluster_shape_mn, @@ -176,9 +232,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.preferred_cluster_count = self.mixed_cga_config.preferred_cluster_count self.fallback_cluster_count = self.mixed_cga_config.fallback_cluster_count self.resolved_fallback_cluster_shape_mn = ( - self.cluster_shape_mn - if self.fallback_cluster_shape_mn is None - else self.fallback_cluster_shape_mn + self.cluster_shape_mn if self.fallback_cluster_shape_mn is None else self.fallback_cluster_shape_mn ) if self.launch_cluster_count != self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred: raise ValueError( @@ -187,6 +241,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) self._validate_geometry() + self._validate_cta_partition() self._resolve_schedule_hint() mma_cta_count = 2 if self.use_2cta_instrs else 1 mma_instruction = Tcgen05MmaInstruction( @@ -200,35 +255,44 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: sf_vec_size=self.sf_vec_size, ) tmem_plan = make_tcgen05_tmem_plan(mma_instruction, self.architecture, self.mma_tiler_mnk) - tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] - instruction_cta_m = self.mma_instruction_mnk[0] // mma_cta_count - instruction_m_repetitions = self.mma_tiler_mnk[0] // self.mma_instruction_mnk[0] - hidden_per_fc2_cluster_tile = ( - instruction_cta_m * instruction_m_repetitions * self.cluster_shape_mn[0] - ) - fc2_done_signals_per_token_tile = ( - ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) - * self.cluster_shape_mn[0] - * self.cluster_shape_mn[1] - ) + if self.token_back_ready_granularity == "token_tile": # nosec B105 + tokens_per_fc1_ready_slot, fc2_done_signals_per_token_tile = derive_fc2_ready_geometry( + hidden_size=self.hidden_size, + mma_tiler_mnk=self.mma_tiler_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_2cta_instrs=self.use_2cta_instrs, + fallback_cluster_shape_mn=self.fallback_cluster_shape_mn, + ) + else: + tokens_per_fc1_ready_slot = self.mma_tiler_mnk[1] * self.cluster_shape_mn[1] + instruction_cta_m = self.mma_instruction_mnk[0] // mma_cta_count + instruction_m_repetitions = self.mma_tiler_mnk[0] // self.mma_instruction_mnk[0] + hidden_per_fc2_cluster_tile = instruction_cta_m * instruction_m_repetitions * self.cluster_shape_mn[0] + fc2_done_signals_per_token_tile = ( + ceil_div(self.hidden_size, hidden_per_fc2_cluster_tile) + * self.cluster_shape_mn[0] + * self.cluster_shape_mn[1] + ) token_comm_impl_desc = ImplDesc( { **impl_desc, + **self._transfer_impl_desc_fields(), "fallback_cluster_shape_mn": self.fallback_cluster_shape_mn, "preferred_cluster_count": self.preferred_cluster_count, "fallback_cluster_count": self.fallback_cluster_count, + "token_back_ready_granularity": self.token_back_ready_granularity, "tokens_per_fc1_ready_slot": tokens_per_fc1_ready_slot, "fc2_done_signals_per_token_tile": fc2_done_signals_per_token_tile, "promised_launchable_sm_count": self.mixed_cga_config.total_cta_cnt, - "token_back_schedule_mode": ( - "atomic_counter" if self.work_id_mode == "atomic_counter" else "static" - ), - "router_smem_limit_bytes": cutlass.memory.get_smem_capacity_in_bytes( - self.architecture - ), + "token_back_schedule_mode": self.token_back_schedule_mode, + "router_smem_limit_bytes": utils.get_smem_capacity_in_bytes(self.architecture), } ) - self.token_comm = TokenCommNonDeterministic(problem_desc, token_comm_impl_desc) + # One top-level assignment on purpose: tests/test_dlb_contracts.py replays the constructor's + # descriptor-forwarding statements and only sees top-level assignments to self.token_comm. + self.token_comm = (TokenCommDedup if self.combine_dedup else TokenCommNonDeterministic)( + problem_desc, token_comm_impl_desc + ) self.max_tokens = self.token_comm.worst_case_token_count fc12_problem_desc = ProblemDesc({**problem_desc, "expert_count": self.local_expert_count}) @@ -248,15 +312,32 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "token_back_push_data": self.token_comm.token_back_push_data, "fc1_epi_flag_batch": self.epi_flag_batches[0], "fc2_epi_flag_batch": self.epi_flag_batches[1], + "compute_cluster_count": self.compute_cluster_count, + "token_back_ready_granularity": self.token_back_ready_granularity, + # Combine dedup: the dedup FC2 epilogue scatters rows into the local order buffer + # [T, R, chunks_per_row, K, chunk] whose geometry TokenCommDedup owns. + **( + { + "dedup_world_size": self.world_size, + "dedup_topk": self.topk, + "dedup_chunk_bytes": self.token_comm.order_chunk_bytes, + "dedup_chunks_per_row": self.token_comm.order_chunks_per_row, + } + if self.combine_dedup + else {} + ), } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) - self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) + if self.combine_dedup: + self.epilogue = SwapABGatedActEpilogueDedup(fc12_problem_desc, resolved_impl_desc) + else: + self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) + if self.token_back_ready_granularity == "token_tile": # nosec B105 + self._validate_token_tile_ready_contract() if self.epilogue.fc1_output_dtype is not self.b_dtype: raise ValueError("Epilogue FC1 output dtype must match the Mainloop B dtype.") if self.epilogue.fc1_output_sf_dtype is not self.sf_dtype: @@ -264,16 +345,153 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: if self.epilogue.sf_vec_size != self.sf_vec_size: raise ValueError("Epilogue and Mainloop scale vector sizes must match.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) - self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + self._comm_cta_mainloop = None + self._comm_smem_workspace = None + if self.comm_clusters > 0: + self._comm_cta_mainloop, self._comm_smem_workspace = self._build_comm_cta_and_smem() + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) + # Every CTA allocates the same arena and the two role tables are laid over it, each CTA reading only + # the one for the role it took. The arena therefore has to be the larger role's, and the workspace + # that owns it is whichever that is. + self._arena_workspace = self._smem_workspace + if ( + self._comm_smem_workspace is not None + and self._comm_smem_workspace.total_bytes > self._smem_workspace.total_bytes + ): + self._arena_workspace = self._comm_smem_workspace + self.smem_arena_bytes = self._arena_workspace.total_bytes + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + if self.smem_arena_bytes > smem_limit: + comm_bytes = 0 if self._comm_smem_workspace is None else self._comm_smem_workspace.total_bytes + raise ValueError( + f"CTA shared-memory arena requires {self.smem_arena_bytes} bytes " + f"(compute={self._smem_workspace.total_bytes}, comm={comm_bytes}), exceeding the " + f"{smem_limit}-byte limit for {self.architecture} at occupancy {self.occupancy}." + ) + if self._comm_smem_workspace is not None: + self._comm_smem_workspace.assert_fits_arena(self.smem_arena_bytes) + self._smem_workspace.assert_fits_arena(self.smem_arena_bytes) + if self.reduce_topk_in_kernel: + self._topk_reduce = None + elif self.combine_dedup: + self._topk_reduce = TopkReduceDedup(self.hidden_size, self.world_size, self.combine_format) + else: + self._topk_reduce = TopkReduce(self.hidden_size, self.topk, self.combine_format) + + def _validate_cta_partition(self) -> None: + for name, value in (("launch_cluster_count", self.launch_cluster_count), ("comm_clusters", self.comm_clusters)): + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{name} must be a Python int, got {value!r}.") + if not 0 <= self.comm_clusters < self.launch_cluster_count: + raise ValueError("comm_clusters must leave at least one compute cluster.") + if len(self.comm_stage_depths) != 2 or any( + isinstance(depth, bool) or not isinstance(depth, int) or not 1 <= depth <= 32 + for depth in self.comm_stage_depths + ): + raise ValueError("comm_stage_depths must be two ints in [1, 32], one phase bit per stage.") + if self.token_back_chunk_bytes <= 0 or self.token_back_chunk_bytes % 16: + raise ValueError("token_back_chunk_bytes must be a positive multiple of 16.") + if self.dispatch_sf_group not in (1, 4): + raise ValueError("dispatch_sf_group must be 1 or 4.") + if self.comm_clusters == 0: + if self.token_back_on_compute_ctas: + raise ValueError("token_back_on_compute_ctas requires a nonzero comm_clusters split.") + return + warps = getattr(self, "comm_warps", None) + if ( + warps is None + or len(warps) != 2 + or any(isinstance(w, bool) or not isinstance(w, int) or w <= 0 for w in warps) + ): + raise ValueError(f"A split needs comm_warps as two positive ints, got {warps}.") + if self.token_back_mode != "standalone_warps": # nosec B105 + raise ValueError("A CTA split requires token_back_mode='standalone_warps'.") + dispatch_warps, back_warps = warps + # Check the warps that actually coexist in each concrete CTA branch. + required_warps = ( + max(dispatch_warps, self.transfer_warp_idx_start + back_warps) + if self.token_back_on_compute_ctas + else dispatch_warps + back_warps ) + if required_warps > self.threads_per_cta // 32: + raise ValueError( + f"comm_warps={warps} requires {required_warps} warps per CTA, " + f"but the launch provides only {self.threads_per_cta // 32}." + ) + dispatch_depth = self.comm_stage_depths[0] + if dispatch_depth > 1 and (self.dispatch_sf_group != 4 or dispatch_depth % 4): + raise ValueError("Multistage token-in needs dispatch_sf_group=4 and a depth divisible by four.") + + def _transfer_impl_desc_fields(self) -> dict: + return { + "comm_cta_count": self.comm_clusters * self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + "comm_stage_depths": self.comm_stage_depths, + "token_back_chunk_bytes": self.token_back_chunk_bytes, + "dispatch_sf_group": self.dispatch_sf_group, + } + + def _is_dispatch_cta(self, cluster_idx): + # Partition preferred-shaped slots, keeping each actual MMA cluster intact. + return (cluster_idx * self.comm_clusters) % self.launch_cluster_count < self.comm_clusters + + def _compute_cluster_ordinal(self, cluster_idx): + preceding_dispatch = ( + cluster_idx * self.comm_clusters + self.launch_cluster_count - 1 + ) // self.launch_cluster_count + return cluster_idx - preceding_dispatch + + def _dispatch_cta_ordinal(self, cluster_idx, physical_cta_offset): + # An active fallback cluster rank restarts within a preferred slot; blockIdx.x does not. + slot = (cluster_idx * self.comm_clusters) // self.launch_cluster_count + return slot * self.cluster_shape_mn[0] + physical_cta_offset + + def _combine_cta_ordinal(self, cluster_idx, physical_cta_offset): + if self.token_back_on_compute_ctas: + return self._compute_cluster_ordinal(cluster_idx) * self.cluster_shape_mn[0] + physical_cta_offset + return self._dispatch_cta_ordinal(cluster_idx, physical_cta_offset) + + def _validate_token_tile_ready_contract(self) -> None: + """Check the consumer target against the actual scheduler and publisher. + + Both preferred clusters and a complete fallback group cover the same + logical CTA coordinates. Using active-cluster dimensions here would + permit token-back before the other fallback clusters finish their data. + """ + if not self.epilogue.fc2_use_tma: + raise ValueError("token_tile return-ready requires FC2 TMA stores.") + mapping_tile = self.scheduler.mapping_cta_tile_shape_mnk + mapping_cluster = self.scheduler.mapping_cluster_shape_mn + if ( + mapping_cluster != (1, self.cluster_shape_mn[0]) + or mapping_tile[:2] != (self.epilogue.cta_tile_n, self.epilogue.cta_tile_m) + ): + raise ValueError("token_tile scheduler and epilogue must share the preferred logical CTA geometry.") + expected_signals = ceil_div(self.hidden_size, mapping_tile[1] * mapping_cluster[1]) * mapping_cluster[1] + if ( + self.token_comm.tokens_per_fc1_ready_slot != mapping_tile[0] + or self.token_comm.fc2_done_signals_per_token_tile != expected_signals + ): + raise ValueError("token_tile producer and consumer must share logical ready slots and completion counts.") def _validate_geometry(self) -> None: + if self.token_back_ready_granularity not in ("expert", "token_tile"): + raise ValueError(f"Unsupported token_back_ready_granularity {self.token_back_ready_granularity!r}.") + if self.token_back_ready_granularity == "token_tile": # nosec B105 + if not ( + self.schedule_mode in ("grouped", "phase_interleave") + and self.work_id_mode == "atomic_counter" + and self.token_back_mode in ("standalone_warps", "reuse_dispatch_warps") + and self.use_2cta_instrs + and self.fc2_use_bulk + and self.quant_kind == QuantKind.nvfp4 + and not self.reduce_topk_in_kernel + ): + raise ValueError( + "token_tile return-ready requires grouped or phase_interleave/atomic_counter, " + "standalone_warps or reuse_dispatch_warps, " + "two-CTA MMA with hidden-only preferred/fallback clusters, " + "FC2 TMA stores, NVFP4, and separate top-k reduction." + ) static_expert_dimensions = ( isinstance(self.expert_count, int), isinstance(self.intermediate_gateup_size, int), @@ -282,9 +500,7 @@ def _validate_geometry(self) -> None: if any(static_expert_dimensions) and not all(static_expert_dimensions): raise ValueError("MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The MegaMoE Kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The MegaMoE Kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -292,11 +508,38 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") + if self.topk > TokenSrcMetadata.MAX_TOPK or self.world_size > TokenSrcMetadata.MAX_RANKS: raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." + f"TokenSrcMetadata keeps topk and rank in {TokenSrcMetadata.FIELD_BITS} bits each; " + f"topk must be <= {TokenSrcMetadata.MAX_TOPK} and world_size <= {TokenSrcMetadata.MAX_RANKS}." + ) + if self.max_tokens_per_rank > TokenSrcMetadata.MAX_TOKENS: + raise ValueError( + f"TokenSrcMetadata keeps the token index in {TokenSrcMetadata.TOKEN_BITS} bits; " + f"max_tokens_per_rank must be <= {TokenSrcMetadata.MAX_TOKENS}." ) if self.reduce_topk_in_kernel and not self.apply_topk_at_fc1: raise ValueError("In-kernel top-k reduction requires apply_topk_at_fc1=True.") + if self.combine_dedup: + if not self.apply_topk_at_fc1: + raise ValueError("combine_dedup applies top-k scores at FC1 only (pre_fc1_quant).") + if self.combine_format.is_quantized or self.combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("combine_dedup supports the BF16 combine format only.") + if self.world_size > 32: + raise ValueError("combine_dedup keeps the destination-rank mask in an Int32; world_size must be <= 32.") + if self.token_back_mode == "epi_warps": # nosec B105 + raise ValueError( + "combine_dedup pushes the reduced rank-major rows from the combine warps; use " + "reuse_dispatch_warps or standalone_warps." + ) + if not self.fc2_use_bulk: + raise ValueError("combine_dedup stages FC2 subtiles in SMEM for per-token bulk stores; fc2_use_bulk must be on.") + if self.comm_clusters > 0: + raise ValueError("combine_dedup does not support the communication CTA split (comm_clusters must be 0).") + if self.token_back_ready_granularity != "expert": # nosec B105 + # The dedup combine warps poll one fc2_done counter per local expert (expert-prefix wait). + raise ValueError("combine_dedup signals FC2 completion per expert; token_back_ready_granularity must be 'expert'.") if self.intermediate_gateup_size % 2 != 0: raise ValueError("The SwiGLU intermediate dimension must be even.") intermediate_downproj = self.intermediate_gateup_size // 2 @@ -312,12 +555,8 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: @@ -335,20 +574,14 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") if self.quant_kind == QuantKind.mxfp4_mxfp8: # Rubin consumes the mixed FP4 weight from native packed SMEM. Its K-major source only # needs ordinary 16-byte TMA alignment and complete SFVec32 blocks. - for name, extent in ( - ("hidden_size", self.hidden_size), - ("intermediate_downproj", intermediate_downproj), - ): + for name, extent in (("hidden_size", self.hidden_size), ("intermediate_downproj", intermediate_downproj)): if extent * int(self.a_dtype.width) % 128 != 0 or extent % self.sf_vec_size != 0: raise ValueError( f"Rubin {self.quant_kind} requires {name} to satisfy ordinary TMA and " @@ -372,9 +605,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin MegaMoE does not implement M/N instruction repetition or B-reuse." - ) + raise NotImplementedError("Rubin MegaMoE does not implement M/N instruction repetition or B-reuse.") if mma_n not in (64, 128, 256): raise ValueError(f"mma_tiler N must be 64, 128, or 256, got {mma_n}.") if mma_k % instruction_k != 0: @@ -393,9 +624,7 @@ def _validate_geometry(self) -> None: f"four-vector SF atom ({sf_atom_size} elements)." ) elif self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if cluster_n != 1: raise ValueError(f"The swap-AB FC12 path requires cluster N=1, got {cluster_n}.") if cluster_m <= 0 or cluster_m > 16 or cluster_m & (cluster_m - 1): @@ -405,14 +634,8 @@ def _validate_geometry(self) -> None: if self.fallback_cluster_shape_mn is not None: fallback_cluster_m, fallback_cluster_n = self.fallback_cluster_shape_mn if fallback_cluster_n != 1: - raise ValueError( - f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}." - ) - if ( - fallback_cluster_m <= 0 - or fallback_cluster_m > 16 - or fallback_cluster_m & (fallback_cluster_m - 1) - ): + raise ValueError(f"The swap-AB FC12 fallback path requires cluster N=1, got {fallback_cluster_n}.") + if fallback_cluster_m <= 0 or fallback_cluster_m > 16 or fallback_cluster_m & (fallback_cluster_m - 1): raise ValueError("fallback cluster M must be a power of two no greater than 16.") if self.use_2cta_instrs and fallback_cluster_m % 2 != 0: raise ValueError("Two-CTA MMA requires an even fallback cluster M.") @@ -428,7 +651,7 @@ def _resolve_schedule_hint(self) -> None: raw_minimum_hint = minimum_phase_interleave_hint( blocks_fc1=blocks_fc1, blocks_fc2=blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.mixed_cga_config.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) minimum_safe_hint = round_up(raw_minimum_hint, self.epi_flag_batches[0]) if self.hint is None: @@ -508,6 +731,7 @@ def dtype_name(dtype: type) -> str: f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" f"e{self.expert_count}_ep{self.world_size}_topk{self.topk}_" + f"helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -526,6 +750,10 @@ def dtype_name(dtype: type) -> str: f"situ{self.situ_beta}x{self.situ_linear_beta}_" f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_redg' if self.reduce_topk_in_kernel else 'separate_reduce'}" + "_dlb_v1" + + ("_fc2readytile_v1" if self.token_back_ready_granularity == "token_tile" else "") # nosec B105 + + f"_{'dedup' if self.combine_dedup else 'nodedup'}" + + ("_diagskiptb" if self._diag_skip_token_back else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -547,14 +775,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -563,20 +787,14 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged, not a kernel choice. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) # TopkReduce and the in-kernel REDG both take their output dtype from this tensor. output_dtype = cutlass.BFloat16 fake_arguments = dict( # Swap-AB puts the weights on the MMA A operand, so the token wire dtype is TokenComm's, not a_dtype. - activation=fake_tensor( - self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16 - ), + activation=fake_tensor(self.token_comm.activation_dtype, (tokens, hidden), (1, 0), {0}, 16), activation_sf=fake_tensor( self.token_comm.activation_sf_dtype, (tokens, self.token_comm.activation_sf_hidden_padded), @@ -586,18 +804,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {0}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {0}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {0, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {0}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {0, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {0}, 16), output_activation=fake_tensor(output_dtype, (tokens, hidden), (1, 0), {0}, 16), # Opaque byte workspaces: a bare base pointer, like to_cute_ptr. The 128-byte promise matches the # alignment the registered regions assume, so the caller must hand over 128-byte-aligned bases. @@ -612,6 +822,16 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr( + cutlass.Uint64, + 0, + AddressSpace.gmem, + assumed_align=8, + ), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: # nvfp4 carries per-expert dequant scalars. Under an e8m0 scale they do not exist at # all, and declaring them here would put three tensors in the ABI that the caller has @@ -626,22 +846,40 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): if out_path is None: return compiled compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True + out_path, + function_name=self._aot_symbol_name(), + export_only_tvm_ffi_symbols=True, ) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) + + def _build_comm_cta_and_smem(self) -> Tuple[TokenCommCtaMainloop, SmemWorkspace]: + mainloop = TokenCommCtaMainloop( + self.token_comm, token_back_on_compute_ctas=self.token_back_on_compute_ctas + ) + workspace = SmemWorkspace() + mainloop.register_smem_regions(workspace) + workspace.finalize() + return mainloop, workspace def _build_mainloop_and_smem( self, problem_desc: ProblemDesc, resolved_impl_desc: ImplDesc ) -> Tuple[BlockScaledSwapAbFc12Mainloop, SmemWorkspace]: - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() - self.token_comm.register_smem_regions(smem_workspace) + if self.comm_clusters == 0: + self.token_comm.register_smem_regions(smem_workspace) + elif self.token_back_on_compute_ctas: + self._comm_cta_mainloop.register_token_back_smem(smem_workspace) self.scheduler.register_smem_regions(smem_workspace) self.epilogue.register_smem_regions(smem_workspace) fixed_component_bytes = smem_workspace.estimate_total_bytes() @@ -650,15 +888,14 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12Mainloop(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) smem_workspace.finalize(max_bytes=smem_limit) return mainloop, smem_workspace + @cute.jit def __call__( self, @@ -678,16 +915,34 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch router, fused MegaMoE compute, and optional TopK reduce.""" + helper_ready_bundle = ( + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, + ) + if cutlass.const_expr( + ( + self.helper_expert_count > 0 + and not all(value is not None for value in helper_ready_bundle) + ) + or ( + self.helper_expert_count == 0 + and any(value is not None for value in helper_ready_bundle) + ) + ): + raise ValueError( + "helper READY flags/generation must both be present iff helper_expert_count > 0" + ) + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -776,19 +1031,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * padded_hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(padded_hidden), - cutlass.Int32(experts), - ), + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(padded_hidden), cutlass.Int32(experts)), sf_vec_size, ), ) @@ -806,31 +1056,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) if cutlass.const_expr(self.reduce_topk_in_kernel): @@ -838,28 +1078,24 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) pre_reduced_activation_sf = None else: - pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor( - self._device_workspace - ) - pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor( - self._device_workspace - ) + pre_reduced_activation = self.token_comm.pre_reduced_activation_tensor(self._device_workspace) + pre_reduced_activation_sf = self.token_comm.pre_reduced_activation_sf_tensor(self._device_workspace) + dedup_valid_rank_mask = None + if cutlass.const_expr(self.combine_dedup): + dedup_valid_rank_mask = self.token_comm.valid_rank_mask_tensor(self._device_workspace) if cutlass.const_expr(self.token_comm.token_back_push_data): + # Local FC2 output the combine warps push from: the local-expert-major pool, or under dedup the + # (T*R*K, 1, H) order buffer the STG store scatters into by metadata. fc2_output = self.token_comm.fc2_activation_tensor(self._device_workspace) else: fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) fc2_output_sf = self.token_comm.fc2_activation_sf_tensor(self._device_workspace) @@ -898,7 +1134,11 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: grid = self.scheduler.get_grid_shape(max_active_clusters=self.launch_cluster_count) self._device_workspace.remove_device_members() - kernel = self._kernel( + if cutlass.const_expr(self.comm_clusters > 0): + kernel_entry = self._kernel_cta_specialized + else: + kernel_entry = self._kernel + kernel = kernel_entry( fc1_tma_a_tensor, fc1_tma_a_atom, fc1_tma_sfa_tensor, @@ -933,6 +1173,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ) if cutlass.const_expr(self.mixed_cga_config.is_mixed): kernel.launch( @@ -954,14 +1196,15 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: use_pdl=True, ) if cutlass.const_expr(not self.reduce_topk_in_kernel): - reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, - pre_reduced_activation_sf, - output_activation, - reduce_scores, - stream, - ) + if cutlass.const_expr(self.combine_dedup): + # Weights are already folded (FC1 pre-multiply or FC2-epilogue multiply); sum the rank slots. + self._topk_reduce(pre_reduced_activation, dedup_valid_rank_mask, output_activation, stream) + else: + reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores + self._topk_reduce( + pre_reduced_activation, pre_reduced_activation_sf, output_activation, reduce_scores, stream + ) + @cute.kernel def _kernel( @@ -1000,6 +1243,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose TokenComm, Scheduler, Mainloop, and Epilogue.""" self._mainloop.materialize_codegen_members() @@ -1021,8 +1266,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -1033,9 +1277,7 @@ def _kernel( is_fallback_cluster = Boolean(False) if cutlass.const_expr(self.mixed_cga_config.is_mixed): active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() - is_fallback_cluster = active_cluster_m == Int32( - self.resolved_fallback_cluster_shape_mn[0] - ) + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) cta_rank_in_cluster = cute.arch.block_idx_in_cluster() cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) @@ -1058,12 +1300,7 @@ def _kernel( pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) self._mainloop.assign_device_members( - self._smem_workspace, - smem_base, - cta_coord_in_cluster, - is_fallback_cluster, - hidden, - intermediate_gateup, + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup ) ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) tma_a_pipeline_state = pipeline.make_pipeline_state( @@ -1091,39 +1328,51 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer( - self._device_workspace - ), - ) - optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), ) + # FC1 pre-multiplies under apply_topk_at_fc1 (the only position combine_dedup supports). + fc1_pool_topk_scores = pool_topk_scores if cutlass.const_expr(self.apply_topk_at_fc1) else None + if cutlass.const_expr(self.combine_dedup): + optional_epilogue_args = GatedActEpilogueArgsDedup( + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=fc1_pool_topk_scores, + fc2_remote_output=pre_reduced_activation, + local_rank=local_rank, + ) + else: + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, + fc2_alpha=fc2_alpha, + fc1_norm_const=fc1_norm_const, + topk_scores=fc1_pool_topk_scores, + ) pipeline_init_arrive_mixed_cga( - self.cluster_shape_mn, - self.resolved_fallback_cluster_shape_mn, - is_fallback_cluster, - is_relaxed=True, + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True ) pipeline_init_wait_mixed_cga( self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster ) finalize_barrier_thread_count = ( - 12 * 32 if self.token_back_mode == "standalone_warps" else 8 * 32 # nosec B105 - ) - finalize_barrier = pipeline.NamedBarrier( - barrier_id=13, num_threads=finalize_barrier_thread_count + 12 * 32 if self.token_back_mode in ("standalone_warps", "hybrid_warps") else 8 * 32 ) + finalize_barrier = pipeline.NamedBarrier(barrier_id=13, num_threads=finalize_barrier_thread_count) if warp_idx == self.scheduler_warp_id: cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) @@ -1138,9 +1387,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1235,9 +1490,7 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) iket.range_pop() finalize_barrier.arrive_unaligned() @@ -1247,9 +1500,15 @@ def _kernel( iket.range_push("mega.token_in") self.token_comm.token_in(self._smem_workspace, smem_base) iket.range_pop() + # epi_warps pushes the data from the epilogue, but a quantized wire still needs its scale rows pushed + # token-contiguously (token_back_push_sf); the dispatch warps do that after token_in. if cutlass.const_expr( self.token_comm.token_back_enabled - and self.token_back_mode != "standalone_warps" # nosec B105 + and ( + self.token_back_mode in ("reuse_dispatch_warps", "hybrid_warps") + or (self.token_back_mode == "epi_warps" and not self.token_comm.token_back_push_data) # nosec B105 + ) + and not self._diag_skip_token_back ): iket.range_push("mega.token_back") self.token_comm.token_back(self._smem_workspace, smem_base) @@ -1258,10 +1517,309 @@ def _kernel( iket.range_push("mega.tail_reset") self.token_comm.reset_tail() iket.range_pop() - elif cutlass.const_expr(self.token_back_mode == "standalone_warps"): # nosec B105 + elif cutlass.const_expr(self.token_back_mode in ("standalone_warps", "hybrid_warps")): iket.range_push("mega.token_back_standalone") - self.token_comm.token_back(self._smem_workspace, smem_base) + if cutlass.const_expr(not self._diag_skip_token_back): + self.token_comm.token_back(self._smem_workspace, smem_base) iket.range_pop() finalize_barrier.arrive_unaligned() self.token_comm.remove_device_members() self._device_workspace.remove_device_members() + + + @cute.kernel + def _kernel_cta_specialized( + self, + fc1_tma_a_tensor: cute.Tensor, + fc1_tma_a_atom: cute.CopyAtom, + fc1_tma_sfa_tensor: cute.Tensor, + fc1_tma_sfa_atom: cute.CopyAtom, + fc2_tma_a_tensor: cute.Tensor, + fc2_tma_a_atom: cute.CopyAtom, + fc2_tma_sfa_tensor: cute.Tensor, + fc2_tma_sfa_atom: cute.CopyAtom, + fc1_tma_b_tensor: cute.Tensor, + fc1_tma_b_atom_or_pair: TmaAtomOrPair, + fc1_tma_sfb_tensor: cute.Tensor, + fc1_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc2_tma_b_tensor: cute.Tensor, + fc2_tma_b_atom_or_pair: TmaAtomOrPair, + fc2_tma_sfb_tensor: cute.Tensor, + fc2_tma_sfb_atom_or_pair: TmaAtomOrPair, + fc1_output_tma_atom: cute.CopyAtom, + fc1_output_tma_tensor: cute.Tensor, + fc2_output_tma_atom: Optional[cute.CopyAtom], + fc2_output_tma_tensor: Optional[cute.Tensor], + fc2_output: cute.Tensor, + fc2_output_sf: Optional[cute.Tensor], + actual_expert_shape: Tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32], + activation: cute.Tensor, + activation_sf: cute.Tensor, + pre_reduced_activation: cute.Tensor, + pre_reduced_activation_sf: Optional[cute.Tensor], + peer_rank_ptr_mapper, + local_rank: cutlass.Int32, + local_workspace: cute.Pointer, + shared_workspace: cute.Pointer, + fc1_alpha: Optional[cute.Tensor], + fc2_alpha: Optional[cute.Tensor], + fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], + ): + """Run communication and FC12 on distinct, whole-cluster CTA sets.""" + storage_type = self._arena_workspace.storage_class() + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(storage_type) + smem_base = storage.buffer.data_ptr() + workspace = self._device_workspace + workspace.assign_device_members(local_workspace, shared_workspace) + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = thread_idx // 32 + block_idx = cute.arch.block_idx() + active_cluster_m = Int32(self.cluster_shape_mn[0]) + is_fallback_cluster = Boolean(False) + if cutlass.const_expr(self.mixed_cga_config.is_mixed): + active_cluster_m, _, _ = cute.arch.block_in_cluster_dim() + is_fallback_cluster = active_cluster_m == Int32(self.resolved_fallback_cluster_shape_mn[0]) + cta_rank_in_cluster = cute.arch.block_idx_in_cluster() + cta_coord_in_cluster = (cta_rank_in_cluster, Int32(0), Int32(0)) + linear_cta_idx = block_idx[0] + block_idx[2] * Int32(self.cluster_shape_mn[0]) + token_comm_args = TokenCommArgs( + activation=activation, + activation_sf=activation_sf, + pre_reduced_activation=pre_reduced_activation, + pre_reduced_activation_sf=pre_reduced_activation_sf, + peer_rank_ptr_mapper=peer_rank_ptr_mapper, + ) + self.token_comm.assign_device_members( + device_workspace=workspace, + token_comm_args=token_comm_args, + local_rank=local_rank, + linear_cta_idx=linear_cta_idx, + dispatch_cta_ordinal=self._dispatch_cta_ordinal(block_idx[2], block_idx[0]), + combine_cta_ordinal=self._combine_cta_ordinal(block_idx[2], block_idx[0]), + ) + if self._is_dispatch_cta(block_idx[2]): + self._comm_cta_mainloop.run(self._comm_smem_workspace, smem_base) + else: + compute_block_idx = (block_idx[0], block_idx[1], self._compute_cluster_ordinal(block_idx[2])) + self._mainloop.materialize_codegen_members() + fc1_done_counter = self._device_workspace.tensor(self.fc1_done_counter_region) + fc1_output_sf_storage = self._device_workspace.tensor(self.fc1_output_sf_region) + if cutlass.const_expr(isinstance(self.hidden_size, int)): + hidden = self.hidden_size + else: + hidden = actual_expert_shape[2] + if cutlass.const_expr(isinstance(self.intermediate_gateup_size, int)): + intermediate_gateup = self.intermediate_gateup_size + else: + intermediate_gateup = actual_expert_shape[1] + fc1_output_sf = cute.make_tensor( + fc1_output_sf_storage.iterator, + blockscaled_utils.tile_atom_to_shape_SF( + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size + ), + ) + + expert_token_sizes = self.token_comm.local_expert_sizes(self._device_workspace, local_rank) + token_src_metadata = self.token_comm.token_src_metadata_tensor(self._device_workspace) + fc2_done_counter = self.token_comm.fc2_done_counter_tensor(self._device_workspace) + pool_topk_scores = self.token_comm.fc1_topk_scores_tensor(self._device_workspace) + + self._mainloop.assign_device_members( + self._smem_workspace, smem_base, cta_coord_in_cluster, is_fallback_cluster, hidden, intermediate_gateup + ) + ab_pipeline = self._mainloop.create_ab_pipeline(self._smem_workspace, smem_base) + tma_a_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + tma_b_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self._mainloop.num_ab_pipeline_stages + ) + mma_pipeline_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self._mainloop.num_ab_pipeline_stages + ) + acc_pipeline = self._mainloop.create_acc_pipeline(self._smem_workspace, smem_base) + tmem_allocator = self._mainloop.create_tmem_allocator( + self._smem_workspace, smem_base, allocator_warp_id=self.epilogue_warp_ids[0] + ) + self.scheduler.assign_device_members( + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=None, + actual_expert_shape=actual_expert_shape, + block_idx=compute_block_idx, + smem_workspace=self._smem_workspace, + smem_base=smem_base, + device_workspace=self._device_workspace, + is_fallback_cluster=is_fallback_cluster, + ) + scheduler = self.scheduler + + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) + kernel_extension = BlockScaledSwapAbFc12Extension( + sf_vec_size=self.sf_vec_size, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + fc1_ready_counter_pointer=self.token_comm.fc1_ready_counter_pointer(self._device_workspace), + ) + optional_epilogue_args = GatedActEpilogueArgs( + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores + ) + + pipeline_init_arrive_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster, is_relaxed=True + ) + pipeline_init_wait_mixed_cga( + self.cluster_shape_mn, self.resolved_fallback_cluster_shape_mn, is_fallback_cluster + ) + + if warp_idx == self.scheduler_warp_id: + cute.arch.warpgroup_reg_dealloc(self.scheduler_register_count) + iket.range_push("mega.scheduler") + if cutlass.const_expr(hasattr(scheduler, "initialize_fallback_group")): + iket.range_push("scheduler.register_group") + scheduler.initialize_fallback_group() + iket.range_pop() + iket.range_push("scheduler.wait_sizes_ready") + self.token_comm.wait_for_sizes_ready(self._device_workspace) + iket.range_pop() + iket.range_push("scheduler.gen_work") + work_tile = scheduler.gen_next_work() + iket.range_pop() + helper_ready_latched = Int32(0) + while work_tile.is_valid_tile: + iket.range_push("scheduler.publish_work") + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) + iket.range_pop() + iket.range_push("scheduler.gen_work") + work_tile = scheduler.gen_next_work() + iket.range_pop() + iket.range_push("scheduler.publish_tail") + scheduler.publish_work(work_tile) + scheduler.produce_tail() + iket.range_pop() + iket.range_pop() + + if warp_idx == self.tma_a_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.tma_a") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_a( + fc1_tma_a_tensor=fc1_tma_a_tensor, + fc1_tma_a_atom=fc1_tma_a_atom, + fc1_tma_sfa_tensor=fc1_tma_sfa_tensor, + fc1_tma_sfa_atom=fc1_tma_sfa_atom, + fc2_tma_a_tensor=fc2_tma_a_tensor, + fc2_tma_a_atom=fc2_tma_a_atom, + fc2_tma_sfa_tensor=fc2_tma_sfa_tensor, + fc2_tma_sfa_atom=fc2_tma_sfa_atom, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_a_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + ) + iket.range_pop() + + if warp_idx == self.tma_b_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.tma_b") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_tma_b( + fc1_tma_b_tensor=fc1_tma_b_tensor, + fc1_tma_b_atom_or_pair=fc1_tma_b_atom_or_pair, + fc1_tma_sfb_tensor=fc1_tma_sfb_tensor, + fc1_tma_sfb_atom_or_pair=fc1_tma_sfb_atom_or_pair, + fc2_tma_b_tensor=fc2_tma_b_tensor, + fc2_tma_b_atom_or_pair=fc2_tma_b_atom_or_pair, + fc2_tma_sfb_tensor=fc2_tma_sfb_tensor, + fc2_tma_sfb_atom_or_pair=fc2_tma_sfb_atom_or_pair, + ab_pipeline=ab_pipeline, + ab_pipeline_state=tma_b_pipeline_state, + sched_consumer=sched_consumer, + kernel_extension=kernel_extension, + fc1_done_counter_pointer=fc1_done_counter.iterator, + fc2_spin_threshold=fc2_spin_threshold, + ) + iket.range_pop() + + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + iket.range_push("mega.mma") + sched_consumer = scheduler.make_consumer() + self._mainloop.run_mma( + tmem_allocator=tmem_allocator, + ab_pipeline=ab_pipeline, + ab_pipeline_state=mma_pipeline_state, + acc_pipeline=acc_pipeline, + sched_consumer=sched_consumer, + ) + iket.range_pop() + + if warp_idx < len(self.epilogue_warp_ids): + cute.arch.warpgroup_reg_alloc(self.epilogue_register_count) + iket.range_push("mega.epilogue") + sched_consumer = scheduler.make_consumer() + tmem_allocator.allocate(self._mainloop.num_tmem_alloc_cols) + tmem_allocator.wait_for_alloc() + tmem_pointer = tmem_allocator.retrieve_ptr(self.acc_dtype) + self.epilogue.run( + self._smem_workspace, + smem_base, + tmem_pointer, + acc_pipeline, + sched_consumer, + kernel_extension, + fc1_output_tma_atom, + fc1_output_tma_tensor, + fc1_output_sf, + fc2_output_tma_atom, + fc2_output_tma_tensor, + fc2_output, + fc1_done_counter, + thread_idx, + token_src_metadata, + fc2_done_counter, + fc2_output_sf, + peer_rank_ptr_mapper, + optional_epilogue_args, + ) + tmem_allocator.relinquish_alloc_permit() + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) + iket.range_pop() + + # Donate registers from complete warpgroups, including idle warps. + if warp_idx >= Int32(self.transfer_warp_idx_start): + cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) + if cutlass.const_expr(self.token_back_on_compute_ctas): + if (warp_idx >= self.token_comm.combine_warp_start) & ( + warp_idx < self.token_comm.combine_warp_start + self.token_comm.combine_warp_count + ): + self.token_comm.token_back(self._smem_workspace, smem_base) + + # Every physical CTA participates in the common tail, including idle transfer warps. + # Each operation has already drained its own asynchronous copies before this CTA barrier. + iket.range_push("mega.kernel_tail_wait") + cute.arch.sync_threads() + iket.range_pop() + if (warp_idx >= Int32(self.transfer_warp_idx_start)) & (warp_idx < Int32(self.token_in_end_warp_idx)): + iket.range_push("mega.tail_reset") + self.token_comm.reset_tail() + iket.range_pop() + + self.token_comm.remove_device_members() + workspace.remove_device_members() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py index 2c864e63cc38..9d99afc3344a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/block_scaled_swap_ab_mega_moe_kernel_gen_specialized.py @@ -1,7 +1,5 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause """Generation-phase MegaMoE kernel composition for Rubin. Specialised for few tokens per rank and large expert parallelism. The narrowing @@ -21,16 +19,17 @@ and the FC2 epilogue warps send results to their owning rank themselves through the bulk reduce-add. Twelve warps cover every role. -The kernel owns the workspace tail reset. Because it outlives the communication -kernel, it is the only participant that can safely zero the counters, which it -does between two NVLink barriers so that no rank observes a half-cleared -workspace or exits before the clearing is visible. +The kernel owns the workspace tail reset. All warp roles finish before a grid +barrier allows the epilogue to clear the counters. A final NVLink barrier makes +the reset visible to every rank before the next launch reuses the workspace. """ +import hashlib import math from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline @@ -40,36 +39,28 @@ from cutlass.cutlass_dsl import Int32, Int64 from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait -from .....api import ( - ImplDesc, - KernelClass, - OptionalRequirement, - ProblemDesc, - StaticOrRuntimeIntegerType, -) +from .....api import ImplDesc, KernelClass, OptionalRequirement, ProblemDesc, StaticOrRuntimeIntegerType from .....communication.nvlink_domain.symmetric_buffer import SymmetricBufferDevice -from .....helpers.cute_py_helpers import ( - Tcgen05MmaInstruction, - make_tcgen05_tmem_plan, - tcgen05_block_scaled_acc_dtype, -) +from .....helpers.helper_weight_ready_gate import derive_helper_weight_ready_topology, make_helper_weight_ready_generation_gate +from .....helpers.cute_py_helpers import Tcgen05MmaInstruction, make_tcgen05_tmem_plan, tcgen05_block_scaled_acc_dtype from .....helpers.device_workspace import DeviceWorkspace +from .....helpers.dlb_expert_count import resolve_dlb_expert_counts from .....helpers.dsl_helpers import spin_wait from .....helpers.iket_compat import iket +from .....helpers.megamoe_aot import make_megamoe_aot_callable from .....helpers.smem_workspace import SmemWorkspace from .....helpers.software_sync import NvlinkBarrier +from .....helpers.topk_reduce import TopkReduce from .....helpers.utils import ceil_div, round_up from .....quant_def import CombineFormat, QuantKind -from ....blackwell.inference.mega.topk_reduce import TopkReduce from ....schedulers.base import WorkIdAcquisitionMode from ....schedulers.fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler from .block_scaled_swap_ab_fc12_epilogue import GatedActEpilogueArgs, SwapABGatedActEpilogue from .block_scaled_swap_ab_fc12_extension import BlockScaledSwapAbFc12Extension -from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import ( - BlockScaledSwapAbFc12MainloopGenSpecialized, -) +from .block_scaled_swap_ab_fc12_mainloop_gen_specialized import BlockScaledSwapAbFc12MainloopGenSpecialized from .token_comm_gen_specialized import GenphaseTokenComm, GenphaseTokenCommArgs + _aot_symbol_prefix = "rubin_genphase_mega_moe_aot" @@ -140,10 +131,15 @@ def impl_desc_require(cls) -> dict[str, object]: "epi_flag_batches": tuple, "launch_cluster_count": int, "reduce_topk_in_kernel": bool, + "helper_expert_count": OptionalRequirement(int), + "total_helper_slots": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self._validate_desc_inputs(problem_desc, impl_desc) + problem_desc, impl_desc = resolve_dlb_expert_counts( + problem_desc, impl_desc, world_size=problem_desc["world_size"] + ) for field_name, fixed_value, supplied_value in ( ("cluster_shape_mn", self.cluster_shape_mn, impl_desc.get("cluster_shape_mn")), @@ -157,10 +153,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: ) fallback_cluster_shape_mn = impl_desc.get("fallback_cluster_shape_mn") - if ( - fallback_cluster_shape_mn is not None - and fallback_cluster_shape_mn != self.cluster_shape_mn - ): + if fallback_cluster_shape_mn is not None and fallback_cluster_shape_mn != self.cluster_shape_mn: raise ValueError( f"The generation-phase kernel launches a single {self.cluster_shape_mn} cluster shape, so " f"fallback_cluster_shape_mn must be unset or equal to it, got {fallback_cluster_shape_mn}." @@ -199,11 +192,19 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.epi_flag_batches = impl_desc["epi_flag_batches"] self.launch_cluster_count = impl_desc["launch_cluster_count"] self.reduce_topk_in_kernel = impl_desc["reduce_topk_in_kernel"] - + self.helper_expert_count = impl_desc.get("helper_expert_count", 0) + if type(self.helper_expert_count) is not int or self.helper_expert_count < 0: + raise ValueError("helper_expert_count must be a non-negative exact int") self.occupancy = 1 self.architecture = "sm_107" self.threads_per_cta = 12 * 32 self.local_expert_count = self.expert_count // self.world_size + if self.helper_expert_count > 0: + self.home_expert_count = derive_helper_weight_ready_topology( + self.local_expert_count, self.helper_expert_count + ).home_expert_count + else: + self.home_expert_count = self.local_expert_count self.cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] self.promised_launchable_sm_count = self.launch_cluster_count * self.cluster_size @@ -235,9 +236,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.token_comm = GenphaseTokenComm(problem_desc, token_comm_impl_desc) self.max_routed_rows = self.token_comm.worst_case_data_rows self._validate_communication_geometry() - self.tail_barrier = NvlinkBarrier( - world_size=self.world_size, barrier_id=self.tail_barrier_id - ) + self.tail_barrier = NvlinkBarrier(world_size=self.world_size, barrier_id=self.tail_barrier_id) fc12_problem_desc = ProblemDesc({**problem_desc, "expert_count": self.local_expert_count}) resolved_impl_desc = ImplDesc( @@ -263,9 +262,7 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: } ) scheduler_type = ( - BlackwellFusedFc12Scheduler - if self.schedule_mode == "grouped" - else PhaseInterleavedFc12Scheduler + BlackwellFusedFc12Scheduler if self.schedule_mode == "grouped" else PhaseInterleavedFc12Scheduler ) self.scheduler = scheduler_type(fc12_problem_desc, resolved_impl_desc) self.epilogue = SwapABGatedActEpilogue(fc12_problem_desc, resolved_impl_desc) @@ -281,17 +278,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: "selects the peer-facing UBLK store." ) if self.epilogue.token_back_enabled: - raise ValueError( - "The generation-phase communication component provides no token-back path." - ) + raise ValueError("The generation-phase communication component provides no token-back path.") self._device_workspace = self._build_device_workspace() - self._mainloop, self._smem_workspace = self._build_mainloop_and_smem( - fc12_problem_desc, resolved_impl_desc - ) + self._mainloop, self._smem_workspace = self._build_mainloop_and_smem(fc12_problem_desc, resolved_impl_desc) self._topk_reduce = ( - None - if self.reduce_topk_in_kernel - else TopkReduce(self.hidden_size, self.topk, self.combine_format) + None if self.reduce_topk_in_kernel else TopkReduce(self.hidden_size, self.topk, self.combine_format) ) # ------------------------------------------------------------------ @@ -305,13 +296,9 @@ def _validate_geometry(self) -> None: isinstance(self.hidden_size, int), ) if any(static_expert_dimensions) and not all(static_expert_dimensions): - raise ValueError( - "Genphase MegaMoE expert dimensions must be either all static or all runtime." - ) + raise ValueError("Genphase MegaMoE expert dimensions must be either all static or all runtime.") if not all(static_expert_dimensions): - raise NotImplementedError( - "The Genphase MegaMoE kernel currently requires static expert dimensions." - ) + raise NotImplementedError("The Genphase MegaMoE kernel currently requires static expert dimensions.") if self.expert_count <= 0 or self.intermediate_gateup_size <= 0 or self.hidden_size <= 0: raise ValueError("Genphase MegaMoE expert dimensions must be positive.") if self.world_size <= 0 or self.expert_count % self.world_size != 0: @@ -319,9 +306,7 @@ def _validate_geometry(self) -> None: if self.topk <= 0 or self.topk > self.expert_count: raise ValueError("topk must be positive and no greater than expert_count.") if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") # The FC2 epilogue carries no per-row top-k weight, so folding the weight # in at FC1 is the only way an in-kernel reduction can be correct. The # alternative costs a (token, topk, hidden) symmetric plane and a second @@ -352,20 +337,14 @@ def _validate_geometry(self) -> None: raise ValueError( f"schedule_policy mode must be 'grouped' or 'phase_interleave', got {self.schedule_mode!r}." ) - if self.hint is not None and ( - isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0 - ): - raise ValueError( - f"schedule_policy hint must be a positive Python int, got {self.hint!r}." - ) + if self.hint is not None and (isinstance(self.hint, bool) or not isinstance(self.hint, int) or self.hint <= 0): + raise ValueError(f"schedule_policy hint must be a positive Python int, got {self.hint!r}.") if self.schedule_mode == "phase_interleave" and self.work_id_mode != "atomic_counter": raise ValueError("phase_interleave currently requires work_id_mode='atomic_counter'.") if len(self.epi_flag_batches) != 2: raise ValueError("epi_flag_batches must contain FC1 and FC2 batch sizes.") if any(batch < 1 or batch > len(self.epilogue_warp_ids) for batch in self.epi_flag_batches): - raise ValueError( - f"epi_flag_batches values must be in [1, {len(self.epilogue_warp_ids)}]." - ) + raise ValueError(f"epi_flag_batches values must be in [1, {len(self.epilogue_warp_ids)}].") if self.sf_vec_size <= 0: raise ValueError("sf_vec_size must be positive.") @@ -375,10 +354,7 @@ def _validate_geometry(self) -> None: ): if len(dimensions) != expected_rank: raise ValueError(f"{field_name} must contain {expected_rank} dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in dimensions - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in dimensions): raise TypeError(f"{field_name} dimensions must be Python integers.") if any(dimension <= 0 for dimension in dimensions): raise ValueError(f"{field_name} dimensions must be positive.") @@ -398,9 +374,7 @@ def _validate_geometry(self) -> None: f"{self.quant_kind} requires Rubin 2x instruction K={expected_instruction_k}, got {instruction_k}." ) if (instruction_m, instruction_n) != (mma_m, mma_n): - raise NotImplementedError( - "Rubin Genphase MegaMoE does not implement M/N instruction repetition." - ) + raise NotImplementedError("Rubin Genphase MegaMoE does not implement M/N instruction repetition.") # Tile N 64 is excluded because it turns on the halved SFB tile index and # the SFB TMEM shift, and 512 is excluded because the FC1 gather needs one # B row per 128-byte swizzle atom (see the mainloop's gather geometry). @@ -411,13 +385,9 @@ def _validate_geometry(self) -> None: if mma_k % (self.sf_vec_size * 4) != 0: raise ValueError("mma_tiler K must be divisible by four scale-factor vectors.") if self.intermediate_gateup_size % (self.sf_vec_size * 4) != 0: - raise ValueError( - "The intermediate dimension must be divisible by four scale-factor vectors." - ) + raise ValueError("The intermediate dimension must be divisible by four scale-factor vectors.") if mma_m // mma_cta_count != 128: - raise ValueError( - f"The epilogue fixes the per-CTA output tile at 128, got {mma_m // mma_cta_count}." - ) + raise ValueError(f"The epilogue fixes the per-CTA output tile at 128, got {mma_m // mma_cta_count}.") # The epilogue picks the TMA owner warp with `warp_idx == subtile_idx` and # does not check the bound itself, so a wider token tile would drop stores. subtile_count = mma_n // 64 @@ -453,9 +423,7 @@ def _resolve_schedule_hint(self) -> None: blocks_fc1 = ceil_div(self.intermediate_gateup_size, cluster_feature_tile) blocks_fc2 = ceil_div(self.hidden_size, cluster_feature_tile) # Cover the worst token-block alignment of one full persistent-cluster FC2 claim wave. - max_dependent_token_blocks = ceil_div( - self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2 - ) + max_dependent_token_blocks = ceil_div(self.launch_cluster_count + blocks_fc2 - 1, blocks_fc2) required_fc1_work = max_dependent_token_blocks * blocks_fc1 raw_minimum_hint = max(1, ceil_div(required_fc1_work, self.launch_cluster_count)) minimum_safe_hint = round_up(raw_minimum_hint, self.epi_flag_batches[0]) @@ -477,9 +445,7 @@ def _build_device_workspace(self) -> DeviceWorkspace: max_sf_rows = self.token_comm.worst_case_sf_rows # One slot per token block across every local expert, which is the space # `cumulative_token_block_count + tile_n_idx` indexes. - counter_slot_count = ( - self.token_comm.worst_case_padded_rows(self.cta_tile_n) // self.cta_tile_n - ) + counter_slot_count = self.token_comm.worst_case_padded_rows(self.cta_tile_n) // self.cta_tile_n device_workspace = DeviceWorkspace() device_workspace.register( @@ -529,7 +495,7 @@ def _build_mainloop_and_smem( ) -> Tuple[BlockScaledSwapAbFc12MainloopGenSpecialized, SmemWorkspace]: # The communication component builds its own SMEM workspace for its own # launch, so nothing of it is charged against this kernel's budget. - smem_limit = cutlass.memory.get_smem_capacity_in_bytes(self.architecture) // self.occupancy + smem_limit = utils.get_smem_capacity_in_bytes(self.architecture) // self.occupancy smem_workspace = SmemWorkspace() self.scheduler.register_smem_regions(smem_workspace) self.epilogue.register_smem_regions(smem_workspace) @@ -539,9 +505,7 @@ def _build_mainloop_and_smem( if mainloop_smem_budget_bytes <= 0: raise ValueError("Fixed kernel components leave no SMEM budget for the Mainloop.") - mainloop_impl_desc = ImplDesc( - {**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes} - ) + mainloop_impl_desc = ImplDesc({**resolved_impl_desc, "mainloop_smem_budget_bytes": mainloop_smem_budget_bytes}) mainloop = BlockScaledSwapAbFc12MainloopGenSpecialized(problem_desc, mainloop_impl_desc) mainloop.register_smem_regions(smem_workspace) @@ -573,7 +537,7 @@ def dtype_name(dtype: type) -> str: f"sm107_block_scaled_swap_ab_gen_phase_specialized_moe_{self.quant_kind}_" f"{dtype_name(self.a_dtype)}_{dtype_name(self.b_dtype)}_{dtype_name(self.acc_dtype)}_" f"sfvec{self.sf_vec_size}_a{self.a_major_mode.name.lower()}_b{self.b_major_mode.name.lower()}_" - f"w{self.world_size}_e{self.expert_count}_topk{self.topk}_" + f"w{self.world_size}_e{self.expert_count}_topk{self.topk}_helpers{self.helper_expert_count}_" f"topkidx{dtype_name(self.topk_index_dtype)}_" f"h{self.hidden_size}_i{self.intermediate_gateup_size}_maxtoken{self.max_tokens_per_rank}_" f"inst{instruction}_{self.mma_k_mode}_tile{tile}_" @@ -589,6 +553,7 @@ def dtype_name(dtype: type) -> str: f"{'apply_topk_fc1' if self.apply_topk_at_fc1 else 'apply_topk_fc2'}_" f"{'inkernel_reduce' if self.reduce_topk_in_kernel else 'separate_reduce'}_" f"abstages{self._mainloop.num_weight_ab_stages}x{self._mainloop.num_token_ab_stages}" + + ("_dlb_v1" if self.helper_expert_count > 0 else "") ) def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): @@ -606,14 +571,10 @@ def aot_compile(self, out_path: Optional[str] = None, **_compile_kwargs): def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): extents = tuple( - sym_int64(divisibility=math.gcd(int(extent), 128)) - if axis in dynamic_axes - else int(extent) + sym_int64(divisibility=math.gcd(int(extent), 128)) if axis in dynamic_axes else int(extent) for axis, extent in enumerate(shape) ) - return make_fake_compact_tensor( - dtype, extents, stride_order=stride_order, assumed_align=alignment - ) + return make_fake_compact_tensor(dtype, extents, stride_order=stride_order, assumed_align=alignment) tokens = self.max_tokens_per_rank hidden = self.hidden_size @@ -622,12 +583,8 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): experts = self.local_expert_count sf_vec_size = self.sf_vec_size # Weight SF is atom-swizzled and opaque, so its extent is whatever the caller staged. - fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up( - hidden // sf_vec_size, 4 - ) - fc2_weight_sf_columns = round_up(hidden, 128) * round_up( - intermediate_downproj // sf_vec_size, 4 - ) + fc1_weight_sf_columns = round_up(intermediate_gateup, 128) * round_up(hidden // sf_vec_size, 4) + fc2_weight_sf_columns = round_up(hidden, 128) * round_up(intermediate_downproj // sf_vec_size, 4) fake_arguments = dict( activation=fake_tensor( @@ -646,18 +603,10 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), topk_indices=fake_tensor(self.topk_index_dtype, (tokens, self.topk), (1, 0), {}, 16), topk_scores=fake_tensor(cutlass.Float32, (tokens, self.topk), (1, 0), {}, 4), - fc1_weight=fake_tensor( - self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {1, 2}, 16 - ), - fc1_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {1}, 16 - ), - fc2_weight=fake_tensor( - self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {1, 2}, 16 - ), - fc2_weight_sf=fake_tensor( - self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {1}, 16 - ), + fc1_weight=fake_tensor(self.a_dtype, (experts, hidden, intermediate_gateup), (2, 0, 1), {1, 2}, 16), + fc1_weight_sf=fake_tensor(self.sf_dtype, (experts, fc1_weight_sf_columns), (1, 0), {1}, 16), + fc2_weight=fake_tensor(self.a_dtype, (experts, intermediate_downproj, hidden), (2, 0, 1), {1, 2}, 16), + fc2_weight_sf=fake_tensor(self.sf_dtype, (experts, fc2_weight_sf_columns), (1, 0), {1}, 16), output_activation=fake_tensor(cutlass.BFloat16, (tokens, hidden), (1, 0), {}, 16), local_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), shared_workspace=make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), @@ -672,6 +621,11 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): ), stream=make_fake_stream(), ) + if self.helper_expert_count > 0: + fake_arguments.update( + hot_expert_weight_ready_flags=make_ptr(cutlass.Uint64, 0, AddressSpace.gmem, assumed_align=8), + hot_expert_weight_ready_generation=cutlass.Uint64(1), + ) if self.quant_kind.uses_global_scale: fake_arguments.update( fc1_alpha=fake_tensor(cutlass.Float32, (experts,), (0,), set(), 4), @@ -682,16 +636,18 @@ def fake_tensor(dtype, shape, stride_order, dynamic_axes, alignment): compiled = cute.compile[cute.EnableTVMFFI(True)](self, **fake_arguments) if out_path is None: return compiled - compiled.export_to_c( - out_path, function_name=_aot_symbol_prefix, export_only_tvm_ffi_symbols=True - ) + compiled.export_to_c(out_path, function_name=self._aot_symbol_name(), export_only_tvm_ffi_symbols=True) return out_path - @staticmethod - def load_compiled(path: str): + def _aot_symbol_name(self) -> str: + identity = hashlib.sha256(self.name().encode()).hexdigest()[:16] + return f"{_aot_symbol_prefix}_{identity}" + + def load_compiled(self, path: str): from cutlass.cute.runtime import load_module - return load_module(path, enable_tvm_ffi=True)[_aot_symbol_prefix] + function = getattr(load_module(path, enable_tvm_ffi=True), self._aot_symbol_name()) + return make_megamoe_aot_callable(function, self) # ------------------------------------------------------------------ # Host entry @@ -716,16 +672,23 @@ def __call__( fc1_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc2_alpha: Optional[cute.Tensor] = None, # (local_experts,) fc1_norm_const: Optional[cute.Tensor] = None, # (local_experts,) + hot_expert_weight_ready_flags: Optional[cute.Pointer] = None, # uint64[EP] + hot_expert_weight_ready_generation: Optional[cutlass.Uint64] = None, ) -> None: """Launch generation-phase dispatch, fused FC12, and the optional top-k reduce.""" + helper_ready_bundle = (hot_expert_weight_ready_flags, hot_expert_weight_ready_generation) + if cutlass.const_expr( + (self.helper_expert_count > 0 and not all(value is not None for value in helper_ready_bundle)) + or (self.helper_expert_count == 0 and any(value is not None for value in helper_ready_bundle)) + ): + raise ValueError("helper READY flags/generation must both be present iff helper_expert_count > 0") + def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: return cute.make_tensor(tensor.iterator, cute.make_layout(shape, stride=tensor.stride)) if cutlass.const_expr(topk_indices.element_type is not self.topk_index_dtype): - raise TypeError( - f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}." - ) + raise TypeError(f"topk_indices must be the declared {self.topk_index_dtype}, got {topk_indices.dtype}.") for operand_name, operand, expected_dtype in ( ("activation", activation, self.b_dtype), ("activation_sf", activation_sf, self.sf_dtype), @@ -767,12 +730,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: self._device_workspace.assign_device_members(local_workspace, shared_workspace) peer_rank_ptr_mapper = peer_rank_ptr_mapper_host.make_device_object() dense_activation = self.token_comm.dense_activation_tensor(self._device_workspace) - canonical_activation_sf = self.token_comm.canonical_activation_sf_tensor( - self._device_workspace - ) - scheduler_expert_sizes = self.token_comm.scheduler_expert_sizes_tensor( - self._device_workspace - ) + canonical_activation_sf = self.token_comm.canonical_activation_sf_tensor(self._device_workspace) + scheduler_expert_sizes = self.token_comm.scheduler_expert_sizes_tensor(self._device_workspace) fc1_output = self._device_workspace.tensor(self.fc1_output_region) fc1_output_sf = self._device_workspace.tensor(self.fc1_output_sf_region) @@ -821,20 +780,14 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_intermediate_gateup = cute.round_up(intermediate_gateup, sf_vec_size * 4) expected_fc1_weight_sf_columns = padded_intermediate_gateup * hidden // sf_vec_size if cutlass.const_expr( - isinstance(fc1_weight_sf.shape[1], int) - and isinstance(expected_fc1_weight_sf_columns, int) + isinstance(fc1_weight_sf.shape[1], int) and isinstance(expected_fc1_weight_sf_columns, int) ): if cutlass.const_expr(fc1_weight_sf.shape[1] != expected_fc1_weight_sf_columns): raise ValueError("fc1_weight_sf has an incompatible column count.") fc1_sfa = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_intermediate_gateup), - cutlass.Int32(hidden), - cutlass.Int32(experts), - ), - sf_vec_size, + (cutlass.Int32(padded_intermediate_gateup), cutlass.Int32(hidden), cutlass.Int32(experts)), sf_vec_size ), ) @@ -851,31 +804,21 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: padded_fc2_intermediate = cute.round_up(intermediate_fc2, sf_vec_size * 4) expected_fc2_weight_sf_columns = padded_fc2_hidden * padded_fc2_intermediate // sf_vec_size if cutlass.const_expr( - isinstance(fc2_weight_sf.shape[1], int) - and isinstance(expected_fc2_weight_sf_columns, int) + isinstance(fc2_weight_sf.shape[1], int) and isinstance(expected_fc2_weight_sf_columns, int) ): if cutlass.const_expr(fc2_weight_sf.shape[1] != expected_fc2_weight_sf_columns): raise ValueError("fc2_weight_sf has an incompatible column count.") fc2_sfa = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - cutlass.Int32(padded_fc2_hidden), - cutlass.Int32(padded_fc2_intermediate), - cutlass.Int32(experts_fc2), - ), + (cutlass.Int32(padded_fc2_hidden), cutlass.Int32(padded_fc2_intermediate), cutlass.Int32(experts_fc2)), sf_vec_size, ), ) fc2_sfb = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - ( - fc1_output_sf.shape[0], - cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), - singleton, - ), - sf_vec_size, + (fc1_output_sf.shape[0], cutlass.Int32(fc1_output_sf.shape[1] * sf_vec_size), singleton), sf_vec_size ), ) @@ -887,20 +830,13 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: output_activation.iterator, cute.make_layout( (self.max_tokens_per_rank, 1, hidden), - stride=( - output_activation.stride[0], - output_activation.stride[0], - output_activation.stride[1], - ), + stride=(output_activation.stride[0], output_activation.stride[0], output_activation.stride[1]), ), ) else: - pre_reduced_activation = self._device_workspace.tensor( - self.pre_reduced_activation_region - ) + pre_reduced_activation = self._device_workspace.tensor(self.pre_reduced_activation_region) fc2_output = rewrite_tensor_shape( - pre_reduced_activation, - (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden), + pre_reduced_activation, (pre_reduced_activation.shape[0], pre_reduced_activation.shape[1], hidden) ) self._mainloop.materialize_codegen_members() @@ -966,6 +902,8 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: fc1_alpha, fc2_alpha, fc1_norm_const, + hot_expert_weight_ready_flags, + hot_expert_weight_ready_generation, ).launch( grid=grid, block=[self.threads_per_cta, 1, 1], @@ -976,9 +914,7 @@ def rewrite_tensor_shape(tensor: cute.Tensor, shape: Tuple) -> cute.Tensor: ) if cutlass.const_expr(not self.reduce_topk_in_kernel): reduce_scores = None if cutlass.const_expr(self.apply_topk_at_fc1) else topk_scores - self._topk_reduce( - pre_reduced_activation, None, output_activation, reduce_scores, stream - ) + self._topk_reduce(pre_reduced_activation, None, output_activation, reduce_scores, stream) # ------------------------------------------------------------------ # Device entry @@ -1015,6 +951,8 @@ def _kernel( fc1_alpha: Optional[cute.Tensor], fc2_alpha: Optional[cute.Tensor], fc1_norm_const: Optional[cute.Tensor], + hot_expert_weight_ready_flags: Optional[cute.Pointer], + hot_expert_weight_ready_generation: Optional[cutlass.Uint64], ): """Compose scheduler, mainloop, epilogue, and the workspace tail reset.""" self._mainloop.materialize_codegen_members() @@ -1037,8 +975,7 @@ def _kernel( fc1_output_sf = cute.make_tensor( fc1_output_sf_storage.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), - self.epilogue.sf_vec_size, + (fc1_output_sf_storage.shape[0], intermediate_gateup // 2, cutlass.Int32(1)), self.epilogue.sf_vec_size ), ) @@ -1097,9 +1034,7 @@ def _kernel( ) scheduler = self.scheduler - fc2_spin_threshold = ( - intermediate_gateup + self._mainloop.cta_tile_m - 1 - ) // self._mainloop.cta_tile_m + fc2_spin_threshold = (intermediate_gateup + self._mainloop.cta_tile_m - 1) // self._mainloop.cta_tile_m # The scheduler enumerates slots, so the tile it produces has to be # rewritten into an expert index before any consumer sees it: the weights, # their scale factors, the per-expert alphas and the per-expert @@ -1108,6 +1043,15 @@ def _kernel( expert_remap = self.token_comm.slot_to_expert_tensor(self._device_workspace) # No FC1 ready counter: readiness is published per rank, not per tile, and # the waits below cover it once per warp. + helper_weight_ready_gate = None + if cutlass.const_expr(self.helper_expert_count > 0): + helper_weight_ready_gate = make_helper_weight_ready_generation_gate( + memory_slot_count=self.local_expert_count, + helper_count=self.helper_expert_count, + source_count=self.world_size, + terminal_flags=hot_expert_weight_ready_flags, + expected_generation=hot_expert_weight_ready_generation, + ) kernel_extension = BlockScaledSwapAbFc12Extension( sf_vec_size=self.sf_vec_size, fc1_done_counter_pointer=fc1_done_counter.iterator, @@ -1116,10 +1060,7 @@ def _kernel( expert_remap=expert_remap, ) optional_epilogue_args = GatedActEpilogueArgs( - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - fc1_norm_const=fc1_norm_const, - topk_scores=pool_topk_scores, + fc1_alpha=fc1_alpha, fc2_alpha=fc2_alpha, fc1_norm_const=fc1_norm_const, topk_scores=pool_topk_scores ) pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) @@ -1143,9 +1084,15 @@ def _kernel( iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() iket.range_pop() + helper_ready_latched = Int32(0) while work_tile.is_valid_tile: iket.range_push("scheduler.publish_work") - scheduler.publish_work(kernel_extension.prepare_work_tile(work_tile)) + prepared_tile = kernel_extension.prepare_work_tile(work_tile) + if cutlass.const_expr(helper_weight_ready_gate is not None): + helper_ready_latched = helper_weight_ready_gate.wait_before_first_helper_publish( + prepared_tile.expert_idx, helper_ready_latched + ) + scheduler.publish_work(prepared_tile) iket.range_pop() iket.range_push("scheduler.gen_work") work_tile = scheduler.gen_next_work() @@ -1193,9 +1140,7 @@ def _kernel( kernel_extension=kernel_extension, fc1_done_counter_pointer=fc1_done_counter.iterator, fc2_spin_threshold=fc2_spin_threshold, - fc1_sf_ready_pointer=self._device_workspace.ptr( - self.token_comm.fc1_sf_ready_region - ), + fc1_sf_ready_pointer=self._device_workspace.ptr(self.token_comm.fc1_sf_ready_region), expert_sizes=self._device_workspace.tensor(self.token_comm.expert_sizes_region), ) iket.range_pop() @@ -1262,47 +1207,11 @@ def _kernel( optional_epilogue_args, ) tmem_allocator.relinquish_alloc_permit() - tmem_allocator.free( - tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols - ) - - epilogue_thread_count = 32 * len(self.epilogue_warp_ids) - - iket.range_push("tail.grid_sync_before_reset") - self.tail_barrier.sync( - epilogue_thread_count, - Int32(self.promised_launchable_sm_count), - linear_cta_idx, - thread_idx, - ) - iket.range_pop() + tmem_allocator.free(tmem_allocator.retrieve_ptr(self.acc_dtype), self._mainloop.num_tmem_alloc_cols) - iket.range_push("tail.reset_workspace") - total_reset_threads = self.promised_launchable_sm_count * epilogue_thread_count - global_reset_thread = linear_cta_idx * Int32(epilogue_thread_count) + thread_idx - self._device_workspace.reset_tail_space( - "shared", global_reset_thread, total_reset_threads - ) - self._device_workspace.reset_tail_space( - "local", global_reset_thread, total_reset_threads - ) iket.range_pop() - iket.range_push("tail.nvlink_publish") - self.tail_barrier.arrive_and_wait( - epilogue_thread_count, - Int32(self.promised_launchable_sm_count), - linear_cta_idx, - thread_idx, - prologue_grid_sync=True, - epilogue_grid_sync=False, - ) - iket.range_pop() - iket.range_pop() - - if (warp_idx >= Int32(self.gather_b_warp_ids[0])) & ( - warp_idx <= Int32(self.gather_b_warp_ids[-1]) - ): + if (warp_idx >= Int32(self.gather_b_warp_ids[0])) & (warp_idx <= Int32(self.gather_b_warp_ids[-1])): cute.arch.warpgroup_reg_dealloc(self.other_warp_register_count) iket.range_push("genphase.gather_b1") iket.range_push("gather_b1.wait_input_ready") @@ -1330,6 +1239,36 @@ def _kernel( ) iket.range_pop() + # CTAs with no work can finish the epilogue before gather warps pass + # their readiness waits. Join every role before clearing those flags. + cute.arch.sync_threads() + if warp_idx < len(self.epilogue_warp_ids): + epilogue_thread_count = 32 * len(self.epilogue_warp_ids) + + iket.range_push("tail.grid_sync_before_reset") + self.tail_barrier.sync( + epilogue_thread_count, Int32(self.promised_launchable_sm_count), linear_cta_idx, thread_idx + ) + iket.range_pop() + + iket.range_push("tail.reset_workspace") + total_reset_threads = self.promised_launchable_sm_count * epilogue_thread_count + global_reset_thread = linear_cta_idx * Int32(epilogue_thread_count) + thread_idx + self._device_workspace.reset_tail_space("shared", global_reset_thread, total_reset_threads) + self._device_workspace.reset_tail_space("local", global_reset_thread, total_reset_threads) + iket.range_pop() + + iket.range_push("tail.nvlink_publish") + self.tail_barrier.arrive_and_wait( + epilogue_thread_count, + Int32(self.promised_launchable_sm_count), + linear_cta_idx, + thread_idx, + prologue_grid_sync=True, + epilogue_grid_sync=False, + ) + iket.range_pop() + self.tail_barrier.remove_device_members() self._device_workspace.remove_device_members() diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py index 56e119bbab85..4041a1fda3d8 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/dynamic_mainloop.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Raw SM107 dynamic-N block-scaled MMA emission for Rubin MegaMoE.""" from typing import Optional @@ -12,6 +9,7 @@ from .....quant_def import QuantKind + _bit_k_size_upper = 3 _bit_b_sf_id = 4 _bit_a_format = 7 @@ -56,9 +54,7 @@ def build_static_idesc_base( raise ValueError(f"SM107 MegaMoE requires instruction M 128 or 256, got {instruction_m}.") instruction_n_granularity = 16 if instruction_m == 256 else 8 if instruction_n <= 0 or instruction_n > 256 or instruction_n % instruction_n_granularity != 0: - raise ValueError( - f"Invalid SM107 instruction N {instruction_n} for instruction M {instruction_m}." - ) + raise ValueError(f"Invalid SM107 instruction N {instruction_n} for instruction M {instruction_m}.") expected_instruction_k = quant_kind.instruction_k("2x") if instruction_k != expected_instruction_k: @@ -146,15 +142,7 @@ def _tcgen05_mma_block_scaled( llvm.inline_asm( None, - [ - d_tmem_i32, - a_desc_i64, - b_desc_i64, - idesc_i32, - enable_input_d_i32, - sfa_tmem_i32, - sfb_tmem_i32, - ], + [d_tmem_i32, a_desc_i64, b_desc_i64, idesc_i32, enable_input_d_i32, sfa_tmem_i32, sfb_tmem_i32], "{\n\t" ".reg .pred p;\n\t" "setp.ne.b32 p, $4, 0;\n\t" diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py new file mode 100644 index 000000000000..37d164b4a6e7 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/tma_gather.py @@ -0,0 +1,7 @@ +"""Rubin Genphase MegaMoE shim for the shared Gather4 primitive.""" + +# <<>> +from ..local_mega.tma_gather import sm107_tma_gather4_load + + +__all__ = ["sm107_tma_gather4_load"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py index 99b5f26b170e..5adf1c789837 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/token_comm_gen_specialized.py @@ -47,9 +47,10 @@ """ import dataclasses -from typing import ClassVar, Optional +from typing import ClassVar, Optional, Tuple import cuda.bindings.driver as cuda + import cutlass import cutlass.cute as cute import cutlass.utils @@ -310,22 +311,16 @@ def _validate_configuration(self) -> None: if self.expert_count > 16384: raise NotImplementedError("Genphase token comm supports at most 16384 global experts.") if self.topk > self.expert_count: - raise ValueError( - f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}." - ) + raise ValueError(f"topk must not exceed expert_count, got {self.topk} and {self.expert_count}.") if self.max_tokens_per_rank > 1024: raise ValueError( f"Genphase token comm requires max_tokens_per_rank <= 1024, got {self.max_tokens_per_rank}." ) if self.topk_index_dtype not in (cutlass.Int32, cutlass.Int64): - raise ValueError( - f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}." - ) + raise ValueError(f"topk_index_dtype must be Int32 or Int64, got {self.topk_index_dtype}.") if self.sf_padding_block != 128: - raise ValueError( - f"sf_padding_block must be the 128-row SF atom, got {self.sf_padding_block}." - ) + raise ValueError(f"sf_padding_block must be the 128-row SF atom, got {self.sf_padding_block}.") if self.sf_padding_block % self.token_padding_block != 0: raise ValueError("sf_padding_block must be divisible by token_padding_block.") @@ -333,9 +328,7 @@ def _validate_configuration(self) -> None: # four of them along K, so the hidden extent has to fill whole atoms. elements_per_sf_atom = self.quant_kind.sf_vec_size * 4 if self.hidden_size % elements_per_sf_atom != 0: - raise ValueError( - f"{self.quant_kind} requires hidden_size divisible by {elements_per_sf_atom}." - ) + raise ValueError(f"{self.quant_kind} requires hidden_size divisible by {elements_per_sf_atom}.") # The caller's top-k plane is exactly `routes_per_rank` elements with no # slack, so the bulk copy that stages it has to land on that boundary @@ -363,9 +356,7 @@ def _validate_configuration(self) -> None: ) # One thread stages one row of a 128-row block. if self.refine_warps_per_group * 32 != 128: - raise ValueError( - "A refine group must be exactly four warps so one thread owns one row." - ) + raise ValueError("A refine group must be exactly four warps so one thread owns one row.") # The staging load reads a run-aligned row straight out of the linear # plane, so the plane's row has to be at least that wide. Both extents # round the same value up to sixteen, from a payload alignment and from an @@ -540,9 +531,7 @@ def worst_case_padded_rows(self, block: int) -> int: routes_per_source_token = min(self.topk, self.experts_per_rank) route_capacity = source_token_capacity * routes_per_source_token active_expert_capacity = min(self.experts_per_rank, route_capacity) - route_budget_blocks = ( - active_expert_capacity + (route_capacity - active_expert_capacity) // block - ) + route_budget_blocks = active_expert_capacity + (route_capacity - active_expert_capacity) // block expert_bound_blocks = active_expert_capacity * int(ceil_div(source_token_capacity, block)) return min(route_budget_blocks, expert_bound_blocks) * block @@ -634,23 +623,14 @@ def _register_symmetric_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=16, ) workspace.register( - self.topk_ready_region, - cutlass.Int32, - (self.world_size,), - buffer_space="shared", - reset="tail_reset", + self.topk_ready_region, cutlass.Int32, (self.world_size,), buffer_space="shared", reset="tail_reset" ) # The shape is the logical column extent and the row padding is a stride. # Bulk copies move whole padded rows, so the trailing row's padding has to # be inside the region even though no index reaches it; the workspace sizes # a region by what its strides occupy for exactly this reason. for region_name, dtype, columns, row_elements in ( - ( - self.plain_activation_region, - self.activation_dtype, - self.hidden_size, - self.plain_activation_row_elements, - ), + (self.plain_activation_region, self.activation_dtype, self.hidden_size, self.plain_activation_row_elements), ( self.plain_activation_sf_region, self.activation_sf_dtype, @@ -667,9 +647,7 @@ def _register_symmetric_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=128, ) for region_name in (self.token_data_ready_region, self.token_sf_plain_ready_region): - workspace.register( - region_name, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset" - ) + workspace.register(region_name, cutlass.Int32, (1,), buffer_space="shared", reset="tail_reset") def _register_local_regions(self, workspace: DeviceWorkspace) -> None: workspace.register( @@ -681,11 +659,7 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: byte_alignment=16, ) workspace.register( - self.expert_sizes_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.expert_sizes_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" ) workspace.register( self.source_expert_base_region, @@ -712,17 +686,9 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: # Sized to the padded bucket count rather than the live expert count: the # scan that fills them runs over padded buckets, and letting the tail # entries exist removes every boundary special case downstream. - for region_name in ( - self.data_expert_base_region, - self.sf_expert_base_region, - self.refine_block_base_region, - ): + for region_name in (self.data_expert_base_region, self.sf_expert_base_region, self.refine_block_base_region): workspace.register( - region_name, - cutlass.Int32, - (self.local_expert_bucket_count,), - buffer_space="local", - byte_alignment=16, + region_name, cutlass.Int32, (self.local_expert_bucket_count,), buffer_space="local", byte_alignment=16 ) workspace.register( self.token_src_metadata_region, @@ -771,16 +737,10 @@ def _register_local_regions(self, workspace: DeviceWorkspace) -> None: self.refine_block_total_region, self.refine_work_counter_region, ): - workspace.register( - region_name, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset" - ) + workspace.register(region_name, cutlass.Int32, (1,), buffer_space="local", reset="tail_reset") # One slot per local expert workspace.register( - self.fc1_sf_ready_region, - cutlass.Int32, - (self.experts_per_rank,), - buffer_space="local", - reset="tail_reset", + self.fc1_sf_ready_region, cutlass.Int32, (self.experts_per_rank,), buffer_space="local", reset="tail_reset" ) # ------------------------------------------------------------------ @@ -798,9 +758,7 @@ def _build_smem_workspace(self) -> SmemWorkspace: warp_count = self.threads_per_cta // 32 workspace = SmemWorkspace() workspace.register_mbarrier(self.exchange_mbarrier_region, 1) - workspace.register_mbarrier( - self.pusher_mbarrier_region, warp_count * self.pusher_pipeline_depth - ) + workspace.register_mbarrier(self.pusher_mbarrier_region, warp_count * self.pusher_pipeline_depth) workspace.register_mbarrier(self.refine_mbarrier_region, self.refine_participant_groups) overlay = workspace.create_overlay("genphase.token_comm.smem.role") @@ -810,37 +768,22 @@ def _build_smem_workspace(self) -> SmemWorkspace: # what the transfer reads. exchange = overlay.add_lifetime("exchange") exchange.register_tensor( - self.exchange_topk_stage_region, - self.topk_index_dtype, - (self.topk_slab_elements,), - byte_alignment=16, + self.exchange_topk_stage_region, self.topk_index_dtype, (self.topk_slab_elements,), byte_alignment=16 ) if self.apply_topk_at_fc1: exchange.register_tensor( - self.exchange_score_stage_region, - cutlass.Float32, - (self.topk_slab_elements,), - byte_alignment=16, + self.exchange_score_stage_region, cutlass.Float32, (self.topk_slab_elements,), byte_alignment=16 ) sort = overlay.add_lifetime("sort") sort.register_tensor( - self.sort_histogram_region, - cutlass.Int32, - (self.histogram_bucket_count,), - byte_alignment=16, + self.sort_histogram_region, cutlass.Int32, (self.histogram_bucket_count,), byte_alignment=16 ) sort.register_tensor( - self.sort_prefix_region, - cutlass.Int32, - (self.local_expert_bucket_count,), - byte_alignment=16, + self.sort_prefix_region, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 ) sort.register_tensor( - self.sort_warp_totals_region, - cutlass.Int32, - (self.threads_per_cta // 32,), - byte_alignment=16, + self.sort_warp_totals_region, cutlass.Int32, (self.threads_per_cta // 32,), byte_alignment=16 ) helper = overlay.add_lifetime("helper") @@ -851,44 +794,22 @@ def _build_smem_workspace(self) -> SmemWorkspace: stride=(self.local_expert_bucket_count, 1), byte_alignment=16, ) - for region_name in ( - self.helper_totals_region, - self.helper_scan_input_region, - self.helper_prefix_region, - ): - helper.register_tensor( - region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 - ) + for region_name in (self.helper_totals_region, self.helper_scan_input_region, self.helper_prefix_region): + helper.register_tensor(region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16) helper.register_tensor( - self.helper_warp_totals_region, - cutlass.Int32, - (self.threads_per_cta // 32,), - byte_alignment=16, + self.helper_warp_totals_region, cutlass.Int32, (self.threads_per_cta // 32,), byte_alignment=16 ) for region_name in (self.helper_slot_to_expert_region, self.helper_totals_by_slot_region): - helper.register_tensor( - region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16 - ) + helper.register_tensor(region_name, cutlass.Int32, (self.local_expert_bucket_count,), byte_alignment=16) if not self.ranks_experts_in_warp: helper.register_tensor( - self.helper_load_bin_region, - cutlass.Int32, - (self.load_bin_count,), - byte_alignment=16, + self.helper_load_bin_region, cutlass.Int32, (self.load_bin_count,), byte_alignment=16 ) pusher = overlay.add_lifetime("pusher") for region_name, dtype, row_elements in ( - ( - self.pusher_data_stage_region, - self.activation_dtype, - self.plain_activation_row_elements, - ), - ( - self.pusher_sf_stage_region, - self.activation_sf_dtype, - self.plain_activation_sf_row_elements, - ), + (self.pusher_data_stage_region, self.activation_dtype, self.plain_activation_row_elements), + (self.pusher_sf_stage_region, self.activation_sf_dtype, self.plain_activation_sf_row_elements), ): pusher.register_tensor( region_name, @@ -926,11 +847,7 @@ def _build_smem_workspace(self) -> SmemWorkspace: refine.register_tensor( self.refine_output_region, self.activation_sf_dtype, - ( - self.refine_participant_groups, - self.refine_output_stages, - self.refine_output_stage_elements, - ), + (self.refine_participant_groups, self.refine_output_stages, self.refine_output_stage_elements), stride=( self.refine_output_stages * self.refine_output_stage_elements, self.refine_output_stage_elements, @@ -965,9 +882,7 @@ def topk_score_inbox_tensor(self, device_workspace: DeviceWorkspace) -> Optional return device_workspace.tensor(self.topk_score_inbox_region) @cute.jit - def _payload_plane_by_dense_row( - self, device_workspace: DeviceWorkspace, region_name: str - ) -> cute.Tensor: + def _payload_plane_by_dense_row(self, device_workspace: DeviceWorkspace, region_name: str) -> cute.Tensor: """``(world * tokens, columns)`` view of a payload plane. The registered shape keeps the source rank as its own mode, which is how @@ -982,9 +897,7 @@ def _payload_plane_by_dense_row( row_elements = device_workspace.stride(region_name)[1] return cute.make_tensor( device_workspace.ptr(region_name), - cute.make_layout( - (self.world_size * self.max_tokens_per_rank, columns), stride=(row_elements, 1) - ), + cute.make_layout((self.world_size * self.max_tokens_per_rank, columns), stride=(row_elements, 1)), ) @cute.jit @@ -1010,12 +923,9 @@ def canonical_activation_sf_tensor(self, device_workspace: DeviceWorkspace) -> c ``tile_atom_to_shape_SF`` takes the *data* shape, not the SF shape: the broadcast mode it inserts already absorbs ``sf_vec_size``. """ - layout = tile_atom_to_shape_SF( - (self.worst_case_sf_rows, self.hidden_size, 1), self.quant_kind.sf_vec_size - ) + layout = tile_atom_to_shape_SF((self.worst_case_sf_rows, self.hidden_size, 1), self.quant_kind.sf_vec_size) return cute.make_tensor( - device_workspace.ptr(self.canonical_activation_sf_region), - cute.select(layout, mode=[0, 1]), + device_workspace.ptr(self.canonical_activation_sf_region), cute.select(layout, mode=[0, 1]) ) @cute.jit @@ -1077,13 +987,9 @@ def launch( f"topk_indices must be the declared {self.topk_index_dtype}, got {args.topk_indices.element_type}." ) if cutlass.const_expr(args.activation.element_type is not self.activation_dtype): - raise TypeError( - f"activation must be {self.activation_dtype}, got {args.activation.element_type}." - ) + raise TypeError(f"activation must be {self.activation_dtype}, got {args.activation.element_type}.") if cutlass.const_expr(args.activation_sf.element_type is not self.activation_sf_dtype): - raise TypeError( - f"activation_sf must be {self.activation_sf_dtype}, got {args.activation_sf.element_type}." - ) + raise TypeError(f"activation_sf must be {self.activation_sf_dtype}, got {args.activation_sf.element_type}.") # Row addressing uses each plane's own layout, so the padded stride is a # requirement rather than something this component can paper over: a whole # padded row is bulk-copied, and a tighter caller stride would be read past @@ -1095,9 +1001,7 @@ def launch( ): if cutlass.const_expr(isinstance(plane.stride[0], int)): if cutlass.const_expr(plane.stride[0] != row_elements): - raise ValueError( - f"{plane_name} needs the padded row stride {row_elements}, got {plane.stride[0]}." - ) + raise ValueError(f"{plane_name} needs the padded row stride {row_elements}, got {plane.stride[0]}.") self._kernel( args.topk_indices, @@ -1234,19 +1138,10 @@ def _run_exchange( if thread_idx == Int32(0): cute.arch.mbarrier_arrive_and_expect_tx(load_mbarrier, Int32(staged_bytes)) - tma_load_1d( - index_stage.iterator, topk_indices.iterator, load_mbarrier, Int32(index_plane_bytes) - ) + tma_load_1d(index_stage.iterator, topk_indices.iterator, load_mbarrier, Int32(index_plane_bytes)) if cutlass.const_expr(self.apply_topk_at_fc1): - score_stage = self._smem_workspace.tensor( - self.exchange_score_stage_region, smem_base - ) - tma_load_1d( - score_stage.iterator, - topk_scores.iterator, - load_mbarrier, - Int32(score_plane_bytes), - ) + score_stage = self._smem_workspace.tensor(self.exchange_score_stage_region, smem_base) + tma_load_1d(score_stage.iterator, topk_scores.iterator, load_mbarrier, Int32(score_plane_bytes)) cute.arch.mbarrier_wait(load_mbarrier, 0) self._canonicalize_wire_indices(index_stage, thread_idx) @@ -1265,10 +1160,7 @@ def _run_exchange( if thread_idx == Int32(0): cp_async_bulk_s2g( cute.make_ptr( - cutlass.Int32, - own_index_slab.iterator.toint() + peer_offset, - AddressSpace.gmem, - assumed_align=16, + cutlass.Int32, own_index_slab.iterator.toint() + peer_offset, AddressSpace.gmem, assumed_align=16 ), wire_indices.iterator, Int32(self.topk_slab_elements * 4), @@ -1283,9 +1175,7 @@ def _run_exchange( AddressSpace.gmem, assumed_align=16, ), - self._smem_workspace.tensor( - self.exchange_score_stage_region, smem_base - ).iterator, + self._smem_workspace.tensor(self.exchange_score_stage_region, smem_base).iterator, Int32(self.topk_slab_elements * 4), ) cute.arch.cp_async_bulk_commit_group() @@ -1311,8 +1201,7 @@ def _wire_index_view(self, index_stage: cute.Tensor) -> cute.Tensor: if cutlass.const_expr(self.topk_index_dtype is cutlass.Int32): return index_stage return cute.make_tensor( - cute.recast_ptr(index_stage.iterator, dtype=cutlass.Int32), - cute.make_layout((self.topk_slab_elements,)), + cute.recast_ptr(index_stage.iterator, dtype=cutlass.Int32), cute.make_layout((self.topk_slab_elements,)) ) @cute.jit @@ -1345,9 +1234,7 @@ def _is_unowned_expert(self, raw_index) -> cutlass.Boolean: return cutlass.Uint64(raw_index) >= cutlass.Uint64(self.expert_count) @cute.jit - def _classify_tokens( - self, topk_indices: cute.Tensor, global_warp_idx: Int32, lane_idx: Int32 - ) -> cute.Tensor: + def _classify_tokens(self, topk_indices: cute.Tensor, global_warp_idx: Int32, lane_idx: Int32) -> cute.Tensor: """One bit per (token, peer): does this warp owe that peer this token? Both payload passes need the same answer from the same input, so deriving @@ -1388,13 +1275,9 @@ def _classify_tokens( if lane_owns_slot: raw_index = topk_indices[token, slot] if not self._is_unowned_expert(raw_index): - lane_rank = cutlass.Uint32(raw_index) // cutlass.Uint32( - self.experts_per_rank - ) + lane_rank = cutlass.Uint32(raw_index) // cutlass.Uint32(self.experts_per_rank) if cutlass.const_expr(mask_words == 1): - masks[token_round, 0] = masks[token_round, 0] | ( - cutlass.Uint32(1) << lane_rank - ) + masks[token_round, 0] = masks[token_round, 0] | (cutlass.Uint32(1) << lane_rank) else: lane_word = lane_rank // cutlass.Uint32(32) lane_bit = cutlass.Uint32(1) << (lane_rank % cutlass.Uint32(32)) @@ -1404,9 +1287,7 @@ def _classify_tokens( for token_round in cutlass.range_constexpr(tokens_per_warp): for word in cutlass.range_constexpr(mask_words): - masks[token_round, word] = cutlass.Uint32( - cute.arch.warp_redux_sync(masks[token_round, word], "or") - ) + masks[token_round, word] = cutlass.Uint32(cute.arch.warp_redux_sync(masks[token_round, word], "or")) return masks @cute.jit @@ -1451,9 +1332,7 @@ def _run_pusher( iket.range_push("pusher.classify") destination_masks = self._classify_tokens(topk_indices, global_warp_idx, lane_idx) iket.range_pop() - stage_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.pusher_pipeline_depth - ) + stage_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.pusher_pipeline_depth) iket.range_push("pusher.sf") stage_state = self._push_one_plane( @@ -1538,10 +1417,7 @@ def _push_one_plane( with cute.arch.elect_one(): cute.arch.mbarrier_arrive_and_expect_tx(slot_mbarrier, Int32(row_bytes)) tma_load_1d( - slot_stage.iterator, - source_plane[token, None].iterator, - slot_mbarrier, - Int32(row_bytes), + slot_stage.iterator, source_plane[token, None].iterator, slot_mbarrier, Int32(row_bytes) ) cute.arch.mbarrier_wait(slot_mbarrier, stage_state.phase) iket.range_pop() @@ -1550,9 +1426,7 @@ def _push_one_plane( for destination_rank in cutlass.range_constexpr(self.world_size): mask_word = destination_rank // 32 peer_bit = cutlass.Uint32(1 << (destination_rank % 32)) - peer_offset = peer_rank_ptr_mapper.map( - Int64(0), Int32(destination_rank), Int64(0) - ) + peer_offset = peer_rank_ptr_mapper.map(Int64(0), Int32(destination_rank), Int64(0)) destination_row = destination_plane[local_rank, token, None] if (destination_masks[token_round, mask_word] & peer_bit) != cutlass.Uint32(0): with cute.arch.elect_one(): @@ -1575,10 +1449,7 @@ def _push_one_plane( @cute.jit def _publish_payload_flags( - self, - device_workspace: DeviceWorkspace, - peer_rank_ptr_mapper: SymmetricBufferDevice, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, peer_rank_ptr_mapper: SymmetricBufferDevice, thread_idx: Int32 ) -> None: """Drain both payload passes, then tell every peer this pusher is finished. @@ -1599,9 +1470,7 @@ def _publish_payload_flags( if thread_idx < Int32(self.world_size): for region_name in (self.token_sf_plain_ready_region, self.token_data_ready_region): counter_address = device_workspace.ptr(region_name).toint() - red_add_relaxed_sys_s32( - peer_rank_ptr_mapper.map(counter_address, thread_idx, Int64(0)), Int32(1) - ) + red_add_relaxed_sys_s32(peer_rank_ptr_mapper.map(counter_address, thread_idx, Int64(0)), Int32(1)) @cute.jit def _run_sort( @@ -1624,9 +1493,7 @@ def _run_sort( histogram = self._smem_workspace.tensor(self.sort_histogram_region, smem_base) base_row = self._smem_workspace.tensor(self.sort_prefix_region, smem_base) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.histogram_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.histogram_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.histogram_bucket_count): histogram[bucket] = Int32(0) @@ -1661,25 +1528,18 @@ def _run_sort( if route < Int32(self.routes_per_rank): # One unsigned compare covers "not this rank's" and "invalid": # the wire marker and any wild value both land outside the window. - local_expert = ( - index_slab[route // Int32(self.topk), route % Int32(self.topk)] - - local_expert_begin - ) + local_expert = index_slab[route // Int32(self.topk), route % Int32(self.topk)] - local_expert_begin if cutlass.Uint32(local_expert) < cutlass.Uint32(self.experts_per_rank): buckets[route_round] = local_expert for route_round in cutlass.range_constexpr(route_rounds): offsets[route_round] = Int32( - cute.arch.atomic_add( - histogram.iterator + buckets[route_round], Int32(1), sem="relaxed", scope="cta" - ) + cute.arch.atomic_add(histogram.iterator + buckets[route_round], Int32(1), sem="relaxed", scope="cta") ) cute.arch.sync_threads() expert_counts = device_workspace.tensor(self.expert_counts_by_rank_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): expert_counts[source_rank, bucket] = histogram[bucket] @@ -1688,14 +1548,8 @@ def _run_sort( iket.range_push("sort.wait_prefix") if thread_idx == Int32(0): - red_async_add_release_gpu_s32( - device_workspace.ptr(self.rank_count_ready_region), Int32(1) - ) - spin_wait( - device_workspace.ptr(self.prefix_ready_region), - lambda value: value == Int32(1), - scope="gpu", - ) + red_async_add_release_gpu_s32(device_workspace.ptr(self.rank_count_ready_region), Int32(1)) + spin_wait(device_workspace.ptr(self.prefix_ready_region), lambda value: value == Int32(1), scope="gpu") cute.arch.sync_threads() iket.range_pop() @@ -1703,9 +1557,7 @@ def _run_sort( # This source's segment start inside each expert, staged once: every thread # needs it for up to `route_rounds` unrelated buckets. source_expert_base = device_workspace.tensor(self.source_expert_base_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): base_row[bucket] = source_expert_base[source_rank, bucket] @@ -1715,9 +1567,7 @@ def _run_sort( token_metadata = device_workspace.tensor(self.token_src_metadata_region) routed_scores = self.fc1_topk_scores_tensor(device_workspace) if cutlass.const_expr(self.apply_topk_at_fc1): - score_slab = device_workspace.tensor(self.topk_score_inbox_region)[ - source_rank, None, None - ] + score_slab = device_workspace.tensor(self.topk_score_inbox_region)[source_rank, None, None] dense_row_begin = source_rank * Int32(self.max_tokens_per_rank) for route_round in cutlass.range_constexpr(route_rounds): @@ -1735,18 +1585,12 @@ def _run_sort( cute.arch.sync_threads() if thread_idx == Int32(0): - red_async_add_release_gpu_s32( - device_workspace.ptr(self.metadata_ready_region), Int32(1) - ) + red_async_add_release_gpu_s32(device_workspace.ptr(self.metadata_ready_region), Int32(1)) iket.range_pop() @cute.jit def _run_helper( - self, - device_workspace: DeviceWorkspace, - local_rank: Int32, - smem_base: cute.Pointer, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, local_rank: Int32, smem_base: cute.Pointer, thread_idx: Int32 ) -> None: """Fold the W count rows into per-expert sizes and both row-space prefixes. @@ -1814,12 +1658,8 @@ def _run_helper( # themselves stay keyed by expert, because every reader of them -- the # source segment bases, the refine expert table, and the main kernel's # gather -- addresses an expert rather than a slot. - loads_in_scan_order = self._smem_workspace.tensor( - self.helper_totals_by_slot_region, smem_base - ) - scan_order_experts = self._smem_workspace.tensor( - self.helper_slot_to_expert_region, smem_base - ) + loads_in_scan_order = self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base) + scan_order_experts = self._smem_workspace.tensor(self.helper_slot_to_expert_region, smem_base) for padding_block, region_name in ( (self.token_padding_block, self.data_expert_base_region), @@ -1881,9 +1721,7 @@ def _rank_experts_by_load( if cutlass.const_expr(self.ranks_experts_in_warp): self._rank_by_counting_in_warp(device_workspace, smem_base, lane_idx, warp_idx) else: - self._rank_by_load_histogram( - device_workspace, smem_base, thread_idx, lane_idx, warp_idx - ) + self._rank_by_load_histogram(device_workspace, smem_base, thread_idx, lane_idx, warp_idx) @cute.jit def _place_expert_at_slot( @@ -1907,20 +1745,14 @@ def _place_expert_at_slot( kernel never translates it. """ self._smem_workspace.tensor(self.helper_slot_to_expert_region, smem_base)[slot] = expert - self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base)[slot] = ( - routed_count - ) + self._smem_workspace.tensor(self.helper_totals_by_slot_region, smem_base)[slot] = routed_count device_workspace.tensor(self.slot_to_expert_region)[slot] = expert if slot < Int32(self.experts_per_rank): device_workspace.tensor(self.expert_sizes_by_slot_region)[slot] = routed_count @cute.jit def _rank_by_counting_in_warp( - self, - device_workspace: DeviceWorkspace, - smem_base: cute.Pointer, - lane_idx: Int32, - warp_idx: Int32, + self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, lane_idx: Int32, warp_idx: Int32 ) -> None: """Rank the experts inside one warp by counting how many outrank each. @@ -1993,12 +1825,7 @@ def _rank_by_load_histogram( binade = Int32(32) - Int32(cute.arch.clz(totals[bucket])) keys[bucket_round] = Int32(self.load_bin_count - 1) - binade offsets_in_bin[bucket_round] = Int32( - cute.arch.atomic_add( - load_bins.iterator + keys[bucket_round], - Int32(1), - sem="relaxed", - scope="cta", - ) + cute.arch.atomic_add(load_bins.iterator + keys[bucket_round], Int32(1), sem="relaxed", scope="cta") ) cute.arch.sync_threads() @@ -2015,9 +1842,7 @@ def _rank_by_load_histogram( bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): slot = load_bins[keys[bucket_round]] + offsets_in_bin[bucket_round] - self._place_expert_at_slot( - device_workspace, smem_base, slot, bucket, totals[bucket] - ) + self._place_expert_at_slot(device_workspace, smem_base, slot, bucket, totals[bucket]) cute.arch.sync_threads() @cute.jit @@ -2051,14 +1876,10 @@ def _publish_padded_prefix( for bucket_round in cutlass.range_constexpr(bucket_rounds): scan_entry = Int32(bucket_round * self.threads_per_cta) + thread_idx if scan_entry < Int32(self.local_expert_bucket_count): - scan_input[scan_entry] = padded_expert_rows( - loads_in_scan_order[scan_entry], padding_block - ) + scan_input[scan_entry] = padded_expert_rows(loads_in_scan_order[scan_entry], padding_block) cute.arch.sync_threads() - smem_exclusive_prefix( - scan_input, prefix, warp_totals, self.threads_per_cta, thread_idx, lane_idx, warp_idx - ) + smem_exclusive_prefix(scan_input, prefix, warp_totals, self.threads_per_cta, thread_idx, lane_idx, warp_idx) cute.arch.sync_threads() destination = device_workspace.tensor(region_name) @@ -2100,9 +1921,7 @@ def _publish_refine_block_prefix( for bucket_round in cutlass.range_constexpr(bucket_rounds): scan_entry = Int32(bucket_round * self.threads_per_cta) + thread_idx if scan_entry < Int32(self.local_expert_bucket_count): - scan_input[scan_entry] = ceil_div( - loads_in_scan_order[scan_entry], Int32(self.sf_padding_block) - ) + scan_input[scan_entry] = ceil_div(loads_in_scan_order[scan_entry], Int32(self.sf_padding_block)) cute.arch.sync_threads() block_total = smem_exclusive_prefix( @@ -2121,11 +1940,7 @@ def _publish_refine_block_prefix( @cute.jit def _publish_source_segment_bases( - self, - device_workspace: DeviceWorkspace, - count_matrix: cute.Tensor, - local_rank: Int32, - thread_idx: Int32, + self, device_workspace: DeviceWorkspace, count_matrix: cute.Tensor, local_rank: Int32, thread_idx: Int32 ) -> None: """Where each (source rank, expert) segment starts in the data row space. @@ -2136,9 +1951,7 @@ def _publish_source_segment_bases( data_expert_base = device_workspace.tensor(self.data_expert_base_region) source_expert_base = device_workspace.tensor(self.source_expert_base_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): segment_begin = data_expert_base[bucket] @@ -2149,9 +1962,7 @@ def _publish_source_segment_bases( cute.arch.sync_threads() @cute.jit - def _run_refine( - self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, thread_idx: Int32 - ) -> None: + def _run_refine(self, device_workspace: DeviceWorkspace, smem_base: cute.Pointer, thread_idx: Int32) -> None: """Rewrite the received scale factors into the layout the MMA reads.""" group_threads = self.refine_warps_per_group * 32 group_idx = thread_idx // Int32(group_threads) @@ -2189,21 +2000,12 @@ def _run_refine( work_counter = device_workspace.ptr(self.refine_work_counter_region) block_total = device_workspace.tensor(self.refine_block_total_region)[0] - claim_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + claim_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) claimed = cute.make_rmem_tensor((self.refine_claim_fields,), cutlass.Int32) fc1_sf_ready = device_workspace.ptr(self.fc1_sf_ready_region) owning_expert = self._claim_next_block( - claim, - expert_table, - work_counter, - block_total, - group_idx, - thread_in_group, - group_barrier, - group_threads, + claim, expert_table, work_counter, block_total, group_idx, thread_in_group, group_barrier, group_threads ) cute.copy(claim_atom, claim[group_idx, None], claimed) @@ -2227,14 +2029,7 @@ def _run_refine( red_async_add_release_gpu_s32(fc1_sf_ready + owning_expert, Int32(1)) iket.range_pop() owning_expert = self._claim_next_block( - claim, - expert_table, - work_counter, - block_total, - group_idx, - thread_in_group, - group_barrier, - group_threads, + claim, expert_table, work_counter, block_total, group_idx, thread_in_group, group_barrier, group_threads ) cute.copy(claim_atom, claim[group_idx, None], claimed) @@ -2269,9 +2064,7 @@ def _claim_next_block( if thread_in_group < Int32(32): claimed_block = Int32(0) if thread_in_group == Int32(0): - claimed_block = Int32( - cute.arch.atomic_add(work_counter, Int32(1), sem="relaxed", scope="gpu") - ) + claimed_block = Int32(cute.arch.atomic_add(work_counter, Int32(1), sem="relaxed", scope="gpu")) # Broadcast rather than staged through shared memory: the ballot below # needs the whole warp to agree on the block, and a shuffle gets it # there without a second barrier. @@ -2284,25 +2077,17 @@ def _claim_next_block( resolved[self.refine_claim_live_rows] = Int32(0) if claimed_block < block_total: block_rows = Int32(self.sf_padding_block) - owning_expert = self._resolve_block_expert( - expert_table, claimed_block, thread_in_group - ) - rows_before = ( - claimed_block - expert_table[self.refine_table_block_base, owning_expert] - ) * block_rows + owning_expert = self._resolve_block_expert(expert_table, claimed_block, thread_in_group) + rows_before = (claimed_block - expert_table[self.refine_table_block_base, owning_expert]) * block_rows expert_size = expert_table[self.refine_table_expert_size, owning_expert] data_base = expert_table[self.refine_table_data_base, owning_expert] sf_base = expert_table[self.refine_table_sf_base, owning_expert] resolved[self.refine_claim_data_row] = data_base + rows_before resolved[self.refine_claim_sf_row] = sf_base + rows_before - resolved[self.refine_claim_live_rows] = cutlass.min( - expert_size - rows_before, block_rows - ) + resolved[self.refine_claim_live_rows] = cutlass.min(expert_size - rows_before, block_rows) if thread_in_group == Int32(0): - claim_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128 - ) + claim_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) cute.copy(claim_atom, resolved, claim[group_idx, None]) cute.arch.barrier(barrier_id=group_barrier, number_of_threads=group_threads) return owning_expert @@ -2322,9 +2107,7 @@ def _stage_refine_expert_table( sf_base = device_workspace.tensor(self.sf_expert_base_region) expert_sizes = device_workspace.tensor(self.expert_sizes_region) - for bucket_round in cutlass.range_constexpr( - ceil_div(self.local_expert_bucket_count, self.threads_per_cta) - ): + for bucket_round in cutlass.range_constexpr(ceil_div(self.local_expert_bucket_count, self.threads_per_cta)): bucket = Int32(bucket_round * self.threads_per_cta) + thread_idx if bucket < Int32(self.local_expert_bucket_count): expert_table[self.refine_table_block_base, bucket] = block_base[bucket] @@ -2377,21 +2160,13 @@ def _refine_one_block( iket.range_push("refine.out_of_smem") self._scatter_block_atoms( - device_workspace, - smem_base, - sf_row_begin, - group_idx, - thread_in_group, - group_barrier, - group_threads, + device_workspace, smem_base, sf_row_begin, group_idx, thread_in_group, group_barrier, group_threads ) iket.range_pop() return stage_state @cute.jit - def _resolve_block_expert( - self, expert_table: cute.Tensor, block_index: Int32, lane_idx: Int32 - ) -> Int32: + def _resolve_block_expert(self, expert_table: cute.Tensor, block_index: Int32, lane_idx: Int32) -> Int32: """Which local expert owns a linear refine block. One lane per bucket, so the answer comes out of a ballot rather than a scan @@ -2411,12 +2186,9 @@ def _resolve_block_expert( if bucket < Int32(self.local_expert_bucket_count): bucket_base = expert_table[self.refine_table_block_base, bucket] bucket_blocks = ceil_div( - expert_table[self.refine_table_expert_size, bucket], - Int32(self.sf_padding_block), - ) - bucket_owns_block = (bucket_base <= block_index) & ( - block_index < bucket_base + bucket_blocks + expert_table[self.refine_table_expert_size, bucket], Int32(self.sf_padding_block) ) + bucket_owns_block = (bucket_base <= block_index) & (block_index < bucket_base + bucket_blocks) owning_lane = _first_matching_lane(bucket_owns_block) if owning_lane >= Int32(0): resolved = Int32(bucket_round * 32) + owning_lane @@ -2507,51 +2279,32 @@ def _scatter_block_atoms( # The runs tile the row exactly: the staging tile's column extent is the # run-aligned one, so a partial last run reads columns past `hidden_sf` # that belong to this row rather than to the next. - staged_runs = cute.zipped_divide( - staging[group_idx, thread_in_group, None], (columns_per_run,) - ) + staged_runs = cute.zipped_divide(staging[group_idx, thread_in_group, None], (columns_per_run,)) load_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - self.activation_sf_dtype, - num_bits_per_copy=8 * columns_per_run, - ) - store_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=32 - ) - staged_values = cute.make_rmem_tensor( - cute.make_layout((columns_per_run,)), self.activation_sf_dtype + cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=8 * columns_per_run ) + store_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.activation_sf_dtype, num_bits_per_copy=32) + staged_values = cute.make_rmem_tensor(cute.make_layout((columns_per_run,)), self.activation_sf_dtype) value_atoms = cute.zipped_divide(staged_values, (4,)) for run in cutlass.range_constexpr(self.refine_run_count): stage_slot = run % self.refine_output_stages - live_atoms = min( - self.refine_atoms_per_stage, - self.refine_total_atoms - run * self.refine_atoms_per_stage, - ) + live_atoms = min(self.refine_atoms_per_stage, self.refine_total_atoms - run * self.refine_atoms_per_stage) # One run of atoms, addressed the way the MMA reads it: 32 rows by 4 # row-groups against 4 scale factors by 4 atoms. run_output = cute.make_tensor( output[group_idx, stage_slot, None].iterator, - cute.make_layout( - ((32, 4), (4, self.refine_atoms_per_stage)), stride=((16, 4), (1, 512)) - ), + cute.make_layout(((32, 4), (4, self.refine_atoms_per_stage)), stride=((16, 4), (1, 512))), ) cute.copy(load_atom, staged_runs[None, run], staged_values) for atom in cutlass.range_constexpr(self.refine_atoms_per_stage): - cute.copy( - store_atom, - value_atoms[None, atom], - run_output[(atom_row, atom_row_group), (None, atom)], - ) + cute.copy(store_atom, value_atoms[None, atom], run_output[(atom_row, atom_row_group), (None, atom)]) cute.arch.barrier(barrier_id=group_barrier, number_of_threads=group_threads) # `sf_row_begin` is atom-aligned and the run starts on an atom # boundary, so the tensor's own layout hands back the run base. - destination = cute.domain_offset( - (sf_row_begin, Int32(run * columns_per_run * sf_vec_size)), canonical - ) + destination = cute.domain_offset((sf_row_begin, Int32(run * columns_per_run * sf_vec_size)), canonical) # Elected inside the group's first warp, not in every warp: the stage # slot is written cooperatively by all `group_threads` threads, so it # leaves as one store. A bare `elect_one` would issue the same address @@ -2562,9 +2315,7 @@ def _scatter_block_atoms( if thread_in_group // Int32(32) == Int32(0): with cute.arch.elect_one(): cp_async_bulk_s2g( - destination.iterator, - output[group_idx, stage_slot, None].iterator, - Int32(live_atoms * 512), + destination.iterator, output[group_idx, stage_slot, None].iterator, Int32(live_atoms * 512) ) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(self.refine_output_stages - 1, read=True) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py new file mode 100644 index 000000000000..9e43d0dd1dda --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/rubin/inference/mega/topk_reduce_dedup.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # ctm-lint: allow-internal-marker +# SPDX-License-Identifier: BSD-3-Clause +"""Combine reduce for the dedup path: sums the ``rank_slots[token][rank]`` rows named by ``valid_rank_mask`` in +ascending rank order into ``reduced_output[token]``; top-k weights are already folded in.""" + +from typing import ClassVar + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import Float32, Int32 + +from .....helpers.dsl_helpers import mark_alignment +from .....quant_def import CombineFormat + + +class TopkReduceDedup: + _threads: ClassVar[int] = 128 + _hidden_per_thread: ClassVar[int] = 16 # one 32 B (ld/st.256) access of bf16 per slot + + def __init__(self, hidden: int, world_size: int, combine_format: CombineFormat) -> None: + if combine_format.is_quantized or combine_format.act_dtype is not cutlass.BFloat16: + raise ValueError("TopkReduceDedup supports the BF16 combine format only.") + if world_size > 32: + raise ValueError("TopkReduceDedup keeps the rank mask in an Int32; world_size must be <= 32.") + self.hidden = int(hidden) + self.world_size = int(world_size) + self.combine_format = combine_format + if self.hidden % self._hidden_per_thread != 0: + raise ValueError(f"hidden ({self.hidden}) must be divisible by {self._hidden_per_thread}.") + self.hidden_tiles = self.hidden // self._hidden_per_thread + self.require_predicate = self.hidden_tiles % self._threads != 0 + + @cute.jit + def __call__( + self, + rank_slots: cute.Tensor, # (token, world_size, hidden) bf16 + valid_rank_mask: cute.Tensor, # (token,) i32 bitmask over ranks + reduced_output: cute.Tensor, # (token, hidden) + stream: cuda.CUstream, + ): + threads = self._threads + total_workers = reduced_output.shape[0] * self.hidden_tiles + grid = [(total_workers + threads - 1) // threads, 1, 1] + block = [threads, 1, 1] + rank_slots = cute.make_tensor( + rank_slots.iterator, + cute.make_layout((rank_slots.shape[0], self.world_size, self.hidden), stride=rank_slots.stride), + ) + reduced_output = cute.make_tensor( + reduced_output.iterator, + cute.make_layout((reduced_output.shape[0], self.hidden), stride=reduced_output.stride), + ) + self._reduce_bf16(rank_slots, valid_rank_mask, reduced_output).launch(grid=grid, block=block, stream=stream) + + @cute.kernel + def _reduce_bf16(self, rank_slots: cute.Tensor, valid_rank_mask: cute.Tensor, reduced_output: cute.Tensor): + threads = self._threads + hidden_per_thread = self._hidden_per_thread + hidden_tiles = self.hidden_tiles + world_size: cutlass.Constexpr[int] = self.world_size + needs_guard = self.require_predicate + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32(threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, rank, hidden) -> (rank, hidden_per_thread) slice owned by this worker + slots = cute.zipped_divide(rank_slots[token_idx, None, None], (world_size, hidden_per_thread))[ + (None, None), (0, hidden_tile_idx) + ] + dst = cute.zipped_divide(reduced_output[token_idx, None], (hidden_per_thread,))[(None,), (hidden_tile_idx,)] + rank_mask = valid_rank_mask[token_idx] + + access_bits: cutlass.Constexpr[int] = hidden_per_thread * cutlass.BFloat16.width + load_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=access_bits) + acc = cute.make_rmem_tensor((hidden_per_thread,), cutlass.Float32) + for i in cutlass.range_constexpr(hidden_per_thread): + acc[i] = Float32(0.0) + one_pair = (Float32(1.0), Float32(1.0)) + + # Issue every rank slot's load first (predicated, one register buffer per rank) so all of them are in + # flight together, then accumulate ranks ascending; fma by 1.0 is an exact fp32 add. + terms = [cute.make_rmem_tensor((hidden_per_thread,), cutlass.BFloat16) for _ in range(world_size)] + for rank in cutlass.range_constexpr(world_size): + if ((rank_mask >> Int32(rank)) & Int32(1)) != Int32(0): + cute.copy( + load_atom, + mark_alignment(slots[rank, None], hidden_per_thread * cutlass.BFloat16.width // 8), + terms[rank], + ) + for rank in cutlass.range_constexpr(world_size): + if ((rank_mask >> Int32(rank)) & Int32(1)) != Int32(0): + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + value_pair = (Float32(terms[rank][i]), Float32(terms[rank][i + 1])) + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2(value_pair, one_pair, (acc[i], acc[i + 1])) + + out = cute.make_rmem_tensor((hidden_per_thread,), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), out_dtype, num_bits_per_copy=hidden_per_thread * out_dtype.width), + out, + mark_alignment(dst, hidden_per_thread * out_dtype.width // 8), + ) + + +__all__ = ["TopkReduceDedup"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py index 52a7a9daf028..70999789b802 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/__init__.py @@ -1,2 +1,31 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 +"""Scheduler protocols and implementations.""" + +from .base import SchedulerBase, SchedulerConsumer, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .clc_mixed_cga import ClcMixedCgaSchedulerWorker +from .fc12_mapping import ( + BlockPhase, + Fc12WorkTileState, + NonSwapAbFc12WorkTileInfo, + SwapAbFc12WorkTileInfo, + peek_ready_bit, +) +from .fc12_scheduler import BlackwellFusedFc12Scheduler, PhaseInterleavedFc12Scheduler +from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker + + +__all__ = [ + "BlackwellFusedFc12Scheduler", + "BlockPhase", + "ClcMixedCgaSchedulerWorker", + "Fc12WorkTileState", + "NonSwapAbFc12WorkTileInfo", + "NonClcMixedCgaConfig", + "NonClcMixedCgaSchedulerWorker", + "PhaseInterleavedFc12Scheduler", + "SchedulerBase", + "SchedulerConsumer", + "SchedulerWorkTileBase", + "SwapAbFc12WorkTileInfo", + "WorkIdAcquisitionMode", + "peek_ready_bit", +] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py index d6e22636108e..f4a1f40efda1 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/base.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Scheduler façade and architecture-independent work-tile transport.""" from abc import ABC, abstractmethod @@ -15,6 +12,7 @@ from ...helpers.device_workspace import DeviceWorkspace from ...helpers.smem_workspace import SmemWorkspace + WorkIdAcquisitionMode = Literal["grid_stride", "atomic_counter", "cluster_launch_control"] @@ -55,9 +53,7 @@ def __init__( self._pipeline = scheduler_pipeline self._smem_buffer = smem_buffer self._work_tile_type = work_tile_type - self._consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, num_stages - ) + self._consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, num_stages) def __extract_mlir_values__(self) -> list: return extract_mlir_values(self._consumer_state) @@ -122,9 +118,7 @@ def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: ) work_tile_type = self.work_tile_type work_tile_field_count = work_tile_type.storage_field_count - smem_workspace.register_mbarrier( - self.pipeline_mbarriers_region, self.num_scheduler_stages * 2 - ) + smem_workspace.register_mbarrier(self.pipeline_mbarriers_region, self.num_scheduler_stages * 2) smem_workspace.register_tensor( self.work_tiles_region, work_tile_type.storage_dtype, @@ -138,16 +132,12 @@ def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: pass @cute.jit - def create_scheduler_pipelines( - self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer - ) -> None: + def create_scheduler_pipelines(self, smem_workspace: SmemWorkspace, smem_base: cute.Pointer) -> None: """Create CTA-lifetime scheduler pipelines and transport state.""" self._pipeline = pipeline.PipelineAsync.create( num_stages=self.num_scheduler_stages, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, self.num_scheduler_consumer_threads - ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_scheduler_consumer_threads), barrier_storage=smem_workspace.ptr(self.pipeline_mbarriers_region, smem_base), defer_sync=True, ) @@ -168,13 +158,9 @@ def make_consumer(self) -> SchedulerConsumer: @cute.jit def publish_work(self, work_tile: SchedulerWorkTileBase) -> None: """Publish one work tile through the common transport pipeline.""" - copy_atom = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128 - ) + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), work_tile.storage_dtype, num_bits_per_copy=128) self._pipeline.producer_acquire(self._producer_state) - cute.copy( - copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)] - ) + cute.copy(copy_atom, work_tile.to_rmem(), self._smem_buffer[(None, self._producer_state.index)]) cute.arch.fence_proxy("async.shared", space="cta") self._pipeline.producer_commit(self._producer_state) self._producer_state.advance() @@ -202,9 +188,7 @@ def __new_from_mlir_values__(self, values: list) -> "SchedulerBase": return result @abstractmethod - def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None - ): + def get_grid_shape(self, *, max_active_clusters: Optional[int] = None, problem_desc: Any = None): """Return the launch grid from static scheduler policy.""" ... diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py new file mode 100644 index 000000000000..896f1d15989d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/clc_mixed_cga.py @@ -0,0 +1,162 @@ +"""Hardware CLC work-ID acquisition for preferred/fallback (mixed-CGA) cluster launches.""" + +from typing import List, Optional + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass._mlir import ir +from cutlass.cutlass_dsl import Boolean, Int32, extract_mlir_values, new_from_mlir_values + +from ...helpers.device_workspace import DeviceWorkspace +from ...helpers.smem_workspace import SmemWorkspace +from .non_clc_mixed_cga import NonClcMixedCgaConfig +from .work_id_claim import ClcMixedCgaWorkIdState, claim_work_id + + +class ClcMixedCgaSchedulerWorker: + """Claim linear work IDs from hardware cluster launch control on a mixed-CGA launch. + + Same surface as ``NonClcMixedCgaSchedulerWorker`` plus ``claimed_is_valid``: a failed + ``try_cancel`` means the grid is exhausted and the caller must emit the DONE tile. + """ + + cluster_pipeline_mbarriers_region = "clc_mixed_cga.cluster_pipeline_mbarriers" + clc_response_region = "clc_mixed_cga.clc_response" + + def __init__(self, *, config: NonClcMixedCgaConfig, stream_count: int = 1) -> None: + if stream_count != 1: + raise ValueError("Cluster launch control exposes exactly one hardware work-ID stream.") + self.config = config + self.work_id_mode = "cluster_launch_control" + self.stream_count = stream_count + + def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: + """Register the cluster claim barriers and the 16-byte CLC response slot.""" + smem_workspace.register_mbarrier(self.cluster_pipeline_mbarriers_region, 2) + smem_workspace.register_tensor(self.clc_response_region, cutlass.Int32, (4,), byte_alignment=16) + + def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: + """CLC keeps no GMEM scheduler state.""" + pass + + @cute.jit + def assign_device_members( + self, + *, + is_fallback_cluster: Optional[Boolean], + block_idx, + smem_workspace: SmemWorkspace, + smem_base: cute.Pointer, + device_workspace: DeviceWorkspace, + ) -> None: + """Bind this CTA's active cluster to the CLC claim pipeline and seed the bootstrap ticket.""" + if cutlass.const_expr(self.config.is_mixed and is_fallback_cluster is None): + raise ValueError("is_fallback_cluster is required for a true mixed-CGA launch.") + + preferred_m, preferred_n = self.config.preferred_cluster_shape + active_cluster_m = preferred_m + active_cluster_n = preferred_n + if cutlass.const_expr(self.config.is_mixed): + if is_fallback_cluster: + active_cluster_m = Int32(self.config.fallback_cluster_shape[0]) + active_cluster_n = Int32(self.config.fallback_cluster_shape[1]) + + cta_coord_in_active_cluster = ( + Int32(block_idx[0]) % Int32(active_cluster_m), + Int32(block_idx[1]) % Int32(active_cluster_n), + Int32(0), + ) + active_cluster_size = active_cluster_m * active_cluster_n + cluster_pipeline = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), + defer_sync=True, + ) + self._work_id_state = ClcMixedCgaWorkIdState( + response_pending=Boolean(True), + grid_m=Int32(block_idx[0]), + grid_n=Int32(block_idx[1]), + grid_l=Int32(block_idx[2]), + response_is_valid=Boolean(True), + cta_coord_in_cluster=cta_coord_in_active_cluster, + cluster_pipeline=cluster_pipeline, + producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), + consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), + response_pointer=smem_workspace.ptr(self.clc_response_region, smem_base), + cluster_size=active_cluster_size, + ) + self.cta_coord_in_preferred_cluster = ( + Int32(block_idx[0]) % Int32(preferred_m), + Int32(block_idx[1]) % Int32(preferred_n), + Int32(0), + ) + self.claimed_stream_index = Int32(0) + self.claimed_is_valid = Boolean(True) + + @cute.jit + def initialize_fallback_group(self) -> None: + """No fallback registration: hardware assigns fallback clusters to slots.""" + pass + + @cute.jit + def claim_next_work(self, stream_index=0) -> Int32: + """Claim the next ticket; publish its validity and this CTA's preferred-cluster coordinate.""" + work_id, self._work_id_state = claim_work_id(self._work_id_state) + preferred_m, preferred_n = self.config.preferred_cluster_shape + self.cta_coord_in_preferred_cluster = ( + work_id.grid_m % Int32(preferred_m), + work_id.grid_n % Int32(preferred_n), + Int32(0), + ) + self.claimed_is_valid = work_id.is_valid + self.claimed_stream_index = Int32(0) + # Exhausted grid: hand the mapper a harmless in-range ID; the caller substitutes the DONE tile. + linear_work_id = Int32(0) + if work_id.is_valid: + linear_work_id = work_id.grid_l + return linear_work_id + + @cute.jit + def produce_tail(self) -> None: + """Leader waits for the cluster's final release, since peers release on rank 0's empty barrier.""" + work_id_state = self._work_id_state + if self._work_id_state.is_leader_cta: + self._work_id_state.cluster_pipeline.producer_tail(self._work_id_state.producer_state) + else: + self._work_id_state = work_id_state + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + values.extend(extract_mlir_values(self._work_id_state)) + values.extend(extract_mlir_values(self.cta_coord_in_preferred_cluster)) + values.extend(extract_mlir_values(self.claimed_stream_index)) + values.extend(extract_mlir_values(self.claimed_is_valid)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClcMixedCgaSchedulerWorker": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return result + + result = type(self)(config=self.config, stream_count=self.stream_count) + result._work_id_state = rebuild(self._work_id_state) + result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) + result.claimed_stream_index = rebuild(self.claimed_stream_index) + result.claimed_is_valid = rebuild(self.claimed_is_valid) + if value_index != len(values): + raise ValueError( + f"ClcMixedCgaSchedulerWorker MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + +__all__ = ["ClcMixedCgaSchedulerWorker"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py index 14d186d45fe9..0bdd9b9ce720 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_mapping.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """FC12 work-tile ABI and grouped or phase-interleaved task mapping.""" import dataclasses @@ -16,11 +13,16 @@ from ...helpers.utils import padded_expert_rows from .base import SchedulerWorkTileBase + phase_bits = 16 phase_mask = (1 << phase_bits) - 1 peek_ready_bit = 1 << phase_bits + + + + class Fc12WorkTileState(IntEnum): """Sentinel values carried in the expert index field.""" @@ -269,9 +271,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_Fc12TaskCursorSt rebuilt = [] for field in self._fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -341,11 +341,7 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] if isinstance(self.expert_count, Int32): values.extend(extract_mlir_values(self.expert_count)) - token_counts = ( - self.expert_token_sizes - if self.expert_token_sizes is not None - else self.expert_token_prefix_sum - ) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum values.extend(extract_mlir_values(token_counts)) values.extend(extract_mlir_values(self.cursor_state)) if isinstance(self.num_fc1_intermediate_blocks, Int32): @@ -360,17 +356,11 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "Fc12TaskMappingSt def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result - expert_count = ( - rebuild(self.expert_count) - if isinstance(self.expert_count, Int32) - else self.expert_count - ) + expert_count = rebuild(self.expert_count) if isinstance(self.expert_count, Int32) else self.expert_count if self.expert_token_sizes is not None: expert_token_sizes = rebuild(self.expert_token_sizes) expert_token_prefix_sum = None @@ -422,7 +412,29 @@ def create_fc12_task_mapping_state( expert_token_prefix_sum: Optional[cute.Tensor], ) -> Fc12TaskMappingState: """Create the register-resident state for one CTA's FC12 mapper.""" - cursor_state = _Fc12TaskCursorState( + cursor_state = _make_initial_fc12_cursor() + mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + return Fc12TaskMappingState( + expert_count=expert_count, + mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, + mapping_cluster_shape_mn=mapping_cluster_shape_mn, + group_hint=group_hint, + token_padding_block=token_padding_block, + sf_padding_block=sf_padding_block, + is_swap_ab=is_swap_ab, + expert_token_sizes=expert_token_sizes, + expert_token_prefix_sum=expert_token_prefix_sum, + cursor_state=cursor_state, + num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, + num_fc2_hidden_blocks=num_fc2_hidden_blocks, + ) + + +def _make_initial_fc12_cursor() -> _Fc12TaskCursorState: + """Cursor positioned before the first group; the mapper seeks forward from here.""" + return _Fc12TaskCursorState( current_group_first_expert=Int32(0), current_group_last_expert_exclusive=Int32(0), current_phase=Int32(BlockPhase.Linear1), @@ -442,25 +454,13 @@ def create_fc12_task_mapping_state( current_token_block_count=Int32(0), current_expert_token_count=Int32(0), ) - mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] - num_fc1_intermediate_blocks = ( - intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n - num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n - return Fc12TaskMappingState( - expert_count=expert_count, - mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, - mapping_cluster_shape_mn=mapping_cluster_shape_mn, - group_hint=group_hint, - token_padding_block=token_padding_block, - sf_padding_block=sf_padding_block, - is_swap_ab=is_swap_ab, - expert_token_sizes=expert_token_sizes, - expert_token_prefix_sum=expert_token_prefix_sum, - cursor_state=cursor_state, - num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, - num_fc2_hidden_blocks=num_fc2_hidden_blocks, - ) + + +@cute.jit +def rewind_fc12_mapping_cursor(mapping_state: Fc12TaskMappingState) -> Fc12TaskMappingState: + """Restart the forward-only cursor so a lower linear ID (a non-monotonic CLC ticket) can be mapped.""" + mapping_state.cursor_state = _make_initial_fc12_cursor() + return mapping_state @cute.jit @@ -486,11 +486,7 @@ def _first_matching_lane(predicate) -> Int32: @cute.jit def _load_expert_batch_metrics( - mapping_state: Fc12TaskMappingState, - batch_base: Int32, - active_begin: Int32, - active_end: Int32, - lane_idx: Int32, + mapping_state: Fc12TaskMappingState, batch_base: Int32, active_begin: Int32, active_end: Int32, lane_idx: Int32 ) -> Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32]: expert_idx = batch_base + lane_idx token_count = Int32(0) @@ -582,10 +578,7 @@ def _sum_expert_range( @cute.jit def _build_group_range( - mapping_state: Fc12TaskMappingState, - group_first_expert: Int32, - base_fc1_tiles: Int32, - base_fc2_tiles: Int32, + mapping_state: Fc12TaskMappingState, group_first_expert: Int32, base_fc1_tiles: Int32, base_fc2_tiles: Int32 ) -> Tuple[Int32, Int32, Int32]: lane_idx = Int32(cute.arch.lane_idx()) group_threshold = base_fc1_tiles + Int32(mapping_state.group_hint) @@ -596,11 +589,7 @@ def _build_group_range( while batch_base < mapping_state.expert_count and cumulative_fc1_tiles < group_threshold: (lane_expert_idx, _, _, _, _, lane_fc1_tiles, lane_fc2_tiles) = _load_expert_batch_metrics( - mapping_state, - batch_base, - group_first_expert, - Int32(mapping_state.expert_count), - lane_idx, + mapping_state, batch_base, group_first_expert, Int32(mapping_state.expert_count), lane_idx ) fc1_prefix = _warp_inclusive_sum(lane_fc1_tiles, lane_idx) reaches_threshold = ( @@ -613,14 +602,10 @@ def _build_group_range( included_fc2_tiles = Int32(0) fc2_batch_tiles = Int32(cute.arch.warp_redux_sync(included_fc2_tiles, "add")) if selected_lane >= Int32(0): - cumulative_fc1_tiles = cumulative_fc1_tiles + Int32( - cute.arch.shuffle_sync(fc1_prefix, selected_lane) - ) + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, selected_lane)) group_last_expert = batch_base + selected_lane + Int32(1) else: - cumulative_fc1_tiles = cumulative_fc1_tiles + Int32( - cute.arch.shuffle_sync(fc1_prefix, Int32(31)) - ) + cumulative_fc1_tiles = cumulative_fc1_tiles + Int32(cute.arch.shuffle_sync(fc1_prefix, Int32(31))) batch_base = batch_base + Int32(32) group_last_expert = cutlass.min(batch_base, Int32(mapping_state.expert_count)) cumulative_fc2_tiles = cumulative_fc2_tiles + fc2_batch_tiles @@ -644,9 +629,7 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: iket.range_pop() cursor.current_data_cumulative = cursor.current_data_cumulative + residual_data_rows cursor.current_sf_cumulative = cursor.current_sf_cumulative + residual_sf_rows - cursor.current_token_block_cumulative = ( - cursor.current_token_block_cumulative + residual_token_blocks - ) + cursor.current_token_block_cumulative = cursor.current_token_block_cumulative + residual_token_blocks cursor.group_start_data_cumulative = cursor.current_data_cumulative cursor.group_start_sf_cumulative = cursor.current_sf_cumulative @@ -657,19 +640,15 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: cursor.current_group_first_expert = cursor.current_group_last_expert_exclusive iket.range_push("scheduler.group_scan") - (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = ( - _build_group_range( - mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles - ) + (cursor.current_group_last_expert_exclusive, cumulative_fc1_tiles, cumulative_fc2_tiles) = _build_group_range( + mapping_state, cursor.current_group_first_expert, base_fc1_tiles, base_fc2_tiles ) iket.range_pop() cursor.cumulative_fc1_tiles_at_group_end = cumulative_fc1_tiles cursor.cumulative_fc2_tiles_at_group_end = cumulative_fc2_tiles group_start_tile = cursor.current_group_end cursor.current_group_fc1_subphase_end = group_start_tile + cumulative_fc1_tiles - base_fc1_tiles - cursor.current_group_end = ( - cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles - ) + cursor.current_group_end = cursor.current_group_fc1_subphase_end + cumulative_fc2_tiles - base_fc2_tiles cursor.current_phase = Int32(BlockPhase.Linear1) cursor.current_expert_idx = cursor.current_group_first_expert - Int32(1) @@ -682,9 +661,7 @@ def _advance_group(mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: @cute.jit -def _seek_expert_for_work_id( - linear_work_id: Int32, mapping_state: Fc12TaskMappingState -) -> _Fc12TaskCursorState: +def _seek_expert_for_work_id(linear_work_id: Int32, mapping_state: Fc12TaskMappingState) -> _Fc12TaskCursorState: cursor = mapping_state.cursor_state base_tile_end = cursor.current_expert_tile_end @@ -701,9 +678,7 @@ def _seek_expert_for_work_id( ) base_token_block_cumulative = base_token_block_cumulative + cursor.current_token_block_count - search_begin = cutlass.max( - cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert - ) + search_begin = cutlass.max(cursor.current_expert_idx + Int32(1), cursor.current_group_first_expert) batch_base = (search_begin // Int32(32)) * Int32(32) selected_expert = Int32(-1) selected_token_count = Int32(0) @@ -726,11 +701,7 @@ def _seek_expert_for_work_id( lane_fc1_tiles, lane_fc2_tiles, ) = _load_expert_batch_metrics( - mapping_state, - batch_base, - search_begin, - cursor.current_group_last_expert_exclusive, - lane_idx, + mapping_state, batch_base, search_begin, cursor.current_group_last_expert_exclusive, lane_idx ) lane_phase_tiles = lane_fc1_tiles if cursor.current_phase == Int32(BlockPhase.Linear2): @@ -760,17 +731,13 @@ def _seek_expert_for_work_id( selected_token_count = Int32(cute.arch.shuffle_sync(lane_token_count, selected_lane)) selected_token_blocks = Int32(cute.arch.shuffle_sync(lane_token_blocks, selected_lane)) selected_phase_tiles = Int32(cute.arch.shuffle_sync(lane_phase_tiles, selected_lane)) - selected_tile_end = base_tile_end + Int32( - cute.arch.shuffle_sync(tile_prefix, selected_lane) - ) + selected_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, selected_lane)) selected_tile_start = selected_tile_end - selected_phase_tiles selected_data_rows = Int32(cute.arch.shuffle_sync(lane_data_rows, selected_lane)) selected_sf_rows = Int32(cute.arch.shuffle_sync(lane_sf_rows, selected_lane)) selected_data_cumulative = base_data_cumulative + batch_data_rows - selected_data_rows selected_sf_cumulative = base_sf_cumulative + batch_sf_rows - selected_sf_rows - selected_token_block_cumulative = ( - base_token_block_cumulative + batch_token_blocks - selected_token_blocks - ) + selected_token_block_cumulative = base_token_block_cumulative + batch_token_blocks - selected_token_blocks else: base_tile_end = base_tile_end + Int32(cute.arch.shuffle_sync(tile_prefix, Int32(31))) base_data_cumulative = base_data_cumulative + batch_data_rows @@ -793,9 +760,7 @@ def _seek_expert_for_work_id( @cute.jit def _decode_inside_expert( - linear_work_id: Int32, - cta_id_in_mapping_cluster: cute.Coord, - mapping_state: Fc12TaskMappingState, + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState ) -> SchedulerWorkTileBase: cursor = mapping_state.cursor_state cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] @@ -805,22 +770,16 @@ def _decode_inside_expert( cluster_output_block_idx = Int32(0) if cursor.current_phase == Int32(BlockPhase.Linear1): cluster_token_block_idx = local_work_id // mapping_state.num_fc1_intermediate_blocks - cluster_output_block_idx = ( - local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks - ) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc1_intermediate_blocks else: cluster_token_block_idx = local_work_id // mapping_state.num_fc2_hidden_blocks - cluster_output_block_idx = ( - local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks - ) + cluster_output_block_idx = local_work_id - cluster_token_block_idx * mapping_state.num_fc2_hidden_blocks cta_token_block_idx = ( - cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] - + cta_id_in_mapping_cluster[0] + cluster_token_block_idx * mapping_state.mapping_cluster_shape_mn[0] + cta_id_in_mapping_cluster[0] ) cta_output_block_idx = ( - cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] - + cta_id_in_mapping_cluster[1] + cluster_output_block_idx * mapping_state.mapping_cluster_shape_mn[1] + cta_id_in_mapping_cluster[1] ) token_start = cta_token_block_idx * Int32(cta_tile_m) remaining_tokens = cutlass.max(cursor.current_expert_token_count - token_start, Int32(0)) @@ -840,13 +799,9 @@ def _decode_inside_expert( cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) - remaining_cluster_tokens = cutlass.max( - cursor.current_expert_token_count - cluster_token_start, Int32(0) - ) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) - valid_tokens_in_cta_cluster_tile = ( - valid_tokens_in_cta_tile << Int32(16) - ) | valid_tokens_in_cluster_tile + valid_tokens_in_cta_cluster_tile = (valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile return NonSwapAbFc12WorkTileInfo( expert_idx=cursor.current_expert_idx, tile_m_idx=cta_token_block_idx, @@ -861,9 +816,7 @@ def _decode_inside_expert( @cute.jit def map_fc12_linear_work_id( - linear_work_id: Int32, - cta_id_in_mapping_cluster: cute.Coord, - mapping_state: Fc12TaskMappingState, + linear_work_id: Int32, cta_id_in_mapping_cluster: cute.Coord, mapping_state: Fc12TaskMappingState ) -> Tuple[SchedulerWorkTileBase, Fc12TaskMappingState]: """Map one monotonically increasing scalar ID to an FC12 work tile.""" cursor = mapping_state.cursor_state @@ -939,9 +892,7 @@ def _runtime_fields(self) -> Tuple: @cute.jit def clone(self) -> "_PhaseFc12CursorState": """Fork this cursor while sharing its immutable initial SSA values.""" - return type(self)( - *self._runtime_fields(), blocks_per_token_block=self.blocks_per_token_block - ) + return type(self)(*self._runtime_fields(), blocks_per_token_block=self.blocks_per_token_block) def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] @@ -954,9 +905,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "_PhaseFc12CursorS rebuilt_fields = [] for field in self._runtime_fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt_fields.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -982,7 +931,11 @@ def __init__( fc2_cursor: _PhaseFc12CursorState, num_fc1_intermediate_blocks: int, num_fc2_hidden_blocks: int, + dfc1_m_group: int = 1, ) -> None: + if type(dfc1_m_group) is not int or dfc1_m_group not in (1, 16): + raise ValueError("dfc1_m_group must be 1 or 16") + self.dfc1_m_group = dfc1_m_group self.expert_count = expert_count self.mapping_cta_tile_shape_mnk = mapping_cta_tile_shape_mnk self.mapping_cluster_shape_mn = mapping_cluster_shape_mn @@ -995,6 +948,12 @@ def __init__( self.fc2_cursor = fc2_cursor self.num_fc1_intermediate_blocks = num_fc1_intermediate_blocks self.num_fc2_hidden_blocks = num_fc2_hidden_blocks + self.fc2_claims_per_token_block = num_fc2_hidden_blocks + if fc2_cursor.blocks_per_token_block != self.fc2_claims_per_token_block: + raise ValueError( + "FC2 cursor blocks_per_token_block does not match the selected " + "work-ID domain." + ) @property def mapping_cluster_tile_m(self) -> int: @@ -1002,27 +961,19 @@ def mapping_cluster_tile_m(self) -> int: def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] - token_counts = ( - self.expert_token_sizes - if self.expert_token_sizes is not None - else self.expert_token_prefix_sum - ) + token_counts = self.expert_token_sizes if self.expert_token_sizes is not None else self.expert_token_prefix_sum values.extend(extract_mlir_values(token_counts)) for field in (self.fc1_cursor, self.fc2_cursor): values.extend(extract_mlir_values(field)) return values - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "PhaseInterleavedFc12MappingState": + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "PhaseInterleavedFc12MappingState": value_index = 0 def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -1045,6 +996,7 @@ def rebuild(field): fc2_cursor=rebuild(self.fc2_cursor), num_fc1_intermediate_blocks=self.num_fc1_intermediate_blocks, num_fc2_hidden_blocks=self.num_fc2_hidden_blocks, + dfc1_m_group=self.dfc1_m_group, ) if value_index != len(values): raise ValueError( @@ -1081,13 +1033,13 @@ def create_phase_interleaved_fc12_mapping_state( is_swap_ab: bool, expert_token_sizes: Optional[cute.Tensor], expert_token_prefix_sum: Optional[cute.Tensor], + dfc1_m_group: int = 1, ) -> PhaseInterleavedFc12MappingState: """Create independent monotonic mapping cursors for the FC1 and FC2 streams.""" mapping_cluster_tile_n = mapping_cluster_shape_mn[1] * mapping_cta_tile_shape_mnk[1] - num_fc1_intermediate_blocks = ( - intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n + num_fc1_intermediate_blocks = (intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n num_fc2_hidden_blocks = (hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + fc2_claims_per_token_block = num_fc2_hidden_blocks return PhaseInterleavedFc12MappingState( expert_count=expert_count, mapping_cta_tile_shape_mnk=mapping_cta_tile_shape_mnk, @@ -1098,9 +1050,10 @@ def create_phase_interleaved_fc12_mapping_state( expert_token_sizes=expert_token_sizes, expert_token_prefix_sum=expert_token_prefix_sum, fc1_cursor=_make_phase_cursor(num_fc1_intermediate_blocks), - fc2_cursor=_make_phase_cursor(num_fc2_hidden_blocks), + fc2_cursor=_make_phase_cursor(fc2_claims_per_token_block), num_fc1_intermediate_blocks=num_fc1_intermediate_blocks, num_fc2_hidden_blocks=num_fc2_hidden_blocks, + dfc1_m_group=dfc1_m_group, ) @@ -1129,9 +1082,9 @@ def _advance_phase_cursor( token_count = prefix_end - prefix_begin cursor.current_expert_token_count = token_count - cursor.current_token_block_count = ( - token_count + Int32(mapping_state.mapping_cluster_tile_m - 1) - ) // Int32(mapping_state.mapping_cluster_tile_m) + cursor.current_token_block_count = (token_count + Int32(mapping_state.mapping_cluster_tile_m - 1)) // Int32( + mapping_state.mapping_cluster_tile_m + ) cursor.expert_tile_start = cursor.expert_tile_end cursor.expert_tile_end = cursor.expert_tile_start + cursor.current_token_block_count * Int32( cursor.blocks_per_token_block @@ -1141,9 +1094,7 @@ def _advance_phase_cursor( @cute.jit def _seek_phase_cursor( - linear_work_id: Int32, - cursor: _PhaseFc12CursorState, - mapping_state: PhaseInterleavedFc12MappingState, + linear_work_id: Int32, cursor: _PhaseFc12CursorState, mapping_state: PhaseInterleavedFc12MappingState ) -> _PhaseFc12CursorState: expert_tile_end = cursor.expert_tile_end next_expert_idx = cursor.expert_idx + Int32(1) @@ -1160,9 +1111,7 @@ def resolve_phase_interleaved_fc1_claim_target( ) -> Tuple[Int32, Boolean]: """Resolve the static FC1 watermark or the smaller runtime stream extent.""" target_work_id = minimum_claim_count - Int32(1) - probe_cursor = _seek_phase_cursor( - target_work_id, mapping_state.fc1_cursor.clone(), mapping_state - ) + probe_cursor = _seek_phase_cursor(target_work_id, mapping_state.fc1_cursor.clone(), mapping_state) stream_ends_before_target = target_work_id >= probe_cursor.expert_tile_end claim_target = minimum_claim_count if stream_ends_before_target: @@ -1179,17 +1128,37 @@ def _decode_phase_work_id( mapping_state: PhaseInterleavedFc12MappingState, ) -> SchedulerWorkTileBase: local_work_id = linear_work_id - cursor.expert_tile_start - cluster_token_block_idx = local_work_id // Int32(cursor.blocks_per_token_block) + cluster_token_block_idx = local_work_id // Int32( + cursor.blocks_per_token_block + ) cluster_output_block_idx = local_work_id - cluster_token_block_idx * Int32( cursor.blocks_per_token_block ) + if cutlass.const_expr(mapping_state.dfc1_m_group > 1): + if phase == Int32(BlockPhase.Linear2): + # Within each small M strip, visit its rows for a fixed N weight + # tile before advancing N. Only the per-expert rectangle is + # permuted: expert ranges, real row/ready indices, and task counts + # remain unchanged. A short final strip uses its actual row count + # and therefore adds neither padding work nor duplicate tiles. + strip_span = Int32(mapping_state.dfc1_m_group * cursor.blocks_per_token_block) + strip = local_work_id // strip_span + strip_m_begin = strip * Int32(mapping_state.dfc1_m_group) + strip_rows = cutlass.max(Int32(1), cutlass.min( + Int32(mapping_state.dfc1_m_group), + cursor.current_token_block_count - strip_m_begin, + )) + in_strip = local_work_id - strip * strip_span + cluster_output_block_idx = in_strip // strip_rows + cluster_token_block_idx = ( + strip_m_begin + in_strip - cluster_output_block_idx * strip_rows + ) + cta_token_block_idx = ( - cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) - + cta_id_in_mapping_cluster[0] + cluster_token_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[0]) + cta_id_in_mapping_cluster[0] ) cta_output_block_idx = ( - cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) - + cta_id_in_mapping_cluster[1] + cluster_output_block_idx * Int32(mapping_state.mapping_cluster_shape_mn[1]) + cta_id_in_mapping_cluster[1] ) cta_tile_m = mapping_state.mapping_cta_tile_shape_mnk[0] @@ -1211,9 +1180,7 @@ def _decode_phase_work_id( cluster_tile_m = mapping_state.mapping_cluster_shape_mn[0] * cta_tile_m cluster_token_start = cluster_token_block_idx * Int32(cluster_tile_m) - remaining_cluster_tokens = cutlass.max( - cursor.current_expert_token_count - cluster_token_start, Int32(0) - ) + remaining_cluster_tokens = cutlass.max(cursor.current_expert_token_count - cluster_token_start, Int32(0)) valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster_tokens, Int32(cluster_tile_m)) return NonSwapAbFc12WorkTileInfo( expert_idx=cursor.expert_idx, @@ -1222,8 +1189,7 @@ def _decode_phase_work_id( cumulative_data_physical_row=cursor.data_cumulative, cumulative_sf_physical_row=cursor.sf_cumulative, cumulative_token_block_count=cursor.token_block_cumulative, - valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) - | valid_tokens_in_cluster_tile, + valid_tokens_in_cta_cluster_tile=(valid_tokens_in_cta_tile << Int32(16)) | valid_tokens_in_cluster_tile, phase_and_flags=phase, ) @@ -1261,6 +1227,8 @@ def map_phase_interleaved_fc12_work_id( return work_tile, stream_has_work, mapping_state + + __all__ = [ "BlockPhase", "Fc12TaskMappingState", @@ -1275,4 +1243,5 @@ def map_phase_interleaved_fc12_work_id( "map_phase_interleaved_fc12_work_id", "peek_ready_bit", "resolve_phase_interleaved_fc1_claim_target", + "rewind_fc12_mapping_cursor", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py index 75ef3eccaabd..9fd349cd8048 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/fc12_scheduler.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Composable grouped and phase-interleaved FC12 schedulers.""" import math @@ -16,6 +13,7 @@ from ...helpers.smem_workspace import SmemWorkspace from ...helpers.utils import ceil_div from .base import SchedulerBase, SchedulerWorkTileBase, WorkIdAcquisitionMode +from .clc_mixed_cga import ClcMixedCgaSchedulerWorker from .fc12_mapping import ( BlockPhase, NonSwapAbFc12WorkTileInfo, @@ -26,6 +24,7 @@ map_fc12_linear_work_id, map_phase_interleaved_fc12_work_id, resolve_phase_interleaved_fc1_claim_target, + rewind_fc12_mapping_cursor, ) from .non_clc_mixed_cga import NonClcMixedCgaConfig, NonClcMixedCgaSchedulerWorker @@ -70,16 +69,10 @@ def minimum_phase_interleave_hint( return max(1, ceil_div(required_fc1_claims, launch_cluster_cnt_merge_as_preferred)) -def _to_fc12_mapping_cta_coord( - cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool -) -> cute.Coord: +def _to_fc12_mapping_cta_coord(cta_coord_in_preferred_cluster: cute.Coord, is_swap_ab: bool) -> cute.Coord: if cutlass.const_expr(not is_swap_ab): return cta_coord_in_preferred_cluster - return ( - cta_coord_in_preferred_cluster[1], - cta_coord_in_preferred_cluster[0], - cta_coord_in_preferred_cluster[2], - ) + return (cta_coord_in_preferred_cluster[1], cta_coord_in_preferred_cluster[0], cta_coord_in_preferred_cluster[2]) class BlackwellFusedFc12Scheduler(SchedulerBase): @@ -109,6 +102,7 @@ def impl_desc_require(cls) -> dict[str, type]: "work_id_mode": str, "is_swap_ab": bool, **_mixed_cga_impl_requirements(), + "compute_cluster_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -127,14 +121,11 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] self.is_swap_ab = impl_desc["is_swap_ab"] self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) - self.launch_cluster_cnt_merge_as_preferred = ( - self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred - ) - self.work_tile_type = ( - SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo - ) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.compute_cluster_count = impl_desc.get("compute_cluster_count", self.launch_cluster_cnt_merge_as_preferred) + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo if self.hint is None: - self.hint = self.launch_cluster_cnt_merge_as_preferred + self.hint = self.compute_cluster_count self.group_hint = self.hint self._validate_configuration() @@ -156,10 +147,14 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.mapping_cluster_shape_mn = self.cluster_shape_mn if self.work_id_mode == "cluster_launch_control": - raise NotImplementedError("cluster_launch_control is not implemented for FC12.") - self._work_id_worker = NonClcMixedCgaSchedulerWorker( - config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=1 - ) + self._work_id_worker = ClcMixedCgaSchedulerWorker(config=self.non_clc_mixed_cga_config) + else: + self._work_id_worker = NonClcMixedCgaSchedulerWorker( + config=self.non_clc_mixed_cga_config, + work_id_mode=self.work_id_mode, + stream_count=1, + work_id_stride=self.compute_cluster_count, + ) def _validate_configuration(self) -> None: if len(self.mma_tiler_mnk) != 3: @@ -181,21 +176,29 @@ def _validate_configuration(self) -> None: raise ValueError("sf_padding_block must be positive.") if self.launch_cluster_cnt_merge_as_preferred <= 0: raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") + if ( + isinstance(self.compute_cluster_count, bool) + or not isinstance(self.compute_cluster_count, int) + or not 0 < self.compute_cluster_count <= self.launch_cluster_cnt_merge_as_preferred + ): + raise ValueError( + f"compute_cluster_count must be a Python int in (0, {self.launch_cluster_cnt_merge_as_preferred}], got " + f"{self.compute_cluster_count}." + ) if self.work_id_mode not in ("grid_stride", "atomic_counter", "cluster_launch_control"): raise ValueError( "work_id_mode must be 'grid_stride', 'atomic_counter', or " f"'cluster_launch_control', got {self.work_id_mode!r}." ) cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] - if self.work_id_mode == "atomic_counter" and cluster_size > 32: - raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") + uses_peer_lane_protocol = self.work_id_mode in ("atomic_counter", "cluster_launch_control") + if uses_peer_lane_protocol and cluster_size > 32: + raise ValueError("The per-lane peer barrier protocol supports at most 32 CTAs per cluster.") fallback_cluster_shape = self.non_clc_mixed_cga_config.fallback_cluster_shape if fallback_cluster_shape is not None: fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] - if self.work_id_mode == "atomic_counter" and fallback_cluster_size > 32: - raise ValueError( - "The atomic broadcast protocol supports at most 32 CTAs per fallback cluster." - ) + if uses_peer_lane_protocol and fallback_cluster_size > 32: + raise ValueError("The per-lane peer barrier protocol supports at most 32 CTAs per fallback cluster.") for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): value = getattr(self, field_name) if isinstance(value, int) and value <= 0: @@ -222,10 +225,48 @@ def initialize_fallback_group(self) -> None: """Register this physical fallback cluster with its fixed logical group.""" self._work_id_worker.initialize_fallback_group() + # CUDA caps gridDim.z at 65535; CLC tickets ride on blockIdx.z. + _clc_grid_z_limit = 65535 + + def clc_ticket_count(self, total_token_rows): + """Return the CLC ticket count, an upper bound on FC12 tasks (surplus tickets decode to DONE).""" + if not isinstance(self.expert_count, int): + raise ValueError("cluster_launch_control needs a static expert_count to size the ticket grid.") + if not isinstance(self.intermediate_gateup_size, int) or not isinstance(self.hidden_size, int): + raise ValueError("cluster_launch_control needs static intermediate/hidden sizes to size the ticket grid.") + if isinstance(total_token_rows, bool) or not isinstance(total_token_rows, int) or total_token_rows < 0: + raise TypeError( + "cluster_launch_control needs a host-known non-negative int total_token_rows; the grid " + "must be validated against the gridDim.z limit before launch." + ) + mapping_cluster_tile_m = self.mapping_cta_tile_shape_mnk[0] * self.mapping_cluster_shape_mn[0] + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + num_fc1_blocks = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + num_fc2_blocks = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n + token_blocks = (total_token_rows + mapping_cluster_tile_m - 1) // mapping_cluster_tile_m + ( + self.expert_count - 1 + ) + # Zero rows still needs one cluster to bootstrap, observe DONE and exit (gridDim.z >= 1). + return max(1, token_blocks * (num_fc1_blocks + num_fc2_blocks)) + def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc=None + self, + *, + max_active_clusters: Optional[int] = None, + problem_desc=None, + total_token_rows=None, ) -> Tuple[int, int, int]: - """Return the persistent launch grid in GEMM-domain orientation.""" + """Return the launch grid in GEMM-domain orientation (full ticket grid under cluster launch control).""" + if self.work_id_mode == "cluster_launch_control": + if total_token_rows is None: + raise ValueError("cluster_launch_control requires total_token_rows to size the ticket grid.") + ticket_count = self.clc_ticket_count(total_token_rows) + if ticket_count > self._clc_grid_z_limit: + raise ValueError( + f"cluster_launch_control ticket count {ticket_count} exceeds gridDim.z limit " + f"{self._clc_grid_z_limit}." + ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], ticket_count) if ( not self.non_clc_mixed_cga_config.is_mixed and max_active_clusters is not None @@ -236,11 +277,7 @@ def get_grid_shape( "launch_cluster_cnt_merge_as_preferred " f"({self.launch_cluster_cnt_merge_as_preferred})." ) - return ( - self.cluster_shape_mn[0], - self.cluster_shape_mn[1], - self.launch_cluster_cnt_merge_as_preferred, - ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) @cute.jit def assign_device_members( @@ -257,9 +294,7 @@ def assign_device_members( ) -> None: """Initialize all FC12 scheduler state rooted for one CTA lifetime.""" if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): - raise ValueError( - "Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided." - ) + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") needs_actual_shape = not all( isinstance(dimension, int) for dimension in (self.expert_count, self.intermediate_gateup_size, self.hidden_size) @@ -311,11 +346,30 @@ def gen_next_work(self) -> SchedulerWorkTileBase: cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( self._work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab ) + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + # CLC tickets on a mixed-cluster launch are not monotonic; the mapper is forward-only, so + # rewind when the ticket falls before the cursor's current segment. + if self._work_id_worker.claimed_is_valid: + if work_id < self._task_mapping_state.cursor_state.current_expert_tile_start: + self._task_mapping_state = rewind_fc12_mapping_cursor(self._task_mapping_state) work_tile, self._task_mapping_state = map_fc12_linear_work_id( work_id, cta_id_in_mapping_cluster, self._task_mapping_state ) + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + # A failed try_cancel means no cluster is left to steal: terminate this CTA's stream. + mapped_tile = work_tile + work_tile = make_fc12_done_tile(self.is_swap_ab) + if self._work_id_worker.claimed_is_valid: + work_tile = mapped_tile return work_tile + @cute.jit + def produce_tail(self) -> None: + """Drain the work transport, then the CLC leader waits for the cluster's final release.""" + super().produce_tail() + if cutlass.const_expr(self.work_id_mode == "cluster_launch_control"): + self._work_id_worker.produce_tail() + def __extract_mlir_values__(self) -> list: values = super().__extract_mlir_values__() values.extend(extract_mlir_values(self._work_id_worker)) @@ -335,9 +389,7 @@ def __new_from_mlir_values__(self, values: list) -> "BlackwellFusedFc12Scheduler def rebuild(state): nonlocal value_index state_value_count = len(extract_mlir_values(state)) - rebuilt_state = new_from_mlir_values( - state, values[value_index : value_index + state_value_count] - ) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) value_index += state_value_count return rebuilt_state @@ -362,6 +414,7 @@ def rebuild(state): "work_id_mode", "is_swap_ab", "launch_cluster_cnt_merge_as_preferred", + "compute_cluster_count", "non_clc_mixed_cga_config", "mapping_cta_tile_shape_mnk", "mapping_cluster_shape_mn", @@ -407,9 +460,7 @@ def __new_from_mlir_values__(self, values: list) -> "_PhaseInterleaveControlStat rebuilt_fields = [] for field in self._fields(): field_value_count = len(extract_mlir_values(field)) - rebuilt_fields.append( - new_from_mlir_values(field, values[value_index : value_index + field_value_count]) - ) + rebuilt_fields.append(new_from_mlir_values(field, values[value_index : value_index + field_value_count])) value_index += field_value_count if value_index != len(values): raise ValueError( @@ -440,7 +491,11 @@ def impl_desc_require(cls) -> dict[str, type]: "sf_padding_block": int, "work_id_mode": str, "is_swap_ab": bool, + "dfc1_m_group": OptionalRequirement(int), + "defer_consumer_until_full_ready": OptionalRequirement(bool), + "fuse_ready_probe_and_linear1_claim": OptionalRequirement(bool), **_mixed_cga_impl_requirements(), + "compute_cluster_count": OptionalRequirement(int), } def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: @@ -458,13 +513,17 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: self.sf_padding_block = impl_desc["sf_padding_block"] self.work_id_mode: WorkIdAcquisitionMode = impl_desc["work_id_mode"] self.is_swap_ab = impl_desc["is_swap_ab"] - self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) - self.launch_cluster_cnt_merge_as_preferred = ( - self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.dfc1_m_group = impl_desc.get("dfc1_m_group", 1) + self.defer_consumer_until_full_ready = impl_desc.get( + "defer_consumer_until_full_ready", False ) - self.work_tile_type = ( - SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo + self.fuse_ready_probe_and_linear1_claim = impl_desc.get( + "fuse_ready_probe_and_linear1_claim", False ) + self.non_clc_mixed_cga_config = _make_non_clc_mixed_cga_config(impl_desc) + self.launch_cluster_cnt_merge_as_preferred = self.non_clc_mixed_cga_config.launch_cluster_cnt_merge_as_preferred + self.compute_cluster_count = impl_desc.get("compute_cluster_count", self.launch_cluster_cnt_merge_as_preferred) + self.work_tile_type = SwapAbFc12WorkTileInfo if self.is_swap_ab else NonSwapAbFc12WorkTileInfo self._validate_configuration() mma_cta_count = 2 if self.use_2cta_instrs else 1 @@ -487,30 +546,39 @@ def __init__(self, problem_desc: ProblemDesc, impl_desc: ImplDesc) -> None: config=self.non_clc_mixed_cga_config, work_id_mode=self.work_id_mode, stream_count=2 ) - mapping_cluster_tile_n = ( - self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] - ) - self.blocks_fc1 = ( - self.intermediate_gateup_size + mapping_cluster_tile_n - 1 - ) // mapping_cluster_tile_n + mapping_cluster_tile_n = self.mapping_cta_tile_shape_mnk[1] * self.mapping_cluster_shape_mn[1] + self.blocks_fc1 = (self.intermediate_gateup_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n self.blocks_fc2 = (self.hidden_size + mapping_cluster_tile_n - 1) // mapping_cluster_tile_n - interleave_gcd = math.gcd(self.blocks_fc1, self.blocks_fc2) - self.interleave_fc2_slots = self.blocks_fc2 // interleave_gcd - self.interleave_cycle_length = (self.blocks_fc1 + self.blocks_fc2) // interleave_gcd + self.fc2_claims_per_token_block = self.blocks_fc2 + interleave_gcd = math.gcd( + self.blocks_fc1, self.fc2_claims_per_token_block + ) + self.interleave_fc2_slots = ( + self.fc2_claims_per_token_block // interleave_gcd + ) + self.interleave_cycle_length = ( + self.blocks_fc1 + self.fc2_claims_per_token_block + ) // interleave_gcd self.minimum_global_fc1_claims = minimum_phase_interleave_fc1_claims( blocks_fc1=self.blocks_fc1, blocks_fc2=self.blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) minimum_hint = minimum_phase_interleave_hint( blocks_fc1=self.blocks_fc1, blocks_fc2=self.blocks_fc2, - launch_cluster_cnt_merge_as_preferred=self.launch_cluster_cnt_merge_as_preferred, + launch_cluster_cnt_merge_as_preferred=self.compute_cluster_count, ) - if self.fc1_prologue_tiles < minimum_hint: + if ( + self.fc1_prologue_tiles < minimum_hint + and not ( + self.fc1_prologue_tiles == 0 + and self.defer_consumer_until_full_ready + ) + ): raise ValueError( f"phase_interleave hint {self.fc1_prologue_tiles} cannot cover a " - f"{self.launch_cluster_cnt_merge_as_preferred}-cluster FC2 claim wave; " + f"{self.compute_cluster_count}-cluster FC2 claim wave; " f"raise hint to at least {minimum_hint}." ) @@ -519,14 +587,10 @@ def _validate_configuration(self) -> None: raise ValueError("mma_tiler_mnk must contain three dimensions.") if len(self.cluster_shape_mn) != 2: raise ValueError("cluster_shape_mn must contain two dimensions.") - if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in self.mma_tiler_mnk - ): + if not all(isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.mma_tiler_mnk): raise TypeError("mma_tiler_mnk dimensions must be Python ints.") if not all( - isinstance(dimension, int) and not isinstance(dimension, bool) - for dimension in self.cluster_shape_mn + isinstance(dimension, int) and not isinstance(dimension, bool) for dimension in self.cluster_shape_mn ): raise TypeError("cluster_shape_mn dimensions must be Python ints.") if any(dimension <= 0 for dimension in self.mma_tiler_mnk): @@ -539,21 +603,61 @@ def _validate_configuration(self) -> None: if ( isinstance(self.fc1_prologue_tiles, bool) or not isinstance(self.fc1_prologue_tiles, int) - or self.fc1_prologue_tiles <= 0 + or self.fc1_prologue_tiles < 0 ): raise ValueError( - "fc1_prologue_tiles must be a positive Python int resolved by the kernel frontend." + "fc1_prologue_tiles must be a non-negative Python int " + "resolved by the kernel frontend." ) + if type(self.dfc1_m_group) is not int or self.dfc1_m_group not in (1, 16): + raise ValueError("dfc1_m_group must be 1 or 16") + if self.dfc1_m_group > 1 and ( + self.is_swap_ab + or not self.defer_consumer_until_full_ready + ): + raise ValueError( + "Grouped dFC1 M traversal requires non-swap singleton consumers, " + "ordinary producer order and full-ready admission" + ) + if not isinstance(self.defer_consumer_until_full_ready, bool): + raise TypeError("defer_consumer_until_full_ready must be a bool.") + if not isinstance(self.fuse_ready_probe_and_linear1_claim, bool): + raise TypeError( + "fuse_ready_probe_and_linear1_claim must be a bool." + ) + if ( + self.defer_consumer_until_full_ready + and self.non_clc_mixed_cga_config.is_mixed + ): + raise NotImplementedError( + "defer_consumer_until_full_ready does not yet support fixed " + "preferred/fallback cluster groups." + ) + if self.fuse_ready_probe_and_linear1_claim: + if not self.defer_consumer_until_full_ready: + raise ValueError( + "fuse_ready_probe_and_linear1_claim requires " + "defer_consumer_until_full_ready=True." + ) + if self.fc1_prologue_tiles == 0 and not self.defer_consumer_until_full_ready: + raise ValueError("fc1_prologue_tiles=0 requires deferred full-readiness.") if self.token_padding_block <= 0: raise ValueError("token_padding_block must be positive.") if self.sf_padding_block <= 0: raise ValueError("sf_padding_block must be positive.") if self.launch_cluster_cnt_merge_as_preferred <= 0: raise ValueError("launch_cluster_cnt_merge_as_preferred must be positive.") - if self.work_id_mode != "atomic_counter": + if ( + isinstance(self.compute_cluster_count, bool) + or not isinstance(self.compute_cluster_count, int) + or not 0 < self.compute_cluster_count <= self.launch_cluster_cnt_merge_as_preferred + ): raise ValueError( - "Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'." + f"compute_cluster_count must be a Python int in (0, {self.launch_cluster_cnt_merge_as_preferred}], got " + f"{self.compute_cluster_count}." ) + if self.work_id_mode != "atomic_counter": + raise ValueError("Phase-interleaved FC12 scheduling currently requires work_id_mode='atomic_counter'.") cluster_size = self.cluster_shape_mn[0] * self.cluster_shape_mn[1] if cluster_size > 32: raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per cluster.") @@ -561,9 +665,7 @@ def _validate_configuration(self) -> None: if fallback_cluster_shape is not None: fallback_cluster_size = fallback_cluster_shape[0] * fallback_cluster_shape[1] if fallback_cluster_size > 32: - raise ValueError( - "The atomic broadcast protocol supports at most 32 CTAs per fallback cluster." - ) + raise ValueError("The atomic broadcast protocol supports at most 32 CTAs per fallback cluster.") maximum_int32 = (1 << 31) - 1 for field_name in ("expert_count", "intermediate_gateup_size", "hidden_size"): value = getattr(self, field_name) @@ -587,7 +689,10 @@ def initialize_fallback_group(self) -> None: self._work_id_worker.initialize_fallback_group() def get_grid_shape( - self, *, max_active_clusters: Optional[int] = None, problem_desc=None + self, + *, + max_active_clusters: Optional[int] = None, + problem_desc=None, ) -> Tuple[int, int, int]: """Return the statically configured persistent launch grid.""" if ( @@ -600,11 +705,7 @@ def get_grid_shape( "launch_cluster_cnt_merge_as_preferred " f"({self.launch_cluster_cnt_merge_as_preferred})." ) - return ( - self.cluster_shape_mn[0], - self.cluster_shape_mn[1], - self.launch_cluster_cnt_merge_as_preferred, - ) + return (self.cluster_shape_mn[0], self.cluster_shape_mn[1], self.launch_cluster_cnt_merge_as_preferred) @cute.jit def assign_device_members( @@ -621,9 +722,7 @@ def assign_device_members( ) -> None: """Initialize phase-local mapping, cadence, and counter state.""" if cutlass.const_expr((expert_token_sizes is None) == (expert_token_prefix_sum is None)): - raise ValueError( - "Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided." - ) + raise ValueError("Exactly one of expert_token_sizes and expert_token_prefix_sum must be provided.") self.create_scheduler_pipelines(smem_workspace, smem_base) self._work_id_worker.assign_device_members( is_fallback_cluster=is_fallback_cluster, @@ -643,14 +742,153 @@ def assign_device_members( is_swap_ab=self.is_swap_ab, expert_token_sizes=expert_token_sizes, expert_token_prefix_sum=expert_token_prefix_sum, + dfc1_m_group=self.dfc1_m_group, ) self._control_state = _PhaseInterleaveControlState( prologue_remaining=Int32(self.fc1_prologue_tiles), - cycle_position=Int32(0), + # Stagger the steady-state cadence by physical cluster. Starting + # every persistent cluster at slot zero turns an intended rolling + # mix into chip-wide FC1/FC2 bursts; the phase offset keeps both + # streams represented in every claim wave after the prologue. + cycle_position=( + Int32(block_idx[2]) % Int32(self.interleave_cycle_length) + ), fc1_exhausted=Boolean(False), fc2_exhausted=Boolean(False), - prologue_claims_synchronized=Boolean(False), + prologue_claims_synchronized=Boolean( + self.defer_consumer_until_full_ready + ), + ) + + @cute.jit + def cluster_uniform_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Boolean: + """Observe one counter on the cluster leader and broadcast readiness.""" + if cutlass.const_expr(not self.defer_consumer_until_full_ready): + raise ValueError( + "cluster readiness polling requires " + "defer_consumer_until_full_ready=True." + ) + work_id_worker = self._work_id_worker + ready = work_id_worker.cluster_uniform_counter_ready( + counter_pointer, ready_threshold, ready_is_mask ) + self._work_id_worker = work_id_worker + return ready + + @cute.jit + def gen_next_linear1_work_unless_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Tuple[Boolean, SchedulerWorkTileBase]: + """Probe full readiness and otherwise claim one Linear1 tile. + + When Linear1 still has unclaimed work, the readiness load and atomic + claim share one cluster broadcast transaction. Once the stream is + exhausted this falls back to a probe-only transaction, avoiding + repeated tail over-claims while the held consumer drains. + """ + if cutlass.const_expr(not self.fuse_ready_probe_and_linear1_claim): + raise ValueError( + "fused readiness/Linear1 claims require " + "fuse_ready_probe_and_linear1_claim=True." + ) + + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + fc1_exhausted = control_state.fc1_exhausted + ready = Boolean(False) + + if fc1_exhausted: + ready = work_id_worker.cluster_uniform_counter_ready( + counter_pointer, ready_threshold, ready_is_mask + ) + else: + ready, linear_work_id = ( + work_id_worker.cluster_uniform_counter_ready_or_claim( + counter_pointer, + ready_threshold, + claim_stream_index=Int32(0), + ready_is_mask=ready_is_mask, + ) + ) + if not ready: + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, + self.is_swap_ab, + ) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + Int32(BlockPhase.Linear1), + cta_id_in_mapping_cluster, + task_mapping_state, + ) + ) + if not stream_has_work: + fc1_exhausted = Boolean(True) + + control_state.fc1_exhausted = fc1_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return ready, work_tile + + + @cute.jit + def gen_next_linear1_work(self) -> SchedulerWorkTileBase: + """Force one dependency-free Linear1 claim while a Linear2 tile is held. + + This corrective claim is outside the normal rolling cadence: the + deferred Linear2 claim already consumed its cadence slot. The caller + must publish a returned Linear1 tile before rechecking or publishing + the held consumer. Consequently no producer descriptor can become + trapped behind an unready consumer. + """ + if cutlass.const_expr(not self.defer_consumer_until_full_ready): + raise ValueError( + "forced Linear1 claims require " + "defer_consumer_until_full_ready=True." + ) + + work_tile = make_fc12_done_tile(self.is_swap_ab) + work_id_worker = self._work_id_worker + task_mapping_state = self._task_mapping_state + control_state = self._control_state + fc1_exhausted = control_state.fc1_exhausted + + if not fc1_exhausted: + linear_work_id = work_id_worker.claim_next_work( + Int32(0) + ) + cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( + work_id_worker.cta_coord_in_preferred_cluster, + self.is_swap_ab, + ) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + Int32(BlockPhase.Linear1), + cta_id_in_mapping_cluster, + task_mapping_state, + ) + ) + if not stream_has_work: + fc1_exhausted = Boolean(True) + + control_state.fc1_exhausted = fc1_exhausted + self._work_id_worker = work_id_worker + self._task_mapping_state = task_mapping_state + self._control_state = control_state + return work_tile @cute.jit def _wait_for_fc1_prologue_claims( @@ -664,10 +902,8 @@ def _wait_for_fc1_prologue_claims( fc1_counter_pointer = work_id_worker.get_atomic_counter_pointer(Int32(0)) minimum_claim_count = Int32(self.minimum_global_fc1_claims) if not spin_peek(fc1_counter_pointer, lambda value: value >= minimum_claim_count): - claim_target, stream_ends_before_target = ( - resolve_phase_interleaved_fc1_claim_target( - minimum_claim_count, task_mapping_state - ) + claim_target, stream_ends_before_target = resolve_phase_interleaved_fc1_claim_target( + minimum_claim_count, task_mapping_state ) spin_wait(fc1_counter_pointer, lambda value: value >= claim_target) return fc1_exhausted @@ -683,6 +919,7 @@ def gen_next_work(self) -> SchedulerWorkTileBase: cycle_position = control_state.cycle_position fc1_exhausted = control_state.fc1_exhausted fc2_exhausted = control_state.fc2_exhausted + stream_has_work = Boolean(False) prologue_claims_synchronized = control_state.prologue_claims_synchronized resolved = Boolean(False) @@ -690,7 +927,8 @@ def gen_next_work(self) -> SchedulerWorkTileBase: if fc1_exhausted and fc2_exhausted: work_tile = make_fc12_done_tile(self.is_swap_ab) resolved = Boolean(True) - else: + + if not resolved: want_fc1 = Boolean(True) if prologue_remaining <= Int32(0): is_fc2_slot = (cycle_position * Int32(self.interleave_fc2_slots)) % Int32( @@ -707,14 +945,19 @@ def gen_next_work(self) -> SchedulerWorkTileBase: atomic_counter_index = Int32(0) linear_work_id = work_id_worker.claim_next_work(atomic_counter_index) want_fc1 = work_id_worker.claimed_stream_index == Int32(0) - phase = Int32(BlockPhase.Linear2) - if want_fc1: - phase = Int32(BlockPhase.Linear1) cta_id_in_mapping_cluster = _to_fc12_mapping_cta_coord( work_id_worker.cta_coord_in_preferred_cluster, self.is_swap_ab ) - work_tile, stream_has_work, task_mapping_state = map_phase_interleaved_fc12_work_id( - linear_work_id, phase, cta_id_in_mapping_cluster, task_mapping_state + phase = Int32(BlockPhase.Linear2) + if want_fc1: + phase = Int32(BlockPhase.Linear1) + work_tile, stream_has_work, task_mapping_state = ( + map_phase_interleaved_fc12_work_id( + linear_work_id, + phase, + cta_id_in_mapping_cluster, + task_mapping_state, + ) ) if stream_has_work: if (not want_fc1) and (not prologue_claims_synchronized): @@ -725,9 +968,7 @@ def gen_next_work(self) -> SchedulerWorkTileBase: if prologue_remaining > Int32(0): prologue_remaining = prologue_remaining - Int32(1) else: - cycle_position = (cycle_position + Int32(1)) % Int32( - self.interleave_cycle_length - ) + cycle_position = (cycle_position + Int32(1)) % Int32(self.interleave_cycle_length) resolved = Boolean(True) else: if want_fc1: @@ -764,9 +1005,7 @@ def __new_from_mlir_values__(self, values: list) -> "PhaseInterleavedFc12Schedul def rebuild(state): nonlocal value_index state_value_count = len(extract_mlir_values(state)) - rebuilt_state = new_from_mlir_values( - state, values[value_index : value_index + state_value_count] - ) + rebuilt_state = new_from_mlir_values(state, values[value_index : value_index + state_value_count]) value_index += state_value_count return rebuilt_state @@ -792,11 +1031,16 @@ def rebuild(state): "work_id_mode", "is_swap_ab", "launch_cluster_cnt_merge_as_preferred", + "compute_cluster_count", "non_clc_mixed_cga_config", "mapping_cta_tile_shape_mnk", "mapping_cluster_shape_mn", "blocks_fc1", "blocks_fc2", + "dfc1_m_group", + "defer_consumer_until_full_ready", + "fuse_ready_probe_and_linear1_claim", + "fc2_claims_per_token_block", "interleave_fc2_slots", "interleave_cycle_length", "minimum_global_fc1_claims", @@ -805,8 +1049,4 @@ def rebuild(state): return result -__all__ = [ - "BlackwellFusedFc12Scheduler", - "PhaseInterleavedFc12Scheduler", - "minimum_phase_interleave_hint", -] +__all__ = ["BlackwellFusedFc12Scheduler", "PhaseInterleavedFc12Scheduler", "minimum_phase_interleave_hint"] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py index 90f340877572..cad226f2dd3a 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/non_clc_mixed_cga.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Reusable preferred/fallback cluster scheduling without hardware CLC.""" import dataclasses @@ -19,6 +16,8 @@ AtomicCounterWorkIdState, FixedGroupMixedCgaAtomicCounterWorkIdState, GridStrideWorkIdState, + broadcast_atomic_counter_ready, + broadcast_atomic_counter_ready_or_claim, claim_work_id, initialize_fixed_group_mixed_cga_work_id_state, ) @@ -49,19 +48,14 @@ def __post_init__(self) -> None: and not isinstance(self.fallback_cluster_count, bool) and self.fallback_cluster_count == 0 ) - if ( - self.fallback_cluster_shape == self.preferred_cluster_shape - or has_no_fallback_clusters - ): + if self.fallback_cluster_shape == self.preferred_cluster_shape or has_no_fallback_clusters: object.__setattr__(self, "fallback_cluster_shape", None) object.__setattr__(self, "preferred_cluster_count", None) object.__setattr__(self, "fallback_cluster_count", None) if self.fallback_cluster_shape is None: if not self._is_positive_int(self.launch_cluster_count): - raise ValueError( - "launch_cluster_count must be positive when fallback_cluster_shape is absent." - ) + raise ValueError("launch_cluster_count must be positive when fallback_cluster_shape is absent.") mn_split_factors = (1, 1) split_factor = 1 launch_cluster_cnt_merge_as_preferred = self.launch_cluster_count @@ -70,9 +64,7 @@ def __post_init__(self) -> None: if len(self.fallback_cluster_shape) != len(self.preferred_cluster_shape): raise ValueError("preferred and fallback cluster shapes must have equal rank.") if not self._is_positive_int(self.preferred_cluster_count): - raise ValueError( - "preferred_cluster_count must be positive when fallback_cluster_shape is present." - ) + raise ValueError("preferred_cluster_count must be positive when fallback_cluster_shape is present.") if ( isinstance(self.fallback_cluster_count, bool) or not isinstance(self.fallback_cluster_count, int) @@ -92,16 +84,12 @@ def __post_init__(self) -> None: self.preferred_cluster_shape, self.fallback_cluster_shape ) ): - raise ValueError( - "Every preferred cluster dimension must be divisible by its fallback dimension." - ) + raise ValueError("Every preferred cluster dimension must be divisible by its fallback dimension.") split_factor = self._shape_size(mn_split_factors) is_mixed = self.fallback_cluster_shape != self.preferred_cluster_shape if is_mixed: if split_factor > 16 or split_factor & (split_factor - 1): - raise ValueError( - "The preferred/fallback cluster split factor must be a power of two at most 16." - ) + raise ValueError("The preferred/fallback cluster split factor must be a power of two at most 16.") if self.fallback_cluster_count % split_factor != 0: raise ValueError( "fallback_cluster_count must be divisible by the preferred/fallback cluster split factor." @@ -110,22 +98,14 @@ def __post_init__(self) -> None: self.preferred_cluster_count + self.fallback_cluster_count // split_factor ) else: - launch_cluster_cnt_merge_as_preferred = ( - self.preferred_cluster_count + self.fallback_cluster_count - ) + launch_cluster_cnt_merge_as_preferred = self.preferred_cluster_count + self.fallback_cluster_count if launch_cluster_cnt_merge_as_preferred <= 0: - raise ValueError( - "The resolved launch must contain at least one cluster merged as preferred." - ) + raise ValueError("The resolved launch must contain at least one cluster merged as preferred.") object.__setattr__(self, "mn_split_factors", mn_split_factors) object.__setattr__(self, "split_factor", split_factor) - object.__setattr__( - self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred - ) - object.__setattr__( - self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size - ) + object.__setattr__(self, "launch_cluster_cnt_merge_as_preferred", launch_cluster_cnt_merge_as_preferred) + object.__setattr__(self, "total_cta_cnt", launch_cluster_cnt_merge_as_preferred * preferred_cluster_size) object.__setattr__(self, "is_mixed", is_mixed) @staticmethod @@ -162,16 +142,27 @@ def __init__( config: NonClcMixedCgaConfig, work_id_mode: WorkIdAcquisitionMode, stream_count: int, + work_id_stride: Optional[int] = None, ) -> None: if work_id_mode not in ("grid_stride", "atomic_counter"): - raise ValueError( - "Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs." - ) + raise ValueError("Non-CLC mixed-CGA scheduling supports grid_stride or atomic_counter work IDs.") if isinstance(stream_count, bool) or not isinstance(stream_count, int) or stream_count <= 0: raise ValueError("stream_count must be a positive Python int.") + if work_id_stride is None: + work_id_stride = config.launch_cluster_cnt_merge_as_preferred + if isinstance(work_id_stride, bool) or not isinstance(work_id_stride, int) or work_id_stride <= 0: + raise ValueError("work_id_stride must be a positive Python int.") self.config = config self.work_id_mode = work_id_mode self.stream_count = stream_count + # The work stream can have fewer participants than the physical launch + # when some clusters only communicate. Atomic claims do not use this. + self.work_id_stride = work_id_stride + + @property + def max_fallback_cluster_count(self) -> int: + """Upper bound on physical fallback clusters: every preferred-cluster slot of the launch split.""" + return self.config.launch_cluster_cnt_merge_as_preferred * self.config.split_factor def register_smem_regions(self, smem_workspace: SmemWorkspace) -> None: """Register the cluster-local claim broadcast channel.""" @@ -197,10 +188,12 @@ def register_device_workspace(self, device_workspace: DeviceWorkspace) -> None: buffer_space="local", reset="tail_reset", ) + # Under PDL more slots can split into fallback clusters than the host plan assumed: size the group + # table for every slot splitting. device_workspace.register( self.fallback_group_token_region, cutlass.Int64, - (self.config.fallback_cluster_count,), + (self.max_fallback_cluster_count,), buffer_space="local", reset="tail_reset", ) @@ -240,9 +233,9 @@ def assign_device_members( for dimension_idx in cutlass.range_constexpr(len(self.config.mn_split_factors)): preferred_dimension = self.config.preferred_cluster_shape[dimension_idx] fallback_dimension = self.config.fallback_cluster_shape[dimension_idx] - inner_coordinate = ( - Int32(block_idx[dimension_idx]) % Int32(preferred_dimension) - ) // Int32(fallback_dimension) + inner_coordinate = (Int32(block_idx[dimension_idx]) % Int32(preferred_dimension)) // Int32( + fallback_dimension + ) flattened_index = flattened_index + inner_coordinate * Int32(dimension_stride) dimension_stride *= self.config.mn_split_factors[dimension_idx] cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( @@ -255,20 +248,15 @@ def assign_device_members( cluster_pipeline = pipeline.PipelineAsync.create( num_stages=1, producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), - consumer_group=pipeline.CooperativeGroup( - pipeline.Agent.Thread, 32 * active_cluster_size - ), - barrier_storage=smem_workspace.ptr( - self.cluster_pipeline_mbarriers_region, smem_base - ), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * active_cluster_size), + barrier_storage=smem_workspace.ptr(self.cluster_pipeline_mbarriers_region, smem_base), defer_sync=True, ) atomic_counter_state = AtomicCounterWorkIdState( counter_pointer=device_workspace.ptr(self.work_id_counter_region), counter_count=self.stream_count, broadcast_pointer=smem_workspace.ptr(self.cluster_broadcast_region, smem_base), - is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) - == Int32(0), + is_leader_cta=(cta_coord_in_active_cluster[0] + cta_coord_in_active_cluster[1]) == Int32(0), cluster_pipeline=cluster_pipeline, producer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1), consumer_state=pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1), @@ -278,12 +266,10 @@ def assign_device_members( self._cta_coord_in_active_cluster = cta_coord_in_active_cluster self._work_id_state = FixedGroupMixedCgaAtomicCounterWorkIdState( atomic_counter_state=atomic_counter_state, - registration_counter_pointer=device_workspace.ptr( - self.fallback_registration_counter_region - ), + registration_counter_pointer=device_workspace.ptr(self.fallback_registration_counter_region), group_token_pointer=device_workspace.ptr(self.fallback_group_token_region), split_factor=self.config.split_factor, - fallback_cluster_count=self.config.fallback_cluster_count, + fallback_cluster_count=self.max_fallback_cluster_count, is_fallback_cluster=is_fallback_cluster, fallback_group_idx=Int32(0), in_group_idx=Int32(0), @@ -296,7 +282,7 @@ def assign_device_members( else: self._work_id_state = GridStrideWorkIdState( next_work_id=Int32(block_idx[2]), - work_id_stride=Int32(self.config.launch_cluster_cnt_merge_as_preferred), + work_id_stride=Int32(self.work_id_stride), ) self.claimed_stream_index = Int32(0) @@ -304,9 +290,7 @@ def assign_device_members( def initialize_fallback_group(self) -> None: """Register one physical fallback cluster with its fixed logical group.""" if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): - self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state( - self._work_id_state - ) + self._work_id_state = initialize_fixed_group_mixed_cga_work_id_state(self._work_id_state) if self._work_id_state.is_fallback_cluster: self.cta_coord_in_preferred_cluster = self._preferred_cluster_cta_coord( self._cta_coord_in_active_cluster, self._work_id_state.in_group_idx @@ -321,10 +305,8 @@ def _preferred_cluster_cta_coord( self.config.mn_split_factors[1] ) return ( - cta_coord_in_active_cluster[0] - + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), - cta_coord_in_active_cluster[1] - + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), + cta_coord_in_active_cluster[0] + inner_cluster_m * Int32(self.config.fallback_cluster_shape[0]), + cta_coord_in_active_cluster[1] + inner_cluster_n * Int32(self.config.fallback_cluster_shape[1]), Int32(0), ) @@ -334,18 +316,14 @@ def get_atomic_counter_pointer(self, stream_index: Int32) -> cute.Pointer: if cutlass.const_expr(self.work_id_mode != "atomic_counter"): raise ValueError("Work-ID counters are only available in atomic_counter mode.") work_id_state = self._work_id_state - if cutlass.const_expr( - isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState) - ): + if cutlass.const_expr(isinstance(work_id_state, FixedGroupMixedCgaAtomicCounterWorkIdState)): work_id_state = work_id_state.atomic_counter_state return work_id_state.counter_pointer + stream_index @cute.jit def claim_next_work(self, stream_index=0) -> Int32: """Claim the next canonical work ID and update the preferred CTA coordinate.""" - claimed_work_id, self._work_id_state = claim_work_id( - self._work_id_state, atomic_counter_index=stream_index - ) + claimed_work_id, self._work_id_state = claim_work_id(self._work_id_state, atomic_counter_index=stream_index) canonical_work_id = claimed_work_id if cutlass.const_expr(self.config.is_mixed and self.work_id_mode == "atomic_counter"): cta_coord_in_preferred_cluster = self._cta_coord_in_active_cluster @@ -359,6 +337,69 @@ def claim_next_work(self, stream_index=0) -> Int32: self.claimed_stream_index = Int32(stream_index) return canonical_work_id + @cute.jit + def cluster_uniform_counter_ready( + self, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, + ) -> Boolean: + """Return one leader-loaded readiness predicate to every active CTA. + + V1 deliberately excludes fixed fallback groups. Their several + physical clusters must rendezvous for every logical work-ID claim, so + independently polling readiness could make group members take + different claim paths. Homogeneous preferred clusters have one shared + claim/broadcast channel and therefore one uniform decision. + """ + if cutlass.const_expr(self.work_id_mode != "atomic_counter"): + raise ValueError( + "cluster_uniform_counter_ready requires atomic-counter work IDs." + ) + if cutlass.const_expr(self.config.is_mixed): + raise NotImplementedError( + "cluster-uniform readiness polling does not yet support fixed " + "preferred/fallback groups." + ) + ready, self._work_id_state = broadcast_atomic_counter_ready( + self._work_id_state, + counter_pointer, + ready_threshold, + ready_is_mask, + ) + return ready + + @cute.jit + def cluster_uniform_counter_ready_or_claim( + self, + ready_counter_pointer: cute.Pointer, + ready_threshold: Int32, + claim_stream_index=0, + ready_is_mask: bool = False, + ) -> Tuple[Boolean, Int32]: + """Probe readiness or claim one ID, uniformly across one cluster.""" + if cutlass.const_expr(self.work_id_mode != "atomic_counter"): + raise ValueError( + "cluster_uniform_counter_ready_or_claim requires " + "atomic-counter work IDs." + ) + if cutlass.const_expr(self.config.is_mixed): + raise NotImplementedError( + "fused cluster readiness/claim does not support fixed " + "preferred/fallback groups." + ) + ready, claimed_work_id, self._work_id_state = ( + broadcast_atomic_counter_ready_or_claim( + self._work_id_state, + ready_counter_pointer, + ready_threshold, + atomic_counter_index=claim_stream_index, + ready_is_mask=ready_is_mask, + ) + ) + self.claimed_stream_index = Int32(claim_stream_index) + return ready, claimed_work_id + def __extract_mlir_values__(self) -> List[ir.Value]: values: List[ir.Value] = [] values.extend(extract_mlir_values(self._work_id_state)) @@ -374,14 +415,15 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "NonClcMixedCgaSch def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result result = type(self)( - config=self.config, work_id_mode=self.work_id_mode, stream_count=self.stream_count + config=self.config, + work_id_mode=self.work_id_mode, + stream_count=self.stream_count, + work_id_stride=self.work_id_stride, ) result._work_id_state = rebuild(self._work_id_state) result.cta_coord_in_preferred_cluster = rebuild(self.cta_coord_in_preferred_cluster) diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py index f9e7b799aed5..846fbe346061 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/kernel_src/schedulers/work_id_claim.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Blackwell persistent work-ID claim backends.""" import dataclasses @@ -10,13 +7,10 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline from cutlass._mlir import ir +from cutlass._mlir.dialects import nvvm as _nvvm from cutlass.cutlass_dsl import Boolean, Int32, Int64, extract_mlir_values, new_from_mlir_values -from ...helpers.ptx_helpers import ( - mbarrier_arrive_expect_tx_on_peer, - nanosleep, - store_i32_to_peer_cluster_smem_async, -) +from ...helpers.ptx_helpers import mbarrier_arrive_expect_tx_on_peer, nanosleep, store_i32_to_peer_cluster_smem_async class GridStrideWorkIdState: @@ -42,9 +36,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridStrideWorkIdS ) return type(self)( next_work_id=new_from_mlir_values(self.next_work_id, values[:next_work_id_value_count]), - work_id_stride=new_from_mlir_values( - self.work_id_stride, values[next_work_id_value_count:] - ), + work_id_stride=new_from_mlir_values(self.work_id_stride, values[next_work_id_value_count:]), ) @@ -62,11 +54,7 @@ def __init__( consumer_state, cluster_size: int | Int32, ) -> None: - if ( - isinstance(counter_count, bool) - or not isinstance(counter_count, int) - or counter_count <= 0 - ): + if isinstance(counter_count, bool) or not isinstance(counter_count, int) or counter_count <= 0: raise ValueError("counter_count must be a positive Python int.") self.counter_pointer = counter_pointer self.counter_count = counter_count @@ -97,9 +85,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "AtomicCounterWork def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -111,9 +97,7 @@ def rebuild(field): cluster_pipeline=self.cluster_pipeline, producer_state=rebuild(self.producer_state), consumer_state=rebuild(self.consumer_state), - cluster_size=rebuild(self.cluster_size) - if isinstance(self.cluster_size, Int32) - else self.cluster_size, + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, ) if value_index != len(values): raise ValueError( @@ -180,17 +164,13 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values.extend(extract_mlir_values(field)) return values - def __new_from_mlir_values__( - self, values: List[ir.Value] - ) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "FixedGroupMixedCgaAtomicCounterWorkIdState": value_index = 0 def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - result = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + result = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return result @@ -234,9 +214,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GridWorkId": if len(values) != 4: raise ValueError(f"GridWorkId expects four MLIR values, got {len(values)}.") fields = (self.grid_m, self.grid_n, self.grid_l, self.is_valid) - return type(self)( - *(new_from_mlir_values(field, [value]) for field, value in zip(fields, values)) - ) + return type(self)(*(new_from_mlir_values(field, [value]) for field, value in zip(fields, values))) class ClusterLaunchControlWorkIdState: @@ -291,9 +269,7 @@ def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClusterLaunchCont def rebuild(field): nonlocal value_index field_value_count = len(extract_mlir_values(field)) - rebuilt_field = new_from_mlir_values( - field, values[value_index : value_index + field_value_count] - ) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) value_index += field_value_count return rebuilt_field @@ -317,10 +293,93 @@ def rebuild(field): return result +class ClcMixedCgaWorkIdState: + """Cluster-wide state for hardware CLC claims on a runtime-sized (mixed-CGA) cluster. + + ``PipelineClcFetchAsync`` needs a static cluster layout, so this drives a plain ``PipelineAsync`` + with the atomic broadcast's peer barrier protocol instead. + """ + + def __init__( + self, + response_pending: Boolean, + grid_m: Int32, + grid_n: Int32, + grid_l: Int32, + response_is_valid: Boolean, + cta_coord_in_cluster: cute.Coord, + cluster_pipeline: pipeline.PipelineAsync, + producer_state, + consumer_state, + is_leader_cta: Boolean, + response_pointer: cute.Pointer, + cluster_size: int | Int32, + ) -> None: + self.response_pending = response_pending + self.grid_m = grid_m + self.grid_n = grid_n + self.grid_l = grid_l + self.response_is_valid = response_is_valid + self.cta_coord_in_cluster = cta_coord_in_cluster + self.cluster_pipeline = cluster_pipeline + self.producer_state = producer_state + self.consumer_state = consumer_state + self.is_leader_cta = is_leader_cta + self.response_pointer = response_pointer + self.cluster_size = cluster_size + + def __extract_mlir_values__(self) -> List[ir.Value]: + values: List[ir.Value] = [] + for field in ( + self.response_pending, + self.grid_m, + self.grid_n, + self.grid_l, + self.response_is_valid, + self.cta_coord_in_cluster, + self.producer_state, + self.consumer_state, + self.is_leader_cta, + self.response_pointer, + ): + values.extend(extract_mlir_values(field)) + if isinstance(self.cluster_size, Int32): + values.extend(extract_mlir_values(self.cluster_size)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "ClcMixedCgaWorkIdState": + value_index = 0 + + def rebuild(field): + nonlocal value_index + field_value_count = len(extract_mlir_values(field)) + rebuilt_field = new_from_mlir_values(field, values[value_index : value_index + field_value_count]) + value_index += field_value_count + return rebuilt_field + + result = type(self)( + response_pending=rebuild(self.response_pending), + grid_m=rebuild(self.grid_m), + grid_n=rebuild(self.grid_n), + grid_l=rebuild(self.grid_l), + response_is_valid=rebuild(self.response_is_valid), + cta_coord_in_cluster=rebuild(self.cta_coord_in_cluster), + cluster_pipeline=self.cluster_pipeline, + producer_state=rebuild(self.producer_state), + consumer_state=rebuild(self.consumer_state), + is_leader_cta=rebuild(self.is_leader_cta), + response_pointer=rebuild(self.response_pointer), + cluster_size=rebuild(self.cluster_size) if isinstance(self.cluster_size, Int32) else self.cluster_size, + ) + if value_index != len(values): + raise ValueError( + f"ClcMixedCgaWorkIdState MLIR value count mismatch: consumed {value_index}, got {len(values)}." + ) + return result + + @cute.jit -def _claim_grid_stride_work_id( - work_id_state: GridStrideWorkIdState, -) -> Tuple[Int32, GridStrideWorkIdState]: +def _claim_grid_stride_work_id(work_id_state: GridStrideWorkIdState) -> Tuple[Int32, GridStrideWorkIdState]: """Claim the next ID from one monotonic grid-stride stream.""" linear_work_id = work_id_state.next_work_id work_id_state.next_work_id = linear_work_id + work_id_state.work_id_stride @@ -345,17 +404,13 @@ def _claim_atomic_counter_work_id( if work_id_state.is_leader_cta: cluster_pipeline.producer_acquire(work_id_state.producer_state) - full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( - work_id_state.producer_state.index - ) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) atomic_work_id = Int32(0) if lane_idx == Int32(0): atomic_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1)) - atomic_work_id = cute.arch.shuffle_sync( - atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + atomic_work_id = cute.arch.shuffle_sync(atomic_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) if lane_idx < Int32(work_id_state.cluster_size): store_i32_to_peer_cluster_smem_async( work_id_state.broadcast_pointer, atomic_work_id, full_barrier_pointer, lane_idx @@ -371,6 +426,165 @@ def _claim_atomic_counter_work_id( return linear_work_id, work_id_state +@cute.jit +def broadcast_atomic_counter_ready( + work_id_state: AtomicCounterWorkIdState, + counter_pointer: cute.Pointer, + ready_threshold: Int32, + ready_is_mask: bool = False, +) -> Tuple[Boolean, AtomicCounterWorkIdState]: + """Load one global readiness counter on the leader and broadcast the predicate. + + The scheduler makes publication decisions at preferred-cluster granularity, + so a CTA-local peek is not sufficient: the two scheduler warps must either + both defer or both publish their CTA halves of a work tile. Reusing the + atomic work-ID broadcast channel gives the readiness decision the same + cluster-uniform ordering as a work-ID claim. + """ + if cutlass.const_expr(not isinstance(ready_is_mask, bool)): + raise TypeError("ready_is_mask must be a Python bool.") + + broadcast_tensor = cute.make_tensor( + work_id_state.broadcast_pointer, cute.make_layout((1,)) + ) + cluster_pipeline = work_id_state.cluster_pipeline + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + work_id_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + ready_payload = Int32(0) + if lane_idx == Int32(0): + observed = cute.arch.load( + counter_pointer, Int32, sem="acquire", scope="gpu" + ) + if cutlass.const_expr(ready_is_mask): + if (observed & ready_threshold) == ready_threshold: + ready_payload = Int32(1) + else: + if observed >= ready_threshold: + ready_payload = Int32(1) + ready_payload = cute.arch.shuffle_sync( + ready_payload, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 + ) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, + ready_payload, + full_barrier_pointer, + lane_idx, + ) + mbarrier_arrive_expect_tx_on_peer( + full_barrier_pointer, Int32(4), lane_idx + ) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + ready_payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive( + work_id_state.consumer_state.index, Int32(0) + ) + work_id_state.consumer_state.advance() + return ready_payload != Int32(0), work_id_state + + +@cute.jit +def broadcast_atomic_counter_ready_or_claim( + work_id_state: AtomicCounterWorkIdState, + ready_counter_pointer: cute.Pointer, + ready_threshold: Int32, + atomic_counter_index=0, + ready_is_mask: bool = False, +) -> Tuple[Boolean, Int32, AtomicCounterWorkIdState]: + """Probe readiness or claim one work ID in one cluster transaction. + + The leader first performs the same acquire readiness load as + :func:`broadcast_atomic_counter_ready`. When the counter is not ready it + immediately claims one ID from ``atomic_counter_index`` before broadcasting + the result. A negative payload denotes readiness; valid work IDs are + non-negative, so the one-word cluster broadcast ABI remains unchanged. + + This primitive is intentionally limited to a single claimed ID. It is the + bundle-1 fast path used while a full-readiness-gated consumer is held + outside the execution FIFO. + """ + if cutlass.const_expr(not isinstance(ready_is_mask, bool)): + raise TypeError("ready_is_mask must be a Python bool.") + invalid_static_index = isinstance(atomic_counter_index, int) and ( + atomic_counter_index < 0 + or atomic_counter_index >= work_id_state.counter_count + ) + if cutlass.const_expr(invalid_static_index): + raise ValueError( + "atomic_counter_index must be in [0, " + f"{work_id_state.counter_count}), got {atomic_counter_index}." + ) + + broadcast_tensor = cute.make_tensor( + work_id_state.broadcast_pointer, cute.make_layout((1,)) + ) + cluster_pipeline = work_id_state.cluster_pipeline + selected_counter_pointer = ( + work_id_state.counter_pointer + Int32(atomic_counter_index) + ) + + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( + work_id_state.producer_state.index + ) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + # -1 is the ready sentinel. A failed probe replaces it with the + # non-negative ID returned by the Linear1 atomic counter. + payload = Int32(-1) + if lane_idx == Int32(0): + observed = cute.arch.load( + ready_counter_pointer, Int32, sem="acquire", scope="gpu" + ) + ready = Boolean(False) + if cutlass.const_expr(ready_is_mask): + ready = (observed & ready_threshold) == ready_threshold + else: + ready = observed >= ready_threshold + if not ready: + payload = cute.arch.atomic_add( + selected_counter_pointer, Int32(1) + ) + payload = Int32( + cute.arch.shuffle_sync( + payload, + offset=0, + mask=0xFFFFFFFF, + mask_and_clamp=31, + ) + ) + if lane_idx < Int32(work_id_state.cluster_size): + store_i32_to_peer_cluster_smem_async( + work_id_state.broadcast_pointer, + payload, + full_barrier_pointer, + lane_idx, + ) + mbarrier_arrive_expect_tx_on_peer( + full_barrier_pointer, Int32(4), lane_idx + ) + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + payload = broadcast_tensor[0] + cute.arch.fence_acq_rel_cta() + cluster_pipeline.sync_object_empty.arrive( + work_id_state.consumer_state.index, Int32(0) + ) + work_id_state.consumer_state.advance() + return payload == Int32(-1), payload, work_id_state + + @cute.jit def initialize_fixed_group_mixed_cga_work_id_state( work_id_state: FixedGroupMixedCgaAtomicCounterWorkIdState, @@ -378,9 +592,7 @@ def initialize_fixed_group_mixed_cga_work_id_state( """Register one physical fallback cluster and broadcast its fixed group coordinates.""" atomic_counter_state = work_id_state.atomic_counter_state if work_id_state.is_fallback_cluster: - broadcast_tensor = cute.make_tensor( - atomic_counter_state.broadcast_pointer, cute.make_layout((1,)) - ) + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) cluster_pipeline = atomic_counter_state.cluster_pipeline if atomic_counter_state.is_leader_cta: cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) @@ -395,16 +607,11 @@ def initialize_fixed_group_mixed_cga_work_id_state( work_id_state.registration_counter_pointer, Int32(1), sem="relaxed", scope="gpu" ) fallback_ordinal = Int32( - cute.arch.shuffle_sync( - fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + cute.arch.shuffle_sync(fallback_ordinal, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) ) if lane_idx < Int32(atomic_counter_state.cluster_size): store_i32_to_peer_cluster_smem_async( - atomic_counter_state.broadcast_pointer, - fallback_ordinal, - full_barrier_pointer, - lane_idx, + atomic_counter_state.broadcast_pointer, fallback_ordinal, full_barrier_pointer, lane_idx ) mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) atomic_counter_state.producer_state.advance() @@ -412,16 +619,12 @@ def initialize_fixed_group_mixed_cga_work_id_state( cluster_pipeline.consumer_wait(atomic_counter_state.consumer_state) fallback_ordinal = broadcast_tensor[0] cute.arch.fence_acq_rel_cta() - cluster_pipeline.sync_object_empty.arrive( - atomic_counter_state.consumer_state.index, Int32(0) - ) + cluster_pipeline.sync_object_empty.arrive(atomic_counter_state.consumer_state.index, Int32(0)) atomic_counter_state.consumer_state.advance() fallback_group_idx = fallback_ordinal // Int32(work_id_state.split_factor) work_id_state.fallback_group_idx = fallback_group_idx - work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32( - work_id_state.split_factor - ) + work_id_state.in_group_idx = fallback_ordinal - fallback_group_idx * Int32(work_id_state.split_factor) work_id_state.atomic_counter_state = atomic_counter_state return work_id_state @@ -440,17 +643,13 @@ def _claim_fixed_group_fallback_work_id( f"atomic_counter_index must be in [0, {atomic_counter_state.counter_count}), got {atomic_counter_index}." ) - broadcast_tensor = cute.make_tensor( - atomic_counter_state.broadcast_pointer, cute.make_layout((1,)) - ) + broadcast_tensor = cute.make_tensor(atomic_counter_state.broadcast_pointer, cute.make_layout((1,))) cluster_pipeline = atomic_counter_state.cluster_pipeline selected_counter_pointer = atomic_counter_state.counter_pointer + Int32(atomic_counter_index) if atomic_counter_state.is_leader_cta: cluster_pipeline.producer_acquire(atomic_counter_state.producer_state) - full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier( - atomic_counter_state.producer_state.index - ) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(atomic_counter_state.producer_state.index) thread_idx, _, _ = cute.arch.thread_idx() lane_idx = thread_idx % Int32(32) group_base_offset = work_id_state.fallback_group_idx * Int32(work_id_state.split_factor) @@ -462,9 +661,7 @@ def _claim_fixed_group_fallback_work_id( while not all_members_consumed: observed_token = work_id_state.previous_token if lane_idx < Int32(work_id_state.split_factor): - observed_token = cute.arch.load( - group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu" - ) + observed_token = cute.arch.load(group_token_pointer + lane_idx, Int64, sem="acquire", scope="gpu") lane_is_ready = (lane_idx >= Int32(work_id_state.split_factor)) | ( observed_token == work_id_state.previous_token ) @@ -475,18 +672,12 @@ def _claim_fixed_group_fallback_work_id( claimed_work_id = Int32(0) if lane_idx == Int32(0): - claimed_work_id = cute.arch.atomic_add( - selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu" - ) + claimed_work_id = cute.arch.atomic_add(selected_counter_pointer, Int32(1), sem="relaxed", scope="gpu") claimed_work_id = Int32( - cute.arch.shuffle_sync( - claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31 - ) + cute.arch.shuffle_sync(claimed_work_id, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) ) claimed_payload = claimed_work_id | (Int32(atomic_counter_index) << Int32(31)) - token = (Int64(work_id_state.next_generation) << Int64(32)) | ( - Int64(claimed_payload) & Int64(0xFFFFFFFF) - ) + token = (Int64(work_id_state.next_generation) << Int64(32)) | (Int64(claimed_payload) & Int64(0xFFFFFFFF)) if lane_idx == Int32(0): cute.arch.store(group_token_pointer, token, sem="relaxed", scope="gpu") work_id_state.previous_token = token @@ -497,36 +688,22 @@ def _claim_fixed_group_fallback_work_id( token_high = Int32(0) token_low = Int32(0) if lane_idx == Int32(0): - observed_token = cute.arch.load( - group_token_pointer, Int64, sem="relaxed", scope="gpu" - ) + observed_token = cute.arch.load(group_token_pointer, Int64, sem="relaxed", scope="gpu") token_high = Int32(observed_token >> Int64(32)) token_low = Int32(observed_token & Int64(0xFFFFFFFF)) - token_high = Int32( - cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) - ) - token_low = Int32( - cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31) - ) + token_high = Int32(cute.arch.shuffle_sync(token_high, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) + token_low = Int32(cute.arch.shuffle_sync(token_low, offset=0, mask=0xFFFFFFFF, mask_and_clamp=31)) token = (Int64(token_high) << Int64(32)) | (Int64(token_low) & Int64(0xFFFFFFFF)) if token == work_id_state.previous_token: nanosleep(500) claimed_payload = Int32(token & Int64(0xFFFFFFFF)) if lane_idx == Int32(0): - cute.arch.store( - group_token_pointer + work_id_state.in_group_idx, - token, - sem="relaxed", - scope="gpu", - ) + cute.arch.store(group_token_pointer + work_id_state.in_group_idx, token, sem="relaxed", scope="gpu") work_id_state.previous_token = token if lane_idx < Int32(atomic_counter_state.cluster_size): store_i32_to_peer_cluster_smem_async( - atomic_counter_state.broadcast_pointer, - claimed_payload, - full_barrier_pointer, - lane_idx, + atomic_counter_state.broadcast_pointer, claimed_payload, full_barrier_pointer, lane_idx ) mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(4), lane_idx) atomic_counter_state.producer_state.advance() @@ -555,9 +732,7 @@ def _claim_fixed_group_mixed_cga_work_id( work_id_state.atomic_counter_state = atomic_counter_state work_id_state.claimed_counter_index = Int32(atomic_counter_index) else: - linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id( - work_id_state, atomic_counter_index - ) + linear_work_id, work_id_state = _claim_fixed_group_fallback_work_id(work_id_state, atomic_counter_index) return linear_work_id, work_id_state @@ -578,9 +753,7 @@ def _claim_cluster_launch_control_work_id( state_before_leader_query = work_id_state if work_id_state.is_leader_cta: work_id_state.cluster_pipeline.producer_acquire(work_id_state.producer_state) - response_barrier = work_id_state.cluster_pipeline.producer_get_barrier( - work_id_state.producer_state - ) + response_barrier = work_id_state.cluster_pipeline.producer_get_barrier(work_id_state.producer_state) with cute.arch.elect_one(): cute.arch.issue_clc_query(response_barrier, work_id_state.response_pointer) else: @@ -613,6 +786,68 @@ def _claim_cluster_launch_control_work_id( ) +@cute.jit +def _claim_clc_mixed_cga_work_id( + work_id_state: ClcMixedCgaWorkIdState, +) -> Tuple[GridWorkId, ClcMixedCgaWorkIdState]: + """Claim the next canceled cluster (first claim: own blockIdx) and return this CTA's grid coordinate.""" + use_bootstrap = work_id_state.response_pending + state_before_bootstrap = work_id_state + if use_bootstrap: + work_id_state.response_pending = Boolean(False) + else: + work_id_state = state_before_bootstrap + + state_before_query = work_id_state + if not use_bootstrap: + cluster_pipeline = work_id_state.cluster_pipeline + state_before_leader_query = work_id_state + if work_id_state.is_leader_cta: + cluster_pipeline.producer_acquire(work_id_state.producer_state) + full_barrier_pointer = cluster_pipeline.sync_object_full.get_barrier(work_id_state.producer_state.index) + thread_idx, _, _ = cute.arch.thread_idx() + lane_idx = thread_idx % Int32(32) + if lane_idx < Int32(work_id_state.cluster_size): + # 16-byte CLC response. + mbarrier_arrive_expect_tx_on_peer(full_barrier_pointer, Int32(16), lane_idx) + with cute.arch.elect_one(): + cute.arch.issue_clc_query(full_barrier_pointer, work_id_state.response_pointer) + else: + work_id_state = state_before_leader_query + work_id_state.producer_state.advance() + + cluster_pipeline.consumer_wait(work_id_state.consumer_state) + (cluster_origin_m, cluster_origin_n, grid_l, response_is_valid) = cute.arch.clc_response( + work_id_state.response_pointer + ) + cute.arch.fence_acq_rel_cta() + cute.arch.fence_proxy("async.shared", space="cta") + # Cluster-scoped release orders this CTA's read of the response before the leader's next overwrite. + cute.arch.mbarrier_arrive( + cluster_pipeline.sync_object_empty.get_barrier(work_id_state.consumer_state.index), + Int32(0), + scope=_nvvm.MemScopeKind.CLUSTER, + ) + work_id_state.consumer_state.advance() + + work_id_state.grid_m = cluster_origin_m + work_id_state.cta_coord_in_cluster[0] + work_id_state.grid_n = cluster_origin_n + work_id_state.cta_coord_in_cluster[1] + work_id_state.grid_l = grid_l + work_id_state.response_is_valid = response_is_valid != Int32(0) + else: + work_id_state = state_before_query + + return ( + GridWorkId( + grid_m=work_id_state.grid_m, + grid_n=work_id_state.grid_n, + grid_l=work_id_state.grid_l, + is_valid=work_id_state.response_is_valid, + ), + work_id_state, + ) + + @cute.jit def claim_work_id(work_id_state, atomic_counter_index=0): """Claim the next work ID using the backend encoded by the state type.""" @@ -624,15 +859,20 @@ def claim_work_id(work_id_state, atomic_counter_index=0): return _claim_fixed_group_mixed_cga_work_id(work_id_state, atomic_counter_index) if cutlass.const_expr(isinstance(work_id_state, ClusterLaunchControlWorkIdState)): return _claim_cluster_launch_control_work_id(work_id_state) + if cutlass.const_expr(isinstance(work_id_state, ClcMixedCgaWorkIdState)): + return _claim_clc_mixed_cga_work_id(work_id_state) raise TypeError(f"Unsupported work-ID state: {type(work_id_state).__name__}.") __all__ = [ "AtomicCounterWorkIdState", + "ClcMixedCgaWorkIdState", "ClusterLaunchControlWorkIdState", "FixedGroupMixedCgaAtomicCounterWorkIdState", "GridStrideWorkIdState", "GridWorkId", + "broadcast_atomic_counter_ready", + "broadcast_atomic_counter_ready_or_claim", "claim_work_id", "initialize_fixed_group_mixed_cga_work_id_state", ] diff --git a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py index 82d5d0c7d11b..7ff5730d73c4 100644 --- a/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py +++ b/tensorrt_llm/_torch/moe/kernels/cutedsl_megamoe/quant_def.py @@ -56,10 +56,7 @@ def umma_kind(self) -> str: pair -- falls back to mxf8f6f4. Keeping the two in sync matters because we build the tiled MMA through that helper but emit the instruction ourselves. """ - both_fp4 = ( - self.weight_dtype is cutlass.Float4E2M1FN - and self.activation_dtype is cutlass.Float4E2M1FN - ) + both_fp4 = self.weight_dtype is cutlass.Float4E2M1FN and self.activation_dtype is cutlass.Float4E2M1FN if not both_fp4: return "mxf8f6f4" return "mxf4nvf4" if self.sf_vec_size == 16 else "mxf4" @@ -105,10 +102,7 @@ def needs_unpack_tma(self, architecture: str) -> bool: normalized_architecture = normalized_architecture[2:] if normalized_architecture not in ("100", "103", "107"): raise ValueError(f"Unsupported TCGen05 architecture {architecture!r}.") - return ( - normalized_architecture in ("100", "103") - and self.weight_dtype.width != self.activation_dtype.width - ) + return normalized_architecture in ("100", "103") and self.weight_dtype.width != self.activation_dtype.width @property def uses_global_scale(self) -> bool: @@ -145,23 +139,15 @@ def _instruction_descriptor_format_code(umma_kind: str, dtype: Type[cutlass.Nume try: return codes[dtype] except KeyError as error: - raise ValueError( - f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}." - ) from error + raise ValueError(f"{dtype} has no instruction-descriptor format code under kind::{umma_kind}.") from error @dataclasses.dataclass(frozen=True) class CombineFormat: """Data and scale representation of one cross-rank FC2 payload.""" - _act_by_tag: ClassVar[Dict[str, type]] = { - "e2m1": cutlass.Float4E2M1FN, - "e4m3": cutlass.Float8E4M3FN, - } - _scale_by_tag: ClassVar[Dict[str, type]] = { - "bf16": cutlass.BFloat16, - "e8m0": cutlass.Float8E8M0FNU, - } + _act_by_tag: ClassVar[Dict[str, type]] = {"e2m1": cutlass.Float4E2M1FN, "e4m3": cutlass.Float8E4M3FN} + _scale_by_tag: ClassVar[Dict[str, type]] = {"bf16": cutlass.BFloat16, "e8m0": cutlass.Float8E8M0FNU} _rejection_reason: ClassVar[Dict[str, str]] = { "32e5m2xe8m0": ( "e5m2 costs the same 8.25 bits per element as e4m3 and trades a mantissa bit (6 dB of " @@ -200,9 +186,7 @@ def name(self) -> str: if not self.is_quantized: return "bf16" act_tag = next(tag for tag, dtype in self._act_by_tag.items() if dtype is self.act_dtype) - scale_tag = next( - tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype - ) + scale_tag = next(tag for tag, dtype in self._scale_by_tag.items() if dtype is self.scale_dtype) return f"{self.scale_block}{act_tag}x{scale_tag}" def __str__(self) -> str: @@ -217,9 +201,7 @@ def parse(cls, text: str) -> "CombineFormat": } token = text.strip().lower() if token in cls._rejection_reason: - raise ValueError( - f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}" - ) + raise ValueError(f"Combine format {token!r} is deliberately unsupported: {cls._rejection_reason[token]}") if token not in specs: raise ValueError(f"Invalid combine format {text!r}; expected one of {tuple(specs)}.") act_dtype, scale_dtype, scale_block = specs[token] @@ -252,15 +234,12 @@ def _verify_pinned_hardware_encoding() -> None: kind.scale_format_code, ) if derived != expected: - raise AssertionError( - f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}." - ) + raise AssertionError(f"QuantKind.{kind.name} derives {derived}, pinned encoding is {expected}.") expected_instruction_k_2x = expected[2] * 2 instruction_k_2x = kind.instruction_k("2x") if instruction_k_2x != expected_instruction_k_2x: raise AssertionError( - f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, " - f"expected {expected_instruction_k_2x}." + f"QuantKind.{kind.name} derives 2x instruction K {instruction_k_2x}, expected {expected_instruction_k_2x}." ) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py b/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py new file mode 100644 index 000000000000..abcc241122b6 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_kernels.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MegaMoE KernelClass entry points from the vendored ``cutedsl_megamoe`` package.""" + +from .cutedsl_megamoe.kernel_src.blackwell.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel as BlackwellInferenceMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.local_mega.block_scaled_swap_ab_local_mega_moe_kernel import ( + BlockScaledSwapAbLocalMegaMoeKernel as RubinInferenceLocalMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel import ( + BlockScaledSwapAbMegaMoeKernel as RubinInferenceMegaMoE, +) +from .cutedsl_megamoe.kernel_src.rubin.inference.mega.block_scaled_swap_ab_mega_moe_kernel_gen_specialized import ( + BlockScaledSwapAbGenphaseMoeKernel as RubinInferenceGenphaseMegaMoE, +) + +__all__ = [ + "BlackwellInferenceMegaMoE", + "RubinInferenceGenphaseMegaMoE", + "RubinInferenceLocalMegaMoE", + "RubinInferenceMegaMoE", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py new file mode 100644 index 000000000000..efa1c69eff92 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/__init__.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MegaMoE physical-slot scheduler public API.""" + +from .cuda_scheduler import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + CudaSchedulerOutputs, + recommend_cuda_scheduler_ctas, +) +from .sami import ( + BoundHierarchicalLiveWeightArena, + HierarchicalLiveWeightArena, + HierarchicalLiveWeightArenaProvider, + HierarchicalSamiWeightBroadcast, + BoundLiveWeightArena, + BundleLayout, + Comm, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + LiveWeightArenaProvider, + hierarchy_group_sizes, +) + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", + "BoundLiveWeightArena", + "BundleLayout", + "Comm", + "LivePlaneView", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py new file mode 100644 index 000000000000..90410ba4fbc1 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/__init__.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure-CUDA physical-slot scheduler. + +HALO-M hierarchical placement followed by HALO-Q quota repair, behind one +four-tensor placement ABI. Copy routing belongs to the in-switch copy layer. The CPU oracle that states the same placement lives in +``research/halo_q_oracle.py`` and is not imported from here. +""" + +from .runtime import ( + CudaPhysicalSlotScheduler, + CudaSchedulerConfig, + CudaSchedulerOutputs, + recommend_cuda_scheduler_ctas, +) + +__all__ = [ + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py new file mode 100644 index 000000000000..7074513cb639 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/multirank_dist.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA fabric allocation shared across expert-parallel processes.""" + +from __future__ import annotations + +import torch +from cuda.bindings import driver as cuda + + +def _check(result: object) -> object: + if isinstance(result, tuple): + error, *values = result + else: + error, values = result, [] + if int(error) != int(cuda.CUresult.CUDA_SUCCESS): + _, name = cuda.cuGetErrorName(error) + _, description = cuda.cuGetErrorString(error) + if isinstance(name, bytes): + name = name.decode() + if isinstance(description, bytes): + description = description.decode() + raise RuntimeError(f"CUDA driver call failed: {name}: {description}") + if not values: + return None + return values[0] if len(values) == 1 else tuple(values) + + +class FabricSymmetricBuffer: + """One CUDA-fabric allocation mapped read/write into every EP process.""" + + def __init__(self, nbytes: int, device: int): + self.device = device + properties = cuda.CUmemAllocationProp() + properties.type = cuda.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + properties.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + properties.location.id = device + properties.requestedHandleTypes = ( + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ) + self._properties = properties + self.granularity = int( + _check( + cuda.cuMemGetAllocationGranularity( + properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ) + ) + ) + self.size = (nbytes + self.granularity - 1) // self.granularity + self.size *= self.granularity + self._mappings: list[int] = [] + self._reservations: list[int] = [] + self._imported: list[object] = [] + self._closed = False + self._handle = None + self.ptr = 0 + self.shareable = b"" + try: + self._handle = _check(cuda.cuMemCreate(self.size, properties, 0)) + self.ptr = self._map(self._handle) + _check(cuda.cuMemsetD8(cuda.CUdeviceptr(self.ptr), 0, self.size)) + exported = _check( + cuda.cuMemExportToShareableHandle( + self._handle, + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + 0, + ) + ) + self.shareable = bytes(exported.data) + except BaseException: + self.close() + raise + + def _map(self, handle: object) -> int: + address = int( + _check(cuda.cuMemAddressReserve(self.size, self.granularity, 0, 0)) + ) + self._reservations.append(address) + _check(cuda.cuMemMap(address, self.size, 0, handle, 0)) + self._mappings.append(address) + access = cuda.CUmemAccessDesc() + access.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + access.location.id = self.device + access.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _check(cuda.cuMemSetAccess(address, self.size, [access], 1)) + return int(address) + + def import_peer(self, shareable: bytes) -> int: + handle = _check( + cuda.cuMemImportFromShareableHandle( + shareable, + cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + ) + ) + self._imported.append(handle) + return self._map(handle) + + def close(self) -> None: + """Release every mapping and fabric handle after device work is drained.""" + if self._closed: + return + with torch.cuda.device(self.device): + while self._mappings: + address = self._mappings[-1] + _check(cuda.cuMemUnmap(address, self.size)) + self._mappings.pop() + while self._reservations: + address = self._reservations[-1] + _check(cuda.cuMemAddressFree(address, self.size)) + self._reservations.pop() + while self._imported: + handle = self._imported[-1] + _check(cuda.cuMemRelease(handle)) + self._imported.pop() + if self._handle is not None: + _check(cuda.cuMemRelease(self._handle)) + self._handle = None + self.ptr = 0 + self.shareable = b"" + self._closed = True + + +__all__ = ["FabricSymmetricBuffer"] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py new file mode 100644 index 000000000000..671fece56427 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host runtime for the HALO-Q scheduler, launched with its GPU-direct TMA copy as one kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass +import threading + +import torch + +from ..geometry import MAX_EP, MIN_EP, hierarchy_group_sizes + + +MAX_EXPERTS = 384 +MAX_CTAS = 128 +MAX_BROADCASTS = MAX_EXPERTS +# Nominal clock64 rate that converts a timeout in seconds to SM cycles, as the +# MoE all-to-all kernels assume. +ASSUMED_CLOCK_HZ = 2_000_000_000 +# Default bound, in SM clock cycles, of the scheduler's grid barrier and +# cross-rank waits; a timeout records a status flag. Each launch takes its bound +# as an argument, so callers can allow more peer skew while ranks warm up. The +# default, 300 s at ASSUMED_CLOCK_HZ, matches NVLinkOneSided's serving timeout, +# so a transient stall on one rank does not trap every EP rank; +# TRTLLM_MOE_COMM_TIMEOUT_SEC can set a shorter bound. +SPIN_CYCLES = 600_000_000_000 + + +def spin_cycles_for_timeout(seconds: int | None) -> int: + """Spin bound for a timeout in seconds; ``None`` selects ``SPIN_CYCLES``.""" + if seconds is None: + return SPIN_CYCLES + if type(seconds) is not int or seconds <= 0: + raise ValueError("scheduler timeout must be a positive exact int of seconds") + return seconds * ASSUMED_CLOCK_HZ + + +def recommend_cuda_scheduler_ctas( + ep_size: int, + max_tokens_per_rank: int, + topk: int = 6, + threads: int = 512, +) -> int: + """Return the bounded pure-CUDA CTA policy without importing CuTe.""" + + for name, value in { + "ep_size": ep_size, + "max_tokens_per_rank": max_tokens_per_rank, + "topk": topk, + "threads": threads, + }.items(): + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive exact int") + if threads != 512: + raise ValueError("the pure-CUDA scheduler currently requires threads=512") + routes = max_tokens_per_rank * topk + if ep_size == 8: + # Bound scheduler occupancy so independent same-stream work retains + # launch capacity. + workers = min(31, max(1, (routes + 3 * threads - 1) // (3 * threads))) + return 1 + workers + target_rounds = max(3, (96 + (ep_size + 6) // 2) // (ep_size + 6)) + workers = max( + 1, (max_tokens_per_rank * topk) // (threads * target_rounds) + ) + workers = min(workers, MAX_CTAS - 1) + rounds = (routes + workers * threads - 1) // (workers * threads) + workers = (routes + rounds * threads - 1) // (rounds * threads) + return 1 + workers + + +def symmetric_buffer_ints(ep_size: int, expert_count: int) -> int: + # Two reset-free arrival generations plus a local epoch and two count slabs. + return 64 + 2 * ep_size * expert_count + + +@dataclass(frozen=True) +class CudaSchedulerConfig: + ep_size: int + logical_expert_count: int + extra_slots_per_rank: int + max_tokens_per_rank: int + topk: int = 6 + local_rank: int = 0 + threads: int = 512 + ctas: int | None = None + + def __post_init__(self) -> None: + if self.ctas is None: + object.__setattr__( + self, + "ctas", + recommend_cuda_scheduler_ctas( + self.ep_size, + self.max_tokens_per_rank, + self.topk, + self.threads, + ), + ) + self.validate() + + @property + def home_experts_per_rank(self) -> int: + return self.logical_expert_count // self.ep_size + + @property + def local_physical_slot_count(self) -> int: + return self.home_experts_per_rank + self.extra_slots_per_rank + + @property + def route_count(self) -> int: + return self.max_tokens_per_rank * self.topk + + @property + def world_size(self) -> int: + return self.ep_size + + @property + def global_expert_count(self) -> int: + return self.logical_expert_count + + @property + def helper_count(self) -> int: + return self.extra_slots_per_rank + + @property + def group_sizes(self) -> tuple[int, ...]: + return hierarchy_group_sizes(self.ep_size) + + def validate(self) -> None: + values = { + "ep_size": self.ep_size, + "logical_expert_count": self.logical_expert_count, + "extra_slots_per_rank": self.extra_slots_per_rank, + "max_tokens_per_rank": self.max_tokens_per_rank, + "topk": self.topk, + "local_rank": self.local_rank, + "threads": self.threads, + "ctas": self.ctas, + } + if any(type(value) is not int for value in values.values()): + raise ValueError("all CUDA scheduler dimensions must be exact ints") + if not MIN_EP <= self.ep_size <= MAX_EP or self.ep_size % 2: + raise ValueError( + f"pure-CUDA scheduler supports even EP in [{MIN_EP},{MAX_EP}]") + if not 0 <= self.local_rank < self.ep_size: + raise ValueError("local_rank must be in [0,ep_size)") + if not 1 <= self.logical_expert_count <= MAX_EXPERTS: + raise ValueError(f"logical_expert_count must be in [1,{MAX_EXPERTS}]") + if self.logical_expert_count % self.ep_size: + raise ValueError("logical_expert_count must be EP-divisible") + if self.extra_slots_per_rank <= 0: + raise ValueError("extra_slots_per_rank must be positive") + if self.ep_size * self.local_physical_slot_count > 2**31 - 1: + raise ValueError("global physical slot addresses must fit int32") + if self.max_tokens_per_rank <= 0 or self.topk <= 0: + raise ValueError("token and topk dimensions must be positive") + if self.ep_size * self.route_count > 2**31 - 1: + raise ValueError("global route count must fit int32 histogram totals") + if self.threads != 512: + raise ValueError("the pure-CUDA scheduler currently requires threads=512") + if not 1 <= self.ctas <= MAX_CTAS: + raise ValueError(f"ctas must be in [1,{MAX_CTAS}]") + + +class CudaSchedulerOutputs: + """The exact four-tensor hierarchical scheduler ABI.""" + + __slots__ = ( + "physical_slot_ids", + "hot_expert_ids", + "hot_expert_group_level", + "hot_expert_source_ranks", + ) + + def __init__(self, config: CudaSchedulerConfig, device: torch.device): + self.physical_slot_ids = torch.empty( + (config.max_tokens_per_rank, config.topk), + dtype=torch.int32, + device=device, + ) + shape = (config.extra_slots_per_rank,) + self.hot_expert_ids = torch.empty(shape, dtype=torch.int32, device=device) + self.hot_expert_group_level = torch.empty( + shape, dtype=torch.int32, device=device + ) + self.hot_expert_source_ranks = torch.empty( + shape, dtype=torch.int32, device=device + ) + + +def _stream_handle(stream: object) -> int: + if isinstance(stream, bool): + raise TypeError("stream must be a CUDA stream handle") + handle = getattr(stream, "cuda_stream", stream) + try: + value = int(handle) + except (TypeError, ValueError) as error: + raise TypeError("stream must be a CUDA stream handle") from error + if value < 0: + raise ValueError("stream handle cannot be negative") + return value + + +class CudaPhysicalSlotScheduler: + """Preallocated HALO-Q outputs and workspace for the fused scheduler + copy kernel. + + The GPU-direct copy endpoint binds it once (``bind_gpu_direct``), binds + the scheduler's arguments from ``fused_launch_args`` into its native + context, and then launches it with the copy + (``moe_rebalance_tma_submit_fused``), passing the routes, the route + target and the spin bound. + """ + + def __init__(self, config: CudaSchedulerConfig, device: str | torch.device = "cuda"): + config.validate() + scheduler_device = torch.device(device) + if scheduler_device.type != "cuda": + raise ValueError("CudaPhysicalSlotScheduler requires a CUDA device") + if scheduler_device.index is None: + scheduler_device = torch.device("cuda", torch.cuda.current_device()) + properties = torch.cuda.get_device_properties(scheduler_device) + if config.ctas > properties.multi_processor_count: + raise ValueError("ctas cannot exceed the CUDA device SM count") + + self.cfg = config + self.device = scheduler_device + self.outputs = CudaSchedulerOutputs(config, scheduler_device) + self.out = self.outputs + self.status = torch.zeros(8, dtype=torch.int32, device=scheduler_device) + # These tensors and the config values are fixed for the endpoint's + # lifetime, so the copy binds them once; see fused_launch_args(). + self.partial = torch.zeros( + config.ctas * config.logical_expert_count, + dtype=torch.int32, + device=scheduler_device, + ) + warps = config.threads // 32 + bins = config.logical_expert_count + 1 + route_aux_ints = config.ctas * warps * bins + if config.ctas > 1: + route_aux_ints = max( + route_aux_ints, + 2 + (config.ctas - 1) * bins + config.route_count, + ) + self.route_aux = torch.zeros( + route_aux_ints, + dtype=torch.int32, + device=scheduler_device, + ) + # The main barrier uses [arrival, epoch]; word 2 publishes route readiness. + self.grid_sync = torch.zeros(4, dtype=torch.int32, device=scheduler_device) + # Each expert is broadcast at most once. The private plan's fixed + # field stride is bounded by E, independently of requested slot capacity. + self.max_broadcasts = MAX_BROADCASTS + self.plan_workspace = torch.full( + (2 + 6 * self.max_broadcasts,), + -1, + dtype=torch.int32, + device=scheduler_device, + ) + self.route_prefix = torch.zeros( + self.max_broadcasts * (config.ep_size + 1), + dtype=torch.int32, + device=scheduler_device, + ) + self.sym = torch.zeros( + symmetric_buffer_ints( + config.ep_size, config.logical_expert_count + ), + dtype=torch.int32, + device=scheduler_device, + ) + self.peer_base = torch.full( + (config.ep_size,), + int(self.sym.data_ptr()), + dtype=torch.int64, + device=scheduler_device, + ) + self._connected = config.ep_size == 1 + self._lock = threading.Lock() + # The bound copy's static stream: every fused launch shares the + # workspace and reusable barrier epochs, so all of them use it. + self._bound_stream: int | None = None + + def connect_peer_bases(self, peer_bases: object) -> None: + if self._connected: + raise RuntimeError("peer base binding is single-assignment") + + bases = tuple(peer_bases) # type: ignore[arg-type] + if len(bases) != self.cfg.ep_size or any( + type(pointer) is not int or pointer <= 0 for pointer in bases + ): + raise ValueError("peer_bases must contain EP positive exact addresses") + self.peer_base.copy_( + torch.tensor(bases, dtype=torch.int64, device=self.device) + ) + torch.cuda.synchronize(self.device) + self._connected = True + + def bind_gpu_direct(self, stream: object) -> None: + """Bind the GPU-direct copy that launches this scheduler with its copy. + + The copy reads the outputs and plan workspace in place, in the same + kernel, and every fused launch uses this static stream. No additional + plan allocation or GPU/host transfer is performed. + """ + with self._lock: + if self._bound_stream is not None: + raise RuntimeError("GPU-direct binding is single-assignment") + if not self._connected: + raise RuntimeError("EP peer bases are not connected") + self._bound_stream = _stream_handle(stream) + + def fused_launch_args(self) -> tuple: + """Scheduler arguments that every fused launch of the bound copy + shares, ``out_slots`` through ``threads`` of + moe_rebalance_tma_bind_fused. + + A launch adds its route target, its spin bound and its routes: a + caller-owned CUDA int32[T,K] tensor, read without input staging. T may + be zero and may differ across ranks, up to max_tokens_per_rank. The + kernel never reads outside those T rows and writes -1 to the unused + output tail. All ranks, including empty + ranks, must launch each generation in the same collective order. The + caller keeps the routes alive and unmodified until the scheduler has + read them; the kernel reports invalid expert values through + check_status(). + """ + if self._bound_stream is None: + raise RuntimeError("a fused launch requires a bound GPU-direct copy") + return ( + self.outputs.physical_slot_ids, + self.outputs.hot_expert_ids, + self.outputs.hot_expert_group_level, + self.outputs.hot_expert_source_ranks, + self.peer_base, + self.status, + self.partial, + self.route_aux, + self.grid_sync, + self.plan_workspace, + self.route_prefix, + self.cfg.ep_size, + self.cfg.logical_expert_count, + self.cfg.extra_slots_per_rank, + self.cfg.local_rank, + self.cfg.topk, + self.cfg.route_count, + self.cfg.ctas, + self.cfg.threads, + ) + + def check_status(self) -> None: + """Raise if any generation faulted since the last `reset_status()`. + + Each flag holds the unsigned modular epoch that wrote it, not a bare 1. + Zero means never set; an error at wrapped epoch zero stores UINT32_MAX. + That lets the per-generation status clear stay out of the submission path. + """ + torch.cuda.synchronize(self.device) + flags = [ + flag & 0xFFFFFFFF for flag in self.status[:6].tolist() + ] # one D2H, not six + labels = ( + "cross-rank histogram rendezvous timeout", + "CTA grid barrier timeout", + "plan publication timeout", + "invalid logical expert ID", + "HALO helper coloring failure", + "HALO quota/host invariant failure", + ) + failures = [ + f"{label} (generation {flag})" + for label, flag in zip(labels, flags) + if flag + ] + if failures: + raise RuntimeError("; ".join(failures)) + + def reset_status(self) -> None: + """Acknowledge and clear faults. Not needed per generation.""" + self.status.zero_() + + +__all__ = [ + "ASSUMED_CLOCK_HZ", + "SPIN_CYCLES", + "CudaPhysicalSlotScheduler", + "CudaSchedulerConfig", + "CudaSchedulerOutputs", + "recommend_cuda_scheduler_ctas", + "spin_cycles_for_timeout", + "symmetric_buffer_ints", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py new file mode 100644 index 000000000000..f1fb15739e9d --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/geometry.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The R4 rank hierarchy shared by the CUDA scheduler and the SAMI copy.""" + +from __future__ import annotations + +__all__ = ["MAX_EP", "MIN_EP", "MIN_GROUP_SIZE", "hierarchy_group_sizes"] + +MIN_EP = 2 +MAX_EP = 32 +MIN_GROUP_SIZE = 4 + + +def hierarchy_group_sizes(world: int) -> tuple[int, ...]: + """Return HALO-M's R4 aligned hierarchy, coarsest to finest. + + Level 0 is the whole EP world. Later levels halve the group size only + while the world divides exactly and every group still holds at least + ``MIN_GROUP_SIZE`` ranks. ``halo_q_scheduler.cu::hierarchy_depth`` is the + device-side statement of the same ladder; the two must agree for every + world this package accepts, so there is deliberately one host definition + rather than one per subpackage. + """ + + if (type(world) is not int or world < MIN_EP or world > MAX_EP + or world % 2): + raise ValueError( + f"hierarchical SAMI supports even EP in [{MIN_EP},{MAX_EP}]") + result = [world] + groups = 2 + while world % groups == 0 and world // groups >= MIN_GROUP_SIZE: + result.append(world // groups) + groups *= 2 + return tuple(result) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py new file mode 100644 index 000000000000..ae2b7d7c8f7b --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/__init__.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Production SAMI resources for an external single-bank live arena.""" + +from .arena import ( + BoundLiveWeightArena, + LivePlaneView, + LiveTerminalView, + LiveWeightArena, + LiveWeightArenaProvider, +) +from ._util import Comm +from .geometry import BundleLayout, Plane, PlaneViewLayout +from .hierarchical import ( + BoundHierarchicalLiveWeightArena, + HierarchicalLiveWeightArena, + HierarchicalLiveWeightArenaProvider, + HierarchicalSamiWeightBroadcast, + hierarchy_group_sizes, +) + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "BoundLiveWeightArena", + "BundleLayout", + "Comm", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "LivePlaneView", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", + "Plane", + "PlaneViewLayout", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py new file mode 100644 index 000000000000..a97ad6b6692a --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/_util.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The driver-call wrapper and setup collectives production SAMI needs. + +Two small helpers rather than two small modules. They are unrelated to each +other, but both are leaf utilities with a single consumer set, and neither is +part of the package's public API. + +``check_cuda`` deliberately does not touch torch. That separation is not +load-bearing for import cost -- ``sami/__init__.py`` already imports ``Comm`` +and ``hierarchical``, so torch is on the path for any ``sami.*`` import -- but +it does keep the driver wrapper usable from a torch-free context if the +package layout ever changes. +""" + +from __future__ import annotations + +from cuda.bindings import driver as cuda +import torch.distributed as dist + +__all__ = ["Comm", "check_cuda"] + + +def check_cuda(result: object, operation: str = "CUDA") -> object: + """Unpack a cuda-python driver result and raise on failure.""" + if isinstance(result, tuple): + error, *values = result + else: + error, values = result, [] + if int(error) != int(cuda.CUresult.CUDA_SUCCESS): + _, name = cuda.cuGetErrorName(error) + _, description = cuda.cuGetErrorString(error) + if isinstance(name, bytes): + name = name.decode() + if isinstance(description, bytes): + description = description.decode() + raise RuntimeError( + f"{operation} failed: {name}: {description} (code {int(error)})" + ) + if not values: + return None + return values[0] if len(values) == 1 else tuple(values) + + +class Comm: + """The rank identity and setup collectives required by SAMI.""" + + def __init__(self) -> None: + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before SAMI") + self.rank = dist.get_rank() + self.world = dist.get_world_size() + + def barrier(self) -> None: + dist.barrier() + + def bcast(self, payload: object, root: int) -> object: + values = [payload] + dist.broadcast_object_list(values, src=root) + return values[0] + + def allgather(self, value: object) -> list[object]: + values: list[object] = [None] * self.world + dist.all_gather_object(values, value) + return values diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py new file mode 100644 index 000000000000..20aef7bbaba3 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/arena.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""External single-bank MegaMoE weight arena consumed directly by SAMI. + +The production arena is plane-major. Each of the seven tensors has local +physical slots ``[0, H)`` for canonical experts and ``[H, H + S)`` for live +helpers. A plane exposes every rank's unicast mapping plus the multicast +mapping of the same backing allocation. This is the shape naturally exposed +by TEKit ``McastGPUBuffer.get_uc_buffer``/``get_mc_buffer``; callers without +that class can provide equivalent typed CUDA tensor views. + +This module performs no allocation and deliberately does not import TEKit. +It validates geometry, pointer and object identity once at bind time and +resolves the pointer tables the copy backend uses from then on; identity is +not re-checked per submit. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Protocol, runtime_checkable + +from .geometry import BundleLayout + + +def _exact_positive_int(name: str, value: object) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive exact int") + return value + + +def _pointer(value: object, *, name: str) -> int: + accessor = getattr(value, "data_ptr", None) + pointer = accessor() if callable(accessor) else accessor + if type(pointer) is not int or pointer <= 0: + raise TypeError(f"{name} must expose a positive integer data_ptr") + return pointer + + +def _shape(value: object, *, name: str) -> tuple[int, ...]: + try: + result = tuple(int(extent) for extent in value.shape) # type: ignore[attr-defined] + except (AttributeError, TypeError, ValueError) as exc: + raise TypeError(f"{name} must expose an integer shape") from exc + if not result or any(extent <= 0 for extent in result): + raise ValueError(f"{name} must have a non-empty positive shape") + return result + + +def _stride(value: object, *, name: str) -> tuple[int, ...]: + accessor = getattr(value, "stride", None) + raw = accessor() if callable(accessor) else accessor + try: + result = tuple(int(step) for step in raw) + except (TypeError, ValueError) as exc: + raise TypeError(f"{name} must expose an integer stride") from exc + if any(step <= 0 for step in result): + raise ValueError(f"{name} must have positive strides") + return result + + +def _dtype(value: object) -> str: + return str(getattr(value, "dtype", None)).lower().removeprefix("torch.") + + +def _element_size(value: object, *, name: str) -> int: + accessor = getattr(value, "element_size", None) + raw = accessor() if callable(accessor) else accessor + if type(raw) is not int or raw <= 0: + raise TypeError(f"{name} must expose a positive integer element_size") + return raw + + +def _device_index(value: object, *, name: str) -> int: + if getattr(value, "is_cuda", True) is not True: + raise ValueError(f"{name} must reside on CUDA memory") + device = getattr(value, "device", None) + text = str(device).lower() + if not text.startswith("cuda"): + raise ValueError(f"{name} must reside on CUDA memory") + index = getattr(device, "index", None) + if type(index) is not int: + if ":" not in text: + raise ValueError(f"{name} must expose an explicit CUDA device ordinal") + try: + index = int(text.split(":", 1)[1]) + except ValueError as exc: + raise ValueError(f"{name} exposes an invalid CUDA device") from exc + if index < 0: + raise ValueError(f"{name} exposes an invalid CUDA device") + return index + + +@dataclass(frozen=True) +class LivePlaneView: + """One live plane's rank-ordered UC views and its MC alias. + + ``aliases_same_backing`` is an explicit provider attestation. Pointer + arithmetic alone cannot prove a CUDA UC/MC alias relationship, so an + un-attested collection of unrelated tensors is rejected. + """ + + name: str + uc_views: tuple[object, ...] = field(repr=False, compare=False) + mc_view: object = field(repr=False, compare=False) + aliases_same_backing: bool = False + backing_owner: object | None = field(default=None, repr=False, compare=False) + + +@dataclass(frozen=True) +class LiveTerminalView: + """Local ``uint64[EP]`` terminals and the matching MC alias.""" + + local_view: object = field(repr=False, compare=False) + mc_view: object = field(repr=False, compare=False) + aliases_same_backing: bool = False + backing_owner: object | None = field(default=None, repr=False, compare=False) + + +@dataclass(frozen=True) +class _TensorIdentity: + object_id: int + pointer: int + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + element_size: int + device: int + name: str + + @classmethod + def capture(cls, value: object, *, name: str) -> "_TensorIdentity": + shape = _shape(value, name=name) + stride = _stride(value, name=name) + if len(shape) != len(stride): + raise ValueError(f"{name} shape/stride rank mismatch") + return cls( + object_id=id(value), + pointer=_pointer(value, name=name), + shape=shape, + stride=stride, + dtype=_dtype(value), + element_size=_element_size(value, name=name), + device=_device_index(value, name=name), + name=name, + ) + + @property + def metadata(self) -> tuple: + """The layout a view must match; deliberately excludes the pointer.""" + return (self.shape, self.stride, self.dtype, self.element_size, + self.device) + + +@dataclass(frozen=True) +class BoundLiveWeightArena: + """Validated pointer tables and strong references for one SAMI rank.""" + + arena: "LiveWeightArena" = field(repr=False, compare=False) + world: int + rank: int + home_count: int + helper_count: int + device: int + plane_names: tuple[str, ...] + plane_bytes: tuple[int, ...] + local_plane_views: tuple[object, ...] = field(repr=False, compare=False) + owner_source_ptrs: tuple[int, ...] + scatter_source_ptrs: tuple[int, ...] + destination_ptrs: tuple[int, ...] + terminal_local_ptr: int + terminal_mc_ptr: int + + @property + def local_slot_count(self) -> int: + return self.home_count + self.helper_count + + +@dataclass(frozen=True) +class LiveWeightArena: + """Seven plane-major live weight views plus one terminal allocation.""" + + planes: tuple[LivePlaneView, ...] + terminals: LiveTerminalView + + def bind( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + ) -> BoundLiveWeightArena: + world = _exact_positive_int("world", world) + home_count = _exact_positive_int("home_count", home_count) + helper_count = _exact_positive_int("helper_count", helper_count) + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + + expected_names = tuple(plane.name for plane in bundle.planes) + observed_names = tuple(plane.name for plane in self.planes) + if observed_names != expected_names: + raise ValueError( + f"live arena plane order mismatch: {observed_names} != {expected_names}" + ) + if len(self.planes) != 7: + raise ValueError("live arena must contain exactly seven planes") + + slot_count = home_count + helper_count + local_planes: list[object] = [] + plane_bytes: list[int] = [] + owner_sources: list[int] = [] + scatter_sources: list[int] = [] + destinations: list[int] = [] + uc_ranges: list[list[tuple[int, int, str]]] = [ + [] for _ in range(world) + ] + mc_ranges: list[tuple[int, int, str]] = [] + + plane_identities: list[tuple[_TensorIdentity, ...]] = [] + mc_identities: list[_TensorIdentity] = [] + for plane_index, (view, layout) in enumerate(zip(self.planes, bundle.planes)): + if view.aliases_same_backing is not True: + raise ValueError(f"{view.name} UC/MC alias is not provider-attested") + if len(view.uc_views) != world: + raise ValueError(f"{view.name} must expose exactly world UC views") + + captured = tuple( + _TensorIdentity.capture(value, name=f"{view.name}.uc[{peer}]") + for peer, value in enumerate(view.uc_views) + ) + mc = _TensorIdentity.capture(view.mc_view, name=f"{view.name}.mc") + reference = captured[rank] + ( + expected_shape, + expected_stride, + expected_dtype, + expected_element_size, + expected_pointer_alignment, + ) = layout.kernel_metadata(slot_count) + expected_metadata = ( + expected_shape, + expected_stride, + expected_dtype, + expected_element_size, + device, + ) + for identity in (*captured, mc): + if identity.metadata != expected_metadata: + raise ValueError( + f"{view.name} must have exact MegaMoE " + f"shape={expected_shape}, stride={expected_stride}, " + f"dtype={expected_dtype} on cuda:{device}" + ) + if identity.pointer % expected_pointer_alignment: + raise ValueError( + f"{view.name} pointer must be " + f"{expected_pointer_alignment}-byte aligned" + ) + if len({identity.pointer for identity in captured}) != world: + raise ValueError(f"{view.name} UC rank mappings must have distinct pointers") + if mc.pointer in {identity.pointer for identity in captured}: + raise ValueError(f"{view.name} MC and UC mappings must be distinct") + if layout.nbytes % reference.element_size: + raise ValueError(f"{view.name} bytes/expert are not dtype aligned") + + expert_elements = layout.nbytes // reference.element_size + tail_span = 1 + sum( + (extent - 1) * stride + for extent, stride in zip(reference.shape[1:], reference.stride[1:]) + ) + if reference.stride[0] != expert_elements or tail_span != expert_elements: + raise ValueError( + f"{view.name} must be dense per expert with byte stride {layout.nbytes}" + ) + + plane_identities.append(captured) + mc_identities.append(mc) + local_planes.append(view.uc_views[rank]) + plane_bytes.append(layout.nbytes) + for peer, identity in enumerate(captured): + uc_ranges[peer].append( + ( + identity.pointer, + identity.pointer + slot_count * layout.nbytes, + view.name, + ) + ) + mc_ranges.append( + (mc.pointer, mc.pointer + slot_count * layout.nbytes, view.name) + ) + + weight_ranges = [ + (begin, end, f"{name}.uc[{peer}]") + for peer, peer_ranges in enumerate(uc_ranges) + for begin, end, name in peer_ranges + ] + weight_ranges.extend( + (begin, end, f"{name}.mc") for begin, end, name in mc_ranges + ) + self._reject_overlaps(weight_ranges) + + for local_expert in range(home_count): + for plane_index, layout in enumerate(bundle.planes): + owner_sources.append( + plane_identities[plane_index][rank].pointer + + local_expert * layout.nbytes + ) + for expert in range(world * home_count): + owner, local_expert = divmod(expert, home_count) + for plane_index, layout in enumerate(bundle.planes): + scatter_sources.append( + plane_identities[plane_index][owner].pointer + + local_expert * layout.nbytes + ) + for helper_index in range(helper_count): + physical_slot = home_count + helper_index + for plane_index, layout in enumerate(bundle.planes): + destinations.append( + mc_identities[plane_index].pointer + + physical_slot * layout.nbytes + ) + + if self.terminals.aliases_same_backing is not True: + raise ValueError("terminal UC/MC alias is not provider-attested") + terminal_local = _TensorIdentity.capture( + self.terminals.local_view, name="terminal.local" + ) + terminal_mc = _TensorIdentity.capture(self.terminals.mc_view, name="terminal.mc") + for terminal in (terminal_local, terminal_mc): + if terminal.shape != (world,) or terminal.stride != (1,): + raise ValueError("terminal views must have contiguous shape [EP]") + if terminal.dtype != "uint64" or terminal.element_size != 8: + raise TypeError("terminal views must have dtype uint64") + if terminal.device != device: + raise ValueError("terminal views are on the wrong CUDA device") + if terminal.pointer % 8: + raise ValueError("terminal pointers must be 8-byte aligned") + self._reject_terminal_overlap(terminal_local, weight_ranges, "local") + self._reject_terminal_overlap(terminal_mc, weight_ranges, "multicast") + terminal_bytes = world * 8 + if ( + terminal_local.pointer < terminal_mc.pointer + terminal_bytes + and terminal_mc.pointer < terminal_local.pointer + terminal_bytes + ): + raise ValueError("terminal local and multicast mappings overlap") + + return BoundLiveWeightArena( + arena=self, + world=world, + rank=rank, + home_count=home_count, + helper_count=helper_count, + device=device, + plane_names=expected_names, + plane_bytes=tuple(plane_bytes), + local_plane_views=tuple(local_planes), + owner_source_ptrs=tuple(owner_sources), + scatter_source_ptrs=tuple(scatter_sources), + destination_ptrs=tuple(destinations), + terminal_local_ptr=terminal_local.pointer, + terminal_mc_ptr=terminal_mc.pointer, + ) + + @staticmethod + def _reject_overlaps(ranges: list[tuple[int, int, str]]) -> None: + ordered = sorted(ranges) + for (_, previous_end, previous), (begin, _, current) in zip( + ordered, ordered[1:] + ): + if begin < previous_end: + raise ValueError( + f"live arena mappings overlap: {previous} and {current}" + ) + + @staticmethod + def _reject_terminal_overlap( + terminal: _TensorIdentity, + ranges: list[tuple[int, int, str]], + label: str, + ) -> None: + begin, end = terminal.pointer, terminal.pointer + terminal.shape[0] * 8 + for plane_begin, plane_end, name in ranges: + if begin < plane_end and plane_begin < end: + raise ValueError( + f"{label} terminal overlaps live weight plane {name}" + ) + + +@runtime_checkable +class LiveWeightArenaProvider(Protocol): + """Provider boundary implemented by TEKit or an equivalent allocator.""" + + def build_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + comm: object, + ) -> LiveWeightArena: + """Return one fully constructed and provider-attested live arena.""" + + +__all__ = [ + "BoundLiveWeightArena", + "LivePlaneView", + "LiveTerminalView", + "LiveWeightArena", + "LiveWeightArenaProvider", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py new file mode 100644 index 000000000000..a785bc2690ba --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/fabric.py @@ -0,0 +1,900 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CUDA Fabric and NVSwitch multicast memory. + +This module is in the runtime dependency closure. It is imported lazily by the +live-weight arena to avoid eager native initialization at module import time. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Callable + +from cuda.bindings import driver as cuda + +from ._util import check_cuda + + +FABRIC = cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _read_write_access(device: int) -> cuda.CUmemAccessDesc: + descriptor = cuda.CUmemAccessDesc() + descriptor.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + descriptor.location.id = int(device) + descriptor.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + return descriptor + + +def _allocation_properties(device: int) -> cuda.CUmemAllocationProp: + properties = cuda.CUmemAllocationProp() + properties.type = cuda.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + properties.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + properties.location.id = int(device) + properties.requestedHandleTypes = FABRIC + return properties + + +@dataclass +class FabricAllocation: + """One exported device allocation plus the peer mappings imported from it.""" + + device: int + size: int + granularity: int + handle: object + local_ptr: int + shareable: bytes + _imported_handles: list[object] = field(default_factory=list, repr=False) + + @classmethod + def create( + cls, nbytes: int, device: int, *, extra_alignment: int = 1 + ) -> "FabricAllocation": + if nbytes <= 0: + raise ValueError("allocation size must be positive") + properties = _allocation_properties(device) + granularity = int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + properties, + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ), + "cuMemGetAllocationGranularity", + ) + ) + alignment = math.lcm(granularity, max(1, int(extra_alignment))) + size = _align_up(nbytes, alignment) + handle = check_cuda(cuda.cuMemCreate(size, properties, 0), "cuMemCreate") + allocation = cls(device, size, alignment, handle, 0, b"") + allocation.local_ptr = allocation._map(handle) + check_cuda( + cuda.cuMemsetD8(cuda.CUdeviceptr(allocation.local_ptr), 0, size), + "cuMemsetD8", + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle", + ) + allocation.shareable = bytes(exported.data) + if len(allocation.shareable) != 64: + raise RuntimeError( + f"fabric handle has {len(allocation.shareable)} bytes; expected 64" + ) + return allocation + + def _map(self, handle: object) -> int: + address = check_cuda( + cuda.cuMemAddressReserve(self.size, self.granularity, 0, 0), + "cuMemAddressReserve", + ) + check_cuda(cuda.cuMemMap(address, self.size, 0, handle, 0), "cuMemMap") + check_cuda( + cuda.cuMemSetAccess( + address, + self.size, + [_read_write_access(self.device)], + 1, + ), + "cuMemSetAccess", + ) + return int(address) + + def import_peer(self, shareable: bytes) -> int: + if len(shareable) != 64: + raise ValueError("peer fabric handle must contain exactly 64 bytes") + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(shareable), FABRIC), + "cuMemImportFromShareableHandle", + ) + self._imported_handles.append(handle) + return self._map(handle) + + +@dataclass(frozen=True) +class MulticastGroup: + """All-member NVSwitch multicast aperture bound to each rank's backing.""" + + device: int + size: int + granularity: int + handle: object + mc_ptr: int + creator_rank: int + backing_owner: FabricAllocation = field(repr=False, compare=False) + + alias_set: "MulticastTeamPair | None" = field(default=None, repr=False, compare=False) + + @property + def aliases(self) -> tuple: + return self.alias_set.aliases if self.alias_set is not None else () + + @property + def closed(self) -> bool: + return self.alias_set.closed if self.alias_set is not None else False + + def metadata(self) -> dict: + if self.alias_set is not None: + return self.alias_set.metadata() + return {"device": self.device, "size": self.size, + "creator_rank": self.creator_rank, "mc_ptr": self.mc_ptr, + "alias_count": 1, "dual_creator": False, "closed": False} + + def close_collectively(self) -> bool: + """Close a managed default group on every EP rank via its owners list. + + Explicit-creator legacy groups retain their original lifetime contract. + """ + if self.alias_set is None: + raise RuntimeError("explicit-creator groups have no managed pair lifecycle") + success = self.alias_set.close_collectively() + if success: + object.__setattr__(self, "mc_ptr", 0) + object.__setattr__(self, "handle", None) + return success + + @staticmethod + def granularity_for(world: int, nbytes: int, device: int) -> tuple[int, int]: + properties = cuda.CUmulticastObjectProp() + properties.numDevices = world + properties.size = nbytes + properties.handleTypes = int(FABRIC) + properties.flags = 0 + multicast = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity", + ) + ) + allocation = int( + check_cuda( + cuda.cuMemGetAllocationGranularity( + _allocation_properties(device), + cuda.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED, + ), + "cuMemGetAllocationGranularity", + ) + ) + return multicast, allocation + + @classmethod + def create( + cls, + *, + world: int, + rank: int, + creator_rank: int, + device: int, + size: int, + backing: FabricAllocation, + broadcast: Callable[[object, int], object], + barrier: Callable[[], None], + ) -> "MulticastGroup": + properties = cuda.CUmulticastObjectProp() + properties.numDevices = world + properties.size = size + properties.handleTypes = int(FABRIC) + properties.flags = 0 + granularity = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity", + ) + ) + if size % granularity: + raise ValueError("multicast size is not aligned to its granularity") + if backing.size < size: + raise ValueError("backing allocation is smaller than multicast object") + + serialized = None + if rank == creator_rank: + handle = check_cuda( + cuda.cuMulticastCreate(properties), "cuMulticastCreate" + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle(multicast)", + ) + serialized = bytes(exported.data) + serialized = broadcast(serialized, creator_rank) + if rank != creator_rank: + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(serialized), FABRIC), + "cuMemImportFromShareableHandle(multicast)", + ) + + check_cuda(cuda.cuMulticastAddDevice(handle, device), "cuMulticastAddDevice") + barrier() + check_cuda( + cuda.cuMulticastBindMem_v2( + handle, device, 0, backing.handle, 0, size, 0 + ), + "cuMulticastBindMem_v2", + ) + address = check_cuda( + cuda.cuMemAddressReserve(size, granularity, 0, 0), + "cuMemAddressReserve(multicast)", + ) + check_cuda(cuda.cuMemMap(address, size, 0, handle, 0), "cuMemMap(multicast)") + check_cuda( + cuda.cuMemSetAccess( + address, + size, + [_read_write_access(device)], + 1, + ), + "cuMemSetAccess(multicast)", + ) + barrier() + return cls( + device, + size, + granularity, + handle, + int(address), + creator_rank, + backing, + ) + + @classmethod + def create_for_members( + cls, *, world: int, members: tuple[int, ...], rank: int, + creator_rank: int | None = None, device: int, size: int, + backing: FabricAllocation, broadcast: Callable[[object, int], object], + barrier: Callable[[], None], all_gather: Callable | None = None, + owners: list | None = None, + ) -> "MulticastGroup | None": + """Default to two creators sharing backing; preserve explicit placement. + + The first member uses the second member's alias. Every other member + uses the first member's alias. A singleton creates one object. Pass an + ``owners`` list on every EP rank to retain even nonmember lifecycle + tokens; close these pairs in reverse global creation order. + """ + if creator_rank is not None: + return cls._create_single_for_members(world=world, members=members, + rank=rank, creator_rank=creator_rank, device=device, size=size, + backing=backing, broadcast=broadcast, barrier=barrier) + # These two values define the fallback control communicator itself. + if type(world) is not int or world <= 0: + raise ValueError("world must be a positive exact int") + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if (type(members) is not tuple or not members or + any(type(member) is not int for member in members)): + raise ValueError("multicast members must be a non-empty tuple of ranks") + count, begin = len(members), members[0] + if members != tuple(range(begin, begin + count)): + raise ValueError("multicast members must be contiguous sorted ranks") + if begin < 0 or members[-1] >= world: + raise ValueError("multicast members must be in [0, world)") + if world % count or begin % count: + raise ValueError("multicast members must form an aligned EP subgroup") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if type(size) is not int or size <= 0: + raise ValueError("multicast size must be a positive exact int") + if owners is not None and type(owners) is not list: + raise TypeError("owners must be a list") + if not callable(broadcast) or not callable(barrier): + raise TypeError("broadcast and barrier must be callable") + if all_gather is None: + def all_gather(value): + return [broadcast(value if rank == root else None, root) + for root in range(world)] + if not callable(all_gather): + raise TypeError("all_gather must be callable") + + def validate_aligned_group(): + if not isinstance(backing, FabricAllocation): + raise TypeError("backing must be a FabricAllocation") + if backing.device != device or backing.size < size: + raise ValueError("backing device or size differs from multicast") + if rank in members: + granularity = cls.granularity_for(count, size, device)[0] + if granularity <= 0 or size % granularity: + raise ValueError("subgroup multicast size is not granular") + + try: + pair = MulticastTeam.create_dual(world=world, rank=rank, device=device, + members=members, size=size, all_gather=all_gather, + backing=backing if rank in members else None, + _validate=validate_aligned_group) + except MulticastTeamError as exc: + if owners is not None: + owners.append(exc.team) + raise + if owners is not None: + owners.append(pair) + if rank not in members: + return None + return cls(device, size, pair.granularity, pair.handle, pair.mc_ptr, + pair.creator_rank, backing, pair) + + @classmethod + def _create_single_for_members( + cls, + *, + world: int, + members: tuple[int, ...], + rank: int, + creator_rank: int | None = None, + device: int, + size: int, + backing: FabricAllocation, + broadcast: Callable[[object, int], object], + barrier: Callable[[], None], + ) -> "MulticastGroup | None": + """Collectively create one aligned subgroup multicast alias. + + Every EP rank calls this method in the same order so setup may keep its + existing world control communicator. Only ``members`` add a device, + bind their local backing, and map the multicast virtual address. + An omitted creator selects the second member (the sole member for a + singleton); an explicit creator preserves caller placement. This is + a deterministic placement policy, not a topology-independent guarantee. + """ + + if type(world) is not int or world <= 0: + raise ValueError("world must be a positive exact int") + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("rank must be in [0, world)") + if ( + not isinstance(members, tuple) + or not members + or any(type(member) is not int for member in members) + ): + raise ValueError("multicast members must be a non-empty tuple of ranks") + group_size = len(members) + group_begin = members[0] + if members != tuple(range(group_begin, group_begin + group_size)): + raise ValueError("multicast members must be contiguous sorted ranks") + if group_begin < 0 or members[-1] >= world: + raise ValueError("multicast members must be in [0, world)") + if world % group_size or group_begin % group_size: + raise ValueError("multicast members must form an aligned EP subgroup") + if creator_rank is None: + creator_rank = members[1] if group_size > 1 else members[0] + if type(creator_rank) is not int: + raise ValueError("multicast creator must be an exact rank") + if creator_rank not in members: + raise ValueError("multicast creator must belong to the subgroup") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if type(size) is not int or size <= 0: + raise ValueError("multicast size must be a positive exact int") + if not isinstance(backing, FabricAllocation): + raise TypeError("backing must be a FabricAllocation") + if backing.device != device: + raise ValueError("backing and multicast devices differ") + if not callable(broadcast) or not callable(barrier): + raise TypeError("broadcast and barrier must be callable") + properties = cuda.CUmulticastObjectProp() + properties.numDevices = len(members) + properties.size = size + properties.handleTypes = int(FABRIC) + properties.flags = 0 + granularity = int( + check_cuda( + cuda.cuMulticastGetGranularity( + properties, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, + ), + "cuMulticastGetGranularity(subgroup)", + ) + ) + if size % granularity: + raise ValueError("subgroup multicast size is not granular") + if backing.size < size: + raise ValueError("backing allocation is smaller than subgroup multicast") + + serialized = None + handle = None + if rank == creator_rank: + handle = check_cuda( + cuda.cuMulticastCreate(properties), + "cuMulticastCreate(subgroup)", + ) + exported = check_cuda( + cuda.cuMemExportToShareableHandle(handle, FABRIC, 0), + "cuMemExportToShareableHandle(subgroup)", + ) + serialized = bytes(exported.data) + serialized = broadcast(serialized, creator_rank) + try: + serialized = bytes(serialized) + except (TypeError, ValueError) as exc: + raise TypeError("broadcast returned an invalid multicast handle") from exc + if len(serialized) != 64: + raise RuntimeError( + f"subgroup fabric handle has {len(serialized)} bytes; expected 64" + ) + is_member = rank in members + if is_member and rank != creator_rank: + handle = check_cuda( + cuda.cuMemImportFromShareableHandle(bytearray(serialized), FABRIC), + "cuMemImportFromShareableHandle(subgroup)", + ) + if is_member: + check_cuda( + cuda.cuMulticastAddDevice(handle, device), + "cuMulticastAddDevice(subgroup)", + ) + barrier() + address = 0 + if is_member: + check_cuda( + cuda.cuMulticastBindMem_v2( + handle, device, 0, backing.handle, 0, size, 0 + ), + "cuMulticastBindMem_v2(subgroup)", + ) + address = check_cuda( + cuda.cuMemAddressReserve(size, granularity, 0, 0), + "cuMemAddressReserve(subgroup)", + ) + check_cuda( + cuda.cuMemMap(address, size, 0, handle, 0), + "cuMemMap(subgroup)", + ) + check_cuda( + cuda.cuMemSetAccess(address, size, [_read_write_access(device)], 1), + "cuMemSetAccess(subgroup)", + ) + barrier() + if not is_member: + return None + return cls( + device, + size, + granularity, + handle, + int(address), + creator_rank, + backing, + ) + + +def pick_multicast_creator(rank_to_node: list[int], source_rank: int = 0) -> int: + """Select a creator off source rank's node when the group spans nodes.""" + source_node = rank_to_node[source_rank] + for rank, node in enumerate(rank_to_node): + if node != source_node: + return rank + return source_rank + + +__all__ = [ + "FabricAllocation", + "MulticastGroup", + "check_cuda", + "pick_multicast_creator", +] + + +_RETAINED_MULTICAST_TEAMS: list[object] = [] + + +class MulticastTeamError(RuntimeError): + """A collective team operation failed; keep its partial resource graph.""" + + def __init__(self, message: str, team: "MulticastTeam", failures: list) -> None: + super().__init__(message) + self.team, self.failures = team, failures + + +class MulticastTeam: + """Explicit arbitrary-member multicast team with collective teardown. + + Every EP rank participates in identical create/close calls. Members with a + backing borrow it; members without one own a fresh dummy. Nonmembers never + allocate or map. Callers must release consumers before collective close. + There is deliberately no destructor that starts a distributed collective. + """ + + def __init__(self, *, world, rank, device, members, creator_rank, size, + all_gather, backing) -> None: + self.world, self.rank, self.device = world, rank, device + self.members, self.creator_rank, self.size = members, creator_rank, size + self.all_gather, self.backing_owner = all_gather, backing + self._borrowed = backing is not None + self.is_member = self.bound = self.mc_mapped = self._uc_mapped = False + self.handle = self._owned_handle = None + self.mc_ptr = self.local_ptr = self._owned_ptr = 0 + self.closed = self.retained_after_cleanup_failure = False + self.granularities = None + + @property + def owns_backing(self) -> bool: + return self.is_member and not self._borrowed + + def _call(self, name, *args): + return check_cuda(getattr(cuda, name)(*args), name) + + def _phase(self, label, operation, *, cleanup=False): + value, error = None, None + try: + value = operation() + except Exception as exc: + error = {"rank": self.rank, "type": type(exc).__name__, "message": str(exc)} + rows = self.all_gather({"rank": self.rank, "world": self.world, + "phase": label, "error": error}) + failures = [row["error"] for row in rows if row["error"] is not None] + if ([row["rank"] for row in rows] != list(range(len(rows))) or + any(row["world"] != len(rows) or row["phase"] != label for row in rows)): + failures.append({"message": "multicast team collective phase/rank mismatch"}) + if failures and not cleanup: + raise MulticastTeamError(label + ": " + repr(failures), self, failures) + return not failures if cleanup else value + + @staticmethod + def granularities_for(member_count: int, nbytes: int, device: int) -> dict: + """Query before allocation; sizes are never silently enlarged by create.""" + if any(type(x) is not int for x in (member_count, nbytes, device)) or min(member_count, nbytes) <= 0 or device < 0: + raise ValueError("team geometry requires positive exact counts and a nonnegative device") + supported = check_cuda(cuda.cuDeviceGetAttribute( + cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, device), + "cuDeviceGetAttribute(multicast)") + if int(supported) != 1: + raise RuntimeError("team device does not support multicast") + prop = cuda.CUmulticastObjectProp() + prop.numDevices, prop.size, prop.handleTypes, prop.flags = member_count, nbytes, int(FABRIC), 0 + cm, ca = cuda.CUmulticastGranularity_flags, cuda.CUmemAllocationGranularity_flags + values = {"mc_minimum": int(check_cuda(cuda.cuMulticastGetGranularity(prop, cm.CU_MULTICAST_GRANULARITY_MINIMUM))), + "mc_recommended": int(check_cuda(cuda.cuMulticastGetGranularity(prop, cm.CU_MULTICAST_GRANULARITY_RECOMMENDED))), + "allocation_minimum": int(check_cuda(cuda.cuMemGetAllocationGranularity(_allocation_properties(device), ca.CU_MEM_ALLOC_GRANULARITY_MINIMUM))), + "allocation_recommended": int(check_cuda(cuda.cuMemGetAllocationGranularity(_allocation_properties(device), ca.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)))} + if any(x <= 0 or x & (x - 1) for x in values.values()): + raise ValueError("invalid CUDA team granularity") + return values + + def _validate(self): + if type(self.world) is not int or self.world <= 0: + raise ValueError("world must be a positive exact int") + if type(self.rank) is not int or not 0 <= self.rank < self.world: + raise ValueError("rank must be an exact EP rank") + if type(self.device) is not int or self.device < 0: + raise ValueError("device must be a nonnegative exact ordinal") + if (type(self.members) is not tuple or not self.members or + any(type(x) is not int for x in self.members) or + self.members != tuple(sorted(set(self.members))) or + any(not 0 <= x < self.world for x in self.members)): + raise ValueError("members must be sorted unique exact ranks within EP") + if type(self.creator_rank) is not int or self.creator_rank not in self.members: + raise ValueError("creator must be an exact member rank") + if type(self.size) is not int or self.size <= 0: + raise ValueError("size must be a positive exact byte count") + self.is_member = self.rank in self.members + if not self.is_member and self.backing_owner is not None: + raise ValueError("nonmembers cannot provide backing") + if self.backing_owner is not None: + if not isinstance(self.backing_owner, FabricAllocation): + raise TypeError("borrowed backing must be a FabricAllocation") + if self.backing_owner.device != self.device or self.backing_owner.size < self.size or not self.backing_owner.local_ptr: + raise ValueError("borrowed backing device, size or UC mapping mismatch") + self.local_ptr = int(self.backing_owner.local_ptr) + + def _initialize(self): + self._phase("validate team arguments", self._validate) + configs = self.all_gather((self.world, self.members, self.creator_rank, self.size)) + def uniform(): + if any(row != configs[0] for row in configs): + raise ValueError("multicast team configuration differs across ranks") + self._phase("uniform team configuration", uniform) + def granularity(): + if self.is_member: + self.granularities = self.granularities_for(len(self.members), self.size, self.device) + if self.size % self.granularities["mc_minimum"] or self.size % self.granularities["allocation_minimum"]: + raise ValueError("explicit size is not granular; no automatic enlargement") + return self.granularities + values = self.all_gather(self._phase("team capability and granularity", granularity)) + self.granularity = max(max(row.values()) for row in values if row is not None) + def own_backing(): + if self.owns_backing: + self._owned_handle = self._call("cuMemCreate", self.size, _allocation_properties(self.device), 0) + self.backing_owner = FabricAllocation(self.device, self.size, self.granularity, self._owned_handle, 0, b"") + self._phase("fresh or borrowed backing", own_backing) + def reserve_uc(): + if self.owns_backing: + self._owned_ptr = int(self._call("cuMemAddressReserve", self.size, self.granularity, 0, 0)) + self._phase("owned UC reservation", reserve_uc) + def map_uc(): + if self.owns_backing: + self._call("cuMemMap", self._owned_ptr, self.size, 0, self._owned_handle, 0) + self._uc_mapped = True + self._phase("owned UC mapping", map_uc) + def access_uc(): + if self.owns_backing: + self._call("cuMemSetAccess", self._owned_ptr, self.size, [_read_write_access(self.device)], 1) + self.local_ptr = self._owned_ptr + self.backing_owner.local_ptr = self.local_ptr + self._phase("owned UC access", access_uc) + def create(): + if self.rank == self.creator_rank: + prop = cuda.CUmulticastObjectProp() + prop.numDevices, prop.size, prop.handleTypes, prop.flags = len(self.members), self.size, int(FABRIC), 0 + self.handle = self._call("cuMulticastCreate", prop) + self._phase("creator MC creation", create) + def export(): + if self.rank == self.creator_rank: + data = bytes(self._call("cuMemExportToShareableHandle", self.handle, FABRIC, 0).data) + if len(data) != 64: + raise ValueError("FABRIC handle must contain exactly 64 bytes") + return data + exports = self.all_gather(self._phase("creator MC export", export)) + def import_member(): + data = exports[self.creator_rank] + if type(data) is not bytes or len(data) != 64: + raise ValueError("invalid collective FABRIC handle") + if self.is_member and self.rank != self.creator_rank: + self.handle = self._call("cuMemImportFromShareableHandle", bytearray(data), FABRIC) + self._phase("member MC import", import_member) + self._phase("all AddDevice before bind/map", lambda: + self._call("cuMulticastAddDevice", self.handle, self.device) if self.is_member else None) + def bind(): + if self.is_member: + self._call("cuMulticastBindMem_v2", self.handle, self.device, 0, self.backing_owner.handle, 0, self.size, 0) + self.bound = True + self._phase("member backing bind", bind) + def reserve_mc(): + if self.is_member: + self.mc_ptr = int(self._call("cuMemAddressReserve", self.size, self.granularity, 0, 0)) + self._phase("member MC reservation", reserve_mc) + def map_mc(): + if self.is_member: + self._call("cuMemMap", self.mc_ptr, self.size, 0, self.handle, 0) + self.mc_mapped = True + self._phase("member MC mapping", map_mc) + self._phase("member MC access", lambda: + self._call("cuMemSetAccess", self.mc_ptr, self.size, [_read_write_access(self.device)], 1) if self.is_member else None) + return self + + @classmethod + def create(cls, *, world, rank, device, members, creator_rank, size, + all_gather, backing=None) -> "MulticastTeam": + if not callable(all_gather): + raise TypeError("all_gather must be callable") + team = cls(world=world, rank=rank, device=device, members=members, + creator_rank=creator_rank, size=size, all_gather=all_gather, backing=backing) + try: + return team._initialize() + except Exception as exc: + try: + team.close_collectively() + except Exception: + team._retain() + if isinstance(exc, MulticastTeamError): + raise + raise MulticastTeamError(str(exc), team, [{"message": str(exc)}]) from exc + + @classmethod + def create_dual(cls, *, world, rank, device, members, size, all_gather, + backing=None, _validate=None) -> "MulticastTeamPair": + """Create first/second-member aliases, with at most one owned backing. + + Explicit ``create(..., creator_rank=...)`` semantics are unchanged. + """ + return MulticastTeamPair.create(world=world, rank=rank, device=device, + members=members, size=size, all_gather=all_gather, backing=backing, + team_factory=cls, validate=_validate) + + def _retain(self): + if not self.retained_after_cleanup_failure: + _RETAINED_MULTICAST_TEAMS.append(self) + self.retained_after_cleanup_failure = True + + def close_collectively(self) -> bool: + """Drain and free in dependency order; stop collectively on any error.""" + if self.closed: + return True + def unmap_mc(): + if self.mc_mapped: + self._call("cuMemUnmap", self.mc_ptr, self.size) + self.mc_mapped = False + if self.mc_ptr: + self._call("cuMemAddressFree", self.mc_ptr, self.size) + self.mc_ptr = 0 + def unbind(): + if self.bound: + self._call("cuMulticastUnbind", self.handle, self.device, 0, self.size) + self.bound = False + def release_imports(): + if self.rank != self.creator_rank and self.handle is not None: + self._call("cuMemRelease", self.handle) + self.handle = None + def release_creator(): + if self.rank == self.creator_rank and self.handle is not None: + self._call("cuMemRelease", self.handle) + self.handle = None + def unmap_uc(): + if self._uc_mapped: + self._call("cuMemUnmap", self._owned_ptr, self.size) + self._uc_mapped = False + if self._owned_ptr: + self._call("cuMemAddressFree", self._owned_ptr, self.size) + self._owned_ptr = self.local_ptr = 0 + def release_backing(): + if self._owned_handle is not None: + self._call("cuMemRelease", self._owned_handle) + self._owned_handle = None + self.backing_owner = None + try: + for label, fn in (("drain context", lambda: self._call("cuCtxSynchronize")), + ("unmap MC", unmap_mc), ("unbind MC", unbind), + ("release imports", release_imports), ("release creator", release_creator), + ("unmap owned UC", unmap_uc), ("release owned backing", release_backing)): + if not self._phase("cleanup " + label, fn, cleanup=True): + self._retain() + return False + except Exception: + self._retain() + raise + self.closed = True + return True + + def metadata(self) -> dict: + return {"rank": self.rank, "device": self.device, "members": list(self.members), + "creator_rank": self.creator_rank, "size": self.size, "is_member": self.is_member, + "owns_backing": self.owns_backing, "borrowed_backing": self._borrowed, + "mc_ptr": self.mc_ptr, "local_ptr": self.local_ptr, "bound": self.bound, + "granularities": self.granularities, "closed": self.closed, + "retained_after_cleanup_failure": self.retained_after_cleanup_failure} + + +class MulticastTeamPair: + """Two MC objects over one backing, exposing the rank-selected aperture. + + The first alias owns any dummy; the second borrows that exact allocation. + Every EP rank retains this object, including nonmembers, and calls close in + the same global order. No destructor starts a collective. + """ + + _SELECTED_FIELDS = frozenset(("handle", "mc_ptr", "creator_rank", "bound", + "mc_mapped", "granularity", "granularities")) + _PRIMARY_FIELDS = frozenset(("backing_owner", "owns_backing", "local_ptr", + "_owned_handle", "_owned_ptr", "_uc_mapped", "_borrowed", "is_member")) + + def __init__(self, *, world, rank, device, members, size, all_gather): + self.world, self.rank, self.device = world, rank, device + self.members, self.size, self.all_gather = members, size, all_gather + self._aliases = [] + self.closed = self.retained_after_cleanup_failure = False + + @property + def aliases(self) -> tuple[MulticastTeam, ...]: + return tuple(self._aliases) + + @property + def selected_alias(self) -> MulticastTeam: + if not self._aliases: + raise RuntimeError("multicast pair has no initialized alias") + index = 1 if len(self._aliases) > 1 and self.rank == self.members[0] else 0 + return self._aliases[index] + + def __getattr__(self, name): + if name in self._SELECTED_FIELDS: + return getattr(self.selected_alias, name) + if name in self._PRIMARY_FIELDS and self._aliases: + return getattr(self._aliases[0], name) + raise AttributeError(name) + + def _phase(self, label, operation, *, cleanup=False): + return MulticastTeam._phase(self, label, operation, cleanup=cleanup) + + def _retain(self): + MulticastTeam._retain(self) + + @classmethod + def create(cls, *, world, rank, device, members, size, all_gather, + backing=None, team_factory=MulticastTeam, validate=None): + if not callable(all_gather): + raise TypeError("all_gather must be callable") + pair = cls(world=world, rank=rank, device=device, members=members, + size=size, all_gather=all_gather) + try: + creator = members[0] if type(members) is tuple and members else None + probe = team_factory(world=world, rank=rank, device=device, + members=members, creator_rank=creator, size=size, + all_gather=all_gather, backing=backing) + def validate_arguments(): + probe._validate() + if validate is not None: + validate() + pair._phase("validate dual multicast arguments", validate_arguments) + primary = team_factory.create(world=world, rank=rank, device=device, + members=members, creator_rank=members[0], size=size, + all_gather=all_gather, backing=backing) + pair._aliases.append(primary) + if len(members) > 1: + secondary = team_factory.create(world=world, rank=rank, device=device, + members=members, creator_rank=members[1], size=size, + all_gather=all_gather, + backing=primary.backing_owner if primary.is_member else None) + pair._aliases.append(secondary) + return pair + except Exception as exc: + if isinstance(exc, MulticastTeamError) and exc.team is not pair: + if all(alias is not exc.team for alias in pair._aliases): + pair._aliases.append(exc.team) + # A failed secondary cleanup may still reference primary's dummy. + # Never free the primary allocation while any such alias survives. + if any(alias.retained_after_cleanup_failure for alias in pair._aliases): + pair._retain() + else: + try: + if not pair.close_collectively(): + pair._retain() + except Exception: + pair._retain() + failures = exc.failures if isinstance(exc, MulticastTeamError) else [{"message": str(exc)}] + raise MulticastTeamError(str(exc), pair, failures) from exc + + def close_collectively(self) -> bool: + """Free the borrowed secondary first; the primary owns the only dummy.""" + if self.closed: + return True + if self.retained_after_cleanup_failure: + return False + try: + for alias in reversed(self._aliases): + if alias.retained_after_cleanup_failure or not alias.close_collectively(): + self._retain() + return False + except Exception: + self._retain() + raise + self.closed = True + return True + + def metadata(self) -> dict: + selected = self.selected_alias.metadata() if self._aliases else {} + selected.update({"rank": self.rank, "device": self.device, + "members": list(self.members) if type(self.members) is tuple else [], + "size": self.size, "dual_creator": len(self.members) > 1 if type(self.members) is tuple else False, + "selection_policy": "first member uses second creator; all others use first creator", + "alias_count": len(self._aliases), + "aliases": [alias.metadata() for alias in self._aliases], + "closed": self.closed, + "retained_after_cleanup_failure": self.retained_after_cleanup_failure, + "owns_backing": self._aliases[0].owns_backing if self._aliases else False, + "allocated_bytes": sum(alias.size for alias in self._aliases if alias._owned_handle is not None), + "mapping_count": sum(bool(alias.mc_ptr) for alias in self._aliases), + "reserved_va_bytes": sum((alias.size if alias.mc_ptr else 0) + + (alias.size if alias._owned_ptr else 0) for alias in self._aliases)}) + return selected + + +__all__.extend(["MulticastTeam", "MulticastTeamPair", "MulticastTeamError"]) diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py new file mode 100644 index 000000000000..2cdc1caaa014 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/geometry.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Frozen seven-plane MegaMoE weight geometry. + +Storage placement is owned by :mod:`megamoe_scheduler.sami.arena`; production +uses one external plane-major live bank without internal receive storage. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +CANONICAL_WEIGHT_PLANE_NAMES = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", +) + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +@dataclass(frozen=True) +class PlaneViewLayout: + """One per-expert tensor view of a live weight plane. + + Shape and stride omit the leading physical-slot axis. The descriptor is + Torch-independent, so validation does not resolve optional FP4/FP8 dtypes. + """ + + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: str + element_size: int + pointer_alignment: int + + def __post_init__(self) -> None: + if len(self.shape) != len(self.stride): + raise ValueError("plane view shape and stride must have equal rank") + if any(type(value) is not int or value <= 0 for value in self.shape): + raise ValueError("plane view extents must be positive exact integers") + if any(type(value) is not int or value <= 0 for value in self.stride): + raise ValueError("plane view strides must be positive exact integers") + if type(self.dtype) is not str or not self.dtype: + raise ValueError("plane view dtype must be a non-empty string") + if type(self.element_size) is not int or self.element_size <= 0: + raise ValueError("plane view element_size must be a positive exact integer") + if type(self.pointer_alignment) is not int or self.pointer_alignment <= 0: + raise ValueError("plane view pointer_alignment must be a positive exact integer") + + @property + def elements_per_slot(self) -> int: + if not self.shape: + return 1 + return 1 + sum((extent - 1) * stride for extent, stride in zip(self.shape, self.stride)) + + @property + def nbytes(self) -> int: + return self.elements_per_slot * self.element_size + + def batched_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if type(slot_count) is not int or slot_count <= 0: + raise ValueError("slot_count must be a positive exact integer") + return ( + (slot_count, *self.shape), + (self.elements_per_slot, *self.stride), + self.dtype, + self.element_size, + self.pointer_alignment, + ) + + +@dataclass(frozen=True) +class Plane: + """One canonical MegaMoE weight plane within an expert bundle. + + The first four fields preserve the original public constructor. Layouts + created by BundleLayout.create additionally carry the storage view used by + the loader and the typed view consumed by the kernel. + """ + + name: str + offset: int + nbytes: int + alignment: int + storage: Optional[PlaneViewLayout] = None + kernel: Optional[PlaneViewLayout] = None + + def __post_init__(self) -> None: + if (self.storage is None) != (self.kernel is None): + raise ValueError("plane storage and kernel layouts must be provided together") + if self.storage is not None: + if self.storage.nbytes != self.nbytes: + raise ValueError( + f"{self.name} storage view spans {self.storage.nbytes} bytes; " + f"expected {self.nbytes}" + ) + assert self.kernel is not None + if self.kernel.nbytes != self.nbytes: + raise ValueError( + f"{self.name} kernel view spans {self.kernel.nbytes} bytes; " + f"expected {self.nbytes}" + ) + + def storage_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if self.storage is None: + raise ValueError(f"{self.name} has no canonical storage view") + return self.storage.batched_metadata(slot_count) + + def kernel_metadata( + self, slot_count: int + ) -> tuple[tuple[int, ...], tuple[int, ...], str, int, int]: + if self.kernel is None: + raise ValueError(f"{self.name} has no canonical kernel view") + return self.kernel.batched_metadata(slot_count) + + +@dataclass(frozen=True) +class BundleLayout: + """NVFP4 bytes/expert for the seven plane-major live tensors. + + ``offset`` and ``stride_bytes`` describe a canonical packed serialization + only; direct-live SAMI uses each plane's ``nbytes`` as its physical-slot + byte stride and never allocates that serialization. + """ + + hidden: int + intermediate: int + planes: tuple[Plane, ...] + total_bytes: int + stride_bytes: int + alignment: int + expand_intermediate: Optional[int] = None + + @classmethod + def create( + cls, + hidden: int = 7168, + intermediate: int = 3072, + *, + expand_intermediate: Optional[int] = None, + weight_bits: int = 4, + sf_vec_size: int = 16, + sf_bits: int = 8, + plane_alignment: int = 512, + ) -> "BundleLayout": + dimensions = { + "hidden": hidden, + "intermediate": intermediate, + "weight_bits": weight_bits, + "sf_vec_size": sf_vec_size, + "sf_bits": sf_bits, + "plane_alignment": plane_alignment, + } + if any(type(value) is not int or value <= 0 for value in dimensions.values()): + raise ValueError("all bundle dimensions must be positive exact integers") + if expand_intermediate is None: + expand_intermediate = 2 * intermediate + if type(expand_intermediate) is not int or expand_intermediate <= 0: + raise ValueError("all bundle dimensions must be positive exact integers") + if (weight_bits, sf_vec_size, sf_bits) != (4, 16, 8): + raise ValueError( + "production bundle format is fixed to NVFP4 weights and " + "FP8 block-16 scale factors" + ) + if hidden % sf_vec_size or intermediate % sf_vec_size: + raise ValueError("hidden/intermediate must be block-scale aligned") + if expand_intermediate != 2 * intermediate: + raise ValueError( + "production bundle requires expand_intermediate == 2 * intermediate" + ) + + # These are the exact TEKit transformed-tensor formulas. The simpler + # elements/sf_vec_size expression only coincides at some model shapes. + fc1_sf_elements = _align_up(expand_intermediate, 128) * _align_up( + hidden // sf_vec_size, 4 + ) + fc2_sf_elements = _align_up(hidden, 128) * _align_up( + intermediate // sf_vec_size, 4 + ) + + def view( + shape: tuple[int, ...], + dtype: str, + element_size: int, + pointer_alignment: int, + *, + stride: Optional[tuple[int, ...]] = None, + ) -> PlaneViewLayout: + if stride is None: + steps = [] + running = 1 + for extent in reversed(shape): + steps.append(running) + running *= extent + stride = tuple(reversed(steps)) + return PlaneViewLayout( + shape=shape, + stride=stride, + dtype=dtype, + element_size=element_size, + pointer_alignment=pointer_alignment, + ) + + gate_up = expand_intermediate + specs = ( + ( + "mega_fc1_weight", + plane_alignment, + view((gate_up, hidden // 2), "uint8", 1, 16), + view( + (hidden // 2, gate_up), + "float4_e2m1fn_x2", + 1, + 16, + stride=(1, hidden // 2), + ), + ), + ( + "mega_fc1_weight_sf", + plane_alignment, + view((fc1_sf_elements,), "uint8", 1, 16), + view((fc1_sf_elements,), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + plane_alignment, + view((hidden, intermediate // 2), "uint8", 1, 16), + view( + (intermediate // 2, hidden), + "float4_e2m1fn_x2", + 1, + 16, + stride=(1, intermediate // 2), + ), + ), + ( + "mega_fc2_weight_sf", + plane_alignment, + view((fc2_sf_elements,), "uint8", 1, 16), + view((fc2_sf_elements,), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ( + "fc2_alpha", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 16, + view((), "float32", 4, 4), + view((), "float32", 4, 4), + ), + ) + + planes: list[Plane] = [] + cursor = 0 + for name, alignment, storage, kernel in specs: + nbytes = storage.nbytes + if kernel.nbytes != nbytes: + raise ValueError(f"{name} storage and kernel views span different bytes") + cursor = _align_up(cursor, alignment) + planes.append(Plane(name, cursor, nbytes, alignment, storage=storage, kernel=kernel)) + cursor += nbytes + return cls( + hidden=hidden, + intermediate=intermediate, + planes=tuple(planes), + total_bytes=sum(plane.nbytes for plane in planes), + stride_bytes=_align_up(cursor, plane_alignment), + alignment=plane_alignment, + expand_intermediate=expand_intermediate, + ) + + @property + def plane_count(self) -> int: + return len(self.planes) + + +__all__ = [ + "BundleLayout", + "CANONICAL_WEIGHT_PLANE_NAMES", + "Plane", + "PlaneViewLayout", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py new file mode 100644 index 000000000000..c4159e593d7b --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_scheduler/sami/hierarchical.py @@ -0,0 +1,716 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hierarchical HALO-Q live arena and low-overhead SAMI submission path.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import threading +from typing import Protocol, runtime_checkable + +import torch +from cuda.bindings import driver as cuda + +from .arena import BoundLiveWeightArena, LiveWeightArena, _TensorIdentity +from ._util import check_cuda +from ..geometry import hierarchy_group_sizes +from .geometry import BundleLayout + +@dataclass(frozen=True) +class BoundHierarchicalLiveWeightArena: + """Validated base arena plus this rank's multicast alias at every level.""" + + arena: "HierarchicalLiveWeightArena" = field(repr=False, compare=False) + base: BoundLiveWeightArena = field(repr=False, compare=False) + group_sizes: tuple[int, ...] + group_destination_ptrs: tuple[int, ...] + group_mc_views: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + group_mc_owners: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + + +@dataclass(frozen=True) +class HierarchicalLiveWeightArena: + """One external live bank with a multicast alias for each HALO level. + + ``group_mc_views`` and ``group_mc_owners`` are plane-major: + ``[plane][level]``. At a given level, each rank exposes only the alias and + strong VMM owner token for its own aligned group. Level zero is the + full-world alias already present in ``base.planes[*].mc_view``. + """ + + base: LiveWeightArena + group_sizes: tuple[int, ...] + group_mc_views: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + group_mc_owners: tuple[tuple[object, ...], ...] = field( + repr=False, compare=False + ) + aliases_same_backing: bool = False + + def bind( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + ) -> BoundHierarchicalLiveWeightArena: + if self.aliases_same_backing is not True: + raise ValueError("hierarchical multicast aliases are not provider-attested") + expected_sizes = hierarchy_group_sizes(world) + if self.group_sizes != expected_sizes: + raise ValueError( + f"hierarchy group sizes {self.group_sizes} != {expected_sizes}" + ) + base = self.base.bind( + world=world, + rank=rank, + home_count=home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + ) + if len(self.group_mc_views) != len(bundle.planes): + raise ValueError("hierarchical arena must expose every weight plane") + if len(self.group_mc_owners) != len(bundle.planes): + raise ValueError("hierarchical arena must retain every VMM owner") + if self.base.terminals.backing_owner is None: + raise ValueError("hierarchical terminal backing owner is not retained") + + identities: list[tuple[_TensorIdentity, ...]] = [] + for views, owners, base_view in zip( + self.group_mc_views, self.group_mc_owners, self.base.planes + ): + if base_view.backing_owner is None: + raise ValueError(f"{base_view.name} backing owner is not retained") + if len(views) != len(expected_sizes): + raise ValueError( + f"{base_view.name} has the wrong hierarchy depth" + ) + if len(owners) != len(expected_sizes) or any( + owner is None for owner in owners + ): + raise ValueError( + f"{base_view.name} must retain every subgroup VMM owner" + ) + captured = tuple( + _TensorIdentity.capture( + view, name=f"{base_view.name}.group_mc[{level}]" + ) + for level, view in enumerate(views) + ) + reference = _TensorIdentity.capture( + base_view.mc_view, name=f"{base_view.name}.mc" + ) + for identity in captured: + if identity.metadata != reference.metadata: + raise ValueError( + f"{base_view.name} group multicast metadata changed" + ) + if captured[0].pointer != reference.pointer: + raise ValueError( + f"{base_view.name} level-zero alias is not the base MC view" + ) + if len({identity.pointer for identity in captured}) != len(captured): + raise ValueError( + f"{base_view.name} hierarchy levels need distinct virtual mappings" + ) + identities.append(captured) + + destinations: list[int] = [] + for level in range(len(expected_sizes)): + for helper in range(helper_count): + physical_slot = home_count + helper + for plane_index, layout in enumerate(bundle.planes): + destinations.append( + identities[plane_index][level].pointer + + physical_slot * layout.nbytes + ) + return BoundHierarchicalLiveWeightArena( + arena=self, + base=base, + group_sizes=expected_sizes, + group_destination_ptrs=tuple(destinations), + group_mc_views=self.group_mc_views, + group_mc_owners=self.group_mc_owners, + ) + +# runtime_checkable because __init__ isinstance()-checks the provider below. +# Without it that check raises TypeError instead of validating anything, and +# only at runtime -- no import, compile or signature check sees it. +@runtime_checkable +class HierarchicalLiveWeightArenaProvider(Protocol): + """Provider boundary implemented by TEKit or a Fabric test allocator.""" + + def build_hierarchical_live_arena( + self, + *, + world: int, + rank: int, + home_count: int, + helper_count: int, + bundle: BundleLayout, + device: int, + comm: object, + ) -> HierarchicalLiveWeightArena: + """Construct the external live bank and all aligned-group aliases.""" + + +def _stream_handle(stream: object) -> int: + if isinstance(stream, bool): + raise TypeError("stream must be a CUDA stream handle") + handle = getattr(stream, "cuda_stream", stream) + value = int(handle) + if value <= 0: + raise ValueError("hierarchical copy requires a positive non-default stream") + return value + + +def _event_handle(event: object | None) -> int: + if event is None: + return 0 + if isinstance(event, bool): + raise TypeError("event must be a CUDA event handle") + return int(getattr(event, "cuda_event", event)) + + +class HierarchicalSamiWeightBroadcast: + """One-rank HALO-Q copy endpoint with a single C hot-path context.""" + + partition = "bundle" + remote_visibility_release = True + live_bank_count = 1 + + def __init__( + self, + *, + comm: object, + device: int, + bundle: BundleLayout, + helper_count: int, + global_expert_count: int, + arena_provider: HierarchicalLiveWeightArenaProvider, + tma_sm_count: int | None = None, + tma_warps: int | None = None, + tma_route: str = "plan", + tma_source_load_factor: float = 1.0, + helper_residency: bool = False, + ) -> None: + for name, value in (("tma_sm_count", tma_sm_count), ("tma_warps", tma_warps)): + if value is not None and (type(value) is not int or value <= 0): + raise ValueError(f"{name} must be a positive exact int or None") + if tma_sm_count is not None and tma_sm_count > 8: + raise ValueError("tma_sm_count must be at most 8") + if tma_warps is not None and tma_warps > 32: + raise ValueError("tma_warps must be at most 32; device shared-memory limits also apply") + if tma_route not in ("plan", "source", "scatter"): + raise ValueError("tma_route must be 'plan', 'source', or 'scatter'") + if type(helper_residency) is not bool: + raise ValueError("helper_residency must be a bool") + # Residency skips copies into helper rows that already hold the planned + # expert, so only this layer's copies may write those rows. + self.helper_residency = helper_residency + if (type(tma_source_load_factor) not in (int, float) or + not 0.01 <= tma_source_load_factor < 2.0 or + abs(100 * tma_source_load_factor - round(100 * tma_source_load_factor)) > 1e-9): + raise ValueError("tma_source_load_factor must be 0.01..1.99 in steps of 0.01") + self.tma_route = tma_route + self.tma_source_load_factor = round(100 * tma_source_load_factor) / 100 + world = getattr(comm, "world", None) + rank = getattr(comm, "rank", None) + if type(world) is not int: + raise ValueError("comm.world must be an exact int") + hierarchy_group_sizes(world) + if type(rank) is not int or not 0 <= rank < world: + raise ValueError("comm.rank must be in [0, world)") + if type(device) is not int or device < 0: + raise ValueError("device must be a non-negative exact int") + if ( + type(global_expert_count) is not int + or global_expert_count <= 0 + or global_expert_count % world + ): + raise ValueError("global_expert_count must be positive and EP-divisible") + if type(helper_count) is not int or helper_count <= 0: + raise ValueError("helper_count must be a positive exact int") + if global_expert_count + world * helper_count > (1 << 31) - 1: + raise ValueError("physical slots exceed public int32 indexing") + if 4 + 7 * ((helper_count + 3) // 4 * 4) > (1 << 31) - 1: + raise ValueError("PlanChannel word count exceeds native int32 indexing") + if not isinstance(arena_provider, HierarchicalLiveWeightArenaProvider): + raise TypeError("arena_provider lacks the hierarchical arena protocol") + self.comm = comm + self.device = device + self.bundle = bundle + self.helper_count = helper_count + self.global_expert_count = global_expert_count + self.home_count = global_expert_count // world + self.home_begin = rank * self.home_count + self.provider = arena_provider + self.cuda = cuda + torch.cuda.set_device(device) + # Fail launch-geometry checks before allocating collective Fabric resources. + device_sms = torch.cuda.get_device_properties(device).multi_processor_count + if tma_sm_count is not None and tma_sm_count > device_sms: + raise ValueError(f"tma_sm_count must be at most {device_sms} on this GPU") + self._validate_tma_warp_capacity(tma_warps) + # Unlike CTA/warp counts, the ownership rule must agree across ranks. + # This is one cold-path collective, outside every copy submission. + factors = comm.allgather(self.tma_source_load_factor) + if any(value != self.tma_source_load_factor for value in factors): + raise ValueError("tma_source_load_factor must agree across all ranks") + + live = arena_provider.build_hierarchical_live_arena( + world=world, + rank=rank, + home_count=self.home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + comm=comm, + ) + if not isinstance(live, HierarchicalLiveWeightArena): + raise TypeError("provider returned a non-hierarchical live arena") + self.live_arena = live + self.arena = live.bind( + world=world, + rank=rank, + home_count=self.home_count, + helper_count=helper_count, + bundle=bundle, + device=device, + ) + self.live_plane_tensors = self.arena.base.local_plane_views + self.terminal_flags = self.live_arena.base.terminals.local_view + self.flags_local = self.arena.base.terminal_local_ptr + self.flags_mc = self.arena.base.terminal_mc_ptr + 8 * rank + check_cuda( + cuda.cuMemsetD8( + cuda.CUdeviceptr(self.flags_local), 0, 8 * world + ), + "cuMemsetD8(hierarchical terminals)", + ) + torch.cuda.synchronize(device) + comm.barrier() + + # GPU-direct consumes the scheduler's device storage through the + # statically linked THOP; no mapped channel or runtime compilation exists. + plan_stride = (helper_count + 3) // 4 * 4 + plan_words = 4 + 7 * plan_stride + planes = len(bundle.planes) + ctx = torch.ops.trtllm.moe_rebalance_tma_create( + helper_count * planes, tma_sm_count or 0, tma_warps or 0 + ) + try: + target_count = sum(world // size for size in self.arena.group_sizes[1:]) + torch.ops.trtllm.moe_rebalance_tma_configure_gpu_plan( + ctx, world, rank, helper_count, planes, + global_expert_count, self.home_count, plan_stride, + len(self.arena.group_sizes), target_count, 6, plan_words, 0, + {"plan": 0, "source": 1, "scatter": 2}[tma_route], + round(100 * self.tma_source_load_factor), + list(self.arena.group_sizes), + list(self.arena.base.scatter_source_ptrs), + list(self.arena.group_destination_ptrs), [], + [plane.nbytes for plane in bundle.planes], + ) + config_names = ( + "abi_version", "device", "device_sm_count", "sms", "warps", + "threads_per_cta", "slots_per_warp", "bank0_slots_per_warp", + "bank1_slots_per_warp", "slice_bytes", "dynamic_shared_bytes", + "device_optin_shared_bytes", "device_shared_bytes_per_sm", + "max_active_ctas_per_sm", "compute_major", "compute_minor", + "total_slots", "max_slots_per_warp", "extra_slot_warps", + "max_warps", "max_segments", + ) + self.tma_config = dict(zip( + config_names, torch.ops.trtllm.moe_rebalance_tma_config(ctx) + )) + if int(torch.ops.trtllm.moe_rebalance_tma_current_gen(ctx)) != 1: + raise RuntimeError("hierarchical SAMI context is not fresh") + # The per-layer launch and release ops take this handle of ctx; + # passing the custom-class object costs microseconds per call. + handle = int(torch.ops.trtllm.moe_rebalance_tma_handle(ctx)) + except BaseException: + torch.ops.trtllm.moe_rebalance_tma_destroy(ctx) + raise + self.ctx = ctx + self._handle = handle + + self._lock = threading.Lock() + self._scheduler: object | None = None + self._outputs: object | None = None + self._stream = 0 + self._stream_owner: object | None = None + self._generation = 1 + self._reuse_authority: object | None = None + self._active_generation = 0 + self._local_release_generation = 0 + self._collective_safe_generation = 0 + self._handoff_aborted = False + + def _validate_tma_warp_capacity(self, warps: int | None) -> None: + # Reject impossible device geometry before building a collective arena. + # Native creation additionally checks the compiled kernel's own static + # shared memory, thread limit and occupancy. + device = check_cuda(self.cuda.cuDeviceGet(self.device), "cuDeviceGet(TMA geometry)") + shared_bytes = int(check_cuda(self.cuda.cuDeviceGetAttribute( + self.cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + device, + ), "cuDeviceGetAttribute(TMA opt-in shared memory)")) + max_threads = int(check_cuda(self.cuda.cuDeviceGetAttribute( + self.cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, + device, + ), "cuDeviceGetAttribute(TMA block threads)")) + slot_bytes = 8208 + 32 # Logical 8 KiB plus aligned-load padding and SlotControl. + slots = max(0, shared_bytes // slot_bytes) + while slots and ((slots * slot_bytes + 127) & ~127) > shared_bytes: + slots -= 1 + max_warps = min(32, slots // 2, max_threads // 64) + if (warps or 7) > max_warps: + raise ValueError( + f"tma_warps must be at most {max_warps} on this GPU " + "(at least two 8 KiB slices per worker pair)" + ) + + def _validate_stream_context(self, handle: int) -> None: + get_current = getattr(self.cuda, "cuCtxGetCurrent", None) + get_device = getattr(self.cuda, "cuCtxGetDevice", None) + get_stream_context = getattr(self.cuda, "cuStreamGetCtx", None) + if not callable(get_current) or not callable(get_device): + raise RuntimeError("CUDA driver binding cannot validate the current context") + current_context = check_cuda(get_current(), "cuCtxGetCurrent") + if current_context is None: + raise RuntimeError("no current CUDA context for hierarchical scheduler") + current_device = int(check_cuda(get_device(), "cuCtxGetDevice")) + if current_device != self.device: + raise ValueError("current CUDA context and SAMI devices differ") + if handle: + if not callable(get_stream_context): + raise RuntimeError("CUDA driver binding cannot validate stream context") + stream_context = check_cuda( + get_stream_context(handle), "cuStreamGetCtx(hierarchical scheduler)" + ) + if stream_context != current_context: + raise ValueError("scheduler stream is from a different CUDA context") + + def bind_scheduler( + self, scheduler: object, stream: object, route_epoch: torch.Tensor + ) -> None: + """Bind the four-field CUDA scheduler and one static stream for life. + + The native context also binds every argument that the fused launches + share, including ``route_epoch``, the caller-owned int64 that each + launch raises once its routes are written. + """ + + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if self._scheduler is not None: + raise RuntimeError("hierarchical scheduler binding is single-assignment") + config = getattr(scheduler, "cfg", None) + outputs = getattr(scheduler, "outputs", None) + if config is None or outputs is None: + raise TypeError("scheduler lacks cfg/outputs") + + expected = ( + ("ep_size", int(self.comm.world)), + ("logical_expert_count", self.global_expert_count), + ("extra_slots_per_rank", self.helper_count), + ("local_rank", int(self.comm.rank)), + ) + for name, value in expected: + observed = getattr(config, name, None) + if type(observed) is not int or observed != value: + raise ValueError(f"scheduler {name} is incompatible with SAMI") + for name in ("max_tokens_per_rank", "topk"): + observed = getattr(config, name, None) + if type(observed) is not int or observed <= 0: + raise ValueError(f"scheduler {name} must be a positive exact int") + fields = ( + "physical_slot_ids", + "hot_expert_ids", + "hot_expert_group_level", + "hot_expert_source_ranks", + ) + tensors = tuple(getattr(outputs, name, None) for name in fields) + if any(not isinstance(tensor, torch.Tensor) for tensor in tensors): + raise TypeError("scheduler does not expose the four-tensor HALO ABI") + expected_shapes = ( + (config.max_tokens_per_rank, config.topk), + (self.helper_count,), + (self.helper_count,), + (self.helper_count,), + ) + for name, tensor, expected_shape in zip(fields, tensors, expected_shapes): + if ( + tensor.dtype != torch.int32 + or not tensor.is_cuda + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous CUDA int32 storage" + ) + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} must have shape {expected_shape}") + + scheduler_device = getattr(scheduler, "device", None) + try: + scheduler_device = torch.device(scheduler_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("scheduler must expose one explicit CUDA device") from exc + if scheduler_device.type != "cuda" or scheduler_device.index is None: + raise ValueError("scheduler must expose one explicit CUDA device") + if scheduler_device.index != self.device: + raise ValueError("scheduler and SAMI devices differ") + if any(tensor.device != scheduler_device for tensor in tensors): + raise ValueError("HALO outputs must share the scheduler CUDA device") + + handle = _stream_handle(stream) + stream_device = getattr(stream, "device", None) + if stream_device is not None: + try: + stream_device = torch.device(stream_device) + except (TypeError, RuntimeError) as exc: + raise TypeError("scheduler stream exposes an invalid device") from exc + if stream_device != scheduler_device: + raise ValueError("scheduler stream and output devices differ") + bind = getattr(scheduler, "bind_gpu_direct", None) + fused_args = getattr(scheduler, "fused_launch_args", None) + if not callable(bind) or not callable(fused_args): + raise TypeError("scheduler lacks the fused copy launch lifecycle") + workspace = getattr(scheduler, "plan_workspace", None) + status = getattr(scheduler, "status", None) + grid_sync = getattr(scheduler, "grid_sync", None) + capacity = getattr(scheduler, "max_broadcasts", None) + if type(capacity) is not int or capacity <= 0: + raise ValueError("GPU-direct requires the scheduler plan capacity") + if (not isinstance(workspace, torch.Tensor) or + workspace.device != scheduler_device or workspace.dtype != torch.int32 or + not workspace.is_contiguous() or tuple(workspace.shape) != (2 + 6 * capacity,)): + raise ValueError("GPU-direct requires the original CUDA scheduler plan workspace") + if (not isinstance(status, torch.Tensor) or status.device != scheduler_device or + status.dtype != torch.int32 or not status.is_contiguous() or status.numel() < 6): + raise ValueError("GPU-direct requires the CUDA scheduler status workspace") + if (not isinstance(grid_sync, torch.Tensor) or grid_sync.device != scheduler_device or + grid_sync.dtype != torch.int32 or not grid_sync.is_contiguous() or grid_sync.numel() < 2): + raise ValueError("GPU-direct requires the CUDA scheduler epoch workspace") + self._validate_stream_context(handle) + try: + bind(handle) + torch.ops.trtllm.moe_rebalance_tma_bind_gpu_direct( + self.ctx, + *tensors[1:], + workspace, + status, + grid_sync, + capacity, + self.helper_residency, + ) + torch.ops.trtllm.moe_rebalance_tma_bind_fused( + self.ctx, self.flags_mc, route_epoch, *fused_args(), handle + ) + except BaseException: + self._handoff_aborted = True + raise + self._scheduler = scheduler + self._outputs = outputs + self._stream = handle + # The native context launches on this stream by handle; keep it alive. + self._stream_owner = stream + + def _validate_reuse_authority(self, authority: object) -> None: + if authority is None: + raise TypeError("generation reuse authority cannot be None") + if getattr(authority, "collective_safe", False) is not True: + raise ValueError("generation reuse authority lacks collective safety") + if getattr(authority, "producer_reuse_guarded", False) is not True: + raise ValueError("generation reuse authority does not guard producer reuse") + if getattr(authority, "live_bank_count", None) != self.live_bank_count: + raise ValueError("generation reuse authority must target one live bank") + copy_module = getattr(authority, "bound_copy_module", None) + if copy_module is not self and getattr(copy_module, "backend", None) is not self: + raise ValueError( + "generation reuse authority is not bound to this hierarchical endpoint" + ) + if copy_module is not self: + config = getattr(copy_module, "config", None) + expected = { + "ep_size": int(self.comm.world), + "logical_expert_count": self.global_expert_count, + "extra_slots_per_rank": self.helper_count, + "local_rank": int(self.comm.rank), + } + if config is None or any( + type(getattr(config, name, None)) is not int + or getattr(config, name) != value + for name, value in expected.items() + ): + raise ValueError("generation reuse authority geometry is inconsistent") + + def bind_generation_reuse_authority(self, authority: object) -> None: + """Bind the exact collective single-bank lease provider once.""" + + self._validate_reuse_authority(authority) + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if self._reuse_authority is not None: + raise RuntimeError("generation reuse authority is single-assignment") + if self._active_generation: + raise RuntimeError("generation reuse authority must bind before submit") + self._reuse_authority = authority + + def release_generation_after( + self, generation: int, completion_event: object | None = None + ) -> None: + """Order the copy stream after this rank's live-bank consumer. + + The consumer is ``completion_event``, or without it the work already + queued on the current stream. That stream must be the one the next + fused launch follows: the launch orders the copy stream after the work + queued on its current stream (see ``submit_fused``), which then + includes this consumer, so no separate device wait is queued for it. + This establishes only local stream ordering. The bound reuse authority + must separately attest its bounded all-rank barrier before another + generation may overwrite the single live bank. + """ + + if type(generation) is not int or generation <= 0: + raise ValueError("generation must be a positive exact int") + event_handle = 0 + if completion_event is not None: + event_handle = _event_handle(completion_event) + if event_handle <= 0: + raise ValueError("completion_event must expose a positive CUDA event") + with self._lock: + if self._handoff_aborted or self._scheduler is None: + raise RuntimeError("hierarchical scheduler handoff is not operational") + if generation != self._active_generation: + raise ValueError("generation is not the active live bank") + if self._local_release_generation >= generation: + raise RuntimeError("generation already has a local release event") + if completion_event is not None: + check_cuda( + self.cuda.cuStreamWaitEvent(self._stream, event_handle, 0), + "cuStreamWaitEvent(live-bank consumer)", + ) + self._local_release_generation = generation + + def mark_collective_reuse_safe( + self, authority: object, generation: int + ) -> None: + """Accept the bound provider's all-rank proof for one generation.""" + + if type(generation) is not int or generation <= 0: + raise ValueError("generation must be a positive exact int") + with self._lock: + if self._handoff_aborted or self._scheduler is None: + raise RuntimeError("hierarchical scheduler handoff is not operational") + # bind_generation_reuse_authority() validated this single-assignment + # authority, so identity suffices on the per-layer path. + if authority is not self._reuse_authority: + raise PermissionError("generation reuse authority identity mismatch") + if generation != self._active_generation: + raise ValueError("generation is not the active live bank") + if self._local_release_generation != generation: + raise RuntimeError("local consumer release is not registered") + if self._collective_safe_generation >= generation: + raise RuntimeError("generation is already collectively reusable") + self._collective_safe_generation = generation + + def _require_live_bank_reusable(self) -> None: + # The bound authority was validated once, when it was bound. + if self._reuse_authority is None: + raise RuntimeError("single live bank has no bound collective reuse authority") + if ( + self._active_generation + and self._collective_safe_generation < self._active_generation + ): + raise RuntimeError( + "single live bank is still leased by the previous generation" + ) + + def submit_fused( + self, + outputs: object, + routes: torch.Tensor, + route_target: int, + *, + spin_cycles: int, + ) -> int: + """Enqueue the bound scheduler and this GPU-direct copy as one kernel. + + The kernel runs on the bound stream after the work already queued on + the current stream, which must produce ``routes``. Each scheduler CTA + adds one to the bound route epoch once its routes are written, which + makes it reach ``route_target``; the copy and its READY publication + follow in the same CTAs. ``spin_cycles`` bounds the launch's + cross-rank waits and its wait for the route epoch. Returns the + launched generation. + """ + + with self._lock: + if self._handoff_aborted: + raise RuntimeError( + "hierarchical handoff aborted; quiesced group rebuild required" + ) + if outputs is not self._outputs or self._scheduler is None: + raise RuntimeError("outputs are not from the bound CUDA scheduler") + self._require_live_bank_reusable() + try: + torch.ops.trtllm.moe_rebalance_tma_submit_fused( + self._handle, routes, route_target, spin_cycles + ) + self._active_generation = self._generation + except BaseException: + self._handoff_aborted = True + raise + generation = self._generation + self._generation = generation + 1 + return generation + + def close(self) -> None: + """Release the native copy context after all device work is drained.""" + with self._lock, torch.cuda.device(self.device): + if self.ctx is None: + return + self._handle = 0 + torch.ops.trtllm.moe_rebalance_tma_destroy(self.ctx) + self.ctx = None + self._scheduler = None + self._outputs = None + self._stream_owner = None + + def last_mode_counts(self) -> tuple[int, int, int]: + """Return diagnostic ``(direct_slots, scatter_slots, commands)``. + + GPU-direct mode explicitly waits for the copy and reads GPU counters. + This opt-in diagnostic must stay outside the submission/timing path. + """ + + result = torch.ops.trtllm.moe_rebalance_tma_result(self.ctx) + direct, scatter, commands = result[1], result[2], result[7] + return int(direct), int(scatter), int(commands) + + +__all__ = [ + "BoundHierarchicalLiveWeightArena", + "HierarchicalLiveWeightArena", + "HierarchicalLiveWeightArenaProvider", + "HierarchicalSamiWeightBroadcast", + "hierarchy_group_sizes", +] diff --git a/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py new file mode 100644 index 000000000000..a7b7e477a972 --- /dev/null +++ b/tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py @@ -0,0 +1,706 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Single-expert use of MegaMoE's complete fused FC1/activation/FC2 kernel.""" + +import weakref +from itertools import product +from threading import Condition, RLock + +import torch + +from ...autotuner import AutoTuner, DynamicTensorSpec, TunableRunner, TuningConfig + +__all__ = [ + "release_shared_fc12_cache", + "run_shared_fc12", + "shared_fc12_cache_info", + "shared_fc12_candidate_tactics", + "shared_fc12_tactic_descriptor", +] + +_SHARED_FC12_OP = "trtllm::megamoe_shared_fc12" +_CATALOG_VERSION = 1 +_SharedFc12Tactic = tuple[int, int, int, str, bool] +_CANDIDATES = tuple( + product((1, 2), (64, 128, 256), (128, 256), ("grid_stride", "atomic_counter"), (False, True)) +) + + +def shared_fc12_candidate_tactics() -> tuple[_SharedFc12Tactic, ...]: + """Finite common search space; larger clusters/K and extra TMA stages are not searched.""" + return _CANDIDATES + + +def shared_fc12_tactic_descriptor(tactic: _SharedFc12Tactic) -> dict: + tactic = tuple(tactic) + if tactic not in _CANDIDATES: + raise ValueError(f"Unsupported shared FC12 tactic: {tactic}") + ctas, tile_n, tile_k, work_id_mode, fc2_use_bulk = tactic + return { + "tactic": list(tactic), + "mma_tiler_mnk": [128 * ctas, tile_n, tile_k], + "cluster_shape_mn": [ctas, 1], + "use_2cta_instrs": ctas == 2, + "work_id_mode": work_id_mode, + "fc2_use_bulk": fc2_use_bulk, + "fc2_tma_stages": 1, + "fc1_epi_flag_batch": 1, + "fc2_epi_flag_batch": 1, + } + + +def _token_buckets(max_tokens: int) -> tuple[int, ...]: + if max_tokens <= 0: + raise ValueError("shared FC12 capacity must be positive") + return tuple( + sorted( + {n for n in (1, 128, 256, 512, 1024, 2048, 4096, 8192) if n < max_tokens} | {max_tokens} + ) + ) + + +def _make_kernel( + device_index, hidden_size, intermediate_size, max_tokens, sm_count, swiglu_limit, tactic +): + import cutlass.utils as utils + from cutlass.cute.nvgpu import OperandMajorMode + + from .cutedsl_megamoe.api import ImplDesc, ProblemDesc + from .cutedsl_megamoe.kernel_src.blackwell.inference.mega.block_scaled_swap_ab_fc12_kernel import ( + BlockScaledSwapAbFc12Kernel, + ) + from .cutedsl_megamoe.quant_def import CombineFormat + + properties = torch.cuda.get_device_properties(device_index) + if not 0 < sm_count <= properties.multi_processor_count: + raise ValueError("shared FC12 sm_count must be within the device SM count") + descriptor = shared_fc12_tactic_descriptor(tactic) + cluster_size = descriptor["cluster_shape_mn"][0] + hardware_clusters = utils.HardwareInfo(device_index).get_max_active_clusters(cluster_size) + launch_clusters = min(sm_count // cluster_size, hardware_clusters) + if launch_clusters <= 0: + raise ValueError("shared FC12 requires at least one active cluster") + problem = ProblemDesc( + { + "expert_count": 1, + "intermediate_gateup_size": 2 * intermediate_size, + "hidden_size": hidden_size, + "quant_kind": "mxfp8_e4m3", + "a_major_mode": OperandMajorMode.K, + "b_major_mode": OperandMajorMode.K, + "combine_format": CombineFormat.parse("bf16"), + "gate_up_clamp": swiglu_limit, + } + ) + implementation = ImplDesc( + { + **{key: value for key, value in descriptor.items() if key != "tactic"}, + "mma_tiler_mnk": tuple(descriptor["mma_tiler_mnk"]), + "cluster_shape_mn": tuple(descriptor["cluster_shape_mn"]), + "group_hint": launch_clusters, + "token_padding_block": 64, + "sf_padding_block": 128, + "max_tokens": max_tokens, + "launch_cluster_count": launch_clusters, + } + ) + kernel = BlockScaledSwapAbFc12Kernel(problem, implementation) + return kernel, implementation, properties, hardware_clusters, launch_clusters + + +class _SharedFc12Runner: + def __init__( + self, + device_index: int, + hidden_size: int, + intermediate_size: int, + max_tokens: int, + sm_count: int, + swiglu_limit: float | None, + stream_handle: int, + *, + tactic: _SharedFc12Tactic | None = None, + ) -> None: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_ptr + from cutlass.cute.typing import AddressSpace + + device = torch.device("cuda", device_index) + properties = torch.cuda.get_device_properties(device) + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Create the shared FC12 runner before CUDA graph capture") + self.hidden_size = hidden_size + self.max_tokens = max_tokens + self.reserved_sms = properties.multi_processor_count - sm_count + self.stream_handle = stream_handle + if tactic is None: + tactic = ( + (1, 256, 128, "atomic_counter", True) + if self.reserved_sms > 0 + else (2, 128, 128, "grid_stride", True) + ) + self.tactic = tuple(tactic) + kernel, implementation, properties, hardware_clusters, launch_clusters = _make_kernel( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + self.tactic, + ) + local_bytes, shared_bytes = kernel.get_workspace_sizes() + if shared_bytes: + raise RuntimeError("shared FC12 must not allocate a communication workspace") + self.workspace = torch.empty(local_bytes, dtype=torch.uint8, device=device) + local_zero_bytes, shared_zero_bytes = kernel.require_zero_workspace_leading_bytes + if shared_zero_bytes: + raise RuntimeError("shared FC12 must not initialize communication state") + self.workspace_pointer = self.workspace.data_ptr() + self.workspace_reset_bytes = local_zero_bytes + self._memset_async = cuda.cuMemsetD8Async + self._cuda_success = cuda.CUresult.CUDA_SUCCESS + self._cuda_stream = cuda.CUstream(stream_handle) + # A slice supplies the current E=1 token count without a per-call GPU write. + self.token_counts = torch.arange(max_tokens + 1, dtype=torch.int32, device=device) + self._token_count_views: dict[int, torch.Tensor] = {} + + token_rows = cute.sym_int64(divisibility=1) + scale_rows = cute.sym_int64(divisibility=128) + padded_gateup = (2 * intermediate_size + 127) // 128 * 128 + fc1_sf_columns = padded_gateup * (hidden_size // 32) + fc2_sf_columns = hidden_size * (intermediate_size // 32) + fake_arguments = { + "activation": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (token_rows, hidden_size), + stride_order=(1, 0), + assumed_align=16, + ), + "fc1_weight": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (1, hidden_size, 2 * intermediate_size), + stride_order=(2, 0, 1), + assumed_align=16, + ), + "activation_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, + (scale_rows, hidden_size // 32), + stride_order=(1, 0), + assumed_align=16, + ), + "fc1_weight_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (1, fc1_sf_columns), stride_order=(1, 0), assumed_align=16 + ), + "fc2_weight": make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (1, intermediate_size, hidden_size), + stride_order=(2, 0, 1), + assumed_align=16, + ), + "fc2_weight_sf": make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (1, fc2_sf_columns), stride_order=(1, 0), assumed_align=16 + ), + "fc2_output": make_fake_compact_tensor( + cutlass.BFloat16, + (token_rows, 1, hidden_size), + stride_order=(2, 1, 0), + assumed_align=16, + ), + "local_workspace": make_ptr(cutlass.Uint8, 0, AddressSpace.gmem, assumed_align=128), + "max_active_clusters": hardware_clusters, + "stream": make_fake_stream(), + "expert_token_sizes": make_fake_compact_tensor( + cutlass.Int32, (1,), stride_order=(0,), assumed_align=4 + ), + "expert_token_prefix_sum": None, + "topk_scores": None, + "fc1_alpha": None, + "fc2_alpha": None, + "fc1_norm_const": None, + } + self.compiled = cute.compile[cute.EnableTVMFFI(True)](kernel, **fake_arguments) + + def __call__( + self, + activation: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_sf: torch.Tensor, + ) -> torch.Tensor: + token_rows = activation.shape[0] + token_count = self._token_count_views.get(token_rows) + if token_count is None: + token_count = self._token_count_views[token_rows] = self.token_counts[ + token_rows : token_rows + 1 + ] + quantized, scales = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + activation, reserved_sms=self.reserved_sms + ) + output = torch.empty( + (token_rows, 1, self.hidden_size), dtype=torch.bfloat16, device=activation.device + ) + # The standalone FC12 readiness counters require reset before every invocation. + (status,) = self._memset_async( + self.workspace_pointer, 0, self.workspace_reset_bytes, self._cuda_stream + ) + if status != self._cuda_success: + raise RuntimeError(f"shared FC12 asynchronous counter reset failed: {status}") + self.compiled( + activation=quantized, + fc1_weight=fc1_weight.transpose(1, 2), + activation_sf=scales.view(torch.float8_e8m0fnu).view(-1, self.hidden_size // 32), + fc1_weight_sf=fc1_weight_sf.view(torch.float8_e8m0fnu), + fc2_weight=fc2_weight.transpose(1, 2), + fc2_weight_sf=fc2_weight_sf.view(torch.float8_e8m0fnu), + fc2_output=output, + local_workspace=self.workspace_pointer, + stream=self.stream_handle, + expert_token_sizes=token_count, + expert_token_prefix_sum=None, + topk_scores=None, + fc1_alpha=None, + fc2_alpha=None, + fc1_norm_const=None, + ) + return output.view(token_rows, self.hidden_size) + + +_SHARED_FC12_CACHE_LOCK = RLock() +_SHARED_FC12_CACHE_CONDITION = Condition(_SHARED_FC12_CACHE_LOCK) +_ACTIVE_FC12_CALLS: dict[tuple[int, int], int] = {} +_RELEASING_FC12_DEVICES: set[int] = set() +_RELEASING_FC12_STREAMS: set[tuple[int, int]] = set() +_RUNNER_CACHE: dict[tuple, _SharedFc12Runner] = {} +_RUNNER_CACHE_HITS = 0 +_RUNNER_CACHE_MISSES = 0 + + +def _begin_shared_fc12_release( + device_index: int, stream_handle: int | None +) -> tuple[int, int] | None: + stream_key = (device_index, stream_handle) if stream_handle is not None else None + with _SHARED_FC12_CACHE_CONDITION: + while ( + device_index in _RELEASING_FC12_DEVICES + or ( + stream_key is None + and any(device == device_index for device, _ in _RELEASING_FC12_STREAMS) + ) + or (stream_key is not None and stream_key in _RELEASING_FC12_STREAMS) + ): + _SHARED_FC12_CACHE_CONDITION.wait() + if stream_key is None: + _RELEASING_FC12_DEVICES.add(device_index) + else: + _RELEASING_FC12_STREAMS.add(stream_key) + while any( + count and device == device_index and (stream_handle is None or stream == stream_handle) + for (device, stream), count in _ACTIVE_FC12_CALLS.items() + ): + _SHARED_FC12_CACHE_CONDITION.wait() + return stream_key + + +def _end_shared_fc12_release(device_index: int, stream_key: tuple[int, int] | None) -> None: + with _SHARED_FC12_CACHE_CONDITION: + if stream_key is None: + _RELEASING_FC12_DEVICES.discard(device_index) + else: + _RELEASING_FC12_STREAMS.discard(stream_key) + _SHARED_FC12_CACHE_CONDITION.notify_all() + + +def _begin_shared_fc12_call(device_index: int, stream_handle: int) -> tuple[int, int]: + key = (int(device_index), int(stream_handle)) + with _SHARED_FC12_CACHE_CONDITION: + while key[0] in _RELEASING_FC12_DEVICES or key in _RELEASING_FC12_STREAMS: + _SHARED_FC12_CACHE_CONDITION.wait() + _ACTIVE_FC12_CALLS[key] = _ACTIVE_FC12_CALLS.get(key, 0) + 1 + return key + + +def _end_shared_fc12_call(key: tuple[int, int]) -> None: + with _SHARED_FC12_CACHE_CONDITION: + remaining = _ACTIVE_FC12_CALLS[key] - 1 + if remaining: + _ACTIVE_FC12_CALLS[key] = remaining + else: + del _ACTIVE_FC12_CALLS[key] + _SHARED_FC12_CACHE_CONDITION.notify_all() + + +def _get_runner( + device_index: int, + hidden_size: int, + intermediate_size: int, + max_tokens: int, + sm_count: int, + swiglu_limit: float | None, + stream_handle: int, + tactic: _SharedFc12Tactic | None = None, +) -> _SharedFc12Runner: + global _RUNNER_CACHE_HITS, _RUNNER_CACHE_MISSES + key = ( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + tactic, + ) + with _SHARED_FC12_CACHE_LOCK: + runner = _RUNNER_CACHE.get(key) + if runner is not None: + _RUNNER_CACHE_HITS += 1 + return runner + _RUNNER_CACHE_MISSES += 1 + with torch.cuda.device(device_index): + runner = _SharedFc12Runner( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + tactic=tactic, + ) + _RUNNER_CACHE[key] = runner + return runner + + +def shared_fc12_cache_info() -> dict[str, int | None]: + """Return process-cache counters for diagnostics.""" + with _SHARED_FC12_CACHE_LOCK: + return { + "hits": _RUNNER_CACHE_HITS, + "misses": _RUNNER_CACHE_MISSES, + "maxsize": None, + "currsize": len(_RUNNER_CACHE), + } + + +class SharedFc12TunableRunner(TunableRunner): + # AutoTuner uses this dictionary to prime persisted winners during model warmup. + kernel_cache: dict[tuple, _SharedFc12Runner] = {} + + def __init__( + self, + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + input_signature, + ): + properties = torch.cuda.get_device_properties(device_index) + self.device_index = device_index + self.sm_count = sm_count + self.reserved_sms = properties.multi_processor_count - sm_count + self.max_tokens = max_tokens + self.stream_handle = stream_handle + self._runner_args = ( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + ) + # The AutoTuner keys its cache by str(unique_id()) on every dispatch, + # so build that string once. + self._unique_id = str( + ( + _CATALOG_VERSION, + device_index, + properties.major, + properties.minor, + properties.multi_processor_count, + sm_count, + hidden_size, + intermediate_size, + max_tokens, + swiglu_limit, + "K128_input_K32_intermediate_R128c4_gateup16", + input_signature, + ) + ) + self.buckets = _token_buckets(max_tokens) + self.tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, + dim_idx=0, + gen_tuning_buckets=self.buckets, + map_to_tuning_buckets=self.token_bucket, + ), + ), + tune_max_num_tokens=max_tokens, + use_cold_l2_cache=True, + use_cuda_graph=False, + ) + self._valid_tactics = None + self._compiled_runners = {} + self._compile_lock = RLock() + + def unique_id(self): + return self._unique_id + + def token_bucket(self, tokens: int) -> int: + if not 0 < tokens <= self.max_tokens: + raise ValueError("shared FC12 token count exceeds its configured capacity") + return next(bucket for bucket in self.buckets if tokens <= bucket) + + def get_valid_tactics(self, inputs, profile, **kwargs): + if self._valid_tactics is None: + valid = [] + with torch.cuda.device(self.device_index): + for tactic in shared_fc12_candidate_tactics(): + try: + _make_kernel(*self._runner_args[:6], tactic) + except (ValueError, AssertionError): + continue + else: + valid.append(tactic) + self._valid_tactics = tuple(valid) + return list(self._valid_tactics) + + def _cached_winner_tactics(self, tuner, inputs): + """Return every profiled bucket winner, or None if any is missing.""" + if tuner.skip_dynamic_tuning_buckets: + return None + winners = set() + profiles = tuner._optimization_profiles(self.tuning_config, inputs) + if not profiles: + return None + for profile in profiles: + hit, runner_id, tactic, _ = tuner.profiling_cache.search_cache( + _SHARED_FC12_OP, + [self], + profile.get_opt_shapes(), + self.tuning_config, + apply_map_to_tuning_buckets=False, + ) + if not hit or runner_id != 0 or tactic == -1: + return None + tactic = tuple(tactic) + if tactic not in _CANDIDATES: + return None + winners.add(tactic) + return winners + + def _retain_cached_winners(self, winners) -> None: + """Drop profiled loser workspaces after all bucket winners are known.""" + winners = {tuple(tactic) for tactic in winners} + if not winners or not winners.issubset(_CANDIDATES): + return + with self._compile_lock: + stale = set(self._compiled_runners) - winners + if not stale: + return + # Every cached runner is bound to this stream. Profiling has + # finished, but synchronize before releasing its workspace. + torch.cuda.ExternalStream(self.stream_handle, device=self.device_index).synchronize() + with _SHARED_FC12_CACHE_LOCK: + for tactic in stale: + self._compiled_runners.pop(tactic, None) + key = (*self._runner_args, tactic) + self.kernel_cache.pop(key, None) + _RUNNER_CACHE.pop(key, None) + # Return the released loser workspaces to CUDA before the routed + # MegaMoE autotuner allocates its symmetric profiling scratch. + with torch.cuda.device(self.device_index): + torch.cuda.empty_cache() + + def forward(self, inputs, *, tactic=-1, do_preparation=False, **kwargs): + if do_preparation: + return None + if tactic == -1: + if AutoTuner.get().is_tuning_mode: + raise RuntimeError("shared FC12 profiling did not select a valid tactic") + tactic = ( + (1, 256, 128, "atomic_counter", True) + if self.reserved_sms > 0 + else (2, 128, 128, "grid_stride", True) + ) + tactic = tuple(tactic) + with self._compile_lock: + runner = self._compiled_runners.get(tactic) + if runner is None: + runner = _get_runner(*self._runner_args, tactic) + self._compiled_runners[tactic] = runner + self.kernel_cache[(*self._runner_args, tactic)] = runner + return runner(*inputs) + + +_tunable_runners: dict[tuple, SharedFc12TunableRunner] = {} + + +def release_shared_fc12_cache(device_index: int, stream_handle: int | None = None) -> None: + """Drain and release shared-FC12 runners for one executor stream. + + Dispatches already in progress finish before teardown. New dispatches for + the same stream wait, while unrelated streams and devices remain usable. + Passing no stream handle retains the process-exit fallback that releases + every runner on the device. + """ + device_index = int(device_index) + stream_handle = None if stream_handle is None else int(stream_handle) + stream_key = _begin_shared_fc12_release(device_index, stream_handle) + + def matches(key: tuple) -> bool: + return int(key[0]) == device_index and ( + stream_handle is None or int(key[6]) == stream_handle + ) + + try: + with torch.cuda.device(device_index): + if stream_handle is None: + torch.cuda.synchronize(device_index) + else: + torch.cuda.ExternalStream(stream_handle, device=device_index).synchronize() + with _SHARED_FC12_CACHE_CONDITION: + tunable_keys = [key for key in _tunable_runners if matches(key)] + for key in tunable_keys: + _tunable_runners.pop(key)._compiled_runners.clear() + kernel_keys = [key for key in SharedFc12TunableRunner.kernel_cache if matches(key)] + for key in kernel_keys: + SharedFc12TunableRunner.kernel_cache.pop(key, None) + runner_keys = [key for key in _RUNNER_CACHE if matches(key)] + for key in runner_keys: + _RUNNER_CACHE.pop(key, None) + finally: + _end_shared_fc12_release(device_index, stream_key) + + +# Validated weight layouts by the identity of the four weight tensors. A +# layer's weights stay fixed between loads, so each set is checked once. +_WEIGHT_LAYOUTS: dict[tuple[int, int, int, int], tuple] = {} + + +def _weight_layout(weights: tuple, hidden_size: int) -> tuple[int, torch.device, tuple]: + """Validate FC12 weights once per weight set. + + Returns the intermediate size, the weights' device, and the input + signature of a contiguous BF16 ``[T, hidden_size]`` activation with them. + """ + key = (id(weights[0]), id(weights[1]), id(weights[2]), id(weights[3])) + entry = _WEIGHT_LAYOUTS.get(key) + if entry is not None: + refs, layout_hidden_size, intermediate_size, device, signature = entry + if ( + layout_hidden_size == hidden_size + and refs[0]() is weights[0] + and refs[1]() is weights[1] + and refs[2]() is weights[2] + and refs[3]() is weights[3] + ): + return intermediate_size, device, signature + fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf = weights + if fc2_weight.ndim != 3 or fc2_weight.shape[:2] != (1, hidden_size): + raise ValueError("shared FC12 weight must have shape [1, H, I]") + intermediate_size = fc2_weight.shape[2] + if hidden_size % 128 or intermediate_size % 128: + raise ValueError("shared FC12 hidden and intermediate dimensions must be multiples of 128") + if fc1_weight.shape != (1, 2 * intermediate_size, hidden_size): + raise ValueError("shared FC1 weight must have shape [1, 2*I, H]") + for weight in (fc1_weight, fc2_weight): + if weight.dtype != torch.float8_e4m3fn or not weight.is_contiguous(): + raise ValueError("shared FC12 weights must be contiguous E4M3 tensors") + for scale in (fc1_weight_sf, fc2_weight_sf): + if scale.dtype != torch.uint8 or not scale.is_contiguous(): + raise ValueError("shared FC12 scales must be contiguous uint8 tensors") + device = fc1_weight.device + if any(t.device != device for t in weights): + raise ValueError("shared FC12 inputs and weights must be on the same CUDA device") + signature = ( + (str(torch.bfloat16), (None, hidden_size), (hidden_size, 1)), + *((str(t.dtype), tuple(t.shape), tuple(t.stride())) for t in weights), + ) + with _SHARED_FC12_CACHE_LOCK: + stale = [k for k, e in _WEIGHT_LAYOUTS.items() if any(ref() is None for ref in e[0])] + for k in stale: + del _WEIGHT_LAYOUTS[k] + _WEIGHT_LAYOUTS[key] = ( + tuple(weakref.ref(t) for t in weights), + hidden_size, + intermediate_size, + device, + signature, + ) + return intermediate_size, device, signature + + +def run_shared_fc12( + activation: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_weight_sf: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_weight_sf: torch.Tensor, + *, + swiglu_limit: float | None, + sm_count: int, + max_tokens: int = 8192, +) -> torch.Tensor: + """Run one complete FC12 kernel after the existing shared-expert input quantizer. + + ``activation`` is contiguous BF16 ``[T, H]``. FP8 weights have shapes + ``[1, 2*I, H]`` (gate/up interleaved in groups of 16) and ``[1, H, I]``. + Their uint8 scales are flattened R128c4 UE8M0 ``[1, padded_rows * K/32]``. + ``H`` and ``I`` must be divisible by 128; output is BF16 ``[T, H]``. + ``max_tokens`` is a fixed model capacity, not the current batch size. + """ + if activation.dtype != torch.bfloat16 or not activation.is_cuda: + raise ValueError("shared FC12 requires CUDA BF16 activations") + if activation.ndim != 2 or not activation.is_contiguous(): + raise ValueError("shared FC12 requires contiguous [T, H] activations") + tokens, hidden_size = activation.shape + if max_tokens <= 0 or tokens > max_tokens: + raise ValueError("shared FC12 token count exceeds its configured capacity") + weights = (fc1_weight, fc1_weight_sf, fc2_weight, fc2_weight_sf) + intermediate_size, device, input_signature = _weight_layout(weights, hidden_size) + if activation.device != device: + raise ValueError("shared FC12 inputs and weights must be on the same CUDA device") + if tokens == 0: + return torch.empty_like(activation) + device_index = device.index + stream_handle = torch._C._cuda_getCurrentRawStream(device_index) + inputs = [activation, *weights] + runner_key = ( + device_index, + hidden_size, + intermediate_size, + max_tokens, + sm_count, + swiglu_limit, + stream_handle, + input_signature, + ) + active_key = _begin_shared_fc12_call(device_index, stream_handle) + try: + runner = _tunable_runners.get(runner_key) + if runner is None: + with _SHARED_FC12_CACHE_LOCK, torch.cuda.device(device_index): + runner = _tunable_runners.get(runner_key) + if runner is None: + runner = _tunable_runners[runner_key] = SharedFc12TunableRunner(*runner_key) + tuner = AutoTuner.get() + selected_runner, tactic = tuner.choose_one( + _SHARED_FC12_OP, [runner], runner.tuning_config, inputs + ) + if tuner.is_tuning_mode: + winners = runner._cached_winner_tactics(tuner, inputs) + if winners is not None: + runner._retain_cached_winners(winners) + return selected_runner(inputs, tactic=tactic) + finally: + _end_shared_fc12_call(active_key) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 07fb708341f5..96c9923a08d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -41,7 +41,7 @@ from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric from ..moe.expert_statistic import ExpertStatistic from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup -from ..moe.fused_moe.moe_load_balancer import MoeLoadBalancer +from ..moe.fused_moe.moe_load_balancer import BaseMoeLoadBalancer from ..nccl_window_tensor_scope import install_eager_nccl_window_tensor_scopes from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture from ..speculative import SpecMetadata, update_spec_config_from_loaded_model @@ -255,7 +255,7 @@ def _load_model( model_weights_restore_mode, ) -> None: """Load the model, or take the given one, at the requested capacity.""" - self.moe_load_balancer: Optional[MoeLoadBalancer] = None + self.moe_load_balancer: Optional[BaseMoeLoadBalancer] = None if model is None: lora_config: Optional[LoraConfig] = self.llm_args.lora_config # Keep the model_loader to support reloading the model weights later @@ -273,7 +273,7 @@ def _load_model( # Open checkpoint and load the LLM module object. self.model, moe_load_balancer = self.model_loader.load( checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader) - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): self.moe_load_balancer = moe_load_balancer else: self.model = model @@ -1000,80 +1000,93 @@ def invalidate_multimodal_encoder_cache(self) -> None: if encoder_cache is not None: encoder_cache.clear() - def cleanup(self) -> None: - """Release resources owned by this model engine. - - Tears down, in order: - - 1. The optional ``ModelLoader`` (which in turn releases any - GMS client; see :meth:`ModelLoader.cleanup`). - 2. Runner resources and executor-owned CUDA Graph captures (via - :meth:`_release_cuda_graphs`). - 3. The runner, model caller, MM item scheduler, and model references. - 4. Input processors. - - Idempotency: - Subsequent calls are no-ops (guarded by ``_cleanup_done``). - The flag is set only at the end, so a partial cleanup that - raises mid-way will be retried on the next call. - - Called from: - :meth:`__del__`, and only from there. ``PyExecutor.shutdown`` - deliberately does *not* call this: it is also invoked mid-init by - ``configure_kv_cache_capacity``, which reads ``model`` right - afterwards, so clearing ``model`` here would break it. That path - uses an explicit temporary-executor shutdown and then drops its - reference instead. - """ + def prepare_cleanup(self, + *, + release_nccl_window_owners: bool = True) -> None: + """Release local graph captures before collective model teardown.""" + if self._cleanup_done: + return + self._release_cuda_graphs( + release_nccl_window_owners=release_nccl_window_owners) + + def shutdown_moe_load_balancer(self, *, local_safe: bool) -> bool: + """Run the manager's rank-aligned teardown preflight and collectives.""" + if self._cleanup_done: + return True + moe_load_balancer = getattr(self, "moe_load_balancer", None) + if moe_load_balancer is None: + return bool(local_safe) + return moe_load_balancer.shutdown( + local_safe=bool(local_safe)) is not False + + def finish_cleanup(self) -> None: + """Release rank-local model resources after collective teardown.""" if self._cleanup_done: return - # Cleanup is not truly atomic: released CUDA/GMS resources cannot be - # rolled back. Keep each handle live until its own release succeeds, - # so a failed cleanup can be retried without double-freeing resources - # that were already released. model_loader = self.model_loader if model_loader is not None: model_loader.cleanup() self.model_loader = None - # Release runner-owned graphs before dropping the runner. Keep the - # handle available if graph release fails and cleanup is retried. - # Destruction timing is not rank-symmetric; only explicit shutdown releases owners. - self._release_cuda_graphs(release_nccl_window_owners=False) - - # The runner, caller and scheduler retain the model, so - # clearing the engine's attribute alone would leave the weights - # reachable past `release_gc()` below. + self.moe_load_balancer = None + # The runner, caller and scheduler retain the model, so clearing the + # engine's attribute alone would leave the weights reachable. self._runner = None self._model_caller = None self._mm_item_scheduler = None self._lora_model = None self._lora_owner = None self.model = None - self.input_processor = None - - # Release model weights. release_gc() self._cleanup_done = True - def __del__(self) -> None: - """Best-effort cleanup during garbage collection. - - Delegates to :meth:`cleanup`. Catches ``RuntimeError`` (which a - release step such as :meth:`_release_cuda_graphs` or - ``ModelLoader.cleanup`` may raise) and ``AttributeError`` (typical - on partially-initialized engines torn down during interpreter - shutdown when module references have already been cleared); both - are logged and swallowed because destructors cannot reliably - surface exceptions. - - This is the only production caller of :meth:`cleanup` -- see the - note there on why ``PyExecutor.shutdown`` must not call it. + def cleanup(self, *, collective: bool = True) -> None: + """Release model resources without rank-asymmetric collectives. + + Explicit terminal cleanup first releases local graph captures on every + rank, then passes that local result into the load-balancer's collective + preflight, and only then drops rank-local model state. Python finalizers + perform local cleanup only; worker-owned terminal cleanup is responsible + for collective resources. """ + if self._cleanup_done: + return + + if not collective: + # Destruction timing is not rank-symmetric; only explicit shutdown + # releases NCCL window owners. + self.prepare_cleanup(release_nccl_window_owners=False) + moe_load_balancer = getattr(self, "moe_load_balancer", None) + if moe_load_balancer is not None and not getattr( + moe_load_balancer, "shutdown_on_finalize", False): + logger.warning( + "Skipping collective MoE load-balancer shutdown from Python finalization; " + "terminal executor shutdown must release it explicitly") + self.finish_cleanup() + return + + prepare_error = None + try: + self.prepare_cleanup() + except BaseException as error: # Preserve rank alignment before surfacing it. + prepare_error = error + + manager_safe = self.shutdown_moe_load_balancer( + local_safe=prepare_error is None) + if prepare_error is not None: + raise prepare_error + if not manager_safe: + raise RuntimeError( + "MoE load-balancer collective shutdown was not safe on every rank" + ) + self.finish_cleanup() + + def __del__(self) -> None: + """Best-effort rank-local cleanup during garbage collection.""" try: - self.cleanup() + self.cleanup(collective=False) except (RuntimeError, AttributeError) as e: logger.warning("ModelEngine cleanup failed during destruction: %s", e) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f967e92d7069..933ae2afc772 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,7 @@ get_registered_model_class, timing_metric) from ..modules.embedding import Embedding, LMHead from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm -from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, +from ..moe.fused_moe.moe_load_balancer import (BaseMoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope @@ -1045,7 +1045,7 @@ def load( self, checkpoint_dir: str, checkpoint_loader: BaseCheckpointLoader, - ) -> tuple[DecoderModelForCausalLM, MoeLoadBalancer | None]: + ) -> tuple[DecoderModelForCausalLM, BaseMoeLoadBalancer | None]: """ Loads the model, its weights, and applies necessary configurations. @@ -1569,7 +1569,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor): # `tensorrt_llm/llmapi/llm_args.py`. When implementing the # follow-up, drop the validator gate AFTER moving these calls # into the pool scope. - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): moe_load_balancer.register_weight_slots_after_to_cuda() logger.info("moe_load_balancer finalizing model...") moe_load_balancer.finalize_model() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..81fbab38f593 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -433,6 +433,9 @@ def _request_flag(request, name: str) -> bool: class PyExecutor: + # Shutdown publication is leader-only, while join/finalize runs on every rank. + shutdown_all_ranks = True + # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. # If the number of micro batches is too small, the executor will spend too much time in synchronization. @@ -2045,11 +2048,28 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: Model-owned graphs are terminal model resources, unlike executor-owned graphs that may reference the resource managers torn down below. """ + # Capture model owners before any rank-local shutdown step can fail; + # outer workers still need them for an all-rank unsafe preflight. + terminal_engines = list(getattr(self, "_terminal_model_engines", ())) + for engine in (getattr(self, "model_engine", + None), getattr(self, "draft_model_engine", + None)): + if engine is not None and all(engine is not item + for item in terminal_engines): + terminal_engines.append(engine) + self._terminal_model_engines = tuple(terminal_engines) + if getattr(self, "_shutdown_runtime_complete", False): + self._finish_shutdown() + return + node_server = getattr(self, "_openengine_node_server", None) if node_server is not None: node_server.close() self._openengine_node_server = None - self.executor_request_queue.enqueue_shutdown_request() + # Rank zero owns the request queue, but every rank receives the + # broadcast shutdown sentinel and must execute the common finalize path. + if self.can_shutdown(): + self.executor_request_queue.enqueue_shutdown_request() self.shutdown_event.wait() if (self.enable_attention_dp and self.kv_connector_manager is not None and self.kv_connector_manager.get_pending_transfer_requests()): @@ -2073,9 +2093,11 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: logger.error("Hang detected, shutting down immediately.") # Stop the relay thread and the native progress thread from # probing while the hang detector aborts MPI. A relay-thread probe - # already in flight is not interrupted. + # already in flight is not interrupted. Do not let the outer worker + # enter model collectives on a broken rank set. self._pp_mpi_progress_stop.set() self._stop_mpi_progress_thread() + self._terminal_model_engines = () return self.worker_thread.join() if self.kv_connector_manager is not None and not self._kv_connector_failed: @@ -2100,10 +2122,13 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: encoder_launch_executor.shutdown(wait=True) self.encoder_launch_executor = None - # Stop the sampler's async worker, if it was used - if (isinstance(self.sampler, AsyncWorkerMixin) - and self.sampler.async_worker_enabled()): - self.sampler.async_worker_stop() + # Stop the sampler's async worker, if it was used. A retried + # shutdown stops it only once. + if not getattr(self, "_shutdown_sampler_complete", False): + if (isinstance(self.sampler, AsyncWorkerMixin) + and self.sampler.async_worker_enabled()): + self.sampler.async_worker_stop() + self._shutdown_sampler_complete = True self.worker_started = False if self._kv_connector_failed: @@ -2118,10 +2143,9 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: # empty_cache), the subsequent CUDA graph teardown can trigger a # device-wide cudaErrorIllegalAddress when the driver touches metadata # for the now-freed memory regions. - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None: - engine.release_executor_resources( - release_model_owned_graphs=release_model_owned_graphs) + for engine in self._terminal_model_engines: + engine.release_executor_resources( + release_model_owned_graphs=release_model_owned_graphs) # Ensure graph destruction has fully completed on device before # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): @@ -2129,26 +2153,108 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: for manager in self.resource_manager.resource_managers.values(): if manager: manager.shutdown() - # Note: do NOT call engine.cleanup() here. PyExecutor.shutdown() is - # paired with shutdown_for_kv_cache_estimation() mid-init in - # configure_kv_cache_capacity(); that path immediately reads - # model_engine.model.model_config - # to compute kv_cache_max_memory. cleanup() would set - # model_engine.model = None, breaking that read with - # `'NoneType' object has no attribute 'model_config'`. - # The engine's __del__ still calls cleanup() at terminal teardown - # (when the executor's reference is dropped), which is sufficient for - # the GMS daemon registry eviction the cleanup hook was added for. - del self.model_engine - if self.draft_model_engine is not None: - del self.draft_model_engine - if self.virtual_memory_pools is not None: - keys = list(self.virtual_memory_pools.keys()) + # Retry a later rank-local failure without repeating request + # publication, graph release, or resource-manager shutdown. + self._shutdown_runtime_complete = True + self._finish_shutdown() + + def _finish_shutdown(self): + virtual_memory_pools = getattr(self, "virtual_memory_pools", None) + if virtual_memory_pools is not None: + keys = list(virtual_memory_pools.keys()) for key in keys: - del self.virtual_memory_pools[key] - if self.dwdp_manager is not None: - self.dwdp_manager.__exit__(None, None, None) + del virtual_memory_pools[key] + dwdp_manager = getattr(self, "dwdp_manager", None) + if dwdp_manager is not None: + dwdp_manager.__exit__(None, None, None) self.dwdp_manager = None + # Keep model owners reachable until every rank-local step succeeds. The + # outer worker runs collective terminal cleanup from the snapshot above. + self.model_engine = None + self.draft_model_engine = None + + def terminal_cleanup(self, *, local_safe: bool = True) -> None: + """Collectively release model resources after all worker threads stop. + + Every rank first completes local graph preparation for every engine, + then participates in every manager preflight/collective in the same + order. Rank-local model and FC12 cache release happens only afterwards. + """ + engines = getattr(self, "_terminal_model_engines", ()) + if not engines: + return + + prepare_errors = [] + if local_safe: + for engine in engines: + try: + engine.prepare_cleanup() + except BaseException as error: + prepare_errors.append(error) + + collective_errors = [] + collective_safe = bool(local_safe and not prepare_errors) + for engine in engines: + try: + manager_safe = engine.shutdown_moe_load_balancer( + local_safe=collective_safe) + except BaseException as error: + collective_errors.append(error) + collective_safe = False + else: + collective_safe = collective_safe and manager_safe + + if prepare_errors: + raise prepare_errors[0] + if collective_errors: + raise collective_errors[0] + if not collective_safe: + raise RuntimeError( + "MoE load-balancer collective shutdown was not safe on every rank" + ) + + local_errors = [] + for engine in engines: + try: + engine.finish_cleanup() + except BaseException as error: + local_errors.append(error) + + # FC12 runners are rank-local and may be released only after every + # model-owned collective resource has been closed. + try: + self.release_rank_local_caches() + except BaseException as error: + local_errors.append(error) + + if local_errors: + raise local_errors[0] + self._terminal_model_engines = () + + def local_cleanup(self) -> None: + """Rank-local teardown for executors without collective EP resources. + + Drops the terminal model snapshot so the engines release their model + state through their own finalizers, as before collective cleanup. + """ + try: + self.release_rank_local_caches() + finally: + self._terminal_model_engines = () + + def release_rank_local_caches(self) -> None: + """Release this executor stream's shared-FC12 runners, if any were built.""" + shared_fc12 = sys.modules.get( + "tensorrt_llm._torch.moe.kernels.megamoe_shared_fc12") + if shared_fc12 is None: + return + + device = self.execution_stream.device + device_index = device.index if hasattr(device, "index") else int(device) + shared_fc12.release_shared_fc12_cache( + device_index, + stream_handle=int(self.execution_stream.cuda_stream), + ) def can_enqueue_requests(self) -> bool: """ diff --git a/tensorrt_llm/_torch/weight_sharing/source_identity.py b/tensorrt_llm/_torch/weight_sharing/source_identity.py index 4fa8305d5f92..7ede164b78c6 100644 --- a/tensorrt_llm/_torch/weight_sharing/source_identity.py +++ b/tensorrt_llm/_torch/weight_sharing/source_identity.py @@ -376,7 +376,8 @@ def _build_backend_fingerprint(model_config: "ModelConfig") -> str: Returns: A hex digest covering attention/MoE backends, allowed GEMM - backends, fusion flags, and the all-reduce strategy. + backends, fusion flags, the all-reduce strategy, and -- only when + one is configured -- the MoE load-balancer config. """ payload = { "attn_backend": getattr(model_config, "attn_backend", None), @@ -401,6 +402,9 @@ def _build_backend_fingerprint(model_config: "ModelConfig") -> str: "use_cute_dsl_bf16_bmm": getattr(model_config, "use_cute_dsl_bf16_bmm", False), "use_cute_dsl_bf16_gemm": getattr(model_config, "use_cute_dsl_bf16_gemm", False), } + load_balancer = getattr(model_config, "moe_load_balancer", None) + if load_balancer is not None: + payload["moe_load_balancer"] = load_balancer.model_dump(mode="python") return _canonical_hash(payload) @staticmethod diff --git a/tensorrt_llm/bench/benchmark/low_latency.py b/tensorrt_llm/bench/benchmark/low_latency.py index 25fa2ee9e828..96a142ad2932 100644 --- a/tensorrt_llm/bench/benchmark/low_latency.py +++ b/tensorrt_llm/bench/benchmark/low_latency.py @@ -27,7 +27,8 @@ generate_json_report, get_general_cli_options, get_llm) from tensorrt_llm.bench.benchmark.utils.asynchronous import async_benchmark -from tensorrt_llm.bench.benchmark.utils.general import generate_warmup_dataset +from tensorrt_llm.bench.benchmark.utils.general import (generate_warmup_dataset, + measurement_gc) from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -339,7 +340,7 @@ def latency_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(): + with iteration_writer.capture(), measurement_gc(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/throughput.py b/tensorrt_llm/bench/benchmark/throughput.py index a28d6e00a2b2..22a8f55fb53e 100755 --- a/tensorrt_llm/bench/benchmark/throughput.py +++ b/tensorrt_llm/bench/benchmark/throughput.py @@ -36,7 +36,8 @@ get_settings) # isort: on from tensorrt_llm.bench.benchmark.utils.general import ( - generate_warmup_dataset, update_sampler_args_with_extra_options) + generate_warmup_dataset, measurement_gc, + update_sampler_args_with_extra_options) from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataclasses.reporting import ReportUtility @@ -494,7 +495,7 @@ def throughput_command( logger.info("Warmup done.") iteration_writer = options.iteration_writer - with iteration_writer.capture(): + with iteration_writer.capture(), measurement_gc(): statistics = asyncio.run( async_benchmark(llm, sampling_params, diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index cf567b89a68b..e98f90337c48 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -1,8 +1,11 @@ from __future__ import annotations +import gc +import os +from contextlib import contextmanager from pathlib import Path from random import choices, shuffle -from typing import Dict, List, Union +from typing import Dict, Iterator, List, Union import yaml @@ -168,6 +171,32 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, } +@contextmanager +def measurement_gc() -> Iterator[None]: + """Keep Python's cyclic GC out of the measured run if TRTLLM_BENCH_DISABLE_GC=1. + + The benchmark process submits every request and handles every response. A + full collection over its in-flight and finished requests can stop it long + enough for the executors to run out of requests. This collects once, + freezes the surviving objects and disables the GC until the measured run + ends, as TRTLLM_SERVER_DISABLE_GC and TRTLLM_WORKER_DISABLE_GC disable it + in the server and the workers. + """ + if os.getenv("TRTLLM_BENCH_DISABLE_GC", "0") != "1": + yield + return + gc.collect() + gc.freeze() + gc.disable() + logger.info("Python GC disabled for the measured run " + "(TRTLLM_BENCH_DISABLE_GC=1).") + try: + yield + finally: + gc.enable() + gc.unfreeze() + + def generate_warmup_dataset(requests, steps) -> List[InferenceRequest]: """Warm up the benchmarker.""" warm_up_dataset = choices(requests, k=steps) diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 4cf9d669ef72..4d8f7450c3e0 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -16,6 +16,7 @@ import enum import gc import json +import os import time import traceback import uuid @@ -35,6 +36,9 @@ from .._utils import (global_mpi_rank, global_mpi_size, mpi_rank, nvtx_range_debug) from ..bindings import executor as tllm +from ..llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, + uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase from ..llmapi.tracer import global_tracer @@ -97,6 +101,18 @@ def __init__( tokenizer: Optional[TokenizerBase] = None, llm_args: Optional[BaseLlmArgs] = None, ) -> None: + # Direct worker construction also bypasses BaseLLM/worker_main. + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) + if llm_args is not None and llm_args.env_overrides: + os.environ.update(llm_args.env_overrides) + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + postproc_config = postproc_worker_config or PostprocWorkerConfig() super().__init__( num_postprocess_workers=postproc_config.num_postprocess_workers, @@ -1061,8 +1077,60 @@ def wakeup(self, wakeup_tags: list[str]) -> None: def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return + self.doing_shutdown = True + + try: + if self.engine is not None: + engine = self.engine + engine_shutdown_completed = bool( + getattr(self, "_engine_shutdown_completed", False)) + can_shutdown = getattr(engine, "can_shutdown", None) + if can_shutdown is None: + can_shutdown = engine.can_enqueue_requests + # PyExecutor separates leader-only shutdown publication from the + # all-rank join/finalize path. Other backends retain their gate. + # A retry must bypass that gate once shutdown already completed. + if (engine_shutdown_completed + or getattr(engine, "shutdown_all_ranks", False) + or can_shutdown()): + shutdown_error = None + local_safe = engine_shutdown_completed + if not local_safe: + try: + engine.shutdown() + self._engine_shutdown_completed = True + local_safe = True + except BaseException as error: + shutdown_error = error + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + try: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + if shutdown_error is not None: + raise shutdown_error + self.engine = None + self._engine_shutdown_completed = False + except BaseException: + # Keep the engine reachable so a caller can retry a failed + # collective terminal cleanup in this long-lived worker process. + # Only an explicit call retries it; see __del__. + self.doing_shutdown = False + self._collective_shutdown_failed = True + raise + self._collective_shutdown_failed = False + + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True if self.engine is not None: can_shutdown = getattr(self.engine, "can_shutdown", None) @@ -1070,6 +1138,9 @@ def shutdown(self): can_shutdown = self.engine.can_enqueue_requests if can_shutdown(): self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() self.engine = None def get_disaggregated_params(self) -> dict: @@ -1247,6 +1318,10 @@ def __exit__(self, exc_type, exc_value, traceback) -> bool: return exc_type is None or exc_type == self.WorkerExit def __del__(self): + # Finalizers run at a different time on each rank, so they must not + # retry a failed collective shutdown. + if getattr(self, "_collective_shutdown_failed", False): + return self.shutdown() diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ca28c4960bbe..ee3da08c4977 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -13,15 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. import atexit +import collections import concurrent.futures import json import os +import pickle # nosec B403 import shutil +import sys import tempfile import threading import time import weakref -from queue import Empty +from queue import Empty, Queue from typing import Dict, List, Optional, Union import torch @@ -100,9 +103,132 @@ def _check_collective_rpc_guard( "this shim only reaches rank-0.") +class _RequestSender: + """Batch a proxy's requests to the rank-0 worker on one background thread. + + Used only with ``TRTLLM_PROXY_BATCH_REQUESTS=1``. Callers pickle each + request themselves and ``send`` the bytes, so a request is captured as it + is when it is submitted, and one that cannot be pickled fails in its + caller. The sender thread sends every payload queued since its previous + send as one list in one HMAC-signed message, so a burst of requests costs + the rank-0 worker one receive and one HMAC check instead of one per + request; the worker then unpickles each payload. The thread starts on + the first ``send``. From then on it is the only thread that writes to the + request queue's socket, which is not thread-safe, and it also sends the + worker's shutdown sentinel. + + Trade-off: the sender thread must take the GIL before it can send. While + the submitting thread keeps running Python, for example a busy + trtllm-serve event loop, a lone request can wait up to the interpreter's + switch interval (5 ms by default) before it is sent. + """ + + MAX_BATCH_SIZE = 1024 + + def __init__(self, request_queue: IpcQueue, error_queue: Queue): + self._request_queue = request_queue + self._error_queue = error_queue + self._pending = collections.deque() + self._cond = threading.Condition() + self._closed = False + self._send_sentinel = False + # Set by the thread once it has drained the queue and will stop. + self._stopping = False + self._thread: Optional[threading.Thread] = None + # Bound here because ``close`` can run from a finalizer during + # interpreter teardown, when module globals may already be None. + self._is_finalizing = sys.is_finalizing + + def send(self, payload: bytes) -> None: + """Queue a pickled request for the worker. + + Payloads sent after ``close`` are dropped: the proxy is shutting down. + """ + with self._cond: + if self._closed: + return + if self._thread is None: + self._thread = threading.Thread(target=self._run, + daemon=True, + name="proxy_request_sender") + self._thread.start() + self._pending.append(payload) + self._cond.notify() + + def close(self, + timeout: Optional[float] = None, + *, + send_sentinel: bool = False) -> bool: + """Send everything queued, then stop the sender thread. + + With ``send_sentinel``, the thread also sends the worker's shutdown + sentinel after the last payload, so the request queue keeps a single + writer. Calls after the thread has stopped return at once. During + interpreter finalization this returns without doing anything: the + daemon thread can no longer run, the ZMQ context may be gone, and + joining a thread there fails. + + Returns: + Whether the caller still has to send the sentinel: true only if + ``send_sentinel`` is set and the thread did not send it, because it + never started or had already stopped. + """ + if self._is_finalizing(): + return False + with self._cond: + self._closed = True + thread = self._thread + caller_sends_sentinel = send_sentinel and (thread is None + or self._stopping) + self._send_sentinel = self._send_sentinel or send_sentinel + self._cond.notify() + if thread is not None: + thread.join(timeout) + if not thread.is_alive(): + self._thread = None + return caller_sends_sentinel + + def _run(self) -> None: + while True: + with self._cond: + while not self._pending and not self._closed: + self._cond.wait() + if not self._pending: + self._stopping = True + send_sentinel = self._send_sentinel + break + batch = [ + self._pending.popleft() + for _ in range(min(len(self._pending), self.MAX_BATCH_SIZE)) + ] + try: + self._request_queue.put(batch) + except Exception as e: + # The payloads are already pickled, so this is a transport + # error, which the proxy treats as fatal. + self._error_queue.put_nowait(e) + if send_sentinel: + # A bare None, not a batch: the worker stops at a top-level None. + try: + self._request_queue.put_noblock(None, retry=4) + except Exception as e: + self._error_queue.put_nowait(e) + + +def _create_request_sender(request_queue: IpcQueue, + error_queue: Queue) -> Optional[_RequestSender]: + """Return a batching request sender if TRTLLM_PROXY_BATCH_REQUESTS=1.""" + if os.getenv("TRTLLM_PROXY_BATCH_REQUESTS", "0") != "1": + return None + return _RequestSender(request_queue, error_queue) + + class GenerationExecutorProxy(GenerationExecutor): READY_SIGNAL = b"READY" WORKER_PROCESS_IDENTITIES_SIGNAL = b"WORKER_PROCESS_IDENTITIES" + # Set when TRTLLM_PROXY_BATCH_REQUESTS=1; requests are put on + # ``request_queue`` directly otherwise. + _request_sender: Optional[_RequestSender] = None # Total wait for the worker futures in an orderly shutdown. Generous # enough for the teardown of a large engine. _WORKER_SHUTDOWN_TIMEOUT_S = 180.0 @@ -258,6 +384,9 @@ def __init__( # ``abort_request()`` on user-caller threads — that pre-existing # cross-thread access is out of scope for this change. The fix # here narrowly targets the new profile-control path. + # With TRTLLM_PROXY_BATCH_REQUESTS=1, the ``_request_sender`` thread + # writes the queue instead once it has started, shutdown sentinel + # included. self._profile_control_executor = concurrent.futures.ThreadPoolExecutor( max_workers=1, thread_name_prefix="proxy_profile_control") @@ -487,6 +616,8 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: is_server=False, socket_type=zmq.PUSH, name="proxy_request_queue") + self._request_sender = _create_request_sender( + self.request_queue, self._error_queue) self.result_queue = FusedIpcQueue(frontend_result_addrs[0], is_server=True, fuse_message=False, @@ -496,6 +627,8 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: request_addr = None self.request_queue = IpcQueue(is_server=True, name="proxy_request_queue") + self._request_sender = _create_request_sender( + self.request_queue, self._error_queue) # TODO[chunweiy]: Unify IpcQueue and FusedIpcQueue # Use PULL mode when enable_postprocess_parallel as there are # multiple senders from multiple processes. @@ -584,7 +717,7 @@ def abort_request(self, request_id: int) -> None: # NOTE, it just sends a cancelling request to the request queue, but it # may take a while for the request to be cancelled in the worker and # send back a finished result. - self.request_queue.put(CancellingRequest(request_id)) + self._send_request(CancellingRequest(request_id)) # Per-call timeouts for the synchronous worker round-trip. The stop # path can take up to 30s because PyExecutor.stop_profile itself @@ -662,10 +795,12 @@ def _profile_control_call(self, kind: str, request, timeout: float) -> None: Runs on ``_profile_control_executor`` so all profile-related ZMQ socket ops (``request_queue.put`` and ``profile_ack_queue.get`` - via ``_wait_profile_ack``) happen on a single owning thread. + via ``_wait_profile_ack``) happen on a single owning thread. With + TRTLLM_PROXY_BATCH_REQUESTS=1 the request is put by + ``_request_sender`` instead. """ request.ack_addr = self._ensure_profile_ack_queue() - self.request_queue.put(request) + self._send_request(request) self._wait_profile_ack(kind, timeout) def _submit_profile_control(self, kind: str, request, @@ -1047,7 +1182,35 @@ def pre_shutdown(self): # empty list and therefore covered case (a) by accident. Switching # to bare ``any(...)`` to better handle partial crashes regressed # case (a); keep both behaviours explicit. - if not self.mpi_futures or any(not f.done() for f in self.mpi_futures): + send_sentinel = not self.mpi_futures or any(not f.done() + for f in self.mpi_futures) + self._close_request_sender(send_sentinel=send_sentinel) + + def _send_request(self, request) -> None: + """Send a request to the rank-0 worker. + + With TRTLLM_PROXY_BATCH_REQUESTS=1 the request is pickled here, on the + caller's thread, and ``_request_sender`` batches the bytes; otherwise + it is put on ``request_queue`` directly. Either way the worker receives + the request as it is now, and a request that cannot be pickled raises + here. + """ + if self._request_sender is None: + self.request_queue.put(request) + else: + self._request_sender.send(pickle.dumps(request)) + + def _close_request_sender(self, send_sentinel: bool = False) -> None: + """Stop the request sender, if any, then send the shutdown sentinel if asked. + + The sentinel follows every queued request. It is sent by the thread + that has been writing ``request_queue``: the sender thread once it has + started, the caller otherwise. + """ + if self._request_sender is not None: + send_sentinel = self._request_sender.close( + timeout=5.0, send_sentinel=send_sentinel) + if send_sentinel: self.request_queue.put_noblock(None, retry=4) def _get_next_client_id(self) -> int: @@ -1072,6 +1235,7 @@ def _cleanup_multi_frontend_ipc_dir(self): def shutdown(self): if not self.workers_started: + self._close_request_sender() self._cleanup_multi_frontend_ipc_dir() return @@ -1210,7 +1374,7 @@ def submit(self, request: GenerationRequest) -> GenerationResult: raise EngineDeadError(self._fatal_error) with nvtx_range_debug("request_queue.put"): - self.request_queue.put(request) + self._send_request(request) self._handle_background_error() @@ -1526,6 +1690,8 @@ def __init__( is_server=False, socket_type=zmq.PUSH, name=f"frontend_{frontend_id}_request_queue") + self._request_sender = _create_request_sender(self.request_queue, + self._error_queue) self.result_queue = FusedIpcQueue( (attach_info["result_addrs"][frontend_id], hmac_key), is_server=True, @@ -1575,6 +1741,8 @@ def pre_shutdown(self): def shutdown(self): self.pre_shutdown() + # Send the cancels queued by pre_shutdown before the lane goes idle. + self._close_request_sender() if self.rpc_client is not None: self.rpc_client.close() self.rpc_client = None diff --git a/tensorrt_llm/executor/ray/executor.py b/tensorrt_llm/executor/ray/executor.py index 37fa717ec3e8..844f978d14b0 100644 --- a/tensorrt_llm/executor/ray/executor.py +++ b/tensorrt_llm/executor/ray/executor.py @@ -37,6 +37,7 @@ from tensorrt_llm.executor.ray.utils import unwrap_ray_errors from tensorrt_llm.logger import logger +from ...llmapi._load_balance_env import uses_collective_shutdown from ...llmapi.utils import logger_debug from ..executor import GenerationExecutor from ..postproc_worker import PostprocWorkerConfig @@ -364,17 +365,22 @@ def shutdown(self): # Then, shutdown the workers if hasattr(self, 'workers') and self.workers is not None: - try: - shutdown_refs = [ - worker.shutdown.remote() for worker in self.workers - ] - # Add timeout to prevent indefinite hanging - ray.get(shutdown_refs, timeout=30.0) - except ray.exceptions.GetTimeoutError: - logger.warning( - "Timeout waiting for workers to shutdown after 30 seconds") - except Exception as e: - logger.warning(f"Error shutting down: {e}") + if uses_collective_shutdown( + getattr(self, "worker_kwargs", {}).get("llm_args")): + self._shutdown_workers_collectively() + else: + try: + shutdown_refs = [ + worker.shutdown.remote() for worker in self.workers + ] + # Add timeout to prevent indefinite hanging + ray.get(shutdown_refs, timeout=30.0) + except ray.exceptions.GetTimeoutError: + logger.warning( + "Timeout waiting for workers to shutdown after 30 seconds" + ) + except Exception as e: + logger.warning(f"Error shutting down: {e}") # The engines are already stopped by the shutdown RPC above, so # kill the actor processes explicitly instead of relying on @@ -412,6 +418,52 @@ def shutdown(self): logger.debug("Shutting down Ray cluster") ray.shutdown() + def _shutdown_workers_collectively(self): + """Shut down workers that must release EP resources collectively. + + Raises if a live worker failed, keeping the actors and placement group + for an explicit retry. A dead actor can never join the collective + release, so it falls through to forced teardown instead. + """ + shutdown_refs = getattr(self, "_pending_shutdown_refs", None) + if shutdown_refs is None: + shutdown_refs = [ + worker.shutdown.remote() for worker in self.workers + ] + self._pending_shutdown_refs = shutdown_refs + + # Converge every rank before deciding whether actor teardown is safe. + # Keep timed-out refs for the next explicit shutdown() call: issuing a + # second collective shutdown while the first is still running can + # deadlock the rank set. + shutdown_errors = [] + timed_out = False + actor_died = False + deadline = time.monotonic() + 30.0 + for rank, shutdown_ref in enumerate(shutdown_refs): + try: + ray.get(shutdown_ref, + timeout=max(0.0, deadline - time.monotonic())) + except ray.exceptions.GetTimeoutError as error: + timed_out = True + shutdown_errors.append((rank, error)) + except ray.exceptions.RayActorError as error: + actor_died = True + shutdown_errors.append((rank, error)) + except Exception as error: + shutdown_errors.append((rank, error)) + + if not timed_out or actor_died: + self._pending_shutdown_refs = None + for rank, error in shutdown_errors: + logger.warning(f"Ray worker {rank} shutdown failed: {error}") + if shutdown_errors and not actor_died: + # Preserve the actors and placement group. RayGPUWorker resets its + # own shutdown guard on failure, so a later explicit call can retry + # after all pending refs have converged. + self._shutdown_event.clear() + raise shutdown_errors[0][1] + def _wait_for_cluster_resource_release(self, timeout: float = 30.0) -> None: """Block until Ray returns the workers' resources to the cluster. diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index 25b124fb21ac..a2d8b34f7e65 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -30,6 +30,9 @@ control_action_decorator) from ... import TorchLlmArgs +from ...llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, + uses_collective_shutdown) from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ...llmapi.tokenizer import TokenizerBase from ...llmapi.utils import configure_cpu_affinity @@ -56,6 +59,18 @@ def resolve_obj_by_qualname(qualname: str) -> Any: class RayWorkerWrapper: def __init__(self, worker_cls, worker_kwargs, world_size, rank): + llm_args = worker_kwargs.get("llm_args") + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) + if llm_args is not None and llm_args.env_overrides: + os.environ.update(llm_args.env_overrides) + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + self.master_address = os.environ["MASTER_ADDR"] self.world_size = world_size self.rank = rank @@ -378,8 +393,98 @@ def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return + self.doing_shutdown = True + shutdown_succeeded = False + + try: + logger.debug(f'Worker {self.rank} shutting down...') + + if hasattr(self, 'shutdown_event'): + self.shutdown_event.set() + + shutdown_error = None + local_safe = True + if getattr(self, '_postproc_pool', None) is not None: + try: + self.shutdown_postproc_workers() + except BaseException as error: + shutdown_error = error + local_safe = False + + if hasattr(self, 'rpc_server') and self.rpc_server is not None: + logger.info( + f"[Rank {self.global_rank}] Shutting down RPC server") + try: + self.rpc_server.shutdown() + except Exception as e: + # Suppress errors during RPC server shutdown. These can occur + # if the server is already closed or during cleanup. + logger.debug( + f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" + ) + self.rpc_server = None + + if self.engine is not None: + engine = self.engine + engine_shutdown_completed = bool( + getattr(self, "_engine_shutdown_completed", False)) + if not engine_shutdown_completed: + try: + engine.shutdown() + self._engine_shutdown_completed = True + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + local_safe = False + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + try: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + + if (getattr(getattr(self, "llm_args", None), "backend", None) + == "pytorch" and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None): + try: + self.checkpoint_loader.cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + else: + self.checkpoint_loader = None + + # Check if there are any errors from the threads before shutdown. + try: + self._handle_background_error() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + + if shutdown_error is not None: + raise shutdown_error + self.engine = None + self._engine_shutdown_completed = False + shutdown_succeeded = True + logger.debug(f"Worker {self.rank} shutdown done.") + finally: + # A failed collective shutdown can be retried only explicitly. + self._collective_shutdown_failed = not shutdown_succeeded + if not shutdown_succeeded: + # Preserve the engine and allow an explicit retry to finish any + # model-owned collective teardown that could not complete. + self.doing_shutdown = False + + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True logger.debug(f'Worker {self.rank} shutting down...') @@ -403,6 +508,9 @@ def shutdown(self): if self.engine is not None: self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() self.engine = None if (self.llm_args.backend == "pytorch" diff --git a/tensorrt_llm/executor/rpc/rpc_server.py b/tensorrt_llm/executor/rpc/rpc_server.py index 07e20b5cd729..95a65407a6da 100644 --- a/tensorrt_llm/executor/rpc/rpc_server.py +++ b/tensorrt_llm/executor/rpc/rpc_server.py @@ -417,6 +417,20 @@ async def _process_requests(self) -> None: logger_debug( f"[server] RPC Server sent response for request {req}" ) + # Worker cleanup and server exit use a two-phase + # handshake. Wake RpcWorker.main_task only after the + # successful commit response is on the wire; otherwise + # main_task could cancel this coroutine before the client + # receives its acknowledgement. + if (req.method_name == "commit_shutdown" + and response.error is None): + complete_commit = getattr( + self._instance, + "_complete_shutdown_commit", + None, + ) + if callable(complete_commit): + complete_commit() # Decrement pending count if req.method_name not in ["_rpc_shutdown", "shutdown"]: diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index bb9e37295f0d..3a29a64610b5 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -16,6 +16,7 @@ import threading from typing import List, Optional, Union +from ..llmapi._load_balance_env import uses_collective_shutdown from ..llmapi.mpi_session import (MpiPoolSession, MpiSession, validate_session_world_size) from ..llmapi.utils import logger_debug, print_colored @@ -32,6 +33,9 @@ class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor): + # Bound on each acknowledged RPC of the collective shutdown handshake. + _COLLECTIVE_SHUTDOWN_TIMEOUT_S = 300.0 + def __init__( self, worker_kwargs: dict, @@ -293,9 +297,23 @@ def get_data_transceiver_state(self) -> bytes: def setup_engine_remote(self): return self.rpc_client.setup_engine().remote(need_response=True) + def _uses_collective_shutdown(self) -> bool: + return uses_collective_shutdown( + getattr(self, "worker_kwargs", {}).get("llm_args")) + def shutdown_remote(self): logger_debug("Shutting down rpc remote", color="yellow") - self.rpc_client.shutdown().remote(need_response=False) + if not self._uses_collective_shutdown(): + self.rpc_client.shutdown().remote(need_response=False) + return + # Cleanup must be acknowledged before main_task is allowed to close the + # server. The second RPC commits that acknowledged all-rank result; its + # response is sent before RpcWorker wakes main_task. The timeout keeps + # a dead rank zero from hanging the caller, e.g. an atexit shutdown. + timeout = self._COLLECTIVE_SHUTDOWN_TIMEOUT_S + self.rpc_client.shutdown().remote(timeout=timeout, need_response=True) + self.rpc_client.commit_shutdown().remote(timeout=timeout, + need_response=True) def abort_request(self, request_id: int) -> None: return self.rpc_client.abort_request(request_id).remote() @@ -317,6 +335,67 @@ def stop_profile(self) -> None: def shutdown(self): if self._shutdown_event.is_set(): return + if not self._uses_collective_shutdown(): + self._shutdown_rank_local() + return + self._shutdown_event.set() + logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow") + + # Keep the server, MPI workers, and retry gate live until every rank has + # acknowledged cleanup. RpcWorker resets its engine shutdown guard on a + # failed attempt, so the caller may invoke shutdown() again explicitly. + try: + self.shutdown_remote() + except BaseException: + self._shutdown_event.clear() + raise + + # Stop the response loop only after cleanup and commit are acknowledged. + if self.main_loop and self.main_loop_task_obj: + logger_debug("Cancelling main loop task.", color="yellow") + try: + self.main_loop.call_soon_threadsafe( + self.main_loop_task_obj.cancel) + except Exception as e: + logger_debug(f"Error cancelling main loop task: {e}", + color="yellow") + + if self.main_loop_thread and threading.current_thread( + ) != self.main_loop_thread: + self.main_loop_thread.join(timeout=2.0) + if self.main_loop_thread.is_alive(): + logger.warning("Main loop thread did not exit gracefully") + + shutdown_errors = [] + # Worker futures complete only after followers observe the commit and + # rank zero has joined its RPC server. Always converge them, including + # for an externally owned pool, before allowing that pool to be reused. + for future in getattr(self, "worker_futures", []): + try: + future.result() + except BaseException as error: + shutdown_errors.append(error) + + if self.mpi_session is not None: + if self._owns_mpi_session: + logger_debug("Shutting down mpi session", color="yellow") + try: + self.mpi_session.shutdown() + except BaseException as error: + shutdown_errors.append(error) + logger_debug("Mpi session shutdown", color="yellow") + self.mpi_session = None + + try: + self.rpc_client.close() + except BaseException as error: + shutdown_errors.append(error) + + if shutdown_errors: + raise shutdown_errors[0] + + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" self._shutdown_event.set() logger_debug("Shutting down GenerationExecutorRpcProxy", color="yellow") diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py index 273192943105..6f25cf71d240 100644 --- a/tensorrt_llm/executor/rpc_worker.py +++ b/tensorrt_llm/executor/rpc_worker.py @@ -11,6 +11,7 @@ from tensorrt_llm.llmapi.utils import enable_llm_debug, logger_debug from .._utils import mpi_rank +from ..llmapi._load_balance_env import uses_collective_shutdown from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.tokenizer import TokenizerBase from ..logger import set_level @@ -19,6 +20,9 @@ from .rpc import RPCServer from .rpc_worker_mixin import RpcWorkerMixin +_RPC_SHUTDOWN_RETRY = "retry" +_RPC_SHUTDOWN_COMMIT = "commit" + class RpcWorker(RpcWorkerMixin, BaseWorker): """ @@ -75,6 +79,10 @@ def __init__( and hasattr(llm_args, 'garbage_collection_gen0_threshold') else None) self.shutdown_event = Event() + self._rpc_collective_shutdown = False + self._rpc_shutdown_attempted = False + self._rpc_shutdown_ready = False + self._rpc_shutdown_committed = False self._response_queue = Queue() self.set_result_queue(self._response_queue) @@ -92,13 +100,67 @@ def setup_engine(self): super().setup_engine() def shutdown(self): + if getattr(self, "_rpc_shutdown_committed", False): + return + logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutting down", color="yellow") - self.shutdown_event.set() - super().shutdown() + if not getattr(self, "_rpc_collective_shutdown", False): + self.shutdown_event.set() + super().shutdown() + return + + # Followers wait between attempts. Wake them only when rank zero is + # explicitly retrying a failed, fully-converged shutdown attempt. + if mpi_rank() == 0 and self._rpc_shutdown_attempted: + mpi_comm().bcast(_RPC_SHUTDOWN_RETRY, root=0) + self._rpc_shutdown_attempted = True + self._rpc_shutdown_ready = False + + local_error = None + try: + super().shutdown() + except BaseException as error: + local_error = error + + # Do not acknowledge rank zero until every rank has completed cleanup. + # A peer-only failure must be reported to the client as well, otherwise + # a shared MPI pool could be reused while that peer is still tearing down. + local_failure = (None if local_error is None else + f"{type(local_error).__name__}: {local_error}") + failures = mpi_comm().allgather(local_failure) + failures = [ + f"rank {rank}: {failure}" for rank, failure in enumerate(failures) + if failure is not None + ] + if failures: + if local_error is not None: + raise local_error + raise RuntimeError("RPC worker shutdown failed: " + + "; ".join(failures)) + + self._rpc_shutdown_ready = True logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutdown", color="yellow") + def commit_shutdown(self) -> None: + """Commit an all-rank cleanup after its RPC response was acknowledged.""" + if self._rpc_shutdown_committed: + return + if mpi_rank() != 0: + raise RuntimeError("Only RPC rank zero may commit shutdown") + if not self._rpc_shutdown_ready: + raise RuntimeError( + "RPC worker cleanup has not completed on every rank") + mpi_comm().bcast(_RPC_SHUTDOWN_COMMIT, root=0) + self._rpc_shutdown_committed = True + + def _complete_shutdown_commit(self) -> None: + """Wake main_task only after RPCServer has sent the commit response.""" + if not self._rpc_shutdown_committed: + raise RuntimeError("RPC worker shutdown was not committed") + self.shutdown_event.set() + def start(self): pass @@ -126,6 +188,9 @@ def main_task( hf_model_dir=hf_model_dir, tokenizer=tokenizer, ) + # Only the per-layer dynamic EPLB fabric arena needs the all-rank + # two-phase shutdown; other configurations keep rank-local teardown. + worker._rpc_collective_shutdown = uses_collective_shutdown(llm_args) if mpi_rank() != 0: # The non-leader worker will setup the engine immediately. @@ -135,6 +200,28 @@ def main_task( f"[worker] Worker {mpi_rank()} is setting up the engine", color="yellow") worker.setup_engine() + # Rank zero publishes the shutdown request. Every follower waits + # for its broadcast sentinel, then runs the same finalize path so + # model-owned EP collectives remain rank-aligned. + if worker._rpc_collective_shutdown: + worker.engine.wait_shutdown() + while worker._rpc_collective_shutdown: + shutdown_error = None + try: + worker.shutdown() + except BaseException as error: + shutdown_error = error + command = mpi_comm().bcast(None, root=0) + if command == _RPC_SHUTDOWN_COMMIT: + if shutdown_error is not None: + raise shutdown_error + worker._rpc_shutdown_committed = True + worker.shutdown_event.set() + break + if command != _RPC_SHUTDOWN_RETRY: + raise RuntimeError( + f"Unexpected RPC shutdown command: {command!r}" + ) from shutdown_error else: logger_debug( diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index d017baa0cf6c..602438cf13a7 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -15,6 +15,7 @@ import gc import os +import pickle # nosec B403 import sys import threading import time @@ -28,6 +29,9 @@ from tensorrt_llm.logger import logger from .._utils import mpi_comm, mpi_rank, print_all_stacks +from ..llmapi._load_balance_env import (LAUNCH_QUEUES_ENV, + configure_moe_launch_queues, + uses_collective_shutdown) from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp from ..llmapi.tokenizer import TokenizerBase @@ -153,8 +157,139 @@ def shutdown(self): if self.doing_shutdown: return - else: - self.doing_shutdown = True + if not uses_collective_shutdown(getattr(self, "llm_args", None)): + self._shutdown_rank_local() + return + self.doing_shutdown = True + shutdown_succeeded = False + + try: + logger_debug(f'Worker {mpi_rank()} shutdown...\n', "yellow") + + shutdown_error = None + # The server references this worker; a long-lived worker process would + # otherwise keep one server, its threads and sockets per LLM it hosted. + if self.rpc_server is not None: + try: + self.rpc_server.shutdown() + self.rpc_server = None + except BaseException as error: + shutdown_error = error + if self.engine is not None: + engine = self.engine + local_safe = bool( + getattr(self, "_engine_shutdown_completed", False)) + if not local_safe: + try: + if engine.can_enqueue_requests(): + if self.await_response_thread.is_alive(): + self.await_response_thread.stop() + self.await_response_thread.join() + + engine.shutdown() + self._engine_shutdown_completed = True + local_safe = True + except BaseException as error: + shutdown_error = error + try: + terminal_cleanup = getattr(engine, "terminal_cleanup", None) + if terminal_cleanup is not None: + if getattr(engine, "shutdown_all_ranks", False): + terminal_cleanup(local_safe=local_safe) + elif local_safe: + terminal_cleanup() + except BaseException as error: + if shutdown_error is None: + shutdown_error = error + + if (self.llm_args.backend == "pytorch" + and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None + and shutdown_error is None): + try: + self.checkpoint_loader.cleanup() + except BaseException as error: + shutdown_error = error + else: + self.checkpoint_loader = None + + # Keep every rank aligned before entering the next collective phase. + # A rank-local engine or checkpoint failure must prevent peers from + # advancing to symmetric-buffer release and process-group teardown. + shutdown_failures = mpi_comm().allgather( + None if shutdown_error is + None else f"{type(shutdown_error).__name__}: {shutdown_error}") + shutdown_failures = [ + f"rank {rank}: {failure}" + for rank, failure in enumerate(shutdown_failures) + if failure is not None + ] + if shutdown_failures: + if shutdown_error is not None: + raise shutdown_error + raise RuntimeError("Worker shutdown failed: " + + "; ".join(shutdown_failures)) + + # MegaMoE's NVLink symmetric-memory activation workspaces are + # rendezvoused over the EP group, so they must go before the + # destroy_process_group() below. + mega_moe = sys.modules.get(_MEGA_MOE_DEEPGEMM_MODULE) + symm_release_error = None + if mega_moe is not None: + try: + mega_moe.release_symm_buffer_cache() + except BaseException as error: + symm_release_error = error + # No rank may destroy its process group while a peer still needs that + # group to retry symmetric-buffer destruction. + symm_release_failures = mpi_comm().allgather( + None if symm_release_error is None else + f"{type(symm_release_error).__name__}: {symm_release_error}") + symm_release_failures = [ + f"rank {rank}: {failure}" + for rank, failure in enumerate(symm_release_failures) + if failure is not None + ] + if symm_release_failures: + if symm_release_error is not None: + raise symm_release_error + raise RuntimeError("MegaMoE symm-buffer release failed: " + + "; ".join(symm_release_failures)) + + # Destroy torch distributed process groups so that NCCL communicators + # are torn down cleanly before MPI session shutdown and process exit. + # This is done here (not in PyExecutor.shutdown()) because the MPI + # worker owns the process group. In the Ray path the process group + # belongs to RayWorkerWrapper and must not be destroyed by the engine. + import torch.distributed + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + # Return this rank's GPU memory to the driver. Under an external MPI + # launch (mpirun/srun, e.g. CI), the worker process is long-lived and + # shared across successive LLM instances: a new GenerationExecutorWorker + # is built for each LLM, but the OS process -- and with it the CUDA + # context and PyTorch caching allocator -- persists. + gc.collect() + torch.cuda.empty_cache() + + self.engine = None + self._engine_shutdown_completed = False + shutdown_succeeded = True + logger_debug(f"Worker {mpi_rank()} shutdown done.\n", "yellow") + + # Raise background-thread errors only after the commit above, which + # every rank has reached, so a retry cannot re-enter collectives. + self._handle_background_error() + finally: + # A failed collective shutdown can be retried only explicitly. + self._collective_shutdown_failed = not shutdown_succeeded + if not shutdown_succeeded: + self.doing_shutdown = False + + def _shutdown_rank_local(self): + """Best-effort teardown for configurations without collective EP resources.""" + self.doing_shutdown = True logger_debug(f'Worker {mpi_rank()} shutdown...\n', "yellow") @@ -171,6 +306,9 @@ def shutdown(self): self.await_response_thread.join() self.engine.shutdown() + local_cleanup = getattr(self.engine, "local_cleanup", None) + if local_cleanup is not None: + local_cleanup() self.engine = None if (self.llm_args.backend == "pytorch" @@ -196,7 +334,7 @@ def shutdown(self): # Destroy torch distributed process groups so that NCCL communicators # are torn down cleanly before MPI session shutdown and process exit. # This is done here (not in PyExecutor.shutdown()) because the MPI - # worker owns the process group. In the Ray path the process group + # worker owns the process group. In the Ray path the process group # belongs to RayWorkerWrapper and must not be destroyed by the engine. import torch.distributed if torch.distributed.is_initialized(): @@ -204,15 +342,8 @@ def shutdown(self): # Return this rank's GPU memory to the driver. Under an external MPI # launch (mpirun/srun, e.g. CI), the worker process is long-lived and - # shared across successive LLM instances: a new GenerationExecutorWorker - # is built for each LLM, but the OS process -- and with it the CUDA - # context and PyTorch caching allocator -- persists. Setting - # `self.engine = None` above is not enough to free the GPU: reference - # cycles keep the model tensors alive until a later GC, and the allocator - # holds freed blocks as "reserved" instead of returning them. Without - # this, the previous model's ~weights-sized reservation carries into the - # next LLM built in this process and can OOM its load (e.g. back-to-back - # tests in one CI shard). + # shared across successive LLM instances, so the previous model's + # reservation would otherwise carry into the next LLM in this process. gc.collect() torch.cuda.empty_cache() @@ -228,6 +359,21 @@ def block_subordinates(self): self.engine.wait_shutdown() +def _incoming_requests(request_queue: IpcQueue): + """Yield the proxy's requests until its shutdown sentinel arrives. + + With TRTLLM_PROXY_BATCH_REQUESTS=1 the proxy sends a burst of requests as + one list of pickled requests. ``request_queue.get()`` has already + authenticated that message, and with it every payload inside. + """ + while (msg := request_queue.get()) is not None: + if isinstance(msg, list): + for payload in msg: + yield pickle.loads(payload) # nosec B301 + else: + yield msg + + @print_traceback_on_error def worker_main( worker_queues: WorkerCommIpcAddrs, @@ -247,6 +393,21 @@ def worker_main( hmac_key: bytes = b"", ) -> None: + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) + if llm_args is not None and llm_args.env_overrides: + # MPI may have cached the parent's environment before this worker was + # spawned, so replay overrides before any rebalance state is inspected. + os.environ.update(llm_args.env_overrides) + + # Configure per-layer dynamic EPLB before MPI or CUDA state is created. + if llm_args is not None and llm_args.backend == "pytorch": + queue_overrides = configure_moe_launch_queues( + getattr(llm_args, "moe_config", None), + llm_args.env_overrides, + initial_setting=launch_queues) + if queue_overrides is not llm_args.env_overrides: + llm_args.env_overrides = queue_overrides + def _print_stacks(): counter = 0 while True: @@ -264,14 +425,6 @@ def _print_stacks(): mpi_comm().barrier() - if llm_args is not None and llm_args.env_overrides: - # this is needed because MPI_Init seems to cache the env at import time. - # The cached env snapshot is used to spawn workers. - # Any env overrides to the main process after tensorrt_llm import - # may not get reflected in the spawned worker process, no matter how early, - # unless we update it explicitly here. - os.environ.update(llm_args.env_overrides) - if llm_args is not None and llm_args.trust_remote_code: _init_hf_modules() @@ -493,7 +646,7 @@ def notify_proxy_threads_to_quit(): worker.engine.set_resource_governor_queue( resource_governor_queue) - while (req := request_queue.get()) is not None: + for req in _incoming_requests(request_queue): if isinstance(req, CancellingRequest): worker.abort_request(req.id) elif isinstance(req, StartProfileRequest): diff --git a/tensorrt_llm/llmapi/_load_balance_env.py b/tensorrt_llm/llmapi/_load_balance_env.py new file mode 100644 index 000000000000..8951b6b7cac8 --- /dev/null +++ b/tensorrt_llm/llmapi/_load_balance_env.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configure process-start CUDA settings for opt-in MoE load balancing. + +Keep this module free of GPU imports: it runs before argument validation and +worker initialization are allowed to probe CUDA. +""" + +from __future__ import annotations + +import os +import sys +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .llm_args import MoeConfig + +LAUNCH_QUEUES_ENV = "CUDA_SCALE_LAUNCH_QUEUES" +_FROM_ENVIRON = object() + + +def configure_moe_launch_queues( + moe_config: MoeConfig | None, + env_overrides: dict[str, Any] | None = None, + *, + initial_setting: Any = _FROM_ENVIRON, +) -> dict[str, Any] | None: + """Configure the per-layer dynamic EPLB launch queue before CUDA setup. + + ``moe_config`` must already be validated. Standard EPLB and configurations + without a load balancer leave the environment and overrides untouched. + CUDA has no supported operation to resize existing launch queues here. + The late-init guard detects PyTorch initialization; external CUDA users + must still arrange for this setting before initializing their contexts. + + ``initial_setting`` is this process's ``CUDA_SCALE_LAUNCH_QUEUES`` before + the caller replayed ``env_overrides``: an existing CUDA context was created + with that value, not with the replayed one. It defaults to the current + environment. + """ + if moe_config is not None: + moe_config.resolve_load_balancer_compatibility() + load_balancer = getattr(moe_config, "load_balancer", None) + if getattr(load_balancer, "mode", None) != "per_layer_dynamic": + return env_overrides + + name = LAUNCH_QUEUES_ENV + if initial_setting is _FROM_ENVIRON: + initial_setting = os.environ.get(name) + torch = sys.modules.get("torch") + if initial_setting != "4x" and torch is not None and torch.cuda.is_initialized(): + raise RuntimeError( + "Per-layer dynamic EPLB requires CUDA_SCALE_LAUNCH_QUEUES=4x " + "before CUDA initialization. Restart the process with that " + "environment setting or construct the LLM before using CUDA." + ) + + overrides = dict(env_overrides or {}) + overrides[name] = "4x" + os.environ[name] = "4x" + return overrides + + +def uses_collective_shutdown(llm_args: Any) -> bool: + """Return whether executor teardown must release EP resources collectively. + + Only the per-layer dynamic EPLB fabric arena needs an all-rank release. + Every other configuration keeps the best-effort, rank-local teardown. + """ + moe_config = getattr(llm_args, "moe_config", None) + load_balancer = getattr(moe_config, "load_balancer", None) + return getattr(load_balancer, "mode", None) == "per_layer_dynamic" diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index ca6bec788a23..486cade00086 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -55,10 +55,11 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams +from ._load_balance_env import LAUNCH_QUEUES_ENV, configure_moe_launch_queues from .llm_args import (ENCODER_RUNNER_MANAGED_INPUTS, TORCH_LLMARGS_EXPLICIT_DOCSTRING, TORCH_LLMARGS_REMOVED_ARGS, ExecutorMemoryType, - RuntimeMemoryStatus, TorchLlmArgs, + MoeConfig, RuntimeMemoryStatus, TorchLlmArgs, validate_token_encoder_bucket_config) from .llm_utils import CachedModelLoader, KvCacheRetentionConfig, ModelLoader from .mpi_session import MpiPoolSession, external_mpi_comm_available @@ -384,7 +385,21 @@ def __init__(self, try: env_overrides = kwargs.get("env_overrides", None) + launch_queues = os.environ.get(LAUNCH_QUEUES_ENV) + # Rebalance validation and launch-queue setup read process state. + # Apply caller overrides before either one observes that state. self._process_env_overrides(env_overrides) + if kwargs.get("backend") == "pytorch": + # This must precede TorchLlmArgs validators and GPU probes. + moe_config = kwargs.get("moe_config") + if isinstance(moe_config, dict): + moe_config = MoeConfig(**moe_config) + kwargs["moe_config"] = moe_config + queue_overrides = configure_moe_launch_queues( + moe_config, env_overrides, initial_setting=launch_queues) + if queue_overrides is not env_overrides: + kwargs["env_overrides"] = queue_overrides + env_overrides = queue_overrides backend = kwargs.get('backend', None) if backend == "pytorch": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..fdfa6fcbccd5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -33,7 +33,8 @@ from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, - PrivateAttr, StrictInt, field_validator, model_validator) + PrivateAttr, StrictInt, TypeAdapter, field_validator, + model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -1851,61 +1852,42 @@ def to_sparse_params(self, **kwargs): uses_spcompress=self.uses_spcompress) -class MoeLoadBalancerConfig(StrictBaseModel): - """Pydantic configuration model for the Mixture of Experts (MoE) load balancer. +class _MoeLoadBalancerConfigBase(StrictBaseModel): + """Common slot layout and expert-parallel runtime state for EPLB.""" - This model holds configuration data (`num_slots`, etc.) as well as - runtime state (`_ep_rank`, `_ep_size`) which must be set via the - `setup()` method before use. - """ - - num_slots: Optional[int] = None - initial_global_assignments: Optional[Dict[int, List[int]]] = Field( + num_slots: Optional[int] = Field( default=None, - repr=False # Exclude this large dict from model representation + description= + "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots." ) - layer_updates_per_iter: int = 0 _ep_rank: Optional[int] = PrivateAttr(default=None) _ep_size: Optional[int] = PrivateAttr(default=None) - # --- Methods --- - def setup(self, ep_rank: int, ep_size: int) -> None: - """Initializes the runtime state of the configuration. - This must be called before accessing properties like `num_local_slots`. - """ + """Bind the process-local expert-parallel rank and size.""" self._ep_rank = ep_rank self._ep_size = ep_size - - # This assertion was in the original and is critical. if self.num_slots is None: raise ValueError("`num_slots` cannot be None when calling setup().") - if self.num_slots % ep_size != 0: raise ValueError( f"`num_slots` ({self.num_slots}) must be divisible by `ep_size` ({ep_size})." ) - # --- Computed Properties --- - # These properties depend on the runtime state set by setup() - @property def ep_rank(self) -> int: - """Public accessor for the private expert parallel rank.""" if self._ep_rank is None: raise AttributeError("ep_rank is not set. Call setup() first.") return self._ep_rank @property def ep_size(self) -> int: - """Public accessor for the private expert parallel size.""" if self._ep_size is None: raise AttributeError("ep_size is not set. Call setup() first.") return self._ep_size @property def num_local_slots(self) -> int: - """Calculates the number of slots local to this rank.""" if self.num_slots is None or self._ep_size is None: raise ValueError( "Cannot calculate `num_local_slots`. " @@ -1914,7 +1896,6 @@ def num_local_slots(self) -> int: @property def slot_start(self) -> int: - """Calculates the starting global slot index for this rank.""" if self._ep_rank is None: raise ValueError( "Cannot calculate `slot_start`. Call setup() first.") @@ -1922,34 +1903,139 @@ def slot_start(self) -> int: @property def slot_end(self) -> int: - """Calculates the ending global slot index (exclusive) for this rank.""" return self.slot_start + self.num_local_slots + +class MoeLoadBalancerConfig(_MoeLoadBalancerConfigBase): + """Configure the standard asynchronous EPLB implementation.""" + + mode: Literal["standard"] = Field( + default="standard", + description= + "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + status="prototype") + initial_global_assignments: Optional[Dict[int, List[int]]] = Field( + default=None, + repr=False, + description= + "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert." + ) + layer_updates_per_iter: int = Field( + default=0, + description= + "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB." + ) + def get_layer_initial_global_assignments( self, layer_idx: int) -> Optional[List[int]]: - """Retrieves the initial global assignments for a specific layer.""" if self.initial_global_assignments is None: return None - if layer_idx not in self.initial_global_assignments: raise KeyError( f"layer_idx {layer_idx} not found in `initial_global_assignments`." ) - assignments = self.initial_global_assignments[layer_idx] - if self.num_slots is None: raise ValueError( "`num_slots` is not set, cannot verify assignment length.") - if len(assignments) != self.num_slots: raise ValueError( f"Assignment length ({len(assignments)}) for layer {layer_idx} " f"does not match `num_slots` ({self.num_slots}).") - return assignments +class MoePerLayerDynamicLoadBalancerConfig(_MoeLoadBalancerConfigBase): + """Configure per-layer dynamic EPLB with Halo-Q scheduling.""" + + mode: Literal["per_layer_dynamic"] = Field( + default="per_layer_dynamic", + description= + "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + status="prototype") + num_slots: PositiveInt = Field( + description= + "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", + status="prototype") + auxiliary_sms: int = Field( + default=4, + ge=1, + le=8, + description= + "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", + status="prototype") + helper_slot_mode: Literal["per_layer_resident", "shared_banks"] = Field( + default="per_layer_resident", + description= + "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + status="prototype") + + def validate_expert_capacity(self, num_experts: int) -> None: + """Validate that every EP rank has at least one helper slot.""" + if self._ep_size is None: + raise ValueError("Call setup() before validating expert capacity.") + if num_experts <= 0: + raise ValueError("`num_experts` must be positive.") + if num_experts % self._ep_size != 0: + raise ValueError( + f"`num_experts` ({num_experts}) must be divisible by " + f"`ep_size` ({self._ep_size}) for per-layer dynamic EPLB.") + if self.num_slots <= num_experts: + raise ValueError( + f"`num_slots` ({self.num_slots}) must exceed `num_experts` " + f"({num_experts}) for per-layer dynamic EPLB.") + + +MoeLoadBalancerConfigType = Annotated[ + Union[MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig], + PydanticField(discriminator="mode"), +] +_MOE_LOAD_BALANCER_ADAPTER = TypeAdapter(MoeLoadBalancerConfigType) + + +def _normalize_load_balancer_config(value: Any) -> Any: + """Preserve legacy standard-EPLB dictionaries that predate ``mode``.""" + if isinstance(value, dict) and "mode" not in value: + return {"mode": "standard", **value} + return value + + +def _parse_load_balancer_config(value: Any) -> MoeLoadBalancerConfigType: + return _MOE_LOAD_BALANCER_ADAPTER.validate_python( + _normalize_load_balancer_config(value)) + + +def _resolve_load_balancer_config(value: Any) -> Any: + """Resolve inline or YAML-backed EPLB input to its typed config.""" + if value is None or isinstance( + value, + (MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig)): + return value + + if isinstance(value, str): + config_path = value + if not os.path.exists(config_path): + raise FileNotFoundError( + f"MoE load balancer config file not found: {config_path}") + try: + with open(config_path) as f: + value = yaml.safe_load(f) + return _parse_load_balancer_config(value) + except Exception as e: + raise ValueError( + f"Failed to load MoE load balancer config file: {config_path}" + ) from e + + if isinstance(value, dict): + try: + return _parse_load_balancer_config(value) + except Exception as e: + raise ValueError( + f"Failed to load MoE load balancer config: {value}") from e + + return value + + _MoeBackend = Literal["AUTO", "CUTLASS", "CUTEDSL", "CUTEDSL_FC12", "TRTLLM", "DEEPGEMM", "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"] @@ -1969,10 +2055,9 @@ class MoeConfig(StrictBaseModel): "If set, at most max_num_tokens tokens will be sent to torch.ops.trtllm.fused_moe at the same time. If the number of tokens exceeds max_num_tokens, the input tensors will be split into chunks and a for loop will be used." ) - load_balancer: Optional[Union[object, str]] = Field( + load_balancer: Optional[Union[MoeLoadBalancerConfigType, str]] = Field( default=None, - description="Configuration for MoE load balancing.", - json_schema_extra={"type": "Union[MoeLoadBalancerConfig, dict, str]"}) + description="Configuration for standard or per-layer dynamic EPLB.") disable_finalize_fusion: bool = Field( default=False, @@ -1986,6 +2071,23 @@ class MoeConfig(StrictBaseModel): "Use low precision combine in MoE operations (only for NVFP4 quantization). When enabled, uses lower precision for combining expert outputs to improve performance." ) + @field_validator("load_balancer", mode="before") + @classmethod + def resolve_load_balancer_config(cls, value: Any) -> Any: + return _resolve_load_balancer_config(value) + + @model_validator(mode="after") + def resolve_load_balancer_compatibility(self) -> "MoeConfig": + """Resolve the backend required by per-layer dynamic EPLB.""" + if getattr(self.load_balancer, "mode", None) != "per_layer_dynamic": + return self + if self.backend == "AUTO": + self.backend = "MEGAMOE_CUTEDSL" + elif self.backend != "MEGAMOE_CUTEDSL": + raise ValueError("per-layer dynamic EPLB requires " + "backend='MEGAMOE_CUTEDSL'") + return self + Nvfp4Backend = Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'] @@ -7866,69 +7968,40 @@ def validate_gms_locality_domain_compat(self) -> 'TorchLlmArgs': @model_validator(mode="after") def validate_gms_moe_compat(self) -> 'TorchLlmArgs': - """Reject ``LoadFormat.GMS`` combined with a MoE load balancer. - - The ``MoeLoadBalancer``'s ``register_weight_slots_after_to_cuda`` - and ``finalize_model`` run AFTER the GMS RW pool is closed and - ``finalize_write`` has committed, so any CUDA allocations they - make land in non-GMS memory and are NOT part of the committed - layout that RO peers receive. The result is "wrong inference, - no error" on RO peers (broken MoE routing state). Failing at - config-validation time is strictly better than that silent - miscompute. - - The fix for this gap (running the MoE finalize work INSIDE - ``mem_pool_scope`` and BEFORE ``finalize_write`` so MoE - allocations are part of the committed layout) is tracked as - the (MoE, GMS) follow-up; see ``model_loader.py``'s - ``TODO(GMS-MOE-LB)`` comment. - - Returns: - ``self`` (Pydantic ``model_validator`` contract). + """Reject GMS when MoE load balancing owns extra device state. - Raises: - ValueError: When ``load_format == LoadFormat.GMS`` and - ``moe_config.load_balancer`` is set. + Load-balancer state is currently allocated after GMS commits its + shared layout, so read-only peers cannot safely restore that state. """ if (self.load_format == LoadFormat.GMS and self.moe_config is not None and self.moe_config.load_balancer is not None): raise ValueError( "LoadFormat.GMS is incompatible with moe_config.load_balancer " - "in this PR. The MoE load balancer's " - "register_weight_slots_after_to_cuda and finalize_model run " - "after the GMS pool closes and finalize_write commits, so " - "their allocations land outside the committed layout. RO " - "peers would receive a broken MoE routing state. Either " - "disable moe_config.load_balancer or use LoadFormat.AUTO. " - "Tracked as the (MoE, GMS) follow-up at " - "tensorrt_llm/_torch/pyexecutor/model_loader.py " - "(see TODO(GMS-MOE-LB)).") + "because load-balancer device state is not included in the " + "committed GMS layout. Disable moe_config.load_balancer or " + "use LoadFormat.AUTO.") return self @model_validator(mode="after") def validate_load_balancer(self) -> 'TorchLlmArgs': - if isinstance(self.moe_config.load_balancer, str): - if not os.path.exists(self.moe_config.load_balancer): - raise FileNotFoundError( - f"MoE load balancer config file not found: {self.moe_config.load_balancer}" - ) - try: - with open(self.moe_config.load_balancer) as f: - moe_load_balancer_config = yaml.safe_load(f) - self.moe_config.load_balancer = MoeLoadBalancerConfig( - **moe_load_balancer_config) - except Exception as e: - raise ValueError( - f"Failed to load MoE load balancer config file: {self.moe_config.load_balancer}" - ) from e - elif isinstance(self.moe_config.load_balancer, dict): - try: - self.moe_config.load_balancer = MoeLoadBalancerConfig( - **self.moe_config.load_balancer) - except Exception as e: - raise ValueError( - f"Failed to load MoE load balancer config: {self.moe_config.load_balancer}" - ) from e + self.moe_config.load_balancer = _resolve_load_balancer_config( + self.moe_config.load_balancer) + self.moe_config.resolve_load_balancer_compatibility() + return self + + @model_validator(mode="after") + def validate_per_layer_dynamic_eplb_cuda_graph(self) -> 'TorchLlmArgs': + """Per-layer dynamic EPLB plans cannot be captured in CUDA graphs.""" + if getattr(self.moe_config.load_balancer, "mode", + None) != "per_layer_dynamic": + return self + if self.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED: + raise ValueError( + "per-layer dynamic EPLB does not support prefill CUDA graphs") + if self.cuda_graph_config is not None: + logger.warning("per-layer dynamic EPLB does not support CUDA " + "graphs; setting cuda_graph_config to None.") + self.cuda_graph_config = None return self @model_validator(mode='after') diff --git a/tensorrt_llm/llmapi/mpi_session.py b/tensorrt_llm/llmapi/mpi_session.py index 190a4723a0f0..8e60054ca13d 100644 --- a/tensorrt_llm/llmapi/mpi_session.py +++ b/tensorrt_llm/llmapi/mpi_session.py @@ -36,8 +36,13 @@ # Environment that MpiPoolSession forwards from the parent to its workers. _FORWARDED_ENV_PREFIXES = ("TRTLLM", "TLLM") -_FORWARDED_ENV_NAMES = (*_UNIFIED_CACHE_ENV_VARS, "FLASHINFER_CUBIN_DIR", - "MASTER_ADDR", "MASTER_PORT") +_FORWARDED_ENV_NAMES = ( + *_UNIFIED_CACHE_ENV_VARS, + "FLASHINFER_CUBIN_DIR", + "MASTER_ADDR", + "MASTER_PORT", + # MPI may use an environment snapshot from MPI_Init. + "CUDA_SCALE_LAUNCH_QUEUES") def _is_forwarded_env(name: str) -> bool: diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index 9a9188a18cd8..139b0af15be4 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -8,7 +8,7 @@ import torch -from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import MoeLoadBalancer +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import BaseMoeLoadBalancer from tensorrt_llm._torch.utils import get_device_uuid from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import serialization @@ -211,7 +211,7 @@ def update_weights(self, ipc_handles: Optional[dict] = None): logger.info("Finalize update weights") self.finalize_weight_update() moe_load_balancer = getattr(self.engine.model_engine, "moe_load_balancer", None) - if isinstance(moe_load_balancer, MoeLoadBalancer): + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): moe_load_balancer.register_weight_slots_after_to_cuda() logger.info("moe_load_balancer finalizing model...") moe_load_balancer.finalize_model() diff --git a/tensorrt_llm/llmapi/serialization.py b/tensorrt_llm/llmapi/serialization.py index 83e6f2af0252..f9fa2affaa88 100644 --- a/tensorrt_llm/llmapi/serialization.py +++ b/tensorrt_llm/llmapi/serialization.py @@ -99,7 +99,8 @@ "tensorrt_llm.executor.worker": ["GenerationExecutorWorker", "worker_main"], "tensorrt_llm.llmapi.llm_args": [ "_ParallelConfig", "CalibConfig", "CapacitySchedulerPolicy", - "KvCacheConfig", "SchedulerConfig", "LoadFormat", "DynamicBatchConfig" + "KvCacheConfig", "SchedulerConfig", "LoadFormat", "DynamicBatchConfig", + "MoeLoadBalancerConfig", "MoePerLayerDynamicLoadBalancerConfig" ], "tensorrt_llm.llmapi.mpi_session": ["RemoteTask"], "tensorrt_llm.llmapi.tokenizer": ["TransformersTokenizer"], diff --git a/tensorrt_llm/logger.py b/tensorrt_llm/logger.py index 74253479a9b8..35fa5b6d077e 100644 --- a/tensorrt_llm/logger.py +++ b/tensorrt_llm/logger.py @@ -296,6 +296,10 @@ def _percent_format(cls, msg) -> Optional[str]: return None def log(self, severity, *msg): + # Without per-module levels the caller's module only matters for an emitted message; + # a filtered one returns before the stack walk. + if not self._module_levels and not self.is_severity_enabled(severity): + return module = _get_caller_module() if not self.is_severity_enabled(severity, module): return @@ -364,6 +368,8 @@ def is_debug_enabled(self) -> bool: ``TLLM_LOG_LEVEL_BY_MODULE`` override can enable or disable the caller's module independently of the global level. """ + if not self._module_levels: + return self.is_severity_enabled(self.VERBOSE) return self.is_severity_enabled(self.VERBOSE, _get_caller_module()) @property diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index c133f71946f4..5ddcf62a18ff 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -12,6 +12,7 @@ # 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. +import functools import json import os from abc import ABC, abstractmethod @@ -669,7 +670,7 @@ def _get_sampling_config(self) -> tllme.SamplingConfig: # | Sampling | use_beam_search | beam_width == 1 | # | Sampling | n | num_return_sequences | # | Sampling | best_of | no corresponding param | - fields = {f for f in dir(tllme.SamplingConfig) if not f.startswith("__")} + fields = _sampling_config_fields() unmatched_params = [ "num_return_sequences", "beam_width", @@ -688,12 +689,7 @@ def _get_sampling_config(self) -> tllme.SamplingConfig: return tllme.SamplingConfig(**llmapi_to_rt_param_map) def _get_output_config(self) -> tllme.OutputConfig: - sampling_param_fields = set(dir(SamplingParams)) - fields = [ - f - for f in dir(tllme.OutputConfig) - if not f.startswith("__") and f in sampling_param_fields - ] + fields = _output_config_fields() config_kwargs = {f: getattr(self, f) for f in fields} @@ -751,3 +747,21 @@ def _get_guided_decoding_params(self) -> tllme.GuidedDecodingParams: ) else: return None + + +@functools.cache +def _sampling_config_fields() -> frozenset[str]: + """Return the ``tllme.SamplingConfig`` field names. + + They are fixed by the bindings, so the ``dir()`` scan runs once instead of once per request. + """ + return frozenset(f for f in dir(tllme.SamplingConfig) if not f.startswith("__")) + + +@functools.cache +def _output_config_fields() -> tuple[str, ...]: + """Return the ``tllme.OutputConfig`` field names that ``SamplingParams`` also defines.""" + sampling_param_fields = set(dir(SamplingParams)) + return tuple( + f for f in dir(tllme.OutputConfig) if not f.startswith("__") and f in sampling_param_fields + ) diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 3db659f73b23..574ee0cad605 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4044,47 +4044,13 @@ "load_balancer": { "anyOf": [ { - "additionalProperties": false, - "description": "Pydantic configuration model for the Mixture of Experts (MoE) load balancer.\n\nThis model holds configuration data (`num_slots`, etc.) as well as\nruntime state (`_ep_rank`, `_ep_size`) which must be set via the\n`setup()` method before use.", - "properties": { - "initial_global_assignments": { - "anyOf": [ - { - "additionalProperties": { - "items": { - "type": "integer" - }, - "type": "array" - }, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Initial Global Assignments" - }, - "layer_updates_per_iter": { - "default": 0, - "title": "Layer Updates Per Iter", - "type": "integer" - }, - "num_slots": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Num Slots" - } - }, - "title": "MoeLoadBalancerConfig", - "type": "object" + "$ref": "#/$defs/MoeLoadBalancerConfig" + }, + { + "$ref": "#/$defs/MoePerLayerDynamicLoadBalancerConfig", + "required": [ + "mode" + ] }, { "type": [ @@ -4092,7 +4058,10 @@ "null" ] } - ] + ], + "default": null, + "description": "Configuration for standard or per-layer dynamic EPLB.", + "title": "Load Balancer" }, "max_num_tokens": { "anyOf": [ @@ -4117,6 +4086,106 @@ "title": "MoeConfig", "type": "object" }, + "MoeLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure the standard asynchronous EPLB implementation.", + "properties": { + "initial_global_assignments": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "integer" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert.", + "title": "Initial Global Assignments" + }, + "layer_updates_per_iter": { + "default": 0, + "description": "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB.", + "title": "Layer Updates Per Iter", + "type": "integer" + }, + "mode": { + "const": "standard", + "default": "standard", + "description": "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots.", + "title": "Num Slots" + } + }, + "title": "MoeLoadBalancerConfig", + "type": "object" + }, + "MoePerLayerDynamicLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", + "properties": { + "auxiliary_sms": { + "default": 4, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", + "maximum": 8, + "minimum": 1, + "status": "prototype", + "title": "Auxiliary Sms", + "type": "integer" + }, + "helper_slot_mode": { + "default": "per_layer_resident", + "description": "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + "enum": [ + "per_layer_resident", + "shared_banks" + ], + "status": "prototype", + "title": "Helper Slot Mode", + "type": "string" + }, + "mode": { + "const": "per_layer_dynamic", + "default": "per_layer_dynamic", + "description": "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Num Slots", + "type": "integer" + } + }, + "required": [ + "num_slots" + ], + "title": "MoePerLayerDynamicLoadBalancerConfig", + "type": "object" + }, "MooncakeStoreConfig": { "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 371b27ff1e08..4e4b8e7dd434 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5608,47 +5608,13 @@ "load_balancer": { "anyOf": [ { - "additionalProperties": false, - "description": "Pydantic configuration model for the Mixture of Experts (MoE) load balancer.\n\nThis model holds configuration data (`num_slots`, etc.) as well as\nruntime state (`_ep_rank`, `_ep_size`) which must be set via the\n`setup()` method before use.", - "properties": { - "initial_global_assignments": { - "anyOf": [ - { - "additionalProperties": { - "items": { - "type": "integer" - }, - "type": "array" - }, - "type": "object" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Initial Global Assignments" - }, - "layer_updates_per_iter": { - "default": 0, - "title": "Layer Updates Per Iter", - "type": "integer" - }, - "num_slots": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Num Slots" - } - }, - "title": "MoeLoadBalancerConfig", - "type": "object" + "$ref": "#/$defs/MoeLoadBalancerConfig" + }, + { + "$ref": "#/$defs/MoePerLayerDynamicLoadBalancerConfig", + "required": [ + "mode" + ] }, { "type": [ @@ -5656,7 +5622,10 @@ "null" ] } - ] + ], + "default": null, + "description": "Configuration for standard or per-layer dynamic EPLB.", + "title": "Load Balancer" }, "max_num_tokens": { "anyOf": [ @@ -5681,6 +5650,106 @@ "title": "MoeConfig", "type": "object" }, + "MoeLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure the standard asynchronous EPLB implementation.", + "properties": { + "initial_global_assignments": { + "anyOf": [ + { + "additionalProperties": { + "items": { + "type": "integer" + }, + "type": "array" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Initial expert placement keyed by MoE layer index. Each value has `num_slots` entries, gives the expert ID held by each global slot, and covers every expert.", + "title": "Initial Global Assignments" + }, + "layer_updates_per_iter": { + "default": 0, + "description": "Number of MoE layers whose expert placement is updated per iteration. 0 keeps the initial placement (static EPLB); a positive value enables online EPLB.", + "title": "Layer Updates Per Iter", + "type": "integer" + }, + "mode": { + "const": "standard", + "default": "standard", + "description": "EPLB mode. `standard` is static EPLB when `layer_updates_per_iter` is 0 and online EPLB otherwise. Configs without `mode` use `standard`.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size`; each rank hosts `num_slots / ep_size` slots.", + "title": "Num Slots" + } + }, + "title": "MoeLoadBalancerConfig", + "type": "object" + }, + "MoePerLayerDynamicLoadBalancerConfig": { + "additionalProperties": false, + "description": "Configure per-layer dynamic EPLB with Halo-Q scheduling.", + "properties": { + "auxiliary_sms": { + "default": 4, + "description": "Number of SMs that the routed-input quantization and shared-expert FC12 kernels leave free for the kernel that runs the HALO-Q scheduler and then the in-switch copy of its plan on the copy stream, one CTA per SM. Must be smaller than the device SM count.", + "maximum": 8, + "minimum": 1, + "status": "prototype", + "title": "Auxiliary Sms", + "type": "integer" + }, + "helper_slot_mode": { + "default": "per_layer_resident", + "description": "Storage of the helper slots. `per_layer_resident` gives every MoE layer its own helper slots and skips copying an expert that a helper slot already holds from that layer's previous forward pass; it costs one helper set per MoE layer, about 6.4 GiB per GPU for DeepSeek-V4-Pro with 3 helper slots per rank. `shared_banks` has all MoE layers alternate between two shared helper banks and copy every helper expert in every forward pass.", + "enum": [ + "per_layer_resident", + "shared_banks" + ], + "status": "prototype", + "title": "Helper Slot Mode", + "type": "string" + }, + "mode": { + "const": "per_layer_dynamic", + "default": "per_layer_dynamic", + "description": "EPLB mode. `per_layer_dynamic` replicates each MoE layer's hot experts into helper slots in every forward pass, planned from that layer's current routing.", + "status": "prototype", + "title": "Mode", + "type": "string" + }, + "num_slots": { + "description": "Total number of expert slots across all EP ranks. Must be divisible by `ep_size` and greater than the number of experts. Each rank keeps its `num_experts / ep_size` resident experts and adds `(num_slots - num_experts) / ep_size` helper slots. Unlike standard EPLB, the helper slots are not static replicas: every forward pass refills them from that layer's routing; see `helper_slot_mode` for their storage.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Num Slots", + "type": "integer" + } + }, + "required": [ + "num_slots" + ], + "title": "MoePerLayerDynamicLoadBalancerConfig", + "type": "object" + }, "MooncakeStoreConfig": { "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 4382827aad45..c48a96278a2a 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -9,7 +9,7 @@ import weakref from dataclasses import replace from enum import IntEnum -from typing import Literal, Optional +from typing import Any, Dict, Literal, Optional, Union import torch @@ -22,7 +22,12 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, remove_weights, skip_forward from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.modules.mhc.hyper_connection import HCState from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( + BaseMoeLoadBalancer, + MoeLoadBalancerIterContext, +) from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor._util import _mamba_conv_layout_kwargs, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import ( @@ -411,6 +416,7 @@ def __init__( use_cuda_graph: bool, spec_config: Optional[DecodingBaseConfig] = None, vision_config: Optional[str] = None, + moe_load_balancer: Optional[Union[str, Dict[str, Any]]] = None, ) -> None: super().__init__() @@ -438,12 +444,29 @@ def __init__( max_num_tokens=moe_max_num_tokens, disable_finalize_fusion=False, use_low_precision_moe_combine=use_low_precision_moe_combine, + load_balancer=moe_load_balancer, ), attn_backend="TRTLLM", kv_cache_config=KvCacheConfig( dtype=kv_cache_dtype, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype ), ) + # Reject load balancers that cannot run on a layer slice driven by this + # harness. The validator above has resolved a YAML path to its config. + load_balancer_config = llm_args.moe_config.load_balancer + load_balancer_mode = getattr(load_balancer_config, "mode", None) + if load_balancer_mode == "per_layer_dynamic" and use_cuda_graph: + raise ValueError( + "Per-layer dynamic EPLB does not support CUDA graphs: it binds one submission" + " stream and keeps host-side plan state that a graph replay cannot advance." + " Rerun with --no-use-cuda-graph." + ) + if load_balancer_mode == "standard" and load_balancer_config.layer_updates_per_iter > 0: + raise ValueError( + "Layer-wise benchmarks support per_layer_dynamic EPLB and static standard EPLB" + " (layer_updates_per_iter 0) only: online standard EPLB synchronizes every MoE" + " layer of the model, but only a slice of the layers runs here." + ) with ( self.scaled_from_ctx(scaled_from, mapping), self.vision_config_ctx(vision_config), @@ -464,7 +487,7 @@ def __init__( max_num_tokens=max_num_tokens, max_seq_len=max_seq_len, ) - model, _ = model_loader.load( + model, load_balancer = model_loader.load( checkpoint_dir=pretrained_model_name_or_path, checkpoint_loader=checkpoint_loader ) @@ -520,6 +543,11 @@ def forward_block_residual( self.model_config = model.model_config self.model = model + # The executor drives the load balancer's iterations and teardown; this + # harness does it in `create_run_pack` and `shutdown_moe_load_balancer`. + self.moe_load_balancer = ( + load_balancer if isinstance(load_balancer, BaseMoeLoadBalancer) else None + ) self.layer_indices = layer_indices @staticmethod @@ -774,6 +802,11 @@ def create_run_pack( hc_mult = getattr(pretrained_config, "hc_mult", None) if hc_mult is not None: hidden_states = hidden_states.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() + # Only layer 0 takes the plain residual. A slice starting mid-model + # enters with the resolved state that the previous layer returns. + first_layer = self.model.model.layers[self.layer_indices[0]] + if not getattr(first_layer, "is_first_layer", True): + hidden_states = HCState.resolved(hidden_states) kwargs["input_ids"] = torch.randint( 0, pretrained_config.vocab_size, @@ -801,7 +834,7 @@ def run_pack(*, check=False): discard_nccl_window_tensor_outputs((position_ids, hidden_states, residual)), ): get_model_extra_attrs()["attention_metadata"] = weakref.ref(attn_metadata) - with torch.inference_mode(): + with torch.inference_mode(), MoeLoadBalancerIterContext(self.moe_load_balancer): hidden_states_out, residual_out = self.model( position_ids, hidden_states, attn_metadata, residual, **kwargs ) @@ -816,6 +849,12 @@ def run_pack(*, check=False): return run_pack + def shutdown_moe_load_balancer(self) -> None: + """Release the MoE load balancer's resources; collective over the EP ranks.""" + if self.moe_load_balancer is not None: + self.moe_load_balancer.shutdown(local_safe=True) + self.moe_load_balancer = None + @contextlib.contextmanager def replace_routing_method_ctx(self, balance_method: BalanceMethod, balance_ratio: float): if balance_method == BalanceMethod.NotModified: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f940f2a80937..3529dd5e756c 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1061,6 +1061,39 @@ "kind": "value", "path": "moe_config.disable_finalize_fusion" }, + { + "capture_policy": "int", + "kind": "value", + "path": "moe_config.load_balancer.auxiliary_sms" + }, + { + "allowed_values": [ + "per_layer_resident", + "shared_banks" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "moe_config.load_balancer.helper_slot_mode" + }, + { + "capture_policy": "int", + "kind": "value", + "path": "moe_config.load_balancer.layer_updates_per_iter" + }, + { + "allowed_values": [ + "standard", + "per_layer_dynamic" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "moe_config.load_balancer.mode" + }, + { + "capture_policy": "int|none", + "kind": "value", + "path": "moe_config.load_balancer.num_slots" + }, { "capture_policy": "int|none", "kind": "value", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b597efc96bc0..ee8b77eefe0b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -143,6 +143,7 @@ l0_b200: - unittest/_torch/modules/test_engram.py - unittest/_torch/modules/test_gate_up_swiglu_quack.py - unittest/_torch/modules/test_multi_stream_utils.py + - unittest/_torch/modules/test_megamoe_shared_mlp_scales.py # ------------- DeepSeek-V4.1 model, Engram and DSpark --------------- - unittest/_torch/modeling/test_modeling_deepseekv41.py - unittest/_torch/modules/test_engram_decode_hash.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 73c43d3b8af4..9891e4cbd207 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -189,10 +189,12 @@ l0_cpu: - unittest/executor/test_ipc.py - unittest/executor/test_rpc.py - unittest/executor/test_multi_frontend_routing.py + - unittest/executor/test_proxy_request_sender.py - unittest/executor/test_event_loop_error_broadcast.py - unittest/executor/test_forward_error_survival.py - unittest/executor/test_rpc_worker_mixin.py - unittest/executor/test_sleep_collective_rpc_guards.py + - unittest/executor/test_shutdown_retry_contracts.py - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_stop_word_side_output_trim.py @@ -262,6 +264,7 @@ l0_cpu: - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py + - unittest/llmapi/test_gms_args.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py # Reads the Muse Spark Open tokenizer files from LLM_MODELS_ROOT. diff --git a/tests/unittest/_torch/kernels/parallel/test_quantization_sm_budget.py b/tests/unittest/_torch/kernels/parallel/test_quantization_sm_budget.py new file mode 100644 index 000000000000..9e77ff37b6ce --- /dev/null +++ b/tests/unittest/_torch/kernels/parallel/test_quantization_sm_budget.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from utils.util import isSM100Family + +pytestmark = pytest.mark.skipif( + not isSM100Family(), reason="Requires supported SM100-family quantization kernels" +) + + +def _assert_same_bytes(actual, expected): + for actual_tensor, expected_tensor in zip(actual, expected): + assert actual_tensor.shape == expected_tensor.shape + assert actual_tensor.dtype == expected_tensor.dtype + torch.testing.assert_close( + actual_tensor.view(torch.uint8), expected_tensor.view(torch.uint8), atol=0, rtol=0 + ) + + +def _input(rows): + generator = torch.Generator(device="cuda").manual_seed(17) + return torch.randn((rows, 7168), device="cuda", dtype=torch.bfloat16, generator=generator) + + +def _rows_for_case(case, reserved_sms): + if case == "empty": + return 0 + if case == "small": + return 129 + if case == "capacity": + return 8192 + available_sms = torch.cuda.get_device_properties("cuda").multi_processor_count - reserved_sms + grid_edge = 4 * available_sms + return grid_edge if case == "grid_edge" else grid_edge + 1 + + +@pytest.mark.parametrize("row_case", ["small", "grid_edge", "grid_edge_plus_one", "capacity"]) +@pytest.mark.parametrize("swizzled", [False, True]) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_fp4_quantize_sm_budget_matches_default(row_case, swizzled, reserved_sms): + x = _input(_rows_for_case(row_case, reserved_sms)) + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + expected = torch.ops.trtllm.fp4_quantize(x, scale, 16, False, swizzled) + actual = torch.ops.trtllm.fp4_quantize.sm_budget( + x, scale, 16, False, swizzled, reserved_sms=reserved_sms + ) + _assert_same_bytes(actual, expected) + + +@pytest.mark.parametrize( + "row_case", ["empty", "small", "grid_edge", "grid_edge_plus_one", "capacity"] +) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_shared_fp8_quantize_sm_budget_matches_default(row_case, reserved_sms): + x = _input(_rows_for_case(row_case, reserved_sms)) + expected = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(x) + actual = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + x, reserved_sms=reserved_sms + ) + _assert_same_bytes(actual, expected) + + +@pytest.mark.parametrize("op", ["fp4", "fp8"]) +@pytest.mark.parametrize("budget", ["negative", "all_sms"]) +def test_quantize_sm_budget_rejects_invalid_budget(op, budget): + x = _input(129) + reserved_sms = ( + -1 + if budget == "negative" + else torch.cuda.get_device_properties(x.device).multi_processor_count + ) + with pytest.raises(RuntimeError, match="reserved_sms"): + if op == "fp4": + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + torch.ops.trtllm.fp4_quantize.sm_budget( + x, scale, 16, False, False, reserved_sms=reserved_sms + ) + else: + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget(x, reserved_sms=reserved_sms) + + +def test_shared_fp8_quantize_sm_budget_rejects_legacy_layout(): + with pytest.raises(RuntimeError, match="R128c4"): + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget( + _input(129), False, reserved_sms=8 + ) + + +@pytest.mark.parametrize("op", ["fp4", "fp8"]) +@pytest.mark.parametrize("reserved_sms", [0, 8]) +def test_quantize_sm_budget_fake_and_cuda_graph(op, reserved_sms): + x = _input(129) + if op == "fp4": + scale = torch.ones((1,), device="cuda", dtype=torch.float32) + quantize = torch.ops.trtllm.fp4_quantize.sm_budget + args = (x, scale, 16, False, False) + else: + quantize = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0.sm_budget + args = (x, True) + kwargs = {"reserved_sms": reserved_sms} + torch.library.opcheck(quantize, args, kwargs, test_utils=("test_schema", "test_faketensor")) + expected = quantize(*args, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = quantize(*args, **kwargs) + graph.replay() + _assert_same_bytes(actual, expected) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index d962962761d2..4cccfddfd060 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -46,6 +46,7 @@ DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, DeepseekV4Gate, + DeepseekV4MoE, DeepseekV4MTP, DeepseekV4WeightLoader, _copy_deepseek_v4_fused_a_weight_scale, @@ -678,6 +679,124 @@ def test_deepseek_v4_nvfp4_mixed_precision_config(): ) +@pytest.mark.parametrize( + ("quant_config_dict", "exclude_modules", "expected"), + [ + (None, None, True), + ( + {"model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None)}, + None, + False, + ), + ( + {"model.layers.4.mlp.shared_experts.down_proj": QuantConfig(quant_algo=None)}, + None, + False, + ), + (None, ["*shared_experts.gate_up_proj"], False), + (None, ["*shared_experts.down_proj"], False), + ( + { + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + "model.layers.4.mlp.shared_experts.up_proj": QuantConfig(quant_algo=None), + }, + None, + True, + ), + # The actual shared gate/up module is one fused Linear: its first + # gate_proj/gate_up_proj override wins and up-only overrides are ignored. + ( + { + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None), + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + }, + None, + False, + ), + ( + { + "model.layers.4.mlp.shared_experts.gate_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128 + ), + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig(quant_algo=None), + }, + None, + True, + ), + (None, ["*shared_experts.gate_proj"], False), + (None, ["*shared_experts.up_proj"], False), + ( + { + "model.layers.4.mlp.shared_experts.gate_up_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=64 + ), + "model.layers.4.mlp.shared_experts.down_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=64 + ), + }, + None, + False, + ), + ], +) +def test_deepseek_v4_shared_fc12_uses_final_projection_quantization( + quant_config_dict, exclude_modules, expected +): + model_config = ModelConfig( + pretrained_config=DeepseekV4Config(), + quant_config=QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + group_size=128, + exclude_modules=exclude_modules, + ), + quant_config_dict=quant_config_dict, + ) + + gate_up_quant, down_quant = DeepseekV4MoE._get_shared_expert_projection_quant_configs( + model_config, layer_idx=4 + ) + + assert DeepseekV4MoE._shared_fc12_quantization_supported(gate_up_quant, down_quant) is expected + + +@pytest.mark.parametrize(("supported", "expected"), [(False, False), (True, True)]) +def test_deepseek_v4_shared_fc12_uses_backend_capability(supported, expected): + backend = SimpleNamespace( + capabilities=SimpleNamespace(supports_per_layer_dynamic_eplb=supported) + ) + + assert DeepseekV4MoE._shared_fc12_backend_supported(backend) is expected + + +def test_deepseek_v4_shared_fc12_rejects_backend_without_capabilities(): + assert not DeepseekV4MoE._shared_fc12_backend_supported(SimpleNamespace()) + + +def test_deepseek_v4_shared_fc12_ignores_up_only_layerwise_override(): + global_quant_config = QuantConfig(quant_algo=None) + fp8_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) + model_config = ModelConfig( + pretrained_config=DeepseekV4Config(), + quant_config=global_quant_config, + quant_config_dict={ + "model.layers.4.mlp.shared_experts.up_proj": fp8_quant_config, + "model.layers.4.mlp.shared_experts.down_proj": fp8_quant_config, + }, + ) + + gate_up_quant, down_quant = DeepseekV4MoE._get_shared_expert_projection_quant_configs( + model_config, layer_idx=4 + ) + + assert gate_up_quant is global_quant_config + assert down_quant is fp8_quant_config + assert not DeepseekV4MoE._shared_fc12_quantization_supported(gate_up_quant, down_quant) + + def test_deepseek_v4_routed_moe_quant_config_from_mxfp4_header(tmp_path, monkeypatch): monkeypatch.setattr("tensorrt_llm._torch.model_config.get_sm_version", lambda: 100) tensor_name = "layers.0.ffn.experts.0.w1.weight" diff --git a/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py b/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py new file mode 100644 index 000000000000..702d62003101 --- /dev/null +++ b/tests/unittest/_torch/modules/test_megamoe_shared_mlp_scales.py @@ -0,0 +1,101 @@ +# 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. +"""UE8M0 block scales that the fused shared FC12 kernel reads on SM100 and SM103.""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.megamoe_shared_mlp import _unpack_r128c4 +from tensorrt_llm._utils import get_sm_version, is_sm_100f +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm.quantization.utils.fp8_utils import transform_k128_scales_to_cutedsl_mxfp8_layout + +pytestmark = pytest.mark.skipif( + not is_sm_100f() or get_sm_version() == 107, + reason="SM107 converts these scales to the CuTe DSL MXFP8 layout instead", +) + + +def _block_scaled_weights(out_features: int, in_features: int) -> dict[str, torch.Tensor]: + weight = torch.randn(out_features, in_features) / in_features**0.5 + blocks = weight.reshape(out_features // 128, 128, in_features // 128, 128) + scale = blocks.abs().amax(dim=(1, 3)) / torch.finfo(torch.float8_e4m3fn).max + quantized = (blocks / scale[:, None, :, None]).reshape(out_features, in_features) + return {"weight": quantized.to(torch.float8_e4m3fn), "weight_scale": scale.float()} + + +def _dequantize(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + expanded = scale.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1) + return weight.float() * expanded + + +@torch.inference_mode() +def test_linear_rounds_block_scales_to_ue8m0_when_asked(): + out_features, in_features = 256, 512 + weights = _block_scaled_weights(out_features, in_features) + linear = Linear( + in_features, + out_features, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES), + use_cute_dsl_blockscaling_mm=True, + ) + linear.round_block_scales_to_ue8m0 = True + linear.load_weights([weights]) + linear.post_load_weights() + linear.cuda() + + scale = linear.weight_scale + assert scale.dtype == torch.float32 + assert tuple(scale.shape) == (out_features // 128, in_features // 128) + assert linear.weight.dtype == torch.float8_e4m3fn + torch.testing.assert_close(scale, torch.exp2(torch.round(torch.log2(scale))), rtol=0, atol=0) + # Rounding a block's scale up to a power of two rescales its FP8 values, so the + # dequantized weight changes by at most FP8 rounding. + torch.testing.assert_close( + _dequantize(linear.weight, scale), + _dequantize(weights["weight"].cuda(), weights["weight_scale"].cuda()), + rtol=0.07, + atol=1e-3, + ) + # The fused kernel reads these scales through the CuTe DSL MXFP8 layout. + packed = transform_k128_scales_to_cutedsl_mxfp8_layout(scale, mn=out_features, k=in_features) + unpacked = _unpack_r128c4(packed, out_features, in_features // 32) + expected = scale.repeat_interleave(128, dim=0).repeat_interleave(4, dim=1) + torch.testing.assert_close( + unpacked, expected.to(torch.float8_e8m0fnu).view(torch.uint8), rtol=0, atol=0 + ) + + +@torch.inference_mode() +def test_linear_keeps_block_scales_without_the_request(): + out_features, in_features = 256, 512 + weights = _block_scaled_weights(out_features, in_features) + linear = Linear( + in_features, + out_features, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES), + use_cute_dsl_blockscaling_mm=True, + ) + linear.load_weights([weights]) + linear.post_load_weights() + linear.cuda() + + torch.testing.assert_close(linear.weight_scale.cpu(), weights["weight_scale"], rtol=0, atol=0) diff --git a/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py new file mode 100644 index 000000000000..0111b2a3f9e7 --- /dev/null +++ b/tests/unittest/_torch/moe/test_dynamic_eplb_contracts.py @@ -0,0 +1,2894 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fast CPU contracts for the per-layer dynamic EPLB serving path.""" + +from __future__ import annotations + +import ast +import dataclasses +import functools +import importlib.util +import math +import os +import re +import sys +import threading +from contextlib import contextmanager, nullcontext +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any, List, Optional, Tuple + +import pytest + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(__file__).resolve().parents[4] +# Worker shutdown contracts below cover the collective path, which only +# per-layer dynamic EPLB takes. +_PER_LAYER_DYNAMIC_SHUTDOWN = {"uses_collective_shutdown": lambda llm_args: True} +_TORCH_ROOT = _ROOT / "tensorrt_llm/_torch" +_QUEUE = "CUDA_SCALE_LAUNCH_QUEUES" +_SCHEDULER_RUNTIME = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/cuda_scheduler/runtime.py" +# The scheduler's native bound and the bound of the default 1800 s MoE +# communication warmup timeout at the nominal 2 GHz clock. +_SERVING_SPIN_CYCLES = 300 * 2_000_000_000 +_WARMUP_SPIN_CYCLES = 1800 * 2_000_000_000 + + +def _extract_class( + path: Path, + name: str, + methods: set[str] | None, + namespace: dict, + bases: Tuple[str, ...] = (), +): + tree = ast.parse(path.read_text()) + # Breadth-first, so a module-level class wins over a nested one of the same name. + original = next( + node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == name + ) + if methods is None: + selected = original + else: + body = [ + node + for node in original.body + if isinstance(node, ast.FunctionDef) and node.name in methods + ] + assert {node.name for node in body} == methods + selected = ast.ClassDef( + name=name, + bases=[ast.Name(id=base, ctx=ast.Load()) for base in bases], + keywords=[], + body=body, + decorator_list=[], + ) + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + selected, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace[name] + + +def _extract_function(path: Path, name: str, namespace: dict): + tree = ast.parse(path.read_text()) + function = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name + ) + function.decorator_list = [] + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + function, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + return namespace[name] + + +def test_shared_slot_provider_uses_model_owned_pool_state(): + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__"}, + {"Any": Any, "Optional": Optional, "Tuple": Tuple}, + ) + first_owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + second_owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + + first = provider_cls(first_owner) + same_model = provider_cls(first_owner) + second = provider_cls(second_owner) + + assert first._pools is same_model._pools + assert first._pool_order is same_model._pool_order + assert first._pools is not second._pools + with pytest.raises(TypeError, match="model-local pool state"): + provider_cls(SimpleNamespace()) + + +def test_shared_slot_provider_pool_key_includes_complete_bundle(): + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" + + class FakePool: + instances = [] + + def __init__(self, **kwargs): + self.bundle = kwargs["bundle"] + self._closed = False + self.group_sizes = () + self.instances.append(self) + + def build_layer(self, comm): + return object(), {}, object(), 0 + + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__", "build_hierarchical_live_arena"}, + { + "Any": Any, + "Dict": dict, + "Optional": Optional, + "Tuple": Tuple, + "_SharedSlotPool": FakePool, + }, + ) + + @dataclasses.dataclass(frozen=True) + class Plane: + name: str + nbytes: int + descriptor: str + + @dataclasses.dataclass(frozen=True) + class Bundle: + hidden: int + intermediate: int + planes: tuple[Plane, ...] + + # Equal names and byte sizes but different view descriptors must not share. + first_bundle = Bundle(160, 80, (Plane("weight", 12800, "storage-a"),)) + second_bundle = Bundle(160, 80, (Plane("weight", 12800, "storage-b"),)) + owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + kwargs = dict( + world=8, + rank=0, + home_count=4, + helper_count=1, + device=0, + comm=object(), + ) + first_provider = provider_cls(owner) + first_provider.build_hierarchical_live_arena(bundle=first_bundle, **kwargs) + with pytest.raises(RuntimeError, match="single-geometry"): + first_provider.build_hierarchical_live_arena(bundle=second_bundle, **kwargs) + provider_cls(owner).build_hierarchical_live_arena(bundle=second_bundle, **kwargs) + + assert len(owner._rebalance_shared_slot_pools) == 2 + assert [pool.bundle for pool in FakePool.instances] == [first_bundle, second_bundle] + + +@pytest.mark.parametrize("per_layer_helper_sets", [False, True]) +def test_shared_slot_provider_pool_key_includes_helper_set_mode(per_layer_helper_sets): + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" + + class FakePool: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self._closed = False + self.group_sizes = () + self.instances.append(self) + + def build_layer(self, comm): + return object(), {}, object(), 0 + + provider_cls = _extract_class( + source, + "SharedSlotArenaProvider", + {"__init__", "build_hierarchical_live_arena"}, + { + "Any": Any, + "Dict": dict, + "Optional": Optional, + "Tuple": Tuple, + "_SharedSlotPool": FakePool, + }, + ) + owner = SimpleNamespace( + _rebalance_shared_slot_pools={}, + _rebalance_shared_slot_pool_order=[], + ) + kwargs = dict(world=8, rank=0, home_count=4, helper_count=1, bundle="bundle", device=0) + provider_cls(owner, per_layer_helper_sets=per_layer_helper_sets).build_hierarchical_live_arena( + comm=object(), **kwargs + ) + provider_cls( + owner, per_layer_helper_sets=not per_layer_helper_sets + ).build_hierarchical_live_arena(comm=object(), **kwargs) + + # One pool per mode: a per-layer pool never hands out a shared bank. + assert [pool.kwargs["per_layer_sets"] for pool in FakePool.instances] == [ + per_layer_helper_sets, + not per_layer_helper_sets, + ] + assert len(owner._rebalance_shared_slot_pools) == 2 + + +def test_hierarchical_fused_submit_launches_bound_scheduler_and_copy(): + source = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/hierarchical.py" + native = [] + + def submit_fused(handle, routes, route_target, spin_cycles): + native.append(("launch", handle, routes, route_target, spin_cycles)) + + torch = SimpleNamespace( + ops=SimpleNamespace(trtllm=SimpleNamespace(moe_rebalance_tma_submit_fused=submit_fused)), + ) + event_handle = _extract_function(source, "_event_handle", {}) + broadcast_cls = _extract_class( + source, + "HierarchicalSamiWeightBroadcast", + { + "submit_fused", + "release_generation_after", + "_require_live_bank_reusable", + "mark_collective_reuse_safe", + }, + { + "torch": torch, + "_event_handle": event_handle, + "check_cuda": lambda result, operation: result, + }, + ) + outputs = object() + + broadcast = broadcast_cls.__new__(broadcast_cls) + broadcast._lock = nullcontext() + broadcast._handoff_aborted = False + broadcast._outputs = outputs + # The native context binds every scheduler argument; a launch passes only its routes. + broadcast._scheduler = SimpleNamespace() + authority = object() + broadcast._reuse_authority = authority + + def revalidate(_authority): + raise AssertionError("the bound authority was validated when it was bound") + + # The per-layer path checks the bound authority by identity only. + broadcast._validate_reuse_authority = revalidate + broadcast._active_generation = 0 + broadcast._collective_safe_generation = 0 + broadcast._local_release_generation = 0 + broadcast._generation = 1 + broadcast._stream = 19 + broadcast.device = 0 + broadcast.ctx = "ctx" + # The per-layer ops take the context's handle, not the custom-class object. + broadcast._handle = 4242 + broadcast.terminal_flags = "terminals" + broadcast.cuda = SimpleNamespace( + cuStreamWaitEvent=lambda stream, event, flags: native.append(("wait", stream, event)) + ) + + generation = broadcast.submit_fused(outputs, "routes", 8, spin_cycles=12345) + + # One native call orders the launch after the current stream and enqueues + # it on the bound copy stream; Python switches no stream. + assert native == [("launch", 4242, "routes", 8, 12345)] + assert generation == 1 + assert (broadcast._active_generation, broadcast._generation) == (1, 2) + # The bank stays leased until the consumer of generation 1 is proven done. + with pytest.raises(RuntimeError, match="still leased"): + broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) + with pytest.raises(RuntimeError, match="bound CUDA scheduler"): + broadcast.submit_fused(object(), "routes", 16, spin_cycles=12345) + + # Only the bound authority, after the local consumer release, frees the bank. + with pytest.raises(PermissionError, match="identity mismatch"): + broadcast.mark_collective_reuse_safe(object(), 1) + with pytest.raises(RuntimeError, match="local consumer release"): + broadcast.mark_collective_reuse_safe(authority, 1) + native.clear() + # Without an event, the consumer is the work queued on the current stream. + # The next launch orders the copy stream after that stream, so the release + # queues no device work. + broadcast.release_generation_after(1) + assert native == [] + with pytest.raises(RuntimeError, match="already has a local release"): + broadcast.release_generation_after(1) + broadcast.mark_collective_reuse_safe(authority, 1) + assert broadcast._collective_safe_generation == 1 + assert broadcast.submit_fused(outputs, "routes", 16, spin_cycles=12345) == 2 + + # A caller-recorded consumer event is waited for on the copy stream instead. + native.clear() + broadcast.release_generation_after(2, SimpleNamespace(cuda_event=23)) + assert native == [("wait", 19, 23)] + with pytest.raises(ValueError, match="positive CUDA event"): + broadcast.release_generation_after(3, SimpleNamespace(cuda_event=0)) + + +@pytest.mark.parametrize( + ( + "hidden", + "intermediate", + "expected_total_bytes", + "expected_stride_bytes", + "expected_planes", + ), + [ + ( + 7168, + 3072, + 37158924, + 37159424, + ( + ( + "mega_fc1_weight", + 0, + 22020096, + 512, + ((3, 6144, 3584), (22020096, 3584, 1), "uint8", 1, 16), + ( + (3, 3584, 6144), + (22020096, 1, 3584), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc1_weight_sf", + 22020096, + 2752512, + 512, + ((3, 2752512), (2752512, 1), "uint8", 1, 16), + ((3, 2752512), (2752512, 1), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + 24772608, + 11010048, + 512, + ((3, 7168, 1536), (11010048, 1536, 1), "uint8", 1, 16), + ( + (3, 1536, 7168), + (11010048, 1, 1536), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc2_weight_sf", + 35782656, + 1376256, + 512, + ((3, 1376256), (1376256, 1), "uint8", 1, 16), + ((3, 1376256), (1376256, 1), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 37158912, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc2_alpha", + 37158928, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 37158944, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ), + ), + ( + 160, + 80, + 24332, + 25088, + ( + ( + "mega_fc1_weight", + 0, + 12800, + 512, + ((3, 160, 80), (12800, 80, 1), "uint8", 1, 16), + ( + (3, 80, 160), + (12800, 1, 80), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc1_weight_sf", + 12800, + 3072, + 512, + ((3, 3072), (3072, 1), "uint8", 1, 16), + ((3, 3072), (3072, 1), "float8_e4m3fn", 1, 16), + ), + ( + "mega_fc2_weight", + 15872, + 6400, + 512, + ((3, 160, 40), (6400, 40, 1), "uint8", 1, 16), + ( + (3, 40, 160), + (6400, 1, 40), + "float4_e2m1fn_x2", + 1, + 16, + ), + ), + ( + "mega_fc2_weight_sf", + 22528, + 2048, + 512, + ((3, 2048), (2048, 1), "uint8", 1, 16), + ((3, 2048), (2048, 1), "float8_e4m3fn", 1, 16), + ), + ( + "fc31_alpha", + 24576, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc2_alpha", + 24592, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ( + "fc1_norm_const", + 24608, + 4, + 16, + ((3,), (1,), "float32", 4, 4), + ((3,), (1,), "float32", 4, 4), + ), + ), + ), + ], + ids=["production", "padding-sensitive"], +) +def test_vendored_live_plane_abi_matches_independent_goldens( + hidden, + intermediate, + expected_total_bytes, + expected_stride_bytes, + expected_planes, +): + path = _TORCH_ROOT / "moe/kernels/megamoe_scheduler/sami/geometry.py" + module_name = f"_live_plane_geometry_{hidden}_{intermediate}" + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(module_name, None) + + bundle = module.BundleLayout.create(hidden=hidden, intermediate=intermediate) + observed = tuple( + ( + plane.name, + plane.offset, + plane.nbytes, + plane.alignment, + plane.storage_metadata(3), + plane.kernel_metadata(3), + ) + for plane in bundle.planes + ) + assert observed == expected_planes + assert bundle.total_bytes == expected_total_bytes + assert bundle.stride_bytes == expected_stride_bytes + + +def test_ep_collective_adapter_uses_the_resolved_process_group(): + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/arena.py" + tree = ast.parse(source.read_text()) + selected = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "_EpComm" + ) + + calls = [] + + class FakeProcessGroup: + def __init__(self, rank, world, global_ranks): + self.rank = rank + self.world = world + self.global_ranks = global_ranks + + class FakeDistributed: + @staticmethod + def get_rank(*, group): + calls.append(("rank", group)) + return group.rank + + @staticmethod + def get_world_size(*, group): + calls.append(("world", group)) + return group.world + + @staticmethod + def get_global_rank(group, group_rank): + calls.append(("global_rank", group, group_rank)) + return group.global_ranks[group_rank] + + @staticmethod + def barrier(*, group): + calls.append(("barrier", group)) + + @staticmethod + def broadcast_object_list(values, *, src, group): + calls.append(("bcast", group, src)) + values[0] = ("from", src) + + @staticmethod + def all_gather_object(values, value, *, group): + calls.append(("allgather", group, value)) + for rank in range(len(values)): + values[rank] = (rank, value) + + namespace = { + "Any": Any, + "_disable_current_modes": nullcontext, + "torch": SimpleNamespace(distributed=FakeDistributed), + } + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + selected, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(source), "exec"), namespace) + process_group = FakeProcessGroup(rank=1, world=3, global_ranks=(11, 13, 17)) + comm = namespace["_EpComm"](process_group) + comm.validate_geometry(expected_rank=1, expected_world=3) + assert comm.allgather("handle") == [(0, "handle"), (1, "handle"), (2, "handle")] + assert comm.bcast("payload", 2) == ("from", 17) + comm.barrier() + assert all(call[1] is process_group for call in calls) + with pytest.raises(ValueError, match="rank/size differs"): + comm.validate_geometry(expected_rank=0, expected_world=3) + with pytest.raises(ValueError, match="broadcast root"): + comm.bcast("payload", 3) + + +def test_model_engine_warmup_uses_load_balancer_public_api(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "ModelEngine", + {"_set_warmup_effects", "warmup_scope"}, + { + "contextmanager": contextmanager, + "set_moe_comm_warmup": lambda value: events.append(("comm", value)), + }, + ) + manager = SimpleNamespace( + enable_statistic=True, + enable_update_weights=True, + set_iter_info=lambda **kwargs: events.append(("iter", kwargs)), + ) + engine = engine_cls() + engine.moe_load_balancer = manager + engine._warmup_depth = 0 + + with engine.warmup_scope(): + assert events == [ + ("comm", True), + ("iter", {"enable_statistic": False, "enable_update_weights": False}), + ] + + assert events[2:] == [ + ("comm", False), + ("iter", {"enable_statistic": True, "enable_update_weights": True}), + ] + + +def test_model_engine_cleanup_shuts_manager_after_graphs_while_model_is_alive(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "ModelEngine", + { + "prepare_cleanup", + "shutdown_moe_load_balancer", + "finish_cleanup", + "cleanup", + }, + {"release_gc": lambda: events.append("release_gc")}, + ) + engine = engine_cls() + model = object() + + class Manager: + def shutdown(self, *, local_safe): + assert engine.model is model + events.append(("manager.shutdown", local_safe)) + return local_safe + + engine._cleanup_done = False + engine.model_loader = None + engine._release_cuda_graphs = lambda **_: events.append("release_graphs") + engine.moe_load_balancer = Manager() + engine._runner = object() + engine._model_caller = object() + engine._mm_item_scheduler = object() + engine.model = model + engine.input_processor = object() + + engine.cleanup() + engine.cleanup() + + assert events == [ + "release_graphs", + ("manager.shutdown", True), + "release_gc", + ] + assert engine.moe_load_balancer is None + assert engine.model is None + + +def test_model_engine_prepare_failure_participates_in_unsafe_preflight(): + events = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "ModelEngine", + { + "prepare_cleanup", + "shutdown_moe_load_balancer", + "finish_cleanup", + "cleanup", + }, + {"release_gc": lambda: events.append("release_gc")}, + ) + + class Manager: + def shutdown(self, *, local_safe): + events.append(("manager.shutdown", local_safe)) + return local_safe + + engine = engine_cls() + engine._cleanup_done = False + engine.model_loader = None + engine._release_cuda_graphs = lambda **_: (_ for _ in ()).throw( + RuntimeError("graph cleanup failed") + ) + engine.moe_load_balancer = Manager() + engine._runner = object() + engine._model_caller = object() + engine._mm_item_scheduler = object() + engine.model = object() + engine.input_processor = object() + + with pytest.raises(RuntimeError, match="graph cleanup failed"): + engine.cleanup() + + assert events == [("manager.shutdown", False)] + assert engine.moe_load_balancer is not None + assert engine.model is not None + + +def test_model_engine_finalizer_never_enters_collective_cleanup(): + calls = [] + engine_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/model_engine.py", + "ModelEngine", + {"__del__"}, + {"logger": SimpleNamespace(warning=lambda *args: calls.append(("warning", args)))}, + ) + engine = engine_cls() + engine.cleanup = lambda *, collective: calls.append(("cleanup", collective)) + + engine.__del__() + + assert calls == [("cleanup", False)] + + +def test_executor_terminal_cleanup_is_three_phase_and_stream_scoped(monkeypatch): + events = [] + module_name = "tensorrt_llm._torch.moe.kernels.megamoe_shared_fc12" + monkeypatch.setitem( + sys.modules, + module_name, + SimpleNamespace( + release_shared_fc12_cache=lambda device, stream_handle=None: events.append( + ("cache", device, stream_handle) + ) + ), + ) + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"terminal_cleanup", "release_rank_local_caches"}, + { + "__name__": "tensorrt_llm._torch.pyexecutor.py_executor", + "__package__": "tensorrt_llm._torch.pyexecutor", + "sys": sys, + }, + ) + + class Engine: + def __init__(self, name): + self.name = name + + def prepare_cleanup(self): + events.append(("prepare", self.name)) + + def shutdown_moe_load_balancer(self, *, local_safe): + events.append(("collective", self.name, local_safe)) + return local_safe + + def finish_cleanup(self): + events.append(("finish", self.name)) + + executor = executor_cls() + executor.execution_stream = SimpleNamespace(device=SimpleNamespace(index=3), cuda_stream=99) + executor._terminal_model_engines = (Engine("first"), Engine("second")) + + executor.terminal_cleanup() + executor.terminal_cleanup() + + assert events == [ + ("prepare", "first"), + ("prepare", "second"), + ("collective", "first", True), + ("collective", "second", True), + ("finish", "first"), + ("finish", "second"), + ("cache", 3, 99), + ] + assert executor._terminal_model_engines == () + + +def test_executor_terminal_cleanup_propagates_prepare_failure_to_all_managers(monkeypatch): + events = [] + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"terminal_cleanup"}, + { + "__name__": "tensorrt_llm._torch.pyexecutor.py_executor", + "__package__": "tensorrt_llm._torch.pyexecutor", + }, + ) + + class Engine: + def __init__(self, name, fail=False): + self.name = name + self.fail = fail + + def prepare_cleanup(self): + events.append(("prepare", self.name)) + if self.fail: + raise RuntimeError(f"{self.name} graph failure") + + def shutdown_moe_load_balancer(self, *, local_safe): + events.append(("collective", self.name, local_safe)) + return local_safe + + def finish_cleanup(self): + events.append(("finish", self.name)) + + executor = executor_cls() + executor.execution_stream = SimpleNamespace(device=SimpleNamespace(index=3), cuda_stream=99) + executor._terminal_model_engines = (Engine("first", fail=True), Engine("second")) + + with pytest.raises(RuntimeError, match="first graph failure"): + executor.terminal_cleanup() + + assert events == [ + ("prepare", "first"), + ("prepare", "second"), + ("collective", "first", False), + ("collective", "second", False), + ] + assert len(executor._terminal_model_engines) == 2 + + +def test_pyexecutor_shutdown_tail_failure_is_retryable(): + events = [] + + class AsyncWorkerMixin: + pass + + class RetryingPool(dict): + def __init__(self): + super().__init__(pool=object()) + self.delete_attempts = 0 + + def __delitem__(self, key): + self.delete_attempts += 1 + events.append(("pool_delete", self.delete_attempts)) + if self.delete_attempts == 1: + raise RuntimeError("pool delete failed") + super().__delitem__(key) + + class Sampler(AsyncWorkerMixin): + def __init__(self): + self.stop_attempts = 0 + + def async_worker_enabled(self): + return True + + def async_worker_stop(self): + self.stop_attempts += 1 + events.append(("sampler_stop", self.stop_attempts)) + + class Engine: + def release_executor_resources(self, *, release_model_owned_graphs): + events.append("release_graphs") + + class DwdpManager: + def __init__(self): + self.exit_attempts = 0 + + def __exit__(self, *args): + self.exit_attempts += 1 + events.append(("dwdp_exit", self.exit_attempts)) + if self.exit_attempts == 1: + raise RuntimeError("dwdp exit failed") + + executor_cls = _extract_class( + _TORCH_ROOT / "pyexecutor/py_executor.py", + "PyExecutor", + {"shutdown", "_shutdown", "_finish_shutdown"}, + { + "AsyncWorkerMixin": AsyncWorkerMixin, + "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)), + }, + ) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"shutdown"}, + _PER_LAYER_DYNAMIC_SHUTDOWN, + ) + + engine = Engine() + executor = executor_cls() + executor.model_engine = engine + executor.draft_model_engine = None + executor.shutdown_all_ranks = True + executor.can_shutdown = lambda: True + executor.executor_request_queue = SimpleNamespace( + enqueue_shutdown_request=lambda: events.append("enqueue") + ) + executor.shutdown_event = SimpleNamespace(wait=lambda: events.append("wait")) + executor._profile_manager = SimpleNamespace(cleanup=lambda: events.append("profile_cleanup")) + executor.hang_detector = SimpleNamespace(detected=lambda: False) + executor.worker_thread = SimpleNamespace(join=lambda: events.append("worker_join")) + executor.kv_connector_manager = None + executor.enable_attention_dp = False + executor._kv_connector_failed = False + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = lambda: events.append("listener_shutdown") + executor.encoder_launch_executor = None + executor.worker_started = True + executor.resource_manager = SimpleNamespace( + resource_managers={ + "manager": SimpleNamespace(shutdown=lambda: events.append("resource_shutdown")) + } + ) + executor.virtual_memory_pools = RetryingPool() + executor.sampler = Sampler() + executor.dwdp_manager = DwdpManager() + + def terminal_cleanup(*, local_safe): + events.append(("terminal_cleanup", local_safe)) + if local_safe: + assert executor.model_engine is None + assert executor._terminal_model_engines == (engine,) + executor._terminal_model_engines = () + + executor.terminal_cleanup = terminal_cleanup + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = executor + + for message in ("pool delete failed", "dwdp exit failed"): + with pytest.raises(RuntimeError, match=message): + worker.shutdown() + assert executor.model_engine is engine + assert executor._terminal_model_engines == (engine,) + assert executor._shutdown_runtime_complete + assert worker.engine is executor + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "enqueue", + "wait", + "profile_cleanup", + "worker_join", + "listener_shutdown", + ("sampler_stop", 1), + "release_graphs", + "resource_shutdown", + ("pool_delete", 1), + ("terminal_cleanup", False), + ("pool_delete", 2), + ("dwdp_exit", 1), + ("terminal_cleanup", False), + ("dwdp_exit", 2), + ("terminal_cleanup", True), + ] + assert executor.virtual_memory_pools == {} + assert executor.model_engine is None + assert executor.draft_model_engine is None + assert executor.dwdp_manager is None + assert executor._shutdown_sampler_complete + assert executor._terminal_model_engines == () + assert worker.engine is None + + +class _Routes: + def __init__( + self, + rows, + trace, + *, + topk=6, + dtype="int32", + device="cuda:0", + contiguous=True, + ): + self.shape = (rows, topk) + self.ndim = 2 + self.dtype = dtype + self.device = device + self._contiguous = contiguous + self.trace = trace + + def is_contiguous(self): + return self._contiguous + + def get_device(self): + return int(self.device.split(":")[1]) if self.device.startswith("cuda:") else -1 + + def record_stream(self, stream): + self.trace.append(("lifetime", stream.cuda_stream)) + + def __getitem__(self, rows): + return _Routes(rows.stop, self.trace) + + +class _StageTensor: + def __init__(self, name, trace, shape=(8, 6)): + self.name = name + self.trace = trace + self.shape = shape + + def __getitem__(self, index): + return _StageTensor(self.name, self.trace, self.shape) + + def view(self, dtype): + return self + + def copy_(self, other, **kwargs): + self.trace.append("copy:" + self.name) + + def zero_(self): + self.trace.append("zero:" + self.name) + + def fill_(self, value): + self.trace.append("fill:" + self.name) + + +def test_launch_queue_is_enabled_only_for_per_layer_dynamic_mode(monkeypatch): + path = _ROOT / "tensorrt_llm/llmapi/_load_balance_env.py" + spec = importlib.util.spec_from_file_location("load_balance_env_under_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + initialized = False + fake_torch = SimpleNamespace( + cuda=SimpleNamespace(is_initialized=lambda: initialized), + ) + process_queue = os.environ.get(_QUEUE) + test_environ = os.environ.copy() + test_environ.pop(_QUEUE, None) + monkeypatch.setattr(module, "os", SimpleNamespace(environ=test_environ)) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + def config(mode="per_layer_dynamic"): + return SimpleNamespace( + load_balancer=SimpleNamespace(mode=mode), + resolve_load_balancer_compatibility=lambda: None, + ) + + original = {"KEEP": "value"} + configured = module.configure_moe_launch_queues(config(), original) + assert configured == {"KEEP": "value", _QUEUE: "4x"} + assert configured is not original + assert original == {"KEEP": "value"} + assert test_environ[_QUEUE] == "4x" + assert module.configure_moe_launch_queues(config(), configured) == configured + + for off in (None, config(mode="standard")): + test_environ.pop(_QUEUE, None) + before = {"KEEP": "value"} + assert module.configure_moe_launch_queues(off, before) is before + assert _QUEUE not in test_environ + + initialized = True + before = {"KEEP": "value"} + assert module.configure_moe_launch_queues(config(mode="standard"), before) is before + assert _QUEUE not in test_environ + + with pytest.raises(RuntimeError, match="before CUDA initialization"): + module.configure_moe_launch_queues(config(), original) + assert original == {"KEEP": "value"} + assert _QUEUE not in test_environ + + # A replayed override is too late for a context that already exists: + # the check reads the setting the process had before the replay. + test_environ[_QUEUE] = "4x" + with pytest.raises(RuntimeError, match="before CUDA initialization"): + module.configure_moe_launch_queues(config(), original, initial_setting=None) + assert module.configure_moe_launch_queues(config(), original, initial_setting="4x") == { + "KEEP": "value", + _QUEUE: "4x", + } + assert os.environ.get(_QUEUE) == process_queue + + +class _RuntimeInitReached(Exception): + """Stops an entrypoint at its first runtime-initialization call.""" + + +class _RecordingEnviron(dict): + def __init__(self, events, **values): + super().__init__(values) + self._events = events + + def __setitem__(self, key, value): + self._events.append(("env", key, value)) + super().__setitem__(key, value) + + def update(self, values): + for key, value in values.items(): + self[key] = value + + +def test_launch_queue_setup_sees_env_overrides_before_runtime_initialization(monkeypatch): + events = [] + environ = _RecordingEnviron(events, MASTER_ADDR="127.0.0.1") + fake_os = SimpleNamespace( + environ=environ, + getenv=lambda key, default=None: environ.get(key, default), + ) + logger = SimpleNamespace( + level="warning", + set_level=lambda level: None, + info=lambda *args, **kwargs: None, + error=lambda *args, **kwargs: None, + debug=lambda *args, **kwargs: None, + ) + user_overrides = {"USER_SETTING": "1"} + configured = {**user_overrides, _QUEUE: "4x"} + + def configure_moe_launch_queues(moe_config, env_overrides, *, initial_setting): + # The override must already be visible when the launch queue is configured, + # while the late-init check sees the setting from before the replay. + events.append(("configure", environ.get("USER_SETTING"), initial_setting)) + return {**(env_overrides or {}), _QUEUE: "4x"} + + def runtime_init(name): + def call(*args, **kwargs): + events.append(name) + raise _RuntimeInitReached(name) + + return call + + def run(entrypoint, stop): + events.clear() + environ.pop("USER_SETTING", None) + with pytest.raises(_RuntimeInitReached, match=stop): + entrypoint() + return list(events) + + def worker_args(): + return SimpleNamespace( + backend="pytorch", env_overrides=dict(user_overrides), moe_config=object() + ) + + expected = [("env", "USER_SETTING", "1"), ("configure", "1", None)] + common = { + "os": fake_os, + "logger": logger, + "configure_moe_launch_queues": configure_moe_launch_queues, + "LAUNCH_QUEUES_ENV": _QUEUE, + } + + # LLM API: argument validation is the first step that may probe CUDA. + class TorchLlmArgs: + model_fields = {"env_overrides": None, "moe_config": None} + + def __init__(self, **kwargs): + events.append(("llm_args", kwargs["env_overrides"])) + raise _RuntimeInitReached("llm_args") + + llm_cls = _extract_class( + _ROOT / "tensorrt_llm/llmapi/llm.py", + "BaseLLM", + {"__init__", "_process_env_overrides"}, + { + **common, + "GenerationExecutor": object, + "MoeConfig": lambda **kwargs: SimpleNamespace(**kwargs), + "TorchLlmArgs": TorchLlmArgs, + "TORCH_LLMARGS_REMOVED_ARGS": set(), + "mpi_disabled": lambda: False, + }, + ) + assert run( + lambda: llm_cls( + "model", backend="pytorch", env_overrides=dict(user_overrides), moe_config={} + ), + "llm_args", + ) == expected + [("llm_args", configured)] + + # MPI worker process: MPI is the first runtime initialization. + worker_main = _extract_function( + _ROOT / "tensorrt_llm/executor/worker.py", + "worker_main", + {**common, "GenerationExecutorWorker": object, "mpi_comm": runtime_init("mpi_comm")}, + ) + args = worker_args() + assert run( + lambda: worker_main(worker_queues=None, log_level="info", llm_args=args), + "mpi_comm", + ) == expected + ["mpi_comm"] + assert args.env_overrides == configured + + # Direct worker construction: the executor base initializes the runtime. + class GenerationExecutor: + __init__ = runtime_init("executor_init") + + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"__init__"}, + { + **common, + "GenerationExecutor": GenerationExecutor, + "PostprocWorkerConfig": lambda: SimpleNamespace( + num_postprocess_workers=0, postprocess_tokenizer_dir=None + ), + }, + bases=("GenerationExecutor",), + ) + args = worker_args() + assert run(lambda: worker_cls(llm_args=args), "executor_init") == expected + ["executor_init"] + assert args.env_overrides == configured + + # Ray worker: the first CUDA query follows launch-queue setup. + monkeypatch.setitem(sys.modules, "tensorrt_llm.logger", SimpleNamespace(logger=logger)) + ray_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayWorkerWrapper", + {"__init__"}, + { + **common, + "torch": SimpleNamespace( + cuda=SimpleNamespace(device_count=runtime_init("torch.cuda.device_count")) + ), + }, + ) + args = worker_args() + assert run(lambda: ray_cls(object, {"llm_args": args}, 1, 0), "device_count") == expected + [ + "torch.cuda.device_count" + ] + assert args.env_overrides == configured + + +@pytest.mark.parametrize("auxiliary_sms", [4, 8]) +def test_fused_submit_orders_generations_and_defers_the_route_wait(auxiliary_sms): + trace = [] + state = SimpleNamespace(device=0) + main = SimpleNamespace(cuda_stream=44, priority=0) + copy = SimpleNamespace(cuda_stream=19, priority=-1) + cuda = SimpleNamespace( + current_device=lambda: state.device, + current_stream=lambda _: main, + nvtx=SimpleNamespace(range=lambda _: nullcontext()), + ) + raw = SimpleNamespace( + _cuda_getDevice=lambda: state.device, + _cuda_getCurrentRawStream=lambda _: main.cuda_stream, + ) + torch = SimpleNamespace(cuda=cuda, _C=raw, Tensor=_Routes, int32="int32") + source = _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py" + namespace = {"torch": torch} + group_cls = _extract_class(source, "RebalanceSlotSchedulerGroup", None, namespace) + lease_cls = _extract_class(source, "_LiveBankLeaseProvider", None, namespace) + group = group_cls.__new__(group_cls) + for name in ( + "_generation", + "_finished_generation", + "_scheduled_generation", + "_route_wait_generation", + "_route_target", + "plan_calls", + ): + setattr(group, name, 0) + group._plan_part = None + group._pending_release = None + group._routes_in_flight = None + group._lent_rows = -1 + group._lent_routes = None + group._closed = False + group._execution_stream_handle = None + group.device = 0 + group.topk = 6 + group.max_tokens_per_rank = 8192 + group.auxiliary_sms = auxiliary_sms + group.copy_stream = group.scheduler_stream = copy + group._stream_handle = 19 + group._route_epoch_ptr = 77 + # Generation 1 launches during warmup, the later ones while serving. + bounds = iter((_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, _SERVING_SPIN_CYCLES)) + group._spin_bound = SimpleNamespace(cycles=lambda: next(bounds)) + + def driver_call(*entry): + trace.append(entry) + return (0,) + + # The planner records no event itself: the native launch and release + # order the copy stream after MAIN. + group._driver = SimpleNamespace( + CUresult=SimpleNamespace(CUDA_SUCCESS=0), + CUstreamWaitValue_flags=SimpleNamespace(CU_STREAM_WAIT_VALUE_GEQ="geq"), + cuStreamWaitValue64=lambda stream, address, value, flags: driver_call( + "wait_value", stream, address, value, flags + ), + ) + group._wait_value_geq = group._driver.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_GEQ + group._cuda_success = group._driver.CUresult.CUDA_SUCCESS + outputs = SimpleNamespace(physical_slot_ids=_Routes(8192, trace)) + + launch_bounds = [] + + def submit_fused(result, routes, target, *, spin_cycles): + assert result is outputs + generation = group._scheduled_generation + trace.append(("FUSED", routes.shape[0], target, generation)) + launch_bounds.append(spin_cycles) + return generation + + group.scheduler = SimpleNamespace(device="cuda:0", outputs=outputs) + group.broadcaster = SimpleNamespace( + submit_fused=submit_fused, + release_generation_after=lambda generation, event: trace.append( + ("consumer_release", generation, event) + ), + mark_collective_reuse_safe=lambda provider, generation: trace.append(("lease", generation)), + ) + group.lease = lease_cls(group) + + lent = {} + for generation, rows in enumerate((16, 16, 0), 1): + trace.clear() + routes = _Routes(rows, trace) + target = generation * auxiliary_sms + route_wait = ("wait_value", 44, 77, target, "geq") + part = group.plan_schedule(routes) + physical, actual_generation = group.plan_finish(part, defer_wait=True) + assert physical.shape == (rows, 6) + assert actual_generation == generation + # The view of the fixed output buffer is reused while the token count repeats. + assert lent.setdefault(rows, physical) is physical + assert route_wait not in trace + # The group, not an allocator stream event, keeps the routes alive + # until MAIN waits for every HALO-Q read of them. + assert group._routes_in_flight is routes + group.wait_for_routes() + assert group._routes_in_flight is None + group.wait_for_routes() + group.finish() + assert trace.count(route_wait) == 1 + assert not any(entry[0] == "lifetime" for entry in trace) + assert trace.index(("FUSED", rows, target, generation)) < trace.index(route_wait) + # Without an event, the release follows the MAIN work queued so far. + assert trace[-1] == ("consumer_release", generation, None) + if generation > 1: + # The previous consumer is proven reusable before the next launch. + assert trace.index(("lease", generation - 1)) < trace.index( + ("FUSED", rows, target, generation) + ) + assert launch_bounds == [_WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, _SERVING_SPIN_CYCLES] + + # The bound-stream fast path still rejects another stream or device. + main.cuda_stream = 45 + with pytest.raises(RuntimeError, match="MAIN stream cannot change"): + group.plan_schedule(_Routes(8, trace)) + main.cuda_stream = 44 + state.device = 1 + with pytest.raises(RuntimeError, match="bound CUDA device"): + group.plan_schedule(_Routes(8, trace)) + state.device = 0 + + trace.clear() + with pytest.raises(ValueError, match="contiguous CUDA int32"): + group.plan_schedule(_Routes(8, trace, dtype="int64")) + assert trace == [] + + # A submission that fails may have launched, so it keeps its routes alive, + # and the group refuses further plans rather than let the allocator reuse them. + group._pending_release = None + group._spin_bound = SimpleNamespace(cycles=lambda: _SERVING_SPIN_CYCLES) + group.broadcaster.submit_fused = lambda *args, **kwargs: 99 + stuck = _Routes(8, trace) + with pytest.raises(RuntimeError, match="generation mismatch"): + group.plan_schedule(stuck) + assert group._routes_in_flight is stuck + with pytest.raises(RuntimeError, match="never ordered on MAIN"): + group.plan_schedule(_Routes(8, trace)) + + # A closed group rejects every route wait, including one of a waited generation. + group._closed = True + with pytest.raises(RuntimeError, match="closed"): + group.wait_for_routes() + + +def test_rebalance_submission_binds_one_main_stream_for_every_thread(): + priorities = {19: -1, 44: 0, 45: 0, 46: -1} + current = SimpleNamespace(stream=46) + torch = SimpleNamespace( + _C=SimpleNamespace( + _cuda_getDevice=lambda: 0, + _cuda_getCurrentRawStream=lambda _: current.stream, + ), + cuda=SimpleNamespace( + current_stream=lambda _: SimpleNamespace( + cuda_stream=current.stream, priority=priorities[current.stream] + ) + ), + ) + group_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py", + "RebalanceSlotSchedulerGroup", + {"_check_owner"}, + {"torch": torch}, + ) + group = group_cls() + group._closed = False + group.device = 0 + group.copy_stream = SimpleNamespace(cuda_stream=19, priority=-1) + group._stream_handle = 19 + group._execution_stream_handle = None + + # Neither COPY nor a stream of its priority can become MAIN. + for stream in (46, 19): + current.stream = stream + with pytest.raises(RuntimeError, match="distinct from the higher-priority COPY"): + group._check_owner() + assert group._execution_stream_handle is None + + # The first other stream binds as MAIN, and any thread may submit on it: + # warmup and serving run on different threads but share one stream. + current.stream = 44 + group._check_owner() + assert group._execution_stream_handle == 44 + errors = [] + + def submit(): + try: + group._check_owner() + except BaseException as error: # noqa: BLE001 + errors.append(error) + + worker = threading.Thread(target=submit) + worker.start() + worker.join() + assert errors == [] + + current.stream = 45 + with pytest.raises(RuntimeError, match="MAIN stream cannot change"): + group._check_owner() + assert group._execution_stream_handle == 44 + + +def _module_constants(path: Path, names: set[str]) -> dict: + tree = ast.parse(path.read_text()) + constants = { + node.targets[0].id: ast.literal_eval(node.value) + for node in tree.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in names + } + assert set(constants) == names + return constants + + +def _scheduler_spin_namespace() -> dict: + namespace = _module_constants(_SCHEDULER_RUNTIME, {"SPIN_CYCLES", "ASSUMED_CLOCK_HZ"}) + _extract_function(_SCHEDULER_RUNTIME, "spin_cycles_for_timeout", namespace) + return namespace + + +def _moe_comm_timeout_guard_module(monkeypatch) -> ModuleType: + """Execute the timeout guard module without its torch and logger imports.""" + path = _TORCH_ROOT / "moe/fused_moe/moe_comm_timeout_guard.py" + tree = ast.parse(path.read_text()) + tree.body = [ + node + for node in tree.body + if not ( + (isinstance(node, ast.Import) and node.names[0].name == "torch") + or (isinstance(node, ast.ImportFrom) and node.module == "tensorrt_llm.logger") + ) + ] + module = ModuleType("moe_comm_timeout_guard_under_test") + module.logger = SimpleNamespace(info=lambda *args, **kwargs: None) + module.torch = None + # Dataclasses resolve their module through sys.modules. + monkeypatch.setitem(sys.modules, module.__name__, module) + exec(compile(tree, str(path), "exec"), module.__dict__) + return module + + +def test_scheduler_spin_cycles_convert_timeouts_at_the_nominal_clock(): + namespace = _scheduler_spin_namespace() + to_cycles = namespace["spin_cycles_for_timeout"] + + assert namespace["SPIN_CYCLES"] == _SERVING_SPIN_CYCLES + assert to_cycles(None) == _SERVING_SPIN_CYCLES + assert to_cycles(1800) == _WARMUP_SPIN_CYCLES + for seconds in (0, -1, 2.5, True): + with pytest.raises(ValueError, match="positive exact int"): + to_cycles(seconds) + + +@pytest.mark.parametrize( + ("environ", "warmup_cycles", "serving_cycles"), + [ + pytest.param({}, _WARMUP_SPIN_CYCLES, _SERVING_SPIN_CYCLES, id="defaults"), + pytest.param( + { + "TRTLLM_MOE_COMM_WARMUP_TIMEOUT_SEC": "600", + "TRTLLM_MOE_COMM_TIMEOUT_SEC": "30", + }, + 600 * 2_000_000_000, + 30 * 2_000_000_000, + id="configured", + ), + ], +) +def test_scheduler_spin_bound_follows_warmup_transitions_and_autotuning( + monkeypatch, environ, warmup_cycles, serving_cycles +): + guard_module = _moe_comm_timeout_guard_module(monkeypatch) + guard = guard_module.MoECommTimeoutGuard(environ=environ, is_capturing=lambda: False) + runtime = SimpleNamespace(**_scheduler_spin_namespace()) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm._torch.moe.kernels.megamoe_scheduler.cuda_scheduler.runtime", + runtime, + ) + bound_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/planner.py", + "SchedulerSpinBound", + None, + {"Callable": Any, "Optional": Optional}, + ) + tuning = [False] + bound = bound_cls(warmup_seconds=guard.budgets().warmup_seconds, is_tuning=lambda: tuning[0]) + guard.register(bound) + + # The engine starts in serving; its warmup transitions switch the bound. + assert bound.cycles() == serving_cycles + guard.set_warmup(True) + assert bound.cycles() == warmup_cycles + guard.set_warmup(False) + assert bound.cycles() == serving_cycles + # AutoTuner profiling outside warmup gets the warmup bound too. + tuning[0] = True + assert bound.cycles() == warmup_cycles + + +def _tma_op_parameters(name: str) -> List[str]: + """Parameter names of a registered moe_rebalance_tma operator schema.""" + op_source = (_ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp").read_text() + schemas = ["".join(re.findall(r'"([^"]*)"', block)) for block in op_source.split("m.def(")[1:]] + schema = next(text for text in schemas if text.startswith(f"{name}(")) + parameters = re.fullmatch(rf"{name}\((.*)\) -> \(\)", schema).group(1) + return [parameter.split()[-1] for parameter in parameters.split(",")] + + +def test_fused_launch_binds_shared_arguments_and_passes_only_routes(): + # Map the scheduler arguments onto the registered operator schemas rather than hard-coding them. + bind_names = _tma_op_parameters("moe_rebalance_tma_bind_fused") + scheduler_names = bind_names[bind_names.index("out_slots") : bind_names.index("copy_stream")] + # A launch passes its routes, route target and spin bound; the bind holds the rest. + # They take the context's handle, since a custom-class argument costs + # microseconds of dispatch per call. + assert _tma_op_parameters("moe_rebalance_tma_submit_fused") == [ + "handle", + "routes", + "route_target", + "spin_cycles", + ] + assert _tma_op_parameters("moe_rebalance_tma_release") == ["handle"] + + scheduler_cls = _extract_class( + _SCHEDULER_RUNTIME, + "CudaPhysicalSlotScheduler", + {"fused_launch_args"}, + {}, + ) + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.cfg = SimpleNamespace( + ep_size=8, + logical_expert_count=384, + extra_slots_per_rank=4, + local_rank=0, + route_count=8192 * 6, + ctas=4, + threads=512, + topk=6, + max_tokens_per_rank=8192, + ) + scheduler.outputs = SimpleNamespace( + physical_slot_ids="slots", + hot_expert_ids="ids", + hot_expert_group_level="levels", + hot_expert_source_ranks="owners", + ) + for name in ( + "peer_base", + "status", + "partial", + "route_aux", + "grid_sync", + "plan_workspace", + "route_prefix", + ): + setattr(scheduler, name, name) + scheduler._bound_stream = None + with pytest.raises(RuntimeError, match="bound GPU-direct copy"): + scheduler.fused_launch_args() + scheduler._bound_stream = 19 + + args = scheduler.fused_launch_args() + assert len(args) == len(scheduler_names) + expected = { + "out_slots": "slots", + "out_ids": "ids", + "out_levels": "levels", + "out_owners": "owners", + "peer_bases": "peer_base", + "route_capacity": 8192 * 6, + "topk": 6, + "ctas": 4, + "threads": 512, + "ep": 8, + "experts": 384, + "helpers": 4, + "local_rank": 0, + } + for name, value in expected.items(): + assert args[scheduler_names.index(name)] == value + + +def _cpp_method_body(source: str, signature: str) -> str: + """Body of the C++ function whose definition starts with ``signature``.""" + open_brace = source.index("{", source.index(signature)) + depth = 0 + for index in range(open_brace, len(source)): + if source[index] == "{": + depth += 1 + elif source[index] == "}": + depth -= 1 + if depth == 0: + return source[open_brace + 1 : index] + raise AssertionError(f"unbalanced body of {signature}") + + +def _assert_in_order(body: str, calls: Tuple[str, ...]) -> None: + positions = [body.index(call) for call in calls] + assert positions == sorted(positions), calls + + +def test_native_launch_and_release_order_the_copy_stream_after_main(): + # The cross-stream ordering lives in the native context, so pin its call order: + # HALO-Q must not read routes before MAIN wrote them, and the next generation's + # copy must not overwrite the live bank before MAIN's consumer of it. + source = (_ROOT / "cpp/tensorrt_llm/thop/moe/loadBalance/moeRebalanceTmaOp.cpp").read_text() + submit = _cpp_method_body(source, "void submitFused(torch::Tensor const& routes") + _assert_in_order( + submit, + ( + "TORCH_CHECK_VALUE(routes.is_cuda()", + "TORCH_CHECK_VALUE(spinCycles > 0", + "cudaStream_t const execution = currentExecutionStream();", + "cudaEventRecord(mInputReady, execution)", + "cudaStreamWaitEvent(mCopyStream, mInputReady, 0)", + "megamoe_tma_copy_submit_fused(", + ), + ) + launch = submit[submit.index("megamoe_tma_copy_submit_fused(") :] + assert "reinterpret_cast(mCopyStream)" in launch + release = _cpp_method_body(source, "void releaseAfterCurrentStream()") + _assert_in_order( + release, + ( + "cudaStream_t const execution = currentExecutionStream();", + "cudaEventRecord(mConsumerDone, execution)", + "cudaStreamWaitEvent(mCopyStream, mConsumerDone, 0)", + ), + ) + # The stream a launch or release follows can never be the copy stream itself. + current = _cpp_method_body(source, "cudaStream_t currentExecutionStream() const") + assert "TORCH_CHECK(stream != mCopyStream" in current + + +def test_megamoe_cute_weight_cache_keys_every_view_field(): + conversions = [] + + def to_cute(t, assumed_align): + conversions.append((t, assumed_align)) + return ("cute", len(conversions)) + + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + namespace = { + "torch": SimpleNamespace(Tensor=object), + "Tuple": Tuple, + "_cute_launch_helpers": lambda: (to_cute,), + } + _extract_function(source, "_cute_view_key", namespace) + to_cute_cached = _extract_function(source, "_to_cute_cached", namespace) + + def weight(ptr=4096, shape=(8, 16), stride=(16, 1), dtype="fp4", device=0): + return SimpleNamespace( + data_ptr=lambda: ptr, + shape=shape, + stride=lambda: stride, + dtype=dtype, + device=SimpleNamespace(index=device), + ) + + cache = {} + first = to_cute_cached(cache, weight()) + # A new view object of the same memory and layout reuses the cute view. + assert to_cute_cached(cache, weight()) is first + assert len(conversions) == 1 + for changed in ( + weight(ptr=8192), + weight(shape=(16, 8)), + weight(stride=(1, 8)), + weight(dtype="fp8"), + weight(device=1), + ): + assert to_cute_cached(cache, changed) is not first + assert to_cute_cached(cache, weight(), assumed_align=4) is not first + assert len(conversions) == 7 + # The cache belongs to the caller: another module's dict starts empty. + assert to_cute_cached({}, weight()) is not first + + # A slot keeps only the latest view of a buffer whose live rows change. + to_cute_slot = _extract_function(source, "_to_cute_slot", namespace) + slots = {} + conversions.clear() + rows_8 = to_cute_slot(slots, "output", weight(shape=(8, 16))) + assert to_cute_slot(slots, "output", weight(shape=(8, 16))) is rows_8 + rows_4 = to_cute_slot(slots, "output", weight(shape=(4, 16))) + assert rows_4 is not rows_8 and len(slots) == 1 + assert to_cute_slot(slots, "output", weight(shape=(8, 16))) is not rows_8 + assert len(conversions) == 3 + + +def test_megamoe_cute_pointer_cache_keys_address_and_alignment(): + made = [] + + def to_cute_ptr(t, assumed_align): + made.append((t.data_ptr(), assumed_align)) + return ("ptr", len(made)) + + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + cache = {} + namespace = { + "torch": SimpleNamespace(Tensor=object), + "_cute_launch_helpers": lambda: (None, to_cute_ptr), + "_MEGAMOE_CUTE_PTR_CACHE": cache, + "_MEGAMOE_CUTE_PTR_CACHE_MAX_ENTRIES": 2, + } + to_cute_ptr_cached = _extract_function(source, "_to_cute_ptr_cached", namespace) + + def buffer(ptr): + return SimpleNamespace(data_ptr=lambda: ptr) + + first = to_cute_ptr_cached(buffer(4096)) + # Another tensor at the same address gets the same pointer. + assert to_cute_ptr_cached(buffer(4096)) is first + assert to_cute_ptr_cached(buffer(4096), assumed_align=4) is not first + assert len(made) == 2 + # The cache is bounded: a third entry starts it over. + to_cute_ptr_cached(buffer(8192)) + assert len(cache) == 1 + assert to_cute_ptr_cached(buffer(4096)) is not first + + +def test_executors_of_one_engine_share_the_execution_stream(): + # The planner binds the first MAIN stream for good, so the KV-cache + # estimation executor and the final executor must share one stream. + tree = ast.parse((_TORCH_ROOT / "pyexecutor/py_executor_creator.py").read_text()) + create = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_create_py_executor" + ) + calls = [ + node + for node in ast.walk(create) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "create_py_executor_instance" + ] + assert len(calls) >= 2 + for call in calls: + stream = next( + (keyword.value for keyword in call.keywords if keyword.arg == "execution_stream"), + None, + ) + assert isinstance(stream, ast.Name) and stream.id == "execution_stream" + + +def test_route_wait_occurs_after_independent_staging(): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" + impl = _extract_class( + source, + "TrtllmCutedslMegaMoeNvfp4Impl", + {"_stage_inputs"}, + {"torch": SimpleNamespace(uint8="uint8")}, + ) + + def run(): + trace = [] + owner = impl() + owner._last_staged_T = {8: 8} + owner._wait_rebalance_routes = lambda replica_plan=None: trace.append("wait") + bufs = SimpleNamespace( + **{ + key: _StageTensor(key, trace) + for key in ("topk_idx_local", "activation", "activation_sf", "topk_weights") + } + ) + owner._stage_inputs( + bufs=bufs, + x=_StageTensor("x", trace), + x_sf=_StageTensor("sf", trace), + topk_idx=_StageTensor("routes", trace), + topk_weights=_StageTensor("weights", trace), + num_tokens=4, + top_k=6, + stage_activation=True, + ) + return trace, owner + + trace, owner = run() + assert trace.index("copy:activation") < trace.index("wait") + assert trace.index("copy:topk_weights") < trace.index("wait") + assert trace.index("wait") < trace.index("copy:topk_idx_local") + assert owner._last_staged_T[8] == 4 + + +@pytest.mark.parametrize("has_sm_budget", [True, False]) +def test_megamoe_quantize_reserves_sms_only_for_a_replica_plan(has_sm_budget): + source = _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py" + calls = [] + + class Activation: + def __init__(self, rows, cols): + self.shape = (rows, cols) + + def to(self, dtype): + return self + + def contiguous(self): + return self + + def view(self, *shape): + return shape + + def fp4_quantize(*args): + calls.append(("default", *args[1:])) + return "x_fp4", Activation(3, 4) + + def sm_budget(*args): + calls.append(("sm_budget", *args[1:])) + return "x_fp4", Activation(3, 4) + + if has_sm_budget: + fp4_quantize.sm_budget = sm_budget + namespace = { + "torch": SimpleNamespace( + bfloat16="bf16", ops=SimpleNamespace(trtllm=SimpleNamespace(fp4_quantize=fp4_quantize)) + ), + "logger": SimpleNamespace(warning=lambda message: None), + "Fp4QuantizedTensor": type("Fp4QuantizedTensor", (), {}), + "megamoe_activation_sf_bytes_per_row": lambda hidden: hidden // 16, + "ceil_div": lambda a, b: -(-a // b), + } + _extract_function(source, "_fp4_quantize_has_sm_budget", namespace) + impl = _extract_class(source, "TrtllmCutedslMegaMoeNvfp4Impl", {"quantize_input"}, namespace) + owner = impl() + owner.fc31_input_scale = "scale" + + owner.quantize_input(Activation(3, 64)) + owner.quantize_input(Activation(3, 64), reserved_sms=4) + + default = ("default", "scale", 16, False, False) + budget = ("sm_budget", "scale", 16, False, False, 4) if has_sm_budget else default + assert calls == [default, budget] + + +def test_megamoe_launch_buffers_are_built_once_per_bucket_and_mode(): + inference = SimpleNamespace(enabled=True) + fake_torch = SimpleNamespace(is_inference_mode_enabled=lambda: inference.enabled) + impl = _extract_class( + _TORCH_ROOT / "moe/fused_moe/mega_moe/mega_moe_cute_dsl.py", + "TrtllmCutedslMegaMoeNvfp4Impl", + {"_acquire_buffers"}, + {"torch": fake_torch}, + ) + built = [] + owner = impl() + owner._launch_buffers = {} + + def build_buffers(**kwargs): + built.append((kwargs["launch_max_T"], inference.enabled)) + return object() + + owner._build_buffers = build_buffers + args = dict(top_k=6, hidden=7168, device="cuda:0", output_dtype="bf16") + + first = owner._acquire_buffers(**args, launch_max_T=8192) + assert owner._acquire_buffers(**args, launch_max_T=8192) is first + assert owner._acquire_buffers(**args, launch_max_T=4096) is not first + # Views made under inference mode are kept apart from normal-mode views. + inference.enabled = False + assert owner._acquire_buffers(**args, launch_max_T=8192) is not first + assert built == [(8192, True), (4096, True), (8192, False)] + + +def test_megamoe_symmetric_regions_are_built_once_per_mode(): + inference = SimpleNamespace(enabled=True) + provider_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "MegaMoeSymmMemProvider", + {"get_regions"}, + {"torch": SimpleNamespace(is_inference_mode_enabled=lambda: inference.enabled)}, + ) + provider = provider_cls() + provider._regions = {} + built = [] + provider._build_regions = lambda: built.append(inference.enabled) or object() + + regions = provider.get_regions() + assert provider.get_regions() is regions + inference.enabled = False + assert provider.get_regions() is not regions + assert built == [True, False] + + +def test_megamoe_tuning_config_is_built_once_per_runner(): + configs = [] + + def tuning_config(**kwargs): + configs.append(kwargs) + return SimpleNamespace(**kwargs) + + runner_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "Sm100MegaMoENvfp4Runner", + {"get_tuning_config", "_autotuner_inputs_pre_hook"}, + { + "TuningConfig": tuning_config, + "DynamicTensorSpec": lambda *args: args, + "ConstraintSpec": lambda *args: args, + "DistributedTuningStrategy": SimpleNamespace(MERGE="merge"), + "get_last_power_of_2_num_tokens_buckets": None, + "last_positive_power_of_2": None, + "_megamoe_autotune_num_tokens": None, + }, + ) + first, second = runner_cls(), runner_cls() + for runner in (first, second): + runner.max_tokens_per_rank = 8192 + + config = first.get_tuning_config() + assert first.get_tuning_config() is config + # The profiling pre-hook stays bound to the runner that built the config. + assert config.inputs_pre_hook.__self__ is first + assert second.get_tuning_config().inputs_pre_hook.__self__ is second + assert len(configs) == 2 + + +def test_megamoe_runner_identity_is_built_once(): + runner_cls = _extract_class( + _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py", + "Sm100MegaMoENvfp4Runner", + {"unique_id"}, + {}, + ) + runner = runner_cls() + builds = [] + runner._build_unique_id = lambda: builds.append(1) or (100, 8, 6) + + assert runner.unique_id() == (100, 8, 6) + assert runner.unique_id() == (100, 8, 6) + assert builds == [1] + + +def _autotune_helpers(torch): + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + functions = { + "_autotune_pl_alpha", + "synthesize_profiling_topk", + "_token_back_ready_granularity_choices", + "_megamoe_tuning_op_name", + "_unpack_tactic", + "_is_pow2_in_range", + "validate_megamoe_tactic", + "_epi_flag_batch_for_tokens", + "_expand_megamoe_ready_tactics", + "enumerate_megamoe_candidate_tactics", + } + constants = { + "_NVFP4_BLOCK_SIZE", + "_SUPPORTED_MMA_TILE_M", + "_SUPPORTED_MMA_TILE_N", + "_AUTOTUNE_PL_SEED", + "_DEFAULT_TOKEN_BACK_READY_GRANULARITY", + "_TOKEN_BACK_READY_GRANULARITIES", + "_TACTIC_LEN", + "_TACTIC_LEN_V4", + "_LEGACY_TACTIC_LEN", + "_FLAG_BATCH_MAX", + "_GEOMETRIES", + "_TOKEN_BACK_MODES", + "_TOKEN_BACK_STORE_BINDING", + "_WORK_ID_MODE_CANDIDATES", + "_GROUP_HINTS", + "_FLAG_BATCHES", + "_EPI_FLAG_BATCH_SMALL", + "_EPI_FLAG_BATCH_LARGE", + "_EPI_FLAG_BATCH_TOKEN_THRESHOLD", + } + tree = ast.parse(source.read_text()) + nodes = [ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)] + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name in functions: + nodes.append(node) + elif isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + if any(isinstance(target, ast.Name) and target.id in constants for target in targets): + nodes.append(node) + namespace = { + "torch": torch, + "math": math, + "functools": functools, + "Tuple": Tuple, + "List": List, + "Optional": Optional, + "Any": Any, + "logger": SimpleNamespace(debug=lambda *args: None), + } + exec( + compile( + ast.fix_missing_locations(ast.Module(body=nodes, type_ignores=[])), str(source), "exec" + ), + namespace, + ) + return namespace + + +def test_megamoe_runner_cache_reuses_only_production_runners(): + source = _TORCH_ROOT / "moe/custom_ops/cute_dsl_megamoe_custom_op.py" + tree = ast.parse(source.read_text()) + factory_node = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_megamoe_get_runner" + ) + + created = [] + + class FakeRunner: + def __init__(self, **kwargs): + self.kwargs = kwargs + self._profiling_scratch = object() + self._profiling_scratch_factory = object() + created.append(self) + + tuner = SimpleNamespace(is_tuning_mode=False) + namespace = { + "AutoTuner": SimpleNamespace(get=lambda: tuner), + "Sm100MegaMoENvfp4Runner": FakeRunner, + "_MEGAMOE_RUNNER_CACHE": {}, + } + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + factory_node, + ], + type_ignores=[], + ) + exec(compile(ast.fix_missing_locations(module), str(source), "exec"), namespace) + get_runner = namespace["_megamoe_get_runner"] + kwargs = dict( + world_size=8, + local_rank=2, + num_topk=6, + num_experts_per_rank=32, + hidden_size=7168, + intermediate_size_per_partition=2048, + expand_intermediate_size_per_partition=4096, + max_tokens_per_rank=8192, + output_dtype="bf16", + apply_topk_in_fc1=True, + swiglu_alpha=None, + swiglu_beta=None, + gate_up_clamp=None, + situ_beta=None, + situ_linear_beta=None, + in_kernel_fc2_reduce=False, + combine_format="bf16", + helper_expert_count=4, + ) + + production = get_runner(**kwargs, tactic_autotune=False) + production._profiling_scratch = "stale" + production._profiling_scratch_factory = "stale" + assert get_runner(**kwargs, tactic_autotune=False) is production + assert production._profiling_scratch is None + assert production._profiling_scratch_factory is None + + # While the AutoTuner profiles, an autotune opt-in gets a fresh runner + # that owns its profiling scratch. + tuner.is_tuning_mode = True + first_tuning = get_runner(**kwargs, tactic_autotune=True) + second_tuning = get_runner(**kwargs, tactic_autotune=True) + assert first_tuning is not second_tuning + assert first_tuning is not production + assert len(created) == 3 + + # Outside tuning mode it only looks up its tactic: one runner is reused, + # separate from the opted-out production runner. + tuner.is_tuning_mode = False + serving = get_runner(**kwargs, tactic_autotune=True) + serving._profiling_scratch = "stale" + assert get_runner(**kwargs, tactic_autotune=True) is serving + assert serving is not production + assert serving._profiling_scratch is None + assert get_runner(**kwargs, tactic_autotune=False) is production + assert len(created) == 4 + + +def test_on_autotune_uses_distinct_tactics(): + helpers = _autotune_helpers(None) + + off_key = helpers["_megamoe_tuning_op_name"](0) + on_key = helpers["_megamoe_tuning_op_name"](4) + assert off_key == "trtllm::cute_dsl_megamoe_nvfp4_blackwell" + assert on_key != off_key and "alpha=0.8" in on_key + + off = helpers["enumerate_megamoe_candidate_tactics"](32, 107) + on = helpers["enumerate_megamoe_candidate_tactics"](32, 107, load_balance=True) + assert all(helpers["_unpack_tactic"](tactic)[10] == "expert" for tactic in off) + assert {helpers["_unpack_tactic"](tactic)[10] for tactic in on} == { + "expert", + "token_tile", + } + assert any(tactic[2] is not None for tactic in on) + + +def test_on_autotune_builds_balanced_powerlaw_input(request): + torch = pytest.importorskip("torch") + original_num_threads = torch.get_num_threads() + request.addfinalizer(lambda: torch.set_num_threads(original_num_threads)) + torch.set_num_threads(4) + helpers = _autotune_helpers(torch) + kwargs = dict( + num_tokens=8192, + num_topk=6, + num_experts_per_rank=52, + world_size=8, + alpha=0.8, + device=torch.device("cpu"), + source_rank=0, + ) + routes = helpers["synthesize_profiling_topk"](**kwargs) + assert torch.equal(routes, helpers["synthesize_profiling_topk"](**kwargs)) + counts = torch.bincount(routes.reshape(-1), minlength=8 * 52).view(8, 52) + assert bool((counts[:, 1:] >= counts[:, :-1]).all()) + per_rank = counts.sum(1) + assert int(per_rank.max() - per_rank.min()) <= 1 + assert bool((counts[:, -1] > 0).all()) + + uniform = helpers["synthesize_profiling_topk"](**{**kwargs, "alpha": 0.0}) + expected = (torch.arange(8192 * 6) % (8 * 52)).view(8192, 6) + assert torch.equal(uniform, expected) + + +def test_per_layer_dynamic_shutdown_reports_in_iter_through_collective_preflight(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/manager.py", + "PerLayerDynamicMoeLoadBalancer", + {"shutdown"}, + {}, + ) + calls = [] + + class Comm: + def allgather(self, state): + calls.append(("allgather", state)) + return [state, {"local_safe": True, "pool_count": 1}] + + class Pool: + def close_collectively(self, *, local_safe): + calls.append(("pool", local_safe)) + return True + + manager = manager_cls() + manager.is_shutdown = False + manager.in_iter = True + manager._ep_comm = Comm() + manager._rebalance_shared_slot_pool_order = [Pool()] + + assert manager.shutdown(local_safe=True) is False + assert calls == [ + ("allgather", {"local_safe": False, "pool_count": 1}), + ] + + +def test_post_moe_hook_releases_a_replica_plan_every_chunk(): + moe_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/interface.py", + "MoE", + {"_load_balancer_start_set_cpu_stage", "_load_balancer_uses_replica_plan"}, + {}, + ) + moe_cls._using_dynamic_load_balancer = _extract_function( + _TORCH_ROOT / "moe/fused_moe/impl_blocks.py", "_using_dynamic_load_balancer", {} + ) + started = [] + + class Layer: + def __init__(self, *, produces_replica_plan, dynamic_routing): + self.produces_replica_plan = produces_replica_plan + self.dynamic_routing = dynamic_routing + + def is_dynamic_routing(self): + return self.dynamic_routing + + def start_set_cpu_stage(self): + started.append(self) + + cases = [ + # A per-layer dynamic layer disables standard host migration. + (Layer(produces_replica_plan=True, dynamic_routing=False), [True, True]), + (Layer(produces_replica_plan=False, dynamic_routing=True), [False, True]), + (Layer(produces_replica_plan=False, dynamic_routing=False), [False, False]), + (None, [False, False]), + ] + for layer, expected in cases: + moe = moe_cls() + moe.layer_load_balancer = layer + observed = [] + for is_last_call in (False, True): + started.clear() + moe._load_balancer_start_set_cpu_stage(is_last_call) + observed.append(started == [layer]) + assert observed == expected + + +def test_iteration_context_aborts_the_iteration_of_a_failed_forward(): + source = _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py" + context_cls = _extract_class(source, "MoeLoadBalancerIterContext", None, {}) + base_cls = _extract_class(source, "BaseMoeLoadBalancer", {"abort_iter"}, {}) + events = [] + + class Manager: + static = False + + def is_static_routing(self): + return self.static + + def set_iter_info(self, enable_statistic, enable_update_weights): + events.append("set_iter_info") + + def start_iter(self): + events.append("start_iter") + + def end_iter(self): + events.append("end_iter") + + def abort_iter(self, error): + events.append(("abort_iter", error)) + + manager = Manager() + with context_cls(manager): + pass + assert events == ["set_iter_info", "start_iter", "end_iter"] + + events.clear() + error = RuntimeError("forward failed") + with pytest.raises(RuntimeError, match="forward failed") as caught: + with context_cls(manager): + raise error + assert caught.value is error + assert events == ["set_iter_info", "start_iter", ("abort_iter", error)] + + # Static placement runs no iteration hook. + events.clear() + manager.static = True + with pytest.raises(RuntimeError, match="forward failed"): + with context_cls(manager): + raise error + assert events == [] + + # By default an aborted iteration simply ends, as standard EPLB always did. + class Standard(base_cls): + def end_iter(self): + events.append("end_iter") + + Standard().abort_iter(error) + assert events == ["end_iter"] + + +def test_per_layer_dynamic_abort_releases_pending_plans_and_keeps_the_error(): + namespace = {} + _extract_function( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/layer.py", "_append_note", namespace + ) + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/per_layer_dynamic_eplb/manager.py", + "PerLayerDynamicMoeLoadBalancer", + {"abort_iter", "shutdown"}, + namespace, + ) + released = [] + + class Layer: + def __init__(self, layer_idx, plan=None, cleanup_error=None): + self.layer_idx = layer_idx + self.replica_plan = plan + self.cleanup_error = cleanup_error + + def finish_replica_plan(self): + released.append(self.layer_idx) + if self.cleanup_error is not None: + raise self.cleanup_error + self.replica_plan = None + + manager = manager_cls() + manager.in_iter = True + manager.iter_id = 3 + manager.single_layer_load_balancers = [ + Layer(0), + Layer(1, plan="first", cleanup_error=RuntimeError("route wait failed")), + Layer(2, plan="second"), + ] + error = ValueError("forward failed") + + manager.abort_iter(error) + + # Every pending plan is released; a failed release only annotates the + # forward's error, which the iteration context then re-raises. + assert released == [1, 2] + assert [layer.replica_plan for layer in manager.single_layer_load_balancers] == [ + None, + "first", + None, + ] + assert len(error.__notes__) == 1 + assert "layer 1" in error.__notes__[0] + assert "route wait failed" in error.__notes__[0] + assert (manager.in_iter, manager.iter_id) == (False, 4) + + # The plan left pending makes this rank unsafe for collective teardown. + states = [] + + class Comm: + def allgather(self, state): + states.append(state) + return [state, {"local_safe": True, "pool_count": 0}] + + manager.is_shutdown = False + manager._ep_comm = Comm() + manager._rebalance_shared_slot_pool_order = [] + assert manager.shutdown(local_safe=True) is False + assert states == [{"local_safe": False, "pool_count": 0}] + + # The iteration ends even if the release loop itself fails. + class Broken: + @property + def replica_plan(self): + raise RuntimeError("layer state unavailable") + + manager.in_iter = True + manager.single_layer_load_balancers = [Broken()] + with pytest.raises(RuntimeError, match="layer state unavailable"): + manager.abort_iter(ValueError("forward failed")) + assert (manager.in_iter, manager.iter_id) == (False, 5) + + +class _ChunkRoutes: + dtype = "int32" + + def __init__(self, name): + self.name = name + + +@pytest.mark.parametrize("mode", ["per_layer_dynamic", "standard", "off"]) +def test_fused_comm_chunk_routes_once_after_recording_the_planner_input(mode): + source = _TORCH_ROOT / "moe/fused_moe/moe_scheduler.py" + trace = [] + + class ExpertStatistic: + @staticmethod + def set_layer(layer_idx): + assert layer_idx == 7 + + @staticmethod + def maybe_add_info(num_slots, routes): + trace.append(("statistic", routes.name)) + + def maybe_collect_or_replay_slots(num_slots, routes): + trace.append(("calibrator", routes.name)) + return _ChunkRoutes("replayed " + routes.name) + + namespace = { + "ExpertStatistic": ExpertStatistic, + "get_calibrator": lambda: SimpleNamespace( + maybe_collect_or_replay_slots=maybe_collect_or_replay_slots + ), + "get_active_route_capture": lambda: None, + "torch": SimpleNamespace(int32="int32", float32="float32"), + "Fp4QuantizedTensor": type("Fp4QuantizedTensor", (), {}), + "MoERunContext": lambda **kwargs: SimpleNamespace(**kwargs), + } + _extract_function(source, "_record_routes", namespace) + _extract_function(source, "_leading_rows", namespace) + scheduler_cls = _extract_class(source, "FusedCommMoEScheduler", {"_forward_chunk"}, namespace) + + plan = SimpleNamespace(reserved_sms=4) + layer = ( + None + if mode == "off" + else SimpleNamespace( + produces_replica_plan=mode != "standard", + require_replica_plan=lambda: plan, + ) + ) + + def route(routes, use_dp): + trace.append(("route", routes.name)) + return _ChunkRoutes("physical" if mode == "per_layer_dynamic" else "slots") + + def run_moe(ctx): + trace.append(("run_moe", ctx.token_selected_experts.name, ctx.replica_plan)) + return "out" + + moe = SimpleNamespace( + apply_router_weight_on_input=False, + mapping=SimpleNamespace(moe_ep_rank=0), + layer_idx=7, + num_slots=24, + use_dp=True, + layer_load_balancer=layer, + routing_method=SimpleNamespace( + apply=lambda logits, input_ids: ( + _ChunkRoutes("logical"), + SimpleNamespace(dtype="float32"), + ) + ), + _maybe_get_perfect_router_expert_ids=lambda routes: routes, + _load_balancer_start_wait_gpu_stage=lambda is_first_call: trace.append("start_wait_gpu"), + _load_balancer_done_wait_gpu_stage=lambda is_first_call: trace.append("done_wait_gpu"), + _load_balancer_update_statistic=lambda routes, *args, **kwargs: trace.append( + ("update_statistic", routes.name) + ), + _load_balancer_uses_replica_plan=lambda: mode == "per_layer_dynamic", + _load_balancer_route=route, + _load_balancer_start_set_cpu_stage=lambda is_last_call: trace.append("start_set_cpu"), + _load_balancer_done_set_cpu_stage=lambda is_last_call: trace.append("done_set_cpu"), + backend=SimpleNamespace( + quantize_input=lambda x, **kwargs: trace.append(("quantize", kwargs)) or ("x", "sf"), + set_adaptive_launch_tokens=lambda tokens: trace.append(("bucket", tokens)), + run_moe=run_moe, + ), + ) + scheduler = scheduler_cls() + scheduler.moe = moe + x = SimpleNamespace(shape=(4, 16)) + + out = scheduler._forward_chunk( + x, + SimpleNamespace(shape=(4, 24)), + output_dtype="bf16", + all_rank_num_tokens=[4, 6], + do_finalize=True, + ) + + assert out == "out" + tail = [("quantize", {}), ("bucket", 6)] + if mode == "per_layer_dynamic": + # The planner consumes the logical routes the calibrator records or + # replays, and quantization leaves the plan's SMs to its copy kernel. + # The layer has no CPU stage or host statistics, so only the post-MoE + # hook that releases its plan runs. + expected = [ + ("statistic", "logical"), + ("calibrator", "logical"), + ("route", "replayed logical"), + ("quantize", {"reserved_sms": 4}), + ("bucket", 6), + ("run_moe", "physical", plan), + "start_set_cpu", + ] + elif mode == "standard": + expected = [ + "start_wait_gpu", + "done_wait_gpu", + ("update_statistic", "logical"), + ("route", "logical"), + ("statistic", "slots"), + ("calibrator", "slots"), + *tail, + ("run_moe", "replayed slots", None), + "start_set_cpu", + "done_set_cpu", + ] + else: + expected = [ + "start_wait_gpu", + ("statistic", "logical"), + ("calibrator", "logical"), + *tail, + ("run_moe", "replayed logical", None), + "start_set_cpu", + "done_set_cpu", + ] + assert trace == expected + + +def test_base_worker_finalizes_rank_follower_pyexecutor(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", + "BaseWorker", + {"shutdown"}, + _PER_LAYER_DYNAMIC_SHUTDOWN, + ) + events = [] + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_shutdown=lambda: False, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is None + + events.clear() + cleanup_attempts = 0 + gate_attempts = 0 + + def can_shutdown(): + nonlocal gate_attempts + gate_attempts += 1 + return gate_attempts == 1 + + def terminal_cleanup(): + nonlocal cleanup_attempts + cleanup_attempts += 1 + events.append(("terminal_cleanup", True)) + if cleanup_attempts == 1: + raise RuntimeError("terminal cleanup failed") + + engine = SimpleNamespace( + shutdown_all_ranks=False, + can_shutdown=can_shutdown, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=terminal_cleanup, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.engine = engine + + with pytest.raises(RuntimeError, match="terminal cleanup failed"): + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "shutdown", + ("terminal_cleanup", True), + ("terminal_cleanup", True), + ] + assert worker.engine is None + assert gate_attempts == 1 + + +def test_mpi_worker_terminal_cleanup_failure_is_retryable(monkeypatch): + import types + + events = [] + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_process_group") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + **_PER_LAYER_DYNAMIC_SHUTDOWN, + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_test_megamoe_module", + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + cleanup_attempts = 0 + + def terminal_cleanup(*, local_safe): + nonlocal cleanup_attempts + cleanup_attempts += 1 + events.append(("terminal_cleanup", local_safe)) + if cleanup_attempts == 1: + raise RuntimeError("terminal cleanup failed") + + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=terminal_cleanup, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="terminal cleanup failed"): + worker.shutdown() + + assert events == ["shutdown", ("terminal_cleanup", True)] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + + assert events == [ + "shutdown", + ("terminal_cleanup", True), + ("terminal_cleanup", True), + "destroy_process_group", + "gc", + "empty_cache", + "background", + ] + assert worker.engine is None + + +def test_ray_worker_postproc_failure_still_runs_unsafe_terminal_cleanup(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + **_PER_LAYER_DYNAMIC_SHUTDOWN, + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + + def fail_postproc(): + events.append("postproc") + raise RuntimeError("postproc cleanup failed") + + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker._postproc_pool = object() + worker.shutdown_postproc_workers = fail_postproc + engine = SimpleNamespace( + shutdown_all_ranks=True, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker.engine = engine + worker.llm_args = SimpleNamespace(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="postproc cleanup failed"): + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "background", + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker._postproc_pool = None + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "background", + ("terminal_cleanup", True), + "background", + ] + assert worker.engine is None + + +def test_ray_worker_shutdown_preserves_first_error_and_attempts_all_cleanup(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + **_PER_LAYER_DYNAMIC_SHUTDOWN, + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + + def fail(name): + def raise_error(): + events.append(name) + raise RuntimeError(f"{name} failed") + + return raise_error + + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker._postproc_pool = object() + worker.shutdown_postproc_workers = fail("postproc") + worker.engine = SimpleNamespace( + shutdown_all_ranks=True, + shutdown=lambda: events.append("shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker.llm_args = SimpleNamespace(backend="pytorch") + worker.checkpoint_loader = SimpleNamespace(cleanup=fail("checkpoint")) + worker._handle_background_error = fail("background") + + with pytest.raises(RuntimeError, match="postproc failed"): + worker.shutdown() + + assert events == [ + "postproc", + "shutdown", + ("terminal_cleanup", False), + "checkpoint", + "background", + ] + assert worker.engine is not None + assert not worker.doing_shutdown + assert worker.checkpoint_loader is not None + + +def test_ray_worker_partial_init_without_llm_args_still_checks_background(): + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/gpu_worker.py", + "RayGPUWorker", + {"shutdown"}, + { + **_PER_LAYER_DYNAMIC_SHUTDOWN, + "logger": SimpleNamespace( + debug=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + events = [] + worker = worker_cls() + worker.doing_shutdown = False + worker.rank = 0 + worker.engine = None + worker._handle_background_error = lambda: events.append("background") + + worker.shutdown() + + assert events == ["background"] + + +def test_standard_load_balancer_unsafe_preflight_skips_destructive_cleanup(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py", + "MoeLoadBalancer", + {"shutdown"}, + {}, + ) + calls = [] + + class Comm: + def allgather(self, local_safe): + calls.append(("allgather", local_safe)) + return [local_safe, True] + + def barrier(self): + calls.append("barrier") + + manager = manager_cls() + manager.is_shutdown = False + manager.in_iter = True + manager.shared_mpi_comm = Comm() + manager.single_layer_load_balancers = [] + manager.load_balancer_impl = SimpleNamespace(shutdown=lambda: calls.append("shutdown")) + + assert manager.shutdown(local_safe=True) is False + assert calls == [("allgather", False)] + + +def test_standard_load_balancer_finalizer_releases_it_as_on_main(): + manager_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/moe_load_balancer.py", + "MoeLoadBalancer", + {"__del__"}, + {}, + ) + calls = [] + manager = manager_cls() + manager.shutdown = lambda: calls.append("shutdown") + + manager.is_shutdown = True + manager.__del__() + assert calls == [] + + manager.is_shutdown = False + manager.__del__() + assert calls == ["shutdown"] + manager.is_shutdown = True + + +class _FakePlane: + """Stand-in for an arena alias: records fills, passes the layout checks.""" + + is_cuda = True + dtype = "float32" + + def __init__(self, name, address, fills): + self.name = name + self.shape = (6,) + self._address = address + self._fills = fills + + def stride(self): + return (1,) + + def data_ptr(self): + return self._address + + def fill_(self, value): + self._fills.append((self.name, value)) + + +def test_live_plane_binding_restores_per_expert_scalar_defaults(): + names = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + fills = [] + aliases = [_FakePlane(name, 0x1000 + index, fills) for index, name in enumerate(names)] + arena = SimpleNamespace( + bundle="bundle", + plane_names=names, + local_plane_views=aliases, + tekit_alias_views=aliases, + home_experts=4, + ) + method_cls = _extract_class( + _TORCH_ROOT / "moe/fused_moe/quantization.py", + "NVFP4MegaMoECuteDslMethod", + {"_bind_live_weight_planes"}, + {"torch": SimpleNamespace(nn=SimpleNamespace(Parameter=None))}, + ) + method = method_cls() + method.live_weight_plane_spec = lambda module: SimpleNamespace( + planes=[SimpleNamespace(name=name) for name in names] + ) + arena.bundle = method.live_weight_plane_spec(None) + method.assert_live_weight_aliases = lambda module, bound: None + module = SimpleNamespace( + **{ + name: SimpleNamespace( + is_meta=False, shape=(6,), stride=lambda: (1,), dtype="float32", data=None + ) + for name in names + } + ) + + method._bind_live_weight_planes(module, arena) + + # The arena starts zeroed; an expert the checkpoint leaves unset must keep + # the create-time scalars (1.0) instead of silently computing zeros. + assert sorted(fills) == sorted( + [("fc31_alpha", 1.0), ("fc2_alpha", 1.0), ("fc1_norm_const", 1.0)] + ) + assert all(getattr(module, name).data is alias for name, alias in zip(names, aliases)) diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index ee1dbf8df39f..aa99dda21ce5 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -2543,6 +2543,25 @@ def _run(lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, to ) whole = _run(0, num_slots) + bundle = method._live_weight_bundle_layout(hidden, intermediate, expand_intermediate) + planes = {plane.name: plane for plane in bundle.planes} + for tensor, name in zip( + whole, + ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + ), + ): + shape, stride, dtype_name, element_size, pointer_alignment = planes[name].storage_metadata( + num_slots + ) + assert tuple(tensor.shape) == shape + assert tuple(tensor.stride()) == stride + assert tensor.dtype == getattr(torch, dtype_name) + assert tensor.element_size() == element_size + assert tensor.data_ptr() % pointer_alignment == 0 for chunk in (1, 3, 16): # 16 > num_slots: the single-call degenerate case pieces = [_run(lo, min(lo + chunk, num_slots)) for lo in range(0, num_slots, chunk)] diff --git a/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py new file mode 100644 index 000000000000..14a0535444e7 --- /dev/null +++ b/tests/unittest/_torch/moe/test_megamoe_shared_fc12_autotune.py @@ -0,0 +1,516 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contract tests; GPU numerical/concurrent admission runs separately.""" + +import __future__ + +import ast +import inspect +import os +import unittest +import weakref +from contextlib import nullcontext +from copy import deepcopy +from itertools import product +from pathlib import Path +from threading import Condition, Event, RLock, Thread +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +_ROOT = Path(__file__).resolve().parents[4] +_SOURCE = _ROOT / "tensorrt_llm/_torch/moe/kernels/megamoe_shared_fc12.py" +_AUTOTUNER = _ROOT / "tensorrt_llm/_torch/autotuner.py" + +pytestmark = pytest.mark.cpu_only + + +class Tensor: + def __init__(self, shape, dtype="bf16"): + self.shape = tuple(shape) + self.dtype = dtype + self.device = SimpleNamespace(index=0) + self.is_cuda = True + self.ndim = len(shape) + + def is_contiguous(self): + return True + + def stride(self): + strides = [] + n = 1 + for extent in reversed(self.shape): + strides.append(n) + n *= extent + return tuple(reversed(strides)) + + +def inputs(tokens): + return [ + Tensor((tokens, 7168)), + Tensor((1, 6144, 7168), "e4m3"), + Tensor((1, 1376256), "uint8"), + Tensor((1, 7168, 3072), "e4m3"), + Tensor((1, 688128), "uint8"), + ] + + +class TuningConfig(SimpleNamespace): + def __init__(self, **kwargs): + super().__init__(constraint_specs=(), distributed_tuning_strategy="independent", **kwargs) + self.distributed_tuning_strategy = SimpleNamespace(value="independent") + + +class TunableRunner: + def __call__(self, inputs, **kwargs): + return self.forward(inputs, **kwargs) + + +class ProfilingCache: + def __init__(self): + self.cache = {} + + def get_cache_key(self, op, runner, shapes, config, apply_map_to_tuning_buckets=True): + shapes = list(shapes) + if apply_map_to_tuning_buckets: + shapes[0] = (runner.token_bucket(shapes[0][0]), shapes[0][1]) + return op, type(runner).__name__, str(runner.unique_id()), tuple(shapes) + + def search_cache(self, op, runners, shapes, config, apply_map_to_tuning_buckets=True): + key = self.get_cache_key(op, runners[0], shapes, config, apply_map_to_tuning_buckets) + return (True, *self.cache[key]) if key in self.cache else (False, 0, -1, float("inf")) + + +class CpuContractTests(unittest.TestCase): + def setUp(self): + self.tuner = SimpleNamespace( + is_tuning_mode=False, + skip_dynamic_tuning_buckets=False, + _active_capture=None, + profiling_cache=ProfilingCache(), + ) + properties = SimpleNamespace(multi_processor_count=80, major=10, minor=7) + self.device_syncs = [] + self.stream_syncs = [] + self.empty_cache_calls = 0 + + def external_stream(stream_handle, *, device): + return SimpleNamespace( + synchronize=lambda: self.stream_syncs.append((device, stream_handle)) + ) + + torch = SimpleNamespace( + Tensor=Tensor, + bfloat16="bf16", + float8_e4m3fn="e4m3", + uint8="uint8", + cuda=SimpleNamespace( + get_device_properties=lambda device: properties, + device=lambda device: nullcontext(), + ExternalStream=external_stream, + synchronize=lambda device: self.device_syncs.append(device), + empty_cache=lambda: setattr(self, "empty_cache_calls", self.empty_cache_calls + 1), + ), + empty_like=lambda x: Tensor(x.shape, x.dtype), + _C=SimpleNamespace(_cuda_getCurrentRawStream=lambda device_index: 17), + ) + self.ns = dict( + torch=torch, + AutoTuner=SimpleNamespace(get=lambda: self.tuner), + DynamicTensorSpec=SimpleNamespace, + TunableRunner=TunableRunner, + TuningConfig=TuningConfig, + Condition=Condition, + RLock=RLock, + deepcopy=deepcopy, + product=product, + weakref=weakref, + ) + module = ast.parse(_SOURCE.read_text()) + module.body = [ + node for node in module.body if not isinstance(node, (ast.Import, ast.ImportFrom)) + ] + exec( + compile(module, str(_SOURCE), "exec", flags=__future__.annotations.compiler_flag), + self.ns, + ) + self.creations = [] + self.launches = [] + + def factory(*args): + self.creations.append(args) + launches = self.launches + + class Compiled: + def __call__(self, *tensors): + launches.append((args[-1], tensors[0].shape[0])) + return tensors[0] + + return Compiled() + + self.ns["_get_runner"] = factory + + def runner(self, sms=72, capacity=8192, clamp=10.0, stream=17): + return self.ns["SharedFc12TunableRunner"]( + 0, 7168, 3072, capacity, sms, clamp, stream, ("fixed_weight_layout",) + ) + + def test_catalog_includes_both_baselines_and_resource_decisions(self): + catalog = self.ns["shared_fc12_candidate_tactics"]() + self.assertEqual(len(catalog), 48) + self.assertEqual(len(set(catalog)), 48) + self.assertIn((1, 256, 128, "atomic_counter", True), catalog) + self.assertIn((2, 128, 128, "grid_stride", True), catalog) + rejected = catalog[0] + calls = [] + + def resource_check(*args): + calls.append(args) + if args[-1] == rejected: + raise ValueError("insufficient test resource") + + self.ns["_make_kernel"] = resource_check + for sms in (72, 80): + runner = self.runner(sms) + valid = runner.get_valid_tactics([], None) + self.assertEqual(set(valid), set(catalog) - {rejected}) + self.assertEqual(runner.get_valid_tactics([], None), valid) + self.assertEqual(len(calls), 96) + with self.assertRaises(ValueError): + self.ns["shared_fc12_tactic_descriptor"]((1, 256, 512, "atomic_counter", True)) + + def test_bucket_and_persistent_key_separation(self): + r = self.runner() + self.assertEqual( + [r.token_bucket(x) for x in (1, 2, 128, 129, 8191, 8192)], + [1, 128, 128, 256, 8192, 8192], + ) + self.assertEqual(self.runner(capacity=300).buckets, (1, 128, 256, 300)) + self.assertNotEqual(r.unique_id(), self.runner(80).unique_id()) + self.assertNotEqual(r.unique_id(), self.runner(clamp=None).unique_id()) + self.assertNotEqual(r.unique_id(), self.runner(capacity=4096).unique_id()) + self.assertEqual(r.unique_id(), self.runner(stream=18).unique_id()) + for invalid in (0, 8193): + with self.assertRaises(ValueError): + r.token_bucket(invalid) + + def test_preparation_is_not_fallback_and_dynamic_t_reuses_compilation(self): + r = self.runner() + self.tuner.is_tuning_mode = True + self.assertIsNone(r.forward(inputs(8192), do_preparation=True)) + self.assertEqual(self.creations, []) + with self.assertRaises(RuntimeError): + r.forward(inputs(8192), tactic=-1) + tactic = (2, 128, 128, "atomic_counter", True) + for tokens in (8192, 129, 4096, 256, 8192): + x = inputs(tokens) + self.assertIs(r.forward(x, tactic=tactic), x[0]) + self.assertEqual(len(self.creations), 1) + self.assertEqual([tokens for _, tokens in self.launches], [8192, 129, 4096, 256, 8192]) + self.tuner.is_tuning_mode = False + r.forward(inputs(256), tactic=-1) + self.assertEqual(self.creations[-1][-1], (1, 256, 128, "atomic_counter", True)) + + def test_complete_profile_cache_prunes_only_loser_workspaces(self): + r = self.runner(capacity=300) + winners = ( + (1, 256, 128, "atomic_counter", True), + (2, 128, 128, "grid_stride", True), + ) + loser = (1, 64, 256, "grid_stride", False) + unrelated = (*r._runner_args[:-1], 18, loser) + for tactic in (*winners, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + self.ns["_RUNNER_CACHE"][unrelated] = object() + r.kernel_cache[unrelated] = object() + + profiles = [] + for index, bucket in enumerate(r.buckets): + shapes = tuple(t.shape for t in inputs(bucket)) + tactic = winners[index % len(winners)] + key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", + r, + shapes, + r.tuning_config, + apply_map_to_tuning_buckets=False, + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.1) + profiles.append(SimpleNamespace(get_opt_shapes=lambda shapes=shapes: shapes)) + self.tuner._optimization_profiles = lambda config, tensors: profiles + + retained = r._cached_winner_tactics(self.tuner, inputs(300)) + self.assertEqual(retained, set(winners)) + r._retain_cached_winners(retained) + self.assertEqual(set(r._compiled_runners), set(winners)) + for tactic in winners: + key = (*r._runner_args, tactic) + self.assertIn(key, self.ns["_RUNNER_CACHE"]) + self.assertIn(key, r.kernel_cache) + loser_key = (*r._runner_args, loser) + self.assertNotIn(loser_key, self.ns["_RUNNER_CACHE"]) + self.assertNotIn(loser_key, r.kernel_cache) + self.assertIn(unrelated, self.ns["_RUNNER_CACHE"]) + self.assertIn(unrelated, r.kernel_cache) + self.assertEqual(self.stream_syncs, [(0, 17)]) + self.assertEqual(self.empty_cache_calls, 1) + r._retain_cached_winners(retained) + self.assertEqual(self.empty_cache_calls, 1) + + def test_incomplete_profile_cache_does_not_prune(self): + r = self.runner(capacity=300) + winner = (1, 256, 128, "atomic_counter", True) + loser = (1, 64, 256, "grid_stride", False) + for tactic in (winner, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + profiles = [] + for bucket in r.buckets: + shapes = tuple(t.shape for t in inputs(bucket)) + profiles.append(SimpleNamespace(get_opt_shapes=lambda shapes=shapes: shapes)) + first_key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", + r, + profiles[0].get_opt_shapes(), + r.tuning_config, + apply_map_to_tuning_buckets=False, + ) + self.tuner.profiling_cache.cache[first_key] = (0, winner, 0.1) + self.tuner._optimization_profiles = lambda config, tensors: profiles + + before_compiled = dict(r._compiled_runners) + before_runners = dict(self.ns["_RUNNER_CACHE"]) + before_kernels = dict(r.kernel_cache) + self.assertIsNone(r._cached_winner_tactics(self.tuner, inputs(300))) + self.assertEqual(r._compiled_runners, before_compiled) + self.assertEqual(self.ns["_RUNNER_CACHE"], before_runners) + self.assertEqual(r.kernel_cache, before_kernels) + self.assertEqual(self.stream_syncs, []) + self.assertEqual(self.empty_cache_calls, 0) + + def test_skipped_dynamic_buckets_do_not_prune(self): + r = self.runner(capacity=300) + winner = (1, 256, 128, "atomic_counter", True) + loser = (1, 64, 256, "grid_stride", False) + for tactic in (winner, loser): + cached = object() + r._compiled_runners[tactic] = cached + key = (*r._runner_args, tactic) + self.ns["_RUNNER_CACHE"][key] = cached + r.kernel_cache[key] = cached + self.tuner.skip_dynamic_tuning_buckets = True + + self.assertIsNone(r._cached_winner_tactics(self.tuner, inputs(300))) + self.assertEqual(set(r._compiled_runners), {winner, loser}) + self.assertEqual(self.stream_syncs, []) + self.assertEqual(self.empty_cache_calls, 0) + + def test_public_dispatch_obeys_chooser_and_reuses_compilation(self): + tactic = (2, 64, 256, "grid_stride", False) + choices = [] + + def choose(op, runners, config, tensors): + runner = runners[0] + key = self.tuner.profiling_cache.get_cache_key( + op, runner, tuple(t.shape for t in tensors), config + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.25) + choices.append(key) + return runner, tactic + + self.tuner.choose_one = choose + outputs = [ + self.ns["run_shared_fc12"](*inputs(tokens), swiglu_limit=10.0, sm_count=72) + for tokens in (129, 200) + ] + + self.assertEqual(len(choices), 2) + self.assertEqual(choices[0], choices[1]) + self.assertEqual(len(self.creations), 1) + self.assertEqual([output.shape[0] for output in outputs], [129, 200]) + self.assertEqual([tokens for _, tokens in self.launches], [129, 200]) + + def test_weight_layout_is_validated_once_per_weight_set(self): + checks = [] + + class CountedWeight(Tensor): + @property + def ndim(self): + checks.append(id(self)) + return len(self.shape) + + @ndim.setter + def ndim(self, value): + pass + + self.tuner.choose_one = lambda op, runners, config, tensors: (runners[0], -1) + activation, fc1, fc1_sf, fc2, _ = inputs(129) + fc2 = CountedWeight(fc2.shape, fc2.dtype) + fc2_sf = Tensor((1, 688128), "uint8") + run = self.ns["run_shared_fc12"] + for tokens in (129, 200): + run(Tensor((tokens, 7168)), fc1, fc1_sf, fc2, fc2_sf, swiglu_limit=10.0, sm_count=72) + # The second call of the same weight set skips the weight checks. + self.assertEqual(len(checks), 1) + layouts = self.ns["_WEIGHT_LAYOUTS"] + self.assertEqual(len(layouts), 1) + + # New weight tensors are validated again; entries of freed weights are pruned. + fc2_new = CountedWeight(fc2.shape, fc2.dtype) + del fc2 + run(activation, fc1, fc1_sf, fc2_new, fc2_sf, swiglu_limit=10.0, sm_count=72) + self.assertEqual(len(checks), 2) + self.assertEqual(len(layouts), 1) + + def test_tuning_key_text_is_unchanged(self): + # Persisted AutoTuner caches key the runner by str(unique_id()); the prebuilt + # string must stay the text of the identity tuple. + r = self.runner() + self.assertEqual( + r.unique_id(), + str( + ( + self.ns["_CATALOG_VERSION"], + 0, + 10, + 7, + 80, + 72, + 7168, + 3072, + 8192, + 10.0, + "K128_input_K32_intermediate_R128c4_gateup16", + ("fixed_weight_layout",), + ) + ), + ) + + def test_terminal_release_drops_all_device_runner_references(self): + first = self.runner(stream=17) + second = self.runner(stream=18) + first._compiled_runners[(1, 256, 128, "atomic_counter", True)] = object() + second._compiled_runners[(2, 128, 128, "grid_stride", True)] = object() + first_key = (0, 7168, 3072, 8192, 72, 10.0, 17, ("layout",)) + second_key = (0, 7168, 3072, 8192, 72, 10.0, 18, ("layout",)) + self.ns["_tunable_runners"].update({first_key: first, second_key: second}) + self.ns["_RUNNER_CACHE"].update({first_key: object(), second_key: object()}) + first.kernel_cache.update({first_key: object(), second_key: object()}) + + self.ns["release_shared_fc12_cache"](0) + + self.assertEqual(self.ns["_tunable_runners"], {}) + self.assertEqual(self.ns["_RUNNER_CACHE"], {}) + self.assertEqual(first.kernel_cache, {}) + self.assertEqual(first._compiled_runners, {}) + self.assertEqual(second._compiled_runners, {}) + self.assertEqual(self.device_syncs, [0]) + self.assertEqual(self.stream_syncs, []) + + def test_terminal_release_is_scoped_to_the_executor_stream(self): + first = self.runner(stream=17) + second = self.runner(stream=18) + first._compiled_runners[(1, 256, 128, "atomic_counter", True)] = object() + second._compiled_runners[(2, 128, 128, "grid_stride", True)] = object() + first_key = (0, 7168, 3072, 8192, 72, 10.0, 17, ("layout",)) + second_key = (0, 7168, 3072, 8192, 72, 10.0, 18, ("layout",)) + self.ns["_tunable_runners"].update({first_key: first, second_key: second}) + self.ns["_RUNNER_CACHE"].update({first_key: object(), second_key: object()}) + first.kernel_cache.update({first_key: object(), second_key: object()}) + + self.ns["release_shared_fc12_cache"](0, 17) + + self.assertNotIn(first_key, self.ns["_tunable_runners"]) + self.assertIn(second_key, self.ns["_tunable_runners"]) + self.assertNotIn(first_key, self.ns["_RUNNER_CACHE"]) + self.assertIn(second_key, self.ns["_RUNNER_CACHE"]) + self.assertNotIn(first_key, first.kernel_cache) + self.assertIn(second_key, first.kernel_cache) + self.assertEqual(first._compiled_runners, {}) + self.assertNotEqual(second._compiled_runners, {}) + self.assertEqual(self.device_syncs, []) + self.assertEqual(self.stream_syncs, [(0, 17)]) + + def test_overlapping_terminal_releases_are_serialized(self): + first = self.ns["_begin_shared_fc12_release"](0, 17) + acquired = Event() + + def acquire_same_scope(): + second = self.ns["_begin_shared_fc12_release"](0, 17) + acquired.set() + self.ns["_end_shared_fc12_release"](0, second) + + thread = Thread(target=acquire_same_scope) + thread.start() + blocked = not acquired.wait(0.05) + self.ns["_end_shared_fc12_release"](0, first) + + self.assertTrue(blocked) + self.assertTrue(acquired.wait(1.0)) + thread.join(timeout=1.0) + self.assertFalse(thread.is_alive()) + + def test_original_autotuner_primes_each_cached_tactic_before_serving(self): + source = ast.parse(_AUTOTUNER.read_text()) + autotuner_class = next( + n for n in source.body if isinstance(n, ast.ClassDef) and n.name == "AutoTuner" + ) + method = deepcopy( + next( + n + for n in autotuner_class.body + if isinstance(n, ast.FunctionDef) and n.name == "_prime_cached_tactics" + ) + ) + namespace = dict( + os=os, + inspect=inspect, + nvtx_range=lambda name: nullcontext(), + DistributedTuningStrategy=SimpleNamespace(MERGE=object()), + ) + exec( + compile( + ast.Module(body=[method], type_ignores=[]), + str(_AUTOTUNER), + "exec", + flags=__future__.annotations.compiler_flag, + ), + namespace, + ) + r = self.runner() + tactics = [(2, 128, 128, "grid_stride", True), (1, 256, 128, "atomic_counter", True)] + tensors = [inputs(256), inputs(8192)] + profiles = [ + SimpleNamespace(get_opt_shapes=lambda x=x: tuple(t.shape for t in x)) for x in tensors + ] + for x, tactic in zip(tensors, tactics): + key = self.tuner.profiling_cache.get_cache_key( + "trtllm::megamoe_shared_fc12", r, tuple(t.shape for t in x), r.tuning_config + ) + self.tuner.profiling_cache.cache[key] = (0, tactic, 0.1) + self.tuner._optimization_profiles = lambda config, x: profiles + self.tuner._prepare_input_tensors = lambda profile, x: inputs( + profile.get_opt_shapes()[0][0] + ) + self.tuner._primed_cached_tactics = set() + self.tuner.is_tuning_mode = True + r.tuning_config.inputs_pre_hook = None + with patch.dict(os.environ, {"TLLM_AUTOTUNER_PRIME_CACHED_TACTICS": "1"}): + for _ in range(2): + namespace["_prime_cached_tactics"]( + self.tuner, "trtllm::megamoe_shared_fc12", [r], r.tuning_config, inputs(8192) + ) + self.assertEqual(len(self.creations), 2) + self.assertEqual([t for t, _ in self.launches], tactics) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/moe/test_megamoe_streaming_load.py b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py index 2c6bb606de21..776af12b1ab1 100644 --- a/tests/unittest/_torch/moe/test_megamoe_streaming_load.py +++ b/tests/unittest/_torch/moe/test_megamoe_streaming_load.py @@ -51,7 +51,7 @@ def _load_classes() -> tuple[type[MoEWeightLoadingMode], type[NVFP4MegaMoECuteDs class _StreamingMoEModule(nn.Module): """Minimal single-rank module for the quant-method load path.""" - def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: + def __init__(self, weight_loading_mode: MoEWeightLoadingMode, helper_slots: int = 0) -> None: from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig super().__init__() @@ -61,6 +61,7 @@ def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: self.intermediate_size_per_partition = INTERMEDIATE_SIZE self.expand_intermediate_size_per_partition = 2 * INTERMEDIATE_SIZE self.expert_size_per_partition = NUM_EXPERTS + self.helper_slots = int(helper_slots) self.initial_local_expert_ids = list(range(NUM_EXPERTS)) self.tp_size = 1 self.tp_rank = 0 @@ -72,11 +73,38 @@ def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: # No EPLB in this test: need_load_shared_weights() must be False. self.layer_load_balancer = None + @property + def compute_slot_count(self) -> int: + return int(self.expert_size_per_partition) + self.helper_slots + def _add_raw_shared_weights_for_unmap(self, weight_tensors: list[torch.Tensor]) -> None: # Only forwards to the dynamic load balancer in production; no-op here. del weight_tensors +class _SentinelAlphaMoEModule(_StreamingMoEModule): + def __init__(self, weight_loading_mode: MoEWeightLoadingMode, helper_slots: int) -> None: + super().__init__(weight_loading_mode, helper_slots) + self.alpha_sentinels: dict[str, torch.Tensor] = {} + + def register_parameter(self, name, parameter): + if ( + name in ("fc31_alpha", "fc2_alpha") + and parameter is not None + and tuple(parameter.shape) == (NUM_EXPERTS,) + and name not in self.alpha_sentinels + ): + base = 11 if name == "fc31_alpha" else 21 + sentinel = torch.arange( + NUM_EXPERTS, + device=parameter.device, + dtype=parameter.dtype, + ).add_(base) + parameter.data.copy_(sentinel) + self.alpha_sentinels[name] = sentinel.clone() + super().register_parameter(name, parameter) + + def _w13_input_scale(expert_id: int) -> float: return 0.5 + 0.125 * expert_id @@ -178,6 +206,87 @@ def _expected_fc1_norm_const() -> torch.Tensor: ) +def test_helper_slots_widen_only_seven_kernel_planes() -> None: + mode_cls, method_cls = _load_classes() + helper_slots = 2 + module = _SentinelAlphaMoEModule(mode_cls.VANILLA, helper_slots=helper_slots) + method = method_cls() + with torch.device("cuda"): + method.create_weights(module) + + assert module.expert_size_per_partition == NUM_EXPERTS + for name in _STREAMED_PARAMS: + assert module.rebuild_tensor_metadata[name]["meta"].shape[0] == NUM_EXPERTS + + compute_slots = NUM_EXPERTS + helper_slots + kernel_planes = ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + assert all(getattr(module, name).shape[0] == compute_slots for name in kernel_planes) + + bundle = method.live_weight_plane_spec(module) + for plane in bundle.planes: + parameter = getattr(module, plane.name) + shape, stride, dtype_name, element_size, pointer_alignment = plane.storage_metadata( + compute_slots + ) + assert tuple(parameter.shape) == shape + assert tuple(parameter.stride()) == stride + assert parameter.dtype == getattr(torch, dtype_name) + assert parameter.element_size() == element_size + assert parameter.data_ptr() % pointer_alignment == 0 + + kernel_shape, kernel_stride, kernel_dtype, kernel_element_size, _ = plane.kernel_metadata( + compute_slots + ) + kernel_view = parameter.view(getattr(torch, kernel_dtype)) + if plane.name in ("mega_fc1_weight", "mega_fc2_weight"): + kernel_view = kernel_view.transpose(1, 2) + assert tuple(kernel_view.shape) == kernel_shape + assert tuple(kernel_view.stride()) == kernel_stride + assert kernel_view.dtype == getattr(torch, kernel_dtype) + assert kernel_view.element_size() == kernel_element_size + + for name, sentinel in module.alpha_sentinels.items(): + actual = getattr(module, name) + assert torch.equal(actual[:NUM_EXPERTS], sentinel) + assert torch.equal( + actual[NUM_EXPERTS:], + torch.ones(helper_slots, dtype=actual.dtype, device=actual.device), + ) + assert set(module.alpha_sentinels) == {"fc31_alpha", "fc2_alpha"} + assert module.quant_scales.fc1_global is module.fc31_alpha + assert module.quant_scales.fc2_global is module.fc2_alpha + + +@pytest.mark.parametrize("helper_slots", [0, 2]) +def test_create_weights_supports_meta_init(helper_slots: int) -> None: + """The model loader builds MoE layers under MetaInitMode. + + Any op that touches a meta tensor raises there and makes the loader rebuild + the whole model on the host, so the large derived planes must stay meta. + """ + from tensorrt_llm._torch.models.modeling_utils import MetaInitMode + + mode_cls, method_cls = _load_classes() + module = _StreamingMoEModule(mode_cls.VANILLA, helper_slots=helper_slots) + with MetaInitMode(): + method_cls().create_weights(module) + + for name in ("mega_fc1_weight", "mega_fc1_weight_sf", "mega_fc2_weight", "mega_fc2_weight_sf"): + assert getattr(module, name).is_meta + norm = module.fc1_norm_const + assert not norm.is_meta + assert norm.shape[0] == NUM_EXPERTS + helper_slots + assert torch.equal(norm, torch.ones_like(norm)) + + def test_initial_streaming_load_layer_atomic() -> None: method, module, weights = _fresh() diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index fc562dbcbbbd..6b18124cd41b 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -105,6 +105,7 @@ MoEProblem, MoERejection, MoERejectReason, + MoEReplicaPlan, MoEResolutionReport, MoERunContext, MoEStaticCapability, @@ -132,6 +133,10 @@ ExternalCommMoEScheduler, FusedCommMoEScheduler, ) +from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb.planner import ( + RebalanceSlotSchedulerGroup, + _LiveBankLeaseProvider, +) from tensorrt_llm._torch.moe.fused_moe.quantization import ( FusedMoEMethodBase, NVFP4FusedMoEMethod, @@ -366,6 +371,53 @@ def test_fused_comm_scheduler_indexes_full_ep_token_metadata(): assert had_meta +def _split_chunk(tokens, chunks): + return [tokens // chunks + 1] * (tokens % chunks) + [tokens // chunks] * ( + chunks - tokens % chunks + ) + + +@pytest.mark.cpu_only +def test_fused_comm_scheduler_passes_unpadded_single_chunk_inputs_through(): + moe = SimpleNamespace( + mapping=SimpleNamespace(moe_ep_rank=1, moe_ep_size=2), + moe_max_num_tokens=8, + split_chunk=_split_chunk, + ) + scheduler = FusedCommMoEScheduler(moe) + x = torch.randn(6, 8) + router_logits = torch.randn(6, 4) + input_ids = torch.arange(6) + + x_real, rl_real, ids_real, counts, ep_rank, had_meta = scheduler._strip_adp_padding( + x, router_logits, input_ids, [4, 6] + ) + # This rank carries no ADP padding and fits one chunk: no view is taken. + assert x_real is x and rl_real is router_logits and ids_real is input_ids + assert (counts, ep_rank, had_meta) == ([4, 6], 1, True) + num_chunks, x_chunks, rl_chunks, chunk_sizes = scheduler._compute_chunk_layout( + x_real, rl_real, counts, ep_rank + ) + assert num_chunks == 1 + assert x_chunks[0] is x and rl_chunks[0] is router_logits + assert chunk_sizes == [[4], [6]] + + # Padding is still stripped and several chunks are still split. + x_real, rl_real, _, counts, ep_rank, _ = scheduler._strip_adp_padding( + x, router_logits, None, [4, 5] + ) + assert x_real.shape[0] == 5 and torch.equal(x_real, x[:5]) + moe.moe_max_num_tokens = 2 + num_chunks, x_chunks, rl_chunks, chunk_sizes = scheduler._compute_chunk_layout( + x_real, rl_real, counts, ep_rank + ) + assert num_chunks == 3 + assert [chunk.shape[0] for chunk in x_chunks] == [2, 2, 1] + assert torch.equal(torch.cat(x_chunks), x[:5]) + assert torch.equal(torch.cat(rl_chunks), router_logits[:5]) + assert chunk_sizes == [[2, 1, 1], [2, 2, 1]] + + def test_fused_comm_scheduler_rejects_partial_ep_token_metadata(): mapping = SimpleNamespace(moe_ep_rank=1, moe_ep_size=4) scheduler = FusedCommMoEScheduler(SimpleNamespace(mapping=mapping)) @@ -2038,6 +2090,429 @@ def test_nvfp4_trtllm_gen_tp_accepts_logical_group_aligned_shard() -> None: assert backend.intermediate_size_per_partition % backend.quant_method.weight_alignment != 0 +@pytest.mark.cpu_only +def test_megamoe_quant_method_defines_canonical_live_plane_spec(): + method = NVFP4MegaMoECuteDslMethod() + module = SimpleNamespace( + hidden_size=7168, + intermediate_size_per_partition=3072, + expand_intermediate_size_per_partition=6144, + ) + + spec = method.live_weight_plane_spec(module) + + assert (spec.hidden, spec.intermediate, spec.expand_intermediate) == (7168, 3072, 6144) + assert tuple(plane.name for plane in spec.planes) == ( + "mega_fc1_weight", + "mega_fc1_weight_sf", + "mega_fc2_weight", + "mega_fc2_weight_sf", + "fc31_alpha", + "fc2_alpha", + "fc1_norm_const", + ) + + # Exercise the storage-to-kernel dtype/view bridge on a small legal + # geometry; production metadata is covered by the independent goldens. + bridge = method._live_weight_bundle_layout(160, 80, 160) + for plane in bridge.planes: + shape, stride, dtype_name, element_size, pointer_alignment = plane.storage_metadata(3) + tensor = torch.empty(shape, dtype=getattr(torch, dtype_name)) + assert tuple(tensor.stride()) == stride + assert tensor.element_size() == element_size + assert tensor.data_ptr() % pointer_alignment == 0 + + kernel_shape, kernel_stride, _, kernel_element_size, _ = plane.kernel_metadata(3) + kernel_view = tensor + if plane.name in ("mega_fc1_weight", "mega_fc2_weight"): + kernel_view = kernel_view.transpose(1, 2) + assert tuple(kernel_view.shape) == kernel_shape + assert tuple(kernel_view.stride()) == kernel_stride + assert kernel_view.element_size() == kernel_element_size + + +@pytest.mark.cpu_only +def test_megamoe_live_plane_spec_rejects_expand_before_allocation(): + method = NVFP4MegaMoECuteDslMethod() + module = SimpleNamespace( + hidden_size=160, + intermediate_size_per_partition=80, + expand_intermediate_size_per_partition=176, + ) + + with pytest.raises(ValueError, match="expand_intermediate == 2"): + method.live_weight_plane_spec(module) + + +@pytest.mark.cpu_only +def test_megamoe_rejects_expand_before_arena_allocation(): + allocate_arena = MagicMock() + module = SimpleNamespace( + hidden_size=160, + intermediate_size_per_partition=80, + expand_intermediate_size_per_partition=176, + layer_load_balancer=SimpleNamespace( + produces_replica_plan=True, allocate_arena=allocate_arena + ), + ) + + with pytest.raises(ValueError, match="expand_intermediate == 2"): + NVFP4MegaMoECuteDslMethod().create_weights(module) + + allocate_arena.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("produces_replica_plan", [True, False, None]) +def test_megamoe_quant_method_allocates_and_binds_the_arena_through_the_layer_handle( + monkeypatch, produces_replica_plan +): + method = NVFP4MegaMoECuteDslMethod() + bundle = method._live_weight_bundle_layout(160, 80, 160) + resident, helpers = 2, 1 if produces_replica_plan else 0 + arena = object() + events = [] + + def create_parent_weights(self, module, *args): + events.append("parent") + for name in NVFP4MegaMoECuteDslMethod._STREAMED_SOURCE_PARAMS: + module.register_parameter( + name, torch.nn.Parameter(torch.empty(1, dtype=torch.uint8), requires_grad=False) + ) + for name in ("fc31_alpha", "fc2_alpha"): + shape, _, dtype_name, _, _ = method._live_weight_plane(bundle, name).storage_metadata( + resident + ) + module.register_parameter( + name, + torch.nn.Parameter( + torch.ones(shape, dtype=getattr(torch, dtype_name)), requires_grad=False + ), + ) + + def allocate_arena(**kwargs): + events.append(("allocate", kwargs)) + return arena + + monkeypatch.setattr(NVFP4FusedMoEMethod, "create_weights", create_parent_weights) + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "setup_quant_scales", lambda self, module: None) + monkeypatch.setattr( + NVFP4MegaMoECuteDslMethod, + "_bind_live_weight_planes", + lambda self, module, bound: events.append(("bind", bound)), + ) + module = torch.nn.Module() + module.hidden_size = 160 + module.intermediate_size_per_partition = 80 + module.expand_intermediate_size_per_partition = 160 + module.bias = False + module.expert_size_per_partition = resident + module.compute_slot_count = resident + helpers + module.rebuild_tensor_metadata = {} + module.ep_process_group = "ep_group" + module.mapping = SimpleNamespace(local_rank=3) + module.layer_idx = 5 + module.layer_load_balancer = ( + None + if produces_replica_plan is None + else SimpleNamespace( + produces_replica_plan=produces_replica_plan, allocate_arena=allocate_arena + ) + ) + + method.create_weights(module) + + assert module.mega_fc1_weight.shape[0] == resident + helpers + if not produces_replica_plan: + assert events == ["parent"] + return + # Every rank allocates the arena before any Parameter exists and binds the + # seven planes once all of them are registered. + assert events == [ + ( + "allocate", + dict( + bundle=bundle, + ep_process_group="ep_group", + mapping=module.mapping, + device=3, + layer_idx=5, + ), + ), + "parent", + ("bind", arena), + ] + + +@pytest.mark.cpu_only +def test_megamoe_quant_method_builds_the_producer_after_checking_the_planes(monkeypatch): + events = [] + monkeypatch.setattr( + NVFP4FusedMoEMethod, "cache_derived_state", lambda self, module: events.append("parent") + ) + monkeypatch.setattr( + NVFP4MegaMoECuteDslMethod, + "assert_live_weight_aliases", + lambda self, module, arena: events.append(("aliases", arena)), + ) + arena = SimpleNamespace(assert_identity=lambda: events.append("identity")) + handle = SimpleNamespace( + produces_replica_plan=True, + live_weight_arena=arena, + bind_backend=lambda module: events.append(("bind_backend", module)), + ) + module = SimpleNamespace(layer_load_balancer=handle) + method = NVFP4MegaMoECuteDslMethod() + + method.cache_derived_state(module) + + assert events == ["parent", "identity", ("aliases", arena), ("bind_backend", module)] + + # Standard EPLB and EPLB-off layers only refresh the parent state. + events.clear() + method.cache_derived_state(SimpleNamespace(layer_load_balancer=None)) + method.cache_derived_state( + SimpleNamespace(layer_load_balancer=SimpleNamespace(produces_replica_plan=False)) + ) + assert events == ["parent", "parent"] + + handle.live_weight_arena = None + with pytest.raises(RuntimeError, match="no live-weight arena"): + method.cache_derived_state(module) + + +@pytest.mark.cpu_only +def test_megamoe_live_plane_identity_guard_detects_parameter_replacement(): + class Tensor: + def __init__(self, pointer): + self.shape = (2, 3) + self.dtype = "u8" + self.device = "cuda:0" + self.is_meta = False + self._pointer = pointer + + def stride(self): + return (3, 1) + + def data_ptr(self): + return self._pointer + + arena_view = Tensor(17) + alias = Tensor(17) + plane = SimpleNamespace(name="weight") + bundle = SimpleNamespace(planes=(plane,)) + arena = SimpleNamespace( + bundle=bundle, + local_plane_views=(arena_view,), + tekit_alias_views=(alias,), + ) + module = SimpleNamespace(weight=Tensor(17)) + method = NVFP4MegaMoECuteDslMethod() + method.live_weight_plane_spec = MagicMock(return_value=bundle) + + method.assert_live_weight_aliases(module, arena) + + module.weight = Tensor(23) + with pytest.raises(RuntimeError, match="changed storage"): + method.assert_live_weight_aliases(module, arena) + + +def _per_layer_dynamic_handle(helper_slots_per_rank): + return SimpleNamespace(produces_replica_plan=True, helper_slots_per_rank=helper_slots_per_rank) + + +@pytest.mark.cpu_only +def test_megamoe_compute_slots_come_from_the_layer_handle(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe.layer_load_balancer = None + assert moe.compute_slot_count == 8 + + # Standard EPLB produces no replica plan and adds no helper slot. + moe.layer_load_balancer = SimpleNamespace(produces_replica_plan=False) + assert moe.compute_slot_count == 8 + + moe.layer_load_balancer = _per_layer_dynamic_handle(3) + assert moe.compute_slot_count == 11 + assert moe.expert_size_per_partition == 8 + + +@pytest.mark.cpu_only +def test_megamoe_replica_plan_supplies_launch_facts_and_route_wait(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe.layer_load_balancer = _per_layer_dynamic_handle(3) + wait_for_routes = MagicMock() + ready_flags = torch.empty(8, dtype=torch.int64) + plan = MoEReplicaPlan( + resident_experts_per_rank=8, + compute_slots_per_rank=11, + ready_flags=ready_flags, + ready_generation=7, + reserved_sms=4, + _wait_for_routes=wait_for_routes, + ) + + compute_slots, helper_slots, actual_flags, generation, reserved_sms = ( + moe._replica_plan_launch_args(plan) + ) + assert (compute_slots, helper_slots, generation, reserved_sms) == (11, 3, 7, 4) + assert actual_flags is ready_flags + moe._wait_rebalance_routes(plan) + wait_for_routes.assert_called_once_with() + + moe._wait_rebalance_routes() + wait_for_routes.assert_called_once_with() + + +@pytest.mark.cpu_only +def test_megamoe_staging_widens_int32_routes_after_the_route_wait(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe._last_staged_T = {8: 6} + trace = [] + moe._wait_rebalance_routes = lambda replica_plan=None: trace.append(("wait", replica_plan)) + bufs = SimpleNamespace( + activation=torch.zeros(8, 16, dtype=torch.uint8), + activation_sf=torch.zeros(8, 4, dtype=torch.uint8), + topk_weights=torch.ones(8, 4, dtype=torch.float32), + topk_idx_local=torch.full((8, 4), 5, dtype=torch.int64), + ) + routes = torch.tensor([[0, 3, -1, 2**31 - 1], [7, 6, 5, 4], [1, 1, 1, 1]], dtype=torch.int32) + weights = torch.full((3, 4), 0.25, dtype=torch.float32) + + moe._stage_inputs( + bufs=bufs, + x=torch.ones(3, 16, dtype=torch.uint8), + x_sf=torch.ones(3, 4, dtype=torch.uint8), + topk_idx=routes, + topk_weights=weights, + num_tokens=3, + top_k=4, + stage_activation=True, + replica_plan="plan", + ) + + # One copy widens the int32 routes into the int64 kernel index buffer. + assert bufs.topk_idx_local.dtype == torch.int64 + assert torch.equal(bufs.topk_idx_local[:3], routes.to(torch.int64)) + # Rows a previous launch wrote, and only those, are reset to -1. + assert torch.equal(bufs.topk_idx_local[3:6], torch.full((3, 4), -1, dtype=torch.int64)) + assert torch.equal(bufs.topk_idx_local[6:], torch.full((2, 4), 5, dtype=torch.int64)) + assert torch.equal(bufs.topk_weights[:3], weights) + assert not bufs.topk_weights[3:].any() + assert trace == [("wait", "plan")] + + +@pytest.mark.cpu_only +def test_megamoe_run_moe_hands_routes_to_staging_unconverted(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.ep_size = 1 + moe._weights_created = True + moe.quant_method = object() + moe.hidden_size = 16 + moe.dtype = torch.bfloat16 + moe._topk_idx_dtype = torch.int64 + moe._build_weight_view = lambda: "weights" + moe._wait_rebalance_routes = MagicMock() + captured = {} + moe._run_moe = lambda **kwargs: captured.update(kwargs) or "out" + routes = torch.zeros(3, 4, dtype=torch.int32) + scales = torch.ones(3, 4, dtype=torch.float32) + + out = moe.run_moe( + MoERunContext( + token_selected_experts=routes, + token_final_scales=scales, + x=torch.zeros(3, 8, dtype=torch.uint8), + x_sf=torch.zeros(3, 4, dtype=torch.uint8), + output_dtype=torch.bfloat16, + ) + ) + + assert out == "out" + # No conversion kernel and no early route wait: staging reads the routes. + assert captured["topk_idx"] is routes + assert captured["topk_weights"] is scales + moe._wait_rebalance_routes.assert_not_called() + + half_scales = scales.to(torch.float16) + moe.run_moe( + MoERunContext( + token_selected_experts=routes, + token_final_scales=half_scales, + x=torch.zeros(3, 8, dtype=torch.uint8), + x_sf=torch.zeros(3, 4, dtype=torch.uint8), + output_dtype=torch.bfloat16, + ) + ) + assert captured["topk_weights"].dtype == torch.float32 + assert torch.equal(captured["topk_weights"], scales) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("has_sm_budget", [True, False]) +def test_megamoe_quantize_reserves_sms_only_for_a_replica_plan(monkeypatch, request, has_sm_budget): + from tensorrt_llm._torch.moe.fused_moe.mega_moe import mega_moe_cute_dsl + + calls = [] + + def outputs(x): + rows = x.shape[0] + return ( + torch.zeros(rows, x.shape[1] // 2, dtype=torch.uint8), + torch.zeros(rows * x.shape[1] // 16, dtype=torch.uint8), + ) + + def fp4_quantize(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled): + calls.append(("default", sf_vec_size, sf_use_ue8m0, is_sf_swizzled)) + return outputs(x) + + def sm_budget(x, scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled, reserved_sms): + calls.append(("sm_budget", sf_vec_size, sf_use_ue8m0, is_sf_swizzled, reserved_sms)) + return outputs(x) + + if has_sm_budget: + fp4_quantize.sm_budget = sm_budget + monkeypatch.setattr(torch.ops.trtllm, "fp4_quantize", fp4_quantize, raising=False) + mega_moe_cute_dsl._fp4_quantize_has_sm_budget.cache_clear() + request.addfinalizer(mega_moe_cute_dsl._fp4_quantize_has_sm_budget.cache_clear) + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.fc31_input_scale = torch.ones(1) + x = torch.zeros(3, 64, dtype=torch.bfloat16) + + moe.quantize_input(x) + moe.quantize_input(x, reserved_sms=4) + + default = ("default", 16, False, False) + budget = ("sm_budget", 16, False, False, 4) if has_sm_budget else default + assert calls == [default, budget] + + +@pytest.mark.cpu_only +def test_megamoe_replica_plan_rejects_unallocated_geometry(): + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe.expert_size_per_partition = 8 + moe.layer_load_balancer = _per_layer_dynamic_handle(3) + plan = MoEReplicaPlan( + resident_experts_per_rank=8, + compute_slots_per_rank=12, + ready_flags=torch.empty(8, dtype=torch.int64), + ready_generation=7, + reserved_sms=8, + _wait_for_routes=lambda: None, + ) + + with pytest.raises(RuntimeError, match="geometry does not match"): + moe._replica_plan_launch_args(plan) + + +@pytest.mark.cpu_only def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) @@ -2096,6 +2571,147 @@ def test_megamoe_cutedsl_mpi_bootstrap_keeps_executor_cuda_device(monkeypatch): ) +@pytest.mark.cpu_only +def test_megamoe_cutedsl_cache_derived_state_reuses_rebalance_binding( + monkeypatch, +): + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import ( + PerLayerDynamicMoeLoadBalancer, + planner, + ) + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module + + config = SimpleNamespace( + num_slots=24, + auxiliary_sms=8, + validate_expert_capacity=lambda num_experts: None, + ) + manager = PerLayerDynamicMoeLoadBalancer( + ep_rank=0, + ep_size=4, + config=config, + ) + layer = manager.add_layer( + expert_count=16, + top_k=2, + slot_count_per_rank=6, + ) + bundle = NVFP4MegaMoECuteDslMethod._live_weight_bundle_layout(160, 80, 160) + arena = MagicMock(home_experts=4, helper_slots=2, bundle=bundle) + arena.assert_identity = MagicMock() + allocate = MagicMock(return_value=arena) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", allocate) + monkeypatch.setattr(manager, "get_ep_comm", lambda process_group: "ep_comm") + monkeypatch.setattr(manager, "create_arena_provider", lambda: "provider") + moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) + torch.nn.Module.__init__(moe) + moe._ep_pg = object() + moe._maxt_buckets = [64, 128] + moe.mapping = SimpleNamespace(local_rank=0) + moe.layer_idx = 5 + moe.layer_load_balancer = layer + moe.initial_global_assignments = list(range(16)) + # The quant method's create_weights allocates through the handle; a second + # call with the same plane spec returns the same arena. + for _ in range(2): + assert ( + layer.allocate_arena( + bundle=bundle, + ep_process_group=moe.ep_process_group, + mapping=moe.mapping, + device=0, + layer_idx=5, + ) + is arena + ) + allocate.assert_called_once() + moe.quant_method = NVFP4MegaMoECuteDslMethod() + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "setup_quant_scales", lambda self, module: None) + assert_aliases = MagicMock() + monkeypatch.setattr(NVFP4MegaMoECuteDslMethod, "assert_live_weight_aliases", assert_aliases) + builder = MagicMock(return_value=SimpleNamespace(home_experts=4, helper_slots=2)) + monkeypatch.setattr( + planner, + "build_rebalance_slot_scheduler_group", + builder, + ) + + moe.cache_derived_state() + moe.cache_derived_state() + + assert arena.assert_identity.call_count == 2 + assert assert_aliases.call_count == 2 + assert layer.live_weight_arena is arena + builder.assert_called_once_with( + arena=arena, + manager=manager, + mapping=moe.mapping, + ep_process_group=moe._ep_pg, + home_experts=4, + helper_slots=2, + topk=2, + max_tokens_per_rank=128, + layer_idx=5, + ) + + +@pytest.mark.cpu_only +def test_rebalance_discard_plan_orders_release_and_allows_next_reuse(): + trace = [] + outputs = SimpleNamespace(physical_slot_ids=[3, 4]) + part = (outputs, 2) + copy_stream = object() + + group = RebalanceSlotSchedulerGroup.__new__(RebalanceSlotSchedulerGroup) + group._check_owner = MagicMock() + group._plan_part = part + group._generation = 1 + group._finished_generation = 0 + group._scheduled_generation = 1 + group._route_wait_generation = 0 + group._pending_release = None + group.plan_calls = 0 + group._execution_stream_handle = 23 + group.copy_stream = copy_stream + group.scheduler_stream = copy_stream + + def wait_for_routes(): + trace.append("wait_routes") + group._route_wait_generation = group._generation + + def release_generation_after(generation, event): + # Without an event, the release follows the MAIN work queued so far. + assert (generation, event) == (1, None) + trace.append("release_generation") + + broadcaster = SimpleNamespace( + release_generation_after=MagicMock(side_effect=release_generation_after), + mark_collective_reuse_safe=MagicMock( + side_effect=lambda authority, generation: trace.append("prove_reuse") + ), + ) + group.wait_for_routes = wait_for_routes + group.broadcaster = broadcaster + + group.discard_plan(part) + + assert group._plan_part is None + assert group.plan_calls == 1 + assert group._finished_generation == 1 + assert group._pending_release == (1, None) + assert trace == ["wait_routes", "release_generation"] + + # The next schedule advances this counter before proving that generation 1 + # can be reused. Exercise the real lease check to show cleanup left a valid + # transition rather than a permanently poisoned group. + group._scheduled_generation = 2 + lease = _LiveBankLeaseProvider(group) + lease.prove(1) + + assert lease.proved_generation == 1 + assert trace[-1] == "prove_reuse" + + def test_megamoe_load_weights_invalidates_cached_deepgemm_views(): method = W4A8MXFP4MXFP8MegaMoEDeepGemmMethod() hidden_size = 128 @@ -3326,29 +3942,32 @@ def test_enumerate_megamoe_candidate_tactics_curated_space( decode = megamoe_op.enumerate_megamoe_candidate_tactics(1024, sm_version=sm_version) prefill = megamoe_op.enumerate_megamoe_candidate_tactics(16384, sm_version=sm_version) assert (len(decode), len(prefill)) == (decode_count, prefill_count) - assert all(len(t) == 10 for t in decode + prefill) - assert {t[-1] for t in decode} == {(1, 1)} - assert {t[-1] for t in prefill} == {(2, 4)} + assert all(len(t) == 11 for t in decode + prefill) + assert {t[-2] for t in decode} == {(1, 1)} + assert {t[-2] for t in prefill} == {(2, 4)} + assert {t[-1] for t in decode + prefill} == {"expert"} for tactic in decode + prefill: megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) - for num_tokens in (64, 4096, 16384): - megamoe_op.validate_megamoe_tactic( - megamoe_op.default_megamoe_tactic(num_tokens), sm_version=sm_version - ) + default_tactics = [ + megamoe_op.default_megamoe_tactic(num_tokens) for num_tokens in (64, 4096, 16384) + ] + assert all(len(tactic) == 11 for tactic in default_tactics) + assert {tactic[-1] for tactic in default_tactics} == {"expert"} + for tactic in default_tactics: + megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) if sm_version == 107: for bucket, tactic in megamoe_op._SM107_GENPHASE_TACTICS.items(): megamoe_op.validate_megamoe_tactic(tactic, sm_version=sm_version) - assert ( - megamoe_op._default_megamoe_tactic_for_problem( - sm_version=sm_version, - max_tokens_per_rank=bucket, - num_tokens=bucket, - apply_topk_in_fc1=False, - in_kernel_fc2_reduce=False, - combine_format="bf16", - ) - == tactic + selected = megamoe_op._default_megamoe_tactic_for_problem( + sm_version=sm_version, + max_tokens_per_rank=bucket, + num_tokens=bucket, + apply_topk_in_fc1=False, + in_kernel_fc2_reduce=False, + combine_format="bf16", ) + assert len(selected) == 11 + assert selected == megamoe_op._unpack_tactic(tactic) invalid_tactic = list(megamoe_op.default_megamoe_tactic(64)) invalid_tactic[3] = ("grouped", 0) with pytest.raises(ValueError, match=r"schedule_policy hint must be a positive int or None"): @@ -3356,6 +3975,9 @@ def test_enumerate_megamoe_candidate_tactics_curated_space( legacy = ([256, 128, 256], [2, 1, 1], 512, "static", "epi_warps", True, 1, (1, 1)) megamoe_op.validate_megamoe_tactic(legacy, sm_version=sm_version) assert megamoe_op._unpack_tactic(legacy) == megamoe_op.default_megamoe_tactic(64) + v4 = megamoe_op.default_megamoe_tactic(64)[:-1] + megamoe_op.validate_megamoe_tactic(v4, sm_version=sm_version) + assert megamoe_op._unpack_tactic(v4) == megamoe_op.default_megamoe_tactic(64) @pytest.mark.gpu diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index b89f8476b958..070c273f3443 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -424,6 +424,35 @@ def test_pinned_megamoe_cutedsl_identity_resolves_to_the_leaf(): assert not report.degraded +def test_dynamic_eplb_is_a_declared_backend_capability(): + assert TrtllmCutedslMegaMoeNvfp4Impl.capabilities.supports_per_layer_dynamic_eplb + assert not DeepgemmCudaW4a8Mxfp4Mxfp8Impl.capabilities.supports_per_layer_dynamic_eplb + + +def test_per_layer_dynamic_eplb_filters_backends_by_declared_capability(): + deployment = dataclasses.replace( + _megamoe_cutedsl_deployment(), + eplb_enabled=True, + eplb_mode="per_layer_dynamic", + ) + accepted = resolve_moe_impl( + _megamoe_cutedsl_model_config(), + problem=_megamoe_cutedsl_problem(), + deployment=deployment, + impl_id=_MEGAMOE_CUTEDSL_IMPL_ID, + ) + assert impl_class_for(accepted) is TrtllmCutedslMegaMoeNvfp4Impl + + rejected = resolve_moe_impl( + _deepgemm_model_config(QuantAlgo.W4A8_MXFP4_MXFP8), + problem=_megamoe_cutedsl_problem(QuantAlgo.W4A8_MXFP4_MXFP8), + deployment=deployment, + impl_id="deepgemm.cuda.mega_moe.w4a8_mxfp4_mxfp8", + ) + assert rejected.winner is None + assert [item.reason for item in rejected.rejected] == [MoERejectReason.EPLB_UNSUPPORTED] + + def test_pinned_megamoe_cutedsl_identity_fails_hard_on_another_format(): """A pin that cannot run must fail, not widen to the sibling that can. diff --git a/tests/unittest/_torch/moe/test_moe_routing.py b/tests/unittest/_torch/moe/test_moe_routing.py index 31503d512dbd..babdd4d24e60 100644 --- a/tests/unittest/_torch/moe/test_moe_routing.py +++ b/tests/unittest/_torch/moe/test_moe_routing.py @@ -22,8 +22,8 @@ RenormalizeNaiveMoeRoutingMethod, SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod, create_load_balanced_logits, create_moe) from tensorrt_llm._torch.moe.fused_moe import routing as moe_routing -from tensorrt_llm._torch.moe.fused_moe.routing import \ - get_cached_perfect_router_logits +from tensorrt_llm._torch.moe.fused_moe.routing import ( + get_cached_perfect_router_expert_ids, get_cached_perfect_router_logits) from tensorrt_llm._torch.nccl_window_tensor_scope import \ discard_nccl_window_tensor_outputs from tensorrt_llm._utils import mpi_rank @@ -636,6 +636,66 @@ def test_perfect_router_get_cached_logits_e2e(routing_name, shape): assert torch.equal(logits, logits_2) +@pytest.mark.parametrize( + "num_experts,top_k,moe_ep_size", + [ + pytest.param(64, 6, 4, id="small"), + pytest.param(384, 6, 8, id="dsv4_pro_ep8"), + ], +) +def test_perfect_router_expert_ids_balanced(num_experts, top_k, moe_ep_size): + """ENABLE_PERFECT_ROUTER=expert_ids: balanced ids for any token count. + + Checks, on CPU, for every sender rank and a range of token counts: + - every EP rank and every expert receives an equal share of the routes, + within one; + - a token's routes land on distinct EP ranks and distinct experts; + - every call returns a view of the same table, sized to the token bound; + - a call above the bound raises instead of growing the table. + """ + device = torch.device("cpu") + experts_per_rank = num_experts // moe_ep_size + max_num_tokens = 8192 + moe_routing._PERFECT_ROUTER_EXPERT_IDS_CACHE.clear() + + def get_ids(num_tokens, ep_rank): + return get_cached_perfect_router_expert_ids( + num_tokens=num_tokens, + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=moe_ep_size, + ep_rank=ep_rank, + max_num_tokens=max_num_tokens, + device=device) + + for ep_rank in range(moe_ep_size): + table_ptr = None + for num_tokens in [1, 3, 7, 100, 1000, max_num_tokens]: + ids = get_ids(num_tokens, ep_rank) + assert ids.shape == (num_tokens, top_k) + assert ids.dtype == torch.int32 + assert ids.is_contiguous() + if table_ptr is None: + table_ptr = ids.data_ptr() + assert ids.data_ptr() == table_ptr + + rank_counts = torch.bincount(ids.flatten() // experts_per_rank, + minlength=moe_ep_size) + assert rank_counts.max() - rank_counts.min() <= 1 + expert_counts = torch.bincount(ids.flatten(), minlength=num_experts) + assert expert_counts.max() - expert_counts.min() <= 1 + + if top_k <= moe_ep_size: + ranks = ids // experts_per_rank + assert all(len(set(row)) == top_k for row in ranks.tolist()) + assert all(len(set(row)) == top_k for row in ids.tolist()) + + with pytest.raises(ValueError, match="max_num_tokens"): + get_ids(max_num_tokens + 1, ep_rank) + + assert len(moe_routing._PERFECT_ROUTER_EXPERT_IDS_CACHE) == moe_ep_size + + # ----------------------------------------------------------------- # Unified routing-method builder. Covers every RoutingMethodType in # tensorrt_llm._torch.moe.fused_moe.routing: diff --git a/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py new file mode 100644 index 000000000000..672d4dfbcc29 --- /dev/null +++ b/tests/unittest/_torch/moe/test_per_layer_dynamic_eplb.py @@ -0,0 +1,574 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from tensorrt_llm._torch.moe.fused_moe.interface import MoE +from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( + BaseMoeLoadBalancer, + MoeLoadBalancer, + MoeLoadBalancerIterContext, + get_moe_load_balancer, + maybe_create_moe_load_balancer, +) +from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import PerLayerDynamicMoeLoadBalancer + +pytestmark = pytest.mark.cpu_only + + +class _Config: + mode = "per_layer_dynamic" + num_slots = 24 + auxiliary_sms = 8 + + def setup(self, ep_rank, ep_size): + self.ep_rank = ep_rank + self.ep_size = ep_size + + def validate_expert_capacity(self, num_experts): + if self.num_slots <= num_experts: + raise ValueError("num_slots must exceed num_experts") + + +class _StandardConfig: + mode = "standard" + layer_updates_per_iter = 0 + + def setup(self, ep_rank, ep_size): + raise AssertionError("unsupported standard EPLB must remain disabled") + + +class _Arena: + home_experts = 4 + helper_slots = 2 + + def __init__(self, bundle="bundle"): + self.bundle = bundle + self.flags = torch.empty(4, dtype=torch.uint64) + self.generation = 0 + self.provider = SimpleNamespace(pool=None) + + def publish_ready_generation(self, generation): + self.generation = generation + + def terminal_flags_tensor(self): + return self.flags + + def assert_identity(self): + pass + + +class _Backend: + def __init__(self): + self.tactic_autotune = False + self.mapping = SimpleNamespace(local_rank=0) + self.ep_process_group = object() + self.max_tokens_per_rank_capacity = 256 + self.layer_idx = 9 + + +class _Group: + home_experts = 4 + helper_slots = 2 + + def __init__(self): + self.wait_calls = 0 + self.finish_calls = [] + self.events = [] + + def plan(self, routes, *, defer_wait): + assert defer_wait + self.events.append("plan") + return routes + 4, 7 + + def plan_schedule(self, routes): + self.events.append("schedule") + return routes + + def plan_finish(self, part, *, defer_wait): + assert defer_wait + self.events.append("finish") + return part + 4, 7 + + def discard_plan(self, part): + self.events.append("discard") + + def wait_for_routes(self): + self.wait_calls += 1 + + def finish(self, completion_event=None): + self.finish_calls.append(completion_event) + + def close(self): + pass + + +def _manager(): + return PerLayerDynamicMoeLoadBalancer( + ep_rank=0, + ep_size=4, + config=_Config(), + ) + + +def _allocate_arena(layer, monkeypatch, arena): + """Allocate ``arena`` through the layer handle, as the quant method does.""" + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module + + monkeypatch.setattr(layer.manager, "get_ep_comm", lambda process_group: object()) + monkeypatch.setattr(layer.manager, "create_arena_provider", lambda: object()) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", MagicMock(return_value=arena)) + return layer.allocate_arena( + bundle=arena.bundle, + ep_process_group=object(), + mapping=object(), + device=0, + layer_idx=layer.layer_idx, + ) + + +def test_layer_capacity_uses_configured_compute_slots(): + manager = _manager() + with pytest.raises(ValueError, match="configured local slot count"): + manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=4) + + +def test_layer_produces_and_finishes_replica_plan(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + group = _Group() + layer.bind_runtime(arena=arena, scheduler_group=group) + + logical = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + physical = layer.route(logical) + plan = layer.require_replica_plan() + + torch.testing.assert_close(physical, logical + 4) + assert plan.resident_experts_per_rank == 4 + assert plan.compute_slots_per_rank == 6 + assert plan.helper_slots_per_rank == 2 + assert plan.ready_generation == arena.generation == 7 + assert plan.ready_flags is arena.flags + assert plan.reserved_sms == 8 + plan.wait_for_routes() + assert group.wait_calls == 1 + + layer.finish_replica_plan(plan) + assert group.finish_calls == [None] + assert layer.replica_plan is None + + +def test_post_moe_hook_releases_the_chunk_plan(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + layer.route(torch.tensor([[0, 1]], dtype=torch.int32)) + layer.start_set_cpu_stage() + + assert group.wait_calls == 1 + assert group.finish_calls == [None] + assert layer.replica_plan is None + + +def test_overlap_fn_runs_between_schedule_and_plan_finish(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + group = _Group() + layer.bind_runtime(arena=arena, scheduler_group=group) + + logical = torch.tensor([[0, 1]], dtype=torch.int32) + physical = layer.route(logical, overlap_fn=lambda: group.events.append("overlap")) + torch.testing.assert_close(physical, logical + 4) + assert group.events == ["schedule", "overlap", "finish"] + layer.finish_replica_plan() + + +def test_overlap_fn_failure_discards_generation_and_adds_context(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + def fail(): + raise RuntimeError("shared expert failed") + + with pytest.raises(RuntimeError, match="shared expert failed") as caught: + layer.route(torch.tensor([[0, 1]], dtype=torch.int32), overlap_fn=fail) + assert group.events == ["schedule", "discard"] + assert "safely released" in caught.value.__notes__[0] + assert layer.replica_plan is None + + +def test_load_balancer_overlap_is_consumed_once_and_always_cleared(): + calls = [] + + class _Layer: + produces_replica_plan = True + + def route(self, routes, offset_by_ep_rank, *, overlap_fn=None): + calls.append(overlap_fn) + return routes + + moe = SimpleNamespace(layer_load_balancer=_Layer(), _load_balancer_overlap_fn=None) + routes = torch.tensor([[0, 1]], dtype=torch.int32) + + def overlap(): + pass + + with MoE.load_balancer_overlap(moe, overlap): + MoE._load_balancer_route(moe, routes, False) + MoE._load_balancer_route(moe, routes, False) + assert calls == [overlap, None] + assert moe._load_balancer_overlap_fn is None + + with pytest.raises(RuntimeError, match="routed path failed"): + with MoE.load_balancer_overlap(moe, overlap): + raise RuntimeError("routed path failed") + assert moe._load_balancer_overlap_fn is None + + with pytest.raises(TypeError, match="overlap_fn must be callable"): + with MoE.load_balancer_overlap(moe, None): + pass + assert MoE._load_balancer_uses_replica_plan(moe) + + +def test_layer_disables_standard_host_migration_and_checks_geometry(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + assert not layer.is_dynamic_routing() + assert not layer.need_load_shared_weights() + assert layer.get_load_expert_ids() == [0, 1, 2, 3] + assert layer.get_local_statistic_tensor() is None + layer.set_initial_weight_assignments(range(16)) + with pytest.raises(ValueError, match="custom initial_global_assignments"): + layer.set_initial_weight_assignments([0]) + with pytest.raises(ValueError, match="arena H/S geometry"): + layer.bind_runtime( + arena=SimpleNamespace(home_experts=5, helper_slots=1), + scheduler_group=_Group(), + ) + + +def test_bind_backend_is_idempotent_before_scheduler_construction(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import planner + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + assert _allocate_arena(layer, monkeypatch, arena) is arena + backend = _Backend() + group = _Group() + builder = MagicMock(return_value=group) + monkeypatch.setattr( + planner, + "build_rebalance_slot_scheduler_group", + builder, + ) + + layer.bind_backend(backend) + layer.bind_backend(backend) + + # The group is built from the layer's geometry, the manager's resources and + # the backend's public facts; layer_idx stays the backend's model index. + builder.assert_called_once_with( + arena=arena, + manager=manager, + mapping=backend.mapping, + ep_process_group=backend.ep_process_group, + home_experts=4, + helper_slots=2, + topk=2, + max_tokens_per_rank=256, + layer_idx=9, + ) + assert layer._arena is arena + assert layer._scheduler_group is group + + +def test_bind_backend_requires_the_layer_arena_and_rejects_another_backend( + monkeypatch, +): + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import planner + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + backend = _Backend() + builder = MagicMock(return_value=_Group()) + monkeypatch.setattr( + planner, + "build_rebalance_slot_scheduler_group", + builder, + ) + with pytest.raises(RuntimeError, match="no live-weight arena"): + layer.bind_backend(backend) + + arena = _Arena() + _allocate_arena(layer, monkeypatch, arena) + # The runtime binds only the arena this handle allocated. + with pytest.raises(RuntimeError, match="did not allocate"): + layer.bind_runtime(arena=_Arena(), scheduler_group=_Group()) + layer.bind_backend(backend) + + with pytest.raises(RuntimeError, match="different backend"): + layer.bind_backend(_Backend()) + builder.assert_called_once() + assert layer.live_weight_arena is arena + + +def test_iteration_context_reuses_common_lifecycle(): + manager = _manager() + with manager: + assert get_moe_load_balancer() is manager + with MoeLoadBalancerIterContext(manager, True, True): + assert manager.in_iter + assert not manager.in_iter + assert manager.iter_id == 1 + assert get_moe_load_balancer() is None + + +def test_iteration_context_releases_the_plan_of_a_failed_forward(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + group = _Group() + layer.bind_runtime(arena=_Arena(), scheduler_group=group) + + with pytest.raises(RuntimeError, match="MoE kernel failed"): + with MoeLoadBalancerIterContext(manager, True, True): + layer.route(torch.tensor([[0, 1]], dtype=torch.int32)) + raise RuntimeError("MoE kernel failed") + + assert group.finish_calls == [None] + assert layer.replica_plan is None + assert not manager.in_iter + assert manager.iter_id == 1 + + +def test_model_owned_shutdown(): + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + arena = _Arena() + pool = MagicMock() + pool.close_collectively.return_value = True + arena.provider.pool = pool + manager._rebalance_shared_slot_pools[("test",)] = pool + manager._rebalance_shared_slot_pool_order.append(pool) + layer.bind_runtime(arena=arena, scheduler_group=_Group()) + + assert manager.shutdown() + pool.close_collectively.assert_called_once_with(local_safe=True) + assert manager.is_shutdown + + +def test_ep_communicator_is_model_owned(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module + + creations = [] + + class Comm: + def __init__(self, process_group): + self.process_group = process_group + self.validations = [] + creations.append(self) + + def validate_geometry(self, *, expected_rank, expected_world): + self.validations.append((expected_rank, expected_world)) + + monkeypatch.setattr(arena_module, "_EpComm", Comm) + manager = _manager() + process_group = object() + + first = manager.get_ep_comm(process_group) + second = manager.get_ep_comm(process_group) + + assert first is second + assert len(creations) == 1 + assert first.validations == [(0, 4), (0, 4)] + with pytest.raises(RuntimeError, match="cannot span EP ProcessGroups"): + manager.get_ep_comm(object()) + + +def test_layer_allocates_arena_from_model_owned_resources(monkeypatch): + from tensorrt_llm._torch.moe.fused_moe.per_layer_dynamic_eplb import arena as arena_module + + manager = _manager() + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + order = [] + process_group, comm, provider, bundle, mapping = (object() for _ in range(5)) + arena = SimpleNamespace(bundle=bundle) + + def get_ep_comm(group): + assert group is process_group + order.append("ep_comm") + return comm + + def create_arena_provider(): + order.append("provider") + return provider + + monkeypatch.setattr(manager, "get_ep_comm", get_ep_comm) + monkeypatch.setattr(manager, "create_arena_provider", create_arena_provider) + allocate = MagicMock(return_value=arena) + monkeypatch.setattr(arena_module, "allocate_rebalance_arena", allocate) + kwargs = dict(ep_process_group=process_group, mapping=mapping, device=1, layer_idx=7) + + allocated = layer.allocate_arena(bundle=bundle, **kwargs) + + assert allocated is arena + assert layer.live_weight_arena is arena + assert order == ["ep_comm", "provider"] + allocate.assert_called_once_with( + home_experts=4, + helper_slots=2, + bundle=bundle, + ep_comm=comm, + mapping=mapping, + provider=provider, + device=1, + layer_idx=7, + ) + + # Allocation happens once per layer; only the same plane spec gets the arena back. + assert layer.allocate_arena(bundle=bundle, **kwargs) is arena + with pytest.raises(RuntimeError, match="different live-weight plane spec"): + layer.allocate_arena(bundle=object(), **kwargs) + allocate.assert_called_once() + assert order == ["ep_comm", "provider"] + + +@pytest.mark.parametrize( + ("helper_slot_mode", "per_layer_sets"), + [(None, True), ("per_layer_resident", True), ("shared_banks", False)], +) +def test_helper_slot_mode_selects_helper_residency(helper_slot_mode, per_layer_sets): + config = _Config() + if helper_slot_mode is not None: + config.helper_slot_mode = helper_slot_mode + manager = PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, config=config) + + assert manager.helper_residency is per_layer_sets + provider = manager.create_arena_provider() + assert provider._per_layer_helper_sets is per_layer_sets + + +def test_helper_slot_mode_rejects_unknown_mode(): + config = _Config() + config.helper_slot_mode = "private" + with pytest.raises(ValueError, match="helper_slot_mode"): + PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, config=config) + + +def test_scheduler_spin_bound_follows_the_warmup_guard_until_shutdown(monkeypatch): + from tensorrt_llm._torch.autotuner import AutoTuner + from tensorrt_llm._torch.moe.fused_moe import moe_comm_timeout_guard + + guard = moe_comm_timeout_guard.MoECommTimeoutGuard(environ={}, is_capturing=lambda: False) + monkeypatch.setattr(moe_comm_timeout_guard, "_DEFAULT_GUARD", guard) + # The real AutoTuner reads the CUDA device name; this test runs without a GPU. + tuner = SimpleNamespace(is_tuning_mode=False) + monkeypatch.setattr(AutoTuner, "get", lambda: tuner) + manager = _manager() + + bound = manager.get_scheduler_spin_bound() + + assert manager.get_scheduler_spin_bound() is bound + assert bound.cycles() == 300 * 2_000_000_000 + tuner.is_tuning_mode = True + assert bound.cycles() == 1800 * 2_000_000_000 + tuner.is_tuning_mode = False + guard.set_warmup(True) + assert bound.cycles() == 1800 * 2_000_000_000 + guard.set_warmup(False) + assert bound.cycles() == 300 * 2_000_000_000 + assert manager.shutdown() + guard.set_warmup(True) + # A released bound no longer follows the guard. + assert bound.cycles() == 300 * 2_000_000_000 + + +def test_per_layer_dynamic_manager_shares_the_base_lifecycle_only(): + manager = _manager() + assert isinstance(manager, BaseMoeLoadBalancer) + assert not isinstance(manager, MoeLoadBalancer) + assert manager.is_dynamic_routing() + assert repr(manager) == ( + "PerLayerDynamicMoeLoadBalancer(ep_rank=0, ep_size=4, num_slots=24, auxiliary_sms=8)" + ) + layer = manager.add_layer(expert_count=16, top_k=2, slot_count_per_rank=6) + assert layer.produces_replica_plan + + +def test_factory_selects_per_layer_dynamic_manager(): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=4, + moe_cluster_size=1, + ) + config = _Config() + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), + moe_load_balancer=config, + ) + + with maybe_create_moe_load_balancer(model_config, mapping) as manager: + assert isinstance(manager, PerLayerDynamicMoeLoadBalancer) + assert config.ep_rank == 0 + assert config.ep_size == 4 + + +@pytest.mark.parametrize( + "architecture, ep_size, cluster_size, message", + [ + ("UnsupportedForCausalLM", 4, 1, "does not support model architecture"), + ("DeepseekV4ForCausalLM", 1, 1, "requires expert parallelism"), + ("DeepseekV4ForCausalLM", 4, 2, "incompatible with smart routing"), + ], +) +def test_per_layer_dynamic_factory_fails_fast_for_unsupported_topology( + architecture, ep_size, cluster_size, message +): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=ep_size, + moe_cluster_size=cluster_size, + ) + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=[architecture]), + moe_load_balancer=_Config(), + ) + + with pytest.raises(ValueError, match=message): + maybe_create_moe_load_balancer(model_config, mapping) + + +@pytest.mark.parametrize( + "architecture, ep_size, cluster_size", + [ + ("UnsupportedForCausalLM", 4, 1), + ("DeepseekV4ForCausalLM", 1, 1), + ("DeepseekV4ForCausalLM", 4, 2), + ], +) +def test_standard_factory_keeps_unsupported_topology_disabled(architecture, ep_size, cluster_size): + mapping = SimpleNamespace( + moe_ep_rank=0, + moe_ep_size=ep_size, + moe_cluster_size=cluster_size, + ) + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(architectures=[architecture]), + moe_load_balancer=_StandardConfig(), + ) + + with maybe_create_moe_load_balancer(model_config, mapping) as manager: + assert manager is None diff --git a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py index e3c31c33712f..b992d39812f1 100644 --- a/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py +++ b/tests/unittest/_torch/weight_sharing/_source_identity_fakes.py @@ -151,6 +151,7 @@ def __init__( quant_config=_UNSET, attn_backend: str = "TRTLLM", moe_backend: str = "CUTLASS", + moe_load_balancer=None, ): self.mapping = mapping or FakeMapping() self.pretrained_config = ( @@ -161,6 +162,7 @@ def __init__( self.force_dynamic_quantization = False self.attn_backend = attn_backend self.moe_backend = moe_backend + self.moe_load_balancer = moe_load_balancer self.nvfp4_gemm_allowed_backends = ["cutlass", "cublaslt", "cuda_core"] self.moe_disable_finalize_fusion = False self.use_low_precision_moe_combine = False diff --git a/tests/unittest/_torch/weight_sharing/test_source_identity.py b/tests/unittest/_torch/weight_sharing/test_source_identity.py index d03964feaed8..9b8128b05c85 100644 --- a/tests/unittest/_torch/weight_sharing/test_source_identity.py +++ b/tests/unittest/_torch/weight_sharing/test_source_identity.py @@ -44,6 +44,7 @@ SourceIdentityMismatchError, check_weight_sharing_compatibility, ) +from tensorrt_llm.llmapi.llm_args import MoeLoadBalancerConfig, MoePerLayerDynamicLoadBalancerConfig pytestmark = pytest.mark.cpu_only @@ -120,6 +121,21 @@ def test_backend_mismatch_flags_global(): assert "backend_fingerprint" in result.mismatched_fields +def test_moe_load_balancer_config_controls_backend_identity(): + base = identity_from(FakeModelConfig()) + standard = identity_from(FakeModelConfig(moe_load_balancer=MoeLoadBalancerConfig(num_slots=64))) + per_layer_dynamic = identity_from( + FakeModelConfig(moe_load_balancer=MoePerLayerDynamicLoadBalancerConfig(num_slots=72)) + ) + same_per_layer_dynamic = identity_from( + FakeModelConfig(moe_load_balancer=MoePerLayerDynamicLoadBalancerConfig(num_slots=72)) + ) + + assert not base.matches(standard).matched + assert not standard.matches(per_layer_dynamic).matched + assert per_layer_dynamic.matches(same_per_layer_dynamic).matched + + def test_model_layout_field_mismatch_flags_global_and_shard(): a = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(hidden_size=4096))) b = identity_from(FakeModelConfig(pretrained_config=FakePretrainedConfig(hidden_size=8192))) diff --git a/tests/unittest/executor/test_proxy_request_sender.py b/tests/unittest/executor/test_proxy_request_sender.py new file mode 100644 index 000000000000..7dbeb489063d --- /dev/null +++ b/tests/unittest/executor/test_proxy_request_sender.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for request batching between the proxy and the rank-0 worker. + +``TRTLLM_PROXY_BATCH_REQUESTS=1`` makes the proxy pickle each request on the +submitting thread and send bursts of payloads as one list; the worker unpickles +each payload. +""" + +import pickle +import queue +import threading +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.executor.proxy import ( + GenerationExecutorProxy, + _create_request_sender, + _RequestSender, +) +from tensorrt_llm.executor.worker import _incoming_requests + +pytestmark = pytest.mark.cpu_only + + +class _RecordingQueue: + """Request-queue stand-in that records each message; the first put can block. + + Like ``ZeroMqQueue`` with ``TLLM_LLMAPI_ZMQ_DEBUG=1``, it rejects a write from + any thread other than the first one that wrote. + """ + + def __init__(self, block_first_put: bool = False): + self.messages = [] + self.writer = None + self.first_put_started = threading.Event() + self.release_first_put = threading.Event() + if not block_first_put: + self.release_first_put.set() + + def _check_writer(self): + if self.writer is None: + self.writer = threading.get_ident() + elif self.writer != threading.get_ident(): + raise RuntimeError("request queue written from a second thread") + + def put(self, msg): + self._check_writer() + if not self.messages: + self.first_put_started.set() + self.release_first_put.wait(timeout=10) + self.messages.append(msg) + + def put_noblock(self, msg, *, retry=1, wait_time=0.001): + self._check_writer() + self.messages.append(msg) + + +def _payloads(*items): + return [pickle.dumps(item) for item in items] + + +def _bare_proxy(): + proxy = GenerationExecutorProxy.__new__(GenerationExecutorProxy) + # Keeps the __del__ -> shutdown() path a no-op at GC time. + proxy.workers_started = False + return proxy + + +def _pre_shutdown(proxy): + """Run ``pre_shutdown`` on a bare proxy with live workers, on a new thread.""" + proxy.workers_started = True + proxy.doing_shutdown = False + proxy._worker_process_monitor = SimpleNamespace(close=lambda: None) + proxy._results = {} + proxy.mpi_futures = [] + shutdown = threading.Thread(target=proxy.pre_shutdown) + shutdown.start() + return shutdown + + +def test_batching_is_opt_in(monkeypatch): + monkeypatch.delenv("TRTLLM_PROXY_BATCH_REQUESTS", raising=False) + assert _create_request_sender(_RecordingQueue(), queue.Queue()) is None + monkeypatch.setenv("TRTLLM_PROXY_BATCH_REQUESTS", "1") + assert isinstance(_create_request_sender(_RecordingQueue(), queue.Queue()), _RequestSender) + + +def test_lone_request_is_sent_as_one_payload_list(): + request_queue = _RecordingQueue() + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert request_queue.messages == [_payloads("a")] + + +def test_burst_is_sent_as_one_list_in_order(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps(0)) + assert request_queue.first_put_started.wait(timeout=10) + # Everything queued while the first send is in flight goes out as one list. + for item in range(1, 6): + sender.send(pickle.dumps(item)) + request_queue.release_first_put.set() + sender.close(timeout=10) + assert request_queue.messages == [_payloads(0), _payloads(1, 2, 3, 4, 5)] + + +def test_close_flushes_queued_requests_and_drops_later_ones(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + closer = threading.Thread(target=sender.close, kwargs={"timeout": 10}) + closer.start() + request_queue.release_first_put.set() + closer.join(timeout=10) + assert not closer.is_alive() + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + sender.send(pickle.dumps("late")) + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + + +def test_close_is_idempotent(): + request_queue = _RecordingQueue() + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert sender._thread is None + sender.close(timeout=10) + assert request_queue.messages == [_payloads("a")] + + +def test_sender_thread_sends_the_sentinel_after_the_queued_requests(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + result = [] + closer = threading.Thread( + target=lambda: result.append(sender.close(timeout=10, send_sentinel=True)) + ) + closer.start() + request_queue.release_first_put.set() + closer.join(timeout=10) + assert result == [False] + # A bare None after the payload lists, written by the sender thread. + assert request_queue.messages == [_payloads("first"), _payloads("queued"), None] + + +def test_caller_sends_the_sentinel_when_the_sender_thread_did_not(): + # The thread never started, so nothing has written the queue yet. + assert _RequestSender(_RecordingQueue(), queue.Queue()).close(send_sentinel=True) + # The thread stopped before the sentinel was asked for. + sender = _RequestSender(_RecordingQueue(), queue.Queue()) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + assert sender.close(timeout=10, send_sentinel=True) + + +def test_close_during_interpreter_finalization_does_not_join(): + request_queue = _RecordingQueue(block_first_put=True) + sender = _RequestSender(request_queue, queue.Queue()) + sender.send(pickle.dumps("first")) + assert request_queue.first_put_started.wait(timeout=10) + sender.send(pickle.dumps("queued")) + # Joining would wait for the blocked send; at teardown close() must return at + # once, and must not ask anyone to send the sentinel. + sender._is_finalizing = lambda: True + assert sender.close(send_sentinel=True) is False + assert sender._thread is not None + # Outside finalization the same sender still flushes and stops. + sender._is_finalizing = lambda: False + request_queue.release_first_put.set() + sender.close(timeout=10) + assert sender._thread is None + assert request_queue.messages == [_payloads("first"), _payloads("queued")] + + +def test_pre_shutdown_flushes_batched_requests_before_the_sentinel(): + request_queue = _RecordingQueue(block_first_put=True) + proxy = _bare_proxy() + proxy.request_queue = request_queue + proxy._request_sender = _RequestSender(request_queue, queue.Queue()) + proxy._send_request("first") + assert request_queue.first_put_started.wait(timeout=10) + proxy._send_request("queued") + + shutdown = _pre_shutdown(proxy) + request_queue.release_first_put.set() + shutdown.join(timeout=10) + assert not shutdown.is_alive() + proxy.workers_started = False + # The sender thread wrote everything, the sentinel included. + assert request_queue.messages == [_payloads("first"), _payloads("queued"), None] + assert request_queue.writer != shutdown.ident + + +@pytest.mark.parametrize("batching", [True, False]) +def test_pre_shutdown_sends_the_sentinel_itself_if_nothing_else_wrote(batching): + request_queue = _RecordingQueue() + proxy = _bare_proxy() + proxy.request_queue = request_queue + if batching: + proxy._request_sender = _RequestSender(request_queue, queue.Queue()) + + shutdown = _pre_shutdown(proxy) + shutdown.join(timeout=10) + assert not shutdown.is_alive() + proxy.workers_started = False + assert request_queue.messages == [None] + assert request_queue.writer == shutdown.ident + + +def test_transport_error_reaches_error_queue(): + class _FailingQueue: + def put(self, msg): + raise RuntimeError("socket closed") + + error_queue = queue.Queue() + sender = _RequestSender(_FailingQueue(), error_queue) + sender.send(pickle.dumps("a")) + sender.close(timeout=10) + error = error_queue.get_nowait() + assert isinstance(error, RuntimeError) + assert str(error) == "socket closed" + + +def test_batched_request_is_captured_when_sent(): + proxy = _bare_proxy() + sent = [] + proxy._request_sender = SimpleNamespace(send=sent.append) + + request = {"max_tokens": 1} + proxy._send_request(request) + # A change after submission does not reach the worker. + request["max_tokens"] = 2 + assert pickle.loads(sent[0]) == {"max_tokens": 1} + + # A request that cannot be pickled fails in its caller, and nothing is queued. + with pytest.raises(TypeError): + proxy._send_request({"lock": threading.Lock()}) + assert len(sent) == 1 + + +def test_unbatched_request_is_put_directly(): + proxy = _bare_proxy() + request_queue = _RecordingQueue() + proxy.request_queue = request_queue + proxy._send_request("a") + assert request_queue.messages == ["a"] + + +def test_worker_unpacks_batched_requests_in_order(): + class _ListQueue: + def __init__(self, messages): + self._messages = list(messages) + + def get(self): + return self._messages.pop(0) + + request_queue = _ListQueue(["a", _payloads("b", "c"), "d", None, "after sentinel"]) + assert list(_incoming_requests(request_queue)) == ["a", "b", "c", "d"] diff --git a/tests/unittest/executor/test_shutdown_retry_contracts.py b/tests/unittest/executor/test_shutdown_retry_contracts.py new file mode 100644 index 000000000000..f9bb440e5235 --- /dev/null +++ b/tests/unittest/executor/test_shutdown_retry_contracts.py @@ -0,0 +1,736 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only lifecycle contracts for retryable Ray and RPC shutdown.""" + +import ast +import sys +import threading +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(__file__).resolve().parents[3] + + +def _extract_function(source, function_name, namespace): + tree = ast.parse(source.read_text()) + function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == function_name + ) + module = ast.fix_missing_locations(ast.Module(body=[function], type_ignores=[])) + exec(compile(module, str(source), "exec"), namespace) + return namespace[function_name] + + +def _extract_class(source, class_name, methods, namespace, bases=()): + tree = ast.parse(source.read_text()) + original = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name + ) + body = [ + node + for node in original.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in methods + ] + derived = ast.ClassDef( + name=class_name, + bases=[ast.Name(id=base, ctx=ast.Load()) for base in bases], + keywords=[], + body=body, + decorator_list=[], + ) + module = ast.fix_missing_locations(ast.Module(body=[derived], type_ignores=[])) + exec(compile(module, str(source), "exec"), namespace) + return namespace[class_name] + + +_uses_collective_shutdown = _extract_function( + _ROOT / "tensorrt_llm/llmapi/_load_balance_env.py", "uses_collective_shutdown", {"Any": object} +) + + +def _llm_args(per_layer_dynamic=True, **fields): + load_balancer = SimpleNamespace(mode="per_layer_dynamic") if per_layer_dynamic else None + return SimpleNamespace(moe_config=SimpleNamespace(load_balancer=load_balancer), **fields) + + +def test_collective_shutdown_is_limited_to_per_layer_dynamic_eplb(): + assert _uses_collective_shutdown(_llm_args()) + assert not _uses_collective_shutdown(_llm_args(per_layer_dynamic=False)) + assert not _uses_collective_shutdown(SimpleNamespace(backend="pytorch")) + assert not _uses_collective_shutdown(None) + + +class _Timeout(Exception): + pass + + +class _ActorDied(Exception): + pass + + +class _Ref: + def __init__(self, outcomes): + self.outcomes = list(outcomes) + + def get(self): + outcome = self.outcomes.pop(0) if len(self.outcomes) > 1 else self.outcomes[0] + if isinstance(outcome, BaseException): + raise outcome + return outcome + + +class _RemoteShutdown: + def __init__(self, refs): + self.refs = list(refs) + self.calls = 0 + + def remote(self): + ref = self.refs[min(self.calls, len(self.refs) - 1)] + self.calls += 1 + return ref + + +class _RayWorker: + def __init__(self, refs): + self.shutdown = _RemoteShutdown(refs) + + +class _FakeRay: + def __init__(self): + self.exceptions = SimpleNamespace(GetTimeoutError=_Timeout, RayActorError=_ActorDied) + self.util = SimpleNamespace(remove_placement_group=lambda group: None) + self.killed = [] + + @staticmethod + def get(ref, timeout=None): + return ref.get() + + def kill(self, worker, no_restart=True): + self.killed.append(worker) + + @staticmethod + def is_initialized(): + return False + + @staticmethod + def shutdown(): + return None + + +def _ray_executor(ray, per_layer_dynamic=True): + cls = _extract_class( + _ROOT / "tensorrt_llm/executor/ray/executor.py", + "RayExecutor", + {"shutdown", "_shutdown_workers_collectively"}, + { + "ray": ray, + "uses_collective_shutdown": _uses_collective_shutdown, + "time": __import__("time"), + "logger": SimpleNamespace( + warning=lambda *args, **kwargs: None, debug=lambda *args, **kwargs: None + ), + "logger_debug": lambda *args, **kwargs: None, + }, + ) + executor = cls() + executor.worker_kwargs = {"llm_args": _llm_args(per_layer_dynamic)} + executor._shutdown_event = threading.Event() + executor.main_loop = None + executor.main_loop_task_obj = None + executor.rpc_client = SimpleNamespace(close=lambda: None) + executor.placement_group = None + executor.bundle_indices = None + executor.has_start_local_cluser = False + executor._wait_for_cluster_resource_release = lambda timeout: None + return executor + + +def test_ray_shutdown_failure_preserves_actors_for_retry(): + ray = _FakeRay() + worker = _RayWorker([_Ref([RuntimeError("cleanup failed")]), _Ref([None])]) + executor = _ray_executor(ray) + executor.workers = [worker] + + with pytest.raises(RuntimeError, match="cleanup failed"): + executor.shutdown() + + assert executor.workers == [worker] + assert not executor._shutdown_event.is_set() + assert ray.killed == [] + + executor.shutdown() + assert worker.shutdown.calls == 2 + assert ray.killed == [worker] + assert executor.workers is None + + +def test_ray_shutdown_timeout_reuses_pending_ref(): + ray = _FakeRay() + pending = _Ref([_Timeout("still running"), None]) + worker = _RayWorker([pending]) + executor = _ray_executor(ray) + executor.workers = [worker] + + with pytest.raises(_Timeout, match="still running"): + executor.shutdown() + assert executor._pending_shutdown_refs == [pending] + + executor.shutdown() + assert worker.shutdown.calls == 1 + assert executor.workers is None + + +def test_ray_dead_actor_falls_through_to_forced_teardown(): + ray = _FakeRay() + dead = _RayWorker([_Ref([_ActorDied("actor died")])]) + live = _RayWorker([_Ref([None])]) + executor = _ray_executor(ray) + executor.workers = [dead, live] + + executor.shutdown() + + assert ray.killed == [dead, live] + assert executor.workers is None + assert executor._pending_shutdown_refs is None + + +def test_ray_shutdown_without_per_layer_eplb_tears_down_best_effort(): + ray = _FakeRay() + worker = _RayWorker([_Ref([RuntimeError("cleanup failed")])]) + executor = _ray_executor(ray, per_layer_dynamic=False) + executor.workers = [worker] + + executor.shutdown() + + assert ray.killed == [worker] + assert executor.workers is None + + +class _BaseWorker: + def shutdown(self): + outcome = self.base_outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + + +class _Comm: + def __init__(self, gathered): + self.gathered = list(gathered) + self.broadcasts = [] + + def allgather(self, value): + return self.gathered.pop(0) + + def bcast(self, value, root): + self.broadcasts.append(value) + return value + + +def _rpc_worker(comm): + cls = _extract_class( + _ROOT / "tensorrt_llm/executor/rpc_worker.py", + "RpcWorker", + {"shutdown", "commit_shutdown", "_complete_shutdown_commit"}, + { + "Base": _BaseWorker, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: comm, + "logger_debug": lambda *args, **kwargs: None, + "_RPC_SHUTDOWN_RETRY": "retry", + "_RPC_SHUTDOWN_COMMIT": "commit", + }, + bases=("Base",), + ) + worker = cls() + worker._rpc_collective_shutdown = True + worker._rpc_shutdown_attempted = False + worker._rpc_shutdown_ready = False + worker._rpc_shutdown_committed = False + worker.shutdown_event = threading.Event() + return worker + + +def test_rpc_cleanup_retry_and_commit_are_rank_aligned_and_idempotent(): + comm = _Comm([["RuntimeError: local failure", None], [None, None]]) + worker = _rpc_worker(comm) + worker.base_outcomes = [RuntimeError("local failure"), None] + + with pytest.raises(RuntimeError, match="local failure"): + worker.shutdown() + assert not worker._rpc_shutdown_ready + assert comm.broadcasts == [] + + worker.shutdown() + assert worker._rpc_shutdown_ready + assert comm.broadcasts == ["retry"] + + worker.commit_shutdown() + worker.commit_shutdown() + assert comm.broadcasts == ["retry", "commit"] + assert not worker.shutdown_event.is_set() + + worker._complete_shutdown_commit() + assert worker.shutdown_event.is_set() + + +def test_rpc_peer_failure_is_propagated_to_rank_zero(): + comm = _Comm([[None, "RuntimeError: peer failure"]]) + worker = _rpc_worker(comm) + worker.base_outcomes = [None] + + with pytest.raises(RuntimeError, match="rank 1: RuntimeError: peer failure"): + worker.shutdown() + assert not worker._rpc_shutdown_ready + + +class _RemoteCall: + def __init__(self, name, calls, error=None): + self.name = name + self.calls = calls + self.error = error + + def remote(self, *, need_response, timeout=None): + self.calls.append((self.name, need_response, timeout)) + if self.error is not None: + raise self.error + + +class _RpcClient: + def __init__(self, calls, shutdown_error=None): + self.calls = calls + self.shutdown_error = shutdown_error + + def shutdown(self): + return _RemoteCall("cleanup", self.calls, self.shutdown_error) + + def commit_shutdown(self): + return _RemoteCall("commit", self.calls) + + def close(self): + self.calls.append(("close", True, None)) + + +def _rpc_proxy(per_layer_dynamic=True): + cls = _extract_class( + _ROOT / "tensorrt_llm/executor/rpc_proxy.py", + "GenerationExecutorRpcProxy", + {"_uses_collective_shutdown", "shutdown_remote", "shutdown", "_shutdown_rank_local"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "logger": SimpleNamespace(warning=lambda *args, **kwargs: None), + "threading": threading, + }, + ) + cls._COLLECTIVE_SHUTDOWN_TIMEOUT_S = 300.0 + proxy = cls() + proxy.worker_kwargs = {"llm_args": _llm_args(per_layer_dynamic)} + return proxy + + +def test_rpc_proxy_requires_cleanup_and_commit_responses(): + calls = [] + proxy = _rpc_proxy() + proxy.rpc_client = _RpcClient(calls) + + proxy.shutdown_remote() + + assert calls == [("cleanup", True, 300.0), ("commit", True, 300.0)] + + +def test_rpc_proxy_without_per_layer_eplb_sends_one_unacknowledged_shutdown(): + calls = [] + proxy = _rpc_proxy(per_layer_dynamic=False) + proxy.rpc_client = _RpcClient(calls) + + proxy.shutdown_remote() + + assert calls == [("cleanup", False, None)] + + +def test_rpc_proxy_cleanup_failure_keeps_retry_gate_and_resources_live(): + calls = [] + proxy = _rpc_proxy() + proxy._shutdown_event = threading.Event() + proxy.rpc_client = _RpcClient(calls, RuntimeError("remote cleanup failed")) + proxy.main_loop = None + proxy.main_loop_task_obj = None + proxy.main_loop_thread = None + proxy.mpi_session = object() + + with pytest.raises(RuntimeError, match="remote cleanup failed"): + proxy.shutdown() + + assert not proxy._shutdown_event.is_set() + assert proxy.mpi_session is not None + assert calls == [("cleanup", True, 300.0)] + + +def test_rpc_server_completes_commit_only_after_response_send(): + source = _ROOT / "tensorrt_llm/executor/rpc/rpc_server.py" + tree = ast.parse(source.read_text()) + server = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "RPCServer" + ) + process = next( + node + for node in server.body + if isinstance(node, ast.AsyncFunctionDef) and node.name == "_process_requests" + ) + send_line = min( + node.lineno + for node in ast.walk(process) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_send_response" + ) + complete_line = min( + node.lineno + for node in ast.walk(process) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "complete_commit" + ) + method_source = ast.get_source_segment(source.read_text(), process) + + assert send_line < complete_line + assert 'req.method_name == "commit_shutdown"' in method_source + assert "response.error is None" in method_source + + +def test_symm_buffer_release_retains_failures_for_retry(): + cache = { + "failed": ("bad", object()), + "released": ("good", object()), + } + attempts = {"bad": 0} + + def free(buffer): + if buffer == "bad" and attempts["bad"] == 0: + attempts["bad"] += 1 + raise RuntimeError("destroy failed") + return 16 + + release = _extract_function( + _ROOT / "tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py", + "release_symm_buffer_cache", + { + "_MEGA_MOE_SYMM_BUFFER_CACHE": cache, + "_free_symm_buffer": free, + "logger": SimpleNamespace( + error=lambda *args, **kwargs: None, info=lambda *args, **kwargs: None + ), + }, + ) + + with pytest.raises(RuntimeError, match="failed to release 1"): + release() + assert list(cache) == ["failed"] + + release() + assert cache == {} + + +def test_mpi_worker_retries_symm_release_before_destroying_process_group(): + events = [] + module_name = "test_retryable_megamoe_symm_release" + release_attempts = 0 + + def release_symm_buffer_cache(): + nonlocal release_attempts + release_attempts += 1 + events.append("release_symm") + if release_attempts == 1: + raise RuntimeError("symm release failed") + + mega_moe = types.ModuleType(module_name) + mega_moe.release_symm_buffer_cache = release_symm_buffer_cache + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + + previous = {name: sys.modules.get(name) for name in (module_name, "torch", "torch.distributed")} + sys.modules[module_name] = mega_moe + sys.modules["torch"] = fake_torch + sys.modules["torch.distributed"] = fake_distributed + try: + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": module_name, + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = engine + worker.llm_args = _llm_args(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="symm release failed"): + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + "release_symm", + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + "release_symm", + ("terminal_cleanup", True), + "release_symm", + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert worker.engine is None + finally: + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_mpi_worker_peer_cleanup_failure_blocks_later_collectives(): + events = [] + outcomes = [ + [None, "RuntimeError: peer cleanup failed"], + [None, None], + [None, None], + ] + + class Comm: + def allgather(self, value): + events.append(("allgather", value)) + return outcomes.pop(0) + + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + + previous = {name: sys.modules.get(name) for name in ("torch", "torch.distributed")} + sys.modules["torch"] = fake_torch + sys.modules["torch.distributed"] = fake_distributed + try: + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: Comm(), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_peer_failure_megamoe", + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + terminal_cleanup=lambda *, local_safe: events.append(("terminal_cleanup", local_safe)), + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = engine + worker.llm_args = _llm_args(backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + with pytest.raises(RuntimeError, match="rank 1.*peer cleanup failed"): + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + ("allgather", None), + ] + assert worker.engine is engine + assert not worker.doing_shutdown + + worker.shutdown() + assert events == [ + "engine_shutdown", + ("terminal_cleanup", True), + ("allgather", None), + ("terminal_cleanup", True), + ("allgather", None), + ("allgather", None), + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert outcomes == [] + assert worker.engine is None + finally: + for name, module in previous.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +def test_mpi_worker_background_error_keeps_shutdown_committed(monkeypatch): + cleanups = [] + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: None + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: None) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "mpi_rank": lambda: 0, + "mpi_comm": lambda: SimpleNamespace(allgather=lambda value: [value]), + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": "missing_background_error_megamoe", + "gc": SimpleNamespace(collect=lambda: None), + }, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = SimpleNamespace( + shutdown_all_ranks=True, + can_enqueue_requests=lambda: False, + shutdown=lambda: None, + terminal_cleanup=lambda *, local_safe: cleanups.append(local_safe), + ) + worker.llm_args = _llm_args(backend="pytorch") + + def background_error(): + raise RuntimeError("background") + + worker._handle_background_error = background_error + + with pytest.raises(RuntimeError, match="background"): + worker.shutdown() + # Peers have committed teardown; a retry must not re-enter its collectives. + assert worker.doing_shutdown and worker.engine is None + worker.shutdown() + assert cleanups == [True] + + +def test_mpi_worker_without_per_layer_eplb_keeps_best_effort_teardown(monkeypatch): + events = [] + module_name = "test_best_effort_megamoe_symm_release" + + def release_symm_buffer_cache(): + events.append("release_symm") + raise RuntimeError("symm release failed") + + def no_collectives(): + raise AssertionError("rank-local teardown must not enter MPI collectives") + + mega_moe = types.ModuleType(module_name) + mega_moe.release_symm_buffer_cache = release_symm_buffer_cache + fake_distributed = types.ModuleType("torch.distributed") + fake_distributed.is_initialized = lambda: True + fake_distributed.destroy_process_group = lambda: events.append("destroy_pg") + fake_torch = types.ModuleType("torch") + fake_torch.distributed = fake_distributed + fake_torch.cuda = SimpleNamespace(empty_cache=lambda: events.append("empty_cache")) + monkeypatch.setitem(sys.modules, module_name, mega_moe) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "torch.distributed", fake_distributed) + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/worker.py", + "GenerationExecutorWorker", + {"shutdown", "_shutdown_rank_local"}, + { + "uses_collective_shutdown": _uses_collective_shutdown, + "logger_debug": lambda *args, **kwargs: None, + "logger": SimpleNamespace(error=lambda *args, **kwargs: events.append("logged")), + "mpi_rank": lambda: 0, + "mpi_comm": no_collectives, + "sys": sys, + "_MEGA_MOE_DEEPGEMM_MODULE": module_name, + "gc": SimpleNamespace(collect=lambda: events.append("gc")), + }, + ) + worker = worker_cls() + worker.doing_shutdown = False + worker.rpc_server = None + worker.engine = SimpleNamespace( + can_enqueue_requests=lambda: False, + shutdown=lambda: events.append("engine_shutdown"), + local_cleanup=lambda: events.append("local_cleanup"), + terminal_cleanup=lambda **kwargs: events.append("terminal_cleanup"), + ) + worker.llm_args = _llm_args(per_layer_dynamic=False, backend="pytorch") + worker._handle_background_error = lambda: events.append("background") + + worker.shutdown() + + assert events == [ + "engine_shutdown", + "local_cleanup", + "release_symm", + "logged", + "destroy_pg", + "gc", + "empty_cache", + "background", + ] + assert worker.doing_shutdown and worker.engine is None + + +def test_finalizer_does_not_retry_a_failed_collective_shutdown(): + calls = [] + worker_cls = _extract_class( + _ROOT / "tensorrt_llm/executor/base_worker.py", "BaseWorker", {"__del__"}, {} + ) + worker = worker_cls() + worker.shutdown = lambda: calls.append("shutdown") + + worker._collective_shutdown_failed = True + worker.__del__() + assert calls == [] + + worker._collective_shutdown_failed = False + worker.__del__() + assert calls == ["shutdown"] + worker.shutdown = lambda: None diff --git a/tests/unittest/llmapi/test_gms_args.py b/tests/unittest/llmapi/test_gms_args.py index d2f3e162b12e..dd8cf5f51a9a 100644 --- a/tests/unittest/llmapi/test_gms_args.py +++ b/tests/unittest/llmapi/test_gms_args.py @@ -127,6 +127,51 @@ def test_non_default_config_warns_when_gms_inactive( assert relevant == [] +class TestGmsMoeCompatibility: + @pytest.mark.parametrize( + "load_balancer", + [ + pytest.param({"num_slots": 64}, id="standard"), + pytest.param( + {"mode": "per_layer_dynamic", "num_slots": 72}, + id="per-layer-dynamic", + ), + ], + ) + def test_rejects_gms_with_load_balancer(self, load_balancer): + with pytest.raises( + ValueError, + match=r"LoadFormat\.GMS.*moe_config\.load_balancer", + ): + _make_args( + load_format=LoadFormat.GMS, + moe_config={"load_balancer": load_balancer}, + ) + + @pytest.mark.parametrize( + "load_format, load_balancer", + [ + pytest.param(LoadFormat.GMS, None, id="gms-without-eplb"), + pytest.param( + LoadFormat.AUTO, + {"num_slots": 64}, + id="standard-without-gms", + ), + pytest.param( + LoadFormat.AUTO, + {"mode": "per_layer_dynamic", "num_slots": 72}, + id="per-layer-dynamic-without-gms", + ), + ], + ) + def test_accepts_compatible_load_balancer(self, load_format, load_balancer): + args = _make_args( + load_format=load_format, + moe_config={"load_balancer": load_balancer}, + ) + assert args.load_format == load_format + + class TestLoadFormatGms: def test_gms_enum_present(self): assert LoadFormat.GMS.name == "GMS" diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..d668adcfd1c7 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -55,8 +55,9 @@ ExecutorMemoryType, ExtendedRuntimePerfKnobConfig, KvCacheConfig, MambaStateConfig, - MoeConfig, MTPDecodingConfig, - MultimodalConfig, + MoeConfig, MoeLoadBalancerConfig, + MoePerLayerDynamicLoadBalancerConfig, + MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, NGramDecodingConfig, PeftCacheConfig, PrefillCudaGraphBackend, PybindMirror, @@ -251,6 +252,230 @@ def test_ignores_override_without_neural_drafter(self) -> None: assert llm_args.speculative_config.moe_backend == "CUTLASS" +@pytest.mark.cpu_only +class TestMoeConfigPerLayerDynamicEplb: + + @staticmethod + def _per_layer_dynamic() -> dict[str, object]: + return {"mode": "per_layer_dynamic", "num_slots": 72} + + @pytest.mark.parametrize("backend", ["CUTLASS", "DEEPGEMM"]) + def test_rejects_explicit_incompatible_backend(self, backend: str) -> None: + with pytest.raises(ValidationError, + match=r"EPLB requires backend='MEGAMOE_CUTEDSL'"): + MoeConfig(backend=backend, load_balancer=self._per_layer_dynamic()) + + @pytest.mark.parametrize("backend", [None, "AUTO", "MEGAMOE_CUTEDSL"]) + def test_resolves_auto_and_accepts_compatible_backend( + self, backend: str | None) -> None: + kwargs = {} if backend is None else {"backend": backend} + + config = MoeConfig(load_balancer=self._per_layer_dynamic(), **kwargs) + + assert isinstance(config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert config.backend == "MEGAMOE_CUTEDSL" + assert config.load_balancer.auxiliary_sms == 4 + assert config.load_balancer.helper_slot_mode == "per_layer_resident" + + @pytest.mark.parametrize( + "num_experts, message", + [ + pytest.param(0, "must be positive", id="zero-experts"), + pytest.param(65, "must be divisible", id="non-divisible-experts"), + pytest.param(72, "must exceed", id="no-helper-capacity"), + pytest.param(80, "must exceed", id="negative-helper-capacity"), + ], + ) + def test_validates_physical_slot_capacity_after_setup( + self, num_experts: int, message: str) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + config.setup(ep_rank=0, ep_size=8) + + with pytest.raises(ValueError, match=message): + config.validate_expert_capacity(num_experts) + + def test_accepts_positive_helper_capacity_after_setup(self) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + config.setup(ep_rank=0, ep_size=8) + + config.validate_expert_capacity(num_experts=64) + assert config.num_local_slots == 9 + + def test_expert_capacity_requires_setup(self) -> None: + config = MoePerLayerDynamicLoadBalancerConfig(num_slots=72) + with pytest.raises(ValueError, match=r"Call setup\(\)"): + config.validate_expert_capacity(num_experts=64) + + def test_legacy_dictionary_defaults_to_standard_mode(self) -> None: + config = MoeConfig( + backend="CUTLASS", + load_balancer={ + "num_slots": 64, + "layer_updates_per_iter": 2 + }, + ) + + assert isinstance(config.load_balancer, MoeLoadBalancerConfig) + assert config.load_balancer.mode == "standard" + assert config.load_balancer.layer_updates_per_iter == 2 + assert config.backend == "CUTLASS" + + @pytest.mark.parametrize("field", ["num_slots", "auxiliary_sms"]) + def test_per_layer_dynamic_positive_fields(self, field: str) -> None: + values = self._per_layer_dynamic() + values[field] = 0 + with pytest.raises(ValidationError, match=field): + MoeConfig(load_balancer=values) + + @pytest.mark.parametrize("auxiliary_sms", [1, 8]) + def test_per_layer_dynamic_accepts_other_auxiliary_sm_budgets( + self, auxiliary_sms: int) -> None: + values = self._per_layer_dynamic() + values["auxiliary_sms"] = auxiliary_sms + config = MoeConfig(load_balancer=values) + assert config.load_balancer.auxiliary_sms == auxiliary_sms + + def test_per_layer_dynamic_rejects_more_than_fused_copy_sm_limit( + self) -> None: + values = self._per_layer_dynamic() + values["auxiliary_sms"] = 9 + with pytest.raises(ValidationError, match="auxiliary_sms"): + MoeConfig(load_balancer=values) + + @pytest.mark.parametrize("mode", ["per_layer_resident", "shared_banks"]) + def test_per_layer_dynamic_accepts_helper_slot_modes(self, + mode: str) -> None: + values = self._per_layer_dynamic() + values["helper_slot_mode"] = mode + config = MoeConfig(load_balancer=values) + assert config.load_balancer.helper_slot_mode == mode + + def test_per_layer_dynamic_rejects_unknown_helper_slot_mode(self) -> None: + values = self._per_layer_dynamic() + values["helper_slot_mode"] = "per_layer" + with pytest.raises(ValidationError, match="helper_slot_mode"): + MoeConfig(load_balancer=values) + + def test_standard_mode_rejects_helper_slot_mode(self) -> None: + with pytest.raises(ValidationError, match="helper_slot_mode"): + MoeConfig(backend="CUTLASS", + load_balancer={ + "mode": "standard", + "num_slots": 64, + "helper_slot_mode": "shared_banks", + }) + + def test_per_layer_dynamic_disables_decode_cuda_graph(self) -> None: + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": self._per_layer_dynamic()}) + assert args.cuda_graph_config is None + + def test_per_layer_dynamic_rejects_prefill_cuda_graph(self) -> None: + with pytest.raises(ValidationError, match="prefill CUDA graphs"): + TorchLlmArgs( + model=llama_model_path, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE, + moe_config={"load_balancer": self._per_layer_dynamic()}) + + def test_legacy_yaml_defaults_to_standard_mode(self, + tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("num_slots: 64\nlayer_updates_per_iter: 2\n") + + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": str(config_path)}, + ) + + assert isinstance(args.moe_config.load_balancer, MoeLoadBalancerConfig) + assert args.moe_config.load_balancer.mode == "standard" + + def test_per_layer_dynamic_yaml_sets_backend(self, tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + + args = TorchLlmArgs( + model=llama_model_path, + moe_config={"load_balancer": str(config_path)}, + ) + + assert isinstance(args.moe_config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert args.moe_config.backend == "MEGAMOE_CUTEDSL" + + def test_per_layer_dynamic_yaml_is_typed_by_moe_config_before_llm_args( + self, tmp_path: Path) -> None: + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + + config = MoeConfig(load_balancer=str(config_path)) + + assert isinstance(config.load_balancer, + MoePerLayerDynamicLoadBalancerConfig) + assert config.backend == "MEGAMOE_CUTEDSL" + + def test_per_layer_dynamic_yaml_configures_queue_before_cuda_init( + self, tmp_path: Path, monkeypatch) -> None: + from tensorrt_llm.llmapi._load_balance_env import \ + configure_moe_launch_queues + + config_path = tmp_path / "eplb.yaml" + config_path.write_text("mode: per_layer_dynamic\nnum_slots: 72\n") + # The code under test writes os.environ. Seed the key first so that + # monkeypatch records the caller's state and removes "4x" on teardown; + # delenv alone records nothing when the key is absent. + monkeypatch.setenv("CUDA_SCALE_LAUNCH_QUEUES", "") + monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES") + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False) + + config = MoeConfig(load_balancer=str(config_path)) + overrides = configure_moe_launch_queues(config, {"KEEP": "value"}) + + assert overrides == { + "KEEP": "value", + "CUDA_SCALE_LAUNCH_QUEUES": "4x", + } + assert config.load_balancer.mode == "per_layer_dynamic" + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + assert configure_moe_launch_queues(config, overrides) == overrides + + def test_standard_yaml_does_not_configure_launch_queues( + self, tmp_path: Path, monkeypatch) -> None: + from tensorrt_llm.llmapi._load_balance_env import \ + configure_moe_launch_queues + + config_path = tmp_path / "eplb.yaml" + config_path.write_text("num_slots: 64\nlayer_updates_per_iter: 2\n") + monkeypatch.delenv("CUDA_SCALE_LAUNCH_QUEUES", raising=False) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + + config = MoeConfig(load_balancer=str(config_path)) + original = {"KEEP": "value"} + + assert configure_moe_launch_queues(config, original) is original + assert "CUDA_SCALE_LAUNCH_QUEUES" not in original + + def test_yaml_resolution_preserves_missing_file_error( + self, tmp_path: Path) -> None: + config_path = tmp_path / "missing.yaml" + + with pytest.raises(FileNotFoundError, + match=r"MoE load balancer config file not found:"): + MoeConfig(load_balancer=str(config_path)) + + def test_yaml_resolution_preserves_invalid_file_error( + self, tmp_path: Path) -> None: + config_path = tmp_path / "invalid.yaml" + config_path.write_text("mode: per_layer_dynamic\n") + + with pytest.raises( + ValidationError, + match=r"Failed to load MoE load balancer config file:"): + MoeConfig(load_balancer=str(config_path)) + + @pytest.mark.cpu_only def test_rejection_sampling_allows_attention_dp(monkeypatch): """ADP (incl. ADP+LM-head-TP) supports rejection sampling. @@ -4144,7 +4369,6 @@ class TestPydanticBestPractices: "drafter", # abstract base class type "resource_manager", # abstract base class type ], - MoeConfig: ["load_balancer"], # allows multiple types including dict RayPlacementConfig: ["placement_groups"], # contains Ray-specific types } diff --git a/tests/unittest/llmapi/test_serialization.py b/tests/unittest/llmapi/test_serialization.py index cbf59b545d19..f8a637fab6eb 100644 --- a/tests/unittest/llmapi/test_serialization.py +++ b/tests/unittest/llmapi/test_serialization.py @@ -24,6 +24,33 @@ def test_serialization_allowed_class(): assert type(obj) == type(b) and obj.name == b.name +@pytest.mark.parametrize( + "config", + [ + pytest.param( + __import__("tensorrt_llm.llmapi.llm_args", + fromlist=["MoeLoadBalancerConfig" + ]).MoeLoadBalancerConfig(num_slots=8), + id="standard", + ), + pytest.param( + __import__( + "tensorrt_llm.llmapi.llm_args", + fromlist=["MoePerLayerDynamicLoadBalancerConfig"], + ).MoePerLayerDynamicLoadBalancerConfig(num_slots=16), + id="per-layer-dynamic", + ), + ], +) +def test_moe_load_balancer_configs_use_restricted_serialization(config): + restored = serialization.loads( + serialization.dumps(config), + approved_imports=serialization.BASE_EXAMPLE_CLASSES, + ) + assert type(restored) is type(config) + assert restored == config + + def test_serialization_disallowed_class(): obj = TestClass("test") a = serialization.dumps(obj) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 38cf3edc82ba..39420bb22730 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -60,6 +60,7 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.SERVE_SCHEMA, {"moe_config": {"backend": "CUTLASS"}}, True), (generator.SERVE_SCHEMA, {"moe_config": {"backend_typo": "CUTLASS"}}, False), (generator.SERVE_SCHEMA, {"moe_config": {"backend": "not-a-backend"}}, False), + (generator.SERVE_SCHEMA, {"moe_config": {"load_balancer": "eplb.yaml"}}, True), (generator.SERVE_SCHEMA, {"backend": "pytorch"}, True), (generator.SERVE_SCHEMA, {"backend": "_autodeploy"}, False), ( @@ -200,6 +201,40 @@ def test_branch_snapshot_field_types( assert (not errors) == valid, [error.message for error in errors] +@pytest.mark.parametrize( + ("load_balancer", "valid"), + [ + ({"num_slots": 288}, True), + ({"mode": "standard", "num_slots": 288, "layer_updates_per_iter": 1}, True), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 4}, True), + ( + {"mode": "per_layer_dynamic", "num_slots": 320, "helper_slot_mode": "shared_banks"}, + True, + ), + # A mapping without mode is standard EPLB, which has no auxiliary_sms. + ({"num_slots": 320, "auxiliary_sms": 4}, False), + ({"mode": "per_layer_dynamic"}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "layer_updates_per_iter": 1}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "auxiliary_sms": 9}, False), + ({"mode": "per_layer_dynamic", "num_slots": 320, "helper_slot_mode": "private"}, False), + ({"mode": "dynamic", "num_slots": 320}, False), + ], +) +def test_load_balancer_schema_matches_runtime( + schemas: dict[str, dict], load_balancer: dict, valid: bool +) -> None: + from tensorrt_llm.llmapi.llm_args import MoeConfig + + config = {"moe_config": {"load_balancer": load_balancer}} + errors = list(Draft202012Validator(schemas[generator.SERVE_SCHEMA]).iter_errors(config)) + assert (not errors) == valid, [error.message for error in errors] + if valid: + MoeConfig(load_balancer=load_balancer) + else: + with pytest.raises(ValueError): + MoeConfig(load_balancer=load_balancer) + + def _assert_serving_fields_in_schema(source: str, function_name: str, properties: dict) -> None: """Check literal top-level config accesses, without tracing aliases or dynamic keys.""" function = next( diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 13af3416689e..d9dbc29f3043 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -884,15 +884,18 @@ class _RawExcludeConfig(StrictBaseModel): assert config == {"kept": 1} -def test_manifest_excludes_loosely_typed_model_children(): - # B-1 regression: moe_config.load_balancer is Optional[Union[object, str]]; - # a validator coerces it into a MoeLoadBalancerConfig at runtime, but the - # annotation names no BaseModel, so its children must NOT be capturable. +def test_manifest_includes_safe_load_balancer_fields(): from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.usage.llmapi_config import build_capture_manifest - paths = {e.path for e in build_capture_manifest(TorchLlmArgs)} - assert not any(p.startswith("moe_config.load_balancer.") for p in paths) + paths = {entry.path for entry in build_capture_manifest(TorchLlmArgs)} + assert { + "moe_config.load_balancer.auxiliary_sms", + "moe_config.load_balancer.layer_updates_per_iter", + "moe_config.load_balancer.mode", + "moe_config.load_balancer.num_slots", + } <= paths + assert "moe_config.load_balancer.initial_global_assignments" not in paths def test_collect_llm_api_config_caps_total_payload_size(monkeypatch): diff --git a/tests/unittest/utils/test_logger.py b/tests/unittest/utils/test_logger.py index 27a209e9d801..7d65a1182432 100644 --- a/tests/unittest/utils/test_logger.py +++ b/tests/unittest/utils/test_logger.py @@ -14,6 +14,7 @@ # limitations under the License. """Unit tests for the modular logger (auto module detection, per-module filtering).""" +import importlib import logging from io import StringIO @@ -239,6 +240,26 @@ def test_unmatched_module_uses_global(self): # Cleanup singleton._module_levels = {} + def test_filtered_message_skips_caller_lookup(self, monkeypatch): + # The package re-exports the Logger singleton as ``tensorrt_llm.logger``, + # shadowing the submodule, so fetch the module object explicitly. + logger_module = importlib.import_module("tensorrt_llm.logger") + + singleton = Logger() + monkeypatch.setattr(singleton, "_min_severity", "warning") + monkeypatch.setattr(singleton, "_module_levels", {}) + + def fail(): + raise AssertionError("caller module looked up for a filtered message") + + monkeypatch.setattr(logger_module, "_get_caller_module", fail) + singleton.info("filtered") + assert singleton.is_debug_enabled() is False + # A per-module override still needs the caller's module. + monkeypatch.setattr(singleton, "_module_levels", {"runtime": 10}) + with pytest.raises(AssertionError, match="caller module"): + singleton.info("module filtered") + class TestIsDebugEnabled: """Tests for is_debug_enabled, the guard for expensive debug arguments.