Optimisation TileIR Triton

Optimisez les kernels Triton EXISTANTS pour le backend TileIR d'NVIDIA sur les GPUs Blackwell. Cette skill ne rédige PAS de kernels à partir de zéro -- c'est le travail du Triton Specialist.

Principes

Backend TileIR vs PTX

TileIR est le backend compilateur d'NVIDIA pour Triton qui génère du code CUDA optimisé en utilisant des représentations de tuiles au niveau CGA (Cooperative Grid Array). Différences critiques :

Implication clé : ne tunez pas num_warps pour TileIR -- concentrez-vous sur occupancy à la place.

Paysage des packages Triton

Trois packages partagent import triton :

Un seul package triton doit être installé à la fois. « Convertir en TileIR » signifie ajouter des configs spécifiques à TileIR, NON changer les imports. TileIR s'active via ENABLE_TILE=1.

Quand TileIR s'applique

TileIR cible Blackwell (sm_100+). Sans nvtriton ou matériel Blackwell, le specialist ajoute quand même des configs optimisés pour TileIR que triton standard ignore en toute sécurité, permettant un déploiement futur.

Accélérations attendues (avec nvtriton sur Blackwell) :

Workflow

Workflow à cinq phases : compatibilité, classification, transformation, validation, benchmark.

Phase 1 : Test de compatibilité (ENABLE_TILE=0)

Vérifiez que le kernel fonctionne en mode PTX avant d'appliquer les optimisations TileIR.

python scripts/tileir_check.py

Puis utilisez la skill kernel-triton-writing avec verify_kernel.py pour vérifier avec ENABLE_TILE=0 :

python scripts/verify_kernel.py --kernel path/to/kernel.py --reference 'torch reference' --shapes '{"x": [32, 512, 4096]}' --dtypes '{"x": "bfloat16"}'

Phase 2 : Classification du kernel

Déterminez le type de kernel pour sélectionner la stratégie d'optimisation.

python scripts/classify_kernel.py --file kernel.py

Arbre de décision de classification :

Contient tl.dot()?
  OUI --> dot-related: TMA + 2CTA + occupancy + plus grands blocs
  NON  --> A réduction + normalisation?
            OUI --> norm-like: occupancy élevée (2, 4) + num_warps (4, 8)
            NON  --> Seulement point-wise?
                      OUI --> element-wise: occupancy (1-16) + num_stages (2-4)
                      NON  --> reduction: occupancy élevée + num_warps

Phase 3 : Appliquer les transformations

Classifiez et appliquez les optimisations en une seule étape :

python scripts/classify_kernel.py --file kernel.py --apply-optimizations

Le JSON de sortie inclut les champs optimized_code et changes_applied.

Transformations spécifiques au type :

Dot-related (priorité la plus haute) :

1. Convertir tl.load/tl.store en descripteurs TMA (OBLIGATOIRE). Voir references/tma-conversion.md.
2. Ajouter des configs 2CTA (num_ctas=2) avec protection de sursouscription SM dans pre-hook.
3. Ajouter occupancy (1, 2, 4) et num_stages étendu (4, 6).
4. Utiliser des tailles de blocs plus grandes (256x256, 256x128).

Norm-like (LayerNorm, Softmax, RMSNorm) :

• Ajouter occupancy (2, 4), num_warps (4, 8). Pas besoin de TMA.

Element-wise (ReLU, GELU, Add, Mul, Exp) :

• Ajouter occupancy (1, 2, 4, 16), num_stages (2, 3, 4). Inclure des configs extrêmes pour petites entrées.

Réduction (Sum, Mean, Max) :

• Même stratégie que norm-like : occupancy élevée (2, 4), num_warps (4, 8).

Gâtez les configs spécifiques à TileIR pour sm_100+ :

import torch

def get_configs_with_gating(pre_hook=None):
    configs = get_baseline_configs()
    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10:
        configs.extend(get_tileir_specific_configs(pre_hook))
    return configs

Voir references/config-templates.md pour les templates de config complets par type de kernel.

Phase 4 : Validation TileIR (ENABLE_TILE=1)

Utilisez la skill kernel-triton-writing avec verify_kernel.py pour vérifier le kernel optimisé avec le backend TileIR :

python scripts/verify_kernel.py --kernel path/to/optimized_kernel.py --reference 'torch reference' --shapes '{"x": [32, 512, 4096]}' --dtypes '{"x": "bfloat16"}'

Définissez ENABLE_TILE=1 avant d'exécuter. Vérifiez : correctness numérique, pas d'erreurs de compilation, les patterns TMA/2CTA compilent avec succès.

