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/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"); 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(