Merge branch 'main' into feature/tle_fix_extract_tile

Author: lzllx123

.github/workflows/enflame3.6-gcu400-build-and-test.yml

@@ -106,7 +106,8 @@ jobs:
           python3 -m pytest python/test/tle \
               --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \
               --ignore=python/test/tle/unit/test_tle_get_local_pe.py \
-              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py \
+              --ignore=python/test/tle/unit/test_tle_get_node_rank.py
 
           ## tle raw test
           python3 -m pytest third_party/enflame/python/test/tle/raw

.github/workflows/new-prs.yml

@@ -40,7 +40,7 @@ jobs:
       pull-requests: write
     steps:
       - id: labeler
-        uses: actions/labeler@v6
+        uses: actions/labeler@v7
         with:
           repo-token: "${{ secrets.GITHUB_TOKEN }}"
           sync-labels: true

.github/workflows/nvidia3.6-build-and-test.yml

@@ -130,7 +130,8 @@ jobs:
           python3 -m pytest -s python/test/tle/unit \
               --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \
               --ignore=python/test/tle/unit/test_tle_get_local_pe.py \
-              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py \
+              --ignore=python/test/tle/unit/test_tle_get_node_rank.py
           ## flagtree hints python tutorials
           python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
           # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test

packaging/README.md

@@ -38,6 +38,9 @@ A 36-minute compile in `override_dh_auto_build` is slow, hard to debug, and
 collides with `dh-python`'s expectations for a pure-Python `pyproject` build.
 The two-stage Docker approach builds the wheel once, then wraps it cheaply
 into a `.deb`/`.rpm` — keeping the heavy lifting in a controlled container.
+The wheel builders use the distro nlohmann-json headers through
+`JSON_SYSPATH=/usr`; backend-specific LLVM and device tools remain pinned to
+the versions expected by each backend.
 
 ## Build locally
 

packaging/debian/build-helpers/Dockerfile.deb

@@ -59,7 +59,7 @@ ENV PIP_ROOT_USER_ACTION=ignore
 RUN apt-get update && apt-get install -y --no-install-recommends \
         python3 python3-pip python3-dev python3-venv \
         build-essential cmake ninja-build git \
