哲见WISE AITHE INTELLIGENCE JOURNAL
Participer
Voir les rubriques
vLLM BlogPublication de la source

vLLM Team

Migration de vLLM vers l'ABI Stable de LibTorch

vLLM repose fortement sur PyTorch, avec des appels fréquents à LibTorch dans son backend C++. Par conséquent, les modifications de l'interface binaire de LibTorch peuvent nécessiter que vLLM soit recompilé et réédité, même lorsque les API restent inchangées.

Introduction

vLLM repose fortement sur PyTorch, avec des appels fréquents à LibTorch dans son backend C++. Par conséquent, les modifications de l'interface binaire de LibTorch peuvent nécessiter que vLLM soit recompilé et réédité, même lorsque les API restent inchangées. Et si nous pouvions mettre à jour LibTorch sans recompiler vLLM ? C'est là qu'intervient l'ABI Stable de LibTorch. Avec une ABI Stable de LibTorch, la mise à jour de LibTorch ne nécessite pas de recompiler ou de mettre à jour vLLM pour continuer à l'utiliser. Au lieu de cela, nous pourrions remplacer la wheel de torch et l'importer directement dans notre paquet vLLM. À partir de la v0.33.0, vLLM publie des wheels CUDA avec ABI Stable de LibTorch.

Import instable
Import stable

À gauche : ABI instable — le nouveau torch nécessite que vLLM soit recompilé. À droite : ABI stable — vLLM peut importer directement le nouveau torch.

Remarque : La compatibilité entre vLLM et PyTorch dépend toujours d'autres facteurs, notamment les API Python, les exigences liées au GPU (telles que la version minimale de CUDA prise en charge) et les dépendances tierces. L'ABI Stable de LibTorch traite la compatibilité binaire mais n'élimine pas ces autres contraintes.

Qu'est-ce que l'ABI Stable de LibTorch ?

Une interface binaire applicative (ABI) définit la manière dont le code compilé interagit avec une bibliothèque : comment les fonctions sont appelées, comment les données sont disposées et comment les symboles sont nommés. Même si une API semble identique dans le code source, un changement d'ABI peut forcer la recompilation des extensions. Une ABI stable préserve un ensemble pris en charge d'interfaces binaires entre les versions compatibles, de sorte que les extensions qui les utilisent peuvent continuer à fonctionner sans recompilation.

L'ABI Stable de LibTorch de PyTorch comporte trois couches : un shim C stable exposant certaines API de LibTorch, des wrappers C++ qui rendent le shim plus facile à utiliser, et des utilitaires headeronly dans torch::headeronly qui ne sont pas liés à libtorch.so.

LibTorch Stable

Pourquoi l'ABI Stable est importante

Avant l'ABI Stable, une nouvelle version de PyTorch pouvait nécessiter une recompilation de vLLM, retardant la prise en charge de la nouvelle version et de ses fonctionnalités. Avec l'ABI Stable, la même wheel de vLLM peut fonctionner sur plusieurs versions compatibles de PyTorch, permettant aux utilisateurs de mettre à niveau plus tôt et réduisant les recompilations pour les distributeurs de paquets. Pour les développeurs, cela signifie que vous pouvez mettre à jour et recompiler PyTorch, puis tester ces changements dans vLLM sans recompiler vLLM.

Diagramme de l'ABI instable
Diagramme de l'ABI stable

À gauche : vLLM doit évoluer au fur et à mesure des changements de LibTorch. À droite : vLLM peut évoluer indépendamment de LibTorch.

Compromis

L'ABI Stable de LibTorch est un sous-ensemble limité de l'ensemble de LibTorch. Bien que de nombreuses opérations courantes soient déjà prises en charge via le shim C et via les wrappers C++ de plus haut niveau, toutes ne le sont pas. Par exemple, is_pinned() ne disposait pas d'équivalent stable pendant la migration. Heureusement, il existe un shim C torch_call_dispatcher à usage général qui nous permet d'appeler les opérations ATen via le dispatcher de manière stable au niveau de l'ABI ! Cela nous a permis de maintenir le comportement tout en atteignant la stabilité de l'ABI. Cette migration a également aidé PyTorch à identifier les API couramment utilisées, et dans la version 2.14, ils ont ajouté une méthode C++ stable is_pinned() .

