Skip to content

[pull] main from NVIDIA:main #318

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pull merged 8 commits into from
Mar 24, 2026
Merged

[pull] main from NVIDIA:main #318

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ffb1fed
[TRTLLM-11471][feat] Add safe version of allgather with chunking (#12…
chienchunhung Mar 24, 2026
c42e86e
[None][perf] add Dynamic SMEM block routing in MOE (#12456)
jiahanc Mar 24, 2026
6c0328c
[TRTLLM-11544][feat] Add Qwen 3.5 supporting(NVFP4). (#12302)
nv-guomingz Mar 24, 2026
00cd284
[https://nvbugs/5997090][fix] Add Disagg Perf Test back as MPI Issue …
chenfeiz0326 Mar 24, 2026
d6b8e6f
[https://nvbugs/5841976][fix] Remove test_fused_moe_alltoall_fp4[Deep…
xxi-nv Mar 24, 2026
6667779
[None][infra] Waive 2 failed cases for main in post-merge 2613 (#12473)
ZhanruiSunCh Mar 24, 2026
c816f8b
[None][infra] Check in most recent lock file from nightly pipeline
tensorrt-cicd Mar 24, 2026
14a75ee
[https://nvbugs/5866619][test] Add unit test for load_state_dict safe…
crazydemo Mar 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ void run(Data const& data, void* stream)
TLLM_CHECK_WITH_INFO(
data.mNumExperts % 4 == 0, "Routing kernel expects #experts %d to be a multiple of 4.", data.mNumExperts);

bool const useSingleBlock = data.mNumTokens <= BlockKernelMaxNumTokens;
bool const useSingleBlock = data.mNumTokens <= BlockKernelMaxNumTokens
|| (data.mNumTokens <= DynBlockKernelMaxNumTokens && data.mNumExperts <= DynBlockKernelMaxNumExperts);

bool const useSingleCluster = data.mNumTokens <= ((data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr)
? MaxNumTokensSingleClusterScores
Expand Down
2 changes: 2 additions & 0 deletions ...lm/kernels/trtllmGenKernels/blockScaleMoe/routingRenormalize/RoutingRenormalizeCommon.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ static constexpr int MaxNumTokensSingleCluster = NumBlocksPerCluster * NumThread
static constexpr int MaxNumTokensSingleClusterScores = NumBlocksPerCluster * NumWarps;

static constexpr int BlockKernelMaxNumTokens = 4;
static constexpr int DynBlockKernelMaxNumTokens = 16;
static constexpr int DynBlockKernelMaxNumExperts = 512;

template <typename DataType, typename InputType, int VecSize, int K, bool DoSoftmaxBeforeTopK>
__forceinline__ __device__ void routingTopKExperts(cg::thread_block_tile<WarpSize> const& warp,
Expand Down
382 changes: 379 additions & 3 deletions ...nsorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routingRenormalize/launchBlockKernel.cu
View file
Open in desktop

Large diffs are not rendered by default.

66 changes: 66 additions & 0 deletions cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,72 @@ TYPED_TEST(RoutingRenormalizeKernelTest, DeviceLevelParallelizationTop4)
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockBasic)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/8,
/*numExperts=*/128, /*topK=*/8,
/*expertParallelization=*/1, /*expertParallelizationId=*/0, /*tileTokensDim=*/256,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*hasInvalidTopKInput=*/false,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockMaxTokens)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/16,
/*numExperts=*/512, /*topK=*/8,
/*expertParallelization=*/1, /*expertParallelizationId=*/0, /*tileTokensDim=*/256,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*hasInvalidTopKInput=*/false,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockWithExpertParallelization)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/12,
/*numExperts=*/512, /*topK=*/8,
/*expertParallelization=*/2, /*expertParallelizationId=*/1, /*tileTokensDim=*/192,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*hasInvalidTopKInput=*/false,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockWithTopKAsInput)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/8,
/*numExperts=*/128, /*topK=*/8,
/*expertParallelization=*/1, /*expertParallelizationId=*/0, /*tileTokensDim=*/256,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*hasInvalidTopKInput=*/false,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockWithInvalidTopKInput)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/10,
/*numExperts=*/512, /*topK=*/8,
/*expertParallelization=*/2, /*expertParallelizationId=*/0, /*tileTokensDim=*/256,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*hasInvalidTopKInput=*/true,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, DynBlockWithRenormalizeNaive)
{
RoutingKernelTestParam param(RoutingMethodType::RenormalizeNaive, /*numTokens=*/16,
/*numExperts=*/512, /*topK=*/8,
/*expertParallelization=*/1, /*expertParallelizationId=*/0, /*tileTokensDim=*/256,
/*paddingLog2=*/3, /*localExpertsStrideLog2=*/0,
/*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*hasInvalidTopKInput=*/false,
/*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9);
this->runTest(param);
};

TYPED_TEST(RoutingRenormalizeKernelTest, BlockLevelParallelizationLargeN)
{
RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/4,
Expand Down
Loading
Loading