Nvfp4 A16 Blackwell Moe Dispatch Policy

namespace trt_edgellm

Argument bundles for the CuTe DSL FMHA launchers.

Each public run* entry point fills one of these once from its own arguments and members, and the launcher then spreads it across the generated descriptors. One struct per wrapper shape rather than a single superset: a superset would let a field that matters on one path (kvCacheCapacity on the dense path, tokensPerPage on the paged one) sit silently at zero on another.

These carry no generated types, so this header is safe to include from translation units that use the required FMHA-v2 runner, the optional CUTE_DSL_FMHA_BLACKWELL_ENABLED runner, or both. The descriptor-filling machinery itself lives in cuteDslTensorDescriptors.h, which stays free of any FMHA concept.

Helpers for populating the tensor descriptors emitted by the CuTe DSL C exporter.

Every AOT variant exports its own nominally distinct but layout-identical descriptor structs (fmha_d64_Tensor_q_tensor_t vs fmha_d128_Tensor_q_tensor_t) plus a cute_dsl_<variant>_wrapper entry point. There is no umbrella C type, so descriptor types are recovered here from the signature of the wrapper that consumes them: a call site names only the wrapper and the kernel module, and pairing a descriptor with the wrong variant is not expressible.

The exporter (cutlass/cute/export/c_header_generator.py) always names the members data, dynamic_shapes and dynamic_strides, but emits each array only when its dynamic mask is non-empty. A rank-1 descriptor therefore has no dynamic_strides member at all and needs makeCuSeqLenTensor() rather than the strided builders below.

Small-vocabulary specialisations of the DSpark sampling and verification kernels.

They differ in how the top-k set is found. DSpark selects it in topK passes over the row, each a block-wide max-scan, and drops to a single-threaded walk when top-p is used without top-k or topK exceeds its parallel bound; that cost scales with topK. A residual-VQ codebook row is small enough to stage in shared memory, so these kernels instead run one MSB-radix select plus a bitonic sort over the surviving candidates — a single pass whose cost is independent of topK.

Semantics, tensor layouts, and accept/residual/bonus behaviour are identical to the DSpark entry points named in each declaration, which remain the fallback above kCpSpecMaxVocab.

namespace kernel

Per-position sum across codec groups (Qwen3.5-Omni speaker codec embedding fold).

Mirrors HF _get_codec_input_embeddings: concatenate per-group embeddings then sum across the group axis. Implemented here as a single-row reduction so the caller only invokes one kernel per row of the speaker codec template.

namespace nvfp4_a16_blackwell_moe

Enums

enum class Backend : int32_t

Backend selection for Nvfp4A16BlackwellMoePlugin (plugin attribute backend).

Values:

enumerator kAuto = 0

token-count policy below

enumerator kDecode = 1

force the CUDA-core decode kernels

enumerator kPrefill = 2

force the tcgen05 grouped GEMM

enum class TokenTile : int32_t

Token (MMA-N) tiles baked into the nvfp4_a16_blackwell_moe AOT group. The tile doubles as the per-expert row padding granularity of the permuted activation buffer.

Values:

enumerator kTn8 = 8

enumerator kTn16 = 16

enumerator kTn32 = 32

enumerator kTn64 = 64

enumerator kTn128 = 128

enum class DecodeDtype : int32_t

Activation dtypes of the CUDA-core kernels (part of the JIT key, like the decode constants below and their EDGELLM_MOE_DECODE_* overrides, all fixed when the plugin compiles its bundle at engine build).

Values:

enumerator kFP16 = 0

enumerator kBF16 = 1

Functions

Backend resolveBackend(

Backend const requested,

int32_t const numTokens

) noexcept

TokenTile selectTokenTile(int32_t const numTokens) noexcept

int64_t maxRowsPadded(

int64_t const numTokens,

int64_t const topK,

int64_t const numExperts,

int64_t const tile

) noexcept

Conservative padded-row capacity of the permuted activation buffer for a given tile: every expert may waste tile-1 rows. Multiple of tile.

Variables

int32_t kLargestTokenTile = {128}

int32_t kMaxTokenTiles = {1024}

Upper bound on token tiles per grouped-GEMM launch: the kernel stages tile_group_idx into shared memory once per CTA (MAX_TOKEN_TILES in nvfp4_a16_blackwell_moe_gemm.py). validShape rejects profiles whose padded row count at the selected tile would exceed it (T > ~19k at E=128, topK=6).

int32_t kDecodeMaxTokens = {1}

int32_t kDecodeFc1SplitK = {2}

Decode FC1 split-K: FC1 has only 90 row tiles at T=1 (2.25 waves of 40 resident CTAs on 20 SMs); split-K 2 doubles the CTA count and was the best of {1,2,4,8} in the engine (decode step 11.41 ms vs 11.57 / 11.49 / 11.63). The fp32 partials cost one ~3 us reduce; decode stays deterministic.

int32_t kDecodeFc1MaxSplitK = {8}

Largest FC1 split-K the benchmark override may select; the decode workspace is sized for it so the size recorded at engine build never depends on the environment.

int32_t kDecodeFc2SplitK = {8}

int32_t kDecodeFc2PrefetchSlots = {0}

Decode FC2 slots whose weight tiles are staged into shared memory before the kernel’s PDL wait. Sealed to 0: the larger shared-memory carve-out (18 KB per staged slot for Nemotron) keeps FC2 from co-residing with the shared-expert GEMV on TensorRT’s auxiliary stream, which cost more than the staging saved (decode step 11.44 / 11.51 / 11.52 ms for 0 / 1 / 2 slots). The path stays available through EDGELLM_MOE_DECODE_FC2_PREFETCH or the runner parameter and is exercised by the unit tests.

int32_t kDecodeFc2MaxPrefetchSlots = {2}