Avant : appel direct de méthode LibTorch

if (cpu_tensor.is_pinned())

Après : appel via le dispatcher stable au niveau de l'ABI

std::array<StableIValue, 2> is_pinned_stack{
    torch::stable::detail::from(cpu_tensor),
    torch::stable::detail::from(std::nullopt)};
TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
    "aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
if (torch::stable::detail::to<bool>(is_pinned_stack[0]))

Le second compromis est une légère surcharge supplémentaire côté hôte. Le shim C stable ajoute une surcharge de conversion de types et, pour certains wrappers qui reposent sur torch_call_dispatcher, une recherche d'opérateur et du boxing. Cette surcharge est généralement négligeable et disparaît après l'échauffement sous CUDA Graphs, mais elle peut avoir de l'importance en mode eager lors du lancement répété de kernels très rapides. Nous n'avons constaté aucune régression notable des temps d'exécution lors de la migration.

Migrer vLLM vers la Stable ABI

Nous avons décomposé la migration vers la Stable ABI en une série d'étapes :

  1. Évaluer la situation

Tout d'abord, nous avons évalué si vLLM pouvait migrer en utilisant les API stables disponibles à l'époque. Nous avons recensé ses appels à l'API PyTorch, identifié les manques et planifié les API stables nécessaires pour les combler. Ce travail est décrit dans le RFC original visant à mettre à jour le build CUDA de vLLM pour utiliser la LibTorch Stable ABI.

  1. Migrer les AutogradFunctions et les méta-fonctions vers Python

À ce jour, la LibTorch Stable ABI ne prend pas en charge SymInt de manière stable, donc les méta-fonctions doivent être écrites en Python plutôt qu'en C++. Heureusement, vLLM avait déjà la plupart de ses méta-fonctions en Python, si bien que migrer les quelques-unes qui ne l'étaient pas a été relativement simple et a rendu le projet plus cohérent. De même, la Stable ABI n'expose pas la machinerie autograd de PyTorch, donc les AutogradFunctions doivent également être écrites en Python. Cela importe moins pour un moteur d'inférence comme vLLM, mais si une fonction backward est implémentée, elle devra être en Python.

  1. Mettre en place l'échafaudage nécessaire

