ITADN

[BUG] TMA cache policy requires materialized Int64

#3252Closedgau-nernst 创建于 2026-05-20
bug? - Needs TriageCuTe DSL
G
gau-nernstcommented
### Which component has the problem? CuTe DSL ### Bug Report **Describe the bug** A clear and concise description of what the bug is. Update: looks like there is something strange with my environment. Will clear the cache and re-install the env **Steps/Code to reproduce bug** Follow this guide http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports to craft a minimal bug report. This helps us reproduce the issue you're having and resolve the issue more quickly. ```python import traceback import cutlass from cuda.bindings.driver import CUstream from cutlass import BFloat16, Int64, cute from cutlass.cute.nvgpu import cpasync from quack.compile_utils import make_fake_tensor EVICT_FIRST = Int64(0x12F0000000000000) class TmaCachePolicyRepro: def __init__(self, cache_policy_kind: str) -> None: self.cache_policy_kind = cache_policy_kind @cute.jit def __call__(self, src: cute.Tensor, stream: CUstream): slayout = cute.make_layout(128) atom, tma_tensor = cpasync.make_tiled_tma_atom( cpasync.CopyBulkTensorTileG2SOp(), src, slayout, cta_tiler=(128,), ) self.kernel(atom, tma_tensor, slayout).launch( grid=(1, 1, 1), block=(32, 1, 1), stream=stream ) @cute.kernel def kernel( self, atom: cute.CopyAtom, tma_tensor: cute.Tensor, slayout: cute.Layout ): smem = cutlass.utils.SmemAllocator() smem_tensor = smem.allocate_tensor(BFloat16, slayout, byte_alignment=128) mbar = smem.allocate_array(Int64, 1) with cute.arch.elect_one(): cute.arch.mbarrier_init(mbar, 1) cute.arch.mbarrier_init_fence() cute.arch.sync_threads() with cute.arch.elect_one(): cute.arch.mbarrier_arrive_and_expect_tx(mbar, 64 * 128 * 2) if cutlass.const_expr(self.cache_policy_kind == "materialized"): cache_policy = Int64(EVICT_FIRST.ir_value()) else: cache_policy = EVICT_FIRST s_part, g_part = cpasync.tma_partition( atom, 0, cute.make_layout(1), cute.group_modes(smem_tensor, 0), cute.group_modes(tma_tensor, 0), ) cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy) def main() -> None: for kind, label in ( ("plain", "plain Int64"), ("materialized", "materialized Int64"), ): print(f"\nCompiling with {label} cache_policy...") try: src = make_fake_tensor(BFloat16, (128,), divisibility=16) stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) cute.compile( TmaCachePolicyRepro(kind), src, stream, options="--enable-tvm-ffi", ) except Exception as exc: print(f"FAILED: {type(exc).__name__}: {exc}") for frame in traceback.extract_tb(exc.__traceback__)[-8:]: print(f" {frame.filename}:{frame.lineno} in {frame.name}") print(f" {frame.line}") else: print("PASSED") if __name__ == "__main__": main() ``` On 4.5.1 ```python Compiling with plain Int64 cache_policy... FAILED: AssertionError: /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/base_dsl/_mlir_helpers/op.py:154 in wrapper raise e /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/base_dsl/_mlir_helpers/op.py:121 in wrapper res_or_list = opFunc(*args, **kwargs, loc=loc) /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/cute/algorithm.py:609 in copy value = atom._unpack(loc=loc, ip=ip, **kwargs) /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/cute/atom.py:306 in _unpack return self._trait.unpack(loc=loc, ip=ip, **kwargs) /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/cute/nvgpu/cpasync/copy.py:261 in unpack exec_value = _cute_nvgpu_ir.atom_set_value( /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/_mlir/dialects/_cute_nvgpu_ops_gen.py:3279 in atom_set_value return AtomSetValueOp(atom=atom, field=field, value=value, loc=loc, ip=ip).result /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/_mlir/dialects/_cute_nvgpu_ops_gen.py:3247 in __init__ operands.append(_get_op_result_or_value(value)) /home/thien/vllm/gdn/.venv/lib/python3.12/site-packages/nvidia_cutlass_dsl/python_packages/cutlass/_mlir/dialects/_ods_common.py:104 in get_op_result_or_value assert isinstance(arg, _cext.ir.Value) Compiling with materialized Int64 cache_policy... PASSED ``` This does not happen on 4.5.0 **Expected behavior** A clear and concise description of what you expected to happen. Passing `cache_value = Int64(xxx)` should work **Environment details (please complete the following information):** - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)] **Additional context** Add any other context about the problem here.
关闭于 2026-05-20 1 条评论