도구 호출 대기와 캐시를 함께 고려하는 에이전트 추론 스케줄러 TAO

배경과 문제

에이전트 세션은 LLM 추론 요청과 외부 도구 호출을 번갈아 수행한다. 동작을 설명하는 예로, 코드 에이전트가 컴파일러를 부르거나 검색 도구를 실행하는 동안 GPU는 그 세션을 위해 일하지 않아도 된다. 이 글은 GPU가 여러 노드에 모여 있고 도구는 분산 CPU에서 실행되는 'Clustered-GPUs with Distributed-CPUs(CGDC)' 시나리오를 정의한다. 도구 실행 시간이 길고 변동적이기 때문에, 언제 어떤 세션이 다시 GPU를 필요로 하는지가 전체 응답성과 처리량을 좌우한다.

문제는 두 갈래로 나뉜다. 첫째, 세션의 다음 추론 요청에 가치를 매기는 평가 단계에서 논문이 비교하는 기존 방식은 도구 호출 지연을 무시하거나 캐시 재사용·토큰 수의 대리 변수로만 사용한다. 둘째, 어떤 세션을 어느 노드에 넣을지 결정하는 단계에서 세션별 휴리스틱을 적용해 근시안적이고 차선의 배치를 만든다. 한 세션을 배치하면 다른 세션들의 가능한 조합을 막을 수 있는데, 기존 방식은 이 경쟁을 공동으로 고려하지 않는다.

기존 접근의 한계

논문이 분류한 캐시 인지 시스템은 KV 프리픽스를 재사용하고 추론 부하를 분산하지만 도구 타이밍을 평가에 넣지 않는다. 워크플로 인지 시스템은 미래 호출을 예측하지만 도구 복귀 시점을 직접 다루지 않는다. 지연 인지 스케줄러는 남은 작업 추정으로 우선순위를 정하지만 도구 실행 중인 세션과 곧 돌아올 세션을 구분하지 못한다. 일부 연구는 도구 지속 시간을 예측해 KV 캐시 보존 시간을 정하거나 토큰 수요를 할인하지만, 이런 특징은 고정 점수 규칙이나 수동 임계값에 들어간다.

의사 결정 측면에서 llm-d와 SGLang 같은 시스템은 세션별 점수와 라우팅 규칙으로 추론 기회를 개별 할당한다. 그러나 세션들은 같은 노드 용량을 두고 경쟁하므로, 한 세션을 넣으면 더 가치 있는 다른 조합이 배제될 수 있다. 휴리스틱 선택과 배치 규칙도 차선책을 낳는다. 이 글은 평가와 결정을 분리하고, 결정 단계에서 대기 풀의 모든 세션을 공동으로 최적화하는 방향을 제안한다.

연구방법을 단계별로

TAO는 먼저 각 세션-노드 쌍에 대해 세 가지 값으로 이루어진 삼중항을 만든다. 첫째, 시간 감쇠 캐시 크레딧 ki,b는 해당 노드에서 재사용할 가능성이 있는 KV 프리픽스의 관측량을 나타내며, 관측이 오래될수록 반감기 10초로 지수 감쇠한다. 둘째, 이진 준비도 ri는 도구가 곧 끝나 1초 안에 추론이 필요할지를 나타낸다. 셋째, 토큰 수 ti는 누적 컨텍스트 크기로, 앞의 두 값이 같을 때 더 큰 컨텍스트를 우선한다.

준비도 예측은 여러 신호를 결합한다. 도구 완료가 관측되면 즉시 준비 상태가 된다. 도구 인터페이스가 노출하는 진행 이벤트나 반환 데드라인, sleep 명령 분석, 도구별·전역 역사 평균을 사용한다. 예측 시각이 30초 이상 오래되면 무효화하고, 예측이 지났는데도 완료가 없으면 일시적 probe 기회만 준다. 마지막 폴백은 5초다. 이 신호는 도구 실행 중인 세션과 곧 돌아올 세션을 구분해 GPU admission을 준비시킨다.