Comme vLLM est un projet de si grande envergure, avec plus de 175 fichiers source .cpp/.cu plus 120 autres en-têtes rien que dans son build CUDA, tout migrer en une seule fois aurait été ingérable et difficile à relire. À la place, nous avons créé une bibliothèque stable séparée (voir les mises à jour de CMakeLists.txt dans la première PR #31509de migration), et nous avons migré les kernels vers cette bibliothèque par morceaux gérables. Le TORCH_TARGET_VERSION de build aide à maintenir la stabilité de la nouvelle bibliothèque vis-à-vis de l'ABI LibTorch en provoquant l'échec du build si elle inclut des en-têtes LibTorch instables.

Migration de vLLM de la bibliothèque instable vers la bibliothèque stable

À mesure que la bibliothèque stable grandit, la bibliothèque instable rétrécit puis disparaît.

  1. Assurez-vous d'enregistrer les ops via le dispatcher

À ce jour, l'enregistrement de vos ops LibTorch Stable requiert un enregistrement via STABLE_TORCH_LIBRARY. Cela vous sera très familier si vous enregistrez déjà vos ops via TORCH_LIBRARY, ce que fait vLLM. La plus grande différence avec le nouvel enregistrement stable est que l'implémentation doit être enveloppée avec TORCH_BOX. Cela génère un adaptateur qui encadre et décadre les types de données C++ dans les types génériques StableIValue capables de traverser l'ABI Stable de PyTorch.

Si vos kernels sont enregistrés via un autre mécanisme, tel que pybind11, nous recommandons de d'abord convertir l'enregistrement en TORCH_LIBRARY avant d'enregistrer via STABLE_TORCH_LIBRARY (plus de détails ci-dessous dans la section Challenges Along the Way ).

  1. Migrer chaque fichier source vers les API stables

Une fois les étapes 1 à 4 en place, nous pouvons migrer chaque fichier source vers les API stables. Tous les en-têtes PyTorch que nous incluons doivent provenir des wrappers stables C++, des utilitaires header-only, ou du shim C lui-même :

  • torch/csrc/stable
  • torch/headeronly
  • torch/csrc/stable/c/shim.h (ou torch/csrc/inductor/aoti_torch/c/shim.h)

Les modifications des fichiers sources sont pour la plupart mécaniques. Par exemple :

  • torch::Tensor → torch::stable::Tensor
  • at::kInt → torch::headeronly::ScalarType::Int
  • t.view({...}) → torch::stable::view(t, {...})

Pour faciliter la migration des fichiers sources, nous avons créé des skills pour agents d'IA qui associent les API existantes à leurs équivalents stables et signalent les erreurs courantes, telles que des modifications sans rapport avec la migration. Ces skills prennent également en charge les étapes précédentes (1 à 4). Ils se sont révélés utiles dans notre travail, mais nous n'avons pas encore développé d'évaluations pour mesurer leur efficacité (voir abi-stable-skills).

  1. Vérifiez !

vLLM est utilisé par des millions d'utilisateurs dans plusieurs secteurs, et nous ne voulons pas interrompre aucun de leurs processus avec la migration. Par conséquent, l'étape la plus importante de cette liste est de vérifier que la migration n'a apporté aucune modification significative. Cela signifie exécuter les tests unitaires et de performance appropriés et examiner manuellement chaque diff de migration ligne par ligne pour s'assurer qu'il n'y a eu aucune mise à jour involontaire.

En plus de vérifier l'exactitude, nous devons également vérifier que la nouvelle bibliothèque est réellement LibTorch ABI Stable. Pour cela, nous utilisons torch-abi-audit (merci Ralf Gommers !). Cet outil inspecte les fichiers .so et catégorise les symboles PyTorch stables et instables. S'il n'y a pas de symboles instables, la bibliothèque est LibTorch ABI Stable. Il signale également la version minimale de LibTorch qui fournit tous les symboles stables requis (par exemple, min-torch=2.11.0). Nous l'avons utilisé pour suivre la progression de la migration et vérifier l'absence de régressions, et nous l'utilisons désormais dans la CI de vLLM pour garantir que la wheel reste LibTorch ABI Stable.

torch-abi-audit vllm
Package: vllm
  Root: /workspaces/vllm_container/vllm/vllm
  Torch ABI:   UNSTABLE
  Min torch:   2.11.0
  CPython ABI: mixed
  Extensions:  10
  Bundled libs: 0
  -- extensions --
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=88, unstable=0)  min-torch=2.11.0
    [UNSTABLE] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=0, unstable=70)
    [UNSTABLE] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=0, unstable=68)
    [STABLE  ] [abi3-ok               ] _moe_C_stable_libtorch.abi3.so  (stable_shim=72, unstable=0)  min-torch=2.11.0
    [UNSTABLE] [abi3-ok               ] _qutlass_C.abi3.so  (stable_shim=0, unstable=59)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so  (stable_shim=0, unstable=59)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=0, unstable=85)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=0, unstable=81)

torch-abi-audit avant que les dépendances ne soient migrées

Dépendances de build de vLLM

La migration des sources de vLLM seule ne permet pas d'établir que la wheel entière de vLLM est ABI stable. Pour que la wheel soit complètement stable, ses dépendances compilées doivent également être LibTorch ABI Stable. vLLM a actuellement sept dépendances qui dépendent également de LibTorch. Ce sont :

