[_sunrise] update commits of sunrise backend. [till 20260812] (#5478)

* #19382 tang -> ptpu

* add sunrise's mode. use sunrise's sort

* [#19375] update attention

* #11009, #11010, #11012, #10976 optimize performence

Author: Dayrker

src/flag_gems/runtime/backend/_sunrise/__init__.py

@@ -23,7 +23,7 @@
     vendor_name="sunrise",
     device_name="ptpu",
     device_query_cmd="pt_smi",
-    triton_extra_name="tang",
+    triton_extra_name="ptpu",
     dispatch_key="PrivateUse1",
     fp64_enabled=False,
     bf16_enabled=False,
@@ -102,6 +102,8 @@ def _sunrise_monkey_patch_enabled():
     "log_sigmoid",
     "margin_ranking_loss",
     "max_pool3d_with_indices_backward",
+    "mm",
+    "mm_out",
     "nonzero_numpy",
     "pad",
     "prelu",
@@ -111,6 +113,7 @@ def _sunrise_monkey_patch_enabled():
     "repeat_interleave.self_Tensor",
     "resolve_conj",
     "resolve_neg",
+    "router_gemm",
     "selu",
     "selu_",
     "square",

src/flag_gems/runtime/backend/_sunrise/fused/sparse_attention.py

@@ -134,7 +134,10 @@ def sparse_attn_triton(
     b, m, h, d = q.shape
     topk = topk_idxs.shape[-1]
     o = torch.empty_like(q)
-    BLOCK = 64
+    # A 64-wide KV tile exceeds the 128 KiB shared-memory limit on PTPU
+    # for the D=512 sparse-attention shapes. The single pipeline stage is
+    # also required for correct results with the smaller tile on PTPU.
+    BLOCK = 16
 
     grid = (m, b)  # each program handles ALL h heads
     if h < 8:
@@ -169,6 +172,7 @@ def sparse_attn_triton(
             D=d,
             H=8,
             num_warps=8,  # 256 threads, matching tilelang
+            num_stages=1,
         )
         o = o_new[:, :, :h].contiguous()
         return o
@@ -199,5 +203,6 @@ def sparse_attn_triton(
         D=d,
         H=h,
         num_warps=8,  # 256 threads, matching tilelang
+        num_stages=1,
     )
     return o

src/flag_gems/runtime/backend/_sunrise/monkey_patch.py

@@ -637,6 +637,78 @@ def randn_with_ptpu_complex_cpu_fallback(*args, **kwargs):
     setattr(torch, patched_attr, True)
 
 
+def _patch_torch_abs_long_runtime_error():
+    """Run ``torch.abs`` on CPU for PTPU int64 eager helper tensors.
+
+    Sunrise's eager ``UNARY_ABS`` kernel raises a plain ``RuntimeError`` for
+    ``torch.long`` inputs. Keep the workaround outside ``use_gems()`` so it
+    only covers setup/reference helpers such as flash-attention ALiBi bias
+    construction and cannot hide a missing FlagGems ``abs`` implementation.
+    """
+    patched_attr = "_flag_gems_sunrise_abs_long_runtime_error_patched"
+    if getattr(torch, patched_attr, False):
+        return
+
+    original_fn = torch.abs
+    runtime_marker = "unary_op<ptpu_kernel::UNARY_ABS>"
+    dtype_marker = "failed to dispatch data type Long"
+
+    @functools.wraps(original_fn)
+    def abs_with_ptpu_long_cpu_fallback(*args, **kwargs):
+        tensor = args[0] if args else kwargs.get("input")
+        try:
+            return original_fn(*args, **kwargs)
+        except RuntimeError as exc:
+            message = str(exc)
+            if (
+                _flag_gems_use_gems_active()
+                or not _is_ptpu_tensor(tensor)
+                or tensor.dtype != torch.long
+                or runtime_marker not in message
+                or dtype_marker not in message
+            ):
+                raise
+            return _torch_function_cpu_fallback(tensor, args, kwargs, original_fn)
+
+    torch.abs = abs_with_ptpu_long_cpu_fallback
+    setattr(torch, patched_attr, True)
+
+
+def _patch_torch_all_keepdim_runtime_error():
+    """Run unsupported PTPU bool ``torch.all(..., keepdim=True)`` on CPU.
+
+    Sunrise eager reduction raises a plain ``RuntimeError`` for this form.
+    Limit the fallback to the exact runtime message and reference/setup code
+    outside ``use_gems()`` so FlagGems' real ``all`` kernel remains visible.
+    """
+    patched_attr = "_flag_gems_sunrise_all_keepdim_runtime_error_patched"
+    if getattr(torch, patched_attr, False):
+        return
+
+    original_fn = torch.all
+    runtime_marker = "all_out with keepdim true is not implemented yet."
+
+    @functools.wraps(original_fn)
+    def all_with_ptpu_keepdim_cpu_fallback(*args, **kwargs):
+        tensor = args[0] if args else kwargs.get("input")
+        keepdim = args[2] if len(args) > 2 else kwargs.get("keepdim", False)
+        try:
+            return original_fn(*args, **kwargs)
+        except RuntimeError as exc:
+            if (
+                _flag_gems_use_gems_active()
+                or not _is_ptpu_tensor(tensor)
+                or tensor.dtype != torch.bool
+                or keepdim is not True
+                or runtime_marker not in str(exc)
+            ):
+                raise
+            return _torch_function_cpu_fallback(tensor, args, kwargs, original_fn)
+
+    torch.all = all_with_ptpu_keepdim_cpu_fallback
+    setattr(torch, patched_attr, True)
+
+
 def _patch_torch_cudnn_convolution():
     """Run `torch.cudnn_convolution(...)` on CPU via `F.conv{1,2,3}d` for PTPU.
 
@@ -3365,6 +3437,7 @@ def apply_sunrise_monkey_patches():
     _patch_torch_function("logsumexp", "aten::amax.out")
     _patch_tensor_method("mean", "aten::mean")
     _patch_torch_function("mean", "aten::mean")
+    _patch_torch_all_keepdim_runtime_error()
     _patch_torch_function("norm", "aten::linalg_vector_norm.out")
     _patch_torch_linalg_function("vector_norm", "aten::linalg_vector_norm.out")
     _patch_torch_linalg_function("qr", "aten::linalg_qr.out")
@@ -3403,6 +3476,7 @@ def apply_sunrise_monkey_patches():
     _patch_torch_nn_functional("logsigmoid", "aten::log_sigmoid_forward")
     _patch_torch_nn_functional_one_hot_cpu_reference()
     _patch_torch_randn_complex_dtype()
+    _patch_torch_abs_long_runtime_error()
     _patch_torch_cudnn_convolution()
     _patch_conv_depthwise2d_cpu_reference()
     _patch_thnn_fused_lstm_cell_cpu_reference()

src/flag_gems/runtime/backend/_sunrise/ops/__init__.py

@@ -127,6 +127,8 @@
 from .max_pool3d_with_indices import max_pool3d_backward, max_pool3d_with_indices
 from .mean import mean, mean_dim
 from .median import median, median_dim, median_dim_values, median_out
+from .mm import mm, mm_out, router_gemm
+from .mode import mode
 from .mul import mul, mul_
 from .multinomial import multinomial
 from .multiply_ import multiply_
@@ -329,6 +331,9 @@
     "median_dim",
     "median_dim_values",
     "median_out",
+    "mode",
+    "mm",
+    "mm_out",
     "mul",
     "mul_",
     "multiply_",
@@ -360,6 +365,7 @@
     "rms_norm",
     "rms_norm_forward",
     "rms_norm_backward",
+    "router_gemm",
     "scaled_grouped_mm",
     "scaled_dot_product_attention",
     "scaled_dot_product_attention_backward",