Detect tcgen05 load/store API by signature

cula/ops/_cutedsl_compat.py

@@ -0,0 +1,65 @@
+# Copyright 2025-2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compatibility helpers for low-level CuTeDSL bindings."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from inspect import Parameter, signature
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class Tcgen05LdStApi:
+    """Detected keyword interface of the generated tcgen05 load/store ops."""
+
+    ld_has_num: bool
+    st_has_num: bool
+    st_value_keyword: Literal["r", "val"]
+
+
+def _parameter_names(op: Callable, op_name: str) -> set[str]:
+    try:
+        parameters = signature(op).parameters
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError(f"Unable to inspect the CuTeDSL {op_name} binding") from exc
+
+    if any(parameter.kind is Parameter.VAR_KEYWORD for parameter in parameters.values()):
+        raise RuntimeError(f"Unsupported CuTeDSL {op_name} signature: variadic keyword arguments are ambiguous")
+    return set(parameters)
+
+
+def detect_tcgen05_ldst_api(ld_op: Callable, st_op: Callable) -> Tcgen05LdStApi:
+    """Detect supported tcgen05 load/store keyword variants from their signatures."""
+
+    ld_parameters = _parameter_names(ld_op, "tcgen05_ld")
+    st_parameters = _parameter_names(st_op, "tcgen05_st")
+
+    missing_ld = {"res", "shape", "tmem_addr"} - ld_parameters
+    if missing_ld:
+        raise RuntimeError(f"Unsupported CuTeDSL tcgen05_ld signature: missing {sorted(missing_ld)}")
+
+    missing_st = {"shape", "tmem_addr"} - st_parameters
+    if missing_st:
+        raise RuntimeError(f"Unsupported CuTeDSL tcgen05_st signature: missing {sorted(missing_st)}")
+
+    value_keywords = {"r", "val"} & st_parameters
+    if len(value_keywords) != 1:
+        raise RuntimeError("Unsupported CuTeDSL tcgen05_st signature: expected exactly one value keyword from ['r', 'val']")
+
+    return Tcgen05LdStApi(
+        ld_has_num="num" in ld_parameters,
+        st_has_num="num" in st_parameters,
+        st_value_keyword=value_keywords.pop(),
+    )

cula/ops/sm100/ptx.py

@@ -64,9 +64,16 @@
 from cutlass.cute.typing import Int32
 from cutlass.cutlass_dsl import dsl_user_op
 
+from cula.ops._cutedsl_compat import detect_tcgen05_ldst_api
+
 CollectorBBuffer = _nvvm.Tcgen05MMACollectorBBuffer
 CollectorOp = _nvvm.Tcgen05MMACollectorOp
 
+_tcgen05_ldst_api = detect_tcgen05_ldst_api(_nvvm.tcgen05_ld, _nvvm.tcgen05_st)
+_TCGEN05_LD_HAS_NUM = _tcgen05_ldst_api.ld_has_num
+_TCGEN05_ST_HAS_NUM = _tcgen05_ldst_api.st_has_num
+_TCGEN05_ST_USES_R = _tcgen05_ldst_api.st_value_keyword == "r"
+
 
 def _to_ir(val, loc=None, ip=None):
     return val.ir_value(loc=loc, ip=ip) if hasattr(val, "ir_value") else val
@@ -168,14 +175,24 @@ def _do(addr_val, *, loc=None, ip=None):
         ptr6_ty = llvm.PointerType.get(address_space=6)
         tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip)
         vec_i32_ty = ir.VectorType.get([num], i32_ty)
-        return _nvvm.tcgen05_ld(
-            res=vec_i32_ty,
-            shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
-            num=num,
-            tmem_addr=tmem_ptr,
-            loc=loc,
-            ip=ip,
-        )
+        if cutlass.const_expr(_TCGEN05_LD_HAS_NUM):
+            result = _nvvm.tcgen05_ld(
+                res=vec_i32_ty,
+                shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                num=num,
+                tmem_addr=tmem_ptr,
+                loc=loc,
+                ip=ip,
+            )
+        else:
+            result = _nvvm.tcgen05_ld(
+                res=vec_i32_ty,
+                shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                tmem_addr=tmem_ptr,
+                loc=loc,
+                ip=ip,
+            )
+        return result
 
     return _do(Int32(taddr))
 
@@ -188,14 +205,42 @@ def tcgen05_st_32x32b(num: int, taddr: int, vec):
     def _do(addr_val, vec_val, *, loc=None, ip=None):
         ptr6_ty = llvm.PointerType.get(address_space=6)
         tmem_ptr = llvm.inttoptr(ptr6_ty, _to_ir(addr_val, loc, ip), loc=loc, ip=ip)