DépendanceMigrée sur le fork vLLMMigrée en amont
FlashMLA✅✅
DeepGEMM✅🚧 Travail en cours
FlashAttention-2✅🚧 PR ouverte pour le rendre stable
FlashAttention-3✅✅
DeepSelect✅🚧 Travail en cours
FlashKDA✅❌
QuTLASSN/A✅

Après que le code vLLM intégré dans l'arbre a été migré, chacune de ces dépendances a également été migrée vers la Stable ABI de LibTorch, rendant la wheel CUDA de vLLM entièrement Stable au sens de l'ABI LibTorch.

torch-abi-audit vllm
Package: vllm
  Root: /workspaces/vllm_container/vllm/vllm
  Torch ABI:   STABLE
  Min torch:   2.13.0
  CPython ABI: abi3
  Extensions:  13
  Bundled libs: 0
  -- extensions --
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=88, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] _deepselect_C.abi3.so  (stable_shim=40, unstable=0)  min-torch=2.10.0
    [STABLE  ] [abi3-ok               ] _flashkda_C.abi3.so  (stable_shim=43, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=66, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=63, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] _moe_C_stable_libtorch.abi3.so  (stable_shim=72, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] _qutlass_C.abi3.so  (stable_shim=63, unstable=0)  min-torch=2.13.0
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] fs_io_C.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [STABLE  ] [abi3-ok               ] third_party/deep_gemm/_C_extension.abi3.so  (stable_shim=85, unstable=0)  min-torch=2.11.0
    [STABLE  ] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=69, unstable=0)  min-torch=2.10.0
    [STABLE  ] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=63, unstable=0)  min-torch=2.10.0

torch-abi-audit après la migration des dépendances

Les défis rencontrés en cours de route

Le plus grand défi était de suivre le rythme d'un projet vaste et en évolution rapide. Les fichiers amont changeaient souvent pendant que nos correctifs de migration étaient en cours, et nous devions donc mettre à jour ces correctifs avant de pouvoir les fusionner. Les dépendances externes de vLLM signifiaient aussi une coordination entre projets et la construction d'un support pour la migration.

DeepGEMM : une étude de cas de migration

DeepGEMM a réuni plusieurs de ces défis. Pendant que notre migration était en cours, DeepGEMM est passé de la version 1 à la version 2 et a ajouté DeepJIT, un autre projet dépendant de LibTorch qui devait également être migré. vLLM a ensuite basculé du dépôt amont deepseek-ai/DeepGEMM vers son propre vllm-project/DeepGEMM fork, nous obligeant à refaire la migration sur le nouveau fork.

Ces changements amont ne constituaient qu'une partie du défi. Comme DeepGEMM enregistrait ses opérateurs via pybind11, nous devions également déplacer son enregistrement vers TORCH_LIBRARY avant d'adopter la Stable ABI. Par exemple, la liaison originale de fp8_fp4_paged_mqa_logits était :

m.def("fp8_fp4_paged_mqa_logits", &fp8_fp4_paged_mqa_logits,
      py::arg("q"), ...,
      py::arg("max_context_len"), ...);

où q est un std::tuple<Tensor, std::optional<Tensor>>& et max_context_len est un const int&. Pour rendre ces arguments compatibles avec TORCH_LIBRARY, nous avons ajouté un wrapper qui accepte les éléments du tuple séparément et utilise int64_t dans sa signature C++ :

static torch::Tensor fp8_fp4_paged_mqa_logits(
    const torch::Tensor& q, const c10::optional<torch::Tensor>& q_sf, ...
    const int64_t& max_context_len, ...) {
  return attention::fp8_fp4_paged_mqa_logits(
      std::make_tuple(q, q_sf), ...,
      static_cast<int>(max_context_len), ...);
}

Ceci peut maintenant être enregistré avec TORCH_LIBRARY comme :

TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
    m.def("fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, ..., int max_context_len, ...) -> Tensor");
}
 
TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
    m.impl("fp8_fp4_paged_mqa_logits", TORCH_FN(fp8_fp4_paged_mqa_logits));
}

L'interface Python déballe le tuple avant d'appeler l'opérateur enregistré :

def fp8_fp4_paged_mqa_logits(q, ..., max_context_len, ...):
    q_fp, q_sf = q[0], q[1]
    return _torch_ops.fp8_fp4_paged_mqa_logits(
        q_fp, q_sf, ..., max_context_len, ...)

Cela permet à l'enregistrement d'utiliser TORCH_LIBRARY sans modifier l'API exposée aux utilisateurs.

Après la migration de DeepGEMM vers TORCH_LIBRARY, la migration vers la Stable ABI a suivi un processus similaire aux autres migrations. Nous avons remplacé ses API instables et mis à jour les enregistrements d'opérateurs :

STABLE_TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
    m.def("fp8_fp4_paged_mqa_logits(Tensor q, Tensor? q_sf, ..., int max_context_len, ...) -> Tensor");
}
 
STABLE_TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
    m.impl("fp8_fp4_paged_mqa_logits", TORCH_BOX(fp8_fp4_paged_mqa_logits));
}

DeepGEMM étant la dernière dépendance de build de vLLM à migrer, l'achèvement de ce travail a rendu la wheel vLLM entière Stable au sens de l'ABI LibTorch !

Exemple d'exécution sur une version plus récente de PyTorch

# stable_abi_demo.py
import os
 
os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
 
import torch
import vllm
from vllm import LLM, SamplingParams
 
print(f"PyTorch runtime: {torch.__version__}")
print(f"vLLM: {vllm.__version__}")
 
llm = LLM(
    model="HuggingFaceTB/SmolLM2-135M-Instruct",
    max_model_len=128,
    gpu_memory_utilization=0.1,
    enforce_eager=True,
)
outputs = llm.generate(
    ["The capital of France is"],
    SamplingParams(temperature=0, max_tokens=8),
)
print(outputs[0].outputs[0].text)

Compilez vLLM et exécutez la démo avec PyTorch 2.14.1 :

> python stable_abi_demo.py
 
PyTorch runtime: 2.14.1+cu132
vLLM: 0.1.dev24716+gcb95d994f.d20261105
Paris. Paris is a city known for

Mettez à jour la wheel PyTorch vers torch==2.15.0:

uv pip install torch==2.15.0 \
  --index-url https://download.pytorch.org/whl/cu132

Réexécutez la démo sans recompiler vLLM :

> python stable_abi_demo.py
 
PyTorch runtime: 2.15.0+cu132
vLLM: 0.1.dev24716+gcb95d994f.d20261105
Paris. Paris is a city known for

Remarque : les journaux d'initialisation courants de vLLM, les barres de progression et les avertissements de dépréciation sans rapport ont été omis des sorties ci-dessus pour plus de clarté.

Travaux futurs

Maintenant que la version CUDA de vLLM utilise l'ABI stable de LibTorch, nos prochaines étapes consistent à rendre ABI stable toutes les versions de vLLM, y compris les versions CPU et ROCm. Ces migrations sont plus petites que la migration CUDA mais restent importantes pour l'écosystème vLLM. La version ROCm partage de nombreux fichiers sources avec la version CUDA : hipify convertit .cu fichiers à .hip fichiers, ces fichiers utilisent donc déjà l'ABI stable de LibTorch. Seuls quelques fichiers spécifiques à ROCm restent à migrer.

Nous élargissons également la surface de l'API stable afin que vLLM et ses dépendances n'aient plus besoin de wrappers. À mesure que la surface de l'API s'accroît, nous mettrons également à jour les tests amont de PyTorch pour garantir que les API restent stables au sens de l'ABI LibTorch.

Remerciements aux principaux contributeurs

Source originale

vLLM Blog

À propos du contenu

La publication originale et les droits appartiennent à la source.

Traduction automatique · Consultez l’original