결정 단계는 동적 계획법(DP)으로 대기 풀의 모든 세션을 공동 선택하고 노드에 배치한다. 용량을 양자 ∆로 이산화하고, 각 후보 세션을 건너뛰거나 적격 노드에 할당하는 선택을 누적해 최적 목적을 찾는다. 목적은 삼중항의 사전식 순서로, 캐시 크레딧이 가장 우선하고 그다음 준비도, 마지막으로 토큰 수다. 계산은 비동기 작업자에서 수행되고, 적용 전에 실제 상태와 서명을 재검증해 변경된 세션이나 용량 초과를 걸러낸다. 탐색 상태가 예산을 넘으면 best-fit-decreasing(BFD) 휴리스틱으로 전환한다.

결과와 숫자의 의미

주 평가는 DeepSeek-V4 Flash를 2노드×노드당 4 H100, Qwen3-32B를 4노드×노드당 2 H100에서 실행했다. Dynamo, vLLM RR, vLLM Cache-aware, ThunderAgent 네 방식과 비교했다. DeepSeek SWE-Lite의 동시 실행 작업자 64개에서 작업 지연은 vLLM Cache-aware의 253.228초에서 185.410초로 26.8% 줄었다. 첫 토큰 대기는 Dynamo의 0.351초에서 0.224초로 36.3%, 생성 처리량은 vLLM Cache-aware의 초당 1561.307토큰에서 2021.523토큰으로 29.5% 개선됐다. 지표별 최상 기준선이 다르며, 작업 지연에는 모델과 도구 시간이 함께 들어간다. 생성 처리량은 과제 정답률과 다른 지표다.

모든 설정에서 일관되게 좋은 것은 아니다. Terminal-Bench 16/64 workers에서 디코드 처리량은 각각 8.8%, 6.2% 낮고, Qwen SWE-Lite 96 workers에서 5.1% 낮다. Qwen SWE-verified 128 workers에서는 지연이 ThunderAgent보다 2.5% 높다. GLM-4.7 단일 8 H100 노드의 별도 6개 설정에서는 SWE-verified 96 workers에서 TTFT가 38.4% 개선되지만, 이는 주 평가의 36.3%와 다른 실험이다. 표본 수에도 불일치가 있다. 본문은 SWE-verified를 'all 500 tasks'라 하지만 표2와 캡션은 300 tasks/run으로 표기한다.

세부 효과는 별도의 DeepSeek 32-worker SWE-Lite 재생 기록에서 분석했다. 캐시 인지 평가의 적중률 +0.19퍼센트포인트와 대기 시간 −4.5%, 감쇠까지 적용한 적중률 +0.31퍼센트포인트와 생성 처리량 +1.68%는 모두 ThunderAgent 기준과의 비교다. 두 향상분을 순서대로 더하는 값이 아니다. 준비도를 켠 대응 실험에서는 처리량이 4.03% 늘고 도구 종료부터 첫 토큰까지 시간이 10.88% 줄었다. 완료 도구 10,967건의 최종 시간 예측 오차는 평균 절대값 0.410초였고 82.9%가 0.5초 이내였다. 공동 배치의 목적값은 복원 가능한 35개 스냅샷 중 6개에서 단순 휴리스틱보다 좋았다. 후보가 여럿인 23개만 보면 6개, 즉 26.1%다.

구체적인 예시

32-worker DeepSeek 재현 스냅샷에서 DP는 캐시 크레딧 16,791토큰을 보존해 short-first 휴리스틱의 13,624토큰보다 23.2% 더 많았다. 다른 스냅샷에서는 DP가 긴 컨텍스트 세션 177을 포함해 7개 세션을 수용했지만, 휴리스틱은 짧은 세션으로 용량을 채우고 177을 미뤄 6개만 수용했다. 이는 개별 세션의 즉각적 가치가 아니라 조합의 총 가치를 보는 효과를 보여준다.

