Autodeploy : Ajouter une passe de transformation Fusion

Où cette skill s'applique

Ce fichier vit dans le plugin trtllm-agent-toolkit. Les chemins tels que tensorrt_llm/..., examples/auto_deploy/... et tests/... sont relatifs à un checkout source TensorRT-LLM sur la machine de l'utilisateur — et non à l'arborescence du plugin.

Après l'installation du plugin (voir le README.md du toolkit), les skills utilisent le préfixe trtllm-agent-toolkit: (par exemple trtllm-agent-toolkit:ad-add-fusion-transformation).

Skills connexes dans ce plugin

Utilisez cette skill quand vous savez déjà quel sous-graphe ou motif vous visez (à partir de dumps graphe, de logs, ou de lecture de code). Pour la capture de dump et la sémantique des fichiers, suivez d'abord ad-graph-dump.

Quand utiliser cette skill

• Ajouter, étendre ou relire une fusion sous les transformations AutoDeploy dans un arbre TensorRT-LLM.

Workflow (concis)

1. Confirmer le motif dans les dumps graphe actuels (voir ad-graph-dump).
2. Chercher un kernel existant ou un chemin custom-op avant nouveau Triton ou CUDA.
3. Implémenter le changement minimal qui prouve la correction et la correspondance ; ajouter des tests.
4. Relancer les dumps et tests ; si les sorties dérivent, séparez les problèmes de correspondance, la perte de métadonnées, et les différences numériques.

Trouver des candidats fusion (léger)

Faites-le avant d'écrire une nouvelle passe pour travailler sur la structure graphe réelle.

Entrées

• Répertoire de dump graphe d'une exécution avec AD_DUMP_GRAPHS_DIR défini (voir ad-graph-dump).
• ID modèle et config AutoDeploy active (YAML registry, overlays default.yaml).
• Arbre source TensorRT-LLM pour la recherche de kernel et transformation.

Sorties

• Liste ordonnée de candidats avec : preuve graphe, recherche kernel existant (found / not_found), recommandation (use_existing_kernel, needs_triton_fallback, defer), et compromis (complexité, risque de correction).

Workflow de découverte

1. Parser les dumps pour les motifs non fusionnés répétés (chaînes élément-par-élément, chaînes norm, epilogues, ops adjacentes à attention).
2. Chercher dans l'arbre les transformations ou ops custom équivalentes ; enregistrer la preuve fichier/symbole.
3. Si rien ne convient, marquer le travail Triton ou autre kernel comme fallback délibéré.
4. Préférer les candidats avec récurrence claire, support existant, et risque numérique plus faible.

Template par candidat

Candidate: <short-name>
Affected graph pattern: <pattern>
Existing kernel lookup: <found|not_found>
Evidence: <path/symbol>
Recommendation: <use_existing_kernel|needs_triton_fallback|defer>
Strengths / weaknesses / risks:
- ...

Garde-fous

• Ne pas sauter la recherche kernel existant.
• Ne pas se tourner par défaut vers Triton quand un op viable existant existe déjà.
• En cas d'incertitude, defer et affiner la question avec un dump ou test supplémentaire.

Entrées (implémentation)

• Candidat choisi ou sous-graphe concret.
• Fichiers modèle et config actifs.
• Dumps graphe frais quand disponibles.
• Baseline actuelle : comptes de correspondance dans les logs, statut test unitaire, notes de précision que vous maintenez déjà.

Sorties (implémentation)

• Conception de passe ou patch : transformation enregistrée, entrée default.yaml, YAML model-registry optionnel.
• Décision de chemin : existing_kernel_path vs triton_fallback_path (ou autre stack kernel).
• Notes de validation : preuve graphe, [SUMMARY] matches=... avant/après des logs AutoDeploy, résultats de test.

Workflow d'implémentation

1. Aligner la passe avec la structure graphe observée des dumps — pas des noms d'op supposés à partir de docs seules.
2. Chercher dans transform/library/, custom_ops/, torch.ops.auto_deploy.*, et tests connexes pour réutiliser.
3. Intégrer un op existant quand possible ; sinon déléguer le travail kernel à la skill appropriée (triton-kernel-writing, cuda-kernel-writing, etc.).
4. Garder un changement logique par patch ; étendre les tests dans le même changement.
5. Relire les dumps après la modification ; si les comptes de correspondance s'effondrent, soupçonnez la disponibilité du motif ou la propagation des métadonnées.

Où vivent les passes fusion

