Backends rapides

Portabilité de TileLang

Les cinq kernels de laya/tl_kernels.py peuvent s’exécuter sur la cible CPU (c) de TileLang via compile_cpu. C’est une spécialisation scalaire fp32 explicite, pas un backend d’accélération CPU pour Agent.accelerate. Elle nécessite TileLang et un compilateur C++ local. Elle n’introduit aucune dépendance à un service d’exécution.

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs) prend les mêmes options de forme et d’opération que les cinq fabriques GPU. Il sélectionne cpu=True, dtype="float32", la cible c, et désactive la vectorisation. Toutes les entrées et sorties flottantes doivent être des tenseurs float32 CPU contigus ; les longueurs d’attention restent en int32. Convertis les activations 16 bits avec .cpu().float().contiguous() avant de l’appeler. LayerNorm met toujours à jour son tenseur résiduel sur place, et RoPE met toujours à jour les colonnes Q/K sur place. Conserve le kernel compilé pour réutiliser ses dimensions dynamiques.

La spécialisation CPU remplace les allocations de fragments et de shared par des tampons locaux, utilise des boucles T.grid série, et omet l’annotation swizzle GPU de l’attention. Les GEMM utilisent l’implémentation CPU scalaire de TileLang et les réductions utilisent des tampons locaux. Les algorithmes tuilés d’origine, les prédicats de padding, les masques d’attention et le softmax en ligne restent partagés avec l’implémentation GPU. Les intermédiaires CPU restent en fp32, y compris les probabilités d’attention ; les intermédiaires GPU conservent leur dtype d’origine. Aucune accélération CPU ni backend CPU pour le modèle complet n’est revendiqué.

Nouvelle sonde sur TileLang 0.1.14

Observé sous Linux, Python 3.12.13, torch 2.11.0+cu130, TileLang 0.1.14 et une RTX 4070 Ti SUPER. La préoccupation de portabilité précédente concerne la compilation de la spécialisation GPU inchangée, pas la possibilité d’une compilation vers le CPU. Au commit de base fa9a2a7, cette page de documentation était absente du checkout.

Le test compile factory.get_tir(...) avec target="c" et la configuration de passe FAST du GPU. Ce sont les textes de diagnostic exacts (emplacements de source et traces de pile omis). bf16 et fp16 ont tous deux été sondés pour chaque kernel.

Les quatre premiers échecs de fragment/réduction sont tvm.error.InternalError. L’échec de génération de code bf16 est aussi un InternalError. RoPE en FP16 lève RuntimeError: Compilation Failed! suivi de l’invocation du compilateur et de la source ; le diagnostic ci-dessus est émis sur stderr. Ses conversions vectorielles générées échouent aussi à reconvertir des vecteurs float en vecteurs half.

Pour isoler le support dtype du support fragment, les tests compilent chaque kernel à nouveau avec cpu=True (tampons locaux et boucles série), conservent bf16 et désactivent la vectorisation. Les cinq échouent alors exactement avec :

Cannot convert type bfloat16 to C type

Ainsi, les fragments sont le premier obstacle pour GEMM, GEGLU et l’attention, les réductions de fragments pour LayerNorm, et bf16 est indépendamment un obstacle pour les cinq. RoPE n’a ni allocation de fragment ni réduction ; GEMM et GEGLU n’ont pas d’opération T.reduce_* explicite. Les réductions de l’attention sont d’abord masquées par son échec de disposition. Des sondes de fragments fp32 uniquement, séparées, reproduisent l’erreur de remplissage et l’erreur de réduction à la fois pour reduce_sum et reduce_max, sans aucun GEMM ni bf16. Remplacer les scopes seuls a aussi produit cette erreur sémantique pour le tampon GEMM (les autres tampons touchés étaient Ci, x et s) :

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

La spécialisation fp32 série supprime ces obstacles. Les assertions de diagnostic sont épinglées à la version 0.1.14 et ignorées sur une autre version, où les erreurs devraient être sondées à nouveau. Les tests numériques CPU continuent de s’exécuter sur d’autres versions.

Vérifications numériques

Exécute python -m pytest tests/test_fast_cpu.py -q -s. Sur l’environnement ci-dessus : 58 réussis. Les vérifications CPU uniquement ne nécessitent pas CUDA ; seules les comparaisons GPU sont ignorées sans lui. La couverture inclut des tuiles GEMM M/N/K irrégulières, tous les épilogues GEMM, GEGLU, toutes les combinaisons résidu/biais, de grandes valeurs résiduelles, le rebouclage de position RoPE et les colonnes V intactes, les formes d’attention statiques/dynamiques, les fenêtres glissantes, les tuiles partielles, les longueurs inégales, les séquences vides et les sorties de padding finies.