-        _nvvm.tcgen05_st(
-            shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
-            num=num,
-            tmem_addr=tmem_ptr,
-            r=_to_ir(vec_val, loc, ip),
-            loc=loc,
-            ip=ip,
-        )
+        if cutlass.const_expr(_TCGEN05_ST_HAS_NUM):
+            if cutlass.const_expr(_TCGEN05_ST_USES_R):
+                _nvvm.tcgen05_st(
+                    shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                    num=num,
+                    tmem_addr=tmem_ptr,
+                    r=_to_ir(vec_val, loc, ip),
+                    loc=loc,
+                    ip=ip,
+                )
+            else:
+                _nvvm.tcgen05_st(
+                    shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                    num=num,
+                    tmem_addr=tmem_ptr,
+                    val=_to_ir(vec_val, loc, ip),
+                    loc=loc,
+                    ip=ip,
+                )
+        else:
+            if cutlass.const_expr(_TCGEN05_ST_USES_R):
+                _nvvm.tcgen05_st(
+                    shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                    tmem_addr=tmem_ptr,
+                    r=_to_ir(vec_val, loc, ip),
+                    loc=loc,
+                    ip=ip,
+                )
+            else:
+                _nvvm.tcgen05_st(
+                    shape=_nvvm.Tcgen05LdStShape.SHAPE_32X32B,
+                    tmem_addr=tmem_ptr,
+                    val=_to_ir(vec_val, loc, ip),
+                    loc=loc,
+                    ip=ip,
+                )
 
     _do(Int32(taddr), vec)
 

tests/test_cutedsl_compat.py

@@ -0,0 +1,65 @@
+# Copyright 2025-2026 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from cula.ops._cutedsl_compat import Tcgen05LdStApi, detect_tcgen05_ldst_api
+
+
+def _legacy_ld(res, shape, num, tmem_addr, *, pack=None, half_split_offset=None):
+    pass
+
+
+def _legacy_st(shape, num, tmem_addr, r, *, unpack=None, half_split_offset=None):
+    pass
+
+
+def _inferred_ld(res, shape, tmem_addr, *, pack=None, offset=None):
+    pass
+
+
+def _inferred_st(shape, tmem_addr, val, *, unpack=None, offset=None):
+    pass
+
+
+def test_detect_tcgen05_ldst_legacy_api():
+    assert detect_tcgen05_ldst_api(_legacy_ld, _legacy_st) == Tcgen05LdStApi(
+        ld_has_num=True,
+        st_has_num=True,
+        st_value_keyword="r",
+    )
+
+
+def test_detect_tcgen05_ldst_inferred_api():
+    assert detect_tcgen05_ldst_api(_inferred_ld, _inferred_st) == Tcgen05LdStApi(
+        ld_has_num=False,
+        st_has_num=False,
+        st_value_keyword="val",
+    )
+
+
+def test_detect_tcgen05_ldst_mixed_api():
+    assert detect_tcgen05_ldst_api(_legacy_ld, _inferred_st) == Tcgen05LdStApi(
+        ld_has_num=True,
+        st_has_num=False,
+        st_value_keyword="val",
+    )
+
+
+def test_detect_tcgen05_ldst_rejects_unknown_store_value_keyword():
+    def unsupported_st(shape, tmem_addr, value):
+        pass
+
+    with pytest.raises(RuntimeError, match="expected exactly one value keyword"):
+        detect_tcgen05_ldst_api(_inferred_ld, unsupported_st)

tests/test_kda_sm100_intracard_cp.py

@@ -91,6 +91,7 @@ def make_varlen_inputs(seq_lens, H, *, use_gk=False, use_h0=False, seed=42):
 
 
 def run_cula_no_cp(k, w, u, gk, h0, cu, **kw):
+    kw["use_intracard_cp"] = False
     return chunk_gated_delta_rule_fwd_h(
         k=k,
         w=w,
@@ -99,7 +100,6 @@ def run_cula_no_cp(k, w, u, gk, h0, cu, **kw):
         initial_state=h0,
         chunk_size=BT,
         cu_seqlens=cu,
-        _no_cp=True,
         **kw,
     )
 
@@ -160,7 +160,7 @@ def run_intracard_direct(k, w, u, gk, h0, cu, *, output_final_state=True, save_n
             chunk_size=BT,
             save_new_value=save_new_value,
             cu_seqlens=cu,
-            _no_cp=True,
+            use_intracard_cp=False,
         )