• Transformations : tensorrt_llm/_torch/auto_deploy/transform/library/
• Registry / comportement de base : tensorrt_llm/_torch/auto_deploy/transform/interface.py
• Liste de transformation par défaut : tensorrt_llm/_torch/auto_deploy/config/default.yaml
• Aide au dump : tensorrt_llm/_torch/auto_deploy/utils/graph_writer.py
• Utilitaires graphe : tensorrt_llm/_torch/auto_deploy/utils/node_utils.py, tensorrt_llm/_torch/auto_deploy/utils/_graph.py
• Ops personnalisées : tensorrt_llm/_torch/auto_deploy/custom_ops/

Tests (typiques) :

• tests/unittest/auto_deploy/singlegpu/transformations/library/
• tests/integration/defs/accuracy/test_llm_api_autodeploy.py (quand le comportement ou les numériques peuvent changer)

Comment ajouter une transformation

Implémenter la passe

Créer ou mettre à jour un module sous transform/library/ et enregistrer la classe :

@TransformRegistry.register("my_transform_key")
class MyTransform(BaseTransform):
    @classmethod
    def get_config_class(cls):
        return MyTransformConfig

Utiliser une classe config dédiée uniquement quand la passe a besoin de paramètres au-delà de la config de transformation de base.

Enregistrer dans default.yaml

Ajouter une clé sous transforms: dans tensorrt_llm/_torch/auto_deploy/config/default.yaml. Copier l'ensemble des champs de la transformation existante la plus proche dans la même section du fichier (les clés requises dépendent de la classe config de transformation et de comment les pairs sont déclarés). Les passes expérimentales nouvelles doivent rester enabled: false jusqu'à être couvertes par des tests et dumps.

Activer pour un modèle spécifique

Pour un déploiement ciblé, ajuster les YAMLs registry sous examples/auto_deploy/model_registry/configs/ plutôt que d'activer des passes non prouvées globalement.

Règles d'implémentation

• Préférer les ops AutoDeploy / TRT-LLM existantes et les entrées torch.ops.auto_deploy.
• Préférer les contrats graphe stables, neutres au backend ; éviter de cacher le flux de données réel dans node.meta quand une arête devrait le porter.
• Utiliser les métadonnées pour les faits tenseur observables (shape, dtype) et les préserver lors des réécritures quand les remplaçants doivent rester traçables.
• Une hypothèse par patch — ne pas mélanger les fusions non connexes.

Kernel existant d'abord, Triton deuxièmement

Avant Triton :

1. Chercher dans transform/library/ et custom_ops/.
2. Chercher torch.ops.auto_deploy.* et définitions d'op custom TRT-LLM.
3. Lire les tests pour des intégrations similaires.

Utiliser triton-kernel-writing uniquement quand aucun op adapté n'existe et vous acceptez de posséder kernel + travail d'intégration.

Ordre de validation

1. Dumps graphe — motif présent, réécriture visible (voir ad-graph-dump).
2. Tests unitaires pour la transformation.
3. Vérifications d'intégration ou de précision quand les numériques ou le comportement bout-en-bout peuvent changer.

Comptes de correspondance

AutoDeploy enregistre [SUMMARY] matches=<n> (ou skipped / disabled) par transformation. Comparer avant et après votre changement ; une grosse baisse indique généralement des problèmes de motif ou métadonnées, pas "des runs lents."

Attentes de test

Suivre trtllm-code-contribution pour les conventions du repo. Couvrir :

• Réécritures micrographe ou graphe exporté en cas d'usage nominal.
• Modes d'échec qui ne doivent pas fusionner (multiples consommateurs, consommateurs mixtes).
• Préservation des métadonnées quand une passe upstream alimente votre motif.

Localisation unittest primaire pour les transformations library :

• tests/unittest/auto_deploy/singlegpu/transformations/library/

Checklist de review

• Structure cible apparaît dans les dumps actuels.
• Transformation enregistrée et listée dans default.yaml de manière cohérente avec les entrées pairs.
• Les bascules model-registry sont intentionnelles.
• matches non-zéro où attendu, ou skipped est expliqué.
• Snippets ou diffs dump avant/après sauvegardés pour le fil de review.
• Tests couvrent à la fois les cas de succès et les cas de non-correspondance intentionnels.
• Si les sorties changent, classer la perte de correspondance vs perte de métadonnées vs dérive numérique acceptable.

Garde-fous

• Ne pas grouper des passes non connexes dans un changement.
• Si les dumps contredisent les attentes, documenter ce que vous avez observé avant de poursuivre des hypothèses non connexes.

Note d'itération (template)

Candidate: <name>
Path: <existing_kernel_path|triton_fallback_path|other>
Rationale:
- ...
Graph validation: <pass|fail — what files / ops>
Summary logs: <matches before / after>
Tests: <what ran>
Open risks:
- ...