demo131 - patch lua flags and includes

src/infiniop/ops/flash_attention/ninetoothed/build.py

@@ -6,16 +6,14 @@ import infiniop.ninetoothed.build
 
 import torch
 
+import os
 
-def build():
 
-    if torch.cuda.is_available():
-        device_count = torch.cuda.device_count()
-        for i in range(device_count):
-            device_name = torch.cuda.get_device_name(i).lower()
+def build():
 
-            if "metax" in device_name:
-                return
+    env_vars_to_check = ["MACA_HOME", "MACA_PATH", "MACA_ROOT"]
+    if any(var in os.environ for var in env_vars_to_check):
+        return
 
     with_kv_cache_values = (0,)
     emb_dim_values = (16, 32, 64, 128, 256)

xmake/hygon.lua

@@ -77,7 +77,7 @@ target("infiniop-hygon")
     add_files("../src/infiniop/ops/swiglu/nvidia/*.cu")
 
     if has_config("ninetoothed") then
-        add_files("../build/ninetoothed/*.c", {cxxflags = {"-Wno-return-type"}})
+        add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp", {cxxflags = {"-Wno-return-type"}})
     end
 target_end()
 

xmake/iluvatar.lua

@@ -58,7 +58,7 @@ target("infiniop-iluvatar")
     add_files("../src/infiniop/ops/dequantize_awq/iluvatar/*.cu")
 
     if has_config("ninetoothed") then
-        add_files("../build/ninetoothed/*.c", {cxxflags = {"-Wno-return-type"}})
+        add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp", {cxxflags = {"-Wno-return-type"}})
     end
 target_end()
 

xmake/metax.lua

@@ -54,7 +54,7 @@ target("infiniop-metax")
 
     if has_config("ninetoothed") then
         add_includedirs(MACA_ROOT .. "/include/mcr")
-        add_files("../build/ninetoothed/*.c", {
+        add_files("../build/ninetoothed/*.c", "../build/ninetoothed/*.cpp", {
             cxflags = {
                 "-include stdlib.h", 
                 "-Wno-return-type",