Netflix: практическое применение Vector API, которое спасло CPU / Хабр

Иногда одна незаметная фича может сжигать гигантский объём ресурсов. В Netflix именно так и вышло: скоринг серендипности в Ranker оказался дорогой горячей точкой, а попытка слегка его ускорить в итоге привела к большой инженерной перестройке — от батчинга до SIMD через JDK Vector API. В новом переводе от команды Spring АйО разберем, как SIMD AVX инструкции на практике позволили снизить потребление CPU.

Ranker — один из крупнейших и самых сложных сервисов Netflix. Помимо прочего, он формирует персонализированные ряды, которые вы видите на главной странице Netflix, и работает в гигантском масштабе. Когда мы изучили профили CPU для этого сервиса, одна особенность постоянно бросалась в глаза: скоринг «серендипности» (serendipity scoring) видео — логика, которая отвечает на простой вопрос:

«Насколько этот новый тайтл отличается от того, что вы смотрели до сих пор, и при этом релевантен лично для вас?»

Одна эта функция потребляла около 7,5% общего CPU на каждом узле, где работал сервис. То, что начиналось как простая идея — «просто пакетировать (batch) фичу скоринга видео», — превратилось в более глубокое путешествие по оптимизации. По пути мы внедрили батчинг, переосмыслили архитектуру раскладки данных в памяти и попробовали разные библиотеки для выполнения вычислительных ядер скоринга.

Проблема: «горячая точка» в Ranker

В общих чертах скоринг серендипности устроен так: кандидатный тайтл и каждый элемент истории просмотров пользователя представляются эмбеддингами в векторном пространстве. Для каждого кандидата мы вычисляем его сходство с эмбеддингами истории, находим максимальное значение сходства и преобразуем его в оценку «новизны». Эта оценка становится входной фичей для последующей логики рекомендаций.

Изначальная реализация была простой, но дорогой. Для каждого кандидата мы извлекаем его эмбеддинг, проходим циклом по истории, вычисляя косинусное сходство для каждой пары по отдельности, и отслеживаем максимальное значение.

Комментарий от Михаила Поливаха

Я как-то делал доклад, который как раз детально разбирает некоторые нюансы векторных БД. Я просто буквально пару слов скажу, чтобы люди не глубоко знакомые с Векторными БД поняли, что происходит

В векторных БД данные хранятся в виде векторов (алгебраически это просто массив чисел с плавающей точкой). И поиск в векторной БД, которая хранит эти самые векторы, почти всегда используется с одной целью, а именно ответить на вопрос:

Насколько вот этот вот объект А релевантен для вот этого вот объекта В?

Это часто нужно как раз в поисковых системах, например если представить себе маркетплейс, то там задача отобразить именно те товары пользователю, которые, как предполагается, для него релевантны, т.е. которые он скорее всего купит.

Вот эту вот "релевантность" определяют расстоянием между векторами (т.е. расстояниями между массивами чисел), для этого есть разные способы, все со своими нюансами, но чаще всего, гораздо чаще, используют косинусное расстояние между векторами (что кстати довольно CPU intensive), что и использовал Netflix.

Это легко понять и поддерживать, но в масштабе Ranker приводит к значительному последовательному объёму вычислений, повторяющимся обращениям за эмбеддингами, разрозненному доступу к памяти и слабой локальности кэша. Профилирование это подтвердило.

Флеймграф показал это предельно наглядно: одной из главных «горячих точек» сервиса оказались Java-вычисления скалярных произведений внутри энкодера серендипности. С точки зрения алгоритма эта горячая точка представляла собой вложенную структуру циклов M кандидатов × N элементов истории, где для каждой пары вычислялось своё косинусное сходство, то есть O(M×N) отдельных операций скалярного произведения.

Решение

Изначальная реализация: цикл косинусного сходства для одного видео

В упрощённом виде код выглядел так:

for (Video candidate : candidates) {
  Vector c = embedding(candidate); // D-dimensional
  double maxSim = -1.0;

  for (Video h : history) {
    Vector v = embedding(h); // D-dimensional
    double sim = cosine(c, v); // dot(c, v) / (||c|| * ||v||)
    maxSim = Math.max(maxSim, sim);
  }

  double serendipity = 1.0 - maxSim;
  emitFeature(candidate, serendipity);
}

