From 0f0600152b289a54bcf5ebbcc6bacc9fe12e3d61 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Sat, 19 Sep 2026 09:41:08 +0200 Subject: [PATCH 1/3] refactor(ops): let the sparse-MoE decode codecs own their matrix addressing The four decode kernels addressed a matrix from the call site: a flat group index built from the row, and three plane pointers picked by hand. That holds while every codec keeps one scale per group per row. It does not generalise - a codec whose scale plane is swizzled needs the row to form an address, and the call site cannot compute it without already knowing the codec. A codec now receives the matrix as a plane set plus a row, with the column count as a template argument. The shared expert becomes a codec parameter in the same move: every registered profile stores it as Q8 today, but that is a property of the profiles rather than of the kernels, and the kernels no longer assert it. No behaviour change. The arithmetic, the accumulation order and the launch geometry are untouched for all four registered codecs. --- .../decode/sparse_moe_decode_kernels.cu | 481 +++++++++--------- 1 file changed, 240 insertions(+), 241 deletions(-) diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 8aa0e1bc62..f5dd00f701 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -104,98 +104,128 @@ __global__ void sparse_moe_d2_warp_kernel(const float* __restrict__ scores, int* sparse_moe_select_top8_warp(scores, ids, alpha, shared_scale, selected_logits); } +// The stored planes of one matrix. They travel together because a dot product takes one matrix at +// a time and each codec reads a different subset of them: a call site handing them over one by one +// had to know that subset before it could name the codec. +struct SparseMoePlanes { + const std::uint8_t* __restrict__ codes = nullptr; + const std::uint8_t* __restrict__ high = nullptr; + const std::uint8_t* __restrict__ scales = nullptr; +}; + +// Codecs come in two lane ownerships. Under `kPackedWord8` a lane owns eight consecutive K values +// and a single decode feeds eight FP32 FMAs, so four adjacent groups of eight lanes each form one +// 128-byte warp transaction. The other ownership is scalar - one value per lane in D3, a pair in +// D4 - which is what a 32-wide group leaves room for. Either way the codec takes the matrix's +// column count as a template argument, because it walks its own row and only it knows how its +// planes are laid out. + struct Q4Codec { - static constexpr int kGroupK = 64; - static constexpr bool kD3PackedWord8 = true; + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, + float (&weights)[8]) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; + const std::uint32_t packed = *reinterpret_cast( + planes.codes + index * Q4RowSplitStorage::kCodeBytesPerGroup + lane_in_group * 4); + const auto scale_bits = *reinterpret_cast( + planes.scales + index * Q4RowSplitStorage::kScaleBytesPerGroup); + Q4SimtDecodeAtom::decode_eight(packed, scale_bits, weights); + } }; struct Q5Codec { - static constexpr int kGroupK = 64; - static constexpr bool kPackedWord8 = true; - - __device__ static __forceinline__ void - load_eight(const std::uint8_t* codes, const std::uint8_t* high, const std::uint8_t* scales, - std::int64_t group_index, int lane_in_group, float (&weights)[8]) { + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, + float (&weights)[8]) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( - codes + group_index * Q5RowSplitStorage::kCodeBytesPerGroup + lane_in_group * 4); + planes.codes + index * Q5RowSplitStorage::kCodeBytesPerGroup + lane_in_group * 4); const std::uint8_t high_bits = - high[group_index * Q5RowSplitStorage::kHighBytesPerGroup + lane_in_group]; + planes.high[index * Q5RowSplitStorage::kHighBytesPerGroup + lane_in_group]; const auto scale_bits = *reinterpret_cast( - scales + group_index * Q5RowSplitStorage::kScaleBytesPerGroup); + planes.scales + index * Q5RowSplitStorage::kScaleBytesPerGroup); Q5SimtDecodeAtom::decode_eight(packed, high_bits, scale_bits, weights); } }; struct Q6Codec { - static constexpr int kGroupK = 64; - static constexpr bool kPackedWord8 = true; - - __device__ static __forceinline__ void - load_eight(const std::uint8_t* codes, const std::uint8_t* high, const std::uint8_t* scales, - std::int64_t group_index, int lane_in_group, float (&weights)[8]) { + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, + float (&weights)[8]) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( - codes + group_index * Q6RowSplitStorage::kCodeBytesPerGroup + lane_in_group * 4); + planes.codes + index * Q6RowSplitStorage::kCodeBytesPerGroup + lane_in_group * 4); const std::uint16_t high_bits = *reinterpret_cast( - high + group_index * Q6RowSplitStorage::kHighBytesPerGroup + lane_in_group * 2); + planes.high + index * Q6RowSplitStorage::kHighBytesPerGroup + lane_in_group * 2); const auto scale_bits = *reinterpret_cast( - scales + group_index * Q6RowSplitStorage::kScaleBytesPerGroup); + planes.scales + index * Q6RowSplitStorage::kScaleBytesPerGroup); Q6SimtDecodeAtom::decode_eight(packed, high_bits, scale_bits, weights); } }; struct Q8Codec { - static constexpr int kGroupK = 32; - static constexpr bool kD3SingleValuePerLane = true; - static constexpr bool kD3PackedWord8 = false; - static constexpr bool kPackedWord8 = false; - - __device__ static __forceinline__ float load_one(const std::uint8_t* codes, - const std::uint8_t* scales, - std::int64_t group_index, int lane) { - const float scale = __half2float( - __ushort_as_half(*reinterpret_cast(scales + group_index * 2))); - return static_cast(static_cast(codes[group_index * kGroupK + lane])) * + static constexpr int kGroupK = 32; + static constexpr bool kPackedWord8 = false; + static constexpr bool kSingleValuePerLane = true; + + template + __device__ static __forceinline__ float load_one(const SparseMoePlanes& planes, int row, + int group, int lane) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; + const float scale = __half2float(__ushort_as_half(*reinterpret_cast( + planes.scales + index * Q8RowSplitStorage::kScaleBytesPerGroup))); + return static_cast(static_cast( + planes.codes[index * Q8RowSplitStorage::kCodeBytesPerGroup + lane])) * scale; } - __device__ static __forceinline__ void - load_pair(const std::uint8_t* codes, const std::uint8_t* high, const std::uint8_t* scales, - std::int64_t group_index, int lane, float& w0, float& w1) { - Q8ScalarDecodeAtom::load_pair(codes, high, scales, group_index, lane, w0, w1); + template + __device__ static __forceinline__ void load_pair(const SparseMoePlanes& planes, int row, + int group, int lane, float& w0, float& w1) { + Q8ScalarDecodeAtom::load_pair(planes.codes, planes.high, planes.scales, + static_cast(row) * (K / kGroupK) + group, lane, + w0, w1); } }; +// A codec declares exactly one lane ownership. The trap below catches a codec that declares +// neither, which would otherwise compile and reduce a zero accumulator. +template +inline constexpr bool kNoLaneOwnership = false; + template -__device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const std::uint8_t* high, - const std::uint8_t* scales, int row0, int row1, +__device__ __forceinline__ void dot_two_rows(const SparseMoePlanes& planes, int row0, int row1, const __nv_bfloat16* x, int k_begin, int k_end, float& result0, float& result1) { - constexpr int kGroups = K / Codec::kGroupK; const int lane = static_cast(threadIdx.x) & 31; float acc0 = 0.0f; float acc1 = 0.0f; const int first_group = k_begin / Codec::kGroupK; const int last_group = k_end / Codec::kGroupK; - if constexpr (Codec::kD3PackedWord8) { - // Four adjacent Q4 groups form one 128-byte warp transaction. Each lane owns eight - // consecutive K values, so one mantissa decode feeds eight FP32 FMAs instead of issuing - // four scalar code-pair/decode iterations. + if constexpr (Codec::kPackedWord8) { const int lane_group = lane >> 3; const int lane_in_group = lane & 7; for (int group_base = first_group; group_base < last_group; group_base += 4) { - const int group = group_base + lane_group; - const std::int64_t index0 = static_cast(row0) * kGroups + group; - const std::int64_t index1 = static_cast(row1) * kGroups + group; - const std::uint32_t packed0 = - *reinterpret_cast(codes + index0 * 32 + lane_in_group * 4); - const std::uint32_t packed1 = - *reinterpret_cast(codes + index1 * 32 + lane_in_group * 4); - const auto scale0 = *reinterpret_cast(scales + index0 * 2); - const auto scale1 = *reinterpret_cast(scales + index1 * 2); + const int group = group_base + lane_group; float weights0[8]; float weights1[8]; - Q4SimtDecodeAtom::decode_eight(packed0, scale0, weights0); - Q4SimtDecodeAtom::decode_eight(packed1, scale1, weights1); + Codec::template load_eight(planes, row0, group, lane_in_group, weights0); + Codec::template load_eight(planes, row1, group, lane_in_group, weights1); const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); const float2 x0 = bf16x2_bits_to_float2(input.x); const float2 x1 = bf16x2_bits_to_float2(input.y); @@ -208,41 +238,25 @@ __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const st acc1 = fmaf(weights1[item], values[item], acc1); } } - } else if constexpr (Codec::kD3SingleValuePerLane) { - for (int group = first_group; group < last_group; ++group) { - const std::int64_t index0 = static_cast(row0) * kGroups + group; - const std::int64_t index1 = static_cast(row1) * kGroups + group; - const float w0 = Codec::load_one(codes, scales, index0, lane); - const float w1 = Codec::load_one(codes, scales, index1, lane); - const float xv = __bfloat162float(x[group * Codec::kGroupK + lane]); - acc0 = fmaf(w0, xv, acc0); - acc1 = fmaf(w1, xv, acc1); - } - } else if (lane < Codec::kGroupK / 2) { + } else if constexpr (Codec::kSingleValuePerLane) { for (int group = first_group; group < last_group; ++group) { - float w00, w01, w10, w11; - Codec::load_pair(codes, high, scales, static_cast(row0) * kGroups + group, - lane, w00, w01); - Codec::load_pair(codes, high, scales, static_cast(row1) * kGroups + group, - lane, w10, w11); - const int k = group * Codec::kGroupK + lane * 2; - const float2 xv = __bfloat1622float2(load_vec<__nv_bfloat162>(x + k)); - acc0 = fmaf(w00, xv.x, acc0); - acc0 = fmaf(w01, xv.y, acc0); - acc1 = fmaf(w10, xv.x, acc1); - acc1 = fmaf(w11, xv.y, acc1); + const float w0 = Codec::template load_one(planes, row0, group, lane); + const float w1 = Codec::template load_one(planes, row1, group, lane); + const float xv = __bfloat162float(x[group * Codec::kGroupK + lane]); + acc0 = fmaf(w0, xv, acc0); + acc1 = fmaf(w1, xv, acc1); } + } else { + static_assert(kNoLaneOwnership, "dot_two_rows: codec declares no lane ownership"); } result0 = warp_reduce_sum(acc0); result1 = warp_reduce_sum(acc1); } -template -__global__ void sparse_moe_d3_nine_warp_kernel( - const __nv_bfloat16* __restrict__ x, const int* __restrict__ ids, - const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, - const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, float* __restrict__ act) { +template +__global__ void sparse_moe_d3_nine_warp_kernel(const __nv_bfloat16* __restrict__ x, + const int* __restrict__ ids, SparseMoePlanes routed, + SparseMoePlanes shared, float* __restrict__ act) { __shared__ __align__(16) __nv_bfloat16 x_shared[kHidden]; const int tid = static_cast(threadIdx.x); const int warp = tid >> 5; @@ -258,23 +272,21 @@ __global__ void sparse_moe_d3_nine_warp_kernel( pdl::wait_for_dependencies(); const int expert = ids[warp]; const int row_base = expert * 1024; - dot_two_rows(routed_codes, routed_high, routed_scales, row_base + j, - row_base + kIntermediate + j, x_shared, 0, kHidden, gate, - up); + dot_two_rows(routed, row_base + j, row_base + kIntermediate + j, + x_shared, 0, kHidden, gate, up); } else { - dot_two_rows(shared_codes, nullptr, shared_scales, j, kIntermediate + j, - x_shared, 0, kHidden, gate, up); + dot_two_rows(shared, j, kIntermediate + j, x_shared, 0, kHidden, gate, + up); } if (lane == 0) { act[static_cast(warp) * kIntermediate + j] = silu(gate) * up; } } -template -__global__ void sparse_moe_d3_path_tiled_kernel( - const __nv_bfloat16* __restrict__ x, const int* __restrict__ token_ids, - const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, - const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, float* __restrict__ token_activations, - int tokens, const int* __restrict__ adaptive_route_jobs) { +template +__global__ void sparse_moe_d3_path_tiled_kernel(const __nv_bfloat16* __restrict__ x, + const int* __restrict__ token_ids, + SparseMoePlanes routed, SparseMoePlanes shared, + float* __restrict__ token_activations, int tokens, + const int* __restrict__ adaptive_route_jobs) { // Three path CTAs per token/output row expose enough blocks for the 170-SM target and keep the // heavier shared Q8 path from holding eight completed routed warps resident. static_assert(PathsPerBlock > 0 && (kTopK + 1) % PathsPerBlock == 0); @@ -317,12 +329,11 @@ __global__ void sparse_moe_d3_path_tiled_kernel( if (path < kTopK) { const int expert = token_ids[token * kTopK + path]; const int row_base = expert * (2 * kIntermediate); - dot_two_rows(routed_codes, routed_high, routed_scales, - row_base + j, row_base + kIntermediate + j, x_shared, - 0, kHidden, gate, up); + dot_two_rows(routed, row_base + j, row_base + kIntermediate + j, + x_shared, 0, kHidden, gate, up); } else { - dot_two_rows(shared_codes, nullptr, shared_scales, j, - kIntermediate + j, x_shared, 0, kHidden, gate, up); + dot_two_rows(shared, j, kIntermediate + j, x_shared, 0, kHidden, + gate, up); } if (lane == 0) { token_activations[(static_cast(token) * (kTopK + 1) + path) * @@ -333,19 +344,18 @@ __global__ void sparse_moe_d3_path_tiled_kernel( } } -template -__device__ __forceinline__ void dot_fp32_rows(const std::uint8_t* codes, const std::uint8_t* high, - const std::uint8_t* scales, int row_base, +template +__device__ __forceinline__ void dot_fp32_rows(const SparseMoePlanes& planes, int row_base, const float* x, int first_group, int last_group, float (&result)[Rows]) { - constexpr int kGroups = kIntermediate / Codec::kGroupK; - const int lane = static_cast(threadIdx.x) & 31; + const int lane = static_cast(threadIdx.x) & 31; float acc[Rows]; #pragma unroll for (int row = 0; row < Rows; ++row) { acc[row] = 0.0f; } if constexpr (Codec::kPackedWord8) { - // Q5/Q6 use the same eight-value lane ownership as D3. The high plane and FP16 scale are - // decoded exactly from their registered row-split codec before FP32 accumulation. + // The same eight-value lane ownership as D3, over the FP32 SwiGLU result rather than the + // BF16 hidden state. Each row decodes from its own stored planes before the FP32 + // accumulation. const int lane_group = lane >> 3; const int lane_in_group = lane & 7; for (int group_base = first_group; group_base < last_group; group_base += 4) { @@ -355,42 +365,43 @@ __device__ __forceinline__ void dot_fp32_rows(const std::uint8_t* codes, const s const float values[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w}; #pragma unroll for (int row = 0; row < Rows; ++row) { - const std::int64_t group_index = - static_cast(row_base + row) * kGroups + group; float weights[8]; - Codec::load_eight(codes, high, scales, group_index, lane_in_group, weights); + Codec::template load_eight(planes, row_base + row, group, lane_in_group, + weights); #pragma unroll for (int item = 0; item < 8; ++item) { acc[row] = fmaf(weights[item], values[item], acc[row]); } } } - } else if (lane < Codec::kGroupK / 2) { - for (int group = first_group; group < last_group; ++group) { - const int k = group * Codec::kGroupK + lane * 2; - const float2 xv = load_vec(x + k); + } else if constexpr (Codec::kSingleValuePerLane) { + // Two values a lane here rather than one: the FP32 input is half the width of the BF16 + // hidden state, so the same group needs half the lanes. + if (lane < Codec::kGroupK / 2) { + for (int group = first_group; group < last_group; ++group) { + const int k = group * Codec::kGroupK + lane * 2; + const float2 xv = load_vec(x + k); #pragma unroll - for (int row = 0; row < Rows; ++row) { - float w0, w1; - Codec::load_pair(codes, high, scales, - static_cast(row_base + row) * kGroups + group, lane, - w0, w1); - acc[row] = fmaf(w0, xv.x, acc[row]); - acc[row] = fmaf(w1, xv.y, acc[row]); + for (int row = 0; row < Rows; ++row) { + float w0, w1; + Codec::template load_pair(planes, row_base + row, group, lane, w0, w1); + acc[row] = fmaf(w0, xv.x, acc[row]); + acc[row] = fmaf(w1, xv.y, acc[row]); + } } } + } else { + static_assert(kNoLaneOwnership, "dot_fp32_rows: codec declares no lane ownership"); } #pragma unroll for (int row = 0; row < Rows; ++row) { result[row] = warp_reduce_sum(acc[row]); } } -template +template __global__ void sparse_moe_d4_nine_warp_kernel( const int* __restrict__ ids, const float* __restrict__ alpha, - const float* __restrict__ shared_scale, const float* __restrict__ act, - const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, - const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination, + const float* __restrict__ shared_scale, const float* __restrict__ act, SparseMoePlanes routed, + SparseMoePlanes shared, __nv_bfloat16* __restrict__ destination, const char* __restrict__ prefetch_data, unsigned long long prefetch_bytes) { __shared__ float paths[kTopK + 1][Rows]; pdl::wait_for_dependencies(); @@ -400,19 +411,19 @@ __global__ void sparse_moe_d4_nine_warp_kernel( if (warp < kTopK) { const int expert = ids[warp]; float dot[Rows]; - dot_fp32_rows(routed_codes, routed_high, routed_scales, - expert * kHidden + row_base, - act + static_cast(warp) * kIntermediate, 0, - kIntermediate / RoutedCodec::kGroupK, dot); + dot_fp32_rows( + routed, expert * kHidden + row_base, + act + static_cast(warp) * kIntermediate, 0, + kIntermediate / RoutedCodec::kGroupK, dot); if (lane == 0) { #pragma unroll for (int row = 0; row < Rows; ++row) { paths[warp][row] = alpha[warp] * dot[row]; } } } else { float dot[Rows]; - dot_fp32_rows(shared_codes, nullptr, shared_scales, row_base, - act + static_cast(kTopK) * kIntermediate, 0, - kIntermediate / Q8Codec::kGroupK, dot); + dot_fp32_rows( + shared, row_base, act + static_cast(kTopK) * kIntermediate, 0, + kIntermediate / SharedCodec::kGroupK, dot); if (lane == 0) { #pragma unroll for (int row = 0; row < Rows; ++row) { paths[kTopK][row] = *shared_scale * dot[row]; } @@ -437,14 +448,13 @@ __global__ void sparse_moe_d4_nine_warp_kernel( } } -template -__global__ void sparse_moe_d4_token_kernel( - const int* __restrict__ token_ids, const float* __restrict__ token_alpha, - const float* __restrict__ shared_scale, const float* __restrict__ token_activations, - const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, - const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination, - int tokens, const int* __restrict__ adaptive_route_jobs) { +template +__global__ void +sparse_moe_d4_token_kernel(const int* __restrict__ token_ids, const float* __restrict__ token_alpha, + const float* __restrict__ shared_scale, + const float* __restrict__ token_activations, SparseMoePlanes routed, + SparseMoePlanes shared, __nv_bfloat16* __restrict__ destination, + int tokens, const int* __restrict__ adaptive_route_jobs) { // Token is a grid dimension rather than an in-CTA serial loop. Rows lets one routed-weight // stream serve adjacent outputs while retaining the deterministic rank-order FP32 epilogue. __shared__ float paths[kTopK + 1][Rows]; @@ -468,10 +478,10 @@ __global__ void sparse_moe_d4_token_kernel( if (warp < kTopK) { const int expert = token_ids[token * kTopK + warp]; float dot[Rows]; - dot_fp32_rows(routed_codes, routed_high, routed_scales, - expert * kHidden + row_base, - act + static_cast(warp) * kIntermediate, - 0, kIntermediate / RoutedCodec::kGroupK, dot); + dot_fp32_rows( + routed, expert * kHidden + row_base, + act + static_cast(warp) * kIntermediate, 0, + kIntermediate / RoutedCodec::kGroupK, dot); if (lane == 0) { #pragma unroll for (int row = 0; row < Rows; ++row) { @@ -480,9 +490,9 @@ __global__ void sparse_moe_d4_token_kernel( } } else { float dot[Rows]; - dot_fp32_rows(shared_codes, nullptr, shared_scales, row_base, - act + static_cast(kTopK) * kIntermediate, 0, - kIntermediate / Q8Codec::kGroupK, dot); + dot_fp32_rows( + shared, row_base, act + static_cast(kTopK) * kIntermediate, 0, + kIntermediate / SharedCodec::kGroupK, dot); if (lane == 0) { #pragma unroll for (int row = 0; row < Rows; ++row) { @@ -514,20 +524,21 @@ void launch_d1(const Tensor& x, const SparseMoeWeights& weights, CUDA_CHECK(cudaGetLastError()); } -template +SparseMoePlanes matrix_planes(const Weight& weight) { + return {static_cast(weight.qdata), + static_cast(weight.qhigh), + static_cast(weight.scales)}; +} + +template void launch_d3_dependent_codec(const Tensor& x, const SparseMoeWeights& weights, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { - const auto* input = static_cast(x.data); - const auto* ids = static_cast(workspace.ids.data); - auto* act = static_cast(workspace.scratch.data); - const auto* routed_codes = static_cast(weights.routed_gate_up.qdata); - const auto* routed_high = static_cast(weights.routed_gate_up.qhigh); - const auto* routed_scales = static_cast(weights.routed_gate_up.scales); - const auto* shared_codes = static_cast(weights.shared_gate_up.qdata); - const auto* shared_scales = static_cast(weights.shared_gate_up.scales); CUDA_CHECK(pdl::launch_dependent( - {dim3(kIntermediate), dim3(9 * 32), 0, stream}, sparse_moe_d3_nine_warp_kernel, - input, ids, routed_codes, routed_high, routed_scales, shared_codes, shared_scales, act)); + {dim3(kIntermediate), dim3(9 * 32), 0, stream}, + sparse_moe_d3_nine_warp_kernel, + static_cast(x.data), static_cast(workspace.ids.data), + matrix_planes(weights.routed_gate_up), matrix_planes(weights.shared_gate_up), + static_cast(workspace.scratch.data))); } void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, @@ -541,35 +552,30 @@ void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, switch (weights.routed_gate_up.qtype) { case QType::Q4_G64_FP16: - launch_d3_dependent_codec(x, weights, workspace, stream); + launch_d3_dependent_codec(x, weights, workspace, stream); return; case QType::Q8_G32_FP16: - launch_d3_dependent_codec(x, weights, workspace, stream); + launch_d3_dependent_codec(x, weights, workspace, stream); return; default: throw std::invalid_argument("sparse_moe: unsupported D3 codec"); } } -template +template void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destination, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, const void* prefetch_data, std::size_t prefetch_bytes) { - const auto* ids = static_cast(workspace.ids.data); - const auto* alpha = static_cast(workspace.alpha.data); - const auto* shared_scale = static_cast(workspace.shared_scale.data); - const auto* act = static_cast(workspace.scratch.data); - const auto* routed_codes = static_cast(weights.routed_down.qdata); - const auto* routed_high = static_cast(weights.routed_down.qhigh); - const auto* routed_scales = static_cast(weights.routed_down.scales); - const auto* shared_codes = static_cast(weights.shared_down.qdata); - const auto* shared_scales = static_cast(weights.shared_down.scales); - auto* output = static_cast<__nv_bfloat16*>(destination.data); + constexpr int kRows = 1; CUDA_CHECK(pdl::launch_dependent( - {dim3(kHidden), dim3(9 * 32), 0, stream}, sparse_moe_d4_nine_warp_kernel, ids, - alpha, shared_scale, act, routed_codes, routed_high, routed_scales, shared_codes, - shared_scales, output, static_cast(prefetch_data), - static_cast(prefetch_bytes))); + {dim3(kHidden / kRows), dim3(9 * 32), 0, stream}, + sparse_moe_d4_nine_warp_kernel, + static_cast(workspace.ids.data), + static_cast(workspace.alpha.data), + static_cast(workspace.shared_scale.data), + static_cast(workspace.scratch.data), matrix_planes(weights.routed_down), + matrix_planes(weights.shared_down), static_cast<__nv_bfloat16*>(destination.data), + static_cast(prefetch_data), static_cast(prefetch_bytes))); } void launch_d4_dependent(const SparseMoeWeights& weights, Tensor& destination, @@ -577,89 +583,81 @@ void launch_d4_dependent(const SparseMoeWeights& weights, Tensor& destination, const void* prefetch_data, std::size_t prefetch_bytes) { switch (weights.routed_down.qtype) { case QType::Q5_G64_FP16: - launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, - prefetch_bytes); + launch_d4_dependent_codec(weights, destination, workspace, stream, + prefetch_data, prefetch_bytes); return; case QType::Q6_G64_FP16: - launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, - prefetch_bytes); + launch_d4_dependent_codec(weights, destination, workspace, stream, + prefetch_data, prefetch_bytes); return; case QType::Q8_G32_FP16: - launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, - prefetch_bytes); + launch_d4_dependent_codec(weights, destination, workspace, stream, + prefetch_data, prefetch_bytes); return; default: throw std::invalid_argument("sparse_moe: unsupported D4 codec"); } } -template +template void launch_d3_small_t_paths(const Tensor& x, const SparseMoeWeights& weights, const int* token_ids, float* token_activations, std::int32_t tokens, cudaStream_t stream, const int* adaptive_route_jobs) { - constexpr int kPathBlocks = (kTopK + 1) / PathsPerBlock; - const auto* input = static_cast(x.data); - const auto* routed_codes = static_cast(weights.routed_gate_up.qdata); - const auto* routed_high = static_cast(weights.routed_gate_up.qhigh); - const auto* routed_scales = static_cast(weights.routed_gate_up.scales); - const auto* shared_codes = static_cast(weights.shared_gate_up.qdata); - const auto* shared_scales = static_cast(weights.shared_gate_up.scales); + constexpr int kPathBlocks = (kTopK + 1) / PathsPerBlock; + const auto* input = static_cast(x.data); + const SparseMoePlanes routed = matrix_planes(weights.routed_gate_up); + const SparseMoePlanes shared = matrix_planes(weights.shared_gate_up); if constexpr (Adaptive) { - sparse_moe_d3_path_tiled_kernel + sparse_moe_d3_path_tiled_kernel <<>>( - input, token_ids, routed_codes, routed_high, routed_scales, shared_codes, - shared_scales, token_activations, tokens, adaptive_route_jobs); + input, token_ids, routed, shared, token_activations, tokens, adaptive_route_jobs); CUDA_CHECK(cudaGetLastError()); } else { CUDA_CHECK(pdl::launch_dependent( {dim3(kIntermediate, tokens * kPathBlocks), dim3(PathsPerBlock * 32), 0, stream}, - sparse_moe_d3_path_tiled_kernel, input, token_ids, - routed_codes, routed_high, routed_scales, shared_codes, shared_scales, - token_activations, tokens, nullptr)); + sparse_moe_d3_path_tiled_kernel, input, + token_ids, routed, shared, token_activations, tokens, nullptr)); } } -template +template void launch_d3_small_t_codec(const Tensor& x, const SparseMoeWeights& weights, const int* token_ids, float* token_activations, std::int32_t tokens, SparseMoeSmallTD3Schedule schedule, cudaStream_t stream, const int* adaptive_route_jobs) { switch (schedule) { case SparseMoeSmallTD3Schedule::Paths1: - launch_d3_small_t_paths(x, weights, token_ids, token_activations, - tokens, stream, adaptive_route_jobs); + launch_d3_small_t_paths( + x, weights, token_ids, token_activations, tokens, stream, adaptive_route_jobs); return; case SparseMoeSmallTD3Schedule::Paths3: - launch_d3_small_t_paths(x, weights, token_ids, token_activations, - tokens, stream, adaptive_route_jobs); + launch_d3_small_t_paths( + x, weights, token_ids, token_activations, tokens, stream, adaptive_route_jobs); return; case SparseMoeSmallTD3Schedule::Paths9: - launch_d3_small_t_paths(x, weights, token_ids, token_activations, - tokens, stream, adaptive_route_jobs); + launch_d3_small_t_paths( + x, weights, token_ids, token_activations, tokens, stream, adaptive_route_jobs); return; } throw std::logic_error("sparse_moe: unknown small-T D3 schedule"); } -template +template void launch_d4_small_t_rows(const SparseMoeWeights& weights, Tensor& destination, const int* token_ids, const float* token_alpha, const float* shared_scale, const float* token_activations, std::int32_t tokens, cudaStream_t stream, const int* adaptive_route_jobs) { const dim3 grid = Adaptive ? dim3(kAdaptiveD4Blocks) : dim3(kHidden / Rows, tokens); - sparse_moe_d4_token_kernel<<>>( - token_ids, token_alpha, shared_scale, token_activations, - static_cast(weights.routed_down.qdata), - static_cast(weights.routed_down.qhigh), - static_cast(weights.routed_down.scales), - static_cast(weights.shared_down.qdata), - static_cast(weights.shared_down.scales), - static_cast<__nv_bfloat16*>(destination.data), tokens, adaptive_route_jobs); + sparse_moe_d4_token_kernel + <<>>( + token_ids, token_alpha, shared_scale, token_activations, + matrix_planes(weights.routed_down), matrix_planes(weights.shared_down), + static_cast<__nv_bfloat16*>(destination.data), tokens, adaptive_route_jobs); CUDA_CHECK(cudaGetLastError()); } -template +template void launch_d4_small_t_codec(const SparseMoeWeights& weights, Tensor& destination, const int* token_ids, const float* token_alpha, const float* shared_scale, const float* token_activations, @@ -667,19 +665,19 @@ void launch_d4_small_t_codec(const SparseMoeWeights& weights, Tensor& destinatio cudaStream_t stream, const int* adaptive_route_jobs) { switch (schedule) { case SparseMoeSmallTD4Schedule::Rows1: - launch_d4_small_t_rows(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, stream, - adaptive_route_jobs); + launch_d4_small_t_rows( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + stream, adaptive_route_jobs); return; case SparseMoeSmallTD4Schedule::Rows2: - launch_d4_small_t_rows(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, stream, - adaptive_route_jobs); + launch_d4_small_t_rows( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + stream, adaptive_route_jobs); return; case SparseMoeSmallTD4Schedule::Rows4: - launch_d4_small_t_rows(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, stream, - adaptive_route_jobs); + launch_d4_small_t_rows( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + stream, adaptive_route_jobs); return; } throw std::logic_error("sparse_moe: unknown small-T D4 schedule"); @@ -694,16 +692,17 @@ void sparse_moe_decode_launch_d3_small_t(const Tensor& x, const SparseMoeWeights switch (weights.routed_gate_up.qtype) { case QType::Q4_G64_FP16: if (adaptive_route_jobs == nullptr) { - launch_d3_small_t_codec(x, weights, token_ids, token_activations, - tokens, schedule, stream, nullptr); + launch_d3_small_t_codec( + x, weights, token_ids, token_activations, tokens, schedule, stream, nullptr); } else { - launch_d3_small_t_codec(x, weights, token_ids, token_activations, tokens, - schedule, stream, adaptive_route_jobs); + launch_d3_small_t_codec(x, weights, token_ids, + token_activations, tokens, schedule, + stream, adaptive_route_jobs); } return; case QType::Q8_G32_FP16: - launch_d3_small_t_codec(x, weights, token_ids, token_activations, tokens, - schedule, stream, nullptr); + launch_d3_small_t_codec(x, weights, token_ids, token_activations, + tokens, schedule, stream, nullptr); return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D3 codec"); @@ -718,30 +717,30 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor switch (weights.routed_down.qtype) { case QType::Q5_G64_FP16: if (adaptive_route_jobs == nullptr) { - launch_d4_small_t_codec(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, - schedule, stream, nullptr); + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, + tokens, schedule, stream, nullptr); } else { - launch_d4_small_t_codec(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, - schedule, stream, adaptive_route_jobs); + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, + tokens, schedule, stream, adaptive_route_jobs); } return; case QType::Q6_G64_FP16: if (adaptive_route_jobs == nullptr) { - launch_d4_small_t_codec(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, - schedule, stream, nullptr); + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, + tokens, schedule, stream, nullptr); } else { - launch_d4_small_t_codec(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, - schedule, stream, adaptive_route_jobs); + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, + tokens, schedule, stream, adaptive_route_jobs); } return; case QType::Q8_G32_FP16: - launch_d4_small_t_codec(weights, destination, token_ids, token_alpha, - shared_scale, token_activations, tokens, schedule, - stream, nullptr); + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + schedule, stream, nullptr); return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D4 codec"); From a83221a5e15c2e5e483ac34e3fe6148e06ba6e69 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Sat, 19 Sep 2026 09:41:08 +0200 Subject: [PATCH 2/3] feat(artifact): carry one NVFP4 divisor per source matrix in a stacked tensor A tensor object holds one NVFP4 weight divisor, which is right for a plane quantised as one matrix and wrong for a plane assembled from several that were quantised apart. Published NVFP4 checkpoints of MoE models are the second case: each expert matrix carries its own weight_global_scale, so the converter has to refuse the pack. The tensor object gains an optional `divisors`: the number of separately quantised source matrices stacked into the plane, default 1. The payload keeps that many FP32 words immediately after the scale plane, in stacking order, so payload_bytes = weight_divisor_offset + 4 * divisors, and row r of an N-row plane uses word r / (N / divisors). A divisor's share must be a whole number of 128-row scale tiles, so a boundary never splits one. Absent means one word and the present layout exactly, so every existing artifact stays valid and no reader changes behaviour. The count is decided where the parent's shape is chosen, in the recipe plan, rather than being discovered at encode time; the converter refuses it outside NVFP4, and `validate_nvfp4_weight` refuses a plane whose divisor count does not divide its rows. --- bench/ops/quantized_weight.cuh | 28 ++++- docs/maintainer/artifact-container.md | 3 +- docs/maintainer/storage-layouts.md | 13 ++- docs/maintainer/tensor-formats.md | 25 +++-- src/artifact/layouts.cpp | 4 +- src/artifact/materializer.cpp | 24 +++-- src/artifact/schema.cpp | 6 +- src/artifact/schema.h | 3 + src/core/weight.h | 10 ++ src/core/weight_view.cpp | 25 ++++- src/core/weight_view.h | 4 +- src/ops/linear/nvfp4/nvfp4_format.cpp | 30 ++++-- src/ops/linear/nvfp4/nvfp4_format.h | 6 +- tests/convert/test_recipe.py | 149 ++++++++++++++++++++++++-- tests/ops/quantized_weight.h | 15 ++- tools/artifact/layouts.py | 22 +++- tools/artifact/schema.py | 22 +++- tools/artifact/tensor_output.py | 19 ++-- tools/convert/methods.py | 50 +++++++-- tools/convert/recipe.py | 34 +++++- 20 files changed, 403 insertions(+), 89 deletions(-) diff --git a/bench/ops/quantized_weight.cuh b/bench/ops/quantized_weight.cuh index 3e6e1ecad3..ed9aad72ce 100644 --- a/bench/ops/quantized_weight.cuh +++ b/bench/ops/quantized_weight.cuh @@ -10,6 +10,7 @@ #include #include #include +#include namespace ninfer::bench { @@ -163,10 +164,16 @@ inline PackedQuantizedWeight make_row_split_weight(QType qtype, std::int32_t n, return result; } -inline PackedQuantizedWeight make_nvfp4_weight(std::int32_t n, std::int32_t k) { - if (n <= 0 || k <= 0 || (n % 128) != 0 || (k % 64) != 0) { +// `divisor_rows` is how many consecutive rows one stored divisor covers: the whole plane for a +// matrix quantised on its own, and one source matrix's share for a plane stacked from several. +inline PackedQuantizedWeight make_nvfp4_weight(std::int32_t n, std::int32_t k, + std::int32_t divisor_rows = 0) { + if (divisor_rows == 0) { divisor_rows = n; } + if (n <= 0 || k <= 0 || (n % 128) != 0 || (k % 64) != 0 || divisor_rows <= 0 || + (n % divisor_rows) != 0 || (divisor_rows % 128) != 0) { throw std::invalid_argument("invalid benchmark NVFP4 weight shape"); } + const std::int32_t divisors = n / divisor_rows; const std::uint64_t elements = detail::checked_mul(static_cast(n), static_cast(k), "benchmark NVFP4 element count overflow"); @@ -176,7 +183,8 @@ inline PackedQuantizedWeight make_nvfp4_weight(std::int32_t n, std::int32_t k) { const std::uint64_t divisor_offset = detail::checked_add(scale_offset, scale_bytes, "benchmark NVFP4 divisor offset overflow"); const std::uint64_t payload_bytes = - detail::checked_add(divisor_offset, sizeof(float), "benchmark NVFP4 payload size overflow"); + detail::checked_add(divisor_offset, static_cast(divisors) * sizeof(float), + "benchmark NVFP4 payload size overflow"); if (payload_bytes > std::numeric_limits::max()) { throw std::overflow_error("benchmark NVFP4 payload does not fit size_t"); } @@ -194,9 +202,17 @@ inline PackedQuantizedWeight make_nvfp4_weight(std::int32_t n, std::int32_t k) { CUDA_CHECK(cudaMemset(result.storage.p, 0x22, code_bytes)); CUDA_CHECK( cudaMemset(static_cast(result.storage.p) + scale_offset, 0x38, scale_bytes)); + // A published checkpoint quantises every source matrix on its own, so the divisors differ. + // Equal ones would let a per-row lookup pass while reading only the first word. constexpr float kWeightDivisor = 0.125F; + std::vector divisor_words(static_cast(divisors)); + for (std::int32_t index = 0; index < divisors; ++index) { + divisor_words[static_cast(index)] = + kWeightDivisor * (1.0F + 0.75F * static_cast(index % 5)); + } CUDA_CHECK(cudaMemcpy(static_cast(result.storage.p) + divisor_offset, - &kWeightDivisor, sizeof(kWeightDivisor), cudaMemcpyHostToDevice)); + divisor_words.data(), divisor_words.size() * sizeof(float), + cudaMemcpyHostToDevice)); Weight& weight = result.weight; weight.payload = result.storage.p; @@ -216,8 +232,10 @@ inline PackedQuantizedWeight make_nvfp4_weight(std::int32_t n, std::int32_t k) { weight.scales = static_cast(result.storage.p) + scale_offset; weight.n = n; weight.k = k; - weight.weight_scale_divisor = kWeightDivisor; + weight.weight_scale_divisor = divisor_words[0]; weight.input_scale_divisor = 3.5F; + weight.weight_divisors = static_cast(result.storage.p) + divisor_offset; + weight.weight_divisor_rows = divisor_rows; return result; } diff --git a/docs/maintainer/artifact-container.md b/docs/maintainer/artifact-container.md index 44149b7637..f8a7c65381 100644 --- a/docs/maintainer/artifact-container.md +++ b/docs/maintainer/artifact-container.md @@ -237,6 +237,7 @@ Target 引用须能找到对应组件，其数学关联由架构 binder 检查 | layout | ID | 第 6 节的布局 | | offset | U64 | 对象起点在逻辑 payload 中的字节偏移 | | bytes | PositiveU64 | 该对象完整编码的字节数 | +| divisors | PositiveU64 | 可选，缺省为 1。堆叠进本平面且各自独立量化的源矩阵数量；每个源矩阵各持一个 NVFP4 权重除数，按平面行数均分。仅 `nvfp4` 与 `block_scale_k16_m128x4_v1` 允许大于 1 | ```json { @@ -322,7 +323,7 @@ group size、scale 类型和解码规则直接由该 codec 定义。 量化名字末尾的 FP16/BF16 表示 scale 类型。激活计算许可在 uses 中表达。 Code 范围、特殊浮点值、舍入与精确重建按[数值合同](tensor-formats.md)解释。 -尤其是 NVFP4 的重建采用 `code_value * block_scale / weight_divisor`，逐行 FP8 采用其既定的 +尤其是 NVFP4 的重建采用 `code_value * block_scale / weight_divisor[floor(row / (N / divisors))]`，逐行 FP8 采用其既定的 `code_value * row_scale` 重建规则。 同一数值含义更换 encoder 或校准过程时，format 名保持相同，生成方法记录在 recipe/provenance。 diff --git a/docs/maintainer/storage-layouts.md b/docs/maintainer/storage-layouts.md index a1b88a8f62..9880238f21 100644 --- a/docs/maintainer/storage-layouts.md +++ b/docs/maintainer/storage-layouts.md @@ -266,11 +266,16 @@ code_plane_bytes = N * K / 2 scale_plane_offset = align_up(code_plane_bytes, 256) scale_plane_bytes = N * K / 16 weight_divisor_offset = scale_plane_offset + scale_plane_bytes -payload_bytes = weight_divisor_offset + 4 +payload_bytes = weight_divisor_offset + 4 * divisors ``` +`divisors` is the tensor object's member of that name, one for a plane quantized as a single +matrix. A plane stacked from several separately quantized source matrices holds one divisor per +source, in the order the sources are stacked; each covers `N / divisors` consecutive rows, and that +share is a whole number of 128-row scale tiles. + The payload is a row-major E2M1 packed-code plane, zero padding to `scale_plane_offset`, a -swizzled E4M3FN scale plane, and the little-endian FP32 weight-divisor word. Within each packed code +swizzled E4M3FN scale plane, and `divisors` little-endian FP32 weight-divisor words. Within each packed code byte, the low nibble is the smaller K coordinate and the high nibble is the next coordinate. For logical row `n`, scale-group coordinate `g=floor(k/16)`, and `K_tiles=K/64`, define: @@ -292,7 +297,7 @@ The scale word's byte offset within the scale plane is: ``` Layout decoding must recover the original packed E2M1 words, natural `[N,K/16]` E4M3FN scale-word -matrix, and exact divisor word. It never decodes and re-encodes either floating-point format. +matrix, and exact divisor words. It never decodes and re-encodes either floating-point format. ## 5. `row_scale_v1` @@ -344,7 +349,7 @@ Layout decoding yields only persistent logical words: - `row_split_k128_v1` yields the grouped signed codes and binary16 scales for logical columns `0..K-1`, discarding physical columns `K..K_pad-1`; - `block_scale_k16_m128x4_v1` yields the packed E2M1 words, natural E4M3FN group-scale words, and - matrix-level FP32 weight divisor; + FP32 weight divisor of each stacked source matrix; - `row_scale_v1` yields the natural row-major E4M3FN code words and one BF16 multiplier per logical row; - `raw_bytes_v1` yields the enclosing resource bytes. diff --git a/docs/maintainer/tensor-formats.md b/docs/maintainer/tensor-formats.md index 3fbcc1f48b..6fbac1a7ca 100644 --- a/docs/maintainer/tensor-formats.md +++ b/docs/maintainer/tensor-formats.md @@ -31,7 +31,7 @@ The block-scaled floating-point weight format is: | Canonical name | Code | K group | Block scale | Global field | |---|---|---:|---|---| -| `nvfp4` | E2M1, 4 bits/weight | 16 | one E4M3FN word/group | one positive FP32 weight divisor | +| `nvfp4` | E2M1, 4 bits/weight | 16 | one E4M3FN word/group | one positive FP32 weight divisor per stacked source matrix | The row-scaled floating-point weight format is: @@ -89,7 +89,7 @@ may preserve an already encoded source or quantize floating-point values. The built-in `grouped_absmax` method implements the reference encoder in Section 7 for all four grouped integer formats. `fp8_row_maxabs` rounds source values to BF16 and quantizes each row to E4M3FN codes with a BF16 multiplier. `import_encoded` preserves compatible FP8 or NVFP4 codes, -scales, and, for NVFP4, the matrix weight divisor. NInfer currently provides no built-in +scales, and, for NVFP4, the weight divisor of the row's own source matrix. NInfer currently provides no built-in floating-point-to-NVFP4 quantizer. A recipe can supply a Python callable as its method. Different methods can produce different @@ -226,8 +226,11 @@ abbreviations; artifacts store the complete canonical names. `nvfp4` is a block-scaled floating-point weight representation, not a signed-integer `QuantFormat`. For a logical matrix `[N,K]`, every K-axis group contains 16 E2M1 code words and one -E4M3FN scale word. The representation also contains one FP32 serialized weight divisor `d_w` for -the complete matrix. +E4M3FN scale word. The representation also contains one FP32 serialized weight divisor `d_w` per +source matrix stacked into the plane: a plane quantized as one matrix holds one, and a plane +assembled from several matrices that were quantized apart holds one for each, covering +`N / divisors` consecutive rows in the order the sources are stacked. `divisors` is the tensor +object's member of that name. An E2M1 word has sign bit 3, exponent bits 2:1, and mantissa bit 0. Positive code words `0..7` decode to: @@ -250,13 +253,13 @@ e == 15, m == 7: NaN ``` Stored NVFP4 weight scales admit only sign-zero finite words, including positive zero. Negative -values, negative zero, and both NaN words are invalid. The serialized binary32 word `d_w` must be +values, negative zero, and both NaN words are invalid. Every serialized binary32 word of `d_w` must be finite and strictly positive. For code `c[n,k]`, scale word `s[n,g]`, and `g=floor(k/16)`, the exact represented weight is: ```text -W[n,k] = decode_e2m1(c[n,k]) * decode_e4m3fn(s[n,g]) / d_w +W[n,k] = decode_e2m1(c[n,k]) * decode_e4m3fn(s[n,g]) / d_w[floor(n / (N / divisors))] ``` `import_encoded` copies all three fields without requantizing or canonicalizing them. Activation @@ -552,8 +555,8 @@ A conforming producer must: - for a quantized format, preserve the logical shape and last-axis group rule; - for a grouped signed-integer format, emit one valid binary16 scale per logical group and only legal signed codes, including never emitting Q8 `-128`; -- for `nvfp4`, emit only valid E2M1 code words, nonnegative finite E4M3FN scale words, and one finite - positive FP32 weight divisor under Section 3.3; +- for `nvfp4`, emit only valid E2M1 code words, nonnegative finite E4M3FN scale words, and one + finite positive FP32 weight divisor per stacked source matrix under Section 3.3; - for `fp8_e4m3fn_row_bf16`, emit only finite E4M3FN code words and valid BF16 row multipliers, with signed-zero codes as the only legal codes in a positive-zero-scale row under Section 3.4; - record enough conversion provenance for the artifact producer to identify how the values @@ -580,8 +583,8 @@ The `.ninfer` container and each registered storage layout must: - for grouped signed-integer formats, make the number and ownership of logical groups unambiguous and reconstruct every signed code and binary16 scale without inference from a kernel implementation; -- for `nvfp4`, reconstruct every E2M1 code word, natural E4M3FN scale word, and the matrix FP32 - divisor under Section 3.3; +- for `nvfp4`, reconstruct every E2M1 code word, natural E4M3FN scale word, and the FP32 divisor + of the row's own source matrix under Section 3.3; - for `fp8_e4m3fn_row_bf16`, reconstruct every E4M3FN code word and its owning BF16 row multiplier under Section 3.4; - define its canonical physical-padding contents and producer responsibilities, if it materializes @@ -632,7 +635,7 @@ enum spellings or private kernel layout. The retained codec and encoder evidence - Q4, Q5, Q6, and Q8 plane bit order, legal interval endpoints, encoded-size geometry, partial-K zero padding, consecutive row views, and arbitrary row gathers; - all 16 E2M1 words, all 256 E4M3FN words, NVFP4 scale/divisor validity, the exact divisor-based - reconstruction equation, and known block-scale swizzle offsets; + reconstruction equation for a plane of one source, and known block-scale swizzle offsets; - finite E4M3FN weight-code validity, BF16 row-scale validity, signed-zero rows, exact code/scale plane round trips, and the row-multiplier reconstruction equation for `fp8_e4m3fn_row_bf16`; diff --git a/src/artifact/layouts.cpp b/src/artifact/layouts.cpp index 2ecc1950fa..e7b4114ee2 100644 --- a/src/artifact/layouts.cpp +++ b/src/artifact/layouts.cpp @@ -7,8 +7,8 @@ namespace ninfer::artifact { WeightGeometry describe_tensor(const TensorObject& object) { try { - auto geometry = - weight_geometry(parse_format(object.format), parse_layout(object.layout), object.shape); + auto geometry = weight_geometry(parse_format(object.format), parse_layout(object.layout), + object.shape, object.divisors); if (geometry.bytes != object.bytes || object.offset % geometry.alignment) { throw ArtifactError("encoded size or object alignment differs from layout"); } diff --git a/src/artifact/materializer.cpp b/src/artifact/materializer.cpp index 6517af2104..8097f94d60 100644 --- a/src/artifact/materializer.cpp +++ b/src/artifact/materializer.cpp @@ -7,11 +7,11 @@ #include #include -#include #include #include #include #include +#include #include #include @@ -83,21 +83,25 @@ struct ReadSpan { float read_divisor(const Reader& reader, ObjectHandle handle, const WeightGeometry& geometry, std::span host, MaterializationStats& stats) { if (geometry.format != QType::NVFP4) { return 0.0F; } - std::array word{}; + const auto bytes = checked_mul(geometry.divisor_count, sizeof(float), "weight divisors"); + std::vector words(bytes); if (!host.empty()) { - std::copy_n(host.data() + geometry.divisor_offset, word.size(), word.data()); + std::copy_n(host.data() + geometry.divisor_offset, bytes, words.data()); } else { const auto& object = reader.directory().tensor(handle); reader.read_into(checked_add(object.offset, geometry.divisor_offset, "weight divisor"), - word); - stats.read_bytes = checked_add(stats.read_bytes, word.size(), "read bytes"); + words); + stats.read_bytes = checked_add(stats.read_bytes, bytes, "read bytes"); } - const auto value = std::bit_cast(read_u32_le(word.data())); - if (!std::isfinite(value) || value <= 0) { - throw ArtifactError(reader.directory().tensor(handle).id + - ": invalid NVFP4 weight divisor"); + // Every divisor is read by a kernel, so every one is checked before anything binds. + for (std::uint64_t index = 0; index < geometry.divisor_count; ++index) { + const auto value = std::bit_cast(read_u32_le(words.data() + index * sizeof(float))); + if (!std::isfinite(value) || value <= 0) { + throw ArtifactError(reader.directory().tensor(handle).id + + ": invalid NVFP4 weight divisor"); + } } - return value; + return std::bit_cast(read_u32_le(words.data())); } } // namespace diff --git a/src/artifact/schema.cpp b/src/artifact/schema.cpp index 71f7b86d1d..89ae3b40d1 100644 --- a/src/artifact/schema.cpp +++ b/src/artifact/schema.cpp @@ -218,7 +218,7 @@ void parse_objects(Directory& out, const Json& objects) { const auto kind = require_id(value.at("kind"), "object kind"); if (kind == "tensor") { require_members(value, {"id", "kind", "shape", "format", "layout", "offset", "bytes"}, - {}, "tensor"); + {"divisors"}, "tensor"); } else if (kind == "resource") { require_members(value, {"id", "kind", "encoding", "offset", "bytes"}, {}, "resource"); } else { @@ -236,9 +236,11 @@ void parse_objects(Directory& out, const Json& objects) { } previous_end = end; if (kind == "tensor") { + const std::uint64_t divisors = + value.contains("divisors") ? require_u64(value.at("divisors"), id, true) : 1; out.objects.emplace_back(TensorObject{ id, parse_shape(value.at("shape"), id), require_id(value.at("format"), id), - require_id(value.at("layout"), id), offset, bytes}); + require_id(value.at("layout"), id), offset, bytes, divisors}); } else { out.objects.emplace_back( ResourceObject{id, require_id(value.at("encoding"), id), offset, bytes}); diff --git a/src/artifact/schema.h b/src/artifact/schema.h index 518933af79..616c395d46 100644 --- a/src/artifact/schema.h +++ b/src/artifact/schema.h @@ -37,6 +37,9 @@ struct TensorObject { std::string layout; std::uint64_t offset = 0; std::uint64_t bytes = 0; + // Source matrices this plane was assembled from, each with its own NVFP4 divisor. Absent means + // one, which is every plane that is not a stack of separately quantised matrices. + std::uint64_t divisors = 1; }; struct ResourceObject { diff --git a/src/core/weight.h b/src/core/weight.h index 9d68447928..e62b13efd7 100644 --- a/src/core/weight.h +++ b/src/core/weight.h @@ -47,6 +47,16 @@ struct Weight { std::int64_t scale_nb[4] = {0, 0, 0, 0}; float weight_scale_divisor = 0.0F; float input_scale_divisor = 0.0F; + + // An NVFP4 plane assembled from several source matrices carries one divisor per source, in the + // payload after the scales. `weight_divisors` addresses them and `weight_divisor_rows` says how + // many consecutive rows each covers, so the divisor of row r is element r / + // weight_divisor_rows. A plane with one source sets the rows to its own row count, which makes + // that index zero for every row, and `weight_scale_divisor` is then the whole story. On a stack + // `weight_scale_divisor` holds only the first source's word, so a route that reads it for any + // other row is silently wrong by a scale factor. + const void* weight_divisors = nullptr; + std::int32_t weight_divisor_rows = 0; }; } // namespace ninfer diff --git a/src/core/weight_view.cpp b/src/core/weight_view.cpp index a625a24f71..34791b7dc0 100644 --- a/src/core/weight_view.cpp +++ b/src/core/weight_view.cpp @@ -77,10 +77,14 @@ std::uint64_t weight_element_count(std::span shape) { } WeightGeometry weight_geometry(QType format, QuantLayout layout, - std::span shape) { + std::span shape, std::uint64_t divisors) { + if (divisors == 0 || (divisors != 1 && format != QType::NVFP4)) { + throw std::invalid_argument("only an NVFP4 plane carries more than one divisor"); + } WeightGeometry out; - out.format = format; - out.layout = layout; + out.format = format; + out.divisor_count = divisors; + out.layout = layout; out.shape.assign(shape.begin(), shape.end()); out.elements = weight_element_count(shape); if (layout == QuantLayout::Contiguous) { @@ -148,8 +152,17 @@ WeightGeometry weight_geometry(QType format, QuantLayout layout, out.scale_bytes = mul(n, out.scale_bytes_per_row); out.bytes = add(out.scale_offset, out.scale_bytes); if (format == QType::NVFP4) { + if (n % divisors) { + throw std::invalid_argument("divisor count must divide the plane's rows"); + } + // Every other structural NVFP4 rule is checked here rather than left to the operator, and + // this one belongs with them: a divisor that covered part of a 128-row scale tile would + // make two rows of one tile disagree about which word they were quantised against. + if (divisors > 1 && (n / divisors) % 128) { + throw std::invalid_argument("each NVFP4 divisor must cover whole 128-row scale tiles"); + } out.divisor_offset = out.bytes; - out.bytes = add(out.bytes, 4); + out.bytes = add(out.bytes, mul(divisors, 4)); } return out; } @@ -288,7 +301,9 @@ Weight native_weight(const WeightView& view, float input_divisor) { out.scale_nb[0] = 2; out.scale_nb[1] = out.scale_nb[2] = out.scale_nb[3] = static_cast(out.n) * 2; } else if (g.layout == QuantLayout::BlockScaleK16M128x4) { - out.scale_dtype = DType::FP8_E4M3FN; + out.scale_dtype = DType::FP8_E4M3FN; + out.weight_divisors = region.parent->data + g.divisor_offset; + out.weight_divisor_rows = dimension(g.shape[0] / g.divisor_count); } return out; } diff --git a/src/core/weight_view.h b/src/core/weight_view.h index b051957d89..d56f5b9077 100644 --- a/src/core/weight_view.h +++ b/src/core/weight_view.h @@ -47,10 +47,12 @@ struct WeightGeometry { std::uint64_t scale_offset = 0; std::uint64_t scale_bytes = 0; std::uint64_t divisor_offset = 0; + std::uint64_t divisor_count = 1; }; [[nodiscard]] WeightGeometry weight_geometry(QType format, QuantLayout layout, - std::span shape); + std::span shape, + std::uint64_t divisors = 1); [[nodiscard]] std::uint64_t weight_element_count(std::span shape); struct WeightParent { diff --git a/src/ops/linear/nvfp4/nvfp4_format.cpp b/src/ops/linear/nvfp4/nvfp4_format.cpp index 1902a54a55..28993cefc7 100644 --- a/src/ops/linear/nvfp4/nvfp4_format.cpp +++ b/src/ops/linear/nvfp4/nvfp4_format.cpp @@ -36,7 +36,8 @@ std::uint64_t align_up(std::uint64_t value, std::uint64_t alignment, const char* } // namespace -Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* operation) { +Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* operation, + bool stacked) { if (weight.n <= 0 || weight.k <= 0 || (weight.n % 128) != 0 || (weight.k % 64) != 0) { throw std::invalid_argument(std::string(operation) + ": NVFP4 requires N%128=0 and K%64=0"); } @@ -44,12 +45,18 @@ Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* oper const std::uint64_t elements = checked_mul(static_cast(weight.n), static_cast(weight.k), operation); Nvfp4WeightGeometry geometry{}; - geometry.code_plane_bytes = elements / 2; - geometry.scale_plane_offset = align_up(geometry.code_plane_bytes, 256, operation); - geometry.scale_plane_bytes = elements / 16; - geometry.required_payload_bytes = - checked_add(checked_add(geometry.scale_plane_offset, geometry.scale_plane_bytes, operation), - sizeof(float), operation); + geometry.code_plane_bytes = elements / 2; + geometry.scale_plane_offset = align_up(geometry.code_plane_bytes, 256, operation); + geometry.scale_plane_bytes = elements / 16; + const std::int32_t rows_per_divisor = weight.weight_divisor_rows; + const bool divisor_shape = rows_per_divisor > 0 && (weight.n % rows_per_divisor) == 0 && + (rows_per_divisor % 128) == 0 && + (stacked || rows_per_divisor == weight.n); + geometry.required_payload_bytes = checked_add( + checked_add(geometry.scale_plane_offset, geometry.scale_plane_bytes, operation), + checked_mul(divisor_shape ? static_cast(weight.n / rows_per_divisor) : 1, + sizeof(float), operation), + operation); if (weight.qtype != QType::NVFP4 || weight.layout != QuantLayout::BlockScaleK16M128x4 || weight.scale_dtype != DType::FP8_E4M3FN || weight.group_size != 16 || weight.group != 16 || @@ -57,10 +64,11 @@ Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* oper weight.padded_shape[0] != weight.n || weight.padded_shape[1] != weight.k || weight.payload == nullptr || weight.qdata == nullptr || weight.scales == nullptr || weight.qhigh != nullptr || weight.high_plane_bytes != 0 || - weight.payload_bytes < geometry.required_payload_bytes || - !std::isfinite(weight.weight_scale_divisor) || weight.weight_scale_divisor <= 0.0F || - !std::isfinite(weight.input_scale_divisor) || weight.input_scale_divisor <= 0.0F || - !aligned_to(weight.qdata, 16) || !aligned_to(weight.scales, 16)) { + weight.payload_bytes < geometry.required_payload_bytes || !divisor_shape || + weight.weight_divisors == nullptr || !std::isfinite(weight.weight_scale_divisor) || + weight.weight_scale_divisor <= 0.0F || !std::isfinite(weight.input_scale_divisor) || + weight.input_scale_divisor <= 0.0F || !aligned_to(weight.qdata, 16) || + !aligned_to(weight.scales, 16)) { throw std::invalid_argument(std::string(operation) + ": invalid NVFP4 weight"); } diff --git a/src/ops/linear/nvfp4/nvfp4_format.h b/src/ops/linear/nvfp4/nvfp4_format.h index feb0c21368..43c17bc389 100644 --- a/src/ops/linear/nvfp4/nvfp4_format.h +++ b/src/ops/linear/nvfp4/nvfp4_format.h @@ -14,6 +14,10 @@ struct Nvfp4WeightGeometry { std::uint64_t required_payload_bytes; }; -Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* operation); +// `stacked` admits a plane assembled from several separately quantised matrices, which carries a +// divisor per source and is read per row. Dense routes leave it false: their kernels read the one +// scalar, so a plane with more than one divisor is not something they can execute. +Nvfp4WeightGeometry validate_nvfp4_weight(const Weight& weight, const char* operation, + bool stacked = false); } // namespace ninfer::ops::detail diff --git a/tests/convert/test_recipe.py b/tests/convert/test_recipe.py index 0429084f24..250148a31d 100644 --- a/tests/convert/test_recipe.py +++ b/tests/convert/test_recipe.py @@ -9,6 +9,7 @@ from tools.artifact.reader import Artifact from tools.artifact.codecs.row_split import decode_row_split_codes from tools.artifact.codecs.nvfp4 import encode_nvfp4 +from tools.artifact.layouts import block_scale_geometry, encoded_size from tools.artifact.schema import binding_parts from tools.artifact.writer import ArtifactWriter from tools.convert.methods import grouped_absmax, import_encoded @@ -154,12 +155,15 @@ def test_shared_weight_keeps_use_independent_and_can_be_overridden(): assert independent.bindings["key"] != independent.bindings["context_key"] -def _encoded_source(name, divisor, shift=0): +def _encoded_source(name, divisor, shift=0, rows=128, activation=None): codes = ( - (torch.arange(128 * 32) + shift).remainder(256).to(torch.uint8).reshape(128, 32) + (torch.arange(rows * 32) + shift) + .remainder(256) + .to(torch.uint8) + .reshape(rows, 32) ) scales = ( - (torch.arange(128 * 4) + shift).remainder(127).to(torch.uint8).reshape(128, 4) + (torch.arange(rows * 4) + shift).remainder(127).to(torch.uint8).reshape(rows, 4) ) word = struct.pack("(device_payload) + high_plane_offset; w.scales = static_cast(device_payload) + scale_plane_offset; + if (w.qtype == QType::NVFP4) { + w.weight_divisors = + static_cast(device_payload) + weight_divisor_offset; + } } else { - w.qdata = nullptr; - w.qhigh = nullptr; - w.scales = nullptr; + w.qdata = nullptr; + w.qhigh = nullptr; + w.scales = nullptr; + w.weight_divisors = nullptr; } return w; } @@ -293,6 +298,8 @@ struct PackedWeight { w.n = row_count; w.shape[0] = row_count; w.padded_shape[0] = row_count; + // A view is one source's worth of rows, so its divisor covers all of them. + w.weight_divisor_rows = row_count; return w; } }; @@ -483,6 +490,8 @@ inline PackedWeight make_patterned_weight(QType qtype, std::int32_t n, std::int3 packed.weight.k = k; packed.weight.weight_scale_divisor = options.weight_scale_divisor; packed.weight.input_scale_divisor = options.input_scale_divisor; + packed.weight.weight_divisors = packed.payload.data() + packed.weight_divisor_offset; + packed.weight.weight_divisor_rows = n; return packed; } if (options.weight_scale_divisor != 0.0F || options.input_scale_divisor != 0.0F) { diff --git a/tools/artifact/layouts.py b/tools/artifact/layouts.py index 4c59059c6f..9ecd5c3585 100644 --- a/tools/artifact/layouts.py +++ b/tools/artifact/layouts.py @@ -64,6 +64,7 @@ class BlockScaleGeometry: scale_plane_offset: int scale_plane_bytes: int weight_divisor_offset: int + weight_divisor_count: int payload_bytes: int @@ -198,7 +199,7 @@ def row_split_geometry( def block_scale_geometry( - format: str | Nvfp4Format, shape: Sequence[int] + format: str | Nvfp4Format, shape: Sequence[int], divisors: int = 1 ) -> BlockScaleGeometry: spec = _format(format) if not isinstance(spec, Nvfp4Format): @@ -209,6 +210,15 @@ def block_scale_geometry( "block_scale_k16_m128x4_v1 requires N divisible by 128 " "and K divisible by 64" ) + if divisors < 1 or n % divisors: + raise ValueError("NVFP4 divisor count must divide the plane's rows") + # Each divisor covers one source matrix, and a 128-row scale tile may not span two of them: + # the swizzled scale plane is addressed in whole tiles and the engine takes the divisor from + # the row. + if divisors > 1 and (n // divisors) % 128: + raise ValueError( + "each NVFP4 divisor must cover a whole number of 128-row scale tiles" + ) code_plane_bytes = n * k // 2 scale_plane_offset = align_up(code_plane_bytes, PLANE_ALIGNMENT) scale_plane_bytes = n * k // spec.group_size @@ -222,7 +232,8 @@ def block_scale_geometry( scale_plane_offset=scale_plane_offset, scale_plane_bytes=scale_plane_bytes, weight_divisor_offset=weight_divisor_offset, - payload_bytes=weight_divisor_offset + 4, + weight_divisor_count=divisors, + payload_bytes=weight_divisor_offset + 4 * divisors, ) @@ -250,9 +261,14 @@ def encoded_size( layout: str | Layout, format: str | NumericFormat, shape: Sequence[int], + divisors: int = 1, ) -> int: layout_spec = _layout(layout) numeric_spec = _format(format) + if divisors != 1 and layout_spec is not BLOCK_SCALE_K16_M128X4_V1: + raise ValueError( + f"layout {layout_spec.name!r} stores one divisor, not {divisors}" + ) if numeric_spec.name not in layout_spec.formats: raise ValueError( f"layout {layout_spec.name!r} does not accept format {numeric_spec.name!r}" @@ -271,7 +287,7 @@ def encoded_size( if layout_spec is BLOCK_SCALE_K16_M128X4_V1: if not isinstance(numeric_spec, Nvfp4Format): raise ValueError("block_scale_k16_m128x4_v1 requires NVFP4") - return block_scale_geometry(numeric_spec, shape).payload_bytes + return block_scale_geometry(numeric_spec, shape, divisors).payload_bytes if layout_spec is ROW_SCALE_V1: if not isinstance(numeric_spec, Fp8RowFormat): raise ValueError("row_scale_v1 requires a row-scaled FP8 format") diff --git a/tools/artifact/schema.py b/tools/artifact/schema.py index b7445f1364..726a2a34a3 100644 --- a/tools/artifact/schema.py +++ b/tools/artifact/schema.py @@ -66,6 +66,9 @@ class TensorSpec: shape: tuple[int, ...] format: str layout: str + # Source matrices stacked into this parent, each with its own NVFP4 divisor. One means the + # parent has a single divisor, which is every parent that is not such a stack. + divisors: int = 1 @dataclass(frozen=True, slots=True) @@ -83,10 +86,16 @@ class TensorObject: layout: str offset: int bytes: int + divisors: int = 1 kind: ClassVar[str] = "tensor" def to_json(self) -> dict: - return {**asdict(self), "shape": list(self.shape), "kind": self.kind} + out = {**asdict(self), "shape": list(self.shape), "kind": self.kind} + # The common parent has one divisor; saying so on every one of a hundred thousand + # objects would be a megabyte of directory that means nothing. + if self.divisors == 1: + del out["divisors"] + return out @dataclass(frozen=True, slots=True) @@ -146,7 +155,7 @@ def validate_encoding(obj: ArtifactObject) -> None: ) return try: - size = encoded_size(obj.layout, obj.format, obj.shape) + size = encoded_size(obj.layout, obj.format, obj.shape, obj.divisors) alignment = get_layout(obj.layout).alignment except (TypeError, ValueError) as error: raise ArtifactError(f"{obj.id}: {error}") from error @@ -171,10 +180,12 @@ def plan_objects(specs: Sequence[ObjectSpec]) -> tuple[ArtifactObject, ...]: dims = shape(spec.shape, name) try: start = align_up(offset, get_layout(spec.layout).alignment) - size = encoded_size(spec.layout, spec.format, dims) + size = encoded_size(spec.layout, spec.format, dims, spec.divisors) except (TypeError, ValueError) as error: raise ArtifactError(f"{name}: {error}") from error - obj = TensorObject(name, dims, spec.format, spec.layout, start, size) + obj = TensorObject( + name, dims, spec.format, spec.layout, start, size, spec.divisors + ) elif isinstance(spec, ResourceSpec): obj = ResourceObject(name, spec.encoding, offset, spec.bytes) else: @@ -196,7 +207,7 @@ def _parse_object(value: object) -> ArtifactObject: common = {"id", "kind", "offset", "bytes"} kind = value.get("kind") if kind == "tensor": - members(value, common | {"shape", "format", "layout"}, set(), "tensor") + members(value, common | {"shape", "format", "layout"}, {"divisors"}, "tensor") elif kind == "resource": members(value, common | {"encoding"}, set(), "resource") else: @@ -214,6 +225,7 @@ def _parse_object(value: object) -> ArtifactObject: identifier(value["layout"], name), offset, size, + integer(value.get("divisors", 1), f"{name} divisors", positive=True), ) diff --git a/tools/artifact/tensor_output.py b/tools/artifact/tensor_output.py index bb67761dbd..aa66296651 100644 --- a/tools/artifact/tensor_output.py +++ b/tools/artifact/tensor_output.py @@ -37,7 +37,7 @@ def __init__(self, writer: ArtifactWriter, object_id: str): self.object = obj self.format = get_format(obj.format) self._padding_initialized = False - self._divisor: bytes | None = None + self._divisors: list[bytes | None] = [None] * obj.divisors def write_bytes(self, offset: int, data: bytes | memoryview) -> None: self.writer.write_region(self.object.id, offset, data) @@ -110,7 +110,7 @@ def write_codes( raise ValueError( f"{obj.id}: NVFP4 output needs whole 128-row tiles and weight divisor" ) - g = block_scale_geometry(self.format, obj.shape) + g = block_scale_geometry(self.format, obj.shape, len(self._divisors)) local = block_scale_geometry(self.format, (rows, k)) block = memoryview(encode_nvfp4(codes, scales, weight_divisor, (rows, k))) self.write_bytes(row_begin * (k // 2), block[: local.code_plane_bytes]) @@ -118,10 +118,17 @@ def write_codes( g.scale_plane_offset + row_begin * (k // 16), block[local.scale_plane_offset : local.weight_divisor_offset], ) - if self._divisor is None: - self.write_bytes(g.weight_divisor_offset, weight_divisor) - self._divisor = bytes(weight_divisor) - elif self._divisor != weight_divisor: + # A row block belongs to exactly one source, so it names exactly one divisor. + rows_per_divisor = g.n // len(self._divisors) + index, offset_in_source = divmod(row_begin, rows_per_divisor) + if offset_in_source + rows > rows_per_divisor: + raise ValueError( + f"{obj.id}: a row block may not span two separately quantised sources" + ) + if self._divisors[index] is None: + self.write_bytes(g.weight_divisor_offset + index * 4, weight_divisor) + self._divisors[index] = bytes(weight_divisor) + elif self._divisors[index] != weight_divisor: raise ValueError(f"{obj.id}: weight divisor changed between row blocks") else: raise TypeError(f"{obj.id}: direct format does not accept quantized codes") diff --git a/tools/convert/methods.py b/tools/convert/methods.py index 02c07d1a2d..557f20960d 100644 --- a/tools/convert/methods.py +++ b/tools/convert/methods.py @@ -249,7 +249,26 @@ def import_encoded(request: PrepareRequest) -> PreparedMethod: raise ValueError("import_encoded requires a known encoded matrix target") _preflight(request, values=False) auxiliaries = {} - weight_divisor = None + # The sources of one parent read one activation tensor, which is quantised once, and the + # consumer holds one `input_scale_divisor` for the whole plane - so exactly one of their + # calibrated divisors can survive. It cancels in the GEMM's alpha, so the choice only decides + # where a block scale lands on the e4m3 grid; the smallest is taken, the one direction that + # cannot saturate another source's blocks upward. + # + # This is decided by the number of sources, not by `divisors`: the two are calibrated apart, so + # sources that agree on their weight divisor - which collapses `divisors` to one - can still + # disagree here, and keeping both words has the bank refused at bind. + activation_divisor = None + if request.target.format == "nvfp4" and len(request.inputs) > 1: + words = [ + item.source.input_divisor() + for item in request.inputs + if item.source.input_divisor is not None + ] + if words: + activation_divisor = min( + words, key=lambda word: struct.unpack(" PreparedMethod: f"{item.parameter}: source {first.format} differs from target {request.target.format}" ) if first.format == "nvfp4": - if weight_divisor is None: - weight_divisor = first.weight_divisor - elif first.weight_divisor != weight_divisor: - raise ValueError( - "NVFP4 weight divisors differ; choose separate parents or a conversion method" - ) for parameter, input_name in item.uses: key = (parameter, input_name, "activation_input_divisor") if key in request.auxiliary_overrides: @@ -276,7 +289,9 @@ def import_encoded(request: PrepareRequest) -> PreparedMethod: f"{parameter}: supply an activation divisor for AllowA4" ) auxiliaries[key] = AuxiliaryValue.activation_divisor( - source.input_divisor() + activation_divisor + if activation_divisor is not None + else source.input_divisor() ) n = request.target.shape[0] chunk = ( @@ -286,9 +301,22 @@ def import_encoded(request: PrepareRequest) -> PreparedMethod: ) def produce(output): - for begin in range(0, n, chunk): - words = request.encoded_rows(begin, min(n, begin + chunk)) - output.write_codes(begin, words.codes, words.scales, words.weight_divisor) + # A stacked plane starts a new run of chunks at every source boundary, because a chunk that + # straddled two sources would carry two divisors and a row block carries one. A plane of one + # source keeps the single run it always had, so its chunking is untouched. + bounds = [n] + if request.target.divisors > 1: + bounds = [] + edge = 0 + for item in request.inputs: + edge += item.source.shape[0] + bounds.append(edge) + cursor = 0 + for edge in bounds: + for begin in range(cursor, edge, chunk): + words = request.encoded_rows(begin, min(edge, begin + chunk)) + output.write_codes(begin, words.codes, words.scales, words.weight_divisor) + cursor = edge return request.job(produce=produce, auxiliaries=auxiliaries) diff --git a/tools/convert/recipe.py b/tools/convert/recipe.py index 46181454ba..0db60c22a1 100644 --- a/tools/convert/recipe.py +++ b/tools/convert/recipe.py @@ -303,6 +303,35 @@ def parent_shape(items, chosen): "grouped inputs need compatible rows or an explicit parent shape" ) + def parent_divisors(sources, selection) -> int: + """How many NVFP4 divisors this parent stores: one per source, or one for all. + + A parent whose sources were quantised apart has to keep each source's divisor, and the + stored plane addresses them by an equal share of its rows. That is a partition only if + the sources have equal row counts, so a parent that needs several divisors and cannot be + divided equally is refused here, where its shape is chosen, rather than part way through + writing it. + """ + if selection.format != "nvfp4" or len(sources) < 2: + return 1 + words = [ + source.weight_divisor() if source.weight_divisor else None + for source in sources + ] + # A stack is only chosen when the sources are known to disagree. A divisor no source + # states at plan time is unknown, not different, and answering the source count there + # would mark a plane a stack that no dense route can execute. + if None in words or len(set(words)) == 1: + return 1 + if any(len(source.shape) != 2 for source in sources) or ( + len({source.shape[0] for source in sources}) != 1 + ): + raise ValueError( + "NVFP4 sources quantised against different divisors must be complete " + "matrices with equal row counts" + ) + return len(sources) + def emit(items, chosen=None): if len({self.model.parameters[name].residency for name, _ in items}) > 1: raise ValueError( @@ -315,9 +344,10 @@ def emit(items, chosen=None): dims, sources = parent_shape(items, chosen) selection = items[0][1] layout = selection.layout or default_layout(selection.format) - encoded_size(layout, selection.format, dims) + divisors = parent_divisors(sources, selection) + encoded_size(layout, selection.format, dims, divisors) spec = TensorSpec( - f"weight/{len(weights):06d}", dims, selection.format, layout + f"weight/{len(weights):06d}", dims, selection.format, layout, divisors ) inputs = tuple( MethodInput( From a7d1d5211a5f9e8cd3cf243219fda673bab0602b Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Sat, 19 Sep 2026 09:41:08 +0200 Subject: [PATCH 3/3] feat(ops): add an NVFP4 expert profile to sparse_moe The operation admitted expert banks as Q4+Q5/Q6 or Q8+Q8 and nothing else, so an artifact whose experts are NVFP4 had no route at any token count. This registers the profile across all three: decode, small-token and prefill. Decode reads NVFP4 against represented BF16 activations. The kernels are CUDA-core dot products, so four-bit activations buy no arithmetic there and would cost a quantiser of their own; keeping the activation represented puts rel_l2 at T=1 at 1.60e-3 instead of 1.70e-1. The D4 block takes two hidden rows, because the scale plane puts rows r and r+1 in one 32-byte sector. Prefill quantises the chunk once and runs W4A4 over the existing work list, and the gate/up epilogue writes the SwiGLU intermediate already encoded, so `down` needs no separate quantiser. The per-source divisor is applied in that epilogue rather than inside the MMA loop, by a shift where the stride is a power of two. The frontier between the small-token and prefill routes is 13 for this profile, swept cold over [2,46] against its own routes rather than inherited from the groupwise ones, and the small-token schedule is chosen by the same sweep. The staging buffers the route needs are charged to this profile alone. Three launches written for it - the small-token router, the fused selection/scan/index block and a narrower route-job bound - are gated on it for the same reason: they would otherwise change what a groupwise call runs on the ragged tail of a sliced chunk. The wrapper also requires the routed and shared banks to agree on their activation divisor, because the route encodes the chunk once and both gate/up GEMMs read that one plane. The operator test walks the profile at T = 1, 2, 12, 13, 64, 768, 4097, with a separate tolerance either side of the frontier, and its fixture stacks gate and up under divisors that differ, as the published checkpoint does. --- bench/ops/sparse_moe_bench.cu | 76 +- src/ops/linear/nvfp4/nvfp4_codec.cuh | 58 +- src/ops/linear/nvfp4/nvfp4_gemv.cuh | 9 +- src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh | 77 +- .../decode/sparse_moe_decode_kernels.cu | 121 ++- .../decode/sparse_moe_decode_plan.cpp | 6 +- .../sparse_moe/prefill/sparse_moe_prefill.h | 62 +- .../prefill/sparse_moe_prefill_kernels.cu | 749 +++++++++++++++++- .../prefill/sparse_moe_prefill_plan.cpp | 17 +- .../small_t/sparse_moe_small_t_plan.cpp | 15 +- src/ops/wrapper/sparse_moe.cpp | 60 +- tests/ops/test_sparse_moe.cpp | 404 +++++++++- 12 files changed, 1515 insertions(+), 139 deletions(-) diff --git a/bench/ops/sparse_moe_bench.cu b/bench/ops/sparse_moe_bench.cu index 98ffd21e52..a48e855ce7 100644 --- a/bench/ops/sparse_moe_bench.cu +++ b/bench/ops/sparse_moe_bench.cu @@ -46,6 +46,7 @@ enum class CodecProfile : std::uint8_t { Q4Q5, Q4Q6, Q8Q8, + Nvfp4, }; enum class ExpertDistribution : std::uint8_t { @@ -123,12 +124,23 @@ const char* codec_name(CodecProfile profile) { return "q4-q6"; case CodecProfile::Q8Q8: return "q8-q8"; + case CodecProfile::Nvfp4: + return "nvfp4"; } return "unknown"; } QType gate_codec(CodecProfile profile) { - return profile == CodecProfile::Q8Q8 ? QType::Q8_G32_FP16 : QType::Q4_G64_FP16; + switch (profile) { + case CodecProfile::Q4Q5: + case CodecProfile::Q4Q6: + return QType::Q4_G64_FP16; + case CodecProfile::Q8Q8: + return QType::Q8_G32_FP16; + case CodecProfile::Nvfp4: + return QType::NVFP4; + } + throw std::logic_error("unknown SparseMoe codec profile"); } QType down_codec(CodecProfile profile) { @@ -139,10 +151,18 @@ QType down_codec(CodecProfile profile) { return QType::Q6_G64_FP16; case CodecProfile::Q8Q8: return QType::Q8_G32_FP16; + case CodecProfile::Nvfp4: + return QType::NVFP4; } throw std::logic_error("unknown SparseMoe codec profile"); } +// The shared expert follows the routed codec only where a profile stores all four matrices one +// way. The row-split profiles keep it at Q8 whatever they route with. +QType shared_codec(CodecProfile profile) { + return profile == CodecProfile::Nvfp4 ? QType::NVFP4 : QType::Q8_G32_FP16; +} + const char* distribution_name(ExpertDistribution distribution) { switch (distribution) { case ExpertDistribution::TraceLike: @@ -170,6 +190,11 @@ const char* execution_name(Execution execution) { const char* cache_name(CacheState cache) { return cache == CacheState::Cold ? "cold" : "warm"; } std::uint64_t packed_weight_bytes(QType qtype, std::int32_t rows, std::int32_t columns) { + if (qtype == QType::NVFP4) { + // Four bits of code plus one e4m3 byte per sixteen values. + const std::uint64_t elements = static_cast(rows) * columns; + return elements / 2 + elements / 16; + } const std::int32_t group = qtype == QType::Q8_G32_FP16 ? 32 : 64; const std::uint64_t groups = static_cast(rows) * columns / group; const std::uint64_t low = qtype == QType::Q8_G32_FP16 @@ -182,9 +207,10 @@ std::uint64_t packed_weight_bytes(QType qtype, std::int32_t rows, std::int32_t c } double unique_weight_bytes(const Result& result) { + const QType shared = shared_codec(result.codec); const std::uint64_t fixed = static_cast(kRouterRows) * kHidden * 2 + - packed_weight_bytes(QType::Q8_G32_FP16, 1024, kHidden) + - packed_weight_bytes(QType::Q8_G32_FP16, kHidden, kIntermediate); + packed_weight_bytes(shared, 1024, kHidden) + + packed_weight_bytes(shared, kHidden, kIntermediate); const std::uint64_t per_expert = packed_weight_bytes(gate_codec(result.codec), 1024, kHidden) + packed_weight_bytes(down_codec(result.codec), kHidden, kIntermediate); @@ -286,7 +312,8 @@ void usage(const char* argv0) { std::fprintf(stderr, "Usage: %s [options]\n\n" "Public workload:\n" - " --codec q4-q5|q4-q6|q8-q8|all Routed weight profile (default q4-q5).\n" + " --codec q4-q5|q4-q6|q8-q8|nvfp4|all Expert weight profile (default " + "q4-q5).\n" " --tokens T Exact token extent (default 1).\n" " --sweep START:END[:STEP] Public token-extent sweep.\n" " --distribution trace-like|independent|same\n" @@ -356,8 +383,8 @@ Options parse_options(int argc, char** argv) { throw std::invalid_argument("--tokens and --sweep are mutually exclusive"); } if (options.codec != "q4-q5" && options.codec != "q4-q6" && options.codec != "q8-q8" && - options.codec != "all") { - throw std::invalid_argument("--codec must be q4-q5, q4-q6, q8-q8, or all"); + options.codec != "nvfp4" && options.codec != "all") { + throw std::invalid_argument("--codec must be q4-q5, q4-q6, q8-q8, nvfp4, or all"); } if (options.repeat <= 0) { throw std::invalid_argument("--repeat must be positive"); } if (options.flush_bytes > std::numeric_limits::max()) { @@ -367,10 +394,13 @@ Options parse_options(int argc, char** argv) { } std::vector selected_profiles(const std::string& codec) { - if (codec == "all") { return {CodecProfile::Q4Q5, CodecProfile::Q4Q6, CodecProfile::Q8Q8}; } + if (codec == "all") { + return {CodecProfile::Q4Q5, CodecProfile::Q4Q6, CodecProfile::Q8Q8, CodecProfile::Nvfp4}; + } if (codec == "q4-q5") return {CodecProfile::Q4Q5}; if (codec == "q4-q6") return {CodecProfile::Q4Q6}; - return {CodecProfile::Q8Q8}; + if (codec == "q8-q8") return {CodecProfile::Q8Q8}; + return {CodecProfile::Nvfp4}; } std::vector selected_tokens(const TokenSweep& sweep) { @@ -501,20 +531,30 @@ Weight dense_weight(void* data, std::int32_t rows, std::int32_t columns) { return result; } +// `divisor_rows` is the artifact's own stride for that bank: gate and up are quantised apart, so +// the routed gate/up plane carries one divisor per 512 rows, and routed down one per 2048. A shared +// bank is one matrix and keeps a single divisor. Only NVFP4 stores divisors at all. +bench::PackedQuantizedWeight make_expert_plane(QType qtype, std::int32_t n, std::int32_t k, + bench::QuantizedWeightFill fill, + std::int32_t divisor_rows = 0) { + return qtype == QType::NVFP4 ? bench::make_nvfp4_weight(n, k, divisor_rows) + : bench::make_row_split_weight(qtype, n, k, k, fill); +} + class BenchmarkWeights { public: BenchmarkWeights(CodecProfile profile, std::uint32_t seed, std::size_t flush_bytes) : router_(static_cast(kRouterRows) * kHidden * 2), - routed_gate_(bench::make_row_split_weight( - gate_codec(profile), kExperts * 1024, kHidden, kHidden, - {static_cast(0x31U ^ seed), 0xa5, 0x1401})), - routed_down_(bench::make_row_split_weight( - down_codec(profile), kExperts * kHidden, kIntermediate, kIntermediate, - {static_cast(0x59U ^ (seed >> 8)), 0x6d, 0x1403})), - shared_gate_(bench::make_row_split_weight(QType::Q8_G32_FP16, 1024, kHidden, kHidden, - {0x27, 0x00, 0x1405})), - shared_down_(bench::make_row_split_weight(QType::Q8_G32_FP16, kHidden, kIntermediate, - kIntermediate, {0x73, 0x00, 0x1407})), + routed_gate_(make_expert_plane(gate_codec(profile), kExperts * 1024, kHidden, + {static_cast(0x31U ^ seed), 0xa5, 0x1401}, + 512)), + routed_down_(make_expert_plane( + down_codec(profile), kExperts * kHidden, kIntermediate, + {static_cast(0x59U ^ (seed >> 8)), 0x6d, 0x1403}, kIntermediate * 4)), + shared_gate_( + make_expert_plane(shared_codec(profile), 1024, kHidden, {0x27, 0x00, 0x1405})), + shared_down_(make_expert_plane(shared_codec(profile), kHidden, kIntermediate, + {0x73, 0x00, 0x1407})), flush_(flush_bytes) { std::vector router(static_cast(kRouterRows) * kHidden, bench::f32_to_bf16(0.0F)); diff --git a/src/ops/linear/nvfp4/nvfp4_codec.cuh b/src/ops/linear/nvfp4/nvfp4_codec.cuh index b56419b066..9db06234f4 100644 --- a/src/ops/linear/nvfp4/nvfp4_codec.cuh +++ b/src/ops/linear/nvfp4/nvfp4_codec.cuh @@ -23,6 +23,20 @@ __device__ __forceinline__ float decode_nvfp4_e4m3(std::uint8_t storage) { return static_cast(value).x; } +// One e4m3 block scale in the K16M128x4 layout. A 512-byte tile holds 128 rows by four groups, +// with the rows interleaved so that four rows 32 apart share one 16-byte line. `ScaleTilesPerRow` +// is the matrix's column count over 64; a caller that owns rows of a larger plane than a +// registered linear shape supplies it directly. +template +__device__ __forceinline__ std::int64_t nvfp4_scale_byte(int row, int group) { + const int m_tile = row / 128; + const int row_inner = row - m_tile * 128; + const int scale_tile = group / 4; + const int scale_lane = group & 3; + return (static_cast(m_tile) * ScaleTilesPerRow + scale_tile) * 512 + + (row_inner & 31) * 16 + (row_inner >> 5) * 4 + scale_lane; +} + struct alignas(8) Nvfp4QuantizedK16 { std::uint32_t codes_lo; std::uint32_t codes_hi; @@ -60,14 +74,34 @@ pack_nvfp4_e2m1x16(const float2 (&values)[8], std::uint32_t& codes_lo, std::uint "f"(values[6].x), "f"(values[6].y), "f"(values[7].x), "f"(values[7].y)); } -__device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloat16* source, - float input_scale_divisor) { - const uint4 packed0 = load_vec(source); - const uint4 packed1 = load_vec(source + 8); - const std::uint32_t represented[8] = { - packed0.x, packed0.y, packed0.z, packed0.w, packed1.x, packed1.y, packed1.z, packed1.w, - }; +// Half a group, packed by one lane. A fused epilogue splits the sixteen values of a group across +// two neighbouring lanes so that every thread of the block takes part. +__device__ __forceinline__ std::uint32_t pack_nvfp4_e2m1x8(const float2 (&values)[4]) { + std::uint32_t codes = 0; + asm volatile("{\n" + ".reg .b8 b0;\n" + ".reg .b8 b1;\n" + ".reg .b8 b2;\n" + ".reg .b8 b3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + "mov.b32 %0, {b0,b1,b2,b3};\n" + "}\n" + : "=r"(codes) + : "f"(values[0].x), "f"(values[0].y), "f"(values[1].x), "f"(values[1].y), + "f"(values[2].x), "f"(values[2].y), "f"(values[3].x), "f"(values[3].y)); + return codes; +} +// Sixteen bf16 values already in registers, for a caller that holds them there rather than writing +// them out for a separate pass to read back. Note that this divides where the sparse-MoE fused +// epilogue multiplies by a reciprocal, so the two can differ by one code at an exact tie. +// Nothing in the engine compares the two planes: each route is checked end to end against an +// oracle that quantises neither. +__device__ __forceinline__ Nvfp4QuantizedK16 +quantize_nvfp4_k16_bits(const std::uint32_t (&represented)[8], float input_scale_divisor) { float2 values[8]; float max_abs = 0.0F; #pragma unroll @@ -92,4 +126,14 @@ __device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloa return result; } +__device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloat16* source, + float input_scale_divisor) { + const uint4 packed0 = load_vec(source); + const uint4 packed1 = load_vec(source + 8); + const std::uint32_t represented[8] = { + packed0.x, packed0.y, packed0.z, packed0.w, packed1.x, packed1.y, packed1.z, packed1.w, + }; + return quantize_nvfp4_k16_bits(represented, input_scale_divisor); +} + } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_gemv.cuh b/src/ops/linear/nvfp4/nvfp4_gemv.cuh index 2164c4121f..de5c5df3b6 100644 --- a/src/ops/linear/nvfp4/nvfp4_gemv.cuh +++ b/src/ops/linear/nvfp4/nvfp4_gemv.cuh @@ -85,14 +85,7 @@ stage_nvfp4_scales(const std::uint8_t* __restrict__ scales, template __device__ __forceinline__ std::int64_t nvfp4_scale_offset(int parent_row, int group) { - const int m_tile = parent_row / 128; - const int row_inner = parent_row - m_tile * 128; - const int scale_tile = group / 4; - const int scale_lane = group & 3; - const int row_mod32 = row_inner & 31; - const int row_quartile = row_inner >> 5; - return static_cast(m_tile * Geometry::kScaleTilesPerRow + scale_tile) * 512 + - row_mod32 * 16 + row_quartile * 4 + scale_lane; + return nvfp4_scale_byte(parent_row, group); } template diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh index 2f8abbc227..8914f32389 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh @@ -22,7 +22,9 @@ struct Nvfp4W4a4MmaSchedule { static_assert(WarpsM > 0 && WarpsN > 0); static_assert((BlockM % WarpsM) == 0 && ((BlockM / WarpsM) % 16) == 0); static_assert((BlockN % WarpsN) == 0 && ((BlockN / WarpsN) % 8) == 0); - static_assert(Stages >= 1 && Stages <= 4); + // Six is what the narrow shared-expert schedules take: a tile of sixteen tokens has little + // else to hide its weight stream behind, and the shared-memory budget admits the depth. + static_assert(Stages >= 1 && Stages <= 6); static_assert(MinBlocksPerSm > 0); static constexpr int kBlockM = BlockM; @@ -56,6 +58,32 @@ struct Nvfp4W4a4IdentityRows { } }; +// The activation plane is addressed by token. A dense operator consumes it in order; a routed +// one consumes the same plane through its own column order, so the row a tile stages is a +// policy rather than the tile index. +struct Nvfp4W4a4IdentityTokens { + __device__ __forceinline__ int source_token(int token) const { return token; } + + // A dense operator knows its extent at launch. A routed one learns it per group, so the + // extent is read through the same policy instead of being a second launch argument. + __device__ __forceinline__ int active_tokens(int launched) const { return launched; } +}; + +// Which tile of the grid a CTA takes. A dense operator reads it straight off the block index; +// a routed one folds its own work list into it, so the tile a CTA serves is a policy rather than +// the block index itself. +struct Nvfp4W4a4MmaRasterRowFast { + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = static_cast(blockIdx.x); + block_token = static_cast(blockIdx.y); + } + + // A routed launch sizes its grid from a host-side bound on the work list and learns the real + // count only on the device, so the tiles past the end have to leave before they read a weight + // plane. A dense launch has no such tiles. + __device__ __forceinline__ bool live() const { return true; } +}; + template struct Nvfp4W4a4SharedStorage { alignas( @@ -77,18 +105,18 @@ __device__ __forceinline__ int nvfp4_w4a4_swizzled_byte(int row, int logical_byt return physical_segment * 16 + byte_in_segment; } -template +template __device__ __forceinline__ void stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, Nvfp4W4a4SharedStorage& shared, int stage, int k_tile, - int token_begin, int active_tokens) { + int token_begin, int active_tokens, TokenPolicy token_policy) { constexpr int kCodeTasks = Schedule::kBlockM * Schedule::kSegmentsPerRow; for (int task = static_cast(threadIdx.x); task < kCodeTasks; task += Schedule::kThreads) { const int row = task / Schedule::kSegmentsPerRow; const int logical_segment = task - row * Schedule::kSegmentsPerRow; const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; const int physical_byte = nvfp4_w4a4_swizzled_byte(row, logical_segment * 16); auto* destination = shared.a_codes[stage] + row * Schedule::kCodeRowBytes + physical_byte; const auto* input = source.codes + @@ -103,7 +131,7 @@ stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, row += Schedule::kThreads) { const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; auto* destination = &shared.a_scale4[stage][row * Schedule::kK64PerStage]; const auto* input = source.scales + (static_cast(source_token) * Geometry::kScaleTilesPerRow + @@ -121,7 +149,7 @@ stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, const int segment = task - row * kSegmentsPerRow; const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; const int local_k64 = segment * 4; auto* destination = &shared.a_scale4[stage][row * Schedule::kK64PerStage + local_k64]; const auto* input = source.scales + (static_cast(source_token) * @@ -203,21 +231,31 @@ __device__ __forceinline__ void stage_nvfp4_w4a4_weight(const std::uint8_t* __re } template + class RowPolicy = Nvfp4W4a4IdentityRows, bool PairRows = false, + class TokenPolicy = Nvfp4W4a4IdentityTokens, + class RasterPolicy = Nvfp4W4a4MmaRasterRowFast> __global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4_mma_kernel( Nvfp4W4a4MaterializedActivation activation, const std::uint8_t* __restrict__ weight_codes, const std::uint8_t* __restrict__ weight_scales, std::int32_t tokens, float alpha, - Epilogue epilogue, OutputPolicy output, RowPolicy row_policy = {}) { + Epilogue epilogue, OutputPolicy output, RowPolicy row_policy = {}, + TokenPolicy token_policy = {}, RasterPolicy raster_policy = {}) { static_assert((Geometry::kInputRows % Schedule::kBlockK) == 0); static_assert((Geometry::kOutputRows % Schedule::kBlockN) == 0); static_assert(!PairRows || (Schedule::kBlockN % 2) == 0); static_assert(!PairRows || ((Geometry::kOutputRows / 2) % (Schedule::kBlockN / 2)) == 0); __shared__ Nvfp4W4a4SharedStorage shared; - const int token_begin = static_cast(blockIdx.y) * Schedule::kBlockM; + int block_row = 0; + int block_token = 0; + // Liveness first: a dead tile of a routed launch has no work-list entry to read, and reading + // one would touch arena bytes nothing has written. + if (!raster_policy.live()) { return; } + raster_policy.blocks(block_row, block_token); + const int token_begin = block_token * Schedule::kBlockM; constexpr int kRowsPerBlock = PairRows ? Schedule::kBlockN / 2 : Schedule::kBlockN; - const int row_begin = static_cast(blockIdx.x) * kRowsPerBlock; + const int row_begin = block_row * kRowsPerBlock; + const int active = token_policy.active_tokens(tokens); constexpr int kKTiles = Geometry::kInputRows / Schedule::kBlockK; constexpr int kWaitGroups = kKTiles < Schedule::kStages ? kKTiles - 1 : Schedule::kStages - 1; @@ -225,7 +263,7 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 for (int stage = 0; stage < Schedule::kStages; ++stage) { if (stage < kKTiles) { stage_nvfp4_w4a4_activation(activation, shared, stage, stage, - token_begin, tokens); + token_begin, active, token_policy); stage_nvfp4_w4a4_weight(weight_codes, weight_scales, shared, stage, stage, row_begin, row_policy); cp_commit(); @@ -320,7 +358,7 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 const int next_k_tile = k_tile + Schedule::kStages; if (next_k_tile < kKTiles) { stage_nvfp4_w4a4_activation(activation, shared, stage, next_k_tile, - token_begin, tokens); + token_begin, active, token_policy); stage_nvfp4_w4a4_weight(weight_codes, weight_scales, shared, stage, next_k_tile, row_begin, row_policy); } @@ -350,11 +388,11 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 float value01 = accumulators[mma_m][mma_n][1] * alpha; float value10 = accumulators[mma_m][mma_n][2] * alpha; float value11 = accumulators[mma_m][mma_n][3] * alpha; - if (token0 < tokens) { + if (token0 < active) { value00 = epilogue.apply(parent_row0, token0, value00); value01 = epilogue.apply(parent_row1, token0, value01); } - if (token1 < tokens) { + if (token1 < active) { value10 = epilogue.apply(parent_row0, token1, value10); value11 = epilogue.apply(parent_row1, token1, value11); } @@ -365,13 +403,22 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 __syncthreads(); constexpr int kStoredRows = PairRows ? Schedule::kBlockN / 2 : Schedule::kBlockN; constexpr int kVectorsPerRow = kStoredRows / 8; + // An output policy that completes a value across the neighbouring lane - the sparse-MoE + // gate/up epilogue does - needs both lanes of a pair to hold the same token, and they do only + // while a row's vectors come in pairs. A geometry that broke it would deadlock on the shuffle + // rather than fail a test, so it is refused here. This is stricter than the kernel needs: the + // dense output policies pair no lanes and would be safe with an odd count. Every schedule in + // the tree satisfies it, the narrowest being BlockN=32 with PairRows, which gives two. + static_assert((kVectorsPerRow % 2) == 0, + "a lane pair must stay inside one token: give the tile an even number of " + "output vectors per row"); constexpr int kOutputVectors = Schedule::kBlockM * kVectorsPerRow; for (int task = static_cast(threadIdx.x); task < kOutputVectors; task += Schedule::kThreads) { const int token_local = task / kVectorsPerRow; const int row_vector = task - token_local * kVectorsPerRow; const int token = token_begin + token_local; - if (token < tokens) { + if (token < active) { const uint4 values = load_vec(shared_output + token_local * kOutputStride + row_vector * 8); if constexpr (PairRows) { diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index f5dd00f701..0fc318bc67 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -6,6 +6,7 @@ #include "ops/common/math.cuh" #include "ops/common/memory.cuh" #include "ops/common/warp.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" #include "ops/linear/q4/q4_rowsplit_storage.cuh" #include "ops/linear/q5/q5_rowsplit_storage.cuh" #include "ops/linear/q6/q6_rowsplit_storage.cuh" @@ -111,6 +112,10 @@ struct SparseMoePlanes { const std::uint8_t* __restrict__ codes = nullptr; const std::uint8_t* __restrict__ high = nullptr; const std::uint8_t* __restrict__ scales = nullptr; + // One divisor per source matrix the plane was assembled from, each covering `divisor_rows` + // consecutive rows. Only NVFP4 has them. + const float* __restrict__ divisors = nullptr; + int divisor_rows = 1; }; // Codecs come in two lane ownerships. Under `kPackedWord8` a lane owns eight consecutive K values @@ -125,9 +130,13 @@ struct Q4Codec { static constexpr bool kPackedWord8 = true; static constexpr bool kSingleValuePerLane = false; + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + template __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, - int group, int lane_in_group, + int group, int lane_in_group, float, float (&weights)[8]) { const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( @@ -143,9 +152,13 @@ struct Q5Codec { static constexpr bool kPackedWord8 = true; static constexpr bool kSingleValuePerLane = false; + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + template __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, - int group, int lane_in_group, + int group, int lane_in_group, float, float (&weights)[8]) { const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( @@ -163,9 +176,13 @@ struct Q6Codec { static constexpr bool kPackedWord8 = true; static constexpr bool kSingleValuePerLane = false; + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + template __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, - int group, int lane_in_group, + int group, int lane_in_group, float, float (&weights)[8]) { const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( @@ -183,6 +200,10 @@ struct Q8Codec { static constexpr bool kPackedWord8 = false; static constexpr bool kSingleValuePerLane = true; + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + template __device__ static __forceinline__ float load_one(const SparseMoePlanes& planes, int row, int group, int lane) { @@ -203,6 +224,42 @@ struct Q8Codec { } }; +// NVFP4 packs two e2m1 codes per byte and one e4m3 scale per sixteen values, so a lane's eight +// values are four code bytes and exactly one scale: whichever half of a block the lane owns, its +// eight values lie inside that one block. Only the scale plane is swizzled; the codes are plain +// row-major. The per-tensor divisor is folded into the decoded block scale, once per eight values +// rather than once per value. +struct Nvfp4Codec { + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + // Which stored divisor this row was quantised against. Taken once per row, not per group. + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes& planes, + int row) { + return __frcp_rn(planes.divisors[row / planes.divisor_rows]); + } + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, float row_scale, + float (&weights)[8]) { + const std::uint32_t packed = *reinterpret_cast( + planes.codes + static_cast(row) * (K / 2) + group * 32 + + lane_in_group * 4); + const std::uint8_t scale = + planes.scales[nvfp4_scale_byte(row, group * 4 + (lane_in_group >> 1))]; + const float coefficient = decode_nvfp4_e4m3(scale) * row_scale; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const float2 code = + decode_nvfp4_e2m1x2(static_cast(packed >> (8 * pair))); + weights[pair * 2] = code.x * coefficient; + weights[pair * 2 + 1] = code.y * coefficient; + } + } +}; + // A codec declares exactly one lane ownership. The trap below catches a codec that declares // neither, which would otherwise compile and reduce a zero accumulator. template @@ -220,12 +277,14 @@ __device__ __forceinline__ void dot_two_rows(const SparseMoePlanes& planes, int if constexpr (Codec::kPackedWord8) { const int lane_group = lane >> 3; const int lane_in_group = lane & 7; + const float scale0 = Codec::row_coefficient(planes, row0); + const float scale1 = Codec::row_coefficient(planes, row1); for (int group_base = first_group; group_base < last_group; group_base += 4) { const int group = group_base + lane_group; float weights0[8]; float weights1[8]; - Codec::template load_eight(planes, row0, group, lane_in_group, weights0); - Codec::template load_eight(planes, row1, group, lane_in_group, weights1); + Codec::template load_eight(planes, row0, group, lane_in_group, scale0, weights0); + Codec::template load_eight(planes, row1, group, lane_in_group, scale1, weights1); const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); const float2 x0 = bf16x2_bits_to_float2(input.x); const float2 x1 = bf16x2_bits_to_float2(input.y); @@ -358,6 +417,11 @@ __device__ __forceinline__ void dot_fp32_rows(const SparseMoePlanes& planes, int // accumulation. const int lane_group = lane >> 3; const int lane_in_group = lane & 7; + float row_scale[Rows]; +#pragma unroll + for (int row = 0; row < Rows; ++row) { + row_scale[row] = Codec::row_coefficient(planes, row_base + row); + } for (int group_base = first_group; group_base < last_group; group_base += 4) { const int group = group_base + lane_group; const float4 x0 = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); @@ -367,7 +431,7 @@ __device__ __forceinline__ void dot_fp32_rows(const SparseMoePlanes& planes, int for (int row = 0; row < Rows; ++row) { float weights[8]; Codec::template load_eight(planes, row_base + row, group, lane_in_group, - weights); + row_scale[row], weights); #pragma unroll for (int item = 0; item < 8; ++item) { acc[row] = fmaf(weights[item], values[item], acc[row]); @@ -513,6 +577,14 @@ sparse_moe_d4_token_kernel(const int* __restrict__ token_ids, const float* __res } } +// How many bytes of the shared-expert down codes the D1 prefetch may touch. A four-bit plane is +// half the size of an eight-bit one, and walking the eight-bit length over it addresses memory the +// artifact never allocated. +unsigned long long shared_down_code_bytes(QType qtype) { + const unsigned long long elements = static_cast(kHidden) * kIntermediate; + return qtype == QType::NVFP4 ? elements / 2 : elements; +} + void launch_d1(const Tensor& x, const SparseMoeWeights& weights, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { sparse_moe_d1_kernel<<>>( @@ -520,14 +592,20 @@ void launch_d1(const Tensor& x, const SparseMoeWeights& weights, static_cast(weights.router_shared_gate.qdata), static_cast(workspace.scratch.data), static_cast(weights.shared_down.qdata), - static_cast(kHidden) * kIntermediate); + shared_down_code_bytes(weights.shared_down.qtype)); CUDA_CHECK(cudaGetLastError()); } SparseMoePlanes matrix_planes(const Weight& weight) { + // Only NVFP4 stores its block scale as a quotient of a divisor, and only it has a plane of + // them; for the row-split codecs the field is not part of the representation at all. The row + // stride is one here so that a row's divisor index is well formed whatever the codec, and the + // codecs without divisors never read it. return {static_cast(weight.qdata), static_cast(weight.qhigh), - static_cast(weight.scales)}; + static_cast(weight.scales), + static_cast(weight.weight_divisors), + weight.weight_divisor_rows > 0 ? weight.weight_divisor_rows : 1}; } template @@ -557,16 +635,28 @@ void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, case QType::Q8_G32_FP16: launch_d3_dependent_codec(x, weights, workspace, stream); return; + case QType::NVFP4: + launch_d3_dependent_codec(x, weights, workspace, stream); + return; default: throw std::invalid_argument("sparse_moe: unsupported D3 codec"); } } +// Hidden rows one D4 block takes. NVFP4 interleaves rows inside a 512-byte scale tile so that rows +// r and r+1 share one 32-byte sector: taking both halves the sectors the kernel reads from L2, 398k +// to 326k, and the kernel from 10.2 us to 8.0. A third and fourth row start the next sector and buy +// nothing. The row-split planes store a scale per row, so a second row only costs them grid. +template +inline constexpr int kDecodeDownRows = 1; +template <> +inline constexpr int kDecodeDownRows = 2; + template void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destination, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, const void* prefetch_data, std::size_t prefetch_bytes) { - constexpr int kRows = 1; + constexpr int kRows = kDecodeDownRows; CUDA_CHECK(pdl::launch_dependent( {dim3(kHidden / kRows), dim3(9 * 32), 0, stream}, sparse_moe_d4_nine_warp_kernel, @@ -594,6 +684,10 @@ void launch_d4_dependent(const SparseMoeWeights& weights, Tensor& destination, launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, prefetch_bytes); return; + case QType::NVFP4: + launch_d4_dependent_codec(weights, destination, workspace, stream, + prefetch_data, prefetch_bytes); + return; default: throw std::invalid_argument("sparse_moe: unsupported D4 codec"); } @@ -704,6 +798,10 @@ void sparse_moe_decode_launch_d3_small_t(const Tensor& x, const SparseMoeWeights launch_d3_small_t_codec(x, weights, token_ids, token_activations, tokens, schedule, stream, nullptr); return; + case QType::NVFP4: + launch_d3_small_t_codec( + x, weights, token_ids, token_activations, tokens, schedule, stream, nullptr); + return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D3 codec"); } @@ -742,6 +840,11 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, schedule, stream, nullptr); return; + case QType::NVFP4: + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + schedule, stream, nullptr); + return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D4 codec"); } diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp index 5412838524..240011839d 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp @@ -25,7 +25,11 @@ SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, QType r (routed_down == QType::Q5_G64_FP16 || routed_down == QType::Q6_G64_FP16); const bool mtp_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - if (!main_profile && !mtp_profile) { + // The decode kernels are CUDA-core dot products, so NVFP4 is read here against represented + // BF16 activations: four-bit activations buy no arithmetic on this path and would cost a + // quantiser kernel of their own. + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + if (!main_profile && !mtp_profile && !nvfp4_profile) { throw std::invalid_argument("sparse_moe: unsupported routed codec profile"); } diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill.h b/src/ops/sparse_moe/prefill/sparse_moe_prefill.h index 09a8d6395f..c3e347a601 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill.h +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill.h @@ -18,9 +18,16 @@ inline constexpr std::int32_t kSparseMoePrefillWorkspaceMin = 20; inline constexpr std::int32_t kSparseMoePrefillQ4Q5Min = 47; inline constexpr std::int32_t kSparseMoePrefillQ4Q6Min = 47; inline constexpr std::int32_t kSparseMoePrefillQ8Q8Min = 20; -inline constexpr std::int32_t kSparseMoePrefillWideMin = 768; -inline constexpr std::int32_t kSparseMoePrefillSliceMax = 4096; -inline constexpr std::int32_t kSparseMoeRouteTileTokens = 8; +// Swept cold against this profile's own small-T route over [2,46], each route measured through the +// operator benchmark at its own schedule. Small-T is ahead through T=8 (79.6 us against 81.9), and +// from T=9 to T=12 it is at most 4.5 % behind - the price of keeping the activation represented, +// which is worth paying for the accuracy: rel_l2 at T=1 is 1.60e-3 against 1.70e-1 for a four-bit +// activation. From T=13 the prefill route's lead becomes real (+9 % at T=16, +29 % at T=32) and +// outgrows that. +inline constexpr std::int32_t kSparseMoePrefillNvfp4Min = 13; +inline constexpr std::int32_t kSparseMoePrefillWideMin = 768; +inline constexpr std::int32_t kSparseMoePrefillSliceMax = 4096; +inline constexpr std::int32_t kSparseMoeRouteTileTokens = 8; // 257 logits padded to a 16-byte-aligned per-token stride. inline constexpr std::int32_t kSparseMoeRouterScoreRows = 260; @@ -28,6 +35,9 @@ struct SparseMoePrefillPlan { std::int32_t tokens = 0; std::int32_t slice_tokens = 0; std::size_t workspace_bytes = 0; + // Whether the route stages its operands in NVFP4. The staging buffers are the profile's own, + // so the flag has to reach the allocation and not only the launch. + bool nvfp4 = false; }; struct SparseMoePrefillWorkspace { @@ -59,17 +69,30 @@ struct SparseMoePrefillWorkspace { // router scores FP32 <-> shared SwiGLU BF16 // gathered X BF16 <-> routed down output BF16 <-> adaptive FP32 activations // (a Q4 routed gate/up stages from x and gathers nothing, so it only ever writes here) - // routed SwiGLU BF16 <-> routed FP32 token reduction + // routed SwiGLU BF16 <-> routed FP32 token reduction <-> every NVFP4 staging plane Tensor score_storage; Tensor shared_activation; Tensor grouped_io; Tensor routed_storage; Tensor routed_sum; + + // The NVFP4 route quantises the chunk once and keeps both SwiGLU intermediates encoded, so + // `down` needs no separate quantiser: 288 bytes a row against the bf16 plane's 1024. All six + // are carved out of `routed_storage`, which this profile never writes -- the encoded + // intermediate replaces the routed SwiGLU plane and the shared-down epilogue replaces the + // token reduction, so the branch returns before either view has a reader. 3744 bytes a token + // inside 8192, so the profile costs no workspace at all. + Tensor nvfp4_input_codes; + Tensor nvfp4_input_scales; + Tensor nvfp4_routed_codes; + Tensor nvfp4_routed_scales; + Tensor nvfp4_shared_codes; + Tensor nvfp4_shared_scales; }; template -SparseMoePrefillWorkspace allocate_sparse_moe_prefill_workspace(Arena& arena, - std::int32_t capacity_tokens) { +SparseMoePrefillWorkspace +allocate_sparse_moe_prefill_workspace(Arena& arena, std::int32_t capacity_tokens, bool nvfp4) { SparseMoePrefillWorkspace out; const std::int32_t assignments = 8 * capacity_tokens; const std::int32_t route_tiles = @@ -104,12 +127,37 @@ SparseMoePrefillWorkspace allocate_sparse_moe_prefill_workspace(Arena& arena, out.routed_storage = arena.alloc(DType::BF16, {512, assignments}, 256); out.routed_sum = Tensor(out.routed_storage.data, DType::FP32, {2048, capacity_tokens}); + + if (nvfp4) { + // The layout builder allocates nothing, so an offset is only formed when there is a + // pointer -- the same guard the index maps above use. + auto* const base = static_cast(out.routed_storage.data); + std::int64_t cursor = 0; + const auto carve = [&](std::int32_t rows, std::int32_t columns) { + std::uint8_t* const at = base != nullptr ? base + cursor : nullptr; + cursor += (static_cast(rows) * columns + 255) / 256 * 256; + return Tensor(at, DType::U8, {rows, columns}); + }; + out.nvfp4_input_codes = carve(1024, capacity_tokens); + out.nvfp4_input_scales = carve(128, capacity_tokens); + out.nvfp4_routed_codes = carve(256, assignments); + out.nvfp4_routed_scales = carve(32, assignments); + out.nvfp4_shared_codes = carve(256, capacity_tokens); + out.nvfp4_shared_scales = carve(32, capacity_tokens); + // 3744 bytes a token, and five alignment gaps of under 256 each, inside the 8192 a token + // that `routed_storage` holds for a profile that has a routed SwiGLU plane. Asserted + // rather than trusted: a wider staging plane has to move the union, not overrun it. The + // eight is `assignments / capacity_tokens`, which is kTopK. + static_assert(1024 + 128 + 8 * 256 + 8 * 32 + 256 + 32 + 5 * 255 <= 512 * 8 * 2, + "the NVFP4 staging planes must fit the routed SwiGLU plane they alias"); + static_cast(cursor); + } return out; } [[nodiscard]] bool sparse_moe_uses_prefill(std::int32_t tokens, QType routed_gate_up, QType routed_down) noexcept; -[[nodiscard]] std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens); +[[nodiscard]] std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens, bool nvfp4); [[nodiscard]] SparseMoePrefillPlan resolve_sparse_moe_prefill_plan(std::int32_t tokens, QType routed_gate_up, QType routed_down); diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu index 3d9fea13fa..974cefae9b 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu @@ -6,6 +6,10 @@ #include "ops/common/memory.cuh" #include "ops/common/mma.cuh" #include "ops/common/rowsplit_mma.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" +#include "ops/linear/nvfp4/nvfp4_geometry.h" +#include "ops/linear/nvfp4/nvfp4_output.cuh" +#include "ops/linear/nvfp4/nvfp4_w4a4_mma.cuh" #include "ops/linear/q4/q4_rowsplit_storage.cuh" #include "ops/linear/q5/q5_rowsplit_storage.cuh" #include "ops/linear/q6/q6_rowsplit_storage.cuh" @@ -37,6 +41,77 @@ constexpr int kRouterStages = 2; constexpr int kRouterWarps = 8; constexpr int kRouterThreads = 32 * kRouterWarps; +constexpr int kRouterSimtThreads = 256; +constexpr int kRouterSimtMaxTokens = 16; + +// One CTA per router row. The tile of the MMA router is 16 rows, which is right for a chunk and +// leaves 17 CTAs for a decode step; this leaves 257 and streams one weight row per CTA. +__global__ __launch_bounds__(kRouterSimtThreads, 4) void sparse_moe_prefill_router_simt_kernel( + const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, + float* __restrict__ scores, int tokens, std::uint8_t* __restrict__ quantized_codes, + std::uint8_t* __restrict__ quantized_scales, float input_scale_divisor) { + constexpr int kWarps = kRouterSimtThreads / 32; + const int row = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + const int lane = tid & 31; + const int warp = tid >> 5; + const auto* w = weight + static_cast(row) * kHidden; + __shared__ float partial[kWarps]; + + for (int token = 0; token < tokens; ++token) { + const auto* xt = x + static_cast(token) * kHidden; + float acc = 0.0F; + for (int i = tid * 8; i < kHidden; i += kRouterSimtThreads * 8) { + const uint4 wv = load_vec(w + i); + const uint4 xv = load_vec(xt + i); + const std::uint32_t wb[4]{wv.x, wv.y, wv.z, wv.w}; + const std::uint32_t xb[4]{xv.x, xv.y, xv.z, xv.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 a = bf16x2_bits_to_float2(wb[j]); + const float2 b = bf16x2_bits_to_float2(xb[j]); + acc = fmaf(a.x, b.x, acc); + acc = fmaf(a.y, b.y, acc); + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + acc += __shfl_xor_sync(0xffffffffU, acc, offset); + } + if (lane == 0) { partial[warp] = acc; } + __syncthreads(); + if (tid == 0) { + float total = 0.0F; +#pragma unroll + for (int i = 0; i < kWarps; ++i) { total += partial[i]; } + scores[static_cast(token) * kSparseMoeRouterScoreRows + row] = total; + } + __syncthreads(); + } + + // The NVFP4 route needs this same input as an encoded plane. It is the tensor this kernel has + // just finished reading and nothing between the two depends on the other, so it is written here + // instead of by a launch of its own. + if (quantized_codes != nullptr) { + constexpr int kGroupsPerRow = kHidden / 16; + const int total = tokens * kGroupsPerRow; + const int stride = static_cast(gridDim.x) * kRouterSimtThreads; + for (int task = static_cast(blockIdx.x) * kRouterSimtThreads + tid; task < total; + task += stride) { + const int token_index = task / kGroupsPerRow; + const int group = task - token_index * kGroupsPerRow; + const Nvfp4QuantizedK16 quantized = quantize_nvfp4_k16( + x + static_cast(token_index) * kHidden + group * 16, + input_scale_divisor); + store_vec(quantized_codes + static_cast(token_index) * (kHidden / 2) + + group * 8, + make_uint2(quantized.codes_lo, quantized.codes_hi)); + quantized_scales[static_cast(token_index) * kGroupsPerRow + group] = + quantized.scale; + } + } +} + __global__ __launch_bounds__(kRouterThreads, 2) void sparse_moe_prefill_router_mma_kernel( const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, float* __restrict__ scores, int tokens) { @@ -288,7 +363,107 @@ constexpr int kExpertStages = 2; constexpr int kGateUpNarrowStages = 6; constexpr int kExpertWarps = 8; constexpr int kExpertThreads = 32 * kExpertWarps; -constexpr int kRtx5090SmCount = 170; + +// An inclusive prefix sum over one block of kExperts lanes. `totals` needs one slot per warp. +__device__ __forceinline__ int sparse_moe_block_inclusive_scan(int value, int* totals, int tid) { + constexpr int kScanWarps = kExperts / 32; + const int lane = tid & 31; + const int warp = tid >> 5; + int running = value; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const int carried = __shfl_up_sync(0xffffffffU, running, offset); + if (lane >= offset) { running += carried; } + } + if (lane == 31) { totals[warp] = running; } + __syncthreads(); + if (warp == 0) { + int total = lane < kScanWarps ? totals[lane] : 0; +#pragma unroll + for (int offset = 1; offset < kScanWarps; offset <<= 1) { + const int carried = __shfl_up_sync(0xffffffffU, total, offset); + if (lane >= offset) { total += carried; } + } + if (lane < kScanWarps) { totals[lane] = total; } + } + __syncthreads(); + return running + (warp == 0 ? 0 : totals[warp - 1]); +} + +// Selection, scan and index as one kernel. Only correct for a single route tile, which is what the +// launch checks: with route_tiles == 1 the tile-local ranks this writes are the global ones. +__global__ __launch_bounds__(kExpertThreads, 1) void sparse_moe_prefill_small_route_kernel( + const float* __restrict__ scores, int* __restrict__ ids, float* __restrict__ alpha, + float* __restrict__ shared_scale, int* __restrict__ local_rank, int* __restrict__ tile_counts, + int* __restrict__ tile_bases, int* __restrict__ expert_offsets, + int* __restrict__ route_job_experts, int* __restrict__ route_job_columns, + int* __restrict__ route_job_count, int* __restrict__ packed_index, + int* __restrict__ packed_token, int job_bn, int tokens) { + __shared__ int counts[kExperts]; + __shared__ int prefix[kExperts]; + __shared__ int warp_totals[kExperts / 32]; + __shared__ float selected_logits[kSparseMoeRouteTileTokens][kTopK]; + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + + // Phase one: the top-k of every token and how many rows each expert collected. + if (tid < kExperts) { counts[tid] = 0; } + __syncthreads(); + if (warp < tokens) { + sparse_moe_select_top8_warp( + scores + static_cast(warp) * kSparseMoeRouterScoreRows, + ids + warp * kTopK, alpha + warp * kTopK, shared_scale + warp, selected_logits[warp]); + __syncwarp(); + if (lane == 0) { +#pragma unroll + for (int rank = 0; rank < kTopK; ++rank) { + const int assignment = warp * kTopK + rank; + local_rank[assignment] = atomicAdd(&counts[ids[assignment]], 1); + } + } + } + __syncthreads(); + + // Phase two: the expert prefix and the work list. One tile, so a base is a cursor. + const int expert = tid; + const int count = counts[expert]; + tile_counts[expert] = count; + prefix[expert] = sparse_moe_block_inclusive_scan(count, warp_totals, tid); + __syncthreads(); + const int base = expert == 0 ? 0 : prefix[expert - 1]; + expert_offsets[expert] = base; + if (expert == kExperts - 1) { expert_offsets[kExperts] = prefix[expert]; } + tile_bases[expert] = base; + + __syncthreads(); + prefix[expert] = + sparse_moe_block_inclusive_scan((count + job_bn - 1) / job_bn, warp_totals, tid); + const int touched = __syncthreads_count(count > 0); + const int job_begin = expert == 0 ? 0 : prefix[expert - 1]; + const int jobs = (count + job_bn - 1) / job_bn; + for (int job = 0; job < jobs; ++job) { + route_job_experts[job_begin + job] = expert; + route_job_columns[job_begin + job] = job * job_bn; + } + if (expert == kExperts - 1) { + route_job_count[0] = prefix[expert]; + route_job_count[1] = touched; + } + + // Phase three: the packed map. The bases above are visible through the barrier. + __syncthreads(); + const int assignments = tokens * kTopK; + for (int assignment = tid; assignment < assignments; assignment += kExpertThreads) { + const int token = assignment / kTopK; + const int packed = tile_bases[ids[assignment]] + local_rank[assignment]; + packed_index[assignment] = packed; + packed_token[packed] = token; + } +} + +constexpr int kRtx5090SmCount = 170; // Upper bound on the persistent grid. The routed GEMMs stride their work list by gridDim.x, // so any grid is correct; this caps the launch when the work list is long. constexpr int kPrefillMaxBlocksPerSm = 32; @@ -1130,6 +1305,348 @@ __global__ __launch_bounds__(kExpertThreads, 1) void sparse_moe_prefill_q8_down_ } } +// --------------------------------------------------------------------------------------------- +// The NVFP4 route. +// +// The routed experts of this model are one contiguous weight plane, so the whole block runs on +// nvfp4_w4a4_mma_kernel with four policies: which expert a tile serves (the scan's work list), +// which weight rows it stages, which activation columns it reads, and what the epilogue writes. +// Three things differ from the Q4/Q5 route beyond the codec: +// +// * the chunk's activation plane is quantised once, up front, and both the routed and the +// shared gate/up read it; +// * the gate/up epilogue emits NVFP4 rather than bf16, so `down` needs no separate quantiser +// and the intermediate plane is 288 bytes a row instead of 1024; +// * the same kernels serve every token count this route is entered for, down to the one-token +// ragged tail of a sliced call - decode itself runs elsewhere. They are launched over the +// scan's work list, so a single token touches eight experts and reads eight experts' +// weights -- the tile wastes compute on its empty columns, which decode does not pay for +// because it is bound by the weight stream. +// +// The routed tile schedules are the ones the operator bench settled on: a 128-deep tile, +// because the 256-deep one takes 55 KiB of staging buffers and only one CTA then fits on an +// SM. The dense shared gate/up keeps the 256-deep tile, where one CTA is the right answer. +// --------------------------------------------------------------------------------------------- + +using Nvfp4RoutedGateUpGeometry = Nvfp4Geometry; +using Nvfp4RoutedDownGeometry = Nvfp4Geometry; +using Nvfp4SharedGateUpGeometry = Nvfp4Geometry<2 * kIntermediate, kHidden>; +using Nvfp4SharedDownGeometry = Nvfp4Geometry; +using Nvfp4InputGeometry = Nvfp4ActivationGeometry; + +using Nvfp4RoutedSchedule = Nvfp4W4a4MmaSchedule<64, 128, 128, 2, 4, 2, 3>; +using Nvfp4DenseSchedule = Nvfp4W4a4MmaSchedule<64, 128, 256, 4, 4, 2, 1>; + +template +struct Nvfp4ScheduleTag { + using type = Schedule; +}; + +// The shared expert is dense, so its grid is the output rows over the tile width: at one token the +// schedules above leave 8 CTAs for gate/up and 16 for down on a 170-SM card, and the pair measures +// 162 GB/s where the routed pair reaches 597. A narrower tile multiplies the CTA count without +// changing the bytes a CTA reads, and BlockM 16 stops staging 64 token rows when one exists. +// gate/up reaches 32 because `Nvfp4SharedGateUpRows` is not contiguous and stages scales per +// (row, k64); down stays at 64, the narrowest the contiguous scale path allows. +using Nvfp4SharedGateUpSmallTSchedule = Nvfp4W4a4MmaSchedule<16, 32, 128, 1, 4, 6, 4>; +// Four stages, not six: this geometry has 512 input rows over a 128-deep tile, so there are +// only four k-tiles to prefetch, and the extra two stages would buy nothing and cost 11.5 KiB +// of shared memory a block. +using Nvfp4SharedDownSmallTSchedule = Nvfp4W4a4MmaSchedule<16, 64, 128, 1, 2, 4, 4>; +// Below this many tokens the wide tile cannot fill the machine and the narrow one cannot lose. +constexpr int kNvfp4SharedSmallTokens = 32; + +// The routed pair keeps BlockM at 64 because the work list is cut into column blocks of that +// width; only N narrows, which is what decides how many CTAs a job expands into. +using Nvfp4RoutedGateUpSmallTSchedule = Nvfp4W4a4MmaSchedule<64, 32, 128, 2, 2, 4, 1>; +using Nvfp4RoutedDownSmallTSchedule = Nvfp4W4a4MmaSchedule<64, 64, 128, 2, 2, 2, 4>; +// Reached only by the ragged tail of a sliced call: this profile enters prefill at thirteen +// tokens, so a whole call never has two. +constexpr int kNvfp4RoutedSmallTokens = 2; + +// One divisor per source matrix, taken from the row the value belongs to. The routed planes are +// stacks of separately quantised matrices - gate and up of an expert were quantised apart, so the +// stride is half an expert's gate/up rows - and a stack of one carries its reciprocal directly, +// which is what every plane did before there were stacks. +// +// This runs once per output element, so the index may not cost an integer division: `shift` is the +// stride's base-two logarithm where it has one, and every stride a source can give this +// architecture does. +struct Nvfp4SourceDivisorEpilogue { + const float* __restrict__ divisors; + float uniform; + int divisor_rows; + int shift; + + __device__ __forceinline__ float apply(std::int32_t row, std::int32_t, float value) const { + if (divisors == nullptr) { return value * uniform; } + const int index = shift >= 0 ? (row >> shift) : (row / divisor_rows); + return value * __frcp_rn(divisors[index]); + } +}; + +// A stacked plane reads its divisors per row; a plane with one carries its reciprocal and never +// touches the plane of them. +Nvfp4SourceDivisorEpilogue nvfp4_divisor_epilogue(const Weight& weight) { + const int rows = weight.weight_divisor_rows; + if (rows == weight.n) { return {nullptr, 1.0F / weight.weight_scale_divisor, rows, -1}; } + const int shift = (rows > 0 && (rows & (rows - 1)) == 0) ? __builtin_ctz(rows) : -1; + return {static_cast(weight.weight_divisors), 0.0F, rows, shift}; +} + +// The scan emits one job per column tile of this width, so it has to match the tile the routed +// schedules stage. +static_assert(Nvfp4RoutedSchedule::kBlockM == 64); +// A route job is one column tile of one expert, so every routed schedule's token tile has to be the +// job width exactly: wider, and one job would cover two tiles of which only one would run. +constexpr int kNvfp4JobColumns = Nvfp4RoutedSchedule::kBlockM; +static_assert(Nvfp4RoutedGateUpSmallTSchedule::kBlockM == kNvfp4JobColumns); +static_assert(Nvfp4RoutedDownSmallTSchedule::kBlockM == kNvfp4JobColumns); + +// The routed grid is sized from a host-side bound on the work list; the real count lives on the +// device, so tiles past the end leave before they touch a weight plane. +struct Nvfp4Jobs { + const int* __restrict__ experts; + const int* __restrict__ columns; + const int* __restrict__ count; + + __device__ __forceinline__ int expert() const { return experts[blockIdx.y]; } + + __device__ __forceinline__ int column_base() const { return columns[blockIdx.y]; } + + __device__ __forceinline__ bool live() const { return static_cast(blockIdx.y) < count[0]; } +}; + +template +struct Nvfp4RoutedRaster { + Nvfp4Jobs jobs; + + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = static_cast(blockIdx.x); + block_token = jobs.column_base() / BlockM; + } + + __device__ __forceinline__ bool live() const { return jobs.live(); } +}; + +// `down` weight rows are contiguous inside an expert, so folding the expert's row base into the +// block row lets the stock identity row policy and the stock contiguous scale staging address +// the plane without knowing about routing. +template +struct Nvfp4RoutedDownRaster { + Nvfp4Jobs jobs; + + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = jobs.expert() * (kHidden / BlockN) + static_cast(blockIdx.x); + block_token = jobs.column_base() / BlockM; + } + + __device__ __forceinline__ bool live() const { return jobs.live(); } +}; + +template +struct Nvfp4ExpertGateUpRows { + static constexpr bool kContiguous = false; + static constexpr int kRowsPerBranch = RowsPerBranch; + + Nvfp4Jobs jobs; + + __device__ __forceinline__ int weight_row(int row_begin, int local_row) const { + const int within = row_begin + (local_row & (kRowsPerBranch - 1)) + + (local_row >= kRowsPerBranch ? kIntermediate : 0); + return jobs.expert() * (2 * kIntermediate) + within; + } +}; + +// The chunk's activation plane is shared by every expert; the route decides which rows this +// expert reads and how many there are. +struct Nvfp4RoutedGatherTokens { + const int* __restrict__ packed_token; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ int source_token(int column) const { + return packed_token[expert_offsets[jobs.expert()] + column]; + } + + __device__ __forceinline__ int active_tokens(int) const { + const int expert = jobs.expert(); + return expert_offsets[expert + 1] - expert_offsets[expert]; + } +}; + +// `down` reads what gate/up wrote, which is already in route order, so this is an offset rather +// than a lookup. +struct Nvfp4RoutedPackedTokens { + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ int source_token(int column) const { + return expert_offsets[jobs.expert()] + column; + } + + __device__ __forceinline__ int active_tokens(int) const { + const int expert = jobs.expert(); + return expert_offsets[expert + 1] - expert_offsets[expert]; + } +}; + +union Nvfp4Bf16Pair { + unsigned bits; + __nv_bfloat162 values; +}; + +__device__ __forceinline__ unsigned nvfp4_swiglu_pair(unsigned gate_bits, unsigned up_bits) { + Nvfp4Bf16Pair gate{gate_bits}; + Nvfp4Bf16Pair up{up_bits}; + const float2 g = __bfloat1622float2(gate.values); + const float2 u = __bfloat1622float2(up.values); + Nvfp4Bf16Pair result; + result.values = __floats2bfloat162_rn(silu(g.x) * u.x, silu(g.y) * u.y); + return result.bits; +} + +// One store is half an NVFP4 group, so every thread of the block takes part instead of half of +// them; the group's maximum is completed across the neighbouring lane. The two lanes of a pair +// always share a token, so they are active together. +__device__ __forceinline__ void nvfp4_store_half_group(std::uint8_t* codes, std::uint8_t* scales, + std::int64_t row_stride_codes, + std::int64_t row_stride_scales, + std::int64_t packed, std::int32_t row, + uint4 gate, uint4 up, float divisor) { + const std::uint32_t bits[4] = { + nvfp4_swiglu_pair(gate.x, up.x), + nvfp4_swiglu_pair(gate.y, up.y), + nvfp4_swiglu_pair(gate.z, up.z), + nvfp4_swiglu_pair(gate.w, up.w), + }; + float2 values[4]; + float max_abs = 0.0F; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + values[pair] = bf16x2_bits_to_float2(bits[pair]); + max_abs = fmaxf(max_abs, fabsf(values[pair].x)); + max_abs = fmaxf(max_abs, fabsf(values[pair].y)); + } + // The two lanes of a group are always both here or both gone, because they share a token, so + // the mask is known and does not have to be asked for. `__activemask` is not composable with a + // shuffle: if the compiler reconverges differently it can name a lane that has left. + const unsigned pair_mask = 3U << (threadIdx.x & 31U & ~1U); + max_abs = fmaxf(max_abs, __shfl_xor_sync(pair_mask, max_abs, 1)); + + const float scale_unencoded = __fdiv_rn(divisor * max_abs, 6.0F); + const std::uint8_t scale = __nv_cvt_float_to_fp8(scale_unencoded, __NV_SATFINITE, __NV_E4M3); + std::uint32_t word = 0; + if (scale != 0) { + const float reciprocal = __frcp_rn(decode_nvfp4_e4m3(scale)); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + values[pair].x = values[pair].x * divisor * reciprocal; + values[pair].y = values[pair].y * divisor * reciprocal; + } + word = pack_nvfp4_e2m1x8(values); + } + const int half = (row >> 3) & 1; + const int group_row = row & ~15; + store_vec(codes + packed * row_stride_codes + (group_row >> 1) + half * 4, word); + if (half == 0) { scales[packed * row_stride_scales + (group_row >> 4)] = scale; } +} + +struct Nvfp4RoutedGateUpOutput { + std::uint8_t* codes; + std::uint8_t* scales; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + float divisor; + + __device__ __forceinline__ void store_pair_vector(std::int32_t row, std::int32_t column, + uint4 gate, uint4 up) const { + const std::int64_t packed = expert_offsets[jobs.expert()] + column; + nvfp4_store_half_group(codes, scales, kIntermediate / 2, kIntermediate / 16, packed, row, + gate, up, divisor); + } +}; + +struct Nvfp4RoutedDownOutput { + __nv_bfloat16* data; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ void store_vector(std::int32_t row, std::int32_t column, + uint4 values) const { + const std::int64_t packed = expert_offsets[jobs.expert()] + column; + store_vec(data + packed * kHidden + (row & (kHidden - 1)), values); + } +}; + +template +struct Nvfp4SharedGateUpRows { + static constexpr bool kContiguous = false; + static constexpr int kRowsPerBranch = RowsPerBranch; + + __device__ __forceinline__ int weight_row(int row_begin, int local_row) const { + return row_begin + (local_row & (kRowsPerBranch - 1)) + + (local_row >= kRowsPerBranch ? kIntermediate : 0); + } +}; + +struct Nvfp4SharedGateUpOutput { + std::uint8_t* codes; + std::uint8_t* scales; + float divisor; + + __device__ __forceinline__ void store_pair_vector(std::int32_t row, std::int32_t column, + uint4 gate, uint4 up) const { + nvfp4_store_half_group(codes, scales, kIntermediate / 2, kIntermediate / 16, column, row, + gate, up, divisor); + } +}; + +// The block's last write: the residual it was handed, plus the routed sum the reduce kernel +// built, plus the shared expert scaled by its gate. +struct Nvfp4SharedDownOutput { + __nv_bfloat16* destination; + const __nv_bfloat16* __restrict__ grouped_output; + const int* __restrict__ packed_index; + const float* __restrict__ alpha; + const float* __restrict__ shared_scale; + + // The routed sum is formed here rather than by a kernel of its own. This epilogue already + // visits every (token, hidden) element once and already had to read the sum back, so reading + // the eight expert rows instead costs the same traffic and saves a graph node per layer. + __device__ __forceinline__ void store_vector(std::int32_t row, std::int32_t column, + uint4 values) const { + const std::int64_t base = static_cast(column) * kHidden + row; + const float gain = shared_scale[column]; + float routed[8] = {}; +#pragma unroll + for (int route = 0; route < kTopK; ++route) { + const int packed = packed_index[column * kTopK + route]; + const float weight = alpha[column * kTopK + route]; + const uint4 raw = + load_vec(grouped_output + static_cast(packed) * kHidden + row); + const std::uint32_t words[4]{raw.x, raw.y, raw.z, raw.w}; +#pragma unroll + for (int index = 0; index < 4; ++index) { + const Nvfp4Bf16Pair slot{words[index]}; + const float2 decoded = __bfloat1622float2(slot.values); + routed[2 * index] += weight * decoded.x; + routed[2 * index + 1] += weight * decoded.y; + } + } + const Nvfp4Bf16Pair pair[4] = {{values.x}, {values.y}, {values.z}, {values.w}}; +#pragma unroll + for (int index = 0; index < 4; ++index) { + const float2 shared = __bfloat1622float2(pair[index].values); + const std::int64_t at = base + 2 * index; + destination[at] = __float2bfloat16_rn(__bfloat162float(destination[at]) + + routed[2 * index] + gain * shared.x); + destination[at + 1] = __float2bfloat16_rn(__bfloat162float(destination[at + 1]) + + routed[2 * index + 1] + gain * shared.y); + } + } +}; + union alignas(16) SparseMoeBf16x8 { uint4 raw; __nv_bfloat162 pair[4]; @@ -1173,6 +1690,146 @@ __global__ void sparse_moe_prefill_reduce_kernel(const __nv_bfloat16* __restrict } } +struct Nvfp4PrefillPlanes { + std::uint8_t* input_codes; + std::uint8_t* input_scales; + std::uint8_t* routed_codes; + std::uint8_t* routed_scales; + std::uint8_t* shared_codes; + std::uint8_t* shared_scales; +}; + +void launch_sparse_moe_prefill_nvfp4(const __nv_bfloat16* input, const SparseMoeWeights& weights, + __nv_bfloat16* destination, int tokens, int assignments, + int max_route_jobs, const int* packed_token, + const int* offsets, const int* route_job_experts, + const int* route_job_columns, const int* route_job_count, + const int* packed_index, const float* alpha, + const float* shared_scale, __nv_bfloat16* grouped_io, + float* /*routed_sum*/, Nvfp4PrefillPlanes planes, + cudaStream_t stream) { + const Nvfp4Jobs jobs{route_job_experts, route_job_columns, route_job_count}; + + // One quantisation of the chunk serves both the routed and the shared gate/up. Below the + // small-token threshold the router has already produced it on its way out. + if (tokens > kRouterSimtMaxTokens) { + constexpr int kThreads = 256; + const int groups = tokens * (kHidden / 16); + nvfp4_w4a4_quantize_kernel + <<<(groups + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + input, planes.input_codes, planes.input_scales, tokens, tokens, + weights.routed_gate_up.input_scale_divisor); + CUDA_CHECK(cudaGetLastError()); + } + + const Nvfp4W4a4MaterializedActivation chunk{planes.input_codes, planes.input_scales}; + const Nvfp4W4a4MaterializedActivation routed_middle{planes.routed_codes, planes.routed_scales}; + const Nvfp4W4a4MaterializedActivation shared_middle{planes.shared_codes, planes.shared_scales}; + + { + // Alpha keeps only what the whole plane shares; the row's own divisor is the epilogue's. + const float scale = 1.0F / weights.routed_gate_up.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.routed_gate_up); + const Nvfp4RoutedGatherTokens token_policy{packed_token, offsets, jobs}; + const Nvfp4RoutedGateUpOutput output{planes.routed_codes, planes.routed_scales, offsets, + jobs, weights.routed_down.input_scale_divisor}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + constexpr int kPairRows = Schedule::kBlockN / 2; + using Rows = Nvfp4ExpertGateUpRows; + using Raster = Nvfp4RoutedRaster; + const dim3 grid(kIntermediate / kPairRows, max_route_jobs); + nvfp4_w4a4_mma_kernel<<>>( + chunk, static_cast(weights.routed_gate_up.qdata), + static_cast(weights.routed_gate_up.scales), assignments, scale, + epilogue, output, Rows{jobs}, token_policy, Raster{jobs}); + }; + if (tokens <= kNvfp4RoutedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + { + const float scale = 1.0F / weights.routed_down.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.routed_down); + const Nvfp4RoutedPackedTokens token_policy{offsets, jobs}; + const Nvfp4RoutedDownOutput output{grouped_io, offsets, jobs}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + using Raster = Nvfp4RoutedDownRaster; + const dim3 grid(kHidden / Schedule::kBlockN, max_route_jobs); + nvfp4_w4a4_mma_kernel + <<>>( + routed_middle, static_cast(weights.routed_down.qdata), + static_cast(weights.routed_down.scales), assignments, + scale, epilogue, output, Nvfp4W4a4IdentityRows{}, token_policy, Raster{jobs}); + }; + if (tokens <= kNvfp4RoutedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + + { + const float scale = 1.0F / weights.shared_gate_up.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.shared_gate_up); + const Nvfp4SharedGateUpOutput output{planes.shared_codes, planes.shared_scales, + weights.shared_down.input_scale_divisor}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + constexpr int kPairRows = Schedule::kBlockN / 2; + using Rows = Nvfp4SharedGateUpRows; + const dim3 grid(kIntermediate / kPairRows, + (tokens + Schedule::kBlockM - 1) / Schedule::kBlockM); + nvfp4_w4a4_mma_kernel + <<>>( + chunk, static_cast(weights.shared_gate_up.qdata), + static_cast(weights.shared_gate_up.scales), tokens, scale, + epilogue, output, Rows{}); + }; + if (tokens < kNvfp4SharedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + { + const float scale = 1.0F / weights.shared_down.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.shared_down); + const Nvfp4SharedDownOutput output{destination, grouped_io, packed_index, alpha, + shared_scale}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + const dim3 grid(kHidden / Schedule::kBlockN, + (tokens + Schedule::kBlockM - 1) / Schedule::kBlockM); + nvfp4_w4a4_mma_kernel<<>>( + shared_middle, static_cast(weights.shared_down.qdata), + static_cast(weights.shared_down.scales), tokens, scale, + epilogue, output); + }; + if (tokens < kNvfp4SharedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } +} + } // namespace void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, @@ -1223,41 +1880,71 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, const int route_tiles = (tokens + kSparseMoeRouteTileTokens - 1) / kSparseMoeRouteTileTokens; const int assignments = tokens * kTopK; + const bool nvfp4 = weights.routed_gate_up.qtype == QType::NVFP4; const int adaptive_last = weights.routed_down.qtype == QType::Q5_G64_FP16 ? 51 : weights.routed_down.qtype == QType::Q6_G64_FP16 ? 52 : 0; - const bool adaptive = tokens >= 47 && tokens <= adaptive_last; - - sparse_moe_prefill_router_mma_kernel<<>>(input, router, scores, - tokens); + const bool adaptive = !nvfp4 && tokens >= 47 && tokens <= adaptive_last; + + // Only this profile: the small-token router folds the NVFP4 quantisation of the chunk + // into its epilogue, and its reduction order differs from the MMA router's, so a + // groupwise call would change scores for nothing. + if (nvfp4 && tokens <= kRouterSimtMaxTokens) { + sparse_moe_prefill_router_simt_kernel<<>>( + input, router, scores, tokens, + static_cast(workspace.nvfp4_input_codes.data), + static_cast(workspace.nvfp4_input_scales.data), + weights.routed_gate_up.input_scale_divisor); + } else { + sparse_moe_prefill_router_mma_kernel<<>>(input, router, + scores, tokens); + } CUDA_CHECK(cudaGetLastError()); - sparse_moe_prefill_select_count_kernel<<>>( - scores, ids, alpha, shared_scale, local_rank, tile_counts, tokens); + const bool routed_gate_up_q4 = weights.routed_gate_up.qtype == QType::Q4_G64_FP16; + // Like the Q4 route, NVFP4 stages its activation tile from the chunk through the route + // order, so it needs the packed token map rather than a materialised gather. + const bool needs_packed_token = routed_gate_up_q4 || nvfp4; + // One route tile means one block can carry selection, scan and index together, which is + // two graph nodes fewer per layer. Restricted to this profile: the Q4 route would take it + // on the ragged tail of a sliced call, where it replaces three kernels whose result it has + // to reproduce exactly, and that is a separate claim with a separate gate. + const bool fused_route = nvfp4 && tokens <= kSparseMoeRouteTileTokens; + if (!fused_route) { + sparse_moe_prefill_select_count_kernel<<>>( + scores, ids, alpha, shared_scale, local_rank, tile_counts, tokens); + } CUDA_CHECK(cudaGetLastError()); - const bool wide_plan = tokens >= kSparseMoePrefillWideMin; - const int route_job_bn = wide_plan ? 64 : 32; + const bool wide_plan = tokens >= kSparseMoePrefillWideMin; + // The NVFP4 tile is 64 columns wide at every token count, so its jobs have to be too. + const int route_job_bn = nvfp4 ? kNvfp4JobColumns : (wide_plan ? 64 : 32); // The scan emits one route job per nonempty column tile of an expert, so it cannot // emit more than one job per full tile of assignments plus one tail per expert -- // the same bound the workspace is sized by. Each job expands into row blocks, and // sizing the grid from that product keeps a persistent block on one work item while // there are fewer work items than the cap, instead of a fixed count that has to // iterate. The exact job count only exists on the device. - const int max_route_jobs = assignments / route_job_bn + kExperts; + // A routed token reaches one expert per assignment, so the tail term is bounded by the + // assignment count as well as by the expert count. Narrowing it launches fewer blocks that + // would exit at once; it is numerically inert, and it is applied only here because on the + // groupwise routes it was measured at zero and is not this change's to make. + const int max_route_jobs = + assignments / route_job_bn + (nvfp4 ? std::min(kExperts, assignments) : kExperts); const int routed_gate_work = max_route_jobs * (kIntermediate / (kExpertBM / 2)); const int routed_down_work = max_route_jobs * (kHidden / kExpertBM); const int routed_gate_blocks = std::min(routed_gate_work, kPrefillMaxBlocks); const int routed_down_blocks = std::min(routed_down_work, kPrefillMaxBlocks); - sparse_moe_prefill_scan_kernel<<<1, kExpertThreads, 0, stream>>>( - tile_counts, tile_bases, offsets, route_job_experts, route_job_columns, route_job_count, - route_tiles, route_job_bn, tokens, adaptive); + if (!fused_route) { + sparse_moe_prefill_scan_kernel<<<1, kExpertThreads, 0, stream>>>( + tile_counts, tile_bases, offsets, route_job_experts, route_job_columns, + route_job_count, route_tiles, route_job_bn, tokens, adaptive); + } CUDA_CHECK(cudaGetLastError()); - const bool routed_gate_up_q4 = weights.routed_gate_up.qtype == QType::Q4_G64_FP16; - const int index_blocks = (assignments + kExpertThreads - 1) / kExpertThreads; + const int index_blocks = (assignments + kExpertThreads - 1) / kExpertThreads; if (adaptive) { auto* adaptive_activations = reinterpret_cast(grouped_io); sparse_moe_decode_launch_d3_small_t(input_slice, weights, ids, adaptive_activations, @@ -1271,15 +1958,37 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, sparse_moe_prefill_index_kernel<<>>( ids, local_rank, packed_index, tile_bases, packed_token, assignments, route_job_count); - } else if (routed_gate_up_q4) { - sparse_moe_prefill_index_kernel<<>>( - ids, local_rank, packed_index, tile_bases, packed_token, assignments, nullptr); + } else if (needs_packed_token) { + if (fused_route) { + sparse_moe_prefill_small_route_kernel<<<1, kExpertThreads, 0, stream>>>( + scores, ids, alpha, shared_scale, local_rank, tile_counts, tile_bases, offsets, + route_job_experts, route_job_columns, route_job_count, packed_index, + packed_token, route_job_bn, tokens); + } else { + sparse_moe_prefill_index_kernel<<>>( + ids, local_rank, packed_index, tile_bases, packed_token, assignments, nullptr); + } } else { sparse_moe_prefill_gather_kernel<<>>( input, ids, local_rank, packed_index, tile_bases, grouped_io, nullptr); } CUDA_CHECK(cudaGetLastError()); + if (nvfp4) { + const Nvfp4PrefillPlanes planes{ + static_cast(workspace.nvfp4_input_codes.data), + static_cast(workspace.nvfp4_input_scales.data), + static_cast(workspace.nvfp4_routed_codes.data), + static_cast(workspace.nvfp4_routed_scales.data), + static_cast(workspace.nvfp4_shared_codes.data), + static_cast(workspace.nvfp4_shared_scales.data)}; + launch_sparse_moe_prefill_nvfp4( + input, weights, output, tokens, assignments, max_route_jobs, packed_token, offsets, + route_job_experts, route_job_columns, route_job_count, packed_index, alpha, + shared_scale, grouped_io, routed_sum, planes, stream); + continue; + } + const dim3 routed_gate_grid(kIntermediate / (kExpertBM / 2), kExperts); // Same predicate that decided whether packed_token was written above. if (routed_gate_up_q4) { diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp b/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp index 1bba6ed4d8..0ee4c3c27c 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp @@ -17,6 +17,9 @@ std::int32_t prefill_min_tokens(QType routed_gate_up, QType routed_down) noexcep if (routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16) { return kSparseMoePrefillQ8Q8Min; } + if (routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4) { + return kSparseMoePrefillNvfp4Min; + } return 0; } @@ -28,13 +31,13 @@ bool sparse_moe_uses_prefill(std::int32_t tokens, QType routed_gate_up, return minimum != 0 && tokens >= minimum; } -std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens) { +std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens, bool nvfp4) { if (max_tokens < kSparseMoePrefillWorkspaceMin) { throw std::invalid_argument("sparse_moe prefill: max_tokens must be at least 20"); } const std::int32_t capacity_tokens = std::min(max_tokens, kSparseMoePrefillSliceMax); WorkspaceLayoutBuilder layout; - (void)allocate_sparse_moe_prefill_workspace(layout, capacity_tokens); + (void)allocate_sparse_moe_prefill_workspace(layout, capacity_tokens, nvfp4); return layout.peak_bytes(1); } @@ -48,8 +51,14 @@ SparseMoePrefillPlan resolve_sparse_moe_prefill_plan(std::int32_t tokens, QType throw std::invalid_argument("sparse_moe prefill: unsupported token count"); } - const std::int32_t slice_tokens = std::min(tokens, kSparseMoePrefillSliceMax); - return {tokens, slice_tokens, sparse_moe_prefill_workspace_bytes(tokens)}; + // The workspace query has a floor of twenty tokens, and NVFP4 enters prefill at thirteen, so a + // call in [13,19] has to allocate to that floor or the execution high-water mark falls short of + // what the query promised. Slicing is unaffected: the loop still steps by the real token + // count. + const std::int32_t slice_tokens = + std::max(std::min(tokens, kSparseMoePrefillSliceMax), kSparseMoePrefillWorkspaceMin); + const bool nvfp4 = routed_gate_up == QType::NVFP4; + return {tokens, slice_tokens, sparse_moe_prefill_workspace_bytes(slice_tokens, nvfp4), nvfp4}; } } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp b/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp index 6c483106dc..757609b0d7 100644 --- a/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp +++ b/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp @@ -30,11 +30,24 @@ SparseMoeSmallTPlan resolve_sparse_moe_small_t_plan(std::int32_t tokens, QType r (routed_down == QType::Q5_G64_FP16 || routed_down == QType::Q6_G64_FP16); const bool mtp_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - if (!main_profile && !mtp_profile) { + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + if (!main_profile && !mtp_profile && !nvfp4_profile) { throw std::invalid_argument("sparse_moe small-T: unsupported routed codec profile"); } SparseMoeSmallTPlan plan{tokens, sparse_moe_small_t_workspace_bytes(tokens)}; + if (nvfp4_profile) { + // Swept over the whole [2,46] window against all nine pairs, cold, with the prefill + // frontier held open; the figures below are that sweep's, and the frontier comment in + // sparse_moe_prefill.h quotes a later one, so the two differ in the third digit. One path + // per warp and two output rows per block win everywhere NVFP4 keeps the window: 34.8 us at + // T=2 and 79.9 at T=8, against 38.4 and 90.1 for the three-path, one-row pair. Rows4 + // overtakes Rows2 somewhere between T=16 and T=32, which is beyond the frontier this + // profile stops at, so it has no interval of its own. + plan.d3_schedule = SparseMoeSmallTD3Schedule::Paths9; + plan.d4_schedule = SparseMoeSmallTD4Schedule::Rows2; + return plan; + } if (mtp_profile) { plan.d3_schedule = tokens <= 5 ? SparseMoeSmallTD3Schedule::Paths1 : SparseMoeSmallTD3Schedule::Paths9; diff --git a/src/ops/wrapper/sparse_moe.cpp b/src/ops/wrapper/sparse_moe.cpp index aa8eb1dc65..ae769bc29d 100644 --- a/src/ops/wrapper/sparse_moe.cpp +++ b/src/ops/wrapper/sparse_moe.cpp @@ -2,6 +2,7 @@ #include "ninfer/ops/sparse_moe.h" #include "core/nvtx.h" +#include "ops/linear/nvfp4/nvfp4_format.h" #include "ops/sparse_moe/decode/sparse_moe_decode.h" #include "ops/sparse_moe/prefill/sparse_moe_prefill.h" #include "ops/sparse_moe/small_t/sparse_moe_small_t.h" @@ -137,8 +138,48 @@ void require_quantized(const Weight& weight, std::int32_t n, std::int32_t k, con ranges.push_back(address_range(weight.scales, scale_bytes, std::string(name) + " scales")); } +// The NVFP4 profile is all four expert matrices together: the gate/up epilogue emits the +// encoded intermediate `down` consumes, so the two cannot be mixed with a groupwise codec. +bool nvfp4_profile(const SparseMoeWeights& weights) { + return weights.routed_gate_up.qtype == QType::NVFP4 && + weights.routed_down.qtype == QType::NVFP4 && + weights.shared_gate_up.qtype == QType::NVFP4 && + weights.shared_down.qtype == QType::NVFP4; +} + +void require_nvfp4(const Weight& weight, std::int32_t n, std::int32_t k, const char* name, + std::vector& ranges) { + require_matrix_metadata(weight, n, k, name); + // Every expert matrix was quantised on its own, so each of these planes is a stack. + const detail::Nvfp4WeightGeometry geometry = + detail::validate_nvfp4_weight(weight, (std::string("sparse_moe ") + name).c_str(), true); + ranges.push_back( + address_range(weight.qdata, geometry.code_plane_bytes, std::string(name) + " code")); + ranges.push_back( + address_range(weight.scales, geometry.scale_plane_bytes, std::string(name) + " scales")); +} + void validate_weights(const SparseMoeWeights& weights, std::vector& ranges) { require_router(weights.router_shared_gate, ranges); + if (nvfp4_profile(weights)) { + require_nvfp4(weights.routed_gate_up, kRoutedGateRows, kHidden, "routed_gate_up", ranges); + require_nvfp4(weights.routed_down, kRoutedDownRows, kIntermediate, "routed_down", ranges); + require_nvfp4(weights.shared_gate_up, kSharedGateRows, kHidden, "shared_gate_up", ranges); + require_nvfp4(weights.shared_down, kHidden, kIntermediate, "shared_down", ranges); + // The prefill route encodes the hidden state once and both gate/up GEMMs read that one + // plane, each dividing by its own bank's activation divisor. Two different divisors would + // scale the shared expert's whole contribution by their ratio, silently. + if (weights.routed_gate_up.input_scale_divisor != + weights.shared_gate_up.input_scale_divisor) { + throw std::invalid_argument( + "sparse_moe: routed and shared experts must share one activation divisor"); + } + return; + } + if (weights.routed_gate_up.qtype == QType::NVFP4 || weights.routed_down.qtype == QType::NVFP4 || + weights.shared_gate_up.qtype == QType::NVFP4 || weights.shared_down.qtype == QType::NVFP4) { + throw std::invalid_argument("sparse_moe: NVFP4 requires all four expert matrices"); + } if (weights.routed_gate_up.qtype != QType::Q4_G64_FP16 && weights.routed_gate_up.qtype != QType::Q8_G32_FP16) { throw std::invalid_argument("sparse_moe: routed_gate_up must be Q4 or Q8"); @@ -169,10 +210,12 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout const bool q8_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - const std::int32_t prefill_first = - q8_profile ? detail::kSparseMoePrefillQ8Q8Min - : (routed_down == QType::Q5_G64_FP16 ? detail::kSparseMoePrefillQ4Q5Min - : detail::kSparseMoePrefillQ4Q6Min); + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + const std::int32_t prefill_first = q8_profile ? detail::kSparseMoePrefillQ8Q8Min + : nvfp4_profile ? detail::kSparseMoePrefillNvfp4Min + : routed_down == QType::Q5_G64_FP16 + ? detail::kSparseMoePrefillQ4Q5Min + : detail::kSparseMoePrefillQ4Q6Min; std::size_t required = 0; if (min_tokens == 1) { required = detail::sparse_moe_decode_workspace_bytes(); } @@ -186,7 +229,12 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout const std::int32_t prefill_interval_first = std::max(min_tokens, prefill_first); if (prefill_interval_first <= max_tokens) { - required = std::max(required, detail::sparse_moe_prefill_workspace_bytes(max_tokens)); + // The prefill workspace query has a floor of its own, and the plan allocates to it, so an + // interval whose top is below that floor still has to promise what execution will take. + required = + std::max(required, detail::sparse_moe_prefill_workspace_bytes( + std::max(max_tokens, detail::kSparseMoePrefillWorkspaceMin), + nvfp4_profile)); } return required; } @@ -246,7 +294,7 @@ void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilo const detail::SparseMoePrefillPlan plan = detail::resolve_sparse_moe_prefill_plan( tokens, weights.routed_gate_up.qtype, weights.routed_down.qtype); const detail::SparseMoePrefillWorkspace views = - detail::allocate_sparse_moe_prefill_workspace(workspace, plan.slice_tokens); + detail::allocate_sparse_moe_prefill_workspace(workspace, plan.slice_tokens, plan.nvfp4); detail::sparse_moe_prefill_launch(x, weights, destination, plan, views, stream); return; } diff --git a/tests/ops/test_sparse_moe.cpp b/tests/ops/test_sparse_moe.cpp index 35dc23ca07..da9f8de427 100644 --- a/tests/ops/test_sparse_moe.cpp +++ b/tests/ops/test_sparse_moe.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,52 @@ constexpr ReductionCriterion kSparseMoeA16Tolerance{ /*gross_relative_to_max_reference*/ 0.0, }; +// NVFP4 below its prefill frontier reads represented BF16 activations, and its weights decode into +// exactly the values the oracle is built from, so the remaining error is the BF16 rounding of the +// destination plus the FP32 accumulation of the dot product. Both terms are relative: the absolute +// term above is a number from the row-split fixture, whose references peak near 0.3 where these +// peak near 1154, so it bounds a different quantity. +// +// The two terms are bounded differently, and the second is why the gross bound is not the same +// number as the L2 one. Over a whole vector the roundings are independent and average down, which +// is what 2.5e-3 of the reference RMS describes. One element does not average: the accumulation +// can carry the exact value across a rounding boundary, and the destination then lands on the +// wrong side of it. +// +// BF16 keeps eight significand bits, so inside a binade the spacing is 2^(e-7) and the relative +// ulp runs over (2^-8, 2^-7]. Half an ulp is therefore at most 2^-8 of the value - 3.9e-3 - and +// since no element exceeds the maximum reference, 4.0e-3 of that maximum bounds the rounding of +// any one of them. That is the format's number, not this fixture's. +// +// The case that showed the old bound was below it: at T=1 the reference peaks at 665, so 2.5e-3 +// set the limit at 1.66, and an element whose exact value is 637.701 lands on 636 - the nearest +// BF16 there, the grid stepping by 4 - for an error of 1.70. The destination was right and the +// bound was not. +// +// Measured over the represented cases: rel-L2 2.09e-3 against 2.5e-3, worst element 1.79 against +// the 2.66 this bound allows. The L2 term is the binding one and stays where it was. +constexpr ReductionCriterion kSparseMoeNvfp4Tolerance{ + /*relative_l2*/ 2.5e-3, + /*gross_absolute*/ 0.0, + /*gross_relative_to_max_reference*/ 4.0e-3, +}; + +// The NVFP4 prefill route quantises activations as well as weights, so the A16 bound above does +// not apply to it. Against an oracle built from the same represented weights the remaining error +// is the four-bit activation plane, and a dot product does not average it away: the sum of K terms +// carries the same cancellation as the error does, so the output keeps the per-element relative +// error of e2m1 rather than that error over the square root of K. Both terms are relative for the +// same reason -- an absolute bound would mean nothing across cases whose references differ by an +// order of magnitude. It also depends on how far the experts spread in magnitude, which is a +// property of the fixture rather than of the route: every source matrix here carries its own +// divisor over a fourfold range, narrower than the 3.3x to 11.9x a real plane spans. Measured +// maximum over the prefill cases is 1.011e-1, and the worst element is half the gross bound. +constexpr ReductionCriterion kSparseMoeA4Tolerance{ + /*relative_l2*/ 1.8e-1, + /*gross_absolute*/ 0.0, + /*gross_relative_to_max_reference*/ 2.4e-1, +}; + constexpr std::size_t kOutputGuardBytes = 256; constexpr std::uint8_t kOutputGuardByte = 0xa5; @@ -81,11 +128,6 @@ std::vector bf16_bits(const std::vector& values) { return bits; } -int compare_output(const std::string& label, const std::vector& actual, - const std::vector& reference) { - return verify_reduction(label, actual, reference, kSparseMoeA16Tolerance); -} - class GuardedBf16Output { public: explicit GuardedBf16Output(std::size_t words) @@ -105,6 +147,126 @@ class GuardedBf16Output { std::size_t words_; }; +// NVFP4 keeps its two planes in one payload at a fixed geometry -- codes from the start, +// scales at the 256-aligned end of them, and after those one weight divisor for each separately +// quantised matrix stacked into the plane -- so this cannot reuse the two independent buffers a +// row-split weight has. +class DeviceNvfp4 { +public: + DeviceNvfp4(std::int32_t rows, std::int32_t columns, std::int32_t divisor_rows) + : rows_(rows), columns_(columns), divisor_rows_(checked_divisor_rows(rows, divisor_rows)), + code_plane_bytes_(static_cast(rows) * columns / 2), + scale_offset_((code_plane_bytes_ + 255) / 256 * 256), + scale_plane_bytes_(static_cast(rows) * columns / 16), + storage_(scale_offset_ + scale_plane_bytes_ + + static_cast(rows / divisor_rows) * sizeof(float)) { + storage_.fill(0); + } + + // The member initialiser divides by this, so the check cannot wait for the body. + static std::int32_t checked_divisor_rows(std::int32_t rows, std::int32_t divisor_rows) { + if (divisor_rows <= 0 || rows % divisor_rows != 0 || divisor_rows % 128 != 0) { + throw std::invalid_argument("sparse_moe test: invalid NVFP4 divisor stride"); + } + return divisor_rows; + } + + void copy_rows(const quantized_weight::PackedWeight& source, std::int32_t destination_row) { + const std::int32_t source_rows = source.weight.n; + if (source.weight.qtype != QType::NVFP4 || source.weight.k != columns_ || source_rows < 1 || + (source_rows % 128) != 0 || destination_row < 0 || (destination_row % 128) != 0 || + destination_row > rows_ - source_rows) { + throw std::invalid_argument("sparse_moe test: invalid NVFP4 row copy"); + } + const std::size_t code_bytes = static_cast(source_rows) * columns_ / 2; + const std::size_t scale_bytes = static_cast(source_rows) * columns_ / 16; + storage_.copy_from_host(source.payload.data(), code_bytes, + static_cast(destination_row) * columns_ / 2); + // A scale tile spans 128 rows, and both the source and the destination row are multiples + // of that, so the tile index scales with the row exactly. + storage_.copy_from_host(source.payload.data() + source.scale_plane_offset, scale_bytes, + scale_offset_ + + static_cast(destination_row) * columns_ / 16); + if (source_rows > divisor_rows_ || destination_row % divisor_rows_ != 0) { + throw std::invalid_argument("sparse_moe test: a copy may not span two divisors"); + } + storage_.copy_from_host(source.payload.data() + source.weight_divisor_offset, sizeof(float), + scale_offset_ + scale_plane_bytes_ + + static_cast(destination_row / divisor_rows_) * + sizeof(float)); + if (destination_row == 0) { weight_scale_divisor_ = source.weight.weight_scale_divisor; } + input_scale_divisor_ = source.weight.input_scale_divisor; + } + + Weight weight() const { + Weight result{}; + result.payload = storage_.p; + result.payload_bytes = storage_.bytes; + result.high_plane_bytes = 0; + result.qtype = QType::NVFP4; + result.group_size = 16; + result.qdata = storage_.p; + result.qhigh = nullptr; + result.scales = static_cast(storage_.p) + scale_offset_; + result.n = rows_; + result.k = columns_; + result.group = 16; + result.layout = QuantLayout::BlockScaleK16M128x4; + result.scale_dtype = DType::FP8_E4M3FN; + result.ndim = 2; + result.shape[0] = rows_; + result.shape[1] = columns_; + result.padded_shape[0] = rows_; + result.padded_shape[1] = columns_; + result.weight_scale_divisor = weight_scale_divisor_; + result.input_scale_divisor = input_scale_divisor_; + result.weight_divisors = + static_cast(storage_.p) + scale_offset_ + scale_plane_bytes_; + result.weight_divisor_rows = divisor_rows_; + return result; + } + + int verify_rows(const std::string& label, const quantized_weight::PackedWeight& source, + std::int32_t destination_row) const { + const std::int32_t source_rows = source.weight.n; + const std::size_t code_bytes = static_cast(source_rows) * columns_ / 2; + const std::size_t scale_bytes = static_cast(source_rows) * columns_ / 16; + int failures = 0; + failures += + verify_span(label + " code", source.payload.data(), + static_cast(destination_row) * columns_ / 2, code_bytes); + failures += verify_span( + label + " scale", source.payload.data() + source.scale_plane_offset, + scale_offset_ + static_cast(destination_row) * columns_ / 16, scale_bytes); + failures += verify_span( + label + " divisor", source.payload.data() + source.weight_divisor_offset, + scale_offset_ + scale_plane_bytes_ + + static_cast(destination_row / divisor_rows_) * sizeof(float), + sizeof(float)); + return failures; + } + +private: + int verify_span(const std::string& label, const std::uint8_t* expected, std::size_t offset, + std::size_t bytes) const { + std::vector actual(bytes); + storage_.copy_to_host(actual.data(), bytes, offset); + if (std::memcmp(actual.data(), expected, bytes) == 0) { return 0; } + std::cerr << label << ": persistent weight was modified\n"; + return 1; + } + + std::int32_t rows_; + std::int32_t columns_; + std::int32_t divisor_rows_; + std::size_t code_plane_bytes_; + std::size_t scale_offset_; + std::size_t scale_plane_bytes_; + DeviceBuffer storage_; + float weight_scale_divisor_ = 1.0F; + float input_scale_divisor_ = 1.0F; +}; + class DeviceRowSplit { public: DeviceRowSplit(QType qtype, std::int32_t rows, std::int32_t columns) @@ -323,9 +485,19 @@ std::vector make_down(std::int32_t rows, std::int32_t columns, std::uint3 return source; } +// Gate and up are quantised apart in every published checkpoint, so each is packed on its own +// with its own divisor and the two are stacked into one plane. This takes one of them out of the +// generator's combined matrix. +std::vector take_rows(const std::vector& matrix, std::int32_t columns, + std::int32_t begin, std::int32_t rows) { + const auto first = matrix.begin() + static_cast(begin) * columns; + return std::vector(first, first + static_cast(rows) * columns); +} + struct HostExpert { int id; - quantized_weight::PackedWeight gate_up; + quantized_weight::PackedWeight gate; + quantized_weight::PackedWeight up; quantized_weight::PackedWeight down; }; @@ -373,7 +545,8 @@ std::vector sparse_moe_oracle(const std::vector& input, const std::vector& residual, const std::vector& router, const std::vector& experts, - const quantized_weight::PackedWeight& shared_gate_up, + const quantized_weight::PackedWeight& shared_gate, + const quantized_weight::PackedWeight& shared_up, const quantized_weight::PackedWeight& shared_down, const RoutePattern& intended_route) { const std::vector x(input.begin(), input.end()); @@ -406,16 +579,16 @@ std::vector sparse_moe_oracle(const std::vector& input, routed_activation[route].resize(kIntermediate); const HostExpert& expert = find_expert(experts, selected[route]); parallel_rows(kIntermediate, [&](std::int32_t row) { - const double gate = dot_fp64(expert.gate_up.dequant, row, kHidden, x); - const double up = dot_fp64(expert.gate_up.dequant, kIntermediate + row, kHidden, x); + const double gate = dot_fp64(expert.gate.dequant, row, kHidden, x); + const double up = dot_fp64(expert.up.dequant, row, kHidden, x); routed_activation[route][row] = (gate / (1.0 + std::exp(-gate))) * up; }); } std::vector shared_activation(kIntermediate); parallel_rows(kIntermediate, [&](std::int32_t row) { - const double gate = dot_fp64(shared_gate_up.dequant, row, kHidden, x); - const double up = dot_fp64(shared_gate_up.dequant, kIntermediate + row, kHidden, x); + const double gate = dot_fp64(shared_gate.dequant, row, kHidden, x); + const double up = dot_fp64(shared_up.dequant, row, kHidden, x); shared_activation[row] = (gate / (1.0 + std::exp(-gate))) * up; }); @@ -438,19 +611,121 @@ struct CodecProfile { const char* name; QType routed_gate_up; QType routed_down; + QType shared; + const ReductionCriterion* tolerance; + // The bound from the token count at which this profile starts quantising its activations, and + // that token count. Zero where no route of the profile does. + const ReductionCriterion* quantized_activation_tolerance; + std::int32_t quantized_activation_from; std::span token_cases; bool verify_graph_replay; + + [[nodiscard]] const ReductionCriterion& criterion(std::int32_t tokens) const { + return quantized_activation_from != 0 && tokens >= quantized_activation_from + ? *quantized_activation_tolerance + : *tolerance; + } }; +// The four expert matrices take whichever storage the profile names; nothing else in the +// fixture depends on which. +class ExpertPlane { +public: + ExpertPlane(QType qtype, std::int32_t rows, std::int32_t columns, std::int32_t divisor_rows) { + if (qtype == QType::NVFP4) { + nvfp4_.emplace(rows, columns, divisor_rows); + } else { + row_split_.emplace(qtype, rows, columns); + } + } + + void copy_rows(const quantized_weight::PackedWeight& source, std::int32_t destination_row) { + if (nvfp4_) { + nvfp4_->copy_rows(source, destination_row); + } else { + row_split_->copy_rows(source, destination_row); + } + } + + Weight weight() const { return nvfp4_ ? nvfp4_->weight() : row_split_->weight(); } + + int verify_rows(const std::string& label, const quantized_weight::PackedWeight& source, + std::int32_t destination_row) const { + return nvfp4_ ? nvfp4_->verify_rows(label, source, destination_row) + : row_split_->verify_rows(label, source, destination_row); + } + +private: + std::optional nvfp4_; + std::optional row_split_; +}; + +float decode_e4m3_scale(std::uint8_t bits) { + const int exponent = (bits >> 3) & 15; + const int mantissa = bits & 7; + const float value = exponent == 0 + ? std::ldexp(static_cast(mantissa) / 8.0F, -6) + : std::ldexp(1.0F + static_cast(mantissa) / 8.0F, exponent - 7); + return ((bits >> 7) & 1) != 0 ? -value : value; +} + +// The oracle needs the values the weight represents. Rather than quantise a second time and risk +// disagreeing with the packer, decode the payload the packer produced: by construction those are +// the values, and the decode is the format's own definition. +std::vector decode_nvfp4_payload(const quantized_weight::PackedWeight& packed, + std::int32_t n, std::int32_t k) { + static constexpr float kMagnitude[8]{0.0F, 0.5F, 1.0F, 1.5F, 2.0F, 3.0F, 4.0F, 6.0F}; + std::vector out(static_cast(n) * k); + const std::int32_t k_tiles = k / 64; + const float inverse = 1.0F / packed.weight.weight_scale_divisor; + for (std::int32_t row = 0; row < n; ++row) { + const std::int32_t row_tile = row / 128; + const std::int32_t row_inner = row % 128; + for (std::int32_t column = 0; column < k; ++column) { + const std::uint8_t byte = + packed.payload[static_cast(row) * k / 2 + column / 2]; + const unsigned nibble = (column & 1) != 0 ? (byte >> 4) : (byte & 0x0fU); + const std::int32_t group = column / 16; + const std::size_t scale_at = + static_cast(packed.scale_plane_offset) + + static_cast(row_tile * k_tiles + group / 4) * 512U + + static_cast(row_inner % 32) * 16U + + static_cast(row_inner / 32) * 4U + static_cast(group % 4); + const float magnitude = kMagnitude[nibble & 7U]; + out[static_cast(row) * k + column] = + ((nibble & 8U) != 0 ? -magnitude : magnitude) * + decode_e4m3_scale(packed.payload[scale_at]) * inverse; + } + } + return out; +} + +// A patterned NVFP4 matrix carries its own values; the groupwise path quantises the fixture's. +// Either way the oracle reads the represented values this returns. +quantized_weight::PackedWeight pack_expert(QType qtype, const std::vector& values, + std::int32_t n, std::int32_t k, std::uint32_t seed, + float weight_divisor = 0.125F) { + if (qtype == QType::NVFP4) { + quantized_weight::PatternedWeightOptions options; + options.weight_scale_divisor = weight_divisor; + options.input_scale_divisor = 3.5F; + quantized_weight::PackedWeight packed = + quantized_weight::make_patterned_weight(QType::NVFP4, n, k, seed, options); + packed.dequant = decode_nvfp4_payload(packed, n, k); + return packed; + } + return quantized_weight::pack_row_split_lowbit(values, n, k, qtype); +} + class SparseMoeFixture { public: explicit SparseMoeFixture(const CodecProfile& profile) : profile_(profile), router_(make_router()), router_bits_(bf16_bits(router_)), device_router_(to_device(router_bits_)), - routed_gate_(profile.routed_gate_up, kRoutedGateRows, kHidden), - routed_down_(profile.routed_down, kRoutedDownRows, kIntermediate), - shared_gate_(QType::Q8_G32_FP16, kSharedGateRows, kHidden), - shared_down_device_(QType::Q8_G32_FP16, kHidden, kIntermediate) { + routed_gate_(profile.routed_gate_up, kRoutedGateRows, kHidden, kIntermediate), + routed_down_(profile.routed_down, kRoutedDownRows, kIntermediate, kHidden), + shared_gate_(profile.shared, kSharedGateRows, kHidden, kIntermediate), + shared_down_device_(profile.shared, kHidden, kIntermediate, kHidden) { for (int pattern = 0; pattern < static_cast(kRoutePatterns.size()); ++pattern) { inputs_.push_back(make_input(pattern)); residuals_.push_back(make_residual(pattern)); @@ -465,32 +740,58 @@ class SparseMoeFixture { } } std::sort(expert_ids.begin(), expert_ids.end()); + // A published checkpoint quantises every expert matrix on its own, so no two of them share + // a divisor. Here they differ by construction, which is what makes a lookup that ignores + // the row visible to the oracle. + const auto expert_divisor = [](int expert, int matrix) { + return 0.0625F * (1.0F + static_cast((expert * 7 + matrix * 3) % 13) * 0.25F); + }; for (int expert : expert_ids) { - const float factor = 0.8f + static_cast((expert * 3) % 11) * 0.045f; - auto gate_up = quantized_weight::pack_row_split_lowbit( - make_gate_up(kExpertGateRows, kHidden, 100U + static_cast(expert), - factor), - kExpertGateRows, kHidden, profile.routed_gate_up); - auto down = quantized_weight::pack_row_split_lowbit( - make_down(kHidden, kIntermediate, 300U + static_cast(expert), - factor), - kHidden, kIntermediate, profile.routed_down); - routed_gate_.copy_rows(gate_up, expert * kExpertGateRows); + const float factor = 0.8f + static_cast((expert * 3) % 11) * 0.045f; + const std::uint32_t seed = 100U + static_cast(expert); + const std::vector gate_up_source = + make_gate_up(kExpertGateRows, kHidden, seed, factor); + auto gate = pack_expert(profile.routed_gate_up, + take_rows(gate_up_source, kHidden, 0, kIntermediate), + kIntermediate, kHidden, seed, expert_divisor(expert, 0)); + // A different seed, not just a different divisor: the NVFP4 packer derives its codes + // from the seed, so one seed for both halves would leave them bit-identical and a route + // that read the wrong half would look right. + auto up = + pack_expert(profile.routed_gate_up, + take_rows(gate_up_source, kHidden, kIntermediate, kIntermediate), + kIntermediate, kHidden, seed ^ 0x9e3779b9U, expert_divisor(expert, 1)); + auto down = + pack_expert(profile.routed_down, + make_down(kHidden, kIntermediate, + 300U + static_cast(expert), factor), + kHidden, kIntermediate, 300U + static_cast(expert), + expert_divisor(expert, 2)); + routed_gate_.copy_rows(gate, expert * kExpertGateRows); + routed_gate_.copy_rows(up, expert * kExpertGateRows + kIntermediate); routed_down_.copy_rows(down, expert * kHidden); - experts_.push_back({expert, std::move(gate_up), std::move(down)}); + experts_.push_back({expert, std::move(gate), std::move(up), std::move(down)}); } - shared_gate_host_ = quantized_weight::pack_q8_g32_row_split( - make_gate_up(kSharedGateRows, kHidden, 0x512U, 0.93f), kSharedGateRows, kHidden); - shared_down_host_ = quantized_weight::pack_q8_g32_row_split( - make_down(kHidden, kIntermediate, 0x731U, 0.87f), kHidden, kIntermediate); + const std::vector shared_source = + make_gate_up(kSharedGateRows, kHidden, 0x512U, 0.93f); + shared_gate_host_ = + pack_expert(profile.shared, take_rows(shared_source, kHidden, 0, kIntermediate), + kIntermediate, kHidden, 0x512U, 0.3125F); + shared_up_host_ = pack_expert( + profile.shared, take_rows(shared_source, kHidden, kIntermediate, kIntermediate), + kIntermediate, kHidden, 0x513U, 0.4375F); + shared_down_host_ = + pack_expert(profile.shared, make_down(kHidden, kIntermediate, 0x731U, 0.87f), kHidden, + kIntermediate, 0x731U); shared_gate_.copy_rows(shared_gate_host_, 0); + shared_gate_.copy_rows(shared_up_host_, kIntermediate); shared_down_device_.copy_rows(shared_down_host_, 0); for (int pattern = 0; pattern < static_cast(kRoutePatterns.size()); ++pattern) { references_.push_back(sparse_moe_oracle(inputs_[pattern], residuals_[pattern], router_, - experts_, shared_gate_host_, shared_down_host_, - kRoutePatterns[pattern])); + experts_, shared_gate_host_, shared_up_host_, + shared_down_host_, kRoutePatterns[pattern])); } } @@ -559,7 +860,8 @@ class SparseMoeFixture { cuda_synchronize(); } - int failures = compare_output(label, destination_storage.values(), reference); + int failures = verify_reduction(label, destination_storage.values(), reference, + profile_.criterion(tokens)); failures += destination_storage.verify_guards(label); failures += verify_exact((label + " input preservation").c_str(), @@ -579,13 +881,18 @@ class SparseMoeFixture { for (const HostExpert& expert : experts_) { failures += routed_gate_.verify_rows( std::string(profile_.name) + " routed gate expert " + std::to_string(expert.id), - expert.gate_up, expert.id * kExpertGateRows); + expert.gate, expert.id * kExpertGateRows); + failures += routed_gate_.verify_rows( + std::string(profile_.name) + " routed up expert " + std::to_string(expert.id), + expert.up, expert.id * kExpertGateRows + kIntermediate); failures += routed_down_.verify_rows( std::string(profile_.name) + " routed down expert " + std::to_string(expert.id), expert.down, expert.id * kHidden); } failures += shared_gate_.verify_rows(std::string(profile_.name) + " shared gate", shared_gate_host_, 0); + failures += shared_gate_.verify_rows(std::string(profile_.name) + " shared up", + shared_up_host_, kIntermediate); failures += shared_down_device_.verify_rows(std::string(profile_.name) + " shared down", shared_down_host_, 0); return failures; @@ -596,14 +903,15 @@ class SparseMoeFixture { std::vector router_; std::vector router_bits_; DeviceBuffer device_router_; - DeviceRowSplit routed_gate_; - DeviceRowSplit routed_down_; - DeviceRowSplit shared_gate_; - DeviceRowSplit shared_down_device_; + ExpertPlane routed_gate_; + ExpertPlane routed_down_; + ExpertPlane shared_gate_; + ExpertPlane shared_down_device_; std::vector> inputs_; std::vector> residuals_; std::vector experts_; quantized_weight::PackedWeight shared_gate_host_; + quantized_weight::PackedWeight shared_up_host_; quantized_weight::PackedWeight shared_down_host_; std::vector> references_; }; @@ -645,11 +953,21 @@ int main() { // and one call crossing the 4096-token internal slice without observing any private plan. constexpr std::array kQ4Q5Tokens{{1, 2, 46, 47, 768, 4097}}; constexpr std::array kQ4Q6Tokens{{1, 2, 46, 47, 768}}; - constexpr std::array kQ8Q8Tokens{{1, 2, 19, 20, 768}}; - const std::array profiles{{ - {"sparse_moe q4+q5 a16", QType::Q4_G64_FP16, QType::Q5_G64_FP16, kQ4Q5Tokens, true}, - {"sparse_moe q4+q6 a16", QType::Q4_G64_FP16, QType::Q6_G64_FP16, kQ4Q6Tokens, false}, - {"sparse_moe q8+q8 a16", QType::Q8_G32_FP16, QType::Q8_G32_FP16, kQ8Q8Tokens, false}, + // 4097 ends on a one-token slice, which is the only way a profile whose prefill starts at + // twenty reaches the route's small-token branch. + constexpr std::array kQ8Q8Tokens{{1, 2, 19, 20, 768, 4097}}; + // NVFP4 walks both frontiers: its Small-T edge and first prefill T, where the activation stops + // being represented, and then the prefill tile boundary and the internal slice. + constexpr std::array kNvfp4Tokens{{1, 2, 12, 13, 64, 768, 4097}}; + const std::array profiles{{ + {"sparse_moe q4+q5 a16", QType::Q4_G64_FP16, QType::Q5_G64_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ4Q5Tokens, true}, + {"sparse_moe q4+q6 a16", QType::Q4_G64_FP16, QType::Q6_G64_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ4Q6Tokens, false}, + {"sparse_moe q8+q8 a16", QType::Q8_G32_FP16, QType::Q8_G32_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ8Q8Tokens, false}, + {"sparse_moe nvfp4", QType::NVFP4, QType::NVFP4, QType::NVFP4, &kSparseMoeNvfp4Tolerance, + &kSparseMoeA4Tolerance, 13, kNvfp4Tokens, true}, }}; int failures = 0;