From 115a257171208c31825a8e42ccb16856bfdd3d3c Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 02:04:42 -0700 Subject: [PATCH 1/5] Add --torch-dtype to convert for bf16-safe model loading The convert step hard-coded torch.float16 for HF model loading, which corrupts models whose weights exceed fp16 range (Gemma 4). bf16 loads them safely; quantization math is unchanged. Default stays float16. Signed-off-by: Noah Cylich --- python/cactus/convert/cli.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 68c4f4916..aa7655a24 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -43,7 +43,9 @@ def _hf_cache_dir() -> str | None: return os.environ.get("HF_HUB_CACHE") or os.environ.get("HF_HOME") or None -def _load_hf(model_id_or_path: str, device: str): +def _load_hf(model_id_or_path: str, device: str, torch_dtype: str = "float16"): + load_dtype = {"float16": torch.float16, "float32": torch.float32, "bfloat16": torch.bfloat16}.get( + torch_dtype, torch.float16) if torch is not None else None import logging, warnings logging.getLogger("transformers").setLevel(logging.ERROR) warnings.filterwarnings("ignore", message=".*You are using a model of type.*") @@ -78,7 +80,7 @@ def _load_hf(model_id_or_path: str, device: str): try: model = model_cls.from_pretrained( model_id_or_path, - torch_dtype=torch.float16 if torch is not None else None, + torch_dtype=load_dtype, trust_remote_code=True, low_cpu_mem_usage=True, local_files_only=Path(model_id_or_path).exists(), @@ -87,7 +89,7 @@ def _load_hf(model_id_or_path: str, device: str): try: model = AutoModel.from_pretrained( model_id_or_path, - torch_dtype=torch.float16 if torch is not None else None, + torch_dtype=load_dtype, trust_remote_code=True, low_cpu_mem_usage=True, local_files_only=Path(model_id_or_path).exists(), @@ -393,7 +395,7 @@ def convert(args: argparse.Namespace) -> None: shutil.rmtree(out_dir) out_dir.mkdir(parents=True, exist_ok=True) - cfg, processor, model = _load_hf(args.model, args.device) + cfg, processor, model = _load_hf(args.model, args.device, getattr(args, 'torch_dtype', 'float16')) runtime_source = ensure_parakeet_tdt_nemo_source(args.model, cache_dir=_hf_cache_dir()) or args.model family = detect_family(cfg, args.model_family) adapter = adapter_for_family(family) @@ -584,6 +586,7 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--embedding-bits", type=int, choices=[1, 2, 3, 4]) p.add_argument("--calibration-manifest") p.add_argument("--device", default="auto") + p.add_argument("--torch-dtype", default="float16", choices=["float16", "float32", "bfloat16"]) p.add_argument("--model-family", default="auto", choices=sorted(SUPPORTED_FAMILIES)) p.add_argument("--strict", action="store_true") p.add_argument("--force", action="store_true") From df4b2acc38776eac9fa76c6216897b394dbcdbd2 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Thu, 11 Jun 2026 13:34:41 -0700 Subject: [PATCH 2/5] Disable the dense-MLP TQ fusion by default The fused node costs ~15 ms/decode-step on Tensor G4 (profile-diffed on device) while gaining ~2% on M4. Its gate keys on ".cq" in the weight path, which the packed-panel suffix is about to satisfy; flipping the default keeps panel bundles on the plain matmul path. Re-enable by constructing FusionConfig(enable_dense_mlp_tq_fused=True). Signed-off-by: Noah Cylich --- python/cactus/transpile/optimize_graph.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index c700c7128..c53516d44 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -54,7 +54,8 @@ class FusionConfig: enable_attention_block: bool = True enable_conv_module: bool = True enable_add_clipped: bool = True - enable_dense_mlp_tq_fused: bool = True + # Opt-in: measured -15% decode on Tensor G4 when fused (+2% on M4). + enable_dense_mlp_tq_fused: bool = False def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig | None = None, From 842e644389daa5bae5cc91aa6d639c146cfeb1e4 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 00:04:14 -0700 Subject: [PATCH 3/5] Add packed-panel CQ4 weight format and NEON kernels Introduce a kernel-ready "packed panel" on-disk layout for CQ4 weights that both NEON (this change) and future SME paths consume zero-copy from the mmap, with no runtime cache and no second resident weight copy. Format - Data region holds panels [SB64][num_groups][gs/4][128B]; each 128B kg-block is 256 nibbles (low-first) = the int8 SDOT panel [4 vec][16 ch][4 K]. - Scales region adds folded fp32 panel norms (matmul rescale) alongside the legacy fp16 norms (embedding row decode), plus the transposed rotation for orthogonal lm_heads (virtual 128-wide groups written directly). - New header flag FLAG_PACKED_PANELS (1<<5) and a unique `.cq4p.weights` filename suffix; new bundles ship only that file. Enabled on all aarch64. Kernels (cactus-kernels) - Panel SDOT GEMV (M=1) and a fused panel NEON GEMM (M>1) over the file layout, plus the orthogonal lm_head GEMV with the file's transposed rotation. Mobile thread budget (ceil(SB64 / CACTUS_GEMV_SB_PER_THREAD), default 8). Batched orthogonal embedding row decode reads panels directly. - cactus_quant_build_panels is the byte-exact reference encoder (test/fixture). Loader (cactus-graph) - io.cpp parses the flag and points packed_panels / norm_panels / rotation_t straight into the mmap; the cache/registry/madvise machinery is removed. Panel files expose no legacy packed-indices region. Transpiler (python/cactus) - cq.write_cq_tensor emits panels for CQ4 matmul weights and orthogonal lm_heads/embeddings (chunked for large vocabs), byte-identical to the engine reference encoder. qdq learns to dequantize panels for export round-trips. Tests - Panel GEMV/GEMM/orth/embed vs fp32 oracle; panel file loaded through the graph matches the direct kernel; transpiler panels are byte-identical to the C reference encoder (ctypes); panel qdq round-trips the legacy dequant. Signed-off-by: Noah Cylich --- cactus-graph/cactus_graph.h | 15 +- cactus-graph/src/core.cpp | 12 + cactus-graph/src/io.cpp | 172 +++-- cactus-graph/src/ops_tensor.cpp | 57 ++ cactus-graph/tests/test_nn.cpp | 136 ++++ cactus-kernels/cactus_kernels.h | 52 ++ cactus-kernels/src/matmul.cpp | 618 +++++++++++++++++- cactus-kernels/tests/test_matmul.cpp | 462 +++++++++++-- python/cactus/convert/cli.py | 7 +- python/cactus/convert/export/qdq.py | 79 +++ python/cactus/convert/quantization/cq.py | 142 +++- python/cactus/convert/tests/test_cq.py | 51 +- python/cactus/convert/tests/test_cq_panels.py | 136 ++++ python/cactus/transpile/weight_compat.py | 173 ++++- python/tests/test_transpile_weight_compat.py | 8 +- 15 files changed, 1916 insertions(+), 204 deletions(-) create mode 100644 python/cactus/convert/tests/test_cq_panels.py diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 1cf4e87cc..8247023a3 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -263,6 +263,12 @@ struct BufferDesc { const int8_t* cq_right_signs = nullptr; const uint32_t* cq_permutation = nullptr; const __fp16* cq_rotation = nullptr; + // Panel-format file extras (see CactusQuantMatrix::packed_panels): kernel-ready packed + // panels (the file's data region), folded fp32 panel norms, and the transposed rotation + // (orthogonal lm_head only) — all pointing straight into the weight-file mapping. + const uint8_t* cq_packed_panels = nullptr; + const float* cq_norm_panels = nullptr; + const __fp16* cq_rotation_t = nullptr; uint32_t cq_flags = 0; CactusQuantMatrix to_cq_matrix() const { @@ -277,13 +283,18 @@ struct BufferDesc { .input_scale = cq_input_scale, .input_scale_recip = cq_input_scale_recip, .norms = cq_norms, - .packed_indices = static_cast(get_data()), + // Panel files store panels (not legacy packed indices) in the data region; keep + // packed_indices null there so legacy kernels can never misread panel bytes. + .packed_indices = cq_packed_panels ? nullptr : static_cast(get_data()), .left_signs = cq_left_signs, .right_signs = cq_right_signs, .permutation = cq_permutation, .rotation = cq_rotation, + .rotation_t = cq_rotation_t, .expanded = nullptr, .norm_f32 = nullptr, + .packed_panels = cq_packed_panels, + .norm_panels = cq_norm_panels, }; } @@ -781,6 +792,7 @@ namespace GraphFile { const void* scales_data() const; bool is_orthogonal_rotation() const { return is_orthogonal_rotation_; } bool is_interleaved_4row() const { return is_interleaved_4row_; } + bool is_packed_panels() const { return is_packed_panels_; } size_t original_N() const { return original_N_; } void* data(); const void* data() const; @@ -802,6 +814,7 @@ namespace GraphFile { uint32_t alignment_ = 32; bool is_orthogonal_rotation_ = false; bool is_interleaved_4row_ = false; + bool is_packed_panels_ = false; size_t original_N_ = 0; void parse_header(); diff --git a/cactus-graph/src/core.cpp b/cactus-graph/src/core.cpp index 1ad864bb9..eabb7e5c2 100644 --- a/cactus-graph/src/core.cpp +++ b/cactus-graph/src/core.cpp @@ -96,6 +96,9 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept cq_right_signs(other.cq_right_signs), cq_permutation(other.cq_permutation), cq_rotation(other.cq_rotation), + cq_packed_panels(other.cq_packed_panels), + cq_norm_panels(other.cq_norm_panels), + cq_rotation_t(other.cq_rotation_t), cq_flags(other.cq_flags) { other.total_size = 0; other.byte_size = 0; @@ -113,6 +116,9 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept other.cq_right_signs = nullptr; other.cq_permutation = nullptr; other.cq_rotation = nullptr; + other.cq_packed_panels = nullptr; + other.cq_norm_panels = nullptr; + other.cq_rotation_t = nullptr; other.cq_flags = 0; } @@ -142,6 +148,9 @@ BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { cq_right_signs = other.cq_right_signs; cq_permutation = other.cq_permutation; cq_rotation = other.cq_rotation; + cq_packed_panels = other.cq_packed_panels; + cq_norm_panels = other.cq_norm_panels; + cq_rotation_t = other.cq_rotation_t; cq_flags = other.cq_flags; other.total_size = 0; @@ -160,6 +169,9 @@ BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { other.cq_right_signs = nullptr; other.cq_permutation = nullptr; other.cq_rotation = nullptr; + other.cq_packed_panels = nullptr; + other.cq_norm_panels = nullptr; + other.cq_rotation_t = nullptr; other.cq_flags = 0; } return *this; diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index b7fff474d..edf59c342 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -25,6 +25,7 @@ namespace { constexpr uint32_t FLAG_ORTHOGONAL_ROTATION = 1 << 1; constexpr uint32_t FLAG_INTERLEAVED_4ROW = 1 << 2; constexpr uint32_t FLAG_EXTENDED_SHAPE = 1 << 4; + constexpr uint32_t FLAG_PACKED_PANELS = 1 << 5; constexpr size_t HEADER_SIZE = 84; inline size_t align_offset(size_t offset, size_t alignment) { @@ -41,6 +42,12 @@ namespace { } const std::string stem = filename.substr(0, filename.size() - std::strlen(suffix)); + // Packed-panel CQ4 weights carry their own suffix (new bundles ship ONLY this file, so + // old engines fail loudly on new bundles instead of silently misreading the layout). + const std::string cq4p = stem + ".cq4p.weights"; + if (std::filesystem::exists(cq4p)) { + return cq4p; + } const std::string cq4 = stem + ".cq4.weights"; if (std::filesystem::exists(cq4)) { return cq4; @@ -52,6 +59,67 @@ namespace { return filename; } + // Derive the CQ scale/panel pointers of a mapped CQ weight buffer from the file's scales + // region. Panel-format files share the legacy prefix (codebook, input scales, fp16 norms) + // and then carry the folded fp32 panel norms (4-byte aligned) followed by the rotation + + // transposed rotation (orthogonal) or the Hadamard parts; their data region holds the + // packed panels themselves (see CactusQuantMatrix::packed_panels). + void set_cq_buffer_pointers(BufferDesc& buffer, const GraphFile::MappedFile& mf) { + const char* scales_base = static_cast(mf.scales_data()); + const uint32_t bits = PrecisionTraits::cq_bits(buffer.precision); + const uint32_t cb_size = 1u << bits; + const uint32_t gs = static_cast(mf.group_size()); + const uint32_t K = gs * static_cast(mf.num_groups()); + const uint32_t N = static_cast(mf.shape().size() >= 2 ? mf.shape()[0] : 1); + + size_t off = 0; + buffer.cq_codebook = reinterpret_cast(scales_base + off); + off += cb_size * sizeof(__fp16); + buffer.cq_input_scale = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); + off += K * sizeof(__fp16); + buffer.cq_norms = reinterpret_cast(scales_base + off); + off += static_cast(N) * mf.num_groups() * sizeof(__fp16); + + if (mf.is_packed_panels()) { + const size_t SB64 = (static_cast(N) + 63) / 64; + off = align_offset(off, sizeof(float)); + buffer.cq_norm_panels = reinterpret_cast(scales_base + off); + off += SB64 * mf.num_groups() * 64 * sizeof(float); + if (mf.is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + off += static_cast(K) * K * sizeof(__fp16); + buffer.cq_rotation_t = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + buffer.cq_packed_panels = static_cast(mf.data()); + return; + } + + if (mf.is_orthogonal_rotation()) { + buffer.cq_rotation = reinterpret_cast(scales_base + off); + buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else { + buffer.cq_left_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_right_signs = reinterpret_cast(scales_base + off); + off += gs; + buffer.cq_permutation = reinterpret_cast(scales_base + off); + buffer.cq_flags = 0; + } + if (mf.is_interleaved_4row()) { + buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + } + } + inline void write_u32(std::ostream& out, uint32_t v) { out.write(reinterpret_cast(&v), sizeof(v)); } @@ -388,38 +456,7 @@ size_t CactusGraph::mmap_embeddings(const std::string& filename) { if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { buffer.group_size = mapped_file->group_size(); buffer.num_groups = mapped_file->num_groups(); - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } + set_cq_buffer_pointers(buffer, *mapped_file); } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { buffer.group_size = mapped_file->group_size(); buffer.num_groups = mapped_file->num_groups(); @@ -452,39 +489,7 @@ size_t CactusGraph::mmap_weights(const std::string& filename) { auto& buffer = nodes_[node_index_map_.at(node_id)]->output_buffer; buffer.group_size = mapped_file->group_size(); buffer.num_groups = mapped_file->num_groups(); - - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } + set_cq_buffer_pointers(buffer, *mapped_file); } size_t file_idx = mapped_files_.size(); @@ -531,43 +536,15 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) buffer.cq_right_signs = nullptr; buffer.cq_permutation = nullptr; buffer.cq_rotation = nullptr; + buffer.cq_packed_panels = nullptr; + buffer.cq_norm_panels = nullptr; + buffer.cq_rotation_t = nullptr; buffer.cq_flags = 0; if (PrecisionTraits::is_cq(precision) && mapped_file->group_size() > 0) { buffer.group_size = mapped_file->group_size(); buffer.num_groups = mapped_file->num_groups(); - - const char* scales_base = static_cast(mapped_file->scales_data()); - uint32_t bits = PrecisionTraits::cq_bits(precision); - uint32_t cb_size = 1u << bits; - uint32_t gs = static_cast(mapped_file->group_size()); - uint32_t K = gs * static_cast(mapped_file->num_groups()); - uint32_t N = static_cast(shape.size() >= 2 ? shape[0] : 1); - - size_t off = 0; - buffer.cq_codebook = reinterpret_cast(scales_base + off); - off += cb_size * sizeof(__fp16); - buffer.cq_input_scale = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_input_scale_recip = reinterpret_cast(scales_base + off); - off += K * sizeof(__fp16); - buffer.cq_norms = reinterpret_cast(scales_base + off); - off += static_cast(N) * mapped_file->num_groups() * sizeof(__fp16); - - if (mapped_file->is_orthogonal_rotation()) { - buffer.cq_rotation = reinterpret_cast(scales_base + off); - buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; - } else { - buffer.cq_left_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_right_signs = reinterpret_cast(scales_base + off); - off += gs; - buffer.cq_permutation = reinterpret_cast(scales_base + off); - buffer.cq_flags = 0; - } - if (mapped_file->is_interleaved_4row()) { - buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; - } + set_cq_buffer_pointers(buffer, *mapped_file); } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) { buffer.group_size = mapped_file->group_size(); buffer.num_groups = mapped_file->num_groups(); @@ -963,6 +940,7 @@ void MappedFile::parse_header() { offset += sizeof(uint32_t); is_orthogonal_rotation_ = (flags & FLAG_ORTHOGONAL_ROTATION) != 0; is_interleaved_4row_ = (flags & FLAG_INTERLEAVED_4ROW) != 0; + is_packed_panels_ = (flags & FLAG_PACKED_PANELS) != 0; alignment_ = *reinterpret_cast(ptr + offset); offset += sizeof(uint32_t); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp index 857c677b9..39a60952d 100644 --- a/cactus-graph/src/ops_tensor.cpp +++ b/cactus-graph/src/ops_tensor.cpp @@ -157,6 +157,63 @@ void compute_embedding_node(GraphNode& node, const std::vector 0) { bool orthogonal = (embeddings_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) != 0; + bool interleaved = (embeddings_buffer.cq_flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0; + bool panels = embeddings_buffer.cq_packed_panels != nullptr; + if (orthogonal && (interleaved || panels) && (hidden_dim % 8) == 0) { + // Batched path (production orthogonal embeddings — packed-panel or + // INTERLEAVED_4ROW): dedup indices, then ONE parallel dequant+un-rotation pass over + // the unique rows (the per-row fallback below is a serial scalar K^2 matvec per row + // — it was a top prefill term for per-layer embeddings). Legacy row-major + // orthogonal bundles (a transpiler bug) take the per-row fallback. + std::unordered_map slot; + slot.reserve(std::min(num_indices, 256)); + std::vector uniq; + uniq.reserve(std::min(num_indices, 256)); + std::vector idx_slot(num_indices); + for (size_t i = 0; i < num_indices; i++) { + size_t idx = static_cast(indices_ptr[i]); + if (idx >= vocab_size) { + throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size)); + } + auto it = slot.emplace(idx, static_cast(uniq.size())); + if (it.second) uniq.push_back(static_cast(idx)); + idx_slot[i] = it.first->second; + } + std::vector<__fp16> rows_out(uniq.size() * hidden_dim); + if (panels) { + cactus_quant_dequantize_orthogonal_embedding_rows_panels( + static_cast(hidden_dim), + static_cast(embeddings_buffer.group_size), + static_cast(embeddings_buffer.num_groups), + uniq.data(), + static_cast(uniq.size()), + embeddings_buffer.cq_packed_panels, + embeddings_buffer.cq_codebook, + embeddings_buffer.cq_norms, + embeddings_buffer.cq_input_scale_recip, + embeddings_buffer.cq_rotation, + rows_out.data()); + } else { + cactus_quant_dequantize_orthogonal_embedding_rows( + PrecisionTraits::cq_bits(emb_prec), + static_cast(hidden_dim), + uniq.data(), + static_cast(uniq.size()), + embeddings_buffer.data_as(), + embeddings_buffer.cq_codebook, + embeddings_buffer.cq_norms, + embeddings_buffer.cq_input_scale_recip, + embeddings_buffer.cq_rotation, + embeddings_buffer.cq_flags, + rows_out.data()); + } + for (size_t i = 0; i < num_indices; i++) { + std::memcpy(output + i * hidden_dim, + rows_out.data() + static_cast(idx_slot[i]) * hidden_dim, + hidden_dim * sizeof(__fp16)); + } + return; + } std::unordered_map row_cache; std::vector<__fp16> cached_rows; if (num_indices > 16) { diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index c2843f090..e8fd47a86 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -193,6 +193,141 @@ bool test_matmul_cq() { return true; } +// Writes a packed-panel CQ4 weight file (FLAG_PACKED_PANELS) exactly as the transpiler does: +// scales = codebook, input_scale, recip, fp16 norms, [4B pad], fp32 folded panel norms, then the +// hadamard parts; data = panels from the reference encoder. Exercises the io.cpp panel loader +// (set_cq_buffer_pointers offsets) + the panel GEMV/GEMM through the real graph. +static void write_test_panel_weights( + const std::filesystem::path& path, size_t K, size_t N, size_t gs, + const std::vector& panels, const std::vector& panel_norms, + const std::vector<__fp16>& codebook, const std::vector<__fp16>& input_scale, + const std::vector<__fp16>& input_scale_recip, const std::vector<__fp16>& norms, + const std::vector& left_signs, const std::vector& right_signs, + const std::vector& permutation) { + constexpr uint32_t CACTUS_MAGIC = 0x54434143; + constexpr uint32_t FLAG_PACKED_PANELS = 1u << 5; + constexpr size_t HEADER_SIZE = 84; + constexpr uint32_t alignment = 32; + const size_t ng = K / gs; + + std::string prefix; // codebook + input_scale + recip + fp16 norms + auto append = [&](const void* p, size_t bytes) { + prefix.append(reinterpret_cast(p), bytes); + }; + append(codebook.data(), codebook.size() * sizeof(__fp16)); + append(input_scale.data(), input_scale.size() * sizeof(__fp16)); + append(input_scale_recip.data(), input_scale_recip.size() * sizeof(__fp16)); + append(norms.data(), norms.size() * sizeof(__fp16)); + const size_t pad = align_offset_test(prefix.size(), 4) - prefix.size(); + + std::string scales = prefix; + scales.append(pad, '\0'); + scales.append(reinterpret_cast(panel_norms.data()), panel_norms.size() * sizeof(float)); + scales.append(reinterpret_cast(left_signs.data()), left_signs.size()); + scales.append(reinterpret_cast(right_signs.data()), right_signs.size()); + scales.append(reinterpret_cast(permutation.data()), permutation.size() * sizeof(uint32_t)); + + std::ofstream file(path, std::ios::binary); + if (!file) throw std::runtime_error("cannot open test panel weights"); + auto write_u32 = [&](uint32_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_u64 = [&](uint64_t v) { file.write(reinterpret_cast(&v), sizeof(v)); }; + auto write_padding = [&](size_t bytes) { for (size_t i = 0; i < bytes; ++i) file.put('\0'); }; + + write_u32(CACTUS_MAGIC); + write_u32(FLAG_PACKED_PANELS); + write_u32(alignment); + write_u32(2); + write_u64(N); + write_u64(K); + write_u64(0); + write_u64(0); + write_u32(static_cast(cq_precision_for_bits(4))); + write_u64(panels.size()); + write_u64(scales.size()); + write_u32(static_cast(gs)); + write_u32(static_cast(ng)); + write_u64(N); + size_t aligned_header = align_offset_test(HEADER_SIZE, alignment); + write_padding(aligned_header - HEADER_SIZE); + file.write(scales.data(), scales.size()); + size_t data_start = align_offset_test(aligned_header + scales.size(), alignment); + write_padding(data_start - (aligned_header + scales.size())); + file.write(reinterpret_cast(panels.data()), panels.size()); + if (!file) throw std::runtime_error("failed writing test panel weights"); +} + +bool test_matmul_cq_panels() { + // CQ4 panel file loaded through the graph must match the direct legacy-format kernel for + // GEMV (M=1) and GEMM (M>1). N=80 spans 2 padded super-blocks; K=256 spans 2 groups. + const size_t K = 256, N = 80, gs = 128, ng = K / gs; + const uint32_t bits = 4; + std::mt19937 gen(2024); + std::uniform_real_distribution dist(-1.f, 1.f); + + const uint32_t pgb = cactus_quant_packed_group_bytes(bits, gs); + std::vector packed(N * ng * pgb); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + std::vector<__fp16> codebook(16), input_scale(K), input_scale_recip(K), norms(N * ng); + std::vector left_signs(gs), right_signs(gs); + std::vector permutation(gs); + for (auto& v : codebook) v = static_cast<__fp16>(dist(gen)); + for (size_t i = 0; i < K; i++) { + float s = 0.5f + std::abs(dist(gen)); + input_scale[i] = static_cast<__fp16>(s); + input_scale_recip[i] = static_cast<__fp16>(1.f / s); + } + for (auto& v : norms) v = static_cast<__fp16>(dist(gen) * 0.1f); + for (auto& v : left_signs) v = (gen() & 1) ? 1 : -1; + for (auto& v : right_signs) v = (gen() & 1) ? 1 : -1; + for (uint32_t i = 0; i < gs; i++) permutation[i] = i; + + CactusQuantMatrix mat{ + .bits = bits, .K = static_cast(K), .N = static_cast(N), + .group_size = static_cast(gs), .num_groups = static_cast(ng), + .flags = 0, + .codebook = codebook.data(), .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), .norms = norms.data(), + .packed_indices = packed.data(), .left_signs = left_signs.data(), + .right_signs = right_signs.data(), .permutation = permutation.data(), + .rotation = nullptr, .rotation_t = nullptr, .expanded = nullptr, .norm_f32 = nullptr, + .packed_panels = nullptr, .norm_panels = nullptr, + }; + + const size_t SB64 = (N + 63) / 64; + std::vector panels(SB64 * ng * (gs / 4) * 128); + std::vector panel_norms(SB64 * ng * 64); + cactus_quant_build_panels(&mat, panels.data(), panel_norms.data()); + + auto path = std::filesystem::temp_directory_path() / "cactus_graph_cq4p_matmul.cq4p.weights"; + write_test_panel_weights(path, K, N, gs, panels, panel_norms, codebook, input_scale, + input_scale_recip, norms, left_signs, right_signs, permutation); + + bool ok = true; + for (size_t M : {size_t(1), size_t(20)}) { + std::vector<__fp16> A(M * K); + for (auto& v : A) v = static_cast<__fp16>(dist(gen)); + std::vector<__fp16> direct(M * N, static_cast<__fp16>(0)); + cactus_quant_matmul(&mat, A.data(), static_cast(M), direct.data()); + + CactusGraph g; + size_t ia = g.input({M, K}, Precision::FP16); + size_t iw = g.mmap_weights(path.string()); + size_t out = g.matmul(ia, iw, true); + g.set_input(ia, A.data(), Precision::FP16); + g.execute(); + __fp16* graph_out = static_cast<__fp16*>(g.get_output(out)); + for (size_t i = 0; i < M * N; i++) { + float actual = static_cast(graph_out[i]); + float expected = static_cast(direct[i]); + if (!std::isfinite(actual) || std::abs(actual - expected) > 2e-3f) { ok = false; break; } + } + g.hard_reset(); + if (!ok) break; + } + std::filesystem::remove(path); + return ok; +} + bool test_attention_int8_hybrid() { const size_t b = 1, s = 1, h = 2, kv = 2, d = 16; const size_t cache_len = 4; @@ -683,6 +818,7 @@ int main() { runner.run_test("Matrix Multiplication", test_matrix_multiplication()); runner.run_test("MatMul CQ", test_matmul_cq()); + runner.run_test("MatMul CQ Panels", test_matmul_cq_panels()); runner.run_test("Transpose", test_transpose()); runner.run_test("Reshape", test_reshape()); runner.run_test("RMS Norm", test_rms_norm()); diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index f371cffbe..68f0a3522 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/cactus-kernels/cactus_kernels.h @@ -219,8 +219,15 @@ struct CactusQuantMatrix { const int8_t* right_signs; const uint32_t* permutation; const __fp16* rotation; + // Transposed rotation (rot_t[i][k] = R[k][i], fp16 K x K), stored alongside `rotation` in + // panel-format files so the orthogonal GEMV phase-A needs no derived buffers. + const __fp16* rotation_t; const int8_t* expanded; const float* norm_f32; + // Panels mmap'd from the file: [SB64][ng][gs/4][128B], nibble j of a kg-block = byte j of + // the int8 panel [4 vec][16 ch][4 K]; norms [SB64][ng][64] f32 folded. Null -> legacy. + const uint8_t* packed_panels; + const float* norm_panels; }; uint32_t cactus_quant_packed_group_bytes(uint32_t bits, uint32_t group_size); @@ -281,6 +288,19 @@ void cactus_quant_orthogonal_matmul( uint32_t M, __fp16* C); +// Reference encoder for the packed-panel format (row-major or INTERLEAVED_4ROW source layout). +// panels_out: SB64*num_groups*(group_size/4)*128 bytes (SB64 = ceil(N/64)); norms_out: +// SB64*num_groups*64 floats. See CactusQuantMatrix::packed_panels. Production panel files are +// written by the transpiler; this is the byte-exact reference its output is tested against. +void cactus_quant_build_panels(const CactusQuantMatrix* W, uint8_t* panels_out, float* norms_out); + +// Orthogonal-rotation CQ4 GEMV over packed panels (lm_head). W2 = virtual-group view of the +// orthogonal matrix as stored in panel files (group_size=128, num_groups=K/128, norms replicated +// per group, packed_panels/norm_panels set); rot_t = K x K fp16 rotation TRANSPOSED +// (rot_t[i][k] = R[k][i]). +void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot_t, + const __fp16* input_scale_recip, const __fp16* A, __fp16* C); + void cactus_quant_4bit_gemv_interleaved( const CactusQuantMatrix* W, const uint8_t* packed_interleaved, @@ -336,6 +356,38 @@ void cactus_quant_dequantize_orthogonal_embedding_row( uint32_t flags, __fp16* out_row); +// Batched + parallelized version of the above for num_rows unique rows (same fp32-accumulate +// math; the per-row variant is a serial scalar K^2 matvec). INTERLEAVED_4ROW only (the +// production orthogonal-embedding format; no-op otherwise). out_rows: [num_rows][K]. +void cactus_quant_dequantize_orthogonal_embedding_rows( + uint32_t bits, + uint32_t K, + const uint32_t* rows, + uint32_t num_rows, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + uint32_t flags, + __fp16* out_rows); + +// Panel-format variant of the batched embedding dequant: decodes rows straight from the +// file-borne packed panels with the UNFOLDED fp16 norms (row-major [N][num_groups] in panel +// files), then the same vectorized un-rotation. out_rows: [num_rows][K]. +void cactus_quant_dequantize_orthogonal_embedding_rows_panels( + uint32_t K, + uint32_t group_size, + uint32_t num_groups, + const uint32_t* rows, + uint32_t num_rows, + const uint8_t* packed_panels, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + __fp16* out_rows); + void cactus_rms_norm_f16( const __fp16* input, const __fp16* weight, diff --git a/cactus-kernels/src/matmul.cpp b/cactus-kernels/src/matmul.cpp index 41daa4bde..7effe2949 100644 --- a/cactus-kernels/src/matmul.cpp +++ b/cactus-kernels/src/matmul.cpp @@ -459,7 +459,6 @@ static void cactus_quant_group_gemm( constexpr size_t TILE_N = 16; const size_t n_blocks = (W.N + TILE_N - 1) / TILE_N; - CactusThreading::parallel_gemm_tiles(M, n_blocks, [&, decode_group](size_t block_start, size_t block_end) { thread_local std::vector<__fp16> b_tile; @@ -497,13 +496,19 @@ static void cactus_quant_group_gemm( }); } +// Format-driven dispatch: the loader sets the panel pointers from panel files only. +static inline bool cactus_quant_has_panels(const CactusQuantMatrix* W) { + return W->packed_panels != nullptr && W->norm_panels != nullptr; +} + static bool cactus_quant_valid_common(const CactusQuantMatrix* W, const void* A, void* C) { if (W == nullptr || A == nullptr || C == nullptr) return false; if (W->K == 0 || W->N == 0 || W->group_size == 0 || W->num_groups == 0) return false; if (W->group_size > 256) return false; if ((W->group_size & (W->group_size - 1)) != 0) return false; if (W->K != W->group_size * W->num_groups) return false; - if (W->codebook == nullptr || W->norms == nullptr || W->packed_indices == nullptr) return false; + if (W->codebook == nullptr || W->norms == nullptr) return false; + if (W->packed_indices == nullptr && !cactus_quant_has_panels(W)) return false; return true; } @@ -691,34 +696,13 @@ __attribute__((always_inline)) static inline void tq_expand_i8_32( } template -__attribute__((always_inline)) static inline void cactus_quant_sdot_gemv_int8( - const CactusQuantMatrix* W, - const __fp16* code_basis, - __fp16* y) { - static_assert(Bits >= 1 && Bits <= 4); +static void cactus_quant_sdot_packed_blocks( + const CactusQuantMatrix* W, const int8_t* act_i8, const float* act_scales, + const int8x16_t cb_lut, const float cb_scale, + size_t block_start, size_t block_end, __fp16* y) { + constexpr size_t INT8_TILE_N = 16; const uint32_t gs = W->group_size; const uint32_t pgb = cactus_quant_packed_group_bytes(Bits, gs); - constexpr uint32_t cb_size = 1u << Bits; - - constexpr size_t INT8_TILE_N = 16; - const size_t int8_n_blocks = (W->N + INT8_TILE_N - 1) / INT8_TILE_N; - - thread_local std::vector act_i8_buf; - thread_local std::vector act_scales_buf; - if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); - if (act_scales_buf.size() < W->num_groups) act_scales_buf.resize(W->num_groups); - for (uint32_t g = 0; g < W->num_groups; ++g) { - act_scales_buf[g] = tq_quantize_group_i8( - code_basis + static_cast(g) * gs, - act_i8_buf.data() + static_cast(g) * gs, gs); - } - const int8_t* act_i8 = act_i8_buf.data(); - const float* act_scales = act_scales_buf.data(); - - int8_t cb_i8[16] = {}; - const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, cb_size); - const int8x16_t cb_lut = vld1q_s8(cb_i8); - auto expand_group4 = [&](size_t cache_block, uint32_t g, int8_t* dst, float* norm_scale) { const uint32_t n_vecs = gs / 16; const size_t n_start4 = cache_block * 4; @@ -763,7 +747,6 @@ __attribute__((always_inline)) static inline void cactus_quant_sdot_gemv_int8( } }; - cactus_quant_parallel_ranges(int8_n_blocks, 16, [&](size_t block_start, size_t block_end) { for (size_t block = block_start; block < block_end; ++block) { const size_t n_start = block * INT8_TILE_N; const size_t actual_n = std::min(INT8_TILE_N, static_cast(W->N) - n_start); @@ -877,6 +860,38 @@ __attribute__((always_inline)) static inline void cactus_quant_sdot_gemv_int8( y[n_start + ni] = static_cast<__fp16>(acc[ni]); } } +} + +template +__attribute__((always_inline)) static inline void cactus_quant_sdot_gemv_int8( + const CactusQuantMatrix* W, + const __fp16* code_basis, + __fp16* y) { + static_assert(Bits >= 1 && Bits <= 4); + const uint32_t gs = W->group_size; + constexpr uint32_t cb_size = 1u << Bits; + + constexpr size_t INT8_TILE_N = 16; + const size_t int8_n_blocks = (W->N + INT8_TILE_N - 1) / INT8_TILE_N; + + thread_local std::vector act_i8_buf; + thread_local std::vector act_scales_buf; + if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); + if (act_scales_buf.size() < W->num_groups) act_scales_buf.resize(W->num_groups); + for (uint32_t g = 0; g < W->num_groups; ++g) { + act_scales_buf[g] = tq_quantize_group_i8( + code_basis + static_cast(g) * gs, + act_i8_buf.data() + static_cast(g) * gs, gs); + } + const int8_t* act_i8 = act_i8_buf.data(); + const float* act_scales = act_scales_buf.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, cb_size); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + cactus_quant_parallel_ranges(int8_n_blocks, 16, [&](size_t b0, size_t b1) { + cactus_quant_sdot_packed_blocks(W, act_i8, act_scales, cb_lut, cb_scale, b0, b1, y); }); } @@ -1043,8 +1058,6 @@ void cactus_quant_2bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_ cactus_quant_Nbit_gemm_impl<2, CactusTQ2ScaledDecoder>(W, A, M, C, 8); } - - struct CactusTQ1ScaledDecoder { uint8x8_t cb_bytes; @@ -1162,8 +1175,6 @@ void cactus_quant_1bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_ cactus_quant_Nbit_gemm_impl<1, CactusTQ1ScaledDecoder>(W, A, M, C, 8); } - - static inline uint8x8_t cactus_tq3_unpack_8x3bit(const uint8_t* packed) { uint32_t b0 = packed[0], b1 = packed[1], b2 = packed[2]; @@ -1315,7 +1326,6 @@ static inline float tq_quantize_group_i8(const __fp16* src, int8_t* dst, uint32_ return scale; } - static inline uint32_t tq_extract_idx(const uint8_t* packed, uint32_t k, uint32_t bits) { if (bits == 4) return (packed[k / 2] >> ((k & 1) * 4)) & 0xF; if (bits == 2) return (packed[k / 4] >> ((k & 3) * 2)) & 0x3; @@ -1519,6 +1529,277 @@ static void tq_preexpand_weights( } } +// Byte-exact reference encoder for the panel format (tests/fixtures; production panels come +// from the transpiler). +void cactus_quant_build_panels(const CactusQuantMatrix* W, uint8_t* panels_out, float* norms_out) { + const uint32_t bits = W->bits; + const uint32_t gs = W->group_size; + const uint32_t ng = W->num_groups; + const uint32_t N = W->N; + const uint32_t pgb = cactus_quant_packed_group_bytes(bits, gs); + const size_t N_blocks = (static_cast(N) + 3) / 4; + const size_t SB64 = (static_cast(N) + 63) / 64; + const uint32_t nkg = gs / 4; + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 1u << bits); + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + std::vector w_il(N_blocks * ng * static_cast(gs) * 4); + std::vector n_f32(N_blocks * ng * 4); + if (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) + tq_preexpand_weights_interleaved(W, bits, ng, gs, pgb, cb_lut, cb_scale, + N_blocks, w_il.data(), n_f32.data()); + else + tq_preexpand_weights(W, bits, ng, gs, pgb, cb_lut, cb_scale, + N_blocks, w_il.data(), n_f32.data()); + + uint8_t inv[256]; + bool seen[256] = {}; + std::memset(inv, 0, sizeof inv); + for (uint32_t i = 0; i < (1u << bits); ++i) { + const uint8_t v = static_cast(cb_i8[i]) ^ 0x80u; // value+128 + if (!seen[v]) { seen[v] = true; inv[v] = static_cast(i); } + } + + std::memset(panels_out, 0, SB64 * ng * static_cast(nkg) * 128); + for (size_t sb = 0; sb < SB64; ++sb) + for (uint32_t g = 0; g < ng; ++g) { + uint8_t* dst = panels_out + (sb * ng + g) * static_cast(nkg) * 128; + for (uint32_t kg = 0; kg < nkg; ++kg) + for (uint32_t b = 0; b < 256; ++b) { + const uint32_t v = b / 64, ch = (b % 64) / 4, kc = b % 4; + const size_t n = sb * 64 + v * 16 + ch; + if (n >= N) continue; // padded channel: idx 0, norm 0 + const size_t nb = n / 4, ni = n % 4; + const uint32_t k = kg * 4 + kc; + const uint32_t v2 = k / 16, c = (k % 16) / 4, kc2 = k % 4; + const int8_t val = w_il[((nb * ng + g) * static_cast(gs)) * 4 + + static_cast(v2) * 64 + c * 16 + ni * 4 + kc2]; + const uint8_t idx = inv[static_cast(val) ^ 0x80u]; + uint8_t& d = dst[kg * 128 + b / 2]; + d = (b & 1u) ? static_cast(d | (idx << 4)) + : static_cast(d | idx); + } + float* nd = norms_out + (sb * ng + g) * 64; + for (uint32_t c = 0; c < 64; ++c) { + const size_t n = sb * 64 + c; + nd[c] = (n < N) ? n_f32[(n / 4 * ng + g) * 4 + (n % 4)] : 0.f; + } + } +} + +static void cactus_quant_interleaved4_gemv_blocks( + const CactusQuantMatrix* W, const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, const int8_t* act_i8, const float* act_scales, + const int8x16_t cb_lut, const float cb_scale, + size_t block_start, size_t block_end, __fp16* y); + +// NEON SDOT GEMV over the file-borne panels: per 16B, and/shr to nibbles, 2x vqtbl1q(codebook), +// vzip(lo,hi), 2x vdotq_laneq; 16 independent accumulators. Panel order = [4 vec][16 ch][4 K]. +static void cactus_quant_panel_gemv_blocks( + const CactusQuantMatrix* W, const int8_t* act_i8, const float* act_scales, + const int8x16_t cb_lut, size_t sb_start, size_t sb_end, __fp16* C) { + const uint32_t gs = W->group_size; + const uint32_t ng = W->num_groups; + const uint32_t N = W->N; + const uint32_t nkg = gs / 4; + const size_t g_stride = static_cast(nkg) * 128; + const size_t sb_stride = static_cast(ng) * g_stride; + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + + for (size_t sb = sb_start; sb < sb_end; ++sb) { + alignas(16) float facc[64] = {}; + const uint8_t* sbp = W->packed_panels + sb * sb_stride; + const float* nsb = W->norm_panels + sb * static_cast(ng) * 64; + for (uint32_t g = 0; g < ng; ++g) { + const uint8_t* gp = sbp + static_cast(g) * g_stride; + const int8_t* ag = act_i8 + static_cast(g) * gs; + int32x4_t acc[16]; + for (int v = 0; v < 16; ++v) acc[v] = vdupq_n_s32(0); + for (uint32_t kq = 0; kq < nkg; kq += 4) { // gs % 16 == 0 (cache gate: gs % 32 == 0) + const int8x16_t av = vld1q_s8(ag + static_cast(kq) * 4); + #define CACTUS_PANEL_KG(J) do { \ + const uint8_t* p_ = gp + (static_cast(kq) + (J)) * 128; \ + for (int c2 = 0; c2 < 8; ++c2) { \ + const uint8x16_t raw_ = vld1q_u8(p_ + c2 * 16); \ + const int8x16_t lov_ = vqtbl1q_s8(cb_lut, vandq_u8(raw_, lo_mask)); \ + const int8x16_t hiv_ = vqtbl1q_s8(cb_lut, vshrq_n_u8(raw_, 4)); \ + acc[c2 * 2 + 0] = vdotq_laneq_s32(acc[c2 * 2 + 0], \ + vzip1q_s8(lov_, hiv_), av, (J)); \ + acc[c2 * 2 + 1] = vdotq_laneq_s32(acc[c2 * 2 + 1], \ + vzip2q_s8(lov_, hiv_), av, (J)); \ + } } while (0) + CACTUS_PANEL_KG(0); CACTUS_PANEL_KG(1); CACTUS_PANEL_KG(2); CACTUS_PANEL_KG(3); + #undef CACTUS_PANEL_KG + } + const float32x4_t as = vdupq_n_f32(act_scales[g]); + const float* nf = nsb + static_cast(g) * 64; + for (int v = 0; v < 16; ++v) { + float32x4_t f = vld1q_f32(facc + v * 4); + f = vfmaq_f32(f, vmulq_f32(vcvtq_f32_s32(acc[v]), as), vld1q_f32(nf + v * 4)); + vst1q_f32(facc + v * 4, f); + } + } + const size_t n0 = sb * 64; + const uint32_t valid = static_cast(std::min(64, N - n0)); + for (uint32_t c = 0; c < valid; ++c) C[n0 + c] = static_cast<__fp16>(facc[c]); + } +} + +// Fused panel GEMV (M=1): one pool dispatch = phase A (group-parallel Hadamard + int8 quantize, +// spin barrier) + phase B (dynamic super-block stealing). Main runs as worker 0 and spin-joins. +// Thread budget = ceil(SB64 / CACTUS_GEMV_SB_PER_THREAD), default 8 (mobile: battery > saturation). +static void cactus_quant_panel_gemv(const CactusQuantMatrix* W, const __fp16* A, __fp16* C) { + const uint32_t gs = W->group_size; + const uint32_t num_groups = W->num_groups; + const uint32_t N = W->N; + const size_t SB64 = (static_cast(N) + 63) / 64; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (SB64 + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::min(pool.num_workers(), std::min(nt_budget, SB64)); + + static thread_local std::vector tl_h_act_i8; + static thread_local std::vector tl_h_act_scales; + if (tl_h_act_i8.size() < W->K) tl_h_act_i8.resize(W->K); + if (tl_h_act_scales.size() < num_groups) tl_h_act_scales.resize(num_groups); + int8_t* act_i8 = tl_h_act_i8.data(); + float* act_scales = tl_h_act_scales.data(); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 1u << W->bits); + (void)cb_scale; // folded into norm_panels by the transpiler + const int8x16_t cb_lut = vld1q_s8(cb_i8); + + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[256]; // gs <= 256 (same bound as the transform's tmp) + cactus_quant_transform_hadamard_group(*W, A + static_cast(g) * gs, g, basis); + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; + + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, 0, SB64, C); + return; + } + + std::atomic ga{0}; // phase A group grab + std::atomic a_done{0}; // phase A completion (release/acquire barrier) + std::atomic next{0}; // phase B super-block grab + std::atomic done{0}; // worker completion (spin-join, no cv wake) + auto worker = [&](size_t) { + for (uint32_t g; (g = ga.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { + phase_a_group(g); + a_done.fetch_add(1, std::memory_order_release); + } + while (a_done.load(std::memory_order_acquire) < num_groups) { /* spin */ } + for (;;) { + const uint32_t seen = next.load(std::memory_order_relaxed); + if (seen >= SB64) break; + const uint32_t want = (SB64 - seen > 4u * nt) ? 4u : 1u; + const uint32_t sb = next.fetch_add(want, std::memory_order_relaxed); + if (sb >= SB64) break; + const uint32_t cnt = std::min(want, static_cast(SB64) - sb); + cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, + sb, static_cast(sb) + cnt, C); + } + }; + pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { + worker(wid + 1); + done.fetch_add(1, std::memory_order_release); + }); + worker(0); + while (done.load(std::memory_order_acquire) < nt - 1) { /* spin */ } +} + +// Orthogonal-rotation CQ4 GEMV (the Gemma lm_head). Panel files store the lm_head with virtual +// 128-wide groups already applied (exact regrouping: norms constant across subgroups, panel +// kg-blocks K-chunk-ordered). Phase A rotates+quantizes each group via rot_t (fp32 accumulate); +// phase B is the panel super-block stealing loop. +void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot_t, + const __fp16* input_scale_recip, const __fp16* A, __fp16* C) { + const uint32_t K = W2->K; + const uint32_t N = W2->N; + const uint32_t gs = W2->group_size; + const uint32_t num_groups = W2->num_groups; + const size_t SB64 = (static_cast(N) + 63) / 64; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (SB64 + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::max(1, std::min(pool.num_workers(), std::min(nt_budget, SB64))); + + static thread_local std::vector<__fp16> tl_a_scaled; + static thread_local std::vector tl_o_act_i8; + static thread_local std::vector tl_o_act_scales; + if (tl_a_scaled.size() < K) tl_a_scaled.resize(K); + if (tl_o_act_i8.size() < K) tl_o_act_i8.resize(K); + if (tl_o_act_scales.size() < num_groups) tl_o_act_scales.resize(num_groups); + __fp16* a_scaled = tl_a_scaled.data(); + int8_t* act_i8 = tl_o_act_i8.data(); + float* act_scales = tl_o_act_scales.data(); + for (uint32_t k = 0; k < K; ++k) + a_scaled[k] = static_cast<__fp16>(static_cast(A[k]) * + (input_scale_recip ? static_cast(input_scale_recip[k]) : 1.f)); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W2->codebook, cb_i8, 1u << W2->bits); + (void)cb_scale; // folded into norm_panels by the transpiler + const int8x16_t cb_lut = vld1q_s8(cb_i8); + // rotate-and-quantize one virtual group: A_rot[i] = sum_k a_scaled[k]*R[k][i] = dot(a_scaled, rot_t[i]) + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[128]; + for (uint32_t j = 0; j < gs; ++j) { + const __fp16* rt = rot_t + (static_cast(g) * gs + j) * K; + float32x4_t s0 = vdupq_n_f32(0.f), s1 = vdupq_n_f32(0.f); + for (uint32_t k = 0; k < K; k += 8) { + float16x8_t av = vld1q_f16(a_scaled + k); + float16x8_t rv = vld1q_f16(rt + k); + s0 = vfmaq_f32(s0, vcvt_f32_f16(vget_low_f16(av)), vcvt_f32_f16(vget_low_f16(rv))); + s1 = vfmaq_f32(s1, vcvt_f32_f16(vget_high_f16(av)), vcvt_f32_f16(vget_high_f16(rv))); + } + basis[j] = static_cast<__fp16>(vaddvq_f32(vaddq_f32(s0, s1))); + } + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; + + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, 0, SB64, C); + return; + } + + std::atomic ga{0}; + std::atomic a_done{0}; + std::atomic next{0}; + std::atomic done{0}; + auto worker = [&](size_t) { + for (uint32_t g; (g = ga.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { + phase_a_group(g); + a_done.fetch_add(1, std::memory_order_release); + } + while (a_done.load(std::memory_order_acquire) < num_groups) { /* spin */ } + for (;;) { + const uint32_t seen = next.load(std::memory_order_relaxed); + if (seen >= SB64) break; + const uint32_t want = (SB64 - seen > 4u * nt) ? 4u : 1u; + const uint32_t sb = next.fetch_add(want, std::memory_order_relaxed); + if (sb >= SB64) break; + const uint32_t cnt = std::min(want, static_cast(SB64) - sb); + cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, + sb, static_cast(sb) + cnt, C); + } + }; + // Main participates as worker 0 + spin-join (no wait_all cv sleep) — see the GEMV driver. + // (nt >= 2 here: the nt <= 1 case took the serial path above.) + pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { + worker(wid + 1); + done.fetch_add(1, std::memory_order_release); + }); + worker(0); + while (done.load(std::memory_order_acquire) < nt - 1) { /* spin */ } +} + void cactus_quant_matmul( const CactusQuantMatrix* W, const __fp16* A, @@ -1528,6 +1809,11 @@ void cactus_quant_matmul( cactus_quant_orthogonal_matmul(W, A, M, C); return; } + // Panels take precedence: panel files carry no legacy packed-indices region. + if (W != nullptr && M == 1 && cactus_quant_has_panels(W) && (W->group_size % 32) == 0) { + cactus_quant_panel_gemv(W, A, C); + return; + } if (W != nullptr && M == 1 && (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0) { if (W->bits == 4) { @@ -1717,6 +2003,143 @@ void cactus_quant_matmul( return; } + // Fused panel GEMM: phase A packs activations [kg][16 rows][4 K] per (M-tile, group); + // phase B work-steals (M-tile x super-block) pairs. + if (cactus_quant_has_panels(W)) { + const uint32_t N = W->N; + const uint32_t K = W->K; + const size_t SB64 = (static_cast(N) + 63) / 64; + constexpr uint32_t TILE_M16 = 16; + const size_t M_tiles = (M + TILE_M16 - 1) / TILE_M16; + const size_t total_pairs = M_tiles * SB64; + const uint32_t nkg = gs / 4; + const size_t g_stride = static_cast(nkg) * 128; + + std::vector act_packed(M_tiles * num_groups * static_cast(gs) * 16); + std::vector act_scales(static_cast(M) * num_groups); + + int8_t cb_i8[16] = {}; + const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, codebook_size); + (void)cb_scale; // folded into norm_panels by the transpiler + const int8x16_t cb_lut = vld1q_s8(cb_i8); + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + + auto& pool = CactusThreading::get_thread_pool(); + // Same thread budget as the original M>1 NEON GEMM (one thread per 16-row M-tile). + const size_t nt = std::max(1, + std::min(pool.num_workers(), std::min(M_tiles, total_pairs))); + const uint32_t total_a = static_cast(M_tiles * num_groups); + + auto phase_a_item = [&](uint32_t item) { + const uint32_t mt = item / num_groups; + const uint32_t g = item % num_groups; + int8_t rowq[16][256]; // gs <= 256 + for (uint32_t r = 0; r < 16; ++r) { + const uint32_t m = mt * TILE_M16 + r; + if (m < M) { + __fp16 basis[256]; + cactus_quant_transform_hadamard_group( + *W, A + static_cast(m) * K + static_cast(g) * gs, g, basis); + act_scales[static_cast(m) * num_groups + g] = + tq_quantize_group_i8(basis, rowq[r], gs); + } else { + std::memset(rowq[r], 0, gs); + } + } + int8_t* dst = act_packed.data() + (static_cast(mt) * num_groups + g) * gs * 16; + for (uint32_t kg = 0; kg < gs / 4; ++kg) + for (uint32_t r = 0; r < 16; ++r) + std::memcpy(dst + static_cast(kg) * 64 + r * 4, rowq[r] + kg * 4, 4); + }; + auto run_pair = [&](size_t pair) { + const size_t mt = pair / SB64; + const size_t sb = pair % SB64; + const uint32_t m_rows = std::min(TILE_M16, M - static_cast(mt) * TILE_M16); + const uint32_t valid_n = std::min(64, N - static_cast(sb) * 64); + const int8_t* atile = act_packed.data() + (mt * num_groups) * static_cast(gs) * 16; + const uint8_t* sbp = W->packed_panels + sb * static_cast(num_groups) * g_stride; + const float* nsb = W->norm_panels + sb * static_cast(num_groups) * 64; + const uint32_t vmax = (valid_n + 15) / 16; // 16-channel quarters with output + const uint32_t strips = (m_rows + 3) / 4; // 4-row strips with live rows + for (uint32_t v = 0; v < vmax; ++v) { + for (uint32_t t = 0; t < strips; ++t) { + const uint32_t rmax = std::min(4, m_rows - t * 4); + alignas(16) float facc[64] = {}; // [4 rows][16 channels] + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* gp = sbp + static_cast(g) * g_stride + v * 32; + const int8_t* ag = atile + (static_cast(g) * nkg) * 64 + t * 16; + int32x4_t acc[16]; // [4 rows][4 channel-quartets] + for (int i2 = 0; i2 < 16; ++i2) acc[i2] = vdupq_n_s32(0); + for (uint32_t kg = 0; kg < nkg; ++kg) { + const int8x16_t act16 = vld1q_s8(ag + static_cast(kg) * 64); + const uint8_t* p = gp + static_cast(kg) * 128; + const uint8x16_t raw0 = vld1q_u8(p); + const uint8x16_t raw1 = vld1q_u8(p + 16); + const int8x16_t lo0 = vqtbl1q_s8(cb_lut, vandq_u8(raw0, lo_mask)); + const int8x16_t hi0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(raw0, 4)); + const int8x16_t lo1 = vqtbl1q_s8(cb_lut, vandq_u8(raw1, lo_mask)); + const int8x16_t hi1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(raw1, 4)); + const int8x16_t w0 = vzip1q_s8(lo0, hi0); // ch 16v+0..3 x 4 K + const int8x16_t w1 = vzip2q_s8(lo0, hi0); // ch 16v+4..7 + const int8x16_t w2 = vzip1q_s8(lo1, hi1); // ch 16v+8..11 + const int8x16_t w3 = vzip2q_s8(lo1, hi1); // ch 16v+12..15 + acc[0] = vdotq_laneq_s32(acc[0], w0, act16, 0); + acc[1] = vdotq_laneq_s32(acc[1], w1, act16, 0); + acc[2] = vdotq_laneq_s32(acc[2], w2, act16, 0); + acc[3] = vdotq_laneq_s32(acc[3], w3, act16, 0); + acc[4] = vdotq_laneq_s32(acc[4], w0, act16, 1); + acc[5] = vdotq_laneq_s32(acc[5], w1, act16, 1); + acc[6] = vdotq_laneq_s32(acc[6], w2, act16, 1); + acc[7] = vdotq_laneq_s32(acc[7], w3, act16, 1); + acc[8] = vdotq_laneq_s32(acc[8], w0, act16, 2); + acc[9] = vdotq_laneq_s32(acc[9], w1, act16, 2); + acc[10] = vdotq_laneq_s32(acc[10], w2, act16, 2); + acc[11] = vdotq_laneq_s32(acc[11], w3, act16, 2); + acc[12] = vdotq_laneq_s32(acc[12], w0, act16, 3); + acc[13] = vdotq_laneq_s32(acc[13], w1, act16, 3); + acc[14] = vdotq_laneq_s32(acc[14], w2, act16, 3); + acc[15] = vdotq_laneq_s32(acc[15], w3, act16, 3); + } + const float* nf = nsb + static_cast(g) * 64 + v * 16; + for (uint32_t r = 0; r < rmax; ++r) { + const uint32_t m = static_cast(mt) * TILE_M16 + t * 4 + r; + const float32x4_t as = + vdupq_n_f32(act_scales[static_cast(m) * num_groups + g]); + float* fr = facc + r * 16; + for (int q = 0; q < 4; ++q) { + float32x4_t f = vld1q_f32(fr + q * 4); + f = vfmaq_f32(f, + vmulq_f32(vcvtq_f32_s32(acc[r * 4 + q]), as), + vld1q_f32(nf + q * 4)); + vst1q_f32(fr + q * 4, f); + } + } + } + const uint32_t n0 = static_cast(sb) * 64 + v * 16; + const uint32_t cmax = std::min(16, valid_n - v * 16); + for (uint32_t r = 0; r < rmax; ++r) { + const uint32_t m = static_cast(mt) * TILE_M16 + t * 4 + r; + __fp16* crow = C + static_cast(m) * N + n0; + for (uint32_t c = 0; c < cmax; ++c) + crow[c] = static_cast<__fp16>(facc[r * 16 + c]); + } + } + } + }; + + if (nt <= 1) { + for (uint32_t i2 = 0; i2 < total_a; ++i2) phase_a_item(i2); + for (size_t p2 = 0; p2 < total_pairs; ++p2) run_pair(p2); + return; + } + + cactus_quant_two_phase_run(nt, total_a, static_cast(total_pairs), phase_a_item, + [&](size_t, uint32_t pair, uint32_t cnt) { + for (uint32_t i2 = 0; i2 < cnt; ++i2) run_pair(pair + i2); + }); + return; + } + const size_t act_size = static_cast(M) * W->K; std::vector<__fp16> heap_code_basis(act_size); __fp16* code_basis_ptr = heap_code_basis.data(); @@ -1956,6 +2379,124 @@ void cactus_quant_dequantize_orthogonal_embedding_row( } } +// Batched dequant + un-rotate in one pass (the per-row scalar K^2 matvec was a top Gemma +// prefill term). INTERLEAVED_4ROW only; row-major orthogonal bundles use the per-row fallback. +static void cactus_quant_unrotate_rows(const std::vector& dq, uint32_t K, uint32_t num_rows, + const __fp16* rotation, const __fp16* input_scale_recip, + __fp16* out_rows) { + constexpr uint32_t JBLK = 64; + const uint32_t jblocks = (K + JBLK - 1) / JBLK; + // ~8 j-blocks per thread caps single-row decode lookups at ~3 threads; prefill batches spread. + CactusThreading::parallel_for(static_cast(num_rows) * jblocks, + CactusThreading::ParallelConfig{16, 8}, + [&](size_t start, size_t end) { + for (size_t item = start; item < end; ++item) { + const uint32_t u = static_cast(item / jblocks); + const uint32_t j0 = static_cast(item % jblocks) * JBLK; + const uint32_t j1 = std::min(j0 + JBLK, K); + const float* d = dq.data() + static_cast(u) * K; + __fp16* out = out_rows + static_cast(u) * K; + for (uint32_t j = j0; j < j1; ++j) { + const __fp16* rrow = rotation + static_cast(j) * K; + float32x4_t acc0 = vdupq_n_f32(0.0f), acc1 = vdupq_n_f32(0.0f); + for (uint32_t i = 0; i < K; i += 8) { + float16x8_t r = vld1q_f16(rrow + i); + acc0 = vfmaq_f32(acc0, vcvt_f32_f16(vget_low_f16(r)), vld1q_f32(d + i)); + acc1 = vfmaq_f32(acc1, vcvt_f32_f16(vget_high_f16(r)), vld1q_f32(d + i + 4)); + } + const float scale = input_scale_recip ? static_cast(input_scale_recip[j]) : 1.0f; + out[j] = static_cast<__fp16>(vaddvq_f32(vaddq_f32(acc0, acc1)) * scale); + } + } + }); +} + +void cactus_quant_dequantize_orthogonal_embedding_rows( + uint32_t bits, + uint32_t K, + const uint32_t* rows, + uint32_t num_rows, + const uint8_t* packed_base, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + uint32_t flags, + __fp16* out_rows) { + if (!packed_base || !codebook || !norms || !rotation || !out_rows || !rows) return; + if (bits == 0 || bits > 4 || K == 0 || (K % 8) != 0 || num_rows == 0) return; + + const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(bits, K); + if ((flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) == 0 || bits != 4) return; + + std::vector dq(static_cast(num_rows) * K); + for (uint32_t u = 0; u < num_rows; ++u) { + const size_t row = rows[u]; + const float norm = static_cast(norms[(row / 4) * 4 + (row & 3u)]); + float* d = dq.data() + static_cast(u) * K; + const size_t panel_bytes = static_cast(4) * packed_group_bytes; + const uint8_t* panel = packed_base + (row / 4) * panel_bytes; + for (uint32_t i = 0; i < K; ++i) { + const uint32_t chunk = i / 8; + const uint32_t sub = i & 7u; + const uint8_t b = panel[chunk * 16 + (row & 3u) * 4 + (sub & 3u)]; + const uint32_t idx = (sub & 4u) ? (b >> 4) : (b & 0x0F); + d[i] = static_cast(codebook[idx]) * norm; + } + } + + cactus_quant_unrotate_rows(dq, K, num_rows, rotation, input_scale_recip, out_rows); +} + +// Batched embedding dequant straight from the panels (row n: super-block n/64, vector (n%64)/16, +// channel n%16; nibble j = v*64+ch*4+kc) with unfolded fp16 norms, then the vectorized un-rotation. +void cactus_quant_dequantize_orthogonal_embedding_rows_panels( + uint32_t K, + uint32_t group_size, + uint32_t num_groups, + const uint32_t* rows, + uint32_t num_rows, + const uint8_t* packed_panels, + const __fp16* codebook, + const __fp16* norms, + const __fp16* input_scale_recip, + const __fp16* rotation, + __fp16* out_rows) { + if (!packed_panels || !codebook || !norms || !rotation || !out_rows || !rows) return; + if (K == 0 || (K % 8) != 0 || num_rows == 0) return; + if (group_size == 0 || (group_size % 4) != 0 || group_size * num_groups != K) return; + + const uint32_t nkg = group_size / 4; + const size_t g_stride = static_cast(nkg) * 128; + const size_t sb_stride = static_cast(num_groups) * g_stride; + + std::vector dq(static_cast(num_rows) * K); + for (uint32_t u = 0; u < num_rows; ++u) { + const size_t row = rows[u]; + const size_t sb = row / 64; + const uint32_t v = static_cast((row % 64) / 16); + const uint32_t ch = static_cast(row % 16); + const uint32_t jbase = v * 64 + ch * 4; + float* d = dq.data() + static_cast(u) * K; + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* gp = packed_panels + sb * sb_stride + static_cast(g) * g_stride; + const float norm = static_cast(norms[row * num_groups + g]); + for (uint32_t kg = 0; kg < nkg; ++kg) { + const uint8_t* kgb = gp + static_cast(kg) * 128; + float* dk = d + static_cast(g) * group_size + kg * 4; + for (uint32_t kc = 0; kc < 4; ++kc) { + const uint32_t j = jbase + kc; + const uint8_t b = kgb[j >> 1]; + const uint32_t idx = (j & 1u) ? (b >> 4) : (b & 0x0F); + dk[kc] = static_cast(codebook[idx]) * norm; + } + } + } + } + + cactus_quant_unrotate_rows(dq, K, num_rows, rotation, input_scale_recip, out_rows); +} + static inline uint32_t tq_extract_interleaved_4row_4bit( const uint8_t* packed_base, uint32_t K, @@ -2059,6 +2600,14 @@ void cactus_quant_orthogonal_matmul( uint32_t M, __fp16* C) { if (!W || !A || !C || M == 0) return; + // Panel orthogonal lm_head: dispatch straight to the panel driver per row. + if (cactus_quant_has_panels(W) && W->rotation_t != nullptr) { + for (uint32_t m = 0; m < M; ++m) + cactus_quant_orth_panel_gemv(W, W->rotation_t, W->input_scale_recip, + A + static_cast(m) * W->K, + C + static_cast(m) * W->N); + return; + } if (!W->rotation || !W->packed_indices || !W->codebook || !W->norms) return; const uint32_t K = W->K; @@ -2203,6 +2752,7 @@ void cactus_quant_orthogonal_matmul( }); } +// Core of the IL CQ4 GEMV: pre-quantized activations against the interleaved packed panels. static void cactus_quant_interleaved4_gemv_blocks( const CactusQuantMatrix* W, const uint8_t* packed_interleaved, const __fp16* norms_interleaved, const int8_t* act_i8, const float* act_scales, diff --git a/cactus-kernels/tests/test_matmul.cpp b/cactus-kernels/tests/test_matmul.cpp index 74f917974..34f66be3e 100644 --- a/cactus-kernels/tests/test_matmul.cpp +++ b/cactus-kernels/tests/test_matmul.cpp @@ -54,6 +54,8 @@ struct SyntheticCQ { std::vector expanded_buf; std::vector norm_f32_buf; + std::vector packed_panels_buf; + std::vector norm_panels_buf; void preexpand() { int8_t cb_i8[16] = {}; @@ -132,31 +134,41 @@ struct SyntheticCQ { nd[ni] = (n_start+ni < N) ? static_cast(norms[(n_start+ni)*num_groups+g]) * cb_sc : 0.f; } } - } - CactusQuantMatrix matrix() const { - return CactusQuantMatrix{ - .bits = bits, .K = K, .N = N, - .group_size = group_size, .num_groups = num_groups, - .flags = 0, - .codebook = codebook.data(), - .input_scale = input_scale.data(), - .input_scale_recip = input_scale_recip.data(), - .norms = norms.data(), - .packed_indices = packed.data(), - .left_signs = left_signs.data(), - .right_signs = right_signs.data(), - .permutation = permutation.data(), - .rotation = nullptr, - .expanded = expanded_buf.empty() ? nullptr : expanded_buf.data(), - .norm_f32 = norm_f32_buf.empty() ? nullptr : norm_f32_buf.data(), - }; + // Independent panel packer: nibble j of a kg-block = byte j of the int8 panel + // [4 vec][16 ch][4 K]; norms [SB64][num_groups][64] with cb_scale folded. + const uint32_t nkg = group_size / 4; + const size_t SB64 = (size_t(N) + 63) / 64; + packed_panels_buf.assign(SB64 * num_groups * nkg * 128, 0); + norm_panels_buf.assign(SB64 * num_groups * 64, 0.f); + for (size_t sb = 0; sb < SB64; ++sb) + for (uint32_t g = 0; g < num_groups; ++g) { + uint8_t* wbase = packed_panels_buf.data() + (sb * num_groups + g) * nkg * 128; + float* nbase = norm_panels_buf.data() + (sb * num_groups + g) * 64; + for (uint32_t kg = 0; kg < nkg; ++kg) + for (uint32_t b = 0; b < 256; ++b) { // expanded byte index in the kg-panel + const uint32_t v = b / 64, ch = (b % 64) / 4, kc = b % 4; + const size_t n = sb * 64 + v * 16 + ch; + if (n >= N) continue; + const uint8_t* row = packed.data() + (n * num_groups + g) * pgb; + const uint8_t idx = unpack_index(row, bits, kg * 4 + kc); + uint8_t& dst = wbase[kg * 128 + b / 2]; + dst = (b & 1u) ? static_cast(dst | (idx << 4)) + : static_cast(dst | idx); + } + for (uint32_t c = 0; c < 64; ++c) { + const size_t n = sb * 64 + c; + if (n < N) nbase[c] = static_cast(norms[n * num_groups + g]) * cb_sc; + } + } } // INTERLEAVED_4ROW encoding: exact inverse of the shipped decoder. Per 8-byte half, // low nibbles = k 0-3 and high nibbles = k 4-7 of the half's K-range. std::vector packed_il; std::vector<__fp16> norms_il; + std::vector panels_il_buf; + std::vector npanels_il_buf; void make_interleaved() { if (bits != 4 || (N % 4) != 0) return; @@ -200,6 +212,37 @@ struct SyntheticCQ { .rotation = nullptr, .expanded = nullptr, .norm_f32 = nullptr, + .packed_panels = panels_il_buf.empty() ? nullptr : panels_il_buf.data(), + .norm_panels = npanels_il_buf.empty() ? nullptr : npanels_il_buf.data(), + }; + } + + void preexpand_il() { + CactusQuantMatrix W = matrix_interleaved(); + const size_t SB64 = ((size_t)N + 63) / 64; + panels_il_buf.resize(SB64 * num_groups * (size_t)(group_size / 4) * 128); + npanels_il_buf.resize(SB64 * num_groups * 64); + cactus_quant_build_panels(&W, panels_il_buf.data(), npanels_il_buf.data()); + } + + CactusQuantMatrix matrix() const { + return CactusQuantMatrix{ + .bits = bits, .K = K, .N = N, + .group_size = group_size, .num_groups = num_groups, + .flags = 0, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms.data(), + .packed_indices = packed.data(), + .left_signs = left_signs.data(), + .right_signs = right_signs.data(), + .permutation = permutation.data(), + .rotation = nullptr, + .expanded = expanded_buf.empty() ? nullptr : expanded_buf.data(), + .norm_f32 = norm_f32_buf.empty() ? nullptr : norm_f32_buf.data(), + .packed_panels = packed_panels_buf.empty() ? nullptr : packed_panels_buf.data(), + .norm_panels = norm_panels_buf.empty() ? nullptr : norm_panels_buf.data(), }; } }; @@ -308,7 +351,7 @@ bool test_cq_correctness(uint32_t bits) { double mse = compute_mse(ref.data(), y_f16.data(), N); - double threshold = 0.1; + double threshold = 0.1; if (mse > threshold) { std::cerr << " cq" << bits << " MSE=" << mse << " > " << threshold << "\n"; return false; @@ -316,7 +359,325 @@ bool test_cq_correctness(uint32_t bits) { return true; } +// ── Panel-format GEMV validation ───────────────────────────────────────────────────────────── +// With packed panels present (built here by the independent test packer), the real +// cactus_quant_matmul dispatch takes the panel GEMV kernel; validate it against the SAME FP32 +// reference oracle the legacy kernels are gated on. +static bool test_cq_panel(uint32_t bits, double& mse_out) { + const uint32_t K = 1024, N = 64, gs = 128; + SyntheticCQ cq(bits, K, N, gs, 123); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + + std::mt19937 gen(77); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + std::vector<__fp16> x_f16(K), y_f16(N, static_cast<__fp16>(0)); + for (size_t i = 0; i < K; i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); + + cactus_quant_matmul(&mat, x_f16.data(), 1, y_f16.data()); + + mse_out = compute_mse(ref.data(), y_f16.data(), N); + return mse_out <= 0.1; +} + +// GEMM (M>1) variant: validates the panel NEON GEMM (M and N tails included) vs a per-row FP32 +// reference. +static bool test_cq_panel_gemm(uint32_t bits, uint32_t M, uint32_t N, double& mse_out) { + const uint32_t K = 512, gs = 128; + SyntheticCQ cq(bits, K, N, gs, 321); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + + std::mt19937 gen(91); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector X(static_cast(M) * K); + for (auto& v : X) v = dist(gen); + + std::vector ref(static_cast(M) * N, 0.f); + for (uint32_t m = 0; m < M; m++) + cq_reference_gemv_f32(cq, X.data() + static_cast(m) * K, ref.data() + static_cast(m) * N); + + std::vector<__fp16> Xf(static_cast(M) * K), Yf(static_cast(M) * N, static_cast<__fp16>(0)); + for (size_t i = 0; i < static_cast(M) * K; i++) Xf[i] = static_cast<__fp16>(X[i]); + + cactus_quant_matmul(&mat, Xf.data(), M, Yf.data()); + + mse_out = compute_mse(ref.data(), Yf.data(), static_cast(M) * N); + return mse_out <= 0.1; +} + +// ── Orthogonal-rotation CQ4 (the Gemma lm_head format) ────────────────────────────────────── +// Oracle: y[n] = norm[n] * sum_i cb[idx(n,i)] * (a_scaled @ R)[i], fp32. Gates the incumbent and +// the panel virtual-group driver against the same reference, plus head-to-head agreement. +static bool test_orth_panel(double& mse_inc, double& mse_panel) { + // PRODUCTION lm_head format: orthogonal-rotation CQ4 in the INTERLEAVED_4ROW packed layout + // (gs == K, ng == 1) for legacy bundles; panel files store it with virtual 128-wide groups. + // Row-major orthogonal bundles were a transpiler bug and take the per-row fallback. + const uint32_t K = 1536, N = 1024, VGS = 128, VNG = K / VGS; + std::mt19937 gen(2024); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector<__fp16> codebook(16), isr(K), norms(N), rot((size_t)K * K), rot_t((size_t)K * K); + for (auto& v : codebook) v = (__fp16)dist(gen); + for (auto& v : isr) v = (__fp16)(0.7f + 0.3f * dist(gen)); + for (auto& v : norms) v = (__fp16)(dist(gen) * 0.05f); + for (size_t i = 0; i < rot.size(); i++) rot[i] = (__fp16)(dist(gen) * 0.04f); + for (uint32_t k = 0; k < K; k++) + for (uint32_t i = 0; i < K; i++) rot_t[(size_t)i * K + k] = rot[(size_t)k * K + i]; + std::vector packed((size_t)N * K / 2); + for (auto& v : packed) v = (uint8_t)(gen() & 0xFF); + std::vector<__fp16> x(K); + for (auto& v : x) v = (__fp16)dist(gen); + + // IL panels from row-major nibbles: byte [c*16 + r4*4 + b] = idx(row, 8c+b) | idx(row, 8c+4+b) << 4. + std::vector packed_il((size_t)N * K / 2); + for (uint32_t nb = 0; nb < N / 4; nb++) { + uint8_t* panel = packed_il.data() + (size_t)nb * 4 * (K / 2); + for (uint32_t c = 0; c < K / 8; c++) + for (uint32_t r4 = 0; r4 < 4; r4++) { + const uint8_t* row = packed.data() + (size_t)(nb * 4 + r4) * K / 2; + for (uint32_t b = 0; b < 4; b++) { + uint8_t lo = unpack_index(row, 4, c * 8 + b); + uint8_t hi = unpack_index(row, 4, c * 8 + 4 + b); + panel[c * 16 + r4 * 4 + b] = (uint8_t)(lo | (hi << 4)); + } + } + } + + std::vector ar(K, 0.f), ref(N); + for (uint32_t k = 0; k < K; k++) { + float a = (float)x[k] * (float)isr[k]; + for (uint32_t i = 0; i < K; i++) ar[i] += a * (float)rot[(size_t)k * K + i]; + } + for (uint32_t n = 0; n < N; n++) { + const uint8_t* row = packed.data() + (size_t)n * K / 2; + float acc = 0.f; + for (uint32_t i = 0; i < K; i++) + acc += (float)codebook[unpack_index(row, 4, i)] * ar[i]; + ref[n] = acc * (float)norms[n]; + } + + CactusQuantMatrix Wo{ .bits = 4, .K = K, .N = N, .group_size = K, .num_groups = 1, + .flags = CACTUS_QUANT_FLAG_ORTHOGONAL | CACTUS_QUANT_FLAG_INTERLEAVED_4ROW, + .codebook = codebook.data(), .input_scale = nullptr, + .input_scale_recip = isr.data(), .norms = norms.data(), .packed_indices = packed_il.data(), + .left_signs = nullptr, .right_signs = nullptr, .permutation = nullptr, + .rotation = rot.data(), .rotation_t = nullptr, .expanded = nullptr, .norm_f32 = nullptr, + .packed_panels = nullptr, .norm_panels = nullptr }; + std::vector<__fp16> y_inc(N); + cactus_quant_orthogonal_matmul(&Wo, x.data(), 1, y_inc.data()); + mse_inc = compute_mse(ref.data(), y_inc.data(), N); + + // Panel-file view: virtual 128-wide groups, replicated norms, rotation_t — what the + // loader hands the kernels; routed through the real dispatch. + std::vector<__fp16> norms_rep((size_t)N * VNG); + for (uint32_t nb = 0; nb < N / 4; nb++) + for (uint32_t g = 0; g < VNG; g++) + for (uint32_t ni = 0; ni < 4; ni++) + norms_rep[((size_t)nb * VNG + g) * 4 + ni] = norms[nb * 4 + ni]; + CactusQuantMatrix W2 = Wo; + W2.flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW; + W2.group_size = VGS; W2.num_groups = VNG; W2.norms = norms_rep.data(); + const size_t SB64 = (N + 63) / 64; + std::vector panels(SB64 * VNG * (VGS / 4) * 128); + std::vector npanels(SB64 * VNG * 64); + cactus_quant_build_panels(&W2, panels.data(), npanels.data()); + W2.flags = CACTUS_QUANT_FLAG_ORTHOGONAL; // loader view: panel files carry no IL flag + W2.packed_indices = nullptr; // panel files carry no legacy packed region + W2.packed_panels = panels.data(); W2.norm_panels = npanels.data(); + W2.rotation_t = rot_t.data(); + std::vector<__fp16> y_panel(N); + cactus_quant_orthogonal_matmul(&W2, x.data(), 1, y_panel.data()); + mse_panel = compute_mse(ref.data(), y_panel.data(), N); + return mse_inc <= 0.1 && mse_panel <= 0.1 && mse_panel <= mse_inc * 4 + 1e-4; +} + +// ── Interleaved-4row (legacy PRODUCTION format) tests ──────────────────────────────────────── +// CQ4 interleaved vs the FP32 oracle, through the real dispatch. use_panels=false exercises the +// legacy interleaved NEON kernel (old bundles); use_panels=true builds the panel layout from the +// SAME interleaved fixture through the reference encoder and exercises the panel GEMV +// (multi-super-block stealing included: N=192 = 3 super-blocks). +static bool test_cq4_interleaved(bool use_panels, double& mse_out, + uint32_t K = 1024, uint32_t N = 192, uint32_t gs = 128) { + SyntheticCQ cq(4, K, N, gs, 777); + if (use_panels) cq.preexpand_il(); + CactusQuantMatrix mat = cq.matrix_interleaved(); + + std::mt19937 gen(31); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + std::vector<__fp16> x_f16(K), y(N, (__fp16)0); + for (size_t i = 0; i < K; i++) x_f16[i] = (__fp16)x_f32[i]; + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + mse_out = compute_mse(ref.data(), y.data(), N); + return mse_out <= 0.1; +} + +// Batched orthogonal embedding dequant vs the per-row reference, both packed layouts. +static bool test_orth_embed_rows() { + const uint32_t K = 256, vocab = 64; + std::mt19937 gen(99); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector packed(static_cast(vocab) * K / 2); + for (auto& b : packed) b = static_cast(gen() & 0xFF); + std::vector<__fp16> codebook(16), rotation(static_cast(K) * K), scale_recip(K); + std::vector<__fp16> norms(vocab); + for (auto& v : codebook) v = static_cast<__fp16>(dist(gen)); + for (auto& v : rotation) v = static_cast<__fp16>(dist(gen) * 0.06f); + for (auto& v : scale_recip) v = static_cast<__fp16>(0.9f + 0.2f * std::fabs(dist(gen))); + for (auto& v : norms) v = static_cast<__fp16>(0.5f + std::fabs(dist(gen))); + const uint32_t rows[5] = {5, 9, 63, 0, 17}; + for (uint32_t flags : {(uint32_t)CACTUS_QUANT_FLAG_INTERLEAVED_4ROW}) { + std::vector<__fp16> ref(5 * K), got(5 * K); + for (int u = 0; u < 5; ++u) + cactus_quant_dequantize_orthogonal_embedding_row( + 4, K, rows[u], packed.data(), codebook.data(), norms.data(), + scale_recip.data(), rotation.data(), flags, ref.data() + (size_t)u * K); + cactus_quant_dequantize_orthogonal_embedding_rows( + 4, K, rows, 5, packed.data(), codebook.data(), norms.data(), + scale_recip.data(), rotation.data(), flags, got.data()); + double mx = 0; + for (size_t i = 0; i < ref.size(); ++i) + mx = std::max(mx, (double)std::fabs((float)ref[i] - (float)got[i])); + if (mx > 1e-2) { + std::cerr << " orth_embed_rows flags=" << flags << " max_err=" << mx << "\n"; + return false; + } + } + // Panel batched variant: same weights through the reference encoder, checked against + // the row-major per-row reference. + { + const uint32_t VGS = 128, VNG = K / VGS; + std::vector<__fp16> norms_rep((size_t)vocab * VNG); + for (uint32_t n = 0; n < vocab; ++n) + for (uint32_t g = 0; g < VNG; ++g) norms_rep[(size_t)n * VNG + g] = norms[n]; + CactusQuantMatrix Wv{ .bits = 4, .K = K, .N = vocab, .group_size = VGS, .num_groups = VNG, + .flags = 0, .codebook = codebook.data(), .input_scale = nullptr, + .input_scale_recip = scale_recip.data(), .norms = norms_rep.data(), + .packed_indices = packed.data(), .left_signs = nullptr, .right_signs = nullptr, + .permutation = nullptr, .rotation = rotation.data(), .rotation_t = nullptr, + .expanded = nullptr, .norm_f32 = nullptr, + .packed_panels = nullptr, .norm_panels = nullptr }; + const size_t SB64 = (vocab + 63) / 64; + std::vector panels(SB64 * VNG * (VGS / 4) * 128); + std::vector npanels(SB64 * VNG * 64); + cactus_quant_build_panels(&Wv, panels.data(), npanels.data()); + + std::vector<__fp16> ref(5 * K), got(5 * K); + for (int u = 0; u < 5; ++u) + cactus_quant_dequantize_orthogonal_embedding_row( + 4, K, rows[u], packed.data(), codebook.data(), norms.data(), + scale_recip.data(), rotation.data(), 0, ref.data() + (size_t)u * K); + cactus_quant_dequantize_orthogonal_embedding_rows_panels( + K, VGS, VNG, rows, 5, panels.data(), codebook.data(), norms_rep.data(), + scale_recip.data(), rotation.data(), got.data()); + double mx = 0; + for (size_t i = 0; i < ref.size(); ++i) + mx = std::max(mx, (double)std::fabs((float)ref[i] - (float)got[i])); + if (mx > 1e-2) { + std::cerr << " orth_embed_rows[panel] max_err=" << mx << "\n"; + return false; + } + } + return true; +} + +// The reference encoder is layout-invariant: identical panels from row-major or interleaved +// sources (the test packer keeps original indices, so it is not byte-compared). +static bool test_panel_layout_invariance() { + const uint32_t K = 512, N = 128, gs = 128; + SyntheticCQ cq(4, K, N, gs, 555); + CactusQuantMatrix w_rm = cq.matrix(); + cq.preexpand_il(); // reference encoder over the interleaved fixture + const size_t SB64 = (N + 63) / 64; + std::vector panels_rm(SB64 * cq.num_groups * (gs / 4) * 128); + std::vector npanels_rm(SB64 * cq.num_groups * 64); + cactus_quant_build_panels(&w_rm, panels_rm.data(), npanels_rm.data()); + bool ok = panels_rm == cq.panels_il_buf && npanels_rm == cq.npanels_il_buf; + if (!ok) std::cerr << " panel layout invariance FAILED\n"; + return ok; +} + +// Bench at Gemma shapes, alternating rounds + best-of: run-to-run noise exceeds the deltas. +static void print_panel_comparison() { + std::cout << "── CQ4 GEMV: panel NEON vs legacy file-layout NEON (gemma shapes) ──────────────────\n"; + struct Shape { const char* name; uint32_t K, N; int iters; }; + Shape shapes[] = { + {"cq4 1x1536x6144 (ffn)", 1536, 6144, 30}, + {"cq4 1x1536x12288 (gate_up)", 1536, 12288, 20}, + {"cq4 1x2048x1536 (o_proj)", 2048, 1536, 30}, + {"cq4 1x1536x2048 (q_proj)", 1536, 2048, 30}, + {"cq4 1x6144x1536 (down)", 6144, 1536, 30}, + {"cq4 1x1536x262144 (lm_head)", 1536, 262144, 6}, + {"cq4 1x1536x512 (kv_proj)", 1536, 512, 60}, + }; + for (auto& sh : shapes) { + SyntheticCQ cq(4, sh.K, sh.N, 128); + CactusQuantMatrix mat_file = cq.matrix_interleaved(); // no panels: legacy IL kernel + cq.preexpand_il(); + CactusQuantMatrix mat_panel = cq.matrix_interleaved(); // panels set: panel GEMV + std::vector<__fp16> x(sh.K), y(sh.N); + fill_random_fp16(x, -1.f, 1.f); + auto run_ms = [&](CactusQuantMatrix* m) { + cactus_quant_matmul(m, x.data(), 1, y.data()); // warmup + Timer t; + for (int i = 0; i < sh.iters; i++) cactus_quant_matmul(m, x.data(), 1, y.data()); + return t.elapsed_ms() / sh.iters; + }; + double ms_f = 1e30, ms_p = 1e30; + for (int round = 0; round < 5; round++) { + ms_f = std::min(ms_f, run_ms(&mat_file)); + ms_p = std::min(ms_p, run_ms(&mat_panel)); + } + const double fl = 2.0 * sh.K * sh.N; + double gf = fl / (ms_f * 1e6), gp = fl / (ms_p * 1e6); + std::cout << " " << std::left << std::setw(31) << sh.name + << " file-NEON " << std::fixed << std::setprecision(1) << gf + << " | panel " << gp << " GF" + << " | " << std::setprecision(2) << (gp / gf) + << "x (best of 5x" << sh.iters << ")\n"; + } + // GEMM (M>1): panel GEMM vs the legacy expand-from-packed NEON GEMM (the real-model M>1 + // path for old bundles). + for (uint32_t M : {4u, 16u, 32u, 64u, 128u, 256u}) { + const uint32_t K = 1024, N = 1024, gs = 128; + SyntheticCQ cq(4, K, N, gs); + CactusQuantMatrix mat_legacy = cq.matrix(); // no panels: per-call expand NEON GEMM + cq.preexpand(); + CactusQuantMatrix mat_panel = cq.matrix(); // panels set: panel NEON GEMM + std::vector<__fp16> A(static_cast(M) * K), Cc(static_cast(M) * N); + fill_random_fp16(A, -1.f, 1.f); + auto bench = [&](CactusQuantMatrix* m) -> double { + cactus_quant_matmul(m, A.data(), M, Cc.data()); // warmup + const int iters = 8; + Timer t; + for (int i = 0; i < iters; i++) cactus_quant_matmul(m, A.data(), M, Cc.data()); + return t.elapsed_ms() / iters; + }; + double ml = 1e30, mp = 1e30; + for (int round = 0; round < 3; round++) { + ml = std::min(ml, bench(&mat_legacy)); + mp = std::min(mp, bench(&mat_panel)); + } + double gl = (2.0 * M * K * N) / (ml * 1e6), gp = (2.0 * M * K * N) / (mp * 1e6); + std::string label = "cq4 M" + std::to_string(M) + " 1024x1024"; + std::cout << " " << std::left << std::setw(20) << label + << " legacy " << std::fixed << std::setprecision(1) << gl << " GFLOPS" + << " | panel " << gp << " GFLOPS" + << " | " << std::setprecision(2) << (gp / gl) << "x (best of 3x8)\n"; + } +} + bool run_benchmarks() { + print_panel_comparison(); auto bench = [](const char* label, size_t M, size_t K, size_t N, auto fn) { fn(); Timer t; @@ -554,26 +915,6 @@ void print_mse_report() { } } -// Legacy IL CQ4 vs the FP32 oracle through the real dispatch. -static bool test_cq4_interleaved(double& mse_out, - uint32_t K = 1024, uint32_t N = 192, uint32_t gs = 128) { - SyntheticCQ cq(4, K, N, gs, 777); - CactusQuantMatrix mat = cq.matrix_interleaved(); - - std::mt19937 gen(31); - std::uniform_real_distribution dist(-1.f, 1.f); - std::vector x_f32(K); - for (auto& v : x_f32) v = dist(gen); - std::vector ref(N, 0.f); - cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); - - std::vector<__fp16> x_f16(K), y(N, (__fp16)0); - for (size_t i = 0; i < K; i++) x_f16[i] = (__fp16)x_f32[i]; - cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); - mse_out = compute_mse(ref.data(), y.data(), N); - return mse_out <= 0.1; -} - int main() { TestRunner runner("Matrix Multiplication"); runner.run_test("matmul_f16", test_matmul_f16()); @@ -581,13 +922,48 @@ int main() { runner.run_test("matmul_cq2", test_cq_correctness(2)); runner.run_test("matmul_cq3", test_cq_correctness(3)); runner.run_test("matmul_cq4", test_cq_correctness(4)); + + // ── Panel-format registry: validate the panel kernels through the real dispatch ── + for (uint32_t bits : {1u, 2u, 3u, 4u}) { + double mse_p = 0.0; + bool ok = test_cq_panel(bits, mse_p); + runner.run_test(std::string("matmul_cq") + std::to_string(bits) + "[panel]", ok); + if (!ok) std::cerr << " cq" << bits << "[panel] MSE=" << mse_p << "\n"; + } + // GEMM (M>1) variant: M tail (20 -> 16+4) and N tail (72 -> 4 super-blocks + 8 valid) for CQ4. + struct GemmCase { uint32_t M, N; }; + for (GemmCase gc : {GemmCase{5, 64}, GemmCase{20, 64}, GemmCase{20, 72}, GemmCase{64, 256}}) { + double mse_g = 0.0; + bool ok = test_cq_panel_gemm(4, gc.M, gc.N, mse_g); + runner.run_test(std::string("matmul_cq4_M") + std::to_string(gc.M) + + "_N" + std::to_string(gc.N) + "[panel]", ok); + if (!ok) std::cerr << " cq4 M=" << gc.M << " N=" << gc.N << "[panel] MSE=" << mse_g << "\n"; + } { - double m1 = 0; - runner.run_test("matmul_cq4_il", test_cq4_interleaved(m1)); - // N=4164 -> 66 chunks incl. a 1-block tail: exercises the multi-thread fused driver. + double mi = 0, mp = 0; + bool orth_ok = test_orth_panel(mi, mp); + runner.run_test("matmul_cq4_orth[panel]", orth_ok); + if (!orth_ok) std::cerr << " orth mse_incumbent=" << mi << " mse_panel=" << mp << "\n"; + } + { + double m1 = 0, m2 = 0; + runner.run_test("matmul_cq4_il[file]", test_cq4_interleaved(false, m1)); + // N=4164: 1041 IL blocks -> 66 chunks (16-block + 1-block tail) -> multi-thread fused + // driver (phase-A stealing, spin barrier, 4-chunk grabs); N=192 stays on the serial path. double m_mt = 0; - runner.run_test("matmul_cq4_il_mt", test_cq4_interleaved(m_mt, 1024, 4164, 128)); + runner.run_test("matmul_cq4_il_mt[file]", test_cq4_interleaved(false, m_mt, 1024, 4164, 128)); + runner.run_test("matmul_cq4_il[panel]", test_cq4_interleaved(true, m2)); + runner.run_test("panel_layout_invariance", test_panel_layout_invariance()); + runner.run_test("orth_embed_rows_batched", test_orth_embed_rows()); } + // The panel GEMM is bits-agnostic (panel nibbles are codebook indices) — validate CQ1-3 too. + for (uint32_t b : {1u, 2u, 3u}) { + double mse_g = 0.0; + bool ok = test_cq_panel_gemm(b, 20, 64, mse_g); + runner.run_test(std::string("matmul_cq") + std::to_string(b) + "_gemm[panel]", ok); + if (!ok) std::cerr << " cq" << b << " gemm[panel] MSE=" << mse_g << "\n"; + } + runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); print_mse_report(); diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index aa7655a24..992f04308 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -495,6 +495,7 @@ def convert(args: argparse.Namespace) -> None: precision = emit_policy.precision gptq_used = False hessian_missing_reason = None + written_path = out_path module_name = adapter.module_target_name(name, model) or _module_name(name) try: if not match.recognized and args.strict: @@ -524,7 +525,9 @@ def convert(args: argparse.Namespace) -> None: cq = _scale_cq_norms(cq, adapter.scale_factor(out_path.name)) if getattr(emit_policy, "layout", "row_major") == "interleaved_4row": cq = replace(cq, interleaved_4row=True) - write_cq_tensor(out_path, cq) + # Panel-format CQ4 weights carry their own filename suffix; the manifest + # records the actual written file so bindings resolve to it. + written_path = write_cq_tensor(out_path, cq) gptq_used = cq.gptq_used status = "converted" if match.recognized else "unrecognized" elif emit_policy.action == "fallback" and out_path is not None: @@ -544,7 +547,7 @@ def convert(args: argparse.Namespace) -> None: "source_name": name, "hf_name": match.hf_name or name, "adapter_name": match.adapter_name or name, - "output_file": str(out_path.name) if out_path else None, + "output_file": str(written_path.name) if written_path else None, "shape": list(_tensor_shape(emit_tensor)), "dtype": str(getattr(emit_tensor, "dtype", "")), "component": emit_policy.component, diff --git a/python/cactus/convert/export/qdq.py b/python/cactus/convert/export/qdq.py index 5c73081d6..18b7488c9 100644 --- a/python/cactus/convert/export/qdq.py +++ b/python/cactus/convert/export/qdq.py @@ -25,6 +25,7 @@ FLAG_ORTHOGONAL_ROTATION = 1 << 1 FLAG_INTERLEAVED_4ROW = 1 << 2 FLAG_INTERLEAVED = 1 << 3 +FLAG_PACKED_PANELS = 1 << 5 PRECISION_INT8 = 0 PRECISION_FP16 = 1 @@ -220,6 +221,82 @@ def unpack_interleaved_4row_4bit(packed: np.ndarray, rows: int, cols: int) -> np return out.reshape(rows, cols) +def unpack_panels(packed: np.ndarray, n: int, group_size: int, num_groups: int) -> np.ndarray: + """Inverse of pack_indices_panels: panel bytes [SB64][ng][gs/4][128] -> indices [n][k].""" + nkg = group_size // 4 + sb64 = (n + 63) // 64 + p = packed.reshape(sb64, num_groups, nkg, 128) + nib = np.empty((sb64, num_groups, nkg, 256), dtype=np.uint8) + nib[..., 0::2] = p & 0x0F + nib[..., 1::2] = (p >> 4) & 0x0F + a = nib.reshape(sb64, num_groups, nkg, 64, 4).transpose(0, 3, 1, 2, 4) + return a.reshape(sb64 * 64, num_groups * group_size)[:n] + + +def parse_panel_metadata(blob: bytes, header: CactusHeader): + """Returns (codebook, input_scale, norms[n][ng] fp16, rotation_or_hadamard) for a panel file. + + The folded fp32 panel-norms region and the transposed rotation are skipped — QDQ reconstructs + from the unfolded fp16 norms + the forward rotation (orthogonal) or hadamard parts. + """ + n, k = header.shape + ng = header.num_groups + gs = header.group_size + pos = 0 + codebook = np.frombuffer(blob, dtype=np.float16, count=1 << header.bits, offset=pos).astype(np.float32) + pos += (1 << header.bits) * 2 + input_scale = np.frombuffer(blob, dtype=np.float16, count=k, offset=pos).astype(np.float32) + pos += k * 2 + pos += k * 2 # recip + norms = np.frombuffer(blob, dtype=np.float16, count=n * ng, offset=pos).astype(np.float32).reshape(n, ng) + pos += n * ng * 2 + pos = align_offset(pos, 4) + sb64 = (n + 63) // 64 + pos += sb64 * ng * 64 * 4 # folded fp32 panel norms + if header.flags & FLAG_ORTHOGONAL_ROTATION: + rotation = np.frombuffer(blob, dtype=np.float16, count=k * k, offset=pos).astype(np.float32).reshape(k, k) + return codebook, input_scale, norms, rotation + left = np.frombuffer(blob, dtype=np.int8, count=gs, offset=pos).astype(np.float32) + pos += gs + right = np.frombuffer(blob, dtype=np.int8, count=gs, offset=pos).astype(np.float32) + pos += gs + perm = np.frombuffer(blob, dtype=" torch.Tensor: + n, k = header.shape + ng = header.num_groups + gs = header.group_size + scales_blob, packed = read_cq_payload(path, header) + codebook, input_scale, norms, rot = parse_panel_metadata(scales_blob, header) + idx_all = unpack_panels(packed, n, gs, ng).astype(np.int64) + scale = input_scale[None, :] + is_orth = bool(header.flags & FLAG_ORTHOGONAL_ROTATION) + if is_orth: + rt = rot.T.copy() + else: + left, right, perm = rot + base_h = (hadamard(gs, dtype=float) / math.sqrt(gs)).astype(np.float32) + rotation = (left[:, None] * base_h * right[None, :])[:, perm] + rt = rotation.T.copy() + out = torch.empty(n, k, dtype=out_dtype) + batch = max(1, int(row_batch_size)) + for start in range(0, n, batch): + stop = min(start + batch, n) + cb = codebook[idx_all[start:stop]] # [b, k] + if is_orth: + recon = (cb @ rt) * norms[start:stop, :1] + arr = recon / scale + else: + arr = np.empty((stop - start, k), dtype=np.float32) + for g in range(ng): + lo, hi = g * gs, (g + 1) * gs + arr[:, lo:hi] = ((cb[:, lo:hi] @ rt) * norms[start:stop, g : g + 1]) / scale[:, lo:hi] + out[start:stop] = torch.from_numpy(arr.copy()).to(out_dtype) + return out + + def dequantize_fp_file(path: Path, header: CactusHeader, out_dtype: torch.dtype) -> torch.Tensor: offset = align_offset(HEADER_SIZE, header.alignment) dtype = np.float16 if header.precision == PRECISION_FP16 else np.float32 @@ -322,6 +399,8 @@ def parse_orthogonal_metadata(blob: bytes, header: CactusHeader): def dequantize_cq_file(path: Path, header: CactusHeader, out_dtype: torch.dtype, row_batch_size: int) -> torch.Tensor: if header.ndim != 2: raise ValueError(f"{path}: CQ tensors must be 2D, got shape={header.shape}") + if header.flags & FLAG_PACKED_PANELS: + return dequantize_cq_panels(path, header, out_dtype, row_batch_size) n, k = header.shape bits = header.bits scales_blob, packed = read_cq_payload(path, header) diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py index 2d3942389..c599fb827 100644 --- a/python/cactus/convert/quantization/cq.py +++ b/python/cactus/convert/quantization/cq.py @@ -24,6 +24,12 @@ PRECISION_CQ = {1: 3, 2: 4, 3: 5, 4: 6} FLAG_ORTHOGONAL_ROTATION = 1 << 1 FLAG_INTERLEAVED_4ROW = 1 << 2 +FLAG_PACKED_PANELS = 1 << 5 + +# Filename suffix for packed-panel CQ4 weight files. Panel bundles ship ONLY this file (no +# legacy-layout twin), so pre-panel engines fail loudly on new bundles instead of silently +# misreading the data region. +PANEL_WEIGHTS_SUFFIX = ".cq4p.weights" @dataclass(frozen=True) @@ -201,6 +207,62 @@ def pack_indices_interleaved_4row_3bit(indices: np.ndarray, group_size: int) -> return out.reshape(-1) +def quantize_codebook_i8(codebook_f16: np.ndarray) -> tuple[np.ndarray, np.float32]: + # Bit-exact mirror of the engine's tq_quantize_codebook_i8 (float32 max-abs/127, 1e-10 clamp, + # round-half-away-from-zero) so folded panel norms match the reference encoder byte-for-byte. + cb = codebook_f16.astype(np.float16).astype(np.float32) + max_abs = np.float32(np.max(np.abs(cb))) if cb.size else np.float32(0.0) + scale = max_abs / np.float32(127.0) + if scale < np.float32(1e-10): + scale = np.float32(1e-10) + inv = np.float32(1.0) / scale + # float64 holds every float32 exactly, so floor(|x| + 0.5) in float64 is exact + # round-half-away-from-zero of the float32 products (std::round semantics). + scaled = (cb * inv).astype(np.float64) + cb_i8 = (np.sign(scaled) * np.floor(np.abs(scaled) + 0.5)).astype(np.int8) + return cb_i8, scale + + +def canonicalize_codebook_indices(indices: np.ndarray, cb_i8: np.ndarray) -> np.ndarray: + # Map each index to the first index with an equal int8 value (both dequant identically) — + # mirrors the reference encoder's expand-then-reverse-map so panels are byte-identical. + lut = np.arange(cb_i8.size, dtype=np.uint8) + first: dict[int, int] = {} + for i, v in enumerate(cb_i8.tolist()): + if v in first: + lut[i] = first[v] + else: + first[v] = i + return lut[indices] + + +def pack_indices_panels(indices: np.ndarray, group_size: int) -> np.ndarray: + # 4-bit indices -> [SB64][ng][gs/4][128B] panels; each 128B kg-block = 256 nibbles (low first) + # = int8 panel [4 vec][16 ch][4 K]. Channels past N zero-padded (index 0). + n, k = indices.shape + if group_size % 4 != 0 or k % group_size != 0: + raise ValueError(f"packed-panel constraints not met: gs={group_size}, k={k}") + ng = k // group_size + nkg = group_size // 4 + sb64 = (n + 63) // 64 + pad = sb64 * 64 - n + idx = indices if pad == 0 else np.vstack([indices, np.zeros((pad, k), dtype=indices.dtype)]) + a = idx.reshape(sb64, 64, ng, nkg, 4).transpose(0, 2, 3, 1, 4).reshape(sb64, ng, nkg, 256) + a = a.astype(np.uint8) + lo = a[..., 0::2] & 0x0F + hi = a[..., 1::2] & 0x0F + return ((hi << 4) | lo).astype(np.uint8).reshape(-1) + + +def fold_panel_norms(norms_f16: np.ndarray, cb_scale: np.float32) -> np.ndarray: + # fp32 panel norms [SB64][ng][64] = float(norms fp16) * cb_scale; channels past N -> 0. + n, ng = norms_f16.shape + sb64 = (n + 63) // 64 + out = np.zeros((sb64 * 64, ng), dtype=np.float32) + out[:n] = norms_f16.astype(np.float16).astype(np.float32) * np.float32(cb_scale) + return np.ascontiguousarray(out.reshape(sb64, 64, ng).transpose(0, 2, 1)) + + def pack_indices_interleaved_4row(indices: np.ndarray, group_size: int, bits: int) -> np.ndarray: if bits == 4: return pack_indices_interleaved_4row_4bit(indices, group_size) @@ -318,15 +380,40 @@ def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np }) -def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: +def panel_weights_path(out_path: Path) -> Path: + """Panel files carry their own unique suffix (q_proj.weights -> q_proj.cq4p.weights).""" + name = out_path.name + if not name.endswith(".weights"): + return out_path + stem = name[: -len(".weights")] + if stem.endswith(".cq4"): + stem = stem[: -len(".cq4")] + return out_path.with_name(stem + PANEL_WEIGHTS_SUFFIX) + + +def write_cq_tensor(out_path: Path, cq: CQTensor, *, allow_panels: bool = True) -> Path: + # Returns the path actually written. CQ4 matmul weights and orthogonal lm_heads/embeddings go + # to the packed-panel format (.cq4p.weights, kernel-ready mmap); other tensors keep legacy. out_path.parent.mkdir(parents=True, exist_ok=True) n, k = cq.indices.shape group_size = int(cq.group_size) groups = k // group_size is_embedding = out_path.stem in _EMBEDDING_TENSOR_STEMS + is_orth = cq.rotation_family == "orthogonal" + + # Hadamard embeddings keep the legacy row-major layout (the engine's per-row hadamard + # decoder reads it); everything else CQ4 with a panel-compatible group size goes panel. + packed_panels = ( + allow_panels + and cq.bits == 4 + and ( + (is_orth and group_size == k and k % 128 == 0) + or (not is_orth and group_size % 32 == 0 and group_size <= 256 and not is_embedding) + ) + ) - if cq.interleaved_4row: + if cq.interleaved_4row and not packed_panels: if cq.bits not in (1, 2, 3, 4): raise ValueError(f"INTERLEAVED_4ROW only supports CQ1..CQ4, got CQ{cq.bits}") if n % 4 != 0: @@ -339,7 +426,7 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: elif group_size % 32 != 0 or group_size > 256: raise ValueError(f"INTERLEAVED_4ROW requires group_size % 32 == 0 and group_size <= 256, got {group_size}") - interleaved = ( + interleaved = not packed_panels and ( cq.interleaved_4row or ( cq.bits in (1, 2, 3, 4) @@ -353,7 +440,7 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: codebook_f32 = make_codebook(group_size, cq.bits).astype(np.float32) norms_f32 = cq.norms.astype(np.float32, copy=True) - if interleaved and cq.rotation_family != "orthogonal": + if (interleaved or packed_panels) and not is_orth: cb_max = float(np.max(np.abs(codebook_f32))) cb_factor = cb_max / 127.0 if cb_max > 1e-12 else 1.0 / 127.0 codebook_f32 = codebook_f32 / cb_factor @@ -363,6 +450,41 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: input_scale = cq.input_scale.astype(np.float16, copy=False).reshape(k) recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) + if packed_panels: + # Header stores the kernel-facing group view: orthogonal tensors are written with + # virtual 128-wide groups already applied (an exact regrouping — norms are constant + # across subgroups and panel kg-blocks are K-chunk-ordered). + header_gs = 128 if is_orth else group_size + header_ng = k // header_gs + norms_f16 = norms_f32.astype(np.float16).reshape(n, groups) + if is_orth: + norms_f16 = np.repeat(norms_f16.reshape(n, 1), header_ng, axis=1) + norms_f16 = np.ascontiguousarray(norms_f16) + + cb_i8, cb_scale = quantize_codebook_i8(codebook) + canonical = canonicalize_codebook_indices(cq.indices, cb_i8) + packed = pack_indices_panels(canonical, header_gs) + panel_norms = fold_panel_norms(norms_f16, cb_scale) + + parts = [codebook.tobytes(), input_scale.tobytes(), recip.tobytes(), norms_f16.tobytes()] + prefix_len = sum(len(p) for p in parts) + parts.append(b"\x00" * (align_offset(prefix_len, 4) - prefix_len)) + parts.append(panel_norms.tobytes()) + flags = FLAG_PACKED_PANELS + if is_orth: + flags |= FLAG_ORTHOGONAL_ROTATION + rotation = make_orthogonal_rotation(group_size, cq.seed).astype(np.float16) + parts.append(rotation.tobytes()) + parts.append(np.ascontiguousarray(rotation.T).tobytes()) + else: + left, right, perm = make_hadamard_components(group_size, cq.seed) + parts.extend([left.tobytes(), right.tobytes(), perm.astype(" None: else: left, right, perm = make_hadamard_components(group_size, cq.seed) parts.extend([left.tobytes(), right.tobytes(), perm.astype(" None: scales_bytes = len(scales_blob) data_bytes = packed.size scales_end = align_offset(HEADER_SIZE, CACTUS_ALIGNMENT) + scales_bytes @@ -394,7 +524,7 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: f.write(struct.pack(" ctypes.c_void_p: + return ctypes.c_void_p(arr.ctypes.data) + + +def _reference_panels(bits, k, n, group_size, flags, codebook_f16, norms_f16, packed_indices): + num_groups = k // group_size + nkg = group_size // 4 + sb64 = (n + 63) // 64 + panels = np.zeros(sb64 * num_groups * nkg * 128, dtype=np.uint8) + norms = np.zeros(sb64 * num_groups * 64, dtype=np.float32) + w = _CactusQuantMatrix( + bits=bits, K=k, N=n, group_size=group_size, num_groups=num_groups, flags=flags, + codebook=_ptr(codebook_f16), input_scale=None, input_scale_recip=None, + norms=_ptr(norms_f16), packed_indices=_ptr(packed_indices), + left_signs=None, right_signs=None, permutation=None, rotation=None, rotation_t=None, + expanded=None, norm_f32=None, packed_panels=None, norm_panels=None, + ) + _REFERENCE_ENCODER(ctypes.byref(w), _ptr(panels), _ptr(norms)) + return panels, norms + + +def _python_panels(k, n, group_size, codebook_f16, norms_f16, raw_indices): + cb_i8, cb_scale = quantize_codebook_i8(codebook_f16) + canonical = canonicalize_codebook_indices(raw_indices, cb_i8) + panels = pack_indices_panels(canonical, group_size) + num_groups = k // group_size + panel_norms = fold_panel_norms(norms_f16.reshape(n, num_groups), cb_scale) + return panels, panel_norms.reshape(-1) + + +@_skip +@pytest.mark.parametrize( + "n,k,group_size", + [ + (64, 256, 128), # one super-block, two groups + (192, 512, 128), # three super-blocks + (100, 256, 128), # ragged N (padded super-block) + (256, 256, 256), # single group spanning K + ], +) +def test_transpiler_panels_match_reference_encoder(n, k, group_size): + rng = np.random.default_rng(20240611 + n + k) + bits = 4 + codebook = make_codebook(group_size, bits).astype(np.float16) + raw = rng.integers(0, 1 << bits, size=(n, k), dtype=np.uint8) + norms = (rng.standard_normal((n, k // group_size)) * 0.1).astype(np.float16) + + packed_lsb = np.ascontiguousarray(pack_indices_lsb(raw, group_size, bits)) + ref_panels, ref_norms = _reference_panels( + bits, k, n, group_size, 0, codebook, np.ascontiguousarray(norms.reshape(-1)), packed_lsb + ) + py_panels, py_norms = _python_panels(k, n, group_size, codebook, norms, raw) + + assert np.array_equal(py_panels, ref_panels) + assert np.array_equal(py_norms, ref_norms) diff --git a/python/cactus/transpile/weight_compat.py b/python/cactus/transpile/weight_compat.py index 7b6c7421b..1ea5bc7d7 100644 --- a/python/cactus/transpile/weight_compat.py +++ b/python/cactus/transpile/weight_compat.py @@ -16,14 +16,21 @@ from cactus.transpile.weight_binding import WeightBinding from cactus.convert.quantization.cq import FLAG_ORTHOGONAL_ROTATION from cactus.convert.quantization.cq import FLAG_INTERLEAVED_4ROW +from cactus.convert.quantization.cq import FLAG_PACKED_PANELS from cactus.convert.quantization.cq import GROUP_SIZE as CQ_GROUP_SIZE +from cactus.convert.quantization.cq import PANEL_WEIGHTS_SUFFIX from cactus.convert.quantization.cq import PRECISION_CQ +from cactus.convert.quantization.cq import canonicalize_codebook_indices +from cactus.convert.quantization.cq import fold_panel_norms from cactus.convert.quantization.cq import make_codebook from cactus.convert.quantization.cq import make_hadamard_components from cactus.convert.quantization.cq import make_hadamard_matrix from cactus.convert.quantization.cq import make_orthogonal_rotation from cactus.convert.quantization.cq import pack_indices_interleaved_4row from cactus.convert.quantization.cq import pack_indices_lsb +from cactus.convert.quantization.cq import pack_indices_panels +from cactus.convert.quantization.cq import panel_weights_path +from cactus.convert.quantization.cq import quantize_codebook_i8 from cactus.convert.quantization.cq import quantize_hadamard from cactus.convert.quantization.cq import quantize_orthogonal from cactus.convert.quantization.cq import write_cq_tensor @@ -90,13 +97,23 @@ def ensure_embedding_binding_compatible(binding: WeightBinding) -> WeightBinding if config is None: return binding - compat_path = source_path.with_name(source_path.stem + f".cq{config['bits']}.weights") + bits = int(config["bits"]) + rotation = str(config["rotation"]) + # CQ4 orthogonal embeddings (token_embeddings == the tied lm_head) materialize in the + # packed-panel format under its own suffix; CQ2 per-layer hadamard embeddings keep the + # legacy row-major layout (the engine's hadamard row decoder reads it). + use_panels = bits == 4 and rotation == "orthogonal" + if use_panels: + compat_path = panel_weights_path(source_path) + else: + compat_path = source_path.with_name(source_path.stem + f".cq{bits}.weights") _materialize_cq_embedding_cache( opened, compat_path, - bits=int(config["bits"]), - rotation=str(config["rotation"]), + bits=bits, + rotation=rotation, layout=str(config.get("layout", "row_major")), + packed_panels=use_panels, ) _cleanup_legacy_fp16_cache(source_path) return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) @@ -121,12 +138,14 @@ def _ensure_legacy_int4_weight_binding_compatible( if not _is_legacy_packed_int4_tensor(opened): return binding - compat_path = source_path.with_name(source_path.stem + ".cq4.weights") + # CQ4 companions are written in the packed-panel format (its own suffix); the matmul kernels + # consume the mmap directly. + compat_path = panel_weights_path(source_path) if compat_path.exists(): if _can_materialize_compat_weight(source_tensor): _materialize_source_cq_weight( source_tensor, - compat_path, + source_path, bits=4, rotation="hadamard", scale_factor=_compat_weight_scale_factor(source_path), @@ -142,15 +161,15 @@ def _ensure_legacy_int4_weight_binding_compatible( "or generate them through `cq_convert` first." ) - _materialize_source_cq_weight( + written = _materialize_source_cq_weight( source_tensor, - compat_path, + source_path, bits=4, rotation="hadamard", scale_factor=_compat_weight_scale_factor(source_path), source_mtime_ns=source_path.stat().st_mtime_ns, ) - return WeightBinding(path=str(compat_path), kind=binding.kind, source_name=binding.source_name) + return WeightBinding(path=str(written), kind=binding.kind, source_name=binding.source_name) def _is_legacy_packed_int4_tensor(opened: _OpenedTensor) -> bool: @@ -197,9 +216,12 @@ def _materialize_source_cq_weight( rotation: str, scale_factor: float = 1.0, source_mtime_ns: int | None = None, -) -> None: - if source_mtime_ns is not None and out_path.exists() and out_path.stat().st_mtime_ns >= source_mtime_ns: - return +) -> Path: + # write_cq_tensor picks the panel suffix for CQ4 weights; the freshness check targets that + # actual companion file. + target = panel_weights_path(out_path) if bits == 4 else out_path + if source_mtime_ns is not None and target.exists() and target.stat().st_mtime_ns >= source_mtime_ns: + return target if rotation == "orthogonal": cq = quantize_orthogonal(tensor, bits=bits) @@ -207,7 +229,7 @@ def _materialize_source_cq_weight( cq = quantize_hadamard(tensor, bits=bits, use_gptq=False) if float(scale_factor) != 1.0: cq = replace(cq, norms=(cq.norms.astype(np.float32) * float(scale_factor)).astype(np.float16)) - write_cq_tensor(out_path, cq) + return write_cq_tensor(out_path, cq) def _open_cactus_tensor_file(path: str | Path) -> _OpenedTensor: @@ -290,6 +312,7 @@ def _materialize_cq_embedding_cache( bits: int, rotation: str, layout: str = "row_major", + packed_panels: bool = False, ) -> None: src_mtime_ns = opened.path.stat().st_mtime_ns if out_path.exists() and out_path.stat().st_mtime_ns >= src_mtime_ns: @@ -298,6 +321,10 @@ def _materialize_cq_embedding_cache( if len(opened.shape) != 2: raise RuntimeError(f"expected rank-2 embedding tensor, got shape={opened.shape}") + if packed_panels: + _materialize_cq_embedding_cache_panels(opened, out_path, bits=bits, rotation=rotation) + return + rows = int(opened.original_n or opened.shape[0]) cols = int(opened.shape[1]) interleaved = layout == "interleaved_4row" @@ -391,6 +418,107 @@ def _materialize_cq_embedding_cache( pass +def _materialize_cq_embedding_cache_panels( + opened: _OpenedTensor, + out_path: Path, + *, + bits: int, + rotation: str, +) -> None: + """Materialize a tied orthogonal CQ4 embedding/lm_head as a packed-panel file. + + The data region holds kernel-ready panels [SB64][ng][gs/4][128B] (ng = K/128 virtual + groups). The scales region carries the legacy fp16 norms (unfolded, replicated per virtual + group — the embedding row decoder uses them with the fp16 codebook) AND the fp32 folded + panel norms (the lm_head matmul rescale), then the rotation and its transpose. Rows stream + in 64-aligned chunks so the super-block panel order assembles by append. + """ + if rotation != "orthogonal" or bits != 4: + raise RuntimeError(f"packed-panel embedding cache requires orthogonal CQ4, got rotation={rotation} bits={bits}") + + rows = int(opened.original_n or opened.shape[0]) + cols = int(opened.shape[1]) + if cols % 128 != 0: + raise RuntimeError(f"packed-panel embedding cache requires cols % 128 == 0, got {cols}") + + input_scale = _estimate_embedding_input_scale(opened, rows=rows, cols=cols) + input_scale = np.asarray(input_scale, dtype=np.float16).reshape(cols) + recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) + + codebook = make_codebook(cols, bits).astype(np.float16) + cb_i8, cb_scale = quantize_codebook_i8(codebook) + rotation_mat = make_orthogonal_rotation(cols, seed=1234).astype(np.float16) + + header_gs = 128 + header_ng = cols // header_gs + nkg = header_gs // 4 + sb64 = (rows + 63) // 64 + + base = _recommended_chunk_rows(cols=cols, rotation="orthogonal") + chunk_rows = max(64, ((base + 63) // 64) * 64) # 64-aligned: only the last chunk is partial + + norms_tmp = out_path.with_suffix(out_path.suffix + ".norms.tmp") + panel_norms_tmp = out_path.with_suffix(out_path.suffix + ".pn.tmp") + data_tmp = out_path.with_suffix(out_path.suffix + ".data.tmp") + final_tmp = out_path.with_suffix(out_path.suffix + ".tmp") + + try: + with norms_tmp.open("wb") as norms_handle, \ + panel_norms_tmp.open("wb") as pnorms_handle, \ + data_tmp.open("wb") as data_handle: + for start in range(0, rows, chunk_rows): + stop = min(start + chunk_rows, rows) + chunk = _dequantize_int8_rows(opened, start, stop) + idx, norms = _quantize_orthogonal_chunk_indices(chunk, bits=bits, input_scale=input_scale) + canonical = canonicalize_codebook_indices(idx, cb_i8) + data_handle.write(np.ascontiguousarray(pack_indices_panels(canonical, header_gs)).tobytes()) + norms_rep = np.repeat(norms.astype(np.float16).reshape(-1, 1), header_ng, axis=1) + norms_handle.write(np.ascontiguousarray(norms_rep, dtype=np.float16).tobytes()) + pnorms_handle.write(np.ascontiguousarray(fold_panel_norms(norms_rep, cb_scale), dtype=np.float32).tobytes()) + + norms_bytes = norms_tmp.read_bytes() + panel_norms_bytes = panel_norms_tmp.read_bytes() + prefix = b"".join([codebook.tobytes(), input_scale.tobytes(), recip.tobytes(), norms_bytes]) + pad = b"\x00" * (align_offset(len(prefix), 4) - len(prefix)) + scales_blob = b"".join([ + prefix, + pad, + panel_norms_bytes, + rotation_mat.tobytes(), + np.ascontiguousarray(rotation_mat.T).tobytes(), + ]) + packed_bytes = sb64 * header_ng * nkg * 128 + flags = FLAG_PACKED_PANELS | FLAG_ORTHOGONAL_ROTATION + with final_tmp.open("wb") as handle: + _write_cq_header( + handle, + rows=rows, + cols=cols, + precision=int(PRECISION_CQ[int(bits)]), + packed_bytes=packed_bytes, + scales_bytes=len(scales_blob), + group_size=header_gs, + num_groups=header_ng, + flags=flags, + ) + scales_offset = align_offset(_HEADER_SIZE, CACTUS_ALIGNMENT) + handle.write(b"\x00" * (scales_offset - _HEADER_SIZE)) + handle.write(scales_blob) + data_offset = align_offset(scales_offset + len(scales_blob), CACTUS_ALIGNMENT) + current = handle.tell() + if data_offset > current: + handle.write(b"\x00" * (data_offset - current)) + with data_tmp.open("rb") as packed_handle: + shutil.copyfileobj(packed_handle, handle, length=8 * 1024 * 1024) + final_tmp.replace(out_path) + finally: + for temp_path in (norms_tmp, panel_norms_tmp, data_tmp): + try: + temp_path.unlink() + except FileNotFoundError: + pass + + def _recommended_chunk_rows(*, cols: int, rotation: str) -> int: if rotation == "orthogonal": if cols >= 8192: @@ -470,16 +598,29 @@ def _quantize_orthogonal_chunk( input_scale: np.ndarray, interleaved: bool = False, ) -> tuple[np.ndarray, np.ndarray]: + indices, norms = _quantize_orthogonal_chunk_indices(chunk, bits=bits, input_scale=input_scale) + cols = chunk.shape[1] + norms_f16 = norms.astype(np.float16).reshape(-1, 1) + if interleaved: + return pack_indices_interleaved_4row(indices, cols, bits), norms_f16 + return pack_indices_lsb(indices, cols, bits), norms_f16 + + +def _quantize_orthogonal_chunk_indices( + chunk: np.ndarray, + *, + bits: int, + input_scale: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Raw codebook indices [chunk, cols] (uint8) + per-row norms [chunk] (fp32).""" work = chunk.astype(np.float32, copy=False) * input_scale.astype(np.float32, copy=False)[None, :] - rows, cols = work.shape + cols = work.shape[1] rotation = make_orthogonal_rotation(cols, seed=1234).astype(np.float32) codebook = make_codebook(cols, bits).astype(np.float32) norms = np.linalg.norm(work, axis=1).clip(min=1e-8).astype(np.float32, copy=False) rotated = (work / norms[:, None]) @ rotation indices = np.abs(rotated[..., None] - codebook[None, None, :]).argmin(axis=-1).astype(np.uint8) - if interleaved: - return pack_indices_interleaved_4row(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) - return pack_indices_lsb(indices, cols, bits), norms.astype(np.float16).reshape(rows, 1) + return indices, norms def _write_cq_header( diff --git a/python/tests/test_transpile_weight_compat.py b/python/tests/test_transpile_weight_compat.py index 29fc9c8f1..444e4c848 100644 --- a/python/tests/test_transpile_weight_compat.py +++ b/python/tests/test_transpile_weight_compat.py @@ -33,7 +33,8 @@ def test_token_embedding_binding_upgrades_to_cq4(tmp_path: Path) -> None: ) compat = ensure_embedding_binding_compatible(binding) - assert compat.path.endswith(".cq4.weights") + # Orthogonal CQ4 token embeddings (the tied lm_head) materialize as packed panels. + assert compat.path.endswith(".cq4p.weights") assert compat.kind == "embedding" assert not legacy_fp16.exists() @@ -75,7 +76,8 @@ def test_gemma4_per_layer_projection_binding_upgrades_legacy_int4(tmp_path: Path ) compat = ensure_binding_compatible(binding, source_tensor=tensor) - assert compat.path.endswith(".cq4.weights") + # Legacy INT4 weights upgrade to packed-panel CQ4 companions. + assert compat.path.endswith(".cq4p.weights") opened = _open_cactus_tensor_file(compat.path) assert opened.precision == 6 assert opened.shape == (8, 128) @@ -88,7 +90,7 @@ def test_decoder_binding_prefers_existing_cq4_companion(tmp_path: Path) -> None: tensor = rng.standard_normal((8, 128), dtype=np.float32) save_tensor_with_header(tensor, source, precision="INT4", model_type="generic") - companion = tmp_path / "layer_0_attn_q.cq4.weights" + companion = tmp_path / "layer_0_attn_q.cq4p.weights" save_tensor_with_header(tensor, companion, precision="FP16", model_type="generic") binding = WeightBinding( From 15f9b98ac0d231eefdfb5ea9504e8af1686db060 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 03:34:00 -0700 Subject: [PATCH 4/5] Report converted tensor bytes from the written panel path Signed-off-by: Noah Cylich --- python/cactus/convert/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 992f04308..1883bea8f 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -559,7 +559,7 @@ def convert(args: argparse.Namespace) -> None: "fallback_reason": emit_policy.fallback_reason, "hessian_samples": int(hessian_samples.get(module_name, 0)), "gptq_used": bool(gptq_used), - "bytes": _tensor_bytes(out_path), + "bytes": _tensor_bytes(written_path), "scale_factor": float(adapter.scale_factor(out_path.name)) if out_path else 1.0, "recognized": bool(match.recognized), "adapter_family": family, From 22df501e778b5678ffb3c3b3a434fcb365306f6b Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 02:52:26 -0700 Subject: [PATCH 5/5] Add SME2 streaming kernels over the packed-panel weight format Stacks the SME2 layer on the panel file format: the LUTI4/SMOPA leaves read the mmap'd panels directly (the panel bytes ARE the streaming layout), so there is no cache builder and no second weight copy. Non-SME2 silicon and non-512-bit SVL fall back to the NEON panel path bit-for-bit. Kernels - matmul_sme2_gemv.cpp: LUTI4/ZT0 nibble-expanding GEMV + fused GEMM leaves emitting raw int32 partials (fp32 rescale stays in NEON, numerics identical to the NEON path). Compiled at -O1: clang mis-compiles the vg1x4 ZA dots at -O2+ into encodings that SIGILL on M4. - matmul_sme2.cpp: SMOPA/USMOPA leaves incl. prefill-attention QK/AV passes and the streaming-SVL probe; compiled with armv8.2-a+...+sme2 (armv9-a implies non-streaming SVE2 Apple Silicon lacks). - attention_hybrid.cpp: SME prefill attention over the cached int8 KV segment (flat-scale QK requantize at pre-pack, u8 USMOPA AV with folded V scales). Dispatch - SME workers replace NEON workers inside the existing thread budget (flat k = min(2, nt-1), the measured power/speed frontier; CACTUS_SME_GEMV_WORKERS overrides). The fused SMOPA GEMM takes all M>1 panel matmuls. - CACTUS_QUANT_BACKEND / cactus_quant_set_backend: 0 auto, 1 force NEON, 2 force the streaming leaves (test coverage). - cpu_has_sme2() gates on streaming svcntb() == 64: every layout is pinned to 512-bit SVL; other SME2 implementations fall back to NEON for correctness. Tests - [sme2] variants of every panel test (GEMV/GEMM incl. M/N tails, orthogonal lm_head, interleaved fixtures, CQ1-4) through the real dispatch, plus a NEON<->SME agreement gate on identical fixtures and the SME prefill attention differential. Measured (M4 Pro, Gemma 4 E2B IT panel bundle, ~1k-token prompt, 32 tokens, 6 alternating cycles vs the NEON panel engine): prefill +17.6% (300->353 tok/s), TTFT -14.9% (3278->2788 ms), decode +2.9%, decode energy/token -17.9% (358->294 mJ). Kernel-level: fused SMOPA GEMM 1.5-3.1x the NEON panel GEMM; GEMV with SME workers up to 2.1x on K-heavy shapes. Signed-off-by: Noah Cylich --- cactus-graph/cactus_graph.h | 8 +- cactus-graph/src/io.cpp | 10 +- cactus-graph/src/ops_tensor.cpp | 7 +- cactus-graph/tests/test_nn.cpp | 6 +- cactus-kernels/CMakeLists.txt | 22 + cactus-kernels/cactus_kernels.h | 32 +- cactus-kernels/src/attention_hybrid.cpp | 493 +++++++++++++++++++++++ cactus-kernels/src/matmul.cpp | 228 ++++++++--- cactus-kernels/src/matmul_simd.h | 43 ++ cactus-kernels/src/matmul_sme2.cpp | 133 ++++++ cactus-kernels/src/matmul_sme2_gemv.cpp | 139 +++++++ cactus-kernels/src/threading.h | 30 ++ cactus-kernels/tests/test_attention.cpp | 177 ++++++++ cactus-kernels/tests/test_matmul.cpp | 200 +++++---- python/cactus/convert/cli.py | 3 +- python/cactus/convert/quantization/cq.py | 12 +- python/cactus/transpile/weight_compat.py | 15 +- 17 files changed, 1363 insertions(+), 195 deletions(-) create mode 100644 cactus-kernels/src/matmul_simd.h create mode 100644 cactus-kernels/src/matmul_sme2.cpp create mode 100644 cactus-kernels/src/matmul_sme2_gemv.cpp diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 8247023a3..8e5549395 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -263,9 +263,8 @@ struct BufferDesc { const int8_t* cq_right_signs = nullptr; const uint32_t* cq_permutation = nullptr; const __fp16* cq_rotation = nullptr; - // Panel-format file extras (see CactusQuantMatrix::packed_panels): kernel-ready packed - // panels (the file's data region), folded fp32 panel norms, and the transposed rotation - // (orthogonal lm_head only) — all pointing straight into the weight-file mapping. + // Panel-format file extras: packed panels, folded fp32 norms, transposed rotation — + // all pointing straight into the weight-file mapping. const uint8_t* cq_packed_panels = nullptr; const float* cq_norm_panels = nullptr; const __fp16* cq_rotation_t = nullptr; @@ -283,8 +282,7 @@ struct BufferDesc { .input_scale = cq_input_scale, .input_scale_recip = cq_input_scale_recip, .norms = cq_norms, - // Panel files store panels (not legacy packed indices) in the data region; keep - // packed_indices null there so legacy kernels can never misread panel bytes. + // Panel files keep packed_indices null so legacy kernels can never misread them. .packed_indices = cq_packed_panels ? nullptr : static_cast(get_data()), .left_signs = cq_left_signs, .right_signs = cq_right_signs, diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index edf59c342..db967ee90 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -42,8 +42,7 @@ namespace { } const std::string stem = filename.substr(0, filename.size() - std::strlen(suffix)); - // Packed-panel CQ4 weights carry their own suffix (new bundles ship ONLY this file, so - // old engines fail loudly on new bundles instead of silently misreading the layout). + // Panel bundles ship only the suffixed file, so old engines fail loudly. const std::string cq4p = stem + ".cq4p.weights"; if (std::filesystem::exists(cq4p)) { return cq4p; @@ -59,11 +58,8 @@ namespace { return filename; } - // Derive the CQ scale/panel pointers of a mapped CQ weight buffer from the file's scales - // region. Panel-format files share the legacy prefix (codebook, input scales, fp16 norms) - // and then carry the folded fp32 panel norms (4-byte aligned) followed by the rotation + - // transposed rotation (orthogonal) or the Hadamard parts; their data region holds the - // packed panels themselves (see CactusQuantMatrix::packed_panels). + // Panel files share the legacy scales prefix, then carry folded fp32 panel norms + // (4-byte aligned), then rotation + transposed rotation or the Hadamard parts. void set_cq_buffer_pointers(BufferDesc& buffer, const GraphFile::MappedFile& mf) { const char* scales_base = static_cast(mf.scales_data()); const uint32_t bits = PrecisionTraits::cq_bits(buffer.precision); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp index 39a60952d..71a29cd68 100644 --- a/cactus-graph/src/ops_tensor.cpp +++ b/cactus-graph/src/ops_tensor.cpp @@ -160,11 +160,8 @@ void compute_embedding_node(GraphNode& node, const std::vector slot; slot.reserve(std::min(num_indices, 256)); std::vector uniq; diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index e8fd47a86..9128bae32 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -193,10 +193,8 @@ bool test_matmul_cq() { return true; } -// Writes a packed-panel CQ4 weight file (FLAG_PACKED_PANELS) exactly as the transpiler does: -// scales = codebook, input_scale, recip, fp16 norms, [4B pad], fp32 folded panel norms, then the -// hadamard parts; data = panels from the reference encoder. Exercises the io.cpp panel loader -// (set_cq_buffer_pointers offsets) + the panel GEMV/GEMM through the real graph. +// Writes a panel CQ4 weight file exactly as the transpiler does; exercises the io.cpp panel +// loader offsets and the panel GEMV/GEMM through the real graph. static void write_test_panel_weights( const std::filesystem::path& path, size_t K, size_t N, size_t gs, const std::vector& panels, const std::vector& panel_norms, diff --git a/cactus-kernels/CMakeLists.txt b/cactus-kernels/CMakeLists.txt index bd0f5ff03..902872d46 100644 --- a/cactus-kernels/CMakeLists.txt +++ b/cactus-kernels/CMakeLists.txt @@ -15,6 +15,8 @@ add_library(cactus_kernels STATIC src/image.cpp src/lstm.cpp src/matmul.cpp + src/matmul_sme2.cpp + src/matmul_sme2_gemv.cpp src/nn.cpp src/norms_rope.cpp src/quants.cpp @@ -46,6 +48,26 @@ if(APPLE) target_link_libraries(cactus_kernels PUBLIC "-framework Accelerate") endif() +# SME2 TUs compile at armv8.2-a+...+sme2, NOT armv9-a+sme2 — armv9-a implies non-streaming +# SVE2 that Apple Silicon lacks (runtime SIGILL); the base library stays armv8.2-a. +include(CheckCXXSourceCompiles) +set(_CACTUS_SME2_FLAG "-march=armv8.2-a+fp16+simd+dotprod+i8mm+sme2") +set(CMAKE_REQUIRED_FLAGS "${_CACTUS_SME2_FLAG}") +check_cxx_source_compiles(" +#include +__arm_locally_streaming __arm_new(\"za\") static void f(){ svzero_za(); } +int main(){ return 0; } +" CACTUS_HAS_SME2) +unset(CMAKE_REQUIRED_FLAGS) +if(CACTUS_HAS_SME2) + set_source_files_properties(src/matmul_sme2.cpp PROPERTIES COMPILE_OPTIONS "${_CACTUS_SME2_FLAG}") + # -O1 only: clang 17 miscompiles the vg1x4 ZA dots at -O2/-O3 into encodings that SIGILL. + set_source_files_properties(src/matmul_sme2_gemv.cpp PROPERTIES COMPILE_OPTIONS "${_CACTUS_SME2_FLAG};-O1") + message(STATUS "cactus_kernels: SME2 enabled (matmul_sme2.cpp; matmul_sme2_gemv.cpp at -O1)") +else() + message(STATUS "cactus_kernels: SME2 not supported by toolchain; SME TUs use scalar stubs") +endif() + if((BUILD_TESTING OR CACTUS_BUILD_TESTS) AND CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) file(GLOB TEST_SOURCES "tests/test_*.cpp") foreach(TEST_SRC ${TEST_SOURCES}) diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index 68f0a3522..a4bc5e6a9 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/cactus-kernels/cactus_kernels.h @@ -219,8 +219,7 @@ struct CactusQuantMatrix { const int8_t* right_signs; const uint32_t* permutation; const __fp16* rotation; - // Transposed rotation (rot_t[i][k] = R[k][i], fp16 K x K), stored alongside `rotation` in - // panel-format files so the orthogonal GEMV phase-A needs no derived buffers. + // rot_t[i][k] = R[k][i] (fp16 K x K), stored in panel-format files. const __fp16* rotation_t; const int8_t* expanded; const float* norm_f32; @@ -288,19 +287,24 @@ void cactus_quant_orthogonal_matmul( uint32_t M, __fp16* C); -// Reference encoder for the packed-panel format (row-major or INTERLEAVED_4ROW source layout). -// panels_out: SB64*num_groups*(group_size/4)*128 bytes (SB64 = ceil(N/64)); norms_out: -// SB64*num_groups*64 floats. See CactusQuantMatrix::packed_panels. Production panel files are -// written by the transpiler; this is the byte-exact reference its output is tested against. +// Byte-exact reference encoder for the panel format (production panels come from the transpiler). void cactus_quant_build_panels(const CactusQuantMatrix* W, uint8_t* panels_out, float* norms_out); -// Orthogonal-rotation CQ4 GEMV over packed panels (lm_head). W2 = virtual-group view of the -// orthogonal matrix as stored in panel files (group_size=128, num_groups=K/128, norms replicated -// per group, packed_panels/norm_panels set); rot_t = K x K fp16 rotation TRANSPOSED -// (rot_t[i][k] = R[k][i]). +// Orthogonal CQ4 GEMV over panels (lm_head): W2 is the virtual-128-group view stored in panel +// files, rot_t the transposed rotation. void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot_t, const __fp16* input_scale_recip, const __fp16* A, __fp16* C); +// Backend override for tests/benches: 0 auto, 1 force NEON, 2 force SME2 leaves. +int cactus_quant_set_backend(int backend); +int cactus_quant_sme_available(void); +// 1 when SME runtime paths should be used (SME2 present + SVL == 64, backend != force-NEON). +int cactus_quant_sme_enabled(void); +// 1 when the SME prefill-attention path should be used (sme_enabled + CACTUS_SME_ATTENTION != 0). +int cactus_quant_sme_attn_enabled(void); +// SME GEMV workers: -1 = env CACTUS_SME_GEMV_WORKERS / auto flat k = min(2, nt-1); 0 = NEON only. +int cactus_quant_set_sme_gemv_workers(int n); + void cactus_quant_4bit_gemv_interleaved( const CactusQuantMatrix* W, const uint8_t* packed_interleaved, @@ -356,9 +360,7 @@ void cactus_quant_dequantize_orthogonal_embedding_row( uint32_t flags, __fp16* out_row); -// Batched + parallelized version of the above for num_rows unique rows (same fp32-accumulate -// math; the per-row variant is a serial scalar K^2 matvec). INTERLEAVED_4ROW only (the -// production orthogonal-embedding format; no-op otherwise). out_rows: [num_rows][K]. +// Batched variant; INTERLEAVED_4ROW only (no-op otherwise). out_rows: [num_rows][K]. void cactus_quant_dequantize_orthogonal_embedding_rows( uint32_t bits, uint32_t K, @@ -372,9 +374,7 @@ void cactus_quant_dequantize_orthogonal_embedding_rows( uint32_t flags, __fp16* out_rows); -// Panel-format variant of the batched embedding dequant: decodes rows straight from the -// file-borne packed panels with the UNFOLDED fp16 norms (row-major [N][num_groups] in panel -// files), then the same vectorized un-rotation. out_rows: [num_rows][K]. +// Panel-format variant: decodes rows from the packed panels with the unfolded fp16 norms. void cactus_quant_dequantize_orthogonal_embedding_rows_panels( uint32_t K, uint32_t group_size, diff --git a/cactus-kernels/src/attention_hybrid.cpp b/cactus-kernels/src/attention_hybrid.cpp index 044ff1408..94d4e1d98 100644 --- a/cactus-kernels/src/attention_hybrid.cpp +++ b/cactus-kernels/src/attention_hybrid.cpp @@ -1,5 +1,6 @@ #include "../cactus_kernels.h" #include "threading.h" +#include "matmul_simd.h" #include #include #include @@ -326,6 +327,478 @@ static void cactus_attention_hybrid_int8_fp16_decode_dot( }); } +// SME2 prefill attention over the cached-int8 segment: with position_offset >= cache_len the +// segment is fully causally visible to every query row, so no online softmax is needed there. +static void cactus_attention_sme_prefill( + const __fp16* queries, + const int8_t* keys_cached, + const int8_t* values_cached, + const float* k_scales, + const float* v_scales, + const __fp16* keys_new, + const __fp16* values_new, + __fp16* output, + size_t batch_size, + size_t seq_len, + size_t cache_len, + size_t new_len, + size_t num_q_heads, + size_t num_kv_heads, + size_t head_dim, + float scale, + size_t position_offset, + bool is_causal, + size_t window_size, + size_t quant_group_size +) { + constexpr size_t QTILE = 16; + constexpr size_t KVBLK = 64; + constexpr size_t SINK_SIZE = 4; + const float NEG_INF = -std::numeric_limits::infinity(); + + const size_t kv_seq_len = cache_len + new_len; + const size_t n_qgroups = head_dim / quant_group_size; + const size_t dim_groups = quant_group_size / 4; + const size_t num_passes = head_dim / KVBLK; + const size_t num_blocks = (cache_len + KVBLK - 1) / KVBLK; + const size_t gqa_group_size = num_q_heads / num_kv_heads; + const size_t num_qtiles = (seq_len + QTILE - 1) / QTILE; + const size_t num_accum_slots = head_dim / 8; + + const size_t q_batch_stride = seq_len * num_q_heads * head_dim; + const size_t k_cached_batch_stride = cache_len * num_kv_heads * head_dim; + const size_t v_cached_batch_stride = cache_len * num_kv_heads * head_dim; + const size_t k_new_batch_stride = new_len * num_kv_heads * head_dim; + const size_t v_new_batch_stride = new_len * num_kv_heads * head_dim; + const size_t o_batch_stride = seq_len * num_q_heads * head_dim; + const size_t q_seq_stride = num_q_heads * head_dim; + const size_t k_seq_stride = num_kv_heads * head_dim; + const size_t v_seq_stride = num_kv_heads * head_dim; + const size_t o_seq_stride = num_q_heads * head_dim; + + const size_t cache_abs_offset = position_offset - cache_len; // gate guarantees >= 0 + + // Packed-KV scratch built once per call, shared by every q-head tile in the GQA group; + // tail kv beyond cache_len pack as zeros with zero scales. + const size_t BH = batch_size * num_kv_heads; + const size_t kpack_blk = head_dim * KVBLK; + const size_t vpack_pass_blk = KVBLK * KVBLK; + const size_t scale_blk = n_qgroups * KVBLK; + std::vector kpack_all(BH * num_blocks * kpack_blk); + std::vector vpack_all(BH * num_passes * num_blocks * vpack_pass_blk); + std::vector kst_all(BH * num_blocks * KVBLK); // flat per-kv K scale + std::vector vst_all(BH * num_blocks * scale_blk); + + const size_t total_pack_items = BH * num_blocks; + const size_t n_pack_slots = std::min(total_pack_items, + CactusThreading::get_thread_pool().num_workers()); + CactusThreading::parallel_for(n_pack_slots, CactusThreading::ParallelConfig{1, 1}, + [&](size_t slot_start, size_t slot_end) { + alignas(16) int8_t zero_row[512] = {0}; + alignas(16) int8_t krq[512]; + for (size_t slot = slot_start; slot < slot_end; ++slot) + for (size_t item = slot; item < total_pack_items; item += n_pack_slots) { + const size_t bh = item / num_blocks; + const size_t blk = item % num_blocks; + const size_t batch_idx = bh / num_kv_heads; + const size_t kv_head_idx = bh % num_kv_heads; + const int8_t* K_base = keys_cached + batch_idx * k_cached_batch_stride; + const int8_t* V_base = values_cached + batch_idx * v_cached_batch_stride; + int8_t* kp = kpack_all.data() + (bh * num_blocks + blk) * kpack_blk; + float* kst = kst_all.data() + (bh * num_blocks + blk) * KVBLK; + float* vst = vst_all.data() + (bh * num_blocks + blk) * scale_blk; + + // Requantize each kv row to ONE flat scale (ksflat = max_g ks[g], ratios folded + // into the int8 values), gathered as [dim-quad][vec(c/16)][c%16][4]. + const size_t words = head_dim / 4; + for (size_t c = 0; c < KVBLK; ++c) { + const size_t kvp = blk * KVBLK + c; + const bool valid = kvp < cache_len; + const int8_t* krow = zero_row; + if (valid) { + const int8_t* ksrc = K_base + kvp * k_seq_stride + kv_head_idx * head_dim; + const float* ks = k_scales + (kvp * num_kv_heads + kv_head_idx) * n_qgroups; + const float* vs = v_scales + (kvp * num_kv_heads + kv_head_idx) * n_qgroups; + float ksflat = 0.0f; + for (size_t g = 0; g < n_qgroups; ++g) ksflat = std::max(ksflat, ks[g]); + kst[c] = ksflat; + const float kinv = ksflat > 0.0f ? 1.0f / ksflat : 0.0f; + for (size_t g = 0; g < n_qgroups; ++g) { + vst[g * KVBLK + c] = vs[g]; + const float32x4_t rho = vdupq_n_f32(ks[g] * kinv); + for (size_t i = g * quant_group_size; i < (g + 1) * quant_group_size; i += 8) { + int16x8_t k16 = vmovl_s8(vld1_s8(ksrc + i)); + float32x4_t lo = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(k16))), rho); + float32x4_t hi = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(k16))), rho); + int16x4_t l16 = vqmovn_s32(vcvtnq_s32_f32(lo)); + int16x4_t h16 = vqmovn_s32(vcvtnq_s32_f32(hi)); + vst1_s8(krq + i, vqmovn_s16(vcombine_s16(l16, h16))); + } + } + krow = krq; + } else { + kst[c] = 0.0f; + for (size_t g = 0; g < n_qgroups; ++g) vst[g * KVBLK + c] = 0.0f; + } + int8_t* dstc = kp + (c >> 4) * 64 + (c & 15) * 4; + for (size_t w = 0; w < words; ++w) + memcpy(dstc + w * 256, krow + w * 4, 4); + } + + // V per 64-dim slice: [kg][tile][16 dims][4 kv] via 4-way byte interleave + for (size_t kg = 0; kg < 16; ++kg) { + const int8_t* vrow[4]; + for (size_t c4 = 0; c4 < 4; ++c4) { + const size_t kvp = blk * KVBLK + kg * 4 + c4; + vrow[c4] = (kvp < cache_len) + ? V_base + kvp * v_seq_stride + kv_head_idx * head_dim : zero_row; + } + for (size_t pass = 0; pass < num_passes; ++pass) { + int8_t* vp = vpack_all.data() + + ((bh * num_passes + pass) * num_blocks + blk) * vpack_pass_blk + kg * 256; + for (size_t t = 0; t < 4; ++t) { + const size_t off = pass * 64 + t * 16; + int8x16x4_t q; + q.val[0] = vld1q_s8(vrow[0] + off); + q.val[1] = vld1q_s8(vrow[1] + off); + q.val[2] = vld1q_s8(vrow[2] + off); + q.val[3] = vld1q_s8(vrow[3] + off); + vst4q_s8(vp + t * 64, q); + } + } + } + } + }); + + // Strided round-robin over tiles: parallel_for's static splitter dumps the remainder on + // the last worker (measured ~60% pool idle); workers walk slot, slot+nw, ... + const size_t total_tiles = batch_size * num_q_heads * num_qtiles; + const size_t n_slots = std::min(total_tiles, + CactusThreading::get_thread_pool().num_workers()); + CactusThreading::parallel_for(n_slots, + CactusThreading::ParallelConfig{1, 1}, + [&](size_t slot_start, size_t slot_end) { + std::vector qpack(n_qgroups * dim_groups * 64); + std::vector qs_scaled(QTILE); + std::vector qk_partials(num_blocks * 16 * 64); + const size_t srow_stride = num_blocks * KVBLK; + std::vector scores(QTILE * srow_stride); + std::vector m_c(QTILE), sum_c(QTILE); + std::vector ps_blk(QTILE * n_qgroups * num_blocks); + std::vector inv_ps_blk(QTILE * n_qgroups * num_blocks); + std::vector ppack(num_blocks * 2048); + std::vector av_out(num_blocks * 16 * 64); + std::vector oacc(QTILE * head_dim); + size_t mask_lo[QTILE], mask_hi[QTILE]; + std::vector output_accum_low(num_accum_slots); + std::vector output_accum_high(num_accum_slots); + std::vector block_accum(num_accum_slots); + float block_scores[32]; + + for (size_t slot = slot_start; slot < slot_end; ++slot) + for (size_t work_idx = slot; work_idx < total_tiles; work_idx += n_slots) { + const size_t batch_idx = work_idx / (num_q_heads * num_qtiles); + const size_t rem = work_idx % (num_q_heads * num_qtiles); + const size_t q_head_idx = rem / num_qtiles; + const size_t tile_idx = rem % num_qtiles; + const size_t kv_head_idx = q_head_idx / gqa_group_size; + const size_t bh = batch_idx * num_kv_heads + kv_head_idx; + const size_t tile0 = tile_idx * QTILE; + const size_t rows = std::min(QTILE, seq_len - tile0); + + const __fp16* Q_base = queries + batch_idx * q_batch_stride; + const __fp16* K_new_base = keys_new + batch_idx * k_new_batch_stride; + const __fp16* V_new_base = values_new + batch_idx * v_new_batch_stride; + __fp16* O_base = output + batch_idx * o_batch_stride; + const int8_t* kpack_bh = kpack_all.data() + bh * num_blocks * kpack_blk; + const float* kst_bh = kst_all.data() + bh * num_blocks * KVBLK; + const float* vst_bh = vst_all.data() + bh * num_blocks * scale_blk; + + // 1) quantize + pack Q with one flat scale per row + memset(qpack.data(), 0, qpack.size()); + for (size_t r = 0; r < rows; ++r) { + const __fp16* q_vec = Q_base + (tile0 + r) * q_seq_stride + q_head_idx * head_dim; + float32x4_t amx = vdupq_n_f32(0.0f); + for (size_t i = 0; i < head_dim; i += 8) { + float16x8_t v = vld1q_f16(q_vec + i); + amx = vmaxq_f32(amx, vmaxq_f32( + vabsq_f32(vcvt_f32_f16(vget_low_f16(v))), + vabsq_f32(vcvt_f32_f16(vget_high_f16(v))))); + } + const float amax = vmaxvq_f32(amx); + const float inv = amax > 0.0f ? 127.0f / amax : 0.0f; + qs_scaled[r] = scale * (amax / 127.0f); + const float32x4_t invv = vdupq_n_f32(inv); + int8_t* qdst = qpack.data() + r * 4; + for (size_t i = 0; i < head_dim; i += 8) { + float16x8_t v = vld1q_f16(q_vec + i); + int32x4_t i0 = vcvtnq_s32_f32(vmulq_f32(vcvt_f32_f16(vget_low_f16(v)), invv)); + int32x4_t i1 = vcvtnq_s32_f32(vmulq_f32(vcvt_f32_f16(vget_high_f16(v)), invv)); + int8x8_t s8 = vqmovn_s16(vcombine_s16(vqmovn_s32(i0), vqmovn_s32(i1))); + vst1_lane_s32(reinterpret_cast(qdst + (i / 4) * 64), + vreinterpret_s32_s8(s8), 0); + vst1_lane_s32(reinterpret_cast(qdst + (i / 4 + 1) * 64), + vreinterpret_s32_s8(s8), 1); + } + } + + // 2) per-row contiguous window-mask range [lo, hi) over the cached segment + size_t skip_lo_max = 0, skip_hi_min = 0; + bool any_mask = false; + for (size_t r = 0; r < rows; ++r) { + const size_t absolute_q_pos = position_offset + tile0 + r; + size_t kv_start = 0; + if (window_size > 0 && absolute_q_pos > window_size) + kv_start = absolute_q_pos - window_size; + size_t lo = 0, hi = 0; + if (window_size > 0 && kv_start > 0) { + lo = (cache_abs_offset > 0) ? SINK_SIZE : 0; + hi = (kv_start > cache_abs_offset) ? kv_start - cache_abs_offset : 0; + hi = std::min(hi, cache_len); + if (hi < lo) hi = lo; + } + mask_lo[r] = lo; + mask_hi[r] = hi; + if (r == 0) { skip_lo_max = lo; skip_hi_min = hi; } + else { + skip_lo_max = std::max(skip_lo_max, lo); + skip_hi_min = std::min(skip_hi_min, hi); + } + if (hi > lo) any_mask = true; + } + + // 3) SME QK over the cached segment -> fp32 scores (leaf skips the masked run) + size_t skipA = num_blocks, skipB = num_blocks; + if (any_mask) { + const size_t a = (skip_lo_max + KVBLK - 1) / KVBLK; + const size_t b = skip_hi_min / KVBLK; + if (a < b) { skipA = a; skipB = b; } + } + if (skipA > 0) + cactus_sme_attn_qk_seg(qpack.data(), kpack_bh, qk_partials.data(), + (uint32_t)(head_dim / 4), (uint32_t)skipA); + if (skipB < num_blocks) + cactus_sme_attn_qk_seg(qpack.data(), kpack_bh + skipB * kpack_blk, + qk_partials.data() + skipB * 1024, + (uint32_t)(head_dim / 4), (uint32_t)(num_blocks - skipB)); + for (size_t blk = 0; blk < num_blocks; ++blk) { + const size_t c0 = blk * KVBLK; + if (blk >= skipA && blk < skipB) { + for (size_t r = 0; r < rows; ++r) { + float* srow = scores.data() + r * srow_stride + c0; + for (size_t c = 0; c < KVBLK; ++c) srow[c] = NEG_INF; + } + continue; + } + const float* kstf = kst_bh + blk * KVBLK; + for (size_t r = 0; r < rows; ++r) { + float* srow = scores.data() + r * srow_stride + c0; + const float32x4_t qsv = vdupq_n_f32(qs_scaled[r]); + const int32_t* p64 = qk_partials.data() + blk * 1024 + r * 64; + for (size_t c = 0; c < KVBLK; c += 4) + vst1q_f32(srow + c, vmulq_f32( + vmulq_f32(vcvtq_f32_s32(vld1q_s32(p64 + c)), qsv), + vld1q_f32(kstf + c))); + } + } + for (size_t r = 0; r < rows; ++r) { + float* srow = scores.data() + r * srow_stride; + for (size_t c = cache_len; c < srow_stride; ++c) srow[c] = NEG_INF; + for (size_t c = mask_lo[r]; c < mask_hi[r]; ++c) srow[c] = NEG_INF; + } + + // 4) masked softmax stats per row (exp in place) + for (size_t r = 0; r < rows; ++r) { + float* srow = scores.data() + r * srow_stride; + float32x4_t mv = vdupq_n_f32(NEG_INF); + for (size_t c = 0; c < srow_stride; c += 4) + mv = vmaxq_f32(mv, vld1q_f32(srow + c)); + const float m = vmaxvq_f32(mv); + m_c[r] = m; + if (!(m > NEG_INF)) { + sum_c[r] = 0.0f; + memset(srow, 0, srow_stride * sizeof(float)); + continue; + } + // vectorized exp; masked -inf lanes clamp to exp(-87) and round to 0 in u8 + const float32x4_t mv4 = vdupq_n_f32(m); + float32x4_t sumv = vdupq_n_f32(0.0f); + for (size_t c = 0; c < srow_stride; c += 4) { + const float32x4_t e = fast_exp_f32x4(vsubq_f32(vld1q_f32(srow + c), mv4)); + vst1q_f32(srow + c, e); + sumv = vaddq_f32(sumv, e); + } + sum_c[r] = vaddvq_f32(sumv); + } + + // 5) per-(row, v-group, block) P quantization scales + for (size_t r = 0; r < rows; ++r) { + const float* prow = scores.data() + r * srow_stride; + for (size_t g = 0; g < n_qgroups; ++g) { + float* psd = ps_blk.data() + (r * n_qgroups + g) * num_blocks; + float* ipsd = inv_ps_blk.data() + (r * n_qgroups + g) * num_blocks; + for (size_t blk = 0; blk < num_blocks; ++blk) { + const float* vs64 = vst_bh + blk * scale_blk + g * KVBLK; + const float* p64 = prow + blk * KVBLK; + float32x4_t mx = vdupq_n_f32(0.0f); + for (size_t c = 0; c < KVBLK; c += 4) + mx = vmaxq_f32(mx, vmulq_f32(vld1q_f32(p64 + c), vld1q_f32(vs64 + c))); + const float m = vmaxvq_f32(mx); + psd[blk] = m / 255.0f; + ipsd[blk] = m > 0.0f ? 255.0f / m : 0.0f; + } + } + } + + // 6) SME AV per 64-dim slice: P folded with vs and block scale, u8 via USMOPA + for (size_t pass = 0; pass < num_passes; ++pass) { + memset(ppack.data(), 0, ppack.size()); + for (size_t grp = 0; grp < 2; ++grp) { + const size_t g = pass * 2 + grp; + for (size_t r = 0; r < rows; ++r) { + const float* prow = scores.data() + r * srow_stride; + const float* ipsd = inv_ps_blk.data() + (r * n_qgroups + g) * num_blocks; + for (size_t blk = 0; blk < num_blocks; ++blk) { + const float32x4_t invv = vdupq_n_f32(ipsd[blk]); + const float* vs64 = vst_bh + blk * scale_blk + g * KVBLK; + const float* p64 = prow + blk * KVBLK; + uint8_t* pdst = ppack.data() + blk * 2048 + grp * 1024 + r * 4; + for (size_t kg = 0; kg < 16; ++kg) { + float32x4_t f = vmulq_f32(vmulq_f32( + vld1q_f32(p64 + kg * 4), vld1q_f32(vs64 + kg * 4)), invv); + int16x4_t i16 = vqmovn_s32(vcvtnq_s32_f32(f)); + uint8x8_t u8 = vqmovun_s16(vcombine_s16(i16, i16)); + vst1_lane_u32(reinterpret_cast(pdst + kg * 64), + vreinterpret_u32_u8(u8), 0); + } + } + } + } + cactus_sme_attn_av_pass(ppack.data(), + vpack_all.data() + (bh * num_passes + pass) * num_blocks * vpack_pass_blk, + av_out.data(), (uint32_t)num_blocks); + for (size_t r = 0; r < rows; ++r) { + for (size_t grp = 0; grp < 2; ++grp) { + const size_t g = pass * 2 + grp; + const float* psd = ps_blk.data() + (r * n_qgroups + g) * num_blocks; + float32x4_t acc[8]; + for (size_t i = 0; i < 8; ++i) acc[i] = vdupq_n_f32(0.0f); + for (size_t blk = 0; blk < num_blocks; ++blk) { + const float32x4_t psv = vdupq_n_f32(psd[blk]); + const int32_t* src = av_out.data() + blk * 1024 + r * 64 + grp * 32; + for (size_t d = 0; d < 32; d += 4) + acc[d / 4] = vfmaq_f32(acc[d / 4], + vcvtq_f32_s32(vld1q_s32(src + d)), psv); + } + float* dst = oacc.data() + r * head_dim + pass * 64 + grp * 32; + for (size_t d = 0; d < 32; d += 4) vst1q_f32(dst + d, acc[d / 4]); + } + } + } + + // 7) seed the flash state from the cached segment, then the incumbent fp16 loop + for (size_t r = 0; r < rows; ++r) { + const size_t q_pos = tile0 + r; + const size_t absolute_q_pos = position_offset + q_pos; + const __fp16* q_vec = Q_base + q_pos * q_seq_stride + q_head_idx * head_dim; + __fp16* o_vec = O_base + q_pos * o_seq_stride + q_head_idx * head_dim; + + float running_max = m_c[r]; + float running_sum = sum_c[r]; + const float* orow = oacc.data() + r * head_dim; + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vld1q_f32(orow + i * 8); + output_accum_high[i] = vld1q_f32(orow + i * 8 + 4); + } + + size_t kv_start = 0; + if (window_size > 0 && absolute_q_pos > window_size) + kv_start = absolute_q_pos - window_size; + const size_t kv_end = is_causal + ? std::min(kv_seq_len, cache_len + q_pos + 1) : kv_seq_len; + + for (size_t kv_block_start = cache_len; kv_block_start < kv_end; kv_block_start += 32) { + const size_t kv_block_end = std::min(kv_block_start + 32, kv_end); + const size_t block_size = kv_block_end - kv_block_start; + + float block_max = NEG_INF; + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const size_t kv_pos = kv_block_start + kv_idx; + bool window_masked = false; + if (window_size > 0 && kv_start > 0) + window_masked = (kv_pos + cache_abs_offset < kv_start); + if ((is_causal && kv_pos > absolute_q_pos) || window_masked) { + block_scores[kv_idx] = NEG_INF; + continue; + } + const size_t new_pos = kv_pos - cache_len; + const __fp16* k_vec = K_new_base + new_pos * k_seq_stride + kv_head_idx * head_dim; + float16x8_t s_acc = vdupq_n_f16((__fp16)0.0f); + for (size_t d = 0; d < head_dim; d += 8) + s_acc = vfmaq_f16(s_acc, vld1q_f16(q_vec + d), vld1q_f16(k_vec + d)); + float score = vaddvq_f32(vcvt_f32_f16(vget_low_f16(s_acc))) + + vaddvq_f32(vcvt_f32_f16(vget_high_f16(s_acc))); + score *= scale; + block_scores[kv_idx] = score; + block_max = std::max(block_max, score); + } + + if (block_max > NEG_INF) { + const float scale_correction = expf(running_max - block_max); + running_sum *= scale_correction; + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vmulq_n_f32(output_accum_low[i], scale_correction); + output_accum_high[i] = vmulq_n_f32(output_accum_high[i], scale_correction); + } + running_max = block_max; + } + + float block_sum = 0.0f; + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + if (block_scores[kv_idx] != NEG_INF) { + block_scores[kv_idx] = expf(block_scores[kv_idx] - block_max); + block_sum += block_scores[kv_idx]; + } else { + block_scores[kv_idx] = 0.0f; + } + } + + for (size_t i = 0; i < num_accum_slots; ++i) + block_accum[i] = vdupq_n_f16((__fp16)0.0f); + + for (size_t kv_idx = 0; kv_idx < block_size; ++kv_idx) { + const float attn_weight = block_scores[kv_idx]; + if (attn_weight == 0.0f) continue; + const size_t new_pos = kv_block_start + kv_idx - cache_len; + const __fp16* v_vec = V_new_base + new_pos * v_seq_stride + kv_head_idx * head_dim; + const float16x8_t w_vec = vdupq_n_f16((__fp16)attn_weight); + for (size_t d = 0; d < head_dim; d += 8) + block_accum[d / 8] = vfmaq_f16(block_accum[d / 8], vld1q_f16(v_vec + d), w_vec); + } + + for (size_t i = 0; i < num_accum_slots; ++i) { + output_accum_low[i] = vaddq_f32(output_accum_low[i], vcvt_f32_f16(vget_low_f16(block_accum[i]))); + output_accum_high[i] = vaddq_f32(output_accum_high[i], vcvt_f32_f16(vget_high_f16(block_accum[i]))); + } + running_sum += block_sum; + } + + if (running_sum > 0.0f) { + const float32x4_t inv_sum_vec = vdupq_n_f32(1.0f / running_sum); + for (size_t d = 0; d < head_dim; d += 8) { + const size_t idx = d / 8; + vst1q_f16(o_vec + d, vcombine_f16( + vcvt_f16_f32(vmulq_f32(output_accum_low[idx], inv_sum_vec)), + vcvt_f16_f32(vmulq_f32(output_accum_high[idx], inv_sum_vec)))); + } + } else { + memset(o_vec, 0, head_dim * sizeof(__fp16)); + } + } + } + }); +} + void cactus_attention_hybrid_int8_fp16( const __fp16* queries, const int8_t* keys_cached, @@ -368,6 +841,26 @@ void cactus_attention_hybrid_int8_fp16( return; } + // SME2 prefill path: cached-int8 segment via 16x64 SMOPA tiles, new fp16 segment via the + // incumbent flash loop. Needs position_offset >= cache_len and cache_len <= 8192 (scratch). + if (seq_len >= 2 && + k_scales != nullptr && v_scales != nullptr && + head_dim == v_head_dim && + head_dim % 64 == 0 && head_dim <= 512 && + quant_group_size == 32 && + cache_len >= 64 && cache_len <= 8192 && + position_offset >= cache_len && + num_kv_heads > 0 && num_q_heads % num_kv_heads == 0 && + cactus_quant_sme_attn_enabled()) { + cactus_attention_sme_prefill( + queries, keys_cached, values_cached, k_scales, v_scales, + keys_new, values_new, output, + batch_size, seq_len, cache_len, new_len, + num_q_heads, num_kv_heads, head_dim, + scale, position_offset, is_causal, window_size, quant_group_size); + return; + } + const size_t kv_seq_len = cache_len + new_len; constexpr size_t VECTOR_WIDTH = 8; diff --git a/cactus-kernels/src/matmul.cpp b/cactus-kernels/src/matmul.cpp index 7effe2949..677a6e9f8 100644 --- a/cactus-kernels/src/matmul.cpp +++ b/cactus-kernels/src/matmul.cpp @@ -1,5 +1,6 @@ #include "../cactus_kernels.h" #include "threading.h" +#include "matmul_simd.h" #include #include #include @@ -501,6 +502,27 @@ static inline bool cactus_quant_has_panels(const CactusQuantMatrix* W) { return W->packed_panels != nullptr && W->norm_panels != nullptr; } +// SME2 backend override for tests/benchmarks: 0 = auto, 1 = force NEON, 2 = force SME2 leaves. +static std::atomic g_cactus_quant_backend{[] { + const char* e = getenv("CACTUS_QUANT_BACKEND"); + return e ? atoi(e) : 0; +}()}; +static std::atomic g_sme_gemv_workers{-1}; + +static inline int cactus_quant_backend_now() { + return g_cactus_quant_backend.load(std::memory_order_relaxed); +} +// Explicit SME GEMV worker override (setter wins over env), or -1 for the flat auto policy. +static inline int cactus_quant_sme_gemv_workers() { + const int o = g_sme_gemv_workers.load(std::memory_order_relaxed); + if (o >= 0) return o; + static const int env_override = [] { + const char* e = getenv("CACTUS_SME_GEMV_WORKERS"); + return e ? atoi(e) : -1; + }(); + return env_override; +} + static bool cactus_quant_valid_common(const CactusQuantMatrix* W, const void* A, void* C) { if (W == nullptr || A == nullptr || C == nullptr) return false; if (W->K == 0 || W->N == 0 || W->group_size == 0 || W->num_groups == 0) return false; @@ -514,6 +536,26 @@ static bool cactus_quant_valid_common(const CactusQuantMatrix* W, const void* A, } // namespace +int cactus_quant_set_backend(int backend) { + g_cactus_quant_backend.store(backend, std::memory_order_relaxed); + return cpu_has_sme2() ? 1 : 0; +} +int cactus_quant_sme_available(void) { return cpu_has_sme2() ? 1 : 0; } +int cactus_quant_sme_enabled(void) { + return (cpu_has_sme2() && cactus_quant_backend_now() != 1) ? 1 : 0; +} +int cactus_quant_sme_attn_enabled(void) { + static const int env_on = [] { + const char* e = getenv("CACTUS_SME_ATTENTION"); + return e ? atoi(e) : 1; + }(); + return (env_on && cactus_quant_sme_enabled()) ? 1 : 0; +} +int cactus_quant_set_sme_gemv_workers(int n) { + g_sme_gemv_workers.store(n, std::memory_order_relaxed); + return cpu_has_sme2() ? 1 : 0; +} + static void cactus_quant_dispatch_group_gemm( const CactusQuantMatrix* W, const __fp16* A, @@ -1647,9 +1689,56 @@ static void cactus_quant_panel_gemv_blocks( } } -// Fused panel GEMV (M=1): one pool dispatch = phase A (group-parallel Hadamard + int8 quantize, -// spin barrier) + phase B (dynamic super-block stealing). Main runs as worker 0 and spin-joins. -// Thread budget = ceil(SB64 / CACTUS_GEMV_SB_PER_THREAD), default 8 (mobile: battery > saturation). +// Rescale one super-block of raw int32 partials in NEON: +// C[sb*64+c] = sum_g partials[g][c] * act_scale[g] * norm_panels[sb][g][c]. +static void cactus_quant_panel_rescale_sb(const int32_t* psb, size_t sb, const float* norm_panels, + const float* act_scales, uint32_t num_groups, + uint32_t N, __fp16* C) { + float32x4_t acc[16]; + for (int v = 0; v < 16; ++v) acc[v] = vdupq_n_f32(0.f); + const float* nsb = norm_panels + sb * num_groups * 64; + for (uint32_t g = 0; g < num_groups; ++g) { + const float32x4_t as = vdupq_n_f32(act_scales[g]); + const int32_t* pp = psb + static_cast(g) * 64; + const float* nf = nsb + static_cast(g) * 64; + for (int v = 0; v < 16; ++v) + acc[v] = vmlaq_f32(acc[v], + vmulq_f32(vcvtq_f32_s32(vld1q_s32(pp + v * 4)), as), + vld1q_f32(nf + v * 4)); + } + const size_t n_start = sb * 64; + const uint32_t valid = std::min(64, N - static_cast(n_start)); + float tmp[64]; + for (int v = 0; v < 16; ++v) vst1q_f32(tmp + v * 4, acc[v]); + for (uint32_t c = 0; c < valid; ++c) C[n_start + c] = static_cast<__fp16>(tmp[c]); +} + +// SME workers replace NEON workers inside the budget; flat k=2 is the measured speed/power frontier. +static inline size_t cactus_quant_panel_k_sme(size_t nt, uint32_t gs) { + if (!cactus_quant_sme_enabled() || (gs % 32) != 0) return 0; + const size_t cap = nt > 1 ? nt - 1 : 1; + const int ov = cactus_quant_sme_gemv_workers(); + if (cactus_quant_backend_now() == 2) + return std::max(1, std::min(ov > 0 ? static_cast(ov) : 2, cap)); + if (ov >= 0) return std::min(static_cast(ov), nt > 1 ? nt - 1 : 0); + return std::min(2, nt > 1 ? nt - 1 : 0); +} + +static void cactus_quant_panel_sme_range(const CactusQuantMatrix* W, const int8_t* act_i8, + const float* act_scales, const uint8_t* zt, + uint32_t sb, uint32_t cnt, __fp16* C) { + thread_local std::vector tl_partials; + const uint32_t ng = W->num_groups; + const size_t need = static_cast(cnt) * ng * 64; + if (tl_partials.size() < need) tl_partials.resize(need); + cactus_sme_cq_gemv_luti4_s8(act_i8, W->packed_panels, zt, tl_partials.data(), + ng, W->group_size, sb, cnt); + for (uint32_t s = 0; s < cnt; ++s) + cactus_quant_panel_rescale_sb(tl_partials.data() + static_cast(s) * ng * 64, + static_cast(sb) + s, W->norm_panels, act_scales, + ng, W->N, C); +} + static void cactus_quant_panel_gemv(const CactusQuantMatrix* W, const __fp16* A, __fp16* C) { const uint32_t gs = W->group_size; const uint32_t num_groups = W->num_groups; @@ -1671,6 +1760,9 @@ static void cactus_quant_panel_gemv(const CactusQuantMatrix* W, const __fp16* A, const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 1u << W->bits); (void)cb_scale; // folded into norm_panels by the transpiler const int8x16_t cb_lut = vld1q_s8(cb_i8); + alignas(16) uint8_t zt[64] = {}; + for (int i = 0; i < 16; ++i) zt[4 * i] = static_cast(cb_i8[i]); + const bool sme = cactus_quant_sme_enabled() && (gs % 32) == 0; auto phase_a_group = [&](uint32_t g) { __fp16 basis[256]; // gs <= 256 (same bound as the transform's tmp) @@ -1680,43 +1772,23 @@ static void cactus_quant_panel_gemv(const CactusQuantMatrix* W, const __fp16* A, if (nt <= 1) { for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); - cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, 0, SB64, C); + if (sme) cactus_quant_panel_sme_range(W, act_i8, act_scales, zt, + 0, static_cast(SB64), C); + else cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, 0, SB64, C); return; } - std::atomic ga{0}; // phase A group grab - std::atomic a_done{0}; // phase A completion (release/acquire barrier) - std::atomic next{0}; // phase B super-block grab - std::atomic done{0}; // worker completion (spin-join, no cv wake) - auto worker = [&](size_t) { - for (uint32_t g; (g = ga.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { - phase_a_group(g); - a_done.fetch_add(1, std::memory_order_release); - } - while (a_done.load(std::memory_order_acquire) < num_groups) { /* spin */ } - for (;;) { - const uint32_t seen = next.load(std::memory_order_relaxed); - if (seen >= SB64) break; - const uint32_t want = (SB64 - seen > 4u * nt) ? 4u : 1u; - const uint32_t sb = next.fetch_add(want, std::memory_order_relaxed); - if (sb >= SB64) break; - const uint32_t cnt = std::min(want, static_cast(SB64) - sb); - cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, - sb, static_cast(sb) + cnt, C); - } - }; - pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { - worker(wid + 1); - done.fetch_add(1, std::memory_order_release); - }); - worker(0); - while (done.load(std::memory_order_acquire) < nt - 1) { /* spin */ } + const size_t k_sme = sme ? cactus_quant_panel_k_sme(nt, gs) : 0; + cactus_quant_two_phase_run(nt, num_groups, static_cast(SB64), phase_a_group, + [&](size_t wid, uint32_t sb, uint32_t cnt) { + if (wid < k_sme) cactus_quant_panel_sme_range(W, act_i8, act_scales, zt, sb, cnt, C); + else cactus_quant_panel_gemv_blocks(W, act_i8, act_scales, cb_lut, + sb, static_cast(sb) + cnt, C); + }); } -// Orthogonal-rotation CQ4 GEMV (the Gemma lm_head). Panel files store the lm_head with virtual -// 128-wide groups already applied (exact regrouping: norms constant across subgroups, panel -// kg-blocks K-chunk-ordered). Phase A rotates+quantizes each group via rot_t (fp32 accumulate); -// phase B is the panel super-block stealing loop. +// Orthogonal CQ4 GEMV (lm_head): panel files store virtual 128-wide groups, folded norms, and +// the transposed rotation. void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot_t, const __fp16* input_scale_recip, const __fp16* A, __fp16* C) { const uint32_t K = W2->K; @@ -1746,6 +1818,9 @@ void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot const float cb_scale = tq_quantize_codebook_i8(W2->codebook, cb_i8, 1u << W2->bits); (void)cb_scale; // folded into norm_panels by the transpiler const int8x16_t cb_lut = vld1q_s8(cb_i8); + alignas(16) uint8_t zt[64] = {}; + for (int i = 0; i < 16; ++i) zt[4 * i] = static_cast(cb_i8[i]); + const bool sme = cactus_quant_sme_enabled() && (gs % 32) == 0; // rotate-and-quantize one virtual group: A_rot[i] = sum_k a_scaled[k]*R[k][i] = dot(a_scaled, rot_t[i]) auto phase_a_group = [&](uint32_t g) { __fp16 basis[128]; @@ -1765,39 +1840,19 @@ void cactus_quant_orth_panel_gemv(const CactusQuantMatrix* W2, const __fp16* rot if (nt <= 1) { for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); - cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, 0, SB64, C); + if (sme) cactus_quant_panel_sme_range(W2, act_i8, act_scales, zt, + 0, static_cast(SB64), C); + else cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, 0, SB64, C); return; } - std::atomic ga{0}; - std::atomic a_done{0}; - std::atomic next{0}; - std::atomic done{0}; - auto worker = [&](size_t) { - for (uint32_t g; (g = ga.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { - phase_a_group(g); - a_done.fetch_add(1, std::memory_order_release); - } - while (a_done.load(std::memory_order_acquire) < num_groups) { /* spin */ } - for (;;) { - const uint32_t seen = next.load(std::memory_order_relaxed); - if (seen >= SB64) break; - const uint32_t want = (SB64 - seen > 4u * nt) ? 4u : 1u; - const uint32_t sb = next.fetch_add(want, std::memory_order_relaxed); - if (sb >= SB64) break; - const uint32_t cnt = std::min(want, static_cast(SB64) - sb); - cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, - sb, static_cast(sb) + cnt, C); - } - }; - // Main participates as worker 0 + spin-join (no wait_all cv sleep) — see the GEMV driver. - // (nt >= 2 here: the nt <= 1 case took the serial path above.) - pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { - worker(wid + 1); - done.fetch_add(1, std::memory_order_release); - }); - worker(0); - while (done.load(std::memory_order_acquire) < nt - 1) { /* spin */ } + const size_t k_sme = sme ? cactus_quant_panel_k_sme(nt, gs) : 0; + cactus_quant_two_phase_run(nt, num_groups, static_cast(SB64), phase_a_group, + [&](size_t wid, uint32_t sb, uint32_t cnt) { + if (wid < k_sme) cactus_quant_panel_sme_range(W2, act_i8, act_scales, zt, sb, cnt, C); + else cactus_quant_panel_gemv_blocks(W2, act_i8, act_scales, cb_lut, + sb, static_cast(sb) + cnt, C); + }); } void cactus_quant_matmul( @@ -2023,9 +2078,12 @@ void cactus_quant_matmul( (void)cb_scale; // folded into norm_panels by the transpiler const int8x16_t cb_lut = vld1q_s8(cb_i8); const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + alignas(16) uint8_t zt[64] = {}; + for (int i2 = 0; i2 < 16; ++i2) zt[4 * i2] = static_cast(cb_i8[i2]); + // SMOPA beats the NEON panel GEMM at every measured M>1; backend 1 pins the NEON body. + const bool use_smopa = cactus_quant_sme_enabled() && (gs % 32) == 0; auto& pool = CactusThreading::get_thread_pool(); - // Same thread budget as the original M>1 NEON GEMM (one thread per 16-row M-tile). const size_t nt = std::max(1, std::min(pool.num_workers(), std::min(M_tiles, total_pairs))); const uint32_t total_a = static_cast(M_tiles * num_groups); @@ -2051,7 +2109,45 @@ void cactus_quant_matmul( for (uint32_t r = 0; r < 16; ++r) std::memcpy(dst + static_cast(kg) * 64 + r * 4, rowq[r] + kg * 4, 4); }; + // The leaf emits raw int32 partials [g][16 rows][64 ch]; the NEON rescale matches the + // fp32 body exactly. + auto run_pair_smopa = [&](size_t pair, int32_t* partials) { + const size_t mt = pair / SB64; + const size_t sb = pair % SB64; + cactus_sme_cq_gemm_luti4_s8( + act_packed.data() + (mt * num_groups) * static_cast(gs) * 16, + W->packed_panels + sb * static_cast(num_groups) * g_stride, + zt, partials, num_groups, gs); + const uint32_t m_rows = std::min(TILE_M16, M - static_cast(mt) * TILE_M16); + const uint32_t valid_n = std::min(64, N - static_cast(sb) * 64); + const float* nsb = W->norm_panels + sb * static_cast(num_groups) * 64; + for (uint32_t r = 0; r < m_rows; ++r) { + const uint32_t m = static_cast(mt) * TILE_M16 + r; + float32x4_t acc[16]; + for (int v = 0; v < 16; ++v) acc[v] = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const float32x4_t as = vdupq_n_f32(act_scales[static_cast(m) * num_groups + g]); + const int32_t* pp = partials + (static_cast(g) * 16 + r) * 64; + const float* nf = nsb + static_cast(g) * 64; + for (int v = 0; v < 16; ++v) + acc[v] = vmlaq_f32(acc[v], + vmulq_f32(vcvtq_f32_s32(vld1q_s32(pp + v * 4)), as), + vld1q_f32(nf + v * 4)); + } + float tmp[64]; + for (int v = 0; v < 16; ++v) vst1q_f32(tmp + v * 4, acc[v]); + __fp16* crow = C + static_cast(m) * N + sb * 64; + for (uint32_t c = 0; c < valid_n; ++c) crow[c] = static_cast<__fp16>(tmp[c]); + } + }; auto run_pair = [&](size_t pair) { + if (use_smopa) { + thread_local std::vector tl_partials; + const size_t need = static_cast(num_groups) * 16 * 64; + if (tl_partials.size() < need) tl_partials.resize(need); + run_pair_smopa(pair, tl_partials.data()); + return; + } const size_t mt = pair / SB64; const size_t sb = pair % SB64; const uint32_t m_rows = std::min(TILE_M16, M - static_cast(mt) * TILE_M16); diff --git a/cactus-kernels/src/matmul_simd.h b/cactus-kernels/src/matmul_simd.h new file mode 100644 index 000000000..965238d72 --- /dev/null +++ b/cactus-kernels/src/matmul_simd.h @@ -0,0 +1,43 @@ +#pragma once +// Declarations bridging the NEON TUs and the SME2 TUs; leaves are __arm_locally_streaming +// (self-transition + own ZA), so they call as ordinary functions. Stubbed without SME2. +#include + +extern "C" { + +// Streaming svcntb (0 without SME2); every layout is pinned to SVL 64 and dispatch gates on it. +uint32_t cactus_sme2_svl_bytes(void); + +// M=1 per-group SMOPA over one 4-channel block; raw int32 partials, caller rescales in NEON. +void cactus_sme_cq_gemv_4col_s8(const int8_t* act_i8, const int8_t* w_block, + int32_t* partials_out, uint32_t num_groups, uint32_t gs); + +// M>1 full 16x16 ZA tile; act_packed and w_sme both [num_groups][gs/4][16][4] int8. +void cactus_sme_cq_gemm_16x16_s8(const int8_t* act_packed, const int8_t* w_sme, + int32_t* partials_out, uint32_t num_groups, uint32_t gs, + uint32_t m_rows); + +// LUTI4/ZT0 GEMV over packed panels: raw int32 partials [sb][group][64], zt_table byte 4*i +// = cb_i8[i]. Defined in matmul_sme2_gemv.cpp (-O1: clang miscompiles vg1x4 dots at -O2+). +void cactus_sme_cq_gemv_luti4_s8(const int8_t* act_i8, const uint8_t* esme_packed, + const uint8_t* zt_table, int32_t* partials, + uint32_t num_groups, uint32_t gs, + uint32_t sb_start, uint32_t sb_count); + +// SMOPA GEMM over one (16-row M-tile, 64-channel super-block) pair of packed panels; +// act_packed [num_groups][gs/4][16][4], partials [num_groups][16][64]. +void cactus_sme_cq_gemm_luti4_s8(const int8_t* act_packed, const uint8_t* esme_sb, + const uint8_t* zt_table, int32_t* partials, + uint32_t num_groups, uint32_t gs); + +// Prefill QK: 16 q-rows x whole cached segment with FLAT scales, one ZA readout per 64-kv +// block; qpack [hd/4][16][4], kpack [hd/4][4][16][4] per block, partials [block][16][64]. +void cactus_sme_attn_qk_seg(const int8_t* qpack, const int8_t* kpack, int32_t* partials, + uint32_t dim_quads, uint32_t num_blocks); + +// Prefill AV: one 64-dim slice, u8 P x s8 V (USMOPA — softmax >= 0), per-block readout; +// ppack [block][2][16][16][4] u8, vpack [block][16][4][16][4] s8, out [block][16][64]. +void cactus_sme_attn_av_pass(const uint8_t* ppack, const int8_t* vpack, int32_t* out, + uint32_t num_blocks); + +} // extern "C" diff --git a/cactus-kernels/src/matmul_sme2.cpp b/cactus-kernels/src/matmul_sme2.cpp new file mode 100644 index 000000000..6fe4c5269 --- /dev/null +++ b/cactus-kernels/src/matmul_sme2.cpp @@ -0,0 +1,133 @@ +// SME2 streaming leaves: integer SMOPA accumulate only — transform, quantize, and fp rescale +// stay in NEON, so partials are bit-identical to the SDOT path. Scalar stubs without SME2. +#include "matmul_simd.h" + +#if defined(__ARM_FEATURE_SME2) +#include + +// Dispatch gates on SVL == 64; other SVLs fall back to NEON. +__arm_locally_streaming +uint32_t cactus_sme2_svl_bytes(void) { + return (uint32_t)svcntb(); +} + +// M=1: one act row vs up to 4 channels per group, ZA row 0 only. +__arm_locally_streaming __arm_new("za") +void cactus_sme_cq_gemv_4col_s8(const int8_t* act_i8, const int8_t* w_block, + int32_t* partials_out, uint32_t num_groups, uint32_t gs) { + const svbool_t pn = svwhilelt_b8_u32(0u, 4u); // activation: row 0, 4 K lanes + const svbool_t pm = svwhilelt_b8_u32(0u, 16u); // weights: 4 channels x 4 K + const svbool_t pst = svwhilelt_b32_u32(0u, 4u); // store 4 int32 (the 4 channels) + const uint32_t nkg = gs >> 2; + for (uint32_t g = 0; g < num_groups; ++g) { + const int8_t* aptr = act_i8 + (size_t)g * gs; + const int8_t* wptr = w_block + (size_t)g * gs * 4; + svzero_za(); + for (uint32_t kg = 0; kg < nkg; ++kg) { + svint8_t zn = svld1_s8(pn, aptr + (size_t)kg * 4); + svint8_t zm = svld1_s8(pm, wptr + (size_t)kg * 16); + svmopa_za32_s8_m(0, pn, pm, zn, zm); + } + svst1_hor_za32(0, 0, pst, partials_out + (size_t)g * 4); + } +} + +// M>1: full 16x16 ZA tile over pre-packed [num_groups][gs/4][16][4]. +__arm_locally_streaming __arm_new("za") +void cactus_sme_cq_gemm_16x16_s8(const int8_t* act_packed, const int8_t* w_sme, + int32_t* partials_out, uint32_t num_groups, uint32_t gs, + uint32_t m_rows) { + const svbool_t pn = svwhilelt_b8_u32(0u, m_rows * 4u); // m_rows x 4 K + const svbool_t pm = svptrue_b8(); // all 16 channels x 4 K + const svbool_t pst = svptrue_b32(); // store 16 int32 columns + const uint32_t nkg = gs >> 2; + for (uint32_t g = 0; g < num_groups; ++g) { + const int8_t* a = act_packed + (size_t)g * gs * 16; + const int8_t* w = w_sme + (size_t)g * gs * 16; + svzero_za(); + for (uint32_t kg = 0; kg < nkg; ++kg) { + svint8_t zn = svld1_s8(pn, a + (size_t)kg * 64); + svint8_t zm = svld1_s8(pm, w + (size_t)kg * 64); + svmopa_za32_s8_m(0, pn, pm, zn, zm); + } + for (uint32_t i = 0; i < m_rows; ++i) + svst1_hor_za32(0, i, pst, partials_out + ((size_t)g * 16 + i) * 16); + } +} + +// QK: ZA accumulates the full head_dim per 64-kv block, single readout per block; caller +// rescales by qs[r]*ksflat[c] in NEON. +__arm_locally_streaming __arm_new("za") +void cactus_sme_attn_qk_seg(const int8_t* qpack, const int8_t* kpack, int32_t* partials, + uint32_t dim_quads, uint32_t num_blocks) { + const svbool_t pg = svptrue_b8(); + const svbool_t pg32 = svptrue_b32(); + for (uint32_t b = 0; b < num_blocks; ++b) { + svzero_za(); + const int8_t* kb = kpack + (size_t)b * dim_quads * 256; + for (uint32_t w = 0; w < dim_quads; ++w) { + svint8_t zn = svld1_s8(pg, qpack + (size_t)w * 64); + const int8_t* p = kb + (size_t)w * 256; + svmopa_za32_s8_m(0, pg, pg, zn, svld1_s8(pg, p)); + svmopa_za32_s8_m(1, pg, pg, zn, svld1_s8(pg, p + 64)); + svmopa_za32_s8_m(2, pg, pg, zn, svld1_s8(pg, p + 128)); + svmopa_za32_s8_m(3, pg, pg, zn, svld1_s8(pg, p + 192)); + } + int32_t* ob = partials + (size_t)b * 1024; + for (uint32_t r = 0; r < 16; ++r) { + svst1_hor_za32(0, r, pg32, ob + (size_t)r * 64); + svst1_hor_za32(1, r, pg32, ob + (size_t)r * 64 + 16); + svst1_hor_za32(2, r, pg32, ob + (size_t)r * 64 + 32); + svst1_hor_za32(3, r, pg32, ob + (size_t)r * 64 + 48); + } + } +} + +// AV: one 64-dim slice, per-block readout (P scale is per row/v-group/block); P is u8 via +// USMOPA. zna feeds dim tiles 0-1, znb tiles 2-3. +__arm_locally_streaming __arm_new("za") +void cactus_sme_attn_av_pass(const uint8_t* ppack, const int8_t* vpack, int32_t* out, + uint32_t num_blocks) { + const svbool_t pg = svptrue_b8(); + const svbool_t pg32 = svptrue_b32(); + for (uint32_t b = 0; b < num_blocks; ++b) { + svzero_za(); + const uint8_t* pp = ppack + (size_t)b * 2048; + const int8_t* vp = vpack + (size_t)b * 4096; + for (uint32_t kg = 0; kg < 16; ++kg) { + svuint8_t zna = svld1_u8(pg, pp + (size_t)kg * 64); + svuint8_t znb = svld1_u8(pg, pp + 1024 + (size_t)kg * 64); + const int8_t* v = vp + (size_t)kg * 256; + svusmopa_za32_u8_m(0, pg, pg, zna, svld1_s8(pg, v)); + svusmopa_za32_u8_m(1, pg, pg, zna, svld1_s8(pg, v + 64)); + svusmopa_za32_u8_m(2, pg, pg, znb, svld1_s8(pg, v + 128)); + svusmopa_za32_u8_m(3, pg, pg, znb, svld1_s8(pg, v + 192)); + } + int32_t* ob = out + (size_t)b * 1024; + for (uint32_t r = 0; r < 16; ++r) { + svst1_hor_za32(0, r, pg32, ob + (size_t)r * 64); + svst1_hor_za32(1, r, pg32, ob + (size_t)r * 64 + 16); + svst1_hor_za32(2, r, pg32, ob + (size_t)r * 64 + 32); + svst1_hor_za32(3, r, pg32, ob + (size_t)r * 64 + 48); + } + } +} + +#else // ---- compile-time fallback (toolchain/target without SME2) ---- +uint32_t cactus_sme2_svl_bytes(void) { return 0; } +void cactus_sme_cq_gemv_4col_s8(const int8_t*, const int8_t*, int32_t* p, + uint32_t num_groups, uint32_t) { + for (uint32_t i = 0; i < num_groups * 4u; ++i) p[i] = 0; +} +void cactus_sme_attn_qk_seg(const int8_t*, const int8_t*, int32_t* p, + uint32_t, uint32_t num_blocks) { + for (uint32_t i = 0; i < num_blocks * 16u * 64u; ++i) p[i] = 0; +} +void cactus_sme_attn_av_pass(const uint8_t*, const int8_t*, int32_t* out, uint32_t num_blocks) { + for (uint32_t i = 0; i < num_blocks * 16u * 64u; ++i) out[i] = 0; +} +void cactus_sme_cq_gemm_16x16_s8(const int8_t*, const int8_t*, int32_t* p, + uint32_t num_groups, uint32_t, uint32_t) { + for (uint32_t i = 0; i < num_groups * 16u * 16u; ++i) p[i] = 0; +} +#endif diff --git a/cactus-kernels/src/matmul_sme2_gemv.cpp b/cactus-kernels/src/matmul_sme2_gemv.cpp new file mode 100644 index 000000000..11bb55d17 --- /dev/null +++ b/cactus-kernels/src/matmul_sme2_gemv.cpp @@ -0,0 +1,139 @@ +// LUTI4/ZT0 leaves emitting raw int32 partials to memory: ZA->core reads drain the shared +// in-order SME queue (stores do not). -O1 TU only: clang miscompiles vg1x4 dots at -O2+ (SIGILL). +#include "matmul_simd.h" + +#if defined(__ARM_FEATURE_SME2) +#include + +__arm_locally_streaming __arm_new("za") __arm_new("zt0") +void cactus_sme_cq_gemv_luti4_s8(const int8_t* act_i8, const uint8_t* esme_packed, + const uint8_t* zt_table, int32_t* partials, + uint32_t num_groups, uint32_t gs, + uint32_t sb_start, uint32_t sb_count) { + svldr_zt(0, zt_table); + const svbool_t pg = svptrue_b8(); + const svbool_t pg32 = svptrue_b32(); + const uint32_t nkg = gs >> 2; + const size_t g_stride = (size_t)nkg * 128; // packed bytes per group panel + const size_t sb_stride = (size_t)num_groups * g_stride; + + #define CACTUS_LUTI4_DOT(SL, P, OFF, AV, IDX) do { \ + svint8x2_t e0_ = svluti4_lane_zt_s8_x2(0, svld1_u8(pg, (P) + (OFF)), 0); \ + svint8x2_t e1_ = svluti4_lane_zt_s8_x2(0, svld1_u8(pg, (P) + (OFF) + 64), 0); \ + svdot_lane_za32_s8_vg1x4(SL, \ + svcreate4_s8(svget2_s8(e0_, 0), svget2_s8(e0_, 1), \ + svget2_s8(e1_, 0), svget2_s8(e1_, 1)), AV, IDX); \ + } while (0) + // vg1 ZA vector v stores at hor(tile 0, slice (v&3)*4 + (v>>2)); fire-and-forget. + #define CACTUS_STORE_QUAD(S, DST) do { \ + svst1_hor_za32(0, (((S) + 0) & 3u) * 4u + (((S) + 0) >> 2), pg32, (DST)); \ + svst1_hor_za32(0, (((S) + 1) & 3u) * 4u + (((S) + 1) >> 2), pg32, (DST) + 16); \ + svst1_hor_za32(0, (((S) + 2) & 3u) * 4u + (((S) + 2) >> 2), pg32, (DST) + 32); \ + svst1_hor_za32(0, (((S) + 3) & 3u) * 4u + (((S) + 3) >> 2), pg32, (DST) + 48); \ + } while (0) + + for (uint32_t s = 0; s < sb_count; ++s) { + const uint8_t* wsb = esme_packed + (size_t)(sb_start + s) * sb_stride; + int32_t* psb = partials + (size_t)s * num_groups * 64; + svzero_za(); + uint32_t g = 0; + for (; g + 3 < num_groups; g += 4) { // 4 groups on slice-quads 0/4/8/12 = 4 chains + const int8_t* aA = act_i8 + (size_t)g * gs; + const uint8_t* wA = wsb + (size_t)g * g_stride; + for (uint32_t kk = 0; kk < gs; kk += 16) { + svint8_t avA = svld1rq_s8(pg, aA + kk); + svint8_t avB = svld1rq_s8(pg, aA + gs + kk); + svint8_t avC = svld1rq_s8(pg, aA + 2 * gs + kk); + svint8_t avD = svld1rq_s8(pg, aA + 3 * gs + kk); + const uint8_t* pA = wA + (size_t)(kk >> 2) * 128; + const uint8_t* pB = pA + g_stride; + const uint8_t* pC = pB + g_stride; + const uint8_t* pD = pC + g_stride; + CACTUS_LUTI4_DOT(0, pA, 0, avA, 0); + CACTUS_LUTI4_DOT(4, pB, 0, avB, 0); + CACTUS_LUTI4_DOT(8, pC, 0, avC, 0); + CACTUS_LUTI4_DOT(12, pD, 0, avD, 0); + CACTUS_LUTI4_DOT(0, pA, 128, avA, 1); + CACTUS_LUTI4_DOT(4, pB, 128, avB, 1); + CACTUS_LUTI4_DOT(8, pC, 128, avC, 1); + CACTUS_LUTI4_DOT(12, pD, 128, avD, 1); + CACTUS_LUTI4_DOT(0, pA, 256, avA, 2); + CACTUS_LUTI4_DOT(4, pB, 256, avB, 2); + CACTUS_LUTI4_DOT(8, pC, 256, avC, 2); + CACTUS_LUTI4_DOT(12, pD, 256, avD, 2); + CACTUS_LUTI4_DOT(0, pA, 384, avA, 3); + CACTUS_LUTI4_DOT(4, pB, 384, avB, 3); + CACTUS_LUTI4_DOT(8, pC, 384, avC, 3); + CACTUS_LUTI4_DOT(12, pD, 384, avD, 3); + } + int32_t* pd = psb + (size_t)g * 64; + CACTUS_STORE_QUAD(0, pd); + CACTUS_STORE_QUAD(4, pd + 64); + CACTUS_STORE_QUAD(8, pd + 128); + CACTUS_STORE_QUAD(12, pd + 192); + if (g + 4 < num_groups) svzero_za(); // in-order queue: runs after the stores drain + } + for (uint32_t t = 0; g < num_groups; ++g, ++t) { // <=3 tail groups; slices zeroed above + const int8_t* a = act_i8 + (size_t)g * gs; + const uint8_t* w = wsb + (size_t)g * g_stride; + const uint32_t sl = t * 4u; + for (uint32_t kk = 0; kk < gs; kk += 16) { + svint8_t av = svld1rq_s8(pg, a + kk); + const uint8_t* p = w + (size_t)(kk >> 2) * 128; + CACTUS_LUTI4_DOT(sl, p, 0, av, 0); + CACTUS_LUTI4_DOT(sl, p, 128, av, 1); + CACTUS_LUTI4_DOT(sl, p, 256, av, 2); + CACTUS_LUTI4_DOT(sl, p, 384, av, 3); + } + CACTUS_STORE_QUAD(sl, psb + (size_t)g * 64); + } + } + #undef CACTUS_LUTI4_DOT + #undef CACTUS_STORE_QUAD +} + +// GEMM leaf: per (group, kg) one 64B act load, two luti4_x2, four SMOPA (one per ZA tile); +// act_packed [num_groups][gs/4][16][4], partials [num_groups][16][64], NEON rescales after. +__arm_locally_streaming __arm_new("za") __arm_new("zt0") +void cactus_sme_cq_gemm_luti4_s8(const int8_t* act_packed, const uint8_t* esme_sb, + const uint8_t* zt_table, int32_t* partials, + uint32_t num_groups, uint32_t gs) { + svldr_zt(0, zt_table); + const svbool_t pg = svptrue_b8(); + const svbool_t pg32 = svptrue_b32(); + const uint32_t nkg = gs >> 2; + for (uint32_t g = 0; g < num_groups; ++g) { + const int8_t* ag = act_packed + (size_t)g * gs * 16; + const uint8_t* wg = esme_sb + (size_t)g * nkg * 128; + svzero_za(); + for (uint32_t kg = 0; kg < nkg; ++kg) { + svint8_t zn = svld1_s8(pg, ag + (size_t)kg * 64); + const uint8_t* p = wg + (size_t)kg * 128; + svint8x2_t e0 = svluti4_lane_zt_s8_x2(0, svld1_u8(pg, p), 0); + svint8x2_t e1 = svluti4_lane_zt_s8_x2(0, svld1_u8(pg, p + 64), 0); + svmopa_za32_s8_m(0, pg, pg, zn, svget2_s8(e0, 0)); + svmopa_za32_s8_m(1, pg, pg, zn, svget2_s8(e0, 1)); + svmopa_za32_s8_m(2, pg, pg, zn, svget2_s8(e1, 0)); + svmopa_za32_s8_m(3, pg, pg, zn, svget2_s8(e1, 1)); + } + int32_t* pgrp = partials + (size_t)g * 16 * 64; + for (uint32_t r = 0; r < 16; ++r) { + int32_t* dst = pgrp + (size_t)r * 64; + svst1_hor_za32(0, r, pg32, dst); + svst1_hor_za32(1, r, pg32, dst + 16); + svst1_hor_za32(2, r, pg32, dst + 32); + svst1_hor_za32(3, r, pg32, dst + 48); + } + } +} + +#else // ---- compile-time fallback ---- +void cactus_sme_cq_gemv_luti4_s8(const int8_t*, const uint8_t*, const uint8_t*, int32_t* partials, + uint32_t num_groups, uint32_t, uint32_t, uint32_t sb_count) { + for (uint32_t i = 0; i < sb_count * num_groups * 64u; ++i) partials[i] = 0; +} +void cactus_sme_cq_gemm_luti4_s8(const int8_t*, const uint8_t*, const uint8_t*, int32_t* partials, + uint32_t num_groups, uint32_t) { + for (uint32_t i = 0; i < num_groups * 16u * 64u; ++i) partials[i] = 0; +} +#endif diff --git a/cactus-kernels/src/threading.h b/cactus-kernels/src/threading.h index e5795329f..9224fb2d5 100644 --- a/cactus-kernels/src/threading.h +++ b/cactus-kernels/src/threading.h @@ -74,6 +74,36 @@ inline bool cpu_has_i8mm() { #endif } +// Streaming svcntb; stub returns 0 without SME2. Only call when the SME2 feature bit is set. +extern "C" uint32_t cactus_sme2_svl_bytes(void); + +inline bool cpu_has_sme2() { +#if defined(__aarch64__) + static std::once_flag once; + static bool has = false; + std::call_once(once, []() { +#if defined(__APPLE__) + int ret = 0; + size_t size = sizeof(ret); + if (sysctlbyname("hw.optional.arm.FEAT_SME2", &ret, &size, nullptr, 0) == 0) { + has = (ret == 1); + } +#elif defined(__ANDROID__) + unsigned long hwcap2 = getauxval(AT_HWCAP2); + #ifndef HWCAP2_SME2 + #define HWCAP2_SME2 (1UL << 37) + #endif + has = (hwcap2 & HWCAP2_SME2) != 0; +#endif + // Every shipped SME layout is pinned to SVL 64; other SVLs fall back to NEON for correctness. + if (has) has = (cactus_sme2_svl_bytes() == 64); + }); + return has; +#else + return false; +#endif +} + inline float32x4_t fast_exp_f32x4(float32x4_t x) { const float32x4_t log2e = vdupq_n_f32(1.44269504088896341f); const float32x4_t ln2_hi = vdupq_n_f32(6.93145751953125e-1f); diff --git a/cactus-kernels/tests/test_attention.cpp b/cactus-kernels/tests/test_attention.cpp index e6089ba6a..c4f4fc460 100644 --- a/cactus-kernels/tests/test_attention.cpp +++ b/cactus-kernels/tests/test_attention.cpp @@ -1,6 +1,8 @@ #include "test_utils.h" #include #include +#include +#include using namespace TestUtils; @@ -76,6 +78,150 @@ bool test_attention_f16() { return true; } +// Double-precision reference for cactus_attention_hybrid_int8_fp16, replicating the incumbent's +// mask semantics exactly (causal kv_end clamp, sliding window, sink exemption on rolled caches). +static void ref_hybrid_attention( + const __fp16* Q, const int8_t* Kc, const int8_t* Vc, + const float* ks, const float* vs, const __fp16* Kn, const __fp16* Vn, + double* out, size_t B, size_t seq, size_t cache_len, size_t new_len, + size_t qh, size_t kvh, size_t hd, size_t pos_off, + bool causal, size_t window, size_t qg) +{ + const size_t kv_seq = cache_len + new_len, gqa = qh / kvh, ngk = hd / qg; + const size_t cache_abs_offset = pos_off >= cache_len ? pos_off - cache_len : 0; + const double scale = 1.0 / static_cast(sqrtf(static_cast(hd))); + std::vector sc(kv_seq); + for (size_t b = 0; b < B; b++) + for (size_t h = 0; h < qh; h++) + for (size_t q = 0; q < seq; q++) { + const __fp16* qv = Q + ((b * seq + q) * qh + h) * hd; + double* ov = out + ((b * seq + q) * qh + h) * hd; + const size_t kv_head = h / gqa; + const size_t abs_q = pos_off + q; + const size_t kv_start = (window > 0 && abs_q > window) ? abs_q - window : 0; + const size_t kv_end = causal ? std::min(kv_seq, cache_len + q + 1) : kv_seq; + std::fill(sc.begin(), sc.end(), -1e300); + for (size_t kv = 0; kv < kv_end; kv++) { + bool wm = false; + if (window > 0 && kv_start > 0) { + if (kv < cache_len) { + if (cache_abs_offset == 0 || kv >= 4) wm = (cache_abs_offset + kv < kv_start); + } else { + wm = (kv + cache_abs_offset < kv_start); + } + } + if ((causal && kv > abs_q) || wm) continue; + double s = 0.0; + if (kv < cache_len) { + const int8_t* kr = Kc + ((b * cache_len + kv) * kvh + kv_head) * hd; + const float* kss = ks + (kv * kvh + kv_head) * ngk; // scales: no batch stride (matches kernel) + for (size_t g = 0; g < ngk; g++) + for (size_t d = g * qg; d < (g + 1) * qg; d++) + s += static_cast(static_cast(qv[d])) * kr[d] * kss[g]; + } else { + const __fp16* kr = Kn + ((b * new_len + (kv - cache_len)) * kvh + kv_head) * hd; + for (size_t d = 0; d < hd; d++) + s += static_cast(static_cast(qv[d])) * static_cast(kr[d]); + } + sc[kv] = s * scale; + } + double m = -1e300; + for (size_t kv = 0; kv < kv_end; kv++) m = std::max(m, sc[kv]); + for (size_t d = 0; d < hd; d++) ov[d] = 0.0; + if (m <= -1e300) continue; + double den = 0.0; + for (size_t kv = 0; kv < kv_end; kv++) { + if (sc[kv] <= -1e300) { sc[kv] = 0.0; continue; } + sc[kv] = std::exp(sc[kv] - m); + den += sc[kv]; + } + if (den <= 0.0) continue; + for (size_t kv = 0; kv < kv_end; kv++) { + if (sc[kv] == 0.0) continue; + const double p = sc[kv]; + if (kv < cache_len) { + const int8_t* vr = Vc + ((b * cache_len + kv) * kvh + kv_head) * hd; + const float* vss = vs + (kv * kvh + kv_head) * ngk; + for (size_t g = 0; g < ngk; g++) + for (size_t d = g * qg; d < (g + 1) * qg; d++) + ov[d] += p * vr[d] * vss[g]; + } else { + const __fp16* vr = Vn + ((b * new_len + (kv - cache_len)) * kvh + kv_head) * hd; + for (size_t d = 0; d < hd; d++) + ov[d] += p * static_cast(static_cast(vr[d])); + } + } + for (size_t d = 0; d < hd; d++) ov[d] /= den; + } +} + +// SME2 prefill attention vs the incumbent NEON flash kernel: global and sliding-window cases, +// rolled caches, partial tiles/blocks, head_dim 128, batch 2. +bool test_attention_hybrid_sme() { + struct Case { size_t B, cache_len, seq, window, pos_off, hd, qh, kvh; }; + const Case cases[] = { + {1, 640, 128, 0, 640, 256, 8, 4}, + {1, 603, 128, 512, 768, 256, 8, 4}, + {1, 600, 128, 512, 600, 256, 8, 4}, + {2, 137, 40, 0, 137, 128, 4, 2}, + }; + std::mt19937 gen(42); + std::uniform_int_distribution i8d(-127, 127); + std::uniform_real_distribution scd(0.005f, 0.02f); + bool all_ok = true; + for (const auto& cs : cases) { + const size_t qg = 32, ngk = cs.hd / qg, new_len = cs.seq; + std::vector<__fp16> Q(cs.B * cs.seq * cs.qh * cs.hd); + std::vector<__fp16> Kn(cs.B * new_len * cs.kvh * cs.hd), Vn(cs.B * new_len * cs.kvh * cs.hd); + std::vector Kc(cs.B * cs.cache_len * cs.kvh * cs.hd), Vc(cs.B * cs.cache_len * cs.kvh * cs.hd); + std::vector ks(cs.cache_len * cs.kvh * ngk), vs(cs.cache_len * cs.kvh * ngk); + fill_random_fp16(Q, -1.0f, 1.0f); + fill_random_fp16(Kn, -1.0f, 1.0f); + fill_random_fp16(Vn, -1.0f, 1.0f); + for (auto& x : Kc) x = static_cast(i8d(gen)); + for (auto& x : Vc) x = static_cast(i8d(gen)); + for (auto& x : ks) x = scd(gen); + for (auto& x : vs) x = scd(gen); + + const size_t osz = cs.B * cs.seq * cs.qh * cs.hd; + std::vector<__fp16> out_neon(osz), out_sme(osz); + cactus_quant_set_backend(1); + cactus_attention_hybrid_int8_fp16(Q.data(), Kc.data(), Vc.data(), ks.data(), vs.data(), + Kn.data(), Vn.data(), out_neon.data(), cs.B, cs.seq, cs.cache_len, new_len, + cs.qh, cs.kvh, cs.hd, 0.0f, cs.pos_off, true, cs.window, qg, cs.hd); + cactus_quant_set_backend(0); + cactus_attention_hybrid_int8_fp16(Q.data(), Kc.data(), Vc.data(), ks.data(), vs.data(), + Kn.data(), Vn.data(), out_sme.data(), cs.B, cs.seq, cs.cache_len, new_len, + cs.qh, cs.kvh, cs.hd, 0.0f, cs.pos_off, true, cs.window, qg, cs.hd); + + std::vector ref(osz); + ref_hybrid_attention(Q.data(), Kc.data(), Vc.data(), ks.data(), vs.data(), + Kn.data(), Vn.data(), ref.data(), cs.B, cs.seq, cs.cache_len, new_len, + cs.qh, cs.kvh, cs.hd, cs.pos_off, true, cs.window, qg); + + double ref_amax = 0.0, e_neon = 0.0, e_sme = 0.0, d_ns = 0.0; + for (size_t i = 0; i < osz; i++) { + ref_amax = std::max(ref_amax, std::fabs(ref[i])); + e_neon = std::max(e_neon, std::fabs(static_cast(out_neon[i]) - ref[i])); + e_sme = std::max(e_sme, std::fabs(static_cast(out_sme[i]) - ref[i])); + d_ns = std::max(d_ns, static_cast(std::fabs( + static_cast(out_sme[i]) - static_cast(out_neon[i])))); + } + const double rel_neon = e_neon / std::max(ref_amax, 1e-12); + const double rel_sme = e_sme / std::max(ref_amax, 1e-12); + const bool engaged_expected = cactus_quant_sme_available() != 0; + bool ok = rel_sme < 0.05 && rel_sme < 5.0 * std::max(rel_neon, 0.002); + if (engaged_expected && d_ns == 0.0) ok = false; // SME path silently not taken + std::cout << " hybrid_sme cache=" << cs.cache_len << " seq=" << cs.seq + << " win=" << cs.window << " off=" << cs.pos_off << " hd=" << cs.hd + << " B=" << cs.B << ": rel_neon=" << rel_neon << " rel_sme=" << rel_sme + << " maxdiff=" << d_ns << (ok ? "" : " <-- FAIL") << "\n"; + all_ok = all_ok && ok; + } + cactus_quant_set_backend(0); + return all_ok; +} + bool run_benchmarks() { auto bench = [](const char* label, auto fn) { fn(); @@ -121,6 +267,36 @@ bool run_benchmarks() { nullptr, 0, 0, true, false, false, 0, 0.0f); }); } + // Hybrid int8 prefill attention at Gemma 4 E2B chunked-prefill shapes: NEON vs SME backend + for (size_t window : {size_t(0), size_t(512)}) { + const size_t B = 1, cache_len = 896, seq = 128, qh = 8, kvh = 4, hd = 256, qg = 32; + const size_t ngk = hd / qg, pos_off = window ? cache_len + 128 : cache_len; + std::mt19937 g2(7); + std::uniform_int_distribution i8d(-127, 127); + std::uniform_real_distribution scd(0.005f, 0.02f); + std::vector<__fp16> Q(B * seq * qh * hd), Kn(B * seq * kvh * hd), Vn(B * seq * kvh * hd); + std::vector Kc(B * cache_len * kvh * hd), Vc(B * cache_len * kvh * hd); + std::vector ks(cache_len * kvh * ngk), vs(cache_len * kvh * ngk); + fill_random_fp16(Q, -1.0f, 1.0f); fill_random_fp16(Kn, -1.0f, 1.0f); fill_random_fp16(Vn, -1.0f, 1.0f); + for (auto& x : Kc) x = static_cast(i8d(g2)); + for (auto& x : Vc) x = static_cast(i8d(g2)); + for (auto& x : ks) x = scd(g2); + for (auto& x : vs) x = scd(g2); + std::vector<__fp16> out(B * seq * qh * hd); + auto run = [&]{ + cactus_attention_hybrid_int8_fp16(Q.data(), Kc.data(), Vc.data(), ks.data(), vs.data(), + Kn.data(), Vn.data(), out.data(), B, seq, cache_len, seq, + qh, kvh, hd, 0.0f, pos_off, true, window, qg, hd); + }; + char label[96]; + cactus_quant_set_backend(1); + snprintf(label, sizeof label, "hybrid prefill c=896 win=%zu NEON", window); + bench(label, run); + cactus_quant_set_backend(0); + snprintf(label, sizeof label, "hybrid prefill c=896 win=%zu SME", window); + bench(label, run); + cactus_quant_set_backend(0); + } return true; } @@ -131,6 +307,7 @@ int main() { runner.run_test("softmax", test_softmax()); runner.run_test("rope", test_rope()); runner.run_test("attention_f16", test_attention_f16()); + runner.run_test("attention_hybrid_sme", test_attention_hybrid_sme()); runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); runner.print_summary(); diff --git a/cactus-kernels/tests/test_matmul.cpp b/cactus-kernels/tests/test_matmul.cpp index 34f66be3e..5e9cb7be9 100644 --- a/cactus-kernels/tests/test_matmul.cpp +++ b/cactus-kernels/tests/test_matmul.cpp @@ -3,6 +3,7 @@ #include #include #include +#include using namespace TestUtils; @@ -359,11 +360,13 @@ bool test_cq_correctness(uint32_t bits) { return true; } -// ── Panel-format GEMV validation ───────────────────────────────────────────────────────────── -// With packed panels present (built here by the independent test packer), the real -// cactus_quant_matmul dispatch takes the panel GEMV kernel; validate it against the SAME FP32 -// reference oracle the legacy kernels are gated on. -static bool test_cq_panel(uint32_t bits, double& mse_out) { +// Panel GEMV through the real dispatch vs the FP32 oracle; backend 1 = NEON, 2 = SME2 leaves. +struct BackendGuard { + explicit BackendGuard(int b) { cactus_quant_set_backend(b); } + ~BackendGuard() { cactus_quant_set_backend(0); } +}; + +static bool test_cq_panel(uint32_t bits, int backend, double& mse_out) { const uint32_t K = 1024, N = 64, gs = 128; SyntheticCQ cq(bits, K, N, gs, 123); cq.preexpand(); @@ -379,15 +382,17 @@ static bool test_cq_panel(uint32_t bits, double& mse_out) { std::vector<__fp16> x_f16(K), y_f16(N, static_cast<__fp16>(0)); for (size_t i = 0; i < K; i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); - cactus_quant_matmul(&mat, x_f16.data(), 1, y_f16.data()); + { + BackendGuard bg(backend); + cactus_quant_matmul(&mat, x_f16.data(), 1, y_f16.data()); + } mse_out = compute_mse(ref.data(), y_f16.data(), N); return mse_out <= 0.1; } -// GEMM (M>1) variant: validates the panel NEON GEMM (M and N tails included) vs a per-row FP32 -// reference. -static bool test_cq_panel_gemm(uint32_t bits, uint32_t M, uint32_t N, double& mse_out) { +// Panel GEMM (M and N tails included) vs a per-row FP32 reference. +static bool test_cq_panel_gemm(uint32_t bits, uint32_t M, uint32_t N, int backend, double& mse_out) { const uint32_t K = 512, gs = 128; SyntheticCQ cq(bits, K, N, gs, 321); cq.preexpand(); @@ -405,19 +410,51 @@ static bool test_cq_panel_gemm(uint32_t bits, uint32_t M, uint32_t N, double& ms std::vector<__fp16> Xf(static_cast(M) * K), Yf(static_cast(M) * N, static_cast<__fp16>(0)); for (size_t i = 0; i < static_cast(M) * K; i++) Xf[i] = static_cast<__fp16>(X[i]); - cactus_quant_matmul(&mat, Xf.data(), M, Yf.data()); + { + BackendGuard bg(backend); + cactus_quant_matmul(&mat, Xf.data(), M, Yf.data()); + } mse_out = compute_mse(ref.data(), Yf.data(), static_cast(M) * N); return mse_out <= 0.1; } -// ── Orthogonal-rotation CQ4 (the Gemma lm_head format) ────────────────────────────────────── -// Oracle: y[n] = norm[n] * sum_i cb[idx(n,i)] * (a_scaled @ R)[i], fp32. Gates the incumbent and -// the panel virtual-group driver against the same reference, plus head-to-head agreement. -static bool test_orth_panel(double& mse_inc, double& mse_panel) { - // PRODUCTION lm_head format: orthogonal-rotation CQ4 in the INTERLEAVED_4ROW packed layout - // (gs == K, ng == 1) for legacy bundles; panel files store it with virtual 128-wide groups. - // Row-major orthogonal bundles were a transpiler bug and take the per-row fallback. +// NEON and SME2 accumulate identical int32 partials, so outputs differ only by fp16 rounding. +static bool test_panel_neon_sme_agreement() { + const uint32_t K = 1024, gs = 128; + for (uint32_t M : {1u, 20u}) { + const uint32_t N = (M == 1) ? 192 : 72; + SyntheticCQ cq(4, K, N, gs, 909); + cq.preexpand(); + CactusQuantMatrix mat = cq.matrix(); + std::vector<__fp16> X(static_cast(M) * K); + fill_random_fp16(X, -1.f, 1.f); + std::vector<__fp16> y_neon(static_cast(M) * N), y_sme(static_cast(M) * N); + { + BackendGuard bg(1); + cactus_quant_matmul(&mat, X.data(), M, y_neon.data()); + } + { + BackendGuard bg(2); + cactus_quant_matmul(&mat, X.data(), M, y_sme.data()); + } + double mx = 0; + for (size_t i = 0; i < y_neon.size(); ++i) { + const double d = std::abs(static_cast(y_neon[i]) - static_cast(y_sme[i])) / + std::max(1.0, std::abs(static_cast(y_neon[i]))); + mx = std::max(mx, d); + } + if (mx > 2e-3) { + std::cerr << " neon/sme agreement M=" << M << " rel_err=" << mx << "\n"; + return false; + } + } + return true; +} + +// Orthogonal CQ4 (lm_head) oracle: y[n] = norm[n] * sum_i cb[idx(n,i)] * (a_scaled @ R)[i]. +static bool test_orth_panel(int backend, double& mse_inc, double& mse_panel) { + // Production lm_head: orthogonal CQ4, INTERLEAVED_4ROW, gs == K, ng == 1. const uint32_t K = 1536, N = 1024, VGS = 128, VNG = K / VGS; std::mt19937 gen(2024); std::uniform_real_distribution dist(-1.f, 1.f); @@ -491,17 +528,17 @@ static bool test_orth_panel(double& mse_inc, double& mse_panel) { W2.packed_panels = panels.data(); W2.norm_panels = npanels.data(); W2.rotation_t = rot_t.data(); std::vector<__fp16> y_panel(N); - cactus_quant_orthogonal_matmul(&W2, x.data(), 1, y_panel.data()); + { + BackendGuard bg(backend); + cactus_quant_orthogonal_matmul(&W2, x.data(), 1, y_panel.data()); + } mse_panel = compute_mse(ref.data(), y_panel.data(), N); return mse_inc <= 0.1 && mse_panel <= 0.1 && mse_panel <= mse_inc * 4 + 1e-4; } -// ── Interleaved-4row (legacy PRODUCTION format) tests ──────────────────────────────────────── -// CQ4 interleaved vs the FP32 oracle, through the real dispatch. use_panels=false exercises the -// legacy interleaved NEON kernel (old bundles); use_panels=true builds the panel layout from the -// SAME interleaved fixture through the reference encoder and exercises the panel GEMV -// (multi-super-block stealing included: N=192 = 3 super-blocks). -static bool test_cq4_interleaved(bool use_panels, double& mse_out, +// Legacy IL CQ4 vs the FP32 oracle through real dispatch; use_panels rebuilds the same +// fixture as panels through the reference encoder. +static bool test_cq4_interleaved(bool use_panels, int backend, double& mse_out, uint32_t K = 1024, uint32_t N = 192, uint32_t gs = 128) { SyntheticCQ cq(4, K, N, gs, 777); if (use_panels) cq.preexpand_il(); @@ -516,7 +553,10 @@ static bool test_cq4_interleaved(bool use_panels, double& mse_out, std::vector<__fp16> x_f16(K), y(N, (__fp16)0); for (size_t i = 0; i < K; i++) x_f16[i] = (__fp16)x_f32[i]; - cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + { + BackendGuard bg(backend); + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + } mse_out = compute_mse(ref.data(), y.data(), N); return mse_out <= 0.1; } @@ -626,27 +666,29 @@ static void print_panel_comparison() { CactusQuantMatrix mat_panel = cq.matrix_interleaved(); // panels set: panel GEMV std::vector<__fp16> x(sh.K), y(sh.N); fill_random_fp16(x, -1.f, 1.f); - auto run_ms = [&](CactusQuantMatrix* m) { + auto run_ms = [&](CactusQuantMatrix* m, int backend) { + BackendGuard bg(backend); cactus_quant_matmul(m, x.data(), 1, y.data()); // warmup Timer t; for (int i = 0; i < sh.iters; i++) cactus_quant_matmul(m, x.data(), 1, y.data()); return t.elapsed_ms() / sh.iters; }; - double ms_f = 1e30, ms_p = 1e30; + const bool sme = cactus_quant_sme_available() != 0; + double ms_f = 1e30, ms_p = 1e30, ms_s = 1e30; for (int round = 0; round < 5; round++) { - ms_f = std::min(ms_f, run_ms(&mat_file)); - ms_p = std::min(ms_p, run_ms(&mat_panel)); + ms_f = std::min(ms_f, run_ms(&mat_file, 1)); + ms_p = std::min(ms_p, run_ms(&mat_panel, 1)); + if (sme) ms_s = std::min(ms_s, run_ms(&mat_panel, 0)); } const double fl = 2.0 * sh.K * sh.N; double gf = fl / (ms_f * 1e6), gp = fl / (ms_p * 1e6); std::cout << " " << std::left << std::setw(31) << sh.name << " file-NEON " << std::fixed << std::setprecision(1) << gf - << " | panel " << gp << " GF" - << " | " << std::setprecision(2) << (gp / gf) + << " | panel " << gp << " GF"; + if (sme) std::cout << " | panel-auto(sme) " << fl / (ms_s * 1e6) << " GF"; + std::cout << " | " << std::setprecision(2) << (gp / gf) << "x (best of 5x" << sh.iters << ")\n"; } - // GEMM (M>1): panel GEMM vs the legacy expand-from-packed NEON GEMM (the real-model M>1 - // path for old bundles). for (uint32_t M : {4u, 16u, 32u, 64u, 128u, 256u}) { const uint32_t K = 1024, N = 1024, gs = 128; SyntheticCQ cq(4, K, N, gs); @@ -655,24 +697,28 @@ static void print_panel_comparison() { CactusQuantMatrix mat_panel = cq.matrix(); // panels set: panel NEON GEMM std::vector<__fp16> A(static_cast(M) * K), Cc(static_cast(M) * N); fill_random_fp16(A, -1.f, 1.f); - auto bench = [&](CactusQuantMatrix* m) -> double { + auto bench = [&](CactusQuantMatrix* m, int backend) -> double { + BackendGuard bg(backend); cactus_quant_matmul(m, A.data(), M, Cc.data()); // warmup const int iters = 8; Timer t; for (int i = 0; i < iters; i++) cactus_quant_matmul(m, A.data(), M, Cc.data()); return t.elapsed_ms() / iters; }; - double ml = 1e30, mp = 1e30; + const bool sme = cactus_quant_sme_available() != 0; + double ml = 1e30, mp = 1e30, ms = 1e30; for (int round = 0; round < 3; round++) { - ml = std::min(ml, bench(&mat_legacy)); - mp = std::min(mp, bench(&mat_panel)); + ml = std::min(ml, bench(&mat_legacy, 1)); + mp = std::min(mp, bench(&mat_panel, 1)); + if (sme) ms = std::min(ms, bench(&mat_panel, 2)); } double gl = (2.0 * M * K * N) / (ml * 1e6), gp = (2.0 * M * K * N) / (mp * 1e6); std::string label = "cq4 M" + std::to_string(M) + " 1024x1024"; std::cout << " " << std::left << std::setw(20) << label << " legacy " << std::fixed << std::setprecision(1) << gl << " GFLOPS" - << " | panel " << gp << " GFLOPS" - << " | " << std::setprecision(2) << (gp / gl) << "x (best of 3x8)\n"; + << " | panel " << gp << " GFLOPS"; + if (sme) std::cout << " | panel-sme " << (2.0 * M * K * N) / (ms * 1e6) << " GFLOPS"; + std::cout << " | " << std::setprecision(2) << (gp / gl) << "x (best of 3x8)\n"; } } @@ -923,45 +969,57 @@ int main() { runner.run_test("matmul_cq3", test_cq_correctness(3)); runner.run_test("matmul_cq4", test_cq_correctness(4)); - // ── Panel-format registry: validate the panel kernels through the real dispatch ── - for (uint32_t bits : {1u, 2u, 3u, 4u}) { - double mse_p = 0.0; - bool ok = test_cq_panel(bits, mse_p); - runner.run_test(std::string("matmul_cq") + std::to_string(bits) + "[panel]", ok); - if (!ok) std::cerr << " cq" << bits << "[panel] MSE=" << mse_p << "\n"; - } - // GEMM (M>1) variant: M tail (20 -> 16+4) and N tail (72 -> 4 super-blocks + 8 valid) for CQ4. + // Panel tests: backend 1 = NEON panel kernels, backend 2 = forced SME2 leaves. struct GemmCase { uint32_t M, N; }; - for (GemmCase gc : {GemmCase{5, 64}, GemmCase{20, 64}, GemmCase{20, 72}, GemmCase{64, 256}}) { - double mse_g = 0.0; - bool ok = test_cq_panel_gemm(4, gc.M, gc.N, mse_g); - runner.run_test(std::string("matmul_cq4_M") + std::to_string(gc.M) + - "_N" + std::to_string(gc.N) + "[panel]", ok); - if (!ok) std::cerr << " cq4 M=" << gc.M << " N=" << gc.N << "[panel] MSE=" << mse_g << "\n"; - } - { - double mi = 0, mp = 0; - bool orth_ok = test_orth_panel(mi, mp); - runner.run_test("matmul_cq4_orth[panel]", orth_ok); - if (!orth_ok) std::cerr << " orth mse_incumbent=" << mi << " mse_panel=" << mp << "\n"; + const GemmCase gemm_cases[] = {GemmCase{5, 64}, GemmCase{20, 64}, GemmCase{20, 72}, GemmCase{64, 256}}; + std::vector> backends{{1, "[panel]"}}; + if (cactus_quant_sme_available()) backends.push_back({2, "[sme2]"}); + for (auto [backend, tag] : backends) { + for (uint32_t bits : {1u, 2u, 3u, 4u}) { + double mse_p = 0.0; + bool ok = test_cq_panel(bits, backend, mse_p); + runner.run_test(std::string("matmul_cq") + std::to_string(bits) + tag, ok); + if (!ok) std::cerr << " cq" << bits << tag << " MSE=" << mse_p << "\n"; + } + // GEMM (M>1): M tail (20 -> 16+4) and N tail (72 -> 4 super-blocks + 8 valid) for CQ4. + for (GemmCase gc : gemm_cases) { + double mse_g = 0.0; + bool ok = test_cq_panel_gemm(4, gc.M, gc.N, backend, mse_g); + runner.run_test(std::string("matmul_cq4_M") + std::to_string(gc.M) + + "_N" + std::to_string(gc.N) + tag, ok); + if (!ok) std::cerr << " cq4 M=" << gc.M << " N=" << gc.N << tag << " MSE=" << mse_g << "\n"; + } + { + double mi = 0, mp = 0; + bool orth_ok = test_orth_panel(backend, mi, mp); + runner.run_test(std::string("matmul_cq4_orth") + tag, orth_ok); + if (!orth_ok) std::cerr << " orth" << tag << " mse_incumbent=" << mi << " mse_panel=" << mp << "\n"; + } + { + double m2 = 0; + runner.run_test(std::string("matmul_cq4_il") + tag, test_cq4_interleaved(true, backend, m2)); + } + // The panel GEMM is bits-agnostic (panel nibbles are codebook indices) — validate CQ1-3 too. + for (uint32_t b : {1u, 2u, 3u}) { + double mse_g = 0.0; + bool ok = test_cq_panel_gemm(b, 20, 64, backend, mse_g); + runner.run_test(std::string("matmul_cq") + std::to_string(b) + "_gemm" + tag, ok); + if (!ok) std::cerr << " cq" << b << " gemm" << tag << " MSE=" << mse_g << "\n"; + } } { - double m1 = 0, m2 = 0; - runner.run_test("matmul_cq4_il[file]", test_cq4_interleaved(false, m1)); - // N=4164: 1041 IL blocks -> 66 chunks (16-block + 1-block tail) -> multi-thread fused - // driver (phase-A stealing, spin barrier, 4-chunk grabs); N=192 stays on the serial path. + double m1 = 0; + runner.run_test("matmul_cq4_il[file]", test_cq4_interleaved(false, 1, m1)); + // N=4164 -> 66 chunks incl. a 1-block tail: exercises the multi-thread fused driver. double m_mt = 0; - runner.run_test("matmul_cq4_il_mt[file]", test_cq4_interleaved(false, m_mt, 1024, 4164, 128)); - runner.run_test("matmul_cq4_il[panel]", test_cq4_interleaved(true, m2)); + runner.run_test("matmul_cq4_il_mt[file]", test_cq4_interleaved(false, 1, m_mt, 1024, 4164, 128)); runner.run_test("panel_layout_invariance", test_panel_layout_invariance()); runner.run_test("orth_embed_rows_batched", test_orth_embed_rows()); } - // The panel GEMM is bits-agnostic (panel nibbles are codebook indices) — validate CQ1-3 too. - for (uint32_t b : {1u, 2u, 3u}) { - double mse_g = 0.0; - bool ok = test_cq_panel_gemm(b, 20, 64, mse_g); - runner.run_test(std::string("matmul_cq") + std::to_string(b) + "_gemm[panel]", ok); - if (!ok) std::cerr << " cq" << b << " gemm[panel] MSE=" << mse_g << "\n"; + if (cactus_quant_sme_available()) { + runner.run_test("panel_neon_sme_agreement", test_panel_neon_sme_agreement()); + } else { + std::cout << " (SME2 unavailable on this CPU — [sme2] variants skipped)\n"; } runner.print_benchmarks_header(); diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 1883bea8f..20731de3c 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -525,8 +525,7 @@ def convert(args: argparse.Namespace) -> None: cq = _scale_cq_norms(cq, adapter.scale_factor(out_path.name)) if getattr(emit_policy, "layout", "row_major") == "interleaved_4row": cq = replace(cq, interleaved_4row=True) - # Panel-format CQ4 weights carry their own filename suffix; the manifest - # records the actual written file so bindings resolve to it. + # The manifest records the actual written file so bindings resolve to it. written_path = write_cq_tensor(out_path, cq) gptq_used = cq.gptq_used status = "converted" if match.recognized else "unrecognized" diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py index c599fb827..6069fd869 100644 --- a/python/cactus/convert/quantization/cq.py +++ b/python/cactus/convert/quantization/cq.py @@ -26,9 +26,7 @@ FLAG_INTERLEAVED_4ROW = 1 << 2 FLAG_PACKED_PANELS = 1 << 5 -# Filename suffix for packed-panel CQ4 weight files. Panel bundles ship ONLY this file (no -# legacy-layout twin), so pre-panel engines fail loudly on new bundles instead of silently -# misreading the data region. +# Panel bundles ship only this file, so pre-panel engines fail loudly on new bundles. PANEL_WEIGHTS_SUFFIX = ".cq4p.weights" @@ -216,8 +214,7 @@ def quantize_codebook_i8(codebook_f16: np.ndarray) -> tuple[np.ndarray, np.float if scale < np.float32(1e-10): scale = np.float32(1e-10) inv = np.float32(1.0) / scale - # float64 holds every float32 exactly, so floor(|x| + 0.5) in float64 is exact - # round-half-away-from-zero of the float32 products (std::round semantics). + # float64 holds float32 exactly, so floor(|x| + 0.5) reproduces std::round half-away-from-zero. scaled = (cb * inv).astype(np.float64) cb_i8 = (np.sign(scaled) * np.floor(np.abs(scaled) + 0.5)).astype(np.int8) return cb_i8, scale @@ -451,9 +448,8 @@ def write_cq_tensor(out_path: Path, cq: CQTensor, *, allow_panels: bool = True) recip = np.minimum(1.0 / np.maximum(input_scale.astype(np.float32), 1e-8), 65504.0).astype(np.float16) if packed_panels: - # Header stores the kernel-facing group view: orthogonal tensors are written with - # virtual 128-wide groups already applied (an exact regrouping — norms are constant - # across subgroups and panel kg-blocks are K-chunk-ordered). + # Orthogonal tensors store virtual 128-wide groups already applied (exact regrouping: + # norms constant across subgroups, kg-blocks K-chunk-ordered). header_gs = 128 if is_orth else group_size header_ng = k // header_gs norms_f16 = norms_f32.astype(np.float16).reshape(n, groups) diff --git a/python/cactus/transpile/weight_compat.py b/python/cactus/transpile/weight_compat.py index 1ea5bc7d7..a3a041910 100644 --- a/python/cactus/transpile/weight_compat.py +++ b/python/cactus/transpile/weight_compat.py @@ -99,9 +99,8 @@ def ensure_embedding_binding_compatible(binding: WeightBinding) -> WeightBinding bits = int(config["bits"]) rotation = str(config["rotation"]) - # CQ4 orthogonal embeddings (token_embeddings == the tied lm_head) materialize in the - # packed-panel format under its own suffix; CQ2 per-layer hadamard embeddings keep the - # legacy row-major layout (the engine's hadamard row decoder reads it). + # CQ4 orthogonal token embeddings (the tied lm_head) materialize as packed panels; CQ2 + # hadamard embeddings keep the legacy row-major layout. use_panels = bits == 4 and rotation == "orthogonal" if use_panels: compat_path = panel_weights_path(source_path) @@ -425,14 +424,8 @@ def _materialize_cq_embedding_cache_panels( bits: int, rotation: str, ) -> None: - """Materialize a tied orthogonal CQ4 embedding/lm_head as a packed-panel file. - - The data region holds kernel-ready panels [SB64][ng][gs/4][128B] (ng = K/128 virtual - groups). The scales region carries the legacy fp16 norms (unfolded, replicated per virtual - group — the embedding row decoder uses them with the fp16 codebook) AND the fp32 folded - panel norms (the lm_head matmul rescale), then the rotation and its transpose. Rows stream - in 64-aligned chunks so the super-block panel order assembles by append. - """ + """Materialize a tied orthogonal CQ4 embedding/lm_head as a packed-panel file; rows stream + in 64-aligned chunks so the super-block panel order assembles by append.""" if rotation != "orthogonal" or bits != 4: raise RuntimeError(f"packed-panel embedding cache requires orthogonal CQ4, got rotation={rotation} bits={bits}")