From deefc8fcf9732d013c37d1b7461f919a7406f2ea Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 19 Jul 2026 06:26:03 -0400 Subject: [PATCH] [FEAT][WEB] Support per-parameter tensor cache encoding --- python/tvm/contrib/tvmjs.py | 60 +++++++++---- tests/python/contrib/test_tvmjs.py | 132 +++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 18 deletions(-) diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py index acd86c44455a..e06035bc4ad6 100644 --- a/python/tvm/contrib/tvmjs.py +++ b/python/tvm/contrib/tvmjs.py @@ -17,6 +17,7 @@ # ruff: noqa: E501, F401 """Namespace to store utilities for building web runtime.""" +import copy import hashlib import json import math @@ -26,7 +27,6 @@ # pylint: disable=unused-import import sys from collections.abc import Iterator, Mapping -from types import GeneratorType from typing import Any, Optional, Union import numpy as np @@ -203,7 +203,7 @@ def dump_tensor_cache( params: Mapping[str, np.ndarray | tvm.runtime.Tensor] | Iterator[tuple[str, np.ndarray | tvm.runtime.Tensor]], cache_dir: str, - encode_format="f32-to-bf16", + encode_format: str | Mapping[str, str] = "f32-to-bf16", meta_data=None, shard_cap_mb=32, show_progress: bool = True, @@ -222,8 +222,12 @@ def dump_tensor_cache( cache_dir: str The path to the cache - encode_format: {"f32-to-bf16", "raw"} - Encoding format. + encode_format: Union[ + {"f32-to-bf16", "raw"}, + Mapping[str, {"f32-to-bf16", "raw"}] + ] + Encoding format. A mapping selects the format by parameter name and may + use ``"*"`` as the fallback for names not explicitly listed. meta_data: json-compatible-struct or Callable[[], Any] Extra meta_data to be stored in the cache json file, @@ -239,11 +243,25 @@ def dump_tensor_cache( If the cache already exists, update the cache. When set to False, it will overwrite the existing files. """ - if encode_format not in ("raw", "f32-to-bf16"): - raise ValueError(f"Invalie encode_format {encode_format}") + if isinstance(encode_format, str): + if encode_format not in ("raw", "f32-to-bf16"): + raise ValueError(f"Invalid encode_format {encode_format}") + elif not isinstance(encode_format, Mapping): + raise TypeError("encode_format must be a string or parameter-name mapping") + else: + for name, param_format in encode_format.items(): + if param_format not in ("raw", "f32-to-bf16"): + raise ValueError(f"Invalid encode_format for parameter {name}: {param_format}") + + def resolve_encode_format(name): + if isinstance(encode_format, str): + return encode_format + param_format = encode_format.get(name, encode_format.get("*")) + if param_format not in ("raw", "f32-to-bf16"): + raise ValueError(f"Invalid encode_format for parameter {name}: {param_format}") + return param_format records = [] - from_generator = isinstance(params, GeneratorType) total_bytes = 0 counter = 0 max_out_length = 0 @@ -251,8 +269,6 @@ def dump_tensor_cache( if not os.path.exists(cache_dir): os.makedirs(cache_dir) - f32_to_bf16_triggered = False - print(f"Start storing to cache {cache_dir}") shard_cap_nbytes = shard_cap_mb * (1 << 20) @@ -268,8 +284,9 @@ def dump_tensor_cache( cache_dir, "params_shard", shard_cap_nbytes, initial_shard_records=records ) - param_generator = params.items() if not from_generator else params + param_generator = params.items() if isinstance(params, Mapping) else params for k, origin_v in param_generator: + param_encode_format = resolve_encode_format(k) shape = list(origin_v.shape) v = origin_v if not isinstance(v, np.ndarray): @@ -286,9 +303,8 @@ def dump_tensor_cache( total_bytes += math.prod(v.shape) * np.dtype(v.dtype).itemsize # convert fp32 to bf16 - if encode_format == "f32-to-bf16" and dtype == "float32": + if param_encode_format == "f32-to-bf16" and dtype == "float32": data = _convert_f32_to_bf16(v).tobytes() - f32_to_bf16_triggered = True else: data = v.tobytes() @@ -297,7 +313,7 @@ def dump_tensor_cache( name=k, shape=shape, dtype=dtype, - encode_format=encode_format, + encode_format=param_encode_format, allow_update=update_if_exists, ) @@ -317,17 +333,25 @@ def dump_tensor_cache( f"\nAll finished, {shard_manager.counter} total shards committed, record saved to {nd_cache_json}" ) - if f32_to_bf16_triggered: - for shard in records: + b16_nd_cache_json = os.path.join(cache_dir, "tensor-cache-b16.json") + has_f32_to_bf16 = any( + item["format"] == "f32-to-bf16" and item["dtype"] == "float32" + for shard in records + for item in shard["records"] + ) + if has_f32_to_bf16: + b16_records = copy.deepcopy(records) + for shard in b16_records: for item in shard["records"]: - if item["dtype"] == "float32": + if item["format"] == "f32-to-bf16" and item["dtype"] == "float32": item["format"] = "raw" item["dtype"] = "bfloat16" - b16_nd_cache_json = os.path.join(cache_dir, "tensor-cache-b16.json") # also dump a file that contains bf16 with open(b16_nd_cache_json, "w") as outfile: - json.dump({"metadata": meta_data, "records": records}, outfile, indent=4) + json.dump({"metadata": meta_data, "records": b16_records}, outfile, indent=4) print(f"Also saved a bf16 record to {b16_nd_cache_json}") + elif os.path.exists(b16_nd_cache_json): + os.remove(b16_nd_cache_json) def load_tensor_cache(cachepath: str, device: tvm.runtime.Device): diff --git a/tests/python/contrib/test_tvmjs.py b/tests/python/contrib/test_tvmjs.py index 4de1b6c9850c..41a975c90771 100644 --- a/tests/python/contrib/test_tvmjs.py +++ b/tests/python/contrib/test_tvmjs.py @@ -17,6 +17,8 @@ """Test contrib.tvmjs""" +import json +import os import tempfile import numpy as np @@ -60,5 +62,135 @@ def test_save_load_float8(dtype): np.testing.assert_array_equal(arr, after_roundtrip) +def _records_by_name(manifest_path): + with open(manifest_path, encoding="utf-8") as source: + manifest = json.load(source) + records = { + record["name"]: record for shard in manifest["records"] for record in shard["records"] + } + return records, manifest + + +def test_dump_tensor_cache_supports_per_parameter_encoding_roundtrip(): + raw = np.array([0.1234567, -0.7654321], dtype="float32") + compressed = np.array([0.2345678, -0.8765432], dtype="float32") + + with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir: + tvmjs.dump_tensor_cache( + {"raw": raw, "compressed": compressed}, + temp_dir, + encode_format={"raw": "raw", "*": "f32-to-bf16"}, + ) + cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu()) + b16_cache, _ = tvmjs.load_tensor_cache( + os.path.join(temp_dir, "tensor-cache-b16.json"), tvm.cpu() + ) + + records, _ = _records_by_name(os.path.join(temp_dir, "tensor-cache.json")) + b16_records, _ = _records_by_name(os.path.join(temp_dir, "tensor-cache-b16.json")) + + assert records["raw"]["format"] == "raw" + assert records["raw"]["dtype"] == "float32" + assert records["compressed"]["format"] == "f32-to-bf16" + assert records["compressed"]["dtype"] == "float32" + assert b16_records["raw"]["format"] == "raw" + assert b16_records["raw"]["dtype"] == "float32" + assert b16_records["compressed"]["format"] == "raw" + assert b16_records["compressed"]["dtype"] == "bfloat16" + np.testing.assert_array_equal(cache["raw"].numpy(), raw) + np.testing.assert_allclose(cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3) + np.testing.assert_array_equal(b16_cache["raw"].numpy(), raw) + np.testing.assert_allclose(b16_cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3) + + +def test_dump_tensor_cache_supports_generator_input(): + params = ( + item + for item in [ + ("raw", np.arange(4, dtype="float32")), + ("compressed", np.linspace(-1, 1, 4, dtype="float32")), + ] + ) + + with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir: + tvmjs.dump_tensor_cache( + params, + temp_dir, + encode_format={"raw": "raw", "*": "f32-to-bf16"}, + ) + cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu()) + + np.testing.assert_array_equal(cache["raw"].numpy(), np.arange(4, dtype="float32")) + np.testing.assert_allclose( + cache["compressed"].numpy(), + np.linspace(-1, 1, 4, dtype="float32"), + rtol=4e-3, + atol=1e-3, + ) + + +def test_dump_tensor_cache_updates_mixed_encoding_manifests(): + original_raw = np.array([1.0, 2.0], dtype="float32") + updated_raw = np.array([3.0, 4.0], dtype="float32") + compressed = np.array([0.2345678, -0.8765432], dtype="float32") + + with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir: + tvmjs.dump_tensor_cache( + {"raw": original_raw, "compressed": compressed}, + temp_dir, + encode_format={"raw": "raw", "*": "f32-to-bf16"}, + ) + _, old_manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache.json")) + old_md5 = old_manifest["records"][0]["md5sum"] + + tvmjs.dump_tensor_cache( + iter([("raw", updated_raw)]), + temp_dir, + encode_format={"raw": "raw"}, + update_if_exists=True, + ) + + cache, _ = tvmjs.load_tensor_cache(temp_dir, tvm.cpu()) + b16_cache, _ = tvmjs.load_tensor_cache( + os.path.join(temp_dir, "tensor-cache-b16.json"), tvm.cpu() + ) + _, manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache.json")) + _, b16_manifest = _records_by_name(os.path.join(temp_dir, "tensor-cache-b16.json")) + + np.testing.assert_array_equal(cache["raw"].numpy(), updated_raw) + np.testing.assert_array_equal(b16_cache["raw"].numpy(), updated_raw) + np.testing.assert_allclose(cache["compressed"].numpy(), compressed, rtol=4e-3, atol=1e-3) + assert manifest["records"][0]["md5sum"] != old_md5 + assert manifest["records"][0]["md5sum"] == b16_manifest["records"][0]["md5sum"] + + +def test_dump_tensor_cache_removes_stale_b16_manifest(): + with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir: + tvmjs.dump_tensor_cache( + {"compressed": np.ones(2, dtype="float32")}, + temp_dir, + encode_format="f32-to-bf16", + ) + b16_manifest = os.path.join(temp_dir, "tensor-cache-b16.json") + assert os.path.exists(b16_manifest) + + tvmjs.dump_tensor_cache( + {"raw": np.ones(2, dtype="float32")}, + temp_dir, + encode_format="raw", + ) + assert not os.path.exists(b16_manifest) + + +def test_dump_tensor_cache_requires_a_format_for_every_parameter(): + with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir: + with pytest.raises(ValueError, match="parameter arr"): + tvmjs.dump_tensor_cache( + {"arr": np.ones(2, dtype="float32")}, + temp_dir, + encode_format={"other": "raw"}, + ) + + if __name__ == "__main__": tvm.testing.main()