도구 지연 실험은 고정된 SWE-Lite 기록을 재생하면서 모델 응답은 실제로 생성하고, 도구 복귀에 1~5초·2~10초·4~20초 지연을 추가했다. 이전 구성 대비 전체 완료 시간은 각각 20.72%·16.42%·19.28% 줄었다. 그러나 준비도뿐 아니라 캐시 감쇠와 공유 토큰 갱신도 함께 바뀌어 준비도 하나의 효과로 볼 수 없다. 같은 상태에서 준비도만 제거했을 때 배치가 달라진 비율은 11/668, 53/1,345, 200/2,521로 각각 1.65%·3.94%·7.93%였다. 재생 기록의 정확한 도구 시간을 알려 준 비교와의 완료 시간 차이는 최대 2.24%였다. 각 조건은 한 차례의 탐색적 실행이므로 실제 모든 도구에서 예측이 정확하다는 증거는 아니다.

한계와 오해하기 쉬운 점

TAO는 도구 상태를 관찰해 GPU 추론을 언제 어느 노드에 넣을지 정한다. CPU나 메모리를 쓰는 도구 작업 자체를 스케줄하거나 GPU KV 블록을 직접 보존·이동하지 않는다. 물리적 캐시의 할당과 제거는 기존 추론 엔진이 담당한다. 따라서 캐시 크레딧의 증가가 같은 양의 실제 메모리 보존을 보장하는 것은 아니다. 이 프리프린트의 결과를 독립 재현하거나 동료 심사 완료를 확인한 것도 아니다.

동적 계획법의 exact 보장은 고정 스냅샷의 이산화된 용량과 목적에만 적용된다. 반올림으로 실제 가능한 배치가 제외될 수 있고, 비동기 계산 후 실제 상태를 다시 검증한다. 탐색 상태가 제한을 넘으면 BFD로 전환하므로 실제 모든 순간에 전역 최적 지연을 보장하지 않는다. 준비도는 관측 완료·명시적 시간 힌트·sleep 분석·역사 평균과 여유 시간을 이용하지만 도구 완료를 정확히 아는 것은 아니다. 미래 1초 horizon과 오래된 예측 재검증을 쓰지만 예측 실패 시 폴백이 작동한다.

표 9에서 Qwen 160-worker 조건의 완료된 계획 계산 시간은 평균 1.59ms, 95번째 백분위는 3.68ms였다. 상태 준비와 완료된 계산의 합계 약 4.03초는 실행 시간의 0.11236%지만, 한도를 넘긴 실패 탐색과 단순 배치 방식의 실행 비용은 빠져 있다. 이는 전체 CPU 비용이 아니다. 계획 제출 수를 분모로 한 단순 방식 전환 비율은 DeepSeek 64에서 5.24%, Qwen 96·128·160에서 3.06%·6.74%·8.98%였다. 실험의 반복 시드와 신뢰구간이 없어 관측된 차이 모두가 통계적으로 유의하다고 단정할 수 없다.

실제로 어떻게 활용할까

TAO는 기존 추론 엔진 위에서 세션 평가와 공동 배치를 담당하는 구조다. TokenCake 엔진과 함께 쓴 16-worker 실험에서는 같은 엔진의 일관된 해시 라우팅 기준보다 생성 처리량이 초당 785.159토큰에서 886.772토큰으로 12.9% 높았다. 다만 이 실행에는 프로그램 중단·재개·이동 기록이 없고 최종 오프로드 기록도 없어, 개선을 DP 선택이나 전송 중첩 효과에 귀속할 수 없다. 이 결과가 뒷받침하는 범위는 해당 부하에서의 결합 실행 가능성이다.

낮은 동시성도 별도로 확인했다. 단일 8 H100 노드에서 작업자를 4~64개로 바꾼 SWE-Lite 300과제 실행 중, 4~8개 조건의 평균 작업 지연은 최상 기준선 대비 2.3% 이내였고 처리량은 최상값의 92% 이상이었다. 첫 토큰 대기는 15~20ms 더 길었다. 저자들은 이를 평가 범위에서 기본 스케줄러로 고려할 근거로 보지만, 모든 초기 서비스에 손실 없이 도입할 수 있다는 보장은 아니다. 실제 선택에는 자기 도구의 시간 분포와 캐시 재사용, 노드 규모에서 다시 비교할 필요가 있다.