[Web] Improve large tensor loading in wasm runtime

[MakotoUwu]

This splits the Web runtime portion of #19766 into a focused tensor-cache
chunking change.

What changes

• Plan outer-dimension chunks from both encoded and decoded byte counts.
• Plan CPU decode independently from 4-byte-aligned WebGPU copies.
• Create and dispose chunk tensor views safely.
• Respect cpu_arr->byte_offset when BF16 data is expanded into a chunk view.
• Avoid a record-sized ArrayBuffer.slice copy before Wasm staging.

The planner chunks only when an outer-dimension split can keep both source and
target payloads within the cap and, for WebGPU, keep every copy offset and
length 4-byte aligned. Unsupported layouts retain the existing full-record
path. This PR does not change tensor-cache format semantics or general JS/Wasm
integer and Shape marshalling.

Motivation and failure boundary

The historical Gemma 4 E2B artifact at Hugging Face revision
4e7d43f11998bac8aa25e46bc43d6a16e6d78131
contains this record:

name:   model.embed_tokens_per_layer.q_weight
shape:  [262144, 1120]
dtype:  uint32
format: f32-to-bf16
nbytes: 1174405120 (1120 MiB)

Because the dtype is uint32, this record uses the normal raw-copy branch; it
does not require a native-float32 fallthrough under f32-to-bf16.

On the unchunked path, the final CPU tensor first grows Wasm memory to
1174994944 bytes. JavaScript then creates the record-sized call-stack buffer,
but CachedCallStack.allocThenSetArgBytes fails while allocating a second
record-sized block in Wasm argument space. No bytes are copied and the C++
packed function is never entered. The runtime Wasm memory has a 2 GiB maximum.

With the same artifact and environment:

1. Clean apache/main full-record loading fails at that allocation.
2. A validation-only build with the cap set to Number.MAX_SAFE_INTEGER
   deliberately disables chunking and fails at the same allocation.
3. The unmodified proposed revision with 128 MiB chunks succeeds in 9 decode
   calls.

WebGPU and content validation

The current artifact
shards the historical 1120 MiB weight. The adapter advertises 4294967292
bytes for both buffer and binding limits, while the requested device exposes 1
GiB for maxBufferSize and maxStorageBufferBindingSize. The WebGPU check uses
the current 384 MiB shard, which fits those requested limits:

name:   model.embed_tokens_per_layer.shards.0.q_weight
shape:  [262144, 384]
dtype:  uint32
nbytes: 402653184 (384 MiB)

Chrome 150 on Apple Metal WebGPU loaded that tensor with 128 MiB chunks,
synchronized it, copied it back to CPU, and matched the source SHA-256:

c5ba516f710992810c9511c83c7c8f1d6ea3ce912fb313d6026ebb72c537ab6f

There was no device loss or uncaptured WebGPU error. A separate 256 MiB raw CPU
record also produced identical content through full-record and chunked loading:

405342ee2074eca1d3312ad16473d43f1b7700c91d9d2a7abb74e1139988c625

Chunk cap

The 1120 MiB CPU reproducer produced this diagnostic sweep:

32 MiB    pass, 36 chunks
64 MiB    pass, 18 chunks
128 MiB   pass, 9 chunks
256 MiB   pass, 5 chunks
512 MiB   fail, Wasm allocation request 1073739664 bytes
1024 MiB  fail, Wasm allocation request 1073739592 bytes

CachedCallStack grows geometrically and retains its backing allocation. The
128 MiB cap stays below the largest passing value to leave headroom for that
growth, call metadata, the final tensor, and other Wasm allocations.

Validation

• ESLint: pass
• TypeScript --noEmit: pass
• Planner tests: 8/8 pass
• Planner property check: 18,785 valid plans pass
• Focused planner and tensor-runtime tests: 10/10 pass
• Broader runnable Node tests: 22/22 pass across five suites
• Emscripten 3.1.56 runtime build: pass
• JS bundle build: pass
• Changed-file pre-commit hooks: pass
• git diff --check: pass
• Exact final Chrome/WebGPU 384 MiB round trip: pass
• WebLLM v0.2.82 rebuilt against the exact TVM source, Gemma 4 E2B smoke: 3/3 pass

[gemini-code-assist]

Code Review

This pull request restructures TVM-FFI includes in wasm_runtime.cc to prevent static initialization crashes and updates ArrayDecodeStorage to tolerate uncompressed float32 weights under the 'f32-to-bf16' format. In runtime.ts, it introduces chunked record loading and copying (up to 128MB chunks) to handle large tensors efficiently, and adds support for kTVMFFIShape types. The review feedback suggests optimizing these chunking loops by utilizing the cached makeShapeTuple method on the Instance class rather than invoking the FFI this.ctx.makeShapeTuple repeatedly, which reduces redundant FFI round-trips.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[tqchen]

