[3/5] feat(ops): add an NVFP4 expert profile to sparse_moe

bench/ops/quantized_weight.cuh

@@ -10,6 +10,7 @@
 #include <cstdint>
 #include <limits>
 #include <stdexcept>
+#include <vector>
 
 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<std::uint64_t>(n), static_cast<std::uint64_t>(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<std::uint64_t>(divisors) * sizeof(float),
+                            "benchmark NVFP4 payload size overflow");
     if (payload_bytes > std::numeric_limits<std::size_t>::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<std::uint8_t*>(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<float> divisor_words(static_cast<std::size_t>(divisors));
+    for (std::int32_t index = 0; index < divisors; ++index) {
+        divisor_words[static_cast<std::size_t>(index)] =
+            kWeightDivisor * (1.0F + 0.75F * static_cast<float>(index % 5));
+    }
     CUDA_CHECK(cudaMemcpy(static_cast<std::uint8_t*>(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<std::uint8_t*>(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<std::uint8_t*>(result.storage.p) + divisor_offset;
+    weight.weight_divisor_rows  = divisor_rows;
     return result;
 }
 

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<std::uint64_t>(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<std::uint64_t>(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<std::uint64_t>(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<std::size_t>::max()) {
@@ -367,10 +394,13 @@ Options parse_options(int argc, char** argv) {
 }
 
 std::vector<CodecProfile> 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<std::int32_t> 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<std::size_t>(kRouterRows) * kHidden * 2),
-          routed_gate_(bench::make_row_split_weight(
-              gate_codec(profile), kExperts * 1024, kHidden, kHidden,
-              {static_cast<std::uint8_t>(0x31U ^ seed), 0xa5, 0x1401})),
-          routed_down_(bench::make_row_split_weight(
-              down_codec(profile), kExperts * kHidden, kIntermediate, kIntermediate,
-              {static_cast<std::uint8_t>(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<std::uint8_t>(0x31U ^ seed), 0xa5, 0x1401},
+                                         512)),
+          routed_down_(make_expert_plane(
+              down_codec(profile), kExperts * kHidden, kIntermediate,
+              {static_cast<std::uint8_t>(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<std::uint16_t> router(static_cast<std::size_t>(kRouterRows) * kHidden,
                                           bench::f32_to_bf16(0.0F));

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。

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.

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`;