rednote-machine-learning/RedKnot
Python1 / 0
RedKnot은 헤드 분류 기반 KV 재사용과 SegPagedAttention으로 장문 컨텍스트 추론의 FLOPs를 50–70% 줄이고 TTFT를 1.35×–2.2× 개선한다.
TL;DR
RedKnot은 SGLang 기반의 장문 컨텍스트 LLM 서빙 가속 라이브러리로, 헤드 단위 분류를 통해 각 헤드에 대해 서로 다른 KV 가시 창과 재사용 전략을 적용하여 불필요한 prefill 연산을 줄이는 설계이다. 오프라인 세그먼트 KV 캐시와 RoPE 재배치를 결합하여 저장된 KV를 수치적으로 정렬한 뒤 필요한 토큰만 선택적으로 재계산하고, 토큰 중요도 기반의 Elastic Sparsity로 FFN 연산을 스킵하여 전체 FLOPs를 약 50–70% 절감한다. README에 제시된 벤치마크는 TTFT 1.35×–2.2×의 속도 향상과 일부 시나리오에서 베이스라인과 동등하거나 더 나은 F1 정확도를 보고했으며, transformers 5.x가 필요한 모델군은 별도 가상환경을 통해 실행해야 한다. 이러한 설계는 긴 문맥에서의 추론 비용과 응답 지연을 낮추는 데 기여하나, 일부 MoE 설정에서는 정확도 저하가 발생할 수 있고 대규모 재현에는 고성능 GPU 자원이 요구된다. arXiv 논문과 통합된 벤치마크 스크립트가 있어 실험 재현과 비교 평가가 가능한 구현을 제공한다.
핵심 포인트
- 헤드 분류를 통해 헤드별로 서로 다른 KV 가시 창과 재사용 전략을 적용하는 점이 기존의 모든 헤드에 동일한 전체 KV를 제공하는 접근과 다르다. 이로 인해 불필요한 KV 로딩과 전체 prefill 계산이 회피되어 메모리 사용과 FLOPs가 감소한다. 분류 결과를 JSON으로 설정·로딩할 수 있게 하여 운영 환경에서 유연한 튜닝이 가능하다.
- 세그먼트 단위의 오프라인 KV 캐시와 RoPE 재배치를 결합하여 저장된 KV를 수치적으로 정렬한 뒤 선택적 재계산만 수행하는 점이 핵심적 차별화 요소이다. 이 구조는 긴 문서 반복 조회 시 전체 프리필을 매번 다시 수행하지 않아도 되는 운영적 이점을 제공한다. 저장된 KV의 형식 일관성 유지로 정확도 손실을 최소화했다는 점이 강조된다.
- 토큰 기여도 기반의 Elastic Sparsity로 FFN 계산을 건너뛰는 기법은 토큰 단위로 연산을 줄여 전체 연산량을 더 낮추는 추가 경로를 제공한다. 이 기법은 특히 긴 컨텍스트에서 효과가 커서 SegPagedAttention과 결합 시 상호보완적인 이득을 제공한다. 결과적으로 FLOPs 절감과 TTFT 가속이 함께 달성되도록 설계되었다.
- 헤드별 KV 저장 정책을 적용하여 각 (layer, kv_head)에 대해 global, local, retrieval, dense 네 가지 클래스 중 하나를 할당하고 해당 클래스에 맞는 KV 조회·재사용 경로를 운영한다. 이 설계는 전체 KV를 모든 헤드에 복제하지 않고 필요한 데이터만 로드하여 메모리와 계산을 절감한다. 클래스 구성은 JSON 설정과 프로파일링 결과를 바탕으로 동작한다.
벤치마크
| 벤치마크 | 지표 | 값 | 비교 |
|---|---|---|---|
| Qwen3-32B HotpotQA 16K | TTFT speedup / FLOPs saved | 1.39x / 69.2% | vs dense FlashAttention-2 baseline |
| Qwen3-32B HotpotQA 32K | TTFT speedup / FLOPs saved | 1.93x / 72.2% | vs dense FlashAttention-2 baseline |
| Qwen3.5-35B-A3B LongBench 16K | TTFT speedup / compute saved | 1.87x / 46.4% | vs honest dense baseline |
| Mistral-7B-Instruct HotpotQA 16K | TTFT speedup / FLOPs saved | 1.35x / 51.5% | vs dense baseline |
659
STARS
211
FORKS
+249
TRENDING
1
조회수
관련 토론
아직 관련 토론이 없습니다.
댓글
댓글을 작성하려면 로그인이 필요합니다.