weftcodec 1.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- weftcodec/__init__.py +62 -0
- weftcodec/__main__.py +6 -0
- weftcodec/ai_modes.py +3360 -0
- weftcodec/cli.py +252 -0
- weftcodec/crc32c.py +33 -0
- weftcodec/dcp.py +2108 -0
- weftcodec/dependency.py +224 -0
- weftcodec/errors.py +119 -0
- weftcodec/frame.py +455 -0
- weftcodec/hashing.py +56 -0
- weftcodec/huffman.py +212 -0
- weftcodec/identity.py +141 -0
- weftcodec/limits.py +58 -0
- weftcodec/modes.py +718 -0
- weftcodec/orbax.py +1734 -0
- weftcodec/service.py +959 -0
- weftcodec/torch.py +94 -0
- weftcodec/varint.py +88 -0
- weftcodec/wcb.py +310 -0
- weftcodec-1.2.0.dist-info/METADATA +257 -0
- weftcodec-1.2.0.dist-info/RECORD +24 -0
- weftcodec-1.2.0.dist-info/WHEEL +5 -0
- weftcodec-1.2.0.dist-info/entry_points.txt +2 -0
- weftcodec-1.2.0.dist-info/top_level.txt +1 -0
weftcodec/ai_modes.py
ADDED
|
@@ -0,0 +1,3360 @@
|
|
|
1
|
+
"""AI-mode payload decoders (WEFT-AI-CODEC §5 and §11-§17).
|
|
2
|
+
|
|
3
|
+
Modes implemented here:
|
|
4
|
+
|
|
5
|
+
* 9 ``FLOAT_TENSOR``
|
|
6
|
+
* 10 ``TENSOR_SHARDED``
|
|
7
|
+
* 11 ``WEIGHT_DELTA``
|
|
8
|
+
* 12 ``KV_CACHE``
|
|
9
|
+
* 13 ``ACTIVATION_STREAM``
|
|
10
|
+
* 14 ``SEMANTIC_DELTA``
|
|
11
|
+
* 15 ``TENSOR_TILED``
|
|
12
|
+
* 16 ``SPARSE_DELTA``
|
|
13
|
+
* 17 ``FLOAT_TILED``
|
|
14
|
+
|
|
15
|
+
Every function is an independent decoder written from the normative spec and
|
|
16
|
+
the canonical WCB reader in :mod:`weftcodec.wcb`. No Rust or Go code is loaded
|
|
17
|
+
or shelled out to. Error classes follow the normative validation order; there
|
|
18
|
+
is no approximate or lossy fallback.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import annotations
|
|
22
|
+
|
|
23
|
+
import math
|
|
24
|
+
import struct
|
|
25
|
+
from dataclasses import dataclass
|
|
26
|
+
|
|
27
|
+
from . import huffman, varint, wcb
|
|
28
|
+
from .errors import (
|
|
29
|
+
Corrupt,
|
|
30
|
+
DependencyMismatch,
|
|
31
|
+
DependencyMissing,
|
|
32
|
+
LimitExceeded,
|
|
33
|
+
Malformed,
|
|
34
|
+
NonCanonical,
|
|
35
|
+
)
|
|
36
|
+
from .hashing import (
|
|
37
|
+
DOMAIN_AI_STREAM,
|
|
38
|
+
DOMAIN_TENSOR_LOGICAL,
|
|
39
|
+
DOMAIN_TENSOR_SEMANTIC,
|
|
40
|
+
DOMAIN_TENSOR_SEMANTIC_VALUES,
|
|
41
|
+
DOMAIN_TENSOR_TILE,
|
|
42
|
+
DOMAIN_TENSOR_FLOAT_TILE,
|
|
43
|
+
DOMAIN_TENSOR_FLOAT_TILE_STREAMS,
|
|
44
|
+
DOMAIN_TENSOR_VALUES,
|
|
45
|
+
DOMAIN_TENSOR_SHARDED_PAYLOAD,
|
|
46
|
+
hash_domain,
|
|
47
|
+
)
|
|
48
|
+
from .limits import Limits
|
|
49
|
+
|
|
50
|
+
U64_MAX = (1 << 64) - 1
|
|
51
|
+
U32_MAX = (1 << 32) - 1
|
|
52
|
+
HASH_LEN = 32
|
|
53
|
+
|
|
54
|
+
FLOAT_TENSOR_SCHEMA = "weft.ai.float_tensor.v1"
|
|
55
|
+
FLOAT_TENSOR_CODER = "huffman-canonical-v1"
|
|
56
|
+
TENSOR_SHARDED_SCHEMA = "weft.ai.tensor_sharded.v1"
|
|
57
|
+
WEIGHT_DELTA_SCHEMA = "weft.ai.weight_delta.v1"
|
|
58
|
+
SEMANTIC_DELTA_SCHEMA = "weft.ai.semantic_delta.v1"
|
|
59
|
+
KV_CACHE_SCHEMA = "weft.ai.kv_cache.v1"
|
|
60
|
+
ACTIVATION_STREAM_SCHEMA = "weft.ai.activation_stream.v1"
|
|
61
|
+
TENSOR_TILED_SCHEMA = "weft.ai.tensor_tiled.v1"
|
|
62
|
+
SPARSE_DELTA_SCHEMA = "weft.ai.sparse_delta.v1"
|
|
63
|
+
SEMANTIC_SCHEMA = "weft.tensor.semantic.v1"
|
|
64
|
+
LOGICAL_SCHEMA = "weft.tensor.logical.v1"
|
|
65
|
+
PLANE_CODEC = "huffman-canonical-v1"
|
|
66
|
+
DELTA_OP = "mod-sub"
|
|
67
|
+
SPARSE_CODEC = "zero-run-v1"
|
|
68
|
+
TILE_LAYOUT = "row_major"
|
|
69
|
+
BLOCK_SIZES = (16, 32, 64, 128, 256)
|
|
70
|
+
MAX_BLOCK_LEVELS = 1 << 24
|
|
71
|
+
|
|
72
|
+
FLOAT_FIELD_BITS = {
|
|
73
|
+
"bf16": (1, 8, 7),
|
|
74
|
+
"fp16": (1, 5, 10),
|
|
75
|
+
"fp32": (1, 8, 23),
|
|
76
|
+
"fp8-e4m3": (1, 4, 3),
|
|
77
|
+
"fp8-e5m2": (1, 5, 2),
|
|
78
|
+
}
|
|
79
|
+
|
|
80
|
+
WEIGHT_WORD_BITS = {"bf16": 16, "fp16": 16, "fp32": 32, "q8_0": 8, "q4_0": 4}
|
|
81
|
+
WEIGHT_STORAGE_BYTES = {"bf16": 2, "fp16": 2, "fp32": 4, "q8_0": 1, "q4_0": 1}
|
|
82
|
+
WEIGHT_PLANES = {"bf16": 2, "fp16": 2, "fp32": 4, "q8_0": 1, "q4_0": 1}
|
|
83
|
+
WEIGHT_MANTISSA_BITS = {"bf16": 7, "fp16": 10, "fp32": 23}
|
|
84
|
+
|
|
85
|
+
SHARD_STORAGE_BYTES = {"q8_0": 1, "q4_0": 1}
|
|
86
|
+
|
|
87
|
+
KV_WORD_BYTES = {"bf16": 2, "fp16": 2, "fp32": 4}
|
|
88
|
+
KV_PLANES = {"bf16": 2, "fp16": 2, "fp32": 4}
|
|
89
|
+
|
|
90
|
+
TENSOR_DTYPE_SIZES = {
|
|
91
|
+
"f16": 2,
|
|
92
|
+
"bf16": 2,
|
|
93
|
+
"f32": 4,
|
|
94
|
+
"f64": 8,
|
|
95
|
+
"i8": 1,
|
|
96
|
+
"u8": 1,
|
|
97
|
+
"i16": 2,
|
|
98
|
+
"i32": 4,
|
|
99
|
+
"i64": 8,
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
# ---------------------------------------------------------------------------
|
|
104
|
+
# Small checked-arithmetic and metadata helpers
|
|
105
|
+
# ---------------------------------------------------------------------------
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def _checked_add(left: int, right: int, what: str) -> int:
|
|
109
|
+
value = left + right
|
|
110
|
+
if value > U64_MAX:
|
|
111
|
+
raise LimitExceeded("%s overflows u64" % what)
|
|
112
|
+
return value
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def _checked_mul(left: int, right: int, what: str) -> int:
|
|
116
|
+
value = left * right
|
|
117
|
+
if value > U64_MAX:
|
|
118
|
+
raise LimitExceeded("%s overflows u64" % what)
|
|
119
|
+
return value
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def _check_len(length: int, maximum: int, what: str) -> None:
|
|
123
|
+
if length > maximum:
|
|
124
|
+
raise LimitExceeded("%s length %d exceeds limit %d" % (what, length, maximum))
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def _map(value, what: str) -> dict:
|
|
128
|
+
if not isinstance(value, dict):
|
|
129
|
+
raise Malformed("%s metadata must be a WCB map" % what)
|
|
130
|
+
return value
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _exact_keys(map_value: dict, expected, what: str) -> None:
|
|
134
|
+
for key in map_value:
|
|
135
|
+
if key not in expected:
|
|
136
|
+
raise Malformed("%s metadata has unknown key `%s`" % (what, key))
|
|
137
|
+
for key in expected:
|
|
138
|
+
if key not in map_value:
|
|
139
|
+
raise Malformed("%s metadata is missing key `%s`" % (what, key))
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
def _text(map_value: dict, key: str) -> str:
|
|
143
|
+
value = map_value.get(key)
|
|
144
|
+
if not isinstance(value, str):
|
|
145
|
+
raise Malformed("metadata `%s` must be text" % key)
|
|
146
|
+
return value
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def _u64(map_value: dict, key: str) -> int:
|
|
150
|
+
value = map_value.get(key)
|
|
151
|
+
if not isinstance(value, wcb.U64):
|
|
152
|
+
raise Malformed("metadata `%s` must be u64" % key)
|
|
153
|
+
return value.value
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
def _u32(map_value: dict, key: str) -> int:
|
|
157
|
+
value = _u64(map_value, key)
|
|
158
|
+
if value > U32_MAX:
|
|
159
|
+
raise Malformed("metadata `%s` = %d does not fit u32" % (key, value))
|
|
160
|
+
return value
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def _i64(map_value: dict, key: str) -> int:
|
|
164
|
+
value = map_value.get(key)
|
|
165
|
+
if not isinstance(value, wcb.I64):
|
|
166
|
+
raise Malformed("metadata `%s` must be i64" % key)
|
|
167
|
+
return value.value
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _bytes_of(map_value: dict, key: str) -> bytes:
|
|
171
|
+
value = map_value.get(key)
|
|
172
|
+
if not isinstance(value, (bytes, bytearray)):
|
|
173
|
+
raise Malformed("metadata `%s` must be bytes" % key)
|
|
174
|
+
return bytes(value)
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def _hash_field(map_value: dict, key: str) -> bytes:
|
|
178
|
+
value = _bytes_of(map_value, key)
|
|
179
|
+
if len(value) != HASH_LEN:
|
|
180
|
+
raise Malformed(
|
|
181
|
+
"metadata `%s` must be %d bytes, got %d" % (key, HASH_LEN, len(value))
|
|
182
|
+
)
|
|
183
|
+
return value
|
|
184
|
+
|
|
185
|
+
|
|
186
|
+
def _array_u64(map_value: dict, key: str) -> "list[int]":
|
|
187
|
+
value = map_value.get(key)
|
|
188
|
+
if not isinstance(value, list):
|
|
189
|
+
raise Malformed("metadata `%s` must be an array" % key)
|
|
190
|
+
out = []
|
|
191
|
+
for item in value:
|
|
192
|
+
if not isinstance(item, wcb.U64):
|
|
193
|
+
raise Malformed("metadata `%s` elements must be u64" % key)
|
|
194
|
+
out.append(item.value)
|
|
195
|
+
return out
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def _shape(map_value: dict) -> "list[int]":
|
|
199
|
+
return _array_u64(map_value, "shape")
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
def _split_prefix(body: bytes) -> "tuple[bytes, bytes]":
|
|
203
|
+
"""Split ``u32le(meta_len) || descriptor || payload`` without a length bias."""
|
|
204
|
+
if len(body) < 4:
|
|
205
|
+
raise Malformed("AI body is shorter than its 4-byte metadata length")
|
|
206
|
+
meta_len = int.from_bytes(body[0:4], "little")
|
|
207
|
+
meta_end = 4 + meta_len
|
|
208
|
+
if meta_end > len(body):
|
|
209
|
+
raise Malformed("AI metadata extends past the frame body")
|
|
210
|
+
return body[4:meta_end], body[meta_end:]
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
def _split_exact(body: bytes, decoded_len: int) -> "tuple[bytes, bytes]":
|
|
214
|
+
descriptor, payload = _split_prefix(body)
|
|
215
|
+
if len(payload) != decoded_len:
|
|
216
|
+
raise Malformed(
|
|
217
|
+
"AI payload length %d differs from decoded_len %d"
|
|
218
|
+
% (len(payload), decoded_len)
|
|
219
|
+
)
|
|
220
|
+
return descriptor, payload
|
|
221
|
+
|
|
222
|
+
|
|
223
|
+
def _decode_descriptor(descriptor_bytes: bytes, limits: Limits):
|
|
224
|
+
return wcb.decode(descriptor_bytes, limits)
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
def _u64_vec(values) -> "list[wcb.U64]":
|
|
228
|
+
return [wcb.U64(value) for value in values]
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def _validate_shape(shape: "list[int]", limits: Limits, what: str) -> int:
|
|
232
|
+
if not shape:
|
|
233
|
+
raise Malformed("%s shape rank must be at least 1" % what)
|
|
234
|
+
if len(shape) > limits.max_tensor_rank:
|
|
235
|
+
raise LimitExceeded(
|
|
236
|
+
"%s rank %d exceeds limit %d"
|
|
237
|
+
% (what, len(shape), limits.max_tensor_rank)
|
|
238
|
+
)
|
|
239
|
+
product = 1
|
|
240
|
+
for dim in shape:
|
|
241
|
+
if dim == 0:
|
|
242
|
+
raise Malformed("%s shape dimensions must be greater than zero" % what)
|
|
243
|
+
product = _checked_mul(product, dim, "%s shape product" % what)
|
|
244
|
+
if product > limits.max_tensor_elements:
|
|
245
|
+
raise LimitExceeded(
|
|
246
|
+
"%s element count %d exceeds limit %d"
|
|
247
|
+
% (what, product, limits.max_tensor_elements)
|
|
248
|
+
)
|
|
249
|
+
return product
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
# ---------------------------------------------------------------------------
|
|
253
|
+
# Logical roots and semantic roots
|
|
254
|
+
# ---------------------------------------------------------------------------
|
|
255
|
+
|
|
256
|
+
|
|
257
|
+
def values_hash(data: bytes) -> bytes:
|
|
258
|
+
return hash_domain(DOMAIN_TENSOR_VALUES, data)
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def tensor_logical_root(
|
|
262
|
+
dtype: str, shape: "list[int]", element_count: int, data: bytes, limits: Limits
|
|
263
|
+
) -> bytes:
|
|
264
|
+
record = {
|
|
265
|
+
"schema": LOGICAL_SCHEMA,
|
|
266
|
+
"dtype": dtype,
|
|
267
|
+
"shape": _u64_vec(shape),
|
|
268
|
+
"element_count": wcb.U64(element_count),
|
|
269
|
+
"values_hash": values_hash(data),
|
|
270
|
+
}
|
|
271
|
+
return hash_domain(DOMAIN_TENSOR_LOGICAL, wcb.encode(record, limits))
|
|
272
|
+
|
|
273
|
+
|
|
274
|
+
def sharded_logical_root(
|
|
275
|
+
dtype: str,
|
|
276
|
+
shape: "list[int]",
|
|
277
|
+
words: "list[int]",
|
|
278
|
+
scales: "list[int]",
|
|
279
|
+
zero_points: "list[int]",
|
|
280
|
+
limits: Limits,
|
|
281
|
+
) -> bytes:
|
|
282
|
+
element_count = _validate_shape(shape, limits, "TENSOR_SHARDED")
|
|
283
|
+
if len(words) != element_count:
|
|
284
|
+
raise Malformed(
|
|
285
|
+
"TENSOR_SHARDED tensor has %d words but shape requires %d"
|
|
286
|
+
% (len(words), element_count)
|
|
287
|
+
)
|
|
288
|
+
if len(scales) != len(zero_points):
|
|
289
|
+
raise Malformed("TENSOR_SHARDED scales and zero_points differ in length")
|
|
290
|
+
material = bytearray()
|
|
291
|
+
for word in words:
|
|
292
|
+
if dtype == "q4_0" and not (-8 <= word <= 7):
|
|
293
|
+
raise Malformed("q4_0 word out of signed 4-bit range")
|
|
294
|
+
material.append(word & 0xFF)
|
|
295
|
+
for scale, zero in zip(scales, zero_points):
|
|
296
|
+
material.extend(struct.pack("<I", scale))
|
|
297
|
+
material.extend(struct.pack("<q", zero))
|
|
298
|
+
record = {
|
|
299
|
+
"schema": LOGICAL_SCHEMA,
|
|
300
|
+
"dtype": dtype,
|
|
301
|
+
"shape": _u64_vec(shape),
|
|
302
|
+
"element_count": wcb.U64(element_count),
|
|
303
|
+
"values_hash": values_hash(bytes(material)),
|
|
304
|
+
}
|
|
305
|
+
return hash_domain(DOMAIN_TENSOR_LOGICAL, wcb.encode(record, limits))
|
|
306
|
+
|
|
307
|
+
|
|
308
|
+
def tile_root(dtype: str, dims: "list[int]", tile_bytes: bytes, limits: Limits) -> bytes:
|
|
309
|
+
record = {
|
|
310
|
+
"dtype": dtype,
|
|
311
|
+
"dims": _u64_vec(dims),
|
|
312
|
+
"element_len": wcb.U64(len(tile_bytes)),
|
|
313
|
+
"values_hash": values_hash(tile_bytes),
|
|
314
|
+
}
|
|
315
|
+
return hash_domain(DOMAIN_TENSOR_TILE, wcb.encode(record, limits))
|
|
316
|
+
|
|
317
|
+
|
|
318
|
+
@dataclass(frozen=True)
|
|
319
|
+
class SemanticPolicy:
|
|
320
|
+
scheme: str
|
|
321
|
+
bits: int = 0
|
|
322
|
+
block_size: int = 0
|
|
323
|
+
levels: int = 0
|
|
324
|
+
|
|
325
|
+
def to_wcb(self):
|
|
326
|
+
if self.scheme == "none":
|
|
327
|
+
return {"scheme": "none"}
|
|
328
|
+
if self.scheme == "mantissa_truncate":
|
|
329
|
+
return {"bits": wcb.U64(self.bits), "scheme": "mantissa_truncate"}
|
|
330
|
+
return {
|
|
331
|
+
"block_size": wcb.U64(self.block_size),
|
|
332
|
+
"levels": wcb.U64(self.levels),
|
|
333
|
+
"scheme": "block_quantize",
|
|
334
|
+
}
|
|
335
|
+
|
|
336
|
+
def validate(self, dtype: str, element_count: int) -> None:
|
|
337
|
+
if self.scheme == "none":
|
|
338
|
+
return
|
|
339
|
+
if self.scheme == "mantissa_truncate":
|
|
340
|
+
maximum = WEIGHT_MANTISSA_BITS.get(dtype)
|
|
341
|
+
if maximum is None:
|
|
342
|
+
raise Malformed(
|
|
343
|
+
"mantissa_truncate policy is not defined for dtype %s" % dtype
|
|
344
|
+
)
|
|
345
|
+
if self.bits == 0 or self.bits > maximum:
|
|
346
|
+
raise Malformed(
|
|
347
|
+
"mantissa_truncate bits %d is outside 1..=%d for %s"
|
|
348
|
+
% (self.bits, maximum, dtype)
|
|
349
|
+
)
|
|
350
|
+
return
|
|
351
|
+
if self.block_size == 0:
|
|
352
|
+
raise Malformed("block_quantize block_size must be greater than zero")
|
|
353
|
+
if self.levels == 0 or self.levels > MAX_BLOCK_LEVELS:
|
|
354
|
+
raise Malformed(
|
|
355
|
+
"block_quantize levels %d is outside 1..=%d"
|
|
356
|
+
% (self.levels, MAX_BLOCK_LEVELS)
|
|
357
|
+
)
|
|
358
|
+
if element_count % self.block_size != 0:
|
|
359
|
+
raise Malformed(
|
|
360
|
+
"block_quantize block_size %d does not divide element_count %d"
|
|
361
|
+
% (self.block_size, element_count)
|
|
362
|
+
)
|
|
363
|
+
|
|
364
|
+
|
|
365
|
+
def parse_semantic_policy(value, dtype: str, element_count: int) -> SemanticPolicy:
|
|
366
|
+
map_value = value if isinstance(value, dict) else None
|
|
367
|
+
if map_value is None:
|
|
368
|
+
raise Malformed("semantic policy must be a WCB map")
|
|
369
|
+
scheme = map_value.get("scheme")
|
|
370
|
+
if not isinstance(scheme, str):
|
|
371
|
+
raise Malformed("semantic policy `scheme` must be text")
|
|
372
|
+
if scheme == "none":
|
|
373
|
+
_exact_keys(map_value, ("scheme",), "semantic policy")
|
|
374
|
+
return SemanticPolicy("none")
|
|
375
|
+
if scheme == "mantissa_truncate":
|
|
376
|
+
_exact_keys(map_value, ("bits", "scheme"), "semantic policy")
|
|
377
|
+
policy = SemanticPolicy("mantissa_truncate", bits=_u64(map_value, "bits"))
|
|
378
|
+
policy.validate(dtype, element_count)
|
|
379
|
+
return policy
|
|
380
|
+
if scheme == "block_quantize":
|
|
381
|
+
_exact_keys(map_value, ("block_size", "levels", "scheme"), "semantic policy")
|
|
382
|
+
policy = SemanticPolicy(
|
|
383
|
+
"block_quantize",
|
|
384
|
+
block_size=_u64(map_value, "block_size"),
|
|
385
|
+
levels=_u64(map_value, "levels"),
|
|
386
|
+
)
|
|
387
|
+
policy.validate(dtype, element_count)
|
|
388
|
+
return policy
|
|
389
|
+
raise Malformed(
|
|
390
|
+
"unknown semantic policy scheme `%s` (v1 allows none, mantissa_truncate, "
|
|
391
|
+
"block_quantize)" % scheme
|
|
392
|
+
)
|
|
393
|
+
|
|
394
|
+
|
|
395
|
+
def _signed_word(dtype: str, word: int) -> int:
|
|
396
|
+
bits = WEIGHT_WORD_BITS[dtype]
|
|
397
|
+
sign = 1 << (bits - 1)
|
|
398
|
+
modulus = 1 << bits
|
|
399
|
+
if word >= sign:
|
|
400
|
+
return word - modulus
|
|
401
|
+
return word
|
|
402
|
+
|
|
403
|
+
|
|
404
|
+
def _raw_word(dtype: str, signed: int) -> int:
|
|
405
|
+
return signed & ((1 << WEIGHT_WORD_BITS[dtype]) - 1)
|
|
406
|
+
|
|
407
|
+
|
|
408
|
+
def _exponent_all_ones(dtype: str, word: int) -> bool:
|
|
409
|
+
if dtype in ("bf16", "fp16"):
|
|
410
|
+
return (word & 0x7F80) == 0x7F80
|
|
411
|
+
if dtype == "fp32":
|
|
412
|
+
return (word & 0x7F800000) == 0x7F800000
|
|
413
|
+
return False
|
|
414
|
+
|
|
415
|
+
|
|
416
|
+
def _word_at(dtype: str, data: bytes, index: int) -> int:
|
|
417
|
+
if dtype in ("bf16", "fp16"):
|
|
418
|
+
offset = index * 2
|
|
419
|
+
return int.from_bytes(data[offset : offset + 2], "little")
|
|
420
|
+
if dtype == "fp32":
|
|
421
|
+
offset = index * 4
|
|
422
|
+
return int.from_bytes(data[offset : offset + 4], "little")
|
|
423
|
+
return data[index]
|
|
424
|
+
|
|
425
|
+
|
|
426
|
+
def _push_word(dtype: str, word: int, out: bytearray) -> None:
|
|
427
|
+
if dtype in ("bf16", "fp16"):
|
|
428
|
+
out.extend(int(word & 0xFFFF).to_bytes(2, "little"))
|
|
429
|
+
elif dtype == "fp32":
|
|
430
|
+
out.extend(int(word & 0xFFFFFFFF).to_bytes(4, "little"))
|
|
431
|
+
else:
|
|
432
|
+
out.append(word & 0xFF)
|
|
433
|
+
|
|
434
|
+
|
|
435
|
+
def _transform_bytes(
|
|
436
|
+
dtype: str, data: bytes, element_count: int, policy: SemanticPolicy
|
|
437
|
+
) -> bytes:
|
|
438
|
+
if policy.scheme == "none":
|
|
439
|
+
return bytes(data)
|
|
440
|
+
if policy.scheme == "mantissa_truncate":
|
|
441
|
+
retain = ~((1 << policy.bits) - 1) & ((1 << WEIGHT_WORD_BITS[dtype]) - 1)
|
|
442
|
+
out = bytearray()
|
|
443
|
+
for index in range(element_count):
|
|
444
|
+
word = _word_at(dtype, data, index)
|
|
445
|
+
transformed = (
|
|
446
|
+
word if _exponent_all_ones(dtype, word) else word & retain
|
|
447
|
+
)
|
|
448
|
+
_push_word(dtype, transformed, out)
|
|
449
|
+
return bytes(out)
|
|
450
|
+
block = policy.block_size
|
|
451
|
+
levels = policy.levels
|
|
452
|
+
out = bytearray()
|
|
453
|
+
start = 0
|
|
454
|
+
while start < element_count:
|
|
455
|
+
end = min(start + block, element_count)
|
|
456
|
+
amax = 0
|
|
457
|
+
for index in range(start, end):
|
|
458
|
+
magnitude = abs(_signed_word(dtype, _word_at(dtype, data, index)))
|
|
459
|
+
if magnitude > amax:
|
|
460
|
+
amax = magnitude
|
|
461
|
+
if amax == 0 or levels == 1:
|
|
462
|
+
for _ in range(start, end):
|
|
463
|
+
_push_word(dtype, 0, out)
|
|
464
|
+
else:
|
|
465
|
+
span = 2 * amax
|
|
466
|
+
for index in range(start, end):
|
|
467
|
+
value = _signed_word(dtype, _word_at(dtype, data, index))
|
|
468
|
+
numerator = (value + amax) * (levels - 1) + amax
|
|
469
|
+
k = numerator // span
|
|
470
|
+
if k >= levels:
|
|
471
|
+
k = levels - 1
|
|
472
|
+
level = (span * k + (levels - 1) // 2) // (levels - 1)
|
|
473
|
+
_push_word(dtype, _raw_word(dtype, level - amax), out)
|
|
474
|
+
start = end
|
|
475
|
+
return bytes(out)
|
|
476
|
+
|
|
477
|
+
|
|
478
|
+
def semantic_root(
|
|
479
|
+
dtype: str, shape: "list[int]", data: bytes, policy: SemanticPolicy, limits: Limits
|
|
480
|
+
) -> bytes:
|
|
481
|
+
count = _validate_shape(shape, limits, "semantic tensor")
|
|
482
|
+
expected = _checked_mul(
|
|
483
|
+
count, WEIGHT_STORAGE_BYTES[dtype], "semantic tensor byte length"
|
|
484
|
+
)
|
|
485
|
+
_check_len(expected, limits.max_decoded_bytes, "semantic tensor bytes")
|
|
486
|
+
if expected != len(data):
|
|
487
|
+
raise Malformed(
|
|
488
|
+
"semantic tensor is %d bytes but shape %r x %s requires %d"
|
|
489
|
+
% (len(data), shape, dtype, expected)
|
|
490
|
+
)
|
|
491
|
+
policy.validate(dtype, count)
|
|
492
|
+
if dtype == "q4_0":
|
|
493
|
+
for byte in data:
|
|
494
|
+
if byte > 15:
|
|
495
|
+
raise Corrupt("q4_0 semantic tensor word is outside 0..=15")
|
|
496
|
+
if policy.scheme == "none":
|
|
497
|
+
quantized_values_hash = values_hash(bytes(data))
|
|
498
|
+
else:
|
|
499
|
+
quantized_values_hash = hash_domain(
|
|
500
|
+
DOMAIN_TENSOR_SEMANTIC_VALUES,
|
|
501
|
+
_transform_bytes(dtype, data, count, policy),
|
|
502
|
+
)
|
|
503
|
+
record = {
|
|
504
|
+
"schema": SEMANTIC_SCHEMA,
|
|
505
|
+
"dtype": dtype,
|
|
506
|
+
"shape": _u64_vec(shape),
|
|
507
|
+
"element_count": wcb.U64(count),
|
|
508
|
+
"policy": policy.to_wcb(),
|
|
509
|
+
"quantized_values_hash": quantized_values_hash,
|
|
510
|
+
}
|
|
511
|
+
return hash_domain(DOMAIN_TENSOR_SEMANTIC, wcb.encode(record, limits))
|
|
512
|
+
|
|
513
|
+
|
|
514
|
+
def semantic_record_from_canonical(data: bytes, limits: Limits) -> dict:
|
|
515
|
+
"""Parse and validate a canonical ``weft.tensor.semantic.v1`` record."""
|
|
516
|
+
value = wcb.decode(data, limits)
|
|
517
|
+
map_value = _map(value, "semantic root record")
|
|
518
|
+
_exact_keys(
|
|
519
|
+
map_value,
|
|
520
|
+
("dtype", "element_count", "policy", "quantized_values_hash", "schema", "shape"),
|
|
521
|
+
"semantic root record",
|
|
522
|
+
)
|
|
523
|
+
schema = _text(map_value, "schema")
|
|
524
|
+
if schema != SEMANTIC_SCHEMA:
|
|
525
|
+
raise Malformed("unsupported semantic root schema `%s`" % schema)
|
|
526
|
+
dtype = _weight_dtype(_text(map_value, "dtype"))
|
|
527
|
+
shape = _shape(map_value)
|
|
528
|
+
element_count = _u64(map_value, "element_count")
|
|
529
|
+
policy = parse_semantic_policy(map_value.get("policy"), dtype, element_count)
|
|
530
|
+
quantized_values_hash = _hash_field(map_value, "quantized_values_hash")
|
|
531
|
+
count = _validate_shape(shape, limits, "semantic")
|
|
532
|
+
if element_count != count:
|
|
533
|
+
raise Malformed(
|
|
534
|
+
"semantic element_count %d does not match shape product %d"
|
|
535
|
+
% (element_count, count)
|
|
536
|
+
)
|
|
537
|
+
_checked_mul(count, WEIGHT_STORAGE_BYTES[dtype], "semantic tensor byte length")
|
|
538
|
+
policy.validate(dtype, count)
|
|
539
|
+
record = {
|
|
540
|
+
"schema": schema,
|
|
541
|
+
"dtype": dtype,
|
|
542
|
+
"shape": _u64_vec(shape),
|
|
543
|
+
"element_count": wcb.U64(element_count),
|
|
544
|
+
"policy": policy.to_wcb(),
|
|
545
|
+
"quantized_values_hash": quantized_values_hash,
|
|
546
|
+
}
|
|
547
|
+
if wcb.encode(record, limits) != data:
|
|
548
|
+
raise NonCanonical("semantic root bytes are not in canonical form")
|
|
549
|
+
return record
|
|
550
|
+
|
|
551
|
+
|
|
552
|
+
# ---------------------------------------------------------------------------
|
|
553
|
+
# Mode 9 -- FLOAT_TENSOR
|
|
554
|
+
# ---------------------------------------------------------------------------
|
|
555
|
+
|
|
556
|
+
_FLOAT_KEYS = (
|
|
557
|
+
"schema",
|
|
558
|
+
"dtype",
|
|
559
|
+
"shape",
|
|
560
|
+
"element_count",
|
|
561
|
+
"sign_bits",
|
|
562
|
+
"exponent_bits",
|
|
563
|
+
"mantissa_bits",
|
|
564
|
+
"exponent_coder",
|
|
565
|
+
"exponent_plane_len",
|
|
566
|
+
"mantissa_plane_len",
|
|
567
|
+
"payload_hash",
|
|
568
|
+
)
|
|
569
|
+
|
|
570
|
+
|
|
571
|
+
@dataclass(frozen=True)
|
|
572
|
+
class FloatTensorMeta:
|
|
573
|
+
dtype: str
|
|
574
|
+
shape: "list[int]"
|
|
575
|
+
element_count: int
|
|
576
|
+
sign_bits: int
|
|
577
|
+
exponent_bits: int
|
|
578
|
+
mantissa_bits: int
|
|
579
|
+
exponent_plane_len: int
|
|
580
|
+
mantissa_plane_len: int
|
|
581
|
+
|
|
582
|
+
|
|
583
|
+
def _float_dtype(name: str) -> str:
|
|
584
|
+
if name not in FLOAT_FIELD_BITS:
|
|
585
|
+
raise Malformed(
|
|
586
|
+
"unknown float tensor dtype `%s` (v1 allows bf16, fp16, fp32, "
|
|
587
|
+
"fp8-e4m3, fp8-e5m2)" % name
|
|
588
|
+
)
|
|
589
|
+
return name
|
|
590
|
+
|
|
591
|
+
|
|
592
|
+
def _byte_field(map_value: dict, key: str) -> int:
|
|
593
|
+
value = _u64(map_value, key)
|
|
594
|
+
if value > 0xFF:
|
|
595
|
+
raise Malformed(
|
|
596
|
+
"float tensor `%s` must fit an unsigned byte, got %d" % (key, value)
|
|
597
|
+
)
|
|
598
|
+
return value
|
|
599
|
+
|
|
600
|
+
|
|
601
|
+
def parse_float_descriptor(
|
|
602
|
+
frame, limits: Limits
|
|
603
|
+
) -> "tuple[FloatTensorMeta, bytes, bytes, bytes, bytes]":
|
|
604
|
+
descriptor_bytes, stored = _split_prefix(frame.body)
|
|
605
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "float tensor")
|
|
606
|
+
_exact_keys(map_value, _FLOAT_KEYS, "float tensor")
|
|
607
|
+
schema = _text(map_value, "schema")
|
|
608
|
+
if schema != FLOAT_TENSOR_SCHEMA:
|
|
609
|
+
raise Malformed("unsupported float tensor schema `%s`" % schema)
|
|
610
|
+
dtype = _float_dtype(_text(map_value, "dtype"))
|
|
611
|
+
coder = _text(map_value, "exponent_coder")
|
|
612
|
+
if coder != PLANE_CODEC:
|
|
613
|
+
raise Malformed(
|
|
614
|
+
"unsupported float tensor exponent coder `%s` (v1 allows `%s`)"
|
|
615
|
+
% (coder, PLANE_CODEC)
|
|
616
|
+
)
|
|
617
|
+
sign_bits = _byte_field(map_value, "sign_bits")
|
|
618
|
+
exponent_bits = _byte_field(map_value, "exponent_bits")
|
|
619
|
+
mantissa_bits = _byte_field(map_value, "mantissa_bits")
|
|
620
|
+
if (sign_bits, exponent_bits, mantissa_bits) != FLOAT_FIELD_BITS[dtype]:
|
|
621
|
+
raise Malformed(
|
|
622
|
+
"float tensor sub-field widths %d, %d, %d disagree with dtype `%s`"
|
|
623
|
+
% (sign_bits, exponent_bits, mantissa_bits, dtype)
|
|
624
|
+
)
|
|
625
|
+
element_bits = sign_bits + exponent_bits + mantissa_bits
|
|
626
|
+
if element_bits % 8 != 0:
|
|
627
|
+
raise Malformed("float tensor element width is not byte-aligned")
|
|
628
|
+
element_size = element_bits // 8
|
|
629
|
+
shape = _shape(map_value)
|
|
630
|
+
declared_count = _u64(map_value, "element_count")
|
|
631
|
+
expected_count = _validate_shape(shape, limits, "float tensor")
|
|
632
|
+
if declared_count != expected_count:
|
|
633
|
+
raise Malformed(
|
|
634
|
+
"float tensor element_count %d does not match shape product %d"
|
|
635
|
+
% (declared_count, expected_count)
|
|
636
|
+
)
|
|
637
|
+
decoded_bytes = _checked_mul(
|
|
638
|
+
declared_count, element_size, "float tensor decoded byte count"
|
|
639
|
+
)
|
|
640
|
+
_check_len(decoded_bytes, limits.max_decoded_bytes, "float tensor decoded bytes")
|
|
641
|
+
if decoded_bytes != frame.header.decoded_len:
|
|
642
|
+
raise Corrupt(
|
|
643
|
+
"float tensor decoded_len is %d but the descriptor requires %d"
|
|
644
|
+
% (frame.header.decoded_len, decoded_bytes)
|
|
645
|
+
)
|
|
646
|
+
mantissa_bits_total = _checked_mul(
|
|
647
|
+
declared_count,
|
|
648
|
+
sign_bits + mantissa_bits,
|
|
649
|
+
"float tensor mantissa bit count",
|
|
650
|
+
)
|
|
651
|
+
expected_mantissa = _checked_add(
|
|
652
|
+
mantissa_bits_total, 7, "float tensor mantissa length"
|
|
653
|
+
) // 8
|
|
654
|
+
declared_mantissa = _u64(map_value, "mantissa_plane_len")
|
|
655
|
+
if declared_mantissa != expected_mantissa:
|
|
656
|
+
raise Malformed(
|
|
657
|
+
"float tensor mantissa_plane_len %d does not match element_count x "
|
|
658
|
+
"(sign_bits + mantissa_bits) / 8 %d"
|
|
659
|
+
% (declared_mantissa, expected_mantissa)
|
|
660
|
+
)
|
|
661
|
+
declared_exponent = _u64(map_value, "exponent_plane_len")
|
|
662
|
+
stored_len = _checked_add(
|
|
663
|
+
declared_exponent, declared_mantissa, "float tensor plane lengths"
|
|
664
|
+
)
|
|
665
|
+
if stored_len != len(stored):
|
|
666
|
+
raise Malformed(
|
|
667
|
+
"float tensor stored payload is %d bytes but the descriptor declares %d"
|
|
668
|
+
% (len(stored), stored_len)
|
|
669
|
+
)
|
|
670
|
+
exponent_compressed = stored[:declared_exponent]
|
|
671
|
+
mantissa_plane = stored[declared_exponent:]
|
|
672
|
+
meta = FloatTensorMeta(
|
|
673
|
+
dtype=dtype,
|
|
674
|
+
shape=shape,
|
|
675
|
+
element_count=declared_count,
|
|
676
|
+
sign_bits=sign_bits,
|
|
677
|
+
exponent_bits=exponent_bits,
|
|
678
|
+
mantissa_bits=mantissa_bits,
|
|
679
|
+
exponent_plane_len=declared_exponent,
|
|
680
|
+
mantissa_plane_len=declared_mantissa,
|
|
681
|
+
)
|
|
682
|
+
return meta, stored, exponent_compressed, mantissa_plane, _bytes_of(
|
|
683
|
+
map_value, "payload_hash"
|
|
684
|
+
)
|
|
685
|
+
|
|
686
|
+
|
|
687
|
+
def _join_float_planes(
|
|
688
|
+
exponent_plane: bytes,
|
|
689
|
+
mantissa_plane: bytes,
|
|
690
|
+
dtype: str,
|
|
691
|
+
element_count: int,
|
|
692
|
+
limits: Limits,
|
|
693
|
+
) -> bytes:
|
|
694
|
+
if len(exponent_plane) != element_count:
|
|
695
|
+
raise Corrupt(
|
|
696
|
+
"float tensor exponent plane has %d bytes but element_count is %d"
|
|
697
|
+
% (len(exponent_plane), element_count)
|
|
698
|
+
)
|
|
699
|
+
sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
|
|
700
|
+
element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
|
|
701
|
+
decoded_bytes = _checked_mul(
|
|
702
|
+
element_count, element_size, "float tensor decoded byte count"
|
|
703
|
+
)
|
|
704
|
+
_check_len(decoded_bytes, limits.max_decoded_bytes, "float tensor decoded bytes")
|
|
705
|
+
expected_mantissa_bits = _checked_mul(
|
|
706
|
+
element_count,
|
|
707
|
+
sign_bits + mantissa_bits,
|
|
708
|
+
"float tensor mantissa bit count",
|
|
709
|
+
)
|
|
710
|
+
expected_mantissa_bytes = _checked_add(
|
|
711
|
+
expected_mantissa_bits, 7, "float tensor mantissa length"
|
|
712
|
+
) // 8
|
|
713
|
+
if len(mantissa_plane) != expected_mantissa_bytes:
|
|
714
|
+
raise Corrupt(
|
|
715
|
+
"float tensor mantissa plane has %d bytes but descriptor arithmetic "
|
|
716
|
+
"requires %d" % (len(mantissa_plane), expected_mantissa_bytes)
|
|
717
|
+
)
|
|
718
|
+
|
|
719
|
+
width = sign_bits + mantissa_bits
|
|
720
|
+
bit = 0
|
|
721
|
+
total_bits = len(mantissa_plane) * 8
|
|
722
|
+
out = bytearray(decoded_bytes)
|
|
723
|
+
mantissa_mask = (1 << mantissa_bits) - 1
|
|
724
|
+
for index, exponent_byte in enumerate(exponent_plane):
|
|
725
|
+
if exponent_byte >> exponent_bits:
|
|
726
|
+
raise Corrupt(
|
|
727
|
+
"float tensor exponent byte exceeds the dtype exponent width"
|
|
728
|
+
)
|
|
729
|
+
value = 0
|
|
730
|
+
for _ in range(width):
|
|
731
|
+
if bit >= total_bits:
|
|
732
|
+
# The plane is exactly long enough by construction; this is an
|
|
733
|
+
# internal invariant, reported as corruption rather than panic.
|
|
734
|
+
raise Corrupt("float tensor mantissa plane is truncated")
|
|
735
|
+
byte = mantissa_plane[bit >> 3]
|
|
736
|
+
value = (value << 1) | ((byte >> (7 - (bit & 7))) & 1)
|
|
737
|
+
bit += 1
|
|
738
|
+
sign = value >> mantissa_bits
|
|
739
|
+
mantissa = value & mantissa_mask
|
|
740
|
+
word = (
|
|
741
|
+
(sign << (exponent_bits + mantissa_bits))
|
|
742
|
+
| (exponent_byte << mantissa_bits)
|
|
743
|
+
| mantissa
|
|
744
|
+
)
|
|
745
|
+
offset = index * element_size
|
|
746
|
+
out[offset : offset + element_size] = word.to_bytes(element_size, "little")
|
|
747
|
+
|
|
748
|
+
padding = len(mantissa_plane) * 8 - expected_mantissa_bits
|
|
749
|
+
if padding > 0:
|
|
750
|
+
last = mantissa_plane[-1]
|
|
751
|
+
if last & ((1 << padding) - 1):
|
|
752
|
+
raise Corrupt("float tensor mantissa padding bits must be zero")
|
|
753
|
+
return bytes(out)
|
|
754
|
+
|
|
755
|
+
|
|
756
|
+
def decode_float_tensor(frame, limits: Limits) -> bytes:
|
|
757
|
+
meta, stored, exponent_compressed, mantissa_plane, declared_hash = (
|
|
758
|
+
parse_float_descriptor(frame, limits)
|
|
759
|
+
)
|
|
760
|
+
exponent_plane = huffman.decode(exponent_compressed, meta.element_count)
|
|
761
|
+
out = _join_float_planes(
|
|
762
|
+
exponent_plane,
|
|
763
|
+
mantissa_plane,
|
|
764
|
+
meta.dtype,
|
|
765
|
+
meta.element_count,
|
|
766
|
+
limits,
|
|
767
|
+
)
|
|
768
|
+
if len(declared_hash) != HASH_LEN:
|
|
769
|
+
raise Malformed(
|
|
770
|
+
"float tensor payload_hash must be %d bytes, got %d"
|
|
771
|
+
% (HASH_LEN, len(declared_hash))
|
|
772
|
+
)
|
|
773
|
+
if declared_hash != hash_domain("weft/float-tensor/payload/v1", stored):
|
|
774
|
+
raise Corrupt("float tensor payload_hash mismatch")
|
|
775
|
+
return out
|
|
776
|
+
|
|
777
|
+
|
|
778
|
+
# ---------------------------------------------------------------------------
|
|
779
|
+
# Mode 10 -- TENSOR_SHARDED
|
|
780
|
+
# ---------------------------------------------------------------------------
|
|
781
|
+
|
|
782
|
+
_SHARDED_KEYS = (
|
|
783
|
+
"schema",
|
|
784
|
+
"dtype",
|
|
785
|
+
"shape",
|
|
786
|
+
"element_count",
|
|
787
|
+
"shard_axis",
|
|
788
|
+
"shard_count",
|
|
789
|
+
"shard_index",
|
|
790
|
+
"block_size",
|
|
791
|
+
"scales",
|
|
792
|
+
"zero_points",
|
|
793
|
+
"layout_root",
|
|
794
|
+
"payload_len",
|
|
795
|
+
"payload_hash",
|
|
796
|
+
)
|
|
797
|
+
|
|
798
|
+
|
|
799
|
+
@dataclass
|
|
800
|
+
class ShardedMeta:
|
|
801
|
+
dtype: str
|
|
802
|
+
shape: "list[int]"
|
|
803
|
+
element_count: int
|
|
804
|
+
shard_axis: int
|
|
805
|
+
shard_count: int
|
|
806
|
+
shard_index: int
|
|
807
|
+
block_size: int
|
|
808
|
+
scales: "list[int]"
|
|
809
|
+
zero_points: "list[int]"
|
|
810
|
+
layout_root: bytes
|
|
811
|
+
payload_len: int
|
|
812
|
+
|
|
813
|
+
|
|
814
|
+
@dataclass
|
|
815
|
+
class DecodedShard:
|
|
816
|
+
meta: ShardedMeta
|
|
817
|
+
words: "list[int]"
|
|
818
|
+
planes: bytes
|
|
819
|
+
|
|
820
|
+
|
|
821
|
+
def _sharded_dtype(name: str) -> str:
|
|
822
|
+
if name not in ("q8_0", "q4_0"):
|
|
823
|
+
raise Malformed(
|
|
824
|
+
"unknown TENSOR_SHARDED dtype `%s` (v1 allows q8_0, q4_0)" % name
|
|
825
|
+
)
|
|
826
|
+
return name
|
|
827
|
+
|
|
828
|
+
|
|
829
|
+
def _u32_array(map_value: dict, key: str) -> "list[int]":
|
|
830
|
+
values = _array_u64(map_value, key)
|
|
831
|
+
for value in values:
|
|
832
|
+
if value > U32_MAX:
|
|
833
|
+
raise Malformed(
|
|
834
|
+
"TENSOR_SHARDED scale bits exceed the 32-bit f32 pattern"
|
|
835
|
+
)
|
|
836
|
+
return values
|
|
837
|
+
|
|
838
|
+
|
|
839
|
+
def _i64_array(map_value: dict, key: str) -> "list[int]":
|
|
840
|
+
value = map_value.get(key)
|
|
841
|
+
if not isinstance(value, list):
|
|
842
|
+
raise Malformed("metadata `%s` must be an array" % key)
|
|
843
|
+
out = []
|
|
844
|
+
for item in value:
|
|
845
|
+
if not isinstance(item, wcb.I64):
|
|
846
|
+
raise Malformed("metadata `%s` elements must be i64" % key)
|
|
847
|
+
out.append(item.value)
|
|
848
|
+
return out
|
|
849
|
+
|
|
850
|
+
|
|
851
|
+
def parse_sharded_descriptor(descriptor_bytes: bytes, limits: Limits):
|
|
852
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "TENSOR_SHARDED")
|
|
853
|
+
_exact_keys(map_value, _SHARDED_KEYS, "TENSOR_SHARDED")
|
|
854
|
+
schema = _text(map_value, "schema")
|
|
855
|
+
if schema != TENSOR_SHARDED_SCHEMA:
|
|
856
|
+
raise Malformed("unsupported TENSOR_SHARDED schema `%s`" % schema)
|
|
857
|
+
dtype = _sharded_dtype(_text(map_value, "dtype"))
|
|
858
|
+
shape = _shape(map_value)
|
|
859
|
+
element_count = _u64(map_value, "element_count")
|
|
860
|
+
shard_axis = _u32(map_value, "shard_axis")
|
|
861
|
+
shard_count = _u32(map_value, "shard_count")
|
|
862
|
+
shard_index = _u32(map_value, "shard_index")
|
|
863
|
+
block_size = _u64(map_value, "block_size")
|
|
864
|
+
scales = _u32_array(map_value, "scales")
|
|
865
|
+
zero_points = _i64_array(map_value, "zero_points")
|
|
866
|
+
layout_root = _hash_field(map_value, "layout_root")
|
|
867
|
+
payload_len = _u64(map_value, "payload_len")
|
|
868
|
+
payload_hash = _hash_field(map_value, "payload_hash")
|
|
869
|
+
meta = ShardedMeta(
|
|
870
|
+
dtype=dtype,
|
|
871
|
+
shape=shape,
|
|
872
|
+
element_count=element_count,
|
|
873
|
+
shard_axis=shard_axis,
|
|
874
|
+
shard_count=shard_count,
|
|
875
|
+
shard_index=shard_index,
|
|
876
|
+
block_size=block_size,
|
|
877
|
+
scales=scales,
|
|
878
|
+
zero_points=zero_points,
|
|
879
|
+
layout_root=layout_root,
|
|
880
|
+
payload_len=payload_len,
|
|
881
|
+
)
|
|
882
|
+
return meta, payload_hash
|
|
883
|
+
|
|
884
|
+
|
|
885
|
+
def _validate_sharded_meta(
|
|
886
|
+
meta: ShardedMeta, limits: Limits
|
|
887
|
+
) -> "tuple[int, int, int]":
|
|
888
|
+
element_count = _validate_shape(meta.shape, limits, "TENSOR_SHARDED")
|
|
889
|
+
if meta.element_count != element_count:
|
|
890
|
+
raise Malformed(
|
|
891
|
+
"TENSOR_SHARDED element_count %d does not match shape product %d"
|
|
892
|
+
% (meta.element_count, element_count)
|
|
893
|
+
)
|
|
894
|
+
if meta.shard_count == 0:
|
|
895
|
+
raise Malformed("TENSOR_SHARDED shard_count must be at least 1")
|
|
896
|
+
if meta.shard_axis >= len(meta.shape):
|
|
897
|
+
raise Malformed(
|
|
898
|
+
"TENSOR_SHARDED shard_axis %d is not below rank %d"
|
|
899
|
+
% (meta.shard_axis, len(meta.shape))
|
|
900
|
+
)
|
|
901
|
+
if meta.shard_index >= meta.shard_count:
|
|
902
|
+
raise Malformed(
|
|
903
|
+
"TENSOR_SHARDED shard_index %d is not below shard_count %d"
|
|
904
|
+
% (meta.shard_index, meta.shard_count)
|
|
905
|
+
)
|
|
906
|
+
axis_dim = meta.shape[meta.shard_axis]
|
|
907
|
+
if axis_dim % meta.shard_count != 0:
|
|
908
|
+
raise Malformed(
|
|
909
|
+
"TENSOR_SHARDED shape[%d] = %d is not divisible by shard_count %d"
|
|
910
|
+
% (meta.shard_axis, axis_dim, meta.shard_count)
|
|
911
|
+
)
|
|
912
|
+
if element_count % meta.shard_count != 0:
|
|
913
|
+
raise Malformed(
|
|
914
|
+
"TENSOR_SHARDED element_count %d is not divisible by shard_count %d"
|
|
915
|
+
% (element_count, meta.shard_count)
|
|
916
|
+
)
|
|
917
|
+
shard_elements = element_count // meta.shard_count
|
|
918
|
+
if meta.block_size not in BLOCK_SIZES:
|
|
919
|
+
raise Malformed(
|
|
920
|
+
"TENSOR_SHARDED block_size %d is not one of 16, 32, 64, 128, 256"
|
|
921
|
+
% meta.block_size
|
|
922
|
+
)
|
|
923
|
+
if element_count % meta.block_size != 0:
|
|
924
|
+
raise Malformed(
|
|
925
|
+
"TENSOR_SHARDED element_count %d is not divisible by block_size %d"
|
|
926
|
+
% (element_count, meta.block_size)
|
|
927
|
+
)
|
|
928
|
+
block_count = element_count // meta.block_size
|
|
929
|
+
if block_count > limits.max_structured_items:
|
|
930
|
+
raise LimitExceeded(
|
|
931
|
+
"TENSOR_SHARDED block count %d exceeds WCB item limit %d"
|
|
932
|
+
% (block_count, limits.max_structured_items)
|
|
933
|
+
)
|
|
934
|
+
if len(meta.scales) != block_count or len(meta.zero_points) != block_count:
|
|
935
|
+
raise Malformed(
|
|
936
|
+
"TENSOR_SHARDED descriptor has %d scales and %d zero points but "
|
|
937
|
+
"requires %d"
|
|
938
|
+
% (len(meta.scales), len(meta.zero_points), block_count)
|
|
939
|
+
)
|
|
940
|
+
for bits in meta.scales:
|
|
941
|
+
scale = struct.unpack("<f", struct.pack("<I", bits))[0]
|
|
942
|
+
if not math.isfinite(scale):
|
|
943
|
+
raise Malformed("TENSOR_SHARDED scale bits must be a finite f32")
|
|
944
|
+
if scale <= 0.0:
|
|
945
|
+
raise Malformed("TENSOR_SHARDED scale must be greater than zero")
|
|
946
|
+
if meta.dtype == "q4_0" and shard_elements % 2 != 0:
|
|
947
|
+
raise Malformed(
|
|
948
|
+
"q4_0 shard element count must be even to pack nibble pairs"
|
|
949
|
+
)
|
|
950
|
+
quant_len = shard_elements if meta.dtype == "q8_0" else shard_elements // 2
|
|
951
|
+
planes_len = _checked_add(
|
|
952
|
+
quant_len, _checked_mul(block_count, 4, "TENSOR_SHARDED scale plane"),
|
|
953
|
+
"TENSOR_SHARDED plane length",
|
|
954
|
+
)
|
|
955
|
+
_check_len(
|
|
956
|
+
planes_len, limits.max_decoded_bytes, "TENSOR_SHARDED reconstructed planes"
|
|
957
|
+
)
|
|
958
|
+
return element_count, shard_elements, block_count
|
|
959
|
+
|
|
960
|
+
|
|
961
|
+
def decode_tensor_sharded(frame, limits: Limits) -> DecodedShard:
|
|
962
|
+
descriptor_bytes, stored = _split_prefix(frame.body)
|
|
963
|
+
meta, declared_payload_hash = parse_sharded_descriptor(descriptor_bytes, limits)
|
|
964
|
+
_element_count, shard_elements, block_count = _validate_sharded_meta(meta, limits)
|
|
965
|
+
if meta.payload_len != len(stored):
|
|
966
|
+
raise Malformed(
|
|
967
|
+
"TENSOR_SHARDED payload is %d bytes but payload_len declares %d"
|
|
968
|
+
% (len(stored), meta.payload_len)
|
|
969
|
+
)
|
|
970
|
+
quant_len = shard_elements if meta.dtype == "q8_0" else shard_elements // 2
|
|
971
|
+
planes_len = _checked_add(
|
|
972
|
+
quant_len, _checked_mul(block_count, 4, "TENSOR_SHARDED scale plane"),
|
|
973
|
+
"TENSOR_SHARDED plane length",
|
|
974
|
+
)
|
|
975
|
+
if frame.header.decoded_len != planes_len:
|
|
976
|
+
raise Corrupt(
|
|
977
|
+
"TENSOR_SHARDED decoded_len is %d but the descriptor requires %d"
|
|
978
|
+
% (frame.header.decoded_len, planes_len)
|
|
979
|
+
)
|
|
980
|
+
if declared_payload_hash != hash_domain(DOMAIN_TENSOR_SHARDED_PAYLOAD, stored):
|
|
981
|
+
raise Corrupt("TENSOR_SHARDED payload_hash mismatch")
|
|
982
|
+
_check_len(
|
|
983
|
+
planes_len, limits.max_decoded_bytes, "TENSOR_SHARDED reconstructed planes"
|
|
984
|
+
)
|
|
985
|
+
planes = huffman.decode(stored, planes_len)
|
|
986
|
+
if quant_len > len(planes):
|
|
987
|
+
raise Corrupt(
|
|
988
|
+
"TENSOR_SHARDED reconstructed plane length is inconsistent"
|
|
989
|
+
)
|
|
990
|
+
quant_plane = planes[:quant_len]
|
|
991
|
+
scale_bytes = planes[quant_len:]
|
|
992
|
+
if len(scale_bytes) != block_count * 4:
|
|
993
|
+
raise Corrupt("TENSOR_SHARDED scale plane length is inconsistent")
|
|
994
|
+
for index in range(block_count):
|
|
995
|
+
chunk = scale_bytes[index * 4 : index * 4 + 4]
|
|
996
|
+
actual = int.from_bytes(chunk, "little")
|
|
997
|
+
if actual != meta.scales[index]:
|
|
998
|
+
raise Corrupt(
|
|
999
|
+
"TENSOR_SHARDED scale plane block %d does not match the descriptor"
|
|
1000
|
+
% index
|
|
1001
|
+
)
|
|
1002
|
+
if frame.header.content_hash != hash_domain("weft/chunk/v1", planes):
|
|
1003
|
+
raise Corrupt(
|
|
1004
|
+
"TENSOR_SHARDED content_hash mismatch against the reconstructed planes"
|
|
1005
|
+
)
|
|
1006
|
+
if meta.dtype == "q8_0":
|
|
1007
|
+
words = [byte - 256 if byte >= 128 else byte for byte in quant_plane]
|
|
1008
|
+
else:
|
|
1009
|
+
words = []
|
|
1010
|
+
for byte in quant_plane:
|
|
1011
|
+
low = byte & 0x0F
|
|
1012
|
+
high = byte >> 4
|
|
1013
|
+
words.append(low - 16 if low >= 8 else low)
|
|
1014
|
+
words.append(high - 16 if high >= 8 else high)
|
|
1015
|
+
words = words[:shard_elements]
|
|
1016
|
+
if len(words) != shard_elements:
|
|
1017
|
+
raise Corrupt(
|
|
1018
|
+
"TENSOR_SHARDED reconstructed word count is inconsistent"
|
|
1019
|
+
)
|
|
1020
|
+
return DecodedShard(meta=meta, words=words, planes=planes)
|
|
1021
|
+
|
|
1022
|
+
|
|
1023
|
+
def _same_sharded_tensor(left: ShardedMeta, right: ShardedMeta) -> bool:
|
|
1024
|
+
return (
|
|
1025
|
+
left.dtype == right.dtype
|
|
1026
|
+
and left.shape == right.shape
|
|
1027
|
+
and left.element_count == right.element_count
|
|
1028
|
+
and left.shard_axis == right.shard_axis
|
|
1029
|
+
and left.shard_count == right.shard_count
|
|
1030
|
+
and left.block_size == right.block_size
|
|
1031
|
+
and left.scales == right.scales
|
|
1032
|
+
and left.zero_points == right.zero_points
|
|
1033
|
+
and left.layout_root == right.layout_root
|
|
1034
|
+
)
|
|
1035
|
+
|
|
1036
|
+
|
|
1037
|
+
def _assemble_shard_words(
|
|
1038
|
+
shape: "list[int]", axis: int, shards: "list[list[int]]"
|
|
1039
|
+
) -> "list[int]":
|
|
1040
|
+
count = len(shards)
|
|
1041
|
+
shard_shape = list(shape)
|
|
1042
|
+
shard_shape[axis] = shard_shape[axis] // count
|
|
1043
|
+
total = 1
|
|
1044
|
+
for dim in shape:
|
|
1045
|
+
total *= dim
|
|
1046
|
+
out = [0] * total
|
|
1047
|
+
plane = shape[axis] // count
|
|
1048
|
+
strides = [0] * len(shape)
|
|
1049
|
+
stride = 1
|
|
1050
|
+
for index in range(len(shape) - 1, -1, -1):
|
|
1051
|
+
strides[index] = stride
|
|
1052
|
+
stride *= shape[index]
|
|
1053
|
+
for shard_index, payload in enumerate(shards):
|
|
1054
|
+
# Iterate the shard-local row-major coordinates.
|
|
1055
|
+
def walk(dimension: int, local_coord: "list[int]", local_flat: int) -> int:
|
|
1056
|
+
if dimension == len(shard_shape):
|
|
1057
|
+
full = 0
|
|
1058
|
+
for axis_index, coord in enumerate(local_coord):
|
|
1059
|
+
full_coord = coord
|
|
1060
|
+
if axis_index == axis:
|
|
1061
|
+
full_coord = shard_index * plane + coord
|
|
1062
|
+
full += full_coord * strides[axis_index]
|
|
1063
|
+
out[full] = payload[local_flat]
|
|
1064
|
+
return local_flat + 1
|
|
1065
|
+
extent = shard_shape[dimension]
|
|
1066
|
+
for position in range(extent):
|
|
1067
|
+
local_coord[dimension] = position
|
|
1068
|
+
local_flat = walk(dimension + 1, local_coord, local_flat)
|
|
1069
|
+
return local_flat
|
|
1070
|
+
|
|
1071
|
+
walk(0, [0] * len(shard_shape), 0)
|
|
1072
|
+
return out
|
|
1073
|
+
|
|
1074
|
+
|
|
1075
|
+
def shard_set_root(
|
|
1076
|
+
shards: "list[DecodedShard]", limits: Limits
|
|
1077
|
+
) -> "tuple[bytes, list[int]]":
|
|
1078
|
+
if not shards:
|
|
1079
|
+
raise Malformed("TENSOR_SHARDED shard set is empty")
|
|
1080
|
+
first = shards[0].meta
|
|
1081
|
+
_validate_sharded_meta(first, limits)
|
|
1082
|
+
count = first.shard_count
|
|
1083
|
+
if len(shards) != count:
|
|
1084
|
+
raise Malformed(
|
|
1085
|
+
"TENSOR_SHARDED shard set has %d entries but the descriptor declares %d"
|
|
1086
|
+
% (len(shards), count)
|
|
1087
|
+
)
|
|
1088
|
+
ordered: "list[DecodedShard | None]" = [None] * count
|
|
1089
|
+
for shard in shards:
|
|
1090
|
+
if not _same_sharded_tensor(first, shard.meta):
|
|
1091
|
+
raise NonCanonical(
|
|
1092
|
+
"TENSOR_SHARDED shard descriptors disagree on tensor identity"
|
|
1093
|
+
)
|
|
1094
|
+
index = shard.meta.shard_index
|
|
1095
|
+
if index >= count:
|
|
1096
|
+
raise NonCanonical(
|
|
1097
|
+
"TENSOR_SHARDED shard index %d is not below shard count %d"
|
|
1098
|
+
% (index, count)
|
|
1099
|
+
)
|
|
1100
|
+
if ordered[index] is not None:
|
|
1101
|
+
raise NonCanonical(
|
|
1102
|
+
"duplicate TENSOR_SHARDED shard index %d" % index
|
|
1103
|
+
)
|
|
1104
|
+
ordered[index] = shard
|
|
1105
|
+
for index, entry in enumerate(ordered):
|
|
1106
|
+
if entry is None:
|
|
1107
|
+
raise NonCanonical("missing TENSOR_SHARDED shard index %d" % index)
|
|
1108
|
+
words = _assemble_shard_words(
|
|
1109
|
+
first.shape, first.shard_axis, [entry.words for entry in ordered]
|
|
1110
|
+
)
|
|
1111
|
+
root = sharded_logical_root(
|
|
1112
|
+
first.dtype,
|
|
1113
|
+
first.shape,
|
|
1114
|
+
words,
|
|
1115
|
+
first.scales,
|
|
1116
|
+
first.zero_points,
|
|
1117
|
+
limits,
|
|
1118
|
+
)
|
|
1119
|
+
if root != first.layout_root:
|
|
1120
|
+
raise Corrupt(
|
|
1121
|
+
"TENSOR_SHARDED shard set recomputes a different logical root"
|
|
1122
|
+
)
|
|
1123
|
+
return root, words
|
|
1124
|
+
|
|
1125
|
+
|
|
1126
|
+
# ---------------------------------------------------------------------------
|
|
1127
|
+
# Modes 11/14/16 -- delta modes
|
|
1128
|
+
# ---------------------------------------------------------------------------
|
|
1129
|
+
|
|
1130
|
+
|
|
1131
|
+
def _weight_dtype(name: str) -> str:
|
|
1132
|
+
if name not in WEIGHT_WORD_BITS:
|
|
1133
|
+
raise Malformed(
|
|
1134
|
+
"unknown WEIGHT_DELTA dtype `%s` (v1 allows bf16, fp16, fp32, "
|
|
1135
|
+
"q8_0, q4_0)" % name
|
|
1136
|
+
)
|
|
1137
|
+
return name
|
|
1138
|
+
|
|
1139
|
+
|
|
1140
|
+
@dataclass
|
|
1141
|
+
class WeightDeltaMeta:
|
|
1142
|
+
dtype: str
|
|
1143
|
+
shape: "list[int]"
|
|
1144
|
+
element_count: int
|
|
1145
|
+
word_bits: int
|
|
1146
|
+
base: bytes
|
|
1147
|
+
base_dtype: str
|
|
1148
|
+
base_shape: "list[int]"
|
|
1149
|
+
op: str
|
|
1150
|
+
planes: int
|
|
1151
|
+
plane_codec: str
|
|
1152
|
+
delta_root: bytes
|
|
1153
|
+
payload_len: int
|
|
1154
|
+
|
|
1155
|
+
|
|
1156
|
+
_WEIGHT_KEYS = (
|
|
1157
|
+
"schema",
|
|
1158
|
+
"dtype",
|
|
1159
|
+
"shape",
|
|
1160
|
+
"element_count",
|
|
1161
|
+
"word_bits",
|
|
1162
|
+
"base",
|
|
1163
|
+
"base_dtype",
|
|
1164
|
+
"base_shape",
|
|
1165
|
+
"op",
|
|
1166
|
+
"planes",
|
|
1167
|
+
"plane_codec",
|
|
1168
|
+
"delta_root",
|
|
1169
|
+
"payload_len",
|
|
1170
|
+
)
|
|
1171
|
+
|
|
1172
|
+
|
|
1173
|
+
def parse_weight_descriptor(descriptor_bytes: bytes, limits: Limits) -> WeightDeltaMeta:
|
|
1174
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "WEIGHT_DELTA")
|
|
1175
|
+
_exact_keys(map_value, _WEIGHT_KEYS, "WEIGHT_DELTA")
|
|
1176
|
+
schema = _text(map_value, "schema")
|
|
1177
|
+
if schema != WEIGHT_DELTA_SCHEMA:
|
|
1178
|
+
raise Malformed("unsupported WEIGHT_DELTA schema `%s`" % schema)
|
|
1179
|
+
dtype = _weight_dtype(_text(map_value, "dtype"))
|
|
1180
|
+
shape = _shape(map_value)
|
|
1181
|
+
element_count = _u64(map_value, "element_count")
|
|
1182
|
+
word_bits = _u64(map_value, "word_bits")
|
|
1183
|
+
base = _hash_field(map_value, "base")
|
|
1184
|
+
base_dtype = _weight_dtype(_text(map_value, "base_dtype"))
|
|
1185
|
+
base_shape = _array_u64(map_value, "base_shape")
|
|
1186
|
+
op = _text(map_value, "op")
|
|
1187
|
+
planes = _u64(map_value, "planes")
|
|
1188
|
+
plane_codec = _text(map_value, "plane_codec")
|
|
1189
|
+
delta_root = _hash_field(map_value, "delta_root")
|
|
1190
|
+
payload_len = _u64(map_value, "payload_len")
|
|
1191
|
+
meta = WeightDeltaMeta(
|
|
1192
|
+
dtype=dtype,
|
|
1193
|
+
shape=shape,
|
|
1194
|
+
element_count=element_count,
|
|
1195
|
+
word_bits=word_bits,
|
|
1196
|
+
base=base,
|
|
1197
|
+
base_dtype=base_dtype,
|
|
1198
|
+
base_shape=base_shape,
|
|
1199
|
+
op=op,
|
|
1200
|
+
planes=planes,
|
|
1201
|
+
plane_codec=plane_codec,
|
|
1202
|
+
delta_root=delta_root,
|
|
1203
|
+
payload_len=payload_len,
|
|
1204
|
+
)
|
|
1205
|
+
validate_weight_meta(meta, limits)
|
|
1206
|
+
return meta
|
|
1207
|
+
|
|
1208
|
+
|
|
1209
|
+
def validate_weight_meta(meta: WeightDeltaMeta, limits: Limits) -> int:
|
|
1210
|
+
if meta.op != DELTA_OP:
|
|
1211
|
+
raise Malformed(
|
|
1212
|
+
"WEIGHT_DELTA op `%s` is not implemented (only %s)"
|
|
1213
|
+
% (meta.op, DELTA_OP)
|
|
1214
|
+
)
|
|
1215
|
+
if meta.plane_codec != PLANE_CODEC:
|
|
1216
|
+
raise Malformed(
|
|
1217
|
+
"WEIGHT_DELTA plane_codec `%s` is not implemented (only %s)"
|
|
1218
|
+
% (meta.plane_codec, PLANE_CODEC)
|
|
1219
|
+
)
|
|
1220
|
+
if WEIGHT_WORD_BITS[meta.dtype] != meta.word_bits:
|
|
1221
|
+
raise Malformed(
|
|
1222
|
+
"WEIGHT_DELTA word_bits %d disagrees with dtype %s (%d bits)"
|
|
1223
|
+
% (meta.word_bits, meta.dtype, WEIGHT_WORD_BITS[meta.dtype])
|
|
1224
|
+
)
|
|
1225
|
+
if WEIGHT_PLANES[meta.dtype] != meta.planes:
|
|
1226
|
+
raise Malformed(
|
|
1227
|
+
"WEIGHT_DELTA planes %d disagrees with dtype %s (%d planes)"
|
|
1228
|
+
% (meta.planes, meta.dtype, WEIGHT_PLANES[meta.dtype])
|
|
1229
|
+
)
|
|
1230
|
+
if meta.base_dtype != meta.dtype:
|
|
1231
|
+
raise Malformed(
|
|
1232
|
+
"WEIGHT_DELTA base_dtype `%s` must equal dtype `%s`"
|
|
1233
|
+
% (meta.base_dtype, meta.dtype)
|
|
1234
|
+
)
|
|
1235
|
+
if meta.base_shape != meta.shape:
|
|
1236
|
+
raise Malformed("WEIGHT_DELTA base_shape must equal shape")
|
|
1237
|
+
count = _validate_shape(meta.shape, limits, "WEIGHT_DELTA")
|
|
1238
|
+
if meta.element_count != count:
|
|
1239
|
+
raise Malformed(
|
|
1240
|
+
"WEIGHT_DELTA element_count %d does not match shape product %d"
|
|
1241
|
+
% (meta.element_count, count)
|
|
1242
|
+
)
|
|
1243
|
+
_checked_mul(
|
|
1244
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1245
|
+
"WEIGHT_DELTA element byte length",
|
|
1246
|
+
)
|
|
1247
|
+
return count
|
|
1248
|
+
|
|
1249
|
+
|
|
1250
|
+
def weight_metadata(frame, limits: Limits) -> WeightDeltaMeta:
|
|
1251
|
+
descriptor_bytes, _payload = _split_prefix(frame.body)
|
|
1252
|
+
return parse_weight_descriptor(descriptor_bytes, limits)
|
|
1253
|
+
|
|
1254
|
+
|
|
1255
|
+
def _decode_planes(dtype: str, payload: bytes, element_count: int, what: str):
|
|
1256
|
+
plane_count = WEIGHT_PLANES[dtype]
|
|
1257
|
+
planes = []
|
|
1258
|
+
offset = 0
|
|
1259
|
+
for _ in range(plane_count):
|
|
1260
|
+
if offset >= len(payload):
|
|
1261
|
+
raise Corrupt("%s payload ends before all declared planes" % what)
|
|
1262
|
+
plane, consumed = huffman.decode_prefix(payload[offset:], element_count)
|
|
1263
|
+
offset += consumed
|
|
1264
|
+
planes.append(plane)
|
|
1265
|
+
if offset != len(payload):
|
|
1266
|
+
raise Corrupt(
|
|
1267
|
+
"%s payload has trailing bytes after the declared planes" % what
|
|
1268
|
+
)
|
|
1269
|
+
return planes
|
|
1270
|
+
|
|
1271
|
+
|
|
1272
|
+
def _delta_word_from_planes(dtype: str, planes, index: int) -> int:
|
|
1273
|
+
if dtype in ("bf16", "fp16"):
|
|
1274
|
+
return planes[0][index] | (planes[1][index] << 8)
|
|
1275
|
+
if dtype == "fp32":
|
|
1276
|
+
return (
|
|
1277
|
+
planes[0][index]
|
|
1278
|
+
| (planes[1][index] << 8)
|
|
1279
|
+
| (planes[2][index] << 16)
|
|
1280
|
+
| (planes[3][index] << 24)
|
|
1281
|
+
)
|
|
1282
|
+
if dtype == "q8_0":
|
|
1283
|
+
return planes[0][index]
|
|
1284
|
+
nibble = planes[0][index]
|
|
1285
|
+
if nibble > 15:
|
|
1286
|
+
raise Corrupt(
|
|
1287
|
+
"WEIGHT_DELTA q4_0 delta word %d is %d, outside 0..=15"
|
|
1288
|
+
% (index, nibble)
|
|
1289
|
+
)
|
|
1290
|
+
return nibble
|
|
1291
|
+
|
|
1292
|
+
|
|
1293
|
+
def _build_delta_bytes(dtype: str, planes, count: int) -> bytes:
|
|
1294
|
+
out = bytearray()
|
|
1295
|
+
for index in range(count):
|
|
1296
|
+
_push_word(dtype, _delta_word_from_planes(dtype, planes, index), out)
|
|
1297
|
+
return bytes(out)
|
|
1298
|
+
|
|
1299
|
+
|
|
1300
|
+
def _mod_mask(dtype: str) -> int:
|
|
1301
|
+
return (1 << WEIGHT_WORD_BITS[dtype]) - 1
|
|
1302
|
+
|
|
1303
|
+
|
|
1304
|
+
def _base_frame_descriptor(frame, limits: Limits) -> "tuple[str, list[int]]":
|
|
1305
|
+
from . import modes as _modes
|
|
1306
|
+
|
|
1307
|
+
mode = frame.header.mode
|
|
1308
|
+
if mode == 6:
|
|
1309
|
+
meta, _payload = _modes._parse_tensor_body(frame, limits)
|
|
1310
|
+
mapped = {
|
|
1311
|
+
"bf16": "bf16",
|
|
1312
|
+
"f16": "fp16",
|
|
1313
|
+
"f32": "fp32",
|
|
1314
|
+
"i8": "q8_0",
|
|
1315
|
+
"u8": "q4_0",
|
|
1316
|
+
}.get(meta["dtype"])
|
|
1317
|
+
if mapped is None:
|
|
1318
|
+
raise Malformed(
|
|
1319
|
+
"base frame dtype `%s` cannot back a WEIGHT_DELTA word plane"
|
|
1320
|
+
% meta["dtype"]
|
|
1321
|
+
)
|
|
1322
|
+
return mapped, meta["shape"]
|
|
1323
|
+
if mode == 9:
|
|
1324
|
+
meta = parse_float_descriptor(frame, limits)[0]
|
|
1325
|
+
mapped = {
|
|
1326
|
+
"bf16": "bf16",
|
|
1327
|
+
"fp16": "fp16",
|
|
1328
|
+
"fp32": "fp32",
|
|
1329
|
+
}.get(meta.dtype)
|
|
1330
|
+
if mapped is None:
|
|
1331
|
+
raise Malformed(
|
|
1332
|
+
"base frame dtype `%s` cannot back a WEIGHT_DELTA word plane"
|
|
1333
|
+
% meta.dtype
|
|
1334
|
+
)
|
|
1335
|
+
return mapped, meta.shape
|
|
1336
|
+
if mode == 11:
|
|
1337
|
+
meta = weight_metadata(frame, limits)
|
|
1338
|
+
return meta.dtype, meta.shape
|
|
1339
|
+
if mode == 14:
|
|
1340
|
+
meta = semantic_delta_metadata(frame, limits)
|
|
1341
|
+
return meta.dtype, meta.shape
|
|
1342
|
+
if mode == 16:
|
|
1343
|
+
meta = sparse_delta_metadata(frame, limits)
|
|
1344
|
+
return meta.dtype, meta.shape
|
|
1345
|
+
raise Malformed(
|
|
1346
|
+
"base frame mode `%s` does not carry tensor dtype/shape metadata"
|
|
1347
|
+
% frame.header.mode_name
|
|
1348
|
+
)
|
|
1349
|
+
|
|
1350
|
+
|
|
1351
|
+
def _base_bytes(
|
|
1352
|
+
frame,
|
|
1353
|
+
meta_base: bytes,
|
|
1354
|
+
resolver,
|
|
1355
|
+
limits: Limits,
|
|
1356
|
+
base_frame,
|
|
1357
|
+
what: str,
|
|
1358
|
+
) -> "tuple[bytes, bytes | None]":
|
|
1359
|
+
from .modes import decode_bytes as _decode_bytes
|
|
1360
|
+
|
|
1361
|
+
if base_frame is not None:
|
|
1362
|
+
return _decode_bytes(base_frame, resolver, limits), base_frame
|
|
1363
|
+
if resolver is None:
|
|
1364
|
+
raise DependencyMissing(
|
|
1365
|
+
"frame needs dependency %s but no resolver was supplied"
|
|
1366
|
+
% meta_base.hex()
|
|
1367
|
+
)
|
|
1368
|
+
fetched = resolver.resolve_frame(meta_base)
|
|
1369
|
+
data = resolver.resolve(meta_base, frame.header.decoded_len, 0)
|
|
1370
|
+
if not isinstance(data, (bytes, bytearray, memoryview)):
|
|
1371
|
+
raise DependencyMismatch("resolver returned a non-bytes value")
|
|
1372
|
+
data = bytes(data)
|
|
1373
|
+
if len(data) != frame.header.decoded_len:
|
|
1374
|
+
raise DependencyMismatch(
|
|
1375
|
+
"dependency length %d differs from declared %d"
|
|
1376
|
+
% (len(data), frame.header.decoded_len)
|
|
1377
|
+
)
|
|
1378
|
+
from .hashing import content_root
|
|
1379
|
+
|
|
1380
|
+
if content_root(data) != meta_base:
|
|
1381
|
+
raise DependencyMismatch("dependency content hash mismatch")
|
|
1382
|
+
return data, fetched
|
|
1383
|
+
|
|
1384
|
+
|
|
1385
|
+
def _verify_base_frame(
|
|
1386
|
+
supplied: bytes, meta, limits: Limits, what: str
|
|
1387
|
+
) -> None:
|
|
1388
|
+
from .frame import parse as _parse
|
|
1389
|
+
|
|
1390
|
+
parsed = _parse(supplied, limits)
|
|
1391
|
+
if parsed.header.content_hash != meta.base:
|
|
1392
|
+
raise DependencyMismatch(
|
|
1393
|
+
"%s supplied base frame content hash does not equal the descriptor base"
|
|
1394
|
+
% what
|
|
1395
|
+
)
|
|
1396
|
+
try:
|
|
1397
|
+
base_dtype, base_shape = _base_frame_descriptor(parsed, limits)
|
|
1398
|
+
except (Malformed, LimitExceeded) as exc:
|
|
1399
|
+
raise DependencyMismatch(exc.message)
|
|
1400
|
+
if base_dtype != meta.dtype:
|
|
1401
|
+
raise DependencyMismatch(
|
|
1402
|
+
"%s base frame dtype `%s` does not match %s dtype `%s`"
|
|
1403
|
+
% (what, base_dtype, what, meta.dtype)
|
|
1404
|
+
)
|
|
1405
|
+
if base_shape != meta.shape:
|
|
1406
|
+
raise DependencyMismatch(
|
|
1407
|
+
"%s base frame shape does not match the %s shape" % (what, what)
|
|
1408
|
+
)
|
|
1409
|
+
|
|
1410
|
+
|
|
1411
|
+
def _check_quantized(dtype: str, data: bytes, what: str) -> None:
|
|
1412
|
+
if dtype == "q4_0":
|
|
1413
|
+
for index, byte in enumerate(data):
|
|
1414
|
+
if byte > 15:
|
|
1415
|
+
raise Corrupt(
|
|
1416
|
+
"%s q4_0 word %d is %d, outside 0..=15"
|
|
1417
|
+
% (what, index, byte)
|
|
1418
|
+
)
|
|
1419
|
+
|
|
1420
|
+
|
|
1421
|
+
def decode_weight_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
|
|
1422
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
1423
|
+
meta = parse_weight_descriptor(descriptor_bytes, limits)
|
|
1424
|
+
count = validate_weight_meta(meta, limits)
|
|
1425
|
+
expected_len = _checked_mul(
|
|
1426
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1427
|
+
"WEIGHT_DELTA element byte length",
|
|
1428
|
+
)
|
|
1429
|
+
if frame.header.decoded_len != expected_len:
|
|
1430
|
+
raise Corrupt(
|
|
1431
|
+
"WEIGHT_DELTA decoded_len is %d but the descriptor requires %d"
|
|
1432
|
+
% (frame.header.decoded_len, expected_len)
|
|
1433
|
+
)
|
|
1434
|
+
_check_len(expected_len, limits.max_decoded_bytes, "WEIGHT_DELTA element bytes")
|
|
1435
|
+
if meta.payload_len != len(payload):
|
|
1436
|
+
raise Malformed(
|
|
1437
|
+
"WEIGHT_DELTA payload is %d bytes but payload_len declares %d"
|
|
1438
|
+
% (len(payload), meta.payload_len)
|
|
1439
|
+
)
|
|
1440
|
+
if frame.header.dep_count != 1:
|
|
1441
|
+
raise Malformed("WEIGHT_DELTA requires exactly one dependency")
|
|
1442
|
+
dep = frame.dependencies[0]
|
|
1443
|
+
if dep.hash != meta.base:
|
|
1444
|
+
raise Malformed(
|
|
1445
|
+
"WEIGHT_DELTA dependency hash must equal the descriptor base"
|
|
1446
|
+
)
|
|
1447
|
+
if dep.decoded_len != frame.header.decoded_len:
|
|
1448
|
+
raise Malformed(
|
|
1449
|
+
"WEIGHT_DELTA dependency decoded_len must equal the frame decoded_len"
|
|
1450
|
+
)
|
|
1451
|
+
planes = _decode_planes(meta.dtype, payload, count, "WEIGHT_DELTA")
|
|
1452
|
+
delta_bytes = _build_delta_bytes(meta.dtype, planes, count)
|
|
1453
|
+
if values_hash(delta_bytes) != meta.delta_root:
|
|
1454
|
+
raise Corrupt("WEIGHT_DELTA delta_root mismatch")
|
|
1455
|
+
|
|
1456
|
+
base_bytes, supplied = _base_bytes(
|
|
1457
|
+
frame, meta.base, resolver, limits, base_frame, "WEIGHT_DELTA"
|
|
1458
|
+
)
|
|
1459
|
+
if supplied is not None:
|
|
1460
|
+
_verify_base_frame(supplied, meta, limits, "WEIGHT_DELTA")
|
|
1461
|
+
if len(base_bytes) != expected_len:
|
|
1462
|
+
raise DependencyMismatch(
|
|
1463
|
+
"base frame decodes to %d bytes but the descriptor requires %d"
|
|
1464
|
+
% (len(base_bytes), expected_len)
|
|
1465
|
+
)
|
|
1466
|
+
_check_quantized(meta.dtype, base_bytes, "WEIGHT_DELTA")
|
|
1467
|
+
|
|
1468
|
+
mask = _mod_mask(meta.dtype)
|
|
1469
|
+
target = bytearray()
|
|
1470
|
+
for index in range(count):
|
|
1471
|
+
delta = _delta_word_from_planes(meta.dtype, planes, index)
|
|
1472
|
+
base_word = _word_at(meta.dtype, base_bytes, index)
|
|
1473
|
+
_push_word(meta.dtype, (delta + base_word) & mask, target)
|
|
1474
|
+
if len(target) != expected_len:
|
|
1475
|
+
raise Corrupt(
|
|
1476
|
+
"WEIGHT_DELTA reconstructed element length is inconsistent"
|
|
1477
|
+
)
|
|
1478
|
+
if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
|
|
1479
|
+
raise Corrupt(
|
|
1480
|
+
"WEIGHT_DELTA content_hash mismatch against the reconstructed tensor"
|
|
1481
|
+
)
|
|
1482
|
+
return bytes(target)
|
|
1483
|
+
|
|
1484
|
+
|
|
1485
|
+
# --- mode 14 ---------------------------------------------------------------
|
|
1486
|
+
|
|
1487
|
+
|
|
1488
|
+
@dataclass
|
|
1489
|
+
class SemanticDeltaMeta:
|
|
1490
|
+
dtype: str
|
|
1491
|
+
shape: "list[int]"
|
|
1492
|
+
element_count: int
|
|
1493
|
+
word_bits: int
|
|
1494
|
+
base: bytes
|
|
1495
|
+
base_dtype: str
|
|
1496
|
+
base_shape: "list[int]"
|
|
1497
|
+
base_semantic_root: bytes
|
|
1498
|
+
policy: SemanticPolicy
|
|
1499
|
+
delta_root: bytes
|
|
1500
|
+
plane_codec: str
|
|
1501
|
+
payload_len: int
|
|
1502
|
+
|
|
1503
|
+
|
|
1504
|
+
_SEMANTIC_DELTA_KEYS = (
|
|
1505
|
+
"schema",
|
|
1506
|
+
"dtype",
|
|
1507
|
+
"shape",
|
|
1508
|
+
"element_count",
|
|
1509
|
+
"word_bits",
|
|
1510
|
+
"base",
|
|
1511
|
+
"base_dtype",
|
|
1512
|
+
"base_shape",
|
|
1513
|
+
"base_semantic_root",
|
|
1514
|
+
"policy",
|
|
1515
|
+
"delta_root",
|
|
1516
|
+
"plane_codec",
|
|
1517
|
+
"payload_len",
|
|
1518
|
+
)
|
|
1519
|
+
|
|
1520
|
+
|
|
1521
|
+
def parse_semantic_delta_descriptor(descriptor_bytes: bytes, limits: Limits):
|
|
1522
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "SEMANTIC_DELTA")
|
|
1523
|
+
_exact_keys(map_value, _SEMANTIC_DELTA_KEYS, "SEMANTIC_DELTA")
|
|
1524
|
+
schema = _text(map_value, "schema")
|
|
1525
|
+
if schema != SEMANTIC_DELTA_SCHEMA:
|
|
1526
|
+
raise Malformed("unsupported SEMANTIC_DELTA schema `%s`" % schema)
|
|
1527
|
+
dtype = _weight_dtype(_text(map_value, "dtype"))
|
|
1528
|
+
shape = _shape(map_value)
|
|
1529
|
+
element_count = _u64(map_value, "element_count")
|
|
1530
|
+
word_bits = _u64(map_value, "word_bits")
|
|
1531
|
+
base = _hash_field(map_value, "base")
|
|
1532
|
+
base_dtype = _weight_dtype(_text(map_value, "base_dtype"))
|
|
1533
|
+
base_shape = _array_u64(map_value, "base_shape")
|
|
1534
|
+
base_semantic_root = _hash_field(map_value, "base_semantic_root")
|
|
1535
|
+
policy = parse_semantic_policy(map_value.get("policy"), dtype, element_count)
|
|
1536
|
+
delta_root = _hash_field(map_value, "delta_root")
|
|
1537
|
+
plane_codec = _text(map_value, "plane_codec")
|
|
1538
|
+
payload_len = _u64(map_value, "payload_len")
|
|
1539
|
+
meta = SemanticDeltaMeta(
|
|
1540
|
+
dtype=dtype,
|
|
1541
|
+
shape=shape,
|
|
1542
|
+
element_count=element_count,
|
|
1543
|
+
word_bits=word_bits,
|
|
1544
|
+
base=base,
|
|
1545
|
+
base_dtype=base_dtype,
|
|
1546
|
+
base_shape=base_shape,
|
|
1547
|
+
base_semantic_root=base_semantic_root,
|
|
1548
|
+
policy=policy,
|
|
1549
|
+
delta_root=delta_root,
|
|
1550
|
+
plane_codec=plane_codec,
|
|
1551
|
+
payload_len=payload_len,
|
|
1552
|
+
)
|
|
1553
|
+
validate_semantic_delta_meta(meta, limits)
|
|
1554
|
+
return meta
|
|
1555
|
+
|
|
1556
|
+
|
|
1557
|
+
def validate_semantic_delta_meta(meta: SemanticDeltaMeta, limits: Limits) -> int:
|
|
1558
|
+
if meta.plane_codec != PLANE_CODEC:
|
|
1559
|
+
raise Malformed(
|
|
1560
|
+
"SEMANTIC_DELTA plane_codec `%s` is not implemented (only %s)"
|
|
1561
|
+
% (meta.plane_codec, PLANE_CODEC)
|
|
1562
|
+
)
|
|
1563
|
+
if WEIGHT_WORD_BITS[meta.dtype] != meta.word_bits:
|
|
1564
|
+
raise Malformed(
|
|
1565
|
+
"SEMANTIC_DELTA word_bits %d disagrees with dtype %s (%d bits)"
|
|
1566
|
+
% (meta.word_bits, meta.dtype, WEIGHT_WORD_BITS[meta.dtype])
|
|
1567
|
+
)
|
|
1568
|
+
if meta.base_dtype != meta.dtype:
|
|
1569
|
+
raise Malformed(
|
|
1570
|
+
"SEMANTIC_DELTA base_dtype `%s` must equal dtype `%s`"
|
|
1571
|
+
% (meta.base_dtype, meta.dtype)
|
|
1572
|
+
)
|
|
1573
|
+
if meta.base_shape != meta.shape:
|
|
1574
|
+
raise Malformed("SEMANTIC_DELTA base_shape must equal shape")
|
|
1575
|
+
count = _validate_shape(meta.shape, limits, "SEMANTIC_DELTA")
|
|
1576
|
+
if meta.element_count != count:
|
|
1577
|
+
raise Malformed(
|
|
1578
|
+
"SEMANTIC_DELTA element_count %d does not match shape product %d"
|
|
1579
|
+
% (meta.element_count, count)
|
|
1580
|
+
)
|
|
1581
|
+
_checked_mul(
|
|
1582
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1583
|
+
"SEMANTIC_DELTA element byte length",
|
|
1584
|
+
)
|
|
1585
|
+
meta.policy.validate(meta.dtype, count)
|
|
1586
|
+
return count
|
|
1587
|
+
|
|
1588
|
+
|
|
1589
|
+
def semantic_delta_metadata(frame, limits: Limits) -> SemanticDeltaMeta:
|
|
1590
|
+
descriptor_bytes, _payload = _split_prefix(frame.body)
|
|
1591
|
+
return parse_semantic_delta_descriptor(descriptor_bytes, limits)
|
|
1592
|
+
|
|
1593
|
+
|
|
1594
|
+
def decode_semantic_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
|
|
1595
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
1596
|
+
meta = parse_semantic_delta_descriptor(descriptor_bytes, limits)
|
|
1597
|
+
count = validate_semantic_delta_meta(meta, limits)
|
|
1598
|
+
expected_len = _checked_mul(
|
|
1599
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1600
|
+
"SEMANTIC_DELTA element byte length",
|
|
1601
|
+
)
|
|
1602
|
+
if frame.header.decoded_len != expected_len:
|
|
1603
|
+
raise Corrupt(
|
|
1604
|
+
"SEMANTIC_DELTA decoded_len is %d but the descriptor requires %d"
|
|
1605
|
+
% (frame.header.decoded_len, expected_len)
|
|
1606
|
+
)
|
|
1607
|
+
_check_len(expected_len, limits.max_decoded_bytes, "SEMANTIC_DELTA element bytes")
|
|
1608
|
+
if meta.payload_len != len(payload):
|
|
1609
|
+
raise Malformed(
|
|
1610
|
+
"SEMANTIC_DELTA payload is %d bytes but payload_len declares %d"
|
|
1611
|
+
% (len(payload), meta.payload_len)
|
|
1612
|
+
)
|
|
1613
|
+
if frame.header.dep_count != 1:
|
|
1614
|
+
raise Malformed("SEMANTIC_DELTA requires exactly one dependency")
|
|
1615
|
+
dep = frame.dependencies[0]
|
|
1616
|
+
if dep.hash != meta.base:
|
|
1617
|
+
raise Malformed(
|
|
1618
|
+
"SEMANTIC_DELTA dependency hash must equal the descriptor base"
|
|
1619
|
+
)
|
|
1620
|
+
if dep.decoded_len != frame.header.decoded_len:
|
|
1621
|
+
raise Malformed(
|
|
1622
|
+
"SEMANTIC_DELTA dependency decoded_len must equal the frame decoded_len"
|
|
1623
|
+
)
|
|
1624
|
+
planes = _decode_planes(meta.dtype, payload, count, "SEMANTIC_DELTA")
|
|
1625
|
+
delta_bytes = _build_delta_bytes(meta.dtype, planes, count)
|
|
1626
|
+
if values_hash(delta_bytes) != meta.delta_root:
|
|
1627
|
+
raise Corrupt("SEMANTIC_DELTA delta_root mismatch")
|
|
1628
|
+
|
|
1629
|
+
base_bytes, supplied = _base_bytes(
|
|
1630
|
+
frame, meta.base, resolver, limits, base_frame, "SEMANTIC_DELTA"
|
|
1631
|
+
)
|
|
1632
|
+
if supplied is not None:
|
|
1633
|
+
_verify_base_frame(supplied, meta, limits, "SEMANTIC_DELTA")
|
|
1634
|
+
if len(base_bytes) != expected_len:
|
|
1635
|
+
raise DependencyMismatch(
|
|
1636
|
+
"base frame decodes to %d bytes but the descriptor requires %d"
|
|
1637
|
+
% (len(base_bytes), expected_len)
|
|
1638
|
+
)
|
|
1639
|
+
_check_quantized(meta.dtype, base_bytes, "SEMANTIC_DELTA")
|
|
1640
|
+
|
|
1641
|
+
try:
|
|
1642
|
+
recomputed = semantic_root(
|
|
1643
|
+
meta.dtype, meta.shape, base_bytes, meta.policy, limits
|
|
1644
|
+
)
|
|
1645
|
+
except (Malformed, NonCanonical, LimitExceeded, Corrupt) as exc:
|
|
1646
|
+
raise DependencyMismatch(exc.message)
|
|
1647
|
+
if recomputed != meta.base_semantic_root:
|
|
1648
|
+
raise DependencyMismatch(
|
|
1649
|
+
"base semantic root does not equal the descriptor (no approximate "
|
|
1650
|
+
"reconstruction exists)"
|
|
1651
|
+
)
|
|
1652
|
+
|
|
1653
|
+
mask = _mod_mask(meta.dtype)
|
|
1654
|
+
target = bytearray()
|
|
1655
|
+
for index in range(count):
|
|
1656
|
+
delta = _delta_word_from_planes(meta.dtype, planes, index)
|
|
1657
|
+
base_word = _word_at(meta.dtype, base_bytes, index)
|
|
1658
|
+
_push_word(meta.dtype, (delta + base_word) & mask, target)
|
|
1659
|
+
if len(target) != expected_len:
|
|
1660
|
+
raise Corrupt(
|
|
1661
|
+
"SEMANTIC_DELTA reconstructed element length is inconsistent"
|
|
1662
|
+
)
|
|
1663
|
+
if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
|
|
1664
|
+
raise Corrupt(
|
|
1665
|
+
"SEMANTIC_DELTA content_hash mismatch against the reconstructed tensor"
|
|
1666
|
+
)
|
|
1667
|
+
return bytes(target)
|
|
1668
|
+
|
|
1669
|
+
|
|
1670
|
+
# --- mode 16 ---------------------------------------------------------------
|
|
1671
|
+
|
|
1672
|
+
_SPARSE_KEYS = (
|
|
1673
|
+
"schema",
|
|
1674
|
+
"dtype",
|
|
1675
|
+
"shape",
|
|
1676
|
+
"element_count",
|
|
1677
|
+
"storage_bytes",
|
|
1678
|
+
"op",
|
|
1679
|
+
"sparse_codec",
|
|
1680
|
+
"base_content_root",
|
|
1681
|
+
"delta_root",
|
|
1682
|
+
)
|
|
1683
|
+
|
|
1684
|
+
|
|
1685
|
+
@dataclass
|
|
1686
|
+
class SparseDeltaMeta:
|
|
1687
|
+
dtype: str
|
|
1688
|
+
shape: "list[int]"
|
|
1689
|
+
element_count: int
|
|
1690
|
+
storage_bytes: int
|
|
1691
|
+
op: str
|
|
1692
|
+
sparse_codec: str
|
|
1693
|
+
base: bytes
|
|
1694
|
+
delta_root: bytes
|
|
1695
|
+
|
|
1696
|
+
|
|
1697
|
+
def parse_sparse_descriptor(descriptor_bytes: bytes, limits: Limits) -> SparseDeltaMeta:
|
|
1698
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "SPARSE_DELTA")
|
|
1699
|
+
_exact_keys(map_value, _SPARSE_KEYS, "SPARSE_DELTA")
|
|
1700
|
+
schema = _text(map_value, "schema")
|
|
1701
|
+
if schema != SPARSE_DELTA_SCHEMA:
|
|
1702
|
+
raise Malformed("unsupported SPARSE_DELTA schema `%s`" % schema)
|
|
1703
|
+
dtype = _weight_dtype(_text(map_value, "dtype"))
|
|
1704
|
+
shape = _shape(map_value)
|
|
1705
|
+
element_count = _u64(map_value, "element_count")
|
|
1706
|
+
storage_bytes = _u64(map_value, "storage_bytes")
|
|
1707
|
+
op = _text(map_value, "op")
|
|
1708
|
+
sparse_codec = _text(map_value, "sparse_codec")
|
|
1709
|
+
base = _hash_field(map_value, "base_content_root")
|
|
1710
|
+
delta_root = _hash_field(map_value, "delta_root")
|
|
1711
|
+
meta = SparseDeltaMeta(
|
|
1712
|
+
dtype=dtype,
|
|
1713
|
+
shape=shape,
|
|
1714
|
+
element_count=element_count,
|
|
1715
|
+
storage_bytes=storage_bytes,
|
|
1716
|
+
op=op,
|
|
1717
|
+
sparse_codec=sparse_codec,
|
|
1718
|
+
base=base,
|
|
1719
|
+
delta_root=delta_root,
|
|
1720
|
+
)
|
|
1721
|
+
validate_sparse_meta(meta, limits)
|
|
1722
|
+
return meta
|
|
1723
|
+
|
|
1724
|
+
|
|
1725
|
+
def validate_sparse_meta(meta: SparseDeltaMeta, limits: Limits) -> int:
|
|
1726
|
+
if meta.op != DELTA_OP:
|
|
1727
|
+
raise Malformed(
|
|
1728
|
+
"SPARSE_DELTA op `%s` is not implemented (only %s)"
|
|
1729
|
+
% (meta.op, DELTA_OP)
|
|
1730
|
+
)
|
|
1731
|
+
if meta.sparse_codec != SPARSE_CODEC:
|
|
1732
|
+
raise Malformed(
|
|
1733
|
+
"SPARSE_DELTA sparse_codec `%s` is not implemented (only %s)"
|
|
1734
|
+
% (meta.sparse_codec, SPARSE_CODEC)
|
|
1735
|
+
)
|
|
1736
|
+
if meta.storage_bytes != WEIGHT_STORAGE_BYTES[meta.dtype]:
|
|
1737
|
+
raise Malformed(
|
|
1738
|
+
"SPARSE_DELTA storage_bytes %d disagrees with dtype %s (%d bytes)"
|
|
1739
|
+
% (
|
|
1740
|
+
meta.storage_bytes,
|
|
1741
|
+
meta.dtype,
|
|
1742
|
+
WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1743
|
+
)
|
|
1744
|
+
)
|
|
1745
|
+
count = _validate_shape(meta.shape, limits, "SPARSE_DELTA")
|
|
1746
|
+
if meta.element_count != count:
|
|
1747
|
+
raise Malformed(
|
|
1748
|
+
"SPARSE_DELTA element_count %d does not match shape product %d"
|
|
1749
|
+
% (meta.element_count, count)
|
|
1750
|
+
)
|
|
1751
|
+
_checked_mul(
|
|
1752
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1753
|
+
"SPARSE_DELTA element byte length",
|
|
1754
|
+
)
|
|
1755
|
+
return count
|
|
1756
|
+
|
|
1757
|
+
|
|
1758
|
+
def sparse_delta_metadata(frame, limits: Limits) -> SparseDeltaMeta:
|
|
1759
|
+
descriptor_bytes, _payload = _split_prefix(frame.body)
|
|
1760
|
+
return parse_sparse_descriptor(descriptor_bytes, limits)
|
|
1761
|
+
|
|
1762
|
+
|
|
1763
|
+
def _decode_runs(
|
|
1764
|
+
dtype: str, payload: bytes, element_count: int, limits: Limits
|
|
1765
|
+
) -> bytes:
|
|
1766
|
+
if len(payload) < 4:
|
|
1767
|
+
raise Malformed("SPARSE_DELTA body is shorter than its 4-byte run_count")
|
|
1768
|
+
run_count = int.from_bytes(payload[0:4], "little")
|
|
1769
|
+
if run_count > limits.max_sparse_runs:
|
|
1770
|
+
raise LimitExceeded(
|
|
1771
|
+
"SPARSE_DELTA run_count %d exceeds limit %d"
|
|
1772
|
+
% (run_count, limits.max_sparse_runs)
|
|
1773
|
+
)
|
|
1774
|
+
if run_count > element_count:
|
|
1775
|
+
raise Malformed(
|
|
1776
|
+
"SPARSE_DELTA run_count %d exceeds element_count %d"
|
|
1777
|
+
% (run_count, element_count)
|
|
1778
|
+
)
|
|
1779
|
+
storage = WEIGHT_STORAGE_BYTES[dtype]
|
|
1780
|
+
delta_len = _checked_mul(element_count, storage, "SPARSE_DELTA delta bytes")
|
|
1781
|
+
delta = bytearray(delta_len)
|
|
1782
|
+
position = 4
|
|
1783
|
+
emitted = 0
|
|
1784
|
+
for record in range(run_count):
|
|
1785
|
+
zero_run, position = varint.decode(payload, position)
|
|
1786
|
+
emitted = _checked_add(
|
|
1787
|
+
emitted, zero_run, "SPARSE_DELTA zero-run length"
|
|
1788
|
+
)
|
|
1789
|
+
if emitted >= element_count:
|
|
1790
|
+
raise Malformed(
|
|
1791
|
+
"SPARSE_DELTA run %d skips %d words past element_count %d"
|
|
1792
|
+
% (record, zero_run, element_count)
|
|
1793
|
+
)
|
|
1794
|
+
end = position + storage
|
|
1795
|
+
if end > len(payload):
|
|
1796
|
+
raise Malformed("SPARSE_DELTA run %d record is truncated" % record)
|
|
1797
|
+
word = payload[position:end]
|
|
1798
|
+
if not any(word):
|
|
1799
|
+
raise Corrupt(
|
|
1800
|
+
"SPARSE_DELTA run %d stores a zero word (non-canonical: zeros "
|
|
1801
|
+
"are implicit)" % record
|
|
1802
|
+
)
|
|
1803
|
+
if dtype == "q4_0" and word[0] > 15:
|
|
1804
|
+
raise Corrupt(
|
|
1805
|
+
"SPARSE_DELTA q4_0 run %d word is %d, outside 0..=15"
|
|
1806
|
+
% (record, word[0])
|
|
1807
|
+
)
|
|
1808
|
+
offset = emitted * storage
|
|
1809
|
+
delta[offset : offset + storage] = word
|
|
1810
|
+
position = end
|
|
1811
|
+
emitted += 1
|
|
1812
|
+
if position != len(payload):
|
|
1813
|
+
raise Corrupt(
|
|
1814
|
+
"SPARSE_DELTA payload has %d trailing bytes after the %d declared runs"
|
|
1815
|
+
% (len(payload) - position, run_count)
|
|
1816
|
+
)
|
|
1817
|
+
return bytes(delta)
|
|
1818
|
+
|
|
1819
|
+
|
|
1820
|
+
def decode_sparse_delta(frame, resolver, limits: Limits, base_frame=None) -> bytes:
|
|
1821
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
1822
|
+
meta = parse_sparse_descriptor(descriptor_bytes, limits)
|
|
1823
|
+
count = validate_sparse_meta(meta, limits)
|
|
1824
|
+
expected_len = _checked_mul(
|
|
1825
|
+
count, WEIGHT_STORAGE_BYTES[meta.dtype],
|
|
1826
|
+
"SPARSE_DELTA element byte length",
|
|
1827
|
+
)
|
|
1828
|
+
if frame.header.decoded_len != expected_len:
|
|
1829
|
+
raise Corrupt(
|
|
1830
|
+
"SPARSE_DELTA decoded_len is %d but the descriptor requires %d"
|
|
1831
|
+
% (frame.header.decoded_len, expected_len)
|
|
1832
|
+
)
|
|
1833
|
+
_check_len(expected_len, limits.max_decoded_bytes, "SPARSE_DELTA element bytes")
|
|
1834
|
+
if frame.header.dep_count != 1:
|
|
1835
|
+
raise Malformed("SPARSE_DELTA requires exactly one dependency")
|
|
1836
|
+
dep = frame.dependencies[0]
|
|
1837
|
+
if dep.hash != meta.base:
|
|
1838
|
+
raise Malformed(
|
|
1839
|
+
"SPARSE_DELTA dependency hash must equal the descriptor "
|
|
1840
|
+
"base_content_root"
|
|
1841
|
+
)
|
|
1842
|
+
if dep.decoded_len != frame.header.decoded_len:
|
|
1843
|
+
raise Malformed(
|
|
1844
|
+
"SPARSE_DELTA dependency decoded_len must equal the frame decoded_len"
|
|
1845
|
+
)
|
|
1846
|
+
delta_bytes = _decode_runs(meta.dtype, payload, count, limits)
|
|
1847
|
+
if values_hash(delta_bytes) != meta.delta_root:
|
|
1848
|
+
raise Corrupt("SPARSE_DELTA delta_root mismatch")
|
|
1849
|
+
|
|
1850
|
+
base_bytes, supplied = _base_bytes(
|
|
1851
|
+
frame, meta.base, resolver, limits, base_frame, "SPARSE_DELTA"
|
|
1852
|
+
)
|
|
1853
|
+
if supplied is not None:
|
|
1854
|
+
_verify_base_frame(supplied, meta, limits, "SPARSE_DELTA")
|
|
1855
|
+
if len(base_bytes) != expected_len:
|
|
1856
|
+
raise DependencyMismatch(
|
|
1857
|
+
"base frame decodes to %d bytes but the descriptor requires %d"
|
|
1858
|
+
% (len(base_bytes), expected_len)
|
|
1859
|
+
)
|
|
1860
|
+
_check_quantized(meta.dtype, base_bytes, "SPARSE_DELTA")
|
|
1861
|
+
|
|
1862
|
+
mask = _mod_mask(meta.dtype)
|
|
1863
|
+
target = bytearray()
|
|
1864
|
+
for index in range(count):
|
|
1865
|
+
delta = _word_at(meta.dtype, delta_bytes, index)
|
|
1866
|
+
base_word = _word_at(meta.dtype, base_bytes, index)
|
|
1867
|
+
_push_word(meta.dtype, (delta + base_word) & mask, target)
|
|
1868
|
+
if len(target) != expected_len:
|
|
1869
|
+
raise Corrupt(
|
|
1870
|
+
"SPARSE_DELTA reconstructed element length is inconsistent"
|
|
1871
|
+
)
|
|
1872
|
+
if hash_domain("weft/chunk/v1", bytes(target)) != frame.header.content_hash:
|
|
1873
|
+
raise Corrupt(
|
|
1874
|
+
"SPARSE_DELTA content_hash mismatch against the reconstructed tensor"
|
|
1875
|
+
)
|
|
1876
|
+
return bytes(target)
|
|
1877
|
+
|
|
1878
|
+
|
|
1879
|
+
def activation_stream_root(
|
|
1880
|
+
stream: str,
|
|
1881
|
+
step: int,
|
|
1882
|
+
rank: int,
|
|
1883
|
+
world: int,
|
|
1884
|
+
chunk_count: int,
|
|
1885
|
+
digest_roots: "list[bytes]",
|
|
1886
|
+
limits: Limits,
|
|
1887
|
+
) -> bytes:
|
|
1888
|
+
if len(digest_roots) != chunk_count:
|
|
1889
|
+
raise Malformed(
|
|
1890
|
+
"ACTIVATION_STREAM stream root needs %d digests, got %d"
|
|
1891
|
+
% (chunk_count, len(digest_roots))
|
|
1892
|
+
)
|
|
1893
|
+
record = {
|
|
1894
|
+
"stream": stream,
|
|
1895
|
+
"step": wcb.U64(step),
|
|
1896
|
+
"rank": wcb.U64(rank),
|
|
1897
|
+
"world": wcb.U64(world),
|
|
1898
|
+
"chunk_count": wcb.U64(chunk_count),
|
|
1899
|
+
"digest_roots": [bytes(digest) for digest in digest_roots],
|
|
1900
|
+
}
|
|
1901
|
+
return hash_domain(DOMAIN_AI_STREAM, wcb.encode(record, limits))
|
|
1902
|
+
|
|
1903
|
+
|
|
1904
|
+
# ---------------------------------------------------------------------------
|
|
1905
|
+
# Mode 12 -- KV_CACHE
|
|
1906
|
+
# ---------------------------------------------------------------------------
|
|
1907
|
+
|
|
1908
|
+
_KV_KEYS = (
|
|
1909
|
+
"schema",
|
|
1910
|
+
"model",
|
|
1911
|
+
"layer",
|
|
1912
|
+
"session",
|
|
1913
|
+
"page_index",
|
|
1914
|
+
"page_count",
|
|
1915
|
+
"seq_start",
|
|
1916
|
+
"seq_end",
|
|
1917
|
+
"dtype",
|
|
1918
|
+
"head_count",
|
|
1919
|
+
"head_dim",
|
|
1920
|
+
"shape",
|
|
1921
|
+
"page_bytes",
|
|
1922
|
+
"plane_codec",
|
|
1923
|
+
"payload_len",
|
|
1924
|
+
)
|
|
1925
|
+
|
|
1926
|
+
|
|
1927
|
+
@dataclass
|
|
1928
|
+
class KvMeta:
|
|
1929
|
+
model: str
|
|
1930
|
+
layer: int
|
|
1931
|
+
session: str
|
|
1932
|
+
page_index: int
|
|
1933
|
+
page_count: int
|
|
1934
|
+
seq_start: int
|
|
1935
|
+
seq_end: int
|
|
1936
|
+
dtype: str
|
|
1937
|
+
head_count: int
|
|
1938
|
+
head_dim: int
|
|
1939
|
+
shape: "list[int]"
|
|
1940
|
+
page_bytes: int
|
|
1941
|
+
plane_codec: str
|
|
1942
|
+
payload_len: int
|
|
1943
|
+
|
|
1944
|
+
|
|
1945
|
+
def _kv_dtype(name: str) -> str:
|
|
1946
|
+
if name not in KV_WORD_BYTES:
|
|
1947
|
+
raise Malformed(
|
|
1948
|
+
"unknown KV_CACHE dtype `%s` (v1 allows bf16, fp16, fp32)" % name
|
|
1949
|
+
)
|
|
1950
|
+
return name
|
|
1951
|
+
|
|
1952
|
+
|
|
1953
|
+
def parse_kv_descriptor(descriptor_bytes: bytes, limits: Limits) -> KvMeta:
|
|
1954
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "KV_CACHE")
|
|
1955
|
+
_exact_keys(map_value, _KV_KEYS, "KV_CACHE")
|
|
1956
|
+
schema = _text(map_value, "schema")
|
|
1957
|
+
if schema != KV_CACHE_SCHEMA:
|
|
1958
|
+
raise Malformed("unsupported KV_CACHE schema `%s`" % schema)
|
|
1959
|
+
meta = KvMeta(
|
|
1960
|
+
model=_text(map_value, "model"),
|
|
1961
|
+
layer=_u32(map_value, "layer"),
|
|
1962
|
+
session=_text(map_value, "session"),
|
|
1963
|
+
page_index=_u64(map_value, "page_index"),
|
|
1964
|
+
page_count=_u64(map_value, "page_count"),
|
|
1965
|
+
seq_start=_u64(map_value, "seq_start"),
|
|
1966
|
+
seq_end=_u64(map_value, "seq_end"),
|
|
1967
|
+
dtype=_kv_dtype(_text(map_value, "dtype")),
|
|
1968
|
+
head_count=_u32(map_value, "head_count"),
|
|
1969
|
+
head_dim=_u32(map_value, "head_dim"),
|
|
1970
|
+
shape=_shape(map_value),
|
|
1971
|
+
page_bytes=_u64(map_value, "page_bytes"),
|
|
1972
|
+
plane_codec=_text(map_value, "plane_codec"),
|
|
1973
|
+
payload_len=_u64(map_value, "payload_len"),
|
|
1974
|
+
)
|
|
1975
|
+
validate_kv_meta(meta, limits)
|
|
1976
|
+
return meta
|
|
1977
|
+
|
|
1978
|
+
|
|
1979
|
+
def validate_kv_meta(meta: KvMeta, limits: Limits) -> int:
|
|
1980
|
+
if meta.plane_codec != PLANE_CODEC:
|
|
1981
|
+
raise Malformed(
|
|
1982
|
+
"KV_CACHE plane_codec `%s` is not implemented (only %s)"
|
|
1983
|
+
% (meta.plane_codec, PLANE_CODEC)
|
|
1984
|
+
)
|
|
1985
|
+
if meta.page_count == 0:
|
|
1986
|
+
raise Malformed("KV_CACHE page_count must be at least 1")
|
|
1987
|
+
if meta.page_index >= meta.page_count:
|
|
1988
|
+
raise Malformed(
|
|
1989
|
+
"KV_CACHE page_index %d is not below page_count %d"
|
|
1990
|
+
% (meta.page_index, meta.page_count)
|
|
1991
|
+
)
|
|
1992
|
+
if meta.seq_start >= meta.seq_end:
|
|
1993
|
+
raise Malformed(
|
|
1994
|
+
"KV_CACHE seq_start %d must be less than seq_end %d"
|
|
1995
|
+
% (meta.seq_start, meta.seq_end)
|
|
1996
|
+
)
|
|
1997
|
+
rows = meta.seq_end - meta.seq_start
|
|
1998
|
+
if meta.head_count == 0:
|
|
1999
|
+
raise Malformed("KV_CACHE head_count must be greater than zero")
|
|
2000
|
+
if meta.head_dim == 0:
|
|
2001
|
+
raise Malformed("KV_CACHE head_dim must be greater than zero")
|
|
2002
|
+
expected_shape = [rows, meta.head_count, meta.head_dim]
|
|
2003
|
+
if meta.shape != expected_shape:
|
|
2004
|
+
raise Malformed(
|
|
2005
|
+
"KV_CACHE shape %r must equal [page_rows, head_count, head_dim] = %r"
|
|
2006
|
+
% (meta.shape, expected_shape)
|
|
2007
|
+
)
|
|
2008
|
+
if len(meta.shape) > limits.max_tensor_rank:
|
|
2009
|
+
raise LimitExceeded(
|
|
2010
|
+
"KV_CACHE rank %d exceeds limit %d"
|
|
2011
|
+
% (len(meta.shape), limits.max_tensor_rank)
|
|
2012
|
+
)
|
|
2013
|
+
element_count = _checked_mul(rows, meta.head_count, "KV_CACHE element count")
|
|
2014
|
+
element_count = _checked_mul(
|
|
2015
|
+
element_count, meta.head_dim, "KV_CACHE element count"
|
|
2016
|
+
)
|
|
2017
|
+
if element_count > limits.max_tensor_elements:
|
|
2018
|
+
raise LimitExceeded(
|
|
2019
|
+
"KV_CACHE element count %d exceeds limit %d"
|
|
2020
|
+
% (element_count, limits.max_tensor_elements)
|
|
2021
|
+
)
|
|
2022
|
+
computed = _checked_mul(
|
|
2023
|
+
element_count, KV_WORD_BYTES[meta.dtype], "KV_CACHE page byte count"
|
|
2024
|
+
)
|
|
2025
|
+
if meta.page_bytes != computed:
|
|
2026
|
+
raise Malformed(
|
|
2027
|
+
"KV_CACHE page_bytes %d does not equal page_rows * head_count * "
|
|
2028
|
+
"head_dim * word_bytes = %d" % (meta.page_bytes, computed)
|
|
2029
|
+
)
|
|
2030
|
+
_check_len(meta.page_bytes, limits.max_decoded_bytes, "KV_CACHE page")
|
|
2031
|
+
return computed
|
|
2032
|
+
|
|
2033
|
+
|
|
2034
|
+
def _decode_byte_planes(
|
|
2035
|
+
payload: bytes, word_bytes: int, byte_len: int, what: str
|
|
2036
|
+
) -> bytes:
|
|
2037
|
+
if byte_len % word_bytes != 0:
|
|
2038
|
+
raise Malformed("%s byte length is not a whole number of elements" % what)
|
|
2039
|
+
element_count = byte_len // word_bytes
|
|
2040
|
+
plane_count = word_bytes
|
|
2041
|
+
planes = []
|
|
2042
|
+
offset = 0
|
|
2043
|
+
for _ in range(plane_count):
|
|
2044
|
+
if offset >= len(payload):
|
|
2045
|
+
raise Corrupt("%s payload ends before all declared planes" % what)
|
|
2046
|
+
plane, consumed = huffman.decode_prefix(payload[offset:], element_count)
|
|
2047
|
+
offset += consumed
|
|
2048
|
+
planes.append(plane)
|
|
2049
|
+
if offset != len(payload):
|
|
2050
|
+
raise Corrupt(
|
|
2051
|
+
"%s payload has trailing bytes after the declared planes" % what
|
|
2052
|
+
)
|
|
2053
|
+
out = bytearray()
|
|
2054
|
+
for index in range(element_count):
|
|
2055
|
+
word = 0
|
|
2056
|
+
for plane_index, plane in enumerate(planes):
|
|
2057
|
+
word |= plane[index] << (8 * plane_index)
|
|
2058
|
+
out.extend(word.to_bytes(word_bytes, "little"))
|
|
2059
|
+
return bytes(out)
|
|
2060
|
+
|
|
2061
|
+
|
|
2062
|
+
def decode_kv_cache(frame, limits: Limits) -> bytes:
|
|
2063
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
2064
|
+
meta = parse_kv_descriptor(descriptor_bytes, limits)
|
|
2065
|
+
page_bytes = validate_kv_meta(meta, limits)
|
|
2066
|
+
if frame.header.decoded_len != page_bytes:
|
|
2067
|
+
raise Corrupt(
|
|
2068
|
+
"KV_CACHE decoded_len is %d but the descriptor requires %d"
|
|
2069
|
+
% (frame.header.decoded_len, page_bytes)
|
|
2070
|
+
)
|
|
2071
|
+
if meta.payload_len != len(payload):
|
|
2072
|
+
raise Malformed(
|
|
2073
|
+
"KV_CACHE payload is %d bytes but payload_len declares %d"
|
|
2074
|
+
% (len(payload), meta.payload_len)
|
|
2075
|
+
)
|
|
2076
|
+
out = _decode_byte_planes(
|
|
2077
|
+
payload, KV_WORD_BYTES[meta.dtype], page_bytes, "KV_CACHE"
|
|
2078
|
+
)
|
|
2079
|
+
if len(out) != page_bytes:
|
|
2080
|
+
raise Corrupt("KV_CACHE reconstructed page length is inconsistent")
|
|
2081
|
+
if hash_domain("weft/chunk/v1", out) != frame.header.content_hash:
|
|
2082
|
+
raise Corrupt(
|
|
2083
|
+
"KV_CACHE content_hash mismatch against the reconstructed page"
|
|
2084
|
+
)
|
|
2085
|
+
return out
|
|
2086
|
+
|
|
2087
|
+
|
|
2088
|
+
# ---------------------------------------------------------------------------
|
|
2089
|
+
# Mode 13 -- ACTIVATION_STREAM
|
|
2090
|
+
# ---------------------------------------------------------------------------
|
|
2091
|
+
|
|
2092
|
+
_ACTIVATION_KEYS = (
|
|
2093
|
+
"schema",
|
|
2094
|
+
"stream",
|
|
2095
|
+
"step",
|
|
2096
|
+
"rank",
|
|
2097
|
+
"world",
|
|
2098
|
+
"chunk_index",
|
|
2099
|
+
"chunk_count",
|
|
2100
|
+
"dtype",
|
|
2101
|
+
"shape",
|
|
2102
|
+
"element_count",
|
|
2103
|
+
"digest_root",
|
|
2104
|
+
"plane_codec",
|
|
2105
|
+
"payload_len",
|
|
2106
|
+
)
|
|
2107
|
+
|
|
2108
|
+
|
|
2109
|
+
@dataclass
|
|
2110
|
+
class ActivationMeta:
|
|
2111
|
+
stream: str
|
|
2112
|
+
step: int
|
|
2113
|
+
rank: int
|
|
2114
|
+
world: int
|
|
2115
|
+
chunk_index: int
|
|
2116
|
+
chunk_count: int
|
|
2117
|
+
dtype: str
|
|
2118
|
+
shape: "list[int]"
|
|
2119
|
+
element_count: int
|
|
2120
|
+
digest_root: bytes
|
|
2121
|
+
plane_codec: str
|
|
2122
|
+
payload_len: int
|
|
2123
|
+
|
|
2124
|
+
|
|
2125
|
+
def _stream_dtype(name: str) -> str:
|
|
2126
|
+
if name not in KV_WORD_BYTES:
|
|
2127
|
+
raise Malformed(
|
|
2128
|
+
"unknown ACTIVATION_STREAM dtype `%s` (v1 allows bf16, fp16, fp32)"
|
|
2129
|
+
% name
|
|
2130
|
+
)
|
|
2131
|
+
return name
|
|
2132
|
+
|
|
2133
|
+
|
|
2134
|
+
def _shape_product(shape: "list[int]", limits: Limits) -> int:
|
|
2135
|
+
if not shape:
|
|
2136
|
+
raise Malformed("ACTIVATION_STREAM shape rank must be at least 1")
|
|
2137
|
+
return _validate_shape(shape, limits, "ACTIVATION_STREAM")
|
|
2138
|
+
|
|
2139
|
+
|
|
2140
|
+
def parse_activation_descriptor(
|
|
2141
|
+
descriptor_bytes: bytes, limits: Limits
|
|
2142
|
+
) -> ActivationMeta:
|
|
2143
|
+
map_value = _map(
|
|
2144
|
+
_decode_descriptor(descriptor_bytes, limits), "ACTIVATION_STREAM"
|
|
2145
|
+
)
|
|
2146
|
+
_exact_keys(map_value, _ACTIVATION_KEYS, "ACTIVATION_STREAM")
|
|
2147
|
+
schema = _text(map_value, "schema")
|
|
2148
|
+
if schema != ACTIVATION_STREAM_SCHEMA:
|
|
2149
|
+
raise Malformed("unsupported ACTIVATION_STREAM schema `%s`" % schema)
|
|
2150
|
+
stream = _text(map_value, "stream")
|
|
2151
|
+
if stream not in ("activation", "gradient"):
|
|
2152
|
+
raise Malformed(
|
|
2153
|
+
"unknown ACTIVATION_STREAM stream `%s` (v1 allows activation, "
|
|
2154
|
+
"gradient)" % stream
|
|
2155
|
+
)
|
|
2156
|
+
meta = ActivationMeta(
|
|
2157
|
+
stream=stream,
|
|
2158
|
+
step=_u64(map_value, "step"),
|
|
2159
|
+
rank=_u32(map_value, "rank"),
|
|
2160
|
+
world=_u32(map_value, "world"),
|
|
2161
|
+
chunk_index=_u64(map_value, "chunk_index"),
|
|
2162
|
+
chunk_count=_u64(map_value, "chunk_count"),
|
|
2163
|
+
dtype=_stream_dtype(_text(map_value, "dtype")),
|
|
2164
|
+
shape=_shape(map_value),
|
|
2165
|
+
element_count=_u64(map_value, "element_count"),
|
|
2166
|
+
digest_root=_hash_field(map_value, "digest_root"),
|
|
2167
|
+
plane_codec=_text(map_value, "plane_codec"),
|
|
2168
|
+
payload_len=_u64(map_value, "payload_len"),
|
|
2169
|
+
)
|
|
2170
|
+
validate_activation_meta(meta, limits)
|
|
2171
|
+
return meta
|
|
2172
|
+
|
|
2173
|
+
|
|
2174
|
+
def validate_activation_meta(meta: ActivationMeta, limits: Limits) -> int:
|
|
2175
|
+
if meta.plane_codec != PLANE_CODEC:
|
|
2176
|
+
raise Malformed(
|
|
2177
|
+
"ACTIVATION_STREAM plane_codec `%s` is not implemented (only %s)"
|
|
2178
|
+
% (meta.plane_codec, PLANE_CODEC)
|
|
2179
|
+
)
|
|
2180
|
+
if meta.world == 0:
|
|
2181
|
+
raise Malformed("ACTIVATION_STREAM world must be at least 1")
|
|
2182
|
+
if meta.rank >= meta.world:
|
|
2183
|
+
raise Malformed(
|
|
2184
|
+
"ACTIVATION_STREAM rank %d must be below world %d"
|
|
2185
|
+
% (meta.rank, meta.world)
|
|
2186
|
+
)
|
|
2187
|
+
if meta.chunk_count == 0:
|
|
2188
|
+
raise Malformed("ACTIVATION_STREAM chunk_count must be at least 1")
|
|
2189
|
+
product = _shape_product(meta.shape, limits)
|
|
2190
|
+
byte_len = _checked_mul(
|
|
2191
|
+
product, KV_WORD_BYTES[meta.dtype], "ACTIVATION_STREAM chunk byte count"
|
|
2192
|
+
)
|
|
2193
|
+
_check_len(byte_len, limits.max_decoded_bytes, "ACTIVATION_STREAM chunk")
|
|
2194
|
+
if meta.chunk_index >= meta.chunk_count:
|
|
2195
|
+
raise Malformed(
|
|
2196
|
+
"ACTIVATION_STREAM chunk_index %d is not below chunk_count %d"
|
|
2197
|
+
% (meta.chunk_index, meta.chunk_count)
|
|
2198
|
+
)
|
|
2199
|
+
if meta.element_count != product:
|
|
2200
|
+
raise Malformed(
|
|
2201
|
+
"ACTIVATION_STREAM element_count %d does not match shape product %d"
|
|
2202
|
+
% (meta.element_count, product)
|
|
2203
|
+
)
|
|
2204
|
+
return byte_len
|
|
2205
|
+
|
|
2206
|
+
|
|
2207
|
+
def activation_metadata(frame, limits: Limits) -> ActivationMeta:
|
|
2208
|
+
descriptor_bytes, _payload = _split_prefix(frame.body)
|
|
2209
|
+
return parse_activation_descriptor(descriptor_bytes, limits)
|
|
2210
|
+
|
|
2211
|
+
|
|
2212
|
+
def decode_activation_stream(frame, limits: Limits) -> bytes:
|
|
2213
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
2214
|
+
meta = parse_activation_descriptor(descriptor_bytes, limits)
|
|
2215
|
+
byte_len = validate_activation_meta(meta, limits)
|
|
2216
|
+
if frame.header.decoded_len != byte_len:
|
|
2217
|
+
raise Corrupt(
|
|
2218
|
+
"ACTIVATION_STREAM decoded_len is %d but the descriptor requires %d"
|
|
2219
|
+
% (frame.header.decoded_len, byte_len)
|
|
2220
|
+
)
|
|
2221
|
+
if meta.payload_len != len(payload):
|
|
2222
|
+
raise Malformed(
|
|
2223
|
+
"ACTIVATION_STREAM payload is %d bytes but payload_len declares %d"
|
|
2224
|
+
% (len(payload), meta.payload_len)
|
|
2225
|
+
)
|
|
2226
|
+
out = _decode_byte_planes(
|
|
2227
|
+
payload, KV_WORD_BYTES[meta.dtype], byte_len, "ACTIVATION_STREAM"
|
|
2228
|
+
)
|
|
2229
|
+
if len(out) != byte_len:
|
|
2230
|
+
raise Corrupt(
|
|
2231
|
+
"ACTIVATION_STREAM reconstructed chunk length is inconsistent"
|
|
2232
|
+
)
|
|
2233
|
+
if values_hash(out) != meta.digest_root:
|
|
2234
|
+
raise Corrupt("ACTIVATION_STREAM digest_root mismatch")
|
|
2235
|
+
if hash_domain("weft/chunk/v1", out) != frame.header.content_hash:
|
|
2236
|
+
raise Corrupt(
|
|
2237
|
+
"ACTIVATION_STREAM content_hash mismatch against the reconstructed chunk"
|
|
2238
|
+
)
|
|
2239
|
+
return out
|
|
2240
|
+
|
|
2241
|
+
|
|
2242
|
+
# ---------------------------------------------------------------------------
|
|
2243
|
+
# Mode 15 -- TENSOR_TILED
|
|
2244
|
+
# ---------------------------------------------------------------------------
|
|
2245
|
+
|
|
2246
|
+
_TILED_KEYS = (
|
|
2247
|
+
"schema",
|
|
2248
|
+
"dtype",
|
|
2249
|
+
"shape",
|
|
2250
|
+
"element_count",
|
|
2251
|
+
"layout",
|
|
2252
|
+
"tile_shape",
|
|
2253
|
+
"tile_counts",
|
|
2254
|
+
"tile_codec",
|
|
2255
|
+
"tiles",
|
|
2256
|
+
"payload_len",
|
|
2257
|
+
"logical_root",
|
|
2258
|
+
)
|
|
2259
|
+
|
|
2260
|
+
_TILE_KEYS = (
|
|
2261
|
+
"index",
|
|
2262
|
+
"dims",
|
|
2263
|
+
"offset",
|
|
2264
|
+
"compressed_len",
|
|
2265
|
+
"element_len",
|
|
2266
|
+
"tile_root",
|
|
2267
|
+
)
|
|
2268
|
+
|
|
2269
|
+
|
|
2270
|
+
@dataclass
|
|
2271
|
+
class TileEntry:
|
|
2272
|
+
index: int
|
|
2273
|
+
dims: "list[int]"
|
|
2274
|
+
offset: int
|
|
2275
|
+
compressed_len: int
|
|
2276
|
+
element_len: int
|
|
2277
|
+
tile_root: bytes
|
|
2278
|
+
|
|
2279
|
+
|
|
2280
|
+
@dataclass
|
|
2281
|
+
class TiledMeta:
|
|
2282
|
+
dtype: str
|
|
2283
|
+
shape: "list[int]"
|
|
2284
|
+
element_count: int
|
|
2285
|
+
tile_shape: "list[int]"
|
|
2286
|
+
tile_counts: "list[int]"
|
|
2287
|
+
tiles: "list[TileEntry]"
|
|
2288
|
+
payload_len: int
|
|
2289
|
+
logical_root: bytes
|
|
2290
|
+
|
|
2291
|
+
|
|
2292
|
+
def _tensor_dtype(name: str) -> str:
|
|
2293
|
+
if name not in TENSOR_DTYPE_SIZES:
|
|
2294
|
+
raise Malformed(
|
|
2295
|
+
"unknown tensor dtype `%s` (v1 allows f16, bf16, f32, f64, i8, u8, "
|
|
2296
|
+
"i16, i32, i64)" % name
|
|
2297
|
+
)
|
|
2298
|
+
return name
|
|
2299
|
+
|
|
2300
|
+
|
|
2301
|
+
def _tile_counts(shape: "list[int]", tile_shape: "list[int]") -> "list[int]":
|
|
2302
|
+
if len(shape) != len(tile_shape):
|
|
2303
|
+
raise Malformed(
|
|
2304
|
+
"tensor tile rank %d does not equal shape rank %d"
|
|
2305
|
+
% (len(tile_shape), len(shape))
|
|
2306
|
+
)
|
|
2307
|
+
counts = []
|
|
2308
|
+
for axis, (extent, tile) in enumerate(zip(shape, tile_shape)):
|
|
2309
|
+
if tile == 0:
|
|
2310
|
+
raise Malformed(
|
|
2311
|
+
"tensor tile shape axis %d must be greater than zero" % axis
|
|
2312
|
+
)
|
|
2313
|
+
if extent % tile != 0:
|
|
2314
|
+
raise Malformed(
|
|
2315
|
+
"tensor tile shape axis %d extent %d does not divide shape "
|
|
2316
|
+
"extent %d" % (axis, tile, extent)
|
|
2317
|
+
)
|
|
2318
|
+
counts.append(extent // tile)
|
|
2319
|
+
return counts
|
|
2320
|
+
|
|
2321
|
+
|
|
2322
|
+
def parse_tiled_descriptor(
|
|
2323
|
+
descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
|
|
2324
|
+
) -> TiledMeta:
|
|
2325
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "tiled tensor")
|
|
2326
|
+
_exact_keys(map_value, _TILED_KEYS, "tiled tensor")
|
|
2327
|
+
schema = _text(map_value, "schema")
|
|
2328
|
+
if schema != TENSOR_TILED_SCHEMA:
|
|
2329
|
+
raise Malformed("unsupported tiled tensor schema `%s`" % schema)
|
|
2330
|
+
dtype = _tensor_dtype(_text(map_value, "dtype"))
|
|
2331
|
+
layout = _text(map_value, "layout")
|
|
2332
|
+
if layout != TILE_LAYOUT:
|
|
2333
|
+
raise Malformed(
|
|
2334
|
+
"tiled tensor layout `%s` is specified but not implemented in v1 "
|
|
2335
|
+
"(only %s)" % (layout, TILE_LAYOUT)
|
|
2336
|
+
)
|
|
2337
|
+
tile_codec = _text(map_value, "tile_codec")
|
|
2338
|
+
if tile_codec != PLANE_CODEC:
|
|
2339
|
+
raise Malformed(
|
|
2340
|
+
"unsupported tensor tile codec `%s` (v1 allows `%s`)"
|
|
2341
|
+
% (tile_codec, PLANE_CODEC)
|
|
2342
|
+
)
|
|
2343
|
+
shape = _shape(map_value)
|
|
2344
|
+
element_count = _validate_shape(shape, limits, "tiled tensor")
|
|
2345
|
+
declared_count = _u64(map_value, "element_count")
|
|
2346
|
+
if declared_count != element_count:
|
|
2347
|
+
raise Malformed(
|
|
2348
|
+
"tiled tensor element_count %d does not match shape product %d"
|
|
2349
|
+
% (declared_count, element_count)
|
|
2350
|
+
)
|
|
2351
|
+
element_size = TENSOR_DTYPE_SIZES[dtype]
|
|
2352
|
+
decoded_bytes = _checked_mul(
|
|
2353
|
+
element_count, element_size, "tiled tensor decoded bytes"
|
|
2354
|
+
)
|
|
2355
|
+
_check_len(decoded_bytes, limits.max_decoded_bytes, "tiled tensor decoded bytes")
|
|
2356
|
+
if decoded_len != decoded_bytes:
|
|
2357
|
+
raise Corrupt(
|
|
2358
|
+
"tiled tensor decoded_len is %d but the descriptor requires %d"
|
|
2359
|
+
% (decoded_len, decoded_bytes)
|
|
2360
|
+
)
|
|
2361
|
+
tile_shape = _array_u64(map_value, "tile_shape")
|
|
2362
|
+
counts = _tile_counts(shape, tile_shape)
|
|
2363
|
+
declared_counts = _array_u64(map_value, "tile_counts")
|
|
2364
|
+
if declared_counts != counts:
|
|
2365
|
+
raise Malformed(
|
|
2366
|
+
"tiled tensor tile_counts %r do not match shape/tile_shape %r"
|
|
2367
|
+
% (declared_counts, counts)
|
|
2368
|
+
)
|
|
2369
|
+
tile_elements = _validate_shape(tile_shape, limits, "tensor tile")
|
|
2370
|
+
tile_element_len = _checked_mul(
|
|
2371
|
+
tile_elements, element_size, "tensor tile byte length"
|
|
2372
|
+
)
|
|
2373
|
+
tile_count = 1
|
|
2374
|
+
for count in counts:
|
|
2375
|
+
tile_count = _checked_mul(tile_count, count, "tensor tile count")
|
|
2376
|
+
if tile_count > limits.max_structured_items:
|
|
2377
|
+
raise LimitExceeded(
|
|
2378
|
+
"tensor tile count %d exceeds limit %d"
|
|
2379
|
+
% (tile_count, limits.max_structured_items)
|
|
2380
|
+
)
|
|
2381
|
+
declared_payload_len = _u64(map_value, "payload_len")
|
|
2382
|
+
if declared_payload_len != payload_len:
|
|
2383
|
+
raise Malformed(
|
|
2384
|
+
"tiled tensor payload is %d bytes but metadata declares %d"
|
|
2385
|
+
% (payload_len, declared_payload_len)
|
|
2386
|
+
)
|
|
2387
|
+
entries = map_value.get("tiles")
|
|
2388
|
+
if not isinstance(entries, list):
|
|
2389
|
+
raise Malformed("tiled tensor `tiles` must be an array")
|
|
2390
|
+
if len(entries) != tile_count:
|
|
2391
|
+
raise Malformed(
|
|
2392
|
+
"tiled tensor has %d tile entries but the grid requires %d"
|
|
2393
|
+
% (len(entries), tile_count)
|
|
2394
|
+
)
|
|
2395
|
+
tiles = []
|
|
2396
|
+
expected_offset = 0
|
|
2397
|
+
for position, item in enumerate(entries):
|
|
2398
|
+
entry = _map(item, "tensor tile")
|
|
2399
|
+
_exact_keys(entry, _TILE_KEYS, "tensor tile")
|
|
2400
|
+
index = _u64(entry, "index")
|
|
2401
|
+
if index != position:
|
|
2402
|
+
raise Malformed(
|
|
2403
|
+
"tiled tensor entry %d declares index %d; entries must be in "
|
|
2404
|
+
"canonical order" % (position, index)
|
|
2405
|
+
)
|
|
2406
|
+
dims = _array_u64(entry, "dims")
|
|
2407
|
+
if dims != tile_shape:
|
|
2408
|
+
raise Malformed(
|
|
2409
|
+
"tensor tile %d dims %r do not equal tile_shape %r; partial "
|
|
2410
|
+
"tiles are not defined in v1" % (position, dims, tile_shape)
|
|
2411
|
+
)
|
|
2412
|
+
declared_element_len = _u64(entry, "element_len")
|
|
2413
|
+
if declared_element_len != tile_element_len:
|
|
2414
|
+
raise Malformed(
|
|
2415
|
+
"tensor tile %d element_len %d does not match %d"
|
|
2416
|
+
% (position, declared_element_len, tile_element_len)
|
|
2417
|
+
)
|
|
2418
|
+
offset = _u64(entry, "offset")
|
|
2419
|
+
if offset != expected_offset:
|
|
2420
|
+
raise Malformed(
|
|
2421
|
+
"tensor tile %d offset %d is not the contiguous offset %d"
|
|
2422
|
+
% (position, offset, expected_offset)
|
|
2423
|
+
)
|
|
2424
|
+
compressed_len = _u64(entry, "compressed_len")
|
|
2425
|
+
if compressed_len == 0:
|
|
2426
|
+
raise Malformed(
|
|
2427
|
+
"tensor tile %d has an empty stored range" % position
|
|
2428
|
+
)
|
|
2429
|
+
expected_offset = _checked_add(
|
|
2430
|
+
expected_offset, compressed_len, "tensor tile offset"
|
|
2431
|
+
)
|
|
2432
|
+
root = _bytes_of(entry, "tile_root")
|
|
2433
|
+
if len(root) != HASH_LEN:
|
|
2434
|
+
raise Malformed("tensor tile_root must be 32 bytes")
|
|
2435
|
+
tiles.append(
|
|
2436
|
+
TileEntry(
|
|
2437
|
+
index=index,
|
|
2438
|
+
dims=dims,
|
|
2439
|
+
offset=offset,
|
|
2440
|
+
compressed_len=compressed_len,
|
|
2441
|
+
element_len=declared_element_len,
|
|
2442
|
+
tile_root=root,
|
|
2443
|
+
)
|
|
2444
|
+
)
|
|
2445
|
+
if expected_offset != declared_payload_len:
|
|
2446
|
+
raise Malformed(
|
|
2447
|
+
"tensor tile ranges cover %d bytes but payload_len is %d"
|
|
2448
|
+
% (expected_offset, declared_payload_len)
|
|
2449
|
+
)
|
|
2450
|
+
logical_root = _bytes_of(map_value, "logical_root")
|
|
2451
|
+
if len(logical_root) != HASH_LEN:
|
|
2452
|
+
raise Malformed("tiled tensor logical_root must be 32 bytes")
|
|
2453
|
+
return TiledMeta(
|
|
2454
|
+
dtype=dtype,
|
|
2455
|
+
shape=shape,
|
|
2456
|
+
element_count=element_count,
|
|
2457
|
+
tile_shape=tile_shape,
|
|
2458
|
+
tile_counts=counts,
|
|
2459
|
+
tiles=tiles,
|
|
2460
|
+
payload_len=declared_payload_len,
|
|
2461
|
+
logical_root=logical_root,
|
|
2462
|
+
)
|
|
2463
|
+
|
|
2464
|
+
|
|
2465
|
+
def _tile_start(counts: "list[int]", index: int) -> "list[int]":
|
|
2466
|
+
coords = [0] * len(counts)
|
|
2467
|
+
remainder = index
|
|
2468
|
+
for axis in range(len(counts) - 1, -1, -1):
|
|
2469
|
+
coords[axis] = remainder % counts[axis]
|
|
2470
|
+
remainder //= counts[axis]
|
|
2471
|
+
return coords
|
|
2472
|
+
|
|
2473
|
+
|
|
2474
|
+
def _decode_tile_planes(
|
|
2475
|
+
stored: bytes, element_size: int, element_len: int
|
|
2476
|
+
) -> bytes:
|
|
2477
|
+
if element_size == 0 or element_len % element_size != 0:
|
|
2478
|
+
raise Malformed(
|
|
2479
|
+
"tensor tile element length is not a whole number of elements"
|
|
2480
|
+
)
|
|
2481
|
+
elements = element_len // element_size
|
|
2482
|
+
planes = []
|
|
2483
|
+
cursor = 0
|
|
2484
|
+
for _ in range(element_size):
|
|
2485
|
+
if cursor > len(stored):
|
|
2486
|
+
raise Corrupt("tensor tile plane is truncated")
|
|
2487
|
+
plane, consumed = huffman.decode_prefix(stored[cursor:], elements)
|
|
2488
|
+
cursor += consumed
|
|
2489
|
+
planes.append(plane)
|
|
2490
|
+
if cursor != len(stored):
|
|
2491
|
+
raise Corrupt(
|
|
2492
|
+
"tensor tile planes leave trailing bytes inside the tile range"
|
|
2493
|
+
)
|
|
2494
|
+
out = bytearray(element_len)
|
|
2495
|
+
for plane_index, plane in enumerate(planes):
|
|
2496
|
+
for element, byte in enumerate(plane):
|
|
2497
|
+
out[element * element_size + plane_index] = byte
|
|
2498
|
+
return bytes(out)
|
|
2499
|
+
|
|
2500
|
+
|
|
2501
|
+
def _iter_coords(extent: "list[int]"):
|
|
2502
|
+
if not extent:
|
|
2503
|
+
yield []
|
|
2504
|
+
return
|
|
2505
|
+
def walk(dimension: int, prefix: "list[int]"):
|
|
2506
|
+
if dimension == len(extent):
|
|
2507
|
+
yield list(prefix)
|
|
2508
|
+
return
|
|
2509
|
+
for position in range(extent[dimension]):
|
|
2510
|
+
prefix.append(position)
|
|
2511
|
+
yield from walk(dimension + 1, prefix)
|
|
2512
|
+
prefix.pop()
|
|
2513
|
+
|
|
2514
|
+
yield from walk(0, [])
|
|
2515
|
+
|
|
2516
|
+
|
|
2517
|
+
def _flat_index(coords: "list[int]", shape: "list[int]") -> int:
|
|
2518
|
+
index = 0
|
|
2519
|
+
for coord, dim in zip(coords, shape):
|
|
2520
|
+
index = index * dim + coord
|
|
2521
|
+
return index
|
|
2522
|
+
|
|
2523
|
+
|
|
2524
|
+
def decode_tensor_tiled(frame, limits: Limits) -> bytes:
|
|
2525
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
2526
|
+
meta = parse_tiled_descriptor(
|
|
2527
|
+
descriptor_bytes, frame.header.decoded_len, len(payload), limits
|
|
2528
|
+
)
|
|
2529
|
+
element_size = TENSOR_DTYPE_SIZES[meta.dtype]
|
|
2530
|
+
total = meta.element_count * element_size
|
|
2531
|
+
out = bytearray(total)
|
|
2532
|
+
tiles_bytes = 0
|
|
2533
|
+
for entry in meta.tiles:
|
|
2534
|
+
start = entry.offset
|
|
2535
|
+
end = start + entry.compressed_len
|
|
2536
|
+
if end > len(payload):
|
|
2537
|
+
raise Corrupt("tensor tile range is outside the payload")
|
|
2538
|
+
stored = payload[start:end]
|
|
2539
|
+
if len(stored) != entry.compressed_len:
|
|
2540
|
+
raise Corrupt("tensor tile stored range length is inconsistent")
|
|
2541
|
+
tile = _decode_tile_planes(stored, element_size, entry.element_len)
|
|
2542
|
+
root = tile_root(meta.dtype, entry.dims, tile, limits)
|
|
2543
|
+
if root != entry.tile_root:
|
|
2544
|
+
raise Corrupt(
|
|
2545
|
+
"tensor tile %d root does not match the descriptor" % entry.index
|
|
2546
|
+
)
|
|
2547
|
+
tiles_bytes += len(tile)
|
|
2548
|
+
tile_coords = _tile_start(meta.tile_counts, entry.index)
|
|
2549
|
+
dst_start = [
|
|
2550
|
+
coord * meta.tile_shape[axis]
|
|
2551
|
+
for axis, coord in enumerate(tile_coords)
|
|
2552
|
+
]
|
|
2553
|
+
for local in _iter_coords(entry.dims):
|
|
2554
|
+
src = _flat_index(local, entry.dims) * element_size
|
|
2555
|
+
full = [
|
|
2556
|
+
dst_start[axis] + local[axis]
|
|
2557
|
+
for axis in range(len(meta.shape))
|
|
2558
|
+
]
|
|
2559
|
+
dst = _flat_index(full, meta.shape) * element_size
|
|
2560
|
+
out[dst : dst + element_size] = tile[src : src + element_size]
|
|
2561
|
+
if tiles_bytes != total:
|
|
2562
|
+
raise Corrupt(
|
|
2563
|
+
"tiled tensor decoded to %d bytes but the descriptor requires %d"
|
|
2564
|
+
% (tiles_bytes, total)
|
|
2565
|
+
)
|
|
2566
|
+
logical = tensor_logical_root(
|
|
2567
|
+
meta.dtype, meta.shape, meta.element_count, bytes(out), limits
|
|
2568
|
+
)
|
|
2569
|
+
if logical != meta.logical_root:
|
|
2570
|
+
raise Corrupt(
|
|
2571
|
+
"tiled tensor logical root does not match the descriptor"
|
|
2572
|
+
)
|
|
2573
|
+
return bytes(out)
|
|
2574
|
+
|
|
2575
|
+
|
|
2576
|
+
# ---------------------------------------------------------------------------
|
|
2577
|
+
# Mode 17 -- FLOAT_TILED
|
|
2578
|
+
# ---------------------------------------------------------------------------
|
|
2579
|
+
|
|
2580
|
+
FLOAT_TILED_SCHEMA = "weft.ai.float_tiled.v1"
|
|
2581
|
+
FLOAT_TILED_DTYPES = ("bf16", "fp16", "fp32")
|
|
2582
|
+
# The descriptor uses the AI-codec spelling; the WEFT-TENSOR identity uses `f16`.
|
|
2583
|
+
FLOAT_TILED_LOGICAL_DTYPE = {"bf16": "bf16", "fp16": "f16", "fp32": "f32"}
|
|
2584
|
+
|
|
2585
|
+
_FLOAT_TILED_KEYS = (
|
|
2586
|
+
"schema",
|
|
2587
|
+
"dtype",
|
|
2588
|
+
"shape",
|
|
2589
|
+
"element_count",
|
|
2590
|
+
"layout",
|
|
2591
|
+
"tile_shape",
|
|
2592
|
+
"tile_counts",
|
|
2593
|
+
"exponent_coder",
|
|
2594
|
+
"tiles",
|
|
2595
|
+
"payload_len",
|
|
2596
|
+
"logical_root",
|
|
2597
|
+
)
|
|
2598
|
+
|
|
2599
|
+
_FLOAT_TILE_KEYS = (
|
|
2600
|
+
"index",
|
|
2601
|
+
"dims",
|
|
2602
|
+
"offset",
|
|
2603
|
+
"compressed_len",
|
|
2604
|
+
"exponent_plane_len",
|
|
2605
|
+
"mantissa_plane_len",
|
|
2606
|
+
"element_len",
|
|
2607
|
+
"tile_root",
|
|
2608
|
+
)
|
|
2609
|
+
|
|
2610
|
+
|
|
2611
|
+
@dataclass
|
|
2612
|
+
class FloatTileEntry:
|
|
2613
|
+
index: int
|
|
2614
|
+
dims: "list[int]"
|
|
2615
|
+
offset: int
|
|
2616
|
+
compressed_len: int
|
|
2617
|
+
exponent_plane_len: int
|
|
2618
|
+
mantissa_plane_len: int
|
|
2619
|
+
element_len: int
|
|
2620
|
+
tile_root: bytes
|
|
2621
|
+
|
|
2622
|
+
|
|
2623
|
+
@dataclass
|
|
2624
|
+
class FloatTiledMeta:
|
|
2625
|
+
dtype: str
|
|
2626
|
+
shape: "list[int]"
|
|
2627
|
+
element_count: int
|
|
2628
|
+
tile_shape: "list[int]"
|
|
2629
|
+
tile_counts: "list[int]"
|
|
2630
|
+
tiles: "list[FloatTileEntry]"
|
|
2631
|
+
payload_len: int
|
|
2632
|
+
logical_root: bytes
|
|
2633
|
+
|
|
2634
|
+
|
|
2635
|
+
def float_tile_root(
|
|
2636
|
+
dtype: str,
|
|
2637
|
+
dims: "list[int]",
|
|
2638
|
+
element_len: int,
|
|
2639
|
+
exponent_plane_len: int,
|
|
2640
|
+
mantissa_plane_len: int,
|
|
2641
|
+
tile_bytes: bytes,
|
|
2642
|
+
limits: Limits,
|
|
2643
|
+
) -> bytes:
|
|
2644
|
+
record = {
|
|
2645
|
+
"dtype": dtype,
|
|
2646
|
+
"dims": _u64_vec(dims),
|
|
2647
|
+
"element_len": wcb.U64(element_len),
|
|
2648
|
+
"exponent_plane_len": wcb.U64(exponent_plane_len),
|
|
2649
|
+
"mantissa_plane_len": wcb.U64(mantissa_plane_len),
|
|
2650
|
+
"values_hash": values_hash(tile_bytes),
|
|
2651
|
+
}
|
|
2652
|
+
return hash_domain(DOMAIN_TENSOR_FLOAT_TILE, wcb.encode(record, limits))
|
|
2653
|
+
|
|
2654
|
+
|
|
2655
|
+
def parse_float_tiled_descriptor(
|
|
2656
|
+
descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
|
|
2657
|
+
) -> FloatTiledMeta:
|
|
2658
|
+
map_value = _map(_decode_descriptor(descriptor_bytes, limits), "float tiled tensor")
|
|
2659
|
+
_exact_keys(map_value, _FLOAT_TILED_KEYS, "float tiled tensor")
|
|
2660
|
+
schema = _text(map_value, "schema")
|
|
2661
|
+
if schema != FLOAT_TILED_SCHEMA:
|
|
2662
|
+
raise Malformed("unsupported float tiled tensor schema `%s`" % schema)
|
|
2663
|
+
dtype = _text(map_value, "dtype")
|
|
2664
|
+
if dtype not in FLOAT_FIELD_BITS:
|
|
2665
|
+
raise Malformed(
|
|
2666
|
+
"unknown float tiled tensor dtype `%s` (v1 allows bf16, fp16, "
|
|
2667
|
+
"fp32, fp8-e4m3, fp8-e5m2)" % dtype
|
|
2668
|
+
)
|
|
2669
|
+
if dtype not in FLOAT_TILED_DTYPES:
|
|
2670
|
+
raise UnsupportedFeature(
|
|
2671
|
+
"float tiled tensor dtype `%s` is not implemented in v1" % dtype
|
|
2672
|
+
)
|
|
2673
|
+
layout = _text(map_value, "layout")
|
|
2674
|
+
if layout != TILE_LAYOUT:
|
|
2675
|
+
raise Malformed(
|
|
2676
|
+
"float tiled tensor layout `%s` is specified but not implemented "
|
|
2677
|
+
"in v1 (only %s)" % (layout, TILE_LAYOUT)
|
|
2678
|
+
)
|
|
2679
|
+
coder = _text(map_value, "exponent_coder")
|
|
2680
|
+
if coder != PLANE_CODEC:
|
|
2681
|
+
raise Malformed(
|
|
2682
|
+
"unsupported float tiled tensor exponent coder `%s` (v1 allows `%s`)"
|
|
2683
|
+
% (coder, PLANE_CODEC)
|
|
2684
|
+
)
|
|
2685
|
+
shape = _shape(map_value)
|
|
2686
|
+
element_count = _validate_shape(shape, limits, "float tiled tensor")
|
|
2687
|
+
declared_count = _u64(map_value, "element_count")
|
|
2688
|
+
if declared_count != element_count:
|
|
2689
|
+
raise Malformed(
|
|
2690
|
+
"float tiled tensor element_count %d does not match shape product %d"
|
|
2691
|
+
% (declared_count, element_count)
|
|
2692
|
+
)
|
|
2693
|
+
sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
|
|
2694
|
+
element_bits = sign_bits + exponent_bits + mantissa_bits
|
|
2695
|
+
if element_bits % 8 != 0:
|
|
2696
|
+
raise Malformed("float tiled tensor element width is not byte-aligned")
|
|
2697
|
+
element_size = element_bits // 8
|
|
2698
|
+
decoded_bytes = _checked_mul(
|
|
2699
|
+
element_count, element_size, "float tiled tensor decoded bytes"
|
|
2700
|
+
)
|
|
2701
|
+
_check_len(decoded_bytes, limits.max_decoded_bytes, "float tiled tensor decoded bytes")
|
|
2702
|
+
if decoded_len != decoded_bytes:
|
|
2703
|
+
raise Corrupt(
|
|
2704
|
+
"float tiled tensor decoded_len is %d but the descriptor requires %d"
|
|
2705
|
+
% (decoded_len, decoded_bytes)
|
|
2706
|
+
)
|
|
2707
|
+
tile_shape = _array_u64(map_value, "tile_shape")
|
|
2708
|
+
counts = _tile_counts(shape, tile_shape)
|
|
2709
|
+
declared_counts = _array_u64(map_value, "tile_counts")
|
|
2710
|
+
if declared_counts != counts:
|
|
2711
|
+
raise Malformed(
|
|
2712
|
+
"float tiled tensor tile_counts %r do not match shape/tile_shape %r"
|
|
2713
|
+
% (declared_counts, counts)
|
|
2714
|
+
)
|
|
2715
|
+
tile_elements = _validate_shape(tile_shape, limits, "float tensor tile")
|
|
2716
|
+
tile_element_len = _checked_mul(
|
|
2717
|
+
tile_elements, element_size, "float tensor tile byte length"
|
|
2718
|
+
)
|
|
2719
|
+
tile_count = 1
|
|
2720
|
+
for count in counts:
|
|
2721
|
+
tile_count = _checked_mul(tile_count, count, "float tensor tile count")
|
|
2722
|
+
if tile_count > limits.max_structured_items:
|
|
2723
|
+
raise LimitExceeded(
|
|
2724
|
+
"float tensor tile count %d exceeds limit %d"
|
|
2725
|
+
% (tile_count, limits.max_structured_items)
|
|
2726
|
+
)
|
|
2727
|
+
declared_payload_len = _u64(map_value, "payload_len")
|
|
2728
|
+
if declared_payload_len != payload_len:
|
|
2729
|
+
# The frame body is the authority; a declared length that disagrees is
|
|
2730
|
+
# a truncated or over-long payload, which is corruption.
|
|
2731
|
+
raise Corrupt(
|
|
2732
|
+
"float tiled tensor payload is %d bytes but metadata declares %d"
|
|
2733
|
+
% (payload_len, declared_payload_len)
|
|
2734
|
+
)
|
|
2735
|
+
mantissa_total_bits = _checked_mul(
|
|
2736
|
+
tile_elements, sign_bits + mantissa_bits, "float tensor tile mantissa bits"
|
|
2737
|
+
)
|
|
2738
|
+
expected_tile_mantissa = _checked_add(
|
|
2739
|
+
mantissa_total_bits, 7, "float tensor tile mantissa length"
|
|
2740
|
+
) // 8
|
|
2741
|
+
entries = map_value.get("tiles")
|
|
2742
|
+
if not isinstance(entries, list):
|
|
2743
|
+
raise Malformed("float tiled tensor `tiles` must be an array")
|
|
2744
|
+
if len(entries) != tile_count:
|
|
2745
|
+
raise Malformed(
|
|
2746
|
+
"float tiled tensor has %d tile entries but the grid requires %d"
|
|
2747
|
+
% (len(entries), tile_count)
|
|
2748
|
+
)
|
|
2749
|
+
tiles = []
|
|
2750
|
+
expected_offset = 0
|
|
2751
|
+
for position, item in enumerate(entries):
|
|
2752
|
+
entry = _map(item, "float tensor tile")
|
|
2753
|
+
_exact_keys(entry, _FLOAT_TILE_KEYS, "float tensor tile")
|
|
2754
|
+
index = _u64(entry, "index")
|
|
2755
|
+
if index != position:
|
|
2756
|
+
raise Malformed(
|
|
2757
|
+
"float tiled tensor entry %d declares index %d; entries must be "
|
|
2758
|
+
"in canonical order" % (position, index)
|
|
2759
|
+
)
|
|
2760
|
+
if index >= tile_count:
|
|
2761
|
+
raise Malformed(
|
|
2762
|
+
"float tiled tensor tile index %d is outside the %d-tile grid"
|
|
2763
|
+
% (index, tile_count)
|
|
2764
|
+
)
|
|
2765
|
+
dims = _array_u64(entry, "dims")
|
|
2766
|
+
if dims != tile_shape:
|
|
2767
|
+
raise Malformed(
|
|
2768
|
+
"float tensor tile %d dims %r do not equal tile_shape %r; "
|
|
2769
|
+
"partial tiles are not defined in v1" % (position, dims, tile_shape)
|
|
2770
|
+
)
|
|
2771
|
+
declared_element_len = _u64(entry, "element_len")
|
|
2772
|
+
if declared_element_len != tile_element_len:
|
|
2773
|
+
raise Malformed(
|
|
2774
|
+
"float tensor tile %d element_len %d does not match %d"
|
|
2775
|
+
% (position, declared_element_len, tile_element_len)
|
|
2776
|
+
)
|
|
2777
|
+
offset = _u64(entry, "offset")
|
|
2778
|
+
if offset != expected_offset:
|
|
2779
|
+
raise Malformed(
|
|
2780
|
+
"float tensor tile %d offset %d is not the contiguous offset %d"
|
|
2781
|
+
% (position, offset, expected_offset)
|
|
2782
|
+
)
|
|
2783
|
+
compressed_len = _u64(entry, "compressed_len")
|
|
2784
|
+
if compressed_len == 0:
|
|
2785
|
+
raise Malformed("float tensor tile %d has an empty stored range" % position)
|
|
2786
|
+
exponent_plane_len = _u64(entry, "exponent_plane_len")
|
|
2787
|
+
if exponent_plane_len < 256:
|
|
2788
|
+
raise Malformed(
|
|
2789
|
+
"float tensor tile %d exponent_plane_len %d is below the "
|
|
2790
|
+
"Huffman table minimum" % (position, exponent_plane_len)
|
|
2791
|
+
)
|
|
2792
|
+
mantissa_plane_len = _u64(entry, "mantissa_plane_len")
|
|
2793
|
+
if mantissa_plane_len != expected_tile_mantissa:
|
|
2794
|
+
raise Malformed(
|
|
2795
|
+
"float tensor tile %d mantissa_plane_len %d does not match the "
|
|
2796
|
+
"tile arithmetic %d"
|
|
2797
|
+
% (position, mantissa_plane_len, expected_tile_mantissa)
|
|
2798
|
+
)
|
|
2799
|
+
if exponent_plane_len + mantissa_plane_len != compressed_len:
|
|
2800
|
+
raise Malformed(
|
|
2801
|
+
"float tensor tile %d plane lengths %d + %d do not sum to "
|
|
2802
|
+
"compressed_len %d"
|
|
2803
|
+
% (
|
|
2804
|
+
position,
|
|
2805
|
+
exponent_plane_len,
|
|
2806
|
+
mantissa_plane_len,
|
|
2807
|
+
compressed_len,
|
|
2808
|
+
)
|
|
2809
|
+
)
|
|
2810
|
+
expected_offset = _checked_add(
|
|
2811
|
+
expected_offset, compressed_len, "float tensor tile offset"
|
|
2812
|
+
)
|
|
2813
|
+
root = _bytes_of(entry, "tile_root")
|
|
2814
|
+
if len(root) != HASH_LEN:
|
|
2815
|
+
raise Malformed("float tensor tile_root must be 32 bytes")
|
|
2816
|
+
tiles.append(
|
|
2817
|
+
FloatTileEntry(
|
|
2818
|
+
index=index,
|
|
2819
|
+
dims=dims,
|
|
2820
|
+
offset=offset,
|
|
2821
|
+
compressed_len=compressed_len,
|
|
2822
|
+
exponent_plane_len=exponent_plane_len,
|
|
2823
|
+
mantissa_plane_len=mantissa_plane_len,
|
|
2824
|
+
element_len=declared_element_len,
|
|
2825
|
+
tile_root=root,
|
|
2826
|
+
)
|
|
2827
|
+
)
|
|
2828
|
+
if expected_offset != declared_payload_len:
|
|
2829
|
+
raise Corrupt(
|
|
2830
|
+
"float tensor tile ranges cover %d bytes but payload_len is %d"
|
|
2831
|
+
% (expected_offset, declared_payload_len)
|
|
2832
|
+
)
|
|
2833
|
+
logical_root = _bytes_of(map_value, "logical_root")
|
|
2834
|
+
if len(logical_root) != HASH_LEN:
|
|
2835
|
+
raise Malformed("float tiled tensor logical_root must be 32 bytes")
|
|
2836
|
+
return FloatTiledMeta(
|
|
2837
|
+
dtype=dtype,
|
|
2838
|
+
shape=shape,
|
|
2839
|
+
element_count=element_count,
|
|
2840
|
+
tile_shape=tile_shape,
|
|
2841
|
+
tile_counts=counts,
|
|
2842
|
+
tiles=tiles,
|
|
2843
|
+
payload_len=declared_payload_len,
|
|
2844
|
+
logical_root=logical_root,
|
|
2845
|
+
)
|
|
2846
|
+
|
|
2847
|
+
|
|
2848
|
+
def decode_float_tiled(frame, limits: Limits) -> bytes:
|
|
2849
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
2850
|
+
meta = parse_float_tiled_descriptor(
|
|
2851
|
+
descriptor_bytes, frame.header.decoded_len, len(payload), limits
|
|
2852
|
+
)
|
|
2853
|
+
sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[meta.dtype]
|
|
2854
|
+
element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
|
|
2855
|
+
total = meta.element_count * element_size
|
|
2856
|
+
out = bytearray(total)
|
|
2857
|
+
tiles_bytes = 0
|
|
2858
|
+
for entry in meta.tiles:
|
|
2859
|
+
start = entry.offset
|
|
2860
|
+
end = start + entry.compressed_len
|
|
2861
|
+
if end > len(payload):
|
|
2862
|
+
raise Corrupt("float tensor tile range is outside the payload")
|
|
2863
|
+
stored = payload[start:end]
|
|
2864
|
+
exponent_compressed = stored[: entry.exponent_plane_len]
|
|
2865
|
+
mantissa_plane = stored[entry.exponent_plane_len :]
|
|
2866
|
+
elements = entry.element_len // element_size
|
|
2867
|
+
exponent_plane = huffman.decode(exponent_compressed, elements)
|
|
2868
|
+
tile = _join_float_planes(
|
|
2869
|
+
exponent_plane,
|
|
2870
|
+
mantissa_plane,
|
|
2871
|
+
meta.dtype,
|
|
2872
|
+
elements,
|
|
2873
|
+
limits,
|
|
2874
|
+
)
|
|
2875
|
+
root = float_tile_root(
|
|
2876
|
+
meta.dtype,
|
|
2877
|
+
entry.dims,
|
|
2878
|
+
entry.element_len,
|
|
2879
|
+
entry.exponent_plane_len,
|
|
2880
|
+
entry.mantissa_plane_len,
|
|
2881
|
+
tile,
|
|
2882
|
+
limits,
|
|
2883
|
+
)
|
|
2884
|
+
if root != entry.tile_root:
|
|
2885
|
+
raise Corrupt(
|
|
2886
|
+
"float tensor tile %d root does not match the descriptor"
|
|
2887
|
+
% entry.index
|
|
2888
|
+
)
|
|
2889
|
+
tiles_bytes += len(tile)
|
|
2890
|
+
tile_coords = _tile_start(meta.tile_counts, entry.index)
|
|
2891
|
+
dst_start = [
|
|
2892
|
+
coord * meta.tile_shape[axis] for axis, coord in enumerate(tile_coords)
|
|
2893
|
+
]
|
|
2894
|
+
for local in _iter_coords(entry.dims):
|
|
2895
|
+
src = _flat_index(local, entry.dims) * element_size
|
|
2896
|
+
full = [dst_start[axis] + local[axis] for axis in range(len(meta.shape))]
|
|
2897
|
+
dst = _flat_index(full, meta.shape) * element_size
|
|
2898
|
+
out[dst : dst + element_size] = tile[src : src + element_size]
|
|
2899
|
+
if tiles_bytes != total:
|
|
2900
|
+
raise Corrupt(
|
|
2901
|
+
"float tiled tensor decoded to %d bytes but the descriptor "
|
|
2902
|
+
"requires %d" % (tiles_bytes, total)
|
|
2903
|
+
)
|
|
2904
|
+
logical = tensor_logical_root(
|
|
2905
|
+
FLOAT_TILED_LOGICAL_DTYPE[meta.dtype],
|
|
2906
|
+
meta.shape,
|
|
2907
|
+
meta.element_count,
|
|
2908
|
+
bytes(out),
|
|
2909
|
+
limits,
|
|
2910
|
+
)
|
|
2911
|
+
if logical != meta.logical_root:
|
|
2912
|
+
raise Corrupt(
|
|
2913
|
+
"float tiled tensor logical root does not match the descriptor"
|
|
2914
|
+
)
|
|
2915
|
+
return bytes(out)
|
|
2916
|
+
|
|
2917
|
+
|
|
2918
|
+
# ---------------------------------------------------------------------------
|
|
2919
|
+
# Mode 18 -- FLOAT_TILED_STREAMS
|
|
2920
|
+
# ---------------------------------------------------------------------------
|
|
2921
|
+
|
|
2922
|
+
FLOAT_TILED_STREAMS_SCHEMA = "weft.ai.tensor_tiled_streams.v1"
|
|
2923
|
+
FLOAT_TILED_STREAMS_MIN = 1
|
|
2924
|
+
FLOAT_TILED_STREAMS_MAX = 16
|
|
2925
|
+
|
|
2926
|
+
_FLOAT_TILED_STREAMS_KEYS = (
|
|
2927
|
+
"schema",
|
|
2928
|
+
"dtype",
|
|
2929
|
+
"shape",
|
|
2930
|
+
"element_count",
|
|
2931
|
+
"layout",
|
|
2932
|
+
"tile_shape",
|
|
2933
|
+
"tile_counts",
|
|
2934
|
+
"exponent_coder",
|
|
2935
|
+
"n_streams",
|
|
2936
|
+
"tiles",
|
|
2937
|
+
"payload_len",
|
|
2938
|
+
"logical_root",
|
|
2939
|
+
)
|
|
2940
|
+
|
|
2941
|
+
_FLOAT_TILE_STREAMS_KEYS = (
|
|
2942
|
+
"index",
|
|
2943
|
+
"dims",
|
|
2944
|
+
"offset",
|
|
2945
|
+
"compressed_len",
|
|
2946
|
+
"exponent_plane_len",
|
|
2947
|
+
"mantissa_plane_len",
|
|
2948
|
+
"element_len",
|
|
2949
|
+
"stream_offsets",
|
|
2950
|
+
"tile_root",
|
|
2951
|
+
)
|
|
2952
|
+
|
|
2953
|
+
|
|
2954
|
+
@dataclass
|
|
2955
|
+
class FloatTileStreamsEntry:
|
|
2956
|
+
index: int
|
|
2957
|
+
dims: "list[int]"
|
|
2958
|
+
offset: int
|
|
2959
|
+
compressed_len: int
|
|
2960
|
+
exponent_plane_len: int
|
|
2961
|
+
mantissa_plane_len: int
|
|
2962
|
+
element_len: int
|
|
2963
|
+
stream_offsets: "list[int]"
|
|
2964
|
+
tile_root: bytes
|
|
2965
|
+
|
|
2966
|
+
|
|
2967
|
+
@dataclass
|
|
2968
|
+
class FloatTiledStreamsMeta:
|
|
2969
|
+
dtype: str
|
|
2970
|
+
shape: "list[int]"
|
|
2971
|
+
element_count: int
|
|
2972
|
+
tile_shape: "list[int]"
|
|
2973
|
+
tile_counts: "list[int]"
|
|
2974
|
+
n_streams: int
|
|
2975
|
+
tiles: "list[FloatTileStreamsEntry]"
|
|
2976
|
+
payload_len: int
|
|
2977
|
+
logical_root: bytes
|
|
2978
|
+
|
|
2979
|
+
|
|
2980
|
+
def interleaved_stream_len(total: int, stream: int, n_streams: int) -> int:
|
|
2981
|
+
if n_streams <= 0 or total <= stream:
|
|
2982
|
+
return 0
|
|
2983
|
+
return (total - stream + n_streams - 1) // n_streams
|
|
2984
|
+
|
|
2985
|
+
|
|
2986
|
+
def float_tile_streams_root(
|
|
2987
|
+
dtype: str,
|
|
2988
|
+
dims: "list[int]",
|
|
2989
|
+
element_len: int,
|
|
2990
|
+
exponent_plane_len: int,
|
|
2991
|
+
mantissa_plane_len: int,
|
|
2992
|
+
n_streams: int,
|
|
2993
|
+
stream_offsets: "list[int]",
|
|
2994
|
+
tile_bytes: bytes,
|
|
2995
|
+
limits: Limits,
|
|
2996
|
+
) -> bytes:
|
|
2997
|
+
record = {
|
|
2998
|
+
"dtype": dtype,
|
|
2999
|
+
"dims": _u64_vec(dims),
|
|
3000
|
+
"element_len": wcb.U64(element_len),
|
|
3001
|
+
"exponent_plane_len": wcb.U64(exponent_plane_len),
|
|
3002
|
+
"mantissa_plane_len": wcb.U64(mantissa_plane_len),
|
|
3003
|
+
"n_streams": wcb.U64(n_streams),
|
|
3004
|
+
"stream_offsets": _u64_vec(stream_offsets),
|
|
3005
|
+
"values_hash": values_hash(tile_bytes),
|
|
3006
|
+
}
|
|
3007
|
+
return hash_domain(DOMAIN_TENSOR_FLOAT_TILE_STREAMS, wcb.encode(record, limits))
|
|
3008
|
+
|
|
3009
|
+
|
|
3010
|
+
def parse_float_tiled_streams_descriptor(
|
|
3011
|
+
descriptor_bytes: bytes, decoded_len: int, payload_len: int, limits: Limits
|
|
3012
|
+
) -> FloatTiledStreamsMeta:
|
|
3013
|
+
map_value = _map(
|
|
3014
|
+
_decode_descriptor(descriptor_bytes, limits), "float tiled streams tensor"
|
|
3015
|
+
)
|
|
3016
|
+
_exact_keys(map_value, _FLOAT_TILED_STREAMS_KEYS, "float tiled streams tensor")
|
|
3017
|
+
schema = _text(map_value, "schema")
|
|
3018
|
+
if schema != FLOAT_TILED_STREAMS_SCHEMA:
|
|
3019
|
+
raise Malformed("unsupported float tiled streams schema `%s`" % schema)
|
|
3020
|
+
dtype = _text(map_value, "dtype")
|
|
3021
|
+
if dtype not in FLOAT_FIELD_BITS:
|
|
3022
|
+
raise Malformed("unknown float tiled streams dtype `%s`" % dtype)
|
|
3023
|
+
if dtype not in FLOAT_TILED_DTYPES:
|
|
3024
|
+
raise UnsupportedFeature(
|
|
3025
|
+
"float tiled streams dtype `%s` is not implemented in v1" % dtype
|
|
3026
|
+
)
|
|
3027
|
+
layout = _text(map_value, "layout")
|
|
3028
|
+
if layout != TILE_LAYOUT:
|
|
3029
|
+
raise Malformed(
|
|
3030
|
+
"float tiled streams layout `%s` is specified but not implemented "
|
|
3031
|
+
"in v1 (only %s)" % (layout, TILE_LAYOUT)
|
|
3032
|
+
)
|
|
3033
|
+
coder = _text(map_value, "exponent_coder")
|
|
3034
|
+
if coder != PLANE_CODEC:
|
|
3035
|
+
raise Malformed(
|
|
3036
|
+
"unsupported float tiled streams exponent coder `%s` (v1 allows `%s`)"
|
|
3037
|
+
% (coder, PLANE_CODEC)
|
|
3038
|
+
)
|
|
3039
|
+
n_streams = _u64(map_value, "n_streams")
|
|
3040
|
+
if not FLOAT_TILED_STREAMS_MIN <= n_streams <= FLOAT_TILED_STREAMS_MAX:
|
|
3041
|
+
raise Malformed(
|
|
3042
|
+
"n_streams %d is outside %d..=%d"
|
|
3043
|
+
% (n_streams, FLOAT_TILED_STREAMS_MIN, FLOAT_TILED_STREAMS_MAX)
|
|
3044
|
+
)
|
|
3045
|
+
shape = _shape(map_value)
|
|
3046
|
+
element_count = _validate_shape(shape, limits, "float tiled streams tensor")
|
|
3047
|
+
declared_count = _u64(map_value, "element_count")
|
|
3048
|
+
if declared_count != element_count:
|
|
3049
|
+
raise Malformed(
|
|
3050
|
+
"float tiled streams element_count %d does not match shape product %d"
|
|
3051
|
+
% (declared_count, element_count)
|
|
3052
|
+
)
|
|
3053
|
+
sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[dtype]
|
|
3054
|
+
element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
|
|
3055
|
+
decoded_bytes = _checked_mul(
|
|
3056
|
+
element_count, element_size, "float tiled streams decoded bytes"
|
|
3057
|
+
)
|
|
3058
|
+
_check_len(
|
|
3059
|
+
decoded_bytes, limits.max_decoded_bytes, "float tiled streams decoded bytes"
|
|
3060
|
+
)
|
|
3061
|
+
if decoded_len != decoded_bytes:
|
|
3062
|
+
raise Corrupt(
|
|
3063
|
+
"float tiled streams decoded_len is %d but the descriptor requires %d"
|
|
3064
|
+
% (decoded_len, decoded_bytes)
|
|
3065
|
+
)
|
|
3066
|
+
tile_shape = _array_u64(map_value, "tile_shape")
|
|
3067
|
+
counts = _tile_counts(shape, tile_shape)
|
|
3068
|
+
if _array_u64(map_value, "tile_counts") != counts:
|
|
3069
|
+
raise Malformed(
|
|
3070
|
+
"float tiled streams tile_counts %r do not match shape/tile_shape %r"
|
|
3071
|
+
% (_array_u64(map_value, "tile_counts"), counts)
|
|
3072
|
+
)
|
|
3073
|
+
tile_elements = _validate_shape(tile_shape, limits, "float tensor tile")
|
|
3074
|
+
tile_element_len = _checked_mul(
|
|
3075
|
+
tile_elements, element_size, "float tensor tile byte length"
|
|
3076
|
+
)
|
|
3077
|
+
tile_count = 1
|
|
3078
|
+
for count in counts:
|
|
3079
|
+
tile_count = _checked_mul(tile_count, count, "float tensor tile count")
|
|
3080
|
+
if tile_count > limits.max_structured_items:
|
|
3081
|
+
raise LimitExceeded(
|
|
3082
|
+
"float tensor tile count %d exceeds limit %d"
|
|
3083
|
+
% (tile_count, limits.max_structured_items)
|
|
3084
|
+
)
|
|
3085
|
+
mantissa_total_bits = _checked_mul(
|
|
3086
|
+
tile_elements, sign_bits + mantissa_bits, "float tensor tile mantissa bits"
|
|
3087
|
+
)
|
|
3088
|
+
expected_tile_mantissa = _checked_add(
|
|
3089
|
+
mantissa_total_bits, 7, "float tensor tile mantissa length"
|
|
3090
|
+
) // 8
|
|
3091
|
+
declared_payload_len = _u64(map_value, "payload_len")
|
|
3092
|
+
if declared_payload_len != payload_len:
|
|
3093
|
+
raise Corrupt(
|
|
3094
|
+
"float tiled streams payload is %d bytes but metadata declares %d"
|
|
3095
|
+
% (payload_len, declared_payload_len)
|
|
3096
|
+
)
|
|
3097
|
+
entries = map_value.get("tiles")
|
|
3098
|
+
if not isinstance(entries, list):
|
|
3099
|
+
raise Malformed("float tiled streams `tiles` must be an array")
|
|
3100
|
+
if len(entries) != tile_count:
|
|
3101
|
+
raise Malformed(
|
|
3102
|
+
"float tiled streams has %d tile entries but the grid requires %d"
|
|
3103
|
+
% (len(entries), tile_count)
|
|
3104
|
+
)
|
|
3105
|
+
tiles = []
|
|
3106
|
+
expected_offset = 0
|
|
3107
|
+
for position, item in enumerate(entries):
|
|
3108
|
+
entry = _map(item, "float tensor streams tile")
|
|
3109
|
+
_exact_keys(entry, _FLOAT_TILE_STREAMS_KEYS, "float tensor streams tile")
|
|
3110
|
+
index = _u64(entry, "index")
|
|
3111
|
+
if index != position:
|
|
3112
|
+
raise Malformed(
|
|
3113
|
+
"float tiled streams entry %d declares index %d; entries must be "
|
|
3114
|
+
"in canonical order" % (position, index)
|
|
3115
|
+
)
|
|
3116
|
+
dims = _array_u64(entry, "dims")
|
|
3117
|
+
if dims != tile_shape:
|
|
3118
|
+
raise Malformed(
|
|
3119
|
+
"float tensor tile %d dims %r do not equal tile_shape %r; "
|
|
3120
|
+
"partial tiles are not defined in v1" % (position, dims, tile_shape)
|
|
3121
|
+
)
|
|
3122
|
+
declared_element_len = _u64(entry, "element_len")
|
|
3123
|
+
if declared_element_len != tile_element_len:
|
|
3124
|
+
raise Malformed(
|
|
3125
|
+
"float tensor tile %d element_len %d does not match %d"
|
|
3126
|
+
% (position, declared_element_len, tile_element_len)
|
|
3127
|
+
)
|
|
3128
|
+
offset = _u64(entry, "offset")
|
|
3129
|
+
if offset != expected_offset:
|
|
3130
|
+
raise Malformed(
|
|
3131
|
+
"float tensor tile %d offset %d is not the contiguous offset %d"
|
|
3132
|
+
% (position, offset, expected_offset)
|
|
3133
|
+
)
|
|
3134
|
+
compressed_len = _u64(entry, "compressed_len")
|
|
3135
|
+
exponent_plane_len = _u64(entry, "exponent_plane_len")
|
|
3136
|
+
mantissa_plane_len = _u64(entry, "mantissa_plane_len")
|
|
3137
|
+
if mantissa_plane_len != expected_tile_mantissa:
|
|
3138
|
+
raise Malformed(
|
|
3139
|
+
"float tensor tile %d mantissa_plane_len %d does not match the "
|
|
3140
|
+
"tile arithmetic %d"
|
|
3141
|
+
% (position, mantissa_plane_len, expected_tile_mantissa)
|
|
3142
|
+
)
|
|
3143
|
+
if exponent_plane_len + mantissa_plane_len != compressed_len:
|
|
3144
|
+
raise Malformed(
|
|
3145
|
+
"float tensor tile %d plane lengths do not sum to compressed_len"
|
|
3146
|
+
% position
|
|
3147
|
+
)
|
|
3148
|
+
expected_offset = _checked_add(
|
|
3149
|
+
expected_offset, compressed_len, "float tensor tile offset"
|
|
3150
|
+
)
|
|
3151
|
+
stream_offsets = _array_u64(entry, "stream_offsets")
|
|
3152
|
+
if len(stream_offsets) != 2 * n_streams:
|
|
3153
|
+
raise Malformed(
|
|
3154
|
+
"float tile %d has %d stream offsets but n_streams %d requires %d"
|
|
3155
|
+
% (position, len(stream_offsets), n_streams, 2 * n_streams)
|
|
3156
|
+
)
|
|
3157
|
+
if stream_offsets[0] != 0:
|
|
3158
|
+
raise Malformed(
|
|
3159
|
+
"float tile %d stream_offsets[0] is %d not 0"
|
|
3160
|
+
% (position, stream_offsets[0])
|
|
3161
|
+
)
|
|
3162
|
+
for earlier, later in zip(stream_offsets, stream_offsets[1:]):
|
|
3163
|
+
if later < earlier:
|
|
3164
|
+
raise Malformed(
|
|
3165
|
+
"float tile %d stream_offsets are not monotonic: %d then %d"
|
|
3166
|
+
% (position, earlier, later)
|
|
3167
|
+
)
|
|
3168
|
+
if stream_offsets[-1] > compressed_len:
|
|
3169
|
+
raise Malformed(
|
|
3170
|
+
"float tile %d stream offset %d is past the %d-byte tile payload"
|
|
3171
|
+
% (position, stream_offsets[-1], compressed_len)
|
|
3172
|
+
)
|
|
3173
|
+
if stream_offsets[n_streams] != exponent_plane_len:
|
|
3174
|
+
raise Malformed(
|
|
3175
|
+
"float tile %d stream boundary %d does not equal exponent_plane_len %d"
|
|
3176
|
+
% (position, stream_offsets[n_streams], exponent_plane_len)
|
|
3177
|
+
)
|
|
3178
|
+
for stream in range(n_streams):
|
|
3179
|
+
end = (
|
|
3180
|
+
stream_offsets[stream + 1]
|
|
3181
|
+
if stream + 1 < n_streams
|
|
3182
|
+
else exponent_plane_len
|
|
3183
|
+
)
|
|
3184
|
+
if end - stream_offsets[stream] < 256:
|
|
3185
|
+
raise Malformed(
|
|
3186
|
+
"float tile %d exponent stream %d is shorter than its "
|
|
3187
|
+
"256-byte code-length table" % (position, stream)
|
|
3188
|
+
)
|
|
3189
|
+
for stream in range(n_streams):
|
|
3190
|
+
start = stream_offsets[n_streams + stream] - exponent_plane_len
|
|
3191
|
+
end = (
|
|
3192
|
+
stream_offsets[n_streams + stream + 1] - exponent_plane_len
|
|
3193
|
+
if stream + 1 < n_streams
|
|
3194
|
+
else mantissa_plane_len
|
|
3195
|
+
)
|
|
3196
|
+
expected = interleaved_stream_len(
|
|
3197
|
+
mantissa_plane_len, stream, n_streams
|
|
3198
|
+
)
|
|
3199
|
+
if end - start != expected:
|
|
3200
|
+
raise Malformed(
|
|
3201
|
+
"float tile %d mantissa stream %d is %d bytes but the "
|
|
3202
|
+
"interleave requires %d"
|
|
3203
|
+
% (position, stream, end - start, expected)
|
|
3204
|
+
)
|
|
3205
|
+
root = _bytes_of(entry, "tile_root")
|
|
3206
|
+
if len(root) != HASH_LEN:
|
|
3207
|
+
raise Malformed("float tensor tile_root must be 32 bytes")
|
|
3208
|
+
tiles.append(
|
|
3209
|
+
FloatTileStreamsEntry(
|
|
3210
|
+
index=index,
|
|
3211
|
+
dims=dims,
|
|
3212
|
+
offset=offset,
|
|
3213
|
+
compressed_len=compressed_len,
|
|
3214
|
+
exponent_plane_len=exponent_plane_len,
|
|
3215
|
+
mantissa_plane_len=mantissa_plane_len,
|
|
3216
|
+
element_len=declared_element_len,
|
|
3217
|
+
stream_offsets=stream_offsets,
|
|
3218
|
+
tile_root=root,
|
|
3219
|
+
)
|
|
3220
|
+
)
|
|
3221
|
+
if expected_offset != declared_payload_len:
|
|
3222
|
+
raise Corrupt(
|
|
3223
|
+
"float tensor tile ranges cover %d bytes but payload_len is %d"
|
|
3224
|
+
% (expected_offset, declared_payload_len)
|
|
3225
|
+
)
|
|
3226
|
+
logical_root = _bytes_of(map_value, "logical_root")
|
|
3227
|
+
if len(logical_root) != HASH_LEN:
|
|
3228
|
+
raise Malformed("float tiled streams logical_root must be 32 bytes")
|
|
3229
|
+
return FloatTiledStreamsMeta(
|
|
3230
|
+
dtype=dtype,
|
|
3231
|
+
shape=shape,
|
|
3232
|
+
element_count=element_count,
|
|
3233
|
+
tile_shape=tile_shape,
|
|
3234
|
+
tile_counts=counts,
|
|
3235
|
+
n_streams=n_streams,
|
|
3236
|
+
tiles=tiles,
|
|
3237
|
+
payload_len=declared_payload_len,
|
|
3238
|
+
logical_root=logical_root,
|
|
3239
|
+
)
|
|
3240
|
+
|
|
3241
|
+
|
|
3242
|
+
def _deinterleave(streams: "list[bytes]", total: int, what: str) -> bytes:
|
|
3243
|
+
n_streams = len(streams)
|
|
3244
|
+
if n_streams == 0:
|
|
3245
|
+
raise Malformed("%s declares zero streams" % what)
|
|
3246
|
+
out = bytearray(total)
|
|
3247
|
+
for index in range(total):
|
|
3248
|
+
stream = streams[index % n_streams]
|
|
3249
|
+
position = index // n_streams
|
|
3250
|
+
if position >= len(stream):
|
|
3251
|
+
raise Corrupt("%s stream ended before its interleaved plane" % what)
|
|
3252
|
+
out[index] = stream[position]
|
|
3253
|
+
return bytes(out)
|
|
3254
|
+
|
|
3255
|
+
|
|
3256
|
+
def decode_float_tiled_streams(frame, limits: Limits) -> bytes:
|
|
3257
|
+
descriptor_bytes, payload = _split_prefix(frame.body)
|
|
3258
|
+
meta = parse_float_tiled_streams_descriptor(
|
|
3259
|
+
descriptor_bytes, frame.header.decoded_len, len(payload), limits
|
|
3260
|
+
)
|
|
3261
|
+
sign_bits, exponent_bits, mantissa_bits = FLOAT_FIELD_BITS[meta.dtype]
|
|
3262
|
+
element_size = (sign_bits + exponent_bits + mantissa_bits) // 8
|
|
3263
|
+
total = meta.element_count * element_size
|
|
3264
|
+
out = bytearray(total)
|
|
3265
|
+
tiles_bytes = 0
|
|
3266
|
+
n_streams = meta.n_streams
|
|
3267
|
+
for entry in meta.tiles:
|
|
3268
|
+
start = entry.offset
|
|
3269
|
+
end = start + entry.compressed_len
|
|
3270
|
+
if end > len(payload):
|
|
3271
|
+
raise Corrupt("float tensor tile range is outside the payload")
|
|
3272
|
+
stored = payload[start:end]
|
|
3273
|
+
if entry.exponent_plane_len > len(stored):
|
|
3274
|
+
raise Corrupt("float tile exponent plane is truncated")
|
|
3275
|
+
exponent_stored = stored[: entry.exponent_plane_len]
|
|
3276
|
+
mantissa_stored = stored[entry.exponent_plane_len :]
|
|
3277
|
+
elements = entry.element_len // element_size
|
|
3278
|
+
exponent_parts = []
|
|
3279
|
+
for stream in range(n_streams):
|
|
3280
|
+
s_start = entry.stream_offsets[stream]
|
|
3281
|
+
s_end = (
|
|
3282
|
+
entry.stream_offsets[stream + 1]
|
|
3283
|
+
if stream + 1 < n_streams
|
|
3284
|
+
else entry.exponent_plane_len
|
|
3285
|
+
)
|
|
3286
|
+
if s_end > len(exponent_stored) or s_start > s_end:
|
|
3287
|
+
raise Corrupt("float tile exponent stream is outside the tile")
|
|
3288
|
+
expected = interleaved_stream_len(elements, stream, n_streams)
|
|
3289
|
+
exponent_parts.append(
|
|
3290
|
+
huffman.decode(exponent_stored[s_start:s_end], expected)
|
|
3291
|
+
)
|
|
3292
|
+
exponent_plane = _deinterleave(exponent_parts, elements, "exponent plane")
|
|
3293
|
+
mantissa_parts = []
|
|
3294
|
+
for stream in range(n_streams):
|
|
3295
|
+
s_start = entry.stream_offsets[n_streams + stream] - entry.exponent_plane_len
|
|
3296
|
+
s_end = (
|
|
3297
|
+
entry.stream_offsets[n_streams + stream + 1] - entry.exponent_plane_len
|
|
3298
|
+
if stream + 1 < n_streams
|
|
3299
|
+
else entry.mantissa_plane_len
|
|
3300
|
+
)
|
|
3301
|
+
if s_end > len(mantissa_stored) or s_start > s_end:
|
|
3302
|
+
raise Corrupt("float tile mantissa stream is outside the tile")
|
|
3303
|
+
expected = interleaved_stream_len(
|
|
3304
|
+
entry.mantissa_plane_len, stream, n_streams
|
|
3305
|
+
)
|
|
3306
|
+
if s_end - s_start != expected:
|
|
3307
|
+
raise Corrupt(
|
|
3308
|
+
"float tile %d mantissa stream %d length differs from the "
|
|
3309
|
+
"interleave" % (entry.index, stream)
|
|
3310
|
+
)
|
|
3311
|
+
mantissa_parts.append(mantissa_stored[s_start:s_end])
|
|
3312
|
+
mantissa_plane = _deinterleave(
|
|
3313
|
+
mantissa_parts, entry.mantissa_plane_len, "mantissa plane"
|
|
3314
|
+
)
|
|
3315
|
+
tile = _join_float_planes(
|
|
3316
|
+
exponent_plane, mantissa_plane, meta.dtype, elements, limits
|
|
3317
|
+
)
|
|
3318
|
+
root = float_tile_streams_root(
|
|
3319
|
+
meta.dtype,
|
|
3320
|
+
entry.dims,
|
|
3321
|
+
entry.element_len,
|
|
3322
|
+
entry.exponent_plane_len,
|
|
3323
|
+
entry.mantissa_plane_len,
|
|
3324
|
+
n_streams,
|
|
3325
|
+
entry.stream_offsets,
|
|
3326
|
+
tile,
|
|
3327
|
+
limits,
|
|
3328
|
+
)
|
|
3329
|
+
if root != entry.tile_root:
|
|
3330
|
+
raise Corrupt(
|
|
3331
|
+
"float tensor tile %d stream root does not match the descriptor"
|
|
3332
|
+
% entry.index
|
|
3333
|
+
)
|
|
3334
|
+
tiles_bytes += len(tile)
|
|
3335
|
+
tile_coords = _tile_start(meta.tile_counts, entry.index)
|
|
3336
|
+
dst_start = [
|
|
3337
|
+
coord * meta.tile_shape[axis] for axis, coord in enumerate(tile_coords)
|
|
3338
|
+
]
|
|
3339
|
+
for local in _iter_coords(entry.dims):
|
|
3340
|
+
src = _flat_index(local, entry.dims) * element_size
|
|
3341
|
+
full = [dst_start[axis] + local[axis] for axis in range(len(meta.shape))]
|
|
3342
|
+
dst = _flat_index(full, meta.shape) * element_size
|
|
3343
|
+
out[dst : dst + element_size] = tile[src : src + element_size]
|
|
3344
|
+
if tiles_bytes != total:
|
|
3345
|
+
raise Corrupt(
|
|
3346
|
+
"float tiled streams decoded to %d bytes but the descriptor requires %d"
|
|
3347
|
+
% (tiles_bytes, total)
|
|
3348
|
+
)
|
|
3349
|
+
logical = tensor_logical_root(
|
|
3350
|
+
FLOAT_TILED_LOGICAL_DTYPE[meta.dtype],
|
|
3351
|
+
meta.shape,
|
|
3352
|
+
meta.element_count,
|
|
3353
|
+
bytes(out),
|
|
3354
|
+
limits,
|
|
3355
|
+
)
|
|
3356
|
+
if logical != meta.logical_root:
|
|
3357
|
+
raise Corrupt(
|
|
3358
|
+
"float tiled streams logical root does not match the descriptor"
|
|
3359
|
+
)
|
|
3360
|
+
return bytes(out)
|