[BUG] TMA cache policy requires materialized Int64
bug? - Needs TriageCuTe DSL
### 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 条评论