Skip to content

Packed-panel CQ4 weight format: NEON + SME2 kernels #709

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ncylich wants to merge 5 commits into
base: gemv-dispatch
Choose a base branch
from
Open

Packed-panel CQ4 weight format: NEON + SME2 kernels #709

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
115a257
Add --torch-dtype to convert for bf16-safe model loading
ncylich Jun 10, 2026
df4b2ac
Disable the dense-MLP TQ fusion by default
ncylich Jun 11, 2026
842e644
Add packed-panel CQ4 weight format and NEON kernels
ncylich Jun 10, 2026
15f9b98
Report converted tensor bytes from the written panel path
ncylich Jun 10, 2026
22df501
Add SME2 streaming kernels over the packed-panel weight format
ncylich Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion cactus-graph/cactus_graph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,11 @@ struct BufferDesc {
const int8_t* cq_right_signs = nullptr;
const uint32_t* cq_permutation = nullptr;
const __fp16* cq_rotation = nullptr;
// 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;
uint32_t cq_flags = 0;

CactusQuantMatrix to_cq_matrix() const {
Expand All @@ -277,13 +282,17 @@ struct BufferDesc {
.input_scale = cq_input_scale,
.input_scale_recip = cq_input_scale_recip,
.norms = cq_norms,
.packed_indices = static_cast<const uint8_t*>(get_data()),
// Panel files keep packed_indices null so legacy kernels can never misread them.
.packed_indices = cq_packed_panels ? nullptr : static_cast<const uint8_t*>(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,
};
}

Expand Down Expand Up @@ -781,6 +790,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;
Expand All @@ -802,6 +812,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();
Expand Down
12 changes: 12 additions & 0 deletions cactus-graph/src/core.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}

Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
168 changes: 71 additions & 97 deletions cactus-graph/src/io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -41,6 +42,11 @@ namespace {
}

const std::string stem = filename.substr(0, filename.size() - std::strlen(suffix));
// 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;
}
const std::string cq4 = stem + ".cq4.weights";
if (std::filesystem::exists(cq4)) {
return cq4;
Expand All @@ -52,6 +58,64 @@ namespace {
return filename;
}

// 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<const char*>(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<uint32_t>(mf.group_size());
const uint32_t K = gs * static_cast<uint32_t>(mf.num_groups());
const uint32_t N = static_cast<uint32_t>(mf.shape().size() >= 2 ? mf.shape()[0] : 1);

size_t off = 0;
buffer.cq_codebook = reinterpret_cast<const __fp16*>(scales_base + off);
off += cb_size * sizeof(__fp16);
buffer.cq_input_scale = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_input_scale_recip = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_norms = reinterpret_cast<const __fp16*>(scales_base + off);
off += static_cast<size_t>(N) * mf.num_groups() * sizeof(__fp16);

if (mf.is_packed_panels()) {
const size_t SB64 = (static_cast<size_t>(N) + 63) / 64;
off = align_offset(off, sizeof(float));
buffer.cq_norm_panels = reinterpret_cast<const float*>(scales_base + off);
off += SB64 * mf.num_groups() * 64 * sizeof(float);
if (mf.is_orthogonal_rotation()) {
buffer.cq_rotation = reinterpret_cast<const __fp16*>(scales_base + off);
off += static_cast<size_t>(K) * K * sizeof(__fp16);
buffer.cq_rotation_t = reinterpret_cast<const __fp16*>(scales_base + off);
buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL;
} else {
buffer.cq_left_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_right_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_permutation = reinterpret_cast<const uint32_t*>(scales_base + off);
buffer.cq_flags = 0;
}
buffer.cq_packed_panels = static_cast<const uint8_t*>(mf.data());
return;
}

if (mf.is_orthogonal_rotation()) {
buffer.cq_rotation = reinterpret_cast<const __fp16*>(scales_base + off);
buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL;
} else {
buffer.cq_left_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_right_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_permutation = reinterpret_cast<const uint32_t*>(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<const char*>(&v), sizeof(v));
}
Expand Down Expand Up @@ -388,38 +452,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<const char*>(mapped_file->scales_data());
uint32_t bits = PrecisionTraits::cq_bits(precision);
uint32_t cb_size = 1u << bits;
uint32_t gs = static_cast<uint32_t>(mapped_file->group_size());
uint32_t K = gs * static_cast<uint32_t>(mapped_file->num_groups());
uint32_t N = static_cast<uint32_t>(shape.size() >= 2 ? shape[0] : 1);

size_t off = 0;
buffer.cq_codebook = reinterpret_cast<const __fp16*>(scales_base + off);
off += cb_size * sizeof(__fp16);
buffer.cq_input_scale = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_input_scale_recip = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_norms = reinterpret_cast<const __fp16*>(scales_base + off);
off += static_cast<size_t>(N) * mapped_file->num_groups() * sizeof(__fp16);

if (mapped_file->is_orthogonal_rotation()) {
buffer.cq_rotation = reinterpret_cast<const __fp16*>(scales_base + off);
buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL;
} else {
buffer.cq_left_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_right_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_permutation = reinterpret_cast<const uint32_t*>(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();
Expand Down Expand Up @@ -452,39 +485,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<const char*>(mapped_file->scales_data());
uint32_t bits = PrecisionTraits::cq_bits(precision);
uint32_t cb_size = 1u << bits;
uint32_t gs = static_cast<uint32_t>(mapped_file->group_size());
uint32_t K = gs * static_cast<uint32_t>(mapped_file->num_groups());
uint32_t N = static_cast<uint32_t>(shape.size() >= 2 ? shape[0] : 1);

size_t off = 0;
buffer.cq_codebook = reinterpret_cast<const __fp16*>(scales_base + off);
off += cb_size * sizeof(__fp16);
buffer.cq_input_scale = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_input_scale_recip = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_norms = reinterpret_cast<const __fp16*>(scales_base + off);
off += static_cast<size_t>(N) * mapped_file->num_groups() * sizeof(__fp16);

if (mapped_file->is_orthogonal_rotation()) {
buffer.cq_rotation = reinterpret_cast<const __fp16*>(scales_base + off);
buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL;
} else {
buffer.cq_left_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_right_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_permutation = reinterpret_cast<const uint32_t*>(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();
Expand Down Expand Up @@ -531,43 +532,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<const char*>(mapped_file->scales_data());
uint32_t bits = PrecisionTraits::cq_bits(precision);
uint32_t cb_size = 1u << bits;
uint32_t gs = static_cast<uint32_t>(mapped_file->group_size());
uint32_t K = gs * static_cast<uint32_t>(mapped_file->num_groups());
uint32_t N = static_cast<uint32_t>(shape.size() >= 2 ? shape[0] : 1);

size_t off = 0;
buffer.cq_codebook = reinterpret_cast<const __fp16*>(scales_base + off);
off += cb_size * sizeof(__fp16);
buffer.cq_input_scale = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_input_scale_recip = reinterpret_cast<const __fp16*>(scales_base + off);
off += K * sizeof(__fp16);
buffer.cq_norms = reinterpret_cast<const __fp16*>(scales_base + off);
off += static_cast<size_t>(N) * mapped_file->num_groups() * sizeof(__fp16);

if (mapped_file->is_orthogonal_rotation()) {
buffer.cq_rotation = reinterpret_cast<const __fp16*>(scales_base + off);
buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL;
} else {
buffer.cq_left_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_right_signs = reinterpret_cast<const int8_t*>(scales_base + off);
off += gs;
buffer.cq_permutation = reinterpret_cast<const uint32_t*>(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();
Expand Down Expand Up @@ -963,6 +936,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<const uint32_t*>(ptr + offset);
offset += sizeof(uint32_t);
Expand Down
54 changes: 54 additions & 0 deletions cactus-graph/src/ops_tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,60 @@ void compute_embedding_node(GraphNode& node, const std::vector<std::unique_ptr<G
Precision emb_prec = embeddings_buffer.precision;
if (PrecisionTraits::is_cq(emb_prec) && embeddings_buffer.group_size > 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) {
// Dedup indices, then one parallel dequant+un-rotation pass over the unique rows;
// legacy row-major orthogonal bundles take the per-row fallback.
std::unordered_map<size_t, uint32_t> slot;
slot.reserve(std::min<size_t>(num_indices, 256));
std::vector<uint32_t> uniq;
uniq.reserve(std::min<size_t>(num_indices, 256));
std::vector<uint32_t> idx_slot(num_indices);
for (size_t i = 0; i < num_indices; i++) {
size_t idx = static_cast<size_t>(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<uint32_t>(uniq.size()));
if (it.second) uniq.push_back(static_cast<uint32_t>(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<uint32_t>(hidden_dim),
static_cast<uint32_t>(embeddings_buffer.group_size),
static_cast<uint32_t>(embeddings_buffer.num_groups),
uniq.data(),
static_cast<uint32_t>(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<uint32_t>(hidden_dim),
uniq.data(),
static_cast<uint32_t>(uniq.size()),
embeddings_buffer.data_as<uint8_t>(),
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<size_t>(idx_slot[i]) * hidden_dim,
hidden_dim * sizeof(__fp16));
}
return;
}
std::unordered_map<size_t, size_t> row_cache;
std::vector<__fp16> cached_rows;
if (num_indices > 16) {
Expand Down
Loading
Loading