[transformers] bad_words_ids를 100배 빠르게: 다중 토큰 시퀀스를 prefix 길이별로 묶어 한 번에 매칭하기

PR 링크: huggingface/transformers#50532 상태: Merged | 변경: +40 / -15

들어가며

Hugging Face transformers의 generate()에는 특정 단어를 금지하는 bad_words_ids와 특정 시퀀스에 bias를 더하는 sequence_bias 옵션이 있습니다. 둘 다 내부적으로는 SequenceBiasLogitsProcessor가 처리하며, NoBadWordsLogitsProcessor는 이를 상속해 bias를 -inf로 주는 형태입니다.

문제는 이 processor가 매 generation step마다 multi-token 시퀀스 전체를 Python 루프로 순회했다는 점입니다. 각 반복마다 device 위에 작은 tensor 두 개를 새로 만들고, 비교 연산과 torch.where를 따로 실행합니다. 즉 비용이 시퀀스 개수에 선형으로 비례합니다.

문서가 권장하는 방식대로 단어를 tokenize해서 blocklist를 만들면 대부분 2~4 토큰짜리 엔트리가 됩니다. CUDA에서 시퀀스 하나당 step마다 약 110us가 들었고, 금지어 100개면 토큰 하나 생성할 때마다 약 17ms가 추가되었습니다(관련 이슈 #39512). 모델 forward보다 logits processor가 더 느려지는 상황이 생긴 것이죠.

코드 분석

logits_process.py — __call__의 핫 루프

Before:

for sequence_ids, sequence_bias in self.sequence_bias.items():
    if len(sequence_ids) == 1:  # the sequence is of length 1, already applied
        continue
    prefix_length = len(sequence_ids) - 1
    if prefix_length > input_ids.shape[1]:
        continue
    last_token = sequence_ids[-1]
    matching_rows = torch.eq(
        input_ids[:, -prefix_length:],
        torch.tensor(sequence_ids[:-1], dtype=input_ids.dtype, device=input_ids.device),
    ).prod(dim=1)
    bias[:, last_token] += torch.where(
        matching_rows.bool(),
        torch.tensor(sequence_bias, device=input_ids.device),
        torch.tensor(0.0, device=input_ids.device),
    )

반복마다 torch.tensor(...)가 세 번 호출됩니다. CUDA에서는 각각이 host→device 복사를 유발하고, eq, prod, where, 인덱싱된 +=까지 여러 개의 작은 kernel launch가 이어집니다. 실제 연산량은 미미하지만 launch overhead와 Python 오버헤드가 시퀀스 수만큼 쌓입니다.

After:

for prefixes, last_tokens, sequence_biases in self.multi_token_biases:
    prefix_length = prefixes.shape[1]
    if prefix_length > input_ids.shape[1]:  # the prefix is longer than the context, ignore
        continue
    matching_rows = (input_ids[:, None, -prefix_length:] == prefixes).all(dim=-1)
    bias.index_add_(1, last_tokens, torch.where(matching_rows, sequence_biases, 0.0).to(bias.dtype))

이제 루프는 시퀀스 단위가 아니라 prefix 길이 그룹 단위로 돕니다. 단어 길이가 1~4 토큰이라면 루프는 최대 3번뿐입니다.

핵심은 broadcasting입니다.

• input_ids[:, None, -prefix_length:]의 shape: (batch, 1, L)
• prefixes의 shape: (num_seqs, L)
• 비교 결과: (batch, num_seqs, L) → .all(dim=-1) → (batch, num_seqs)

한 번의 비교로 모든 batch row × 모든 시퀀스의 매칭 여부를 계산합니다. 기존의 .prod(dim=1) + .bool() 조합 대신 .all()을 써서 의도도 더 명확해졌습니다.

bias 누적에는 index_add_를 사용합니다. 여기가 미묘한 부분인데, 서로 다른 시퀀스가 같은 마지막 토큰으로 끝날 수 있습니다(예: (1, 2)와 (3, 1, 2)). bias[:, last_tokens] += values 같은 advanced indexing 할당은 중복 인덱스가 있을 때 값이 누적되지 않고 하나만 반영됩니다. index_add_는 중복 인덱스를 올바르게 합산하므로 기존 루프의 의미를 그대로 보존합니다.

_prepare_bias_variables — 전처리를 첫 호출로 이동

sequences_by_prefix_length = {}
for sequence_ids, bias in self.sequence_bias.items():
    if len(sequence_ids) > 1:
        sequences_by_prefix_length.setdefault(len(sequence_ids) - 1, []).append((sequence_ids, bias))
self.multi_token_biases = [
    (
        torch.tensor([sequence_ids[:-1] for sequence_ids, _ in group], device=scores.device),
        torch.tensor([sequence_ids[-1] for sequence_ids, _ in group], device=scores.device),
        torch.tensor([bias for _, bias in group], dtype=torch.float, device=scores.device),
    )
    for group in sequences_by_prefix_length.values()
]

이 메서드는 원래부터 vocab size 추론을 위해 첫 호출 시 한 번 실행되던 곳입니다. 여기에 그룹화와 tensor 생성을 얹었기 때문에 device로의 tensor 업로드가 매 step이 아니라 단 한 번만 일어납니다. 같은 prefix 길이끼리 묶었기 때문에 prefixes는 ragged가 아닌 정사각 2D tensor가 되어 padding이나 masking 없이 깔끔하게 broadcast할 수 있습니다.

테스트

sequence_bias = {(1, 2): -1.0, (3, 1, 2): -2.0, (0, 1, 2): -4.0, (1, 4): 3.0, (5, 5, 5, 5, 5, 1): 10.0}

같은 마지막 토큰 2를 공유하는 길이 2·3 시퀀스(합산 검증), context보다 긴 prefix(무시 검증), 그리고 두 번째 호출에서 캐시된 tensor 재사용까지 검증합니다. 첫 row [0, 1, 3, 1]은 (1,2)와 (3,1,2)가 매칭되어 토큰 2에 -3.0이, 두 번째 row [2, 2, 0, 1]은 (1,2)와 (0,1,2)가 매칭되어 -5.0이 됩니다.

왜 이게 좋은가

PR 작성자가 측정한 수치(vocab 32k, context 512)는 인상적입니다.

설정	CUDA before	CUDA after	CPU before	CPU after
10 words	940 us	190 us	276 us	79 us
200 words	17.3 ms	0.18 ms	6.05 ms	0.10 ms
generate() 100 bad words (tok/s)	50	242	40	50

200단어 기준 CUDA에서 약 96배, 실제 generate() 처리량은 약 4.8배 향상되었습니다. 이후 비용은 시퀀스 수가 아니라 prefix 길이 종류 수에 비례하므로 사실상 상수에 가깝습니다.

정확성 검증도 꼼꼼합니다. 약 16k개의 랜덤 케이스(-inf/inf/유한 bias, batch 0~9, int32/int64, fp16/bf16 등)에서 old/new 구현을 비교했고, -inf·inf·NaN 위치는 항상 일치했습니다. 같은 토큰에 여러 bias가 합산될 때만 덧셈 순서 차이로 rounding 수준의 차이가 있을 뿐입니다.

일반적 교훈

1. GPU에서 "작은 연산의 반복"은 독이다. 연산량이 아니라 kernel launch와 host→device 전송이 병목이 됩니다. 루프를 broadcast 한 번으로 바꾸는 것이 핵심입니다.
2. 불변 데이터의 tensor화는 hot path 밖으로. 매 step 동일한 상수를 torch.tensor()로 다시 만들지 말고 lazy-init 시점에 캐시하세요.
3. ragged 데이터는 길이별로 bucket. 가변 길이 시퀀스를 padding하는 대신 길이별로 그룹화하면 각 그룹은 dense tensor로 다룰 수 있습니다.
4. scatter 시 중복 인덱스 의미를 확인하라. 벡터화 과정에서 x[idx] += v는 중복 인덱스를 누적하지 않습니다. index_add_나 scatter_add_를 써야 루프 버전과 의미가 같아집니다.

참고 자료

• torch.Tensor.index_add_ — 중복 인덱스 bias 누적에 사용된 API
• torch.where — 매칭 여부에 따라 bias를 선택하는 데 사용
• torch.all — prefix 전체 일치 판정에 사용
• PyTorch Broadcasting semantics — 한 번의 비교로 모든 시퀀스를 매칭하는 원리
• Transformers Generation utilities (Logits processors) — SequenceBiasLogitsProcessor, NoBadWordsLogitsProcessor 문서

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.