Phase 5 : Benchmark

Utilisez triton.testing.do_bench() (comme documenté dans la skill perf-workload-profiling) pour comparer PTX (ENABLE_TILE=0) vs TileIR (ENABLE_TILE=1).

Benchmarkez sur plusieurs tailles d'entrée (128, 1024, 8192) -- les performances varient selon la taille.

Scripts

tileir_check.py

Vérifiez la disponibilité de TileIR (nvtriton, ENABLE_TILE, GPU Blackwell) :

python scripts/tileir_check.py

Retourne JSON : nvtriton_installed, tileir_active, blackwell_gpu, gpu_capability, recommendation.

classify_kernel.py

Classifiez le type de kernel et optionnellement appliquez les optimisations TileIR :

# Classification uniquement
python scripts/classify_kernel.py --file kernel.py

# Classification + application des optimisations
python scripts/classify_kernel.py --file kernel.py --apply-optimizations

# Depuis du code inline
python scripts/classify_kernel.py --code '<kernel_code>'

Retourne JSON : classification, confidence, indicators, recommendations. Avec --apply-optimizations : ajoute optimized_code et changes_applied.

Gestion des erreurs

Pièges courants

Erreurs de descripteur TMA (kernels dot-related) :

• Toujours passer pre_hook=tma_set_block_size_hook à la génération de configs -- sans cela, les descripteurs TMA gardent des tailles de blocs fictives, causant des erreurs runtime ou des résultats incorrects.
• Pour GEMM : passer b.T.contiguous() dans le wrapper et utiliser tl.dot(a, b.T, accumulator) dans le kernel. Un manque de transposition produit des résultats incorrects silencieusement.

Sursouscription 2CTA :

• Ajustez le nombre de SMs dans pre-hook quand vous utilisez num_ctas=2 :

  if "NUM_SMS" in nargs and "NUM_CTAS" in nargs:
      nargs["NUM_SMS"] = nargs["NUM_SMS"] // nargs["NUM_CTAS"]

Signatures de fonction de config :

• TOUTES les fonctions helper de config DOIVENT accepter pre_hook=None, même si non utilisé. Sans cela : TypeError: get_autotune_configs() takes 0 positional arguments.

Gâtage du matériel :

• Gâtez les configs TileIR avec torch.cuda.get_device_capability()[0] >= 10. TMA/2CTA sur des GPUs pré-Blackwell causent des crashes runtime.

Disponibilité de l'API :

• Utilisez 1.0 / (1.0 + tl.exp(-x)) à la place de tl.sigmoid(x) -- non disponible dans toutes les versions de Triton y compris certains builds nvtriton.

Tuning de performance :

• Ne sur-tunez pas num_warps -- TileIR l'ignore. Concentrez-vous sur occupancy.
• Utilisez des tailles de blocs plus grandes (256x256, 256x128) pour TileIR, pas des petits blocs tunés pour PTX.
• Benchmarkez sur de petites/moyennes/grandes entrées ; les configs one-size sous-performent.
• Pour les kernels intensifs en exp/log, activez le calcul approché :

  export TILEIR_ENABLE_APPROX=1
  export TILEIR_ENABLE_FTZ=1

Quand arrêter

Arrêtez et rapportez si :

1. Pas de triton installé -- impossible de procéder.
2. Le test de compatibilité échoue -- le kernel a des erreurs de syntaxe/runtime avant optimisation.
3. La validation TileIR échoue -- le kernel optimisé produit des résultats incorrects.
4. Pas d'accélération -- la version TileIR est plus lente que la baseline PTX (avec nvtriton).
5. Pas de GPU Blackwell -- ajoutez quand même les configs pour un déploiement futur, mais ignorez les tests ENABLE_TILE et le benchmarking.

Format de sortie

Après optimisation, retournez :

## TileIR Optimization: kernel_name

### Classification
- Kernel type: [dot-related | norm-like | element-wise | reduction]
- Strategy: [TMA + 2CTA | High occupancy | Occupancy + num_stages]

### Compatibility Check (ENABLE_TILE=0)
[PASSED | FAILED] — Max difference: X.Xe-Y

### Transformations Applied
- [List of transformations]

### TileIR Validation (ENABLE_TILE=1)
[PASSED | FAILED] — Max difference: X.Xe-Y

### Benchmark Comparison
| Backend | Time (ms) | Speedup |
|---------|-----------|---------|
| PTX (ENABLE_TILE=0) | X.XXX | 1.0x |
| TileIR (ENABLE_TILE=1) | X.XXX | Y.Yx |

### Output
File: kernel_name_tileir.py