russell@unturf.com
81bef63b2e
transformers-0002: MOAD-0004 (CWE-312) regnet convert script logs HF_TOKEN verbatim transformers-0003: MOAD-0001 (CWE-407) convert_tokens_to_string O(T×S) list scan - marian, m2m_100, speech_to_text, siglip, gpt_sw3 all affected - all_special_tokens is list[str]; fix: cache set() before loop; 5x speedup vllm-0002: MOAD-0001 (CWE-407) Grok2Tokenizer O(N×V) dict.values() scan - decode() and convert_ids_to_tokens() use .values() view per token - sibling Mistral tokenizer already uses frozenset correctly - fix: add _special_token_ids frozenset at __init__; 10x speedup at N=2048, V=200 MOADs 0002/0003/0005 CLEAN for both repos
96 lines
3.7 KiB
Java
96 lines
3.7 KiB
Java
import java.util.*;
|
||
|
||
/**
|
||
* vllm-0002: Grok2Tokenizer decode O(N×S) dict.values() linear scan
|
||
*
|
||
* Models the two decode implementations:
|
||
* DEFECTIVE: for each token, scan all V dict values → O(N×V)
|
||
* PATCHED: use a frozenset (HashSet) of special IDs → O(N×1) = O(N)
|
||
*
|
||
* Mirrors vllm/vllm/tokenizers/grok2.py lines 354-363 and 376-382.
|
||
*/
|
||
public class Vllm0002Grok2SpecialTokenTest {
|
||
|
||
// --- Defective implementation ---
|
||
static List<Integer> decodeDefective(List<Integer> ids,
|
||
Map<String, Integer> specialTokens,
|
||
boolean skipSpecial) {
|
||
if (!skipSpecial) return new ArrayList<>(ids);
|
||
List<Integer> result = new ArrayList<>();
|
||
for (int tokenId : ids) {
|
||
// O(V) dict.values() scan per token — mirrors Python dict.values() membership
|
||
if (!specialTokens.containsValue(tokenId)) {
|
||
result.add(tokenId);
|
||
}
|
||
}
|
||
return result;
|
||
}
|
||
|
||
// --- Patched implementation ---
|
||
static List<Integer> decodePatched(List<Integer> ids,
|
||
Set<Integer> specialTokenIds,
|
||
boolean skipSpecial) {
|
||
if (!skipSpecial) return new ArrayList<>(ids);
|
||
List<Integer> result = new ArrayList<>();
|
||
for (int tokenId : ids) {
|
||
// O(1) frozenset/HashSet lookup
|
||
if (!specialTokenIds.contains(tokenId)) {
|
||
result.add(tokenId);
|
||
}
|
||
}
|
||
return result;
|
||
}
|
||
|
||
public static void main(String[] args) {
|
||
// Build a special token dict simulating Grok-2 (200 special tokens)
|
||
int V = 200;
|
||
Map<String, Integer> specialTokens = new LinkedHashMap<>();
|
||
// Regular vocab: 0..99999; special tokens start at 100000
|
||
for (int i = 0; i < V; i++) {
|
||
specialTokens.put("<special_" + i + ">", 100000 + i);
|
||
}
|
||
Set<Integer> specialTokenIds = new HashSet<>(specialTokens.values());
|
||
|
||
// Build a realistic output sequence: N=2048 tokens, 10 special, rest regular
|
||
int N = 2048;
|
||
List<Integer> ids = new ArrayList<>(N);
|
||
for (int i = 0; i < N; i++) {
|
||
if (i % 200 == 0) {
|
||
ids.add(100000); // special token <special_0>
|
||
} else {
|
||
ids.add(i % 50000); // regular token
|
||
}
|
||
}
|
||
|
||
// Correctness check
|
||
List<Integer> defectResult = decodeDefective(ids, specialTokens, true);
|
||
List<Integer> patchResult = decodePatched(ids, specialTokenIds, true);
|
||
assert defectResult.equals(patchResult)
|
||
: "FAIL: defect and patched outputs differ";
|
||
|
||
// Timing comparison
|
||
int REPS = 500;
|
||
|
||
long t0 = System.nanoTime();
|
||
for (int r = 0; r < REPS; r++) {
|
||
decodeDefective(ids, specialTokens, true);
|
||
}
|
||
long defectNs = System.nanoTime() - t0;
|
||
|
||
long t1 = System.nanoTime();
|
||
for (int r = 0; r < REPS; r++) {
|
||
decodePatched(ids, specialTokenIds, true);
|
||
}
|
||
long patchNs = System.nanoTime() - t1;
|
||
|
||
double ratio = (double) defectNs / patchNs;
|
||
System.out.printf("vllm-0002 MOAD-0001 Grok2 decode O(N×S) dict.values() scan%n");
|
||
System.out.printf(" N=%d tokens, V=%d special tokens, %d reps%n", N, V, REPS);
|
||
System.out.printf(" Defective (dict.values scan): %,d ms%n", defectNs / 1_000_000);
|
||
System.out.printf(" Patched (frozenset lookup): %,d ms%n", patchNs / 1_000_000);
|
||
System.out.printf(" Speedup: %.1fx%n", ratio);
|
||
|
||
assert ratio > 5.0 : "FAIL: expected >5x speedup, got " + ratio;
|
||
System.out.println("PASS");
|
||
}
|
||
}
|