Skip to content

Commit 5fe3dba

Browse files
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#24468
1 parent 45e91d0 commit 5fe3dba

2 files changed

Lines changed: 1034 additions & 546 deletions

File tree

0 commit comments

Comments
 (0)