Вложенный цикл for с O(M×N) отдельными скалярными произведениями порождал и собственные накладные расходы. Один любопытный нюанс мы выяснили, проинструментировав формы трафика: большинство запросов (около 98%) приходились на одиночные видео, но оставшиеся 2% были большими пакетными (batch) запросами. Поскольку эти батчи были настолько крупными, общий объём обработанных видео в итоге распределялся примерно 50:50 между одиночными и пакетными задачами. Это сделало батчинг оправданным направлением, даже если он не улучшал медианный запрос.

Шаг 1: батчинг — от вложенных циклов к умножению матриц

Первая идея заключалась в том, чтобы перестать мыслить категориями «множество маленьких скалярных произведений» и вместо этого рассматривать работу как матричную операцию. То есть для пакетных кандидатов реализовать такую раскладку данных, чтобы распараллелить математику в рамках одной операции — умножения матриц. Пусть D �� размерность эмбеддинга:

Упаковать все эмбеддинги кандидатов в матрицу A размера M × D
Упаковать все эмбеддинги истории просмотров в матрицу B размера N × D
Нормализовать все строки до единичной длины.
Вычислить косинусные сходства как:
[ C = A × B^T ]; где C — матрица косинусных сходств размера M × N.

В псевдокоде:

// Build matrices
double[][] A = new double[M][D]; // candidates
double[][] B = new double[N][D]; // history

for (int i = 0; i < M; i++) {
  A[i] = embedding(candidates[i]).toArray();
}
for (int j = 0; j < N; j++) {
  B[j] = embedding(history[j]).toArray();
}

// Normalize rows to unit vectors
normalizeRows(A);
normalizeRows(B);

// Compute C = A * B^T
double[][] C = matmul(A, B);
C[i][j] = cosine(candidates[i], history[j])

// Derive serendipity
for (int i = 0; i < M; i++) {
  double maxSim = max(C[i][0..N-1]);
  double serendipity = 1.0 - maxSim;
  emitFeature(candidates[i], serendipity);
}

Это превращает M×N отдельных скалярных произведений в одно умножение матриц — именно под такие операции и «заточены» CPU и оптимизированные вычислительные ядра. Мы интегрировали это в существующий фреймворк, поддержав оба варианта: encode() для одиночных видео и batchEncode() для батчей, при этом сохранив обратную совместимость. На этом этапе казалось, что мы «закончились», но это было не так.

Шаг 2: когда одного батчинга недостаточно

Получив батчевую реализацию, мы запустили канареечные прогоны и увидели неожиданное: примерно 5% регрессию по производительности. Проблема была не в алгоритме — замена M×N отдельных скалярных произведений на умножение матриц математически корректна. Проблемой оказались накладные расходы, которые мы привнесли в первой реализации.

Наша первоначальная версия на каждом батче строила матрицы double[][] для кандидатов, истории и результатов. Эти крупные, но краткоживущие аллокации создавали давление на GC, а сама раскладка double[][] не является непрерывной в памяти, из-за чего возникали дополнительные переходы по указателям и ухудшалось поведение кэша.

Кроме того, первая версия умножения матриц на Java была прямолинейной скалярной реализацией, поэтому не могла задействовать SIMD. Иными словами, мы заплатили цену батчинга, но не получили вычислительной эффективности, к которой стремились.

Вывод был мгновенным: алгоритмические улучшения не имеют значения, если детали реализации — раскладка в памяти, стратегия аллокаций и вычислительное ядро — работают против вас. Это подготовило почву для следующего шага: сделать раскладку данных дружественной кэшу и устранить аллокации на каждый батч, прежде чем снова возвращаться к ядру умножения матриц.

Шаг 3: плоские буферы и повторное использование через ThreadLocal

Мы переработали раскладку данных так, чтобы она была дружественной кэшу и минимальной по аллокациям. Вместо double[m][n] мы перешли на плоские буферы double[] в порядке row-major. Это дало непрерывную память и предсказуемые паттерны доступа. Затем мы ввели ThreadLocal<BufferHolder>, который владеет переиспользуемыми буферами для кандидатов, истории и любого другого вспомогательного пространства. Буферы при необходимости растут, но никогда не уменьшаются — это избавляет от аллокаций на каждый запрос и при этом держит каждый поток изолированным (без конкуренции). Упрощённый набросок:

class BufferHolder {  
  double[] candidatesFlat = new double[0];  
  double[] historyFlat = new double[0];  

