This is an automated email from the ASF dual-hosted git repository.
akaashrp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 344d7de3bc [Feat][Web] Support per-parameter tensor cache encoding
(#20136)
344d7de3bc is described below
commit 344d7de3bcd352bbfac8e243529739b69b1bfc17
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Wed Aug 26 05:29:32 2026 -0700
[Feat][Web] Support per-parameter tensor cache encoding (#20136)
---
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 acd86c4445..e06035bc4a 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 @@ import shutil
# 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 4de1b6c985..41a975c907 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()
