Update README.md

README.md

@@ -3,7 +3,65 @@ license: apache-2.0
 datasets:
 - ScalingIntelligence/KernelBench
 - LiShangZ/TritonBench
+pipeline_tag: text-generation
+library_name: transformers
+tags:
+- code-generation
 ---
-## AutoTriton
+# AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs
 
-AutoTriton is an 8B parameter model for Triton programming, which is trained based on Seed-Coder-8B-Reasoning via supervised fine-tuning and reinforcement learning sequentially. We leverage TritonBench and KernelBench as our benchmarks. For more details, please see our repo at [https://github.com/AI9Stars/AutoTriton](https://github.com/AI9Stars/AutoTriton).
+This repository contains the **AutoTriton** model, an 8B parameter model for Triton programming, which is trained based on Seed-Coder-8B-Reasoning via supervised fine-tuning and reinforcement learning sequentially.
+
+The model was presented in the paper [AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs](https://huggingface.co/papers/2507.05687).
+
+## Model Overview
+
+AutoTriton is the first model dedicated to Triton programming powered by reinforcement learning (RL). It addresses the complex challenges in deep learning kernel development by automating the optimization of computational units, memory management, parallelism, and hardware-specific parameters that typically require extensive manual tuning.
+
+The model's training process involves two sequential stages:
+1.  **Supervised Fine-Tuning (SFT)**: AutoTriton is first equipped with essential Triton programming expertise using a high-quality data gathering pipeline.
+2.  **Reinforcement Learning (RL)**: It then undergoes RL with the Group Relative Policy Optimization (GRPO) algorithm, combining a rule-based reward and an execution-based reward to further enhance its Triton programming ability.
+
+This approach underscores the promise of RL for automatically generating high-performance kernels, which are core components for building more efficient AI systems.
+
+## Evaluation
+
+Experiments across five evaluation channels of TritonBench and KernelBench illustrate that the 8B AutoTriton model achieves performance comparable to mainstream large models, including Claude-4-Sonnet and DeepSeek-R1-0528. Further analysis highlights the crucial role of each module within AutoTriton, including the SFT stage, the RL stage, and the reward design strategy.
+
+## Usage
+
+This model is compatible with the `transformers` library and can be loaded and used as follows:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "ai9stars/AutoTriton" # Replace with the actual model ID if different
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Example usage for Triton kernel code generation
+prompt = "Use triton language to write an add kernel for me."
+
+inputs = tokenizer(prompt, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=100)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+## GitHub Repository
+
+For more details on the project, including the full code, training scripts, and additional benchmarks, please refer to the official GitHub repository:
+[https://github.com/AI9Stars/AutoTriton](https://github.com/AI9Stars/AutoTriton)
+
+## Citation
+
+If you find this work useful, please consider citing our paper:
+
+```bibtex
+@article{li2025autotriton,
+  title={AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs},
+  author={Li, Shangzhan and Wang, Zefan and He, Ye and Li, Yuxuan and Shi, Qi and Li, Jianling and Hu, Yonggang and Che, Wanxiang and Han, Xu and Liu, Zhiyuan and others},
+  journal={arXiv preprint arXiv:2507.05687},
+  year={2025}
+}
+```