[None][feat] Add DynamicEPLB kernels and runtime integration

.pre-commit-config.yaml

@@ -1463,7 +1463,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:

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 $<TARGET_OBJECTS:moe_dynamic_eplb_src>)
 set(MOE_KERNELS_SRC
     ${MOE_SRC}
     PARENT_SCOPE)

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)

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 <climits>
+
+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

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 <cstdint>
+
+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