how is the number determined, would be good to have a sense of what WebGPU runtime supports, the motivation here is not as clear

[MakotoUwu]

I reran the historical 1120 MiB CPU reproducer with the same artifact and host:

32 MiB    pass, 36 chunks
64 MiB    pass, 18 chunks
128 MiB   pass, 9 chunks
256 MiB   pass, 5 chunks
512 MiB   fail
1024 MiB  fail

CachedCallStack grows geometrically and retains its backing allocation. At
512 and 1024 MiB, that growth drives the next Wasm allocation close to 1 GiB.
The 128 MiB default stays below the largest passing value to leave headroom for
the final tensor, retained staging memory, call metadata, and other Wasm
allocations.

[tqchen]

@akaashrp would be great if you can help review

[tqchen]

Please elaborate on "reorder FFI implementation includes before runtime implementation includes in the wasm single-translation-unit build, avoiding static initialization ordering issues during module startup"

Since static init ordering should not impact the startup, if there is an impact we need to know details

[tqchen]

would be great to get elaboration on "load large tensor-cache records in chunks to avoid oversized JS-to-wasm decode/copy calls", is there a specific scenario that motivate this, generally, we would favor simplicity as long as such case is covered by webgpu runtime

[MakotoUwu]

Thanks for the review. I updated the PR to narrow the scope and make the motivation more explicit:

• Removed the FFI include reordering change from this PR. I agree the previous static-init wording was too strong without a concrete wasm startup trace in this split PR, so this now leaves the existing include order unchanged.
• Kept the large tensor-cache record chunking and clarified the motivation in the PR description and code comment. The concrete case is a single tensor-cache record larger than the conservative WebGPU staging/view size used by the web runtime. The 128 MiB cap follows the existing maxStorageBufferBindingSize fallback in web/src/webgpu.ts; it is only a per-decode/per-copy view cap for loading large records, while smaller records still use the existing full-record path.

Local checks still pass after the update: npm run lint, npx tsc --noEmit --pretty false, and git diff --check.

[tqchen]

Thanks, is there a real usecase for the chunking? .eg. if maxStorageBufferBindingSize limit is set really to 128MB, seems that means we canot allocate tensor larger than that anyway as a result we won't have copy in such shape. And in cases where maxStorageBufferBindingSize is larger, we don't need chunking.

[MakotoUwu]

You are right that maxStorageBufferBindingSize was not the right motivation. I updated the code comment and PR description to avoid that claim.

The concrete use case is a very large tensor-cache record that is valid for the target device but fragile as one JS/Wasm staging call. The Gemma 4 E2B artifact has three records above 128 MiB:

• model.embed_tokens.q_weight: 192 MiB, shape [262144, 192]
• model.embed_tokens_per_layer.q_scale: 140 MiB, shape [262144, 280]
• model.embed_tokens_per_layer.q_weight: 1120 MiB, shape [262144, 1120]

In the local Chrome/WebGPU validation lane, the failure was deterministic on the 1120 MiB record. Instrumentation reached arrayDecodeStorage:start and aborted before arrayDecodeStorage:done, before GPU copy. So the chunking is intended to avoid one multi-hundred-MiB JS-to-wasm byte-array decode/copy call, not to work around a WebGPU allocation/binding limit. Smaller records still use the original full-record path.

[tqchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces chunked loading for large tensor-cache records to avoid massive single JS-to-wasm byte-array calls, and updates the WASM runtime to handle uncompressed float32 payloads labeled as 'f32-to-bf16'. Feedback on these changes highlights several critical issues: the custom f32-to-bf16 decoding path in C++ ignores cpu_arr->byte_offset which causes data corruption during chunked loading; the storageBytes helper in TypeScript fails for "bool" data types; a potential resource leak exists in the GPU chunked copy path if an exception occurs during view creation; and a defensive null check is needed for shapeObjPtr before loading memory.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[MakotoUwu]

Addressed the Gemini follow-up comments in the latest commit:

• ArrayDecodeStorage now honors cpu_arr->byte_offset in the packed bf16 expansion path, so chunk views write to the correct offset.
• storageBytes now handles bool explicitly.
• The GPU chunk copy path now creates both views before detaching them from the scope, avoiding a leak if the second view creation throws.
• kTVMFFIShape callback decoding now guards against a null shape object pointer.

Local checks still pass: npm run lint, npx tsc --noEmit --pretty false, and git diff --check.

[MakotoUwu]

Pushed a small follow-up on the new head 1d2a3f131:

• Made withNewScope close scopes in a finally block, so failed chunk-view creation cannot leave a live scope/view behind.
• Tightened the f32-to-bf16 decode branch to require the exact bf16 payload size before expansion, so malformed payload sizes fall through to the normal size validation path.

Local validation passes: git diff --check, npm run lint, and npx tsc --noEmit --pretty false. Remote CI is in progress on the new head.

[MakotoUwu]

Pushed a small follow-up on head a2359c159 after rerunning the downstream browser smoke with a WebLLM bundle rebuilt from this PR branch.

The smoke first exposed one more scope issue in the chunked path: this.makeShapeTuple(chunkShape) can create a cached TVM object on a cache miss, so calling it before withNewScope(...) failed with Must call beginScope to use functions that returns TVM objects. I moved both CPU and GPU chunk-path makeShapeTuple calls inside the existing withNewScope callbacks.

Local TVM web checks pass on the pushed head:

• git diff --check
• npm run lint
• npx tsc --noEmit --pretty false

Downstream Chrome/WebGPU smoke also passes with a WebLLM JS bundle rebuilt against the pushed TVM source:

• TVM PR head: a2359c159cb3fb94766335beea33f281fb9a7bba
• Rebuilt WebLLM bundle SHA256: c357f887fde2408078e715abe4b7a6b54c719f44d87e4e9ba2cefe58e101e8a5
• TVM wasm_runtime.bc SHA256: 8dbe21eeb30abda125e1b71464d16c8cf1f1c877be08c3a1a7275f0a27701a23
• Model wasm used for the smoke: 70d7295dc91b622b79ceeada2c64b4c20787832631c04e3714de95db04515dfc
• Browser: Chrome 149.0.7827.104, Apple M3 WebGPU/Metal

Prompt smoke results:

• T1 Hi: load 11.6s, generated Hi there! How can I help you, finish length - PASS
• T2 France one-word answer: load 5.3s, generated Paris, finish stop - PASS
• T3 haiku: load 5.7s, generated a non-empty haiku-like response, generation 2.8s, finish stop - PASS

Boundary note: this validates the current PR-head TypeScript/WebLLM loading path with the latest available Apache model wasm. It is not a fresh model-wasm rebuild from a2359c159.

[MakotoUwu]

Hi @tqchen, just checking whether the latest head addresses your concern about the chunking motivation and runtime scope.

One extra bit of context: the large tensor-cache entries here are tied to Gemma 4's documented E2B/E4B Per-Layer Embedding design, not just to an arbitrary local conversion artifact. Google's current Gemma 4 model card notes that the smaller models keep large per-layer token embedding tables while exposing a lower "effective" parameter count. In the WebLLM artifact, that per-layer embedding weight totals 1120 MiB and is currently split into large tensor-cache records, which is the concrete case this runtime staging change handles.

The PR itself remains scoped to the Web runtime path, required CI is green on a2359c159cb3fb94766335beea33f281fb9a7bba, and the latest follow-up also includes the downstream WebLLM smoke fix for the makeShapeTuple scope issue.

Please let me know if you would prefer an additional runtime test or a smaller adjustment before re-review.

[tqchen]

thanks @MakotoUwu i think main thing is to get reviews on the PR to make sure the logics are polished and readable, i don't have full bandwidth atm.

On chunking, i think as long as it is scoped to tensor cache loading it seems to be OK, and we should have comments about the motivation (e.g. it is possible to allocate large tensor but sometimes webgpu runtime do not like copy with too large tensor)

maybe @akaashrp can help a bit

[MakotoUwu]

@tvm-bot rerun

[MakotoUwu]

The remaining red cpu/pr-head appears unrelated to this web-only PR. The failure is in tests/python/relax/test_transform_legalize_ops_nn.py::test_dropout with call_tir() got an unexpected keyword argument "out_sinfo", while this PR only touches web/emcc/wasm_runtime.cc and web/src/runtime.ts. wasm/pr-head, gpu/pr-head, arm/pr-head, docker/pr-head, and cc-reviewers are green on 42ed716d.

I tried @tvm-bot rerun, but the bot rejected my account because I am not in the repository mentionable users list. Could a maintainer please rerun Jenkins when convenient?

[akaashrp]

Hi @MakotoUwu, apologize for the delay. Will try to review this week.

[MakotoUwu]

Hi @tqchen @akaashrp, quick update: the PR is now rebased on current apache/main, and all required checks are green on head aef19199a.

The latest follow-up tightens the generic Web runtime path and adds repository-local regression tests:

• bound both encoded and decoded chunk sizes
• keep WebGPU copy offsets and lengths 4-byte aligned
• preserve JavaScript-safe int64 sizes, offsets, and shapes across JS/Wasm
• keep the chunk planner internal rather than adding public API surface

A clean local Emscripten 3.1.56 Wasm build, bundle, and 22 focused tests also pass.

Could you please take another look when you have bandwidth?

[akaashrp]

Thanks for iterating on this. I think the PR is doing too many unrelated things and should be narrowed to chunking.

The parts that seem directly relevant are:

1. Planning chunks based on both encoded and decoded sizes.
2. Keeping WebGPU offsets and lengths aligned.
3. Cleaning up chunk views correctly.
4. Respecting cpu_arr->byte_offset in ArrayDecodeStorage, since chunks are decoded into views with nonzero offsets.

I don't think we should also change the meaning of f32-to-bf16 here. Accepting a native four-byte float32 payload under that format makes the interface ambiguous, and it would behave differently from the C++ and Python loaders. The original Gemma checkpoint is entirely BF16. From the reported names, shapes, and sizes, I would expect the converted tensors to use the usual q4f16 representation: packed 32-bit q_weight values and 16-bit scales. Could you share the relevant tensor-cache.json entries and identify the record that actually needs the new float32 fallthrough? I would expect it to have metadata like:

dtype  = float32
format = f32-to-bf16
nbytes = product(shape) × 4

If such a record exists, why can't it be tagged as raw? Either way, this seems like a separate format-compatibility change that should have its own discussion. The Shape callback and general int64 marshalling changes also don't appear necessary for chunking and should be split out unless there is a direct dependency I'm missing.

I'm also not yet convinced we know where the original failure occurred. Seeing arrayDecodeStorage:start without arrayDecodeStorage:done only tells us that the overall packed call did not return. The failure could be while growing the JS call stack, allocating Wasm argument space, copying into Wasm memory, entering C++, or performing the actual decode. The current path does create several record-sized temporary allocations, so memory pressure is a reasonable theory, but we should identify the operation that actually fails.

There is also an issue with the reported 1120 MiB tensor. TVM requests only 1 GiB for maxBufferSize and maxStorageBufferBindingSize in web/src/webgpu.ts. Chunking does not split the final GPU tensor and still allocates a full gpu_arr. Could you log both the adapter limits and the actual device limits? If the successful test used a sharded version of this tensor rather than one [262144, 1120] allocation, please update the PR description accordingly.

Could you provide a reproducible comparison showing, with the same artifact and environment:

1. The current full-record path fails at a specific, identified operation.
2. PR head with chunking disabled fails in the same way.
3. PR head with chunking enabled succeeds.
4. The final GPU allocation fits the actual device limits and is successfully used.
5. For a smaller record, such as 256 MiB, full and chunked loading produce identical tensor contents.

Finally, 128 MiB still looks arbitrary. Could you try a few chunk sizes (perhaps 32, 64, 256, 512, 1024 MiB) and report which ones work and where failures begin (if at all)?

[MakotoUwu]

Thanks for the detailed review. I've pushed a narrower revision limited to the
four chunking items you identified: planning from encoded and decoded sizes,
4-byte-aligned WebGPU copy ranges, safe chunk-view cleanup, and honoring
cpu_arr->byte_offset.

I removed the native-float32 f32-to-bf16 fallback, Shape callback handling,
generic int64 marshalling, and the storage-size helper.

I also added the requested same-artifact validation using this
pinned 1120 MiB tensor-cache record.
The record has dtype uint32, so it uses the existing raw-copy path and does
not require a format-semantics change.

Results:

• clean apache/main and a validation-only no-chunk control both fail at
  CachedCallStack.allocThenSetArgBytes, before bytes are copied or the C++
  packed function is entered
• the two failures are controls, not failures of the proposed revision
• the actual proposed revision succeeds by loading the 1120 MiB CPU record in
  9 chunks
• the adapter advertises 4,294,967,292-byte buffer limits and the requested
  device exposes 1 GiB limits
• the current 384 MiB shard fits those requested limits, is successfully
  synchronized and copied back, and matches the source contents
• a 256 MiB record produces identical contents through full-record and chunked
  loading

The chunk-size sweep passes at 32, 64, 128, and 256 MiB. It fails at 512 and
1024 MiB as retained call-stack growth drives the next Wasm allocation close to
1 GiB. I kept 128 MiB as a conservative cap to leave headroom for the final
tensor, retained staging memory, and call metadata.

Fresh CI is green on this head. The updated PR description contains the exact
traces, limits, hashes, sweep results, and validation commands.

Could you please take another look when you have bandwidth?