[tgghzyh]

Embedding4bit already gathers the packed rows and their scales for the requested token IDs. It then deep-copies the full weight quantization state, including the vocabulary-wide absmax tensor, only to overwrite that copy with the gathered scales.

Use a shallow copy of the state. Assigning absmax and shape on the copy leaves the weight's state intact, and dequantization only reads the other fields.

GPU measurements

Single NVIDIA A10, PyTorch 2.13.0+cu130, bitsandbytes at a616886 with the CUDA 13.0 native library from 0.50.2. The NF4 table had 65,536 rows, 4,096 columns, FP16 source weights and a 16 MiB scale tensor.

Token IDs per lookup	Extra peak allocated, before → after	Median time per lookup, before → after
1	16,780,800 → 10,752 bytes	263.19 → 104.13 µs
32	16,851,968 → 335,872 bytes	277.66 → 116.63 µs

Peak is the increase in torch.cuda.max_memory_allocated() over the allocation immediately before one warmed-up lookup, including its output. Timing is the median of seven synchronized batches of 100 calls after 20 warmups. Each configuration matched F.embedding(ids, dequantize_4bit(weight, state)) exactly.

The existing test_embedding_lossless and test_embedding_error cases passed on CUDA (60 passed). All pre-commit hooks passed.