La graine est 1234. Les comparaisons GPU utilisent des valeurs d’entrée identiques arrondies en bf16 ou fp16, puis promues en fp32 pour l’exécution CPU. Ce sont des tolérances absolues pour les fixtures bornées, pas une garantie pour des magnitudes arbitraires ou une profondeur de modèle. Les comparaisons d’attention utilisent des lignes de requête valides, comme dans tests/test_fast.py.

La tolérance CPU/référence est 2e-5 (2e-6 pour RoPE). Les mises à jour du flux résiduel sont exactement égales, y compris le cas sans résidu.

Preuves de préservation du GPU

Le mot-clé cpu vaut false par défaut. Les valeurs par défaut bf16/fp16 existantes, les scopes d’allocation GPU, les boucles parallèles, les swizzles et les options FAST sont inchangés. Les appels GPU fp32 par défaut lèvent toujours ValueError.

Les exécutions avant/après ont utilisé respectivement le module d’origine extrait avec git show fa9a2a7:laya/tl_kernels.py et le module modifié. L’original a été chargé comme laya.tl_kernels via importlib pour les exécutions de référence ; tout le reste du code Laya et l’environnement Python sont restés identiques.

• python -m pytest tests/test_fast.py -q, avec les tests full-forward configurés pour charger le checkpoint anglais en cache identifié plus bas : 13 réussis avant ; 13 réussis après. Au départ, sans checkpoint, c’était 11 réussis / 2 ignorés. Les deux exécutions complètes ont émis l’avertissement existant de limitation de température du checkpoint (choice:11+=0.10058280825614929 -> 0.5).
• Capture avec graine des tests de kernel existants : 24 tenseurs de sortie identiques bit à bit, y compris les mises à jour résiduelles ; différence maximale avant/après 0.
• Code source CUDA généré pour les cinq formes de sonde ci-dessus, dans les deux dtypes : 10/10 identiques octet pour octet. Cela couvre aussi RoPE, absent de la suite fast d’origine.
• python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... et la commande fp16 équivalente : 288 questions sur 60 états par dtype. Tous les enregistrements JSON avant/après (probabilités fp32, stock et fast) se comparent exactement égaux ; différence maximale de probabilité avant/après 0.

MODEL était l’instantané anglais en cache convaiinnovations/laya 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851 ; les exécutions utilisaient HF_HUB_OFFLINE=1. Aucune comparaison avec un checkpoint multilingue ou à décisions typées n’est revendiquée ici.

Vérifications du dépôt

Toutes les commandes ont utilisé /home/ckl/projects/S/laya/.venv/bin/python ; ruff et zensical provenaient du répertoire bin de cet environnement virtuel.

La nouvelle suite est enregistrée comme nécessitant l’extra TileLang optionnel et un compilateur C++ dans les exemptions existantes du test d’empaquetage. Le test de contrat d’API épingle le sélecteur CPU additif à mot-clé uniquement, le dtype GPU par défaut inchangé et la signature de compile_cpu, sans importer TileLang dans l’environnement CI de base.

AOTInductor

DecisionModel peut être exporté, compilé en un paquet .pt2 et chargé avec torch._inductor.aoti_load_package. Le paquet renvoie à la fois les logits de décision et les logits d’action. La tokenisation, le padding, la calibration de température et le formatage des réponses restent à la charge de l’appelant ; cela n’ajoute pas de backend Agent et ne change pas son chemin d’exécution par défaut.

Le PR #472 fermé a documenté l’obstacle dtype antérieur. L’état regroupé et les caractéristiques de confiance de la tête d’action sont calculés en fp32, même lorsque les poids d’un modèle sont explicitement en bf16. Sans autocast, le premier linear d’action recevait donc une entrée fp32 et des poids bf16. Désormais la passe avant convertit l’entrée concaténée vers le dtype des poids de la tête uniquement hors autocast. Softmax, l’entropie et les logits de décision renvoyés conservent leurs calculs en fp32. L’inférence eager existante à paramètres fp32, y compris l’AMP fp16/bf16, conserve son comportement numérique ; les dtypes mixtes poids/autocast conservent eux aussi la conversion d’origine d’autocast.

Exporte une copie d’évaluation distincte convertie en bf16, hors autocast. Déclare l’axe des jetons comme 16 * Dim("tokens16", ...) pour satisfaire les gardes d’alignement de l’attention. Les lignes et les nombres de marqueurs peuvent être dynamiques indépendamment. Ne présume pas qu’une exportation capturée sous autocast peut être empaquetée hors de ce contexte : utilise des dtypes de poids explicites pour cette recette.

Reproduire hors ligne

La vérification basée sur des assertions utilise par défaut un ModernBERT minuscule initialisé aléatoirement, sans accès réseau. Passe un répertoire de checkpoint local pour une mesure avec un vrai modèle. L’absence de CUDA, des API AOTInductor ou d’un compilateur C++ produit un SKIP explicite ; les échecs de compilation et de parité sur une installation prise en charge font échouer la vérification.

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

