TensorGPT: Generative Tensor Program Optimization for Heterogeneous Edge Devices

Summary

• Problem: Existing generative tensor compilers use weak hardware representations, so schedules often fail to transfer across CPUs, GPUs, and edge GPUs.
• Approach: use richer representation to condition schedule generation on both target hardware, rather than relying only on raw target strings.
  • Technique #1: Gives the generator coarse hardware family cues as special tokens and feed fine-grained resource constraints as KV prefixes.
  • Technique #2: Use a mixture of LoRA experts to store hardware-specific knowledge, and route the generator to mix experts based on hardware and workload.
• Evaluation: On held-out Orin and Ryzen targets, TensorGPT improves validity and latency over TLM/TLM-CT, and requires less measurment budget than Ansor.

Problem: Existing Generative Tensor Compilers Model Hardware Poorly

• The same operator needs different schedules on target hardwares for high perf.
  • schedules: tiling, loop order, vectorization, thread binding, and memory placement.
• Existing generative tensor compilers (e.g., TLM) take a raw string as the target hardware representation, which is insufficient to capture the hardware semantics.
  • e.g., llvm -keys=cpu -mcpu=core-avx2

Approach: Richer Harwdware Conditioned Generation

Technique #1: Dual-Channel Hardware Grounding

Challenge #1: How can the generator understand the target hardware beyond a target string?

Add model-facing hardware information through two channels.

1. coarse bucket tokens (e.g., HW_GPU_EDGE)
2. continuous hardware KV prefixes from a 24-D hardware vector

Technique #2: Mixture of Hardware Experts

Challenge #2: How can one generator reuse tuning knowledge from multiple source devices?

• Store source-device knowledge as LoRA residual experts.
• A router mixes experts using hardware and workload evidence.

Evaluation: Key Results Setup

• Source devices: RTX 4090, Jetson Xavier, Xeon.
• Held-out targets: Jetson Orin, Ryzen 5800H.
• Baselines: Ansor, TLM, TLM-CT
  • TLM-CT: TLM further trained with additional data used for TensorGPT training
• Workload: BERT-style Transformers, DenseNet, MobileNetV2, ResNet, ResNet3D
• Metrics: (normalized) latency, validity

Evaluation: Key Results

Interpretation: Improved validity over TLM, up to 1.5x ish speedup for specific workloads, around 1.1x speedup for representative subgraphs

Evaluation: Ablation Study

• Ryzen CPU, normalized wrt Ansor@1000
• Baselines:
  • w/o HG: remove dual-channel hardware-grounding
  • w/o Routing: disable residual expert composition
  • w/o WA: keeps experts but remove workload-aware routing

Interpretation: Hardware grounding mainly restores CPU validity. LoRA experts and workload-aware routing improve latency