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

Ason93 · 2025-11-03T08:30:41Z

PR Category

Type of Change

Description

Issue

Progress

Change is properly reviewed (1 reviewer required, 2 recommended).
Change is responded to an issue.
Change is fully covered by a UT.

Performance

gemini-code-assist · 2025-11-03T08:30:55Z

Summary of Changes

Hello @Ason93, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on enhancing the performance of various comparison operations (greater than or equal to, greater than, less than or equal to, less than, and not equal to) within the Kunlunxin backend. The improvement is achieved by programmatically enabling a specific Triton XPU fusion optimization flag (TRITONXPU_COMPARE_FUSION) during the execution of these operations, which is expected to lead to more efficient computation.

Highlights

Performance Optimization: Introduced the use of the TRITONXPU_COMPARE_FUSION environment variable to enable specific fusion optimizations for comparison operations within the Kunlunxin backend.
Scope of Application: The TRITONXPU_COMPARE_FUSION flag is now dynamically set and unset around the execution of ge, gt, le, lt, and ne operations, covering both their tensor and scalar variants.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request aims to speed up comparison operations on Kunlunxin hardware by enabling TRITONXPU_COMPARE_FUSION. The approach of setting an environment variable is sound, but the implementation has several issues. There is a critical bug in ne.py where os.environ is modified inside a Triton kernel, which is not allowed. Additionally, the method used to set and unset the environment variable is not robust; it's not exception-safe and doesn't correctly restore the previous state of the environment variable. This could lead to hard-to-debug issues. The repeated logic also makes the code harder to maintain. My review includes suggestions to fix these problems by using try...finally blocks for safe cleanup and state restoration, and I recommend abstracting this logic into a reusable utility to improve code quality.

gemini-code-assist · 2025-11-03T08:33:40Z

src/flag_gems/runtime/backend/_kunlunxin/ops/ne.py

 def ne_func_scalar(x, y):
-    return x.to(tl.float32) != y.to(tl.float32)
+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
+    res = x.to(tl.float32) != y.to(tl.float32)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]
+    return res


This change is incorrect. The os module cannot be used inside a Triton JIT-compiled kernel (@triton.jit). This will cause a runtime error. The logic for setting and unsetting the environment variable should be in the host-side Python wrapper function, ne_scalar, not inside the kernel ne_func_scalar.

Please revert the changes in ne_func_scalar. The environment variable logic should be applied in the ne_scalar function, which is currently not modified in this PR. It should be changed to mirror the pattern in other files.

Suggested change

def ne_func_scalar(x, y):

return x.to(tl.float32) != y.to(tl.float32)

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = x.to(tl.float32) != y.to(tl.float32)

del os.environ["TRITONXPU_COMPARE_FUSION"]

return res

def ne_func_scalar(x, y):

return x.to(tl.float32) != y.to(tl.float32)

gemini-code-assist · 2025-11-03T08:33:40Z