Le répertoire de sortie contient decision.pt2, results.json, inputs.pt et eager.pt. Le JSON inclut des instantanés complets de nvidia-smi, les temps d’exportation/empaquetage/chargement, les octets du paquet, la latence et les deltas maximaux absolus de logit/probabilité pour les deux têtes. Les artefacts binaires et les caches du compilateur sont délibérément tenus hors du dépôt.

Mesuré sur une RTX 4070 Ti SUPER

2026-10-03, Linux x86-64, Python 3.12.13, torch 2.11.0+cu130, transformers 5.17.0, pilote NVIDIA 615.71.09, 16,376 MiB de VRAM. Checkpoint : convaiinnovations/laya, sous-répertoire multilingual à la révision 1c5edc17a7acd8701df6fc341c0d179f1c62c982. La référence est main fa9a2a7. Les instantanés complets de la machine et les mesures non arrondies sont dans aoti_multilingual_rtx4070.json.

L’échec reproduit se produit à l’empaquetage, après une exportation réussie : mat1 and mat2 must have the same dtype, but got Float and BFloat16. Chaque exécution a utilisé un cache Inductor distinct, initialement vide ; la compilation AMP s’est exécutée avant l’empaquetage, donc le temps du paquet n’est pas une mesure de démarrage à froid d’un interpréteur neuf. La vérification en processus neuf n’a chargé que le paquet et les entrées enregistrées (pas de checkpoint), avec torch.compile et les points d’entrée d’empaquetage bloqués. Elle a reproduit les réponses des deux têtes. PyTorch a tout de même compilé une petite sonde de capacité AVX CPU dans le cache initialement vide ; aucun graphe de modèle ni kernel CUDA n’a été compilé. Une affirmation générale « aucune activité du compilateur au chargement » serait donc inexacte sur cette exécution.

Les mêmes entrées enregistrées contiennent sept lignes de décision réparties sur trois lots, avec du texte réel tokenisé de facturation/remboursement, des questions choice/score/noul, du padding et des nombres de marqueurs valides inégaux. Chaque latence est la médiane de cinq groupes de 30 passes avant après dix échauffements, avec une mesure de temps murale synchronisée. torch.compile(dynamic=True) utilise le mode par défaut. Aucun graphe CUDA explicite, tokenisation, exportation ou compilation n’est inclus dans ces chiffres de latence.

Ici AMP signifie paramètres fp32 avec autocast bf16. bf16 explicite signifie que les poids et le flux résiduel sont en bf16, sans autocast ; utilise ces colonnes pour la comparaison d’exécution la plus proche du paquet. bf16 explicite eager/compiled et AOTI étaient indisponibles avant ce correctif.

L’argmax de décision comme celui d’action concordent sur 7/7 lignes pour les deux comparaisons AOTI. Les probabilités sont des sorties softmax brutes, sans calibration. La distribution d’action est saturée sur ces entrées, donc son delta de probabilité nul n’implique pas des logits sous-jacents identiques face à AMP. La passe avant compilée en bf16 explicite diffère aussi du bf16 explicite eager (delta maximal de logit de décision 0.5, delta de probabilité 0.00659859). La compilation en précision réduite n’est pas exacte bit à bit.

Le GPU était partagé avec des applications de bureau et d’autres tâches Python ; la compilation CPU l’était aussi, et les fréquences n’étaient pas verrouillées. Ce sont des temps observés, pas une revendication d’accélération isolée. Instantanés nvidia-smi encadrant les exécutions :

Limites restantes

• L’artefact est spécifique à ce checkpoint, à cette précision, à cette pile PyTorch/exécution et à cette cible GPU ; la portabilité vers d’autres matériels ou versions de PyTorch n’a pas été testée.
• Cette vérification exporte les lignes 1–8, les jetons 32–512 par multiples de 16 et les slots de marqueur 2–8. Elle exerce trois formes, dont des formes différentes de l’exemple d’exportation, plutôt que chaque point de ces plages. Une option valide peut utiliser des slots de marqueur avec padding ; un vrai tenseur à un seul slot emprunte la branche séparée à option unique et nécessite une exportation distincte. Les longueurs de jeton arbitraires et une couche de production de bucketing/dispatch sont hors de ce changement.
• Sept lignes vérifient la régression d’empaquetage, pas la précision large du checkpoint ni la parité de confiance calibrée. Convertir une copie d’exportation en bf16 diffère de conserver les poids fp32 sous autocast ; le chemin eager existant lui-même reste inchangé.
• Le script nécessite une compilation PyTorch compatible CUDA et une chaîne d’outils de compilateur locale pour créer l’artefact. Le .pt2 intègre le modèle et les kernels CUDA ; aucun service hébergé n’intervient.