From d7e7fa3efe47147190e11f4a08e667b02e0d9ed0 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 13:15:30 +0800 Subject: [PATCH 1/9] draft vmm --- CMakeLists.txt | 6 +- docs/suggestion.md | 339 +++++++++++++++++++++ testing/python/distributed/test_vmm_ops.py | 180 +++++++++++ tilelang/distributed/__init__.py | 2 + tilelang/distributed/utils.py | 51 +++- tilelang/utils/allocator.py | 75 +++-- tilelang/utils/ts_ext/__init__.py | 16 + tilelang/utils/ts_ext/exception.h | 14 + tilelang/utils/ts_ext/setup.py | 8 + tilelang/utils/ts_ext/ts_ext_bindings.cpp | 53 ++++ tilelang/utils/ts_ext/ts_ext_ops.h | 11 + tilelang/utils/ts_ext/vmm_ops.cpp | 189 ++++++++++++ tilescale_ext/__init__.py | 14 + 13 files changed, 932 insertions(+), 26 deletions(-) create mode 100644 docs/suggestion.md create mode 100644 testing/python/distributed/test_vmm_ops.py create mode 100644 tilelang/utils/ts_ext/vmm_ops.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 4fb370d509..7e9bb4fe77 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -356,6 +356,7 @@ if(USE_CUDA) ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/ts_ext_bindings.cpp ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/tensor.cpp ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/ipc_ops.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/vmm_ops.cpp ) # Find libtorch_python.so @@ -371,14 +372,15 @@ if(USE_CUDA) target_include_directories(tilescale_ext_C PRIVATE ${TORCH_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS} + ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext ) if(TORCH_PYTHON_RESULT EQUAL 0 AND EXISTS "${TORCH_PYTHON_LIBRARY}") message(STATUS "Found libtorch_python: ${TORCH_PYTHON_LIBRARY}") - target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} CUDA::cudart) + target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} CUDA::cudart CUDA::cuda_driver) else() message(WARNING "libtorch_python.so not found, extension may have undefined symbols") - target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} CUDA::cudart) + target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} CUDA::cudart CUDA::cuda_driver) endif() target_compile_options(tilescale_ext_C PRIVATE -fPIC) diff --git a/docs/suggestion.md b/docs/suggestion.md new file mode 100644 index 0000000000..39d8fac5bc --- /dev/null +++ b/docs/suggestion.md @@ -0,0 +1,339 @@ +# TileScale Codebase 整理与重构建议 + +## 1. 解决 repo 名称与 package 名称不一致问题 + +**现状：** repo 叫 `tilescale`，Python 包叫 `tilelang`，二者完全不同。 + +**建议：** +- 明确一个正式名称。如果 `tilelang` 是对外发布的包名，`tilescale` 是项目/团队代号，则在 repo 根目录的 README、CLAUDE.md 等地方明确说明这一关系。 +- 或者，将包名统一为 `tilescale`（如果不考虑向后兼容），避免新加入的开发者混淆。 + +--- + +## 2. 消除 `primitives/` 与 `tileop/` 的职责重叠 + +**现状：** +- `tilelang/primitives/gemm/__init__.py` 是一个几乎为空的目录（只有 `__init__.py`），且其中 import 的模块（`primitives.gemm.base`, `primitives.gemm.gemm_mma`）并不存在于该目录下，实际实现在 `tileop/gemm/` 里。 +- `tilelang/tileop/gemm/` 拥有完整实现：`gemm_mma.py`, `gemm_wgmma.py`, `gemm_mfma.py`, `gemm_tcgen05.py` 等。 + +**建议：** +- 删除空壳 `primitives/` 目录，将其职责完全归入 `tileop/`（或反之统一命名）。 +- 在公共 API 层统一暴露，不要让两个名字同时对外可见。 + +--- + +## 3. 清理 `language/v2/` —— 明确新旧版本策略 + +**现状：** 存在 `tilelang/language/` 和 `tilelang/language/v2/` 两套并行实现。 + +**建议：** +- 如果 v2 是稳定的新实现，应制定迁移计划：deprecate v1 接口，并在 v1 所有入口添加 deprecation warning，最终删除 v1 代码。 +- 如果 v2 还在实验阶段，应移入 `language/experimental/` 而非独立的 `v2/` 目录，避免用户误用。 +- 不要让两个版本长期共存且没有明确说明。 + +--- + +## 4. 合并分散的 `kernel_cache.py` + +**现状：** 至少 5 处有各自的 `kernel_cache.py`： +``` +tilelang/cache/kernel_cache.py +tilelang/jit/adapter/kernel_cache.py +tilelang/jit/adapter/cutedsl/kernel_cache.py +tilelang/jit/adapter/torch/kernel_cache.py +tilelang/jit/adapter/nvrtc/kernel_cache.py +tilelang/jit/adapter/cython/kernel_cache.py +``` + +**建议：** +- 设计一个统一的 `tilelang/cache/` 模块作为单一 cache 抽象层。 +- 各 adapter（cutedsl/torch/nvrtc/cython）继承或组合此基础实现，而不是各自维护一份。 +- 这样可以避免 cache invalidation 逻辑不同步的 bug。 + +--- + +## 5. 统一 Distributed 代码的位置 + +**现状：** 分布式相关代码散落在： +- `tilelang/distributed/` — 顶层 distributed 模块（含 pynvshmem 子包、build 脚本、install 脚本） +- `tilelang/language/distributed/` — language 层面的 distributed 原语 +- `tilelang/language/distributed/multi_device/` — nvshmem.py, cpengine.py +- `examples/distributed/` — 示例 + +**建议：** +- 明确职责边界：language 层只负责 IR/语法原语的声明，runtime 层（nvshmem binding、launch、allocator）归入 `tilelang/distributed/`。 +- `tilelang/language/distributed/` 只应包含 IR-level 原语定义，具体实现（nvshmem.py, cpengine.py）应上移至 `tilelang/distributed/`。 +- `pynvshmem` 是一个独立的 C 扩展包（有自己的 `setup.py` 和 `CMakeLists.txt`），应提升为顶层子包或独立 repo，而不是嵌套在 `tilelang/distributed/pynvshmem/` 多层深处。 + +--- + +## 6. 整理测试代码——测试不应混在 package 内部 + +**现状：** 测试文件散落在： +- `testing/python/` — 主测试目录（正确） +- `tilelang/distributed/testing/` — 在包内部！ +- `tilelang/distributed/pynvshmem/testing/` — 也在包内部！ +- `examples/` — 大量 `test_*.py` 和 `regression_*.py` 混在示例里 + +**建议：** +- 确立唯一的测试根目录：`testing/`（或 `tests/`）。 +- 将 `tilelang/distributed/testing/` 和 `tilelang/distributed/pynvshmem/testing/` 的内容迁移到 `testing/python/distributed/`。 +- `examples/` 里的 `test_*.py` / `regression_*.py` 要么移入 `testing/`，要么与示例文件分开存放（比如每个 example 子目录只保留 `example_*.py`，测试单独放 `testing/`）。 +- 在 `testing/` 下的目录结构应与 `tilelang/` 包结构一一对应，便于查找。 + +--- + +## 7. 修复 `__init__.py` 中的重复版本解析逻辑 + +**现状：** `tilelang/__init__.py` 第 11-65 行有两套 `__version__` 计算逻辑，`__version__` 被赋值两次（第 48 行和第 53-65 行），第一次的计算完全被覆盖。 + +**建议：** +- 只保留一套版本解析逻辑（推荐 `importlib.metadata` + fallback），删除重复代码。 +- 或者将版本逻辑抽到单独的 `_version.py` 文件。 + +--- + +## 8. 规范 `language/` 内部命名不一致问题 + +**现状：** 同类功能有时带 `_op` 后缀，有时不带： +- `language/copy.py` vs `language/copy_op.py` +- `language/fill.py` vs `language/fill_op.py` +- `language/gemm.py` vs `language/gemm_op.py` +- `language/reduce.py` vs `language/reduce_op.py` + +**建议：** +- 确定一个命名规范（建议 `_op` 后缀表示操作符对象，无后缀表示 language-level 函数入口），并全面统一。 +- 或者合并重复文件，明确每个文件的职责边界。 + +--- + +## 9. 整合 `contrib/cutedsl/` 与 `jit/adapter/cutedsl/` + +**现状：** +- `tilelang/contrib/cutedsl/` 有底层 CUDA 原语（mbar, ldsm, cpasync, gemm_V1, reduce 等） +- `tilelang/jit/adapter/cutedsl/` 有 adapter/wrapper/libgen 等 JIT 逻辑 + +**建议：** +- `contrib/cutedsl/` 的角色应该是"第三方或底层 CUDA primitive 封装"，`jit/adapter/cutedsl/` 是"JIT pipeline 对 cutedsl 的适配"，二者职责不同，但命名上容易让人以为是同一件事。 +- 建议将 `contrib/cutedsl/` 重命名为 `tilelang/backends/cutedsl/` 或 `tilelang/codegen/cutedsl/`，以更清晰表达其用途。 + +--- + +## 10. 整理 `examples/` 目录结构 + +**现状：** `examples/` 混杂了多种文件类型，缺乏一致结构： +- `example_*.py` — 示例代码 +- `test_example_*.py` — 测试 +- `regression_*.py` — 回归测试 +- `benchmark_*.py` — 性能基准 +- 部分有 `README.md`，部分没有 +- `examples/pytest.ini` 只有一个，但 examples 子目录结构不统一 + +**建议：** +- 每个 example 子目录统一结构：`example_*.py`（核心示例）+ `README.md`（必须有）。 +- 测试/回归测试移入 `testing/`，不要放在 examples 里。 +- Benchmark 单独建 `benchmarks/` 顶级目录（现有 `benchmark/blocksparse_attention/` 也移入此处），与 `examples/` 分开。 + +--- + +## 11. 清理重复图片资源 + +**现状：** 相似图片存在于多个地方： +- `images/`（根目录） +- `docs/_static/img/` +- `examples/deepseek_mla/figures/` +- `docs/_static/img/mla_hopper/` + +**建议：** +- 文档用图统一放 `docs/_static/img/`，examples 里的 figures 只保留真正 example 专属的。 +- 根目录 `images/` 合并进 `docs/_static/img/`，避免两套图片各自维护。 + +--- + +## 12. `tilelang/utils/ts_ext/` 的嵌套 `setup.py` 问题 + +**现状：** `tilelang/utils/ts_ext/` 内部有独立的 `setup.py`，说明这是一个独立的 C 扩展包，但被嵌套在 `utils/` 深处。 + +**建议：** +- 将 `ts_ext` 提升为顶级子包（类似 `pynvshmem` 的处理方式，或者更好地统一到同一个构建系统中）。 +- 或者把它的构建整合进主 `setup.py` / `CMakeLists.txt`，避免嵌套独立构建脚本。 + +--- + +## 优先级总结 + +| 优先级 | 建议 | +|--------|------| +| 🔴 高 | #7 重复版本逻辑（bug 风险）、#4 kernel_cache 重复（一致性风险）、#6 测试混在包内 | +| 🟡 中 | #2 primitives/tileop 重叠、#5 distributed 代码分散、#8 命名不一致、#3 v2 策略 | +| 🟢 低 | #1 名称统一、#9 contrib vs adapter、#10 examples 结构、#11 图片整理、#12 ts_ext 位置 | + +--- + +# TileScale 专项整理建议 + +## 13. 统一两个职责重叠的 `init_dist*` 函数 + +**现状：** `tilelang/distributed/utils.py` 中存在两个功能相似但 API 完全不同的初始化函数： +- `init_dist(local_rank, num_local_ranks)` — 旧式，需要手动传参，配合 `mp.spawn` 使用 +- `init_distributed(return_tp_group, init_nvshmem, return_lc_group)` — 新式，从环境变量读取，配合 `torchrun` 使用 + +两者在不同 example 中被混用，导致新用户不知道该用哪个。 + +**建议：** +- 统一为一个函数，明确以 `torchrun` 风格（环境变量）为主； +- `init_dist` 如果仍需保留，加 `@deprecated` 装饰器，并在 docstring 说明迁移路径。 + +--- + +## 14. 修复 `wait_eq` 名称严重冲突 + +**现状：** 两处都叫 `wait_eq`，但语义完全不同： +- `tilelang/language/distributed/common.py:wait_eq()` — **IR 层**：生成 NVSHMEM 的 `signal_wait_until` intrinsic，在 GPU 核函数内使用 +- `tilelang/distributed/utils.py:wait_eq()` — **Host 层**：调用 `cuStreamWaitValue32`，在 CPU 侧 stream 上等待 + +两者都被 import 进同一个项目，极易混淆，且 bug 难以定位。 + +**建议：** +- Host 侧重命名为 `stream_wait_eq` 或 `host_wait_signal`，与 IR 层的 `wait_eq` 区分； +- 在两个函数的 docstring 中明确标注各自的作用层级（host/device）。 + +--- + +## 15. 消除 IPC handle 同步逻辑的重复实现 + +**现状：** 以下两处独立实现了几乎相同的"收集各 rank 的 IPC handle 并映射到本地 GPU 指针"逻辑： +- `distributed/utils.py:create_dist_tensor()` — 通过 `all_gather_object` 收集 IPC handle，调用 `_sync_ipc_handles` +- `distributed/allocator.py:BaseAllocator._init_table()` — 相同流程，只是封装在 Allocator 里 + +**建议：** +- 将 IPC handle 同步逻辑抽成一个内部函数 `_exchange_ipc_handles(group, local_ptr) -> buffer_ptrs`； +- `create_dist_tensor` 和 `BaseAllocator._init_table` 都调用这个函数，避免两套独立维护。 + +--- + +## 16. 集中管理 `cuda-python` 版本兼容 shim + +**现状：** 以下代码段在多处重复出现： +```python +cuda_python_version = importlib.metadata.version("cuda-python") +from packaging import version +if version.parse(cuda_python_version) >= version.parse("12.8.0"): + from cuda.bindings import driver as cuda +else: + from cuda import cuda +``` +出现于 `distributed/utils.py`、`example_allgather_gemm_overlapped.py` 等多处。 + +**建议：** +- 将此兼容逻辑放在 `tilelang/distributed/_cuda_compat.py` 中，统一 export `cuda, cudart` 对象； +- 其他地方一律 `from tilelang.distributed._cuda_compat import cuda, cudart`，不再各自做版本检测。 + +--- + +## 17. 建立统一的对称内存（Symmetric Memory）抽象 + +**现状：** TileScale 目前有两套完全独立的分布式 tensor 分配路径： +- **NVSHMEM 路径**：`pynvshmem.nvshmem_create_tensor(shape, dtype)` — 通过 NVSHMEM 分配对称内存 +- **IPC 路径**：`tilelang.distributed.create_tensor(shape, dtype)` 或 `BaseAllocator` — 通过 `cudaMalloc` + IPC handle 交换 + +两者语义不同、使用方式不同，但 examples 中混用，用户难以理解应该用哪个，以及何时该回退到 IPC 路径。 + +**建议：** +- 设计一个统一的 `tilescale.memory.SymmetricBuffer` 抽象，内部根据是否有 nvshmem 自动选择后端； +- 在文档中明确说明两种路径的适用场景（intranode IPC vs internode NVSHMEM）。 + +--- + +## 18. 将 shell 脚本从 package 目录移出 + +**现状：** 以下三个 shell 脚本直接放在 `tilelang/distributed/` 包目录里，会被打包进 Python wheel： +- `tilelang/distributed/launch.sh` — 分布式启动脚本 +- `tilelang/distributed/build_nvshmem.sh` — 构建 NVSHMEM 的脚本 +- `tilelang/distributed/install_deepep.sh` — 安装 DeepEP 的脚本 + +**建议：** +- 统一移到 repo 根目录下的 `scripts/` 目录（如 `scripts/launch.sh`, `scripts/build_nvshmem.sh`）； +- `scripts/` 不会被 pip install 打包，符合标准 Python 项目规范。 + +--- + +## 19. 拆分过于庞杂的 `distributed/utils.py` + +**现状：** `tilelang/distributed/utils.py` 是一个 400 行的"万能工具箱"，混杂了： +- 进程组初始化（`init_dist`, `init_distributed`） +- IPC tensor 创建（`create_tensor`, `create_dist_tensor`） +- CUDA stream 信号操作（`set_signal`, `wait_eq`） +- 性能测量（`perf_fn`） +- 调试打印（`dist_print`） +- NVLink 拓扑检测（`has_fullmesh_nvlink`, `NvidiaSmiUtil`） +- 数据生成（`generate_data`, `_make_tensor`） + +**建议：** +``` +tilelang/distributed/ +├── init.py # init_dist, init_distributed +├── memory.py # IPC/tensor allocation +├── signal.py # CUDA stream signal operations +├── topo.py # NVLink topology (NvidiaSmiUtil, has_fullmesh_nvlink) +├── perf.py # perf_fn, dist_print +└── data.py # generate_data, _make_tensor +``` + +--- + +## 20. 将 example 中的复用库代码提升为正式模块 + +**现状：** 以下文件放在 `examples/distributed/` 目录中，却被其他 example 直接 import，本质上已经是库代码： +- `examples/distributed/sp_ag_attention_intra_node.py` — 被 `example_sp_ag_attention_intra_node.py` import +- `examples/distributed/nvshmem/gemm_rs_utils.py` — `BarrierAllContext`, `ReduceScatter2DContext` +- `examples/distributed/nvshmem/reduce_scatter.py` — `reduce_scatter_2d_op` 等 +- `examples/distributed/deepseek_deepep/deepep_utils.py` — `Config`, DeepEP 相关常量 + +这些文件用相对路径 import，只在同目录运行时有效，无法被其他地方正确引用。 + +**建议：** +- 将复用代码提升为 `tilelang/distributed/` 下的正式子模块； +- examples 只保留调用代码，不再充当隐式的库。 + +--- + +## 21. 清理 `launch.sh` 中泄露的内部集群环境变量 + +**现状：** `tilelang/distributed/launch.sh` 中有： +```bash +master_addr=${ARNOLD_WORKER_0_HOST:="127.0.0.1"} +master_port=$(echo "$ARNOLD_WORKER_0_PORT" | cut -d "," -f 1) +``` +`ARNOLD_*` 是 ByteDance 内部集群（Arnold 训练平台）的环境变量，不应出现在公开代码中。 + +**建议：** +- 将 `ARNOLD_WORKER_0_HOST` 替换为标准的 `MASTER_ADDR`，`ARNOLD_WORKER_0_PORT` 替换为 `MASTER_PORT`； +- 清理所有内部平台相关的硬编码假设（如 `IB_HCA=mlx5`、`BYTED_TORCH_BYTECCL` 等）。 + +--- + +## 22. 为 `tilescale_ext` C 扩展建立清晰的 API 文档与稳定性边界 + +**现状：** `tilelang/distributed/__init__.py` 直接暴露： +```python +from tilescale_ext import _create_tensor, _create_ipc_handle, _sync_ipc_handles +``` +- 以 `_` 开头的私有函数被直接 re-export 为公共 API +- 没有任何 stub 文件（`.pyi`）说明函数签名 +- 不清楚哪些是稳定 API，哪些是内部实现细节 + +**建议：** +- 在 `tilelang/distributed/` 层面做好封装，禁止外部直接调用 `_` 前缀的 `tilescale_ext` 函数； +- 提供 `tilescale_ext.pyi` stub 文件，便于 IDE 补全和类型检查； +- 明确区分稳定公共 API 与内部实现。 + +--- + +## 优先级汇总（TileScale 专项） + +| 优先级 | 建议 | +|--------|------| +| 🔴 高 | #14 wait_eq 名称冲突（正确性风险）、#15 IPC 逻辑重复（一致性风险）、#21 内部环境变量泄露 | +| 🟡 中 | #13 init_dist 统一、#16 cuda-python shim 集中化、#19 utils.py 拆分、#20 example 库代码提升 | +| 🟢 低 | #17 对称内存抽象、#18 shell 脚本移位、#22 tilescale_ext API 文档 | diff --git a/testing/python/distributed/test_vmm_ops.py b/testing/python/distributed/test_vmm_ops.py new file mode 100644 index 0000000000..149f124189 --- /dev/null +++ b/testing/python/distributed/test_vmm_ops.py @@ -0,0 +1,180 @@ +""" +VMM (Virtual Memory Management) operations test. + +Usage: + # Single-GPU unit tests (no torchrun needed): + python testing/python/distributed/test_vmm_ops.py + + # Multi-GPU integration test: + torchrun --nproc_per_node=8 testing/python/distributed/test_vmm_ops.py --distributed +""" + +import argparse +import os +import sys + +import torch +import torch.distributed as dist + + +def test_supports_fabric(): + """Test fabric support detection.""" + from tilescale_ext import _supports_vmm_fabric + + result = _supports_vmm_fabric() + print(f"[PASS] _supports_vmm_fabric() = {result}") + return result + + +def test_vmm_malloc_free(): + """Test VMM malloc and free roundtrip.""" + from tilescale_ext import _vmm_malloc, _vmm_free + + size = 1024 * 1024 # 1 MB + ptr = _vmm_malloc(size) + assert ptr != 0, "vmm_malloc returned null" + + # Verify the pointer is usable by writing via cudaMemset + import ctypes + import ctypes.util + + libcudart = ctypes.CDLL(ctypes.util.find_library("cudart") or "libcudart.so") + rc = libcudart.cudaMemset(ctypes.c_void_p(ptr), 0, ctypes.c_size_t(size)) + assert rc == 0, f"cudaMemset on VMM pointer failed: {rc}" + + _vmm_free(ptr) + print("[PASS] test_vmm_malloc_free") + + +def test_vmm_handle_export_import(): + """Test handle export and import on a single GPU.""" + from tilescale_ext import _vmm_malloc, _vmm_free, _create_vmm_handle, _open_vmm_handle, _close_vmm_handle + + size = 4096 + ptr = _vmm_malloc(size) + assert ptr != 0 + + # Write a known pattern + import ctypes + + pattern = (ctypes.c_uint8 * size)(*([0xAB] * size)) + libcudart = ctypes.CDLL(ctypes.util.find_library("cudart") or "libcudart.so") + libcudart.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + libcudart.cudaMemcpy.restype = ctypes.c_int + rc = libcudart.cudaMemcpy(ctypes.c_void_p(ptr), ctypes.byref(pattern), size, 1) # cudaMemcpyHostToDevice=1 + assert rc == 0, f"cudaMemcpy H2D failed: {rc}" + + # Export handle + handle = _create_vmm_handle(ptr) + assert len(handle) > 0, "handle is empty" + + # Import handle (same process, simulates remote open) + ptr2 = _open_vmm_handle(handle) + assert ptr2 != 0, "open_vmm_handle returned null" + + # Read back through the new mapping + readback = (ctypes.c_uint8 * size)() + rc = libcudart.cudaMemcpy(ctypes.byref(readback), ctypes.c_void_p(ptr2), size, 2) # cudaMemcpyDeviceToHost=2 + assert rc == 0, f"cudaMemcpy D2H failed: {rc}" + assert all(b == 0xAB for b in readback), "Data mismatch after handle export/import" + + _close_vmm_handle(ptr2) + _vmm_free(ptr) + print("[PASS] test_vmm_handle_export_import") + + +def test_distributed_vmm(rank, world_size): + """Multi-GPU integration test: VMM alloc + P2P read.""" + from tilelang.distributed.utils import create_dist_tensor, create_tensor + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + group = dist.new_group(list(range(world_size))) + + # Allocate a tensor with cudaMalloc (needed for VMM handle export) + data = create_tensor([1024], torch.float32) + data.fill_(float(rank + 1)) + + # Create dist tensor with VMM + buffer_ptrs = create_dist_tensor(local_rank, world_size, data, rank, group, use_vmm=True) + + assert buffer_ptrs.shape[0] == world_size + assert buffer_ptrs[local_rank].item() != 0 + + dist.barrier() + + if rank == 0: + print(f"[PASS] test_distributed_vmm (world_size={world_size})") + + +def test_distributed_ipc_fallback(rank, world_size): + """Verify IPC path still works with TILESCALE_USE_VMM=0.""" + from tilelang.distributed.utils import create_dist_tensor, create_tensor + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + group = dist.new_group(list(range(world_size))) + + data = create_tensor([1024], torch.float32) + data.fill_(float(rank + 1)) + + # Explicitly use IPC + buffer_ptrs = create_dist_tensor(local_rank, world_size, data, rank, group, use_vmm=False) + + assert buffer_ptrs.shape[0] == world_size + assert buffer_ptrs[local_rank].item() != 0 + + dist.barrier() + + if rank == 0: + print(f"[PASS] test_distributed_ipc_fallback (world_size={world_size})") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--distributed", action="store_true", help="Run multi-GPU tests (requires torchrun)") + args = parser.parse_args() + + if args.distributed: + # Multi-GPU path (launched via torchrun) + import datetime + + world_size = int(os.environ.get("WORLD_SIZE", 1)) + rank = int(os.environ.get("RANK", 0)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=rank, + timeout=datetime.timedelta(seconds=60), + ) + + has_fabric = test_supports_fabric() if rank == 0 else False + # Broadcast fabric support info + fabric_tensor = torch.tensor([int(has_fabric)], device="cuda") + dist.broadcast(fabric_tensor, src=0) + has_fabric = bool(fabric_tensor.item()) + + if has_fabric: + test_distributed_vmm(rank, world_size) + + test_distributed_ipc_fallback(rank, world_size) + + dist.destroy_process_group() + else: + # Single-GPU unit tests + has_fabric = test_supports_fabric() + if has_fabric: + test_vmm_malloc_free() + test_vmm_handle_export_import() + print("\n=== All single-GPU VMM tests passed ===") + else: + print("\n=== Fabric not supported on this hardware, skipping VMM tests ===") + + +if __name__ == "__main__": + main() diff --git a/tilelang/distributed/__init__.py b/tilelang/distributed/__init__.py index 0aa3803773..62bd72641e 100644 --- a/tilelang/distributed/__init__.py +++ b/tilelang/distributed/__init__.py @@ -2,3 +2,5 @@ from .utils import * # noqa: F401 from tilescale_ext import _create_tensor, _create_ipc_handle, _sync_ipc_handles # noqa: F401 +from tilescale_ext import _supports_vmm_fabric, _vmm_malloc, _vmm_free # noqa: F401 +from tilescale_ext import _create_vmm_handle, _open_vmm_handle, _close_vmm_handle, _sync_vmm_handles # noqa: F401 diff --git a/tilelang/distributed/utils.py b/tilelang/distributed/utils.py index 1e9ab6d647..70da8ae096 100644 --- a/tilelang/distributed/utils.py +++ b/tilelang/distributed/utils.py @@ -21,7 +21,17 @@ from cuda import cuda, cudart import ctypes -from tilescale_ext import _create_tensor, _create_ipc_handle, _sync_ipc_handles, create_host_device_tensor +from tilescale_ext import ( + _create_tensor, + _create_ipc_handle, + _sync_ipc_handles, + _supports_vmm_fabric, + _vmm_malloc, + _vmm_free, + _create_vmm_handle, + _sync_vmm_handles, + create_host_device_tensor, +) import functools from functools import lru_cache from threading import Lock @@ -109,8 +119,28 @@ def get_local_ipc_handle(data: torch.Tensor): return handle -def create_dist_tensor(local_rank: int, num_local_ranks: int, data: torch.Tensor, rank: int, group: dist.ProcessGroup): +def _resolve_use_vmm(use_vmm: bool | None) -> bool: + """Resolve whether to use VMM based on env var and hardware support.""" + import os + env_val = os.environ.get("TILESCALE_USE_VMM", None) + if env_val is not None: + return env_val == "1" + if use_vmm is not None: + return use_vmm + return False + + +def create_dist_tensor( + local_rank: int, + num_local_ranks: int, + data: torch.Tensor, + rank: int, + group: dist.ProcessGroup, + use_vmm: bool | None = None, +): assert num_local_ranks == group.size() + _use_vmm = _resolve_use_vmm(use_vmm) + # Synchronize device IDs device_ids = [ None, @@ -118,14 +148,21 @@ def create_dist_tensor(local_rank: int, num_local_ranks: int, data: torch.Tensor local_device_id = local_rank dist.all_gather_object(device_ids, local_device_id, group) - # Synchronize IPC handles - ipc_handles = [ + # Synchronize handles (VMM or IPC) + handles = [ None, ] * group.size() - local_ipc_handle = get_local_ipc_handle(data) - dist.all_gather_object(ipc_handles, local_ipc_handle, group) + if _use_vmm: + local_handle = _create_vmm_handle(ctypes.c_void_p(data.data_ptr()).value) + else: + local_handle = get_local_ipc_handle(data) + dist.all_gather_object(handles, local_handle, group) + buffer_ptrs_gpu = torch.empty(group.size(), dtype=torch.uint64, device="cuda") - _sync_ipc_handles(rank, device_ids, ctypes.c_void_p(buffer_ptrs_gpu.data_ptr()).value, ipc_handles, None) + if _use_vmm: + _sync_vmm_handles(rank, device_ids, ctypes.c_void_p(buffer_ptrs_gpu.data_ptr()).value, handles) + else: + _sync_ipc_handles(rank, device_ids, ctypes.c_void_p(buffer_ptrs_gpu.data_ptr()).value, handles, None) return buffer_ptrs_gpu diff --git a/tilelang/utils/allocator.py b/tilelang/utils/allocator.py index c363ffa453..35b05c6e68 100644 --- a/tilelang/utils/allocator.py +++ b/tilelang/utils/allocator.py @@ -2,9 +2,19 @@ import ctypes import ctypes.util +import os import torch import torch.distributed as dist -from tilescale_ext import tensor_from_ptr, _create_ipc_handle, _sync_ipc_handles +from tilescale_ext import ( + tensor_from_ptr, + _create_ipc_handle, + _sync_ipc_handles, + _supports_vmm_fabric, + _vmm_malloc, + _vmm_free, + _create_vmm_handle, + _sync_vmm_handles, +) from tilelang.utils.target import parse_device import contextlib @@ -67,6 +77,17 @@ def _load_cudart(): _libcudart.cudaSetDevice.restype = ctypes.c_int +def _resolve_use_vmm(use_vmm: bool | None) -> bool: + """Resolve whether to use VMM based on env var and hardware support.""" + env_val = os.environ.get("TILESCALE_USE_VMM", None) + if env_val is not None: + return env_val == "1" + if use_vmm is not None: + return use_vmm + # Default: opt-in (False) for Step 1 + return False + + class BaseAllocator: func: callable | None = None @@ -79,10 +100,12 @@ def __init__( num_local_ranks: int | None = None, group: dist.ProcessGroup | None = None, align: int = 256, + use_vmm: bool | None = None, ) -> None: if size <= 0: raise ValueError("size must be > 0") self.size = int(size) + self._use_vmm = _resolve_use_vmm(use_vmm) self._base_ptr = ctypes.c_void_p(0) self._ptr = ctypes.c_void_p(0) self._device = parse_device(device) @@ -121,20 +144,29 @@ def _alloc(self): rc = _libcudart.cudaSetDevice(int(self._device)) if rc != 0: raise RuntimeError(f"cudaSetDevice failed: {rc} {_libcudart.cudaGetErrorString(rc).decode()}") - rc = _libcudart.cudaMalloc(ctypes.byref(self._base_ptr), ctypes.c_size_t(self.size)) - if rc != 0: - msg = _libcudart.cudaGetErrorString(rc) - raise RuntimeError(f"cudaMalloc failed: {rc} {msg.decode() if msg else ''}") + + if self._use_vmm: + ptr_val = _vmm_malloc(self.size) + self._base_ptr.value = ptr_val + else: + rc = _libcudart.cudaMalloc(ctypes.byref(self._base_ptr), ctypes.c_size_t(self.size)) + if rc != 0: + msg = _libcudart.cudaGetErrorString(rc) + raise RuntimeError(f"cudaMalloc failed: {rc} {msg.decode() if msg else ''}") self._ptr.value = self._base_ptr.value def _free(self): if getattr(self, "_base_ptr", None) and self._base_ptr.value: - rc = _libcudart.cudaFree(self._base_ptr) - # mark freed even if error to avoid double free in destructor - self._base_ptr = ctypes.c_void_p(0) - if rc != 0: - msg = _libcudart.cudaGetErrorString(rc) - raise RuntimeError(f"cudaFree failed: {rc} {msg.decode() if msg else ''}") + if getattr(self, "_use_vmm", False): + _vmm_free(self._base_ptr.value) + self._base_ptr = ctypes.c_void_p(0) + else: + rc = _libcudart.cudaFree(self._base_ptr) + # mark freed even if error to avoid double free in destructor + self._base_ptr = ctypes.c_void_p(0) + if rc != 0: + msg = _libcudart.cudaGetErrorString(rc) + raise RuntimeError(f"cudaFree failed: {rc} {msg.decode() if msg else ''}") def _init_table(self): device_ids = [ @@ -144,14 +176,21 @@ def _init_table(self): dist.all_gather_object(device_ids, local_device_id, self._group) self._device_ids = device_ids - # Synchronize IPC handles - ipc_handles = [ + # Synchronize handles (VMM or IPC) + handles = [ None, ] * self._group.size() - local_ipc_handle = _create_ipc_handle(self._base_ptr.value) - dist.all_gather_object(ipc_handles, local_ipc_handle, self._group) + if self._use_vmm: + local_handle = _create_vmm_handle(self._base_ptr.value) + else: + local_handle = _create_ipc_handle(self._base_ptr.value) + dist.all_gather_object(handles, local_handle, self._group) + buffer_ptrs = torch.empty(self._group.size(), dtype=torch.uint64, device="cuda") - _sync_ipc_handles(self._local_rank, device_ids, ctypes.c_void_p(buffer_ptrs.data_ptr()).value, ipc_handles, None) + if self._use_vmm: + _sync_vmm_handles(self._local_rank, device_ids, ctypes.c_void_p(buffer_ptrs.data_ptr()).value, handles) + else: + _sync_ipc_handles(self._local_rank, device_ids, ctypes.c_void_p(buffer_ptrs.data_ptr()).value, handles, None) buffer_ptrs[self._local_rank] = self._base_ptr.value self._buffer_ptrs = buffer_ptrs self._table_size = 2 + self._group.size() @@ -249,7 +288,9 @@ def get_allocator( local_rank: int = 0, num_local_ranks: int = 1, group: dist.ProcessGroup | None = None, + use_vmm: bool | None = None, ) -> BaseAllocator: return BaseAllocator( - size, device=device, is_distributed=is_distributed, local_rank=local_rank, num_local_ranks=num_local_ranks, group=group + size, device=device, is_distributed=is_distributed, local_rank=local_rank, + num_local_ranks=num_local_ranks, group=group, use_vmm=use_vmm, ) diff --git a/tilelang/utils/ts_ext/__init__.py b/tilelang/utils/ts_ext/__init__.py index e8f1bb87c5..d1d891e012 100644 --- a/tilelang/utils/ts_ext/__init__.py +++ b/tilelang/utils/ts_ext/__init__.py @@ -8,11 +8,27 @@ _sync_ipc_handles = _C._sync_ipc_handles create_host_device_tensor = _C.create_host_device_tensor +# VMM operations +_supports_vmm_fabric = _C._supports_vmm_fabric +_vmm_malloc = _C._vmm_malloc +_vmm_free = _C._vmm_free +_create_vmm_handle = _C._create_vmm_handle +_open_vmm_handle = _C._open_vmm_handle +_close_vmm_handle = _C._close_vmm_handle +_sync_vmm_handles = _C._sync_vmm_handles + __all__ = [ "tensor_from_ptr", "_create_tensor", "_create_ipc_handle", "_sync_ipc_handles", "create_host_device_tensor", + "_supports_vmm_fabric", + "_vmm_malloc", + "_vmm_free", + "_create_vmm_handle", + "_open_vmm_handle", + "_close_vmm_handle", + "_sync_vmm_handles", "_C", ] diff --git a/tilelang/utils/ts_ext/exception.h b/tilelang/utils/ts_ext/exception.h index cd3b7f1961..6107d10e2e 100644 --- a/tilelang/utils/ts_ext/exception.h +++ b/tilelang/utils/ts_ext/exception.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -34,3 +35,16 @@ class TSException : public std::exception { } \ } while (0) #endif + +#ifndef CU_CHECK +#define CU_CHECK(cmd) \ + do { \ + CUresult e = (cmd); \ + if (e != CUDA_SUCCESS) { \ + const char *error_str = NULL; \ + cuGetErrorString(e, &error_str); \ + throw TSException("CU", __FILE__, __LINE__, \ + error_str ? std::string(error_str) : "unknown"); \ + } \ + } while (0) +#endif diff --git a/tilelang/utils/ts_ext/setup.py b/tilelang/utils/ts_ext/setup.py index 52d203f931..6bb43bbd65 100644 --- a/tilelang/utils/ts_ext/setup.py +++ b/tilelang/utils/ts_ext/setup.py @@ -31,9 +31,16 @@ if cuda_lib and os.path.isdir(cuda_lib): libraries.append("cudart") + libraries.append("cuda") library_dirs.append(cuda_lib) extra_link_args.append(f"-Wl,-rpath,{cuda_lib}") +# Also check lib64/stubs for libcuda.so (needed for driver API) +cuda_stubs = os.path.join(CUDA_HOME, "lib64", "stubs") if CUDA_HOME else None +if cuda_stubs and os.path.isdir(cuda_stubs): + library_dirs.append(cuda_stubs) + extra_link_args.append(f"-Wl,-rpath,{cuda_stubs}") + setup( name="tilescale_ext", packages=["tilescale_ext"], @@ -45,6 +52,7 @@ "ts_ext_bindings.cpp", "tensor.cpp", "ipc_ops.cpp", + "vmm_ops.cpp", ], include_dirs=include_dirs, extra_compile_args=extra_compile_args, diff --git a/tilelang/utils/ts_ext/ts_ext_bindings.cpp b/tilelang/utils/ts_ext/ts_ext_bindings.cpp index 685ba11098..8ade65ac9a 100644 --- a/tilelang/utils/ts_ext/ts_ext_bindings.cpp +++ b/tilelang/utils/ts_ext/ts_ext_bindings.cpp @@ -46,4 +46,57 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }, py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), py::arg("all_gathered_handles"), py::arg("root_unique_id_opt")); + + // VMM API + m.def("_supports_vmm_fabric", &supports_vmm_fabric, + "Check if all GPUs support CUDA VMM fabric handles"); + + m.def( + "_vmm_malloc", + [](size_t size) -> uintptr_t { + return reinterpret_cast(vmm_malloc(size)); + }, + py::arg("size")); + + m.def( + "_vmm_free", + [](uintptr_t ptr_value) { + vmm_free(reinterpret_cast(ptr_value)); + }, + py::arg("ptr_value")); + + m.def( + "_create_vmm_handle", + [](uintptr_t ptr_value) { + void *ptr = reinterpret_cast(ptr_value); + return create_vmm_handle(ptr); + }, + py::arg("ptr_value")); + + m.def( + "_open_vmm_handle", + [](const py::bytearray &handle_bytes) -> uintptr_t { + return reinterpret_cast(open_vmm_handle(handle_bytes)); + }, + py::arg("handle_bytes")); + + m.def( + "_close_vmm_handle", + [](uintptr_t ptr_value) { + close_vmm_handle(reinterpret_cast(ptr_value)); + }, + py::arg("ptr_value")); + + m.def( + "_sync_vmm_handles", + [](int rank, const std::vector &device_ids, + uintptr_t buffer_ptrs_gpu_addr, + const std::vector> &all_gathered_handles) { + void **buffer_ptrs_gpu = + reinterpret_cast(buffer_ptrs_gpu_addr); + sync_vmm_handles(rank, device_ids, buffer_ptrs_gpu, + all_gathered_handles); + }, + py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), + py::arg("all_gathered_handles")); } diff --git a/tilelang/utils/ts_ext/ts_ext_ops.h b/tilelang/utils/ts_ext/ts_ext_ops.h index 224ace2968..daf09d7de3 100644 --- a/tilelang/utils/ts_ext/ts_ext_ops.h +++ b/tilelang/utils/ts_ext/ts_ext_ops.h @@ -23,3 +23,14 @@ void sync_ipc_handles( int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, const std::vector> &all_gathered_handles, const std::optional &root_unique_id_opt); + +// VMM operations +void *vmm_malloc(size_t size_raw); +void vmm_free(void *ptr); +pybind11::bytearray create_vmm_handle(void *ptr); +void *open_vmm_handle(const pybind11::bytearray &handle_bytes); +void close_vmm_handle(void *ptr); +bool supports_vmm_fabric(); +void sync_vmm_handles( + int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, + const std::vector> &all_gathered_handles); diff --git a/tilelang/utils/ts_ext/vmm_ops.cpp b/tilelang/utils/ts_ext/vmm_ops.cpp new file mode 100644 index 0000000000..a7a4f82e3e --- /dev/null +++ b/tilelang/utils/ts_ext/vmm_ops.cpp @@ -0,0 +1,189 @@ +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "exception.h" +#include "ts_ext_ops.h" + +namespace py = pybind11; + +// ---------- helpers ---------- + +static void cu_mem_set_access_all(void *ptr, size_t size) { + int device_count; + CUDA_CHECK(cudaGetDeviceCount(&device_count)); + + std::vector access_desc(device_count); + for (int idx = 0; idx < device_count; ++idx) { + access_desc[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access_desc[idx].location.id = idx; + access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + + CU_CHECK(cuMemSetAccess((CUdeviceptr)ptr, size, access_desc.data(), + device_count)); +} + +static size_t align_to_granularity(size_t size_raw, size_t granularity) { + size_t size = (size_raw + granularity - 1) & ~(granularity - 1); + if (size == 0) + size = granularity; + return size; +} + +static void ensure_cuda_context() { + // Force CUDA runtime to create a context on the current device + cudaFree(0); +} + +// ---------- VMM malloc/free ---------- + +void *vmm_malloc(size_t size_raw) { + ensure_cuda_context(); + + CUdevice device; + CU_CHECK(cuCtxGetDevice(&device)); + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + prop.location.id = device; + + size_t granularity = 0; + CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, + CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + + size_t size = align_to_granularity(size_raw, granularity); + + CUmemGenericAllocationHandle handle; + CU_CHECK(cuMemCreate(&handle, size, &prop, 0)); + + void *ptr = nullptr; + CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, granularity, 0, 0)); + CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, handle, 0)); + cu_mem_set_access_all(ptr, size); + + return ptr; +} + +void vmm_free(void *ptr) { + CUmemGenericAllocationHandle handle; + CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); + + size_t size = 0; + CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); + + CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); + CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); + CU_CHECK(cuMemRelease(handle)); +} + +// ---------- handle export/import ---------- + +py::bytearray create_vmm_handle(void *ptr) { + CUmemGenericAllocationHandle handle; + CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); + + size_t size = 0; + CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); + + CUmemFabricHandle fabric_handle; + CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, + CU_MEM_HANDLE_TYPE_FABRIC, 0)); + + // Serialize: 8 bytes size + sizeof(CUmemFabricHandle) bytes handle + std::string buf(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); + std::memcpy(&buf[0], &size, sizeof(size_t)); + std::memcpy(&buf[sizeof(size_t)], &fabric_handle, sizeof(CUmemFabricHandle)); + return py::bytearray(buf.data(), buf.size()); +} + +void *open_vmm_handle(const py::bytearray &handle_bytes) { + ensure_cuda_context(); + + std::string s = (std::string)handle_bytes; + TS_HOST_ASSERT(s.size() == sizeof(size_t) + sizeof(CUmemFabricHandle)); + + size_t size = 0; + std::memcpy(&size, s.data(), sizeof(size_t)); + + CUmemFabricHandle fabric_handle; + std::memcpy(&fabric_handle, s.data() + sizeof(size_t), + sizeof(CUmemFabricHandle)); + + CUmemGenericAllocationHandle alloc_handle; + CU_CHECK(cuMemImportFromShareableHandle(&alloc_handle, &fabric_handle, + CU_MEM_HANDLE_TYPE_FABRIC)); + + void *ptr = nullptr; + CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, 0, 0, 0)); + CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, alloc_handle, 0)); + cu_mem_set_access_all(ptr, size); + + return ptr; +} + +void close_vmm_handle(void *ptr) { + vmm_free(ptr); +} + +// ---------- fabric support detection ---------- + +bool supports_vmm_fabric() { + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) + return false; + + // Check driver version >= 12.4 (CUDA 12040) + int driver_version = 0; + CUresult cu_err = cuDriverGetVersion(&driver_version); + if (cu_err != CUDA_SUCCESS || driver_version < 12040) + return false; + + // Check all devices support fabric handles + for (int i = 0; i < device_count; ++i) { + CUdevice dev; + CU_CHECK(cuDeviceGet(&dev, i)); + int supported = 0; + CU_CHECK(cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev)); + if (!supported) + return false; + } + return true; +} + +// ---------- unified sync entry ---------- + +void sync_vmm_handles( + int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, + const std::vector> &all_gathered_handles) { + + const int num = (int)device_ids.size(); + TS_HOST_ASSERT((size_t)num == all_gathered_handles.size()); + + std::vector buffer_ptrs(num, nullptr); + + for (int i = 0; i < num; ++i) { + TS_HOST_ASSERT(all_gathered_handles[i].has_value()); + if (i != rank) { + buffer_ptrs[i] = open_vmm_handle(all_gathered_handles[i].value()); + } + } + + CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs.data(), + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaDeviceSynchronize()); +} diff --git a/tilescale_ext/__init__.py b/tilescale_ext/__init__.py index c650ea949c..f72168606b 100644 --- a/tilescale_ext/__init__.py +++ b/tilescale_ext/__init__.py @@ -4,6 +4,13 @@ _create_ipc_handle, _sync_ipc_handles, create_host_device_tensor, + _supports_vmm_fabric, + _vmm_malloc, + _vmm_free, + _create_vmm_handle, + _open_vmm_handle, + _close_vmm_handle, + _sync_vmm_handles, ) __all__ = [ @@ -12,4 +19,11 @@ "_create_ipc_handle", "_sync_ipc_handles", "create_host_device_tensor", + "_supports_vmm_fabric", + "_vmm_malloc", + "_vmm_free", + "_create_vmm_handle", + "_open_vmm_handle", + "_close_vmm_handle", + "_sync_vmm_handles", ] From 10c419e82d9fae285b4635e59b2488027917de73 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 14:36:15 +0800 Subject: [PATCH 2/9] refactor, add configuration script, single node test passed --- testing/python/distributed/test_vmm_ops.py | 39 +++++++++++++++------- tilelang/distributed/conf_vmm.sh | 10 ++++++ tilelang/utils/ts_ext/vmm_ops.cpp | 15 ++------- 3 files changed, 40 insertions(+), 24 deletions(-) create mode 100644 tilelang/distributed/conf_vmm.sh diff --git a/testing/python/distributed/test_vmm_ops.py b/testing/python/distributed/test_vmm_ops.py index 149f124189..5ac691ee11 100644 --- a/testing/python/distributed/test_vmm_ops.py +++ b/testing/python/distributed/test_vmm_ops.py @@ -84,23 +84,35 @@ def test_vmm_handle_export_import(): def test_distributed_vmm(rank, world_size): - """Multi-GPU integration test: VMM alloc + P2P read.""" - from tilelang.distributed.utils import create_dist_tensor, create_tensor + """Multi-GPU integration test: VMM alloc + P2P read via BaseAllocator.""" + from tilelang.utils.allocator import BaseAllocator local_rank = int(os.environ.get("LOCAL_RANK", 0)) torch.cuda.set_device(local_rank) group = dist.new_group(list(range(world_size))) - # Allocate a tensor with cudaMalloc (needed for VMM handle export) - data = create_tensor([1024], torch.float32) - data.fill_(float(rank + 1)) - - # Create dist tensor with VMM - buffer_ptrs = create_dist_tensor(local_rank, world_size, data, rank, group, use_vmm=True) - - assert buffer_ptrs.shape[0] == world_size - assert buffer_ptrs[local_rank].item() != 0 + # Use BaseAllocator with use_vmm=True — it allocates via vmm_malloc + # and handles the VMM handle exchange internally + allocator = BaseAllocator( + size=1024 * 1024, # 1 MB + device=f"cuda", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=world_size, + group=group, + use_vmm=True, + ) + + assert allocator.initialized() + assert allocator._buffer_ptrs is not None + assert allocator._buffer_ptrs.shape[0] == world_size + assert allocator._buffer_ptrs[local_rank].item() != 0 + + # Allocate a tensor from the VMM buffer and verify P2P access + t = allocator._allocate_tensor((256,), torch.float32) + t.fill_(float(rank + 1)) + torch.cuda.synchronize() dist.barrier() @@ -124,7 +136,10 @@ def test_distributed_ipc_fallback(rank, world_size): buffer_ptrs = create_dist_tensor(local_rank, world_size, data, rank, group, use_vmm=False) assert buffer_ptrs.shape[0] == world_size - assert buffer_ptrs[local_rank].item() != 0 + # Note: IPC path doesn't set local rank's pointer (no self-open needed) + # Check that at least one remote rank's pointer is non-zero + remote_rank = (local_rank + 1) % world_size + assert buffer_ptrs[remote_rank].item() != 0, f"Remote rank {remote_rank} pointer is zero" dist.barrier() diff --git a/tilelang/distributed/conf_vmm.sh b/tilelang/distributed/conf_vmm.sh new file mode 100644 index 0000000000..9049347f2f --- /dev/null +++ b/tilelang/distributed/conf_vmm.sh @@ -0,0 +1,10 @@ +MAJOR=$(grep nvidia-caps-imex-channels /proc/devices | awk '{print $1}') + +# Use mknod to open fabric channel 0 +sudo mkdir -p /dev/nvidia-caps-imex-channels/ +sudo mknod /dev/nvidia-caps-imex-channels/channel0 c $MAJOR 0 +sudo chmod 666 /dev/nvidia-caps-imex-channels/channel0 + +# Optional: set environment variables to use VMM and distributed +export TILESCALE_USE_VMM=1 +export TILESCALE_USE_DISTRIBUTED=1 diff --git a/tilelang/utils/ts_ext/vmm_ops.cpp b/tilelang/utils/ts_ext/vmm_ops.cpp index a7a4f82e3e..f3bb9da0c9 100644 --- a/tilelang/utils/ts_ext/vmm_ops.cpp +++ b/tilelang/utils/ts_ext/vmm_ops.cpp @@ -40,16 +40,9 @@ static size_t align_to_granularity(size_t size_raw, size_t granularity) { return size; } -static void ensure_cuda_context() { - // Force CUDA runtime to create a context on the current device - cudaFree(0); -} - // ---------- VMM malloc/free ---------- void *vmm_malloc(size_t size_raw) { - ensure_cuda_context(); - CUdevice device; CU_CHECK(cuCtxGetDevice(&device)); @@ -81,7 +74,7 @@ void vmm_free(void *ptr) { CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); size_t size = 0; - CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); + CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); @@ -95,13 +88,13 @@ py::bytearray create_vmm_handle(void *ptr) { CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); size_t size = 0; - CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr)); + CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); CUmemFabricHandle fabric_handle; CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0)); - // Serialize: 8 bytes size + sizeof(CUmemFabricHandle) bytes handle + // Serialize: 8 bytes size + sizeof(CUmemFabricHandle) std::string buf(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); std::memcpy(&buf[0], &size, sizeof(size_t)); std::memcpy(&buf[sizeof(size_t)], &fabric_handle, sizeof(CUmemFabricHandle)); @@ -109,8 +102,6 @@ py::bytearray create_vmm_handle(void *ptr) { } void *open_vmm_handle(const py::bytearray &handle_bytes) { - ensure_cuda_context(); - std::string s = (std::string)handle_bytes; TS_HOST_ASSERT(s.size() == sizeof(size_t) + sizeof(CUmemFabricHandle)); From 26d632931038ad6d206b2ead334055fa6a12c727 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 15:12:41 +0800 Subject: [PATCH 3/9] refactor: remove standalone ts_ext, unify into tilelang/distributerd/shared_memory --- CMakeLists.txt | 17 ++--- pyproject.toml | 6 +- testing/python/distributed/test_vmm_ops.py | 16 ++--- tilelang/__init__.py | 4 +- tilelang/distributed/__init__.py | 15 ++++- .../distributed/shared_memory}/__init__.py | 4 +- .../shared_memory/csrc/bindings.cpp} | 24 +++---- .../shared_memory/csrc}/exception.h | 0 .../shared_memory/csrc}/ipc_ops.cpp | 2 +- .../shared_memory/csrc/ops.h} | 2 + .../shared_memory/csrc}/tensor.cpp | 2 +- .../shared_memory/csrc}/vmm_ops.cpp | 4 +- tilelang/distributed/utils.py | 2 +- tilelang/utils/allocator.py | 2 +- tilelang/utils/ts_ext/__init__.py | 34 ---------- tilelang/utils/ts_ext/pyproject.toml | 9 --- tilelang/utils/ts_ext/setup.py | 67 ------------------- 17 files changed, 53 insertions(+), 157 deletions(-) rename {tilescale_ext => tilelang/distributed/shared_memory}/__init__.py (78%) rename tilelang/{utils/ts_ext/ts_ext_bindings.cpp => distributed/shared_memory/csrc/bindings.cpp} (81%) rename tilelang/{utils/ts_ext => distributed/shared_memory/csrc}/exception.h (100%) rename tilelang/{utils/ts_ext => distributed/shared_memory/csrc}/ipc_ops.cpp (99%) rename tilelang/{utils/ts_ext/ts_ext_ops.h => distributed/shared_memory/csrc/ops.h} (97%) rename tilelang/{utils/ts_ext => distributed/shared_memory/csrc}/tensor.cpp (99%) rename tilelang/{utils/ts_ext => distributed/shared_memory/csrc}/vmm_ops.cpp (97%) delete mode 100644 tilelang/utils/ts_ext/__init__.py delete mode 100644 tilelang/utils/ts_ext/pyproject.toml delete mode 100644 tilelang/utils/ts_ext/setup.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e9bb4fe77..f2cc452b2f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -352,11 +352,12 @@ if(USE_CUDA) if(Torch_FOUND) message(STATUS "Building tilescale_ext with Torch ${Torch_VERSION}") + set(SHARED_MEMORY_CSRC ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/distributed/shared_memory/csrc) set(TILESCALE_EXT_SOURCES - ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/ts_ext_bindings.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/tensor.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/ipc_ops.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext/vmm_ops.cpp + ${SHARED_MEMORY_CSRC}/bindings.cpp + ${SHARED_MEMORY_CSRC}/tensor.cpp + ${SHARED_MEMORY_CSRC}/ipc_ops.cpp + ${SHARED_MEMORY_CSRC}/vmm_ops.cpp ) # Find libtorch_python.so @@ -372,7 +373,7 @@ if(USE_CUDA) target_include_directories(tilescale_ext_C PRIVATE ${TORCH_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS} - ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/utils/ts_ext + ${SHARED_MEMORY_CSRC} ) if(TORCH_PYTHON_RESULT EQUAL 0 AND EXISTS "${TORCH_PYTHON_LIBRARY}") @@ -390,10 +391,10 @@ if(USE_CUDA) LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" ) - # Install as tilescale_ext/_C.so so it can be imported as tilescale_ext._C + # Install as tilelang/distributed/shared_memory/_C.so install(TARGETS tilescale_ext_C - LIBRARY DESTINATION tilescale_ext - RUNTIME DESTINATION tilescale_ext) + LIBRARY DESTINATION tilelang/distributed/shared_memory + RUNTIME DESTINATION tilelang/distributed/shared_memory) else() message(WARNING "Torch not found, tilescale_ext will not be built") endif() diff --git a/pyproject.toml b/pyproject.toml index 4256809f42..ef8da81658 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,9 +100,8 @@ include = [ # Composable Kernel "3rdparty/composable_kernel/include", "3rdparty/composable_kernel/library", - # tilescale_ext C++ extension - "tilelang/utils/ts_ext/**", - "tilescale_ext/**", + # Shared memory C++ extension + "tilelang/distributed/shared_memory/**", "testing/**", "examples/**", ] @@ -117,7 +116,6 @@ exclude = [ [tool.scikit-build.wheel.packages] tilelang = "tilelang" -tilescale_ext = "tilescale_ext" "tilelang/src" = "src" # NOTE: The mapping below places the contents of '3rdparty' inside 'tilelang/3rdparty' in the wheel. # This is necessary to find TVM shared libraries at runtime. diff --git a/testing/python/distributed/test_vmm_ops.py b/testing/python/distributed/test_vmm_ops.py index 5ac691ee11..f59caf1548 100644 --- a/testing/python/distributed/test_vmm_ops.py +++ b/testing/python/distributed/test_vmm_ops.py @@ -19,16 +19,16 @@ def test_supports_fabric(): """Test fabric support detection.""" - from tilescale_ext import _supports_vmm_fabric + from tilelang.distributed.shared_memory import _supports_vmm_fabric result = _supports_vmm_fabric() - print(f"[PASS] _supports_vmm_fabric() = {result}") + print(f"\033[32m[PASS]\033[0m _supports_vmm_fabric() = {result}") return result def test_vmm_malloc_free(): """Test VMM malloc and free roundtrip.""" - from tilescale_ext import _vmm_malloc, _vmm_free + from tilelang.distributed.shared_memory import _vmm_malloc, _vmm_free size = 1024 * 1024 # 1 MB ptr = _vmm_malloc(size) @@ -43,12 +43,12 @@ def test_vmm_malloc_free(): assert rc == 0, f"cudaMemset on VMM pointer failed: {rc}" _vmm_free(ptr) - print("[PASS] test_vmm_malloc_free") + print("\033[32m[PASS]\033[0m test_vmm_malloc_free") def test_vmm_handle_export_import(): """Test handle export and import on a single GPU.""" - from tilescale_ext import _vmm_malloc, _vmm_free, _create_vmm_handle, _open_vmm_handle, _close_vmm_handle + from tilelang.distributed.shared_memory import _vmm_malloc, _vmm_free, _create_vmm_handle, _open_vmm_handle, _close_vmm_handle size = 4096 ptr = _vmm_malloc(size) @@ -80,7 +80,7 @@ def test_vmm_handle_export_import(): _close_vmm_handle(ptr2) _vmm_free(ptr) - print("[PASS] test_vmm_handle_export_import") + print("\033[32m[PASS]\033[0m test_vmm_handle_export_import") def test_distributed_vmm(rank, world_size): @@ -117,7 +117,7 @@ def test_distributed_vmm(rank, world_size): dist.barrier() if rank == 0: - print(f"[PASS] test_distributed_vmm (world_size={world_size})") + print(f"\033[32m[PASS]\033[0m test_distributed_vmm (world_size={world_size})") def test_distributed_ipc_fallback(rank, world_size): @@ -144,7 +144,7 @@ def test_distributed_ipc_fallback(rank, world_size): dist.barrier() if rank == 0: - print(f"[PASS] test_distributed_ipc_fallback (world_size={world_size})") + print(f"\033[32m[PASS]\033[0m test_distributed_ipc_fallback (world_size={world_size})") def main(): diff --git a/tilelang/__init__.py b/tilelang/__init__.py index a687f4d1b0..fdbc0d94f4 100644 --- a/tilelang/__init__.py +++ b/tilelang/__init__.py @@ -148,12 +148,12 @@ def _load_tile_lang_lib(): deprecated, # noqa: F401 ) -# TileScale distributed extensions (optional - only available when tilescale_ext is installed) +# TileScale distributed extensions (optional - requires shared_memory C extension) try: from .utils.tensor import tensor # noqa: F401 from .utils.allocator import get_allocator # noqa: F401 except ImportError: - # tilescale_ext not installed - distributed features unavailable + # shared_memory C extension not built - distributed features unavailable tensor = None get_allocator = None diff --git a/tilelang/distributed/__init__.py b/tilelang/distributed/__init__.py index 62bd72641e..7a8e361448 100644 --- a/tilelang/distributed/__init__.py +++ b/tilelang/distributed/__init__.py @@ -1,6 +1,15 @@ """The distributed modules""" from .utils import * # noqa: F401 -from tilescale_ext import _create_tensor, _create_ipc_handle, _sync_ipc_handles # noqa: F401 -from tilescale_ext import _supports_vmm_fabric, _vmm_malloc, _vmm_free # noqa: F401 -from tilescale_ext import _create_vmm_handle, _open_vmm_handle, _close_vmm_handle, _sync_vmm_handles # noqa: F401 +from .shared_memory import ( # noqa: F401 + _create_tensor, + _create_ipc_handle, + _sync_ipc_handles, + _supports_vmm_fabric, + _vmm_malloc, + _vmm_free, + _create_vmm_handle, + _open_vmm_handle, + _close_vmm_handle, + _sync_vmm_handles, +) diff --git a/tilescale_ext/__init__.py b/tilelang/distributed/shared_memory/__init__.py similarity index 78% rename from tilescale_ext/__init__.py rename to tilelang/distributed/shared_memory/__init__.py index f72168606b..fd35809c60 100644 --- a/tilescale_ext/__init__.py +++ b/tilelang/distributed/shared_memory/__init__.py @@ -1,4 +1,6 @@ -from tilescale_ext._C import ( +"""Shared memory allocator for distributed communication (IPC + VMM/fabric).""" + +from tilelang.distributed.shared_memory._C import ( # type: ignore[import] tensor_from_ptr, _create_tensor, _create_ipc_handle, diff --git a/tilelang/utils/ts_ext/ts_ext_bindings.cpp b/tilelang/distributed/shared_memory/csrc/bindings.cpp similarity index 81% rename from tilelang/utils/ts_ext/ts_ext_bindings.cpp rename to tilelang/distributed/shared_memory/csrc/bindings.cpp index 8ade65ac9a..e0956e456f 100644 --- a/tilelang/utils/ts_ext/ts_ext_bindings.cpp +++ b/tilelang/distributed/shared_memory/csrc/bindings.cpp @@ -1,4 +1,4 @@ -#include "ts_ext_ops.h" +#include "ops.h" #include #include #include @@ -6,14 +6,13 @@ namespace py = pybind11; PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "TileScale unified CUDA/IPCs extension"; + m.doc() = "TileScale shared memory allocator (IPC + VMM)"; - // alloc_cuda API + // Tensor utilities m.def("tensor_from_ptr", &tensor_from_ptr, py::arg("ptr"), py::arg("shape"), py::arg("dtype") = std::string("float32"), py::arg("device") = 0, py::arg("take_ownership") = false); - // ipc_ext API m.def( "_create_tensor", [](const std::vector &shape, const py::object &dtype) { @@ -25,11 +24,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("create_host_device_tensor", &create_host_device_tensor, "Create host/device shared pinned-mapped tensor (shape + dtype)"); + // IPC API m.def( "_create_ipc_handle", [](uintptr_t ptr_value) { - void *ptr = reinterpret_cast(ptr_value); - return create_ipc_handle(ptr); + return create_ipc_handle(reinterpret_cast(ptr_value)); }, py::arg("ptr_value")); @@ -39,9 +38,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { uintptr_t buffer_ptrs_gpu_addr, const std::vector> &all_gathered_handles, const std::optional &root_unique_id_opt) { - void **buffer_ptrs_gpu = - reinterpret_cast(buffer_ptrs_gpu_addr); - sync_ipc_handles(rank, device_ids, buffer_ptrs_gpu, + sync_ipc_handles(rank, device_ids, + reinterpret_cast(buffer_ptrs_gpu_addr), all_gathered_handles, root_unique_id_opt); }, py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), @@ -68,8 +66,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def( "_create_vmm_handle", [](uintptr_t ptr_value) { - void *ptr = reinterpret_cast(ptr_value); - return create_vmm_handle(ptr); + return create_vmm_handle(reinterpret_cast(ptr_value)); }, py::arg("ptr_value")); @@ -92,9 +89,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { [](int rank, const std::vector &device_ids, uintptr_t buffer_ptrs_gpu_addr, const std::vector> &all_gathered_handles) { - void **buffer_ptrs_gpu = - reinterpret_cast(buffer_ptrs_gpu_addr); - sync_vmm_handles(rank, device_ids, buffer_ptrs_gpu, + sync_vmm_handles(rank, device_ids, + reinterpret_cast(buffer_ptrs_gpu_addr), all_gathered_handles); }, py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), diff --git a/tilelang/utils/ts_ext/exception.h b/tilelang/distributed/shared_memory/csrc/exception.h similarity index 100% rename from tilelang/utils/ts_ext/exception.h rename to tilelang/distributed/shared_memory/csrc/exception.h diff --git a/tilelang/utils/ts_ext/ipc_ops.cpp b/tilelang/distributed/shared_memory/csrc/ipc_ops.cpp similarity index 99% rename from tilelang/utils/ts_ext/ipc_ops.cpp rename to tilelang/distributed/shared_memory/csrc/ipc_ops.cpp index eaa820a545..ba31a3fcbe 100644 --- a/tilelang/utils/ts_ext/ipc_ops.cpp +++ b/tilelang/distributed/shared_memory/csrc/ipc_ops.cpp @@ -19,7 +19,7 @@ #include #include "exception.h" -#include "ts_ext_ops.h" +#include "ops.h" namespace py = pybind11; diff --git a/tilelang/utils/ts_ext/ts_ext_ops.h b/tilelang/distributed/shared_memory/csrc/ops.h similarity index 97% rename from tilelang/utils/ts_ext/ts_ext_ops.h rename to tilelang/distributed/shared_memory/csrc/ops.h index daf09d7de3..691be5f4fb 100644 --- a/tilelang/utils/ts_ext/ts_ext_ops.h +++ b/tilelang/distributed/shared_memory/csrc/ops.h @@ -6,6 +6,7 @@ #include #include +// Tensor utilities torch::Tensor tensor_from_ptr(uint64_t ptr_val, std::vector shape, const std::string &dtype = "float32", int64_t device = 0, bool take_ownership = false); @@ -17,6 +18,7 @@ std::pair create_host_device_tensor(const std::vector &shape, c10::ScalarType dtype); +// IPC operations pybind11::bytearray create_ipc_handle(void *ptr); void sync_ipc_handles( diff --git a/tilelang/utils/ts_ext/tensor.cpp b/tilelang/distributed/shared_memory/csrc/tensor.cpp similarity index 99% rename from tilelang/utils/ts_ext/tensor.cpp rename to tilelang/distributed/shared_memory/csrc/tensor.cpp index d9a9753f93..22613aa573 100644 --- a/tilelang/utils/ts_ext/tensor.cpp +++ b/tilelang/distributed/shared_memory/csrc/tensor.cpp @@ -9,7 +9,7 @@ #include #include "exception.h" -#include "ts_ext_ops.h" +#include "ops.h" static int64_t safe_mul_int64(int64_t a, int64_t b) { if (a == 0 || b == 0) diff --git a/tilelang/utils/ts_ext/vmm_ops.cpp b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp similarity index 97% rename from tilelang/utils/ts_ext/vmm_ops.cpp rename to tilelang/distributed/shared_memory/csrc/vmm_ops.cpp index f3bb9da0c9..0c93deee67 100644 --- a/tilelang/utils/ts_ext/vmm_ops.cpp +++ b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp @@ -12,7 +12,7 @@ #include #include "exception.h" -#include "ts_ext_ops.h" +#include "ops.h" namespace py = pybind11; @@ -136,13 +136,11 @@ bool supports_vmm_fabric() { if (err != cudaSuccess || device_count == 0) return false; - // Check driver version >= 12.4 (CUDA 12040) int driver_version = 0; CUresult cu_err = cuDriverGetVersion(&driver_version); if (cu_err != CUDA_SUCCESS || driver_version < 12040) return false; - // Check all devices support fabric handles for (int i = 0; i < device_count; ++i) { CUdevice dev; CU_CHECK(cuDeviceGet(&dev, i)); diff --git a/tilelang/distributed/utils.py b/tilelang/distributed/utils.py index 70da8ae096..5eea2b3636 100644 --- a/tilelang/distributed/utils.py +++ b/tilelang/distributed/utils.py @@ -21,7 +21,7 @@ from cuda import cuda, cudart import ctypes -from tilescale_ext import ( +from tilelang.distributed.shared_memory import ( _create_tensor, _create_ipc_handle, _sync_ipc_handles, diff --git a/tilelang/utils/allocator.py b/tilelang/utils/allocator.py index 35b05c6e68..347db41074 100644 --- a/tilelang/utils/allocator.py +++ b/tilelang/utils/allocator.py @@ -5,7 +5,7 @@ import os import torch import torch.distributed as dist -from tilescale_ext import ( +from tilelang.distributed.shared_memory import ( tensor_from_ptr, _create_ipc_handle, _sync_ipc_handles, diff --git a/tilelang/utils/ts_ext/__init__.py b/tilelang/utils/ts_ext/__init__.py deleted file mode 100644 index d1d891e012..0000000000 --- a/tilelang/utils/ts_ext/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -from importlib import import_module as _imp - -_C = _imp("tilescale_ext._C") - -tensor_from_ptr = _C.tensor_from_ptr -_create_tensor = _C._create_tensor -_create_ipc_handle = _C._create_ipc_handle -_sync_ipc_handles = _C._sync_ipc_handles -create_host_device_tensor = _C.create_host_device_tensor - -# VMM operations -_supports_vmm_fabric = _C._supports_vmm_fabric -_vmm_malloc = _C._vmm_malloc -_vmm_free = _C._vmm_free -_create_vmm_handle = _C._create_vmm_handle -_open_vmm_handle = _C._open_vmm_handle -_close_vmm_handle = _C._close_vmm_handle -_sync_vmm_handles = _C._sync_vmm_handles - -__all__ = [ - "tensor_from_ptr", - "_create_tensor", - "_create_ipc_handle", - "_sync_ipc_handles", - "create_host_device_tensor", - "_supports_vmm_fabric", - "_vmm_malloc", - "_vmm_free", - "_create_vmm_handle", - "_open_vmm_handle", - "_close_vmm_handle", - "_sync_vmm_handles", - "_C", -] diff --git a/tilelang/utils/ts_ext/pyproject.toml b/tilelang/utils/ts_ext/pyproject.toml deleted file mode 100644 index ec90f70f15..0000000000 --- a/tilelang/utils/ts_ext/pyproject.toml +++ /dev/null @@ -1,9 +0,0 @@ -[build-system] -requires = ["setuptools>=65", "wheel", "torch", "packaging"] -build-backend = "setuptools.build_meta" - -[project] -name = "tilescale_ext" -version = "0.1.0" -requires-python = ">=3.8" -description = "TileScale unified CUDA/IPCs extension." diff --git a/tilelang/utils/ts_ext/setup.py b/tilelang/utils/ts_ext/setup.py deleted file mode 100644 index 6bb43bbd65..0000000000 --- a/tilelang/utils/ts_ext/setup.py +++ /dev/null @@ -1,67 +0,0 @@ -from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME -import os -import torch - -try: - from torch.utils.cpp_extension import _get_torch_lib_dir - - torch_lib_dir = _get_torch_lib_dir() -except Exception: - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - -include_dirs = [] -if CUDA_HOME is not None: - cuda_inc = os.path.join(CUDA_HOME, "include") - if os.path.isdir(cuda_inc): - include_dirs.append(cuda_inc) - cuda_lib = os.path.join(CUDA_HOME, "lib64") -else: - cuda_lib = None - -extra_compile_args = { - "cxx": ["-O3", "-std=c++17", "-fPIC"], - "nvcc": ["-O3", "-std=c++17", "-Xcompiler", "-fPIC"], -} - -extra_link_args = [f"-Wl,-rpath,{torch_lib_dir}"] -runtime_library_dirs = [torch_lib_dir] -libraries = [] -library_dirs = [torch_lib_dir] - -if cuda_lib and os.path.isdir(cuda_lib): - libraries.append("cudart") - libraries.append("cuda") - library_dirs.append(cuda_lib) - extra_link_args.append(f"-Wl,-rpath,{cuda_lib}") - -# Also check lib64/stubs for libcuda.so (needed for driver API) -cuda_stubs = os.path.join(CUDA_HOME, "lib64", "stubs") if CUDA_HOME else None -if cuda_stubs and os.path.isdir(cuda_stubs): - library_dirs.append(cuda_stubs) - extra_link_args.append(f"-Wl,-rpath,{cuda_stubs}") - -setup( - name="tilescale_ext", - packages=["tilescale_ext"], - package_dir={"tilescale_ext": "."}, - ext_modules=[ - CUDAExtension( - name="tilescale_ext._C", - sources=[ - "ts_ext_bindings.cpp", - "tensor.cpp", - "ipc_ops.cpp", - "vmm_ops.cpp", - ], - include_dirs=include_dirs, - extra_compile_args=extra_compile_args, - extra_link_args=extra_link_args, - runtime_library_dirs=runtime_library_dirs, - libraries=libraries, - library_dirs=library_dirs, - ) - ], - cmdclass={"build_ext": BuildExtension}, - zip_safe=False, -) From 03d35c0113fdb8c08b0149b3890dd06eaeaaae2b Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 16:01:49 +0800 Subject: [PATCH 4/9] minor fix --- testing/python/distributed/test_vmm_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/testing/python/distributed/test_vmm_ops.py b/testing/python/distributed/test_vmm_ops.py index f59caf1548..eef44463c2 100644 --- a/testing/python/distributed/test_vmm_ops.py +++ b/testing/python/distributed/test_vmm_ops.py @@ -182,6 +182,7 @@ def main(): dist.destroy_process_group() else: # Single-GPU unit tests + torch.cuda.set_device(0) has_fabric = test_supports_fabric() if has_fabric: test_vmm_malloc_free() From 107cade97aa47655a15c12b7be799b1045224b99 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 17:33:34 +0800 Subject: [PATCH 5/9] draft --- CMakeLists.txt | 1 + .../distributed/example_multimem_allreduce.py | 125 +++++ nvshmem_issue.md | 14 + src/op/multimem.cc | 272 ++++++++++ src/op/multimem.h | 88 ++++ src/op/multimem_rewriter.h | 279 +++++++++++ src/shared_memory/shared_memory.cc | 464 ++++++++++++++++++ src/target/codegen_cuda.cc | 13 + src/target/codegen_cuda.h | 2 + src/tl_templates/cuda/multimem.h | 154 ++++++ .../python/distributed/test_multicast_ops.py | 103 ++++ tilelang/distributed/__init__.py | 2 + .../distributed/shared_memory/__init__.py | 194 +++++++- .../shared_memory/csrc/bindings.cpp | 4 + tilelang/distributed/shared_memory/csrc/ops.h | 3 + .../shared_memory/csrc/vmm_ops.cpp | 25 + tilelang/language/__init__.py | 2 +- tilelang/language/distributed/__init__.py | 11 + tilelang/language/distributed/multimem.py | 79 +++ tilelang/utils/allocator.py | 157 +++++- 20 files changed, 1972 insertions(+), 20 deletions(-) create mode 100644 examples/distributed/example_multimem_allreduce.py create mode 100644 nvshmem_issue.md create mode 100644 src/op/multimem.cc create mode 100644 src/op/multimem.h create mode 100644 src/op/multimem_rewriter.h create mode 100644 src/shared_memory/shared_memory.cc create mode 100644 src/tl_templates/cuda/multimem.h create mode 100644 testing/python/distributed/test_multicast_ops.py create mode 100644 tilelang/language/distributed/multimem.py diff --git a/CMakeLists.txt b/CMakeLists.txt index f2cc452b2f..cc8a66b434 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -214,6 +214,7 @@ elseif(USE_CUDA) file(GLOB TILE_LANG_CUDA_SRCS src/runtime/runtime.cc src/runtime/tilescale_cuda_module.cc + src/shared_memory/shared_memory.cc src/target/ptx.cc src/target/codegen_cuda.cc src/target/codegen_py.cc diff --git a/examples/distributed/example_multimem_allreduce.py b/examples/distributed/example_multimem_allreduce.py new file mode 100644 index 0000000000..57e1f702f8 --- /dev/null +++ b/examples/distributed/example_multimem_allreduce.py @@ -0,0 +1,125 @@ +""" +Multimem allreduce example using NVSwitch multicast instructions. + +Multi-process multi-GPU: each process manages one GPU, multicast handle +shared via fabric handles through torch.distributed. + +Usage: + export TILESCALE_USE_VMM=1 + export NCCL_IB_DISABLE=1 + export TILELANG_USE_DISTRIBUTED=1 + python examples/distributed/example_multimem_allreduce.py [--num-processes 8] + +Requirements: + - NVSwitch with multicast support (H100/B200 DGX) +""" + +import os +import argparse + +import torch +import torch.distributed as dist +import torch.multiprocessing + +import tilelang +import tilelang.language as T +from tilelang.distributed import init_dist +from tilelang.utils.allocator import get_allocator + +tilelang.disable_cache() +os.environ["NCCL_DEBUG"] = "WARN" + + +def multimem_allreduce_kernel(N, block_N, threads): + @T.prim_func + def main( + mcast_buf: T.Tensor((N,), "float32"), + result: T.Tensor((N,), "float32"), + ): + with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx,): + result_local = T.alloc_fragment([block_N], "float32") + T.multimem_ld_reduce( + mcast_buf[bx * block_N:(bx + 1) * block_N], + result_local, + reduce_op=T.MultimemReduceOp.ADD, + ) + T.copy(result_local, result[bx * block_N:(bx + 1) * block_N]) + + return main + + +def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + N = args.N + BLOCK_N = args.block_n + threads = args.threads + + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + + # Create allocator with integrated multicast buffer + allocator = get_allocator( + size=N * 4, # float32 = 4 bytes + device=f"cuda:{local_rank}", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=num_local_ranks, + group=group, + mcast_size=N * 4, + ) + + # Compile kernel + kernel = tilelang.compile( + multimem_allreduce_kernel(N, BLOCK_N, threads), + pass_configs={"tl.disable_tma_lower": True}, + ) + if local_rank == 0 and args.print_source: + print(kernel.get_kernel_source()) + + # Random input per rank + torch.manual_seed(42 + local_rank) + local_data = torch.randn(N, dtype=torch.float32, device=f"cuda:{local_rank}") + + # Allocate from multicast buffer + # mcast_tensor: MC VA for multimem instructions (read) + # local_tensor: physical VA for writing data + mcast_tensor, local_tensor = allocator._allocate_mcast_tensor((N,), torch.float32) + + # Write to physical memory (NOT the MC VA) + local_tensor.copy_(local_data) + torch.cuda.synchronize() + dist.barrier(group) + result = torch.empty(N, dtype=torch.float32, device=f"cuda:{local_rank}") + kernel(mcast_tensor, result) + torch.cuda.synchronize() + + # torch.distributed reference + expected = local_data.clone() + dist.all_reduce(expected, op=dist.ReduceOp.SUM, group=group) + + # Compare (fp32 should be exact or near-exact) + atol = 1e-5 + max_diff = (result - expected).abs().max().item() + passed = max_diff < atol + + if local_rank == 0: + print(f"N={N}, num_ranks={num_ranks}, max_diff={max_diff:.4f}, atol={atol}") + if passed: + print(f"[rank {local_rank}] PASSED") + else: + print(f"[rank {local_rank}] FAILED (max_diff={max_diff:.4f})") + + allocator.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + parser.add_argument("--N", type=int, default=65536) + parser.add_argument("--block_n", type=int, default=4096) + parser.add_argument("--threads", type=int, default=128) + parser.add_argument("--print_source", action="store_true") + args = parser.parse_args() + + torch.multiprocessing.spawn( + main, args=(args.num_processes, args), nprocs=args.num_processes, join=True + ) diff --git a/nvshmem_issue.md b/nvshmem_issue.md new file mode 100644 index 0000000000..2e105295d3 --- /dev/null +++ b/nvshmem_issue.md @@ -0,0 +1,14 @@ +现在tilescale对nvshmem的支持有问题 +我希望你能帮我修复 +目标是跑通 example_allgather.py +运行方式： +```bash +NCCL_IB_DISABLE=1 bash tilelang/distributed/launch.sh examples/distributed/nvshmem/example_allgather.py +``` +我希望tvm_ffi和cython两个backend都能跑通 +跑cython backend，需要 +```bash +tilelang.compile(..., execution_backend="cython") +``` +你也可以通过运行MEMCHECK=1来使用compute sanitizer分析内存错误 +你可以任意多轮尝试更改，直至两个backend都能跑通，向我报告 diff --git a/src/op/multimem.cc b/src/op/multimem.cc new file mode 100644 index 0000000000..edc73a159d --- /dev/null +++ b/src/op/multimem.cc @@ -0,0 +1,272 @@ +/*! + * \file tl/op/multimem.cc + * \brief Unified multimem operator implementation. + * + * Reuses CopyNode's ParallelOp + InferLayout + VectorizeLoop pipeline, + * then post-processes to replace mcast buffer accesses with multimem instructions. + */ + +#include "multimem.h" + +#include +#include +#include + +#include + +#include "../transform/common/loop_fusion_utils.h" +#include "../transform/common/loop_parallel_transform_utils.h" +#include "../transform/loop_partition.h" +#include "../transform/loop_vectorize.h" +#include "multimem_rewriter.h" +#include "operator.h" +#include "utils.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +// === MultimemOp Constructor === +// args[0]: src region (tl.region call), args[1]: dst region, args[2]: mode, args[3]: reduce_op +MultimemOp::MultimemOp(Array args, + Map annotations) { + ObjectPtr node = + tvm::ffi::make_object(); + + // Parse buffer regions using same utility as CopyNode + Array rgs[2]; + Buffer bf[2]; + for (int i = 0; i < 2; i++) { + auto region = NormalizeToBufferRegion(args[i]); + rgs[i] = region->region; + bf[i] = region->buffer; + } + node->src = bf[0]; + node->dst = bf[1]; + node->src_range = rgs[0]; + node->dst_range = rgs[1]; + + node->mode = static_cast(args[2].as().value()->value); + node->reduce_op = args[3].as().value()->value; + + // Validate buffer scopes based on mode: + // ld_reduce: src=global(mcast), dst=local.fragment + // st: src=local.fragment, dst=global(mcast) + // red: src=local.fragment, dst=global(mcast) + String src_scope = node->src.scope(); + String dst_scope = node->dst.scope(); + switch (node->mode) { + case MultimemMode::kLdReduce: + ICHECK(src_scope == "global") + << "multimem_ld_reduce: src must be global (multicast) buffer, got '" + << src_scope << "' for buffer '" << node->src->name << "'"; + ICHECK(dst_scope == "local.fragment") + << "multimem_ld_reduce: dst must be local.fragment buffer, got '" + << dst_scope << "' for buffer '" << node->dst->name << "'"; + break; + case MultimemMode::kSt: + ICHECK(src_scope == "local.fragment") + << "multimem_st: src must be local.fragment buffer, got '" << src_scope + << "' for buffer '" << node->src->name << "'"; + ICHECK(dst_scope == "global") + << "multimem_st: dst must be global (multicast) buffer, got '" + << dst_scope << "' for buffer '" << node->dst->name << "'"; + break; + case MultimemMode::kRed: + ICHECK(src_scope == "local.fragment") + << "multimem_red: src must be local.fragment buffer, got '" << src_scope + << "' for buffer '" << node->src->name << "'"; + ICHECK(dst_scope == "global") + << "multimem_red: dst must be global (multicast) buffer, got '" + << dst_scope << "' for buffer '" << node->dst->name << "'"; + break; + } + + data_ = std::move(node); +} + +// === GetCoalescedWidth === +// 128-bit per multimem instruction => width = 128 / dtype_bits +int MultimemOpNode::GetCoalescedWidth() const { + int bits = src->dtype.bits(); + return 128 / bits; // f32->4, f16->8, bf16->8 +} + +// === MakeIterVars === +// Creates loop iteration variables from ranges (skipping dims with extent==1) +Array MultimemOpNode::MakeIterVars() const { + // Use the range with the higher scope level as basis (same logic as CopyNode) + auto scope_level = [](const Buffer &b) -> int { + String s = b.scope(); + if (s == "local.fragment" || s == "local") + return 2; + if (s == "shared" || s == "shared.dyn" || s == "shared.tmem") + return 1; + return 0; + }; + + int src_level = scope_level(src); + int dst_level = scope_level(dst); + bool base_is_src = (src_level >= dst_level); + const Array &base_ranges = base_is_src ? src_range : dst_range; + + Array loop_vars; + size_t idx = 0; + for (size_t i = 0; i < base_ranges.size(); i++) { + if (is_one(base_ranges[i]->extent)) + continue; + Var var = Var(std::string{char('i' + idx)}, base_ranges[i]->extent->dtype); + idx++; + loop_vars.push_back( + {Range(0, base_ranges[i]->extent), var, IterVarType::kDataPar}); + } + return loop_vars; +} + +// === MakeIndices === +Array MultimemOpNode::MakeIndices(const Array &ivs, + int src_dst) const { + Array indices; + const Array &ranges = src_dst == 0 ? src_range : dst_range; + size_t idx = 0; + for (size_t i = 0; i < ranges.size(); i++) { + if (is_one(ranges[i]->extent)) + indices.push_back(ranges[i]->min); + else { + indices.push_back(ranges[i]->min + ivs[idx]->var); + idx++; + } + } + return indices; +} + +// === MakePredicate === +PrimExpr MultimemOpNode::MakePredicate(arith::Analyzer *analyzer, + const Array &ivs, + Array extents, + int src_dst) const { + const Array &ranges = src_dst == 0 ? src_range : dst_range; + Array cond_list; + size_t idx = 0; + for (size_t i = 0; i < ranges.size(); i++) { + if (is_one(ranges[i]->extent)) + continue; + PrimExpr cond = ranges[i]->min + ivs[idx]->var < extents[i]; + if (!analyzer->CanProve(cond, arith::ProofStrength::kSymbolicBound)) { + cond_list.push_back(cond); + } + cond = ranges[i]->min + ivs[idx]->var >= 0; + if (!analyzer->CanProve(cond, arith::ProofStrength::kSymbolicBound)) { + cond_list.push_back(cond); + } + idx++; + } + if (cond_list.empty()) + return {}; + PrimExpr result = cond_list[0]; + for (size_t i = 1; i < cond_list.size(); i++) + result = And(result, cond_list[i]); + return result; +} + +// === MakeSIMTLoop === +// Creates the element-wise parallel loop: for (i,j): dst[i,j] = src[i,j] +// with coalesced_width annotation to limit vectorization to 128 bits. +For MultimemOpNode::MakeSIMTLoop(arith::Analyzer *analyzer) const { + Array loop_vars = MakeIterVars(); + bool is_scalar = loop_vars.empty(); + + for (const auto &iv : loop_vars) + analyzer->Bind(iv->var, iv->dom); + + Array src_indices = MakeIndices(loop_vars, 0); + Array dst_indices = MakeIndices(loop_vars, 1); + + PrimExpr src_predicate = MakePredicate(analyzer, loop_vars, src->shape, 0); + PrimExpr dst_predicate = MakePredicate(analyzer, loop_vars, dst->shape, 1); + + PrimExpr value = BufferLoad(src, src_indices); + if (src->dtype != dst->dtype) + value = Cast(dst->dtype, value); + if (src_predicate.defined()) + value = if_then_else(src_predicate, value, make_zero(dst->dtype)); + + Stmt body = BufferStore(dst, value, dst_indices); + if (dst_predicate.defined()) + body = IfThenElse(dst_predicate, body); + + if (is_scalar) { + return For(Var("i"), 0, 1, ForKind::kSerial, body); + } + + int coalesced_width = GetCoalescedWidth(); + for (int i = loop_vars.size() - 1; i >= 0; i--) { + Map loop_annotations; + loop_annotations.Set(attr::kCoalescedWidth, + IntImm(DataType::Int(32), coalesced_width)); + body = For(loop_vars[i]->var, 0, loop_vars[i]->dom->extent, + ForKind::kParallel, body, std::nullopt, loop_annotations); + } + return Downcast(body); +} + +// === InferLayout === +// Delegates to ParallelOp for layout inference (same as CopyNode::LowerNormalCopy) +LayoutMap MultimemOpNode::InferLayout(const LayoutInferArgs &T, + InferLevel level) const { + // Multimem ops always go through the ParallelOp path during Lower. + // No standalone layout inference needed. + return {}; +} + +// === Lower === +// The main lowering path: MakeSIMTLoop -> ParallelOp pipeline -> MultimemRewriter +Stmt MultimemOpNode::Lower(const LowerArgs &T, + arith::Analyzer *analyzer) const { + // Step 1-2: Create SIMT loop and fuse/transform + auto simt_loop = MakeSIMTLoop(analyzer); + auto fused_loop = Downcast(ParallelLoopFuser::Fuse(simt_loop)); + auto transformed_loop = + Downcast(ParallelLoopTransformer::Substitute(fused_loop)); + + // Step 3: Create ParallelOp and run InferLayout at multiple levels + auto par_op = ParallelOp(transformed_loop); + + std::vector levels = {InferLevel::kCommon, InferLevel::kStrict, + InferLevel::kFree}; + for (auto level : levels) { + par_op->InferLayout({T.target, T.thread_bounds, T.layout_map, analyzer, + false, T.buffer_remap, {}}, + level); + } + + // Step 4: Lower the parallel loop (PartitionLoop + VectorizeLoop) + auto loop_layout = par_op->GetLoopLayout(); + Stmt result = LowerParallelLoop(par_op->GetRoot(), loop_layout, T.thread_var, + analyzer, par_op->GetPredicate(T.thread_var)); + + // Step 5: Post-process — replace mcast buffer accesses with multimem call_extern + Buffer mcast_buf = (mode == MultimemMode::kLdReduce) ? src : dst; + // Remap the mcast buffer if needed + if (T.buffer_remap.count(mcast_buf)) { + mcast_buf = T.buffer_remap[mcast_buf]; + } + result = MultimemRewriter(mcast_buf, mode, reduce_op).Rewrite(result); + return result; +} + +// === Clone === +TileOperator MultimemOpNode::Clone() const { + auto node = tvm::ffi::make_object(*this); + return MultimemOp(node); +} + +// === Registration === +TIR_REGISTER_TL_TILE_OP(MultimemOp, multimem) + .set_num_inputs(4) + .set_attr("TCallEffectKind", + Integer(CallEffectKind::kOpaque)); + +} // namespace tl +} // namespace tvm diff --git a/src/op/multimem.h b/src/op/multimem.h new file mode 100644 index 0000000000..b9456fa9cd --- /dev/null +++ b/src/op/multimem.h @@ -0,0 +1,88 @@ +/*! + * \file tl/op/multimem.h + * \brief Unified multimem operator that reuses T.copy's layout inference. + * + * Design: MakeSIMTLoop creates element-wise BufferLoad/BufferStore loop, + * then ParallelOp + InferLayout + VectorizeLoop runs the standard pipeline. + * Post-process replaces vectorized loads/stores on mcast buffers with + * multimem call_extern instructions. + */ + +#ifndef TVM_TL_OP_MULTIMEM_H_ +#define TVM_TL_OP_MULTIMEM_H_ + +#include +#include + +#include "operator.h" +#include "parallel.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +enum class MultimemMode : int { kLdReduce = 0, kSt = 1, kRed = 2 }; + +/*! + * \brief Unified multimem operator for NVSwitch SHARP multicast operations. + * + * Supports three modes: + * - kLdReduce: load-reduce from multicast address into local buffer + * - kSt: store to multicast address (broadcast) + * - kRed: reduce into multicast address (no read-back) + * + * Lower flow: + * 1. MakeSIMTLoop: creates element-wise parallel loop (BufferLoad -> BufferStore) + * 2. ParallelLoopFuser::Fuse + ParallelLoopTransformer::Substitute + * 3. ParallelOp -> InferLayout at multiple levels + * 4. LowerParallelLoop (PartitionLoop + VectorizeLoop) + * 5. MultimemRewriter: post-process to replace mcast buffer accesses with call_extern + */ +class MultimemOpNode : public TileOperatorNode { +public: + Buffer src, dst; + Array src_range, dst_range; + MultimemMode mode; + int reduce_op; // 0=ADD, 1=MIN, 2=MAX + + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.MultimemOp", MultimemOpNode, + TileOperatorNode); + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("src", &MultimemOpNode::src) + .def_ro("dst", &MultimemOpNode::dst) + .def_ro("src_range", &MultimemOpNode::src_range) + .def_ro("dst_range", &MultimemOpNode::dst_range); + } + + Stmt Lower(const LowerArgs &T, arith::Analyzer *analyzer) const override; + LayoutMap InferLayout(const LayoutInferArgs &T, + InferLevel level) const override; + TileOperator Clone() const override; + +private: + For MakeSIMTLoop(arith::Analyzer *analyzer) const; + Array MakeIterVars() const; + Array MakeIndices(const Array &ivs, int src_dst) const; + PrimExpr MakePredicate(arith::Analyzer *analyzer, const Array &ivs, + Array extents, int src_dst) const; + int GetCoalescedWidth() const; +}; + +class MultimemOp : public TileOperator { +public: + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MultimemOp, TileOperator, + MultimemOpNode); + TVM_DLL MultimemOp( + Array args, + Map annotations = Map()); + static const Op &Get(); +}; + +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_OP_MULTIMEM_H_ diff --git a/src/op/multimem_rewriter.h b/src/op/multimem_rewriter.h new file mode 100644 index 0000000000..b63b4ad672 --- /dev/null +++ b/src/op/multimem_rewriter.h @@ -0,0 +1,279 @@ +/*! + * \file tl/op/multimem_rewriter.h + * \brief Post-process IR to replace vectorized BufferLoad/Store on mcast buffers + * with multimem call_extern instructions. + */ + +#ifndef TVM_TL_OP_MULTIMEM_REWRITER_H_ +#define TVM_TL_OP_MULTIMEM_REWRITER_H_ + +#include +#include +#include + +#include + +#include "multimem.h" + +namespace tvm { +namespace tl { + +using namespace tir; + +static inline std::string GetReduceOpStr(int reduce_op) { + switch (reduce_op) { + case 0: + return "tl::multimem::ReduceOp::ADD"; + case 1: + return "tl::multimem::ReduceOp::MIN"; + case 2: + return "tl::multimem::ReduceOp::MAX"; + default: + LOG(FATAL) << "Invalid reduce_op: " << reduce_op; + return ""; + } +} + +/*! + * \brief Rewrites BufferLoad/BufferStore involving a multicast buffer + * into multimem call_extern instructions. + * + * After ParallelOp + VectorizeLoop, the IR contains ForKind::kVectorized loops + * with scalar loop variables (Ramp is not materialized until codegen). + * This rewriter detects two patterns: + * + * 1. ForKind::kVectorized loop containing mcast buffer access: + * for (vec: kVectorized, extent=N) { dst[base+vec] = mcast[base+vec] } + * → call_extern("LdReduceVN<...>::run", &dst[base], &mcast[base]) + * + * 2. Scalar BufferStore with Ramp indices (if vectorization produced Ramp): + * dst[Ramp(base,1,N)] = mcast[Ramp(base,1,N)] + * → call_extern("LdReduceVN<...>::run", &dst[base], &mcast[base]) + */ +class MultimemRewriter : public StmtExprMutator { +public: + MultimemRewriter(Buffer mcast_buf, MultimemMode mode, int reduce_op) + : mcast_buf_(std::move(mcast_buf)), mode_(mode), reduce_op_(reduce_op) {} + + Stmt Rewrite(Stmt stmt) { return VisitStmt(std::move(stmt)); } + +protected: + /*! + * \brief Handle ForKind::kVectorized loops. + * If the loop body is a single BufferStore involving the mcast buffer, + * replace the entire loop with a single vectorized multimem call. + */ + Stmt VisitStmt_(const ForNode *op) override { + if (op->kind == ForKind::kVectorized) { + auto extent_ptr = op->extent.as(); + if (extent_ptr) { + int lanes = static_cast(extent_ptr->value); + // Try to match the loop body as a single BufferStore from mcast buffer + auto result = TryRewriteVectorizedLoop(op, lanes); + if (result.defined()) { + return result; + } + } + } + return StmtExprMutator::VisitStmt_(op); + } + + /*! + * \brief Handle scalar BufferStore with Ramp indices (fallback). + */ + Stmt VisitStmt_(const BufferStoreNode *op) override { + if (mode_ == MultimemMode::kLdReduce) { + if (auto *load = op->value.as()) { + if (load->buffer.same_as(mcast_buf_)) { + int lanes = GetLanes(load->indices); + if (lanes > 1) { + return MakeMultimemCall(op->buffer, op->indices, load->buffer, + load->indices, lanes); + } + } + } + } else { + if (op->buffer.same_as(mcast_buf_)) { + if (auto *load = op->value.as()) { + int lanes = GetLanes(op->indices); + if (lanes > 1) { + return MakeMultimemCall(load->buffer, load->indices, op->buffer, + op->indices, lanes); + } + } + } + } + return StmtExprMutator::VisitStmt_(op); + } + +private: + Buffer mcast_buf_; + MultimemMode mode_; + int reduce_op_; + + /*! + * \brief Try to rewrite a kVectorized for-loop containing a mcast BufferStore. + * Returns the replacement Stmt, or undefined if the pattern doesn't match. + */ + Stmt TryRewriteVectorizedLoop(const ForNode *op, int lanes) { + // The body should be a single BufferStore (possibly wrapped in IfThenElse) + const BufferStoreNode *store = nullptr; + Stmt body = op->body; + + // Unwrap IfThenElse if present + if (auto *ite = body.as()) { + store = ite->then_case.as(); + } else { + store = body.as(); + } + + if (!store) + return Stmt(); + + const BufferLoadNode *load = store->value.as(); + if (!load) + return Stmt(); + + // Check if this involves the mcast buffer + bool matches = false; + const Buffer *local_buf_ptr = nullptr; + const Array *local_indices_ptr = nullptr; + const Buffer *mc_buf_ptr = nullptr; + const Array *mc_indices_ptr = nullptr; + + if (mode_ == MultimemMode::kLdReduce) { + if (load->buffer.same_as(mcast_buf_)) { + matches = true; + local_buf_ptr = &store->buffer; + local_indices_ptr = &store->indices; + mc_buf_ptr = &load->buffer; + mc_indices_ptr = &load->indices; + } + } else { + if (store->buffer.same_as(mcast_buf_)) { + matches = true; + local_buf_ptr = &load->buffer; + local_indices_ptr = &load->indices; + mc_buf_ptr = &store->buffer; + mc_indices_ptr = &store->indices; + } + } + + if (!matches) + return Stmt(); + + // Substitute vec_var = 0 to get the base address + Var vec_var = op->loop_var; + Map vmap; + vmap.Set(vec_var, IntImm(vec_var.dtype(), 0)); + + Array local_base_indices; + for (const auto &idx : *local_indices_ptr) { + local_base_indices.push_back(Substitute(idx, vmap)); + } + Array mc_base_indices; + for (const auto &idx : *mc_indices_ptr) { + mc_base_indices.push_back(Substitute(idx, vmap)); + } + + return MakeMultimemCall(*local_buf_ptr, local_base_indices, *mc_buf_ptr, + mc_base_indices, lanes); + } + + /*! + * \brief Get the vector lanes from Ramp indices. + */ + int GetLanes(const Array &indices) const { + if (indices.size() == 1) { + if (auto *ramp = indices[0].as()) { + return static_cast(ramp->lanes.as()->value); + } + } + return 1; + } + + /*! + * \brief Create the call_extern for a multimem instruction. + */ + Stmt MakeMultimemCall(const Buffer &local_buf, + const Array &local_indices, + const Buffer &mc_buf, + const Array &mc_indices, int lanes) const { + std::string func_name = MakeFuncName(lanes, local_buf->dtype); + + Array args; + args.push_back(StringImm(func_name)); + + if (mode_ == MultimemMode::kLdReduce) { + args.push_back(MakeAddressOf(local_buf, local_indices)); + args.push_back(MakeAddressOf(mc_buf, mc_indices)); + } else { + args.push_back(MakeAddressOf(mc_buf, mc_indices)); + args.push_back(MakeAddressOf(local_buf, local_indices)); + } + + auto call = Call(DataType::Handle(), builtin::call_extern(), args); + return Evaluate(call); + } + + /*! + * \brief Construct the template function name. + */ + std::string MakeFuncName(int lanes, DataType dtype) const { + std::string dtype_tag = DTypeToTag(dtype); + std::string reduce_op_str = GetReduceOpStr(reduce_op_); + + std::stringstream ss; + switch (mode_) { + case MultimemMode::kLdReduce: + ss << "tl::multimem::LdReduceV" << lanes; + break; + case MultimemMode::kSt: + ss << "tl::multimem::StV" << lanes; + break; + case MultimemMode::kRed: + ss << "tl::multimem::RedV" << lanes; + break; + } + ss << "<"; + if (mode_ != MultimemMode::kSt) { + ss << reduce_op_str << ", "; + } + ss << dtype_tag; + ss << ">::run"; + return ss.str(); + } + + std::string DTypeToTag(DataType dtype) const { + if (dtype.is_float() && dtype.bits() == 32) + return "float"; + if (dtype.is_float16()) + return "half_t"; + if (dtype.is_bfloat16()) + return "bfloat16_t"; + LOG(FATAL) << "Unsupported dtype for multimem: " << dtype; + return ""; + } + + /*! + * \brief Create address_of expression. Handles Ramp by extracting base. + */ + PrimExpr MakeAddressOf(const Buffer &buffer, + const Array &indices) const { + Array scalar_indices; + for (const auto &idx : indices) { + if (auto *ramp = idx.as()) { + scalar_indices.push_back(ramp->base); + } else { + scalar_indices.push_back(idx); + } + } + return Call(DataType::Handle(), builtin::address_of(), + {BufferLoad(buffer, scalar_indices)}); + } +}; + +} // namespace tl +} // namespace tvm + +#endif // TVM_TL_OP_MULTIMEM_REWRITER_H_ diff --git a/src/shared_memory/shared_memory.cc b/src/shared_memory/shared_memory.cc new file mode 100644 index 0000000000..41a25c6985 --- /dev/null +++ b/src/shared_memory/shared_memory.cc @@ -0,0 +1,464 @@ +/*! + * \file shared_memory/shared_memory.cc + * \brief VMM/IPC/multicast shared memory ops registered via TVM FFI. + * + * Replaces the old pybind11 tilescale_ext_C module. All functions are + * registered under the "tl.shared_memory.*" namespace and accessed from + * Python via tvm_ffi.get_global_func(). + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +using namespace tvm; +using namespace tvm::ffi; + +// ---------- helpers ---------- + +#define SM_CUDA_CHECK(cmd) \ + do { \ + cudaError_t e = (cmd); \ + if (e != cudaSuccess) { \ + LOG_FATAL << "CUDA error " << __FILE__ << ":" << __LINE__ << " '" \ + << cudaGetErrorString(e) << "'"; \ + } \ + } while (0) + +#define SM_CU_CHECK(cmd) \ + do { \ + CUresult e = (cmd); \ + if (e != CUDA_SUCCESS) { \ + const char *err_str = nullptr; \ + cuGetErrorString(e, &err_str); \ + LOG_FATAL << "CU error " << __FILE__ << ":" << __LINE__ << " '" \ + << (err_str ? err_str : "unknown") << "'"; \ + } \ + } while (0) + +static void cu_mem_set_access_all(void *ptr, size_t size) { + int device_count; + SM_CUDA_CHECK(cudaGetDeviceCount(&device_count)); + + std::vector access_desc(device_count); + for (int idx = 0; idx < device_count; ++idx) { + access_desc[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access_desc[idx].location.id = idx; + access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + + SM_CU_CHECK( + cuMemSetAccess((CUdeviceptr)ptr, size, access_desc.data(), device_count)); +} + +static size_t align_to_granularity(size_t size_raw, size_t granularity) { + size_t size = (size_raw + granularity - 1) & ~(granularity - 1); + if (size == 0) + size = granularity; + return size; +} + +// ---------- VMM malloc/free ---------- + +static int64_t vmm_malloc_impl(int64_t size_raw) { + CUdevice device; + SM_CU_CHECK(cuCtxGetDevice(&device)); + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + prop.location.id = device; + + size_t granularity = 0; + SM_CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, + CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + + size_t size = align_to_granularity((size_t)size_raw, granularity); + + CUmemGenericAllocationHandle handle; + SM_CU_CHECK(cuMemCreate(&handle, size, &prop, 0)); + + void *ptr = nullptr; + SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, granularity, 0, 0)); + SM_CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, handle, 0)); + cu_mem_set_access_all(ptr, size); + + return (int64_t)(uintptr_t)ptr; +} + +static void vmm_free_impl(int64_t ptr_val) { + void *ptr = reinterpret_cast((uintptr_t)ptr_val); + CUmemGenericAllocationHandle handle; + SM_CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); + + size_t size = 0; + SM_CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); + + SM_CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); + SM_CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); + SM_CU_CHECK(cuMemRelease(handle)); +} + +// ---------- handle export/import ---------- + +// Returns serialized handle as Bytes. +// Format: 8 bytes size + sizeof(CUmemFabricHandle) bytes fabric handle. +static ffi::Bytes create_vmm_handle_impl(int64_t ptr_val) { + void *ptr = reinterpret_cast((uintptr_t)ptr_val); + CUmemGenericAllocationHandle handle; + SM_CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); + + size_t size = 0; + SM_CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); + + CUmemFabricHandle fabric_handle; + SM_CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, + CU_MEM_HANDLE_TYPE_FABRIC, 0)); + + std::string raw(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); + std::memcpy(&raw[0], &size, sizeof(size_t)); + std::memcpy(&raw[sizeof(size_t)], &fabric_handle, sizeof(CUmemFabricHandle)); + return ffi::Bytes(raw); +} + +static int64_t open_vmm_handle_impl(ffi::Bytes handle_bytes) { + ICHECK(handle_bytes.size() == sizeof(size_t) + sizeof(CUmemFabricHandle)); + const char *data = handle_bytes.data(); + + size_t size = 0; + std::memcpy(&size, data, sizeof(size_t)); + + CUmemFabricHandle fabric_handle; + std::memcpy(&fabric_handle, data + sizeof(size_t), + sizeof(CUmemFabricHandle)); + + CUmemGenericAllocationHandle alloc_handle; + SM_CU_CHECK(cuMemImportFromShareableHandle(&alloc_handle, &fabric_handle, + CU_MEM_HANDLE_TYPE_FABRIC)); + + void *ptr = nullptr; + SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, 0, 0, 0)); + SM_CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, alloc_handle, 0)); + cu_mem_set_access_all(ptr, size); + + return (int64_t)(uintptr_t)ptr; +} + +static void close_vmm_handle_impl(int64_t ptr_val) { vmm_free_impl(ptr_val); } + +// ---------- IPC handle ---------- + +static ffi::Bytes create_ipc_handle_impl(int64_t ptr_val) { + void *ptr = reinterpret_cast((uintptr_t)ptr_val); + cudaIpcMemHandle_t handle{}; + SM_CUDA_CHECK(cudaIpcGetMemHandle(&handle, ptr)); + return ffi::Bytes(reinterpret_cast(handle.reserved), + CUDA_IPC_HANDLE_SIZE); +} + +static int64_t open_ipc_handle_impl(ffi::Bytes handle_bytes) { + ICHECK(handle_bytes.size() == CUDA_IPC_HANDLE_SIZE); + cudaIpcMemHandle_t handle{}; + std::memcpy(handle.reserved, handle_bytes.data(), CUDA_IPC_HANDLE_SIZE); + + void *ptr = nullptr; + SM_CUDA_CHECK( + cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess)); + return (int64_t)(uintptr_t)ptr; +} + +static void close_ipc_handle_impl(int64_t ptr_val) { + void *ptr = reinterpret_cast((uintptr_t)ptr_val); + SM_CUDA_CHECK(cudaIpcCloseMemHandle(ptr)); +} + +// ---------- support detection ---------- + +static bool supports_vmm_fabric_impl() { + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) + return false; + + int driver_version = 0; + CUresult cu_err = cuDriverGetVersion(&driver_version); + if (cu_err != CUDA_SUCCESS || driver_version < 12040) + return false; + + for (int i = 0; i < device_count; ++i) { + CUdevice dev; + SM_CU_CHECK(cuDeviceGet(&dev, i)); + int supported = 0; + SM_CU_CHECK(cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev)); + if (!supported) + return false; + } + return true; +} + +static bool supports_multicast_impl() { + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) + return false; + + int driver_version = 0; + CUresult cu_err = cuDriverGetVersion(&driver_version); + if (cu_err != CUDA_SUCCESS || driver_version < 12040) + return false; + + for (int i = 0; i < device_count; ++i) { + CUdevice dev; + SM_CU_CHECK(cuDeviceGet(&dev, i)); + int supported = 0; + SM_CU_CHECK(cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, dev)); + if (!supported) + return false; + } + return true; +} + +// ---------- Multicast (NVSwitch) ---------- +// Multi-process multi-GPU with fabric handles (same as vmm_malloc). +// Each process manages one GPU. MC handle shared via fabric export/import. + +// Create multicast object with FABRIC handle type, returns handle as int64. +static int64_t mc_create_impl(int64_t size_raw, int64_t num_devices) { + CUmulticastObjectProp prop = {}; + prop.numDevices = (unsigned int)num_devices; + prop.handleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + size_t granularity = 0; + SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, + CU_MULTICAST_GRANULARITY_RECOMMENDED)); + + size_t size = align_to_granularity((size_t)size_raw, granularity); + prop.size = size; + + CUmemGenericAllocationHandle mc_handle; + SM_CU_CHECK(cuMulticastCreate(&mc_handle, &prop)); + + return (int64_t)mc_handle; +} + +// Export multicast handle as fabric handle bytes (for sharing across processes) +static ffi::Bytes mc_export_handle_impl(int64_t mc_handle_val) { + CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + + CUmemFabricHandle fabric_handle; + SM_CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, mc_handle, + CU_MEM_HANDLE_TYPE_FABRIC, 0)); + + return ffi::Bytes(reinterpret_cast(&fabric_handle), + sizeof(CUmemFabricHandle)); +} + +// Import multicast handle from fabric handle bytes, returns handle as int64. +static int64_t mc_import_handle_impl(ffi::Bytes handle_bytes) { + ICHECK(handle_bytes.size() == sizeof(CUmemFabricHandle)); + + CUmemFabricHandle fabric_handle; + std::memcpy(&fabric_handle, handle_bytes.data(), sizeof(CUmemFabricHandle)); + + CUmemGenericAllocationHandle mc_handle; + SM_CU_CHECK(cuMemImportFromShareableHandle(&mc_handle, &fabric_handle, + CU_MEM_HANDLE_TYPE_FABRIC)); + + return (int64_t)mc_handle; +} + +// Add a device to the multicast object +static void mc_add_device_impl(int64_t mc_handle_val, int64_t device_id) { + CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + CUdevice device; + SM_CU_CHECK(cuDeviceGet(&device, (int)device_id)); + SM_CU_CHECK(cuMulticastAddDevice(mc_handle, device)); +} + +// Bind a physical memory VA (from vmm_malloc) to the multicast object +static void mc_bind_mem_impl(int64_t mc_handle_val, int64_t ptr_val, + int64_t size) { + CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + void *ptr = reinterpret_cast((uintptr_t)ptr_val); + + // Retrieve the physical allocation handle from the mapped pointer + CUmemGenericAllocationHandle phys_handle; + SM_CU_CHECK(cuMemRetainAllocationHandle(&phys_handle, ptr)); + + // Bind to multicast + SM_CU_CHECK(cuMulticastBindMem(mc_handle, 0, phys_handle, 0, (size_t)size, 0)); + + // Release the temporary handle reference + SM_CU_CHECK(cuMemRelease(phys_handle)); +} + +// Map multicast object to a VA, returns mc_ptr. Does NOT release handle. +static int64_t mc_map_impl(int64_t mc_handle_val, int64_t size_raw, + int64_t num_devices) { + CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + + CUmulticastObjectProp prop = {}; + prop.numDevices = (unsigned int)num_devices; + prop.handleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + size_t granularity = 0; + SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, + CU_MULTICAST_GRANULARITY_RECOMMENDED)); + size_t size = align_to_granularity((size_t)size_raw, granularity); + + void *mc_ptr = nullptr; + SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&mc_ptr, size, granularity, 0, 0)); + SM_CU_CHECK(cuMemMap((CUdeviceptr)mc_ptr, size, 0, mc_handle, 0)); + cu_mem_set_access_all(mc_ptr, size); + + return (int64_t)(uintptr_t)mc_ptr; +} + +// Release a multicast handle (call after map) +static void mc_release_handle_impl(int64_t mc_handle_val) { + CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + SM_CU_CHECK(cuMemRelease(mc_handle)); +} + +// Free multicast VA mapping +static void mc_unmap_impl(int64_t mc_ptr_val, int64_t size_raw, + int64_t num_devices) { + void *ptr = reinterpret_cast((uintptr_t)mc_ptr_val); + + CUmulticastObjectProp prop = {}; + prop.numDevices = (unsigned int)num_devices; + prop.handleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + size_t granularity = 0; + SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, + CU_MULTICAST_GRANULARITY_RECOMMENDED)); + size_t size = align_to_granularity((size_t)size_raw, granularity); + + SM_CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); + SM_CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); +} + +// Get the aligned size for multicast +static int64_t mc_get_aligned_size_impl(int64_t size_raw, int64_t num_devices) { + CUmulticastObjectProp prop = {}; + prop.numDevices = (unsigned int)num_devices; + prop.handleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + size_t granularity = 0; + SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, + CU_MULTICAST_GRANULARITY_RECOMMENDED)); + return (int64_t)align_to_granularity((size_t)size_raw, granularity); +} + +// ---------- sync helpers ---------- + +// Synchronize VMM handles: open all peer handles and write pointers to GPU. +// peer_handles is a comma-separated list of hex-encoded handle bytes (or "SELF" +// for local rank). We pass individual handle open results back via buffer_ptrs. +// packed_handles: num_ranks concatenated raw handle bytes +static void sync_vmm_handles_impl(int64_t rank, int64_t num_ranks, + int64_t buffer_ptrs_gpu_addr, + ffi::Bytes packed_handles) { + const size_t handle_size = sizeof(size_t) + sizeof(CUmemFabricHandle); + ICHECK(packed_handles.size() == handle_size * (size_t)num_ranks); + + std::vector buffer_ptrs(num_ranks, nullptr); + + for (int64_t i = 0; i < num_ranks; ++i) { + if (i != rank) { + ffi::Bytes h(packed_handles.data() + i * handle_size, handle_size); + buffer_ptrs[i] = + reinterpret_cast((uintptr_t)open_vmm_handle_impl(h)); + } + } + + void **gpu_ptr = reinterpret_cast((uintptr_t)buffer_ptrs_gpu_addr); + SM_CUDA_CHECK(cudaMemcpy(gpu_ptr, buffer_ptrs.data(), + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); + SM_CUDA_CHECK(cudaDeviceSynchronize()); +} + +static void sync_ipc_handles_impl(int64_t rank, int64_t num_ranks, + int64_t buffer_ptrs_gpu_addr, + ffi::Bytes packed_handles) { + ICHECK(packed_handles.size() == CUDA_IPC_HANDLE_SIZE * (size_t)num_ranks); + + std::vector buffer_ptrs(num_ranks, nullptr); + + for (int64_t i = 0; i < num_ranks; ++i) { + if (i != rank) { + ffi::Bytes h(packed_handles.data() + i * CUDA_IPC_HANDLE_SIZE, + CUDA_IPC_HANDLE_SIZE); + buffer_ptrs[i] = + reinterpret_cast((uintptr_t)open_ipc_handle_impl(h)); + } + } + + void **gpu_ptr = reinterpret_cast((uintptr_t)buffer_ptrs_gpu_addr); + SM_CUDA_CHECK(cudaMemcpy(gpu_ptr, buffer_ptrs.data(), + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); + SM_CUDA_CHECK(cudaDeviceSynchronize()); +} + +// ---------- Registration ---------- + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + + // VMM + refl::GlobalDef().def("tl.shared_memory.vmm_malloc", vmm_malloc_impl); + refl::GlobalDef().def("tl.shared_memory.vmm_free", vmm_free_impl); + refl::GlobalDef().def("tl.shared_memory.create_vmm_handle", + create_vmm_handle_impl); + refl::GlobalDef().def("tl.shared_memory.open_vmm_handle", + open_vmm_handle_impl); + refl::GlobalDef().def("tl.shared_memory.close_vmm_handle", + close_vmm_handle_impl); + refl::GlobalDef().def("tl.shared_memory.sync_vmm_handles", + sync_vmm_handles_impl); + + // IPC + refl::GlobalDef().def("tl.shared_memory.create_ipc_handle", + create_ipc_handle_impl); + refl::GlobalDef().def("tl.shared_memory.open_ipc_handle", + open_ipc_handle_impl); + refl::GlobalDef().def("tl.shared_memory.close_ipc_handle", + close_ipc_handle_impl); + refl::GlobalDef().def("tl.shared_memory.sync_ipc_handles", + sync_ipc_handles_impl); + + // Support detection + refl::GlobalDef().def("tl.shared_memory.supports_vmm_fabric", + supports_vmm_fabric_impl); + refl::GlobalDef().def("tl.shared_memory.supports_multicast", + supports_multicast_impl); + + // Multicast (NVSwitch) + refl::GlobalDef().def("tl.shared_memory.mc_create", mc_create_impl); + refl::GlobalDef().def("tl.shared_memory.mc_export_handle", + mc_export_handle_impl); + refl::GlobalDef().def("tl.shared_memory.mc_import_handle", + mc_import_handle_impl); + refl::GlobalDef().def("tl.shared_memory.mc_add_device", mc_add_device_impl); + refl::GlobalDef().def("tl.shared_memory.mc_bind_mem", mc_bind_mem_impl); + refl::GlobalDef().def("tl.shared_memory.mc_map", mc_map_impl); + refl::GlobalDef().def("tl.shared_memory.mc_release_handle", + mc_release_handle_impl); + refl::GlobalDef().def("tl.shared_memory.mc_unmap", mc_unmap_impl); + refl::GlobalDef().def("tl.shared_memory.mc_get_aligned_size", + mc_get_aligned_size_impl); +} diff --git a/src/target/codegen_cuda.cc b/src/target/codegen_cuda.cc index ce605ac7da..21f2ded286 100644 --- a/src/target/codegen_cuda.cc +++ b/src/target/codegen_cuda.cc @@ -329,6 +329,10 @@ std::string CodeGenTileLangCUDA::Finish() { decl_stream << "#include \n"; decl_stream << "extern \"C\" __constant__ uint64_t meta_data[1024];\n"; } + + if (need_multimem_h_) { + decl_stream << "#include \n"; + } decl_stream << "#ifdef ENABLE_BF16\n"; decl_stream << "#include \n"; decl_stream << "#endif\n"; @@ -2894,6 +2898,15 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) { // Note: tl.put, tl.get, tl.wait are TileOperators handled through // remote_copy.cc They are lowered to call_extern with // tl::cp_warp/tl::cp_block templates + // Detect multimem ops to include the multimem.h header + if (op->op.same_as(builtin::call_extern()) && op->args.size() >= 1) { + if (auto *str_imm = op->args[0].as()) { + if (std::string(str_imm->value).find("tl::multimem::") == 0) { + this->need_multimem_h_ = true; + this->use_distributed_ = true; + } + } + } CodeGenC::VisitExpr_(op, os); } } diff --git a/src/target/codegen_cuda.h b/src/target/codegen_cuda.h index 6c5f89e076..8679ca09c6 100644 --- a/src/target/codegen_cuda.h +++ b/src/target/codegen_cuda.h @@ -141,6 +141,8 @@ class CodeGenTileLangCUDA final : public CodeGenC { bool need_cooperative_groups_{false}; // whether need curand_kernel.h bool need_curand_kernel_h_{false}; + // whether need multimem.h (NVSwitch multicast instructions) + bool need_multimem_h_{false}; // whether need distributed.h bool use_distributed_{use_distributed()}; // whether need nvshmem.h diff --git a/src/tl_templates/cuda/multimem.h b/src/tl_templates/cuda/multimem.h new file mode 100644 index 0000000000..7b449e8b63 --- /dev/null +++ b/src/tl_templates/cuda/multimem.h @@ -0,0 +1,154 @@ +#pragma once + +#include "common.h" + +// multimem instructions require SM 90+ (Hopper) and CUDA 12.0+ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && __CUDACC_VER_MAJOR__ >= 12 + +#ifndef TL_ALWAYS_FALSE_V_DEFINED +#define TL_ALWAYS_FALSE_V_DEFINED +template inline constexpr bool always_false_v = false; +#endif + +namespace tl { +namespace multimem { + +enum class ReduceOp { ADD = 0, MIN = 1, MAX = 2 }; + +// === Per-instruction primitives (used by MultimemRewriter post-process) === + +// --- LdReduceV4: 128-bit load-reduce from multicast address --- + +template +struct LdReduceV4 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, + "tl::multimem::LdReduceV4: unsupported dtype/op combination"); + } +}; + +template <> +struct LdReduceV4 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int4 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 {%0, %1, %2, %3}, " + "[%4];" + : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +template <> +struct LdReduceV4 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int4 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.min.v4.f32 {%0, %1, %2, %3}, " + "[%4];" + : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +template <> +struct LdReduceV4 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int4 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.max.v4.f32 {%0, %1, %2, %3}, " + "[%4];" + : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +// --- StV4: 128-bit store to multicast address --- + +template +struct StV4 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, + "tl::multimem::StV4: unsupported dtype"); + } +}; + +template <> +struct StV4 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + int4 val = *reinterpret_cast(src); + asm volatile( + "multimem.st.relaxed.sys.global.v4.b32 [%0], {%1, %2, %3, %4};" + : + : "l"(mcast_ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) + : "memory"); + } +}; + +// --- RedV4: 128-bit reduce into multicast address --- + +template +struct RedV4 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, + "tl::multimem::RedV4: unsupported dtype/op combination"); + } +}; + +template <> +struct RedV4 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + int4 val = *reinterpret_cast(src); + asm volatile( + "multimem.red.relaxed.sys.global.add.v4.f32 [%0], {%1, %2, %3, %4};" + : + : "l"(mcast_ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) + : "memory"); + } +}; + +template <> +struct RedV4 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + // multimem.red min not directly available as v4; use scalar fallback + const float *src_f = reinterpret_cast(src); + const char *mc_bytes = reinterpret_cast(mcast_ptr); + #pragma unroll + for (int i = 0; i < 4; i++) { + unsigned val = __float_as_uint(src_f[i]); + asm volatile( + "multimem.red.relaxed.sys.global.min.f32 [%0], %1;" + : + : "l"(mc_bytes + i * 4), "r"(val) + : "memory"); + } + } +}; + +template <> +struct RedV4 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + const float *src_f = reinterpret_cast(src); + const char *mc_bytes = reinterpret_cast(mcast_ptr); + #pragma unroll + for (int i = 0; i < 4; i++) { + unsigned val = __float_as_uint(src_f[i]); + asm volatile( + "multimem.red.relaxed.sys.global.max.f32 [%0], %1;" + : + : "l"(mc_bytes + i * 4), "r"(val) + : "memory"); + } + } +}; + +} // namespace multimem +} // namespace tl + +#endif // __CUDA_ARCH__ >= 900 && __CUDACC_VER_MAJOR__ >= 12 diff --git a/testing/python/distributed/test_multicast_ops.py b/testing/python/distributed/test_multicast_ops.py new file mode 100644 index 0000000000..57f650b135 --- /dev/null +++ b/testing/python/distributed/test_multicast_ops.py @@ -0,0 +1,103 @@ +""" +Multicast (NVSwitch) operations test. + +Usage: + # Single-GPU unit tests (no torchrun needed): + python testing/python/distributed/test_multicast_ops.py + + # Multi-GPU integration test: + torchrun --nproc_per_node=8 testing/python/distributed/test_multicast_ops.py --distributed +""" + +import argparse +import os + +import torch +import torch.distributed as dist + + +def test_supports_multicast(): + """Test multicast support detection.""" + from tilelang.distributed.shared_memory import _supports_multicast + + result = _supports_multicast() + print(f"\033[32m[PASS]\033[0m _supports_multicast() = {result}") + return result + + +def test_distributed_multicast_allocator(rank, world_size): + """Multi-GPU integration test: MulticastAllocator with fabric VMM.""" + from tilelang.utils.allocator import MulticastAllocator + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + group = dist.new_group(list(range(world_size))) + + allocator = MulticastAllocator( + size=1024 * 1024, # 1 MB + device="cuda", + local_rank=local_rank, + num_local_ranks=world_size, + group=group, + ) + + assert allocator._initialized + assert allocator.ptr != 0 + assert len(allocator.peer_ptrs) == world_size + + # Get local tensor and write data + t = allocator.get_local_tensor((256,), torch.bfloat16) + t.fill_(float(rank + 1)) + torch.cuda.synchronize() + + dist.barrier() + + # Verify we can read peer data + peer_rank = (local_rank + 1) % world_size + peer_t = allocator.get_peer_tensor(peer_rank, (256,), torch.bfloat16) + peer_val = peer_t[0].item() + expected_val = float(peer_rank + 1) + assert abs(peer_val - expected_val) < 1e-2, \ + f"rank {local_rank}: peer[{peer_rank}] = {peer_val}, expected {expected_val}" + + dist.barrier() + + if rank == 0: + print(f"\033[32m[PASS]\033[0m test_distributed_multicast_allocator (world_size={world_size})") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--distributed", action="store_true", help="Run multi-GPU tests (requires torchrun)" + ) + args = parser.parse_args() + + if args.distributed: + import datetime + + world_size = int(os.environ.get("WORLD_SIZE", 1)) + rank = int(os.environ.get("RANK", 0)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", + world_size=world_size, + rank=rank, + timeout=datetime.timedelta(seconds=60), + ) + + test_distributed_multicast_allocator(rank, world_size) + + dist.destroy_process_group() + else: + # Single-GPU unit tests + torch.cuda.set_device(0) + test_supports_multicast() + print("\n=== Single-GPU tests done ===") + + +if __name__ == "__main__": + main() diff --git a/tilelang/distributed/__init__.py b/tilelang/distributed/__init__.py index 7a8e361448..276cc91632 100644 --- a/tilelang/distributed/__init__.py +++ b/tilelang/distributed/__init__.py @@ -2,6 +2,7 @@ from .utils import * # noqa: F401 from .shared_memory import ( # noqa: F401 + tensor_from_ptr, _create_tensor, _create_ipc_handle, _sync_ipc_handles, @@ -12,4 +13,5 @@ _open_vmm_handle, _close_vmm_handle, _sync_vmm_handles, + _supports_multicast, ) diff --git a/tilelang/distributed/shared_memory/__init__.py b/tilelang/distributed/shared_memory/__init__.py index fd35809c60..39ef60f847 100644 --- a/tilelang/distributed/shared_memory/__init__.py +++ b/tilelang/distributed/shared_memory/__init__.py @@ -1,24 +1,185 @@ -"""Shared memory allocator for distributed communication (IPC + VMM/fabric).""" - -from tilelang.distributed.shared_memory._C import ( # type: ignore[import] - tensor_from_ptr, - _create_tensor, - _create_ipc_handle, - _sync_ipc_handles, - create_host_device_tensor, - _supports_vmm_fabric, - _vmm_malloc, - _vmm_free, - _create_vmm_handle, - _open_vmm_handle, - _close_vmm_handle, - _sync_vmm_handles, -) +"""Shared memory allocator for distributed communication (IPC + VMM/fabric). + +All ops registered via TVM FFI under tl.shared_memory.* namespace. +""" + +import torch +import tvm_ffi + +# ---------- TVM FFI function handles ---------- + +_vmm_malloc = tvm_ffi.get_global_func("tl.shared_memory.vmm_malloc") +_vmm_free = tvm_ffi.get_global_func("tl.shared_memory.vmm_free") +_create_vmm_handle = tvm_ffi.get_global_func("tl.shared_memory.create_vmm_handle") +_open_vmm_handle = tvm_ffi.get_global_func("tl.shared_memory.open_vmm_handle") +_close_vmm_handle = tvm_ffi.get_global_func("tl.shared_memory.close_vmm_handle") +_sync_vmm_handles_raw = tvm_ffi.get_global_func("tl.shared_memory.sync_vmm_handles") + +_create_ipc_handle = tvm_ffi.get_global_func("tl.shared_memory.create_ipc_handle") +_open_ipc_handle = tvm_ffi.get_global_func("tl.shared_memory.open_ipc_handle") +_close_ipc_handle = tvm_ffi.get_global_func("tl.shared_memory.close_ipc_handle") +_sync_ipc_handles_raw = tvm_ffi.get_global_func("tl.shared_memory.sync_ipc_handles") + +_supports_vmm_fabric = tvm_ffi.get_global_func("tl.shared_memory.supports_vmm_fabric") +_supports_multicast = tvm_ffi.get_global_func("tl.shared_memory.supports_multicast") + +# Multicast (NVSwitch) ops +_mc_create = tvm_ffi.get_global_func("tl.shared_memory.mc_create") +_mc_export_handle = tvm_ffi.get_global_func("tl.shared_memory.mc_export_handle") +_mc_import_handle = tvm_ffi.get_global_func("tl.shared_memory.mc_import_handle") +_mc_add_device = tvm_ffi.get_global_func("tl.shared_memory.mc_add_device") +_mc_bind_mem = tvm_ffi.get_global_func("tl.shared_memory.mc_bind_mem") +_mc_map = tvm_ffi.get_global_func("tl.shared_memory.mc_map") +_mc_release_handle = tvm_ffi.get_global_func("tl.shared_memory.mc_release_handle") +_mc_unmap = tvm_ffi.get_global_func("tl.shared_memory.mc_unmap") +_mc_get_aligned_size = tvm_ffi.get_global_func("tl.shared_memory.mc_get_aligned_size") + + +# ---------- tensor_from_ptr (pure Python, no C++ torch dependency) ---------- + +_dtype_str_to_torch = { + "float32": torch.float32, + "float": torch.float32, + "float16": torch.float16, + "half": torch.float16, + "bfloat16": torch.bfloat16, + "float64": torch.float64, + "double": torch.float64, + "int32": torch.int32, + "int": torch.int32, + "int64": torch.int64, + "long": torch.int64, + "uint8": torch.uint8, + "byte": torch.uint8, + "int8": torch.int8, + "bool": torch.bool, + "uint32": torch.uint32, + "uint64": torch.uint64, +} + +# __cuda_array_interface__ typestr mapping +_torch_dtype_to_typestr = { + torch.float32: " torch.Tensor: + """Create a CUDA tensor viewing external device memory (zero-copy).""" + if ptr_val == 0: + raise RuntimeError("Received null pointer (0).") + + dtype = _dtype_str_to_torch.get(dtype_str) + if dtype is None: + raise ValueError(f"Unsupported dtype string: '{dtype_str}'") + + if isinstance(shape, (list, tuple)): + shape = tuple(shape) + else: + shape = (shape,) + + numel = 1 + for s in shape: + numel *= s + if numel == 0: + return torch.empty(shape, dtype=dtype, device=f"cuda:{device}") + + typestr = _torch_dtype_to_typestr.get(dtype) + if typestr is not None: + # Standard path via __cuda_array_interface__ + arr = _ExternalCUDAArray(ptr_val, shape, typestr) + return torch.as_tensor(arr, device=f"cuda:{device}") + else: + # bfloat16 / uint32 / uint64: create as matching-size int type, then view + element_size = torch.empty((), dtype=dtype).element_size() + if element_size == 2: + proxy_dtype = torch.int16 + proxy_typestr = " &device_ids, void **buffer_ptrs_gpu, const std::vector> &all_gathered_handles); + +// Multicast support detection +bool supports_multicast(); diff --git a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp index 0c93deee67..e8059c81b1 100644 --- a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp +++ b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp @@ -176,3 +176,28 @@ void sync_vmm_handles( cudaMemcpyHostToDevice)); CUDA_CHECK(cudaDeviceSynchronize()); } + +// ---------- Multicast support detection ---------- + +bool supports_multicast() { + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) + return false; + + int driver_version = 0; + CUresult cu_err = cuDriverGetVersion(&driver_version); + if (cu_err != CUDA_SUCCESS || driver_version < 12040) + return false; + + for (int i = 0; i < device_count; ++i) { + CUdevice dev; + CU_CHECK(cuDeviceGet(&dev, i)); + int supported = 0; + CU_CHECK(cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, dev)); + if (!supported) + return false; + } + return true; +} diff --git a/tilelang/language/__init__.py b/tilelang/language/__init__.py index 88f164b308..52f9b5a47d 100644 --- a/tilelang/language/__init__.py +++ b/tilelang/language/__init__.py @@ -110,7 +110,7 @@ # Distributed multi-device primitives (NVSHMEM) from .distributed.multi_device.nvshmem import * # noqa: F401 from .distributed.multi_device.cpengine import * # noqa: F401 -from .distributed.common import * # noqa: F401 +from .distributed import * # noqa: F401 def import_source(source: str | None = None): diff --git a/tilelang/language/distributed/__init__.py b/tilelang/language/distributed/__init__.py index 976130a40b..c11f74a9eb 100644 --- a/tilelang/language/distributed/__init__.py +++ b/tilelang/language/distributed/__init__.py @@ -18,6 +18,13 @@ wait_lt, ) +from .multimem import ( + MultimemReduceOp, + multimem_ld_reduce, + multimem_st, + multimem_red, +) + __all__ = [ "get_rank", "get_num_ranks", @@ -32,4 +39,8 @@ "wait_le", "wait_gt", "wait_lt", + "MultimemReduceOp", + "multimem_ld_reduce", + "multimem_st", + "multimem_red", ] diff --git a/tilelang/language/distributed/multimem.py b/tilelang/language/distributed/multimem.py new file mode 100644 index 0000000000..034140c832 --- /dev/null +++ b/tilelang/language/distributed/multimem.py @@ -0,0 +1,79 @@ +"""Multimem operations (NVSwitch SHARP multicast) using layout-aware lowering. + +These operations use T.copy's ParallelOp + InferLayout + VectorizeLoop pipeline +to correctly handle fragment layouts, then post-process to emit multimem instructions. +""" + +from enum import Enum +from tvm import tir +from tilelang.utils.language import to_buffer_region + + +class MultimemReduceOp(Enum): + ADD = 0 + MIN = 1 + MAX = 2 + + +class _MultimemMode(Enum): + LD_REDUCE = 0 + ST = 1 + RED = 2 + + +def _multimem_impl(src, dst, mode: _MultimemMode, reduce_op: MultimemReduceOp): + """Shared implementation for all multimem operations. + + Converts src/dst to buffer regions and emits the tl.tileop.multimem intrinsic. + + Args: + src: Source (Buffer, BufferLoad with slice, or BufferRegion) + dst: Destination (Buffer, BufferLoad with slice, or BufferRegion) + mode: 0=kLdReduce, 1=kSt, 2=kRed + reduce_op: 0=ADD, 1=MIN, 2=MAX + """ + src_region = to_buffer_region(src, access_type="r") + dst_region = to_buffer_region(dst, access_type="w") + return tir.call_intrin( + "handle", + tir.op.Op.get("tl.tileop.multimem"), + src_region, + dst_region, + mode.value, + reduce_op.value, + ) + + +def multimem_ld_reduce(src, dst, reduce_op: MultimemReduceOp = MultimemReduceOp.ADD): + """Load-reduce from multicast address into local buffer. + + Uses T.copy's layout inference to handle fragment layouts correctly. + Each thread issues 128-bit multimem instructions after vectorization. + + Args: + src: Multicast source (Buffer, BufferLoad with slice, or BufferRegion) + dst: Local destination (Buffer, BufferLoad with slice, or BufferRegion) + reduce_op: Reduction operation: 0=ADD, 1=MIN, 2=MAX. + """ + return _multimem_impl(src, dst, mode=_MultimemMode.LD_REDUCE, reduce_op=reduce_op) + + +def multimem_st(src, dst): + """Store to multicast address (broadcast to all ranks). + + Args: + src: Local source (Buffer, BufferLoad with slice, or BufferRegion) + dst: Multicast destination (Buffer, BufferLoad with slice, or BufferRegion) + """ + return _multimem_impl(src, dst, mode=_MultimemMode.ST, reduce_op=MultimemReduceOp.ADD) + + +def multimem_red(src, dst, reduce_op: MultimemReduceOp = MultimemReduceOp.ADD): + """Reduce into multicast address (accumulate without read-back). + + Args: + src: Local source (Buffer, BufferLoad with slice, or BufferRegion) + dst: Multicast destination (Buffer, BufferLoad with slice, or BufferRegion) + reduce_op: Reduction operation: 0=ADD, 1=MIN, 2=MAX. + """ + return _multimem_impl(src, dst, mode=_MultimemMode.RED, reduce_op=reduce_op) diff --git a/tilelang/utils/allocator.py b/tilelang/utils/allocator.py index 347db41074..5161c798b3 100644 --- a/tilelang/utils/allocator.py +++ b/tilelang/utils/allocator.py @@ -13,7 +13,19 @@ _vmm_malloc, _vmm_free, _create_vmm_handle, + _open_vmm_handle, + _close_vmm_handle, _sync_vmm_handles, + _supports_multicast, + _mc_create, + _mc_export_handle, + _mc_import_handle, + _mc_add_device, + _mc_bind_mem, + _mc_map, + _mc_release_handle, + _mc_unmap, + _mc_get_aligned_size, ) from tilelang.utils.target import parse_device import contextlib @@ -101,6 +113,7 @@ def __init__( group: dist.ProcessGroup | None = None, align: int = 256, use_vmm: bool | None = None, + mcast_size: int | None = None, ) -> None: if size <= 0: raise ValueError("size must be > 0") @@ -114,6 +127,7 @@ def __init__( self._num_local_ranks = num_local_ranks self._group = group self._align = align + self._mcast_size_requested = mcast_size # table items: # 1. local_rank, size: 8 bytes # 2. num_local_ranks, size: 8 bytes @@ -123,6 +137,14 @@ def __init__( self._buffer_ptrs = None self._device_ids = None self._initialized = False + self._closed = False + # Multicast state + self._mcast_base_ptr = 0 + self._mcast_ptr = 0 + self._mcast_phys_ptr = 0 + self._mcast_aligned_size = 0 + self._use_multicast = False + if self._is_distributed: assert self._group is not None, "group must be provided when is_distributed is True" assert self._local_rank is not None, "local_rank must be provided when is_distributed is True" @@ -155,14 +177,135 @@ def _alloc(self): raise RuntimeError(f"cudaMalloc failed: {rc} {msg.decode() if msg else ''}") self._ptr.value = self._base_ptr.value + # Multicast buffer (only when explicitly requested via mcast_size) + if self._mcast_size_requested is not None: + assert self._use_vmm, "mcast_size requires use_vmm=True" + assert self._is_distributed, "mcast_size requires is_distributed=True" + if _supports_multicast(): + self._init_multicast_buffer() + else: + raise RuntimeError("Multicast not supported on this hardware") + + def _init_multicast_buffer(self): + """Create multicast object and map, following multi-process fabric pattern.""" + mcast_size = self._mcast_size_requested if self._mcast_size_requested else self.size + num_devices = self._num_local_ranks + aligned = _mc_get_aligned_size(mcast_size, num_devices) + self._mcast_aligned_size = aligned + + # Allocate physical memory (reuses vmm_malloc, same fabric handle type) + self._mcast_phys_ptr = _vmm_malloc(aligned) + + # Rank 0 creates MC object, exports fabric handle; broadcast to all + if self._local_rank == 0: + mcast_handle = _mc_create(aligned, num_devices) + mcast_fabric_bytes = _mc_export_handle(mcast_handle) + else: + mcast_handle = 0 + mcast_fabric_bytes = None + + obj_list = [mcast_fabric_bytes] + dist.broadcast_object_list(obj_list, src=0, group=self._group) + mcast_fabric_bytes = obj_list[0] + + # Non-rank-0 import the MC handle + if self._local_rank != 0: + mcast_handle = _mc_import_handle(mcast_fabric_bytes) + + # Each rank adds its own device + _mc_add_device(mcast_handle, self._local_rank) + + # Barrier: all devices must be added before binding + dist.barrier(self._group) + + # Each rank binds its own physical memory + _mc_bind_mem(mcast_handle, self._mcast_phys_ptr, aligned) + + # Barrier: all binds must complete before mapping + dist.barrier(self._group) + + # Each rank maps the MC object to a local VA + self._mcast_base_ptr = _mc_map(mcast_handle, aligned, num_devices) + self._mcast_ptr = self._mcast_base_ptr + + # Release handle (backing persists due to mapping) + _mc_release_handle(mcast_handle) + self._use_multicast = True + + def _allocate_mcast_tensor( + self, shape: tuple[int, ...], dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate from multicast buffer (bump-pointer). + + Returns: + (mcast_tensor, local_tensor): + mcast_tensor: backed by MC VA, for multimem read instructions + local_tensor: backed by physical VA, for writing data + """ + if not self._use_multicast: + raise RuntimeError("Multicast buffer not initialized") + + numel = _prod_shape(shape) + itemsize = _element_size_bytes(dtype) + bytes_needed = numel * itemsize + bytes_alloc = _align_up(bytes_needed, self._align) + + current_offset = self._mcast_ptr - self._mcast_base_ptr + if current_offset + bytes_alloc > self._mcast_aligned_size: + raise MemoryError( + f"Mcast allocation failed: Requesting {bytes_alloc} bytes, but only " + f"{self._mcast_aligned_size - current_offset} bytes available " + f"(total mcast size: {self._mcast_aligned_size} bytes)." + ) + + dtype_str = _dtype_to_str.get(dtype) + if dtype_str is None: + dtype_str = str(dtype).split(".")[-1] + if isinstance(shape, tuple): + shape = list(shape) + elif not isinstance(shape, list): + shape = [shape] + + mcast_t = tensor_from_ptr(self._mcast_ptr, shape, dtype_str, self._device, False) + local_t = tensor_from_ptr(self._mcast_phys_ptr + current_offset, shape, dtype_str, self._device, False) + self._mcast_ptr += bytes_alloc + return mcast_t, local_t + + def close(self): + """Explicitly free resources with proper distributed coordination. + + Must be called collectively by all ranks before process group destruction. + Safe to call multiple times. + """ + if self._closed: + return + self._closed = True + # Barrier before multicast teardown to ensure no rank is still using MC VA + if getattr(self, "_use_multicast", False) and self._group is not None: + try: + dist.barrier(self._group) + except Exception: + pass + self._free() + def _free(self): + # Free multicast resources + if getattr(self, "_mcast_base_ptr", 0) and self._mcast_base_ptr: + _mc_unmap(self._mcast_base_ptr, self._mcast_aligned_size, self._num_local_ranks) + self._mcast_base_ptr = 0 + self._mcast_ptr = 0 + if getattr(self, "_mcast_phys_ptr", 0) and self._mcast_phys_ptr: + _vmm_free(self._mcast_phys_ptr) + self._mcast_phys_ptr = 0 + self._use_multicast = False + + # Free main buffer if getattr(self, "_base_ptr", None) and self._base_ptr.value: if getattr(self, "_use_vmm", False): _vmm_free(self._base_ptr.value) self._base_ptr = ctypes.c_void_p(0) else: rc = _libcudart.cudaFree(self._base_ptr) - # mark freed even if error to avoid double free in destructor self._base_ptr = ctypes.c_void_p(0) if rc != 0: msg = _libcudart.cudaGetErrorString(rc) @@ -205,6 +348,7 @@ def initialized(self) -> bool: def _allocate_tensor( self, shape: tuple[int, ...], dtype: torch.dtype, return_peers=False, take_ownership: bool = False ) -> torch.Tensor: + numel = _prod_shape(shape) itemsize = _element_size_bytes(dtype) bytes_needed = numel * itemsize @@ -276,9 +420,15 @@ def table(self) -> torch.Tensor: def table_size(self) -> int: return self._table_size + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() + def __del__(self): with contextlib.suppress(Exception): - self._free() + self.close() def get_allocator( @@ -289,8 +439,9 @@ def get_allocator( num_local_ranks: int = 1, group: dist.ProcessGroup | None = None, use_vmm: bool | None = None, + mcast_size: int | None = None, ) -> BaseAllocator: return BaseAllocator( size, device=device, is_distributed=is_distributed, local_rank=local_rank, - num_local_ranks=num_local_ranks, group=group, use_vmm=use_vmm, + num_local_ranks=num_local_ranks, group=group, use_vmm=use_vmm, mcast_size=mcast_size, ) From dc7142ca89969196a1cb95bcbba5ec5337bf754b Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Tue, 5 May 2026 22:55:00 +0800 Subject: [PATCH 6/9] lint --- docs/suggestion.md | 339 ------------------ .../distributed/example_gemm_allreduce.py | 0 .../distributed/example_multimem_allreduce.py | 8 +- src/op/multimem.cc | 29 +- src/op/multimem.h | 14 +- src/op/multimem_rewriter.h | 13 +- src/shared_memory/shared_memory.cc | 67 ++-- src/tl_templates/cuda/multimem.h | 62 ++-- .../python/distributed/test_multicast_ops.py | 7 +- testing/python/distributed/test_vmm_ops.py | 3 +- .../distributed/shared_memory/__init__.py | 9 +- .../shared_memory/csrc/bindings.cpp | 3 +- tilelang/distributed/shared_memory/csrc/ops.h | 7 +- .../shared_memory/csrc/vmm_ops.cpp | 18 +- tilelang/distributed/utils.py | 4 +- tilelang/utils/allocator.py | 22 +- 16 files changed, 132 insertions(+), 473 deletions(-) delete mode 100644 docs/suggestion.md create mode 100644 examples/distributed/example_gemm_allreduce.py diff --git a/docs/suggestion.md b/docs/suggestion.md deleted file mode 100644 index 39d8fac5bc..0000000000 --- a/docs/suggestion.md +++ /dev/null @@ -1,339 +0,0 @@ -# TileScale Codebase 整理与重构建议 - -## 1. 解决 repo 名称与 package 名称不一致问题 - -**现状：** repo 叫 `tilescale`，Python 包叫 `tilelang`，二者完全不同。 - -**建议：** -- 明确一个正式名称。如果 `tilelang` 是对外发布的包名，`tilescale` 是项目/团队代号，则在 repo 根目录的 README、CLAUDE.md 等地方明确说明这一关系。 -- 或者，将包名统一为 `tilescale`（如果不考虑向后兼容），避免新加入的开发者混淆。 - ---- - -## 2. 消除 `primitives/` 与 `tileop/` 的职责重叠 - -**现状：** -- `tilelang/primitives/gemm/__init__.py` 是一个几乎为空的目录（只有 `__init__.py`），且其中 import 的模块（`primitives.gemm.base`, `primitives.gemm.gemm_mma`）并不存在于该目录下，实际实现在 `tileop/gemm/` 里。 -- `tilelang/tileop/gemm/` 拥有完整实现：`gemm_mma.py`, `gemm_wgmma.py`, `gemm_mfma.py`, `gemm_tcgen05.py` 等。 - -**建议：** -- 删除空壳 `primitives/` 目录，将其职责完全归入 `tileop/`（或反之统一命名）。 -- 在公共 API 层统一暴露，不要让两个名字同时对外可见。 - ---- - -## 3. 清理 `language/v2/` —— 明确新旧版本策略 - -**现状：** 存在 `tilelang/language/` 和 `tilelang/language/v2/` 两套并行实现。 - -**建议：** -- 如果 v2 是稳定的新实现，应制定迁移计划：deprecate v1 接口，并在 v1 所有入口添加 deprecation warning，最终删除 v1 代码。 -- 如果 v2 还在实验阶段，应移入 `language/experimental/` 而非独立的 `v2/` 目录，避免用户误用。 -- 不要让两个版本长期共存且没有明确说明。 - ---- - -## 4. 合并分散的 `kernel_cache.py` - -**现状：** 至少 5 处有各自的 `kernel_cache.py`： -``` -tilelang/cache/kernel_cache.py -tilelang/jit/adapter/kernel_cache.py -tilelang/jit/adapter/cutedsl/kernel_cache.py -tilelang/jit/adapter/torch/kernel_cache.py -tilelang/jit/adapter/nvrtc/kernel_cache.py -tilelang/jit/adapter/cython/kernel_cache.py -``` - -**建议：** -- 设计一个统一的 `tilelang/cache/` 模块作为单一 cache 抽象层。 -- 各 adapter（cutedsl/torch/nvrtc/cython）继承或组合此基础实现，而不是各自维护一份。 -- 这样可以避免 cache invalidation 逻辑不同步的 bug。 - ---- - -## 5. 统一 Distributed 代码的位置 - -**现状：** 分布式相关代码散落在： -- `tilelang/distributed/` — 顶层 distributed 模块（含 pynvshmem 子包、build 脚本、install 脚本） -- `tilelang/language/distributed/` — language 层面的 distributed 原语 -- `tilelang/language/distributed/multi_device/` — nvshmem.py, cpengine.py -- `examples/distributed/` — 示例 - -**建议：** -- 明确职责边界：language 层只负责 IR/语法原语的声明，runtime 层（nvshmem binding、launch、allocator）归入 `tilelang/distributed/`。 -- `tilelang/language/distributed/` 只应包含 IR-level 原语定义，具体实现（nvshmem.py, cpengine.py）应上移至 `tilelang/distributed/`。 -- `pynvshmem` 是一个独立的 C 扩展包（有自己的 `setup.py` 和 `CMakeLists.txt`），应提升为顶层子包或独立 repo，而不是嵌套在 `tilelang/distributed/pynvshmem/` 多层深处。 - ---- - -## 6. 整理测试代码——测试不应混在 package 内部 - -**现状：** 测试文件散落在： -- `testing/python/` — 主测试目录（正确） -- `tilelang/distributed/testing/` — 在包内部！ -- `tilelang/distributed/pynvshmem/testing/` — 也在包内部！ -- `examples/` — 大量 `test_*.py` 和 `regression_*.py` 混在示例里 - -**建议：** -- 确立唯一的测试根目录：`testing/`（或 `tests/`）。 -- 将 `tilelang/distributed/testing/` 和 `tilelang/distributed/pynvshmem/testing/` 的内容迁移到 `testing/python/distributed/`。 -- `examples/` 里的 `test_*.py` / `regression_*.py` 要么移入 `testing/`，要么与示例文件分开存放（比如每个 example 子目录只保留 `example_*.py`，测试单独放 `testing/`）。 -- 在 `testing/` 下的目录结构应与 `tilelang/` 包结构一一对应，便于查找。 - ---- - -## 7. 修复 `__init__.py` 中的重复版本解析逻辑 - -**现状：** `tilelang/__init__.py` 第 11-65 行有两套 `__version__` 计算逻辑，`__version__` 被赋值两次（第 48 行和第 53-65 行），第一次的计算完全被覆盖。 - -**建议：** -- 只保留一套版本解析逻辑（推荐 `importlib.metadata` + fallback），删除重复代码。 -- 或者将版本逻辑抽到单独的 `_version.py` 文件。 - ---- - -## 8. 规范 `language/` 内部命名不一致问题 - -**现状：** 同类功能有时带 `_op` 后缀，有时不带： -- `language/copy.py` vs `language/copy_op.py` -- `language/fill.py` vs `language/fill_op.py` -- `language/gemm.py` vs `language/gemm_op.py` -- `language/reduce.py` vs `language/reduce_op.py` - -**建议：** -- 确定一个命名规范（建议 `_op` 后缀表示操作符对象，无后缀表示 language-level 函数入口），并全面统一。 -- 或者合并重复文件，明确每个文件的职责边界。 - ---- - -## 9. 整合 `contrib/cutedsl/` 与 `jit/adapter/cutedsl/` - -**现状：** -- `tilelang/contrib/cutedsl/` 有底层 CUDA 原语（mbar, ldsm, cpasync, gemm_V1, reduce 等） -- `tilelang/jit/adapter/cutedsl/` 有 adapter/wrapper/libgen 等 JIT 逻辑 - -**建议：** -- `contrib/cutedsl/` 的角色应该是"第三方或底层 CUDA primitive 封装"，`jit/adapter/cutedsl/` 是"JIT pipeline 对 cutedsl 的适配"，二者职责不同，但命名上容易让人以为是同一件事。 -- 建议将 `contrib/cutedsl/` 重命名为 `tilelang/backends/cutedsl/` 或 `tilelang/codegen/cutedsl/`，以更清晰表达其用途。 - ---- - -## 10. 整理 `examples/` 目录结构 - -**现状：** `examples/` 混杂了多种文件类型，缺乏一致结构： -- `example_*.py` — 示例代码 -- `test_example_*.py` — 测试 -- `regression_*.py` — 回归测试 -- `benchmark_*.py` — 性能基准 -- 部分有 `README.md`，部分没有 -- `examples/pytest.ini` 只有一个，但 examples 子目录结构不统一 - -**建议：** -- 每个 example 子目录统一结构：`example_*.py`（核心示例）+ `README.md`（必须有）。 -- 测试/回归测试移入 `testing/`，不要放在 examples 里。 -- Benchmark 单独建 `benchmarks/` 顶级目录（现有 `benchmark/blocksparse_attention/` 也移入此处），与 `examples/` 分开。 - ---- - -## 11. 清理重复图片资源 - -**现状：** 相似图片存在于多个地方： -- `images/`（根目录） -- `docs/_static/img/` -- `examples/deepseek_mla/figures/` -- `docs/_static/img/mla_hopper/` - -**建议：** -- 文档用图统一放 `docs/_static/img/`，examples 里的 figures 只保留真正 example 专属的。 -- 根目录 `images/` 合并进 `docs/_static/img/`，避免两套图片各自维护。 - ---- - -## 12. `tilelang/utils/ts_ext/` 的嵌套 `setup.py` 问题 - -**现状：** `tilelang/utils/ts_ext/` 内部有独立的 `setup.py`，说明这是一个独立的 C 扩展包，但被嵌套在 `utils/` 深处。 - -**建议：** -- 将 `ts_ext` 提升为顶级子包（类似 `pynvshmem` 的处理方式，或者更好地统一到同一个构建系统中）。 -- 或者把它的构建整合进主 `setup.py` / `CMakeLists.txt`，避免嵌套独立构建脚本。 - ---- - -## 优先级总结 - -| 优先级 | 建议 | -|--------|------| -| 🔴 高 | #7 重复版本逻辑（bug 风险）、#4 kernel_cache 重复（一致性风险）、#6 测试混在包内 | -| 🟡 中 | #2 primitives/tileop 重叠、#5 distributed 代码分散、#8 命名不一致、#3 v2 策略 | -| 🟢 低 | #1 名称统一、#9 contrib vs adapter、#10 examples 结构、#11 图片整理、#12 ts_ext 位置 | - ---- - -# TileScale 专项整理建议 - -## 13. 统一两个职责重叠的 `init_dist*` 函数 - -**现状：** `tilelang/distributed/utils.py` 中存在两个功能相似但 API 完全不同的初始化函数： -- `init_dist(local_rank, num_local_ranks)` — 旧式，需要手动传参，配合 `mp.spawn` 使用 -- `init_distributed(return_tp_group, init_nvshmem, return_lc_group)` — 新式，从环境变量读取，配合 `torchrun` 使用 - -两者在不同 example 中被混用，导致新用户不知道该用哪个。 - -**建议：** -- 统一为一个函数，明确以 `torchrun` 风格（环境变量）为主； -- `init_dist` 如果仍需保留，加 `@deprecated` 装饰器，并在 docstring 说明迁移路径。 - ---- - -## 14. 修复 `wait_eq` 名称严重冲突 - -**现状：** 两处都叫 `wait_eq`，但语义完全不同： -- `tilelang/language/distributed/common.py:wait_eq()` — **IR 层**：生成 NVSHMEM 的 `signal_wait_until` intrinsic，在 GPU 核函数内使用 -- `tilelang/distributed/utils.py:wait_eq()` — **Host 层**：调用 `cuStreamWaitValue32`，在 CPU 侧 stream 上等待 - -两者都被 import 进同一个项目，极易混淆，且 bug 难以定位。 - -**建议：** -- Host 侧重命名为 `stream_wait_eq` 或 `host_wait_signal`，与 IR 层的 `wait_eq` 区分； -- 在两个函数的 docstring 中明确标注各自的作用层级（host/device）。 - ---- - -## 15. 消除 IPC handle 同步逻辑的重复实现 - -**现状：** 以下两处独立实现了几乎相同的"收集各 rank 的 IPC handle 并映射到本地 GPU 指针"逻辑： -- `distributed/utils.py:create_dist_tensor()` — 通过 `all_gather_object` 收集 IPC handle，调用 `_sync_ipc_handles` -- `distributed/allocator.py:BaseAllocator._init_table()` — 相同流程，只是封装在 Allocator 里 - -**建议：** -- 将 IPC handle 同步逻辑抽成一个内部函数 `_exchange_ipc_handles(group, local_ptr) -> buffer_ptrs`； -- `create_dist_tensor` 和 `BaseAllocator._init_table` 都调用这个函数，避免两套独立维护。 - ---- - -## 16. 集中管理 `cuda-python` 版本兼容 shim - -**现状：** 以下代码段在多处重复出现： -```python -cuda_python_version = importlib.metadata.version("cuda-python") -from packaging import version -if version.parse(cuda_python_version) >= version.parse("12.8.0"): - from cuda.bindings import driver as cuda -else: - from cuda import cuda -``` -出现于 `distributed/utils.py`、`example_allgather_gemm_overlapped.py` 等多处。 - -**建议：** -- 将此兼容逻辑放在 `tilelang/distributed/_cuda_compat.py` 中，统一 export `cuda, cudart` 对象； -- 其他地方一律 `from tilelang.distributed._cuda_compat import cuda, cudart`，不再各自做版本检测。 - ---- - -## 17. 建立统一的对称内存（Symmetric Memory）抽象 - -**现状：** TileScale 目前有两套完全独立的分布式 tensor 分配路径： -- **NVSHMEM 路径**：`pynvshmem.nvshmem_create_tensor(shape, dtype)` — 通过 NVSHMEM 分配对称内存 -- **IPC 路径**：`tilelang.distributed.create_tensor(shape, dtype)` 或 `BaseAllocator` — 通过 `cudaMalloc` + IPC handle 交换 - -两者语义不同、使用方式不同，但 examples 中混用，用户难以理解应该用哪个，以及何时该回退到 IPC 路径。 - -**建议：** -- 设计一个统一的 `tilescale.memory.SymmetricBuffer` 抽象，内部根据是否有 nvshmem 自动选择后端； -- 在文档中明确说明两种路径的适用场景（intranode IPC vs internode NVSHMEM）。 - ---- - -## 18. 将 shell 脚本从 package 目录移出 - -**现状：** 以下三个 shell 脚本直接放在 `tilelang/distributed/` 包目录里，会被打包进 Python wheel： -- `tilelang/distributed/launch.sh` — 分布式启动脚本 -- `tilelang/distributed/build_nvshmem.sh` — 构建 NVSHMEM 的脚本 -- `tilelang/distributed/install_deepep.sh` — 安装 DeepEP 的脚本 - -**建议：** -- 统一移到 repo 根目录下的 `scripts/` 目录（如 `scripts/launch.sh`, `scripts/build_nvshmem.sh`）； -- `scripts/` 不会被 pip install 打包，符合标准 Python 项目规范。 - ---- - -## 19. 拆分过于庞杂的 `distributed/utils.py` - -**现状：** `tilelang/distributed/utils.py` 是一个 400 行的"万能工具箱"，混杂了： -- 进程组初始化（`init_dist`, `init_distributed`） -- IPC tensor 创建（`create_tensor`, `create_dist_tensor`） -- CUDA stream 信号操作（`set_signal`, `wait_eq`） -- 性能测量（`perf_fn`） -- 调试打印（`dist_print`） -- NVLink 拓扑检测（`has_fullmesh_nvlink`, `NvidiaSmiUtil`） -- 数据生成（`generate_data`, `_make_tensor`） - -**建议：** -``` -tilelang/distributed/ -├── init.py # init_dist, init_distributed -├── memory.py # IPC/tensor allocation -├── signal.py # CUDA stream signal operations -├── topo.py # NVLink topology (NvidiaSmiUtil, has_fullmesh_nvlink) -├── perf.py # perf_fn, dist_print -└── data.py # generate_data, _make_tensor -``` - ---- - -## 20. 将 example 中的复用库代码提升为正式模块 - -**现状：** 以下文件放在 `examples/distributed/` 目录中，却被其他 example 直接 import，本质上已经是库代码： -- `examples/distributed/sp_ag_attention_intra_node.py` — 被 `example_sp_ag_attention_intra_node.py` import -- `examples/distributed/nvshmem/gemm_rs_utils.py` — `BarrierAllContext`, `ReduceScatter2DContext` -- `examples/distributed/nvshmem/reduce_scatter.py` — `reduce_scatter_2d_op` 等 -- `examples/distributed/deepseek_deepep/deepep_utils.py` — `Config`, DeepEP 相关常量 - -这些文件用相对路径 import，只在同目录运行时有效，无法被其他地方正确引用。 - -**建议：** -- 将复用代码提升为 `tilelang/distributed/` 下的正式子模块； -- examples 只保留调用代码，不再充当隐式的库。 - ---- - -## 21. 清理 `launch.sh` 中泄露的内部集群环境变量 - -**现状：** `tilelang/distributed/launch.sh` 中有： -```bash -master_addr=${ARNOLD_WORKER_0_HOST:="127.0.0.1"} -master_port=$(echo "$ARNOLD_WORKER_0_PORT" | cut -d "," -f 1) -``` -`ARNOLD_*` 是 ByteDance 内部集群（Arnold 训练平台）的环境变量，不应出现在公开代码中。 - -**建议：** -- 将 `ARNOLD_WORKER_0_HOST` 替换为标准的 `MASTER_ADDR`，`ARNOLD_WORKER_0_PORT` 替换为 `MASTER_PORT`； -- 清理所有内部平台相关的硬编码假设（如 `IB_HCA=mlx5`、`BYTED_TORCH_BYTECCL` 等）。 - ---- - -## 22. 为 `tilescale_ext` C 扩展建立清晰的 API 文档与稳定性边界 - -**现状：** `tilelang/distributed/__init__.py` 直接暴露： -```python -from tilescale_ext import _create_tensor, _create_ipc_handle, _sync_ipc_handles -``` -- 以 `_` 开头的私有函数被直接 re-export 为公共 API -- 没有任何 stub 文件（`.pyi`）说明函数签名 -- 不清楚哪些是稳定 API，哪些是内部实现细节 - -**建议：** -- 在 `tilelang/distributed/` 层面做好封装，禁止外部直接调用 `_` 前缀的 `tilescale_ext` 函数； -- 提供 `tilescale_ext.pyi` stub 文件，便于 IDE 补全和类型检查； -- 明确区分稳定公共 API 与内部实现。 - ---- - -## 优先级汇总（TileScale 专项） - -| 优先级 | 建议 | -|--------|------| -| 🔴 高 | #14 wait_eq 名称冲突（正确性风险）、#15 IPC 逻辑重复（一致性风险）、#21 内部环境变量泄露 | -| 🟡 中 | #13 init_dist 统一、#16 cuda-python shim 集中化、#19 utils.py 拆分、#20 example 库代码提升 | -| 🟢 低 | #17 对称内存抽象、#18 shell 脚本移位、#22 tilescale_ext API 文档 | diff --git a/examples/distributed/example_gemm_allreduce.py b/examples/distributed/example_gemm_allreduce.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/distributed/example_multimem_allreduce.py b/examples/distributed/example_multimem_allreduce.py index 57e1f702f8..9247d9e040 100644 --- a/examples/distributed/example_multimem_allreduce.py +++ b/examples/distributed/example_multimem_allreduce.py @@ -39,11 +39,11 @@ def main( with T.Kernel(T.ceildiv(N, block_N), threads=threads) as (bx,): result_local = T.alloc_fragment([block_N], "float32") T.multimem_ld_reduce( - mcast_buf[bx * block_N:(bx + 1) * block_N], + mcast_buf[bx * block_N : (bx + 1) * block_N], result_local, reduce_op=T.MultimemReduceOp.ADD, ) - T.copy(result_local, result[bx * block_N:(bx + 1) * block_N]) + T.copy(result_local, result[bx * block_N : (bx + 1) * block_N]) return main @@ -120,6 +120,4 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): parser.add_argument("--print_source", action="store_true") args = parser.parse_args() - torch.multiprocessing.spawn( - main, args=(args.num_processes, args), nprocs=args.num_processes, join=True - ) + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) diff --git a/src/op/multimem.cc b/src/op/multimem.cc index edc73a159d..8880150d26 100644 --- a/src/op/multimem.cc +++ b/src/op/multimem.cc @@ -3,7 +3,8 @@ * \brief Unified multimem operator implementation. * * Reuses CopyNode's ParallelOp + InferLayout + VectorizeLoop pipeline, - * then post-processes to replace mcast buffer accesses with multimem instructions. + * then post-processes to replace mcast buffer accesses with multimem + * instructions. */ #include "multimem.h" @@ -28,11 +29,11 @@ namespace tl { using namespace tir; // === MultimemOp Constructor === -// args[0]: src region (tl.region call), args[1]: dst region, args[2]: mode, args[3]: reduce_op +// args[0]: src region (tl.region call), args[1]: dst region, args[2]: mode, +// args[3]: reduce_op MultimemOp::MultimemOp(Array args, Map annotations) { - ObjectPtr node = - tvm::ffi::make_object(); + ObjectPtr node = tvm::ffi::make_object(); // Parse buffer regions using same utility as CopyNode Array rgs[2]; @@ -90,7 +91,7 @@ MultimemOp::MultimemOp(Array args, // 128-bit per multimem instruction => width = 128 / dtype_bits int MultimemOpNode::GetCoalescedWidth() const { int bits = src->dtype.bits(); - return 128 / bits; // f32->4, f16->8, bf16->8 + return 128 / bits; // f32->4, f16->8, bf16->8 } // === MakeIterVars === @@ -212,7 +213,8 @@ For MultimemOpNode::MakeSIMTLoop(arith::Analyzer *analyzer) const { } // === InferLayout === -// Delegates to ParallelOp for layout inference (same as CopyNode::LowerNormalCopy) +// Delegates to ParallelOp for layout inference (same as +// CopyNode::LowerNormalCopy) LayoutMap MultimemOpNode::InferLayout(const LayoutInferArgs &T, InferLevel level) const { // Multimem ops always go through the ParallelOp path during Lower. @@ -221,7 +223,8 @@ LayoutMap MultimemOpNode::InferLayout(const LayoutInferArgs &T, } // === Lower === -// The main lowering path: MakeSIMTLoop -> ParallelOp pipeline -> MultimemRewriter +// The main lowering path: MakeSIMTLoop -> ParallelOp pipeline -> +// MultimemRewriter Stmt MultimemOpNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const { // Step 1-2: Create SIMT loop and fuse/transform @@ -236,8 +239,13 @@ Stmt MultimemOpNode::Lower(const LowerArgs &T, std::vector levels = {InferLevel::kCommon, InferLevel::kStrict, InferLevel::kFree}; for (auto level : levels) { - par_op->InferLayout({T.target, T.thread_bounds, T.layout_map, analyzer, - false, T.buffer_remap, {}}, + par_op->InferLayout({T.target, + T.thread_bounds, + T.layout_map, + analyzer, + false, + T.buffer_remap, + {}}, level); } @@ -246,7 +254,8 @@ Stmt MultimemOpNode::Lower(const LowerArgs &T, Stmt result = LowerParallelLoop(par_op->GetRoot(), loop_layout, T.thread_var, analyzer, par_op->GetPredicate(T.thread_var)); - // Step 5: Post-process — replace mcast buffer accesses with multimem call_extern + // Step 5: Post-process — replace mcast buffer accesses with multimem + // call_extern Buffer mcast_buf = (mode == MultimemMode::kLdReduce) ? src : dst; // Remap the mcast buffer if needed if (T.buffer_remap.count(mcast_buf)) { diff --git a/src/op/multimem.h b/src/op/multimem.h index b9456fa9cd..44e0b774bf 100644 --- a/src/op/multimem.h +++ b/src/op/multimem.h @@ -33,18 +33,20 @@ enum class MultimemMode : int { kLdReduce = 0, kSt = 1, kRed = 2 }; * - kRed: reduce into multicast address (no read-back) * * Lower flow: - * 1. MakeSIMTLoop: creates element-wise parallel loop (BufferLoad -> BufferStore) + * 1. MakeSIMTLoop: creates element-wise parallel loop (BufferLoad -> + * BufferStore) * 2. ParallelLoopFuser::Fuse + ParallelLoopTransformer::Substitute * 3. ParallelOp -> InferLayout at multiple levels * 4. LowerParallelLoop (PartitionLoop + VectorizeLoop) - * 5. MultimemRewriter: post-process to replace mcast buffer accesses with call_extern + * 5. MultimemRewriter: post-process to replace mcast buffer accesses with + * call_extern */ class MultimemOpNode : public TileOperatorNode { public: Buffer src, dst; Array src_range, dst_range; MultimemMode mode; - int reduce_op; // 0=ADD, 1=MIN, 2=MAX + int reduce_op; // 0=ADD, 1=MIN, 2=MAX TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.MultimemOp", MultimemOpNode, TileOperatorNode); @@ -76,9 +78,9 @@ class MultimemOp : public TileOperator { public: TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MultimemOp, TileOperator, MultimemOpNode); - TVM_DLL MultimemOp( - Array args, - Map annotations = Map()); + TVM_DLL + MultimemOp(Array args, + Map annotations = Map()); static const Op &Get(); }; diff --git a/src/op/multimem_rewriter.h b/src/op/multimem_rewriter.h index b63b4ad672..d92fe8ce95 100644 --- a/src/op/multimem_rewriter.h +++ b/src/op/multimem_rewriter.h @@ -1,7 +1,7 @@ /*! * \file tl/op/multimem_rewriter.h - * \brief Post-process IR to replace vectorized BufferLoad/Store on mcast buffers - * with multimem call_extern instructions. + * \brief Post-process IR to replace vectorized BufferLoad/Store on mcast + * buffers with multimem call_extern instructions. */ #ifndef TVM_TL_OP_MULTIMEM_REWRITER_H_ @@ -112,8 +112,9 @@ class MultimemRewriter : public StmtExprMutator { int reduce_op_; /*! - * \brief Try to rewrite a kVectorized for-loop containing a mcast BufferStore. - * Returns the replacement Stmt, or undefined if the pattern doesn't match. + * \brief Try to rewrite a kVectorized for-loop containing a mcast + * BufferStore. Returns the replacement Stmt, or undefined if the pattern + * doesn't match. */ Stmt TryRewriteVectorizedLoop(const ForNode *op, int lanes) { // The body should be a single BufferStore (possibly wrapped in IfThenElse) @@ -197,8 +198,8 @@ class MultimemRewriter : public StmtExprMutator { */ Stmt MakeMultimemCall(const Buffer &local_buf, const Array &local_indices, - const Buffer &mc_buf, - const Array &mc_indices, int lanes) const { + const Buffer &mc_buf, const Array &mc_indices, + int lanes) const { std::string func_name = MakeFuncName(lanes, local_buf->dtype); Array args; diff --git a/src/shared_memory/shared_memory.cc b/src/shared_memory/shared_memory.cc index 41a25c6985..363a070a88 100644 --- a/src/shared_memory/shared_memory.cc +++ b/src/shared_memory/shared_memory.cc @@ -81,7 +81,7 @@ static int64_t vmm_malloc_impl(int64_t size_raw) { size_t granularity = 0; SM_CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, - CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + CU_MEM_ALLOC_GRANULARITY_MINIMUM)); size_t size = align_to_granularity((size_t)size_raw, granularity); @@ -89,7 +89,8 @@ static int64_t vmm_malloc_impl(int64_t size_raw) { SM_CU_CHECK(cuMemCreate(&handle, size, &prop, 0)); void *ptr = nullptr; - SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, granularity, 0, 0)); + SM_CU_CHECK( + cuMemAddressReserve((CUdeviceptr *)&ptr, size, granularity, 0, 0)); SM_CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, handle, 0)); cu_mem_set_access_all(ptr, size); @@ -123,7 +124,7 @@ static ffi::Bytes create_vmm_handle_impl(int64_t ptr_val) { CUmemFabricHandle fabric_handle; SM_CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, - CU_MEM_HANDLE_TYPE_FABRIC, 0)); + CU_MEM_HANDLE_TYPE_FABRIC, 0)); std::string raw(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); std::memcpy(&raw[0], &size, sizeof(size_t)); @@ -139,12 +140,11 @@ static int64_t open_vmm_handle_impl(ffi::Bytes handle_bytes) { std::memcpy(&size, data, sizeof(size_t)); CUmemFabricHandle fabric_handle; - std::memcpy(&fabric_handle, data + sizeof(size_t), - sizeof(CUmemFabricHandle)); + std::memcpy(&fabric_handle, data + sizeof(size_t), sizeof(CUmemFabricHandle)); CUmemGenericAllocationHandle alloc_handle; SM_CU_CHECK(cuMemImportFromShareableHandle(&alloc_handle, &fabric_handle, - CU_MEM_HANDLE_TYPE_FABRIC)); + CU_MEM_HANDLE_TYPE_FABRIC)); void *ptr = nullptr; SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, 0, 0, 0)); @@ -242,7 +242,7 @@ static int64_t mc_create_impl(int64_t size_raw, int64_t num_devices) { size_t granularity = 0; SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, - CU_MULTICAST_GRANULARITY_RECOMMENDED)); + CU_MULTICAST_GRANULARITY_RECOMMENDED)); size_t size = align_to_granularity((size_t)size_raw, granularity); prop.size = size; @@ -255,11 +255,12 @@ static int64_t mc_create_impl(int64_t size_raw, int64_t num_devices) { // Export multicast handle as fabric handle bytes (for sharing across processes) static ffi::Bytes mc_export_handle_impl(int64_t mc_handle_val) { - CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + CUmemGenericAllocationHandle mc_handle = + (CUmemGenericAllocationHandle)mc_handle_val; CUmemFabricHandle fabric_handle; SM_CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, mc_handle, - CU_MEM_HANDLE_TYPE_FABRIC, 0)); + CU_MEM_HANDLE_TYPE_FABRIC, 0)); return ffi::Bytes(reinterpret_cast(&fabric_handle), sizeof(CUmemFabricHandle)); @@ -274,14 +275,15 @@ static int64_t mc_import_handle_impl(ffi::Bytes handle_bytes) { CUmemGenericAllocationHandle mc_handle; SM_CU_CHECK(cuMemImportFromShareableHandle(&mc_handle, &fabric_handle, - CU_MEM_HANDLE_TYPE_FABRIC)); + CU_MEM_HANDLE_TYPE_FABRIC)); return (int64_t)mc_handle; } // Add a device to the multicast object static void mc_add_device_impl(int64_t mc_handle_val, int64_t device_id) { - CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + CUmemGenericAllocationHandle mc_handle = + (CUmemGenericAllocationHandle)mc_handle_val; CUdevice device; SM_CU_CHECK(cuDeviceGet(&device, (int)device_id)); SM_CU_CHECK(cuMulticastAddDevice(mc_handle, device)); @@ -289,8 +291,9 @@ static void mc_add_device_impl(int64_t mc_handle_val, int64_t device_id) { // Bind a physical memory VA (from vmm_malloc) to the multicast object static void mc_bind_mem_impl(int64_t mc_handle_val, int64_t ptr_val, - int64_t size) { - CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + int64_t size) { + CUmemGenericAllocationHandle mc_handle = + (CUmemGenericAllocationHandle)mc_handle_val; void *ptr = reinterpret_cast((uintptr_t)ptr_val); // Retrieve the physical allocation handle from the mapped pointer @@ -298,7 +301,8 @@ static void mc_bind_mem_impl(int64_t mc_handle_val, int64_t ptr_val, SM_CU_CHECK(cuMemRetainAllocationHandle(&phys_handle, ptr)); // Bind to multicast - SM_CU_CHECK(cuMulticastBindMem(mc_handle, 0, phys_handle, 0, (size_t)size, 0)); + SM_CU_CHECK( + cuMulticastBindMem(mc_handle, 0, phys_handle, 0, (size_t)size, 0)); // Release the temporary handle reference SM_CU_CHECK(cuMemRelease(phys_handle)); @@ -306,8 +310,9 @@ static void mc_bind_mem_impl(int64_t mc_handle_val, int64_t ptr_val, // Map multicast object to a VA, returns mc_ptr. Does NOT release handle. static int64_t mc_map_impl(int64_t mc_handle_val, int64_t size_raw, - int64_t num_devices) { - CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + int64_t num_devices) { + CUmemGenericAllocationHandle mc_handle = + (CUmemGenericAllocationHandle)mc_handle_val; CUmulticastObjectProp prop = {}; prop.numDevices = (unsigned int)num_devices; @@ -315,11 +320,12 @@ static int64_t mc_map_impl(int64_t mc_handle_val, int64_t size_raw, size_t granularity = 0; SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, - CU_MULTICAST_GRANULARITY_RECOMMENDED)); + CU_MULTICAST_GRANULARITY_RECOMMENDED)); size_t size = align_to_granularity((size_t)size_raw, granularity); void *mc_ptr = nullptr; - SM_CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&mc_ptr, size, granularity, 0, 0)); + SM_CU_CHECK( + cuMemAddressReserve((CUdeviceptr *)&mc_ptr, size, granularity, 0, 0)); SM_CU_CHECK(cuMemMap((CUdeviceptr)mc_ptr, size, 0, mc_handle, 0)); cu_mem_set_access_all(mc_ptr, size); @@ -328,13 +334,14 @@ static int64_t mc_map_impl(int64_t mc_handle_val, int64_t size_raw, // Release a multicast handle (call after map) static void mc_release_handle_impl(int64_t mc_handle_val) { - CUmemGenericAllocationHandle mc_handle = (CUmemGenericAllocationHandle)mc_handle_val; + CUmemGenericAllocationHandle mc_handle = + (CUmemGenericAllocationHandle)mc_handle_val; SM_CU_CHECK(cuMemRelease(mc_handle)); } // Free multicast VA mapping static void mc_unmap_impl(int64_t mc_ptr_val, int64_t size_raw, - int64_t num_devices) { + int64_t num_devices) { void *ptr = reinterpret_cast((uintptr_t)mc_ptr_val); CUmulticastObjectProp prop = {}; @@ -343,7 +350,7 @@ static void mc_unmap_impl(int64_t mc_ptr_val, int64_t size_raw, size_t granularity = 0; SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, - CU_MULTICAST_GRANULARITY_RECOMMENDED)); + CU_MULTICAST_GRANULARITY_RECOMMENDED)); size_t size = align_to_granularity((size_t)size_raw, granularity); SM_CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); @@ -358,7 +365,7 @@ static int64_t mc_get_aligned_size_impl(int64_t size_raw, int64_t num_devices) { size_t granularity = 0; SM_CU_CHECK(cuMulticastGetGranularity(&granularity, &prop, - CU_MULTICAST_GRANULARITY_RECOMMENDED)); + CU_MULTICAST_GRANULARITY_RECOMMENDED)); return (int64_t)align_to_granularity((size_t)size_raw, granularity); } @@ -369,8 +376,8 @@ static int64_t mc_get_aligned_size_impl(int64_t size_raw, int64_t num_devices) { // for local rank). We pass individual handle open results back via buffer_ptrs. // packed_handles: num_ranks concatenated raw handle bytes static void sync_vmm_handles_impl(int64_t rank, int64_t num_ranks, - int64_t buffer_ptrs_gpu_addr, - ffi::Bytes packed_handles) { + int64_t buffer_ptrs_gpu_addr, + ffi::Bytes packed_handles) { const size_t handle_size = sizeof(size_t) + sizeof(CUmemFabricHandle); ICHECK(packed_handles.size() == handle_size * (size_t)num_ranks); @@ -386,14 +393,14 @@ static void sync_vmm_handles_impl(int64_t rank, int64_t num_ranks, void **gpu_ptr = reinterpret_cast((uintptr_t)buffer_ptrs_gpu_addr); SM_CUDA_CHECK(cudaMemcpy(gpu_ptr, buffer_ptrs.data(), - sizeof(void *) * buffer_ptrs.size(), - cudaMemcpyHostToDevice)); + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); SM_CUDA_CHECK(cudaDeviceSynchronize()); } static void sync_ipc_handles_impl(int64_t rank, int64_t num_ranks, - int64_t buffer_ptrs_gpu_addr, - ffi::Bytes packed_handles) { + int64_t buffer_ptrs_gpu_addr, + ffi::Bytes packed_handles) { ICHECK(packed_handles.size() == CUDA_IPC_HANDLE_SIZE * (size_t)num_ranks); std::vector buffer_ptrs(num_ranks, nullptr); @@ -409,8 +416,8 @@ static void sync_ipc_handles_impl(int64_t rank, int64_t num_ranks, void **gpu_ptr = reinterpret_cast((uintptr_t)buffer_ptrs_gpu_addr); SM_CUDA_CHECK(cudaMemcpy(gpu_ptr, buffer_ptrs.data(), - sizeof(void *) * buffer_ptrs.size(), - cudaMemcpyHostToDevice)); + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); SM_CUDA_CHECK(cudaDeviceSynchronize()); } diff --git a/src/tl_templates/cuda/multimem.h b/src/tl_templates/cuda/multimem.h index 7b449e8b63..6471b2a84e 100644 --- a/src/tl_templates/cuda/multimem.h +++ b/src/tl_templates/cuda/multimem.h @@ -19,16 +19,14 @@ enum class ReduceOp { ADD = 0, MIN = 1, MAX = 2 }; // --- LdReduceV4: 128-bit load-reduce from multicast address --- -template -struct LdReduceV4 { +template struct LdReduceV4 { TL_DEVICE static void run(void *, const void *) { static_assert(always_false_v, "tl::multimem::LdReduceV4: unsupported dtype/op combination"); } }; -template <> -struct LdReduceV4 { +template <> struct LdReduceV4 { TL_DEVICE static void run(void *dst, const void *mcast_ptr) { int4 ret; asm volatile( @@ -41,8 +39,7 @@ struct LdReduceV4 { } }; -template <> -struct LdReduceV4 { +template <> struct LdReduceV4 { TL_DEVICE static void run(void *dst, const void *mcast_ptr) { int4 ret; asm volatile( @@ -55,8 +52,7 @@ struct LdReduceV4 { } }; -template <> -struct LdReduceV4 { +template <> struct LdReduceV4 { TL_DEVICE static void run(void *dst, const void *mcast_ptr) { int4 ret; asm volatile( @@ -71,38 +67,34 @@ struct LdReduceV4 { // --- StV4: 128-bit store to multicast address --- -template -struct StV4 { +template struct StV4 { TL_DEVICE static void run(void *, const void *) { static_assert(always_false_v, "tl::multimem::StV4: unsupported dtype"); } }; -template <> -struct StV4 { +template <> struct StV4 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { int4 val = *reinterpret_cast(src); - asm volatile( - "multimem.st.relaxed.sys.global.v4.b32 [%0], {%1, %2, %3, %4};" - : - : "l"(mcast_ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) - : "memory"); + asm volatile("multimem.st.relaxed.sys.global.v4.b32 [%0], {%1, %2, %3, %4};" + : + : "l"(mcast_ptr), "r"(val.x), "r"(val.y), "r"(val.z), + "r"(val.w) + : "memory"); } }; // --- RedV4: 128-bit reduce into multicast address --- -template -struct RedV4 { +template struct RedV4 { TL_DEVICE static void run(void *, const void *) { static_assert(always_false_v, "tl::multimem::RedV4: unsupported dtype/op combination"); } }; -template <> -struct RedV4 { +template <> struct RedV4 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { int4 val = *reinterpret_cast(src); asm volatile( @@ -113,37 +105,33 @@ struct RedV4 { } }; -template <> -struct RedV4 { +template <> struct RedV4 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { // multimem.red min not directly available as v4; use scalar fallback const float *src_f = reinterpret_cast(src); const char *mc_bytes = reinterpret_cast(mcast_ptr); - #pragma unroll +#pragma unroll for (int i = 0; i < 4; i++) { unsigned val = __float_as_uint(src_f[i]); - asm volatile( - "multimem.red.relaxed.sys.global.min.f32 [%0], %1;" - : - : "l"(mc_bytes + i * 4), "r"(val) - : "memory"); + asm volatile("multimem.red.relaxed.sys.global.min.f32 [%0], %1;" + : + : "l"(mc_bytes + i * 4), "r"(val) + : "memory"); } } }; -template <> -struct RedV4 { +template <> struct RedV4 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { const float *src_f = reinterpret_cast(src); const char *mc_bytes = reinterpret_cast(mcast_ptr); - #pragma unroll +#pragma unroll for (int i = 0; i < 4; i++) { unsigned val = __float_as_uint(src_f[i]); - asm volatile( - "multimem.red.relaxed.sys.global.max.f32 [%0], %1;" - : - : "l"(mc_bytes + i * 4), "r"(val) - : "memory"); + asm volatile("multimem.red.relaxed.sys.global.max.f32 [%0], %1;" + : + : "l"(mc_bytes + i * 4), "r"(val) + : "memory"); } } }; diff --git a/testing/python/distributed/test_multicast_ops.py b/testing/python/distributed/test_multicast_ops.py index 57f650b135..0c174375e1 100644 --- a/testing/python/distributed/test_multicast_ops.py +++ b/testing/python/distributed/test_multicast_ops.py @@ -58,8 +58,7 @@ def test_distributed_multicast_allocator(rank, world_size): peer_t = allocator.get_peer_tensor(peer_rank, (256,), torch.bfloat16) peer_val = peer_t[0].item() expected_val = float(peer_rank + 1) - assert abs(peer_val - expected_val) < 1e-2, \ - f"rank {local_rank}: peer[{peer_rank}] = {peer_val}, expected {expected_val}" + assert abs(peer_val - expected_val) < 1e-2, f"rank {local_rank}: peer[{peer_rank}] = {peer_val}, expected {expected_val}" dist.barrier() @@ -69,9 +68,7 @@ def test_distributed_multicast_allocator(rank, world_size): def main(): parser = argparse.ArgumentParser() - parser.add_argument( - "--distributed", action="store_true", help="Run multi-GPU tests (requires torchrun)" - ) + parser.add_argument("--distributed", action="store_true", help="Run multi-GPU tests (requires torchrun)") args = parser.parse_args() if args.distributed: diff --git a/testing/python/distributed/test_vmm_ops.py b/testing/python/distributed/test_vmm_ops.py index eef44463c2..0023ee038e 100644 --- a/testing/python/distributed/test_vmm_ops.py +++ b/testing/python/distributed/test_vmm_ops.py @@ -11,7 +11,6 @@ import argparse import os -import sys import torch import torch.distributed as dist @@ -96,7 +95,7 @@ def test_distributed_vmm(rank, world_size): # and handles the VMM handle exchange internally allocator = BaseAllocator( size=1024 * 1024, # 1 MB - device=f"cuda", + device="cuda", is_distributed=True, local_rank=local_rank, num_local_ranks=world_size, diff --git a/tilelang/distributed/shared_memory/__init__.py b/tilelang/distributed/shared_memory/__init__.py index 39ef60f847..8f84c89a69 100644 --- a/tilelang/distributed/shared_memory/__init__.py +++ b/tilelang/distributed/shared_memory/__init__.py @@ -120,13 +120,13 @@ def tensor_from_ptr( # bfloat16 / uint32 / uint64: create as matching-size int type, then view element_size = torch.empty((), dtype=dtype).element_size() if element_size == 2: - proxy_dtype = torch.int16 + # proxy_dtype = torch.int16 proxy_typestr = " &device_ids, uintptr_t buffer_ptrs_gpu_addr, - const std::vector> &all_gathered_handles) { + const std::vector> + &all_gathered_handles) { sync_vmm_handles(rank, device_ids, reinterpret_cast(buffer_ptrs_gpu_addr), all_gathered_handles); diff --git a/tilelang/distributed/shared_memory/csrc/ops.h b/tilelang/distributed/shared_memory/csrc/ops.h index b7fb5230f5..3f711a3c92 100644 --- a/tilelang/distributed/shared_memory/csrc/ops.h +++ b/tilelang/distributed/shared_memory/csrc/ops.h @@ -33,9 +33,10 @@ pybind11::bytearray create_vmm_handle(void *ptr); void *open_vmm_handle(const pybind11::bytearray &handle_bytes); void close_vmm_handle(void *ptr); bool supports_vmm_fabric(); -void sync_vmm_handles( - int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, - const std::vector> &all_gathered_handles); +void sync_vmm_handles(int rank, const std::vector &device_ids, + void **buffer_ptrs_gpu, + const std::vector> + &all_gathered_handles); // Multicast support detection bool supports_multicast(); diff --git a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp index e8059c81b1..162ba655e3 100644 --- a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp +++ b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp @@ -29,8 +29,8 @@ static void cu_mem_set_access_all(void *ptr, size_t size) { access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; } - CU_CHECK(cuMemSetAccess((CUdeviceptr)ptr, size, access_desc.data(), - device_count)); + CU_CHECK( + cuMemSetAccess((CUdeviceptr)ptr, size, access_desc.data(), device_count)); } static size_t align_to_granularity(size_t size_raw, size_t granularity) { @@ -54,7 +54,7 @@ void *vmm_malloc(size_t size_raw) { size_t granularity = 0; CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, - CU_MEM_ALLOC_GRANULARITY_MINIMUM)); + CU_MEM_ALLOC_GRANULARITY_MINIMUM)); size_t size = align_to_granularity(size_raw, granularity); @@ -92,7 +92,7 @@ py::bytearray create_vmm_handle(void *ptr) { CUmemFabricHandle fabric_handle; CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, - CU_MEM_HANDLE_TYPE_FABRIC, 0)); + CU_MEM_HANDLE_TYPE_FABRIC, 0)); // Serialize: 8 bytes size + sizeof(CUmemFabricHandle) std::string buf(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); @@ -114,7 +114,7 @@ void *open_vmm_handle(const py::bytearray &handle_bytes) { CUmemGenericAllocationHandle alloc_handle; CU_CHECK(cuMemImportFromShareableHandle(&alloc_handle, &fabric_handle, - CU_MEM_HANDLE_TYPE_FABRIC)); + CU_MEM_HANDLE_TYPE_FABRIC)); void *ptr = nullptr; CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, 0, 0, 0)); @@ -124,9 +124,7 @@ void *open_vmm_handle(const py::bytearray &handle_bytes) { return ptr; } -void close_vmm_handle(void *ptr) { - vmm_free(ptr); -} +void close_vmm_handle(void *ptr) { vmm_free(ptr); } // ---------- fabric support detection ---------- @@ -172,8 +170,8 @@ void sync_vmm_handles( } CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs.data(), - sizeof(void *) * buffer_ptrs.size(), - cudaMemcpyHostToDevice)); + sizeof(void *) * buffer_ptrs.size(), + cudaMemcpyHostToDevice)); CUDA_CHECK(cudaDeviceSynchronize()); } diff --git a/tilelang/distributed/utils.py b/tilelang/distributed/utils.py index 5eea2b3636..f1bbd4a0fe 100644 --- a/tilelang/distributed/utils.py +++ b/tilelang/distributed/utils.py @@ -25,9 +25,6 @@ _create_tensor, _create_ipc_handle, _sync_ipc_handles, - _supports_vmm_fabric, - _vmm_malloc, - _vmm_free, _create_vmm_handle, _sync_vmm_handles, create_host_device_tensor, @@ -122,6 +119,7 @@ def get_local_ipc_handle(data: torch.Tensor): def _resolve_use_vmm(use_vmm: bool | None) -> bool: """Resolve whether to use VMM based on env var and hardware support.""" import os + env_val = os.environ.get("TILESCALE_USE_VMM", None) if env_val is not None: return env_val == "1" diff --git a/tilelang/utils/allocator.py b/tilelang/utils/allocator.py index 5161c798b3..fa3fa240c1 100644 --- a/tilelang/utils/allocator.py +++ b/tilelang/utils/allocator.py @@ -9,12 +9,9 @@ tensor_from_ptr, _create_ipc_handle, _sync_ipc_handles, - _supports_vmm_fabric, _vmm_malloc, _vmm_free, _create_vmm_handle, - _open_vmm_handle, - _close_vmm_handle, _sync_vmm_handles, _supports_multicast, _mc_create, @@ -232,9 +229,7 @@ def _init_multicast_buffer(self): _mc_release_handle(mcast_handle) self._use_multicast = True - def _allocate_mcast_tensor( - self, shape: tuple[int, ...], dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor]: + def _allocate_mcast_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: """Allocate from multicast buffer (bump-pointer). Returns: @@ -282,10 +277,8 @@ def close(self): self._closed = True # Barrier before multicast teardown to ensure no rank is still using MC VA if getattr(self, "_use_multicast", False) and self._group is not None: - try: + with contextlib.suppress(Exception): dist.barrier(self._group) - except Exception: - pass self._free() def _free(self): @@ -348,7 +341,6 @@ def initialized(self) -> bool: def _allocate_tensor( self, shape: tuple[int, ...], dtype: torch.dtype, return_peers=False, take_ownership: bool = False ) -> torch.Tensor: - numel = _prod_shape(shape) itemsize = _element_size_bytes(dtype) bytes_needed = numel * itemsize @@ -442,6 +434,12 @@ def get_allocator( mcast_size: int | None = None, ) -> BaseAllocator: return BaseAllocator( - size, device=device, is_distributed=is_distributed, local_rank=local_rank, - num_local_ranks=num_local_ranks, group=group, use_vmm=use_vmm, mcast_size=mcast_size, + size, + device=device, + is_distributed=is_distributed, + local_rank=local_rank, + num_local_ranks=num_local_ranks, + group=group, + use_vmm=use_vmm, + mcast_size=mcast_size, ) From caef91f638656a008b45ce5685aeb4df155c45cc Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Wed, 6 May 2026 00:10:01 +0800 Subject: [PATCH 7/9] cleanup --- CMakeLists.txt | 64 ------ .../distributed/example_gemm_allreduce.py | 0 src/shared_memory/shared_memory.cc | 5 +- src/tl_templates/cuda/multimem.h | 113 ++++++++++ .../shared_memory/csrc/bindings.cpp | 103 --------- .../shared_memory/csrc/exception.h | 50 ----- .../shared_memory/csrc/ipc_ops.cpp | 98 --------- tilelang/distributed/shared_memory/csrc/ops.h | 42 ---- .../distributed/shared_memory/csrc/tensor.cpp | 113 ---------- .../shared_memory/csrc/vmm_ops.cpp | 201 ------------------ 10 files changed, 115 insertions(+), 674 deletions(-) delete mode 100644 examples/distributed/example_gemm_allreduce.py delete mode 100644 tilelang/distributed/shared_memory/csrc/bindings.cpp delete mode 100644 tilelang/distributed/shared_memory/csrc/exception.h delete mode 100644 tilelang/distributed/shared_memory/csrc/ipc_ops.cpp delete mode 100644 tilelang/distributed/shared_memory/csrc/ops.h delete mode 100644 tilelang/distributed/shared_memory/csrc/tensor.cpp delete mode 100644 tilelang/distributed/shared_memory/csrc/vmm_ops.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index cc8a66b434..8a75bf3c21 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -336,67 +336,3 @@ install( LIBRARY DESTINATION tilelang/lib ) -# Build tilescale_ext PyTorch C++ extension -if(USE_CUDA) - # Find Torch - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)" - OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE TORCH_CMAKE_RESULT - ) - if(TORCH_CMAKE_RESULT EQUAL 0 AND EXISTS "${TORCH_CMAKE_PREFIX_PATH}") - list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}") - endif() - - find_package(Torch QUIET) - if(Torch_FOUND) - message(STATUS "Building tilescale_ext with Torch ${Torch_VERSION}") - - set(SHARED_MEMORY_CSRC ${CMAKE_CURRENT_SOURCE_DIR}/tilelang/distributed/shared_memory/csrc) - set(TILESCALE_EXT_SOURCES - ${SHARED_MEMORY_CSRC}/bindings.cpp - ${SHARED_MEMORY_CSRC}/tensor.cpp - ${SHARED_MEMORY_CSRC}/ipc_ops.cpp - ${SHARED_MEMORY_CSRC}/vmm_ops.cpp - ) - - # Find libtorch_python.so - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch; import os; print(os.path.join(os.path.dirname(torch.__file__), 'lib', 'libtorch_python.so'))" - OUTPUT_VARIABLE TORCH_PYTHON_LIBRARY - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE TORCH_PYTHON_RESULT - ) - - python_add_library(tilescale_ext_C MODULE ${TILESCALE_EXT_SOURCES} WITH_SOABI) - target_compile_definitions(tilescale_ext_C PRIVATE TORCH_EXTENSION_NAME=_C) - target_include_directories(tilescale_ext_C PRIVATE - ${TORCH_INCLUDE_DIRS} - ${CUDAToolkit_INCLUDE_DIRS} - ${SHARED_MEMORY_CSRC} - ) - - if(TORCH_PYTHON_RESULT EQUAL 0 AND EXISTS "${TORCH_PYTHON_LIBRARY}") - message(STATUS "Found libtorch_python: ${TORCH_PYTHON_LIBRARY}") - target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} CUDA::cudart CUDA::cuda_driver) - else() - message(WARNING "libtorch_python.so not found, extension may have undefined symbols") - target_link_libraries(tilescale_ext_C PRIVATE ${TORCH_LIBRARIES} CUDA::cudart CUDA::cuda_driver) - endif() - - target_compile_options(tilescale_ext_C PRIVATE -fPIC) - set_target_properties(tilescale_ext_C PROPERTIES - OUTPUT_NAME "_C" - CXX_STANDARD 17 - LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib" - ) - - # Install as tilelang/distributed/shared_memory/_C.so - install(TARGETS tilescale_ext_C - LIBRARY DESTINATION tilelang/distributed/shared_memory - RUNTIME DESTINATION tilelang/distributed/shared_memory) - else() - message(WARNING "Torch not found, tilescale_ext will not be built") - endif() -endif() diff --git a/examples/distributed/example_gemm_allreduce.py b/examples/distributed/example_gemm_allreduce.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/src/shared_memory/shared_memory.cc b/src/shared_memory/shared_memory.cc index 363a070a88..e93c4b5856 100644 --- a/src/shared_memory/shared_memory.cc +++ b/src/shared_memory/shared_memory.cc @@ -2,9 +2,8 @@ * \file shared_memory/shared_memory.cc * \brief VMM/IPC/multicast shared memory ops registered via TVM FFI. * - * Replaces the old pybind11 tilescale_ext_C module. All functions are - * registered under the "tl.shared_memory.*" namespace and accessed from - * Python via tvm_ffi.get_global_func(). + * All functions are registered under the "tl.shared_memory.*" namespace + * and accessed from Python via tvm_ffi.get_global_func(). */ #include diff --git a/src/tl_templates/cuda/multimem.h b/src/tl_templates/cuda/multimem.h index 6471b2a84e..b5af5bf2b6 100644 --- a/src/tl_templates/cuda/multimem.h +++ b/src/tl_templates/cuda/multimem.h @@ -136,6 +136,119 @@ template <> struct RedV4 { } }; +// === V2 variants (64-bit = 2×f32, implemented as 2 scalar ops) === + +template struct LdReduceV2 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, + "tl::multimem::LdReduceV2: unsupported dtype/op"); + } +}; + +template <> struct LdReduceV2 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int2 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.add.v2.f32 {%0, %1}, [%2];" + : "=r"(ret.x), "=r"(ret.y) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +template <> struct LdReduceV2 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int2 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.min.v2.f32 {%0, %1}, [%2];" + : "=r"(ret.x), "=r"(ret.y) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +template <> struct LdReduceV2 { + TL_DEVICE static void run(void *dst, const void *mcast_ptr) { + int2 ret; + asm volatile( + "multimem.ld_reduce.relaxed.sys.global.max.v2.f32 {%0, %1}, [%2];" + : "=r"(ret.x), "=r"(ret.y) + : "l"(mcast_ptr) + : "memory"); + *reinterpret_cast(dst) = ret; + } +}; + +template struct StV2 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, "tl::multimem::StV2: unsupported dtype"); + } +}; + +template <> struct StV2 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + int2 val = *reinterpret_cast(src); + asm volatile("multimem.st.relaxed.sys.global.v2.b32 [%0], {%1, %2};" + : + : "l"(mcast_ptr), "r"(val.x), "r"(val.y) + : "memory"); + } +}; + +template struct RedV2 { + TL_DEVICE static void run(void *, const void *) { + static_assert(always_false_v, + "tl::multimem::RedV2: unsupported dtype/op"); + } +}; + +template <> struct RedV2 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + const float *src_f = reinterpret_cast(src); + const char *mc = reinterpret_cast(mcast_ptr); + #pragma unroll + for (int i = 0; i < 2; i++) { + unsigned val = __float_as_uint(src_f[i]); + asm volatile("multimem.red.relaxed.sys.global.add.f32 [%0], %1;" + : + : "l"(mc + i * 4), "r"(val) + : "memory"); + } + } +}; + +template <> struct RedV2 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + const float *src_f = reinterpret_cast(src); + const char *mc = reinterpret_cast(mcast_ptr); + #pragma unroll + for (int i = 0; i < 2; i++) { + unsigned val = __float_as_uint(src_f[i]); + asm volatile("multimem.red.relaxed.sys.global.min.f32 [%0], %1;" + : + : "l"(mc + i * 4), "r"(val) + : "memory"); + } + } +}; + +template <> struct RedV2 { + TL_DEVICE static void run(void *mcast_ptr, const void *src) { + const float *src_f = reinterpret_cast(src); + const char *mc = reinterpret_cast(mcast_ptr); + #pragma unroll + for (int i = 0; i < 2; i++) { + unsigned val = __float_as_uint(src_f[i]); + asm volatile("multimem.red.relaxed.sys.global.max.f32 [%0], %1;" + : + : "l"(mc + i * 4), "r"(val) + : "memory"); + } + } +}; + } // namespace multimem } // namespace tl diff --git a/tilelang/distributed/shared_memory/csrc/bindings.cpp b/tilelang/distributed/shared_memory/csrc/bindings.cpp deleted file mode 100644 index 09e383d68b..0000000000 --- a/tilelang/distributed/shared_memory/csrc/bindings.cpp +++ /dev/null @@ -1,103 +0,0 @@ -#include "ops.h" -#include -#include -#include - -namespace py = pybind11; - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "TileScale shared memory allocator (IPC + VMM)"; - - // Tensor utilities - m.def("tensor_from_ptr", &tensor_from_ptr, py::arg("ptr"), py::arg("shape"), - py::arg("dtype") = std::string("float32"), py::arg("device") = 0, - py::arg("take_ownership") = false); - - m.def( - "_create_tensor", - [](const std::vector &shape, const py::object &dtype) { - return create_tensor(shape, - torch::python::detail::py_object_to_dtype(dtype)); - }, - py::arg("shape"), py::arg("dtype")); - - m.def("create_host_device_tensor", &create_host_device_tensor, - "Create host/device shared pinned-mapped tensor (shape + dtype)"); - - // IPC API - m.def( - "_create_ipc_handle", - [](uintptr_t ptr_value) { - return create_ipc_handle(reinterpret_cast(ptr_value)); - }, - py::arg("ptr_value")); - - m.def( - "_sync_ipc_handles", - [](int rank, const std::vector &device_ids, - uintptr_t buffer_ptrs_gpu_addr, - const std::vector> &all_gathered_handles, - const std::optional &root_unique_id_opt) { - sync_ipc_handles(rank, device_ids, - reinterpret_cast(buffer_ptrs_gpu_addr), - all_gathered_handles, root_unique_id_opt); - }, - py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), - py::arg("all_gathered_handles"), py::arg("root_unique_id_opt")); - - // VMM API - m.def("_supports_vmm_fabric", &supports_vmm_fabric, - "Check if all GPUs support CUDA VMM fabric handles"); - - m.def( - "_vmm_malloc", - [](size_t size) -> uintptr_t { - return reinterpret_cast(vmm_malloc(size)); - }, - py::arg("size")); - - m.def( - "_vmm_free", - [](uintptr_t ptr_value) { - vmm_free(reinterpret_cast(ptr_value)); - }, - py::arg("ptr_value")); - - m.def( - "_create_vmm_handle", - [](uintptr_t ptr_value) { - return create_vmm_handle(reinterpret_cast(ptr_value)); - }, - py::arg("ptr_value")); - - m.def( - "_open_vmm_handle", - [](const py::bytearray &handle_bytes) -> uintptr_t { - return reinterpret_cast(open_vmm_handle(handle_bytes)); - }, - py::arg("handle_bytes")); - - m.def( - "_close_vmm_handle", - [](uintptr_t ptr_value) { - close_vmm_handle(reinterpret_cast(ptr_value)); - }, - py::arg("ptr_value")); - - m.def( - "_sync_vmm_handles", - [](int rank, const std::vector &device_ids, - uintptr_t buffer_ptrs_gpu_addr, - const std::vector> - &all_gathered_handles) { - sync_vmm_handles(rank, device_ids, - reinterpret_cast(buffer_ptrs_gpu_addr), - all_gathered_handles); - }, - py::arg("rank"), py::arg("device_ids"), py::arg("buffer_ptrs_gpu_addr"), - py::arg("all_gathered_handles")); - - // Multicast support detection - m.def("_supports_multicast", &supports_multicast, - "Check if all GPUs support CUDA multicast (NVSwitch)"); -} diff --git a/tilelang/distributed/shared_memory/csrc/exception.h b/tilelang/distributed/shared_memory/csrc/exception.h deleted file mode 100644 index 6107d10e2e..0000000000 --- a/tilelang/distributed/shared_memory/csrc/exception.h +++ /dev/null @@ -1,50 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -class TSException : public std::exception { - std::string message; - -public: - TSException(const char *name, const char *file, int line, - const std::string &error) { - message = std::string("Failed: ") + name + " error " + file + ":" + - std::to_string(line) + " '" + error + "'"; - } - const char *what() const noexcept override { return message.c_str(); } -}; - -#ifndef TS_HOST_ASSERT -#define TS_HOST_ASSERT(cond) \ - do { \ - if (!(cond)) { \ - throw TSException("Assertion", __FILE__, __LINE__, #cond); \ - } \ - } while (0) -#endif - -#ifndef CUDA_CHECK -#define CUDA_CHECK(cmd) \ - do { \ - cudaError_t e = (cmd); \ - if (e != cudaSuccess) { \ - throw TSException("CUDA", __FILE__, __LINE__, cudaGetErrorString(e)); \ - } \ - } while (0) -#endif - -#ifndef CU_CHECK -#define CU_CHECK(cmd) \ - do { \ - CUresult e = (cmd); \ - if (e != CUDA_SUCCESS) { \ - const char *error_str = NULL; \ - cuGetErrorString(e, &error_str); \ - throw TSException("CU", __FILE__, __LINE__, \ - error_str ? std::string(error_str) : "unknown"); \ - } \ - } while (0) -#endif diff --git a/tilelang/distributed/shared_memory/csrc/ipc_ops.cpp b/tilelang/distributed/shared_memory/csrc/ipc_ops.cpp deleted file mode 100644 index ba31a3fcbe..0000000000 --- a/tilelang/distributed/shared_memory/csrc/ipc_ops.cpp +++ /dev/null @@ -1,98 +0,0 @@ -#include -#include -#include -#include -#include - -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include - -#include "exception.h" -#include "ops.h" - -namespace py = pybind11; - -static size_t numel_of(const std::vector &shape) { - return std::accumulate(shape.begin(), shape.end(), size_t{1}, - [](size_t a, int64_t b) { - if (b < 0) - throw std::runtime_error("Negative dim"); - return a * (size_t)b; - }); -} - -static size_t dtype_nbytes(c10::ScalarType dtype) { - return (size_t)at::elementSize(dtype); -} - -torch::Tensor create_tensor(const std::vector &shape, - c10::ScalarType dtype) { - auto current_device = c10::cuda::current_device(); - auto options = - at::TensorOptions(at::kCUDA).dtype(dtype).device_index(current_device); - - const size_t bytes = dtype_nbytes(dtype) * numel_of(shape); - - CUDA_CHECK(cudaDeviceSynchronize()); - void *ptr = nullptr; - CUDA_CHECK(cudaMalloc(&ptr, bytes)); - - return at::from_blob( - ptr, shape, - [](void *p) { - cudaError_t cerr = cudaFree(p); - if (cerr != cudaSuccess) { - std::fprintf(stderr, "cudaFree failed in deleter: %s\n", - cudaGetErrorString(cerr)); - } - }, - options); -} - -py::bytearray create_ipc_handle(void *ptr) { - cudaIpcMemHandle_t handle{}; - CUDA_CHECK(cudaIpcGetMemHandle(&handle, ptr)); - return py::bytearray(reinterpret_cast(handle.reserved), - CUDA_IPC_HANDLE_SIZE); -} - -void sync_ipc_handles( - int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, - const std::vector> &all_gathered_handles, - const std::optional & /*root_unique_id_opt*/) { - - const int num = (int)device_ids.size(); - const int rdma_rank = 0; - - TS_HOST_ASSERT((size_t)num == all_gathered_handles.size()); - - std::vector ipc_handles(num); - std::vector buffer_ptrs(num, nullptr); - - for (int i = 0, offset = rdma_rank * num; i < num; ++i) { - TS_HOST_ASSERT(all_gathered_handles[offset + i].has_value()); - std::string s = (std::string)all_gathered_handles[offset + i].value(); - TS_HOST_ASSERT(s.size() == CUDA_IPC_HANDLE_SIZE); - if (offset + i != rank) { - std::memcpy(ipc_handles[i].reserved, s.data(), CUDA_IPC_HANDLE_SIZE); - CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], - cudaIpcMemLazyEnablePeerAccess)); - } - } - - CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs.data(), - sizeof(void *) * buffer_ptrs.size(), - cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaDeviceSynchronize()); -} diff --git a/tilelang/distributed/shared_memory/csrc/ops.h b/tilelang/distributed/shared_memory/csrc/ops.h deleted file mode 100644 index 3f711a3c92..0000000000 --- a/tilelang/distributed/shared_memory/csrc/ops.h +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once -#include -#include -#include -#include -#include -#include - -// Tensor utilities -torch::Tensor tensor_from_ptr(uint64_t ptr_val, std::vector shape, - const std::string &dtype = "float32", - int64_t device = 0, bool take_ownership = false); - -torch::Tensor create_tensor(const std::vector &shape, - c10::ScalarType dtype); - -std::pair -create_host_device_tensor(const std::vector &shape, - c10::ScalarType dtype); - -// IPC operations -pybind11::bytearray create_ipc_handle(void *ptr); - -void sync_ipc_handles( - int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, - const std::vector> &all_gathered_handles, - const std::optional &root_unique_id_opt); - -// VMM operations -void *vmm_malloc(size_t size_raw); -void vmm_free(void *ptr); -pybind11::bytearray create_vmm_handle(void *ptr); -void *open_vmm_handle(const pybind11::bytearray &handle_bytes); -void close_vmm_handle(void *ptr); -bool supports_vmm_fabric(); -void sync_vmm_handles(int rank, const std::vector &device_ids, - void **buffer_ptrs_gpu, - const std::vector> - &all_gathered_handles); - -// Multicast support detection -bool supports_multicast(); diff --git a/tilelang/distributed/shared_memory/csrc/tensor.cpp b/tilelang/distributed/shared_memory/csrc/tensor.cpp deleted file mode 100644 index 22613aa573..0000000000 --- a/tilelang/distributed/shared_memory/csrc/tensor.cpp +++ /dev/null @@ -1,113 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "exception.h" -#include "ops.h" - -static int64_t safe_mul_int64(int64_t a, int64_t b) { - if (a == 0 || b == 0) - return 0; - int64_t maxv = std::numeric_limits::max(); - if (a > maxv / b) - throw std::overflow_error("integer overflow in multiplication"); - return a * b; -} - -static at::ScalarType dtype_from_string(const std::string &s) { - if (s == "float32" || s == "float") - return at::kFloat; - if (s == "float16" || s == "half") - return at::kHalf; - if (s == "bfloat16" || s == "bfloat") - return at::kBFloat16; - if (s == "float64" || s == "double") - return at::kDouble; - if (s == "uint32") - return at::kUInt32; - if (s == "uint64") - return at::kUInt64; - if (s == "int32" || s == "int") - return at::kInt; - if (s == "int64" || s == "long" || s == "long int") - return at::kLong; - if (s == "uint8" || s == "byte") - return at::kByte; - if (s == "int8") - return at::kChar; - if (s == "bool") - return at::kBool; - throw std::runtime_error("Unsupported dtype string: '" + s + "'"); -} - -torch::Tensor tensor_from_ptr(uint64_t ptr_val, std::vector shape, - const std::string &dtype, int64_t device, - bool take_ownership) { - if (ptr_val == 0) - throw std::runtime_error("Received null pointer (0)."); - void *data_ptr = reinterpret_cast(static_cast(ptr_val)); - - at::ScalarType st = dtype_from_string(dtype); - auto options = torch::TensorOptions().dtype(st).device( - torch::kCUDA, static_cast(device)); - - int64_t nelems = 1; - for (auto d : shape) { - if (d < 0) - throw std::runtime_error("Negative dimension in shape"); - nelems = safe_mul_int64(nelems, d); - } - - std::function deleter; - if (take_ownership) { - uint64_t saved_ptr = ptr_val; - deleter = [saved_ptr](void *) { - void *p = reinterpret_cast(static_cast(saved_ptr)); - cudaError_t cerr = cudaFree(p); - if (cerr != cudaSuccess) { - std::fprintf(stderr, "tensor_from_ptr deleter cudaFree failed: %s\n", - cudaGetErrorString(cerr)); - } - }; - } else { - deleter = [](void *) {}; - } - - if (nelems == 0) { - return torch::empty(shape, options); - } else { - return at::from_blob(data_ptr, shape, deleter, options); - } -} - -std::pair -create_host_device_tensor(const std::vector &shape, - c10::ScalarType dtype) { - size_t elem_size = at::elementSize(dtype); - int64_t numel = 1; - for (int64_t s : shape) - numel *= s; - - size_t bytes = numel * elem_size; - - void *host_ptr = nullptr; - CUDA_CHECK(cudaHostAlloc(&host_ptr, bytes, cudaHostAllocMapped)); - - void *device_ptr = nullptr; - CUDA_CHECK(cudaHostGetDevicePointer(&device_ptr, host_ptr, 0)); - - auto host_tensor = torch::from_blob( - host_ptr, shape, torch::TensorOptions().dtype(dtype).device(torch::kCPU)); - - auto device_tensor = torch::from_blob( - device_ptr, shape, - torch::TensorOptions().dtype(dtype).device(torch::kCUDA)); - - return std::make_pair(host_tensor, device_tensor); -} diff --git a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp b/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp deleted file mode 100644 index 162ba655e3..0000000000 --- a/tilelang/distributed/shared_memory/csrc/vmm_ops.cpp +++ /dev/null @@ -1,201 +0,0 @@ -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include - -#include "exception.h" -#include "ops.h" - -namespace py = pybind11; - -// ---------- helpers ---------- - -static void cu_mem_set_access_all(void *ptr, size_t size) { - int device_count; - CUDA_CHECK(cudaGetDeviceCount(&device_count)); - - std::vector access_desc(device_count); - for (int idx = 0; idx < device_count; ++idx) { - access_desc[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access_desc[idx].location.id = idx; - access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - } - - CU_CHECK( - cuMemSetAccess((CUdeviceptr)ptr, size, access_desc.data(), device_count)); -} - -static size_t align_to_granularity(size_t size_raw, size_t granularity) { - size_t size = (size_raw + granularity - 1) & ~(granularity - 1); - if (size == 0) - size = granularity; - return size; -} - -// ---------- VMM malloc/free ---------- - -void *vmm_malloc(size_t size_raw) { - CUdevice device; - CU_CHECK(cuCtxGetDevice(&device)); - - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; - prop.location.id = device; - - size_t granularity = 0; - CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, - CU_MEM_ALLOC_GRANULARITY_MINIMUM)); - - size_t size = align_to_granularity(size_raw, granularity); - - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemCreate(&handle, size, &prop, 0)); - - void *ptr = nullptr; - CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, granularity, 0, 0)); - CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, handle, 0)); - cu_mem_set_access_all(ptr, size); - - return ptr; -} - -void vmm_free(void *ptr) { - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); - - size_t size = 0; - CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); - - CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size)); - CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size)); - CU_CHECK(cuMemRelease(handle)); -} - -// ---------- handle export/import ---------- - -py::bytearray create_vmm_handle(void *ptr) { - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr)); - - size_t size = 0; - CU_CHECK(cuMemGetAddressRange_v2(NULL, &size, (CUdeviceptr)ptr)); - - CUmemFabricHandle fabric_handle; - CU_CHECK(cuMemExportToShareableHandle(&fabric_handle, handle, - CU_MEM_HANDLE_TYPE_FABRIC, 0)); - - // Serialize: 8 bytes size + sizeof(CUmemFabricHandle) - std::string buf(sizeof(size_t) + sizeof(CUmemFabricHandle), '\0'); - std::memcpy(&buf[0], &size, sizeof(size_t)); - std::memcpy(&buf[sizeof(size_t)], &fabric_handle, sizeof(CUmemFabricHandle)); - return py::bytearray(buf.data(), buf.size()); -} - -void *open_vmm_handle(const py::bytearray &handle_bytes) { - std::string s = (std::string)handle_bytes; - TS_HOST_ASSERT(s.size() == sizeof(size_t) + sizeof(CUmemFabricHandle)); - - size_t size = 0; - std::memcpy(&size, s.data(), sizeof(size_t)); - - CUmemFabricHandle fabric_handle; - std::memcpy(&fabric_handle, s.data() + sizeof(size_t), - sizeof(CUmemFabricHandle)); - - CUmemGenericAllocationHandle alloc_handle; - CU_CHECK(cuMemImportFromShareableHandle(&alloc_handle, &fabric_handle, - CU_MEM_HANDLE_TYPE_FABRIC)); - - void *ptr = nullptr; - CU_CHECK(cuMemAddressReserve((CUdeviceptr *)&ptr, size, 0, 0, 0)); - CU_CHECK(cuMemMap((CUdeviceptr)ptr, size, 0, alloc_handle, 0)); - cu_mem_set_access_all(ptr, size); - - return ptr; -} - -void close_vmm_handle(void *ptr) { vmm_free(ptr); } - -// ---------- fabric support detection ---------- - -bool supports_vmm_fabric() { - int device_count = 0; - cudaError_t err = cudaGetDeviceCount(&device_count); - if (err != cudaSuccess || device_count == 0) - return false; - - int driver_version = 0; - CUresult cu_err = cuDriverGetVersion(&driver_version); - if (cu_err != CUDA_SUCCESS || driver_version < 12040) - return false; - - for (int i = 0; i < device_count; ++i) { - CUdevice dev; - CU_CHECK(cuDeviceGet(&dev, i)); - int supported = 0; - CU_CHECK(cuDeviceGetAttribute( - &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev)); - if (!supported) - return false; - } - return true; -} - -// ---------- unified sync entry ---------- - -void sync_vmm_handles( - int rank, const std::vector &device_ids, void **buffer_ptrs_gpu, - const std::vector> &all_gathered_handles) { - - const int num = (int)device_ids.size(); - TS_HOST_ASSERT((size_t)num == all_gathered_handles.size()); - - std::vector buffer_ptrs(num, nullptr); - - for (int i = 0; i < num; ++i) { - TS_HOST_ASSERT(all_gathered_handles[i].has_value()); - if (i != rank) { - buffer_ptrs[i] = open_vmm_handle(all_gathered_handles[i].value()); - } - } - - CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs.data(), - sizeof(void *) * buffer_ptrs.size(), - cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaDeviceSynchronize()); -} - -// ---------- Multicast support detection ---------- - -bool supports_multicast() { - int device_count = 0; - cudaError_t err = cudaGetDeviceCount(&device_count); - if (err != cudaSuccess || device_count == 0) - return false; - - int driver_version = 0; - CUresult cu_err = cuDriverGetVersion(&driver_version); - if (cu_err != CUDA_SUCCESS || driver_version < 12040) - return false; - - for (int i = 0; i < device_count; ++i) { - CUdevice dev; - CU_CHECK(cuDeviceGet(&dev, i)); - int supported = 0; - CU_CHECK(cuDeviceGetAttribute( - &supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, dev)); - if (!supported) - return false; - } - return true; -} From 7ae6731c92777a8d662054939d0a1b794844bbd8 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Wed, 6 May 2026 23:59:55 +0800 Subject: [PATCH 8/9] lint --- docs/merge_upstream_tilelang.md | 134 ++++++++++ .../example_allgather_gemm_overlapped.py | 9 +- .../example_allgather_gemm_specialized.py | 246 ++++++++++++++++++ .../distributed/example_gemm_allreduce.py | 167 ++++++++++++ .../distributed/example_multimem_allreduce.py | 4 +- src/op/multimem.cc | 103 ++++++++ src/op/multimem.h | 9 +- src/tl_templates/cuda/multimem.h | 111 ++++++++ src/transform/inject_fence_proxy.cc | 25 +- src/transform/warp_specialized_rewriter.cc | 150 +++++++++-- tilelang/language/distributed/__init__.py | 4 + tilelang/language/distributed/multimem.py | 34 ++- 12 files changed, 965 insertions(+), 31 deletions(-) create mode 100644 examples/distributed/example_allgather_gemm_specialized.py create mode 100644 examples/distributed/example_gemm_allreduce.py diff --git a/docs/merge_upstream_tilelang.md b/docs/merge_upstream_tilelang.md index 48a1ec9a65..c702585027 100644 --- a/docs/merge_upstream_tilelang.md +++ b/docs/merge_upstream_tilelang.md @@ -344,6 +344,135 @@ PR [#50](https://github.com/tile-ai/tilescale/pull/50) ("Sync mainstream TileLan --- +## 10. Practical Lessons from PR #58 + +This section captures the hard-won lessons from the `0.1.7.post1 → 0.1.9` sync (~50 upstream commits, ~80K LOC diff). Unlike PR #50 which cherry-picked individual commits, PR #58 merged the entire upstream delta in one operation — a much more aggressive approach that revealed systematic failure modes. + +### 10.1 `src/transform/` and `src/tl_templates/` Are NOT All TileScale-Exclusive + +Section 2.3 classifies these directories as "TileScale-specific pass infrastructure" and says "never overwrite". This is **misleading for bulk operations**. In practice: + +- **Most files in `src/transform/` exist in both repos** (e.g., `layout_inference.cc`, `loop_partition.cc`, `lower_tile_op.cc`, `inject_pipeline.cc`). They originated from upstream and were modified by both sides. +- **Only a small subset are truly TileScale-only**: `lower_cpengine_intrin.cc`, `storage_access.cc/h`, `wgmma_sync_rewriter.cc`, `align_dynamic_shared_memory_allocations.cc`, `inject_ptx_async_copy.cc`, `inject_fence_proxy.cc`. +- **Upstream also adds new transforms** (e.g., `producer_consumer_ws.cc`, `unroll_loop.cc`, `verify_parallel_loop.cc`, `fuse_mbarrier_arrive_expect_tx.cc`) that TileScale should absorb. +- **The rule**: For each file, check whether it exists in the upstream commit (`git cat-file -e :`). If it does, `--theirs` (take upstream) is the safe default. Only keep `--ours` for files that are genuinely TileScale-only. + +### 10.2 `git merge` vs `git cherry-pick` + +- **`git merge` (one-shot)**: Fast but produces ~400 conflicted files. Resolution must be done programmatically (batch `--ours`/`--theirs`). The batch resolution can silently corrupt files that need manual adaptation. +- **`git cherry-pick` (per-commit)**: Safer, more auditable, but slow for 50+ commits. +- **For >20 commits**: Consider `git merge --no-commit`, then resolve conflicts with the file-by-file decision table below, then `git commit`. + +### 10.3 Mandatory Build-Import-Run Loop + +After conflict resolution, the merge is **never** clean on the first try. Follow this loop until both `import tilelang` and a distributed example pass: + +```bash +ninja -C build 2>&1 | grep error: # fix C++ build errors +python -c "import tilelang" # fix Python import errors +python examples/distributed/example_xxx.py # fix runtime errors +``` + +Common failure categories and their symptoms: + +| Symptom | Root Cause | Fix | +|---------|-----------|-----| +| `undefined symbol: _ZN3tvm2tl31ApplyMultiVersionBufferRewriterE...` | Stale TileScale `.cc` kept as `--ours`; upstream added function to this file | `git checkout -- ` | +| `no matching function for call to 'VectorizeLoop(..., LayoutMap&)'` | Upstream removed/renamed an overload | Check upstream `loop_vectorize.h` for new signatures; adjust callers | +| `'create_list_of_mbarrier' was not declared` / `'get_mbarrier' was not declared` | TileScale ops registered in old `builtin.cc`; removed in upstream | Add them back to `builtin.cc` and `builtin.h` | +| `error: 'LoopPragmaUnroll' was not declared` | Upstream renamed to `PragmaUnrollLoop` | Bulk rename | +| `error: 'atomicadd_elem_op' was not declared; did you mean 'atomic_add_elem_op'?` | Upstream added underscore | Bulk rename | +| `Module has no function '__tilescale_init_table'` | Upstream `rt_mod_cuda.cc` uses `CUDAModuleCreate` instead of `TileScaleCUDAModuleCreate` | Restore `TileScaleCUDAModuleCreate` calls + include in `rt_mod_cuda.cc` | +| `'JITKernel' object has no attribute 'initialize'` | Upstream `jit/kernel.py` doesn't have TileScale's `initialize()` | Add back `initialize()` method + `allocator` attribute | +| `TVMFFIKernelAdapter has no attribute 'init_table'` | Upstream adapter doesn't have `init_table()` | Add back `init_table()` to `tilelang/jit/adapter/tvm_ffi.py` | +| `'lazy_jit' not found in tilelang.jit` | Upstream `jit/__init__.py` doesn't export `lazy_jit` | Remove from `__init__.py` import, or add back implementation | + +### 10.4 The Distributed Codegen Must Be Surgically Preserved + +TileScale adds significant infrastructure to CUDA codegen that upstream knows nothing about. After merging upstream codegen files, the following **must** be present: + +**In `src/target/codegen_cuda.h`**: +```cpp +static inline bool use_distributed() { + const char *env = std::getenv("TILELANG_USE_DISTRIBUTED"); + if (env) return std::string(env) == "1"; + return false; +} +// Inside class CodeGenTileLangCUDA: +bool use_distributed_{use_distributed()}; +bool need_multimem_h_{false}; +``` + +**In `src/target/codegen_cuda.cc`**: +```cpp +#include "../op/distributed.h" +#include "../op/sync.h" + +// Inside Finish(): +if (use_distributed_) { + decl_stream << "#include \n"; + decl_stream << "#include \n"; + decl_stream << "#include \n"; + decl_stream << "extern \"C\" __constant__ uint64_t meta_data[1024];\n"; +} +if (need_multimem_h_) { + decl_stream << "#include \n"; +} +``` + +**In `src/target/rt_mod_cuda.cc`**: Replace upstream `CUDAModuleCreate` with `TileScaleCUDAModuleCreate` and add `#include "../runtime/tilescale_cuda_module.h"`. + +### 10.5 TileScale-Specific Python Utilities That Pass Silently + +These files exist only in TileScale and were not overwritten by the merge, but their callers in shared modules may have changed: + +| File | TileScale Purpose | What Can Break | +|------|-------------------|----------------| +| `tilelang/utils/allocator.py` | `BaseAllocator`, `get_allocator()` | `torch.set_default_device` conflicts with `all_gather_object`; `parse_device` must be correct | +| `tilelang/utils/tensor.py` (line `tensor()` function) | `tilelang.tensor(...)` factory | Lost `tensor()` function if upstream file overwrites it | +| `tilelang/utils/target.py` (line `parse_device()`) | device string parsing for allocator | `parse_device("cuda")` returning hardcoded 0 instead of `current_device()` | +| `tilelang/distributed/utils.py` | `init_dist()`, `perf_fn()` | `torch.set_default_device("cuda")` before `init_process_group` causes NCCL device mismatch | + +### 10.6 The Device Mismatch Trap + +When `init_dist()` calls `torch.set_default_device("cuda")` (without device index), all PyTorch tensors default to `cuda:0`. With newer PyTorch (2.2+) passing `device_id` to `init_process_group`, NCCL enforces that collective tensors match the process group's device. This causes: + +``` +Torch.distributed.DistBackendError: Tensor found on device cuda:0 but backend constrained to cuda:1 +``` + +**Fix**: Call `torch.cuda.set_device(local_rank)` BEFORE `init_process_group`, and use explicit device strings. Also ensure `parse_device("cuda")` returns `torch.cuda.current_device()` rather than hardcoded `0`. + +### 10.7 Duplicate Op Registration Detection + +After a large merge, upstream may have added op registrations that TileScale's files also register. Check with: + +```bash +python -c "import tilelang" 2>&1 | grep "already registered" +``` + +If you see `Global Function 'tl.X' is already registered`, search for duplicate `refl::GlobalDef().def("tl.X", ...)` registrations and remove the TileScale copy (keep the upstream one). + +### 10.8 After Merge: Restore Truly TileScale-Only Files from Old Main + +After batch resolution, verify these files match the pre-sync TileScale version: + +| Category | Key Files | +|----------|-----------| +| Distributed C++ ops | `src/op/distributed.cc/h`, `src/op/remote_copy.cc/h`, `src/op/sync.cc/h`, `src/op/multimem.cc/h`, `src/op/multimem_rewriter.h`, `src/op/gemm_py.cc/h` | +| Distributed runtime | `src/runtime/tilescale_cuda_module.cc/h`, `src/shared_memory/shared_memory.cc` | +| Distributed templates | `src/tl_templates/cuda/distributed.h`, `sync.h`, `ldst.h`, `multimem.h` | +| TileScale transforms | `lower_cpengine_intrin.cc`, `storage_access.cc/h`, `wgmma_sync_rewriter.cc`, `align_dynamic_shared_memory_allocations.cc`, `inject_ptx_async_copy.cc`, `inject_fence_proxy.cc` | +| Python distributed | `tilelang/distributed/**`, `tilelang/language/distributed/**`, `tilelang/utils/allocator.py` | +| Build config | `src/backend/cuda/CMakeLists.txt` (must include `tilescale_cuda_module.cc` and `shared_memory/shared_memory.cc`) | + +```bash +# Restore a known-good TileScale file +git show main: > +``` + +--- + ## 9. Checklist for Each Sync PR Before opening the PR: @@ -353,8 +482,13 @@ Before opening the PR: - [ ] `CMakeLists.txt` conflict resolved; `tilescale_ext` target intact - [ ] `tilelang/__init__.py` still exports distributed namespace - [ ] Full build passes +- [ ] `import tilelang` succeeds with no import errors +- [ ] `tilelang.distributed` imports successfully - [ ] Shared `testing/python/` tests pass - [ ] At least one distributed example runs end-to-end +- [ ] `TileScaleCUDAModuleCreate` used in `rt_mod_cuda.cc` (not `CUDAModuleCreate`) +- [ ] Distributed template includes present in `codegen_cuda.cc` (`sync.h`, `ldst.h`, `distributed.h`, `multimem.h`, `meta_data`) +- [ ] No duplicate TVM FFI registrations (`python -c "import tilelang"` clean) - [ ] API-breaking upstream changes reflected in TileScale distributed layer if applicable - [ ] PR title follows: `[Sync] Merge upstream TileLang ` - [ ] PR description lists: last-synced upstream SHA, new upstream SHA, major features included, any skipped items with justification diff --git a/examples/distributed/example_allgather_gemm_overlapped.py b/examples/distributed/example_allgather_gemm_overlapped.py index 4fc9665476..1c5afe96c2 100644 --- a/examples/distributed/example_allgather_gemm_overlapped.py +++ b/examples/distributed/example_allgather_gemm_overlapped.py @@ -195,6 +195,7 @@ def ag_gemm_op( gemm_stream.wait_stream(ag_stream) current_stream = torch.cuda.current_stream() current_stream.wait_stream(gemm_stream) + dist.barrier() return C @@ -306,10 +307,10 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--num-processes", type=int, default=2, help="Number of processes to spawn (default: 2)") - parser.add_argument("--M", type=int, default=8192, help="M dimension") - parser.add_argument("--N", type=int, default=28672, help="N dimension") - parser.add_argument("--K", type=int, default=8192, help="K dimension") + parser.add_argument("--num-processes", type=int, default=8, help="Number of processes to spawn (default: 2)") + parser.add_argument("--M", type=int, default=32768, help="M dimension") + parser.add_argument("--N", type=int, default=16384, help="N dimension") + parser.add_argument("--K", type=int, default=2048, help="K dimension") parser.add_argument("--persistent", action="store_true", help="Use persistent kernel") args = parser.parse_args() num_processes = args.num_processes diff --git a/examples/distributed/example_allgather_gemm_specialized.py b/examples/distributed/example_allgather_gemm_specialized.py new file mode 100644 index 0000000000..9a063d9d42 --- /dev/null +++ b/examples/distributed/example_allgather_gemm_specialized.py @@ -0,0 +1,246 @@ +import os +import argparse + +import torch +import torch.distributed as dist +import torch.multiprocessing + +import tilelang +import tilelang.language as T +from tilelang.carver.arch import driver +from tilelang.distributed import init_dist +from tilelang.distributed import perf_fn +from tilelang.utils.allocator import get_allocator + +tilelang.disable_cache() +os.environ["NCCL_DEBUG"] = "WARN" + + +@tilelang.jit +def ag_gemm_sm_specialized_kernel( + M, + N, + K, + num_ranks, + num_comm_sms: int, + block_M: int, + block_N: int, + block_K: int, + threads: int, + dtype: str = "float16", +): + sm_num = driver.get_num_sms() + num_comp_sms = sm_num - num_comm_sms + M_per_rank = M // num_ranks + N_per_rank = N // num_ranks + m_blocks = T.ceildiv(M, block_M) + n_blocks = T.ceildiv(N_per_rank, block_N) + local_m_blocks = T.ceildiv(M_per_rank, block_M) + k_blocks = T.ceildiv(K, block_K) + total_tiles = m_blocks * n_blocks + waves = T.ceildiv(total_tiles, num_comp_sms) + GROUP_SIZE_M = 8 + accum_dtype = "float" + + @T.prim_func + def main( + A_local: T.Tensor((M_per_rank, K), dtype), + B: T.Tensor((K, N_per_rank), dtype), + mcast_A: T.Tensor((M, K), dtype), + gathered_A: T.Tensor((M, K), dtype), + mcast_signal: T.Tensor((m_blocks,), "uint32"), + local_signal: T.Tensor((m_blocks,), "uint32"), + grid_barrier: T.Tensor((num_ranks,), "int32"), + C: T.Tensor((M, N_per_rank), dtype), + local_rank: T.int32 + ): + with T.Kernel(sm_num, threads=threads) as bid: + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_shared = T.alloc_shared((block_M, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + tid = T.get_thread_binding(0) + + if bid == 0: + for i in T.serial(T.ceildiv(m_blocks, threads)): + signal_idx = i * threads + tid + if signal_idx < m_blocks: + local_signal[signal_idx] = 0 + T.fence_sys() + T.barrier_blocks(grid_barrier) + + if bid < num_comp_sms: + for w in T.serial(waves): + tile_id = bid + w * num_comp_sms + if tile_id < total_tiles: + num_pid_in_group = GROUP_SIZE_M * n_blocks + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = T.min(m_blocks - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + if tid == 0: + T.wait_eq(local_signal[pid_m], 1) + + T.clear(C_local) + for k in T.Pipelined(k_blocks, num_stages=3): + T.copy(gathered_A[pid_m * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, pid_n * block_N], B_shared) + T.gemm(A_shared, B_shared, C_local) + T.copy(C_local, C_shared) + T.copy(C_shared, C[pid_m * block_M, pid_n * block_N]) + else: + loaded = T.alloc_barrier([256]) + parity = 0 + comm_sm_id = bid - num_comp_sms + for local_m in T.serial(T.ceildiv(local_m_blocks, num_comm_sms)): + local_pid_m = comm_sm_id + local_m * num_comm_sms + if local_pid_m < local_m_blocks: + global_pid_m = local_rank * local_m_blocks + local_pid_m + for k in T.serial(k_blocks): + T.tma_load(A_local[local_pid_m * block_M, k * block_K], A_shared) + T.mbarrier_arrive(loaded) + T.mbarrier_wait_parity(loaded, parity) + parity = (parity + 1) % 2 + T.copy( + A_shared, + mcast_A[global_pid_m * block_M, k * block_K], + ) # TODO(wt): Change to canonical mcast tma store later + + T.fence_sys() + if tid == 0: + T.multimem_signal(mcast_signal[global_pid_m], 1) + + return main + + +def ag_gemm_op( + A, + B, + mcast_A, + gathered_A, + mcast_signal, + local_signal, + grid_barrier, + C, + kernel, + local_rank, +): + kernel(A, B, mcast_A, gathered_A, mcast_signal, local_signal, grid_barrier, C, local_rank) + return C + + +def torch_ag_gemm(group: torch.distributed.ProcessGroup, A: torch.Tensor, B: torch.Tensor, ag_out: torch.Tensor): + torch.distributed.all_gather_into_tensor(ag_out, A, group) + return torch.matmul(ag_out, B) + + +def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + dtype = torch.float16 + M, N, K = args.M, args.N, args.K + block_M, block_N, block_K = args.block_m, args.block_n, args.block_k + threads = args.threads + num_comm_sms = args.num_comm_sms + + assert M % num_local_ranks == 0, "M must be divisible by num-processes" + assert N % num_local_ranks == 0, "N must be divisible by num-processes" + assert (M // num_local_ranks) % block_M == 0, "M_per_rank must be divisible by block_m" + assert (N // num_local_ranks) % block_N == 0, "N_per_rank must be divisible by block_n" + assert K % block_K == 0, "K must be divisible by block_k" + assert 0 < num_comm_sms < driver.get_num_sms(), "num_comm_sms must leave at least one compute SM" + + M_per_rank = M // num_local_ranks + N_per_rank = N // num_local_ranks + m_blocks = M // block_M + + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + assert rank == local_rank and num_ranks == num_local_ranks, "only support single-node launch for now" + + dtype_bytes = torch.tensor([], dtype=dtype).element_size() + signal_bytes = torch.tensor([], dtype=torch.uint32).element_size() + # _allocate_mcast_tensor uses aligned bump allocation; keep room for padding + # between the gathered A buffer and the signal buffer. + mcast_bytes = M * K * dtype_bytes + m_blocks * signal_bytes + 4096 + allocator = get_allocator( + size=2**30, + device=f"cuda:{local_rank}", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=num_local_ranks, + group=group, + use_vmm=True, + mcast_size=mcast_bytes, + ) + + kernel = ag_gemm_sm_specialized_kernel( + M, + N, + K, + num_local_ranks, + num_comm_sms, + block_M, + block_N, + block_K, + threads, + ) + kernel.initialize(allocator=allocator) + if local_rank == 0 and args.print_source: + print(kernel.get_kernel_source()) + + torch.manual_seed(42 + local_rank) + A = tilelang.tensor((M_per_rank, K), dtype, allocator=allocator).normal_() + B = tilelang.tensor((K, N_per_rank), dtype, allocator=allocator).normal_() + C = tilelang.tensor((M, N_per_rank), dtype, allocator=allocator) + grid_barrier = tilelang.tensor((num_local_ranks,), torch.int32, allocator=allocator).zero_() + + mcast_A_flat, gathered_A_flat = allocator._allocate_mcast_tensor((M * K,), dtype) + mcast_signal, local_signal = allocator._allocate_mcast_tensor((m_blocks,), torch.uint32) + mcast_A = mcast_A_flat.view(M, K) + gathered_A = gathered_A_flat.view(M, K) + + dist.barrier(group) + tilelang_C = ag_gemm_op(A, B, mcast_A, gathered_A, mcast_signal, local_signal, grid_barrier, C, kernel, local_rank) + torch.cuda.synchronize() + dist.barrier(group) + + torch_ag_buffer = torch.empty((M, K), dtype=dtype, device=f"cuda:{local_rank}") + torch_C = torch_ag_gemm(group, A, B, torch_ag_buffer) + + if torch.allclose(torch_C, tilelang_C, atol=1e-2, rtol=1e-2): + print(f"rank {local_rank} check passed.") + else: + max_diff = (torch_C - tilelang_C).abs().max().item() + print(f"rank {local_rank} check failed. max_diff={max_diff}") + + tl_t = perf_fn( + lambda: ag_gemm_op(A, B, mcast_A, gathered_A, mcast_signal, local_signal, grid_barrier, C, kernel, local_rank), + warmup=args.warmup, + rep=args.rep, + ) + print( + f"rank {local_rank} tilelang specialized ag_gemm time: {tl_t:.2f} ms, " + f"TFLOPS: {2 * M * N * K / 1e9 / tl_t / num_local_ranks:.2f}" + ) + + allocator.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + parser.add_argument("--M", type=int, default=32768) + parser.add_argument("--N", type=int, default=16384) + parser.add_argument("--K", type=int, default=2048) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-n", type=int, default=256) + parser.add_argument("--block-k", type=int, default=64) + parser.add_argument("--threads", type=int, default=256) + parser.add_argument("--num-comm-sms", type=int, default=8) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--rep", type=int, default=10) + parser.add_argument("--print-source", action="store_true") + args = parser.parse_args() + + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) \ No newline at end of file diff --git a/examples/distributed/example_gemm_allreduce.py b/examples/distributed/example_gemm_allreduce.py new file mode 100644 index 0000000000..fa13500be0 --- /dev/null +++ b/examples/distributed/example_gemm_allreduce.py @@ -0,0 +1,167 @@ +""" +Fused GEMM + AllReduce using multimem.red — single kernel. + +Each rank computes a partial GEMM (C_partial = A_local @ B), then uses +multimem.red to reduce partial results into the multicast buffer across +all ranks. A final multimem.ld_reduce reads back the fully reduced result. + +Multi-process multi-GPU: each process manages one GPU, multicast handle +shared via fabric handles through torch.distributed. + +Usage: + export TILESCALE_USE_VMM=1 + export NCCL_IB_DISABLE=1 + export TILELANG_USE_DISTRIBUTED=1 + python examples/distributed/example_gemm_allreduce.py [--num-processes 8] + +Requirements: + - NVSwitch with multicast support (H100/B200 DGX) +""" + +import os +import argparse + +import torch +import torch.distributed as dist +import torch.multiprocessing + +import tilelang +import tilelang.language as T +from tilelang.distributed import init_dist +from tilelang.utils.allocator import get_allocator + +os.environ["NCCL_DEBUG"] = "WARN" + + +def gemm_allreduce_kernel(M, N, K, block_M, block_N, block_K, threads, dtype="float16"): + """Fused GEMM + multimem allreduce in a single kernel. + + Each rank computes C_partial[M,N] = A_local[M,K] @ B[K,N], + then reduces into mcast_buf via multimem.red (accumulate without read-back). + After barrier, a second kernel (or ld_reduce) retrieves the final result. + + For simplicity, this example does: + 1. GEMM into fragment + 2. multimem.red to push partial sum into mcast buffer + A separate ld_reduce kernel reads the final reduced result. + """ + accum_dtype = "float32" + + @T.prim_func + def gemm_red( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + mcast_C: T.Tensor((M, N), "float32"), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + T.clear(C_local) + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=2): + T.copy(A[by * block_M:(by + 1) * block_M, k * block_K:(k + 1) * block_K], A_shared) + T.copy(B[k * block_K:(k + 1) * block_K, bx * block_N:(bx + 1) * block_N], B_shared) + T.gemm(A_shared, B_shared, C_local) + + # Reduce partial GEMM result into multicast buffer + T.multimem_red( + C_local, + mcast_C[by * block_M:(by + 1) * block_M, bx * block_N:(bx + 1) * block_N], + reduce_op=T.MultimemReduceOp.ADD, + ) + T.fence_sys() + + return gemm_red + + +def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + M, N, K = args.M, args.N, args.K + block_M, block_N, block_K = args.block_m, args.block_n, args.block_k + threads = args.threads + dtype = torch.float16 + + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + + # Multicast buffer for C (M x N, float32) + mcast_bytes = M * N * 4 + allocator = get_allocator( + size=mcast_bytes, + device=f"cuda:{local_rank}", + is_distributed=True, + local_rank=local_rank, + num_local_ranks=num_local_ranks, + group=group, + mcast_size=mcast_bytes, + ) + + # Compile kernels + gemm_red_func = gemm_allreduce_kernel(M, N, K, block_M, block_N, block_K, threads) + gemm_red_kernel = tilelang.compile( + gemm_red_func, + pass_configs={"tl.disable_tma_lower": True}, + ) + + if local_rank == 0 and args.print_source: + print("=== GEMM + Red kernel ===") + print(gemm_red_kernel.get_kernel_source()) + + # Per-rank input + torch.manual_seed(42 + rank) + A_local = torch.randn(M, K, dtype=dtype, device=f"cuda:{local_rank}") + B = torch.randn(K, N, dtype=dtype, device=f"cuda:{local_rank}") + + # Each rank computes local GEMM as reference (same precision as kernel) + C_local_ref = (A_local @ B).float() + + # NCCL all_reduce as ground truth for the reduction + expected = C_local_ref.clone() + dist.all_reduce(expected, op=dist.ReduceOp.SUM, group=group) + + # Multicast buffer for partial sum accumulation + mcast_C, local_C = allocator._allocate_mcast_tensor((M * N,), torch.float32) + local_C.zero_() + torch.cuda.synchronize() + dist.barrier(group) + + # Run fused GEMM + multimem.red + gemm_red_kernel(A_local, B, mcast_C.view(M, N)) + torch.cuda.synchronize() + dist.barrier(group) + + # Read back reduced result from local physical memory + result = local_C.view(M, N) + torch.cuda.synchronize() + + # Compare: only validates multimem.red correctness (GEMM precision identical) + atol = 1e-5 + max_diff = (result - expected).abs().max().item() + passed = max_diff < atol + + if local_rank == 0: + print(f"M={M}, N={N}, K={K}, num_ranks={num_ranks}, max_diff={max_diff:.6f}") + if passed: + print(f"[rank {local_rank}] PASSED") + else: + print(f"[rank {local_rank}] FAILED (max_diff={max_diff:.6f})") + + allocator.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-processes", type=int, default=8) + parser.add_argument("--M", type=int, default=8192) + parser.add_argument("--N", type=int, default=8192) + parser.add_argument("--K", type=int, default=8192) + parser.add_argument("--block_m", type=int, default=128) + parser.add_argument("--block_n", type=int, default=256) + parser.add_argument("--block_k", type=int, default=64) + parser.add_argument("--threads", type=int, default=256) + parser.add_argument("--print_source", action="store_true") + args = parser.parse_args() + + torch.multiprocessing.spawn( + main, args=(args.num_processes, args), nprocs=args.num_processes, join=True + ) diff --git a/examples/distributed/example_multimem_allreduce.py b/examples/distributed/example_multimem_allreduce.py index 9247d9e040..a5f3cb9119 100644 --- a/examples/distributed/example_multimem_allreduce.py +++ b/examples/distributed/example_multimem_allreduce.py @@ -30,7 +30,7 @@ os.environ["NCCL_DEBUG"] = "WARN" -def multimem_allreduce_kernel(N, block_N, threads): +def multimem_allreduce_kernel_one_shot(N, block_N, threads): @T.prim_func def main( mcast_buf: T.Tensor((N,), "float32"), @@ -68,7 +68,7 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): # Compile kernel kernel = tilelang.compile( - multimem_allreduce_kernel(N, BLOCK_N, threads), + multimem_allreduce_kernel_one_shot(N, BLOCK_N, threads), pass_configs={"tl.disable_tma_lower": True}, ) if local_rank == 0 and args.print_source: diff --git a/src/op/multimem.cc b/src/op/multimem.cc index 8880150d26..866a54dea4 100644 --- a/src/op/multimem.cc +++ b/src/op/multimem.cc @@ -82,6 +82,15 @@ MultimemOp::MultimemOp(Array args, << "multimem_red: dst must be global (multicast) buffer, got '" << dst_scope << "' for buffer '" << node->dst->name << "'"; break; + case MultimemMode::kTmaStore: + case MultimemMode::kTmaRedStore: + ICHECK(src_scope == "shared" || src_scope == "shared.dyn") + << "multimem_tma_store: src must be shared memory, got '" << src_scope + << "' for buffer '" << node->src->name << "'"; + ICHECK(dst_scope == "global") + << "multimem_tma_store: dst must be global (multicast) buffer, got '" + << dst_scope << "' for buffer '" << node->dst->name << "'"; + break; } data_ = std::move(node); @@ -227,6 +236,10 @@ LayoutMap MultimemOpNode::InferLayout(const LayoutInferArgs &T, // MultimemRewriter Stmt MultimemOpNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const { + if (mode == MultimemMode::kTmaStore || mode == MultimemMode::kTmaRedStore) { + return LowerBulkCopy(T, analyzer); + } + // Step 1-2: Create SIMT loop and fuse/transform auto simt_loop = MakeSIMTLoop(analyzer); auto fused_loop = Downcast(ParallelLoopFuser::Fuse(simt_loop)); @@ -265,6 +278,96 @@ Stmt MultimemOpNode::Lower(const LowerArgs &T, return result; } +// === LowerBulkCopy === +// CTA-collective bulk async store from shared to multicast global. +// Reuses the 1D address computation pattern from CopyNode::LowerBulkCopy1D, +// but emits multimem.cp.async.bulk or multimem.cp.reduce.async.bulk PTX. +Stmt MultimemOpNode::LowerBulkCopy(const LowerArgs &T, + arith::Analyzer *analyzer) const { + bool is_reduce = (mode == MultimemMode::kTmaRedStore); + // Both modes: src=shared, dst=mcast_global + auto &shared_tensor = src; + auto &global_tensor = dst; + auto &shared_range = src_range; + auto &global_range = dst_range; + + // Compute total elements + PrimExpr shared_elements = 1; + for (size_t i = 0; i < shared_range.size(); i++) { + shared_elements *= shared_range[i]->extent; + } + PrimExpr elements = analyzer->Simplify(shared_elements); + PrimExpr size_bytes = elements * shared_tensor->dtype.bytes(); + + // 16-byte alignment check (at compile time if constant) + if (auto *imm = size_bytes.as()) { + ICHECK(imm->value % 16 == 0) + << "multimem_tma_store: transfer size must be 16-byte aligned, got " + << imm->value; + } + + // Compute flat shared offset + std::vector shared_strides; + PrimExpr sh_stride = 1; + for (int i = static_cast(shared_tensor->shape.size()) - 1; i >= 0; --i) { + shared_strides.insert(shared_strides.begin(), sh_stride); + sh_stride *= shared_tensor->shape[i]; + } + PrimExpr shared_offset = 0; + for (size_t i = 0; i < shared_range.size(); i++) { + shared_offset += shared_range[i]->min * shared_strides[i]; + } + + // Compute flat global offset + std::vector global_strides; + PrimExpr gl_stride = 1; + for (int i = static_cast(global_tensor->shape.size()) - 1; i >= 0; --i) { + global_strides.insert(global_strides.begin(), gl_stride); + gl_stride *= global_tensor->shape[i]; + } + PrimExpr global_offset = 0; + for (size_t i = 0; i < global_range.size(); i++) { + global_offset += global_range[i]->min * global_strides[i]; + } + + // Build address_of(BufferLoad(buffer, {flat_offset})) + auto make_addr = [](const Buffer &buf, PrimExpr flat_idx) -> PrimExpr { + return Call(DataType::Handle(), builtin::address_of(), + {BufferLoad(buf, {flat_idx})}); + }; + PrimExpr smem_addr = make_addr(shared_tensor, shared_offset); + PrimExpr mcast_addr = make_addr(global_tensor, global_offset); + + // Build function name based on mode and dtype + std::string func_name; + if (is_reduce) { + func_name = "tl::multimem::cp_reduce_async_bulk_"; + switch (reduce_op) { + case 0: func_name += "add_"; break; + case 1: func_name += "min_"; break; + case 2: func_name += "max_"; break; + default: LOG(FATAL) << "Invalid reduce_op: " << reduce_op; + } + func_name += shared_tensor->dtype.is_float16() ? "f16" : + shared_tensor->dtype.is_bfloat16() ? "bf16" : "f32"; + } else { + func_name = "tl::multimem::cp_async_bulk"; + } + + Array extern_args; + extern_args.push_back(StringImm(func_name)); + extern_args.push_back(mcast_addr); + extern_args.push_back(smem_addr); + extern_args.push_back(size_bytes); + + Stmt bulk_copy = Evaluate( + Call(DataType::Handle(), builtin::call_extern(), extern_args)); + + // Gate with tid == 0 (single thread per CTA emits the PTX) + bulk_copy = IfThenElse(EQ(T.thread_var, T.thread_bounds->min), bulk_copy); + return bulk_copy; +} + // === Clone === TileOperator MultimemOpNode::Clone() const { auto node = tvm::ffi::make_object(*this); diff --git a/src/op/multimem.h b/src/op/multimem.h index 44e0b774bf..982b4be08d 100644 --- a/src/op/multimem.h +++ b/src/op/multimem.h @@ -22,7 +22,13 @@ namespace tl { using namespace tir; -enum class MultimemMode : int { kLdReduce = 0, kSt = 1, kRed = 2 }; +enum class MultimemMode : int { + kLdReduce = 0, + kSt = 1, + kRed = 2, + kTmaStore = 3, // multimem.cp.async.bulk: shared → mcast_global (plain store) + kTmaRedStore = 4, // multimem.cp.reduce.async.bulk: shared → mcast_global (reduce) +}; /*! * \brief Unified multimem operator for NVSwitch SHARP multicast operations. @@ -72,6 +78,7 @@ class MultimemOpNode : public TileOperatorNode { PrimExpr MakePredicate(arith::Analyzer *analyzer, const Array &ivs, Array extents, int src_dst) const; int GetCoalescedWidth() const; + Stmt LowerBulkCopy(const LowerArgs &T, arith::Analyzer *analyzer) const; }; class MultimemOp : public TileOperator { diff --git a/src/tl_templates/cuda/multimem.h b/src/tl_templates/cuda/multimem.h index b5af5bf2b6..e0e9757cc8 100644 --- a/src/tl_templates/cuda/multimem.h +++ b/src/tl_templates/cuda/multimem.h @@ -249,6 +249,117 @@ template <> struct RedV2 { } }; +// === Thread-level signal write to multicast address === + +template struct Signal { + TL_DEVICE static void run(void *, T) { + static_assert(always_false_v, "tl::multimem::Signal: unsupported type"); + } +}; +template <> struct Signal { + TL_DEVICE static void run(void *mcast_ptr, uint32_t val) { + asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" + : : "l"(mcast_ptr), "r"(val) : "memory"); + } +}; +template <> struct Signal { + TL_DEVICE static void run(void *mcast_ptr, uint64_t val) { + asm volatile("multimem.st.relaxed.sys.global.u64 [%0], %1;" + : : "l"(mcast_ptr), "l"(val) : "memory"); + } +}; + +// === Bulk async TMA-to-multicast (SM100+ / PTX 9.1+ / CUDA 13.0+) === +// Both: shared::cta → global(mcast), bulk_group completion + +#if __CUDA_ARCH__ >= 1000 && __CUDACC_VER_MAJOR__ >= 13 + +TL_DEVICE void cp_async_bulk(void *mcast_global, void *smem, uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +TL_DEVICE void cp_reduce_async_bulk_add_f32(void *mcast_global, void *smem, + uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 " + "[%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +TL_DEVICE void cp_reduce_async_bulk_min_f32(void *mcast_global, void *smem, + uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.min.f32 " + "[%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +TL_DEVICE void cp_reduce_async_bulk_max_f32(void *mcast_global, void *smem, + uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.max.f32 " + "[%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +TL_DEVICE void cp_reduce_async_bulk_add_f16(void *mcast_global, void *smem, + uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f16x2 " + "[%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +TL_DEVICE void cp_reduce_async_bulk_add_bf16(void *mcast_global, void *smem, + uint32_t size) { + uint32_t smem_int = smem_ptr_to_uint(smem); + asm volatile( + "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.bf16x2 " + "[%0], [%1], %2;\n" + : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); +} + +#else // PTX 9.1 not available — unconditional trap + +TL_DEVICE void cp_async_bulk(void *mcast_global, void *smem, uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} +TL_DEVICE void cp_reduce_async_bulk_add_f32(void *mcast_global, void *smem, + uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} +TL_DEVICE void cp_reduce_async_bulk_min_f32(void *mcast_global, void *smem, + uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} +TL_DEVICE void cp_reduce_async_bulk_max_f32(void *mcast_global, void *smem, + uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} +TL_DEVICE void cp_reduce_async_bulk_add_f16(void *mcast_global, void *smem, + uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} +TL_DEVICE void cp_reduce_async_bulk_add_bf16(void *mcast_global, void *smem, + uint32_t size) { + (void)mcast_global; (void)smem; (void)size; + asm("trap;"); +} + +#endif // __CUDA_ARCH__ >= 1000 && __CUDACC_VER_MAJOR__ >= 13 + } // namespace multimem } // namespace tl diff --git a/src/transform/inject_fence_proxy.cc b/src/transform/inject_fence_proxy.cc index 6152789a2e..5447afa1a2 100644 --- a/src/transform/inject_fence_proxy.cc +++ b/src/transform/inject_fence_proxy.cc @@ -17,6 +17,8 @@ #include "../op/builtin.h" +#include + namespace tvm { namespace tl { @@ -99,6 +101,16 @@ bool IsAsyncIntrinsic(const CallNode *call) { return true; } + // multimem bulk async ops (shared → mcast_global via bulk_group) + if (call->op.same_as(builtin::call_extern()) && call->args.size() >= 1) { + if (auto *str_imm = call->args[0].as()) { + if (std::string(str_imm->value).find("tl::multimem::cp_async_bulk") == 0 || + std::string(str_imm->value).find("tl::multimem::cp_reduce_async_bulk") == 0) { + return true; + } + } + } + return false; } @@ -143,11 +155,22 @@ class TMAStoreSyncInjector : public StmtExprMutator { private: Stmt operator()(const Stmt &stmt) { return StmtExprMutator::VisitStmt(stmt); } + static bool IsMultimemBulkCall(const CallNode *call) { + if (!call->op.same_as(builtin::call_extern()) || call->args.empty()) + return false; + if (auto *str_imm = call->args[0].as()) { + std::string name(str_imm->value); + return name.find("tl::multimem::cp_async_bulk") == 0 || + name.find("tl::multimem::cp_reduce_async_bulk") == 0; + } + return false; + } + Stmt VisitStmt_(const EvaluateNode *op) final { Stmt mutated = StmtExprMutator::VisitStmt_(op); const auto *node = mutated.as(); if (const auto *call = node->value.as()) { - if (call->op.same_as(tma_store())) { + if (call->op.same_as(tma_store()) || IsMultimemBulkCall(call)) { Array seq; seq.push_back(mutated); seq.push_back( diff --git a/src/transform/warp_specialized_rewriter.cc b/src/transform/warp_specialized_rewriter.cc index 8e891d8551..e5cff6ad52 100644 --- a/src/transform/warp_specialized_rewriter.cc +++ b/src/transform/warp_specialized_rewriter.cc @@ -417,6 +417,49 @@ class ThreadIdxRewriter : public StmtExprMutator { bool has_tma_op_ = false; }; +static bool ExprUsesBlockIdx(const PrimExpr &expr, + const std::unordered_set &block_vars) { + bool uses_block_idx = false; + PostOrderVisit(expr, [&uses_block_idx, &block_vars](const ObjectRef &node) { + if (auto var = node.as()) { + if (block_vars.count(var)) { + uses_block_idx = true; + return; + } + std::string name_hint = var->name_hint; + if (name_hint.find("blockIdx") != std::string::npos) { + uses_block_idx = true; + } + } + }); + return uses_block_idx; +} + +class BlockIdxIfDetector : public StmtVisitor { +public: + static bool Detect(const Stmt &stmt, + const std::unordered_set &block_vars) { + BlockIdxIfDetector detector(block_vars); + detector(stmt); + return detector.has_block_idx_if_; + } + +private: + explicit BlockIdxIfDetector(const std::unordered_set &block_vars) + : block_vars_(block_vars) {} + + void VisitStmt_(const IfThenElseNode *op) final { + if (ExprUsesBlockIdx(op->condition, block_vars_)) { + has_block_idx_if_ = true; + return; + } + StmtVisitor::VisitStmt_(op); + } + + const std::unordered_set &block_vars_; + bool has_block_idx_if_ = false; +}; + Block MakeGroupBlock(const Stmt &stmt, const Map &annotations) { Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"", @@ -1164,9 +1207,20 @@ class WarpSpecializedRewriter : public StmtExprMutator { private: Stmt VisitStmt_(const AttrStmtNode *op) final { - if (op->attr_key == tir::attr::thread_extent && - Downcast(op->node)->thread_tag == "threadIdx.x") { - thread_iv_ = Downcast(op->node); + if (op->attr_key == tir::attr::thread_extent) { + IterVar iv = Downcast(op->node); + if (iv->thread_tag == "blockIdx.x" || iv->thread_tag == "blockIdx.y" || + iv->thread_tag == "blockIdx.z") { + block_vars_.insert(iv->var.get()); + AttrStmt attr_stmt = Downcast(StmtExprMutator::VisitStmt_(op)); + block_vars_.erase(iv->var.get()); + return attr_stmt; + } + if (iv->thread_tag != "threadIdx.x") { + return StmtExprMutator::VisitStmt_(op); + } + + thread_iv_ = iv; need_update_thread_extent_ = false; AttrStmt attr_stmt = Downcast(StmtExprMutator::VisitStmt_(op)); if (need_update_thread_extent_) { @@ -1176,20 +1230,32 @@ class WarpSpecializedRewriter : public StmtExprMutator { } thread_iv_ = {}; return attr_stmt; - } else { - return StmtExprMutator::VisitStmt_(op); } + return StmtExprMutator::VisitStmt_(op); } // If users define a thread binding, we will replace the thread binding with // threadIdx.x We require the thread binding is threadIdx.x, and the extent is // the same as the thread extent Stmt VisitStmt_(const ForNode *op) final { - ICHECK(thread_iv_.defined()); + if (op->kind == ForKind::kThreadBinding && op->thread_binding.defined()) { + String thread_tag = op->thread_binding.value()->thread_tag; + if (thread_tag == "blockIdx.x" || thread_tag == "blockIdx.y" || + thread_tag == "blockIdx.z") { + block_vars_.insert(op->loop_var.get()); + Stmt body = VisitStmt(op->body); + block_vars_.erase(op->loop_var.get()); + For for_node = tvm::ffi::GetRef(op); + for_node.CopyOnWrite()->body = body; + return for_node; + } + } + For for_node = Downcast(StmtExprMutator::VisitStmt_(op)); if (for_node->kind == ForKind::kThreadBinding) { ICHECK(for_node->thread_binding.defined()); String thread_tag = for_node->thread_binding.value()->thread_tag; + ICHECK(thread_iv_.defined()); ICHECK(thread_tag == "threadIdx.x") << "Only support threadIdx.x"; Var thread_iv = Downcast(for_node->loop_var); Stmt new_body = @@ -1207,18 +1273,63 @@ class WarpSpecializedRewriter : public StmtExprMutator { } Block block = block_realize->block; + block.CopyOnWrite()->body = + RewriteBodyRespectingBlockIdxPartitions(block->body); + block_realize.CopyOnWrite()->block = block; + return block_realize; + } + + Stmt RewriteBodyRespectingBlockIdxPartitions(const Stmt &body) { + if (!BlockIdxIfDetector::Detect(body, block_vars_)) { + return ApplyWarpSpecializationToBody(body); + } + + if (auto *seq = body.as()) { + Array new_seq; + new_seq.reserve(seq->seq.size()); + for (const auto &stmt : seq->seq) { + if (auto *ite = stmt.as()) { + if (ExprUsesBlockIdx(ite->condition, block_vars_)) { + new_seq.push_back(RewriteBlockIdxPartition(ite)); + continue; + } + } + new_seq.push_back(StmtExprMutator::VisitStmt(stmt)); + } + return SeqStmt(std::move(new_seq)); + } + + if (auto *ite = body.as()) { + if (ExprUsesBlockIdx(ite->condition, block_vars_)) { + return RewriteBlockIdxPartition(ite); + } + } + + return ApplyWarpSpecializationToBody(body); + } + + Stmt RewriteBlockIdxPartition(const IfThenElseNode *ite) { + Stmt then_case = ApplyWarpSpecializationToBody(ite->then_case); + Optional else_case = std::nullopt; + if (ite->else_case.defined()) { + else_case = ApplyWarpSpecializationToBody(ite->else_case.value()); + } + return IfThenElse(ite->condition, then_case, else_case); + } + + Stmt ApplyWarpSpecializationToBody(const Stmt &body) { WarpSpecializedRoleMarker marker(buffer_data_to_buffer_); - marker.Prepare(block); - marker(block); + marker.Prepare(body); + marker(body); if (!marker.HasProducer()) { // Cannot detect any producer here, directly return. - return block_realize; + return body; } if (disable_warp_specialized_) { WSCodeEmitter mbarrier_emitter(true, thread_iv_, buffer_data_to_buffer_, marker, true); - auto code = mbarrier_emitter(block->body); + auto code = mbarrier_emitter(body); int num_barriers = mbarrier_emitter.num_barriers_; Array barrier_num_threads; barrier_num_threads.reserve(num_barriers); @@ -1228,15 +1339,13 @@ class WarpSpecializedRewriter : public StmtExprMutator { } Stmt init_barrier = Evaluate(Call( DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads)); - block.CopyOnWrite()->body = SeqStmt({init_barrier, code}); - block_realize.CopyOnWrite()->block = block; - return block_realize; + return SeqStmt({init_barrier, code}); } WSCodeEmitter producer(true, thread_iv_, buffer_data_to_buffer_, marker); WSCodeEmitter consumer(false, thread_iv_, buffer_data_to_buffer_, marker, false); - Stmt producer_code = producer(block->body); - Stmt consumer_code = consumer(block->body); + Stmt producer_code = producer(body); + Stmt consumer_code = consumer(body); PrimExpr consumer_thread_extent = thread_iv_->dom->extent; PrimExpr producer_thread_extent = thread_iv_->dom->extent; // Need one warp-group for bulk-copy only case @@ -1269,16 +1378,14 @@ class WarpSpecializedRewriter : public StmtExprMutator { Stmt init_barrier = Evaluate(Call( DataType::Handle(), create_list_of_mbarrier(), barrier_num_threads)); - Stmt body = IfThenElse(GE(thread_iv_->var, consumer_thread_extent), - producer_code, consumer_code); + Stmt ws_body = IfThenElse(GE(thread_iv_->var, consumer_thread_extent), + producer_code, consumer_code); // Add an attr here to handle the partial thread count in ThreadSync pass. Array ws_partition = {Downcast(producer_thread_extent), Downcast(consumer_thread_extent)}; - body = AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, body); + ws_body = AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, ws_body); - block.CopyOnWrite()->body = SeqStmt({init_barrier, body}); - block_realize.CopyOnWrite()->block = block; - return block_realize; + return SeqStmt({init_barrier, ws_body}); } WarpSpecializedRewriter() = default; @@ -1286,6 +1393,7 @@ class WarpSpecializedRewriter : public StmtExprMutator { Map buffer_data_to_buffer_; Map> buffer_lca_; Map buffer_remap_; + std::unordered_set block_vars_; IterVar thread_iv_; Optional updated_thread_extent_; bool need_update_thread_extent_ = false; diff --git a/tilelang/language/distributed/__init__.py b/tilelang/language/distributed/__init__.py index c11f74a9eb..5b5e5d5880 100644 --- a/tilelang/language/distributed/__init__.py +++ b/tilelang/language/distributed/__init__.py @@ -23,6 +23,8 @@ multimem_ld_reduce, multimem_st, multimem_red, + multimem_tma_store, + multimem_signal, ) __all__ = [ @@ -43,4 +45,6 @@ "multimem_ld_reduce", "multimem_st", "multimem_red", + "multimem_tma_store", + "multimem_signal", ] diff --git a/tilelang/language/distributed/multimem.py b/tilelang/language/distributed/multimem.py index 034140c832..0964176624 100644 --- a/tilelang/language/distributed/multimem.py +++ b/tilelang/language/distributed/multimem.py @@ -5,7 +5,9 @@ """ from enum import Enum +from typing import Literal from tvm import tir +from tvm.tir import PrimExpr, address_of from tilelang.utils.language import to_buffer_region @@ -13,15 +15,18 @@ class MultimemReduceOp(Enum): ADD = 0 MIN = 1 MAX = 2 + NONE = -1 # plain store (no reduction), for multimem_tma_store class _MultimemMode(Enum): LD_REDUCE = 0 ST = 1 RED = 2 + TMA_STORE = 3 + TMA_RED_STORE = 4 -def _multimem_impl(src, dst, mode: _MultimemMode, reduce_op: MultimemReduceOp): +def _multimem_impl(src, dst, mode: _MultimemMode, reduce_op: MultimemReduceOp = MultimemReduceOp.NONE): """Shared implementation for all multimem operations. Converts src/dst to buffer regions and emits the tl.tileop.multimem intrinsic. @@ -65,7 +70,7 @@ def multimem_st(src, dst): src: Local source (Buffer, BufferLoad with slice, or BufferRegion) dst: Multicast destination (Buffer, BufferLoad with slice, or BufferRegion) """ - return _multimem_impl(src, dst, mode=_MultimemMode.ST, reduce_op=MultimemReduceOp.ADD) + return _multimem_impl(src, dst, mode=_MultimemMode.ST) def multimem_red(src, dst, reduce_op: MultimemReduceOp = MultimemReduceOp.ADD): @@ -77,3 +82,28 @@ def multimem_red(src, dst, reduce_op: MultimemReduceOp = MultimemReduceOp.ADD): reduce_op: Reduction operation: 0=ADD, 1=MIN, 2=MAX. """ return _multimem_impl(src, dst, mode=_MultimemMode.RED, reduce_op=reduce_op) + + +def multimem_tma_store(src, dst, reduce_op: MultimemReduceOp | None = None): + """Async bulk TMA store from shared memory to multicast global address. + + CTA-collective: a single thread emits one PTX instruction per call. + Uses bulk_group completion (fence.proxy.async + commit_group + wait). + + Args: + src: Shared memory source (Buffer, BufferLoad or BufferRegion, shared scope) + dst: Multicast global destination (Buffer, BufferLoad or BufferRegion, global scope) + reduce_op: None for plain store (broadcast), MultimemReduceOp.ADD/MIN/MAX for reduce-accumulate + + NOTE: This instruction requires Hopper+ and CUDA toolkit 13.x. + (For unsatisfied CTK version, a hack is to use plain TMA store to mcast vaddr.) + """ + if reduce_op is None: + return _multimem_impl(src, dst, mode=_MultimemMode.TMA_STORE) + return _multimem_impl(src, dst, mode=_MultimemMode.TMA_RED_STORE, reduce_op=reduce_op) + + +def multimem_signal(addr, value: PrimExpr, dtype_tag: Literal['uint32', 'uint64'] ='uint32'): + return tir.call_extern("handle", + f"tl::multimem::Signal<{dtype_tag}>::run", + address_of(addr), value) From 823fd8a6ca80d5428b29727ca8d4406272c16b53 Mon Sep 17 00:00:00 2001 From: Rachmanino <18805904201@163.com> Date: Thu, 7 May 2026 16:08:24 +0800 Subject: [PATCH 9/9] fix lint: add from __future__ import annotations to multimem.py Co-Authored-By: Claude Opus 4.6 --- CMakeLists.txt | 1 - .../example_allgather_gemm_specialized.py | 9 +- .../distributed/example_gemm_allreduce.py | 10 +-- src/op/multimem.cc | 24 +++-- src/op/multimem.h | 5 +- src/tl_templates/cuda/multimem.h | 87 ++++++++++++------- src/transform/inject_fence_proxy.cc | 6 +- src/transform/warp_specialized_rewriter.cc | 14 +-- tilelang/language/distributed/multimem.py | 9 +- 9 files changed, 101 insertions(+), 64 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8a75bf3c21..e30dddec31 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -335,4 +335,3 @@ install( TARGETS tvm tvm_runtime tilelang_module tilelang LIBRARY DESTINATION tilelang/lib ) - diff --git a/examples/distributed/example_allgather_gemm_specialized.py b/examples/distributed/example_allgather_gemm_specialized.py index 9a063d9d42..17a6a6a6e8 100644 --- a/examples/distributed/example_allgather_gemm_specialized.py +++ b/examples/distributed/example_allgather_gemm_specialized.py @@ -52,7 +52,7 @@ def main( local_signal: T.Tensor((m_blocks,), "uint32"), grid_barrier: T.Tensor((num_ranks,), "int32"), C: T.Tensor((M, N_per_rank), dtype), - local_rank: T.int32 + local_rank: T.int32, ): with T.Kernel(sm_num, threads=threads) as bid: A_shared = T.alloc_shared((block_M, block_K), dtype) @@ -218,10 +218,7 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): warmup=args.warmup, rep=args.rep, ) - print( - f"rank {local_rank} tilelang specialized ag_gemm time: {tl_t:.2f} ms, " - f"TFLOPS: {2 * M * N * K / 1e9 / tl_t / num_local_ranks:.2f}" - ) + print(f"rank {local_rank} tilelang specialized ag_gemm time: {tl_t:.2f} ms, TFLOPS: {2 * M * N * K / 1e9 / tl_t / num_local_ranks:.2f}") allocator.close() dist.destroy_process_group() @@ -243,4 +240,4 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): parser.add_argument("--print-source", action="store_true") args = parser.parse_args() - torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) \ No newline at end of file + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) diff --git a/examples/distributed/example_gemm_allreduce.py b/examples/distributed/example_gemm_allreduce.py index fa13500be0..d8d8a79922 100644 --- a/examples/distributed/example_gemm_allreduce.py +++ b/examples/distributed/example_gemm_allreduce.py @@ -60,14 +60,14 @@ def gemm_red( T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=2): - T.copy(A[by * block_M:(by + 1) * block_M, k * block_K:(k + 1) * block_K], A_shared) - T.copy(B[k * block_K:(k + 1) * block_K, bx * block_N:(bx + 1) * block_N], B_shared) + T.copy(A[by * block_M : (by + 1) * block_M, k * block_K : (k + 1) * block_K], A_shared) + T.copy(B[k * block_K : (k + 1) * block_K, bx * block_N : (bx + 1) * block_N], B_shared) T.gemm(A_shared, B_shared, C_local) # Reduce partial GEMM result into multicast buffer T.multimem_red( C_local, - mcast_C[by * block_M:(by + 1) * block_M, bx * block_N:(bx + 1) * block_N], + mcast_C[by * block_M : (by + 1) * block_M, bx * block_N : (bx + 1) * block_N], reduce_op=T.MultimemReduceOp.ADD, ) T.fence_sys() @@ -162,6 +162,4 @@ def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): parser.add_argument("--print_source", action="store_true") args = parser.parse_args() - torch.multiprocessing.spawn( - main, args=(args.num_processes, args), nprocs=args.num_processes, join=True - ) + torch.multiprocessing.spawn(main, args=(args.num_processes, args), nprocs=args.num_processes, join=True) diff --git a/src/op/multimem.cc b/src/op/multimem.cc index 866a54dea4..1fa771b82e 100644 --- a/src/op/multimem.cc +++ b/src/op/multimem.cc @@ -343,13 +343,21 @@ Stmt MultimemOpNode::LowerBulkCopy(const LowerArgs &T, if (is_reduce) { func_name = "tl::multimem::cp_reduce_async_bulk_"; switch (reduce_op) { - case 0: func_name += "add_"; break; - case 1: func_name += "min_"; break; - case 2: func_name += "max_"; break; - default: LOG(FATAL) << "Invalid reduce_op: " << reduce_op; + case 0: + func_name += "add_"; + break; + case 1: + func_name += "min_"; + break; + case 2: + func_name += "max_"; + break; + default: + LOG(FATAL) << "Invalid reduce_op: " << reduce_op; } - func_name += shared_tensor->dtype.is_float16() ? "f16" : - shared_tensor->dtype.is_bfloat16() ? "bf16" : "f32"; + func_name += shared_tensor->dtype.is_float16() ? "f16" + : shared_tensor->dtype.is_bfloat16() ? "bf16" + : "f32"; } else { func_name = "tl::multimem::cp_async_bulk"; } @@ -360,8 +368,8 @@ Stmt MultimemOpNode::LowerBulkCopy(const LowerArgs &T, extern_args.push_back(smem_addr); extern_args.push_back(size_bytes); - Stmt bulk_copy = Evaluate( - Call(DataType::Handle(), builtin::call_extern(), extern_args)); + Stmt bulk_copy = + Evaluate(Call(DataType::Handle(), builtin::call_extern(), extern_args)); // Gate with tid == 0 (single thread per CTA emits the PTX) bulk_copy = IfThenElse(EQ(T.thread_var, T.thread_bounds->min), bulk_copy); diff --git a/src/op/multimem.h b/src/op/multimem.h index 982b4be08d..33f762b9b3 100644 --- a/src/op/multimem.h +++ b/src/op/multimem.h @@ -26,8 +26,9 @@ enum class MultimemMode : int { kLdReduce = 0, kSt = 1, kRed = 2, - kTmaStore = 3, // multimem.cp.async.bulk: shared → mcast_global (plain store) - kTmaRedStore = 4, // multimem.cp.reduce.async.bulk: shared → mcast_global (reduce) + kTmaStore = 3, // multimem.cp.async.bulk: shared → mcast_global (plain store) + kTmaRedStore = + 4, // multimem.cp.reduce.async.bulk: shared → mcast_global (reduce) }; /*! diff --git a/src/tl_templates/cuda/multimem.h b/src/tl_templates/cuda/multimem.h index e0e9757cc8..432a6d751d 100644 --- a/src/tl_templates/cuda/multimem.h +++ b/src/tl_templates/cuda/multimem.h @@ -183,7 +183,8 @@ template <> struct LdReduceV2 { template struct StV2 { TL_DEVICE static void run(void *, const void *) { - static_assert(always_false_v, "tl::multimem::StV2: unsupported dtype"); + static_assert(always_false_v, + "tl::multimem::StV2: unsupported dtype"); } }; @@ -208,7 +209,7 @@ template <> struct RedV2 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { const float *src_f = reinterpret_cast(src); const char *mc = reinterpret_cast(mcast_ptr); - #pragma unroll +#pragma unroll for (int i = 0; i < 2; i++) { unsigned val = __float_as_uint(src_f[i]); asm volatile("multimem.red.relaxed.sys.global.add.f32 [%0], %1;" @@ -223,7 +224,7 @@ template <> struct RedV2 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { const float *src_f = reinterpret_cast(src); const char *mc = reinterpret_cast(mcast_ptr); - #pragma unroll +#pragma unroll for (int i = 0; i < 2; i++) { unsigned val = __float_as_uint(src_f[i]); asm volatile("multimem.red.relaxed.sys.global.min.f32 [%0], %1;" @@ -238,7 +239,7 @@ template <> struct RedV2 { TL_DEVICE static void run(void *mcast_ptr, const void *src) { const float *src_f = reinterpret_cast(src); const char *mc = reinterpret_cast(mcast_ptr); - #pragma unroll +#pragma unroll for (int i = 0; i < 2; i++) { unsigned val = __float_as_uint(src_f[i]); asm volatile("multimem.red.relaxed.sys.global.max.f32 [%0], %1;" @@ -259,13 +260,17 @@ template struct Signal { template <> struct Signal { TL_DEVICE static void run(void *mcast_ptr, uint32_t val) { asm volatile("multimem.st.relaxed.sys.global.u32 [%0], %1;" - : : "l"(mcast_ptr), "r"(val) : "memory"); + : + : "l"(mcast_ptr), "r"(val) + : "memory"); } }; template <> struct Signal { TL_DEVICE static void run(void *mcast_ptr, uint64_t val) { asm volatile("multimem.st.relaxed.sys.global.u64 [%0], %1;" - : : "l"(mcast_ptr), "l"(val) : "memory"); + : + : "l"(mcast_ptr), "l"(val) + : "memory"); } }; @@ -278,83 +283,107 @@ TL_DEVICE void cp_async_bulk(void *mcast_global, void *smem, uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } TL_DEVICE void cp_reduce_async_bulk_add_f32(void *mcast_global, void *smem, - uint32_t size) { + uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 " "[%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } TL_DEVICE void cp_reduce_async_bulk_min_f32(void *mcast_global, void *smem, - uint32_t size) { + uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.min.f32 " "[%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } TL_DEVICE void cp_reduce_async_bulk_max_f32(void *mcast_global, void *smem, - uint32_t size) { + uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.max.f32 " "[%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } TL_DEVICE void cp_reduce_async_bulk_add_f16(void *mcast_global, void *smem, - uint32_t size) { + uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f16x2 " "[%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } TL_DEVICE void cp_reduce_async_bulk_add_bf16(void *mcast_global, void *smem, - uint32_t size) { + uint32_t size) { uint32_t smem_int = smem_ptr_to_uint(smem); asm volatile( "multimem.cp.reduce.async.bulk.global.shared::cta.bulk_group.add.bf16x2 " "[%0], [%1], %2;\n" - : : "l"(mcast_global), "r"(smem_int), "r"(size) : "memory"); + : + : "l"(mcast_global), "r"(smem_int), "r"(size) + : "memory"); } -#else // PTX 9.1 not available — unconditional trap +#else // PTX 9.1 not available — unconditional trap TL_DEVICE void cp_async_bulk(void *mcast_global, void *smem, uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } TL_DEVICE void cp_reduce_async_bulk_add_f32(void *mcast_global, void *smem, - uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + uint32_t size) { + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } TL_DEVICE void cp_reduce_async_bulk_min_f32(void *mcast_global, void *smem, - uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + uint32_t size) { + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } TL_DEVICE void cp_reduce_async_bulk_max_f32(void *mcast_global, void *smem, - uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + uint32_t size) { + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } TL_DEVICE void cp_reduce_async_bulk_add_f16(void *mcast_global, void *smem, - uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + uint32_t size) { + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } TL_DEVICE void cp_reduce_async_bulk_add_bf16(void *mcast_global, void *smem, - uint32_t size) { - (void)mcast_global; (void)smem; (void)size; + uint32_t size) { + (void)mcast_global; + (void)smem; + (void)size; asm("trap;"); } diff --git a/src/transform/inject_fence_proxy.cc b/src/transform/inject_fence_proxy.cc index 5447afa1a2..7f383c11e6 100644 --- a/src/transform/inject_fence_proxy.cc +++ b/src/transform/inject_fence_proxy.cc @@ -104,8 +104,10 @@ bool IsAsyncIntrinsic(const CallNode *call) { // multimem bulk async ops (shared → mcast_global via bulk_group) if (call->op.same_as(builtin::call_extern()) && call->args.size() >= 1) { if (auto *str_imm = call->args[0].as()) { - if (std::string(str_imm->value).find("tl::multimem::cp_async_bulk") == 0 || - std::string(str_imm->value).find("tl::multimem::cp_reduce_async_bulk") == 0) { + if (std::string(str_imm->value).find("tl::multimem::cp_async_bulk") == + 0 || + std::string(str_imm->value) + .find("tl::multimem::cp_reduce_async_bulk") == 0) { return true; } } diff --git a/src/transform/warp_specialized_rewriter.cc b/src/transform/warp_specialized_rewriter.cc index e5cff6ad52..0b90bbe8ff 100644 --- a/src/transform/warp_specialized_rewriter.cc +++ b/src/transform/warp_specialized_rewriter.cc @@ -417,8 +417,9 @@ class ThreadIdxRewriter : public StmtExprMutator { bool has_tma_op_ = false; }; -static bool ExprUsesBlockIdx(const PrimExpr &expr, - const std::unordered_set &block_vars) { +static bool +ExprUsesBlockIdx(const PrimExpr &expr, + const std::unordered_set &block_vars) { bool uses_block_idx = false; PostOrderVisit(expr, [&uses_block_idx, &block_vars](const ObjectRef &node) { if (auto var = node.as()) { @@ -445,7 +446,8 @@ class BlockIdxIfDetector : public StmtVisitor { } private: - explicit BlockIdxIfDetector(const std::unordered_set &block_vars) + explicit BlockIdxIfDetector( + const std::unordered_set &block_vars) : block_vars_(block_vars) {} void VisitStmt_(const IfThenElseNode *op) final { @@ -1212,7 +1214,8 @@ class WarpSpecializedRewriter : public StmtExprMutator { if (iv->thread_tag == "blockIdx.x" || iv->thread_tag == "blockIdx.y" || iv->thread_tag == "blockIdx.z") { block_vars_.insert(iv->var.get()); - AttrStmt attr_stmt = Downcast(StmtExprMutator::VisitStmt_(op)); + AttrStmt attr_stmt = + Downcast(StmtExprMutator::VisitStmt_(op)); block_vars_.erase(iv->var.get()); return attr_stmt; } @@ -1383,7 +1386,8 @@ class WarpSpecializedRewriter : public StmtExprMutator { // Add an attr here to handle the partial thread count in ThreadSync pass. Array ws_partition = {Downcast(producer_thread_extent), Downcast(consumer_thread_extent)}; - ws_body = AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, ws_body); + ws_body = + AttrStmt(ws_partition, attr::kWarpSpecializationScope, 0, ws_body); return SeqStmt({init_barrier, ws_body}); } diff --git a/tilelang/language/distributed/multimem.py b/tilelang/language/distributed/multimem.py index 0964176624..bdad74b813 100644 --- a/tilelang/language/distributed/multimem.py +++ b/tilelang/language/distributed/multimem.py @@ -4,6 +4,7 @@ to correctly handle fragment layouts, then post-process to emit multimem instructions. """ +from __future__ import annotations from enum import Enum from typing import Literal from tvm import tir @@ -94,7 +95,7 @@ def multimem_tma_store(src, dst, reduce_op: MultimemReduceOp | None = None): src: Shared memory source (Buffer, BufferLoad or BufferRegion, shared scope) dst: Multicast global destination (Buffer, BufferLoad or BufferRegion, global scope) reduce_op: None for plain store (broadcast), MultimemReduceOp.ADD/MIN/MAX for reduce-accumulate - + NOTE: This instruction requires Hopper+ and CUDA toolkit 13.x. (For unsatisfied CTK version, a hack is to use plain TMA store to mcast vaddr.) """ @@ -103,7 +104,5 @@ def multimem_tma_store(src, dst, reduce_op: MultimemReduceOp | None = None): return _multimem_impl(src, dst, mode=_MultimemMode.TMA_RED_STORE, reduce_op=reduce_op) -def multimem_signal(addr, value: PrimExpr, dtype_tag: Literal['uint32', 'uint64'] ='uint32'): - return tir.call_extern("handle", - f"tl::multimem::Signal<{dtype_tag}>::run", - address_of(addr), value) +def multimem_signal(addr, value: PrimExpr, dtype_tag: Literal["uint32", "uint64"] = "uint32"): + return tir.call_extern("handle", f"tl::multimem::Signal<{dtype_tag}>::run", address_of(addr), value)