From e02d2ac70fc70bc8882377bb5d8750b8469d3d5d Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 31 Jul 2026 09:00:11 +0000 Subject: [PATCH 1/4] [None][fix] Harden excluded Qwen3Next MoE experts Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../_torch/models/modeling_qwen3_next.py | 49 ++++----- .../models/test_qwen3_next_moe_quant.py | 102 ++++++++++++++++++ .../mamba/test_flashinfer_gdn_verify.py | 84 +++++++++++++++ 3 files changed, 209 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 15047b48a4d5..3e1c06dceacc 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -248,11 +248,22 @@ def __init__( if quant_config_dict and layer_idx is not None: expert_quant_config = quant_config_dict.get( f"model.layers.{layer_idx}.mlp.experts") - # Excluded experts end up bf16 whatever the per-layer entry says, so - # hand create_moe a bf16 config and let it pick a backend serving bf16. + moe_model_config = model_config if _experts_excluded_from_quant(model_config, layer_idx): + # These experts end up bf16 whatever the per-layer entry says, so + # build them on CUTLASS, the only backend that serves bf16 + # unconditionally. expert_quant_config = QuantConfig(kv_cache_quant_algo=model_config. quant_config.kv_cache_quant_algo) + if model_config.moe_backend != "CUTLASS": + moe_model_config = copy.copy(model_config) + moe_model_config._frozen = False + moe_model_config.moe_backend = "CUTLASS" + moe_model_config._frozen = True + logger.warning( + f"Layer {layer_idx} MoE experts are excluded from " + "quantization; using moe_backend=CUTLASS for this layer " + f"(other layers keep {model_config.moe_backend}).") self.experts = create_moe( num_experts=self.num_experts, routing_method=self.gate.routing_method, @@ -261,7 +272,7 @@ def __init__( aux_stream_dict={AuxStreamType.MoeChunkingOverlap: aux_stream}, dtype=config.torch_dtype, reduce_results=False, - model_config=model_config, + model_config=moe_model_config, layer_idx=layer_idx, weight_loading_mode=weight_loading_mode, override_quant_config=expert_quant_config, @@ -807,25 +818,11 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream]): # Some HF checkpoints (e.g. Qwen3.5 NVFP4) keep the whole MTP layer in - # bf16. Given a bf16 config most backends switch to CUTLASS on their - # own; DEEPGEMM and WIDEEP do not, so switch for them here. - mtp_model_config = model_config - if (model_config.moe_backend != "CUTLASS" - and _experts_excluded_from_quant(model_config, layer_idx)): - original_backend = model_config.moe_backend - mtp_model_config = copy.copy(model_config) - mtp_model_config._frozen = False - mtp_model_config.moe_backend = "CUTLASS" - mtp_model_config._frozen = True - logger.warning( - "Qwen3Next MTP layer is unquantized in the checkpoint; " - "falling back to moe_backend=CUTLASS for the MTP layer " - f"(regular layers keep moe_backend={original_backend}).") - - super().__init__(mtp_model_config, layer_idx, + # bf16. Qwen3NextSparseMoeBlock handles that for every layer, MTP + # included, so nothing layer-specific is needed here. + super().__init__(model_config, layer_idx, aux_stream_dict[AuxStreamType.Attention]) - config = mtp_model_config.pretrained_config - self.model_config = mtp_model_config + config = model_config.pretrained_config self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] self.event_dict = { key: torch.cuda.Event() @@ -845,13 +842,13 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig], use_gemma=True, ) - if mtp_model_config.mapping.enable_attention_dp: + if model_config.mapping.enable_attention_dp: self.fc = Linear( config.hidden_size * 2, config.hidden_size, bias=False, dtype=config.torch_dtype, - skip_create_weights_in_init=mtp_model_config. + skip_create_weights_in_init=model_config. skip_create_weights_in_init, use_cute_dsl_blockscaling_mm=False, ) @@ -862,13 +859,13 @@ def __init__(self, model_config: ModelConfig[Qwen3NextConfig], bias=False, dtype=config.torch_dtype, tensor_parallel_mode=TensorParallelMode.ROW, - mapping=mtp_model_config.mapping, + mapping=model_config.mapping, reduce_output=True, - skip_create_weights_in_init=mtp_model_config. + skip_create_weights_in_init=model_config. skip_create_weights_in_init, use_cute_dsl_blockscaling_mm=False, ) - self.shared_head = Qwen3NextMTPHead(mtp_model_config) + self.shared_head = Qwen3NextMTPHead(model_config) # MTP applies shared_head.norm after the base decoder forward, so its # MoE-output all-reduce cannot consume next_layer_layernorm. self.fusion_config.POST_MOE_FUSION = False diff --git a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py index fb4e190f1dd0..d28b6403291c 100644 --- a/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py +++ b/tests/unittest/_torch/models/test_qwen3_next_moe_quant.py @@ -15,8 +15,10 @@ """Unquantized-MoE probe for the Qwen3Next / Qwen3.5 MTP layer.""" from types import SimpleNamespace +from unittest.mock import patch import pytest +import torch from tensorrt_llm._torch.models.modeling_qwen3_5 import _normalize_qwen35_exclude_modules from tensorrt_llm._torch.models.modeling_qwen3_next import _experts_excluded_from_quant @@ -163,3 +165,103 @@ def test_regular_layers_are_covered_too(pattern, layer_idx, expected): def test_missing_layer_idx_is_a_noop(): assert not _experts_excluded_from_quant(_model_config(["model.layers.5*"]), None) + + +class _StopBlockInit(Exception): + """Raised after the test captures the arguments passed to ``create_moe``.""" + + +def _build_moe_block(moe_backend, exclude_modules, layer_idx, quant_config_dict=None): + """Run sparse-MoE initialization through its ``create_moe`` call.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextSparseMoeBlock + + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=32, + intermediate_size=64, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + num_experts=4, + num_experts_per_tok=2, + num_hidden_layers=NUM_HIDDEN_LAYERS, + num_nextn_predict_layers=1, + torch_dtype=torch.bfloat16, + model_type="qwen3_next", + mlp_bias=False, + ), + moe_backend=moe_backend, + quant_config=QuantConfig( + quant_algo=QuantAlgo.NVFP4, + kv_cache_quant_algo=QuantAlgo.FP8, + exclude_modules=exclude_modules, + ), + quant_config_dict=quant_config_dict, + ) + captured = {} + + def _capture(*args, **kwargs): + from tensorrt_llm._torch.modules.fused_moe.create_moe import resolve_moe_cls + + captured["moe_backend"] = kwargs["model_config"].moe_backend + captured["override"] = kwargs["override_quant_config"] + captured["moe_cls"] = resolve_moe_cls( + kwargs["model_config"], + kwargs["routing_method"], + kwargs["dtype"], + kwargs["override_quant_config"], + kwargs["layer_idx"], + ).__name__ + raise _StopBlockInit + + import tensorrt_llm._torch.models.modeling_qwen3_next as qwen3_next + + with patch.object(qwen3_next, "create_moe", _capture), pytest.raises(_StopBlockInit): + Qwen3NextSparseMoeBlock(model_config, aux_stream=None, layer_idx=layer_idx) + return captured + + +@pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "DEEPGEMM", "WIDEEP", "CUTEDSL"]) +@pytest.mark.parametrize("layer_idx", [5, MTP_LAYER_IDX]) +def test_excluded_layer_builds_bf16_on_cutlass(backend, layer_idx): + per_layer_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + captured = _build_moe_block( + backend, + [f"model.layers.{layer_idx}*"], + layer_idx, + quant_config_dict={ + f"model.layers.{layer_idx}.mlp.experts": per_layer_quant_config, + }, + ) + + assert captured["moe_backend"] == "CUTLASS" + assert captured["moe_cls"] == "CutlassFusedMoE" + assert captured["override"] is not per_layer_quant_config + assert not captured["override"].layer_quant_mode.has_any_quant(exclude_kv_cache=True) + assert captured["override"].kv_cache_quant_algo == QuantAlgo.FP8 + + +_UNEXCLUDED_EXPECTED_MOE_CLS = { + "CUTLASS": "CutlassFusedMoE", + "TRTLLM": "TRTLLMGenFusedMoE", + "DEEPGEMM": "DeepGemmFusedMoE", + "WIDEEP": "WideEPMoE", + "CUTEDSL": "CuteDslFusedMoE", +} + + +@pytest.mark.parametrize("backend", sorted(_UNEXCLUDED_EXPECTED_MOE_CLS)) +def test_unexcluded_layer_keeps_configured_backend_and_layer_quant_config(backend): + per_layer_quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + captured = _build_moe_block( + backend, + ["model.layers.7*"], + 5, + quant_config_dict={ + "model.layers.5.mlp.experts": per_layer_quant_config, + }, + ) + + assert captured["moe_backend"] == backend + assert captured["moe_cls"] == _UNEXCLUDED_EXPECTED_MOE_CLS[backend] + assert captured["override"] is per_layer_quant_config diff --git a/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py b/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py index b7b30c6abcac..6cf6db5a83b6 100644 --- a/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py +++ b/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py @@ -230,6 +230,90 @@ def test_fi_mtp_verify_misaligned_index_slice(): torch.testing.assert_close(out_mis.float(), out_ref.float()) +@skip_unsupported +def test_fi_mtp_verify_misaligned_ab_slices(): + """Non-32B-aligned ``a``/``b`` column slices must be realigned. + + ``gdn_mixer._compute_tokenwise_inputs`` splits the fused ``in_proj_ba`` + projection into ``b = projected_states_ba[:, :num_v_heads_per_tp]`` and + ``a = projected_states_ba[:, num_v_heads_per_tp:]``. ``a`` therefore starts + ``num_v_heads_per_tp`` elements in, and that byte offset is not a multiple + of 32 for many TP splits (in bf16, whenever the per-rank v-head count is not + a multiple of 16). The FI kernel then asserts ``Misaligned Tensor data on + argument``. Attention-DP hides this because v-heads are not sharded there. + + The gating inputs are fp32, matching the other tests in this file. ``HV=4`` + gives the requested 16-byte offset. Two fused buffers exercise each guard: + the ordinary split misaligns ``a``; shifting the fused buffer misaligns + ``b`` instead. + + ``.contiguous()`` is not a fix: at ``draft_token_num == 1`` the token dim is + size 1, so the offset view already reports as contiguous. + """ + from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( + _flashinfer_gdn_verify, + ) + + torch.manual_seed(0) + dev = "cuda" + N, T, H, HV, K, V = 2, 3, 4, 4, 128, 128 + q = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + k = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + v = (torch.randn(N, T, HV, V, device=dev) * 0.1).to(torch.bfloat16) + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + state_pool = (torch.randn(N, HV, V, K, device=dev) * 0.1).to(torch.bfloat16) + idx = torch.arange(N, device=dev, dtype=torch.int32) + + # Same column split gdn_mixer performs on the fused in_proj_ba output: + # the base is aligned, while a starts 16 bytes into each row. + ba = torch.randn(N * T, 2 * HV, device=dev) * 0.1 + b_aligned = ba[:, :HV].view(N, T, HV) + a_misaligned = ba[:, HV:].view(N, T, HV) + assert b_aligned.data_ptr() % 32 == 0 + assert a_misaligned.data_ptr() % 32 == 16 + + # Shift a second fused buffer by 16 bytes so b is misaligned and a is + # aligned after the additional 16-byte column offset. + shifted_storage = torch.randn(N * T * 2 * HV + HV, device=dev) * 0.1 + shifted_ba = shifted_storage[HV:].view(N * T, 2 * HV) + b_misaligned = shifted_ba[:, :HV].view(N, T, HV) + a_aligned = shifted_ba[:, HV:].view(N, T, HV) + assert b_misaligned.data_ptr() % 32 == 16 + assert a_aligned.data_ptr() % 32 == 0 + + def _run(a_in, b_in): + return _flashinfer_gdn_verify( + A_log=A_log, + a=a_in, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b_in, + initial_state_source=state_pool, + initial_state_indices=idx, + intermediate_states_buffer=torch.zeros( + N, T, HV, V, K, device=dev, dtype=torch.bfloat16 + ), + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + ) + + for a, b in [ + (a_misaligned, b_aligned), + (a_aligned, b_misaligned), + ]: + out_misaligned = _run(a, b) + out_aligned = _run( + a.clone(memory_format=torch.contiguous_format), + b.clone(memory_format=torch.contiguous_format), + ) + torch.testing.assert_close(out_misaligned.float(), out_aligned.float()) + + @skip_unsupported def test_fi_verify_gate_env_killswitch(monkeypatch): """The dispatch gate honors the disable env vars and shape constraints.""" From 3c9e940cdca8a17fac6dc2c68c06a11732f3066f Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 31 Jul 2026 09:08:16 +0000 Subject: [PATCH 2/4] [None][perf] Optimize DeepSeek FP8 MoE activation Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../blockScaleMoe/DevKernel.cu | 185 +++++++- .../blockScaleMoe/DevKernel.h | 10 +- .../trtllmGenKernels/blockScaleMoe/runner.cu | 5 +- .../trtllmGenKernels/blockScaleMoe/runner.h | 5 +- cpp/tests/unit_tests/kernels/CMakeLists.txt | 2 + .../kernels/blockScaleMoeActivationTest.cu | 438 ++++++++++++++++++ 6 files changed, 639 insertions(+), 6 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu index 22f1d38b83f7..11450f738934 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -22,6 +22,9 @@ #include #include +#include +#include + //////////////////////////////////////////////////////////////////////////////////////////////////// // Helper function for array conversion @@ -231,6 +234,145 @@ struct KernelTraits<1> constexpr int DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA = 128; +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Permuted-space SwiGLU for the DeepSeek-FP8 separate-activation path. +// +// `activationDeepSeekKernel` below grids over the *expanded* index space +// (numTokens x topK) and discovers work by loading expandedIdxToPermutedIdx, +// skipping entries that map to -1. Under expert parallelism only 1/ep_size of +// those entries are local, so most of the launched CTAs do no memory work at +// all -- yet they still run the unconditional cub::BlockReduce. At a large +// context and a high expert-parallel degree the launched CTA count exceeds the +// permuted rows of real work by the ep_size factor, and the achieved bandwidth +// is a small fraction of what the row count alone would need. +// +// Every memory access in that kernel is addressed by (permutedIdx, hiddenIdx) +// only -- the expanded index exists purely to find the work. So grid directly +// over the permuted rows instead and the indirection, the -1 slots and the +// ep_size-fold CTA inflation all disappear together. +// +// Layout: one warp owns exactly one (permutedRow, 128-element scale block). +// 32 lanes x 4 elements = 128 = one scale block, so the amax reduction is a +// single warp shuffle instead of a shared-memory block reduce, and each lane +// moves 4 bytes per load instead of 1. +// +// totalNumPaddedTokens is only known on the device, so the grid is persistent +// and strides over the row space. This visits the per-expert tile padding that +// the expanded-space kernel skips (~4% extra rows at 32 local experts); those +// rows are dropped by the finalize kernel. The arithmetic below deliberately +// preserves the legacy kernel's 0/0 -> NaN behavior for an all-zero block. +constexpr int kDsActWarpSize = 32; +constexpr int kDsActEltsPerSf = 128; +constexpr int kDsActEltsPerThread = kDsActEltsPerSf / kDsActWarpSize; +constexpr int kDsActWarpsPerCta = 4; +constexpr int kDsActPermutedNumThreadsPerCta = kDsActWarpSize * kDsActWarpsPerCta; + +constexpr bool shouldUsePermutedActivation(int innerDim, int numTokens, int topK, int numExperts, int tileTokensDim) +{ + int const outputDim = innerDim / 2; + bool const layoutEligible = outputDim >= kDsActEltsPerSf && outputDim % kDsActEltsPerSf == 0 && innerDim % 8 == 0; + int64_t const realRowsPerExpert = numExperts > 0 ? static_cast(numTokens) * topK / numExperts : 0; + bool const paddingAmortised = tileTokensDim > 0 && realRowsPerExpert >= tileTokensDim; + return layoutEligible && paddingAmortised; +} + +template +__global__ void activationDeepSeekPermutedKernel(KernelParams params) +{ + using Type = typename KernelParams::Type; + using PackedIo = uint32_t; // kDsActEltsPerThread x 8-bit elements + + static_assert(kDsActEltsPerThread == 4, "PackedIo assumes 4 elements per thread"); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (KernelParams::UsePdl) + { + cudaTriggerProgrammaticLaunchCompletion(); + cudaGridDependencySynchronize(); + } +#endif + + float constexpr kE4m3MaxVal{448.F}; + + int const totalNumPaddedTokens = params.totalNumPaddedTokens[0]; + int const outputDim = params.innerDim / 2; + int const numSfBlocks = outputDim / kDsActEltsPerSf; + + bool const hasSwigluLimit = params.hasSwigluLimit; + float const swigluLimit = params.swigluLimit; + + int const lane = threadIdx.x % kDsActWarpSize; + int const warpInCta = threadIdx.x / kDsActWarpSize; + + int64_t const numTasks = static_cast(totalNumPaddedTokens) * numSfBlocks; + int64_t const taskStride = static_cast(gridDim.x) * kDsActWarpsPerCta; + + for (int64_t task = static_cast(blockIdx.x) * kDsActWarpsPerCta + warpInCta; task < numTasks; + task += taskStride) + { + int const permutedIdx = static_cast(task / numSfBlocks); + int const sfBlock = static_cast(task % numSfBlocks); + int const hiddenBase = sfBlock * kDsActEltsPerSf + lane * kDsActEltsPerThread; + + // Both scales are uniform across the warp: one per (row, scale block). + float const scale1 = params.inDqSfsPtr[permutedIdx + totalNumPaddedTokens * sfBlock]; + float const scale2 = params.inDqSfsPtr[permutedIdx + totalNumPaddedTokens * (sfBlock + numSfBlocks)]; + + int64_t const baseIdx = static_cast(permutedIdx) * params.innerDim + hiddenBase; + PackedIo const packed1 = *reinterpret_cast(params.inPtr + baseIdx); + PackedIo const packed2 = *reinterpret_cast(params.inPtr + baseIdx + outputDim); + + Type const* elts1 = reinterpret_cast(&packed1); + Type const* elts2 = reinterpret_cast(&packed2); + + float out[kDsActEltsPerThread]; + float aMax = 0.F; +#pragma unroll + for (int i = 0; i < kDsActEltsPerThread; ++i) + { + float x1 = scale1 * static_cast(elts1[i]); // up (linear) + float x2 = scale2 * static_cast(elts2[i]); // gate (silu input) + if (hasSwigluLimit) + { + x2 = fminf(x2, swigluLimit); + x1 = fmaxf(fminf(x1, swigluLimit), -swigluLimit); + } + out[i] = silu(x2) * x1; + aMax = fmaxf(aMax, fabsf(out[i])); + } + +#pragma unroll + for (int offset = kDsActWarpSize / 2; offset > 0; offset >>= 1) + { + aMax = fmaxf(aMax, __shfl_xor_sync(0xffffffffu, aMax, offset)); + } + + float const scaleOut = aMax / kE4m3MaxVal; + + if (lane == 0) + { + params.outDqSfsPtr[permutedIdx + totalNumPaddedTokens * sfBlock] = scaleOut; + } + + PackedIo packedOut; + Type* outElts = reinterpret_cast(&packedOut); +#pragma unroll + for (int i = 0; i < kDsActEltsPerThread; ++i) + { + // Divide; do NOT hoist a reciprocal. `x / s` and `x * (1/s)` round + // differently, and an equivalence run showed that single ulp flip a + // greedy-decoded token. This must match activationDeepSeekKernel + // bit for bit, including 0/0 -> NaN on an all-zero scale block. + outElts[i] = static_cast(out[i] / scaleOut); + } + *reinterpret_cast(params.outPtr + static_cast(permutedIdx) * outputDim + hiddenBase) + = packedOut; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + template __global__ void activationDeepSeekKernel(KernelParams params) { @@ -435,8 +577,45 @@ void run(Data const& data, void* stream) const dim3 grid(gridSizeX, gridSizeY, data.topK); - LAUNCH_ACTIVATION( - data, activationDeepSeekKernel, numTokensPerCta, grid, DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA, 0, stream); + // The two kernels sweep different spaces, and which one is cheaper flips + // with batch size. + // + // The expanded-space kernel visits numTokens x topK slots and skips the + // ~(1 - 1/ep_size) of them that are not local, so it never touches the + // per-expert tile padding. The permuted-space kernel sweeps + // [0, totalNumPaddedTokens), which *is* padded: each local expert + // contributes up to tileTokensDim-1 rows of padding that carry no real + // tokens but cost a full row of load/compute/store. + // + // At prefill that padding is noise next to the real rows, and the + // permuted sweep wins by the ep_size factor. At decode the ratio + // inverts: a single token leaves well under one real row per expert + // against the same padding, so the permuted kernel does almost nothing + // but padding. Getting this wrong costs more on every decode step than + // the prefill win is worth over a full generation. + // + // So gate on real work per expert. tileTokensDim is exactly the padding + // granularity, which makes it the natural threshold: below it, an + // expert's real rows do not even fill the tile that must be swept for it. + if (shouldUsePermutedActivation(data.innerDim, data.numTokens, data.topK, data.numExperts, data.tileTokensDim)) + { + int64_t const maxTasks = static_cast(data.numTokens) * data.topK * (outputDim / kDsActEltsPerSf); + int64_t const ctasForAllTasks = (maxTasks + kDsActWarpsPerCta - 1) / kDsActWarpsPerCta; + // Persistent grid: totalNumPaddedTokens is a device-side value, so the + // host can only bound it. Cap at a few waves and let the grid stride + // absorb the difference rather than launching the (numTokens x topK) + // worst case that the expanded-space kernel pays unconditionally. + int const numCtas = static_cast(std::min(ctasForAllTasks, int64_t{numSms} * 32)); + dim3 const permutedGrid(std::max(numCtas, 1), 1, 1); + + LAUNCH_ACTIVATION( + data, activationDeepSeekPermutedKernel, 1, permutedGrid, kDsActPermutedNumThreadsPerCta, 0, stream); + } + else + { + LAUNCH_ACTIVATION(data, activationDeepSeekKernel, numTokensPerCta, grid, + DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA, 0, stream); + } } else { diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h index 171ceaaf12ea..a87106e604f2 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -259,6 +259,14 @@ struct Data int32_t topK; int32_t* expandedIdxToPermutedIdx; + // Used only to pick between the two activation kernels; see the dispatch + // note in DevKernel.cu. The permuted-space kernel sweeps + // totalNumPaddedTokens, which carries up to one partial tile of padding per + // local expert, so it only pays off when there is enough real work per + // expert to amortise that sweep. Both are host-side values. + int32_t numExperts{0}; + int32_t tileTokensDim{0}; + int32_t const* totalNumPaddedTokens; // Optional swiglu clamp limit (fp32, uniform across experts on the FP8 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index 635407433c90..d1737854a60c 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -622,6 +622,7 @@ Runner::Runner( : mPermuteGemm1(PermuteGemm1::Runner(dtypeAct, dtypeWeights, useDeepSeekFp8, tileTokensDim, actType)) , mGemm2(Gemm2::Runner(dtypeAct, dtypeWeights, btg::Dtype::Bfloat16, useDeepSeekFp8, tileTokensDim)) , mActType(actType) + , mTileTokensDim(tileTokensDim) { auto const& gemm1PassingIndices = mPermuteGemm1.getPassingConfigIndices(); auto const& gemm2PassingIndices = mGemm2.getPassingConfigIndices(); @@ -670,6 +671,8 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace activationData.topK = args.top_k; activationData.numTokens = args.num_tokens; activationData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx; + activationData.numExperts = args.num_experts; + activationData.tileTokensDim = mTileTokensDim; // For DeepSeek FP8 the activation runs as a separate kernel rather than // fused into the FC1 GEMM cubin; forward the scalar swiglu_limit so it // can honor swiglu_limit (uniform across experts; see DevKernel.h note). diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index 97d77eddad0a..763b3d8f86b4 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -418,6 +418,9 @@ class Runner PermuteGemm1::Runner mPermuteGemm1; Gemm2::Runner mGemm2; ActType mActType; + // Kept so setOpsData can tell the activation launcher how much per-expert + // tile padding the permuted row space carries. + int32_t mTileTokensDim; // This will be the cartesian product of the passing configs for gemm1 and gemm2 // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2 separately diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 5f121d2406a9..6b0e5a118211 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -100,6 +100,8 @@ set(ROUTING_KERNEL_TEST_SRC add_gtest(routingKernelsTest "${ROUTING_KERNEL_TEST_SRC}") target_link_libraries(routingKernelsTest PRIVATE Python3::Python) +add_gtest(blockScaleMoeActivationTest blockScaleMoeActivationTest.cu) + add_gtest(moeLoadBalanceKernelTest moeLoadBalanceKernelTest.cpp) if(USING_OSS_CUTLASS_MOE_GEMM) diff --git a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu new file mode 100644 index 000000000000..4a56dcd8f558 --- /dev/null +++ b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu @@ -0,0 +1,438 @@ +/* + * 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. + */ + +// Bit-exact equivalence between the two DeepSeek-FP8 MoE activation kernels. +// +// `moe::dev::activation::run` picks between +// * `activationDeepSeekKernel` - grids over the expanded (numTokens x +// topK) index space and discovers work through expandedIdxToPermutedIdx, +// * `activationDeepSeekPermutedKernel` - grids directly over the permuted row +// space with one warp per (row, 128-element scale block), +// via `shouldUsePermutedActivation()`. Both must produce *identical bits* for +// every row that carries a real token: DevKernel.cu documents that the permuted +// kernel must not hoist a reciprocal out of `out / scaleOut`, because `x / s` +// and `x * (1/s)` round differently and one ulp was enough to flip a +// greedy-decoded token. An `isClose`-style comparison would not catch that +// regression, so everything below compares raw bit patterns (which also makes +// the NaN cases comparable). +// +// Note on coverage: fp8 e4m3 carries three mantissa bits, so most 1-ulp fp32 +// differences vanish when the result is rounded back down to fp8 -- only values +// sitting on a rounding boundary survive. A single small shape can therefore +// miss the reciprocal regression by chance, which is why several shapes with +// different scale-block counts are instantiated below. +// +// The dispatch inputs `Data::numExperts` and `Data::tileTokensDim` exist *only* +// for that choice: `KernelParams::setKernelParams` does not forward them, so +// neither kernel can observe them. Overriding `tileTokensDim` therefore selects +// the kernel without perturbing a single input byte, which is what lets this +// test run the very same inputs through both paths. + +#include + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h" +#include "tensorrt_llm/runtime/bufferManager.h" +#include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/iBuffer.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::tests::kernels::blockscalemoe +{ + +namespace tc = tensorrt_llm::common; +namespace tg = batchedGemm::trtllm::gen; + +using tensorrt_llm::runtime::BufferManager; +using tensorrt_llm::runtime::CudaStream; +using tensorrt_llm::runtime::ITensor; +using tensorrt_llm::runtime::MemoryType; +using tensorrt_llm::runtime::bufferCast; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// `shouldUsePermutedActivation()` requires +// realRowsPerExpert = numTokens * topK / numExperts >= tileTokensDim. +// A tile of 1 always satisfies it (given numTokens * topK >= numExperts); a +// tile larger than the whole expanded space never can. +constexpr int32_t kTileForcePermuted = 1; +constexpr int32_t kTileForceLegacy = 1 << 20; + +// The activation scale factor is `aMax / 448.f` with `aMax >= 0`, so a negative +// value can never be produced by either kernel and makes an unambiguous +// "this entry was not written" marker. +constexpr float kScaleSentinel = -12345.0F; +constexpr int8_t kDataSentinel = static_cast(0x5A); + +constexpr int32_t kEltsPerSf = 128; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct ActivationEquivParam +{ + std::string name; + int32_t numTokens; + int32_t topK; + int32_t numExperts; // global expert count + int32_t numLocalExperts; // this rank's share, i.e. numExperts / epSize + int32_t intermediateSize; + int32_t paddingTile; // routing tile the permuted layout was built with + bool hasSwigluLimit; + float swigluLimit; + uint32_t seed; +}; + +inline std::ostream& operator<<(std::ostream& os, ActivationEquivParam const& p) +{ + return os << p.name; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Host-side stand-in for what the routing kernel produces: a permuted row space +// grouped by local expert, with each expert's row count padded up to +// `paddingTile`. +struct PermutedLayout +{ + std::vector expandedIdxToPermutedIdx; // numTokens * topK, -1 for non-local + std::vector realRows; // rows carrying a token + std::vector paddingRows; // rows that exist only because of tile padding + int32_t totalNumPaddedTokens{0}; +}; + +// One activation run, read back to the host as raw bytes / raw float bits. +struct ActivationResult +{ + std::vector bytes; + std::vector scales; +}; + +inline PermutedLayout buildPermutedLayout(ActivationEquivParam const& p) +{ + std::mt19937 rng(p.seed); + + // Pick topK distinct experts per token, then keep the ones this rank owns + // (local expert offset 0, i.e. experts [0, numLocalExperts)). + std::vector> slotsPerLocalExpert(p.numLocalExperts); + std::vector experts(p.numExperts); + std::iota(experts.begin(), experts.end(), 0); + for (int32_t token = 0; token < p.numTokens; ++token) + { + std::shuffle(experts.begin(), experts.end(), rng); + for (int32_t k = 0; k < p.topK; ++k) + { + int32_t const expert = experts[k]; + if (expert < p.numLocalExperts) + { + slotsPerLocalExpert[expert].push_back(token * p.topK + k); + } + } + } + + PermutedLayout layout; + layout.expandedIdxToPermutedIdx.assign(p.numTokens * p.topK, -1); + int32_t offset = 0; + for (int32_t expert = 0; expert < p.numLocalExperts; ++expert) + { + auto const& slots = slotsPerLocalExpert[expert]; + auto const numSlots = static_cast(slots.size()); + for (int32_t i = 0; i < numSlots; ++i) + { + layout.expandedIdxToPermutedIdx[slots[i]] = offset + i; + layout.realRows.push_back(offset + i); + } + int32_t const paddedCount = tc::ceilDiv(numSlots, p.paddingTile) * p.paddingTile; + for (int32_t row = numSlots; row < paddedCount; ++row) + { + layout.paddingRows.push_back(offset + row); + } + offset += paddedCount; + } + layout.totalNumPaddedTokens = offset; + return layout; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int8_t toFp8Byte(float value) +{ + cutlass::float_e4m3_t const converted(value); + int8_t byte{}; + std::memcpy(&byte, &converted, sizeof(byte)); + return byte; +} + +inline uint32_t floatBits(float value) +{ + uint32_t bits{}; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class BlockScaleMoeActivationEquivalenceTest : public ::testing::TestWithParam +{ +protected: + void SetUp() override + { + if (tc::getSMVersion() < 90) + { + GTEST_SKIP() << "The trtllm-gen block-scale MoE activation kernels target SM90+."; + } + mStream = std::make_shared(); + mBufferManager = std::make_shared(mStream); + } + + // Allocates device buffers, fills the inputs deterministically and uploads + // them. `zeroedRowBlock`, when set, forces one (row, scale block) pair of + // the input to all-zero so both kernels take the aMax == 0 -> 0/0 -> NaN + // path on exactly the same element. + void setUp(ActivationEquivParam const& param, PermutedLayout const& layout, + std::optional> zeroedRowBlock = std::nullopt) + { + mParam = param; + mInnerDim = 2 * param.intermediateSize; + mOutputDim = param.intermediateSize; + mTotalRows = layout.totalNumPaddedTokens; + mNumOutSfBlocks = mOutputDim / kEltsPerSf; + + auto const numInElts = static_cast(mTotalRows) * mInnerDim; + auto const numInSfElts = static_cast(mInnerDim / kEltsPerSf) * mTotalRows; + mNumOutElts = static_cast(mTotalRows) * mOutputDim; + mNumOutSfElts = static_cast(mNumOutSfBlocks) * mTotalRows; + + std::mt19937 rng(param.seed + 977U); + std::uniform_real_distribution valueDist(-4.F, 4.F); + std::uniform_real_distribution scaleDist(0.05F, 2.F); + + std::vector hostIn(numInElts); + for (auto& byte : hostIn) + { + byte = toFp8Byte(valueDist(rng)); + } + std::vector hostInSf(numInSfElts); + for (auto& scale : hostInSf) + { + scale = scaleDist(rng); + } + + if (zeroedRowBlock.has_value()) + { + auto const [row, sfBlock] = *zeroedRowBlock; + auto const base = static_cast(row) * mInnerDim + static_cast(sfBlock) * kEltsPerSf; + // Both halves: `up` at [base, base+128) and `gate` at [base+outputDim, ...). + std::fill_n(hostIn.begin() + base, kEltsPerSf, int8_t{0}); + std::fill_n(hostIn.begin() + base + mOutputDim, kEltsPerSf, int8_t{0}); + } + + auto upload = [this](auto const& host) { + return mBufferManager->copyFrom( + host, ITensor::makeShape({static_cast(host.size())}), MemoryType::kGPU); + }; + + mInDevice = upload(hostIn); + mInSfDevice = upload(hostInSf); + mExpandedMapDevice = upload(layout.expandedIdxToPermutedIdx); + mTotalPaddedDevice = upload(std::vector{mTotalRows}); + + mOutDevice = mBufferManager->gpu(ITensor::makeShape({mNumOutElts}), tensorrt_llm::DataType::kINT8); + mOutSfDevice = mBufferManager->gpu(ITensor::makeShape({mNumOutSfElts}), tensorrt_llm::DataType::kFLOAT); + + mStream->synchronize(); + } + + // Resets the outputs to sentinels, runs the activation with the requested + // dispatch override, and reads the results back. + ActivationResult runOnce(int32_t tileTokensDimOverride) + { + ActivationResult result; + std::vector const outSentinel(mNumOutElts, kDataSentinel); + std::vector const outSfSentinel(mNumOutSfElts, kScaleSentinel); + mBufferManager->copy(outSentinel.data(), *mOutDevice); + mBufferManager->copy(outSfSentinel.data(), *mOutSfDevice); + + moe::dev::activation::Data data; + data.mDtypeElt = tg::Dtype::E4m3; + data.mUsePdl = false; + data.mUseDeepSeekFp8 = true; + data.inPtr = bufferCast(*mInDevice); + data.outPtr = bufferCast(*mOutDevice); + data.inDqSfsPtr = bufferCast(*mInSfDevice); + data.outDqSfsPtr = bufferCast(*mOutSfDevice); + data.innerDim = mInnerDim; + data.numTokens = mParam.numTokens; + data.topK = mParam.topK; + data.expandedIdxToPermutedIdx = bufferCast(*mExpandedMapDevice); + data.numExperts = mParam.numExperts; + data.tileTokensDim = tileTokensDimOverride; + data.totalNumPaddedTokens = bufferCast(*mTotalPaddedDevice); + data.swigluLimit = mParam.swigluLimit; + data.hasSwigluLimit = mParam.hasSwigluLimit; + + moe::dev::activation::run(data, mStream->get()); + TLLM_CUDA_CHECK(cudaGetLastError()); + + result.bytes.resize(mNumOutElts); + result.scales.resize(mNumOutSfElts); + mBufferManager->copy(*mOutDevice, result.bytes.data()); + mBufferManager->copy(*mOutSfDevice, result.scales.data()); + mStream->synchronize(); + return result; + } + + std::shared_ptr mStream; + std::shared_ptr mBufferManager; + + ITensor::SharedPtr mInDevice; + ITensor::SharedPtr mInSfDevice; + ITensor::SharedPtr mOutDevice; + ITensor::SharedPtr mOutSfDevice; + ITensor::SharedPtr mExpandedMapDevice; + ITensor::SharedPtr mTotalPaddedDevice; + + ActivationEquivParam mParam{}; + int32_t mInnerDim{0}; + int32_t mOutputDim{0}; + int32_t mTotalRows{0}; + int32_t mNumOutSfBlocks{0}; + int64_t mNumOutElts{0}; + int64_t mNumOutSfElts{0}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST_P(BlockScaleMoeActivationEquivalenceTest, BothKernelsAgreeBitForBit) +{ + auto const param = GetParam(); + auto const layout = buildPermutedLayout(param); + + ASSERT_FALSE(layout.realRows.empty()) << "the test config produced no local rows"; + // The dispatch assertions below rely on padding existing. + ASSERT_FALSE(layout.paddingRows.empty()) << "the test config produced no tile padding"; + ASSERT_GE(static_cast(param.numTokens) * param.topK, param.numExperts) + << "kTileForcePermuted only forces the permuted kernel when numTokens * topK >= numExperts"; + + setUp(param, layout); + + auto const legacy = runOnce(kTileForceLegacy); + auto const permuted = runOnce(kTileForcePermuted); + + // Guard the dispatch itself. Only the permuted kernel sweeps the per-expert + // tile padding, so the padding rows tell the two kernels apart without + // reaching into DevKernel.cu internals. If a future refactor made both runs + // take the same branch, this fires instead of the comparison passing + // vacuously. + for (auto const row : layout.paddingRows) + { + for (int32_t sfBlock = 0; sfBlock < mNumOutSfBlocks; ++sfBlock) + { + auto const idx = static_cast(row) + static_cast(mTotalRows) * sfBlock; + ASSERT_EQ(floatBits(legacy.scales[idx]), floatBits(kScaleSentinel)) + << "the expanded-space kernel must not touch padding row " << row; + ASSERT_NE(floatBits(permuted.scales[idx]), floatBits(kScaleSentinel)) + << "the permuted-space kernel must sweep padding row " << row; + } + } + + // Every real row must have been written by both kernels, bit for bit. + // ASSERT (not EXPECT) so a regression that hits every row reports the first + // offending element instead of hundreds of thousands of them. + for (auto const row : layout.realRows) + { + for (int32_t sfBlock = 0; sfBlock < mNumOutSfBlocks; ++sfBlock) + { + auto const idx = static_cast(row) + static_cast(mTotalRows) * sfBlock; + ASSERT_NE(floatBits(legacy.scales[idx]), floatBits(kScaleSentinel)) + << "real row " << row << " scale block " << sfBlock << " was never written"; + ASSERT_EQ(floatBits(legacy.scales[idx]), floatBits(permuted.scales[idx])) + << "scale mismatch at row " << row << " block " << sfBlock; + } + + for (int32_t elt = 0; elt < mOutputDim; ++elt) + { + auto const idx = static_cast(row) * mOutputDim + elt; + ASSERT_EQ(static_cast(legacy.bytes[idx]), static_cast(permuted.bytes[idx])) + << "fp8 mismatch at row " << row << " element " << elt; + } + } +} + +INSTANTIATE_TEST_SUITE_P(BlockScaleMoeActivation, BlockScaleMoeActivationEquivalenceTest, + ::testing::Values( + // Small shape, heavy padding (paddingTile 8 against ~8 rows per expert). + ActivationEquivParam{"small", /*numTokens=*/64, /*topK=*/4, /*numExperts=*/32, /*numLocalExperts=*/8, + /*intermediateSize=*/128, /*paddingTile=*/8, /*hasSwigluLimit=*/false, /*swigluLimit=*/0.F, /*seed=*/13U}, + // Two scale blocks per row, so the warp-per-(row, block) mapping is + // exercised with more than one block per row. + ActivationEquivParam{"two_sf_blocks", 128, 4, 32, 8, 256, 8, false, 0.F, 17U}, + // Production-sized intermediate dims: 4 and 8 scale blocks per row. + ActivationEquivParam{"four_sf_blocks", 96, 8, 64, 16, 512, 16, false, 0.F, 29U}, + ActivationEquivParam{"eight_sf_blocks", 64, 10, 40, 10, 1024, 8, false, 0.F, 31U}, + // The clamped SwiGLU branch (gemm1_clamp_limit) must match too. + ActivationEquivParam{"swiglu_limit", 128, 4, 32, 8, 256, 8, true, 1.5F, 37U}), + [](::testing::TestParamInfo const& info) { return info.param.name; }); + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// An all-zero scale block yields aMax == 0, so the quantization does 0 / 0. The +// resulting NaN encoding is unspecified, but both kernels evaluate the same +// expression and must therefore land on the same bits -- which is exactly what +// would break if one of them replaced the division with a multiply by the +// reciprocal. +TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNaNs) +{ + ActivationEquivParam const param{"zero_block", /*numTokens=*/64, /*topK=*/4, /*numExperts=*/32, + /*numLocalExperts=*/8, /*intermediateSize=*/256, /*paddingTile=*/8, /*hasSwigluLimit=*/false, + /*swigluLimit=*/0.F, /*seed=*/41U}; + auto const layout = buildPermutedLayout(param); + ASSERT_FALSE(layout.realRows.empty()); + + int32_t const zeroedRow = layout.realRows.front(); + int32_t const zeroedBlock = 0; + setUp(param, layout, std::make_pair(zeroedRow, zeroedBlock)); + + auto const legacy = runOnce(kTileForceLegacy); + auto const permuted = runOnce(kTileForcePermuted); + + auto const sfIdx = static_cast(zeroedRow) + static_cast(mTotalRows) * zeroedBlock; + EXPECT_EQ(floatBits(legacy.scales[sfIdx]), 0U) << "an all-zero block must give scaleOut == +0"; + EXPECT_EQ(floatBits(legacy.scales[sfIdx]), floatBits(permuted.scales[sfIdx])); + + for (int32_t elt = 0; elt < kEltsPerSf; ++elt) + { + auto const idx = static_cast(zeroedRow) * mOutputDim + zeroedBlock * kEltsPerSf + elt; + ASSERT_EQ(static_cast(legacy.bytes[idx]), static_cast(permuted.bytes[idx])) + << "0/0 encoding differs at element " << elt; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace tensorrt_llm::tests::kernels::blockscalemoe From 5b787519b1742ac80e4e462a72b856d2462fc4ba Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 31 Jul 2026 09:08:33 +0000 Subject: [PATCH 3/4] [None][fix] Select explicit FP8 MoE fallback tactic Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../custom_ops/trtllm_gen_custom_ops.py | 57 +++++++++++++++++-- .../_torch/modules/moe/test_moe_backend.py | 15 +++++ 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 083892888fbb..0a9c874785f7 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -1,3 +1,18 @@ +# 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. + import os from dataclasses import dataclass, replace from functools import lru_cache @@ -52,6 +67,19 @@ _RANDOM = "random" +def _select_explicit_fallback_tactic( + valid_tactics: List[List[int]]) -> List[int]: + if not valid_tactics: + raise RuntimeError( + "The FP8 block-scale MoE runner has no valid fallback tactic for " + "the current input shape.") + + # Prefer the widest valid token tile so the fallback uses the largest + # routing workspace among the runner's candidates. Pick the lowest config + # index within that tile to keep cache-miss behavior deterministic. + return max(valid_tactics, key=lambda tactic: (tactic[0], -tactic[1])) + + def prepare_dummy_topk_and_hook( topk_weights: Optional[torch.Tensor], topk_ids: Optional[torch.Tensor], @@ -830,6 +858,7 @@ class FP8BlockScaleMoEInputs: class FP8BlockScaleMoERunner(TunableRunner): runner_dict = dict() + fallback_tactic_dict = dict() tuning_config = None def __init__( @@ -882,6 +911,25 @@ def get_runner(self): instance_key] = torch.classes.trtllm.FP8BlockScaleMoERunner() return FP8BlockScaleMoERunner.runner_dict[instance_key] + def get_fallback_tactic(self, hidden_size: int, + num_tokens: int) -> List[int]: + """Deterministic tactic for the AutoTuner cache-miss path. + + Enumerating the valid tactics walks every (tileN, config) pair through + ``isValidConfig``, and this runs on *every* forward call that misses + the profiling cache. The result depends only on the key below, so + memoize it the same way ``runner_dict`` memoizes the runner. + """ + key = (self.top_k, hidden_size, self.intermediate_size, + self.local_num_experts, num_tokens) + tactic = FP8BlockScaleMoERunner.fallback_tactic_dict.get(key) + if tactic is None: + tactic = tuple( + _select_explicit_fallback_tactic( + self.get_runner().get_valid_configs(*key))) + FP8BlockScaleMoERunner.fallback_tactic_dict[key] = tactic + return list(tactic) + def forward( self, inputs: List[torch.Tensor], @@ -1095,15 +1143,16 @@ def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], input_tensors_for_tuner, ) + if best_tactic == -1: + best_tactic = kernel_runner.get_fallback_tactic(hidden_states.shape[1], + hidden_states.shape[0]) + input_tensors = input_tensors_for_tuner input_tensors[ 0] = routing_logits # replace dummy routing logits with actual routing logits input_tensors[-2] = topk_weights # replace dummy topk_weights with actual input_tensors[-1] = topk_ids # replace dummy topk_ids with actual - result = kernel_runner( - input_tensors, - tactic=[-1, -1] if best_tactic == -1 else best_tactic, - output=output) + result = kernel_runner(input_tensors, tactic=best_tactic, output=output) # When output is provided, the result is written in-place to output. # Return empty tensor to avoid aliasing constraint violation in PyTorch 2.9.1+ # (custom op output cannot be the same tensor as input). diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index e8c8d9511e2c..057ee3e502cc 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -52,6 +52,7 @@ from transformers.configuration_utils import PretrainedConfig from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import _select_explicit_fallback_tactic from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.fused_moe import ( DeepSeekV3MoeRoutingMethod, @@ -81,6 +82,20 @@ } +def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): + valid_tactics = [ + [8, 0], + [32, 4], + [16, 0], + [32, 0], + ] + + assert _select_explicit_fallback_tactic(valid_tactics) == [32, 0] + + with pytest.raises(RuntimeError, match="no valid fallback tactic"): + _select_explicit_fallback_tactic([]) + + def _ensure_single_proc_dist_for_megamoe(backend_type: MoeBackendType, rank: int) -> None: """Every MegaMoE backend (DG + CuteDSL) resolves an EP ProcessGroup at construction time via ``_resolve_ep_pg``. Single-process tests From b30713cbd5213f7aef0ba1b67b834f4439ddc848 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Fri, 31 Jul 2026 09:08:49 +0000 Subject: [PATCH 4/4] [None][fix] Normalize one-model draft KV pool ratio Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 12 ++++ .../executor/test_kv_cache_estimation.py | 57 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9c0023cac331..b72708dbd3b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1256,6 +1256,18 @@ def _create_one_model_draft_kv_cache_manager( "would not be partitioned from the target's and would overrun " "GPU memory. Derived draft max_attention_window=" f"{draft_kv_config.max_attention_window}.") + if (draft_kv_config.pool_ratio is not None + and len(draft_kv_config.pool_ratio) != 1): + # pool_ratio describes one manager's pool-group layout. The + # target hybrid manager may have separate recurrent-state and + # attention groups, while today's supported one-model draft + # manager has one non-VSWA attention group. Reusing the target's + # two ratios for that separate manager fails its arity check. + logger.info( + "Normalizing the separate one-model draft KV cache pool_ratio " + f"from {draft_kv_config.pool_ratio} to [1.0] for its single " + "pool group.") + draft_kv_config.pool_ratio = [1.0] if is_vswa_enabled(self._kv_cache_config): logger.info( f"Derived draft KV cache max_attention_window for separate " diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 6a30bf974d68..b5ad697d5796 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -712,3 +712,60 @@ def __init__(self, kv_cache_config, _kv_cache_type, **kwargs) -> None: assert captured_configs[0].avg_seq_len == expected_avg_seq_len assert kv_cache_config.avg_seq_len == 2055 + + +def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: + creator = object.__new__(KvCacheCreator) + target_pool_ratio = [0.32, 0.68] + creator._kv_cache_config = KvCacheConfig( + pool_ratio=target_pool_ratio, + max_attention_window=None, + ) + creator._max_seq_len = 9472 + creator._tokens_per_block = 32 + creator._max_batch_size = 1024 + creator._max_num_tokens = 9472 + creator._max_beam_width = 1 + creator._kv_connector_manager = None + creator._skip_est = False + creator._execution_stream = None + creator._is_disagg = False + creator._mapping = Mock() + creator._speculative_config = Mock() + + effective_draft_config = Mock() + effective_draft_config.pretrained_config.torch_dtype = "bfloat16" + effective_draft_config.sparse_attention_config = None + + with ( + patch.object(creator, "_get_num_draft_layers", return_value=1), + patch.object(creator, "_get_one_model_draft_layer_mask", return_value=[True]), + patch.object( + creator, + "_get_effective_draft_config", + return_value=effective_draft_config, + ), + patch.object(creator, "_enable_kv_cache_stats", return_value=False), + patch.object( + creator, + "_fallback_if_unsupported_kv_cache_manager_v2", + return_value=KVCacheManagerV2, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util._derive_draft_max_attention_window", + return_value=None, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", + return_value=KVCacheManagerV2, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=Mock(), + ) as create_manager, + ): + creator._create_one_model_draft_kv_cache_manager() + + draft_config = create_manager.call_args.kwargs["kv_cache_config"] + assert draft_config.pool_ratio == [1.0] + assert creator._kv_cache_config.pool_ratio == target_pool_ratio