[Web] Decode packed BF16 tensor records in place - #20167
Conversation
Packed BF16 decoding currently uses separate input and output storage. Add a Wasm entry point that expands values backward in the destination buffer after the packed bytes have been copied into its lower half. Use the in-place path when the entry point is available and keep the existing decoder otherwise. Both paths account for DLTensor byte_offset.
acae0af to
d6cbf31
Compare
guan404ming
left a comment
guan404ming
left a comment
There was a problem hiding this comment.
Thanks for the PR! Built with emsdk 4.0.23 and ran tests/node locally: all 83 tests pass, including the 20 new/updated ones. The backward in-place expansion looks correct to me. A few small suggestions below, none blocking.
| }); | ||
| } | ||
|
|
||
| size_t GetCheckedTensorElementCount(const Tensor& tensor) { |
There was a problem hiding this comment.
Maybe a plain product over shape would be enough here. tensor shapes are already validated at allocation time, so the overflow checks feel a little heavier than needed.
There was a problem hiding this comment.
I couldn't find shape validation in the allocation path. ffi::GetDataSize multiplies dimensions as size_t without checking negative values or overflow. This matters on wasm32, where an int64_t dimension could truncate before allocation. Since this function is also callable directly through FFI, I think the checks should remain unless validation is added to the common tensor construction path.
| } | ||
|
|
||
| /** Return the exact byte size of a packed BF16 tensor. */ | ||
| private getPackedBF16Bytes(shape: Array<number>): number { |
There was a problem hiding this comment.
I think we could drop this helper, since ArrayDecodeBF16ToF32Inplace already rejects a size mismatch via TVM_FFI_ICHECK_EQ. Keeping the check in one place would be a bit simpler to maintain.
There was a problem hiding this comment.
The C++ check only runs after storeRawBytes has copied the packed data into tensor memory, so it cannot protect that copy. I think the JS check is needed before the write. Passing the bytes through ArrayDecodeStorage would reintroduce the temporary allocation this path removes. I prefer to keep the C++ check as validation of the newly exposed FFI entry point itself since it can be invoked independently of the tensor-cache loader.
2 participants
Reduce peak WASM memory during BF16 tensor-cache loading by decoding directly in the final
float32CPU tensor allocation. The packed BF16 input is copied into the lower half of the destination tensor and expanded backward, preventing output writes from overwriting unread input. This removes the separate FFI byte-array allocation previously required during decoding. Leave the existing decoder as a fallback when the new WASM entry point is unavailable.