[kevchengcodes]

Type of Change

• Bug fix
• Small feature (agreed in an issue)
• New model (export recipe only)
• New model (export recipe and new runtime support)
• Documentation update

Motivation and Context

This PR fixes the issue found in #313. The per-layer embedding table for Gemma-3n-E4B is 262,144 x 8,960 parameters. Currently, this is stored in one table. The GPU gather can count up to 2^31 parameters, so lookups for token IDs greater than 239,675 (with 8,960 elements per row) silently fail and return rows of zeroes.

This fix checks the table size and creates a ChunkedEmbedding (splits the table in half so each row holds only 4,480 columns) if the number of elements exceeds 2^31. Then, a token lookup is done on both tables and the row is put back together.

This leads to bit-identical results as before (besides fixing the issue), and the extra operation of split-table lookup/concat is negligible.

Changes

• new ChunkedEmbedding table
• a warning print during the export pipeline if the table is larger than 2^31
• new constant: MAX_GATHER_CONSTANT_ELEMENTS
• new tests to verify chunking behavior and results

Testing

• Tested locally
• Added/updated tests
• [] Added/updated docs

The issue contained a repro script. I was able to reproduce the error mode with the existing code, and verified that this fix passes that script now. Essentially, we now properly lookup the table instead of silently returning rows of zeroes.

Follow Up

This may be fixed downstream in a future OS release. At that point, this fix can be removed.