  double[] getCandidatesFlat(int required) {  
    if (candidatesFlat.length < required) {  
      candidatesFlat = new double[required];  
    }  
    return candidatesFlat;  
  }  
  
  double[] getHistoryFlat(int required) {  
    if (historyFlat.length < required) {  
      historyFlat = new double[required];  
    }  
    return historyFlat;  
  }  
}  

private static final ThreadLocal<BufferHolder> threadBuffers =  
    ThreadLocal.withInitial(BufferHolder::new);

Одно это изменение сделало батчевый путь гораздо более предсказуемым: меньше аллокаций, ниже давление на GC и лучшая локальность кэша.

Теперь оставался вопрос, на который мы изначально и думали, что отвечаем: как лучше всего выполнять умножение матриц?

Шаг 4: BLAS — отлично в тестах, но не в продакшене

Очевидным следующим шагом был BLAS (Basic Linear Algebra Subprograms).

Комментарий от Михаила Поливаха

На всякий случай поясню, т.к. тут парни этого не делают. BLAS – это спека такая, которая по сути определяет набор операций с векторами и матрицами. И эту спеку реализуют разные библиотеки на разных языках программирования.

Так получилось, что референсная реализация BLAS была написана на фортране, и, соответственно, большое количество библиотек на других языках программирования (не только на Java кстати) просто не парились и делегировали всё в нативную либу.

В отрыве от контекста микробенчмарки выглядели многообещающе. Но после интеграции в реальный путь пакетного скоринга выигрыша не получилось. Против нас работало несколько факторов:

Путь по умолчанию в netlib-java использовал F2J (Fortran-to-Java) BLAS, а не действительно нативную реализацию.
Даже с нативным BLAS мы платили накладные расходы за подготовку и переходы через JNI.
Раскладка row-major в Java не совпадает с ожиданиями column-major у многих процедур BLAS, что может требовать преобразования и временных буферов.
Эти дополнительные аллокации и копирования имели значение во всём конвейере, особенно рядом с работой TensorFlow по эмбеддингам.

BLAS всё же оказался полезным экспериментом — он прояснил, куда уходит время, но не стал тем «drop-in» улучшением, на которое мы рассчитывали. Нам нужно было решение, остающееся в чистой Java, подходящее под нашу архитектуру плоских буферов и при этом способное задействовать SIMD.

Шаг 5: на помощь приходит JDK Vector API

Коротко о JDK Vector API: JDK Vector API — это инкубируемая возможность, которая даёт переносимый способ выражать в Java операции с параллелизмом по данным — в духе «SIMD без интринсиков». Вы пишете код в терминах векторов и «линий» (lanes), а JIT отображает эти операции на наиболее подходящие SIMD-инструкции, доступные на хостовом CPU (SSE/AVX2/AVX-512), с запасным скалярным вариантом при необходимости. И, что для нас особенно важно, это чистая Java: никаких нативных зависимостей, никаких переходов через JNI и модель разработки, которая выглядит как обычный Java-код, а не платформенно-специфичный ассемблер или интринсики.

Это оказалось особенно удачным совпадением с нашим профилем нагрузки, потому что мы уже перенесли эмбеддинги в плоские, непрерывные буферы double[], а «горячий» цикл в основном состоял из огромного числа скалярных произведений. Финальным шагом стала замена BLAS на SIMD-реализацию на чистой Java с использованием JDK Vector API. К этому моменту у нас уже была правильная «форма» для высокой про��зводительности — батчинг, плоские буферы и повторное использование через ThreadLocal. Поэтому оставшаяся работа сводилась к тому, чтобы заменить вычислительное ядро, не добавляя накладных расходов JNI и не прибегая к платформенно-специфичному коду. Мы сделали это через небольшой фабричный слой. При загрузке класса MatMulFactory выбирает лучшую доступную реализацию:

Если доступен jdk.incubator.vector — использовать реализацию на Vector API.
В противном случае — откатываться на скалярную реализацию с тщательно оптимизированным, развернутым по циклу скалярным произведением (её реализовал мой коллега Патрик Стродерман; она вдохновлена паттернами, используемыми в Lucene).

В реализации на Vector API внутренний цикл вычисляет скалярное произведение, накапливая a * b в векторный аккумулятор с помощью fma() (fused multiply-add). DoubleVector.SPECIES_PREFERRED позволяет рантайму выбрать подходящую ширину «линий» для конкретной машины. Вот упрощённый набросок внутреннего цикла:

