From 0e467955b541693051323cacf5650c1dbc5c73dc Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 5 Aug 2026 14:23:40 +0800 Subject: [PATCH 1/2] Handle tcgen05 API changes by signature Signed-off-by: shuyan.ycf --- cula/ops/_cutedsl_compat.py | 65 ++++++++++++++++++++++++++++++ cula/ops/sm100/ptx.py | 77 ++++++++++++++++++++++++++++-------- tests/test_cutedsl_compat.py | 65 ++++++++++++++++++++++++++++++ 3 files changed, 191 insertions(+), 16 deletions(-) create mode 100644 cula/ops/_cutedsl_compat.py create mode 100644 tests/test_cutedsl_compat.py diff --git a/cula/ops/_cutedsl_compat.py b/cula/ops/_cutedsl_compat.py new file mode 100644 index 00000000..66bbf0df --- /dev/null +++ b/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(), + ) diff --git a/cula/ops/sm100/ptx.py b/cula/ops/sm100/ptx.py index 1e95b412..2823b231 100644 --- a/cula/ops/sm100/ptx.py +++ b/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) diff --git a/tests/test_cutedsl_compat.py b/tests/test_cutedsl_compat.py new file mode 100644 index 00000000..1bca3ec0 --- /dev/null +++ b/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) From 1f4b3a7ce619d1d55c02dbea4618c68aa2adc285 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 5 Aug 2026 17:45:57 +0800 Subject: [PATCH 2/2] Fix SM100 intracard CP test dispatch Signed-off-by: shuyan.ycf --- tests/test_kda_sm100_intracard_cp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_kda_sm100_intracard_cp.py b/tests/test_kda_sm100_intracard_cp.py index bd44c322..3d3f6f98 100644 --- a/tests/test_kda_sm100_intracard_cp.py +++ b/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, )