From dd0e47e1859e8f563fcc3a10a34020b781ebfb15 Mon Sep 17 00:00:00 2001
From: Ying Zhang <yingz@meta.com>
Date: Mon, 11 Sep 2023 17:45:17 -0700
Subject: [PATCH] Add PyTorch fp8 dtypes to Triton

---
 python/triton/runtime/jit.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py
index d306c160e8ee..f664b33c71ab 100644
--- a/python/triton/runtime/jit.py
+++ b/python/triton/runtime/jit.py
@@ -250,6 +250,8 @@ def _type_of(key):
             "float8e5": "fp8e5",
             "float8e4b15": "fp8e4b15",
             "float8e4b15x4": "fp8e4b15x4",
+            "float8_e4m3fn": "fp8e4nv",
+            "float8_e5m2": "fp8e5",
             "float16": "fp16",
             "bfloat16": "bf16",
             "float32": "fp32",