// Vector API path (simplified)  
for (int i = 0; i < M; i++) {  
  for (int j = 0; j < N; j++) {  
  
  DoubleVector acc = DoubleVector.zero(SPECIES);  
    int k = 0;  
    // SPECIES.length() (e.g. often 4 doubles on AVX2 and 8 doubles on AVX-512). 
    for (; k + SPECIES.length() <= D; k += SPECIES.length()) {  
      DoubleVector a = DoubleVector.fromArray(SPECIES, candidatesFlat, i*D + k);  
      DoubleVector b = DoubleVector.fromArray(SPECIES, historyFlat,   j*D + k);
      acc = a.fma(b, acc);  // fused multiply-add  
    }  
    double dot = acc.reduceLanes(VectorOperators.ADD);  
    // handle tail k..D-1  
    similaritiesFlat[i*N + j] = dot;  
  }  
}

На рисунке ниже показано, как Vector API задействует SIMD-аппаратное обеспечение, обрабатывая несколько значений double за одну инструкцию (например, 4 линии на AVX2 и 8 линий на AVX-512). То, что раньше было множеством скалярных операций multiply-add, превращается в меньшее число векторных операций fma() плюс редукцию — тот же алгоритм, но гораздо более эффективное использование векторных блоков CPU.

Фолбэки и надёжность: когда Vector API недоступен

Поскольку Vector API всё ещё находится в статусе incubating, для его включения нужен флаг рантайма: --add-modules=jdk.incubator.vector. Мы не хотели, чтобы корректность или доступность зависели от этого флага. Поэтому поведение фолбэка мы спроектировали явно: при старте мы определяем, поддерживается ли Vector API, и при наличии используем SIMD-реализацию батчевого matmul; в противном случае откатываемся на оптимизированный скалярный путь, при этом запросы на одиночные видео продолжают использовать реализацию «поэлементно».

Это даёт нам понятную операционную модель: сервисы могут включить Vector API ради максимальной производительности, но система остаётся безопасной и предсказуемой и без него.

Результаты в продакшене:

Когда полная конструкция была готова — батчинг, плоские буферы, повторное использование через ThreadLocal и Vector API, — мы запустили канареечные прогоны на продакшен-трафике. Мы увидели снижение утилизации CPU примерно на ~7% и падение средней задержки примерно на ~12%. Чтобы нормализовать результаты с учётом любых небольших различий в пропускной способности, мы также отслеживали CPU/RPS (CPU, потребляемый на запросы в секунду). Этот показатель улучшился примерно на 10%, то есть мы могли обслуживать тот же трафик примерн�� на 10% меньшими затратами CPU, и схожие значения сохранились после полного выката в продакшен.

На уровне операторов функций мы увидели, что потребление CPU упало с исходных 7,5% до всего лишь ~1% после внедрения оптимизации. На уровне ассемблера сдвиг был очевиден: от развернутых по циклу скалярных скалярных произведений — к векторизованному умножению матриц на аппаратуре AVX-512.

Заключительные мысли

Эта оптимизация в итоге оказалась не столько поиском «самой быстрой библиотеки», сколько наведением порядка в базовых вещах: выбором правильной формы вычислений, кэше-дружественной раскладкой данных и устранением накладных расходов, которые способны свести на нет теоретические выигрыши. Когда эти элементы встали на место, JDK Vector API отлично подошёл под задачу: он позволил выразить SIMD-подобную математику на чистой Java, без JNI, при этом сохранив безопасный путь с фолбэком. Ещё одним бонусом стали низкие затраты разработчиков: по сравнению с более низкоуровневыми подходами Vector API позволил заменить куда более крупную и сложную реализацию относительно небольшим объёмом читабельного Java-кода, что упростило ревью, сопровождение и дальнейшие итерации.

Пробовали ли вы уже Vector API в реальном сервисе? Буду рад услышать, для каких нагрузок он помог (или не помог) и какие выводы вы сделали о бенчмаркинге и выкатке в продакшен.

Присоединяйтесь к русскоязычному сообществу разработчиков на Spring Boot в телеграм — Spring АйО, чтобы быть в курсе последних новостей из мира разработки на Spring Boot и всего, что с ним связано.

Netflix: практическое применение Vector API, которое спасло CPU