Skip to content

[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION and TRITONXPU_FP16_FAST #1032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 4, 2025
Merged

[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION and TRITONXPU_FP16_FAST #1032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
409e681
[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION
Ason93 Nov 3, 2025
965ebf6
[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_FP16_FAST
Ason93 Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion src/flag_gems/runtime/backend/_kunlunxin/ops/ge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os

import triton
import triton.language as tl
Expand All @@ -16,7 +17,11 @@ def ge_func(x, y):

def ge(A, B):
logger.debug("GEMS GE")
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = ge_func(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res


Expand All @@ -28,4 +33,9 @@ def ge_func_scalar(x, y):

def ge_scalar(A, B):
logger.debug("GEMS GE SCALAR")
return ge_func_scalar(A, B)
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = ge_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res
12 changes: 11 additions & 1 deletion src/flag_gems/runtime/backend/_kunlunxin/ops/gt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os

import triton
import triton.language as tl
Expand All @@ -16,7 +17,11 @@ def gt_func(x, y):

def gt(A, B):
logger.debug("GEMS GT")
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = gt_func(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res


Expand All @@ -28,4 +33,9 @@ def gt_func_scalar(x, y):

def gt_scalar(A, B):
logger.debug("GEMS GT SCALAR")
return gt_func_scalar(A, B)
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = gt_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res
10 changes: 9 additions & 1 deletion src/flag_gems/runtime/backend/_kunlunxin/ops/le.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os

import triton
import triton.language as tl
Expand All @@ -16,7 +17,11 @@ def le_func(x, y):

def le(A, B):
logger.debug("GEMS LE")
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = le_func(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res


Expand All @@ -28,4 +33,7 @@ def le_func_scalar(x, y):

def le_scalar(A, B):
logger.debug("GEMS LE SCALAR")
return le_func_scalar(A, B)
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
res = le_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
Comment on lines +36 to +38
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This logic for modifying os.environ is not exception-safe and does not properly restore the previous state. The same robust approach suggested for other functions in this PR should be applied here.

Suggested change
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
res = le_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
key = "TRITONXPU_COMPARE_FUSION"
original_value = os.environ.get(key)
os.environ[key] = "1"
try:
res = le_func_scalar(A, B)
finally:
if original_value is None:
os.environ.pop(key, None)
else:
os.environ[key] = original_value

return res
12 changes: 11 additions & 1 deletion src/flag_gems/runtime/backend/_kunlunxin/ops/lt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os

import triton
import triton.language as tl
Expand All @@ -16,7 +17,11 @@ def lt_func(x, y):

def lt(A, B):
logger.debug("GEMS LT")
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = lt_func(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res


Expand All @@ -28,4 +33,9 @@ def lt_func_scalar(x, y):

def lt_scalar(A, B):
logger.debug("GEMS LT SCALAR")
return lt_func_scalar(A, B)
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = lt_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res
12 changes: 11 additions & 1 deletion src/flag_gems/runtime/backend/_kunlunxin/ops/ne.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import os

import triton
import triton.language as tl
Expand All @@ -16,7 +17,11 @@ def ne_func(x, y):

def ne(A, B):
logger.debug("GEMS NE")
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = ne_func(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res


Expand All @@ -28,4 +33,9 @@ def ne_func_scalar(x, y):

def ne_scalar(A, B):
logger.debug("GEMS NE SCALAR")
return ne_func_scalar(A, B)
os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
os.environ["TRITONXPU_FP16_FAST"] = "1"
res = ne_func_scalar(A, B)
del os.environ["TRITONXPU_COMPARE_FUSION"]
del os.environ["TRITONXPU_FP16_FAST"]
return res
Loading