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/bench/ops/sparse_moe_bench.cu b/bench/ops/sparse_moe_bench.cu index 98ffd21e52..a48e855ce7 100644 --- a/bench/ops/sparse_moe_bench.cu +++ b/bench/ops/sparse_moe_bench.cu @@ -46,6 +46,7 @@ enum class CodecProfile : std::uint8_t { Q4Q5, Q4Q6, Q8Q8, + Nvfp4, }; enum class ExpertDistribution : std::uint8_t { @@ -123,12 +124,23 @@ const char* codec_name(CodecProfile profile) { return "q4-q6"; case CodecProfile::Q8Q8: return "q8-q8"; + case CodecProfile::Nvfp4: + return "nvfp4"; } return "unknown"; } QType gate_codec(CodecProfile profile) { - return profile == CodecProfile::Q8Q8 ? QType::Q8_G32_FP16 : QType::Q4_G64_FP16; + switch (profile) { + case CodecProfile::Q4Q5: + case CodecProfile::Q4Q6: + return QType::Q4_G64_FP16; + case CodecProfile::Q8Q8: + return QType::Q8_G32_FP16; + case CodecProfile::Nvfp4: + return QType::NVFP4; + } + throw std::logic_error("unknown SparseMoe codec profile"); } QType down_codec(CodecProfile profile) { @@ -139,10 +151,18 @@ QType down_codec(CodecProfile profile) { return QType::Q6_G64_FP16; case CodecProfile::Q8Q8: return QType::Q8_G32_FP16; + case CodecProfile::Nvfp4: + return QType::NVFP4; } throw std::logic_error("unknown SparseMoe codec profile"); } +// The shared expert follows the routed codec only where a profile stores all four matrices one +// way. The row-split profiles keep it at Q8 whatever they route with. +QType shared_codec(CodecProfile profile) { + return profile == CodecProfile::Nvfp4 ? QType::NVFP4 : QType::Q8_G32_FP16; +} + const char* distribution_name(ExpertDistribution distribution) { switch (distribution) { case ExpertDistribution::TraceLike: @@ -170,6 +190,11 @@ const char* execution_name(Execution execution) { const char* cache_name(CacheState cache) { return cache == CacheState::Cold ? "cold" : "warm"; } std::uint64_t packed_weight_bytes(QType qtype, std::int32_t rows, std::int32_t columns) { + if (qtype == QType::NVFP4) { + // Four bits of code plus one e4m3 byte per sixteen values. + const std::uint64_t elements = static_cast(rows) * columns; + return elements / 2 + elements / 16; + } const std::int32_t group = qtype == QType::Q8_G32_FP16 ? 32 : 64; const std::uint64_t groups = static_cast(rows) * columns / group; const std::uint64_t low = qtype == QType::Q8_G32_FP16 @@ -182,9 +207,10 @@ std::uint64_t packed_weight_bytes(QType qtype, std::int32_t rows, std::int32_t c } double unique_weight_bytes(const Result& result) { + const QType shared = shared_codec(result.codec); const std::uint64_t fixed = static_cast(kRouterRows) * kHidden * 2 + - packed_weight_bytes(QType::Q8_G32_FP16, 1024, kHidden) + - packed_weight_bytes(QType::Q8_G32_FP16, kHidden, kIntermediate); + packed_weight_bytes(shared, 1024, kHidden) + + packed_weight_bytes(shared, kHidden, kIntermediate); const std::uint64_t per_expert = packed_weight_bytes(gate_codec(result.codec), 1024, kHidden) + packed_weight_bytes(down_codec(result.codec), kHidden, kIntermediate); @@ -286,7 +312,8 @@ void usage(const char* argv0) { std::fprintf(stderr, "Usage: %s [options]\n\n" "Public workload:\n" - " --codec q4-q5|q4-q6|q8-q8|all Routed weight profile (default q4-q5).\n" + " --codec q4-q5|q4-q6|q8-q8|nvfp4|all Expert weight profile (default " + "q4-q5).\n" " --tokens T Exact token extent (default 1).\n" " --sweep START:END[:STEP] Public token-extent sweep.\n" " --distribution trace-like|independent|same\n" @@ -356,8 +383,8 @@ Options parse_options(int argc, char** argv) { throw std::invalid_argument("--tokens and --sweep are mutually exclusive"); } if (options.codec != "q4-q5" && options.codec != "q4-q6" && options.codec != "q8-q8" && - options.codec != "all") { - throw std::invalid_argument("--codec must be q4-q5, q4-q6, q8-q8, or all"); + options.codec != "nvfp4" && options.codec != "all") { + throw std::invalid_argument("--codec must be q4-q5, q4-q6, q8-q8, nvfp4, or all"); } if (options.repeat <= 0) { throw std::invalid_argument("--repeat must be positive"); } if (options.flush_bytes > std::numeric_limits::max()) { @@ -367,10 +394,13 @@ Options parse_options(int argc, char** argv) { } std::vector selected_profiles(const std::string& codec) { - if (codec == "all") { return {CodecProfile::Q4Q5, CodecProfile::Q4Q6, CodecProfile::Q8Q8}; } + if (codec == "all") { + return {CodecProfile::Q4Q5, CodecProfile::Q4Q6, CodecProfile::Q8Q8, CodecProfile::Nvfp4}; + } if (codec == "q4-q5") return {CodecProfile::Q4Q5}; if (codec == "q4-q6") return {CodecProfile::Q4Q6}; - return {CodecProfile::Q8Q8}; + if (codec == "q8-q8") return {CodecProfile::Q8Q8}; + return {CodecProfile::Nvfp4}; } std::vector selected_tokens(const TokenSweep& sweep) { @@ -501,20 +531,30 @@ Weight dense_weight(void* data, std::int32_t rows, std::int32_t columns) { return result; } +// `divisor_rows` is the artifact's own stride for that bank: gate and up are quantised apart, so +// the routed gate/up plane carries one divisor per 512 rows, and routed down one per 2048. A shared +// bank is one matrix and keeps a single divisor. Only NVFP4 stores divisors at all. +bench::PackedQuantizedWeight make_expert_plane(QType qtype, std::int32_t n, std::int32_t k, + bench::QuantizedWeightFill fill, + std::int32_t divisor_rows = 0) { + return qtype == QType::NVFP4 ? bench::make_nvfp4_weight(n, k, divisor_rows) + : bench::make_row_split_weight(qtype, n, k, k, fill); +} + class BenchmarkWeights { public: BenchmarkWeights(CodecProfile profile, std::uint32_t seed, std::size_t flush_bytes) : router_(static_cast(kRouterRows) * kHidden * 2), - routed_gate_(bench::make_row_split_weight( - gate_codec(profile), kExperts * 1024, kHidden, kHidden, - {static_cast(0x31U ^ seed), 0xa5, 0x1401})), - routed_down_(bench::make_row_split_weight( - down_codec(profile), kExperts * kHidden, kIntermediate, kIntermediate, - {static_cast(0x59U ^ (seed >> 8)), 0x6d, 0x1403})), - shared_gate_(bench::make_row_split_weight(QType::Q8_G32_FP16, 1024, kHidden, kHidden, - {0x27, 0x00, 0x1405})), - shared_down_(bench::make_row_split_weight(QType::Q8_G32_FP16, kHidden, kIntermediate, - kIntermediate, {0x73, 0x00, 0x1407})), + routed_gate_(make_expert_plane(gate_codec(profile), kExperts * 1024, kHidden, + {static_cast(0x31U ^ seed), 0xa5, 0x1401}, + 512)), + routed_down_(make_expert_plane( + down_codec(profile), kExperts * kHidden, kIntermediate, + {static_cast(0x59U ^ (seed >> 8)), 0x6d, 0x1403}, kIntermediate * 4)), + shared_gate_( + make_expert_plane(shared_codec(profile), 1024, kHidden, {0x27, 0x00, 0x1405})), + shared_down_(make_expert_plane(shared_codec(profile), kHidden, kIntermediate, + {0x73, 0x00, 0x1407})), flush_(flush_bytes) { std::vector router(static_cast(kRouterRows) * kHidden, bench::f32_to_bf16(0.0F)); diff --git a/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_codec.cuh b/src/ops/linear/nvfp4/nvfp4_codec.cuh index b56419b066..9db06234f4 100644 --- a/src/ops/linear/nvfp4/nvfp4_codec.cuh +++ b/src/ops/linear/nvfp4/nvfp4_codec.cuh @@ -23,6 +23,20 @@ __device__ __forceinline__ float decode_nvfp4_e4m3(std::uint8_t storage) { return static_cast(value).x; } +// One e4m3 block scale in the K16M128x4 layout. A 512-byte tile holds 128 rows by four groups, +// with the rows interleaved so that four rows 32 apart share one 16-byte line. `ScaleTilesPerRow` +// is the matrix's column count over 64; a caller that owns rows of a larger plane than a +// registered linear shape supplies it directly. +template +__device__ __forceinline__ std::int64_t nvfp4_scale_byte(int row, int group) { + const int m_tile = row / 128; + const int row_inner = row - m_tile * 128; + const int scale_tile = group / 4; + const int scale_lane = group & 3; + return (static_cast(m_tile) * ScaleTilesPerRow + scale_tile) * 512 + + (row_inner & 31) * 16 + (row_inner >> 5) * 4 + scale_lane; +} + struct alignas(8) Nvfp4QuantizedK16 { std::uint32_t codes_lo; std::uint32_t codes_hi; @@ -60,14 +74,34 @@ pack_nvfp4_e2m1x16(const float2 (&values)[8], std::uint32_t& codes_lo, std::uint "f"(values[6].x), "f"(values[6].y), "f"(values[7].x), "f"(values[7].y)); } -__device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloat16* source, - float input_scale_divisor) { - const uint4 packed0 = load_vec(source); - const uint4 packed1 = load_vec(source + 8); - const std::uint32_t represented[8] = { - packed0.x, packed0.y, packed0.z, packed0.w, packed1.x, packed1.y, packed1.z, packed1.w, - }; +// Half a group, packed by one lane. A fused epilogue splits the sixteen values of a group across +// two neighbouring lanes so that every thread of the block takes part. +__device__ __forceinline__ std::uint32_t pack_nvfp4_e2m1x8(const float2 (&values)[4]) { + std::uint32_t codes = 0; + asm volatile("{\n" + ".reg .b8 b0;\n" + ".reg .b8 b1;\n" + ".reg .b8 b2;\n" + ".reg .b8 b3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n" + "mov.b32 %0, {b0,b1,b2,b3};\n" + "}\n" + : "=r"(codes) + : "f"(values[0].x), "f"(values[0].y), "f"(values[1].x), "f"(values[1].y), + "f"(values[2].x), "f"(values[2].y), "f"(values[3].x), "f"(values[3].y)); + return codes; +} +// Sixteen bf16 values already in registers, for a caller that holds them there rather than writing +// them out for a separate pass to read back. Note that this divides where the sparse-MoE fused +// epilogue multiplies by a reciprocal, so the two can differ by one code at an exact tie. +// Nothing in the engine compares the two planes: each route is checked end to end against an +// oracle that quantises neither. +__device__ __forceinline__ Nvfp4QuantizedK16 +quantize_nvfp4_k16_bits(const std::uint32_t (&represented)[8], float input_scale_divisor) { float2 values[8]; float max_abs = 0.0F; #pragma unroll @@ -92,4 +126,14 @@ __device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloa return result; } +__device__ __forceinline__ Nvfp4QuantizedK16 quantize_nvfp4_k16(const __nv_bfloat16* source, + float input_scale_divisor) { + const uint4 packed0 = load_vec(source); + const uint4 packed1 = load_vec(source + 8); + const std::uint32_t represented[8] = { + packed0.x, packed0.y, packed0.z, packed0.w, packed1.x, packed1.y, packed1.z, packed1.w, + }; + return quantize_nvfp4_k16_bits(represented, input_scale_divisor); +} + } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_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/linear/nvfp4/nvfp4_gemv.cuh b/src/ops/linear/nvfp4/nvfp4_gemv.cuh index 2164c4121f..de5c5df3b6 100644 --- a/src/ops/linear/nvfp4/nvfp4_gemv.cuh +++ b/src/ops/linear/nvfp4/nvfp4_gemv.cuh @@ -85,14 +85,7 @@ stage_nvfp4_scales(const std::uint8_t* __restrict__ scales, template __device__ __forceinline__ std::int64_t nvfp4_scale_offset(int parent_row, int group) { - const int m_tile = parent_row / 128; - const int row_inner = parent_row - m_tile * 128; - const int scale_tile = group / 4; - const int scale_lane = group & 3; - const int row_mod32 = row_inner & 31; - const int row_quartile = row_inner >> 5; - return static_cast(m_tile * Geometry::kScaleTilesPerRow + scale_tile) * 512 + - row_mod32 * 16 + row_quartile * 4 + scale_lane; + return nvfp4_scale_byte(parent_row, group); } template diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh index 2f8abbc227..8914f32389 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh @@ -22,7 +22,9 @@ struct Nvfp4W4a4MmaSchedule { static_assert(WarpsM > 0 && WarpsN > 0); static_assert((BlockM % WarpsM) == 0 && ((BlockM / WarpsM) % 16) == 0); static_assert((BlockN % WarpsN) == 0 && ((BlockN / WarpsN) % 8) == 0); - static_assert(Stages >= 1 && Stages <= 4); + // Six is what the narrow shared-expert schedules take: a tile of sixteen tokens has little + // else to hide its weight stream behind, and the shared-memory budget admits the depth. + static_assert(Stages >= 1 && Stages <= 6); static_assert(MinBlocksPerSm > 0); static constexpr int kBlockM = BlockM; @@ -56,6 +58,32 @@ struct Nvfp4W4a4IdentityRows { } }; +// The activation plane is addressed by token. A dense operator consumes it in order; a routed +// one consumes the same plane through its own column order, so the row a tile stages is a +// policy rather than the tile index. +struct Nvfp4W4a4IdentityTokens { + __device__ __forceinline__ int source_token(int token) const { return token; } + + // A dense operator knows its extent at launch. A routed one learns it per group, so the + // extent is read through the same policy instead of being a second launch argument. + __device__ __forceinline__ int active_tokens(int launched) const { return launched; } +}; + +// Which tile of the grid a CTA takes. A dense operator reads it straight off the block index; +// a routed one folds its own work list into it, so the tile a CTA serves is a policy rather than +// the block index itself. +struct Nvfp4W4a4MmaRasterRowFast { + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = static_cast(blockIdx.x); + block_token = static_cast(blockIdx.y); + } + + // A routed launch sizes its grid from a host-side bound on the work list and learns the real + // count only on the device, so the tiles past the end have to leave before they read a weight + // plane. A dense launch has no such tiles. + __device__ __forceinline__ bool live() const { return true; } +}; + template struct Nvfp4W4a4SharedStorage { alignas( @@ -77,18 +105,18 @@ __device__ __forceinline__ int nvfp4_w4a4_swizzled_byte(int row, int logical_byt return physical_segment * 16 + byte_in_segment; } -template +template __device__ __forceinline__ void stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, Nvfp4W4a4SharedStorage& shared, int stage, int k_tile, - int token_begin, int active_tokens) { + int token_begin, int active_tokens, TokenPolicy token_policy) { constexpr int kCodeTasks = Schedule::kBlockM * Schedule::kSegmentsPerRow; for (int task = static_cast(threadIdx.x); task < kCodeTasks; task += Schedule::kThreads) { const int row = task / Schedule::kSegmentsPerRow; const int logical_segment = task - row * Schedule::kSegmentsPerRow; const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; const int physical_byte = nvfp4_w4a4_swizzled_byte(row, logical_segment * 16); auto* destination = shared.a_codes[stage] + row * Schedule::kCodeRowBytes + physical_byte; const auto* input = source.codes + @@ -103,7 +131,7 @@ stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, row += Schedule::kThreads) { const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; auto* destination = &shared.a_scale4[stage][row * Schedule::kK64PerStage]; const auto* input = source.scales + (static_cast(source_token) * Geometry::kScaleTilesPerRow + @@ -121,7 +149,7 @@ stage_nvfp4_w4a4_activation(Nvfp4W4a4MaterializedActivation source, const int segment = task - row * kSegmentsPerRow; const int token = token_begin + row; const bool valid = token < active_tokens; - const int source_token = valid ? token : 0; + const int source_token = valid ? token_policy.source_token(token) : 0; const int local_k64 = segment * 4; auto* destination = &shared.a_scale4[stage][row * Schedule::kK64PerStage + local_k64]; const auto* input = source.scales + (static_cast(source_token) * @@ -203,21 +231,31 @@ __device__ __forceinline__ void stage_nvfp4_w4a4_weight(const std::uint8_t* __re } template + class RowPolicy = Nvfp4W4a4IdentityRows, bool PairRows = false, + class TokenPolicy = Nvfp4W4a4IdentityTokens, + class RasterPolicy = Nvfp4W4a4MmaRasterRowFast> __global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4_mma_kernel( Nvfp4W4a4MaterializedActivation activation, const std::uint8_t* __restrict__ weight_codes, const std::uint8_t* __restrict__ weight_scales, std::int32_t tokens, float alpha, - Epilogue epilogue, OutputPolicy output, RowPolicy row_policy = {}) { + Epilogue epilogue, OutputPolicy output, RowPolicy row_policy = {}, + TokenPolicy token_policy = {}, RasterPolicy raster_policy = {}) { static_assert((Geometry::kInputRows % Schedule::kBlockK) == 0); static_assert((Geometry::kOutputRows % Schedule::kBlockN) == 0); static_assert(!PairRows || (Schedule::kBlockN % 2) == 0); static_assert(!PairRows || ((Geometry::kOutputRows / 2) % (Schedule::kBlockN / 2)) == 0); __shared__ Nvfp4W4a4SharedStorage shared; - const int token_begin = static_cast(blockIdx.y) * Schedule::kBlockM; + int block_row = 0; + int block_token = 0; + // Liveness first: a dead tile of a routed launch has no work-list entry to read, and reading + // one would touch arena bytes nothing has written. + if (!raster_policy.live()) { return; } + raster_policy.blocks(block_row, block_token); + const int token_begin = block_token * Schedule::kBlockM; constexpr int kRowsPerBlock = PairRows ? Schedule::kBlockN / 2 : Schedule::kBlockN; - const int row_begin = static_cast(blockIdx.x) * kRowsPerBlock; + const int row_begin = block_row * kRowsPerBlock; + const int active = token_policy.active_tokens(tokens); constexpr int kKTiles = Geometry::kInputRows / Schedule::kBlockK; constexpr int kWaitGroups = kKTiles < Schedule::kStages ? kKTiles - 1 : Schedule::kStages - 1; @@ -225,7 +263,7 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 for (int stage = 0; stage < Schedule::kStages; ++stage) { if (stage < kKTiles) { stage_nvfp4_w4a4_activation(activation, shared, stage, stage, - token_begin, tokens); + token_begin, active, token_policy); stage_nvfp4_w4a4_weight(weight_codes, weight_scales, shared, stage, stage, row_begin, row_policy); cp_commit(); @@ -320,7 +358,7 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 const int next_k_tile = k_tile + Schedule::kStages; if (next_k_tile < kKTiles) { stage_nvfp4_w4a4_activation(activation, shared, stage, next_k_tile, - token_begin, tokens); + token_begin, active, token_policy); stage_nvfp4_w4a4_weight(weight_codes, weight_scales, shared, stage, next_k_tile, row_begin, row_policy); } @@ -350,11 +388,11 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 float value01 = accumulators[mma_m][mma_n][1] * alpha; float value10 = accumulators[mma_m][mma_n][2] * alpha; float value11 = accumulators[mma_m][mma_n][3] * alpha; - if (token0 < tokens) { + if (token0 < active) { value00 = epilogue.apply(parent_row0, token0, value00); value01 = epilogue.apply(parent_row1, token0, value01); } - if (token1 < tokens) { + if (token1 < active) { value10 = epilogue.apply(parent_row0, token1, value10); value11 = epilogue.apply(parent_row1, token1, value11); } @@ -365,13 +403,22 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 __syncthreads(); constexpr int kStoredRows = PairRows ? Schedule::kBlockN / 2 : Schedule::kBlockN; constexpr int kVectorsPerRow = kStoredRows / 8; + // An output policy that completes a value across the neighbouring lane - the sparse-MoE + // gate/up epilogue does - needs both lanes of a pair to hold the same token, and they do only + // while a row's vectors come in pairs. A geometry that broke it would deadlock on the shuffle + // rather than fail a test, so it is refused here. This is stricter than the kernel needs: the + // dense output policies pair no lanes and would be safe with an odd count. Every schedule in + // the tree satisfies it, the narrowest being BlockN=32 with PairRows, which gives two. + static_assert((kVectorsPerRow % 2) == 0, + "a lane pair must stay inside one token: give the tile an even number of " + "output vectors per row"); constexpr int kOutputVectors = Schedule::kBlockM * kVectorsPerRow; for (int task = static_cast(threadIdx.x); task < kOutputVectors; task += Schedule::kThreads) { const int token_local = task / kVectorsPerRow; const int row_vector = task - token_local * kVectorsPerRow; const int token = token_begin + token_local; - if (token < tokens) { + if (token < active) { const uint4 values = load_vec(shared_output + token_local * kOutputStride + row_vector * 8); if constexpr (PairRows) { diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 8aa0e1bc62..0fc318bc67 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -6,6 +6,7 @@ #include "ops/common/math.cuh" #include "ops/common/memory.cuh" #include "ops/common/warp.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" #include "ops/linear/q4/q4_rowsplit_storage.cuh" #include "ops/linear/q5/q5_rowsplit_storage.cuh" #include "ops/linear/q6/q6_rowsplit_storage.cuh" @@ -104,98 +105,186 @@ __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; + // One divisor per source matrix the plane was assembled from, each covering `divisor_rows` + // consecutive rows. Only NVFP4 has them. + const float* __restrict__ divisors = nullptr; + int divisor_rows = 1; +}; + +// Codecs come in two lane ownerships. Under `kPackedWord8` a lane owns eight consecutive K values +// 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; + + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, float, + 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; + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } - __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]) { + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, float, + float (&weights)[8]) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( - 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; + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } - __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]) { + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, float, + float (&weights)[8]) { + const std::int64_t index = static_cast(row) * (K / kGroupK) + group; const std::uint32_t packed = *reinterpret_cast( - 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; + + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes&, int) { + return 1.0F; + } + + template + __device__ static __forceinline__ float load_one(const SparseMoePlanes& planes, int row, + int group, int lane) { + 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); } }; +// NVFP4 packs two e2m1 codes per byte and one e4m3 scale per sixteen values, so a lane's eight +// values are four code bytes and exactly one scale: whichever half of a block the lane owns, its +// eight values lie inside that one block. Only the scale plane is swizzled; the codes are plain +// row-major. The per-tensor divisor is folded into the decoded block scale, once per eight values +// rather than once per value. +struct Nvfp4Codec { + static constexpr int kGroupK = 64; + static constexpr bool kPackedWord8 = true; + static constexpr bool kSingleValuePerLane = false; + + // Which stored divisor this row was quantised against. Taken once per row, not per group. + __device__ static __forceinline__ float row_coefficient(const SparseMoePlanes& planes, + int row) { + return __frcp_rn(planes.divisors[row / planes.divisor_rows]); + } + + template + __device__ static __forceinline__ void load_eight(const SparseMoePlanes& planes, int row, + int group, int lane_in_group, float row_scale, + float (&weights)[8]) { + const std::uint32_t packed = *reinterpret_cast( + planes.codes + static_cast(row) * (K / 2) + group * 32 + + lane_in_group * 4); + const std::uint8_t scale = + planes.scales[nvfp4_scale_byte(row, group * 4 + (lane_in_group >> 1))]; + const float coefficient = decode_nvfp4_e4m3(scale) * row_scale; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const float2 code = + decode_nvfp4_e2m1x2(static_cast(packed >> (8 * pair))); + weights[pair * 2] = code.x * coefficient; + weights[pair * 2 + 1] = code.y * coefficient; + } + } +}; + +// A codec declares exactly one lane ownership. The trap below catches a codec that declares +// neither, which would otherwise compile and reduce a zero accumulator. +template +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; + const float scale0 = Codec::row_coefficient(planes, row0); + const float scale1 = Codec::row_coefficient(planes, row1); for (int group_base = first_group; group_base < last_group; group_base += 4) { - const int group = group_base + lane_group; - 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, scale0, weights0); + Codec::template load_eight(planes, row1, group, lane_in_group, scale1, weights1); const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); const float2 x0 = bf16x2_bits_to_float2(input.x); const float2 x1 = bf16x2_bits_to_float2(input.y); @@ -208,41 +297,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 +331,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 +388,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,21 +403,25 @@ __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; + float row_scale[Rows]; +#pragma unroll + for (int row = 0; row < Rows; ++row) { + row_scale[row] = Codec::row_coefficient(planes, row_base + row); + } for (int group_base = first_group; group_base < last_group; group_base += 4) { const int group = group_base + lane_group; const float4 x0 = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); @@ -355,42 +429,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, + row_scale[row], 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 +475,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 +512,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 +542,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 +554,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) { @@ -503,6 +577,14 @@ __global__ void sparse_moe_d4_token_kernel( } } +// How many bytes of the shared-expert down codes the D1 prefetch may touch. A four-bit plane is +// half the size of an eight-bit one, and walking the eight-bit length over it addresses memory the +// artifact never allocated. +unsigned long long shared_down_code_bytes(QType qtype) { + const unsigned long long elements = static_cast(kHidden) * kIntermediate; + return qtype == QType::NVFP4 ? elements / 2 : elements; +} + void launch_d1(const Tensor& x, const SparseMoeWeights& weights, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { sparse_moe_d1_kernel<<>>( @@ -510,24 +592,31 @@ void launch_d1(const Tensor& x, const SparseMoeWeights& weights, static_cast(weights.router_shared_gate.qdata), static_cast(workspace.scratch.data), static_cast(weights.shared_down.qdata), - static_cast(kHidden) * kIntermediate); + shared_down_code_bytes(weights.shared_down.qtype)); CUDA_CHECK(cudaGetLastError()); } -template +SparseMoePlanes matrix_planes(const Weight& weight) { + // Only NVFP4 stores its block scale as a quotient of a divisor, and only it has a plane of + // them; for the row-split codecs the field is not part of the representation at all. The row + // stride is one here so that a row's divisor index is well formed whatever the codec, and the + // codecs without divisors never read it. + return {static_cast(weight.qdata), + static_cast(weight.qhigh), + static_cast(weight.scales), + static_cast(weight.weight_divisors), + weight.weight_divisor_rows > 0 ? weight.weight_divisor_rows : 1}; +} + +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 +630,42 @@ 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; + case QType::NVFP4: + launch_d3_dependent_codec(x, weights, workspace, stream); return; default: throw std::invalid_argument("sparse_moe: unsupported D3 codec"); } } +// Hidden rows one D4 block takes. NVFP4 interleaves rows inside a 512-byte scale tile so that rows +// r and r+1 share one 32-byte sector: taking both halves the sectors the kernel reads from L2, 398k +// to 326k, and the kernel from 10.2 us to 8.0. A third and fourth row start the next sector and buy +// nothing. The row-split planes store a scale per row, so a second row only costs them grid. template +inline constexpr int kDecodeDownRows = 1; +template <> +inline constexpr int kDecodeDownRows = 2; + +template void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destination, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, const void* prefetch_data, std::size_t prefetch_bytes) { - 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 = kDecodeDownRows; 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 +673,85 @@ 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; + case QType::NVFP4: + launch_d4_dependent_codec(weights, destination, workspace, stream, + prefetch_data, prefetch_bytes); return; default: throw std::invalid_argument("sparse_moe: unsupported D4 codec"); } } -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 +759,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 +786,21 @@ 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; + case QType::NVFP4: + launch_d3_small_t_codec( + x, weights, token_ids, token_activations, tokens, schedule, stream, nullptr); return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D3 codec"); @@ -718,30 +815,35 @@ 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; + case QType::NVFP4: + launch_d4_small_t_codec( + weights, destination, token_ids, token_alpha, shared_scale, token_activations, tokens, + schedule, stream, nullptr); return; default: throw std::invalid_argument("sparse_moe: unsupported small-T D4 codec"); diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp index 5412838524..240011839d 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp @@ -25,7 +25,11 @@ SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, QType r (routed_down == QType::Q5_G64_FP16 || routed_down == QType::Q6_G64_FP16); const bool mtp_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - if (!main_profile && !mtp_profile) { + // The decode kernels are CUDA-core dot products, so NVFP4 is read here against represented + // BF16 activations: four-bit activations buy no arithmetic on this path and would cost a + // quantiser kernel of their own. + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + if (!main_profile && !mtp_profile && !nvfp4_profile) { throw std::invalid_argument("sparse_moe: unsupported routed codec profile"); } diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill.h b/src/ops/sparse_moe/prefill/sparse_moe_prefill.h index 09a8d6395f..c3e347a601 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill.h +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill.h @@ -18,9 +18,16 @@ inline constexpr std::int32_t kSparseMoePrefillWorkspaceMin = 20; inline constexpr std::int32_t kSparseMoePrefillQ4Q5Min = 47; inline constexpr std::int32_t kSparseMoePrefillQ4Q6Min = 47; inline constexpr std::int32_t kSparseMoePrefillQ8Q8Min = 20; -inline constexpr std::int32_t kSparseMoePrefillWideMin = 768; -inline constexpr std::int32_t kSparseMoePrefillSliceMax = 4096; -inline constexpr std::int32_t kSparseMoeRouteTileTokens = 8; +// Swept cold against this profile's own small-T route over [2,46], each route measured through the +// operator benchmark at its own schedule. Small-T is ahead through T=8 (79.6 us against 81.9), and +// from T=9 to T=12 it is at most 4.5 % behind - the price of keeping the activation represented, +// which is worth paying for the accuracy: rel_l2 at T=1 is 1.60e-3 against 1.70e-1 for a four-bit +// activation. From T=13 the prefill route's lead becomes real (+9 % at T=16, +29 % at T=32) and +// outgrows that. +inline constexpr std::int32_t kSparseMoePrefillNvfp4Min = 13; +inline constexpr std::int32_t kSparseMoePrefillWideMin = 768; +inline constexpr std::int32_t kSparseMoePrefillSliceMax = 4096; +inline constexpr std::int32_t kSparseMoeRouteTileTokens = 8; // 257 logits padded to a 16-byte-aligned per-token stride. inline constexpr std::int32_t kSparseMoeRouterScoreRows = 260; @@ -28,6 +35,9 @@ struct SparseMoePrefillPlan { std::int32_t tokens = 0; std::int32_t slice_tokens = 0; std::size_t workspace_bytes = 0; + // Whether the route stages its operands in NVFP4. The staging buffers are the profile's own, + // so the flag has to reach the allocation and not only the launch. + bool nvfp4 = false; }; struct SparseMoePrefillWorkspace { @@ -59,17 +69,30 @@ struct SparseMoePrefillWorkspace { // router scores FP32 <-> shared SwiGLU BF16 // gathered X BF16 <-> routed down output BF16 <-> adaptive FP32 activations // (a Q4 routed gate/up stages from x and gathers nothing, so it only ever writes here) - // routed SwiGLU BF16 <-> routed FP32 token reduction + // routed SwiGLU BF16 <-> routed FP32 token reduction <-> every NVFP4 staging plane Tensor score_storage; Tensor shared_activation; Tensor grouped_io; Tensor routed_storage; Tensor routed_sum; + + // The NVFP4 route quantises the chunk once and keeps both SwiGLU intermediates encoded, so + // `down` needs no separate quantiser: 288 bytes a row against the bf16 plane's 1024. All six + // are carved out of `routed_storage`, which this profile never writes -- the encoded + // intermediate replaces the routed SwiGLU plane and the shared-down epilogue replaces the + // token reduction, so the branch returns before either view has a reader. 3744 bytes a token + // inside 8192, so the profile costs no workspace at all. + Tensor nvfp4_input_codes; + Tensor nvfp4_input_scales; + Tensor nvfp4_routed_codes; + Tensor nvfp4_routed_scales; + Tensor nvfp4_shared_codes; + Tensor nvfp4_shared_scales; }; template -SparseMoePrefillWorkspace allocate_sparse_moe_prefill_workspace(Arena& arena, - std::int32_t capacity_tokens) { +SparseMoePrefillWorkspace +allocate_sparse_moe_prefill_workspace(Arena& arena, std::int32_t capacity_tokens, bool nvfp4) { SparseMoePrefillWorkspace out; const std::int32_t assignments = 8 * capacity_tokens; const std::int32_t route_tiles = @@ -104,12 +127,37 @@ SparseMoePrefillWorkspace allocate_sparse_moe_prefill_workspace(Arena& arena, out.routed_storage = arena.alloc(DType::BF16, {512, assignments}, 256); out.routed_sum = Tensor(out.routed_storage.data, DType::FP32, {2048, capacity_tokens}); + + if (nvfp4) { + // The layout builder allocates nothing, so an offset is only formed when there is a + // pointer -- the same guard the index maps above use. + auto* const base = static_cast(out.routed_storage.data); + std::int64_t cursor = 0; + const auto carve = [&](std::int32_t rows, std::int32_t columns) { + std::uint8_t* const at = base != nullptr ? base + cursor : nullptr; + cursor += (static_cast(rows) * columns + 255) / 256 * 256; + return Tensor(at, DType::U8, {rows, columns}); + }; + out.nvfp4_input_codes = carve(1024, capacity_tokens); + out.nvfp4_input_scales = carve(128, capacity_tokens); + out.nvfp4_routed_codes = carve(256, assignments); + out.nvfp4_routed_scales = carve(32, assignments); + out.nvfp4_shared_codes = carve(256, capacity_tokens); + out.nvfp4_shared_scales = carve(32, capacity_tokens); + // 3744 bytes a token, and five alignment gaps of under 256 each, inside the 8192 a token + // that `routed_storage` holds for a profile that has a routed SwiGLU plane. Asserted + // rather than trusted: a wider staging plane has to move the union, not overrun it. The + // eight is `assignments / capacity_tokens`, which is kTopK. + static_assert(1024 + 128 + 8 * 256 + 8 * 32 + 256 + 32 + 5 * 255 <= 512 * 8 * 2, + "the NVFP4 staging planes must fit the routed SwiGLU plane they alias"); + static_cast(cursor); + } return out; } [[nodiscard]] bool sparse_moe_uses_prefill(std::int32_t tokens, QType routed_gate_up, QType routed_down) noexcept; -[[nodiscard]] std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens); +[[nodiscard]] std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens, bool nvfp4); [[nodiscard]] SparseMoePrefillPlan resolve_sparse_moe_prefill_plan(std::int32_t tokens, QType routed_gate_up, QType routed_down); diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu index 3d9fea13fa..974cefae9b 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu @@ -6,6 +6,10 @@ #include "ops/common/memory.cuh" #include "ops/common/mma.cuh" #include "ops/common/rowsplit_mma.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" +#include "ops/linear/nvfp4/nvfp4_geometry.h" +#include "ops/linear/nvfp4/nvfp4_output.cuh" +#include "ops/linear/nvfp4/nvfp4_w4a4_mma.cuh" #include "ops/linear/q4/q4_rowsplit_storage.cuh" #include "ops/linear/q5/q5_rowsplit_storage.cuh" #include "ops/linear/q6/q6_rowsplit_storage.cuh" @@ -37,6 +41,77 @@ constexpr int kRouterStages = 2; constexpr int kRouterWarps = 8; constexpr int kRouterThreads = 32 * kRouterWarps; +constexpr int kRouterSimtThreads = 256; +constexpr int kRouterSimtMaxTokens = 16; + +// One CTA per router row. The tile of the MMA router is 16 rows, which is right for a chunk and +// leaves 17 CTAs for a decode step; this leaves 257 and streams one weight row per CTA. +__global__ __launch_bounds__(kRouterSimtThreads, 4) void sparse_moe_prefill_router_simt_kernel( + const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, + float* __restrict__ scores, int tokens, std::uint8_t* __restrict__ quantized_codes, + std::uint8_t* __restrict__ quantized_scales, float input_scale_divisor) { + constexpr int kWarps = kRouterSimtThreads / 32; + const int row = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + const int lane = tid & 31; + const int warp = tid >> 5; + const auto* w = weight + static_cast(row) * kHidden; + __shared__ float partial[kWarps]; + + for (int token = 0; token < tokens; ++token) { + const auto* xt = x + static_cast(token) * kHidden; + float acc = 0.0F; + for (int i = tid * 8; i < kHidden; i += kRouterSimtThreads * 8) { + const uint4 wv = load_vec(w + i); + const uint4 xv = load_vec(xt + i); + const std::uint32_t wb[4]{wv.x, wv.y, wv.z, wv.w}; + const std::uint32_t xb[4]{xv.x, xv.y, xv.z, xv.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 a = bf16x2_bits_to_float2(wb[j]); + const float2 b = bf16x2_bits_to_float2(xb[j]); + acc = fmaf(a.x, b.x, acc); + acc = fmaf(a.y, b.y, acc); + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + acc += __shfl_xor_sync(0xffffffffU, acc, offset); + } + if (lane == 0) { partial[warp] = acc; } + __syncthreads(); + if (tid == 0) { + float total = 0.0F; +#pragma unroll + for (int i = 0; i < kWarps; ++i) { total += partial[i]; } + scores[static_cast(token) * kSparseMoeRouterScoreRows + row] = total; + } + __syncthreads(); + } + + // The NVFP4 route needs this same input as an encoded plane. It is the tensor this kernel has + // just finished reading and nothing between the two depends on the other, so it is written here + // instead of by a launch of its own. + if (quantized_codes != nullptr) { + constexpr int kGroupsPerRow = kHidden / 16; + const int total = tokens * kGroupsPerRow; + const int stride = static_cast(gridDim.x) * kRouterSimtThreads; + for (int task = static_cast(blockIdx.x) * kRouterSimtThreads + tid; task < total; + task += stride) { + const int token_index = task / kGroupsPerRow; + const int group = task - token_index * kGroupsPerRow; + const Nvfp4QuantizedK16 quantized = quantize_nvfp4_k16( + x + static_cast(token_index) * kHidden + group * 16, + input_scale_divisor); + store_vec(quantized_codes + static_cast(token_index) * (kHidden / 2) + + group * 8, + make_uint2(quantized.codes_lo, quantized.codes_hi)); + quantized_scales[static_cast(token_index) * kGroupsPerRow + group] = + quantized.scale; + } + } +} + __global__ __launch_bounds__(kRouterThreads, 2) void sparse_moe_prefill_router_mma_kernel( const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ weight, float* __restrict__ scores, int tokens) { @@ -288,7 +363,107 @@ constexpr int kExpertStages = 2; constexpr int kGateUpNarrowStages = 6; constexpr int kExpertWarps = 8; constexpr int kExpertThreads = 32 * kExpertWarps; -constexpr int kRtx5090SmCount = 170; + +// An inclusive prefix sum over one block of kExperts lanes. `totals` needs one slot per warp. +__device__ __forceinline__ int sparse_moe_block_inclusive_scan(int value, int* totals, int tid) { + constexpr int kScanWarps = kExperts / 32; + const int lane = tid & 31; + const int warp = tid >> 5; + int running = value; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const int carried = __shfl_up_sync(0xffffffffU, running, offset); + if (lane >= offset) { running += carried; } + } + if (lane == 31) { totals[warp] = running; } + __syncthreads(); + if (warp == 0) { + int total = lane < kScanWarps ? totals[lane] : 0; +#pragma unroll + for (int offset = 1; offset < kScanWarps; offset <<= 1) { + const int carried = __shfl_up_sync(0xffffffffU, total, offset); + if (lane >= offset) { total += carried; } + } + if (lane < kScanWarps) { totals[lane] = total; } + } + __syncthreads(); + return running + (warp == 0 ? 0 : totals[warp - 1]); +} + +// Selection, scan and index as one kernel. Only correct for a single route tile, which is what the +// launch checks: with route_tiles == 1 the tile-local ranks this writes are the global ones. +__global__ __launch_bounds__(kExpertThreads, 1) void sparse_moe_prefill_small_route_kernel( + const float* __restrict__ scores, int* __restrict__ ids, float* __restrict__ alpha, + float* __restrict__ shared_scale, int* __restrict__ local_rank, int* __restrict__ tile_counts, + int* __restrict__ tile_bases, int* __restrict__ expert_offsets, + int* __restrict__ route_job_experts, int* __restrict__ route_job_columns, + int* __restrict__ route_job_count, int* __restrict__ packed_index, + int* __restrict__ packed_token, int job_bn, int tokens) { + __shared__ int counts[kExperts]; + __shared__ int prefix[kExperts]; + __shared__ int warp_totals[kExperts / 32]; + __shared__ float selected_logits[kSparseMoeRouteTileTokens][kTopK]; + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + + // Phase one: the top-k of every token and how many rows each expert collected. + if (tid < kExperts) { counts[tid] = 0; } + __syncthreads(); + if (warp < tokens) { + sparse_moe_select_top8_warp( + scores + static_cast(warp) * kSparseMoeRouterScoreRows, + ids + warp * kTopK, alpha + warp * kTopK, shared_scale + warp, selected_logits[warp]); + __syncwarp(); + if (lane == 0) { +#pragma unroll + for (int rank = 0; rank < kTopK; ++rank) { + const int assignment = warp * kTopK + rank; + local_rank[assignment] = atomicAdd(&counts[ids[assignment]], 1); + } + } + } + __syncthreads(); + + // Phase two: the expert prefix and the work list. One tile, so a base is a cursor. + const int expert = tid; + const int count = counts[expert]; + tile_counts[expert] = count; + prefix[expert] = sparse_moe_block_inclusive_scan(count, warp_totals, tid); + __syncthreads(); + const int base = expert == 0 ? 0 : prefix[expert - 1]; + expert_offsets[expert] = base; + if (expert == kExperts - 1) { expert_offsets[kExperts] = prefix[expert]; } + tile_bases[expert] = base; + + __syncthreads(); + prefix[expert] = + sparse_moe_block_inclusive_scan((count + job_bn - 1) / job_bn, warp_totals, tid); + const int touched = __syncthreads_count(count > 0); + const int job_begin = expert == 0 ? 0 : prefix[expert - 1]; + const int jobs = (count + job_bn - 1) / job_bn; + for (int job = 0; job < jobs; ++job) { + route_job_experts[job_begin + job] = expert; + route_job_columns[job_begin + job] = job * job_bn; + } + if (expert == kExperts - 1) { + route_job_count[0] = prefix[expert]; + route_job_count[1] = touched; + } + + // Phase three: the packed map. The bases above are visible through the barrier. + __syncthreads(); + const int assignments = tokens * kTopK; + for (int assignment = tid; assignment < assignments; assignment += kExpertThreads) { + const int token = assignment / kTopK; + const int packed = tile_bases[ids[assignment]] + local_rank[assignment]; + packed_index[assignment] = packed; + packed_token[packed] = token; + } +} + +constexpr int kRtx5090SmCount = 170; // Upper bound on the persistent grid. The routed GEMMs stride their work list by gridDim.x, // so any grid is correct; this caps the launch when the work list is long. constexpr int kPrefillMaxBlocksPerSm = 32; @@ -1130,6 +1305,348 @@ __global__ __launch_bounds__(kExpertThreads, 1) void sparse_moe_prefill_q8_down_ } } +// --------------------------------------------------------------------------------------------- +// The NVFP4 route. +// +// The routed experts of this model are one contiguous weight plane, so the whole block runs on +// nvfp4_w4a4_mma_kernel with four policies: which expert a tile serves (the scan's work list), +// which weight rows it stages, which activation columns it reads, and what the epilogue writes. +// Three things differ from the Q4/Q5 route beyond the codec: +// +// * the chunk's activation plane is quantised once, up front, and both the routed and the +// shared gate/up read it; +// * the gate/up epilogue emits NVFP4 rather than bf16, so `down` needs no separate quantiser +// and the intermediate plane is 288 bytes a row instead of 1024; +// * the same kernels serve every token count this route is entered for, down to the one-token +// ragged tail of a sliced call - decode itself runs elsewhere. They are launched over the +// scan's work list, so a single token touches eight experts and reads eight experts' +// weights -- the tile wastes compute on its empty columns, which decode does not pay for +// because it is bound by the weight stream. +// +// The routed tile schedules are the ones the operator bench settled on: a 128-deep tile, +// because the 256-deep one takes 55 KiB of staging buffers and only one CTA then fits on an +// SM. The dense shared gate/up keeps the 256-deep tile, where one CTA is the right answer. +// --------------------------------------------------------------------------------------------- + +using Nvfp4RoutedGateUpGeometry = Nvfp4Geometry; +using Nvfp4RoutedDownGeometry = Nvfp4Geometry; +using Nvfp4SharedGateUpGeometry = Nvfp4Geometry<2 * kIntermediate, kHidden>; +using Nvfp4SharedDownGeometry = Nvfp4Geometry; +using Nvfp4InputGeometry = Nvfp4ActivationGeometry; + +using Nvfp4RoutedSchedule = Nvfp4W4a4MmaSchedule<64, 128, 128, 2, 4, 2, 3>; +using Nvfp4DenseSchedule = Nvfp4W4a4MmaSchedule<64, 128, 256, 4, 4, 2, 1>; + +template +struct Nvfp4ScheduleTag { + using type = Schedule; +}; + +// The shared expert is dense, so its grid is the output rows over the tile width: at one token the +// schedules above leave 8 CTAs for gate/up and 16 for down on a 170-SM card, and the pair measures +// 162 GB/s where the routed pair reaches 597. A narrower tile multiplies the CTA count without +// changing the bytes a CTA reads, and BlockM 16 stops staging 64 token rows when one exists. +// gate/up reaches 32 because `Nvfp4SharedGateUpRows` is not contiguous and stages scales per +// (row, k64); down stays at 64, the narrowest the contiguous scale path allows. +using Nvfp4SharedGateUpSmallTSchedule = Nvfp4W4a4MmaSchedule<16, 32, 128, 1, 4, 6, 4>; +// Four stages, not six: this geometry has 512 input rows over a 128-deep tile, so there are +// only four k-tiles to prefetch, and the extra two stages would buy nothing and cost 11.5 KiB +// of shared memory a block. +using Nvfp4SharedDownSmallTSchedule = Nvfp4W4a4MmaSchedule<16, 64, 128, 1, 2, 4, 4>; +// Below this many tokens the wide tile cannot fill the machine and the narrow one cannot lose. +constexpr int kNvfp4SharedSmallTokens = 32; + +// The routed pair keeps BlockM at 64 because the work list is cut into column blocks of that +// width; only N narrows, which is what decides how many CTAs a job expands into. +using Nvfp4RoutedGateUpSmallTSchedule = Nvfp4W4a4MmaSchedule<64, 32, 128, 2, 2, 4, 1>; +using Nvfp4RoutedDownSmallTSchedule = Nvfp4W4a4MmaSchedule<64, 64, 128, 2, 2, 2, 4>; +// Reached only by the ragged tail of a sliced call: this profile enters prefill at thirteen +// tokens, so a whole call never has two. +constexpr int kNvfp4RoutedSmallTokens = 2; + +// One divisor per source matrix, taken from the row the value belongs to. The routed planes are +// stacks of separately quantised matrices - gate and up of an expert were quantised apart, so the +// stride is half an expert's gate/up rows - and a stack of one carries its reciprocal directly, +// which is what every plane did before there were stacks. +// +// This runs once per output element, so the index may not cost an integer division: `shift` is the +// stride's base-two logarithm where it has one, and every stride a source can give this +// architecture does. +struct Nvfp4SourceDivisorEpilogue { + const float* __restrict__ divisors; + float uniform; + int divisor_rows; + int shift; + + __device__ __forceinline__ float apply(std::int32_t row, std::int32_t, float value) const { + if (divisors == nullptr) { return value * uniform; } + const int index = shift >= 0 ? (row >> shift) : (row / divisor_rows); + return value * __frcp_rn(divisors[index]); + } +}; + +// A stacked plane reads its divisors per row; a plane with one carries its reciprocal and never +// touches the plane of them. +Nvfp4SourceDivisorEpilogue nvfp4_divisor_epilogue(const Weight& weight) { + const int rows = weight.weight_divisor_rows; + if (rows == weight.n) { return {nullptr, 1.0F / weight.weight_scale_divisor, rows, -1}; } + const int shift = (rows > 0 && (rows & (rows - 1)) == 0) ? __builtin_ctz(rows) : -1; + return {static_cast(weight.weight_divisors), 0.0F, rows, shift}; +} + +// The scan emits one job per column tile of this width, so it has to match the tile the routed +// schedules stage. +static_assert(Nvfp4RoutedSchedule::kBlockM == 64); +// A route job is one column tile of one expert, so every routed schedule's token tile has to be the +// job width exactly: wider, and one job would cover two tiles of which only one would run. +constexpr int kNvfp4JobColumns = Nvfp4RoutedSchedule::kBlockM; +static_assert(Nvfp4RoutedGateUpSmallTSchedule::kBlockM == kNvfp4JobColumns); +static_assert(Nvfp4RoutedDownSmallTSchedule::kBlockM == kNvfp4JobColumns); + +// The routed grid is sized from a host-side bound on the work list; the real count lives on the +// device, so tiles past the end leave before they touch a weight plane. +struct Nvfp4Jobs { + const int* __restrict__ experts; + const int* __restrict__ columns; + const int* __restrict__ count; + + __device__ __forceinline__ int expert() const { return experts[blockIdx.y]; } + + __device__ __forceinline__ int column_base() const { return columns[blockIdx.y]; } + + __device__ __forceinline__ bool live() const { return static_cast(blockIdx.y) < count[0]; } +}; + +template +struct Nvfp4RoutedRaster { + Nvfp4Jobs jobs; + + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = static_cast(blockIdx.x); + block_token = jobs.column_base() / BlockM; + } + + __device__ __forceinline__ bool live() const { return jobs.live(); } +}; + +// `down` weight rows are contiguous inside an expert, so folding the expert's row base into the +// block row lets the stock identity row policy and the stock contiguous scale staging address +// the plane without knowing about routing. +template +struct Nvfp4RoutedDownRaster { + Nvfp4Jobs jobs; + + __device__ __forceinline__ void blocks(int& block_row, int& block_token) const { + block_row = jobs.expert() * (kHidden / BlockN) + static_cast(blockIdx.x); + block_token = jobs.column_base() / BlockM; + } + + __device__ __forceinline__ bool live() const { return jobs.live(); } +}; + +template +struct Nvfp4ExpertGateUpRows { + static constexpr bool kContiguous = false; + static constexpr int kRowsPerBranch = RowsPerBranch; + + Nvfp4Jobs jobs; + + __device__ __forceinline__ int weight_row(int row_begin, int local_row) const { + const int within = row_begin + (local_row & (kRowsPerBranch - 1)) + + (local_row >= kRowsPerBranch ? kIntermediate : 0); + return jobs.expert() * (2 * kIntermediate) + within; + } +}; + +// The chunk's activation plane is shared by every expert; the route decides which rows this +// expert reads and how many there are. +struct Nvfp4RoutedGatherTokens { + const int* __restrict__ packed_token; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ int source_token(int column) const { + return packed_token[expert_offsets[jobs.expert()] + column]; + } + + __device__ __forceinline__ int active_tokens(int) const { + const int expert = jobs.expert(); + return expert_offsets[expert + 1] - expert_offsets[expert]; + } +}; + +// `down` reads what gate/up wrote, which is already in route order, so this is an offset rather +// than a lookup. +struct Nvfp4RoutedPackedTokens { + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ int source_token(int column) const { + return expert_offsets[jobs.expert()] + column; + } + + __device__ __forceinline__ int active_tokens(int) const { + const int expert = jobs.expert(); + return expert_offsets[expert + 1] - expert_offsets[expert]; + } +}; + +union Nvfp4Bf16Pair { + unsigned bits; + __nv_bfloat162 values; +}; + +__device__ __forceinline__ unsigned nvfp4_swiglu_pair(unsigned gate_bits, unsigned up_bits) { + Nvfp4Bf16Pair gate{gate_bits}; + Nvfp4Bf16Pair up{up_bits}; + const float2 g = __bfloat1622float2(gate.values); + const float2 u = __bfloat1622float2(up.values); + Nvfp4Bf16Pair result; + result.values = __floats2bfloat162_rn(silu(g.x) * u.x, silu(g.y) * u.y); + return result.bits; +} + +// One store is half an NVFP4 group, so every thread of the block takes part instead of half of +// them; the group's maximum is completed across the neighbouring lane. The two lanes of a pair +// always share a token, so they are active together. +__device__ __forceinline__ void nvfp4_store_half_group(std::uint8_t* codes, std::uint8_t* scales, + std::int64_t row_stride_codes, + std::int64_t row_stride_scales, + std::int64_t packed, std::int32_t row, + uint4 gate, uint4 up, float divisor) { + const std::uint32_t bits[4] = { + nvfp4_swiglu_pair(gate.x, up.x), + nvfp4_swiglu_pair(gate.y, up.y), + nvfp4_swiglu_pair(gate.z, up.z), + nvfp4_swiglu_pair(gate.w, up.w), + }; + float2 values[4]; + float max_abs = 0.0F; +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + values[pair] = bf16x2_bits_to_float2(bits[pair]); + max_abs = fmaxf(max_abs, fabsf(values[pair].x)); + max_abs = fmaxf(max_abs, fabsf(values[pair].y)); + } + // The two lanes of a group are always both here or both gone, because they share a token, so + // the mask is known and does not have to be asked for. `__activemask` is not composable with a + // shuffle: if the compiler reconverges differently it can name a lane that has left. + const unsigned pair_mask = 3U << (threadIdx.x & 31U & ~1U); + max_abs = fmaxf(max_abs, __shfl_xor_sync(pair_mask, max_abs, 1)); + + const float scale_unencoded = __fdiv_rn(divisor * max_abs, 6.0F); + const std::uint8_t scale = __nv_cvt_float_to_fp8(scale_unencoded, __NV_SATFINITE, __NV_E4M3); + std::uint32_t word = 0; + if (scale != 0) { + const float reciprocal = __frcp_rn(decode_nvfp4_e4m3(scale)); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + values[pair].x = values[pair].x * divisor * reciprocal; + values[pair].y = values[pair].y * divisor * reciprocal; + } + word = pack_nvfp4_e2m1x8(values); + } + const int half = (row >> 3) & 1; + const int group_row = row & ~15; + store_vec(codes + packed * row_stride_codes + (group_row >> 1) + half * 4, word); + if (half == 0) { scales[packed * row_stride_scales + (group_row >> 4)] = scale; } +} + +struct Nvfp4RoutedGateUpOutput { + std::uint8_t* codes; + std::uint8_t* scales; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + float divisor; + + __device__ __forceinline__ void store_pair_vector(std::int32_t row, std::int32_t column, + uint4 gate, uint4 up) const { + const std::int64_t packed = expert_offsets[jobs.expert()] + column; + nvfp4_store_half_group(codes, scales, kIntermediate / 2, kIntermediate / 16, packed, row, + gate, up, divisor); + } +}; + +struct Nvfp4RoutedDownOutput { + __nv_bfloat16* data; + const int* __restrict__ expert_offsets; + Nvfp4Jobs jobs; + + __device__ __forceinline__ void store_vector(std::int32_t row, std::int32_t column, + uint4 values) const { + const std::int64_t packed = expert_offsets[jobs.expert()] + column; + store_vec(data + packed * kHidden + (row & (kHidden - 1)), values); + } +}; + +template +struct Nvfp4SharedGateUpRows { + static constexpr bool kContiguous = false; + static constexpr int kRowsPerBranch = RowsPerBranch; + + __device__ __forceinline__ int weight_row(int row_begin, int local_row) const { + return row_begin + (local_row & (kRowsPerBranch - 1)) + + (local_row >= kRowsPerBranch ? kIntermediate : 0); + } +}; + +struct Nvfp4SharedGateUpOutput { + std::uint8_t* codes; + std::uint8_t* scales; + float divisor; + + __device__ __forceinline__ void store_pair_vector(std::int32_t row, std::int32_t column, + uint4 gate, uint4 up) const { + nvfp4_store_half_group(codes, scales, kIntermediate / 2, kIntermediate / 16, column, row, + gate, up, divisor); + } +}; + +// The block's last write: the residual it was handed, plus the routed sum the reduce kernel +// built, plus the shared expert scaled by its gate. +struct Nvfp4SharedDownOutput { + __nv_bfloat16* destination; + const __nv_bfloat16* __restrict__ grouped_output; + const int* __restrict__ packed_index; + const float* __restrict__ alpha; + const float* __restrict__ shared_scale; + + // The routed sum is formed here rather than by a kernel of its own. This epilogue already + // visits every (token, hidden) element once and already had to read the sum back, so reading + // the eight expert rows instead costs the same traffic and saves a graph node per layer. + __device__ __forceinline__ void store_vector(std::int32_t row, std::int32_t column, + uint4 values) const { + const std::int64_t base = static_cast(column) * kHidden + row; + const float gain = shared_scale[column]; + float routed[8] = {}; +#pragma unroll + for (int route = 0; route < kTopK; ++route) { + const int packed = packed_index[column * kTopK + route]; + const float weight = alpha[column * kTopK + route]; + const uint4 raw = + load_vec(grouped_output + static_cast(packed) * kHidden + row); + const std::uint32_t words[4]{raw.x, raw.y, raw.z, raw.w}; +#pragma unroll + for (int index = 0; index < 4; ++index) { + const Nvfp4Bf16Pair slot{words[index]}; + const float2 decoded = __bfloat1622float2(slot.values); + routed[2 * index] += weight * decoded.x; + routed[2 * index + 1] += weight * decoded.y; + } + } + const Nvfp4Bf16Pair pair[4] = {{values.x}, {values.y}, {values.z}, {values.w}}; +#pragma unroll + for (int index = 0; index < 4; ++index) { + const float2 shared = __bfloat1622float2(pair[index].values); + const std::int64_t at = base + 2 * index; + destination[at] = __float2bfloat16_rn(__bfloat162float(destination[at]) + + routed[2 * index] + gain * shared.x); + destination[at + 1] = __float2bfloat16_rn(__bfloat162float(destination[at + 1]) + + routed[2 * index + 1] + gain * shared.y); + } + } +}; + union alignas(16) SparseMoeBf16x8 { uint4 raw; __nv_bfloat162 pair[4]; @@ -1173,6 +1690,146 @@ __global__ void sparse_moe_prefill_reduce_kernel(const __nv_bfloat16* __restrict } } +struct Nvfp4PrefillPlanes { + std::uint8_t* input_codes; + std::uint8_t* input_scales; + std::uint8_t* routed_codes; + std::uint8_t* routed_scales; + std::uint8_t* shared_codes; + std::uint8_t* shared_scales; +}; + +void launch_sparse_moe_prefill_nvfp4(const __nv_bfloat16* input, const SparseMoeWeights& weights, + __nv_bfloat16* destination, int tokens, int assignments, + int max_route_jobs, const int* packed_token, + const int* offsets, const int* route_job_experts, + const int* route_job_columns, const int* route_job_count, + const int* packed_index, const float* alpha, + const float* shared_scale, __nv_bfloat16* grouped_io, + float* /*routed_sum*/, Nvfp4PrefillPlanes planes, + cudaStream_t stream) { + const Nvfp4Jobs jobs{route_job_experts, route_job_columns, route_job_count}; + + // One quantisation of the chunk serves both the routed and the shared gate/up. Below the + // small-token threshold the router has already produced it on its way out. + if (tokens > kRouterSimtMaxTokens) { + constexpr int kThreads = 256; + const int groups = tokens * (kHidden / 16); + nvfp4_w4a4_quantize_kernel + <<<(groups + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + input, planes.input_codes, planes.input_scales, tokens, tokens, + weights.routed_gate_up.input_scale_divisor); + CUDA_CHECK(cudaGetLastError()); + } + + const Nvfp4W4a4MaterializedActivation chunk{planes.input_codes, planes.input_scales}; + const Nvfp4W4a4MaterializedActivation routed_middle{planes.routed_codes, planes.routed_scales}; + const Nvfp4W4a4MaterializedActivation shared_middle{planes.shared_codes, planes.shared_scales}; + + { + // Alpha keeps only what the whole plane shares; the row's own divisor is the epilogue's. + const float scale = 1.0F / weights.routed_gate_up.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.routed_gate_up); + const Nvfp4RoutedGatherTokens token_policy{packed_token, offsets, jobs}; + const Nvfp4RoutedGateUpOutput output{planes.routed_codes, planes.routed_scales, offsets, + jobs, weights.routed_down.input_scale_divisor}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + constexpr int kPairRows = Schedule::kBlockN / 2; + using Rows = Nvfp4ExpertGateUpRows; + using Raster = Nvfp4RoutedRaster; + const dim3 grid(kIntermediate / kPairRows, max_route_jobs); + nvfp4_w4a4_mma_kernel<<>>( + chunk, static_cast(weights.routed_gate_up.qdata), + static_cast(weights.routed_gate_up.scales), assignments, scale, + epilogue, output, Rows{jobs}, token_policy, Raster{jobs}); + }; + if (tokens <= kNvfp4RoutedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + { + const float scale = 1.0F / weights.routed_down.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.routed_down); + const Nvfp4RoutedPackedTokens token_policy{offsets, jobs}; + const Nvfp4RoutedDownOutput output{grouped_io, offsets, jobs}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + using Raster = Nvfp4RoutedDownRaster; + const dim3 grid(kHidden / Schedule::kBlockN, max_route_jobs); + nvfp4_w4a4_mma_kernel + <<>>( + routed_middle, static_cast(weights.routed_down.qdata), + static_cast(weights.routed_down.scales), assignments, + scale, epilogue, output, Nvfp4W4a4IdentityRows{}, token_policy, Raster{jobs}); + }; + if (tokens <= kNvfp4RoutedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + + { + const float scale = 1.0F / weights.shared_gate_up.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.shared_gate_up); + const Nvfp4SharedGateUpOutput output{planes.shared_codes, planes.shared_scales, + weights.shared_down.input_scale_divisor}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + constexpr int kPairRows = Schedule::kBlockN / 2; + using Rows = Nvfp4SharedGateUpRows; + const dim3 grid(kIntermediate / kPairRows, + (tokens + Schedule::kBlockM - 1) / Schedule::kBlockM); + nvfp4_w4a4_mma_kernel + <<>>( + chunk, static_cast(weights.shared_gate_up.qdata), + static_cast(weights.shared_gate_up.scales), tokens, scale, + epilogue, output, Rows{}); + }; + if (tokens < kNvfp4SharedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } + + { + const float scale = 1.0F / weights.shared_down.input_scale_divisor; + const Nvfp4SourceDivisorEpilogue epilogue = nvfp4_divisor_epilogue(weights.shared_down); + const Nvfp4SharedDownOutput output{destination, grouped_io, packed_index, alpha, + shared_scale}; + const auto launch = [&](auto tag) { + using Schedule = typename decltype(tag)::type; + const dim3 grid(kHidden / Schedule::kBlockN, + (tokens + Schedule::kBlockM - 1) / Schedule::kBlockM); + nvfp4_w4a4_mma_kernel<<>>( + shared_middle, static_cast(weights.shared_down.qdata), + static_cast(weights.shared_down.scales), tokens, scale, + epilogue, output); + }; + if (tokens < kNvfp4SharedSmallTokens) { + launch(Nvfp4ScheduleTag{}); + } else { + launch(Nvfp4ScheduleTag{}); + } + CUDA_CHECK(cudaGetLastError()); + } +} + } // namespace void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, @@ -1223,41 +1880,71 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, const int route_tiles = (tokens + kSparseMoeRouteTileTokens - 1) / kSparseMoeRouteTileTokens; const int assignments = tokens * kTopK; + const bool nvfp4 = weights.routed_gate_up.qtype == QType::NVFP4; const int adaptive_last = weights.routed_down.qtype == QType::Q5_G64_FP16 ? 51 : weights.routed_down.qtype == QType::Q6_G64_FP16 ? 52 : 0; - const bool adaptive = tokens >= 47 && tokens <= adaptive_last; - - sparse_moe_prefill_router_mma_kernel<<>>(input, router, scores, - tokens); + const bool adaptive = !nvfp4 && tokens >= 47 && tokens <= adaptive_last; + + // Only this profile: the small-token router folds the NVFP4 quantisation of the chunk + // into its epilogue, and its reduction order differs from the MMA router's, so a + // groupwise call would change scores for nothing. + if (nvfp4 && tokens <= kRouterSimtMaxTokens) { + sparse_moe_prefill_router_simt_kernel<<>>( + input, router, scores, tokens, + static_cast(workspace.nvfp4_input_codes.data), + static_cast(workspace.nvfp4_input_scales.data), + weights.routed_gate_up.input_scale_divisor); + } else { + sparse_moe_prefill_router_mma_kernel<<>>(input, router, + scores, tokens); + } CUDA_CHECK(cudaGetLastError()); - sparse_moe_prefill_select_count_kernel<<>>( - scores, ids, alpha, shared_scale, local_rank, tile_counts, tokens); + const bool routed_gate_up_q4 = weights.routed_gate_up.qtype == QType::Q4_G64_FP16; + // Like the Q4 route, NVFP4 stages its activation tile from the chunk through the route + // order, so it needs the packed token map rather than a materialised gather. + const bool needs_packed_token = routed_gate_up_q4 || nvfp4; + // One route tile means one block can carry selection, scan and index together, which is + // two graph nodes fewer per layer. Restricted to this profile: the Q4 route would take it + // on the ragged tail of a sliced call, where it replaces three kernels whose result it has + // to reproduce exactly, and that is a separate claim with a separate gate. + const bool fused_route = nvfp4 && tokens <= kSparseMoeRouteTileTokens; + if (!fused_route) { + sparse_moe_prefill_select_count_kernel<<>>( + scores, ids, alpha, shared_scale, local_rank, tile_counts, tokens); + } CUDA_CHECK(cudaGetLastError()); - const bool wide_plan = tokens >= kSparseMoePrefillWideMin; - const int route_job_bn = wide_plan ? 64 : 32; + const bool wide_plan = tokens >= kSparseMoePrefillWideMin; + // The NVFP4 tile is 64 columns wide at every token count, so its jobs have to be too. + const int route_job_bn = nvfp4 ? kNvfp4JobColumns : (wide_plan ? 64 : 32); // The scan emits one route job per nonempty column tile of an expert, so it cannot // emit more than one job per full tile of assignments plus one tail per expert -- // the same bound the workspace is sized by. Each job expands into row blocks, and // sizing the grid from that product keeps a persistent block on one work item while // there are fewer work items than the cap, instead of a fixed count that has to // iterate. The exact job count only exists on the device. - const int max_route_jobs = assignments / route_job_bn + kExperts; + // A routed token reaches one expert per assignment, so the tail term is bounded by the + // assignment count as well as by the expert count. Narrowing it launches fewer blocks that + // would exit at once; it is numerically inert, and it is applied only here because on the + // groupwise routes it was measured at zero and is not this change's to make. + const int max_route_jobs = + assignments / route_job_bn + (nvfp4 ? std::min(kExperts, assignments) : kExperts); const int routed_gate_work = max_route_jobs * (kIntermediate / (kExpertBM / 2)); const int routed_down_work = max_route_jobs * (kHidden / kExpertBM); const int routed_gate_blocks = std::min(routed_gate_work, kPrefillMaxBlocks); const int routed_down_blocks = std::min(routed_down_work, kPrefillMaxBlocks); - sparse_moe_prefill_scan_kernel<<<1, kExpertThreads, 0, stream>>>( - tile_counts, tile_bases, offsets, route_job_experts, route_job_columns, route_job_count, - route_tiles, route_job_bn, tokens, adaptive); + if (!fused_route) { + sparse_moe_prefill_scan_kernel<<<1, kExpertThreads, 0, stream>>>( + tile_counts, tile_bases, offsets, route_job_experts, route_job_columns, + route_job_count, route_tiles, route_job_bn, tokens, adaptive); + } CUDA_CHECK(cudaGetLastError()); - const bool routed_gate_up_q4 = weights.routed_gate_up.qtype == QType::Q4_G64_FP16; - const int index_blocks = (assignments + kExpertThreads - 1) / kExpertThreads; + const int index_blocks = (assignments + kExpertThreads - 1) / kExpertThreads; if (adaptive) { auto* adaptive_activations = reinterpret_cast(grouped_io); sparse_moe_decode_launch_d3_small_t(input_slice, weights, ids, adaptive_activations, @@ -1271,15 +1958,37 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, sparse_moe_prefill_index_kernel<<>>( ids, local_rank, packed_index, tile_bases, packed_token, assignments, route_job_count); - } else if (routed_gate_up_q4) { - sparse_moe_prefill_index_kernel<<>>( - ids, local_rank, packed_index, tile_bases, packed_token, assignments, nullptr); + } else if (needs_packed_token) { + if (fused_route) { + sparse_moe_prefill_small_route_kernel<<<1, kExpertThreads, 0, stream>>>( + scores, ids, alpha, shared_scale, local_rank, tile_counts, tile_bases, offsets, + route_job_experts, route_job_columns, route_job_count, packed_index, + packed_token, route_job_bn, tokens); + } else { + sparse_moe_prefill_index_kernel<<>>( + ids, local_rank, packed_index, tile_bases, packed_token, assignments, nullptr); + } } else { sparse_moe_prefill_gather_kernel<<>>( input, ids, local_rank, packed_index, tile_bases, grouped_io, nullptr); } CUDA_CHECK(cudaGetLastError()); + if (nvfp4) { + const Nvfp4PrefillPlanes planes{ + static_cast(workspace.nvfp4_input_codes.data), + static_cast(workspace.nvfp4_input_scales.data), + static_cast(workspace.nvfp4_routed_codes.data), + static_cast(workspace.nvfp4_routed_scales.data), + static_cast(workspace.nvfp4_shared_codes.data), + static_cast(workspace.nvfp4_shared_scales.data)}; + launch_sparse_moe_prefill_nvfp4( + input, weights, output, tokens, assignments, max_route_jobs, packed_token, offsets, + route_job_experts, route_job_columns, route_job_count, packed_index, alpha, + shared_scale, grouped_io, routed_sum, planes, stream); + continue; + } + const dim3 routed_gate_grid(kIntermediate / (kExpertBM / 2), kExperts); // Same predicate that decided whether packed_token was written above. if (routed_gate_up_q4) { diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp b/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp index 1bba6ed4d8..0ee4c3c27c 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill_plan.cpp @@ -17,6 +17,9 @@ std::int32_t prefill_min_tokens(QType routed_gate_up, QType routed_down) noexcep if (routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16) { return kSparseMoePrefillQ8Q8Min; } + if (routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4) { + return kSparseMoePrefillNvfp4Min; + } return 0; } @@ -28,13 +31,13 @@ bool sparse_moe_uses_prefill(std::int32_t tokens, QType routed_gate_up, return minimum != 0 && tokens >= minimum; } -std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens) { +std::size_t sparse_moe_prefill_workspace_bytes(std::int32_t max_tokens, bool nvfp4) { if (max_tokens < kSparseMoePrefillWorkspaceMin) { throw std::invalid_argument("sparse_moe prefill: max_tokens must be at least 20"); } const std::int32_t capacity_tokens = std::min(max_tokens, kSparseMoePrefillSliceMax); WorkspaceLayoutBuilder layout; - (void)allocate_sparse_moe_prefill_workspace(layout, capacity_tokens); + (void)allocate_sparse_moe_prefill_workspace(layout, capacity_tokens, nvfp4); return layout.peak_bytes(1); } @@ -48,8 +51,14 @@ SparseMoePrefillPlan resolve_sparse_moe_prefill_plan(std::int32_t tokens, QType throw std::invalid_argument("sparse_moe prefill: unsupported token count"); } - const std::int32_t slice_tokens = std::min(tokens, kSparseMoePrefillSliceMax); - return {tokens, slice_tokens, sparse_moe_prefill_workspace_bytes(tokens)}; + // The workspace query has a floor of twenty tokens, and NVFP4 enters prefill at thirteen, so a + // call in [13,19] has to allocate to that floor or the execution high-water mark falls short of + // what the query promised. Slicing is unaffected: the loop still steps by the real token + // count. + const std::int32_t slice_tokens = + std::max(std::min(tokens, kSparseMoePrefillSliceMax), kSparseMoePrefillWorkspaceMin); + const bool nvfp4 = routed_gate_up == QType::NVFP4; + return {tokens, slice_tokens, sparse_moe_prefill_workspace_bytes(slice_tokens, nvfp4), nvfp4}; } } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp b/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp index 6c483106dc..757609b0d7 100644 --- a/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp +++ b/src/ops/sparse_moe/small_t/sparse_moe_small_t_plan.cpp @@ -30,11 +30,24 @@ SparseMoeSmallTPlan resolve_sparse_moe_small_t_plan(std::int32_t tokens, QType r (routed_down == QType::Q5_G64_FP16 || routed_down == QType::Q6_G64_FP16); const bool mtp_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - if (!main_profile && !mtp_profile) { + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + if (!main_profile && !mtp_profile && !nvfp4_profile) { throw std::invalid_argument("sparse_moe small-T: unsupported routed codec profile"); } SparseMoeSmallTPlan plan{tokens, sparse_moe_small_t_workspace_bytes(tokens)}; + if (nvfp4_profile) { + // Swept over the whole [2,46] window against all nine pairs, cold, with the prefill + // frontier held open; the figures below are that sweep's, and the frontier comment in + // sparse_moe_prefill.h quotes a later one, so the two differ in the third digit. One path + // per warp and two output rows per block win everywhere NVFP4 keeps the window: 34.8 us at + // T=2 and 79.9 at T=8, against 38.4 and 90.1 for the three-path, one-row pair. Rows4 + // overtakes Rows2 somewhere between T=16 and T=32, which is beyond the frontier this + // profile stops at, so it has no interval of its own. + plan.d3_schedule = SparseMoeSmallTD3Schedule::Paths9; + plan.d4_schedule = SparseMoeSmallTD4Schedule::Rows2; + return plan; + } if (mtp_profile) { plan.d3_schedule = tokens <= 5 ? SparseMoeSmallTD3Schedule::Paths1 : SparseMoeSmallTD3Schedule::Paths9; diff --git a/src/ops/wrapper/sparse_moe.cpp b/src/ops/wrapper/sparse_moe.cpp index aa8eb1dc65..ae769bc29d 100644 --- a/src/ops/wrapper/sparse_moe.cpp +++ b/src/ops/wrapper/sparse_moe.cpp @@ -2,6 +2,7 @@ #include "ninfer/ops/sparse_moe.h" #include "core/nvtx.h" +#include "ops/linear/nvfp4/nvfp4_format.h" #include "ops/sparse_moe/decode/sparse_moe_decode.h" #include "ops/sparse_moe/prefill/sparse_moe_prefill.h" #include "ops/sparse_moe/small_t/sparse_moe_small_t.h" @@ -137,8 +138,48 @@ void require_quantized(const Weight& weight, std::int32_t n, std::int32_t k, con ranges.push_back(address_range(weight.scales, scale_bytes, std::string(name) + " scales")); } +// The NVFP4 profile is all four expert matrices together: the gate/up epilogue emits the +// encoded intermediate `down` consumes, so the two cannot be mixed with a groupwise codec. +bool nvfp4_profile(const SparseMoeWeights& weights) { + return weights.routed_gate_up.qtype == QType::NVFP4 && + weights.routed_down.qtype == QType::NVFP4 && + weights.shared_gate_up.qtype == QType::NVFP4 && + weights.shared_down.qtype == QType::NVFP4; +} + +void require_nvfp4(const Weight& weight, std::int32_t n, std::int32_t k, const char* name, + std::vector& ranges) { + require_matrix_metadata(weight, n, k, name); + // Every expert matrix was quantised on its own, so each of these planes is a stack. + const detail::Nvfp4WeightGeometry geometry = + detail::validate_nvfp4_weight(weight, (std::string("sparse_moe ") + name).c_str(), true); + ranges.push_back( + address_range(weight.qdata, geometry.code_plane_bytes, std::string(name) + " code")); + ranges.push_back( + address_range(weight.scales, geometry.scale_plane_bytes, std::string(name) + " scales")); +} + void validate_weights(const SparseMoeWeights& weights, std::vector& ranges) { require_router(weights.router_shared_gate, ranges); + if (nvfp4_profile(weights)) { + require_nvfp4(weights.routed_gate_up, kRoutedGateRows, kHidden, "routed_gate_up", ranges); + require_nvfp4(weights.routed_down, kRoutedDownRows, kIntermediate, "routed_down", ranges); + require_nvfp4(weights.shared_gate_up, kSharedGateRows, kHidden, "shared_gate_up", ranges); + require_nvfp4(weights.shared_down, kHidden, kIntermediate, "shared_down", ranges); + // The prefill route encodes the hidden state once and both gate/up GEMMs read that one + // plane, each dividing by its own bank's activation divisor. Two different divisors would + // scale the shared expert's whole contribution by their ratio, silently. + if (weights.routed_gate_up.input_scale_divisor != + weights.shared_gate_up.input_scale_divisor) { + throw std::invalid_argument( + "sparse_moe: routed and shared experts must share one activation divisor"); + } + return; + } + if (weights.routed_gate_up.qtype == QType::NVFP4 || weights.routed_down.qtype == QType::NVFP4 || + weights.shared_gate_up.qtype == QType::NVFP4 || weights.shared_down.qtype == QType::NVFP4) { + throw std::invalid_argument("sparse_moe: NVFP4 requires all four expert matrices"); + } if (weights.routed_gate_up.qtype != QType::Q4_G64_FP16 && weights.routed_gate_up.qtype != QType::Q8_G32_FP16) { throw std::invalid_argument("sparse_moe: routed_gate_up must be Q4 or Q8"); @@ -169,10 +210,12 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout const bool q8_profile = routed_gate_up == QType::Q8_G32_FP16 && routed_down == QType::Q8_G32_FP16; - const std::int32_t prefill_first = - q8_profile ? detail::kSparseMoePrefillQ8Q8Min - : (routed_down == QType::Q5_G64_FP16 ? detail::kSparseMoePrefillQ4Q5Min - : detail::kSparseMoePrefillQ4Q6Min); + const bool nvfp4_profile = routed_gate_up == QType::NVFP4 && routed_down == QType::NVFP4; + const std::int32_t prefill_first = q8_profile ? detail::kSparseMoePrefillQ8Q8Min + : nvfp4_profile ? detail::kSparseMoePrefillNvfp4Min + : routed_down == QType::Q5_G64_FP16 + ? detail::kSparseMoePrefillQ4Q5Min + : detail::kSparseMoePrefillQ4Q6Min; std::size_t required = 0; if (min_tokens == 1) { required = detail::sparse_moe_decode_workspace_bytes(); } @@ -186,7 +229,12 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout const std::int32_t prefill_interval_first = std::max(min_tokens, prefill_first); if (prefill_interval_first <= max_tokens) { - required = std::max(required, detail::sparse_moe_prefill_workspace_bytes(max_tokens)); + // The prefill workspace query has a floor of its own, and the plan allocates to it, so an + // interval whose top is below that floor still has to promise what execution will take. + required = + std::max(required, detail::sparse_moe_prefill_workspace_bytes( + std::max(max_tokens, detail::kSparseMoePrefillWorkspaceMin), + nvfp4_profile)); } return required; } @@ -246,7 +294,7 @@ void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilo const detail::SparseMoePrefillPlan plan = detail::resolve_sparse_moe_prefill_plan( tokens, weights.routed_gate_up.qtype, weights.routed_down.qtype); const detail::SparseMoePrefillWorkspace views = - detail::allocate_sparse_moe_prefill_workspace(workspace, plan.slice_tokens); + detail::allocate_sparse_moe_prefill_workspace(workspace, plan.slice_tokens, plan.nvfp4); detail::sparse_moe_prefill_launch(x, weights, destination, plan, views, stream); return; } diff --git a/tests/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/tests/ops/test_sparse_moe.cpp b/tests/ops/test_sparse_moe.cpp index 35dc23ca07..da9f8de427 100644 --- a/tests/ops/test_sparse_moe.cpp +++ b/tests/ops/test_sparse_moe.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,52 @@ constexpr ReductionCriterion kSparseMoeA16Tolerance{ /*gross_relative_to_max_reference*/ 0.0, }; +// NVFP4 below its prefill frontier reads represented BF16 activations, and its weights decode into +// exactly the values the oracle is built from, so the remaining error is the BF16 rounding of the +// destination plus the FP32 accumulation of the dot product. Both terms are relative: the absolute +// term above is a number from the row-split fixture, whose references peak near 0.3 where these +// peak near 1154, so it bounds a different quantity. +// +// The two terms are bounded differently, and the second is why the gross bound is not the same +// number as the L2 one. Over a whole vector the roundings are independent and average down, which +// is what 2.5e-3 of the reference RMS describes. One element does not average: the accumulation +// can carry the exact value across a rounding boundary, and the destination then lands on the +// wrong side of it. +// +// BF16 keeps eight significand bits, so inside a binade the spacing is 2^(e-7) and the relative +// ulp runs over (2^-8, 2^-7]. Half an ulp is therefore at most 2^-8 of the value - 3.9e-3 - and +// since no element exceeds the maximum reference, 4.0e-3 of that maximum bounds the rounding of +// any one of them. That is the format's number, not this fixture's. +// +// The case that showed the old bound was below it: at T=1 the reference peaks at 665, so 2.5e-3 +// set the limit at 1.66, and an element whose exact value is 637.701 lands on 636 - the nearest +// BF16 there, the grid stepping by 4 - for an error of 1.70. The destination was right and the +// bound was not. +// +// Measured over the represented cases: rel-L2 2.09e-3 against 2.5e-3, worst element 1.79 against +// the 2.66 this bound allows. The L2 term is the binding one and stays where it was. +constexpr ReductionCriterion kSparseMoeNvfp4Tolerance{ + /*relative_l2*/ 2.5e-3, + /*gross_absolute*/ 0.0, + /*gross_relative_to_max_reference*/ 4.0e-3, +}; + +// The NVFP4 prefill route quantises activations as well as weights, so the A16 bound above does +// not apply to it. Against an oracle built from the same represented weights the remaining error +// is the four-bit activation plane, and a dot product does not average it away: the sum of K terms +// carries the same cancellation as the error does, so the output keeps the per-element relative +// error of e2m1 rather than that error over the square root of K. Both terms are relative for the +// same reason -- an absolute bound would mean nothing across cases whose references differ by an +// order of magnitude. It also depends on how far the experts spread in magnitude, which is a +// property of the fixture rather than of the route: every source matrix here carries its own +// divisor over a fourfold range, narrower than the 3.3x to 11.9x a real plane spans. Measured +// maximum over the prefill cases is 1.011e-1, and the worst element is half the gross bound. +constexpr ReductionCriterion kSparseMoeA4Tolerance{ + /*relative_l2*/ 1.8e-1, + /*gross_absolute*/ 0.0, + /*gross_relative_to_max_reference*/ 2.4e-1, +}; + constexpr std::size_t kOutputGuardBytes = 256; constexpr std::uint8_t kOutputGuardByte = 0xa5; @@ -81,11 +128,6 @@ std::vector bf16_bits(const std::vector& values) { return bits; } -int compare_output(const std::string& label, const std::vector& actual, - const std::vector& reference) { - return verify_reduction(label, actual, reference, kSparseMoeA16Tolerance); -} - class GuardedBf16Output { public: explicit GuardedBf16Output(std::size_t words) @@ -105,6 +147,126 @@ class GuardedBf16Output { std::size_t words_; }; +// NVFP4 keeps its two planes in one payload at a fixed geometry -- codes from the start, +// scales at the 256-aligned end of them, and after those one weight divisor for each separately +// quantised matrix stacked into the plane -- so this cannot reuse the two independent buffers a +// row-split weight has. +class DeviceNvfp4 { +public: + DeviceNvfp4(std::int32_t rows, std::int32_t columns, std::int32_t divisor_rows) + : rows_(rows), columns_(columns), divisor_rows_(checked_divisor_rows(rows, divisor_rows)), + code_plane_bytes_(static_cast(rows) * columns / 2), + scale_offset_((code_plane_bytes_ + 255) / 256 * 256), + scale_plane_bytes_(static_cast(rows) * columns / 16), + storage_(scale_offset_ + scale_plane_bytes_ + + static_cast(rows / divisor_rows) * sizeof(float)) { + storage_.fill(0); + } + + // The member initialiser divides by this, so the check cannot wait for the body. + static std::int32_t checked_divisor_rows(std::int32_t rows, std::int32_t divisor_rows) { + if (divisor_rows <= 0 || rows % divisor_rows != 0 || divisor_rows % 128 != 0) { + throw std::invalid_argument("sparse_moe test: invalid NVFP4 divisor stride"); + } + return divisor_rows; + } + + void copy_rows(const quantized_weight::PackedWeight& source, std::int32_t destination_row) { + const std::int32_t source_rows = source.weight.n; + if (source.weight.qtype != QType::NVFP4 || source.weight.k != columns_ || source_rows < 1 || + (source_rows % 128) != 0 || destination_row < 0 || (destination_row % 128) != 0 || + destination_row > rows_ - source_rows) { + throw std::invalid_argument("sparse_moe test: invalid NVFP4 row copy"); + } + const std::size_t code_bytes = static_cast(source_rows) * columns_ / 2; + const std::size_t scale_bytes = static_cast(source_rows) * columns_ / 16; + storage_.copy_from_host(source.payload.data(), code_bytes, + static_cast(destination_row) * columns_ / 2); + // A scale tile spans 128 rows, and both the source and the destination row are multiples + // of that, so the tile index scales with the row exactly. + storage_.copy_from_host(source.payload.data() + source.scale_plane_offset, scale_bytes, + scale_offset_ + + static_cast(destination_row) * columns_ / 16); + if (source_rows > divisor_rows_ || destination_row % divisor_rows_ != 0) { + throw std::invalid_argument("sparse_moe test: a copy may not span two divisors"); + } + storage_.copy_from_host(source.payload.data() + source.weight_divisor_offset, sizeof(float), + scale_offset_ + scale_plane_bytes_ + + static_cast(destination_row / divisor_rows_) * + sizeof(float)); + if (destination_row == 0) { weight_scale_divisor_ = source.weight.weight_scale_divisor; } + input_scale_divisor_ = source.weight.input_scale_divisor; + } + + Weight weight() const { + Weight result{}; + result.payload = storage_.p; + result.payload_bytes = storage_.bytes; + result.high_plane_bytes = 0; + result.qtype = QType::NVFP4; + result.group_size = 16; + result.qdata = storage_.p; + result.qhigh = nullptr; + result.scales = static_cast(storage_.p) + scale_offset_; + result.n = rows_; + result.k = columns_; + result.group = 16; + result.layout = QuantLayout::BlockScaleK16M128x4; + result.scale_dtype = DType::FP8_E4M3FN; + result.ndim = 2; + result.shape[0] = rows_; + result.shape[1] = columns_; + result.padded_shape[0] = rows_; + result.padded_shape[1] = columns_; + result.weight_scale_divisor = weight_scale_divisor_; + result.input_scale_divisor = input_scale_divisor_; + result.weight_divisors = + static_cast(storage_.p) + scale_offset_ + scale_plane_bytes_; + result.weight_divisor_rows = divisor_rows_; + return result; + } + + int verify_rows(const std::string& label, const quantized_weight::PackedWeight& source, + std::int32_t destination_row) const { + const std::int32_t source_rows = source.weight.n; + const std::size_t code_bytes = static_cast(source_rows) * columns_ / 2; + const std::size_t scale_bytes = static_cast(source_rows) * columns_ / 16; + int failures = 0; + failures += + verify_span(label + " code", source.payload.data(), + static_cast(destination_row) * columns_ / 2, code_bytes); + failures += verify_span( + label + " scale", source.payload.data() + source.scale_plane_offset, + scale_offset_ + static_cast(destination_row) * columns_ / 16, scale_bytes); + failures += verify_span( + label + " divisor", source.payload.data() + source.weight_divisor_offset, + scale_offset_ + scale_plane_bytes_ + + static_cast(destination_row / divisor_rows_) * sizeof(float), + sizeof(float)); + return failures; + } + +private: + int verify_span(const std::string& label, const std::uint8_t* expected, std::size_t offset, + std::size_t bytes) const { + std::vector actual(bytes); + storage_.copy_to_host(actual.data(), bytes, offset); + if (std::memcmp(actual.data(), expected, bytes) == 0) { return 0; } + std::cerr << label << ": persistent weight was modified\n"; + return 1; + } + + std::int32_t rows_; + std::int32_t columns_; + std::int32_t divisor_rows_; + std::size_t code_plane_bytes_; + std::size_t scale_offset_; + std::size_t scale_plane_bytes_; + DeviceBuffer storage_; + float weight_scale_divisor_ = 1.0F; + float input_scale_divisor_ = 1.0F; +}; + class DeviceRowSplit { public: DeviceRowSplit(QType qtype, std::int32_t rows, std::int32_t columns) @@ -323,9 +485,19 @@ std::vector make_down(std::int32_t rows, std::int32_t columns, std::uint3 return source; } +// Gate and up are quantised apart in every published checkpoint, so each is packed on its own +// with its own divisor and the two are stacked into one plane. This takes one of them out of the +// generator's combined matrix. +std::vector take_rows(const std::vector& matrix, std::int32_t columns, + std::int32_t begin, std::int32_t rows) { + const auto first = matrix.begin() + static_cast(begin) * columns; + return std::vector(first, first + static_cast(rows) * columns); +} + struct HostExpert { int id; - quantized_weight::PackedWeight gate_up; + quantized_weight::PackedWeight gate; + quantized_weight::PackedWeight up; quantized_weight::PackedWeight down; }; @@ -373,7 +545,8 @@ std::vector sparse_moe_oracle(const std::vector& input, const std::vector& residual, const std::vector& router, const std::vector& experts, - const quantized_weight::PackedWeight& shared_gate_up, + const quantized_weight::PackedWeight& shared_gate, + const quantized_weight::PackedWeight& shared_up, const quantized_weight::PackedWeight& shared_down, const RoutePattern& intended_route) { const std::vector x(input.begin(), input.end()); @@ -406,16 +579,16 @@ std::vector sparse_moe_oracle(const std::vector& input, routed_activation[route].resize(kIntermediate); const HostExpert& expert = find_expert(experts, selected[route]); parallel_rows(kIntermediate, [&](std::int32_t row) { - const double gate = dot_fp64(expert.gate_up.dequant, row, kHidden, x); - const double up = dot_fp64(expert.gate_up.dequant, kIntermediate + row, kHidden, x); + const double gate = dot_fp64(expert.gate.dequant, row, kHidden, x); + const double up = dot_fp64(expert.up.dequant, row, kHidden, x); routed_activation[route][row] = (gate / (1.0 + std::exp(-gate))) * up; }); } std::vector shared_activation(kIntermediate); parallel_rows(kIntermediate, [&](std::int32_t row) { - const double gate = dot_fp64(shared_gate_up.dequant, row, kHidden, x); - const double up = dot_fp64(shared_gate_up.dequant, kIntermediate + row, kHidden, x); + const double gate = dot_fp64(shared_gate.dequant, row, kHidden, x); + const double up = dot_fp64(shared_up.dequant, row, kHidden, x); shared_activation[row] = (gate / (1.0 + std::exp(-gate))) * up; }); @@ -438,19 +611,121 @@ struct CodecProfile { const char* name; QType routed_gate_up; QType routed_down; + QType shared; + const ReductionCriterion* tolerance; + // The bound from the token count at which this profile starts quantising its activations, and + // that token count. Zero where no route of the profile does. + const ReductionCriterion* quantized_activation_tolerance; + std::int32_t quantized_activation_from; std::span token_cases; bool verify_graph_replay; + + [[nodiscard]] const ReductionCriterion& criterion(std::int32_t tokens) const { + return quantized_activation_from != 0 && tokens >= quantized_activation_from + ? *quantized_activation_tolerance + : *tolerance; + } }; +// The four expert matrices take whichever storage the profile names; nothing else in the +// fixture depends on which. +class ExpertPlane { +public: + ExpertPlane(QType qtype, std::int32_t rows, std::int32_t columns, std::int32_t divisor_rows) { + if (qtype == QType::NVFP4) { + nvfp4_.emplace(rows, columns, divisor_rows); + } else { + row_split_.emplace(qtype, rows, columns); + } + } + + void copy_rows(const quantized_weight::PackedWeight& source, std::int32_t destination_row) { + if (nvfp4_) { + nvfp4_->copy_rows(source, destination_row); + } else { + row_split_->copy_rows(source, destination_row); + } + } + + Weight weight() const { return nvfp4_ ? nvfp4_->weight() : row_split_->weight(); } + + int verify_rows(const std::string& label, const quantized_weight::PackedWeight& source, + std::int32_t destination_row) const { + return nvfp4_ ? nvfp4_->verify_rows(label, source, destination_row) + : row_split_->verify_rows(label, source, destination_row); + } + +private: + std::optional nvfp4_; + std::optional row_split_; +}; + +float decode_e4m3_scale(std::uint8_t bits) { + const int exponent = (bits >> 3) & 15; + const int mantissa = bits & 7; + const float value = exponent == 0 + ? std::ldexp(static_cast(mantissa) / 8.0F, -6) + : std::ldexp(1.0F + static_cast(mantissa) / 8.0F, exponent - 7); + return ((bits >> 7) & 1) != 0 ? -value : value; +} + +// The oracle needs the values the weight represents. Rather than quantise a second time and risk +// disagreeing with the packer, decode the payload the packer produced: by construction those are +// the values, and the decode is the format's own definition. +std::vector decode_nvfp4_payload(const quantized_weight::PackedWeight& packed, + std::int32_t n, std::int32_t k) { + static constexpr float kMagnitude[8]{0.0F, 0.5F, 1.0F, 1.5F, 2.0F, 3.0F, 4.0F, 6.0F}; + std::vector out(static_cast(n) * k); + const std::int32_t k_tiles = k / 64; + const float inverse = 1.0F / packed.weight.weight_scale_divisor; + for (std::int32_t row = 0; row < n; ++row) { + const std::int32_t row_tile = row / 128; + const std::int32_t row_inner = row % 128; + for (std::int32_t column = 0; column < k; ++column) { + const std::uint8_t byte = + packed.payload[static_cast(row) * k / 2 + column / 2]; + const unsigned nibble = (column & 1) != 0 ? (byte >> 4) : (byte & 0x0fU); + const std::int32_t group = column / 16; + const std::size_t scale_at = + static_cast(packed.scale_plane_offset) + + static_cast(row_tile * k_tiles + group / 4) * 512U + + static_cast(row_inner % 32) * 16U + + static_cast(row_inner / 32) * 4U + static_cast(group % 4); + const float magnitude = kMagnitude[nibble & 7U]; + out[static_cast(row) * k + column] = + ((nibble & 8U) != 0 ? -magnitude : magnitude) * + decode_e4m3_scale(packed.payload[scale_at]) * inverse; + } + } + return out; +} + +// A patterned NVFP4 matrix carries its own values; the groupwise path quantises the fixture's. +// Either way the oracle reads the represented values this returns. +quantized_weight::PackedWeight pack_expert(QType qtype, const std::vector& values, + std::int32_t n, std::int32_t k, std::uint32_t seed, + float weight_divisor = 0.125F) { + if (qtype == QType::NVFP4) { + quantized_weight::PatternedWeightOptions options; + options.weight_scale_divisor = weight_divisor; + options.input_scale_divisor = 3.5F; + quantized_weight::PackedWeight packed = + quantized_weight::make_patterned_weight(QType::NVFP4, n, k, seed, options); + packed.dequant = decode_nvfp4_payload(packed, n, k); + return packed; + } + return quantized_weight::pack_row_split_lowbit(values, n, k, qtype); +} + class SparseMoeFixture { public: explicit SparseMoeFixture(const CodecProfile& profile) : profile_(profile), router_(make_router()), router_bits_(bf16_bits(router_)), device_router_(to_device(router_bits_)), - routed_gate_(profile.routed_gate_up, kRoutedGateRows, kHidden), - routed_down_(profile.routed_down, kRoutedDownRows, kIntermediate), - shared_gate_(QType::Q8_G32_FP16, kSharedGateRows, kHidden), - shared_down_device_(QType::Q8_G32_FP16, kHidden, kIntermediate) { + routed_gate_(profile.routed_gate_up, kRoutedGateRows, kHidden, kIntermediate), + routed_down_(profile.routed_down, kRoutedDownRows, kIntermediate, kHidden), + shared_gate_(profile.shared, kSharedGateRows, kHidden, kIntermediate), + shared_down_device_(profile.shared, kHidden, kIntermediate, kHidden) { for (int pattern = 0; pattern < static_cast(kRoutePatterns.size()); ++pattern) { inputs_.push_back(make_input(pattern)); residuals_.push_back(make_residual(pattern)); @@ -465,32 +740,58 @@ class SparseMoeFixture { } } std::sort(expert_ids.begin(), expert_ids.end()); + // A published checkpoint quantises every expert matrix on its own, so no two of them share + // a divisor. Here they differ by construction, which is what makes a lookup that ignores + // the row visible to the oracle. + const auto expert_divisor = [](int expert, int matrix) { + return 0.0625F * (1.0F + static_cast((expert * 7 + matrix * 3) % 13) * 0.25F); + }; for (int expert : expert_ids) { - const float factor = 0.8f + static_cast((expert * 3) % 11) * 0.045f; - auto gate_up = quantized_weight::pack_row_split_lowbit( - make_gate_up(kExpertGateRows, kHidden, 100U + static_cast(expert), - factor), - kExpertGateRows, kHidden, profile.routed_gate_up); - auto down = quantized_weight::pack_row_split_lowbit( - make_down(kHidden, kIntermediate, 300U + static_cast(expert), - factor), - kHidden, kIntermediate, profile.routed_down); - routed_gate_.copy_rows(gate_up, expert * kExpertGateRows); + const float factor = 0.8f + static_cast((expert * 3) % 11) * 0.045f; + const std::uint32_t seed = 100U + static_cast(expert); + const std::vector gate_up_source = + make_gate_up(kExpertGateRows, kHidden, seed, factor); + auto gate = pack_expert(profile.routed_gate_up, + take_rows(gate_up_source, kHidden, 0, kIntermediate), + kIntermediate, kHidden, seed, expert_divisor(expert, 0)); + // A different seed, not just a different divisor: the NVFP4 packer derives its codes + // from the seed, so one seed for both halves would leave them bit-identical and a route + // that read the wrong half would look right. + auto up = + pack_expert(profile.routed_gate_up, + take_rows(gate_up_source, kHidden, kIntermediate, kIntermediate), + kIntermediate, kHidden, seed ^ 0x9e3779b9U, expert_divisor(expert, 1)); + auto down = + pack_expert(profile.routed_down, + make_down(kHidden, kIntermediate, + 300U + static_cast(expert), factor), + kHidden, kIntermediate, 300U + static_cast(expert), + expert_divisor(expert, 2)); + routed_gate_.copy_rows(gate, expert * kExpertGateRows); + routed_gate_.copy_rows(up, expert * kExpertGateRows + kIntermediate); routed_down_.copy_rows(down, expert * kHidden); - experts_.push_back({expert, std::move(gate_up), std::move(down)}); + experts_.push_back({expert, std::move(gate), std::move(up), std::move(down)}); } - shared_gate_host_ = quantized_weight::pack_q8_g32_row_split( - make_gate_up(kSharedGateRows, kHidden, 0x512U, 0.93f), kSharedGateRows, kHidden); - shared_down_host_ = quantized_weight::pack_q8_g32_row_split( - make_down(kHidden, kIntermediate, 0x731U, 0.87f), kHidden, kIntermediate); + const std::vector shared_source = + make_gate_up(kSharedGateRows, kHidden, 0x512U, 0.93f); + shared_gate_host_ = + pack_expert(profile.shared, take_rows(shared_source, kHidden, 0, kIntermediate), + kIntermediate, kHidden, 0x512U, 0.3125F); + shared_up_host_ = pack_expert( + profile.shared, take_rows(shared_source, kHidden, kIntermediate, kIntermediate), + kIntermediate, kHidden, 0x513U, 0.4375F); + shared_down_host_ = + pack_expert(profile.shared, make_down(kHidden, kIntermediate, 0x731U, 0.87f), kHidden, + kIntermediate, 0x731U); shared_gate_.copy_rows(shared_gate_host_, 0); + shared_gate_.copy_rows(shared_up_host_, kIntermediate); shared_down_device_.copy_rows(shared_down_host_, 0); for (int pattern = 0; pattern < static_cast(kRoutePatterns.size()); ++pattern) { references_.push_back(sparse_moe_oracle(inputs_[pattern], residuals_[pattern], router_, - experts_, shared_gate_host_, shared_down_host_, - kRoutePatterns[pattern])); + experts_, shared_gate_host_, shared_up_host_, + shared_down_host_, kRoutePatterns[pattern])); } } @@ -559,7 +860,8 @@ class SparseMoeFixture { cuda_synchronize(); } - int failures = compare_output(label, destination_storage.values(), reference); + int failures = verify_reduction(label, destination_storage.values(), reference, + profile_.criterion(tokens)); failures += destination_storage.verify_guards(label); failures += verify_exact((label + " input preservation").c_str(), @@ -579,13 +881,18 @@ class SparseMoeFixture { for (const HostExpert& expert : experts_) { failures += routed_gate_.verify_rows( std::string(profile_.name) + " routed gate expert " + std::to_string(expert.id), - expert.gate_up, expert.id * kExpertGateRows); + expert.gate, expert.id * kExpertGateRows); + failures += routed_gate_.verify_rows( + std::string(profile_.name) + " routed up expert " + std::to_string(expert.id), + expert.up, expert.id * kExpertGateRows + kIntermediate); failures += routed_down_.verify_rows( std::string(profile_.name) + " routed down expert " + std::to_string(expert.id), expert.down, expert.id * kHidden); } failures += shared_gate_.verify_rows(std::string(profile_.name) + " shared gate", shared_gate_host_, 0); + failures += shared_gate_.verify_rows(std::string(profile_.name) + " shared up", + shared_up_host_, kIntermediate); failures += shared_down_device_.verify_rows(std::string(profile_.name) + " shared down", shared_down_host_, 0); return failures; @@ -596,14 +903,15 @@ class SparseMoeFixture { std::vector router_; std::vector router_bits_; DeviceBuffer device_router_; - DeviceRowSplit routed_gate_; - DeviceRowSplit routed_down_; - DeviceRowSplit shared_gate_; - DeviceRowSplit shared_down_device_; + ExpertPlane routed_gate_; + ExpertPlane routed_down_; + ExpertPlane shared_gate_; + ExpertPlane shared_down_device_; std::vector> inputs_; std::vector> residuals_; std::vector experts_; quantized_weight::PackedWeight shared_gate_host_; + quantized_weight::PackedWeight shared_up_host_; quantized_weight::PackedWeight shared_down_host_; std::vector> references_; }; @@ -645,11 +953,21 @@ int main() { // and one call crossing the 4096-token internal slice without observing any private plan. constexpr std::array kQ4Q5Tokens{{1, 2, 46, 47, 768, 4097}}; constexpr std::array kQ4Q6Tokens{{1, 2, 46, 47, 768}}; - constexpr std::array kQ8Q8Tokens{{1, 2, 19, 20, 768}}; - const std::array profiles{{ - {"sparse_moe q4+q5 a16", QType::Q4_G64_FP16, QType::Q5_G64_FP16, kQ4Q5Tokens, true}, - {"sparse_moe q4+q6 a16", QType::Q4_G64_FP16, QType::Q6_G64_FP16, kQ4Q6Tokens, false}, - {"sparse_moe q8+q8 a16", QType::Q8_G32_FP16, QType::Q8_G32_FP16, kQ8Q8Tokens, false}, + // 4097 ends on a one-token slice, which is the only way a profile whose prefill starts at + // twenty reaches the route's small-token branch. + constexpr std::array kQ8Q8Tokens{{1, 2, 19, 20, 768, 4097}}; + // NVFP4 walks both frontiers: its Small-T edge and first prefill T, where the activation stops + // being represented, and then the prefill tile boundary and the internal slice. + constexpr std::array kNvfp4Tokens{{1, 2, 12, 13, 64, 768, 4097}}; + const std::array profiles{{ + {"sparse_moe q4+q5 a16", QType::Q4_G64_FP16, QType::Q5_G64_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ4Q5Tokens, true}, + {"sparse_moe q4+q6 a16", QType::Q4_G64_FP16, QType::Q6_G64_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ4Q6Tokens, false}, + {"sparse_moe q8+q8 a16", QType::Q8_G32_FP16, QType::Q8_G32_FP16, QType::Q8_G32_FP16, + &kSparseMoeA16Tolerance, nullptr, 0, kQ8Q8Tokens, false}, + {"sparse_moe nvfp4", QType::NVFP4, QType::NVFP4, QType::NVFP4, &kSparseMoeNvfp4Tolerance, + &kSparseMoeA4Tolerance, 13, kNvfp4Tokens, true}, }}; int failures = 0; 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(