diff --git a/.gitignore b/.gitignore index 0fe3c15b2..8837be447 100644 --- a/.gitignore +++ b/.gitignore @@ -63,6 +63,8 @@ c/tests/test_* !c/tests/test_*.cu !c/tests/test_*.mm !c/tests/test_*.py +# scratch fixture directories the test binaries write at run time +c/tests/tmp_* # bench_* and fuzz_* build to extensionless binaries in the same directory and # were not covered here, which is how c/tests/bench_omp_grain reached the tree. c/tests/bench_* diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index ca473c784..7cc297ee7 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -32,6 +32,22 @@ executes an unmodified vLLM checkout (commit `ffd46bfab2128bb84146050e98b51a617c6575ab`) as a behavioural oracle for the native port; no vLLM code is vendored. +## Swiftlet qpack and MLX affine Metal GEMV + +`c/qpack.c` and `c/qpack.h` implement the Swiftlet qpack v1 container schema +documented by `Sources/SwiftletCore/Qpack.swift` at the commit linked below. +The reader is a new portable implementation for Colibri and does not copy +Swift source code. + +The MLX affine Q4/Q8 kernels in `c/backend_metal.mm` are adapted from +`gemv_affine_fast` and `gemv_affine_fast8` in Swiftlet: +https://github.com/leonickson1/Swiftlet at commit +`b3a04676748c7597800c5bcc8b80a32508f9f43d`. + +Swiftlet is licensed under Apache License 2.0. The kernels were modified for +Colibri's checked descriptor, batched dispatch, buffer ownership, and fallback +contract. This repository's `LICENSE` contains the applicable Apache 2.0 text. + ## DeepGEMM sm120 headers (fetched, not vendored: `c/third_party/deepgemm/`) The DeepSeek V4 CUDA tier's DeepGEMM flavour (`make cuda-dsv4-dg-dll`, diff --git a/c/Makefile b/c/Makefile index f7dbd2878..4ae3f220e 100644 --- a/c/Makefile +++ b/c/Makefile @@ -52,6 +52,13 @@ ifneq (,$(DARWIN)) # Niente -march: su arm64 NEON e' baseline (i kernel __ARM_NEON si attivano da soli). CC = clang OMPDIR := $(shell brew --prefix libomp 2>/dev/null) +# CLT-only Macs routinely have Homebrew installed but off the non-interactive +# PATH (`make` from scripts, SSH, CI), which turned this probe into a silent +# single-threaded build with libomp sitting right there on disk. Fall back to +# the two standard Homebrew prefixes before concluding it is absent. +ifeq ($(OMPDIR),) +OMPDIR := $(firstword $(wildcard /opt/homebrew/opt/libomp /usr/local/opt/libomp)) +endif # `brew --prefix libomp` can print the formula's prospective path even when it # is not installed, so verify both artifacts before adding unusable flags. ifneq ($(and $(OMPDIR),$(wildcard $(OMPDIR)/include/omp.h),$(wildcard $(OMPDIR)/lib/libomp.*)),) @@ -838,16 +845,16 @@ cuda-dsv4-dg-dll: backend_cuda_dsv4.cu backend_cuda_dsv4.h dsv4.def $(DEEPGEMM_S dsv4-cuda-loader-test: tests/test_dsv4_decode_cuda.c backend_loader_dsv4.c backend_cuda_dsv4.h $(CC) $(CFLAGS) tests/test_dsv4_decode_cuda.c backend_loader_dsv4.c -o dsv4_decode_loader_test$(EXE) ./dsv4_decode_loader_test$(EXE) -backend_metal.o: backend_metal.mm backend_metal.h +backend_metal.o: backend_metal.mm backend_metal.h affine_quant.h $(METALXX) -c backend_metal.mm -o $@ -metal-test: tests/test_backend_metal.mm backend_metal.mm backend_metal.h +metal-test: tests/test_backend_metal.mm backend_metal.mm backend_metal.h affine_quant.h $(METALXX) tests/test_backend_metal.mm backend_metal.mm -framework Metal -framework Foundation -o backend_metal_test ./backend_metal_test # Standalone large-batch GEMM correctness sweep — reproduces the long-context prefill # corruption in seconds (no model). OMP-parallel CPU reference. -gemm-test: tests/test_gemm_largebatch.mm backend_metal.mm backend_metal.h +gemm-test: tests/test_gemm_largebatch.mm backend_metal.mm backend_metal.h affine_quant.h $(METALXX) $(OMPC) tests/test_gemm_largebatch.mm backend_metal.mm $(OMPL) -framework Metal -framework Foundation -o gemm_largebatch_test ./gemm_largebatch_test @@ -1067,6 +1074,19 @@ tests/test_qwen36_dense_batch$(EXE): tests/test_qwen36_dense_batch.c qwen36.c qw tests/test_qwen36_json_escape$(EXE): tests/test_qwen36_json_escape.c qwen36.c qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ) $(CC) $(QWEN36_CFLAGS) $< $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o $@ $(QWEN36_LDFLAGS) +# Dense-bridge gate: qwen36's loader against an MLX affine +# model.safetensors (the dense half of a Swiftlet qpack container) -- +# `language_model.` name resolution, U32 triple indexing, checked Q4/Q8 +# expansion, refusals. With argv it doubles as the round-trip harness for +# tests/test_make_qwen36_qpack_snap.py. +tests/test_qwen36_dense_affine$(EXE): tests/test_qwen36_dense_affine.c qwen36.c affine_quant.h qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ) + $(CC) $(QWEN36_CFLAGS) $< $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o $@ $(QWEN36_LDFLAGS) + +# Both tokenizer.json merge spellings ("a b" strings and ["a","b"] pairs) +# must index the same merge table; the pair form is what Qwen3.6 ships. +tests/test_qwen36_tok_merges$(EXE): tests/test_qwen36_tok_merges.c qwen36.c qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ) + $(CC) $(QWEN36_CFLAGS) $< $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o $@ $(QWEN36_LDFLAGS) + # Reproducible local timing evidence; intentionally not a noisy CI perf gate. tests/bench_qwen36_dense_batch$(EXE): tests/bench_qwen36_dense_batch.c qwen36.c qwen36_tier.h st.h json.h compat.h $(QWEN36_TIER_SRC) $(CUDA_OBJ) $(CC) $(QWEN36_CFLAGS) $< $(QWEN36_TIER_SRC) $(CUDA_OBJ) -o $@ $(QWEN36_LDFLAGS) @@ -1117,6 +1137,9 @@ portable: iobench$(EXE): iobench.c compat.h $(CC) $(CFLAGS) iobench.c -o iobench$(EXE) $(LDFLAGS) +qpack-inspect$(EXE): qpack_inspect.c qpack.c qpack.h affine_quant.h json.h compat.h + $(CC) $(CFLAGS) qpack_inspect.c qpack.c -o qpack-inspect$(EXE) $(LDFLAGS) + tests/test_serve_sentinel$(EXE): tests/test_serve_sentinel.c compat.h serve_codec.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) @@ -1379,6 +1402,19 @@ tests/test_expert_store_ops$(EXE): tests/test_expert_store_ops.c expert_store.h tests/test_native_quant$(EXE): tests/test_native_quant.c deepseek_v4.c native_quant.h tensor.h quant.h $(CC) $(CFLAGS) -DCOLI_V4_UNIT_NATIVE_QUANT deepseek_v4.c $< -o $@ $(LDFLAGS) +# Portable MLX affine Q4/Q8 format contract and CPU oracle. This stays +# separate from QT fmt=4, whose signed-int4 scale-only layout is incompatible. +tests/test_affine_quant$(EXE): tests/test_affine_quant.c affine_quant.h backend_metal.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + +tests/test_qpack$(EXE): tests/test_qpack.c qpack.c qpack.h affine_quant.h json.h compat.h + $(CC) $(CFLAGS) tests/test_qpack.c qpack.c -o $@ $(LDFLAGS) + +# Loader-side expansion oracle for the same MLX affine contract: nibble/byte +# order, scalar decodings, parity against coli_affine_matmul_ref, refusals. +tests/test_affine_dequant$(EXE): tests/test_affine_dequant.c affine_quant.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + tests/test_edge_runtime$(EXE): tests/test_edge_runtime.c edge_runtime.c edge_runtime.h $(CC) $(CFLAGS) tests/test_edge_runtime.c edge_runtime.c -o $@ $(LDFLAGS) diff --git a/c/affine_quant.h b/c/affine_quant.h new file mode 100644 index 000000000..a763fc8ff --- /dev/null +++ b/c/affine_quant.h @@ -0,0 +1,285 @@ +/* affine_quant.h -- portable MLX affine-quantization contract and CPU oracle. + * + * This format is deliberately separate from QT fmt=4. Colibri's fmt=4 stores + * signed int4 values with multiplicative group scales; MLX stores unsigned + * uint32-packed values and reconstructs each group as scale*q + bias. + */ +#ifndef COLI_AFFINE_QUANT_H +#define COLI_AFFINE_QUANT_H + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef enum { + COLI_AFFINE_MLX_Q4 = 0, + COLI_AFFINE_MLX_Q8 = 1 +} ColiAffineFormat; + +typedef enum { + COLI_AFFINE_SCALAR_F32 = 0, + COLI_AFFINE_SCALAR_F16 = 1, + COLI_AFFINE_SCALAR_BF16 = 2 +} ColiAffineScalarFormat; + +typedef enum { + COLI_AFFINE_OK = 0, + COLI_AFFINE_NULL = -1, + COLI_AFFINE_BAD_FORMAT = -2, + COLI_AFFINE_BAD_SHAPE = -3, + COLI_AFFINE_TRUNCATED = -4, + COLI_AFFINE_OVERFLOW = -5 +} ColiAffineStatus; + +/* MLX affine weights are little-endian uint32 words, packed along input_dim. + * Scales and biases have identical scalar types and [output_dim, groups] + * layout, where groups = input_dim/group_size. Byte lengths are mandatory: + * qpack sections and resident tensors must be checked before GPU registration. + */ +typedef struct { + const void *weights; + const void *scales; + const void *biases; + size_t weight_bytes; + size_t scale_bytes; + size_t bias_bytes; + size_t output_dim; + size_t input_dim; + size_t group_size; + ColiAffineFormat format; + ColiAffineScalarFormat scalar_format; +} ColiAffineQuantizedView; + +static inline const char *coli_affine_status_string(ColiAffineStatus status) { + switch (status) { + case COLI_AFFINE_OK: return "ok"; + case COLI_AFFINE_NULL: return "null affine buffer"; + case COLI_AFFINE_BAD_FORMAT: return "unsupported affine format"; + case COLI_AFFINE_BAD_SHAPE: return "invalid affine shape"; + case COLI_AFFINE_TRUNCATED: return "truncated affine buffer"; + case COLI_AFFINE_OVERFLOW: return "affine size overflow"; + } + return "unknown affine status"; +} + +static inline unsigned coli_affine_bits(ColiAffineFormat format) { + return format == COLI_AFFINE_MLX_Q4 ? 4u : + format == COLI_AFFINE_MLX_Q8 ? 8u : 0u; +} + +static inline size_t coli_affine_scalar_size(ColiAffineScalarFormat format) { + return format == COLI_AFFINE_SCALAR_F32 ? 4u : + (format == COLI_AFFINE_SCALAR_F16 || + format == COLI_AFFINE_SCALAR_BF16) ? 2u : 0u; +} + +static inline int coli_affine_size_mul(size_t left, size_t right, size_t *out) { + if (left && right > SIZE_MAX / left) return 0; + *out = left * right; + return 1; +} + +static inline ColiAffineStatus +coli_affine_validate(const ColiAffineQuantizedView *view) { + size_t packed_words, group_count, need_weights, need_scalars; + unsigned bits, per_word; + size_t scalar_size; + + if (!view || !view->weights || !view->scales || !view->biases) + return COLI_AFFINE_NULL; + bits = coli_affine_bits(view->format); + scalar_size = coli_affine_scalar_size(view->scalar_format); + if (!bits || !scalar_size) return COLI_AFFINE_BAD_FORMAT; + if (!view->output_dim || !view->input_dim || !view->group_size) + return COLI_AFFINE_BAD_SHAPE; + + per_word = 32u / bits; + if (view->input_dim % view->group_size || + view->input_dim % per_word || view->group_size % per_word) + return COLI_AFFINE_BAD_SHAPE; + + packed_words = view->input_dim / per_word; + group_count = view->input_dim / view->group_size; + if (!coli_affine_size_mul(view->output_dim, packed_words, &need_weights) || + !coli_affine_size_mul(need_weights, sizeof(uint32_t), &need_weights) || + !coli_affine_size_mul(view->output_dim, group_count, &need_scalars) || + !coli_affine_size_mul(need_scalars, scalar_size, &need_scalars)) + return COLI_AFFINE_OVERFLOW; + if (view->weight_bytes < need_weights || + view->scale_bytes < need_scalars || view->bias_bytes < need_scalars) + return COLI_AFFINE_TRUNCATED; + return COLI_AFFINE_OK; +} + +static inline uint16_t coli_affine_load_le16(const uint8_t *source) { + return (uint16_t)source[0] | (uint16_t)((uint16_t)source[1] << 8); +} + +static inline uint32_t coli_affine_load_le32(const uint8_t *source) { + return (uint32_t)source[0] | ((uint32_t)source[1] << 8) | + ((uint32_t)source[2] << 16) | ((uint32_t)source[3] << 24); +} + +static inline float coli_affine_f32_from_bits(uint32_t bits) { + float value; + memcpy(&value, &bits, sizeof(value)); + return value; +} + +static inline float coli_affine_f16_to_f32(uint16_t value) { + uint32_t sign = ((uint32_t)value & 0x8000u) << 16; + int exponent = (int)(((uint32_t)value >> 10) & 0x1fu); + uint32_t mantissa = (uint32_t)value & 0x3ffu; + + if (exponent == 0) { + if (!mantissa) return coli_affine_f32_from_bits(sign); + exponent = 1; + while (!(mantissa & 0x400u)) { + mantissa <<= 1; + exponent--; + } + mantissa &= 0x3ffu; + exponent += 127 - 15; + } else if (exponent == 31) { + exponent = 255; + } else { + exponent += 127 - 15; + } + return coli_affine_f32_from_bits(sign | ((uint32_t)exponent << 23) | + (mantissa << 13)); +} + +static inline float coli_affine_load_scalar(const void *values, size_t index, + ColiAffineScalarFormat format) { + const uint8_t *bytes = (const uint8_t *)values; + if (format == COLI_AFFINE_SCALAR_F32) + return coli_affine_f32_from_bits(coli_affine_load_le32(bytes + index * 4)); + if (format == COLI_AFFINE_SCALAR_F16) + return coli_affine_f16_to_f32(coli_affine_load_le16(bytes + index * 2)); + return coli_affine_f32_from_bits( + (uint32_t)coli_affine_load_le16(bytes + index * 2) << 16); +} + +/* Correctness-first y[S,O] = x[S,I] @ dequant(W[O,I])^T. Accumulation is + * grouped exactly like Swiftlet's generic Metal kernel: two f32 reductions + * (q*x and x), followed by scale*qdot + bias*xsum for each group. + */ +static inline ColiAffineStatus +coli_affine_matmul_ref(float *output, const float *input, size_t batch, + const ColiAffineQuantizedView *view) { + ColiAffineStatus status = coli_affine_validate(view); + unsigned bits, per_word; + uint32_t mask; + size_t packed_words, groups, unused; + const uint8_t *weights; + + if (status != COLI_AFFINE_OK) return status; + if (!output || !input) return COLI_AFFINE_NULL; + if (!batch) return COLI_AFFINE_BAD_SHAPE; + if (!coli_affine_size_mul(batch, view->input_dim, &unused) || + !coli_affine_size_mul(batch, view->output_dim, &unused)) + return COLI_AFFINE_OVERFLOW; + + bits = coli_affine_bits(view->format); + per_word = 32u / bits; + mask = (UINT32_C(1) << bits) - 1u; + packed_words = view->input_dim / per_word; + groups = view->input_dim / view->group_size; + weights = (const uint8_t *)view->weights; + + for (size_t row = 0; row < view->output_dim; row++) { + const uint8_t *weight_row = weights + row * packed_words * 4; + for (size_t sample = 0; sample < batch; sample++) { + const float *input_row = input + sample * view->input_dim; + float sum = 0.0f; + for (size_t group = 0; group < groups; group++) { + const size_t word_base = group * (view->group_size / per_word); + const size_t input_base = group * view->group_size; + float quantized_dot = 0.0f, input_sum = 0.0f; + + for (size_t local = 0; local < view->group_size; local++) { + const size_t column = input_base + local; + const size_t word_index = word_base + local / per_word; + const uint32_t word = + coli_affine_load_le32(weight_row + word_index * 4); + const unsigned shift = bits * (unsigned)(local % per_word); + const float x = input_row[column]; + quantized_dot += (float)((word >> shift) & mask) * x; + input_sum += x; + } + const size_t scalar_index = row * groups + group; + const float scale = coli_affine_load_scalar( + view->scales, scalar_index, view->scalar_format); + const float bias = coli_affine_load_scalar( + view->biases, scalar_index, view->scalar_format); + sum += scale * quantized_dot + bias * input_sum; + } + output[sample * view->output_dim + row] = sum; + } + } + return COLI_AFFINE_OK; +} + +/* Row-major logical expansion out[output_dim, input_dim] = dequant(view): + * out[row*I + col] = scale[row, col/gs] * q[row, col] + bias[row, col/gs], + * with the SAME bit conventions as coli_affine_matmul_ref above (little-endian + * uint32 words, lowest bits = lowest logical column). This is the loader-side + * oracle: matmul_ref folds the dequant into its accumulation and never + * materialises the weights, but a dense tensor the ENGINE keeps resident + * (qwen36's attention/DeltaNet/shared-expert set from a Swiftlet container's + * model.safetensors) needs the expanded f32 rows exactly once, at load time. + * Kept next to matmul_ref so the two bit-level readings of the packed words + * cannot drift apart without this file changing. */ +static inline ColiAffineStatus +coli_affine_dequant_ref(const ColiAffineQuantizedView *view, float *output) { + ColiAffineStatus status = coli_affine_validate(view); + unsigned bits, per_word; + uint32_t mask; + size_t groups, unused; + const uint8_t *weights; + + if (status != COLI_AFFINE_OK) return status; + if (!output) return COLI_AFFINE_NULL; + if (!coli_affine_size_mul(view->output_dim, view->input_dim, &unused)) + return COLI_AFFINE_OVERFLOW; + + bits = coli_affine_bits(view->format); + per_word = 32u / bits; + mask = (UINT32_C(1) << bits) - 1u; + groups = view->input_dim / view->group_size; + weights = (const uint8_t *)view->weights; + + for (size_t row = 0; row < view->output_dim; row++) { + const uint8_t *weight_row = + weights + row * (view->input_dim / per_word) * 4; + float *out_row = output + row * view->input_dim; + for (size_t group = 0; group < groups; group++) { + const size_t word_base = group * (view->group_size / per_word); + const size_t input_base = group * view->group_size; + const size_t scalar_index = row * groups + group; + const float scale = coli_affine_load_scalar( + view->scales, scalar_index, view->scalar_format); + const float bias = coli_affine_load_scalar( + view->biases, scalar_index, view->scalar_format); + for (size_t local = 0; local < view->group_size; local++) { + const size_t word_index = word_base + local / per_word; + const uint32_t word = + coli_affine_load_le32(weight_row + word_index * 4); + const unsigned shift = bits * (unsigned)(local % per_word); + out_row[input_base + local] = + scale * (float)((word >> shift) & mask) + bias; + } + } + } + return COLI_AFFINE_OK; +} + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/c/backend_metal.h b/c/backend_metal.h index a8a01be2f..ef41e1f1f 100644 --- a/c/backend_metal.h +++ b/c/backend_metal.h @@ -4,6 +4,8 @@ #include #include +#include "affine_quant.h" + #ifdef __cplusplus extern "C" { #endif @@ -22,7 +24,9 @@ typedef struct ColiMetalTensor ColiMetalTensor; int coli_metal_init(void); void coli_metal_shutdown(void); int coli_metal_available(void); -/* Bytes of unified memory in use by wrapped tensors, and their count. */ +/* Resident handle count and bytes. Existing QT handles retain weight-only byte + * accounting; affine handles include their weight, scale, and bias buffers. + */ void coli_metal_stats(size_t *tensor_count, size_t *tensor_bytes); int coli_metal_mem_info(size_t *used_bytes, size_t *total_bytes); @@ -68,7 +72,57 @@ int coli_metal_matmul(ColiMetalTensor **tensor, const void *weights, const float *scales, int fmt, int S, int I, int O, int gs); +/* Returns whether a validated MLX affine descriptor and batch fit the Metal + * shader's uint32 indexing ABI. This predicate is portable and does not query + * hardware; coli_metal_affine_available() additionally requires initialized + * Q4/Q8 pipelines on a device with 32-lane simdgroups. + */ +static inline int +coli_metal_affine_dispatch_supported(const ColiAffineQuantizedView *view, + int batch) { + size_t elements, groups; + if (batch <= 0 || coli_affine_validate(view) != COLI_AFFINE_OK) + return 0; + if (view->output_dim > UINT32_MAX || view->input_dim > UINT32_MAX || + view->group_size > UINT32_MAX) + return 0; + groups = view->input_dim / view->group_size; + if (!coli_affine_size_mul(view->output_dim, groups, &elements) || + elements > UINT32_MAX) + return 0; + if (!coli_affine_size_mul((size_t)batch, view->input_dim, &elements) || + elements > UINT32_MAX) + return 0; + if (!coli_affine_size_mul((size_t)batch, view->output_dim, &elements) || + elements > UINT32_MAX) + return 0; + return 1; +} + +/* Dedicated MLX affine path. This does not use or extend Colibri's QT fmt + * namespace: `view` names unsigned uint32-packed Q4/Q8 plus scale and bias. + * The first successful call wraps the three stable buffers in `tensor`, as + * coli_metal_matmul does for its existing weight formats. Returns 0 for CPU + * fallback when validation, capability, allocation, or execution fails. + */ +typedef enum { + COLI_METAL_AFFINE_CAP_NOT_INITIALIZED = 0, + COLI_METAL_AFFINE_CAP_PIPELINE_UNAVAILABLE = 1, + COLI_METAL_AFFINE_CAP_SIMD_WIDTH_UNSUPPORTED = 2, + COLI_METAL_AFFINE_CAP_READY = 3 +} ColiMetalAffineCapability; + +/* Reports why the runtime affine path is unavailable so tests and callers can + * distinguish an optional hardware-width fallback from a pipeline failure. + */ +ColiMetalAffineCapability coli_metal_affine_capability(void); +int coli_metal_affine_available(void); +int coli_metal_matmul_affine(ColiMetalTensor **tensor, + float *y, const float *x, int batch, + const ColiAffineQuantizedView *view); + void coli_metal_tensor_free(ColiMetalTensor *tensor); +/* Weight payload bytes only, preserving the established QT API contract. */ size_t coli_metal_tensor_bytes(const ColiMetalTensor *tensor); /* diff --git a/c/backend_metal.mm b/c/backend_metal.mm index e04847a0d..3f37eb923 100644 --- a/c/backend_metal.mm +++ b/c/backend_metal.mm @@ -154,6 +154,107 @@ kernel void mm_gemv(device const uchar* w [[buffer(0)]], // raw weight by if (slane == 0) y[row] = (fmt == 4 || fmt == 8) ? acc : acc * scale[o]; } +// MLX affine packed GEMV, adapted from Swiftlet's gemv_affine_fast Q4/Q8 +// kernels (Apache-2.0, commit b3a04676748c7597800c5bcc8b80a32508f9f43d). +// Modified here for Colibri's batched standalone API and checked descriptor. +// This is a separate representation from mm_gemv fmt=4: q is unsigned and +// every group reconstructs weights as scale*q + bias. +struct AffineGemvParams { + uint output_dim; + uint input_dim; + uint group_size; + uint scalar_format; // ColiAffineScalarFormat: 0=f32, 1=f16, 2=bf16 + uint batch; +}; + +inline uint affine_load_u32(device const uchar *p, ulong offset) { + return uint(p[offset]) | (uint(p[offset+1]) << 8) | + (uint(p[offset+2]) << 16) | (uint(p[offset+3]) << 24); +} + +inline float affine_load_scalar(device const uchar *p, uint index, uint format) { + if (format == 0) return as_type(affine_load_u32(p, ulong(index)*4)); + ulong offset = ulong(index)*2; + ushort value = ushort(p[offset]) | (ushort(p[offset+1]) << 8); + if (format == 1) return float(as_type(value)); + return as_type(uint(value) << 16); +} + +kernel void mm_affine_q4(device const float *x [[buffer(0)]], + device const uchar *weights [[buffer(1)]], + device const uchar *scales [[buffer(2)]], + device const uchar *biases [[buffer(3)]], + device float *y [[buffer(4)]], + constant AffineGemvParams &p [[buffer(5)]], + uint2 position [[thread_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + uint flat_row=position.y, total=p.batch*p.output_dim; + if (flat_row >= total) return; + uint sample=flat_row/p.output_dim, row=flat_row%p.output_dim; + uint packed_cols=p.input_dim/8, groups=p.input_dim/p.group_size; + uint words_per_group=p.group_size/8; + device const uint *wr=(device const uint*)(weights + ulong(row)*packed_cols*4); + float acc=0.f; + for (uint group=lane; group> 4) & 0xFu)*x[xo+1] + + float((word >> 8) & 0xFu)*x[xo+2] + + float((word >> 12) & 0xFu)*x[xo+3] + + float((word >> 16) & 0xFu)*x[xo+4] + + float((word >> 20) & 0xFu)*x[xo+5] + + float((word >> 24) & 0xFu)*x[xo+6] + + float((word >> 28) & 0xFu)*x[xo+7]; + xsum += x[xo]+x[xo+1]+x[xo+2]+x[xo+3] + + x[xo+4]+x[xo+5]+x[xo+6]+x[xo+7]; + } + uint scalar_index=row*groups+group; + acc += affine_load_scalar(scales,scalar_index,p.scalar_format)*qdot + + affine_load_scalar(biases,scalar_index,p.scalar_format)*xsum; + } + acc=simd_sum(acc); + if (lane == 0) y[flat_row]=acc; +} + +kernel void mm_affine_q8(device const float *x [[buffer(0)]], + device const uchar *weights [[buffer(1)]], + device const uchar *scales [[buffer(2)]], + device const uchar *biases [[buffer(3)]], + device float *y [[buffer(4)]], + constant AffineGemvParams &p [[buffer(5)]], + uint2 position [[thread_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + uint flat_row=position.y, total=p.batch*p.output_dim; + if (flat_row >= total) return; + uint sample=flat_row/p.output_dim, row=flat_row%p.output_dim; + uint packed_cols=p.input_dim/4, groups=p.input_dim/p.group_size; + uint words_per_group=p.group_size/4; + device const uint *wr=(device const uint*)(weights + ulong(row)*packed_cols*4); + float acc=0.f; + for (uint group=lane; group> 8) & 0xFFu)*x[xo+1] + + float((word >> 16) & 0xFFu)*x[xo+2] + + float((word >> 24) & 0xFFu)*x[xo+3]; + xsum += x[xo]+x[xo+1]+x[xo+2]+x[xo+3]; + } + uint scalar_index=row*groups+group; + acc += affine_load_scalar(scales,scalar_index,p.scalar_format)*qdot + + affine_load_scalar(biases,scalar_index,p.scalar_format)*xsum; + } + acc=simd_sum(acc); + if (lane == 0) y[flat_row]=acc; +} + // Batched bindless expert GEMV: each row gr belongs to expert erow[gr], whose weight and // scale live at gpuAddresses waddr[e]/saddr[e] (zero-copy in the RAM slab). fmt 1=i8, 2=i4 // per-row, 4=i4 grouped (scale layout [O][ng], ng=ceil(K/qgs) -- same convention as mm_gemv @@ -583,12 +684,23 @@ kernel void kda_state( struct ColiMetalTensor { id w; // weights (wrapped, zero-copy when page-aligned) id s; // scales - int fmt, I, O; size_t wbytes; + id b; // affine bias (nil for existing QT formats) + int fmt, scalar_fmt, affine; + size_t I, O, group_size, wbytes, resident_bytes; + const void *host_w, *host_s, *host_b; }; +static void tensor_destroy(ColiMetalTensor *t) { + if (!t) return; + t->w=nil; t->s=nil; t->b=nil; + delete t; +} + static id g_dev; static id g_queue; static id g_gemv, g_moe_gemv, g_moe_silu, g_moe_fwht; +static id g_affine_q4, g_affine_q8; +static int g_affine_width_ok; // fmt=6: sign-bit buffers for the GPU FWHT, one per tile size, cached forever (a // handful of sizes). The xorshift64* draw replicates quant.h e8_signs exactly — @@ -814,6 +926,7 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { g_moe_silu = [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"moe_silu"] error:&err]; g_moe_fwht = [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"moe_fwht"] error:&err]; auto P=[&](const char*n){ return [g_dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@(n)] error:&err]; }; + g_affine_q4=P("mm_affine_q4"); g_affine_q8=P("mm_affine_q8"); g_a_rms=P("a_rmsnorm"); g_a_rope=P("a_rope"); g_a_copy=P("a_copy"); g_a_qabs=P("a_qabs"); g_a_score=P("a_score"); g_a_smax=P("a_smax"); g_a_clat=P("a_clat"); g_a_ctx=P("a_ctx"); g_a_add=P("a_add"); g_r_router=P("r_router"); g_r_top8=P("r_top8"); g_r_top8p=P("r_top8_par"); @@ -821,6 +934,17 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { g_kda_state=P("kda_state"); g_kda_conv_silu=P("kda_conv_silu"); g_kda_l2_norm=P("kda_l2_norm"); if(!g_a_add||!g_r_router||!g_r_top8||!g_r_top8p||!g_kv_rope||!g_kv_clear||!g_kda_state||!g_kda_conv_silu||!g_kda_l2_norm){ fprintf(stderr,"[metal] tail pipelines failed\n"); g_dev=nil; return 0; } + if (!g_affine_q4 || !g_affine_q8) { + fprintf(stderr,"[metal] MLX affine GEMV disabled: pipeline creation failed\n"); + g_affine_width_ok = 0; + } else { + g_affine_width_ok = [g_affine_q4 threadExecutionWidth] == 32 && + [g_affine_q8 threadExecutionWidth] == 32; + if (!g_affine_width_ok) + fprintf(stderr,"[metal] MLX affine GEMV disabled: simd width Q4=%lu Q8=%lu, expected 32\n", + (unsigned long)[g_affine_q4 threadExecutionWidth], + (unsigned long)[g_affine_q8 threadExecutionWidth]); + } // r_top8_par's reduction hardcodes SIMD width 32 (shuffle-down offsets 16..1, one // 32-thread threadgroup per row). True on all Apple Silicon shipped to date, but a // non-32-width device would reduce wrongly AND race multiple lane-0 writers, so this @@ -940,9 +1064,21 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { } #endif g_resset_obj=nil; g_resset_enabled=false; g_resset_dirty=false; - g_gemv=nil; g_queue=nil; g_dev=nil; g_tensor_count=g_tensor_bytes=0; + g_gemv=nil; g_affine_q4=nil; g_affine_q8=nil; g_affine_width_ok=0; + g_queue=nil; g_dev=nil; g_tensor_count=g_tensor_bytes=0; } extern "C" int coli_metal_available(void) { return g_dev != nil; } +extern "C" ColiMetalAffineCapability coli_metal_affine_capability(void) { + if (!g_dev) return COLI_METAL_AFFINE_CAP_NOT_INITIALIZED; + if (!g_affine_q4 || !g_affine_q8) + return COLI_METAL_AFFINE_CAP_PIPELINE_UNAVAILABLE; + if (!g_affine_width_ok) + return COLI_METAL_AFFINE_CAP_SIMD_WIDTH_UNSUPPORTED; + return COLI_METAL_AFFINE_CAP_READY; +} +extern "C" int coli_metal_affine_available(void) { + return coli_metal_affine_capability() == COLI_METAL_AFFINE_CAP_READY; +} extern "C" void coli_metal_stats(size_t *c, size_t *b) { if(c)*c=g_tensor_count; if(b)*b=g_tensor_bytes; } extern "C" int coli_metal_mem_info(size_t *used, size_t *total) { if (!g_dev) return 0; @@ -960,13 +1096,18 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { if(!weights||!x||!y) return 0; @autoreleasepool { ColiMetalTensor *t = *tp; + if (t && t->affine) return 0; if (!t) { t = new ColiMetalTensor(); - t->fmt = fmt; t->I = I; t->O = O; t->wbytes = fmt_bytes(fmt, I, O); + t->fmt = fmt; t->scalar_fmt = 0; t->affine = 0; + t->I = I; t->O = O; t->group_size = (size_t)gs; + t->host_w = weights; t->host_s = scales; t->host_b = NULL; + t->wbytes = fmt_bytes(fmt, I, O); t->resident_bytes = t->wbytes; + t->b = nil; t->w = wrap(weights, t->wbytes); { size_t sb=fmt_scale_bytes(fmt, I, O, gs); t->s=sb?wrap(scales,sb):0; } *tp = t; - g_tensor_count++; g_tensor_bytes += t->wbytes; + g_tensor_count++; g_tensor_bytes += t->resident_bytes; } id bx = [g_dev newBufferWithBytes:x length:(size_t)S*I*sizeof(float) options:MTLResourceStorageModeShared]; id by = [g_dev newBufferWithLength:(size_t)S*O*sizeof(float) options:MTLResourceStorageModeShared]; @@ -986,6 +1127,100 @@ static size_t fmt_scale_bytes(int fmt, int I, int O, int gs) { return 1; } +struct AffineGemvParamsHost { + uint32_t output_dim; + uint32_t input_dim; + uint32_t group_size; + uint32_t scalar_format; + uint32_t batch; +}; +static_assert(sizeof(AffineGemvParamsHost) == 5*sizeof(uint32_t), + "Metal affine parameter layout drift"); + +extern "C" int coli_metal_matmul_affine(ColiMetalTensor **tp, + float *y, const float *x, int batch, + const ColiAffineQuantizedView *view) { + if (!tp || !y || !x || !coli_metal_affine_available() || + !coli_metal_affine_dispatch_supported(view,batch)) + return 0; + @autoreleasepool { + const unsigned bits=coli_affine_bits(view->format), per_word=32u/bits; + const size_t groups=view->input_dim/view->group_size; + const size_t wbytes=view->output_dim*(view->input_dim/per_word)*sizeof(uint32_t); + const size_t scalar_bytes=view->output_dim*groups* + coli_affine_scalar_size(view->scalar_format); + if (scalar_bytes > (SIZE_MAX-wbytes)/2) + return 0; + const size_t resident_bytes=wbytes+2*scalar_bytes; + ColiMetalTensor *t=*tp; + if (t && (!t->affine || t->fmt != (int)view->format || + t->scalar_fmt != (int)view->scalar_format || + t->I != view->input_dim || t->O != view->output_dim || + t->group_size != view->group_size || + t->host_w != view->weights || t->host_s != view->scales || + t->host_b != view->biases)) + return 0; + const bool fresh=t==NULL; + if (fresh) { + id bw=wrap(view->weights,wbytes); + id bs=wrap(view->scales,scalar_bytes); + id bb=wrap(view->biases,scalar_bytes); + if (!bw || !bs || !bb) return 0; + t=new ColiMetalTensor(); + t->w=bw; t->s=bs; t->b=bb; + t->fmt=(int)view->format; t->scalar_fmt=(int)view->scalar_format; + t->affine=1; t->I=view->input_dim; t->O=view->output_dim; + t->group_size=view->group_size; t->wbytes=wbytes; + t->resident_bytes=resident_bytes; + t->host_w=view->weights; t->host_s=view->scales; t->host_b=view->biases; + } + + const size_t input_bytes=(size_t)batch*view->input_dim*sizeof(float); + const size_t output_bytes=(size_t)batch*view->output_dim*sizeof(float); + id bx=[g_dev newBufferWithBytes:x length:input_bytes + options:MTLResourceStorageModeShared]; + id by=[g_dev newBufferWithLength:output_bytes + options:MTLResourceStorageModeShared]; + if (!bx || !by) { + if (fresh) tensor_destroy(t); + return 0; + } + AffineGemvParamsHost p = { + (uint32_t)view->output_dim, (uint32_t)view->input_dim, + (uint32_t)view->group_size, (uint32_t)view->scalar_format, + (uint32_t)batch + }; + id cb=[g_queue commandBuffer]; + id e=[cb computeCommandEncoder]; + if (!cb || !e) { + if (fresh) tensor_destroy(t); + return 0; + } + [e setComputePipelineState:(view->format == COLI_AFFINE_MLX_Q4 ? + g_affine_q4 : g_affine_q8)]; + [e setBuffer:bx offset:0 atIndex:0]; [e setBuffer:t->w offset:0 atIndex:1]; + [e setBuffer:t->s offset:0 atIndex:2]; [e setBuffer:t->b offset:0 atIndex:3]; + [e setBuffer:by offset:0 atIndex:4]; [e setBytes:&p length:sizeof(p) atIndex:5]; + [e dispatchThreads:MTLSizeMake(32,(NSUInteger)batch*view->output_dim,1) + threadsPerThreadgroup:MTLSizeMake(32,1,1)]; + [e endEncoding]; [cb commit]; [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) { + fprintf(stderr,"[metal] affine GEMV command buffer failed (status=%ld): %s\n", + (long)cb.status, + cb.error ? [[cb.error localizedDescription] UTF8String] : "?"); + if (fresh) tensor_destroy(t); + return 0; + } + memcpy(y,[by contents],output_bytes); + if (fresh) { + *tp=t; + g_tensor_count++; + g_tensor_bytes+=t->resident_bytes; + } + } + return 1; +} + // ---- fused decode attention scratch (GLM-5.2 dims) ---- enum { AH=6144, AHEADS=64, AQLORA=2048, AKVL=512, AROPE=64, AVH=256, AQH=256, ANOPE=192, AROWSH=448, AHQH=AHEADS*AQH, AHVH=AHEADS*AVH, AMAXS=4 }; static id ax_,aqr_,aqf_,acomp_,aqabs_,ascore_,aclat_,actx_,aout_,aqaln_,akvaln_; static size_t ascore_cap; @@ -1369,8 +1604,8 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc, extern "C" void coli_metal_tensor_free(ColiMetalTensor *t) { if (!t) return; - g_tensor_count--; g_tensor_bytes -= t->wbytes; - t->w = nil; t->s = nil; delete t; + g_tensor_count--; g_tensor_bytes -= t->resident_bytes; + tensor_destroy(t); } extern "C" size_t coli_metal_tensor_bytes(const ColiMetalTensor *t) { return t ? t->wbytes : 0; } diff --git a/c/json.h b/c/json.h index 649b92cae..373870424 100644 --- a/c/json.h +++ b/c/json.h @@ -8,6 +8,7 @@ #include #include #include +#include typedef enum { J_NULL, J_BOOL, J_NUM, J_STR, J_ARR, J_OBJ } jtype; @@ -28,6 +29,8 @@ typedef struct { size_t acap, aoff; int depth; /* annidamento corrente: bound contro lo stack-overflow * da JSON malevolo tipo [[[[...]]]] (discesa ricorsiva) */ + int strict; /* grammatica esatta; rifiuta NUL per stringhe C sicure */ + int error; } jparser; /* tetto di annidamento: gli header safetensors / config sono piatti (profondita' @@ -43,7 +46,14 @@ static char *j_dup(jparser *p, const char *b, int n) { return d; } -static void j_ws(jparser *p) { while (*p->s && isspace((unsigned char)*p->s)) p->s++; } +static void j_ws(jparser *p) { + if (p->strict) { + while (*p->s == ' ' || *p->s == '\t' || + *p->s == '\n' || *p->s == '\r') p->s++; + } else { + while (*p->s && isspace((unsigned char)*p->s)) p->s++; + } +} static jval *j_new(jtype t) { jval *v = (jval *)calloc(1, sizeof(jval)); @@ -52,12 +62,30 @@ static jval *j_new(jtype t) { static jval *j_parse_val(jparser *p); +static int j_hex4(const char *s, unsigned *out) { + unsigned value = 0; + for (int i = 0; i < 4; i++) { + unsigned digit; + unsigned char c = (unsigned char)s[i]; + if (c >= '0' && c <= '9') digit = c - '0'; + else if (c >= 'a' && c <= 'f') digit = c - 'a' + 10; + else if (c >= 'A' && c <= 'F') digit = c - 'A' + 10; + else return -1; + value = (value << 4) | digit; + } + *out = value; + return 0; +} + static char *j_parse_str_raw(jparser *p) { /* SEC (GHSA-2qrj): fail closed if not actually at a quote. The old comment * "assume *p->s == '\"'" was violated on the object-key path, and the * unconditional p->s++ would step past the buffer's NUL terminator and scan * adjacent heap (OOB read leaking into tensor names). */ - if (*p->s != '"') return j_dup(p, "", 0); + if (*p->s != '"') { + if (p->strict) p->error = 1; + return j_dup(p, "", 0); + } p->s++; /* buffer su heap che CRESCE: niente troncamento silenzioso a 64KB (le stringhe * lunghe di tokenizer.json/config venivano tagliate) e niente 64KB di stack. */ @@ -67,7 +95,10 @@ static char *j_parse_str_raw(jparser *p) { if (!tmp) { fprintf(stderr, "OOM parsing JSON string\n"); exit(1); } } tmp[n++] = (char)(ch); }while(0) while (*p->s && *p->s != '"') { char c = *p->s++; - if (c == '\\' && *p->s) { + if (p->strict && (unsigned char)c < 0x20) p->error = 1; + if (c == '\\' && !*p->s) { + if (p->strict) p->error = 1; + } else if (c == '\\') { char e = *p->s++; switch (e) { case 'n': c = '\n'; break; case 't': c = '\t'; break; @@ -75,13 +106,28 @@ static char *j_parse_str_raw(jparser *p) { case 'f': c = '\f'; break; case '/': c = '/'; break; case '\\': c = '\\'; break; case '"': c = '"'; break; case 'u': { /* \uXXXX -> codepoint UTF-8 (con coppie surrogate) */ - if (!p->s[0]||!p->s[1]||!p->s[2]||!p->s[3]) { c='?'; break; } /* \u troncato: non leggere oltre il NUL */ + if (!p->s[0]||!p->s[1]||!p->s[2]||!p->s[3]) { + if (p->strict) p->error = 1; + c='?'; break; /* \u troncato: non leggere oltre il NUL */ + } unsigned cp = (unsigned)strtoul((char[]){p->s[0],p->s[1],p->s[2],p->s[3],0}, NULL, 16); + if (p->strict && j_hex4(p->s, &cp)) p->error = 1; p->s += 4; - if (cp >= 0xD800 && cp <= 0xDBFF && p->s[0]=='\\' && p->s[1]=='u' - && p->s[2] && p->s[3] && p->s[4] && p->s[5]) { - unsigned lo = (unsigned)strtoul((char[]){p->s[2],p->s[3],p->s[4],p->s[5],0}, NULL, 16); - if (lo >= 0xDC00 && lo <= 0xDFFF) { cp = 0x10000 + ((cp-0xD800)<<10) + (lo-0xDC00); p->s += 6; } + if (p->strict && cp == 0) p->error = 1; + if (cp >= 0xD800 && cp <= 0xDBFF) { + unsigned lo = 0; + int has_low = p->s[0]=='\\' && p->s[1]=='u' && + p->s[2] && p->s[3] && p->s[4] && p->s[5] && + j_hex4(p->s + 2, &lo) == 0 && + lo >= 0xDC00 && lo <= 0xDFFF; + if (has_low) { + cp = 0x10000 + ((cp-0xD800)<<10) + (lo-0xDC00); + p->s += 6; + } else if (p->strict) { + p->error = 1; + } + } else if (p->strict && cp >= 0xDC00 && cp <= 0xDFFF) { + p->error = 1; } if (cp < 0x80) { J_PUT(cp); } else if (cp < 0x800) { J_PUT(0xC0|(cp>>6)); J_PUT(0x80|(cp&0x3F)); } @@ -89,13 +135,14 @@ static char *j_parse_str_raw(jparser *p) { else { J_PUT(0xF0|(cp>>18)); J_PUT(0x80|((cp>>12)&0x3F)); J_PUT(0x80|((cp>>6)&0x3F)); J_PUT(0x80|(cp&0x3F)); } continue; } - default: c = e; break; + default: if (p->strict) p->error = 1; c = e; break; } } J_PUT(c); } #undef J_PUT if (*p->s == '"') p->s++; + else if (p->strict) p->error = 1; char *out = j_dup(p, tmp, (int)n); free(tmp); return out; } @@ -105,7 +152,11 @@ static jval *j_parse_val(jparser *p) { char c = *p->s; if (c == '"') { jval *v = j_new(J_STR); v->str = j_parse_str_raw(p); return v; } if (c == '{') { - if (++p->depth > J_MAX_DEPTH) { p->depth--; return j_new(J_NULL); } + if (++p->depth > J_MAX_DEPTH) { + if (p->strict) p->error = 1; + p->depth--; + return j_new(J_NULL); + } p->s++; jval *v = j_new(J_OBJ); int cap = 8; v->keys = malloc(cap * sizeof(char*)); @@ -114,11 +165,17 @@ static jval *j_parse_val(jparser *p) { if (!v->kids) { fprintf(stderr, "OOM parsing JSON object\n"); exit(1); } j_ws(p); if (*p->s == '}') { p->s++; p->depth--; return v; } + int closed = 0; for (;;) { j_ws(p); - if (*p->s != '"') break; /* SEC (GHSA-2qrj): object key must be a quoted string; stop on malformed input */ + if (*p->s != '"') { + if (p->strict) p->error = 1; + break; /* SEC (GHSA-2qrj): object key must be a quoted string; stop on malformed input */ + } char *key = j_parse_str_raw(p); - j_ws(p); if (*p->s == ':') p->s++; + j_ws(p); + if (*p->s == ':') p->s++; + else if (p->strict) p->error = 1; jval *val = j_parse_val(p); if (v->len == cap) { cap *= 2; char **nk = (char**)realloc(v->keys, cap*sizeof(char*)); @@ -129,22 +186,45 @@ static jval *j_parse_val(jparser *p) { v->kids = nc; } v->keys[v->len] = key; v->kids[v->len] = val; v->len++; j_ws(p); - if (*p->s == ',') { p->s++; continue; } - if (*p->s == '}') { p->s++; break; } + if (*p->s == ',') { + p->s++; + if (p->strict) { + j_ws(p); + if (*p->s == '}') p->error = 1; + } + continue; + } + if (*p->s == '}') { p->s++; closed = 1; break; } + if (p->strict) p->error = 1; break; } + if (p->strict && !closed) p->error = 1; p->depth--; return v; } if (c == '[') { - if (++p->depth > J_MAX_DEPTH) { p->depth--; return j_new(J_NULL); } + if (++p->depth > J_MAX_DEPTH) { + if (p->strict) p->error = 1; + p->depth--; + return j_new(J_NULL); + } p->s++; jval *v = j_new(J_ARR); int cap = 8; v->kids = malloc(cap * sizeof(jval*)); if (!v->kids) { fprintf(stderr, "OOM parsing JSON array\n"); exit(1); } j_ws(p); if (*p->s == ']') { p->s++; p->depth--; return v; } + int closed = 0; for (;;) { + if (p->strict) { + j_ws(p); + if (*p->s == ']') { + p->error = 1; + p->s++; + closed = 1; + break; + } + } jval *val = j_parse_val(p); if (v->len == cap) { cap *= 2; jval **nc = (jval**)realloc(v->kids, cap*sizeof(jval*)); @@ -153,9 +233,11 @@ static jval *j_parse_val(jparser *p) { v->kids[v->len++] = val; j_ws(p); if (*p->s == ',') { p->s++; continue; } - if (*p->s == ']') { p->s++; break; } + if (*p->s == ']') { p->s++; closed = 1; break; } + if (p->strict) p->error = 1; break; } + if (p->strict && !closed) p->error = 1; p->depth--; return v; } @@ -163,12 +245,41 @@ static jval *j_parse_val(jparser *p) { if (c == 'f' && !strncmp(p->s, "false", 5)) { p->s += 5; jval *v = j_new(J_BOOL); v->boolean = 0; return v; } if (c == 'n' && !strncmp(p->s, "null", 4)) { p->s += 4; return j_new(J_NULL); } /* numero */ + if (p->strict) { + const char *start = p->s; + const char *scan = start; + if (*scan == '-') scan++; + if (*scan == '0') { + scan++; + if (isdigit((unsigned char)*scan)) p->error = 1; + } else if (*scan >= '1' && *scan <= '9') { + do { scan++; } while (isdigit((unsigned char)*scan)); + } else { + p->error = 1; + } + if (*scan == '.') { + scan++; + if (!isdigit((unsigned char)*scan)) p->error = 1; + while (isdigit((unsigned char)*scan)) scan++; + } + if (*scan == 'e' || *scan == 'E') { + scan++; + if (*scan == '+' || *scan == '-') scan++; + if (!isdigit((unsigned char)*scan)) p->error = 1; + while (isdigit((unsigned char)*scan)) scan++; + } + char *end; + double d = strtod(start, &end); + if (end != scan || end == start || !isfinite(d)) p->error = 1; + p->s = scan; + jval *v = j_new(J_NUM); v->num = d; return v; + } { char *end; double d = strtod(p->s, &end); p->s = end; jval *v = j_new(J_NUM); v->num = d; return v; } } /* API */ -static jval *json_parse(const char *text, char **arena_out) { - jparser p = { text, NULL, 0, 0, 0 }; +static inline jval *json_parse(const char *text, char **arena_out) { + jparser p = { text, NULL, 0, 0, 0, 0, 0 }; jval *v = j_parse_val(&p); if (arena_out) *arena_out = p.arena; else free(p.arena); return v; @@ -195,4 +306,21 @@ static void json_free(jval *v) { free(v); } +/* Exact parsing is opt-in so existing safetensors/config callers retain the + * historical permissive behavior. It rejects malformed syntax, trailing + * non-whitespace input, and U+0000, which cannot be represented safely by + * this parser's NUL-terminated string API. */ +static inline jval *json_parse_exact(const char *text, char **arena_out) { + if (!text) return NULL; + jparser p = { text, NULL, 0, 0, 0, 1, 0 }; + jval *v = j_parse_val(&p); + j_ws(&p); + if (p.error || *p.s) { + json_free(v); + v = NULL; + } + if (arena_out) *arena_out = p.arena; else free(p.arena); + return v; +} + #endif diff --git a/c/qpack.c b/c/qpack.c new file mode 100644 index 000000000..d1e8a8056 --- /dev/null +++ b/c/qpack.c @@ -0,0 +1,685 @@ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#include "qpack.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "compat.h" +#include "json.h" + +#define COLI_QPACK_MAX_JSON (16u * 1024u * 1024u) +#define COLI_QPACK_MAX_LAYERS 4096u +#define COLI_QPACK_MAX_EXPERTS 1048576u +#define COLI_QPACK_MAX_SECTIONS 256u +#define COLI_QPACK_MAX_FILES 8192u +#define COLI_QPACK_MAX_RANK 16u +#define COLI_QPACK_MAX_EXPERT_STRIDE (1u << 30) + +static int qpack_fail(char *error, size_t capacity, const char *format, ...) { + if (error && capacity) { + va_list args; + va_start(args, format); + vsnprintf(error, capacity, format, args); + va_end(args); + } + return -1; +} + +static char *qpack_copy_string(const char *source) { + size_t size = strlen(source) + 1; + char *copy = (char *)malloc(size); + if (copy) memcpy(copy, source, size); + return copy; +} + +static char *qpack_join_path(const char *left, const char *right) { + size_t left_size = strlen(left), right_size = strlen(right); + int slash = left_size > 0 && left[left_size - 1] != '/' && + left[left_size - 1] != '\\'; + size_t extra = (size_t)slash + 1; + if (right_size > SIZE_MAX - extra || + left_size > SIZE_MAX - right_size - extra) + return NULL; + char *path = (char *)malloc(left_size + (size_t)slash + right_size + 1); + if (!path) return NULL; + memcpy(path, left, left_size); + if (slash) path[left_size++] = '/'; + memcpy(path + left_size, right, right_size + 1); + return path; +} + +static char *qpack_read_text(const char *path, size_t *size_out, + char *error, size_t error_capacity) { + int fd = open(path, COMPAT_O_RDONLY); + if (fd < 0) + return qpack_fail(error, error_capacity, "open %s: %s", + path, strerror(errno)), NULL; + struct stat st; + if (fstat(fd, &st) != 0) { + int saved = errno; + close(fd); + return qpack_fail(error, error_capacity, "fstat %s: %s", + path, strerror(saved)), NULL; + } + if (!S_ISREG(st.st_mode) || st.st_size <= 0 || + (uint64_t)st.st_size > COLI_QPACK_MAX_JSON) { + close(fd); + return qpack_fail(error, error_capacity, + "%s is not a bounded regular JSON file", path), NULL; + } + size_t size = (size_t)st.st_size; + char *text = (char *)malloc(size + 1); + if (!text) { + close(fd); + return qpack_fail(error, error_capacity, + "out of memory reading %s", path), NULL; + } + ssize_t count; + do { + count = pread(fd, text, size, 0); + } while (count < 0 && errno == EINTR); + int saved = errno; + close(fd); + if (count < 0 || (size_t)count != size) { + free(text); + if (count < 0) + return qpack_fail(error, error_capacity, "pread %s: %s", + path, strerror(saved)), NULL; + return qpack_fail(error, error_capacity, "short read from %s", path), NULL; + } + if (memchr(text, 0, size)) { + free(text); + return qpack_fail(error, error_capacity, + "%s contains an embedded NUL", path), NULL; + } + text[size] = 0; + if (size_out) *size_out = size; + return text; +} + +static int qpack_json_u64(jval *value, uint64_t *output) { + if (!value || value->t != J_NUM || !isfinite(value->num) || + value->num < 0.0 || floor(value->num) != value->num || + value->num > 9007199254740991.0) + return -1; + *output = (uint64_t)value->num; + return 0; +} + +static int qpack_object_has_duplicate_keys(jval *object) { + if (!object || object->t != J_OBJ) return 0; + for (int i = 0; i < object->len; i++) + for (int j = 0; j < i; j++) + if (strcmp(object->keys[i], object->keys[j]) == 0) return 1; + return 0; +} + +static int qpack_dtype_size(const char *dtype, size_t *size) { + static const struct { const char *name; size_t size; } types[] = { + {"BOOL", 1}, {"U8", 1}, {"I8", 1}, {"U16", 2}, {"I16", 2}, + {"F16", 2}, {"BF16", 2}, {"U32", 4}, {"I32", 4}, {"F32", 4}, + {"U64", 8}, {"I64", 8}, {"F64", 8} + }; + for (size_t i = 0; i < sizeof(types) / sizeof(types[0]); i++) { + if (strcmp(dtype, types[i].name) == 0) { + *size = types[i].size; + return 0; + } + } + return -1; +} + +static void qpack_free_layout(ColiQpackLayout *layout) { + if (!layout) return; + for (size_t i = 0; i < layout->section_count; i++) { + free(layout->sections[i].name); + free(layout->sections[i].dtype); + free(layout->sections[i].shape); + } + free(layout->sections); + free(layout->linear_layers); + memset(layout, 0, sizeof(*layout)); +} + +static int qpack_parse_layout(ColiQpackLayout *layout, const char *text, + char *error, size_t error_capacity) { + jval *root = json_parse_exact(text, NULL); + if (!root || root->t != J_OBJ || qpack_object_has_duplicate_keys(root)) { + json_free(root); + return qpack_fail(error, error_capacity, + "layout.json must be an object with unique keys"); + } + + uint64_t expert_count, layer_count, expert_stride; + if (qpack_json_u64(json_get(root, "expertCount"), &expert_count) || + !expert_count || expert_count > COLI_QPACK_MAX_EXPERTS || + qpack_json_u64(json_get(root, "layerCount"), &layer_count) || + !layer_count || layer_count > COLI_QPACK_MAX_LAYERS || + qpack_json_u64(json_get(root, "expertStride"), &expert_stride) || + !expert_stride || expert_stride > COLI_QPACK_MAX_EXPERT_STRIDE || + expert_stride > SIZE_MAX) { + json_free(root); + return qpack_fail(error, error_capacity, + "layout has invalid counts or stride"); + } + if (expert_stride % COLI_QPACK_PAGE_ALIGNMENT != 0) { + json_free(root); + return qpack_fail(error, error_capacity, + "expertStride must be 16 KiB aligned"); + } + + jval *sections = json_get(root, "sections"); + jval *linear = json_get(root, "linearLayers"); + if (!sections || sections->t != J_ARR || sections->len <= 0 || + sections->len > (int)COLI_QPACK_MAX_SECTIONS || + !linear || linear->t != J_ARR || + (uint64_t)linear->len != layer_count) { + json_free(root); + return qpack_fail(error, error_capacity, + "layout sections or linearLayers are invalid"); + } + + layout->expert_count = (size_t)expert_count; + layout->layer_count = (size_t)layer_count; + layout->expert_stride = (size_t)expert_stride; + layout->section_count = (size_t)sections->len; + layout->sections = (ColiQpackSection *)calloc( + layout->section_count, sizeof(ColiQpackSection)); + layout->linear_layers = (unsigned char *)calloc(layout->layer_count, 1); + if (!layout->sections || !layout->linear_layers) { + json_free(root); + qpack_free_layout(layout); + return qpack_fail(error, error_capacity, + "out of memory parsing layout"); + } + + for (size_t i = 0; i < layout->layer_count; i++) { + jval *value = linear->kids[i]; + if (!value || value->t != J_BOOL) { + json_free(root); + qpack_free_layout(layout); + return qpack_fail(error, error_capacity, + "linearLayers[%zu] is not boolean", i); + } + layout->linear_layers[i] = (unsigned char)value->boolean; + } + + for (size_t i = 0; i < layout->section_count; i++) { + jval *object = sections->kids[i]; + jval *name = json_get(object, "name"); + jval *dtype = json_get(object, "dtype"); + jval *shape = json_get(object, "shape"); + uint64_t offset, size; + if (!object || object->t != J_OBJ || + qpack_object_has_duplicate_keys(object) || + !name || name->t != J_STR || !name->str[0] || + !dtype || dtype->t != J_STR || !dtype->str[0] || + !shape || shape->t != J_ARR || shape->len <= 0 || + shape->len > (int)COLI_QPACK_MAX_RANK || + qpack_json_u64(json_get(object, "offset"), &offset) || + qpack_json_u64(json_get(object, "size"), &size) || !size || + offset > expert_stride || size > expert_stride - offset) { + json_free(root); + qpack_free_layout(layout); + return qpack_fail(error, error_capacity, + "section %zu is invalid", i); + } + + ColiQpackSection *section = &layout->sections[i]; + section->name = qpack_copy_string(name->str); + section->dtype = qpack_copy_string(dtype->str); + section->rank = (size_t)shape->len; + section->shape = (size_t *)calloc(section->rank, sizeof(size_t)); + section->offset = offset; + section->size = size; + if (!section->name || !section->dtype || !section->shape) { + json_free(root); + qpack_free_layout(layout); + return qpack_fail(error, error_capacity, + "out of memory parsing section %zu", i); + } + + size_t element_size; + uint64_t element_count = 1; + if (qpack_dtype_size(section->dtype, &element_size)) { + int result = qpack_fail(error, error_capacity, + "section %s has unsupported dtype %s", + section->name, section->dtype); + json_free(root); + qpack_free_layout(layout); + return result; + } + for (size_t d = 0; d < section->rank; d++) { + uint64_t dimension; + if (qpack_json_u64(shape->kids[d], &dimension) || !dimension || + dimension > SIZE_MAX || + element_count > UINT64_MAX / dimension) { + int result = qpack_fail(error, error_capacity, + "section %s has invalid shape", + section->name); + json_free(root); + qpack_free_layout(layout); + return result; + } + section->shape[d] = (size_t)dimension; + element_count *= dimension; + } + if (element_count > UINT64_MAX / element_size || + element_count * element_size != section->size) { + int result = qpack_fail( + error, error_capacity, + "section %s size does not match dtype and shape", section->name); + json_free(root); + qpack_free_layout(layout); + return result; + } + + for (size_t prior = 0; prior < i; prior++) { + const ColiQpackSection *other = &layout->sections[prior]; + if (strcmp(section->name, other->name) == 0) { + int result = qpack_fail(error, error_capacity, + "duplicate section name %s", section->name); + json_free(root); + qpack_free_layout(layout); + return result; + } + uint64_t end = section->offset + section->size; + uint64_t other_end = other->offset + other->size; + if (section->offset < other_end && other->offset < end) { + int result = qpack_fail(error, error_capacity, + "sections %s and %s overlap", + section->name, other->name); + json_free(root); + qpack_free_layout(layout); + return result; + } + } + } + + json_free(root); + return 0; +} + +static int qpack_parse_manifest(const char *text, + int *quant_bits, int *quant_group_size, + jval **files_out, jval **root_out, + char *error, size_t error_capacity) { + jval *root = json_parse_exact(text, NULL); + jval *magic = json_get(root, "magic"); + jval *model = json_get(root, "modelName"); + jval *source = json_get(root, "sourceCheckpoint"); + uint64_t version; + if (!root || root->t != J_OBJ || qpack_object_has_duplicate_keys(root) || + !magic || magic->t != J_STR || strcmp(magic->str, "QPACK") != 0 || + !model || model->t != J_STR || !model->str[0] || + !source || source->t != J_STR || !source->str[0] || + qpack_json_u64(json_get(root, "version"), &version) || + version != COLI_QPACK_MANIFEST_VERSION) { + json_free(root); + return qpack_fail(error, error_capacity, + "manifest is not a supported QPACK v1 container"); + } + jval *files = json_get(root, "files"); + if (!files || files->t != J_OBJ || files->len <= 0 || + files->len > (int)COLI_QPACK_MAX_FILES || + qpack_object_has_duplicate_keys(files)) { + json_free(root); + return qpack_fail(error, error_capacity, + "manifest files table is invalid"); + } + for (int i = 0; i < files->len; i++) { + uint64_t ignored; + if (!files->keys[i][0] || qpack_json_u64(files->kids[i], &ignored)) { + json_free(root); + return qpack_fail(error, error_capacity, + "manifest file size is invalid"); + } + } + + *quant_bits = 0; + *quant_group_size = 0; + jval *bits = json_get(root, "quantBits"); + jval *group = json_get(root, "quantGroupSize"); + uint64_t number; + if (bits && bits->t != J_NULL) { + if (qpack_json_u64(bits, &number) || number > INT_MAX) { + json_free(root); + return qpack_fail(error, error_capacity, + "manifest quantBits is invalid"); + } + *quant_bits = (int)number; + } + if (group && group->t != J_NULL) { + if (qpack_json_u64(group, &number) || number > INT_MAX) { + json_free(root); + return qpack_fail(error, error_capacity, + "manifest quantGroupSize is invalid"); + } + *quant_group_size = (int)number; + } + *files_out = files; + *root_out = root; + return 0; +} + +static int qpack_verify_manifest_size(jval *files, const char *relative, + uint64_t expected_size, + char *error, size_t error_capacity) { + uint64_t manifest_size; + jval *entry = json_get(files, relative); + if (!entry || qpack_json_u64(entry, &manifest_size)) + return qpack_fail(error, error_capacity, + "manifest has no valid size for %s", relative); + if (manifest_size != expected_size) + return qpack_fail(error, error_capacity, + "manifest size mismatch for %s", relative); + return 0; +} + +int coli_qpack_open(ColiQpackReader *reader, const char *container_dir, + char *error, size_t error_capacity) { + if (!reader) + return qpack_fail(error, error_capacity, "qpack reader is null"); + memset(reader, 0, sizeof(*reader)); + if (!container_dir || !container_dir[0]) + return qpack_fail(error, error_capacity, "qpack path is empty"); + + char *manifest_path = qpack_join_path(container_dir, "manifest.json"); + reader->packed_dir = qpack_join_path(container_dir, "packed_experts"); + char *layout_path = reader->packed_dir ? + qpack_join_path(reader->packed_dir, "layout.json") : NULL; + if (!manifest_path || !reader->packed_dir || !layout_path) { + free(manifest_path); + free(layout_path); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "out of memory building qpack paths"); + } + + size_t layout_size; + char *manifest_text = qpack_read_text( + manifest_path, NULL, error, error_capacity); + char *layout_text = qpack_read_text( + layout_path, &layout_size, error, error_capacity); + free(manifest_path); + if (!manifest_text || !layout_text) { + free(manifest_text); + free(layout_text); + free(layout_path); + coli_qpack_close(reader); + return -1; + } + + jval *manifest_root = NULL, *files = NULL; + if (qpack_parse_manifest(manifest_text, &reader->quant_bits, + &reader->quant_group_size, &files, + &manifest_root, error, error_capacity) || + qpack_parse_layout(&reader->layout, layout_text, + error, error_capacity)) { + free(manifest_text); + free(layout_text); + free(layout_path); + json_free(manifest_root); + coli_qpack_close(reader); + return -1; + } + free(manifest_text); + free(layout_text); + + if (qpack_verify_manifest_size(files, "packed_experts/layout.json", + layout_size, error, error_capacity)) { + free(layout_path); + json_free(manifest_root); + coli_qpack_close(reader); + return -1; + } + free(layout_path); + + if (reader->layout.expert_count > + (uint64_t)INT64_MAX / reader->layout.expert_stride) { + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "expert layer size overflows off_t"); + } + uint64_t layer_size = (uint64_t)reader->layout.expert_count * + reader->layout.expert_stride; + reader->layer_fds = (int *)malloc(reader->layout.layer_count * sizeof(int)); + if (!reader->layer_fds) { + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "out of memory allocating layer descriptors"); + } + for (size_t layer = 0; layer < reader->layout.layer_count; layer++) + reader->layer_fds[layer] = -1; + + for (size_t layer = 0; layer < reader->layout.layer_count; layer++) { + char name[64], relative[96]; + int name_count = snprintf(name, sizeof(name), "layer_%02zu.bin", layer); + int relative_count = snprintf(relative, sizeof(relative), + "packed_experts/%s", name); + char *path = qpack_join_path(reader->packed_dir, name); + if (name_count < 0 || (size_t)name_count >= sizeof(name) || + relative_count < 0 || (size_t)relative_count >= sizeof(relative)) { + free(path); + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "layer filename is too long"); + } + if (!path) { + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "out of memory building layer path"); + } + if (qpack_verify_manifest_size(files, relative, layer_size, + error, error_capacity)) { + free(path); + json_free(manifest_root); + coli_qpack_close(reader); + return -1; + } + int fd = open(path, COMPAT_O_RDONLY); + if (fd < 0) { + int saved = errno; + free(path); + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, "open %s: %s", + relative, strerror(saved)); + } + reader->layer_fds[layer] = fd; + struct stat st; + if (fstat(fd, &st) != 0) { + int saved = errno; + free(path); + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, "fstat %s: %s", + relative, strerror(saved)); + } + if (!S_ISREG(st.st_mode) || st.st_size < 0 || + (uint64_t)st.st_size != layer_size) { + free(path); + json_free(manifest_root); + coli_qpack_close(reader); + return qpack_fail(error, error_capacity, + "%s has unexpected size", relative); + } + free(path); + } + + json_free(manifest_root); + return 0; +} + +void coli_qpack_close(ColiQpackReader *reader) { + if (!reader) return; + if (reader->layer_fds) { + for (size_t i = 0; i < reader->layout.layer_count; i++) + if (reader->layer_fds[i] >= 0) close(reader->layer_fds[i]); + } + free(reader->layer_fds); + free(reader->packed_dir); + qpack_free_layout(&reader->layout); + memset(reader, 0, sizeof(*reader)); +} + +const ColiQpackSection * +coli_qpack_find_section(const ColiQpackReader *reader, const char *name) { + if (!reader || !name) return NULL; + for (size_t i = 0; i < reader->layout.section_count; i++) + if (strcmp(reader->layout.sections[i].name, name) == 0) + return &reader->layout.sections[i]; + return NULL; +} + +int coli_qpack_read_expert(const ColiQpackReader *reader, + size_t layer, size_t expert, + void *buffer, size_t buffer_size, + char *error, size_t error_capacity) { + if (!reader || !buffer || !reader->layer_fds) + return qpack_fail(error, error_capacity, "qpack reader is not open"); + if (layer >= reader->layout.layer_count || + expert >= reader->layout.expert_count) + return qpack_fail(error, error_capacity, + "expert coordinate is out of range"); + if (buffer_size < reader->layout.expert_stride) + return qpack_fail(error, error_capacity, + "expert destination buffer is too small"); + + off_t offset = (off_t)((uint64_t)expert * reader->layout.expert_stride); + ssize_t count; + do { + count = pread(reader->layer_fds[layer], buffer, + reader->layout.expert_stride, offset); + } while (count < 0 && errno == EINTR); + if (count < 0) + return qpack_fail(error, error_capacity, + "pread layer %zu expert %zu: %s", + layer, expert, strerror(errno)); + if ((size_t)count != reader->layout.expert_stride) + return qpack_fail(error, error_capacity, + "short read for layer %zu expert %zu", layer, expert); + return 0; +} + +static const ColiQpackSection * +qpack_projection_section(const ColiQpackReader *reader, + const char *projection, const char *suffix) { + char name[256]; + int count = snprintf(name, sizeof(name), "%s.%s", projection, suffix); + if (count < 0 || (size_t)count >= sizeof(name)) return NULL; + return coli_qpack_find_section(reader, name); +} + +static int qpack_scalar_format(const char *dtype, + ColiAffineScalarFormat *format) { + if (strcmp(dtype, "F32") == 0) *format = COLI_AFFINE_SCALAR_F32; + else if (strcmp(dtype, "F16") == 0) *format = COLI_AFFINE_SCALAR_F16; + else if (strcmp(dtype, "BF16") == 0) *format = COLI_AFFINE_SCALAR_BF16; + else return -1; + return 0; +} + +int coli_qpack_affine_view(const ColiQpackReader *reader, + const void *expert_blob, size_t blob_size, + const char *projection, + ColiAffineQuantizedView *view, + char *error, size_t error_capacity) { + if (!reader || !expert_blob || !projection || !projection[0] || !view) + return qpack_fail(error, error_capacity, + "affine descriptor received a null pointer"); + memset(view, 0, sizeof(*view)); + if (blob_size < reader->layout.expert_stride) + return qpack_fail(error, error_capacity, + "expert blob is shorter than expertStride"); + if ((reader->quant_bits != 4 && reader->quant_bits != 8) || + reader->quant_group_size <= 0) + return qpack_fail(error, error_capacity, + "qpack affine quantization metadata is unsupported"); + + const ColiQpackSection *weight = qpack_projection_section( + reader, projection, "weight"); + const ColiQpackSection *scales = qpack_projection_section( + reader, projection, "scales"); + const ColiQpackSection *biases = qpack_projection_section( + reader, projection, "biases"); + if (!weight || !scales || !biases) + return qpack_fail(error, error_capacity, + "projection is missing weight, scales, or biases"); + if (strcmp(weight->dtype, "U32") != 0 || weight->rank < 2) + return qpack_fail(error, error_capacity, + "affine weight must be a rank-2+ U32 tensor"); + if (strcmp(scales->dtype, biases->dtype) != 0 || + scales->rank != weight->rank || biases->rank != weight->rank) + return qpack_fail(error, error_capacity, + "affine scale and bias metadata is incompatible"); + + ColiAffineScalarFormat scalar_format; + if (qpack_scalar_format(scales->dtype, &scalar_format)) + return qpack_fail(error, error_capacity, + "affine scales must be F32, F16, or BF16"); + + size_t output_dim = 1; + for (size_t i = 0; i + 1 < weight->rank; i++) { + if (output_dim > SIZE_MAX / weight->shape[i]) + return qpack_fail(error, error_capacity, + "affine output dimension overflows"); + output_dim *= weight->shape[i]; + if (scales->shape[i] != weight->shape[i] || + biases->shape[i] != weight->shape[i]) + return qpack_fail(error, error_capacity, + "affine scale or bias shape does not match weights"); + } + const unsigned per_word = 32u / (unsigned)reader->quant_bits; + if (weight->shape[weight->rank - 1] > SIZE_MAX / per_word) + return qpack_fail(error, error_capacity, + "affine input dimension overflows"); + size_t input_dim = weight->shape[weight->rank - 1] * per_word; + size_t group_size = (size_t)reader->quant_group_size; + if (!output_dim || !input_dim || input_dim % group_size != 0) + return qpack_fail(error, error_capacity, + "affine dimensions do not match quantGroupSize"); + size_t groups = input_dim / group_size; + if (scales->shape[scales->rank - 1] != groups || + biases->shape[biases->rank - 1] != groups) + return qpack_fail(error, error_capacity, + "affine scale or bias shape does not match weights"); + + const unsigned char *blob = (const unsigned char *)expert_blob; + view->weights = blob + weight->offset; + view->scales = blob + scales->offset; + view->biases = blob + biases->offset; + view->weight_bytes = (size_t)weight->size; + view->scale_bytes = (size_t)scales->size; + view->bias_bytes = (size_t)biases->size; + view->output_dim = output_dim; + view->input_dim = input_dim; + view->group_size = group_size; + view->format = reader->quant_bits == 4 ? + COLI_AFFINE_MLX_Q4 : COLI_AFFINE_MLX_Q8; + view->scalar_format = scalar_format; + ColiAffineStatus status = coli_affine_validate(view); + if (status != COLI_AFFINE_OK) { + memset(view, 0, sizeof(*view)); + return qpack_fail(error, error_capacity, + "affine descriptor is invalid: %s", + coli_affine_status_string(status)); + } + return 0; +} diff --git a/c/qpack.h b/c/qpack.h new file mode 100644 index 000000000..648c2b854 --- /dev/null +++ b/c/qpack.h @@ -0,0 +1,78 @@ +/* qpack.h -- strict reader for Swiftlet qpack v1 expert containers. */ +#ifndef COLI_QPACK_H +#define COLI_QPACK_H + +#include +#include + +#include "affine_quant.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define COLI_QPACK_MANIFEST_VERSION 1 +#define COLI_QPACK_PAGE_ALIGNMENT 16384u + +typedef struct { + char *name; + char *dtype; + size_t *shape; + size_t rank; + uint64_t offset; + uint64_t size; +} ColiQpackSection; + +typedef struct { + size_t expert_count; + size_t layer_count; + size_t expert_stride; + ColiQpackSection *sections; + size_t section_count; + unsigned char *linear_layers; +} ColiQpackLayout; + +typedef struct { + ColiQpackLayout layout; + int quant_bits; + int quant_group_size; + char *packed_dir; + int *layer_fds; +} ColiQpackReader; + +/* Opens and validates manifest.json, packed_experts/layout.json, and every + * fixed-stride layer file. The reader owns all returned metadata until close. + * Do not reopen a live reader. Close is safe after either success or failure. + * + * Every layer descriptor is opened and fstat-validated before success. The + * resulting reader is immutable and concurrent expert reads are safe because + * they use explicit-offset pread. Close must not race with a read. + */ +int coli_qpack_open(ColiQpackReader *reader, const char *container_dir, + char *error, size_t error_capacity); +void coli_qpack_close(ColiQpackReader *reader); + +const ColiQpackSection * +coli_qpack_find_section(const ColiQpackReader *reader, const char *name); + +/* Reads exactly one expert_stride-byte blob with one pread attempt. */ +int coli_qpack_read_expert(const ColiQpackReader *reader, + size_t layer, size_t expert, + void *buffer, size_t buffer_size, + char *error, size_t error_capacity); + +/* Resolves projection.{weight,scales,biases} inside a caller-owned expert + * blob and constructs Colibri's checked MLX affine descriptor. The view and + * its pointers remain valid only as long as the blob remains valid. + */ +int coli_qpack_affine_view(const ColiQpackReader *reader, + const void *expert_blob, size_t blob_size, + const char *projection, + ColiAffineQuantizedView *view, + char *error, size_t error_capacity); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/c/qpack_inspect.c b/c/qpack_inspect.c new file mode 100644 index 000000000..01e78a08e --- /dev/null +++ b/c/qpack_inspect.c @@ -0,0 +1,47 @@ +#include "qpack.h" + +#include +#include + +int main(int argc, char **argv) { + if (argc != 2) { + fprintf(stderr, "usage: %s MODEL.qpack\n", argv[0]); + return 2; + } + + char error[512] = {0}; + ColiQpackReader reader; + if (coli_qpack_open(&reader, argv[1], error, sizeof(error))) { + fprintf(stderr, "qpack validation failed: %s\n", error); + return 1; + } + + size_t linear_layers = 0; + for (size_t layer = 0; layer < reader.layout.layer_count; layer++) + linear_layers += reader.layout.linear_layers[layer] != 0; + + printf("QPACK v1: %zu layers (%zu linear), %zu experts/layer, " + "%zu-byte stride\n", + reader.layout.layer_count, linear_layers, + reader.layout.expert_count, reader.layout.expert_stride); + if ((reader.quant_bits == 4 || reader.quant_bits == 8) && + reader.quant_group_size > 0) { + printf("quantization: %d-bit affine, group %d\n", + reader.quant_bits, reader.quant_group_size); + } else { + printf("quantization: unspecified or unsupported by affine adapter\n"); + } + + for (size_t i = 0; i < reader.layout.section_count; i++) { + const ColiQpackSection *section = &reader.layout.sections[i]; + printf(" %-24s %-5s offset=%" PRIu64 " size=%" PRIu64 " shape=[", + section->name, section->dtype, + section->offset, section->size); + for (size_t dimension = 0; dimension < section->rank; dimension++) + printf("%s%zu", dimension ? "," : "", section->shape[dimension]); + puts("]"); + } + + coli_qpack_close(&reader); + return 0; +} diff --git a/c/qwen36.c b/c/qwen36.c index 6fe13bd93..871bdcf37 100644 --- a/c/qwen36.c +++ b/c/qwen36.c @@ -62,6 +62,7 @@ static int qwen36_max_ctx(void) { #include "st.h" #include "json.h" /* tokenizer.json parsing (reuse minimal parser) */ #include "qwen36_tier.h" /* optional transparent Vulkan compute backend for MoE experts */ +#include "affine_quant.h" /* MLX affine dense triples in model.safetensors */ #ifdef COLI_SEGMENT_ADAPTER #include "segment_runtime.h" #include "segment_adapters.h" @@ -294,11 +295,29 @@ static void load_tokenizer(const char *path){ jval *merges = json_get(model, "merges"); if (merges && merges->t==J_ARR){ for (int r=0;rlen;r++){ - const char *e = merges->kids[r]->str; if(!e) continue; - const char *sp = strchr(e, ' '); if(!sp) continue; - int la=(int)(sp-e), lb=(int)strlen(sp+1); + /* Two on-disk spellings for one merge table: legacy tokenizer.json + * writes "a b" strings, tokenizers >= 0.20 (transformers 4.45+, + * the Qwen3.6 checkpoints included) writes ["a","b"] pairs. The + * string-only reader SILENTLY indexed zero merges from the pair + * form, and encode_text degraded to one token per byte-symbol -- + * 24 tokens for a 24-char prompt, real-model run -- because + * bpe_piece treats an empty merge table as "nothing to merge", + * not as an error. */ + jval *mk = merges->kids[r]; + const char *a, *b; + int la, lb; + if (mk && mk->t==J_STR && mk->str){ + const char *sp = strchr(mk->str, ' '); if(!sp) continue; + a = mk->str; la = (int)(sp - mk->str); + b = sp + 1; lb = (int)strlen(b); + } else if (mk && mk->t==J_ARR && mk->len==2 && + mk->kids[0] && mk->kids[0]->t==J_STR && mk->kids[0]->str && + mk->kids[1] && mk->kids[1]->t==J_STR && mk->kids[1]->str){ + a = mk->kids[0]->str; la = (int)strlen(a); + b = mk->kids[1]->str; lb = (int)strlen(b); + } else continue; char *key=malloc(la+1+lb+1); - memcpy(key,e,la); key[la]=0x1F; memcpy(key+la+1,sp+1,lb); key[la+1+lb]=0; + memcpy(key,a,la); key[la]=0x1F; memcpy(key+la+1,b,lb); key[la+1+lb]=0; smap_put(&g_merge, key, r); } } @@ -578,6 +597,15 @@ typedef struct { int n_group, topk_group; float theta, eps, partial_rotary_factor; int norm_topk, has_qk_norm, has_bias, attn_output_gate; + /* RMSNorm weight dialect for the tensors rmsnorm_row touches (input/post + * layernorms, q/k norms, final norm). 1 = HF Qwen3.6 zero-centered + * storage, apply (1 + w) -- the convention every converted snapshot uses + * and the engine has always assumed. 0 = full-gamma storage (MLX-derived + * containers: mlx-lm materialises the +1 into the weights at conversion), + * unshifted back to zero-centered at LOAD so the forward pass stays one + * convention. The DeltaNet gated norm is full-gamma in BOTH dialects and + * is untouched by this flag. */ + int zero_centered_norms; uint8_t *is_attn; /* [n_layers] 1 if Gated Attention layer, 0 if DeltaNet */ /* Gated DeltaNet (linear_attention) dims, read from qwen36_meta.json. */ int dn_vheads, dn_kheads, dn_kdim, dn_vdim, dn_convk, dn_conv_dim; @@ -1073,7 +1101,7 @@ static void load_cfg(Cfg *c, const char *snap) { c->partial_rotary_factor = 0.25f; c->n_experts = 256; c->topk = 8; c->inter = 512; c->shared_inter = 512; c->n_group = 1; c->topk_group = 1; c->norm_topk = 1; c->has_qk_norm = 1; c->has_bias = 0; - c->attn_output_gate = 1; c->n_active = 0; + c->attn_output_gate = 1; c->n_active = 0; c->zero_centered_norms = 1; if (c->n_layers <= 0 || c->n_layers > 512) { fprintf(stderr, "load_cfg: n_layers=%d out of range 1..512\n", c->n_layers); exit(1); } c->is_attn = calloc((size_t)c->n_layers, sizeof(uint8_t)); for (int i = 0; i < c->n_layers; i++) c->is_attn[i] = (i % 4 == 3) ? 1 : 0; @@ -1186,6 +1214,7 @@ static void load_meta(Cfg *c, const char *snap) { if((v=json_get(r,"rms_eps"))&&v->t==J_NUM) c->eps=(float)v->num; if((v=json_get(r,"attn_output_gate"))&&v->t==J_BOOL) c->attn_output_gate=v->boolean; if((v=json_get(r,"norm_topk_prob"))&&v->t==J_BOOL) c->norm_topk=v->boolean; + if((v=json_get(r,"zero_centered_norms"))&&v->t==J_BOOL) c->zero_centered_norms=v->boolean; if((v=json_get(r,"has_bias"))&&v->t==J_BOOL) c->has_bias=v->boolean; if((v=json_get(r,"has_qk_norm"))&&v->t==J_BOOL) c->has_qk_norm=v->boolean; /* derive rotary_dim from head_dim * partial_rotary_factor (HF formula) */ @@ -1219,23 +1248,148 @@ static void load_meta(Cfg *c, const char *snap) { c->dn_vheads, c->dn_kheads, c->dn_kdim, c->dn_vdim, c->dn_convk, c->dn_conv_dim); } +/* Multimodal Qwen3.5/3.6 checkpoints (a Swiftlet qpack container's + * model.safetensors included) store the text stack under a `language_model.` + * prefix; the engine and the converter both speak unprefixed names. Resolve + * the plain name first so converted snapshots are untouched, then the + * prefixed one; when neither exists return the CANONICAL name so the caller's + * refusal names the tensor the engine actually wanted. */ +#define QW_DENSE_NAME_MAX 288 /* 256-byte call-site buffers + the prefix */ +static const char *dense_resolve(Model *m, const char *name, + char *buf, size_t cap) { + if (st_has(&m->S, name)) return name; + snprintf(buf, cap, "language_model.%s", name); + if (st_has(&m->S, buf)) return buf; + return name; +} +static int dense_has(Model *m, const char *name) { + char rn[QW_DENSE_NAME_MAX]; + return st_has(&m->S, dense_resolve(m, name, rn, sizeof rn)); +} + +/* Dense tensor stored as an MLX affine triple (packed U32 `.weight` + sibling + * `.scales`/`.biases`, the layout Swiftlet writes into a qpack container's + * model.safetensors): expand to f32 rows through the checked affine contract. + * Every dimension is anchored to `want`, the CONFIG-implied element count the + * forward pass will index with -- the same discipline as load_t_n below, so a + * container whose packed geometry disagrees with config.json is a refusal, + * never a plausible heap OOB. bits (Q4/Q8) and the group size are DERIVED + * from the shapes (logical input over packed words, logical input over scale + * groups) rather than parsed from config quantization overrides: the file's + * own byte layout is what the expansion must agree with, and the affine + * validator re-checks the derived geometry against every buffer length. */ +static float *load_t_affine(Model *m, const char *wname, int64_t want) { + st_tensor *w = st_find(&m->S, wname); + if (!w) { fprintf(stderr, "missing %s\n", wname); exit(1); } + size_t wlen = strlen(wname); + char sname[QW_DENSE_NAME_MAX], bname[QW_DENSE_NAME_MAX]; + if (wlen < 7 || strcmp(wname + wlen - 7, ".weight") != 0 || + wlen - 7 + sizeof(".scales") > sizeof(sname)) { + fprintf(stderr, "%s: U32 tensor is not a `.weight` with room for " + "`.scales`/`.biases` siblings -- refusing\n", wname); exit(1); + } + snprintf(sname, sizeof(sname), "%.*s.scales", (int)(wlen - 7), wname); + snprintf(bname, sizeof(bname), "%.*s.biases", (int)(wlen - 7), wname); + st_tensor *s = st_find(&m->S, sname), *b = st_find(&m->S, bname); + if (!s || !b) { + fprintf(stderr, "%s: affine `.scales`/`.biases` siblings are missing " + "-- refusing\n", wname); exit(1); + } + ColiAffineScalarFormat sf; + switch (s->dtype) { + case 0: sf = COLI_AFFINE_SCALAR_BF16; break; + case 1: sf = COLI_AFFINE_SCALAR_F16; break; + case 2: sf = COLI_AFFINE_SCALAR_F32; break; + default: + fprintf(stderr, "%s: scales dtype %s is not BF16/F16/F32 -- " + "refusing\n", sname, st_dtype_name(s->dtype)); exit(1); + } + if (b->dtype != s->dtype || b->numel != s->numel) { + fprintf(stderr, "%s: biases dtype/numel disagree with scales -- " + "refusing\n", bname); exit(1); + } + int64_t packed = w->rank >= 2 ? w->shape[w->rank - 1] : 0; + int64_t rows = packed > 0 ? w->numel / packed : 0; + if (want <= 0 || rows <= 0 || w->numel != rows * packed || + want % rows != 0) { + fprintf(stderr, "%s: packed shape does not divide the config-implied " + "%lld elements -- refusing\n", wname, (long long)want); exit(1); + } + int64_t in = want / rows; + int64_t per_word = packed > 0 && in % packed == 0 ? in / packed : 0; + if (per_word != 8 && per_word != 4) { + fprintf(stderr, "%s: %lld packed words for %lld logical columns is " + "neither Q4 nor Q8 -- refusing\n", + wname, (long long)packed, (long long)in); exit(1); + } + int64_t groups = s->numel % rows == 0 ? s->numel / rows : 0; + int64_t gs = groups > 0 && in % groups == 0 ? in / groups : 0; + if (gs <= 0) { + fprintf(stderr, "%s: %lld scale groups do not tile %lld logical " + "columns -- refusing\n", + sname, (long long)groups, (long long)in); exit(1); + } + void *wraw = malloc((size_t)w->nbytes); + void *sraw = malloc((size_t)s->nbytes); + void *braw = malloc((size_t)b->nbytes); + if (!wraw || !sraw || !braw) { fprintf(stderr, "OOM reading %s\n", wname); exit(1); } + st_read_raw(&m->S, wname, wraw, 1); + st_read_raw(&m->S, sname, sraw, 1); + st_read_raw(&m->S, bname, braw, 1); + ColiAffineQuantizedView view = { + wraw, sraw, braw, + (size_t)w->nbytes, (size_t)s->nbytes, (size_t)b->nbytes, + (size_t)rows, (size_t)in, (size_t)gs, + per_word == 8 ? COLI_AFFINE_MLX_Q4 : COLI_AFFINE_MLX_Q8, sf + }; + float *p = falloc(want); + ColiAffineStatus status = coli_affine_dequant_ref(&view, p); + if (status != COLI_AFFINE_OK) { + fprintf(stderr, "%s: affine expansion refused (%s)\n", + wname, coli_affine_status_string(status)); exit(1); + } + free(wraw); free(sraw); free(braw); + return p; +} + /* `want` is the element count the forward pass will index with. The container * is a file, not an invariant: this used to allocate whatever st_numel reported * while every read afterwards used CONFIG dims, so a short tensor was a plain * heap OOB read (embed is indexed as m->embed + ids[s]*D). The expert path * already refuses a wrong size; this is the same discipline for the dense set. */ static float *load_t_n(Model *m, const char *name, int64_t want) { - int64_t n = st_numel(&m->S, name); - if (n < 0) { fprintf(stderr, "missing %s\n", name); exit(1); } + char rn[QW_DENSE_NAME_MAX]; + const char *nm = dense_resolve(m, name, rn, sizeof rn); + st_tensor *t = st_find(&m->S, nm); + if (!t) { fprintf(stderr, "missing %s\n", name); exit(1); } + if (t->dtype == 7) return load_t_affine(m, nm, want); + int64_t n = t->numel; if (want > 0 && n != want) { fprintf(stderr, "%s: %lld elements, config implies %lld -- refusing\n", - name, (long long)n, (long long)want); exit(1); + nm, (long long)n, (long long)want); exit(1); } float *p = falloc(n); - st_read_f32(&m->S, name, p, 0); + st_read_f32(&m->S, nm, p, 0); return p; } +/* Loader for the RMSNorm weights rmsnorm_row applies as (1 + w). HF Qwen3.6 + * stores them zero-centered and the forward pass is written for that; an + * MLX-derived container stores FULL gamma (mlx-lm materialises the +1 at + * conversion -- measured on the production qpack container: input_layernorm + * mean 1.03, q_norm mean 1.33, where the zero-centered forms centre on 0). + * Feeding full gamma through (1 + w) doubles every normalised activation and + * the model degenerates to noise, so the shift is undone HERE, at load, and + * the forward pass keeps exactly one convention. The DeltaNet gated norm is + * full gamma in both dialects (its forward multiplies plain w) and must NOT + * come through this loader. */ +static float *load_norm_n(Model *m, const char *name, int64_t want) { + float *w = load_t_n(m, name, want); + if (!m->c.zero_centered_norms) + for (int64_t i = 0; i < want; i++) w[i] -= 1.0f; + return w; +} + static void model_init_range(Model *m, const char *snap, int cap, int bits, int layer_begin, int layer_end, int load_boundaries, int allocate_state) { @@ -1264,7 +1418,7 @@ static void model_init_range(Model *m, const char *snap, int cap, int bits, if (load_boundaries) { m->embed = load_t_n(m, "model.embed_tokens.weight", (int64_t)c->vocab * c->hidden); m->lm_head = load_t_n(m, "lm_head.weight", (int64_t)c->vocab * c->hidden); - m->final_norm = load_t_n(m, "model.norm.weight", c->hidden); + m->final_norm = load_norm_n(m, "model.norm.weight", c->hidden); } m->L = calloc((size_t)c->n_layers, sizeof(Layer)); /* Phase 2: the converter stores EVERY layer (Gated-Attention + Gated DeltaNet) @@ -1276,22 +1430,24 @@ static void model_init_range(Model *m, const char *snap, int cap, int bits, for (int i = layer_begin; i < layer_end; i++) { int ai = m->active_of[i]; /* == i for Phase 2 */ Layer *l = &m->L[i]; - /* input/post layernorms + MoE exist for every layer */ - #define LD(field, suffix, want) snprintf(nm,sizeof(nm),"model.layers.%d." suffix,ai); l->field = load_t_n(m,nm,(want)) + /* input/post layernorms + MoE exist for every layer; the layernorms go + * through load_norm_n (rmsnorm_row weights), the router does not. */ + #define LD(field, suffix, want) snprintf(nm,sizeof(nm),"model.layers.%d." suffix,ai); l->field = load_norm_n(m,nm,(want)) LD(in_ln, "input_layernorm.weight", c->hidden); LD(post_ln,"post_attention_layernorm.weight", c->hidden); - LD(gate, "mlp.gate.weight", (int64_t)c->n_experts * c->hidden); #undef LD + snprintf(nm,sizeof(nm),"model.layers.%d.mlp.gate.weight",ai); + l->gate = load_t_n(m, nm, (int64_t)c->n_experts * c->hidden); /* q/k norms are per-head [head_dim]; only on attention layers, load if present */ if (c->has_qk_norm) { snprintf(nm,sizeof(nm),"model.layers.%d.self_attn.q_norm.weight", ai); - l->qn = st_has(&m->S, nm) ? load_t_n(m, nm, c->head_dim) : NULL; + l->qn = dense_has(m, nm) ? load_norm_n(m, nm, c->head_dim) : NULL; snprintf(nm,sizeof(nm),"model.layers.%d.self_attn.k_norm.weight", ai); - l->kn = st_has(&m->S, nm) ? load_t_n(m, nm, c->head_dim) : NULL; + l->kn = dense_has(m, nm) ? load_norm_n(m, nm, c->head_dim) : NULL; } else { l->qn = NULL; l->kn = NULL; } /* router correction bias (optional) */ snprintf(nm,sizeof(nm),"model.layers.%d.mlp.gate.e_score_correction_bias", ai); - if (st_has(&m->S, nm)) { l->gate_bias = falloc(c->n_experts); st_read_f32(&m->S, nm, l->gate_bias, 0); } + if (dense_has(m, nm)) { l->gate_bias = load_t_n(m, nm, c->n_experts); } else l->gate_bias = NULL; /* shared expert (dense f32) */ #define LD2(field, suffix, want) snprintf(nm,sizeof(nm),"model.layers.%d.mlp.shared_expert." suffix,ai); l->field = load_t_n(m,nm,(want)) @@ -1301,7 +1457,7 @@ static void model_init_range(Model *m, const char *snap, int cap, int bits, #undef LD2 /* shared_expert_gate: Linear(hidden -> 1), sigmoid-gated shared expert */ snprintf(nm,sizeof(nm),"model.layers.%d.mlp.shared_expert_gate.weight", ai); - l->sh_gate = st_has(&m->S, nm) ? load_t_n(m, nm, c->hidden) : NULL; + l->sh_gate = dense_has(m, nm) ? load_t_n(m, nm, c->hidden) : NULL; if (c->is_attn[i]) { /* Gated Attention (full_attention) layer */ #define LD3(field, suffix, want) snprintf(nm,sizeof(nm),"model.layers.%d.self_attn." suffix,ai); l->field = load_t_n(m,nm,(want)) @@ -3369,7 +3525,7 @@ static int qwen36_edge_engine_open( engine->model.lm_head = load_t_n( &engine->model, "lm_head.weight", (int64_t)config->vocab * config->hidden); - engine->model.final_norm = load_t_n( + engine->model.final_norm = load_norm_n( &engine->model, "model.norm.weight", config->hidden); engine->model.quant_bits = container_layer_is_int4(&engine->model, 0) ? 4 : 8; char tokenizer_path[4096]; diff --git a/c/st.h b/c/st.h index 76987ada9..4b2f501af 100644 --- a/c/st.h +++ b/c/st.h @@ -33,7 +33,7 @@ typedef struct { int fd; int64_t off; /* offset assoluto del dato dentro al file */ int64_t nbytes; - int dtype; /* 0=BF16 1=F16 2=F32 3=U8/I8 4=F8_E4M3 5=F8_E8M0 6=I64 */ + int dtype; /* 0=BF16 1=F16 2=F32 3=U8/I8 4=F8_E4M3 5=F8_E8M0 6=I64 7=U32/I32 */ int64_t numel; int rank; int64_t shape[ST_MAX_RANK]; @@ -95,6 +95,13 @@ static int st_dtype_code(const char *s) { !strcmp(s, "float8_e4m3fn")) return 4; if (!strcmp(s, "F8_E8M0") || !strcmp(s, "F8_E8M0FNU")) return 5; if (!strcmp(s, "I64") || !strcmp(s, "U64")) return 6; + /* MLX affine-quantized checkpoints (Swiftlet qpack containers included) + * store packed weights as U32 words with BF16/F16 scales+biases siblings. + * INDEXED here like the fp8 additions above so st_init reaches the rest of + * the file; read through st_read_raw by callers that validated the packed + * geometry (qwen36's dense affine loader). The float readers refuse the + * code by value (dtype >= 3 in st_read_f32), same as I64. */ + if (!strcmp(s, "U32") || !strcmp(s, "I32")) return 7; fprintf(stderr, "unsupported dtype: %s\n", s); exit(1); } @@ -106,6 +113,7 @@ static inline int st_dtype_esz(int dtype) { case 2: return 4; /* F32 */ case 3: case 4: case 5: return 1; /* U8/I8, F8_E4M3, F8_E8M0 */ case 6: return 8; /* I64/U64 */ + case 7: return 4; /* U32/I32 (MLX affine packed words) */ default: return 2; /* BF16, F16 */ } } @@ -115,7 +123,7 @@ static inline const char *st_dtype_name(int dtype) { switch (dtype) { case 0: return "BF16"; case 1: return "F16"; case 2: return "F32"; case 3: return "U8/I8"; case 4: return "F8_E4M3"; case 5: return "F8_E8M0"; - case 6: return "I64"; default: return "?"; + case 6: return "I64"; case 7: return "U32/I32"; default: return "?"; } } diff --git a/c/tests/test_affine_dequant.c b/c/tests/test_affine_dequant.c new file mode 100644 index 000000000..cebda56cd --- /dev/null +++ b/c/tests/test_affine_dequant.c @@ -0,0 +1,161 @@ +/* coli_affine_dequant_ref gates: the loader-side expansion oracle must read + * the SAME bits as coli_affine_matmul_ref and refuse the same malformed + * views. + * + * 1. nibble order -- a Q4 word authored as 0x76543210 must expand to + * logical columns 0,1,...,7 in that order (lowest bits = lowest + * column). A most-significant-first reading produces the reversed + * sequence and fails the exact-value checks. + * 2. affine map -- expansion is scale*q + bias per group, with the + * BF16/F16/F32 scalar decodings all exercised (the values chosen are + * exactly representable in every one of the three, so the checks are + * equality, not tolerance). + * 3. Q8 arm -- per_word drops to 4 and the mask widens; a crafted + * word proves the byte order. + * 4. matmul parity -- x @ dequant(W)^T must equal matmul_ref(x, W) on a + * multi-group view, so the two readings of the packed words cannot + * drift apart. + * 5. refusals -- NULL buffers, truncated scales, and a group size + * that does not divide input_dim come back as statuses, never as + * plausible output. + */ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#include +#include +#include +#include + +#include "../affine_quant.h" + +#define CHECK(condition) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: check failed: %s\n", \ + __FILE__, __LINE__, #condition); \ + exit(1); \ + } \ +} while (0) + +static uint16_t f32_to_bf16(float value) { + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + return (uint16_t)(bits >> 16); /* the test values are bf16-exact */ +} + +/* Q4, one row of 16 logical columns in two groups of 8: word 0 counts + * 0..7 LSB-first, word 1 counts 8..15. Group 0 maps q -> 2q - 1, group 1 + * maps q -> 0.5q + 4; every expected value is exact in bf16/f16/f32. */ +static void test_q4_nibble_order_and_affine(void) { + const uint32_t words[2] = { 0x76543210u, 0xFEDCBA98u }; + const uint16_t scales_bf16[2] = { 0x4000u /* 2.0 */, 0x3F00u /* 0.5 */ }; + const uint16_t biases_bf16[2] = { 0xBF80u /* -1.0 */, 0x4080u /* 4.0 */ }; + ColiAffineQuantizedView view = { + words, scales_bf16, biases_bf16, + sizeof(words), sizeof(scales_bf16), sizeof(biases_bf16), + 1, 16, 8, COLI_AFFINE_MLX_Q4, COLI_AFFINE_SCALAR_BF16 + }; + float out[16]; + CHECK(coli_affine_dequant_ref(&view, out) == COLI_AFFINE_OK); + for (int q = 0; q < 8; q++) + CHECK(out[q] == 2.0f * (float)q - 1.0f); + for (int q = 8; q < 16; q++) + CHECK(out[q] == 0.5f * (float)q + 4.0f); +} + +/* Q8: 4 values per word, byte order proven by an asymmetric word. */ +static void test_q8_byte_order(void) { + const uint32_t words[1] = { 0xFF400201u }; /* q = 1, 2, 64, 255 */ + const float scales_f32[1] = { 3.0f }; + const float biases_f32[1] = { -2.0f }; + const float expected[4] = { 1.0f, 4.0f, 190.0f, 763.0f }; + ColiAffineQuantizedView view = { + words, scales_f32, biases_f32, + sizeof(words), sizeof(scales_f32), sizeof(biases_f32), + 1, 4, 4, COLI_AFFINE_MLX_Q8, COLI_AFFINE_SCALAR_F32 + }; + float out[4]; + CHECK(coli_affine_dequant_ref(&view, out) == COLI_AFFINE_OK); + for (int i = 0; i < 4; i++) CHECK(out[i] == expected[i]); +} + +/* F16 scalar decoding through the same path (1.5 and 0.25 are f16-exact). */ +static void test_f16_scalars(void) { + const uint32_t words[1] = { 0x00000021u }; /* q = 1, 2, 0, 0, ... */ + const uint16_t scales_f16[1] = { 0x3E00u /* 1.5 */ }; + const uint16_t biases_f16[1] = { 0x3400u /* 0.25 */ }; + ColiAffineQuantizedView view = { + words, scales_f16, biases_f16, + sizeof(words), sizeof(scales_f16), sizeof(biases_f16), + 1, 8, 8, COLI_AFFINE_MLX_Q4, COLI_AFFINE_SCALAR_F16 + }; + float out[8]; + CHECK(coli_affine_dequant_ref(&view, out) == COLI_AFFINE_OK); + CHECK(out[0] == 1.75f && out[1] == 3.25f); + for (int i = 2; i < 8; i++) CHECK(out[i] == 0.25f); +} + +/* Multi-row multi-group parity: y = matmul_ref(x, W) must equal + * x @ dequant_ref(W)^T, or the two functions read different bits. */ +static void test_matmul_parity(void) { + enum { O = 3, I = 32, GS = 16, WORDS_PER_ROW = I / 8 }; + uint32_t words[O * WORDS_PER_ROW]; + uint16_t scales[O * (I / GS)], biases[O * (I / GS)]; + float x[I], expanded[O * I], want[O], got[O]; + for (int i = 0; i < O * WORDS_PER_ROW; i++) + words[i] = 0x9E3779B9u * (uint32_t)(i + 1) + 0x7F4A7C15u; + for (int i = 0; i < O * (I / GS); i++) { + scales[i] = f32_to_bf16(0.25f + 0.125f * (float)i); + biases[i] = f32_to_bf16(-1.0f + 0.5f * (float)(i % 3)); + } + for (int i = 0; i < I; i++) x[i] = 0.0625f * (float)(i - I / 2); + ColiAffineQuantizedView view = { + words, scales, biases, + sizeof(words), sizeof(scales), sizeof(biases), + O, I, GS, COLI_AFFINE_MLX_Q4, COLI_AFFINE_SCALAR_BF16 + }; + CHECK(coli_affine_matmul_ref(want, x, 1, &view) == COLI_AFFINE_OK); + CHECK(coli_affine_dequant_ref(&view, expanded) == COLI_AFFINE_OK); + for (int o = 0; o < O; o++) { + double sum = 0.0; + for (int i = 0; i < I; i++) sum += (double)expanded[o * I + i] * x[i]; + got[o] = (float)sum; + CHECK(fabsf(got[o] - want[o]) <= 1e-4f * (1.0f + fabsf(want[o]))); + } +} + +static void test_refusals(void) { + const uint32_t words[2] = { 0, 0 }; + const uint16_t scalars[2] = { 0x3F80u, 0x3F80u }; + float out[16]; + ColiAffineQuantizedView view = { + words, scalars, scalars, + sizeof(words), sizeof(scalars), sizeof(scalars), + 1, 16, 8, COLI_AFFINE_MLX_Q4, COLI_AFFINE_SCALAR_BF16 + }; + CHECK(coli_affine_dequant_ref(NULL, out) == COLI_AFFINE_NULL); + CHECK(coli_affine_dequant_ref(&view, NULL) == COLI_AFFINE_NULL); + + ColiAffineQuantizedView truncated = view; + truncated.scale_bytes = 2; /* needs 2 groups * 2 bytes */ + CHECK(coli_affine_dequant_ref(&truncated, out) == COLI_AFFINE_TRUNCATED); + + ColiAffineQuantizedView bad_group = view; + bad_group.group_size = 12; /* 16 % 12 != 0 */ + CHECK(coli_affine_dequant_ref(&bad_group, out) == COLI_AFFINE_BAD_SHAPE); + + ColiAffineQuantizedView bad_format = view; + bad_format.format = (ColiAffineFormat)7; + CHECK(coli_affine_dequant_ref(&bad_format, out) == COLI_AFFINE_BAD_FORMAT); +} + +int main(void) { + test_q4_nibble_order_and_affine(); + test_q8_byte_order(); + test_f16_scalars(); + test_matmul_parity(); + test_refusals(); + printf("test_affine_dequant: OK\n"); + return 0; +} diff --git a/c/tests/test_affine_quant.c b/c/tests/test_affine_quant.c new file mode 100644 index 000000000..35db06011 --- /dev/null +++ b/c/tests/test_affine_quant.c @@ -0,0 +1,244 @@ +#include "../backend_metal.h" + +#include +#include +#include +#include + +enum { OUTPUT_DIM = 3, INPUT_DIM = 32, GROUP_SIZE = 16, BATCH = 2 }; + +static void store_le16(uint8_t *destination, uint16_t value) { + destination[0] = (uint8_t)value; + destination[1] = (uint8_t)(value >> 8); +} + +static void store_le32(uint8_t *destination, uint32_t value) { + destination[0] = (uint8_t)value; + destination[1] = (uint8_t)(value >> 8); + destination[2] = (uint8_t)(value >> 16); + destination[3] = (uint8_t)(value >> 24); +} + +static uint16_t f32_to_f16(float value) { + uint32_t bits; + uint32_t sign, mantissa; + int exponent; + memcpy(&bits, &value, sizeof(bits)); + sign = (bits >> 16) & 0x8000u; + exponent = (int)((bits >> 23) & 0xffu) - 127 + 15; + mantissa = bits & 0x7fffffu; + if (exponent <= 0) return (uint16_t)sign; + if (exponent >= 31) return (uint16_t)(sign | 0x7c00u); + return (uint16_t)(sign | ((uint32_t)exponent << 10) | (mantissa >> 13)); +} + +static void store_scalar(uint8_t *destination, size_t index, + ColiAffineScalarFormat format, float value) { + uint32_t bits; + if (format == COLI_AFFINE_SCALAR_F32) { + memcpy(&bits, &value, sizeof(bits)); + store_le32(destination + index * 4, bits); + } else if (format == COLI_AFFINE_SCALAR_F16) { + store_le16(destination + index * 2, f32_to_f16(value)); + } else { + memcpy(&bits, &value, sizeof(bits)); + store_le16(destination + index * 2, (uint16_t)(bits >> 16)); + } +} + +static uint32_t packed_value(size_t row, size_t column, unsigned bits) { + const uint32_t mask = (UINT32_C(1) << bits) - 1u; + return (uint32_t)(row * 17 + column * 5 + 3) & mask; +} + +static float direct_reference(const ColiAffineQuantizedView *view, + const float *input, size_t sample, size_t row) { + const unsigned bits = coli_affine_bits(view->format); + const unsigned per_word = 32u / bits; + const size_t packed_words = view->input_dim / per_word; + const size_t groups = view->input_dim / view->group_size; + const uint32_t mask = (UINT32_C(1) << bits) - 1u; + const uint8_t *weight_row = + (const uint8_t *)view->weights + row * packed_words * 4; + double result = 0.0; + + for (size_t column = 0; column < view->input_dim; column++) { + const uint32_t word = + coli_affine_load_le32(weight_row + (column / per_word) * 4); + const uint32_t quantized = + (word >> (bits * (unsigned)(column % per_word))) & mask; + const size_t group = row * groups + column / view->group_size; + const float scale = coli_affine_load_scalar( + view->scales, group, view->scalar_format); + const float bias = coli_affine_load_scalar( + view->biases, group, view->scalar_format); + const float x = input[sample * view->input_dim + column]; + result += (double)(scale * (float)quantized + bias) * (double)x; + } + return (float)result; +} + +static int run_parity(ColiAffineFormat format, + ColiAffineScalarFormat scalar_format) { + const unsigned bits = coli_affine_bits(format); + const unsigned per_word = 32u / bits; + const size_t packed_words = INPUT_DIM / per_word; + const size_t scalar_count = OUTPUT_DIM * (INPUT_DIM / GROUP_SIZE); + const size_t scalar_size = coli_affine_scalar_size(scalar_format); + uint8_t weights[OUTPUT_DIM * INPUT_DIM]; + uint8_t scales[OUTPUT_DIM * (INPUT_DIM / GROUP_SIZE) * 4]; + uint8_t biases[OUTPUT_DIM * (INPUT_DIM / GROUP_SIZE) * 4]; + float input[BATCH * INPUT_DIM], output[BATCH * OUTPUT_DIM]; + + memset(weights, 0, sizeof(weights)); + memset(scales, 0, sizeof(scales)); + memset(biases, 0, sizeof(biases)); + for (size_t row = 0; row < OUTPUT_DIM; row++) { + for (size_t word_index = 0; word_index < packed_words; word_index++) { + uint32_t word = 0; + for (unsigned lane = 0; lane < per_word; lane++) { + const size_t column = word_index * per_word + lane; + word |= packed_value(row, column, bits) << (lane * bits); + } + store_le32(weights + (row * packed_words + word_index) * 4, word); + } + } + for (size_t index = 0; index < scalar_count; index++) { + const float scale = (float)((int)(index % 5) - 2) * 0.125f; + const float bias = (float)((int)(index % 3) - 1) * 0.0625f; + store_scalar(scales, index, scalar_format, scale); + store_scalar(biases, index, scalar_format, bias); + } + for (size_t index = 0; index < BATCH * INPUT_DIM; index++) + input[index] = (float)((int)(index % 13) - 6) * 0.125f; + + ColiAffineQuantizedView view = { + weights, scales, biases, + OUTPUT_DIM * packed_words * 4, scalar_count * scalar_size, + scalar_count * scalar_size, + OUTPUT_DIM, INPUT_DIM, GROUP_SIZE, format, scalar_format + }; + ColiAffineStatus status = coli_affine_validate(&view); + if (status != COLI_AFFINE_OK) { + fprintf(stderr, "affine validation failed bits=%u scalar=%d: %s\n", + bits, (int)scalar_format, coli_affine_status_string(status)); + return 1; + } + if (!coli_metal_affine_dispatch_supported(&view, BATCH)) { + fprintf(stderr, "valid Metal dispatch rejected bits=%u scalar=%d\n", + bits, (int)scalar_format); + return 1; + } + status = coli_affine_matmul_ref(output, input, BATCH, &view); + if (status != COLI_AFFINE_OK) { + fprintf(stderr, "affine matmul failed bits=%u scalar=%d: %s\n", + bits, (int)scalar_format, coli_affine_status_string(status)); + return 1; + } + + for (size_t sample = 0; sample < BATCH; sample++) { + for (size_t row = 0; row < OUTPUT_DIM; row++) { + const float expected = direct_reference(&view, input, sample, row); + const float actual = output[sample * OUTPUT_DIM + row]; + const float tolerance = 2e-6f * fmaxf(1.0f, fabsf(expected)); + if (fabsf(actual - expected) > tolerance) { + fprintf(stderr, + "affine parity failed bits=%u scalar=%d sample=%zu row=%zu: " + "got %.9g want %.9g\n", + bits, (int)scalar_format, sample, row, actual, expected); + return 1; + } + } + } + return 0; +} + +static int test_validation(void) { + uint32_t weight = 0; + float scales = 1.0f, bias = 0.0f; + float input[8] = {0}, output = 123.0f; + ColiAffineQuantizedView view = { + &weight, &scales, &bias, sizeof(weight), sizeof(scales), sizeof(bias), + 1, 8, 8, COLI_AFFINE_MLX_Q4, COLI_AFFINE_SCALAR_F32 + }; + +#define EXPECT_STATUS(expression, expected) do { \ + ColiAffineStatus got = (expression); \ + if (got != (expected)) { \ + fprintf(stderr, "validation line %d: got %s, want %s\n", __LINE__, \ + coli_affine_status_string(got), \ + coli_affine_status_string(expected)); \ + return 1; \ + } \ + } while (0) + + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_OK); + view.biases = NULL; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_NULL); + view.biases = &bias; + view.format = (ColiAffineFormat)99; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_BAD_FORMAT); + view.format = COLI_AFFINE_MLX_Q4; + view.input_dim = 7; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_BAD_SHAPE); + view.input_dim = 8; + view.weight_bytes--; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_TRUNCATED); + view.weight_bytes++; + view.scale_bytes--; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_TRUNCATED); + view.scale_bytes++; + view.bias_bytes--; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_TRUNCATED); + view.bias_bytes++; + view.output_dim = SIZE_MAX; + EXPECT_STATUS(coli_affine_validate(&view), COLI_AFFINE_OVERFLOW); + view.output_dim = 1; + EXPECT_STATUS(coli_affine_matmul_ref(&output, input, 0, &view), + COLI_AFFINE_BAD_SHAPE); + EXPECT_STATUS(coli_affine_matmul_ref(&output, input, SIZE_MAX, &view), + COLI_AFFINE_OVERFLOW); +#undef EXPECT_STATUS + if (output != 123.0f) return 1; + + if (!coli_metal_affine_dispatch_supported(&view, 1) || + coli_metal_affine_dispatch_supported(&view, 0) || + coli_metal_affine_dispatch_supported(&view, -1)) + return 1; + view.format = (ColiAffineFormat)99; + if (coli_metal_affine_dispatch_supported(&view, 1)) return 1; + view.format = COLI_AFFINE_MLX_Q4; +#if SIZE_MAX > UINT32_MAX + view.output_dim = (size_t)UINT32_MAX + 1u; + view.weight_bytes = view.scale_bytes = view.bias_bytes = SIZE_MAX; + if (coli_affine_validate(&view) != COLI_AFFINE_OK || + coli_metal_affine_dispatch_supported(&view, 1)) + return 1; + view.output_dim = UINT32_MAX; + if (coli_affine_validate(&view) != COLI_AFFINE_OK || + coli_metal_affine_dispatch_supported(&view, 2)) + return 1; + view.output_dim = UINT32_C(65536); + view.input_dim = UINT32_C(1048576); + view.group_size = 8; + if (coli_affine_validate(&view) != COLI_AFFINE_OK || + coli_metal_affine_dispatch_supported(&view, 1)) + return 1; +#endif + return 0; +} + +int main(void) { + for (int format = COLI_AFFINE_MLX_Q4; + format <= COLI_AFFINE_MLX_Q8; format++) { + for (int scalar = COLI_AFFINE_SCALAR_F32; + scalar <= COLI_AFFINE_SCALAR_BF16; scalar++) { + if (run_parity((ColiAffineFormat)format, + (ColiAffineScalarFormat)scalar)) + return 1; + } + } + if (test_validation()) return 1; + puts("affine quant tests: ok"); + return 0; +} diff --git a/c/tests/test_backend_metal.mm b/c/tests/test_backend_metal.mm index 1c66137d8..2dffcfa1c 100644 --- a/c/tests/test_backend_metal.mm +++ b/c/tests/test_backend_metal.mm @@ -9,6 +9,7 @@ #include #include #include +#include #include enum { F32=0, I8=1, I4=2, I2=3, I4G=4, FP8=8 }; @@ -56,6 +57,103 @@ static int run(int fmt, int O, int I, int S, const char *name) { return ok?0:1; } +static uint16_t affine_f16(float value) { + uint32_t bits; memcpy(&bits,&value,sizeof(bits)); + uint32_t sign=(bits>>16)&0x8000u, mantissa=bits&0x7fffffu; + int exponent=(int)((bits>>23)&0xffu)-127+15; + if(exponent<=0) return (uint16_t)sign; + if(exponent>=31) return (uint16_t)(sign|0x7c00u); + return (uint16_t)(sign|((uint32_t)exponent<<10)|(mantissa>>13)); +} + +static void affine_store_scalar(std::vector& bytes, size_t index, + ColiAffineScalarFormat format, float value) { + uint32_t bits; memcpy(&bits,&value,sizeof(bits)); + if(format==COLI_AFFINE_SCALAR_F32){ + bytes[index*4]=(uint8_t)bits; bytes[index*4+1]=(uint8_t)(bits>>8); + bytes[index*4+2]=(uint8_t)(bits>>16); bytes[index*4+3]=(uint8_t)(bits>>24); + } else { + uint16_t encoded=format==COLI_AFFINE_SCALAR_F16 ? affine_f16(value) + : (uint16_t)(bits>>16); + bytes[index*2]=(uint8_t)encoded; bytes[index*2+1]=(uint8_t)(encoded>>8); + } +} + +static int run_affine(ColiAffineFormat format, ColiAffineScalarFormat scalar_format, + int S, const char *name) { + enum { O=48, I=2048, GROUP=32 }; // 64 groups exercises lane's group+32 stride. + unsigned bits=coli_affine_bits(format), per_word=32u/bits; + size_t packed_words=I/per_word, scalar_count=(size_t)O*(I/GROUP); + size_t scalar_size=coli_affine_scalar_size(scalar_format); + std::vector weights((size_t)O*packed_words); + std::vector scales(scalar_count*scalar_size), biases(scalar_count*scalar_size); + std::vector x((size_t)S*I), expected((size_t)S*O), actual((size_t)S*O); + uint32_t state=0x12345678u+(uint32_t)bits*17u+(uint32_t)scalar_format; + for(auto &word:weights){ state=state*1664525u+1013904223u; word=state; } + for(size_t i=0;i dev = MTLCreateSystemDefaultDevice(); + if (!dev) { printf("Metal device unavailable (skipping)\n"); return 0; } + if (!coli_metal_init()) { + printf("Metal backend initialization failed on device %s\n", + [[dev name] UTF8String]); + return 1; + } + ColiMetalAffineCapability affine_capability=coli_metal_affine_capability(); + if (affine_capability!=COLI_METAL_AFFINE_CAP_READY && + affine_capability!=COLI_METAL_AFFINE_CAP_SIMD_WIDTH_UNSUPPORTED) { + printf("Metal MLX affine packed GEMV setup: FAIL (capability=%d)\n", + (int)affine_capability); + coli_metal_shutdown(); + return 1; + } /* GitHub Actions Apple Silicon runners expose Metal through the Apple * Paravirtual device, where Metal submissions never complete (hangs -- the * #947 CI observation). The shader compile above (coli_metal_init) still * runs, keeping this suite's compile coverage -- the class of bug #940 * shipped -- while GPU execution is honestly skipped here. (#947 review) */ - id dev = MTLCreateSystemDefaultDevice(); - if (dev && [[dev name] containsString:@"Apple Paravirtual device"]) { + if ([[dev name] containsString:@"Apple Paravirtual device"]) { printf("SKIPPED: paravirtual device (%s) -- compile-only\n", [[dev name] UTF8String]); + coli_metal_shutdown(); return 0; } printf("Metal standalone op tests (MAE / maxAbs error vs CPU reference):\n"); @@ -813,6 +925,22 @@ int main(void) { fail |= run(I8, 2048,6144,4, "int8 gate/up S=4"); fail |= run(I4, 2048,6144,7, "int4 gate/up S=7 (odd)"); fail |= run(I4, 2050,6146,3, "int4 non-mult-4 dims"); + if (affine_capability==COLI_METAL_AFFINE_CAP_READY) { + printf("Metal MLX affine packed GEMV tests:\n"); + const char *scalar_name[] = {"f32","f16","bf16"}; + for(int format=COLI_AFFINE_MLX_Q4;format<=COLI_AFFINE_MLX_Q8;format++) + for(int scalar=COLI_AFFINE_SCALAR_F32;scalar<=COLI_AFFINE_SCALAR_BF16;scalar++){ + char name[64]; snprintf(name,sizeof(name),"affine Q%d/%s batched", + (int)coli_affine_bits((ColiAffineFormat)format),scalar_name[scalar]); + fail |= run_affine((ColiAffineFormat)format,(ColiAffineScalarFormat)scalar,3,name); + } + } else if (affine_capability==COLI_METAL_AFFINE_CAP_SIMD_WIDTH_UNSUPPORTED) { + printf("Metal MLX affine packed GEMV tests: SKIPPED (requires 32-lane simdgroups)\n"); + } else { + printf("Metal MLX affine packed GEMV tests: FAIL (affine pipelines unavailable, capability=%d)\n", + (int)affine_capability); + fail=1; + } printf("Metal fmt=4 grouped-int4 tests (coli_metal_matmul vs matmul_i4_grouped semantics):\n"); // I multiple of gs=64, S=1 and S>1 (real g64-checkpoint shapes: gate/up I=6144, down I=2048) fail |= run_grouped(2048,6144,64,1,0, "grouped gate/up I=6144(mult64) S=1"); diff --git a/c/tests/test_json.c b/c/tests/test_json.c index 56941a907..931d0e187 100644 --- a/c/tests/test_json.c +++ b/c/tests/test_json.c @@ -30,6 +30,30 @@ int main(void) { CHECK(values->kids[1]->num == -2.5); CHECK(values->kids[2]->num == 300.0); CHECK(strcmp(json_get(root, "unicode")->str, "λ 🚀") == 0); + json_free(root); + + root = json_parse_exact(" {\"exact\":[0,-2.5,3e2]} \n", NULL); + CHECK(root && root->t == J_OBJ); + json_free(root); + + static const char *malformed[] = { + "{\"a\" 1}", + "{\"a\":\"unterminated}", + "{\"a\":1", + "{\"a\":1} trailing", + "{\"a\":1,}", + "[1,]", + "{\"a\":1e}", + "{\"magic\\u0000suffix\":\"QPACK\"}", + "{\"a\":1}\v" + }; + for (size_t i = 0; i < sizeof(malformed) / sizeof(malformed[0]); i++) + CHECK(json_parse_exact(malformed[i], NULL) == NULL); + + /* The legacy entry point intentionally keeps its historical behavior. */ + root = json_parse("{\"a\" 1}", NULL); + CHECK(root && root->t == J_OBJ); + json_free(root); puts("json tests: ok"); return 0; diff --git a/c/tests/test_make_qwen36_qpack_snap.py b/c/tests/test_make_qwen36_qpack_snap.py new file mode 100644 index 000000000..b49bef421 --- /dev/null +++ b/c/tests/test_make_qwen36_qpack_snap.py @@ -0,0 +1,274 @@ +"""make_qwen36_qpack_snap: snap view over a qpack container, plus the C round-trip. + +Builds a synthetic MLX-affine container (nested multimodal config.json with +per-module quantization overrides, a hand-written model.safetensors with +packed U32 triples) and checks: + + * the view's config.json is the FLATTENED text config (load_cfg reads + hidden_size at the top level); + * qwen36_meta.json derives head dims from the LOGICAL weight shapes -- + o_proj's input axis is stored packed (uint32 words), so a derivation + that forgets to unpack reports o_in 8x too small and fails here; + * layer_types / n_active / DeltaNet dims / rope_theta (nested under + rope_parameters) survive the trip; + * model.safetensors and tokenizer.json are links into the container, and + the container itself is never written; + * refusal when --out points inside the container. + +Round-trip (the independent cross-implementation evidence): the fixture +carries a seeded-random Q4 tensor whose dequantized rows are computed HERE in +pure Python (bf16 scalars decoded exactly, scale*q + bias per group) and +handed to tests/test_qwen36_dense_affine's compare mode, which loads the SAME +bytes through the engine's own st_init + load_t_n path -- through the snap +view's symlink, exactly as a real run reads them. Skips loudly if the C gate +binary is not built (`make test-c` builds it first in a green `make test`). +""" +import json +import os +import random +import struct +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +HERE = Path(__file__).resolve().parent +C_DIR = HERE.parent +TOOL = C_DIR / "tools" / "make_qwen36_qpack_snap.py" +GATE = C_DIR / "tests" / ("test_qwen36_dense_affine.exe" if os.name == "nt" + else "test_qwen36_dense_affine") + + +def bf16_bits(value: float) -> int: + return struct.unpack("> 16 + + +def bf16_value(bits: int) -> float: + return struct.unpack("> (bits * (c % per_word))) & mask + g = r * groups + c // cls.GS + row.append(bf16_value(scales[g]) * q + bf16_value(biases[g])) + logical.append(row) + tensors = [ + (name + ".weight", "U32", [rows, cols // per_word], + struct.pack(f"<{len(words)}I", *words)), + (name + ".scales", "BF16", [rows, groups], + struct.pack(f"<{len(scales)}H", *scales)), + (name + ".biases", "BF16", [rows, groups], + struct.pack(f"<{len(biases)}H", *biases)), + ] + return tensors, logical + + @classmethod + def _write_model(cls): + p = "language_model.model.layers.3.self_attn." + tensors = [] + for proj, rows in (("q_proj", cls.N_Q * cls.Q_HEAD_DIM), + ("k_proj", cls.N_KV * cls.HEAD_DIM), + ("v_proj", cls.N_KV * cls.HEAD_DIM)): + t, _ = cls._quantized(p + proj, rows, cls.HIDDEN, bits=4) + tensors += t + # o_proj input axis = n_q*head_dim = 16 LOGICAL columns -> 2 packed + # words; the o_in derivation must report 16, not 2. + t, _ = cls._quantized(p + "o_proj", cls.HIDDEN, + cls.N_Q * cls.HEAD_DIM, bits=4) + tensors += t + tensors.append((p + "q_norm.weight", "BF16", [cls.HEAD_DIM], + struct.pack(f"<{cls.HEAD_DIM}H", + *[bf16_bits(1.0)] * cls.HEAD_DIM))) + # Q8 override module: 8 experts x hidden router. + t, _ = cls._quantized("language_model.model.layers.3.mlp.gate", + 8, cls.HIDDEN, bits=8) + tensors += t + # Round-trip tensor for the C compare-mode gate. + t, cls.roundtrip = cls._quantized("language_model.model.roundtrip", + 4, cls.HIDDEN, bits=4) + tensors += t + write_safetensors(cls.container / "model.safetensors", tensors) + + # ---- gates ----------------------------------------------------------- + + def test_flat_config(self): + cfg = json.loads((self.view / "config.json").read_text()) + self.assertEqual(cfg["hidden_size"], self.HIDDEN) + self.assertEqual(cfg["num_hidden_layers"], self.N_LAYERS) + self.assertEqual(cfg["vocab_size"], 64) + self.assertNotIn("text_config", cfg) + hf = json.loads((self.view / "config.hf.json").read_text()) + self.assertIn("text_config", hf) + + def test_meta_head_dims_unpacked(self): + meta = json.loads((self.view / "qwen36_meta.json").read_text()) + self.assertEqual(meta["q_heads"], self.N_Q) + self.assertEqual(meta["kv_heads"], self.N_KV) + self.assertEqual(meta["q_head_dim"], self.Q_HEAD_DIM) + self.assertEqual(meta["k_head_dim"], self.HEAD_DIM) + self.assertEqual(meta["v_head_dim"], self.HEAD_DIM) + self.assertEqual(meta["o_in"], self.N_Q * self.HEAD_DIM) # unpacked! + self.assertEqual(meta["head_dim"], self.HEAD_DIM) + self.assertEqual(meta["qk_rope_head_dim"], self.HEAD_DIM) + + def test_meta_layers_dn_rope(self): + meta = json.loads((self.view / "qwen36_meta.json").read_text()) + self.assertEqual(meta["n_layers"], self.N_LAYERS) + self.assertEqual(meta["n_active"], 1) + self.assertEqual(meta["layer_types"], + ["linear_attention"] * 3 + ["full_attention"]) + self.assertEqual(meta["num_experts"], 8) + self.assertEqual(meta["topk"], 2) + self.assertEqual(meta["moe_inter"], 16) + self.assertEqual(meta["dn_vheads"], 2) + self.assertEqual(meta["dn_kheads"], 1) + self.assertEqual(meta["dn_conv_dim"], 2 * 1 * 8 + 2 * 8) + self.assertEqual(meta["rope_theta"], 12345678.0) + self.assertTrue(meta["attn_output_gate"]) + self.assertEqual(meta["expert_gs"], 0) + # MLX-derived containers carry full-gamma norm weights; the engine + # must be told to undo the shift (load_norm_n) or the model is noise. + self.assertIs(meta["zero_centered_norms"], False) + + def test_links_and_readonly_container(self): + for name in ("model.safetensors", "tokenizer.json"): + link = self.view / name + self.assertTrue(link.exists(), name) + self.assertEqual(os.stat(link).st_ino, + os.stat(self.container / name).st_ino, + f"{name} is a copy, not a link") + after = sorted(p.name for p in self.container.iterdir()) + self.assertEqual(self.before, after, "container was written to") + + def test_refuses_out_inside_container(self): + result = subprocess.run( + [sys.executable, str(TOOL), "--container", str(self.container), + "--out", str(self.container / "view")], + capture_output=True, text=True, cwd=C_DIR) + self.assertNotEqual(result.returncode, 0) + self.assertIn("OUTSIDE", result.stderr + result.stdout) + + def test_roundtrip_against_engine_loader(self): + if not GATE.exists(): + self.skipTest(f"{GATE} not built (run `make test-c` first): the " + "C-vs-Python round-trip was NOT checked") + expected = struct.pack( + f"<{4 * self.HIDDEN}f", + *[v for row in self.roundtrip for v in row]) + expected_path = Path(self.tmp.name) / "roundtrip.f32" + expected_path.write_bytes(expected) + result = subprocess.run( + [str(GATE), str(self.view), "model.roundtrip.weight", + str(4 * self.HIDDEN), str(expected_path)], + capture_output=True, text=True, cwd=C_DIR) + self.assertEqual(result.returncode, 0, + result.stderr + result.stdout) + self.assertIn("round-trip", result.stdout) + + +if __name__ == "__main__": + unittest.main() diff --git a/c/tests/test_qpack.c b/c/tests/test_qpack.c new file mode 100644 index 000000000..c9f5f756a --- /dev/null +++ b/c/tests/test_qpack.c @@ -0,0 +1,412 @@ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#include "../qpack.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../compat.h" + +#define CHECK(condition) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: check failed: %s\n", \ + __FILE__, __LINE__, #condition); \ + exit(1); \ + } \ +} while (0) + +static const char valid_layout_bf16[] = + "{\n" + " \"expertCount\":2,\"layerCount\":2,\"expertStride\":16384,\n" + " \"linearLayers\":[true,false],\"sections\":[\n" + " {\"name\":\"gate_proj.weight\",\"dtype\":\"U32\"," + "\"shape\":[2,1],\"offset\":0,\"size\":8},\n" + " {\"name\":\"gate_proj.scales\",\"dtype\":\"BF16\"," + "\"shape\":[2,1],\"offset\":8,\"size\":4},\n" + " {\"name\":\"gate_proj.biases\",\"dtype\":\"BF16\"," + "\"shape\":[2,1],\"offset\":12,\"size\":4}\n" + " ]\n}"; + +static const char valid_layout_f16[] = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"gate_proj.weight\",\"dtype\":\"U32\"," + "\"shape\":[2,1],\"offset\":0,\"size\":8}," + "{\"name\":\"gate_proj.scales\",\"dtype\":\"F16\"," + "\"shape\":[2,1],\"offset\":8,\"size\":4}," + "{\"name\":\"gate_proj.biases\",\"dtype\":\"F16\"," + "\"shape\":[2,1],\"offset\":12,\"size\":4}]}"; + +static const char valid_layout_f32[] = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"gate_proj.weight\",\"dtype\":\"U32\"," + "\"shape\":[2,1],\"offset\":0,\"size\":8}," + "{\"name\":\"gate_proj.scales\",\"dtype\":\"F32\"," + "\"shape\":[2,1],\"offset\":8,\"size\":8}," + "{\"name\":\"gate_proj.biases\",\"dtype\":\"F32\"," + "\"shape\":[2,1],\"offset\":16,\"size\":8}]}"; + +static const char *misaligned_layout = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":8192," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"w\",\"dtype\":\"U32\",\"shape\":[1]," + "\"offset\":0,\"size\":4}]}"; + +static const char *overlap_layout = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"a\",\"dtype\":\"U32\",\"shape\":[2]," + "\"offset\":0,\"size\":8}," + "{\"name\":\"b\",\"dtype\":\"U32\",\"shape\":[2]," + "\"offset\":4,\"size\":8}]}"; + +static const char *duplicate_layout = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"a\",\"dtype\":\"U32\",\"shape\":[1]," + "\"offset\":0,\"size\":4}," + "{\"name\":\"a\",\"dtype\":\"U32\",\"shape\":[1]," + "\"offset\":4,\"size\":4}]}"; + +static const char *overflow_layout = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"w\",\"dtype\":\"U32\"," + "\"shape\":[4294967296,4294967296],\"offset\":0,\"size\":4}]}"; + +static const char *bad_affine_shape_layout = + "{\"expertCount\":2,\"layerCount\":2,\"expertStride\":16384," + "\"linearLayers\":[true,false],\"sections\":[" + "{\"name\":\"gate_proj.weight\",\"dtype\":\"U32\"," + "\"shape\":[2,1],\"offset\":0,\"size\":8}," + "{\"name\":\"gate_proj.scales\",\"dtype\":\"BF16\"," + "\"shape\":[1,2],\"offset\":8,\"size\":4}," + "{\"name\":\"gate_proj.biases\",\"dtype\":\"BF16\"," + "\"shape\":[2,1],\"offset\":12,\"size\":4}]}"; + +static void make_path(char *output, size_t capacity, + const char *root, const char *name) { + int count = snprintf(output, capacity, "%s/%s", root, name); + CHECK(count > 0 && (size_t)count < capacity); +} + +static void make_dir(const char *path) { +#ifdef _WIN32 + CHECK(_mkdir(path) == 0); +#else + CHECK(mkdir(path, 0700) == 0); +#endif +} + +static void remove_dir(const char *path) { +#ifdef _WIN32 + CHECK(_rmdir(path) == 0); +#else + CHECK(rmdir(path) == 0); +#endif +} + +static void must_write(const char *path, const void *data, size_t size) { + FILE *file = fopen(path, "wb"); + CHECK(file != NULL); + CHECK(fwrite(data, 1, size, file) == size); + CHECK(fclose(file) == 0); +} + +static void write_layer(const char *root, size_t layer, size_t size, + ColiAffineScalarFormat scalar_format) { + char path[1024], name[64]; + int count = snprintf(name, sizeof(name), + "packed_experts/layer_%02zu.bin", layer); + CHECK(count > 0 && (size_t)count < sizeof(name)); + make_path(path, sizeof(path), root, name); + unsigned char *bytes = (unsigned char *)calloc(size ? size : 1, 1); + CHECK(bytes != NULL); + if (size >= 32768) { + static const unsigned char weights[] = { + 0x10, 0x32, 0x54, 0x76, 0x11, 0x11, 0x11, 0x11 + }; + static const unsigned char bf16_scalars[] = { + 0x00, 0x3f, 0x80, 0x3f, 0x00, 0x3e, 0x00, 0xbf + }; + static const unsigned char f16_scalars[] = { + 0x00, 0x38, 0x00, 0x3c, 0x00, 0x30, 0x00, 0xb8 + }; + static const unsigned char f32_scalars[] = { + 0x00, 0x00, 0x00, 0x3f, 0x00, 0x00, 0x80, 0x3f, + 0x00, 0x00, 0x00, 0x3e, 0x00, 0x00, 0x00, 0xbf + }; + memcpy(bytes + 16384, weights, sizeof(weights)); + if (scalar_format == COLI_AFFINE_SCALAR_F32) + memcpy(bytes + 16384 + sizeof(weights), f32_scalars, + sizeof(f32_scalars)); + else if (scalar_format == COLI_AFFINE_SCALAR_F16) + memcpy(bytes + 16384 + sizeof(weights), f16_scalars, + sizeof(f16_scalars)); + else + memcpy(bytes + 16384 + sizeof(weights), bf16_scalars, + sizeof(bf16_scalars)); + } + must_write(path, bytes, size); + free(bytes); +} + +static void write_fixture_sized(const char *root, + const void *layout, size_t layout_size, + const char *magic, size_t layer_size, + int quant_bits, int group_size, + ColiAffineScalarFormat scalar_format) { + char packed[1024], path[1024]; + make_path(packed, sizeof(packed), root, "packed_experts"); + make_dir(packed); + make_path(path, sizeof(path), root, "packed_experts/layout.json"); + must_write(path, layout, layout_size); + write_layer(root, 0, layer_size, scalar_format); + write_layer(root, 1, layer_size, scalar_format); + + char manifest[2048]; + int count = snprintf( + manifest, sizeof(manifest), + "{\"magic\":\"%s\",\"version\":1,\"modelName\":\"fixture\"," + "\"sourceCheckpoint\":\"fixture\",\"quantBits\":%d," + "\"quantGroupSize\":%d,\"files\":{" + "\"packed_experts/layout.json\":%zu," + "\"packed_experts/layer_00.bin\":%zu," + "\"packed_experts/layer_01.bin\":%zu}}", + magic, quant_bits, group_size, layout_size, layer_size, layer_size); + CHECK(count > 0 && (size_t)count < sizeof(manifest)); + make_path(path, sizeof(path), root, "manifest.json"); + must_write(path, manifest, (size_t)count); +} + +static void write_fixture(const char *root, const char *layout, + const char *magic, size_t layer_size, + int quant_bits, int group_size, + ColiAffineScalarFormat scalar_format) { + write_fixture_sized(root, layout, strlen(layout), magic, layer_size, + quant_bits, group_size, scalar_format); +} + +static void cleanup(const char *root) { + char path[1024]; + for (size_t layer = 0; layer < 2; layer++) { + char name[64]; + int count = snprintf(name, sizeof(name), + "packed_experts/layer_%02zu.bin", layer); + CHECK(count > 0 && (size_t)count < sizeof(name)); + make_path(path, sizeof(path), root, name); + CHECK(remove(path) == 0); + } + make_path(path, sizeof(path), root, "packed_experts/layout.json"); + CHECK(remove(path) == 0); + make_path(path, sizeof(path), root, "manifest.json"); + CHECK(remove(path) == 0); + make_path(path, sizeof(path), root, "packed_experts"); + remove_dir(path); + remove_dir(root); +} + +static void expect_open_failure(const char *layout, const char *magic, + size_t layer_size, const char *message) { + char root[] = "test_qpack_bad_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + write_fixture(root, layout, magic, layer_size, 4, 8, + COLI_AFFINE_SCALAR_BF16); + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) != 0); + CHECK(strstr(error, message) != NULL); + CHECK(reader.layer_fds == NULL && reader.layout.sections == NULL); + cleanup(root); +} + +static void test_valid_container(const char *layout, + ColiAffineScalarFormat scalar_format, + int quant_bits, int group_size, + size_t expected_input, + float expected0, float expected1) { + char root[] = "test_qpack_ok_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + write_fixture(root, layout, "QPACK", 32768, + quant_bits, group_size, scalar_format); + + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) == 0); + CHECK(reader.layout.layer_count == 2); + CHECK(reader.layout.expert_count == 2); + CHECK(reader.layout.expert_stride == 16384); + CHECK(reader.quant_bits == quant_bits); + CHECK(reader.quant_group_size == group_size); + CHECK(reader.layout.linear_layers[0] == 1); + CHECK(reader.layout.linear_layers[1] == 0); + CHECK(reader.layer_fds[0] >= 0 && reader.layer_fds[1] >= 0); + + const ColiQpackSection *section = + coli_qpack_find_section(&reader, "gate_proj.weight"); + CHECK(section && section->offset == 0 && section->size == 8); + CHECK(section->rank == 2 && section->shape[0] == 2 && + section->shape[1] == 1); + + unsigned char *blob = (unsigned char *)malloc(reader.layout.expert_stride); + CHECK(blob != NULL); + CHECK(coli_qpack_read_expert(&reader, 1, 1, blob, + reader.layout.expert_stride, + error, sizeof(error)) == 0); + CHECK(blob[7] == 0x11); + + ColiAffineQuantizedView view; + CHECK(coli_qpack_affine_view(&reader, blob, reader.layout.expert_stride, + "gate_proj", &view, + error, sizeof(error)) == 0); + size_t scalar_bytes = scalar_format == COLI_AFFINE_SCALAR_F32 ? 8 : 4; + CHECK(view.weights == blob && view.scales == blob + 8 && + view.biases == blob + 8 + scalar_bytes); + CHECK(view.weight_bytes == 8 && view.scale_bytes == scalar_bytes && + view.bias_bytes == scalar_bytes); + CHECK(view.output_dim == 2 && view.input_dim == expected_input && + view.group_size == (size_t)group_size); + CHECK(coli_affine_bits(view.format) == (unsigned)quant_bits); + CHECK(view.scalar_format == scalar_format); + + float input[8] = {1, 1, 1, 1, 1, 1, 1, 1}; + float output[2]; + CHECK(coli_affine_matmul_ref(output, input, 1, &view) == COLI_AFFINE_OK); + CHECK(fabsf(output[0] - expected0) < 1e-5f); + CHECK(fabsf(output[1] - expected1) < 1e-5f); + + CHECK(coli_qpack_affine_view(&reader, blob, reader.layout.expert_stride, + "up_proj", &view, + error, sizeof(error)) != 0); + CHECK(coli_qpack_affine_view(&reader, blob, 16, "gate_proj", &view, + error, sizeof(error)) != 0); + CHECK(coli_qpack_read_expert(&reader, 2, 0, blob, + reader.layout.expert_stride, + error, sizeof(error)) != 0); + CHECK(coli_qpack_read_expert(&reader, 0, 0, blob, 16, + error, sizeof(error)) != 0); + free(blob); + coli_qpack_close(&reader); + CHECK(reader.layer_fds == NULL && reader.layout.sections == NULL); + cleanup(root); +} + +static void test_affine_shape_refusal(void) { + char root[] = "test_qpack_shape_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + write_fixture(root, bad_affine_shape_layout, "QPACK", 32768, 4, 8, + COLI_AFFINE_SCALAR_BF16); + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) == 0); + unsigned char *blob = (unsigned char *)calloc(reader.layout.expert_stride, 1); + CHECK(blob != NULL); + ColiAffineQuantizedView view; + CHECK(coli_qpack_affine_view(&reader, blob, reader.layout.expert_stride, + "gate_proj", &view, + error, sizeof(error)) != 0); + CHECK(strstr(error, "shape does not match") != NULL); + free(blob); + coli_qpack_close(&reader); + cleanup(root); +} + +static void test_post_open_truncation(void) { + char root[] = "test_qpack_short_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + write_fixture(root, valid_layout_bf16, "QPACK", 32768, 4, 8, + COLI_AFFINE_SCALAR_BF16); + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) == 0); + write_layer(root, 1, 16384, COLI_AFFINE_SCALAR_BF16); + unsigned char *blob = (unsigned char *)malloc(reader.layout.expert_stride); + CHECK(blob != NULL); + CHECK(coli_qpack_read_expert(&reader, 1, 1, blob, + reader.layout.expert_stride, + error, sizeof(error)) != 0); + CHECK(strstr(error, "short read") != NULL); + free(blob); + coli_qpack_close(&reader); + cleanup(root); +} + +static void test_eager_layer_validation_cleanup(void) { + char root[] = "test_qpack_eager_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + write_fixture(root, valid_layout_bf16, "QPACK", 32768, 4, 8, + COLI_AFFINE_SCALAR_BF16); + write_layer(root, 1, 16384, COLI_AFFINE_SCALAR_BF16); + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) != 0); + CHECK(strstr(error, "unexpected size") != NULL); + CHECK(reader.layer_fds == NULL && reader.layout.sections == NULL); + cleanup(root); +} + +static void test_embedded_nul_metadata(void) { + char root[] = "test_qpack_nul_XXXXXX"; + CHECK(mkdtemp(root) != NULL); + size_t prefix_size = strlen(valid_layout_bf16); + static const char suffix[] = "trailing"; + size_t layout_size = prefix_size + 1 + sizeof(suffix) - 1; + char *layout = (char *)malloc(layout_size); + CHECK(layout != NULL); + memcpy(layout, valid_layout_bf16, prefix_size); + layout[prefix_size] = 0; + memcpy(layout + prefix_size + 1, suffix, sizeof(suffix) - 1); + write_fixture_sized(root, layout, layout_size, "QPACK", 32768, 4, 8, + COLI_AFFINE_SCALAR_BF16); + free(layout); + + ColiQpackReader reader; + char error[512] = {0}; + CHECK(coli_qpack_open(&reader, root, error, sizeof(error)) != 0); + CHECK(strstr(error, "embedded NUL") != NULL); + CHECK(reader.layer_fds == NULL && reader.layout.sections == NULL); + cleanup(root); +} + +int main(void) { + static const struct { + const char *layout; + ColiAffineScalarFormat scalar_format; + } formats[] = { + {valid_layout_bf16, COLI_AFFINE_SCALAR_BF16}, + {valid_layout_f16, COLI_AFFINE_SCALAR_F16}, + {valid_layout_f32, COLI_AFFINE_SCALAR_F32} + }; + for (size_t i = 0; i < sizeof(formats) / sizeof(formats[0]); i++) { + test_valid_container(formats[i].layout, formats[i].scalar_format, + 4, 8, 8, 15.0f, 4.0f); + test_valid_container(formats[i].layout, formats[i].scalar_format, + 8, 4, 4, 134.5f, 66.0f); + } + expect_open_failure(valid_layout_bf16, "NOTQPACK", 32768, + "not a supported QPACK v1"); + expect_open_failure(misaligned_layout, "QPACK", 16384, + "16 KiB aligned"); + expect_open_failure(overlap_layout, "QPACK", 32768, "overlap"); + expect_open_failure(duplicate_layout, "QPACK", 32768, + "duplicate section"); + expect_open_failure(overflow_layout, "QPACK", 32768, "invalid shape"); + expect_open_failure(valid_layout_bf16, "QPACK", 32767, + "manifest size mismatch"); + test_affine_shape_refusal(); + test_eager_layer_validation_cleanup(); + test_embedded_nul_metadata(); + test_post_open_truncation(); + puts("test_qpack: strict v1 metadata, fixed-stride read, affine views: ok"); + return 0; +} diff --git a/c/tests/test_qwen36_dense_affine.c b/c/tests/test_qwen36_dense_affine.c new file mode 100644 index 000000000..3cb86a2a4 --- /dev/null +++ b/c/tests/test_qwen36_dense_affine.c @@ -0,0 +1,309 @@ +/* Dense-bridge gates for qwen36 against an MLX affine model.safetensors (the + * dense half of a Swiftlet qpack container). What must hold: + * + * 1. name resolution -- the engine speaks unprefixed names + * (`model.layers.N...`); a multimodal container stores the text stack + * under `language_model.`. Both spellings must load, and the plain + * spelling must win when it exists (converted snapshots untouched). + * 2. U32 indexing -- st_init must index the packed-U32 weight tensors a + * quantized container carries (it used to exit(1) at the first one) + * with the correct shape/rank, while st_read_f32 still refuses them. + * 3. affine expansion -- a `.weight` U32 tensor with `.scales`/`.biases` + * siblings expands through the checked affine contract, with bits + * (Q4/Q8) and group size DERIVED from the packed shape against the + * config-implied element count, and the expanded values equal to + * scale*q + bias for the authored bit patterns. + * 4. refusals -- a want that the packed shape cannot tile, an orphan U32 + * weight without siblings, and a plain tensor of the wrong size all + * exit non-zero instead of returning plausible rows (fork gates, + * POSIX only). + * + * With argv, the binary becomes the round-trip harness for the Python tool + * test (tests/test_make_qwen36_qpack_snap.py): + * test_qwen36_dense_affine + * loads through the engine's own load_t_n and compares against a + * reference computed INDEPENDENTLY in Python -- the cross-implementation + * evidence that the C expansion agrees with the format's producer side. + */ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#define main qwen36_main_unused +#include "../qwen36.c" +#undef main + +#include +#ifndef _WIN32 +#include +#endif + +#define CHECK(condition) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: check failed: %s\n", \ + __FILE__, __LINE__, #condition); \ + exit(1); \ + } \ +} while (0) + +/* ---- fixture: a minimal safetensors file with affine triples ------------- */ + +typedef struct { + const char *name, *dtype, *shape; + const void *bytes; + size_t nbytes; +} FixtureTensor; + +static void write_fixture(const char *path, const FixtureTensor *tensors, int n) { + char header[4096]; + size_t hlen = 0, off = 0; + hlen += (size_t)snprintf(header + hlen, sizeof(header) - hlen, "{"); + for (int i = 0; i < n; i++) { + hlen += (size_t)snprintf(header + hlen, sizeof(header) - hlen, + "%s\"%s\":{\"dtype\":\"%s\",\"shape\":%s," + "\"data_offsets\":[%zu,%zu]}", + i ? "," : "", tensors[i].name, tensors[i].dtype, tensors[i].shape, + off, off + tensors[i].nbytes); + off += tensors[i].nbytes; + } + hlen += (size_t)snprintf(header + hlen, sizeof(header) - hlen, "}"); + CHECK(hlen < sizeof(header)); + FILE *f = fopen(path, "wb"); + CHECK(f != NULL); + uint64_t h64 = (uint64_t)hlen; + CHECK(fwrite(&h64, 8, 1, f) == 1); + CHECK(fwrite(header, 1, hlen, f) == hlen); + for (int i = 0; i < n; i++) + CHECK(tensors[i].nbytes == 0 || + fwrite(tensors[i].bytes, 1, tensors[i].nbytes, f) == tensors[i].nbytes); + CHECK(fclose(f) == 0); +} + +static uint16_t bf16_bits(float value) { /* fixture values are bf16-exact */ + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + return (uint16_t)(bits >> 16); +} + +static void write_file(const char *path, const char *text) { + FILE *f = fopen(path, "wb"); + CHECK(f != NULL); + CHECK(fwrite(text, 1, strlen(text), f) == strlen(text)); + CHECK(fclose(f) == 0); +} + +/* Q4 [2,32] logical, gs=16: per (row,group) scale 0.5*(1+idx), bias + * 0.25*idx - 1.0 (all bf16-exact). */ +static const uint32_t q4_words[2][4] = { + { 0x76543210u, 0xFEDCBA98u, 0x00000000u, 0x11111111u }, + { 0xAAAAAAAAu, 0x55555555u, 0x0F0F0F0Fu, 0xF0F0F0F0u }, +}; +static float q4_scale(int r, int g) { return 0.5f * (float)(1 + r * 2 + g); } +static float q4_bias(int r, int g) { return 0.25f * (float)(r * 2 + g) - 1.0f; } + +/* Q8 [3,8] logical, gs=8: one group per row. */ +static const uint32_t q8_words[3][2] = { + { 0x03020100u, 0x07060504u }, + { 0xFF800001u, 0x00000000u }, + { 0x10204080u, 0x01010101u }, +}; +static float q8_scale(int r) { return 1.0f + 0.5f * (float)r; } +static float q8_bias(int r) { return -0.5f * (float)r; } + +static const float plain_bf16_vals[4] = { 1.0f, -2.0f, 0.5f, 4.0f }; +static const float plain_f32_vals[3] = { 0.25f, -0.5f, 8.0f }; + +static void build_fixture_dir(const char *dir) { +#ifdef _WIN32 + _mkdir(dir); +#else + mkdir(dir, 0700); +#endif + static uint16_t q4_scales[4], q4_biases[4], q8_scales[3], q8_biases[3]; + static uint16_t plain_bf16[4]; + for (int r = 0; r < 2; r++) for (int g = 0; g < 2; g++) { + q4_scales[r * 2 + g] = bf16_bits(q4_scale(r, g)); + q4_biases[r * 2 + g] = bf16_bits(q4_bias(r, g)); + } + for (int r = 0; r < 3; r++) { + q8_scales[r] = bf16_bits(q8_scale(r)); + q8_biases[r] = bf16_bits(q8_bias(r)); + } + for (int i = 0; i < 4; i++) plain_bf16[i] = bf16_bits(plain_bf16_vals[i]); + static const uint32_t orphan_word[1] = { 0x12345678u }; + const FixtureTensor tensors[] = { + { "language_model.model.aff_q4.weight", "U32", "[2,4]", + q4_words, sizeof(q4_words) }, + { "language_model.model.aff_q4.scales", "BF16", "[2,2]", + q4_scales, sizeof(q4_scales) }, + { "language_model.model.aff_q4.biases", "BF16", "[2,2]", + q4_biases, sizeof(q4_biases) }, + { "language_model.model.aff_q8.weight", "U32", "[3,2]", + q8_words, sizeof(q8_words) }, + { "language_model.model.aff_q8.scales", "BF16", "[3,1]", + q8_scales, sizeof(q8_scales) }, + { "language_model.model.aff_q8.biases", "BF16", "[3,1]", + q8_biases, sizeof(q8_biases) }, + { "language_model.model.plain.weight", "BF16", "[4]", + plain_bf16, sizeof(plain_bf16) }, + { "model.unprefixed.weight", "F32", "[3]", + plain_f32_vals, sizeof(plain_f32_vals) }, + { "language_model.model.orphan.weight", "U32", "[1,1]", + orphan_word, sizeof(orphan_word) }, + }; + char path[512]; + snprintf(path, sizeof(path), "%s/model.safetensors", dir); + write_fixture(path, tensors, (int)(sizeof(tensors) / sizeof(tensors[0]))); +} + +/* ---- gates --------------------------------------------------------------- */ + +static Model g_dm; + +static void test_u32_indexed_with_shape(void) { + st_tensor *t = st_find(&g_dm.S, "language_model.model.aff_q4.weight"); + CHECK(t != NULL); + CHECK(t->dtype == 7); + CHECK(t->rank == 2 && t->shape[0] == 2 && t->shape[1] == 4); + CHECK(t->numel == 8 && t->nbytes == 32); +} + +static void test_q4_expansion(void) { + float *w = load_t_n(&g_dm, "model.aff_q4.weight", 2 * 32); + for (int r = 0; r < 2; r++) for (int i = 0; i < 32; i++) { + int g = i / 16; + float q = (float)((q4_words[r][i / 8] >> (4 * (i % 8))) & 0xFu); + CHECK(w[r * 32 + i] == q4_scale(r, g) * q + q4_bias(r, g)); + } + free(w); +} + +static void test_q8_expansion(void) { + float *w = load_t_n(&g_dm, "model.aff_q8.weight", 3 * 8); + for (int r = 0; r < 3; r++) for (int i = 0; i < 8; i++) { + float q = (float)((q8_words[r][i / 4] >> (8 * (i % 4))) & 0xFFu); + CHECK(w[r * 8 + i] == q8_scale(r) * q + q8_bias(r)); + } + free(w); +} + +/* MLX norm-weight dialect: with zero_centered_norms=0 (what the snap-view + * tool emits for MLX-derived containers) load_norm_n must undo the + * materialised +1 so rmsnorm_row's (1 + w) sees zero-centered weights again; + * with the default HF dialect it must not touch a value. */ +static void test_norm_unshift(void) { + g_dm.c.zero_centered_norms = 1; + float *w = load_norm_n(&g_dm, "model.plain.weight", 4); + for (int i = 0; i < 4; i++) CHECK(w[i] == plain_bf16_vals[i]); + free(w); + g_dm.c.zero_centered_norms = 0; + w = load_norm_n(&g_dm, "model.plain.weight", 4); + for (int i = 0; i < 4; i++) CHECK(w[i] == plain_bf16_vals[i] - 1.0f); + free(w); + g_dm.c.zero_centered_norms = 1; +} + +/* qwen36_meta.json carries the dialect flag; load_meta must parse it. */ +static void test_meta_flag(void) { + write_file("tests/tmp_dense_affine/qwen36_meta.json", + "{\"zero_centered_norms\": false}"); + Cfg c; + memset(&c, 0, sizeof(c)); + c.n_layers = 1; + c.is_attn = calloc(1, 1); + c.zero_centered_norms = 1; + load_meta(&c, "tests/tmp_dense_affine"); + CHECK(c.zero_centered_norms == 0); + free(c.is_attn); +} + +static void test_prefix_resolution(void) { + CHECK(dense_has(&g_dm, "model.plain.weight")); + CHECK(dense_has(&g_dm, "model.unprefixed.weight")); + CHECK(!dense_has(&g_dm, "model.absent.weight")); + float *p = load_t_n(&g_dm, "model.plain.weight", 4); + for (int i = 0; i < 4; i++) CHECK(p[i] == plain_bf16_vals[i]); + free(p); + /* the plain spelling wins when it exists: this fixture only has the + * unprefixed one, and it must load without any prefix probing */ + float *u = load_t_n(&g_dm, "model.unprefixed.weight", 3); + for (int i = 0; i < 3; i++) CHECK(u[i] == plain_f32_vals[i]); + free(u); +} + +#ifndef _WIN32 +static void expect_load_death(const char *name, int64_t want) { + pid_t pid = fork(); + CHECK(pid >= 0); + if (pid == 0) { + if (!freopen("/dev/null", "w", stderr)) { /* keep the noise then */ } + float *p = load_t_n(&g_dm, name, want); + (void)p; + _exit(0); /* reaching here means the refusal did not happen */ + } + int status = 0; + CHECK(waitpid(pid, &status, 0) == pid); + CHECK(WIFEXITED(status) && WEXITSTATUS(status) != 0); +} + +static void test_refusals(void) { + expect_load_death("model.aff_q4.weight", 2 * 32 + 1); /* untileable want */ + expect_load_death("model.aff_q4.weight", 2 * 12); /* neither Q4 nor Q8 */ + expect_load_death("model.orphan.weight", 8); /* missing siblings */ + expect_load_death("model.plain.weight", 5); /* plain wrong size */ + expect_load_death("model.absent.weight", 4); /* missing tensor */ +} +#endif + +/* ---- Python round-trip harness ------------------------------------------ */ + +static int compare_mode(const char *dir, const char *tensor, int64_t want, + const char *expected_path) { + st_init(&g_dm.S, dir); + float *got = load_t_n(&g_dm, tensor, want); + FILE *f = fopen(expected_path, "rb"); + if (!f) { perror(expected_path); return 1; } + float *expected = falloc(want); + if (fread(expected, sizeof(float), (size_t)want, f) != (size_t)want) { + fprintf(stderr, "%s: short read\n", expected_path); return 1; + } + fclose(f); + double max_diff = 0.0; + for (int64_t i = 0; i < want; i++) { + double diff = fabs((double)got[i] - (double)expected[i]); + if (diff > max_diff) max_diff = diff; + if (diff > 1e-5 * (1.0 + fabs((double)expected[i]))) { + fprintf(stderr, "%s[%lld]: got %.9g expected %.9g\n", + tensor, (long long)i, got[i], expected[i]); + return 1; + } + } + printf("round-trip %s: %lld elements, max |diff| %.3g\n", + tensor, (long long)want, max_diff); + free(got); free(expected); + return 0; +} + +int main(int argc, char **argv) { + if (argc == 5) + return compare_mode(argv[1], argv[2], atoll(argv[3]), argv[4]); + if (argc != 1) { + fprintf(stderr, "usage: %s [ ]\n", + argv[0]); + return 2; + } + const char *dir = "tests/tmp_dense_affine"; + build_fixture_dir(dir); + st_init(&g_dm.S, dir); + test_u32_indexed_with_shape(); + test_q4_expansion(); + test_q8_expansion(); + test_prefix_resolution(); + test_norm_unshift(); + test_meta_flag(); +#ifndef _WIN32 + test_refusals(); +#endif + printf("test_qwen36_dense_affine: OK\n"); + return 0; +} diff --git a/c/tests/test_qwen36_tok_merges.c b/c/tests/test_qwen36_tok_merges.c new file mode 100644 index 000000000..91ec10bbb --- /dev/null +++ b/c/tests/test_qwen36_tok_merges.c @@ -0,0 +1,77 @@ +/* Tokenizer merge-table format gate. tokenizer.json spells its merges two + * ways: legacy files write "a b" strings, tokenizers >= 0.20 (transformers + * 4.45+, the Qwen3.6 checkpoints included) writes ["a","b"] pairs. The + * string-only reader silently indexed ZERO merges from the pair form and + * encode_text degraded to one token per byte-symbol -- observed as a 24-token + * encoding of a 24-char prompt on the first real-model run, with the model + * fed a token stream it never saw in training. bpe_piece treats an empty + * merge table as "nothing to merge", so nothing ever refused. + * + * Both spellings of the SAME tiny model must produce the SAME merged ids: + * "in in" -> [in, Ġin], exercising both a plain merge (i+n) and a chained + * one (Ġ + in, rank-ordered after i+n produces the intermediate). */ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif + +#define main qwen36_main_unused +#include "../qwen36.c" +#undef main + +#include + +#define CHECK(condition) do { \ + if (!(condition)) { \ + fprintf(stderr, "%s:%d: check failed: %s\n", \ + __FILE__, __LINE__, #condition); \ + exit(1); \ + } \ +} while (0) + +static const char *tok_pairs = + "{\"model\":{\"vocab\":{\"i\":0,\"n\":1,\"\\u0120\":2,\"in\":3," + "\"\\u0120in\":4},\"merges\":[[\"i\",\"n\"],[\"\\u0120\",\"in\"]]}}"; +static const char *tok_strings = + "{\"model\":{\"vocab\":{\"i\":0,\"n\":1,\"\\u0120\":2,\"in\":3," + "\"\\u0120in\":4},\"merges\":[\"i n\",\"\\u0120 in\"]}}"; + +static void write_file(const char *path, const char *text) { + FILE *f = fopen(path, "wb"); + CHECK(f != NULL); + CHECK(fwrite(text, 1, strlen(text), f) == strlen(text)); + CHECK(fclose(f) == 0); +} + +static void check_encoding(const char *tok_path, const char *label) { + load_tokenizer(tok_path); + CHECK(g_tok != NULL && g_tok_n == 5); + int *ids = NULL, n = 0; + encode_text("in in", &ids, &n); + if (n != 2 || ids[0] != 3 || ids[1] != 4) { + fprintf(stderr, "%s: expected [3, 4], got %d ids:", label, n); + for (int i = 0; i < n; i++) fprintf(stderr, " %d", ids[i]); + fprintf(stderr, "\n"); + exit(1); + } + free(ids); +} + +int main(void) { + const char *dir = "tests/tmp_tok_merges"; +#ifdef _WIN32 + _mkdir(dir); +#else + mkdir(dir, 0700); +#endif + char pairs_path[256], strings_path[256]; + snprintf(pairs_path, sizeof(pairs_path), "%s/tokenizer_pairs.json", dir); + snprintf(strings_path, sizeof(strings_path), "%s/tokenizer_strings.json", dir); + write_file(pairs_path, tok_pairs); + write_file(strings_path, tok_strings); + /* load_tokenizer is re-entrant enough for a test: smap_init re-allocates + * fresh tables and neither fixture carries added_tokens. */ + check_encoding(strings_path, "string merges"); + check_encoding(pairs_path, "pair merges"); + printf("test_qwen36_tok_merges: OK\n"); + return 0; +} diff --git a/c/tools/make_qwen36_qpack_snap.py b/c/tools/make_qwen36_qpack_snap.py new file mode 100644 index 000000000..4d4237a0d --- /dev/null +++ b/c/tools/make_qwen36_qpack_snap.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +"""Build a qwen36 snap view over a Swiftlet qpack container (READ-ONLY source). + +The engine's contract is the one convert_qwen36.py established: a snapshot +directory with a FLAT config.json, a qwen36_meta.json whose head dims were +derived from the actual weight shapes, the *.safetensors holding the dense +weights, and tokenizer.json. A Swiftlet qpack container already carries all +of the bytes -- model.safetensors (dense half, MLX affine triples that +c/qwen36.c expands at load), tokenizer.json, and a nested multimodal +config.json -- but not the two engine-side metadata files, and the container +itself must never be written to (it is a production artifact shared with +Swiftlet, hash-pinned by hashes.json). + +So the bridge is a VIEW: a fresh directory holding the generated flat +config.json + qwen36_meta.json next to symlinks into the container. No +tensor bytes are copied or converted on disk; st_init follows the symlink and +the engine's affine dense loader does the expansion in memory. + +Usage: + python tools/make_qwen36_qpack_snap.py \ + --container ~/models/qwen3.6-35b.qpack --out ~/build/qwen36-35b-snap + +Run the engine against it: + SNAP= QWEN36_QPACK= ./qwen36 16 4 prompt.txt + +Head-dim derivation mirrors convert_qwen36.py ("derived from the actual +weight shapes, authoritative"), with one extra step: quantized tensors +declare PACKED shapes (uint32 words along the input axis), so logical dims +are unpacked through the config's MLX quantization spec (default bits/group +plus per-module overrides, the same matching Swiftlet's Checkpoint.quantSpec +uses). Stdlib only -- no torch, no numpy, no safetensors dependency; the +safetensors header is 8 bytes of length plus JSON. +""" +import argparse +import json +import os +import struct +import sys +from pathlib import Path + + +def read_safetensors_header(path: Path) -> dict: + with open(path, "rb") as f: + (hlen,) = struct.unpack(" (512 << 20): + sys.exit(f"{path}: implausible safetensors header length {hlen}") + header = json.loads(f.read(hlen)) + header.pop("__metadata__", None) + return header + + +class QuantSpec: + """Config `quantization` block: default bits/group + per-module overrides.""" + + def __init__(self, cfg: dict): + block = cfg.get("quantization") or cfg.get("quantization_config") or {} + self.default = (int(block.get("bits", 4)), int(block.get("group_size", 64))) + self.overrides = { + key: (int(spec["bits"]), int(spec["group_size"])) + for key, spec in block.items() + if isinstance(spec, dict) and "bits" in spec + } + + def for_module(self, module: str): + for key, spec in self.overrides.items(): + if module.endswith(key) or key.endswith(module): + return spec + return self.default + + +class DenseShapes: + """Logical (unpacked) shapes for the dense tensors, prefix-resolved.""" + + def __init__(self, header: dict, quant: QuantSpec): + self.header = header + self.quant = quant + + def resolve(self, name: str): + if name in self.header: + return name + prefixed = "language_model." + name + return prefixed if prefixed in self.header else None + + def logical_shape(self, name: str): + resolved = self.resolve(name) + if resolved is None: + return None + info = self.header[resolved] + shape = list(info["shape"]) + if info["dtype"] in ("U32", "I32") and self.resolve( + name.removesuffix(".weight") + ".scales") is not None: + bits, _ = self.quant.for_module( + resolved.removesuffix(".weight")) + shape[-1] *= 32 // bits + return shape + + +def derive_meta(cfg_full: dict, mcfg: dict, shapes: DenseShapes) -> dict: + n_layers = int(mcfg["num_hidden_layers"]) + layer_types = mcfg.get("layer_types") + if not layer_types: + interval = int(mcfg.get("full_attention_interval", 4)) + layer_types = ["full_attention" if (i + 1) % interval == 0 + else "linear_attention" for i in range(n_layers)] + full_idx = [i for i, t in enumerate(layer_types) if t == "full_attention"] + rope = mcfg.get("rope_parameters") or {} + meta = { + "model_type": cfg_full.get("model_type"), + "hidden": int(mcfg["hidden_size"]), + "n_layers": n_layers, + "n_active": len(full_idx), + "layer_types": layer_types, + "num_experts": int(mcfg["num_experts"]), + "topk": int(mcfg["num_experts_per_tok"]), + "moe_inter": int(mcfg["moe_intermediate_size"]), + "shared_inter": int(mcfg.get("shared_expert_intermediate_size", + mcfg.get("moe_intermediate_size", 0))), + "rms_eps": float(mcfg.get("rms_norm_eps", 1e-6)), + "scoring_func": mcfg.get("scoring_func", "softmax"), + "n_group": int(mcfg.get("n_group", 1)), + "topk_group": int(mcfg.get("topk_group", 1)), + "norm_topk_prob": bool(mcfg.get("norm_topk_prob", False)), + "attn_output_gate": bool(mcfg.get("attn_output_gate", False)), + "rope_theta": float(rope.get("rope_theta", + mcfg.get("rope_theta", 10000000.0))), + "mrope_section": rope.get("mrope_section", [16, 16, 16]), + "partial_rotary_factor": float(mcfg.get("partial_rotary_factor", 0.25)), + "expert_gs": 0, # snapshot-expert group scaling; unused with a qpack store + # HF Qwen3.6 stores the rmsnorm_row weights zero-centered and the + # engine applies (1 + w); mlx-lm materialises the +1 at conversion, so + # an MLX-derived container carries FULL gamma. This flag makes the + # engine undo the shift at load (c/qwen36.c load_norm_n) -- without it + # every normalised activation is ~doubled and the model is noise. + "zero_centered_norms": False, + } + n_q = int(mcfg["num_attention_heads"]) + n_kv = int(mcfg["num_key_value_heads"]) + meta["q_heads"] = n_q + meta["kv_heads"] = n_kv + if not full_idx: + sys.exit("no full_attention layer to derive head dims from") + fi = full_idx[0] + + def logical(proj): + return shapes.logical_shape(f"model.layers.{fi}.self_attn.{proj}.weight") + + qp, kp, vp, op = (logical(p) for p in ("q_proj", "k_proj", "v_proj", "o_proj")) + qn = shapes.logical_shape(f"model.layers.{fi}.self_attn.q_norm.weight") + if qp is None or kp is None or vp is None or op is None: + sys.exit(f"layer {fi}: attention projections missing from model.safetensors") + meta["q_head_dim"] = qp[0] // n_q + meta["k_head_dim"] = kp[0] // n_kv + meta["v_head_dim"] = vp[0] // n_kv + meta["o_in"] = op[1] + if qn is not None: + meta["qk_rope_head_dim"] = qn[0] + meta["head_dim"] = meta["k_head_dim"] + meta["rope_dim"] = meta.get("qk_rope_head_dim", meta["head_dim"] // 4) + # DeltaNet dims are explicit in config (same source convert_qwen36.py uses). + meta["dn_vheads"] = int(mcfg.get("linear_num_value_heads", + mcfg.get("num_value_heads", 0))) + meta["dn_kheads"] = int(mcfg.get("linear_num_key_heads", + mcfg.get("num_key_heads", 0))) + meta["dn_kdim"] = int(mcfg.get("linear_key_head_dim", + mcfg.get("key_head_dim", 0))) + meta["dn_vdim"] = int(mcfg.get("linear_value_head_dim", + mcfg.get("value_head_dim", 0))) + meta["dn_convk"] = int(mcfg.get("linear_conv_kernel_dim", + mcfg.get("conv_kernel_size", 0))) + meta["dn_conv_dim"] = (meta["dn_kheads"] * meta["dn_kdim"] * 2 + + meta["dn_vheads"] * meta["dn_vdim"]) + return meta + + +def place_link(target: Path, link: Path): + """Symlink preferred (zero copies, container stays the single source of + truth); hard link as the fallback for hosts without symlink privileges.""" + if link.is_symlink() or link.exists(): + link.unlink() + try: + link.symlink_to(target) + return "symlink" + except OSError: + os.link(target, link) # let a cross-device failure surface loudly + return "hardlink" + + +def main(): + # __doc__ is None under python -OO (docstrings stripped): fall back to a + # literal rather than crashing before argparse can even print usage. + ap = argparse.ArgumentParser( + description=__doc__.splitlines()[0] if __doc__ else + "Build a qwen36 snap view over a Swiftlet qpack container") + ap.add_argument("--container", required=True, + help="qpack container directory (never written to)") + ap.add_argument("--out", required=True, + help="snap view directory to create") + args = ap.parse_args() + container = Path(args.container).expanduser().resolve() + out = Path(args.out).expanduser().resolve() + if container == out or container in out.parents: + sys.exit("refusing: --out must live OUTSIDE the read-only container") + st_path = container / "model.safetensors" + cfg_path = container / "config.json" + for required in (st_path, cfg_path): + if not required.is_file(): + sys.exit(f"{required}: missing -- not a qpack container?") + manifest = container / "manifest.json" + if manifest.is_file(): + magic = json.loads(manifest.read_text(encoding="utf-8")).get("magic") + if magic != "QPACK": + sys.exit(f"{manifest}: magic {magic!r} is not QPACK") + else: + print(f"[warn] {manifest} missing; proceeding on model.safetensors alone") + + cfg_full = json.loads(cfg_path.read_text(encoding="utf-8")) + mcfg = cfg_full.get("text_config", cfg_full) + quant = QuantSpec(cfg_full) + shapes = DenseShapes(read_safetensors_header(st_path), quant) + meta = derive_meta(cfg_full, mcfg, shapes) + + out.mkdir(parents=True, exist_ok=True) + # The engine reads FLAT keys from config.json (load_cfg), but multimodal + # checkpoints nest them under text_config -- same flattening + # convert_qwen36.py performs, original kept as config.hf.json. + (out / "config.json").write_text(json.dumps(mcfg, indent=2), encoding="utf-8") + (out / "config.hf.json").write_text(json.dumps(cfg_full, indent=2), + encoding="utf-8") + (out / "qwen36_meta.json").write_text(json.dumps(meta, indent=2), + encoding="utf-8") + links = {"model.safetensors": st_path} + tok = container / "tokenizer.json" + if tok.is_file(): + links["tokenizer.json"] = tok + else: + print(f"[warn] {tok} missing; set TOK= when running the engine") + for name, target in links.items(): + kind = place_link(target, out / name) + print(f"[link] {name} -> {target} ({kind})") + print(f"[meta] {out / 'qwen36_meta.json'} " + f"(q_head_dim={meta['q_head_dim']} k/v={meta['k_head_dim']}/" + f"{meta['v_head_dim']} o_in={meta['o_in']} " + f"n_active={meta['n_active']}/{meta['n_layers']})") + print(f"\nRun: SNAP={out} QWEN36_QPACK={container} ./qwen36 16 4 prompt.txt") + + +if __name__ == "__main__": + main()