234 lines
11 KiB
C++
234 lines
11 KiB
C++
// Validation of the 1024_cascade HR stage against the PyTorch reference
|
|
// (scripts/dump_cascade_reference.py -> dumps/reference_cascade.gguf):
|
|
//
|
|
// 1. shape-decoder upsample(x4): the LR slat -> 512^3 candidate coord set,
|
|
// and its quantized 64^3 HR scaffold (subdivision-boundary set tolerance).
|
|
// 2. HR (1024-model) flow forward at t=500 on the HR scaffold (tight gate).
|
|
// 3. final 1024^3 decode of the reference HR slat: per-level features,
|
|
// subdivision, and 7-channel output.
|
|
// (The 12-step HR sampler is env-gated (TRELLIS2_CASCADE_SAMPLE) — 24 forwards
|
|
// at ~40k tokens is impractical on CPU; the shared Euler loop is already
|
|
// validated by test_slat / test_ss_sample.)
|
|
//
|
|
// usage: test_cascade <slat_512_f32> <slat_1024_f32> <shape_dec_f32> <reference_cascade.gguf>
|
|
// exits 77 (ctest SKIP) when inputs are missing.
|
|
|
|
#include "trellis2.h"
|
|
#include "parity.hpp"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <unordered_set>
|
|
#include <vector>
|
|
|
|
static bool file_exists(const std::string & p) {
|
|
std::ifstream f(p);
|
|
return f.good();
|
|
}
|
|
|
|
// [L*4] (batch,x,y,z) reference coords -> [L*3] int32 (x,y,z)
|
|
static std::vector<int32_t> coords_xyz(const std::vector<float> & c4) {
|
|
const size_t L = c4.size() / 4;
|
|
std::vector<int32_t> out(L * 3);
|
|
for (size_t v = 0; v < L; ++v) {
|
|
out[v * 3] = (int32_t) c4[v * 4 + 1];
|
|
out[v * 3 + 1] = (int32_t) c4[v * 4 + 2];
|
|
out[v * 3 + 2] = (int32_t) c4[v * 4 + 3];
|
|
}
|
|
return out;
|
|
}
|
|
|
|
static uint64_t vkey(int32_t x, int32_t y, int32_t z) {
|
|
return ((uint64_t) (uint32_t) x << 40) | ((uint64_t) (uint32_t) y << 20) | (uint64_t) (uint32_t) z;
|
|
}
|
|
|
|
// symmetric-difference fraction of two coord sets ([*3] int32)
|
|
static double set_diff_frac(const std::vector<int32_t> & a, const std::vector<int32_t> & b) {
|
|
std::unordered_set<uint64_t> sa, sb;
|
|
for (size_t i = 0; i < a.size(); i += 3) sa.insert(vkey(a[i], a[i + 1], a[i + 2]));
|
|
for (size_t i = 0; i < b.size(); i += 3) sb.insert(vkey(b[i], b[i + 1], b[i + 2]));
|
|
size_t only = 0;
|
|
for (uint64_t k : sa) if (!sb.count(k)) ++only;
|
|
for (uint64_t k : sb) if (!sa.count(k)) ++only;
|
|
const size_t big = sa.size() > sb.size() ? sa.size() : sb.size();
|
|
return big ? (double) only / (double) big : 0.0;
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc < 5) {
|
|
std::fprintf(stderr,
|
|
"usage: %s <slat_512_f32.gguf> <slat_1024_f32.gguf> <shape_dec_f32.gguf> <reference_cascade.gguf>\n",
|
|
argv[0]);
|
|
return 2;
|
|
}
|
|
const std::string flow512_path = argv[1];
|
|
const std::string flow1024_path = argv[2];
|
|
const std::string dec_path = argv[3];
|
|
const std::string ref_path = argv[4];
|
|
if (!file_exists(flow512_path) || !file_exists(flow1024_path) ||
|
|
!file_exists(dec_path) || !file_exists(ref_path)) {
|
|
std::fprintf(stderr, "missing input file(s), skipping\n");
|
|
return 77;
|
|
}
|
|
|
|
t2_parity::baseline ref;
|
|
if (!ref.open(ref_path)) {
|
|
std::fprintf(stderr, "failed to open %s\n", ref_path.c_str());
|
|
return 1;
|
|
}
|
|
|
|
std::vector<float> cond512, cond1024, coords32f, lr_slat, up_coordsf, hr_coordsf, hr_noise;
|
|
if (!ref.load("cond_512", cond512) || !ref.load("cond_1024", cond1024) ||
|
|
!ref.load("coords32", coords32f) || !ref.load("lr_slat", lr_slat) ||
|
|
!ref.load("up_coords", up_coordsf) || !ref.load("hr_coords", hr_coordsf) ||
|
|
!ref.load("hr_noise", hr_noise)) {
|
|
std::fprintf(stderr, "reference missing required tensors\n");
|
|
return 1;
|
|
}
|
|
const int L32 = (int) (coords32f.size() / 4);
|
|
const int Lhr = (int) (hr_coordsf.size() / 4);
|
|
const int Lkv512 = (int) (cond512.size() / 1024);
|
|
const int Lkv1024 = (int) (cond1024.size() / 1024);
|
|
std::vector<int32_t> coords32 = coords_xyz(coords32f);
|
|
std::vector<int32_t> ref_up_coords = coords_xyz(up_coordsf);
|
|
std::vector<int32_t> hr_coords = coords_xyz(hr_coordsf);
|
|
std::printf("reference: scaffold %d voxels, upsample %zu, HR %d voxels, cond %d/%d tokens\n",
|
|
L32, up_coordsf.size() / 4, Lhr, Lkv512, Lkv1024);
|
|
|
|
std::string err;
|
|
int n_fail = 0;
|
|
|
|
// ── 1. shape-decoder upsample(x4): coord set + quantized 64^3 scaffold ────
|
|
trellis2_shape_dec_model * dec = trellis2_shape_dec_load(dec_path, true, &err, "cpu");
|
|
if (!dec) { std::fprintf(stderr, "dec load failed: %s\n", err.c_str()); return 1; }
|
|
std::printf("dec backend: %s\n", trellis2_shape_dec_backend_name(dec));
|
|
{
|
|
std::vector<int32_t> got_up;
|
|
if (!trellis2_shape_dec_upsample(dec, lr_slat.data(), L32, coords32.data(),
|
|
/*upsample_times*/ 4, got_up, &err)) {
|
|
std::fprintf(stderr, "upsample failed: %s\n", err.c_str());
|
|
trellis2_shape_dec_free(dec);
|
|
return 1;
|
|
}
|
|
const double up_frac = set_diff_frac(got_up, ref_up_coords);
|
|
std::printf("[upsample coords] got %zu vs ref %zu, sym-diff %.4f%% -> %s\n",
|
|
got_up.size() / 3, ref_up_coords.size() / 3, 100.0 * up_frac,
|
|
up_frac <= 5e-4 ? "OK" : "FAIL");
|
|
if (up_frac > 5e-4) ++n_fail;
|
|
|
|
// quantize both to 64^3 and compare the dedup'd set that drives the HR flow
|
|
auto quant = [](const std::vector<int32_t> & c) {
|
|
std::unordered_set<uint64_t> s;
|
|
std::vector<int32_t> out;
|
|
for (size_t i = 0; i < c.size(); i += 3) {
|
|
int32_t x = (int32_t) ((c[i] + 0.5f) / 512.0f * 64.0f);
|
|
int32_t y = (int32_t) ((c[i + 1] + 0.5f) / 512.0f * 64.0f);
|
|
int32_t z = (int32_t) ((c[i + 2] + 0.5f) / 512.0f * 64.0f);
|
|
if (s.insert(vkey(x, y, z)).second) { out.push_back(x); out.push_back(y); out.push_back(z); }
|
|
}
|
|
return out;
|
|
};
|
|
std::vector<int32_t> my_hr = quant(got_up);
|
|
const double hr_frac = set_diff_frac(my_hr, hr_coords);
|
|
std::printf("[hr scaffold] got %zu vs ref %d, sym-diff %.4f%% -> %s\n",
|
|
my_hr.size() / 3, Lhr, 100.0 * hr_frac, hr_frac <= 5e-4 ? "OK" : "FAIL");
|
|
if (hr_frac > 5e-4) ++n_fail;
|
|
}
|
|
|
|
// ── 2. HR flow forward at t=500 on the reference HR scaffold ──────────────
|
|
{
|
|
// Force CPU: at 10k+ HR tokens the attention exceeds the exact-path
|
|
// threshold and uses flash, and GPU flash (F16-MMA) is ~1e-2 vs the
|
|
// exact fp32 reference. CPU flash is exact-matching (~3e-4), so it
|
|
// gives a meaningful tight gate. TRELLIS2_CASCADE_GPU overrides.
|
|
const char * dev = std::getenv("TRELLIS2_CASCADE_GPU") ? nullptr : "cpu";
|
|
trellis2_slat_flow_model * flow = trellis2_slat_flow_load(flow1024_path, true, &err, dev);
|
|
if (!flow) { std::fprintf(stderr, "1024 flow load failed: %s\n", err.c_str()); trellis2_shape_dec_free(dec); return 1; }
|
|
std::printf("1024 flow backend: %s\n", trellis2_slat_flow_backend_name(flow));
|
|
std::vector<float> got((size_t) Lhr * 32), want;
|
|
if (!trellis2_slat_flow_forward(flow, hr_noise.data(), Lhr, hr_coords.data(), 500.0f,
|
|
cond1024.data(), Lkv1024, 1024, got.data(), &err)) {
|
|
std::fprintf(stderr, "HR flow forward failed: %s\n", err.c_str());
|
|
trellis2_slat_flow_free(flow); trellis2_shape_dec_free(dec); return 1;
|
|
}
|
|
ref.load("hr_flow_t500_out", want);
|
|
t2_parity::compare_stats st;
|
|
t2_parity::compare(got, want, "hr_flow_t500_out", 2e-3, 2e-3, &st);
|
|
if (st.rel_l2 > 3e-3) { std::printf(" -> HR forward rel_l2 %.4g > 3e-3, FAIL\n", st.rel_l2); ++n_fail; }
|
|
|
|
// optional: full HR sampler (expensive at ~40k tokens; off by default)
|
|
if (std::getenv("TRELLIS2_CASCADE_SAMPLE")) {
|
|
trellis2_ss_sampler_params P;
|
|
P.steps = 12; P.guidance_strength = 7.5f; P.guidance_rescale = 0.5f;
|
|
P.guidance_interval_min = 0.6f; P.guidance_interval_max = 1.0f; P.rescale_t = 3.0f;
|
|
std::vector<float> sampled((size_t) Lhr * 32), wslat;
|
|
if (trellis2_slat_flow_sample(flow, Lhr, hr_coords.data(), cond1024.data(), Lkv1024, 1024,
|
|
&P, hr_noise.data(), true, sampled.data(), &err)) {
|
|
ref.load("hr_slat", wslat);
|
|
t2_parity::compare_stats ss;
|
|
t2_parity::compare(sampled, wslat, "hr_slat(sampled)", 5e-2, 5e-2, &ss);
|
|
}
|
|
}
|
|
trellis2_slat_flow_free(flow);
|
|
}
|
|
|
|
// ── 3. final 1024^3 decode of the reference HR slat ──────────────────────
|
|
// The 1024^3 decode is the SAME decoder validated exactly at the 512 tier
|
|
// (test_slat, levels 0-4 at rel-L2 5e-7) applied to more voxels, and the
|
|
// end-to-end demo exercises it directly. It also transiently needs ~14 GB of
|
|
// host RAM (the finest up-block's conv output is held in both the graph and
|
|
// the readback), so it is gated behind TRELLIS2_CASCADE_DECODE — enable it
|
|
// on a big-RAM box to also gate out7 here.
|
|
if (!std::getenv("TRELLIS2_CASCADE_DECODE")) {
|
|
trellis2_shape_dec_free(dec);
|
|
std::printf("\n(1024^3 decode gate skipped; set TRELLIS2_CASCADE_DECODE to enable)\n");
|
|
std::printf("total failures: %d\nRESULT: %s\n", n_fail, n_fail ? "FAIL" : "PASS");
|
|
return n_fail ? 1 : 0;
|
|
}
|
|
std::vector<float> hr_slat;
|
|
if (!ref.load("hr_slat", hr_slat)) {
|
|
std::fprintf(stderr, "reference missing hr_slat\n");
|
|
trellis2_shape_dec_free(dec); return 1;
|
|
}
|
|
// Compare only the final 7-channel output (taps=nullptr): the per-level
|
|
// intermediates would hold multiple GB at 1024^3 and the decoder's level
|
|
// logic is already validated exactly at the 512 tier (test_slat, same
|
|
// decoder). out7 is the load-bearing gate.
|
|
std::vector<float> out_feats;
|
|
std::vector<int32_t> out_coords;
|
|
if (!trellis2_shape_dec_decode(dec, hr_slat.data(), Lhr, hr_coords.data(),
|
|
out_feats, out_coords, nullptr, &err)) {
|
|
std::fprintf(stderr, "HR decode failed: %s\n", err.c_str());
|
|
trellis2_shape_dec_free(dec); return 1;
|
|
}
|
|
trellis2_shape_dec_free(dec);
|
|
|
|
std::printf("HR decode: %zu output voxels\n", out_coords.size() / 3);
|
|
std::vector<float> ref_out7;
|
|
ref.load("out7", ref_out7);
|
|
const size_t a = out_feats.size(), b = ref_out7.size();
|
|
const size_t big = a > b ? a : b, sml = a > b ? b : a;
|
|
if (a != b) {
|
|
const double frac = (double) (big - sml) / (double) big;
|
|
if (frac > 5e-4) {
|
|
std::printf("[out7] SIZE MISMATCH %zu vs %zu (%.3f%%) -> FAIL\n", a, b, 100.0 * frac);
|
|
++n_fail;
|
|
} else {
|
|
std::printf("[out7] near-match (%zu vs %zu, %.4f%% subdivision boundary flip) -> OK\n",
|
|
a, b, 100.0 * frac);
|
|
}
|
|
} else {
|
|
t2_parity::compare_stats st;
|
|
t2_parity::compare(out_feats, ref_out7, "out7", 2e-3, 2e-3, &st);
|
|
if (st.rel_l2 > 2e-2) { std::printf(" -> out7 rel_l2 %.4g > 2e-2, FAIL\n", st.rel_l2); ++n_fail; }
|
|
}
|
|
|
|
std::printf("\ntotal failures: %d\n", n_fail);
|
|
std::printf("RESULT: %s\n", n_fail ? "FAIL" : "PASS");
|
|
return n_fail ? 1 : 0;
|
|
}
|