-        zlib1g zlib1g-dev libxml2 libxml2-dev \
+        zlib1g zlib1g-dev libxml2 libxml2-dev nlohmann-json3-dev \
         wget curl ca-certificates \
     && apt-get clean && rm -rf /var/lib/apt/lists/*
 
@@ -96,6 +96,7 @@ COPY CMakeLists.txt pyproject.toml setup.py MANIFEST.in LICENSE README.md /src/
 # tree and the LLVM tarball cache (~/.triton) would otherwise add
 # several GB and push CI runners into ENOSPC.
 RUN unset FLAGTREE_BACKEND && \
+    JSON_SYSPATH=/usr \
     FLAGTREE_DEFAULT_BACKENDS=nvidia,amd \
     MAX_JOBS=4 python -m pip wheel . --no-build-isolation --no-deps -w /wheels -v && \
     ls -lh /wheels/ && \

packaging/rpm/helpers/Dockerfile.rpm

@@ -38,6 +38,8 @@ RUN dnf install -y --setopt=install_weak_deps=False \
         gcc gcc-c++ make cmake ninja-build git \
         zlib-devel libxml2-devel \
         wget curl ca-certificates \
+    && (dnf install -y --setopt=install_weak_deps=False json-devel \
+        || dnf install -y --setopt=install_weak_deps=False nlohmann-json-devel) \
     && dnf clean all
 
 RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel build
@@ -67,6 +69,7 @@ COPY CMakeLists.txt pyproject.toml setup.py MANIFEST.in LICENSE README.md /src/
 # tree and the LLVM tarball cache (~/.triton) would otherwise add
 # several GB and push CI runners into ENOSPC.
 RUN unset FLAGTREE_BACKEND && \
+    JSON_SYSPATH=/usr \
     MAX_JOBS=4 python3 -m pip wheel . --no-build-isolation --no-deps -w /wheels -v && \
     ls -lh /wheels/ && \
     test -n "$(ls /wheels/flagtree-*.whl 2>/dev/null)" && \

python/test/tle/unit/test_tle_distributed.py

@@ -226,6 +226,16 @@ def __init__(self):
         self.builder = _LegacyDistributedBarrierBuilder()
 
 
+class TestShardId:
+
+    @pytest.mark.parametrize("axis", ("device", "node"))
+    def test_rank_axis_requires_device_dptr(self, axis):
+        mesh = tle.device_mesh({"node": 2, "device": 4})
+        semantic = _FakeSemantic()
+        with pytest.raises(ValueError, match=rf"device_dptr is required for axis '{axis}'"):
+            tle.shard_id(mesh, axis, _semantic=semantic)
+
+
 class TestDistributedBarrierScope:
 
     def test_distributed_barrier_full_cluster_mesh(self):

python/test/tle/unit/test_tle_get_node_rank.py

@@ -0,0 +1,77 @@
+import os
+
+import torch
+import torch.distributed as dist
+import triton
+import triton.experimental.tle.language as tle
+import triton.language as tl
+
+LOCAL_WORLD_SIZE = int(os.environ["LOCAL_WORLD_SIZE"])
+WORLD_SIZE = int(os.environ["WORLD_SIZE"])
+if WORLD_SIZE % LOCAL_WORLD_SIZE != 0:
+    raise ValueError("WORLD_SIZE must be divisible by LOCAL_WORLD_SIZE")
+
+DEVICE_MESH = tle.device_mesh(tle.MeshConfig(node=WORLD_SIZE // LOCAL_WORLD_SIZE, device=LOCAL_WORLD_SIZE))
+
+
+@triton.jit
+def _tle_node_rank_kernel(out_ptr, device_dptr: tl.constexpr, mesh: tl.constexpr):
+    pid = tl.program_id(0)
+    node_rank = tle.shard_id(mesh, "node", device_dptr=device_dptr)
+    tl.store(out_ptr + pid, node_rank)
+
+
+def test_tle_get_node_rank():
+    grid = 2
+    with torch.cuda.use_mem_pool(tle.get_mem_pool()):
+        source = torch.empty((1, ), dtype=torch.float32, device="cuda")
+    device_dptr = tle.create_dist_tensor(source)
+    node_rank_out = torch.empty((grid, ), dtype=torch.int32, device="cuda")
+
+    compiled = _tle_node_rank_kernel.warmup(
+        out_ptr=node_rank_out,
+        device_dptr=device_dptr,
+        mesh=DEVICE_MESH,
+        grid=(grid, ),
+        num_ctas=1,
+        num_warps=4,
+    )
+    assert "get_world_rank" in compiled.asm["ttgir"]
+    assert "get_num_pes" in compiled.asm["ttgir"]
+    assert "flagcxDevCommGetRank" in compiled.asm["ptx"]
+    assert "flagcxDevCommGetIntraSize" in compiled.asm["ptx"]
+
+    _tle_node_rank_kernel[(grid, )](
+        out_ptr=node_rank_out,
+        device_dptr=device_dptr,
+        mesh=DEVICE_MESH,
+    )
+    torch.cuda.synchronize()
+
+    rank = dist.get_rank()
+    expected_node_rank = rank // LOCAL_WORLD_SIZE
+    actual_node_ranks = node_rank_out.cpu().tolist()
+    try:
+        torch.testing.assert_close(
+            node_rank_out,
+            torch.full_like(node_rank_out, expected_node_rank),
+        )
+    except AssertionError:
+        print(
+            f"[Rank {rank}] FAILED: node ranks={actual_node_ranks}, "
+            f"expected={expected_node_rank}",
+            flush=True,
+        )
+        raise
+    else:
+        print(
+            f"[Rank {rank}] PASSED: node ranks={actual_node_ranks}, "
+            f"expected={expected_node_rank}",
+            flush=True,
+        )
+    finally:
+        tle.cleanup_communicator()
+
+
+if __name__ == "__main__":
+    test_tle_get_node_rank()

python/test/tle/unit/test_tle_get_node_rank.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+if [ "$1" = "debug" ]; then
+    export NCCL_DEBUG=INFO
+    export NCCL_DEBUG_SUBSYS=all
+else
+    unset NCCL_DEBUG
+    unset NCCL_DEBUG_SUBSYS
+fi
+
+export FLAGCX_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_6,mlx5_7,mlx5_8,mlx5_9
+export FLAGCX_USE_HETERO_COMM=1
+export FLAGCX_MEM_ENABLE=1
+export FLAGCX_VMM_ENABLE=0
+export FLAGCX_P2P_DISABLE=1
+export CUDA_VISIBLE_DEVICES=0,1
+
+nproc_per_node=${NPROC_PER_NODE:-2}
+nnodes=${NNODES:-2}
+node_rank=${NODE_RANK:-0}
+master_addr=${MASTER_ADDR:-10.0.9.3}
+port=${MASTER_PORT:-8335}
+
+if [ "${nnodes}" -eq 1 ]; then
+    while ss -ltn | grep -q ":${port} "; do
+        echo "Port ${port} is occupied, trying next..."
+        port=$((port + 2))
+    done
+fi
+
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+echo "Using master ${master_addr}:${port}, node ${node_rank}/${nnodes}"
+
+torchrun \
+    --nproc_per_node="${nproc_per_node}" \
+    --nnodes="${nnodes}" \
+    --node_rank="${node_rank}" \
+    --master_addr="${master_addr}" \
+    --master_port="${port}" \
+    "${script_dir}/test_tle_get_node_rank.py"

python/test/tle/unit/test_tle_raw_cache_key.py

@@ -1,6 +1,7 @@
 import ast
 import textwrap
 
+import pytest
 import triton.language as tl
 from triton.experimental.tle.raw import dialect
 from triton.experimental.tle.raw.cache_key import (
@@ -127,6 +128,7 @@ def make_cache_key() -> str:
 
 
 def test_mlir_dialect_cache_key_changes_with_edsl_source():
+    pytest.importorskip("mlir", reason="requires the optional MLIR Python bindings")
 
     @dialect(name="mlir")
     def edsl_v1():