src/flag_gems/runtime/backend/_kunlunxin/ops/ge.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
    res = ge_func(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


Modifying os.environ directly like this has a few issues:

Exception Safety: If ge_func(A, B) raises an exception, the del os.environ[...] line will not be executed, leaving the environment variable set. This could affect other parts of the program. This should be wrapped in a try...finally block.

State Restoration: The code deletes the environment variable, but it might have been set to a different value before this function was called. The correct behavior is to restore its original state.

Concurrency: os.environ is a global, process-wide dictionary. Modifying it without locks is not thread-safe and can lead to race conditions if these functions are called from multiple threads.

To address these issues and also reduce code duplication across files, I recommend creating a thread-safe context manager. As a direct fix, you can use a try...finally block and store/restore the previous value.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = ge_func(A, B)

del os.environ["TRITONXPU_COMPARE_FUSION"]

key = "TRITONXPU_COMPARE_FUSION"

original_value = os.environ.get(key)

os.environ[key] = "1"

try:

res = ge_func(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:40Z

src/flag_gems/runtime/backend/_kunlunxin/ops/ge.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
+    res = ge_func_scalar(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


Similar to the comment above, this logic should be made robust against exceptions and should restore the original state of the environment variable. Using a try...finally block is recommended to ensure correctness and prevent side effects.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = ge_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 = ge_func_scalar(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:40Z

src/flag_gems/runtime/backend/_kunlunxin/ops/gt.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
    res = gt_func(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


This logic for modifying os.environ is not exception-safe and does not properly restore the previous state. Please see my comment in ge.py for a detailed explanation and a suggested fix. The same robust approach should be applied here.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = gt_func(A, B)

del os.environ["TRITONXPU_COMPARE_FUSION"]

key = "TRITONXPU_COMPARE_FUSION"

original_value = os.environ.get(key)

os.environ[key] = "1"

try:

res = gt_func(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:40Z

src/flag_gems/runtime/backend/_kunlunxin/ops/gt.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
+    res = gt_func_scalar(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


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 = gt_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 = gt_func_scalar(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:41Z

src/flag_gems/runtime/backend/_kunlunxin/ops/le.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
    res = le_func(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


This logic for modifying os.environ is not exception-safe and does not properly restore the previous state. Please see my comment in ge.py for a detailed explanation and a suggested fix. The same robust approach should be applied here.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = le_func(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(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:41Z

src/flag_gems/runtime/backend/_kunlunxin/ops/le.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
+    res = le_func_scalar(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


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

gemini-code-assist · 2025-11-03T08:33:41Z

src/flag_gems/runtime/backend/_kunlunxin/ops/lt.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
    res = lt_func(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


This logic for modifying os.environ is not exception-safe and does not properly restore the previous state. Please see my comment in ge.py for a detailed explanation and a suggested fix. The same robust approach should be applied here.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = lt_func(A, B)

del os.environ["TRITONXPU_COMPARE_FUSION"]

key = "TRITONXPU_COMPARE_FUSION"

original_value = os.environ.get(key)

os.environ[key] = "1"

try:

res = lt_func(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:41Z

src/flag_gems/runtime/backend/_kunlunxin/ops/lt.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
+    res = lt_func_scalar(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


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 = lt_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 = lt_func_scalar(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

gemini-code-assist · 2025-11-03T08:33:41Z

src/flag_gems/runtime/backend/_kunlunxin/ops/ne.py

+    os.environ["TRITONXPU_COMPARE_FUSION"] = "1"
    res = ne_func(A, B)
+    del os.environ["TRITONXPU_COMPARE_FUSION"]


This logic for modifying os.environ is not exception-safe and does not properly restore the previous state. Please see my comment in ge.py for a detailed explanation and a suggested fix. The same robust approach should be applied here.

Suggested change

os.environ["TRITONXPU_COMPARE_FUSION"] = "1"

res = ne_func(A, B)

del os.environ["TRITONXPU_COMPARE_FUSION"]

key = "TRITONXPU_COMPARE_FUSION"

original_value = os.environ.get(key)

os.environ[key] = "1"

try:

res = ne_func(A, B)

finally:

if original_value is None:

os.environ.pop(key, None)

else:

os.environ[key] = original_value

[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION

409e681

gemini-code-assist bot reviewed Nov 3, 2025

View reviewed changes

[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_FP16_FAST

965ebf6

Ason93 changed the title ~~[KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION~~ [KUNLUNXIN] Speed Up ge/gt/le/lt/ne with TRITONXPU_COMPARE_FUSION and TRITONXPU_FP16_FAST Nov 4, 2025

nianqi-tian self-requested a review November 4, 2025 09:07

nianqi-tian approved these changes Nov 4, 2025

View reviewed changes

DuanYaQi approved these changes Nov 4, 2025

View reviewed changes

nianqi-tian merged commit c0ec9d8 into master Nov 4, 2025
17 of 22 checks passed

nianqi-tian deleted the klx_speed_ge_3 branch November 4, 2025 09:07

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

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

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

Uh oh!

Conversation

Ason93 commented Nov 3, 2025

PR Category

Type of Change

Description

Issue

Progress

Performance

Uh oh!

gemini-code-assist bot commented Nov 3, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants