Merge branch 'flagos-ai:master' into fix/benchmark-bugs

Author: wuhulalala

src/flag_gems/fused/DSA/bin_topk.py

@@ -3,6 +3,7 @@
 from flag_gems.fused.concat_and_cache_mla import concat_and_cache_mla
 from flag_gems.fused.cross_entropy_loss import cross_entropy_loss
 from flag_gems.fused.cutlass_scaled_mm import cutlass_scaled_mm
+from flag_gems.fused.DSA.bin_topk import bucket_sort_topk
 from flag_gems.fused.FLA import (
     chunk_gated_delta_rule_fwd,
     fused_recurrent_gated_delta_rule_fwd,
@@ -43,6 +44,7 @@
     "apply_repetition_penalties",
     "apply_rotary_pos_emb",
     "bincount",
+    "bucket_sort_topk",
     "chunk_gated_delta_rule_fwd",
     "concat_and_cache_mla",
     "cutlass_scaled_mm",

src/flag_gems/fused/__init__.py

@@ -6,27 +6,9 @@
 import triton
 import triton.language as tl
 
-from flag_gems.utils import libentry, libtuner
-
-
-def _triton_version_at_least(major: int, minor: int, patch: int = 0) -> bool:
-    version = str(getattr(triton, "__version__", "0.0.0")).split("+", 1)[0]
-    parts = version.split(".")
-    parsed = []
-    for part in parts[:3]:
-        digits = []
-        for ch in part:
-            if ch.isdigit():
-                digits.append(ch)
-            else:
-                break
-        parsed.append(int("".join(digits)) if digits else 0)
-    while len(parsed) < 3:
-        parsed.append(0)
-    return tuple(parsed) >= (major, minor, patch)
-
-
-if _triton_version_at_least(3, 6, 0):
+from flag_gems.utils import has_triton_tle, libentry, libtuner
+
+if has_triton_tle(3, 6, 0):
     try:
         import triton.experimental.tle.language as tle
         import triton.experimental.tle.language.gpu as tleg

src/flag_gems/fused/moe_align_block_size.py

@@ -14,7 +14,11 @@
 )
 from flag_gems.utils.triton_driver_helper import get_device_properties
 from flag_gems.utils.triton_lang_helper import tl_extra_shim
-from flag_gems.utils.triton_version_utils import HAS_TLE, _triton_version_at_least
+from flag_gems.utils.triton_version_utils import (
+    HAS_TLE,
+    _triton_version_at_least,
+    has_triton_tle,
+)
 
 __all__ = [
     "libentry",
@@ -31,5 +35,6 @@
     "get_device_properties",
     "tl_extra_shim",
     "_triton_version_at_least",
+    "has_triton_tle",
     "HAS_TLE",
 ]

src/flag_gems/utils/__init__.py

@@ -1,28 +1,36 @@
+import re
+
 import triton
+from packaging.version import InvalidVersion, Version
+
+
+def _coerce_triton_version(version: str) -> Version:
+    try:
+        return Version(version)
+    except InvalidVersion:
+        release = []
+        for part in version.split("+", 1)[0].split(".")[:3]:
+            match = re.match(r"\d+", part)
+            release.append(match.group(0) if match else "0")
+        while len(release) < 3:
+            release.append("0")
+        return Version(".".join(release))
 
 
 def _triton_version_at_least(major: int, minor: int, patch: int = 0) -> bool:
-    version = str(getattr(triton, "__version__", "0.0.0")).split("+", 1)[0]
-    parts = version.split(".")
-    parsed = []
-    for part in parts[:3]:
-        digits = []
-        for ch in part:
-            if ch.isdigit():
-                digits.append(ch)
-            else:
-                break
-        parsed.append(int("".join(digits)) if digits else 0)
-    while len(parsed) < 3:
-        parsed.append(0)
-    return tuple(parsed) >= (major, minor, patch)
-
-
-HAS_TLE = False
-if _triton_version_at_least(3, 1, 0):
+    version = str(getattr(triton, "__version__", "0.0.0"))
+    return _coerce_triton_version(version) >= Version(f"{major}.{minor}.{patch}")
+
+
+def has_triton_tle(major: int = 0, minor: int = 0, patch: int = 0) -> bool:
+    if not _triton_version_at_least(major, minor, patch):
+        return False
     try:
         import triton.experimental.tle.language as _tle  # noqa: F401
 
-        HAS_TLE = True
+        return True
     except ImportError:
-        pass
+        return False
+
+
+HAS_TLE = has_triton_tle()

src/flag_gems/utils/triton_version_utils.py

@@ -4,7 +4,9 @@
 from datetime import datetime
 
 import pytest
-import torch
+
+# TODO(Qiming): Try remove this line
+import torch  # noqa: F401
 import yaml
 
 import flag_gems
@@ -44,16 +46,9 @@ def pytest_addoption(parser):
     )
 
     parser.addoption(
-        (
-            "--mode"
-            if not (flag_gems.vendor_name == "kunlunxin" and torch.__version__ < "2.5")
-            else "--fg_mode"
-        ),  # TODO: fix pytest-* common --mode args,
-        action="store",
-        default="normal",
-        required=False,
-        choices=["normal", "quick"],
-        help="run tests on normal or quick mode",
+        "--quick",
+        action="store_true",
+        help="run tests on quick mode",
     )
 
     parser.addoption(
@@ -85,7 +80,7 @@ def pytest_configure(config):
 
     RECORD_LOG = config.getoption("--record") == "log"
     TO_CPU = config.getoption("--ref") == "cpu"
-    QUICK_MODE = config.getoption("--mode") == "quick"
+    QUICK_MODE = config.getoption("--quick") is True
 
     if RECORD_LOG:
         RUNTEST_INFO = {}

tests/conftest.py

@@ -8,6 +8,7 @@
 from flag_gems.fused.DSA.bin_topk import (
     bucket_sort_topk,  # Replace with actual module name
 )
+from flag_gems.fused.DSA.bin_topk import HAS_TLE
 
 
 def assert_set_similar(actual, expected, dtype, equal_nan=False):
@@ -134,6 +135,37 @@ def debug_topk_results(actual, expected, inputs, test_name=""):
         print(f"  Expected top values: {np.sort(expected_values)[-m:][::-1]}")
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device required")
+@pytest.mark.skipif(not HAS_TLE, reason="TLE bucket_sort_topk is unavailable")
+@pytest.mark.bucket_sort_topk
+@pytest.mark.parametrize(
+    ("starts_list", "ends_list"),
+    [
+        ([0, 0, 0], [512, 768, 1024]),
+        ([7, 31, 63], [700, 900, 1024]),
+    ],
+)
+def test_bucket_sort_topk_public_entrypoint_matches_torch_topk(starts_list, ends_list):
+    batch_size = len(starts_list)
+    seq_len = 1024
+    topk = 32
+    dtype = torch.float32
+
+    init_seed(2026)
+    inputs = torch.randn((batch_size, seq_len), dtype=dtype, device=device)
+    starts = torch.tensor(starts_list, dtype=torch.int32, device=device)
+    ends = torch.tensor(ends_list, dtype=torch.int32, device=device)
+
+    ref_indices = reference_topk_implementation(
+        to_reference(inputs), to_reference(starts), to_reference(ends), topk
+    )
+    actual_indices = bucket_sort_topk(inputs, starts, ends, topk)
+
+    assert actual_indices.shape == (batch_size, topk)
+    assert actual_indices.dtype == torch.int32
+    assert_set_similar(actual_indices, ref_indices, dtype)
+
+
 @pytest.mark.skip(
     "RuntimeError: Cannot call @triton.jit'd outside of the scope of a kernel"
 )

tests/test_DSA/test_bin_topk.py

@@ -93,7 +93,6 @@ def ref_paged_attn(
     return torch.cat(outputs, dim=0)
 
 
-@pytest.mark.skipif(True, reason="Dead code on line 221 fails to pass.")
 @pytest.mark.skipif(vendor_name == "kunlunxin", reason="RESULT TODOFIX")
 @pytest.mark.skipif(vendor_name == "hygon", reason="RESULT TODOFIX")
 @pytest.mark.flash_attn_varlen_func
@@ -229,6 +228,10 @@ def test_flash_attn_varlen_func(
                     softmax_scale=scale,
                     causal=causal,
                     window_size=window_size,
+                    block_table=block_tables,
+                    softcap=soft_cap if soft_cap is not None else 0,
+                    alibi_slopes=alibi_slopes,
+                    fa_version=2,
                 )
 
         ref_output = ref_paged_attn(

tests/test_flash_attn_varlen_func.py

@@ -81,7 +81,7 @@ coverage run -m pytest -s ${EXTRA_OPTS} ${TEST_CASES[@]}
 # Run quick-cpu test if necessary
 if [[ ${#TEST_CASES_CPU[@]} -ne 0 ]]; then
   echo "Running quick-cpu mode unit tests for ${TEST_CASES_CPU[@]}"
-  coverage run -m pytest -s ${EXTRA_OPTS} ${TEST_CASES_CPU[@]} --ref=cpu --mode=quick
+  coverage run -m pytest -s ${EXTRA_OPTS} ${TEST_CASES_CPU[@]} --ref=cpu --quick
 fi
 
 # Process coverage data only when full-range testing