Repository navigation
Commit 5fe3dba
authored
Improve performance and reduce memory usage of nvtext::wordpiece_tokenize (NVIDIA#24468)
Reworks `nvtext::wordpiece_tokenize` to reduce intermediate memory, reduce warp divergence in the tokenizer, and fix incorrect results in the `max_words_per_row` path.
**Memory**
- The intermediate token array (4 bytes per input byte) and its fill, segmented-reduce, and `remove_copy` passes are replaced with a per-word layout: a token count and a single token per word, with tokens for multi-token words stored in an overflow buffer.
- Word starts are located exactly using a row-start bitmask and `count_if`/`copy_if` instead of an array sized for the worst case (4x the input size) merged with the row offsets.
- Row token counts are computed with `cub::DeviceSegmentedReduce` over the per-word counts.
**Tokenizer divergence**
Tokenizing is now done in two phases:
1. A thread per word resolves most words with a single vocabulary lookup.
2. The remaining words are deferred to a kernel where a tile of 8 threads cooperates on each word. The greedy longest-prefix search checks 8 prefix sizes per step, so each step produces a token instead of one lookup per character.
**`max_words_per_row` fix**
The previous implementation produced incorrect results for some rows (about 0.5-1% of rows in randomized testing), including rows with a single byte which produced no tokens. This path now locates the first words of each row with a warp-per-row kernel and shares the same tokenizer. New tests verify the results against a host implementation of the tokenizer for several `max_words_per_row` values on full and sliced input.
**Vocabulary lookups**
Profiling showed about half of the deferred kernel time was spent in vocabulary map probes, most of which are for prefixes not in the vocabulary.
- The map keys now include the hash of each entry, so non-matching slots are rejected without loading and comparing the strings.
- The maps use a lower load factor (capacity of 4x the number of entries) to reduce the number of slots checked for each lookup.
**Results** (RTX A6000, realistic benchmark from NVIDIA#24467)
| num_rows | num_words | max_words | before | after | speedup | peak memory |
|---|---|---|---|---|---|---|
| 262144 | 32 | 0 | 9.13 ms | 4.53 ms | 2.02x | 242 → 136 MiB |
| 2097152 | 32 | 0 | 72.12 ms | 35.32 ms | 2.04x | 1.89 → 1.06 GiB |
| 32768 | 2048 | 0 | 77.14 ms | 32.11 ms | 2.40x | 1.90 → 1.06 GiB |
| 262144 | 256 | 20 | 12.98 ms | 3.31 ms | 3.92x | 1.46 GiB → 129 MiB |
| 32768 | 2048 | 20 | 7.61 ms | 0.63 ms | 12.08x | 1.41 GiB → 55 MiB |
| 2097152 | 32 | 40 | 88.09 ms | 36.05 ms | 2.44x | 2.34 → 1.08 GiB |
| 32768 | 2048 | 40 | 8.55 ms | 0.94 ms | 9.07x | 1.42 GiB → 66 MiB |
`num_words` is words per row (~5.6 bytes/word: 32 ≈ 180 B, 256 ≈ 1.4 KB, 2048 ≈ 11.5 KB rows).
`max_words=0` means no limit (all words are tokenized).
Every benchmark configuration is faster (1.92x-12.08x) and uses less memory.
Authors:
- David Wendt (https://github.com/davidwendt)
Approvers:
- Yunsong Wang (https://github.com/PointKernel)
- Basit Ayantunde (https://github.com/lamarrr)
URL: NVIDIA#244681 parent 45e91d0 commit 5fe3dba
2 files changed
Lines changed: 1034 additions & 546 deletions
0 commit comments