tools: Omega-QVLA dit_svdquant_v1 pack → E0M3/UE4M3 converter + format doc + synthetic fixture

.gitignore

@@ -44,3 +44,4 @@ internal-tests/
 internal-docs/
 notes/
 third_party/cutlass
+/artifacts/omega_e0m3/*.pt

docs/omega_pack_e0m3.md

@@ -0,0 +1,160 @@
+# Omega-QVLA pack to FlashRT E0M3 format
+
+This document describes the **Milestone 1 offline toolchain** shipped in this
+repository. It converts rank-0 Omega-QVLA `dit_svdquant_v1` records into the
+packed E0M3/UE4M3 weight format consumed by FlashRT's existing SM110 GEMM.
+
+This change does not connect the artifact to a frontend or pipeline. It adds no
+runtime route, server, CUDA graph, CMake source, binding, or public Python API.
+Runtime-consumer results mentioned in development discussions were produced on
+an external experimental branch and are not capabilities of this repository.
+
+## Included tools
+
+- `tools/convert_omega_pack_e0m3.py`: validates and converts a pack.
+- `tools/check_omega_e0m3_layer.py`: CPU emulation diagnostics and a Thor
+  kernel round-trip that directly consumes converted `packed` and `sfb` data.
+- `tools/gen_omega_pack_fixture.py`: creates a four-record synthetic fixture.
+
+## Input container
+
+The input is a plain `torch.save` mapping loadable with
+`torch.load(..., weights_only=True)`. A full inspected pi0.5 pack contains 252
+records plus `__meta__`: 126 expert records and 126 PaliGemma records.
+
+Each selected record must have `format == "dit_svdquant_v1"` and satisfy:
+
+| Field | Required contract |
+|---|---|
+| `weight_res_q` | finite floating tensor `[out_features, in_features]` |
+| `rank` | exactly `0` |
+| `lowrank_A`, `lowrank_B` | `[out_features, 0]`, `[in_features, 0]` |
+| `weight_bits`, `a_bits` | both `4` |
+| `act_scale_table` | finite positive floating tensor `[steps, in_features]` |
+| `duquant_rotation_blocks` | finite `[in_features / 64, 64, 64]` |
+| `duquant_rotation_perm` | int64 permutation of `[0, in_features)` |
+| `duquant_rotation_out_blocks` | finite `[out_features / 64, 64, 64]` |
+| `in_features`, `out_features` | exact match for the weight tensor |
+
+Both feature dimensions must be positive multiples of 64. The converter
+validates every selected record before importing the CUDA extension or writing
+an artifact. Unsupported records are errors; they are never skipped.
+
+For an unfiltered non-fixture pack, the default coverage gate is 252 records.
+`--layer-regex` and fixture packs expect every selected record. Use
+`--expected-records` when intentionally converting a different complete pack.
+
+## Output container
+
+The output format is `omega_e0m3_v1`, schema version 1:
+
+```text
+{
+  "format": "omega_e0m3_v1",
+  "schema_version": 1,
+  "source_pack_meta": {...},
+  "source_record_count": int,
+  "selected_record_count": int,
+  "selected_layers": [str, ...],
+  "fold": "none" | "mean" | "actnorm",
+  "weights": {
+    layer: {
+      "packed": uint8[N, K / 2],
+      "sfb": uint8[sfa_size_bytes(N, K, True)],
+      "N": int,
+      "K": int,
+    }
+  },
+  "aux": {layer: {...}}
+}
+```
+
+The converter writes through a sibling temporary file and atomically replaces
+the destination only after all selected records convert successfully. The
+`weights`, `aux`, and `selected_layers` sets must be identical.
+
+## Conversion math
+
+Omega stores `weight_res_q` as a dequantized fp16 tensor already in the
+rotated and permuted domain. No GPTQ bitstream decoding is required.
+
+FlashRT converts each `[N, K]` weight with:
+
+```text
+quantize_e0m3_dynamic_sfa_fp16(weight, packed, sfb, N, K, is_sfb=True)
+```
+
+This produces packed 4-bit E0M3 elements and tile-interleaved UE4M3 per-16
+scales. SFB storage is allocated with
+`flash_rt_fp4.sfa_size_bytes(N, K, True)` and zero-initialized because the
+tile-interleaved layout contains padding.
+
+The default `--fold none` strategy intentionally does not fold the
+per-channel activation table into the weight. Thor experiments on four real
+layers measured E0M3-vs-fp16 cosine around 0.9924-0.9932, comparable to the
+source fake-quant path. `mean` and `actnorm` remain ablation modes; they are not
+the recommended artifact format.
+
+## Reproduction
+
+### CPU fixture and emulation
+
+```bash
+python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt
+
+python tools/convert_omega_pack_e0m3.py \
+    --pack /tmp/fixture_pack.pt \
+    --out /tmp/unused.pt \
+    --validate-only
+
+python tools/check_omega_e0m3_layer.py \
+    --pack /tmp/fixture_pack.pt \
+    --mode emulate
+```
+
+The validation command performs no CUDA work and does not create its output.
+
+### Thor artifact round-trip
+
+```bash
+python tools/convert_omega_pack_e0m3.py \
+    --pack /tmp/fixture_pack.pt \
+    --out /tmp/fixture_e0m3.pt \
+    --fold none
+
+python tools/check_omega_e0m3_layer.py \
+    --pack /tmp/fixture_pack.pt \
+    --artifact /tmp/fixture_e0m3.pt \
+    --mode kernel \
+    --min-artifact-cos 0.98
+```
+
+With `--artifact`, the checker quantizes only the activation. It loads the
+weight's `packed` and `sfb` tensors directly from the converted artifact,
+validates their shape and byte count, runs the GEMM, and fails if cosine is
+below the requested threshold. This is the converter round-trip gate.
+
+### Full pack
+
+```bash
+python tools/convert_omega_pack_e0m3.py \
+    --pack /path/to/quantized.pt \
+    --out /path/to/pi05_e0m3.pt \
+    --fold none
+```
+
+Without a layer filter, this command fails unless all 252 records validate and
+convert. A representative artifact layer should then be checked with the same
+`--artifact --mode kernel` command above. Release evidence must include both
+the direct packed/SFB GEMM result and `252/252` coverage.
+
+## Scope and roadmap
+
+Milestone 1 is limited to conversion, schema validation, a synthetic fixture,
+CPU emulation, and direct artifact GEMM verification.
+
+A future runtime PR may load these artifacts in a pi0.5 Thor frontend, apply
+the DuQuant input/output rotations, and establish end-to-end accuracy and
+latency gates. Such a PR must independently add and test its frontend,
+pipeline, graph-lifecycle, and serving contracts. None of that runtime surface
+is provided here.

tests/test_omega_e0m3_tools.py

@@ -0,0 +1,203 @@
+import copy
+import importlib.util
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_tool(name: str):
+    path = ROOT / "tools" / f"{name}.py"
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+converter = _load_tool("convert_omega_pack_e0m3")
+checker = _load_tool("check_omega_e0m3_layer")
+
+
+def _record(n: int = 64, k: int = 64) -> dict:
+    return {
+        "format": "dit_svdquant_v1",
+        "weight_res_q": torch.randn(n, k, dtype=torch.float16),
+        "lowrank_A": torch.empty(n, 0, dtype=torch.float16),
+        "lowrank_B": torch.empty(k, 0, dtype=torch.float16),
+        "act_scale_table": torch.ones(2, k),
+        "duquant_rotation_blocks": torch.eye(64).half().view(1, 64, 64),
+        "duquant_rotation_perm": torch.arange(k, dtype=torch.int64),
+        "duquant_rotation_out_blocks": torch.eye(64).half().view(1, 64, 64),
+        "weight_bits": 4,
+        "a_bits": 4,
+        "rank": 0,
+        "in_features": k,
+        "out_features": n,
+    }
+
+
+def test_converter_accepts_complete_rank_zero_record():
+    assert converter.validate_record("layer", _record()) == (64, 64)
+
+
+@pytest.mark.parametrize(
+    ("mutation", "message"),
+    [
+        (lambda r: r.update(format="other"), "unsupported format"),
+        (lambda r: r.update(rank=1), "only rank=0"),
+        (lambda r: r.update(in_features=128), "metadata shape mismatch"),
+        (lambda r: r.update(act_scale_table=torch.ones(2, 32)),
+         "act_scale_table"),
+        (lambda r: r.update(duquant_rotation_blocks=torch.eye(64).half()),
+         "input rotation shape"),
+        (lambda r: r.update(duquant_rotation_perm=torch.zeros(64, dtype=torch.int64)),
+         "not a permutation"),
+    ],
+)
+def test_converter_rejects_incomplete_or_inconsistent_records(mutation, message):
+    record = copy.deepcopy(_record())
+    mutation(record)
+    with pytest.raises(ValueError, match=message):
+        converter.validate_record("layer", record)
+
+
+def test_full_pack_count_defaults_to_252_but_fixture_and_subset_are_exact():
+    assert converter.expected_record_count(
+        {}, layer_regex="", selected_count=251, explicit=None) == 252
+    assert converter.expected_record_count(
+        {"recipe": "synthetic fixture (test)"}, layer_regex="", selected_count=4,
+        explicit=None) == 4
+    assert converter.expected_record_count(
+        {}, layer_regex="layer.0", selected_count=1, explicit=None) == 1
+
+
+def test_validate_only_runs_without_cuda_and_writes_no_artifact(tmp_path):
+    pack_path = tmp_path / "fixture_pack.pt"
+    out_path = tmp_path / "must_not_exist.pt"
+    torch.save({
+        "__meta__": {"recipe": "synthetic fixture (test)"},
+        "layer": _record(),
+    }, pack_path)
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(ROOT / "tools" / "convert_omega_pack_e0m3.py"),
+            "--pack", str(pack_path),
+            "--out", str(out_path),
+            "--validate-only",
+        ],
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode == 0, result.stderr
+    assert "validated 1/1 records" in result.stdout
+    assert not out_path.exists()
+
+
+def test_validate_only_fails_before_cuda_for_bad_records(tmp_path):
+    pack_path = tmp_path / "bad_pack.pt"
+    out_path = tmp_path / "must_not_exist.pt"
+    record = _record()
+    record["rank"] = 1
+    torch.save({
+        "__meta__": {"recipe": "synthetic fixture (test)"},
+        "layer": record,
+    }, pack_path)
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(ROOT / "tools" / "convert_omega_pack_e0m3.py"),
+            "--pack", str(pack_path),
+            "--out", str(out_path),
+            "--validate-only",
+        ],
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode == 2
+    assert "only rank=0 is supported" in result.stderr
+    assert not out_path.exists()
+
+
+def test_converter_refuses_to_overwrite_source_pack(tmp_path):
+    pack_path = tmp_path / "fixture_pack.pt"
+    torch.save({
+        "__meta__": {"recipe": "synthetic fixture (test)"},
+        "layer": _record(),
+    }, pack_path)
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(ROOT / "tools" / "convert_omega_pack_e0m3.py"),
+            "--pack", str(pack_path),
+            "--out", str(pack_path),
+        ],
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    assert result.returncode == 2
+    assert "must not overwrite" in result.stderr
+    loaded = torch.load(pack_path, map_location="cpu", weights_only=True)
+    assert "layer" in loaded
+
+
+def test_artifact_schema_requires_complete_direct_packed_sfb_coverage():
+    class FakeFp4:
+        @staticmethod
+        def sfa_size_bytes(n, k, is_sfb):
+            assert is_sfb
+            return 32
+
+    artifact = {
+        "format": "omega_e0m3_v1",
+        "schema_version": 1,
+        "fold": "none",
+        "selected_record_count": 1,
+        "selected_layers": ["layer"],
+        "weights": {
+            "layer": {
+                "packed": torch.zeros(64, 32, dtype=torch.uint8),
+                "sfb": torch.zeros(32, dtype=torch.uint8),
+                "N": 64,
+                "K": 64,
+            }
+        },
+        "aux": {"layer": {"fold": "none"}},
+    }
+    entry, aux = checker.validate_artifact(
+        artifact, "layer", 64, 64, FakeFp4())
+    assert entry["packed"].shape == (64, 32)
+    assert aux["fold"] == "none"
+
+    broken = copy.deepcopy(artifact)
+    broken["aux"] = {}
+    with pytest.raises(ValueError, match="do not cover"):
+        checker.validate_artifact(broken, "layer", 64, 64, FakeFp4())
+
+
+def test_docs_and_gitignore_only_describe_milestone_one_surface():
+    docs = (ROOT / "docs" / "omega_pack_e0m3.md").read_text()
+    gitignore = (ROOT / ".gitignore").read_text().splitlines()
+
+    for nonexistent in (
+        "omega_e0m3_linear.py",
+        "check_omega_e0m3_consumer.py",
+        "serve_omega_e0m3.py",
+        "omega_e0m3_graph.py",
+        "check_omega_e0m3_graph_smoke.py",
+        "start_e0m3_server.sh",
+    ):
+        assert nonexistent not in docs
+    assert "*.pt" not in gitignore
+    assert "/artifacts/omega_e0m3/*.pt" in gitignore