Nvfp4 Tensor

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.

Typedefs

using Dim3 = int3

Three tile indices for NVFP4Tensor: x → strides[0] (slowest), y → strides[1], z → strides[2] (fastest).

Functions

__device__ __forceinline__ float nvfp4TensorScaleAt (float const *global_scale, int const idx)

global_scale[idx] if global_scale is non-null, else 1.f.

Variables

int kNvfp4ElemsPerTile = 64

Packed NVFP4 (E2M1) payload for one 64-element tile: 8×uint32 lanes, 8 FP4 values per uint32.

int64_t kNvfp4Int4PerTilePayload = 2

Number of int4 vectors per tile payload (8×uint32 lanes == 2×int4); int64_t for mixed int4 / int64 stride math.

struct NVFP4Tensor

Device view of an NVFP4 tensor stored in 64-element tiles with per-tile FP8 block scales and optional global_scale (applied after block dequantization).

global_scale is a device pointer to E floats from the MoE plugin (per-expert scales) when the tensor is stacked expert weights (slow tile index x is the expert id); kernels use global_scale[tile.x]. nullptr means a factor of 1.f. MoE W4A4 activation: global_scale[0] scales the activation in the up-proj GEMV (after NVFP4 dequant). Down-proj decode does not read the activation tensor. nullptr means 1.f for index 0.

Logical layout is 3D in tiles (each tile is kNvfp4ElemsPerTile NVFP4 values + one packed scale int). Index order is row-major C: z is the fastest-varying tile index, x the slowest.

quantized_data points at the first int4 of the tile at logical origin (0,0,0). Each tile occupies kNvfp4Int4PerTilePayload contiguous int4 vectors. strides[d] is the offset in int4 elements when logical tile index component d increases by 1 (NumPy-style / sizeof(int4)).

block_scale uses atom-layout 128×4 swizzle (matching Blackwell tcgen05.mma block-scaled MMA and TMA scale-factor multicasts). Each FP8 E4M3 byte is stored at a byte offset computed by tileScaleIndex; four consecutive bytes (innerK 0..3) form one int32 word read by readBlockScaleWord.

Atom-layout byte offset for scale factor at logical (mRow, kChunk): innerK = kChunk % 4; innerM = (mRow % 128) / 32; outerM = mRow % 32; kTile = kChunk / 4; mTile = mRow / 128; byteOffset = mTile * numKTiles * 512 + kTile * 512 + outerM * 16 + innerM * 4 + innerK

The int32 word index (for block_scale pointer) drops innerK: mTile * numKTiles * 128 + kTile * 128 + outerM * 4 + innerM

scaleMDimIdx / scaleKDimIdx select which Dim3 component maps to mRow / kChunk. scaleNumKTiles is ceil(numSfCols / 4) where numSfCols = K / 16. scaleExpertStride is the int32 element stride between experts (ceil(M/128) * numKTiles * 128).

Public Functions

inline __device__ __forceinline__ int64_t tileOffsetInt4 (Dim3 const c) const

inline __device__ __forceinline__ int64_t tileScaleIndex (Dim3 const c) const

Atom-layout 128×4 swizzle: returns int32 index into block_scale for the tile at c.

inline __device__ __forceinline__ int readBlockScaleWord (Dim3 const c) const

Read block-scale int32 word with byte-swap for Marlin dequant compatibility. Atom stores sequential FP8 bytes {s0,s1,s2,s3}; Marlin dequant_fp8_scales expects {s0,s2,s1,s3}.

inline __device__ __forceinline__ int readBlockScaleWordLinear (Dim3 const c) const

Read block-scale int32 word using plain linear tile indexing (one scale word per tile payload). The word is already in Marlin-packed byte order — no byte-swap. Used for activation scale factors whose layout matches TRT_DynamicQuantize output (contiguous, not atom-swizzled).

inline __device__ __forceinline__ void loadTileUint4 (Dim3 const tile, int const chunk, uint4 &out) const

Loads one int4 chunk (4× uint32 NVFP4 lane packs) from the tile payload.

Parameters:

chunk – 0 or 1 — tile payload is kNvfp4Int4PerTilePayload int4 vectors.

inline __device__ __forceinline__ uint32_t loadTileUint32Lane (Dim3 const tile, int const chunk, int const lane) const

Loads one uint32 lane (8 packed NVFP4 values) from the tile payload.

Parameters:

• chunk – 0 or 1;

• lane – 0..3 within that chunk (see loadTileUint4).

Public Members

int4 *quantized_data

int *block_scale

float *global_scale

Device pointer to E per-expert scales (MoE plugin inputs), or nullptr. See struct comment.

int64_t strides[3]

Stride in int4 elements per +1 in tile index dim 0..2 (z fastest).

int32_t scaleNumKTiles

Atom-layout scale factor metadata.

K-tiles for atom swizzle: ceil(numSfCols / 4) where numSfCols = K_dim / 16.

int64_t scaleExpertStride

int32 elements between experts: ceil(M/128) * scaleNumKTiles * 128.

int8_t scaleMDimIdx

Which Dim3 component is the M-row (0 for act, 1 for weights).

int8_t scaleKDimIdx

Which Dim3 component is the K-chunk (1 for act, 2 for weights).