diff --git a/.gitignore b/.gitignore index 6e4c2c90..bece6d8d 100644 --- a/.gitignore +++ b/.gitignore @@ -27,4 +27,7 @@ h3_semantic_dit_test h3_real_video_vae_test h3_semantic_vae_test libh3.a +h3_baseline +build_attention_cache +*.cache .DS_Store diff --git a/Makefile b/Makefile index bb202379..c53436d3 100644 --- a/Makefile +++ b/Makefile @@ -27,6 +27,11 @@ h3: $(CLI_OBJ) $(LIB_OBJ) libh3.a: $(LIB_OBJ) $(AR) rcs $@ $^ +# One-time tool: quantizes QKV/attention-output to int8 and writes the cache +# that H3_ATTENTION_CACHE points the runtime at. See h3_build_attention_cache.c. +build_attention_cache: h3_build_attention_cache.o $(LIB_OBJ) + $(CC) -o $@ $^ $(LDLIBS) + h3_tests: tests/test_h3.o $(LIB_OBJ) $(CC) -o $@ $^ $(LDLIBS) diff --git a/README.md b/README.md index 4750ac49..e9331c2d 100644 --- a/README.md +++ b/README.md @@ -554,6 +554,71 @@ interactive DiT is ready for its next denoiser evaluation. Measurements reached about 13--14.6 GiB/s from the internal SSD. `H3_PROFILE=1` reports total bytes, read throughput, and the part of the read wait that was not hidden by GPU work. +### Streamed int8 attention cache + +`H3_ATTENTION_CACHE=path/to/cache` is a lighter-weight alternative to +`--ssd-streaming`: QKV/attention-output stream per block from a cache +pre-quantized to int8 (~147 MiB/layer, versus ~735 MiB for BF16 all four +matrices) through two double-buffered slots, instead of being resident. +The MLP (FC1/FC2) still stays int8-resident by default, same as the plain +resident-int8 path. `H3_INT8_STREAM_MLP=1` streams FC1/FC2 from the same +cache too, dropping DiT weight residency to just the two slots - the +tradeoff a long (~15s/362-frame) run needs, and, on measurement, is +consistently faster than `--ssd-streaming` even on short clips once set +(a short 22-frame/512-square clip: 141s on `--ssd-streaming` versus ~78s +on `H3_ATTENTION_CACHE`+`H3_INT8_STREAM_MLP=1`, both 20 denoising steps). +Without `H3_INT8_STREAM_MLP`, the cache only avoids the resident path's +one-time MLP quantization cost, which mostly shows up on longer runs - +short clips can come out slower than `--ssd-streaming` in that +configuration. Passing `--ssd-streaming` itself always wins over +`H3_ATTENTION_CACHE` if both are set (see below), rather than erroring. + +Build one with `build_attention_cache `. Both paths measured bit-for-bit identical output against +plain resident-int8 at matched seed. + +The cache format's header (v3) tags which transformer directory it was +quantized from - `model_kind` (FL2VA or Ref2VA) and `model_id` (a cheap, +non-cryptographic fingerprint of the checkpoint's own shard paths/sizes/ +mtimes, not a hash of the ~18GB of weight bytes). `H3_ATTENTION_CACHE` +refuses a cache whose `model_kind` does not match the generation actually +running (e.g. an FL2VA cache used once `--ref-image`/`--ref-video` +switches to Ref2VA) with a clear error, rather than silently streaming +structurally-compatible-but-wrong weights - both models share the same +DiT dimensions, so nothing else would have caught this: + +``` +h3: Ref2VA generation cannot use a FL2VA attention cache (dit_int8_v2.cache) - rebuild it against the matching transformer directory +``` + +A `model_id` mismatch (rewritten weights, a LoRA baked in after the cache +was built, or a moved/copied checkpoint) is a warning, not a hard error, +since the fingerprint can occasionally shift for benign reasons (e.g. a +copy that resets mtimes) that `model_kind` never would. + +This is a breaking format change: v2 caches (from before this) fail the +version check and must be rebuilt with the new `build_attention_cache`. + +For a model directory with both FL2VA and Ref2VA (most releases), +`build_attention_cache ` builds +both in one pass - detected by the presence of `/FL2VA/transformer/config.json` - writing `/fl2va.cache` and, +if a Ref2VA transformer is present, `/ref2va.cache` too. Point +`H3_ATTENTION_CACHE_DIR` at that directory instead of `H3_ATTENTION_CACHE` +at a single file, and h3.c auto-selects the matching cache the same way +it already selects between the two transformer directories (by whether +`--ref-image`/`--ref-video`/etc. are present): + +``` +build_attention_cache MiniMax-H3 ./h3-cache +H3_ATTENTION_CACHE_DIR=./h3-cache ./h3 -d MiniMax-H3 -p "..." +``` + +`H3_ATTENTION_CACHE_DIR` needs the standard `FL2VA/transformer`/ +`Ref2VA/transformer` layout to know which file to pick; use +`H3_ATTENTION_CACHE` (a single file) for a non-standard directory +instead - setting both at once is an error. + ### Metal 4 and TensorOps paths M5 GPUs automatically use native BF16 Metal 4/TensorOps for the DiT QKV and diff --git a/h3_build_attention_cache.c b/h3_build_attention_cache.c new file mode 100644 index 00000000..3af22bb3 --- /dev/null +++ b/h3_build_attention_cache.c @@ -0,0 +1,275 @@ +/* Builds a pre-quantized int8 cache of every DiT weight matrix (QKV, + * attention-output, FC1, FC2), so the runtime can either keep the MLP + * portion resident (as before) or stream all four matrices for every + * block - the latter is what a long (~15s/362-frame) run needs, since + * even int8-resident MLP for all 50 blocks (~10.8 GiB) stops being cheap + * once per-sequence activations grow with frame count. See h3_dit.c's + * H3_ATTENTION_CACHE / H3_INT8_STREAM_MLP handling for the reader side. + * + * Cache layout: a 64-byte header, then per block (in block order): + * qkv_int8[INNER*3*HIDDEN] qkv_scales[INNER*3] (f32) + * out_int8[HIDDEN*INNER] out_scales[HIDDEN] (f32) + * fc1_int8[FFN*2*HIDDEN] fc1_scales[FFN*2] (f32) + * fc2_int8[HIDDEN*FFN] fc2_scales[HIDDEN] (f32) + * Quantization uses the exact same GPU routine (h3_gpu_quantize_weight_int8) + * as the existing resident-int8 path, so results match it bit for bit; they + * are not expected to match a BF16-only run. + * + * The v3 header adds model_kind (FL2VA=1/Ref2VA=2, matching h3_dit.c's + * h3_cache_model_kind) and model_id (h3_weight_store_fingerprint() of the + * transformer directory quantized) so h3_dit.c can refuse a cache built + * for the wrong model rather than silently streaming mismatched weights - + * a v2 cache (no such tagging) simply fails h3_dit.c's version check and + * must be rebuilt. + * + * Usage: + * build_attention_cache + * - single-component mode, backward compatible with v2's usage; the + * transformer directory's own path (its last two components, e.g. + * ".../FL2VA/transformer") sets model_kind, same convention h3.c's + * own dit_path selection writes and h3_dit.c reads back. + * build_attention_cache + * - detected when /FL2VA/transformer/config.json + * exists: writes /fl2va.cache, and, if a Ref2VA + * transformer is present too, /ref2va.cache - matching + * H3_ATTENTION_CACHE_DIR's auto-selection by generation mode. + */ +#include "h3_gpu.h" +#include "h3_weights.h" + +#include +#include +#include +#include +#include +#include + +enum { + HIDDEN = 5376, + HEADS = 56, + HEAD_DIM = 128, + INNER = HEADS * HEAD_DIM, + FFN = 14336, + DIT_BLOCKS = 50, +}; + +#define CACHE_MAGIC "H3AC" +#define CACHE_VERSION 3u + +typedef enum { + MODEL_UNKNOWN = 0, + MODEL_FL2VA = 1, + MODEL_REF2VA = 2, +} model_kind; + +typedef struct { + char magic[4]; + uint32_t version; + uint32_t block_count; + uint32_t hidden; + uint32_t inner; + uint32_t ffn; + uint32_t model_kind; + uint8_t model_id[32]; + uint32_t reserved[1]; +} cache_header; + +/* Mirrors h3_dit.c's detect_model_kind(): reads back the layout h3.c's + * own dit_path selection always produces, rather than guessing at one. */ +static model_kind detect_model_kind(const char *transformer_dir) { + size_t length = strlen(transformer_dir); + static const char ref2va_suffix[] = "Ref2VA/transformer"; + static const char fl2va_suffix[] = "FL2VA/transformer"; + if (length >= sizeof(ref2va_suffix) - 1 && + !strcmp(transformer_dir + length - (sizeof(ref2va_suffix) - 1), + ref2va_suffix)) + return MODEL_REF2VA; + if (length >= sizeof(fl2va_suffix) - 1 && + !strcmp(transformer_dir + length - (sizeof(fl2va_suffix) - 1), + fl2va_suffix)) + return MODEL_FL2VA; + return MODEL_UNKNOWN; +} + +static int path_exists(const char *path) { + struct stat status; + return stat(path, &status) == 0; +} + +static char *join_path(const char *base, const char *suffix) { + size_t length = strlen(base) + strlen(suffix) + 2; + char *result = malloc(length); + if (result) snprintf(result, length, "%s/%s", base, suffix); + return result; +} + +static int quantize_and_write(h3_gpu *gpu, h3_weight_store *store, + const char *name, uint32_t rows, + uint32_t columns, FILE *out, + char *error, size_t error_size) { + uint64_t shape[2] = { rows, columns }; + h3_gpu_tensor *bf16 = h3_weight_load_bf16(store, gpu, name, 2, shape, + error, error_size); + if (!bf16) return 0; + + size_t elements = (size_t)rows * columns; + h3_gpu_tensor *i8 = h3_gpu_tensor_new_i8(gpu, elements); + h3_gpu_tensor *scales = h3_gpu_tensor_new_f32(gpu, rows); + int ok = i8 && scales && + h3_gpu_begin(gpu) && + h3_gpu_quantize_weight_int8(gpu, i8, scales, bf16, rows, + columns) && + h3_gpu_submit(gpu); + h3_gpu_tensor_free(bf16); + if (!ok) { + if (error && error_size && !error[0]) + snprintf(error, error_size, "cannot quantize %s: %s", name, + h3_gpu_error(gpu)); + h3_gpu_tensor_free(i8); + h3_gpu_tensor_free(scales); + return 0; + } + + int8_t *i8_host = malloc(elements); + float *scale_host = malloc((size_t)rows * sizeof(float)); + int result = i8_host && scale_host && + h3_gpu_tensor_read_i8(i8, i8_host, elements) && + h3_gpu_tensor_read_f32(scales, scale_host, rows) && + fwrite(i8_host, 1, elements, out) == elements && + fwrite(scale_host, sizeof(float), rows, out) == rows; + free(i8_host); + free(scale_host); + h3_gpu_tensor_free(i8); + h3_gpu_tensor_free(scales); + if (!result && error && error_size) + snprintf(error, error_size, "cannot write cache payload for %s: %s", + name, strerror(errno)); + return result; +} + +/* Builds one complete attention cache from transformer_dir into + * output_path. Returns 1 on success, 0 with a message already printed to + * stderr on failure. */ +static int build_one_cache(h3_gpu *gpu, const char *transformer_dir, + const char *output_path) { + char error[512] = {0}; + h3_weight_store *store = h3_weight_store_open(transformer_dir, error, + sizeof(error)); + if (!store) { + fprintf(stderr, "h3: %s\n", error); + return 0; + } + + FILE *out = fopen(output_path, "wb"); + if (!out) { + fprintf(stderr, "h3: cannot open %s: %s\n", output_path, + strerror(errno)); + h3_weight_store_free(store); + return 0; + } + + cache_header header = {0}; + memcpy(header.magic, CACHE_MAGIC, 4); + header.version = CACHE_VERSION; + header.block_count = DIT_BLOCKS; + header.hidden = HIDDEN; + header.inner = INNER; + header.ffn = FFN; + header.model_kind = (uint32_t)detect_model_kind(transformer_dir); + h3_weight_store_fingerprint(store, header.model_id); + if (fwrite(&header, sizeof(header), 1, out) != 1) { + fprintf(stderr, "h3: cannot write cache header: %s\n", + strerror(errno)); + fclose(out); + h3_weight_store_free(store); + return 0; + } + + for (uint32_t block = 0; block < DIT_BLOCKS; block++) { + char name[160]; + struct { const char *suffix; uint32_t rows, columns; } projections[] = { + {"attn.qkv_proj.weight", INNER * 3, HIDDEN}, + {"attn.out_proj.weight", HIDDEN, INNER}, + {"mlp.fc1.weight", FFN * 2, HIDDEN}, + {"mlp.fc2.weight", HIDDEN, FFN}, + }; + for (size_t p = 0; p < sizeof(projections) / sizeof(*projections); p++) { + error[0] = '\0'; + snprintf(name, sizeof(name), "blocks.%u.%s", block, + projections[p].suffix); + if (!quantize_and_write(gpu, store, name, projections[p].rows, + projections[p].columns, out, error, + sizeof(error))) { + fprintf(stderr, "h3: %s\n", error); + fclose(out); + h3_weight_store_free(store); + return 0; + } + } + fprintf(stderr, "h3: attention cache block %2u/%u\n", block + 1, + DIT_BLOCKS); + } + + fclose(out); + h3_weight_store_free(store); + fprintf(stderr, "h3: wrote attention cache to %s\n", output_path); + return 1; +} + +int main(int argc, char **argv) { + if (argc != 3) { + fprintf(stderr, + "usage: %s \n" + " %s \n", + argv[0], argv[0]); + return 1; + } + char *fl2va_dir = join_path(argv[1], "FL2VA/transformer"); + char *fl2va_marker = fl2va_dir ? join_path(fl2va_dir, "config.json") : NULL; + int model_root_mode = fl2va_marker && path_exists(fl2va_marker); + free(fl2va_marker); + + char error[512] = {0}; + h3_gpu *gpu = h3_gpu_create("h3_shaders.metal", error, sizeof(error)); + if (!gpu) { + fprintf(stderr, "h3: %s\n", error); + free(fl2va_dir); + return 1; + } + if (!h3_gpu_has_int8_mlp(gpu)) { + fprintf(stderr, + "h3: this GPU lacks the int8 path the cache is built for\n"); + h3_gpu_free(gpu); + free(fl2va_dir); + return 1; + } + + int ok; + if (model_root_mode) { + mkdir(argv[2], 0755); /* ignore EEXIST - a pre-existing dir is fine */ + char *fl2va_out = join_path(argv[2], "fl2va.cache"); + char *ref2va_dir = join_path(argv[1], "Ref2VA/transformer"); + char *ref2va_marker = ref2va_dir ? join_path(ref2va_dir, "config.json") : NULL; + ok = fl2va_out && build_one_cache(gpu, fl2va_dir, fl2va_out); + if (ok) { + if (ref2va_marker && path_exists(ref2va_marker)) { + char *ref2va_out = join_path(argv[2], "ref2va.cache"); + ok = ref2va_out && build_one_cache(gpu, ref2va_dir, ref2va_out); + free(ref2va_out); + } else { + fprintf(stderr, + "h3: no Ref2VA transformer under %s - built " + "fl2va.cache only\n", argv[1]); + } + } + free(fl2va_out); + free(ref2va_dir); + free(ref2va_marker); + } else { + ok = build_one_cache(gpu, argv[1], argv[2]); + } + + h3_gpu_free(gpu); + free(fl2va_dir); + return ok ? 0 : 1; +} diff --git a/h3_dit.c b/h3_dit.c index 667e49b2..4e1f2060 100644 --- a/h3_dit.c +++ b/h3_dit.c @@ -3,6 +3,7 @@ #include "h3_dit_schedule.h" #include "h3_weights.h" +#include #include #include #include @@ -67,6 +68,63 @@ typedef struct { h3_dit_stream_source sources[STREAM_MATRICES]; } h3_dit_stream_layer; +/* H3_ATTENTION_CACHE: an alternative to --ssd-streaming. QKV/attention- + * output are always streamed from a pre-quantized int8 cache file (see + * h3_build_attention_cache.c) instead of being resident. Per layer this + * moves ~147 MiB (int8 QKV+OUT) instead of ~735 MiB (BF16 all four + * matrices). + * + * By default the MLP (FC1/FC2) still stays int8-resident for every active + * block, same as the plain resident-int8 path (good for short clips, where + * per-sequence activations are small and the ~10.8 GiB of resident MLP is + * cheap next to that). H3_INT8_STREAM_MLP=1 streams FC1/FC2 from the same + * cache too (an extra ~221 MiB/layer), dropping DiT weight residency to + * just the two double-buffer slots (~0.72 GiB) - the tradeoff a long + * (~15s/362-frame) run needs, since activations grow with sequence length + * and resident MLP for every block stops being the cheap part. */ +typedef struct { + h3_gpu_tensor *qkv_int8; + h3_gpu_tensor *qkv_scales; + h3_gpu_tensor *out_int8; + h3_gpu_tensor *out_scales; + h3_gpu_tensor *fc1_int8; + h3_gpu_tensor *fc1_scales; + h3_gpu_tensor *fc2_int8; + h3_gpu_tensor *fc2_scales; +} h3_dit_attention_slot; + +#define H3_ATTENTION_CACHE_MAGIC "H3AC" +#define H3_ATTENTION_CACHE_VERSION 3u + +/* Which transformer directory (see h3.c's dit_path selection: "FL2VA/ + * transformer" vs "Ref2VA/transformer") a cache's weights were quantized + * from. H3_CACHE_MODEL_UNKNOWN means the cache (or the checkpoint it is + * about to serve) did not come from that standard layout, so the + * model_kind check is skipped rather than guessed at. */ +typedef enum { + H3_CACHE_MODEL_UNKNOWN = 0, + H3_CACHE_MODEL_FL2VA = 1, + H3_CACHE_MODEL_REF2VA = 2, +} h3_cache_model_kind; + +/* v3 adds model_kind + model_id where v2 had unused reserved[10] (40 + * bytes: 4 + 32 + 4 = 40, so the header stays byte-identical in size) - + * a v2 file simply fails the version check below rather than being + * misread, since those bytes used to be zero/undefined. model_id is + * h3_weight_store_fingerprint()'s output, or all-zero if the tool that + * wrote this cache did not have a checkpoint directory to fingerprint. */ +typedef struct { + char magic[4]; + uint32_t version; + uint32_t block_count; + uint32_t hidden; + uint32_t inner; + uint32_t ffn; + uint32_t model_kind; + uint8_t model_id[32]; + uint32_t reserved[1]; +} h3_attention_cache_header; + struct h3_dit { h3_gpu *gpu; h3_weight_store *weights; @@ -145,6 +203,15 @@ struct h3_dit { uint64_t stream_bytes; double stream_read_seconds; double stream_wait_seconds; + char *attention_cache_path; + int attention_stream; + int mlp_stream; + h3_dit_attention_slot attention_slots[2]; + unsigned attn_ready_layer; + unsigned attn_ready_slot; + uint64_t attn_stream_bytes; + double attn_stream_read_seconds; + double attn_stream_wait_seconds; h3_gpu_tensor *final_norm; h3_gpu_tensor *final_video_w; h3_gpu_tensor *final_video_b; @@ -542,6 +609,34 @@ static int load_block_norms(h3_dit *dit, h3_dit_block *block, return 1; } +/* Like load_block(), but for H3_ATTENTION_CACHE: QKV/attention-output come + * from the streamed int8 cache instead, so only the norms and the MLP (kept + * int8-resident, same as the plain resident path) are loaded here. */ +static int load_block_norms_and_mlp(h3_dit *dit, h3_dit_block *block, + const char *prefix, + char *error, size_t error_size) { + char name[160]; +#define LOAD1(field, suffix, width) do { \ + snprintf(name, sizeof(name), "%s%s", prefix, suffix); \ + block->field = bf1(dit, name, width, error, error_size); \ + if (!block->field) return 0; \ +} while (0) +#define LOAD2(field, suffix, rows, columns) do { \ + snprintf(name, sizeof(name), "%s%s", prefix, suffix); \ + block->field = bf2(dit, name, rows, columns, error, error_size); \ + if (!block->field) return 0; \ +} while (0) + LOAD1(norm1, "norm1.weight", HIDDEN); + LOAD1(norm2, "norm2.weight", HIDDEN); + LOAD1(q_norm, "attn.q_norm.weight", HEAD_DIM); + LOAD1(k_norm, "attn.k_norm.weight", HEAD_DIM); + LOAD2(fc1, "mlp.fc1.weight", FFN * 2, HIDDEN); + LOAD2(fc2, "mlp.fc2.weight", HIDDEN, FFN); +#undef LOAD1 +#undef LOAD2 + return 1; +} + static void free_block(h3_dit_block *block) { free_tensor(&block->norm1); free_tensor(&block->norm2); @@ -659,26 +754,73 @@ typedef struct { char error[512]; } h3_dit_stream_job; +typedef struct { + const h3_dit_stream_source *source; + h3_gpu_tensor *target; + int ok; + uint64_t bytes; + char error[512]; +} h3_dit_matrix_job; + +static void *read_stream_matrix_thread(void *opaque) { + h3_dit_matrix_job *job = opaque; + job->error[0] = '\0'; + if (!job->target || !h3_gpu_tensor_stream_file_bf16( + job->target, job->source->path, job->source->file_offset, + job->source->elements, job->error, sizeof(job->error))) { + if (!job->error[0]) + snprintf(job->error, sizeof(job->error), + "invalid BF16 streaming destination"); + job->ok = 0; + return NULL; + } + job->ok = 1; + job->bytes = (uint64_t)job->source->elements * sizeof(uint16_t); + return NULL; +} + +/* The four matrices (QKV/OUT/FC1/FC2) target independent buffers and each + open their own file descriptor, so they can be read concurrently. A + single pread() per matrix only reaches queue depth 1, well under what + Apple Silicon's NVMe controller can sustain - reading all four at once + raises the outstanding request count instead. */ static int read_stream_layer(h3_dit_stream_job *job) { h3_dit_stream_layer *layer = &job->dit->stream_layers[job->layer]; h3_dit_block *slot = &job->dit->stream_slots[job->slot]; double started = stream_now(); + + h3_dit_matrix_job matrix_jobs[STREAM_MATRICES]; + pthread_t matrix_threads[STREAM_MATRICES]; + int started_thread[STREAM_MATRICES] = {0}; + + for (unsigned index = 0; index < STREAM_MATRICES; index++) { + matrix_jobs[index].source = &layer->sources[index]; + matrix_jobs[index].target = + stream_slot_target(slot, layer->sources[index].field); + matrix_jobs[index].ok = 0; + matrix_jobs[index].bytes = 0; + if (pthread_create(&matrix_threads[index], NULL, + read_stream_matrix_thread, + &matrix_jobs[index]) == 0) { + started_thread[index] = 1; + } else { + read_stream_matrix_thread(&matrix_jobs[index]); + } + } + job->ok = 1; job->bytes = 0; job->error[0] = '\0'; for (unsigned index = 0; index < STREAM_MATRICES; index++) { - const h3_dit_stream_source *source = &layer->sources[index]; - h3_gpu_tensor *target = stream_slot_target(slot, source->field); - if (!target || !h3_gpu_tensor_stream_file_bf16( - target, source->path, source->file_offset, source->elements, - job->error, sizeof(job->error))) { - if (!job->error[0]) - snprintf(job->error, sizeof(job->error), - "invalid BF16 streaming destination"); + if (started_thread[index]) + pthread_join(matrix_threads[index], NULL); + if (!matrix_jobs[index].ok) { job->ok = 0; - break; + if (!job->error[0]) + snprintf(job->error, sizeof(job->error), "%s", + matrix_jobs[index].error); } - job->bytes += (uint64_t)source->elements * sizeof(uint16_t); + job->bytes += matrix_jobs[index].bytes; } job->seconds = stream_now() - started; return job->ok; @@ -689,6 +831,236 @@ static void *read_stream_layer_thread(void *opaque) { return NULL; } +/* --- H3_ATTENTION_CACHE: streamed int8 QKV/OUT (+ optionally FC1/FC2) -- */ + +/* Every record holds all four matrices so the same cache file serves both + * modes; attention-only mode simply never touches the FC1/FC2 portion. */ +static uint64_t attention_cache_record_bytes(void) { + return (uint64_t)INNER * 3 * HIDDEN /* qkv_int8 */ + + (uint64_t)INNER * 3 * sizeof(float) /* qkv_scales */ + + (uint64_t)HIDDEN * INNER /* out_int8 */ + + (uint64_t)HIDDEN * sizeof(float) /* out_scales */ + + (uint64_t)FFN * 2 * HIDDEN /* fc1_int8 */ + + (uint64_t)FFN * 2 * sizeof(float) /* fc1_scales */ + + (uint64_t)HIDDEN * FFN /* fc2_int8 */ + + (uint64_t)HIDDEN * sizeof(float); /* fc2_scales */ +} + +typedef struct { + uint64_t qkv_int8, qkv_scales, out_int8, out_scales; + uint64_t fc1_int8, fc1_scales, fc2_int8, fc2_scales; +} h3_attention_cache_offsets; + +static void attention_cache_offsets(unsigned block, + h3_attention_cache_offsets *off) { + uint64_t base = (uint64_t)sizeof(h3_attention_cache_header) + + (uint64_t)block * attention_cache_record_bytes(); + off->qkv_int8 = base; + off->qkv_scales = off->qkv_int8 + (uint64_t)INNER * 3 * HIDDEN; + off->out_int8 = off->qkv_scales + (uint64_t)INNER * 3 * sizeof(float); + off->out_scales = off->out_int8 + (uint64_t)HIDDEN * INNER; + off->fc1_int8 = off->out_scales + (uint64_t)HIDDEN * sizeof(float); + off->fc1_scales = off->fc1_int8 + (uint64_t)FFN * 2 * HIDDEN; + off->fc2_int8 = off->fc1_scales + (uint64_t)FFN * 2 * sizeof(float); + off->fc2_scales = off->fc2_int8 + (uint64_t)HIDDEN * FFN; +} + +/* Infers FL2VA vs Ref2VA from weight_directory's own last two path + * components - h3.c's dit_path selection (h3.c: h3_path(ctx->model_dir, + * ref2va ? "Ref2VA/transformer" : "FL2VA/transformer")) always ends the + * directory this way, so this is reading back a convention h3.c's own + * caller wrote, not guessing at one. Any other layout (test fixtures, a + * directory pointed at directly with a different name) reports UNKNOWN, + * which disables the model_kind check entirely rather than risk a false + * mismatch. */ +static h3_cache_model_kind detect_model_kind(const char *weight_directory) { + if (!weight_directory) return H3_CACHE_MODEL_UNKNOWN; + size_t length = strlen(weight_directory); + static const char fl2va_suffix[] = "FL2VA/transformer"; + static const char ref2va_suffix[] = "Ref2VA/transformer"; + if (length >= sizeof(ref2va_suffix) - 1 && + !strcmp(weight_directory + length - (sizeof(ref2va_suffix) - 1), + ref2va_suffix)) + return H3_CACHE_MODEL_REF2VA; + if (length >= sizeof(fl2va_suffix) - 1 && + !strcmp(weight_directory + length - (sizeof(fl2va_suffix) - 1), + fl2va_suffix)) + return H3_CACHE_MODEL_FL2VA; + return H3_CACHE_MODEL_UNKNOWN; +} + +static const char *cache_model_kind_name(h3_cache_model_kind kind) { + switch (kind) { + case H3_CACHE_MODEL_FL2VA: return "FL2VA"; + case H3_CACHE_MODEL_REF2VA: return "Ref2VA"; + default: return "an unrecognized model layout"; + } +} + +/* expected_kind is H3_CACHE_MODEL_UNKNOWN when the checkpoint directory + * about to be loaded does not match the standard FL2VA/Ref2VA layout (see + * detect_model_kind()) - the model_kind check is then skipped entirely, + * same as a cache whose own model_kind is UNKNOWN (written by a tool that + * predates this field, or given a non-standard transformer directory + * itself). expected_model_id may be NULL to skip the fingerprint check. */ +static int attention_cache_validate(const char *path, int need_mlp, + h3_cache_model_kind expected_kind, + const uint8_t *expected_model_id, + char *error, size_t error_size) { + FILE *file = fopen(path, "rb"); + if (!file) { + fail(error, error_size, "cannot open attention cache %s: %s", path, + strerror(errno)); + return 0; + } + h3_attention_cache_header header; + int ok = fread(&header, sizeof(header), 1, file) == 1; + if (ok) { + ok = memcmp(header.magic, H3_ATTENTION_CACHE_MAGIC, 4) == 0 && + header.version == H3_ATTENTION_CACHE_VERSION && + header.block_count == H3_DIT_BLOCKS && + header.hidden == HIDDEN && header.inner == INNER && + header.ffn == FFN; + } + if (ok) { + if (fseeko(file, 0, SEEK_END) != 0) { + ok = 0; + } else { + off_t size = ftello(file); + uint64_t wanted = (uint64_t)sizeof(header) + + (uint64_t)H3_DIT_BLOCKS * attention_cache_record_bytes(); + ok = size >= 0 && (uint64_t)size >= wanted; + } + } + fclose(file); + (void)need_mlp; /* every v3 cache carries FC1/FC2 too; kept for clarity */ + if (!ok) { + fail(error, error_size, + "attention cache %s does not match this build (wrong model, " + "quantization version, or a truncated file) - rebuild it with " + "build_attention_cache", path); + return 0; + } + if (expected_kind != H3_CACHE_MODEL_UNKNOWN && + header.model_kind != H3_CACHE_MODEL_UNKNOWN && + header.model_kind != (uint32_t)expected_kind) { + fail(error, error_size, + "%s generation cannot use a %s attention cache (%s) - rebuild " + "it against the matching transformer directory", + cache_model_kind_name(expected_kind), + cache_model_kind_name((h3_cache_model_kind)header.model_kind), + path); + return 0; + } + if (expected_model_id && + memcmp(header.model_id, (const uint8_t[32]){0}, 32) != 0 && + memcmp(header.model_id, expected_model_id, 32) != 0) { + fprintf(stderr, + "h3: warning: attention cache %s's model fingerprint does " + "not match the loaded checkpoint (different weights, a " + "LoRA baked in after the cache was built, or weights " + "rewritten since) - results may be wrong; rebuild the " + "cache if unsure\n", path); + } + return 1; +} + +static int allocate_attention_slot(h3_dit *dit, h3_dit_attention_slot *slot, + char *error, size_t error_size) { + slot->qkv_int8 = h3_gpu_tensor_new_i8(dit->gpu, (size_t)INNER * 3 * HIDDEN); + slot->qkv_scales = h3_gpu_tensor_new_f32(dit->gpu, INNER * 3); + slot->out_int8 = h3_gpu_tensor_new_i8(dit->gpu, (size_t)HIDDEN * INNER); + slot->out_scales = h3_gpu_tensor_new_f32(dit->gpu, HIDDEN); + int ok = slot->qkv_int8 && slot->qkv_scales && slot->out_int8 && + slot->out_scales; + if (ok && dit->mlp_stream) { + slot->fc1_int8 = h3_gpu_tensor_new_i8(dit->gpu, (size_t)FFN * 2 * HIDDEN); + slot->fc1_scales = h3_gpu_tensor_new_f32(dit->gpu, FFN * 2); + slot->fc2_int8 = h3_gpu_tensor_new_i8(dit->gpu, (size_t)HIDDEN * FFN); + slot->fc2_scales = h3_gpu_tensor_new_f32(dit->gpu, HIDDEN); + ok = slot->fc1_int8 && slot->fc1_scales && slot->fc2_int8 && + slot->fc2_scales; + } + if (!ok) { + fail(error, error_size, "cannot allocate int8 attention slot: %s", + h3_gpu_error(dit->gpu)); + return 0; + } + return 1; +} + +static void free_attention_slot(h3_dit_attention_slot *slot) { + free_tensor(&slot->qkv_int8); + free_tensor(&slot->qkv_scales); + free_tensor(&slot->out_int8); + free_tensor(&slot->out_scales); + free_tensor(&slot->fc1_int8); + free_tensor(&slot->fc1_scales); + free_tensor(&slot->fc2_int8); + free_tensor(&slot->fc2_scales); +} + +typedef struct { + h3_dit *dit; + unsigned layer; + unsigned slot; + int ok; + uint64_t bytes; + double seconds; + char error[512]; +} h3_dit_attention_job; + +static int read_attention_layer(h3_dit_attention_job *job) { + h3_dit_attention_slot *slot = &job->dit->attention_slots[job->slot]; + const char *path = job->dit->attention_cache_path; + h3_attention_cache_offsets off; + attention_cache_offsets(job->layer, &off); + double started = stream_now(); + job->ok = + h3_gpu_tensor_stream_file_i8(slot->qkv_int8, path, off.qkv_int8, + (size_t)INNER * 3 * HIDDEN, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_f32(slot->qkv_scales, path, off.qkv_scales, + INNER * 3, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_i8(slot->out_int8, path, off.out_int8, + (size_t)HIDDEN * INNER, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_f32(slot->out_scales, path, off.out_scales, + HIDDEN, + job->error, sizeof(job->error)); + uint64_t bytes = job->ok ? + (uint64_t)INNER * 3 * HIDDEN + (uint64_t)INNER * 3 * sizeof(float) + + (uint64_t)HIDDEN * INNER + (uint64_t)HIDDEN * sizeof(float) : 0; + if (job->ok && job->dit->mlp_stream) { + job->ok = + h3_gpu_tensor_stream_file_i8(slot->fc1_int8, path, off.fc1_int8, + (size_t)FFN * 2 * HIDDEN, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_f32(slot->fc1_scales, path, + off.fc1_scales, FFN * 2, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_i8(slot->fc2_int8, path, off.fc2_int8, + (size_t)HIDDEN * FFN, + job->error, sizeof(job->error)) && + h3_gpu_tensor_stream_file_f32(slot->fc2_scales, path, + off.fc2_scales, HIDDEN, + job->error, sizeof(job->error)); + if (job->ok) + bytes += (uint64_t)FFN * 2 * HIDDEN + + (uint64_t)FFN * 2 * sizeof(float) + + (uint64_t)HIDDEN * FFN + (uint64_t)HIDDEN * sizeof(float); + } + job->bytes = job->ok ? bytes : 0; + job->seconds = stream_now() - started; + return job->ok; +} + +static void *read_attention_layer_thread(void *opaque) { + read_attention_layer(opaque); + return NULL; +} + static int quantize_block_mlp(h3_dit *dit, h3_dit_block *block, char *error, size_t error_size) { block->fc1_int8 = h3_gpu_tensor_new_i8( @@ -1253,6 +1625,16 @@ static int load_core(h3_dit *dit, h3_dit_progress progress, void *opaque, error, error_size) || !prepare_stream_layer(dit, index, error, error_size)) return 0; + } else if (dit->attention_stream && dit->mlp_stream) { + /* QKV/OUT/FC1/FC2 are all streamed; only norms stay resident. */ + if (!load_block_norms(dit, &dit->blocks[index], prefix, + error, error_size)) return 0; + } else if (dit->attention_stream) { + if (!load_block_norms_and_mlp(dit, &dit->blocks[index], prefix, + error, error_size)) return 0; + if (dit->int8_mlp && + !quantize_block_mlp(dit, &dit->blocks[index], + error, error_size)) return 0; } else { if (!load_block(dit, &dit->blocks[index], prefix, error, error_size)) return 0; @@ -1292,6 +1674,30 @@ static int load_core(h3_dit *dit, h3_dit_progress progress, void *opaque, dit->stream_bytes += job.bytes; dit->stream_read_seconds += job.seconds; } + if (dit->attention_stream) { + if (!allocate_attention_slot(dit, &dit->attention_slots[0], + error, error_size) || + !allocate_attention_slot(dit, &dit->attention_slots[1], + error, error_size)) return 0; + unsigned first = first_active_block(dit); + if (first == H3_DIT_BLOCKS) { + fail(error, error_size, + "attention cache stream has no active DiT block"); + return 0; + } + h3_dit_attention_job job = { + .dit = dit, .layer = first, .slot = 0 + }; + if (!read_attention_layer(&job)) { + fail(error, error_size, + "cannot prime attention int8 stream: %s", job.error); + return 0; + } + dit->attn_ready_layer = first; + dit->attn_ready_slot = 0; + dit->attn_stream_bytes += job.bytes; + dit->attn_stream_read_seconds += job.seconds; + } dit->video_patch_w = f2(dit, "video_patch_proj.weight", HIDDEN, VIDEO_PATCH, error, error_size); dit->video_patch_b = f1(dit, "video_patch_proj.bias", HIDDEN, @@ -1640,6 +2046,93 @@ static h3_dit *load_dit(const char *weight_directory, !use_slower_bf16_attention_output && dit->sequence >= 128 && h3_gpu_has_int8_mlp(dit->gpu); + const char *attention_cache_path = getenv("H3_ATTENTION_CACHE"); + const char *attention_cache_dir = getenv("H3_ATTENTION_CACHE_DIR"); + h3_cache_model_kind model_kind = detect_model_kind(weight_directory); + char *selected_cache_path = NULL; + if (attention_cache_path && *attention_cache_path && + attention_cache_dir && *attention_cache_dir) { + fail(error, error_size, + "set only one of H3_ATTENTION_CACHE or H3_ATTENTION_CACHE_DIR"); + goto failed; + } + if (attention_cache_dir && *attention_cache_dir) { + if (model_kind == H3_CACHE_MODEL_UNKNOWN) { + fail(error, error_size, + "H3_ATTENTION_CACHE_DIR needs the standard FL2VA/Ref2VA " + "transformer layout to auto-select a cache file - point " + "H3_ATTENTION_CACHE at a specific file instead for a " + "non-standard directory"); + goto failed; + } + const char *name = model_kind == H3_CACHE_MODEL_REF2VA ? + "ref2va.cache" : "fl2va.cache"; + size_t length = strlen(attention_cache_dir) + strlen(name) + 2; + selected_cache_path = malloc(length); + if (!selected_cache_path) { + fail(error, error_size, "out of memory building cache path"); + goto failed; + } + snprintf(selected_cache_path, length, "%s/%s", attention_cache_dir, + name); + attention_cache_path = selected_cache_path; + } + if (attention_cache_path && *attention_cache_path && dit->ssd_streaming) { + /* --ssd-streaming was passed explicitly - H3_ATTENTION_CACHE is + * very likely just left set in the environment for other, longer + * runs, so honor --ssd-streaming instead of forcing the user to + * unset the env var by hand every time for a quick test. The + * other conditions below that disable int8_qkv/int8_attention_out + * (a short sequence, no int8 GPU support, --use-slower-bf16-qkv/ + * -attention-output) stay hard errors - those are configurations + * that can't serve the cache at all, not an explicit alternate + * mode the user chose. */ + fprintf(stderr, + "h3: --ssd-streaming set; ignoring H3_ATTENTION_CACHE\n"); + attention_cache_path = NULL; + free(selected_cache_path); + selected_cache_path = NULL; + } + if (attention_cache_path && *attention_cache_path) { + if (!dit->int8_qkv || !dit->int8_attention_out) { + fail(error, error_size, + "H3_ATTENTION_CACHE needs the int8 QKV/attention-output " + "path (unavailable here: --ssd-streaming, " + "--use-slower-bf16-qkv/-attention-output, a short " + "sequence, or a GPU without the int8 path all disable it)"); + free(selected_cache_path); + goto failed; + } + const char *stream_mlp = getenv("H3_INT8_STREAM_MLP"); + int want_mlp_stream = stream_mlp && *stream_mlp && strcmp(stream_mlp, "0"); + if (want_mlp_stream && !dit->int8_mlp) { + fail(error, error_size, + "H3_INT8_STREAM_MLP needs the int8 MLP path (unavailable " + "here: --ssd-streaming, --use-slower-bf16-mlp, or a GPU " + "without the int8 path disable it)"); + free(selected_cache_path); + goto failed; + } + uint8_t expected_model_id[32]; + h3_weight_store_fingerprint(dit->weights, expected_model_id); + int cache_ok = attention_cache_validate( + attention_cache_path, want_mlp_stream, model_kind, + expected_model_id, error, error_size); + if (!cache_ok) { + free(selected_cache_path); + goto failed; + } + dit->attention_cache_path = selected_cache_path ? selected_cache_path : + strdup(attention_cache_path); + selected_cache_path = NULL; + if (!dit->attention_cache_path) { + fail(error, error_size, "out of memory copying cache path"); + goto failed; + } + dit->attention_stream = 1; + dit->mlp_stream = want_mlp_stream; + } + free(selected_cache_path); dit->use_slower_row_major_attention_output = use_slower_row_major_attention_output; dit->use_slower_unfused_int8_inputs = @@ -2175,7 +2668,7 @@ static int encode_forward(h3_dit *dit, int step, int begin, int submit, if (evaluate_core) { unsigned command_blocks = disable_command_split ? 0 : command_block_interval(dit); - if (dit->ssd_streaming) command_blocks = 0; + if (dit->ssd_streaming || dit->attention_stream) command_blocks = 0; unsigned completed_blocks = 0; int carried_attention_adaln = 0; int carried_attention_input_quantized = 0; @@ -2256,6 +2749,51 @@ static int encode_forward(h3_dit *dit, int step, int begin, int submit, } stream_started = 1; } + h3_dit_attention_job attn_job; + pthread_t attn_thread; + int attn_started = 0; + if (dit->attention_stream) { + if (dit->attn_ready_layer != block || + dit->attn_ready_slot > 1) { + fail(error, error_size, + "attention int8 stream expected block %u, has " + "block %u", block, dit->attn_ready_layer); + return 0; + } + h3_dit_attention_slot *slot = + &dit->attention_slots[dit->attn_ready_slot]; + streamed_weight = dit->blocks[block]; + streamed_weight.qkv_int8 = slot->qkv_int8; + streamed_weight.qkv_scales = slot->qkv_scales; + streamed_weight.out_int8 = slot->out_int8; + streamed_weight.out_scales = slot->out_scales; + if (dit->mlp_stream) { + streamed_weight.fc1_int8 = slot->fc1_int8; + streamed_weight.fc1_scales = slot->fc1_scales; + streamed_weight.fc2_int8 = slot->fc2_int8; + streamed_weight.fc2_scales = slot->fc2_scales; + } + weight = &streamed_weight; + + unsigned future = next_active_block(dit, block); + if (future == H3_DIT_BLOCKS) + future = first_active_block(dit); + attn_job = (h3_dit_attention_job){ + .dit = dit, + .layer = future, + .slot = dit->attn_ready_slot ^ 1u + }; + int thread_error = pthread_create( + &attn_thread, NULL, read_attention_layer_thread, + &attn_job); + if (thread_error) { + fail(error, error_size, + "cannot start attention int8 prefetch for block " + "%u: %s", future, strerror(thread_error)); + return 0; + } + attn_started = 1; + } int block_ok = run_block( dit, block, step, weight, fused_token_adaln, fused_attention_input_quantized, @@ -2265,6 +2803,7 @@ static int encode_forward(h3_dit *dit, int step, int begin, int submit, error, error_size); if (!block_ok) { if (stream_started) (void)pthread_join(stream_thread, NULL); + if (attn_started) (void)pthread_join(attn_thread, NULL); return 0; } completed_blocks++; @@ -2299,6 +2838,33 @@ static int encode_forward(h3_dit *dit, int step, int begin, int submit, OP(h3_gpu_begin(dit->gpu), "continue after streamed DiT block"); } + if (attn_started) { + int gpu_ok = gpu_op(dit, h3_gpu_submit(dit->gpu), + error, error_size, + "submit attention-streamed DiT block"); + double wait_started = stream_now(); + int join_error = pthread_join(attn_thread, NULL); + dit->attn_stream_wait_seconds += stream_now() - wait_started; + if (!gpu_ok) return 0; + if (join_error) { + fail(error, error_size, + "cannot join attention int8 prefetch: %s", + strerror(join_error)); + return 0; + } + dit->attn_stream_bytes += attn_job.bytes; + dit->attn_stream_read_seconds += attn_job.seconds; + if (!attn_job.ok) { + fail(error, error_size, + "cannot stream attention block %u: %s", + attn_job.layer, attn_job.error); + return 0; + } + dit->attn_ready_layer = attn_job.layer; + dit->attn_ready_slot = attn_job.slot; + OP(h3_gpu_begin(dit->gpu), + "continue after attention-streamed DiT block"); + } } if (use_token_reduction && token_reduction_end == H3_DIT_BLOCKS && @@ -3018,6 +3584,9 @@ void h3_dit_free(h3_dit *dit) { free_block(&dit->blocks[block]); free_block(&dit->stream_slots[0]); free_block(&dit->stream_slots[1]); + free_attention_slot(&dit->attention_slots[0]); + free_attention_slot(&dit->attention_slots[1]); + free(dit->attention_cache_path); free_tensor(&dit->final_norm); free_tensor(&dit->final_video_w); free_tensor(&dit->final_video_b); free_tensor(&dit->final_audio_w); free_tensor(&dit->final_audio_b); @@ -3055,6 +3624,17 @@ void h3_dit_free(h3_dit *dit) { ? gib / dit->stream_read_seconds : 0.0, dit->stream_wait_seconds); } + if (dit->attention_stream && getenv("H3_PROFILE")) { + double gib = (double)dit->attn_stream_bytes / + (1024.0 * 1024.0 * 1024.0); + fprintf(stderr, + "h3: attention int8 stream %.3f GiB read in %.3fs " + "(%.3f GiB/s), unhidden wait %.3fs\n", + gib, dit->attn_stream_read_seconds, + dit->attn_stream_read_seconds > 0.0 + ? gib / dit->attn_stream_read_seconds : 0.0, + dit->attn_stream_wait_seconds); + } h3_gpu_free(dit->gpu); h3_weight_store_free(dit->weights); h3_layout_free(&dit->layout); diff --git a/h3_gpu.h b/h3_gpu.h index 3a47cc35..61b95f04 100644 --- a/h3_gpu.h +++ b/h3_gpu.h @@ -65,6 +65,20 @@ int h3_gpu_tensor_read_file_bf16(h3_gpu_tensor *tensor, const char *path, int h3_gpu_tensor_stream_file_bf16(h3_gpu_tensor *tensor, const char *path, uint64_t file_offset, size_t elements, char *error, size_t error_size); +/* I8/F32 counterparts of the two functions above, for a pre-quantized + * attention cache (int8 weights + f32 scales) rather than raw BF16 shards. */ +int h3_gpu_tensor_read_file_i8(h3_gpu_tensor *tensor, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size); +int h3_gpu_tensor_stream_file_i8(h3_gpu_tensor *tensor, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size); +int h3_gpu_tensor_read_file_f32(h3_gpu_tensor *tensor, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size); +int h3_gpu_tensor_stream_file_f32(h3_gpu_tensor *tensor, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size); void h3_gpu_tensor_free(h3_gpu_tensor *tensor); size_t h3_gpu_tensor_elements(const h3_gpu_tensor *tensor); h3_gpu_dtype h3_gpu_tensor_dtype(const h3_gpu_tensor *tensor); @@ -75,6 +89,8 @@ int h3_gpu_tensor_read_f32_range(const h3_gpu_tensor *tensor, size_t elements); int h3_gpu_tensor_read_bf16(const h3_gpu_tensor *tensor, uint16_t *values, size_t elements); +int h3_gpu_tensor_read_i8(const h3_gpu_tensor *tensor, int8_t *values, + size_t elements); int h3_gpu_tensor_write_f32(h3_gpu_tensor *tensor, const float *values, size_t elements); int h3_gpu_tensor_write_f32_range(h3_gpu_tensor *tensor, diff --git a/h3_gpu.m b/h3_gpu.m index c61d04c6..8a5bc944 100644 --- a/h3_gpu.m +++ b/h3_gpu.m @@ -719,24 +719,29 @@ int h3_gpu_has_int8_mlp(const h3_gpu *opaque) { sizeof(float), H3_GPU_F32, "F32"); } -static int h3_gpu_tensor_read_file_bf16_mode( +/* Shared by the BF16/I8/F32 file-read and SSD-stream entry points below: + * fills an EXISTING tensor's buffer from a file, unlike h3_gpu_tensor_load_file + * which allocates a new one. item_size/dtype select which element width and + * H3_GPU_* tag the request must match. */ +static int h3_gpu_tensor_read_file_mode( h3_gpu_tensor *opaque, const char *path, uint64_t file_offset, size_t elements, - int uncached, + size_t item_size, h3_gpu_dtype dtype, + const char *label, int uncached, char *error, size_t error_size) { if (error && error_size) error[0] = '\0'; if (!opaque || !path || !*path || - TENSOR(opaque).dtype != H3_GPU_BF16 || + TENSOR(opaque).dtype != dtype || elements != TENSOR(opaque).elements || - elements > SIZE_MAX / sizeof(uint16_t) || file_offset > INT64_MAX) { + elements > SIZE_MAX / item_size || file_offset > INT64_MAX) { if (error && error_size) - snprintf(error, error_size, "invalid BF16 file read request"); + snprintf(error, error_size, "invalid %s file read request", label); return 0; } - size_t bytes = elements * sizeof(uint16_t); + size_t bytes = elements * item_size; if ((uint64_t)bytes > (uint64_t)INT64_MAX - file_offset) { if (error && error_size) - snprintf(error, error_size, "BF16 file read range overflows"); + snprintf(error, error_size, "%s file read range overflows", label); return 0; } int descriptor = open(path, O_RDONLY | O_CLOEXEC); @@ -761,9 +766,9 @@ static int h3_gpu_tensor_read_file_bf16_mode( if (count <= 0) { int detail = count < 0 ? errno : 0; if (error && error_size) { - snprintf(error, error_size, "cannot read BF16 payload from %s: %s", - path, detail ? strerror(detail) : - "unexpected end of file"); + snprintf(error, error_size, "cannot read %s payload from %s: %s", + label, path, detail ? strerror(detail) : + "unexpected end of file"); } close(descriptor); return 0; @@ -777,15 +782,49 @@ static int h3_gpu_tensor_read_file_bf16_mode( int h3_gpu_tensor_read_file_bf16(h3_gpu_tensor *opaque, const char *path, uint64_t file_offset, size_t elements, char *error, size_t error_size) { - return h3_gpu_tensor_read_file_bf16_mode( - opaque, path, file_offset, elements, 0, error, error_size); + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(uint16_t), H3_GPU_BF16, + "BF16", 0, error, error_size); } int h3_gpu_tensor_stream_file_bf16(h3_gpu_tensor *opaque, const char *path, uint64_t file_offset, size_t elements, char *error, size_t error_size) { - return h3_gpu_tensor_read_file_bf16_mode( - opaque, path, file_offset, elements, 1, error, error_size); + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(uint16_t), H3_GPU_BF16, + "BF16", 1, error, error_size); +} + +int h3_gpu_tensor_read_file_i8(h3_gpu_tensor *opaque, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size) { + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(int8_t), H3_GPU_I8, + "I8", 0, error, error_size); +} + +int h3_gpu_tensor_stream_file_i8(h3_gpu_tensor *opaque, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size) { + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(int8_t), H3_GPU_I8, + "I8", 1, error, error_size); +} + +int h3_gpu_tensor_read_file_f32(h3_gpu_tensor *opaque, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size) { + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(float), H3_GPU_F32, + "F32", 0, error, error_size); +} + +int h3_gpu_tensor_stream_file_f32(h3_gpu_tensor *opaque, const char *path, + uint64_t file_offset, size_t elements, + char *error, size_t error_size) { + return h3_gpu_tensor_read_file_mode( + opaque, path, file_offset, elements, sizeof(float), H3_GPU_F32, + "F32", 1, error, error_size); } void h3_gpu_tensor_free(h3_gpu_tensor *tensor) { @@ -837,6 +876,14 @@ int h3_gpu_tensor_read_bf16(const h3_gpu_tensor *tensor, uint16_t *values, return 1; } +int h3_gpu_tensor_read_i8(const h3_gpu_tensor *tensor, int8_t *values, + size_t elements) { + if (!tensor || !values || TENSOR(tensor).dtype != H3_GPU_I8 || + elements > TENSOR(tensor).elements) return 0; + memcpy(values, TENSOR(tensor).buffer.contents, elements * sizeof(int8_t)); + return 1; +} + int h3_gpu_tensor_write_f32(h3_gpu_tensor *tensor, const float *values, size_t elements) { return h3_gpu_tensor_write_f32_range(tensor, 0, values, elements); diff --git a/h3_weights.c b/h3_weights.c index 9a0224da..ec161e36 100644 --- a/h3_weights.c +++ b/h3_weights.c @@ -5,6 +5,7 @@ #include #include #include +#include struct h3_weight_store { h3_st_header *headers; @@ -126,6 +127,33 @@ size_t h3_weight_store_shards(const h3_weight_store *store) { return store ? store->count : 0; } +static uint64_t fingerprint_fnv1a64(uint64_t hash, const void *data, + size_t bytes) { + const unsigned char *p = data; + for (size_t i = 0; i < bytes; i++) { + hash ^= p[i]; + hash *= 1099511628211ull; + } + return hash; +} + +void h3_weight_store_fingerprint(const h3_weight_store *store, uint8_t out[32]) { + memset(out, 0, 32); + if (!store) return; + uint64_t hash = 1469598103934665603ull; + for (size_t index = 0; index < store->count; index++) { + const char *path = store->headers[index].path; + uint64_t size = store->headers[index].file_size; + int64_t mtime = 0; + struct stat status; + if (path && stat(path, &status) == 0) mtime = (int64_t)status.st_mtime; + hash = fingerprint_fnv1a64(hash, path, path ? strlen(path) : 0); + hash = fingerprint_fnv1a64(hash, &size, sizeof(size)); + hash = fingerprint_fnv1a64(hash, &mtime, sizeof(mtime)); + } + memcpy(out, &hash, sizeof(hash)); +} + const h3_st_tensor *h3_weight_find(const h3_weight_store *store, const char *name, const h3_st_header **header) { diff --git a/h3_weights.h b/h3_weights.h index 89fbecc6..4b71a359 100644 --- a/h3_weights.h +++ b/h3_weights.h @@ -16,6 +16,15 @@ h3_weight_store *h3_weight_store_open(const char *directory, void h3_weight_store_free(h3_weight_store *store); size_t h3_weight_store_shards(const h3_weight_store *store); +/* A lightweight, non-cryptographic fingerprint of this store's shards + * (each shard's path, file size, and mtime, hashed together) - cheap + * enough to recompute at every run (a stat() per shard, no payload + * reads), unlike hashing the ~18GB of actual weight bytes. Meant for + * staleness detection (a cache built against different/rewritten + * weights), not as a security property. Always fills all 32 bytes of + * `out` (zero-padded beyond the 8 bytes the hash occupies). */ +void h3_weight_store_fingerprint(const h3_weight_store *store, uint8_t out[32]); + const h3_st_tensor *h3_weight_find(const h3_weight_store *store, const char *name, const h3_st_header **header);