WebGPU가 브라우저의 경계를 넘어 네이티브 생태계로 확장되는 동안, 드미트리의 역할 역시 미묘하게 변하고 있었다. 그는 더 이상 Dawn의 코드 한 줄 한 줄을 책임지는 실무 엔지니어가 아니었다. 대신 그는 구글 내부와 외부의 수많은 팀들이 WebGPU를 성공적으로 도입할 수 있도록 돕는 기술 고문이자 전략가의 역할을 수행하고 있었다.
그의 새로운 일과 중 하나는, 구글의 또 다른 거대한 프로젝트, 텐서플로우(TensorFlow) 팀과의 정기적인 기술 회의에 참석하는 것이었다.
텐서플로우.
전 세계의 머신러닝 연구와 산업을 지배하는, 구글의 가장 중요한 오픈소스 프로젝트 중 하나.
이 거대한 코끼리가 마침내 웹이라는 새로운 초원을 주목하기 시작한 것이다.
회의실의 분위기는 사뭇 달랐다. 텐서플로우 팀의 엔지니어들은 Dawn 팀과는 다른 종류의 천재들이었다. 그들의 언어는 렌더 파이프라인이나 셰이더가 아니었다. 텐서(Tensor), 연산 그래프(Operation Graph), 그리고 신경망 레이어(Neural Network Layer)였다.
텐서플로우 팀의 리드 아키텍트, 샨카르가 입을 열었다.
“드미트리, 저희는 현재 TensorFlow.js라는 자바스크립트 라이브러리를 통해 웹에서 머신러닝을 지원하고 있습니다. 현재 백엔드는 WebGL을 사용하고 있죠.”
그는 화면에 텐서플로우의 아키텍처 다이어그램을 띄웠다.
“WebGL 백엔드는 훌륭하게 작동해왔습니다. 하지만 저희는 한계에 부딪혔습니다. 특히, 모바일넷(MobileNet) 같은 대규모 모델을 실행할 때, WebGL의 텍스처 기반 연산 방식은 비효율적이고, 복잡한 연산 그래프를 구현하기가 매우 까다롭습니다.”
그는 WebGL을 이용해 행렬 곱셈을 구현하는 코드를 보여주었다. 프래그먼트 셰이더를 이용해 텍스처의 픽셀 값으로 계산을 시뮬레이션하는, 복잡하고 기교에 가까운 코드였다.
“보시다시피, 이건 거의 예술에 가깝습니다. 하지만 유지보수하기는 끔찍하게 어렵죠.”
샨카르는 다음 슬라이드를 넘겼다. 그 슬라이드에는 단 하나의 단어가 크게 쓰여 있었다.
WebGPU
“저희의 목표는 명확합니다. 텐서플로우.js의 차세대 백엔드로 WebGPU를 도입하는 겁니다. WebGL의 편법적인 계산 대신, 컴퓨트 셰이더의 순수한 힘을 이용하고 싶습니다. 이를 통해 웹에서의 머신러닝 추론 성능을 한 차원 끌어올릴 수 있을 거라 믿습니다.”
이것은 Google Meet이나 Figma의 사례와는 차원이 다른 이야기였다.
수백만 명의 AI 개발자들이 사용하는, 세계 최대의 머신러닝 플랫폼이 그들의 운명을 WebGPU에 걸겠다고 선언한 것이다.
드미트리는 이 프로젝트의 엄청난 잠재력과 동시에, 그 안에 도사린 기술적 난관을 즉시 간파했다.
“매우 흥미로운 도전입니다, 샨카르. 하지만 텐서플로우의 수백 가지 연산(Operation)들을 모두 WGSL 컴퓨트 셰이더로 구현하는 것은 엄청난 작업이 될 겁니다. 특히 컨볼루션(Convolution)이나 풀링(Pooling) 같은 복잡한 연산들은 최적화하기가 매우 까다로울 텐데요.”
샨카르가 고개를 끄덕였다.
“알고 있습니다. 그래서 저희는 처음부터 모든 것을 새로 만들지 않을 겁니다.”
그는 놀라운 계획을 밝혔다.
“저희는 텐서플로우의 핵심 C++ 커널을 WebAssembly로 컴파일할 겁니다. 그리고 이 WASM 코드가 연산 그래프를 관리하고 스케줄링하는 역할을 맡게 할 겁니다. 그리고 실제 계산이 일어나는 가장 깊은 부분, 즉 각 연산의 개별 커널들만 WGSL로 작성하여 WebGPU 위에서 실행시키는 하이브리드 접근법을 생각하고 있습니다.”
그것은 유니티가 보여주었던, WebAssembly와 WebGPU의 협력 모델을 머신러닝의 영역에서 그대로 재현하는 것이었다.
그날 이후, 드미트리는 텐서플로우.js 팀의 가장 중요한 기술 자문이 되었다.
그는 텐서플로우 팀이 WGSL의 특성을 최대한 활용하여 행렬 곱셈 커널을 최적화할 수 있도록 도왔다. 메모리 레이아웃을 어떻게 구성해야 캐시 효율성을 높일 수 있는지, 워크그룹(Workgroup)의 크기를 어떻게 조절해야 다양한 GPU에서 최적의 성능을 낼 수 있는지에 대한 노하우를 아낌없이 전수했다.
이 협력은 드미트리에게도 새로운 배움의 기회였다. 그는 텐서플로우 팀을 통해, 머신러닝 워크로드의 독특한 특성을 이해하게 되었다. 수많은 작은 버퍼를 효율적으로 관리해야 하는 문제, 중간 계산 결과를 재사용하기 위한 파이프라인 배리어(Pipeline Barrier)의 중요성 등을 깨달았다. 이 경험들은 Dawn 엔진의 다음 버전에서 컴퓨트 셰이더의 성능을 개선하는 데 귀중한 밑거름이 되었다.
몇 달 후, 텐서플로우 팀은 마침내 WebGPU 백엔드의 첫 번째 프로토타입을 공개했다.
결과는 놀라웠다. 동일한 이미지 분류 모델을 실행했을 때, WebGPU 백엔드는 기존의 WebGL 백엔드보다 평균 3배 이상 빠른 추론 속도를 보여주었다. 특히 모델의 크기가 커질수록 그 격차는 더욱 벌어졌다.
드미트리는 그 결과를 보며, 뉴턴의 말을 떠올렸다.
‘내가 더 멀리 보았다면, 그것은 거인들의 어깨 위에 서 있었기 때문이다.’
WebGPU는 이제 텐서플로우라는 거인에게 자신의 어깨를 빌려주고 있었다. 그리고 그 거인은 이제 웹이라는 새로운 세상을 향해, 이전보다 훨씬 더 멀리, 더 넓게 내다볼 수 있게 된 것이다. 드미트리는 그 거대한 변화의 흐름 속에서, 자신이 단단한 디딤돌 하나를 놓았다는 사실에 조용한 자부심을 느꼈다.
