diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 77c9df153..e364da905 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -69,6 +69,10 @@ ARG TILELANG_REPO="https://github.com/HaiShaw/tilelang.git" ARG TILELANG_BRANCH="dsv32-mi35x" ARG TILELANG_COMMIT="ae938cf885743f165a19656d1122ad42bb0e30b8" +ARG TILELANG_GFX942_REPO="https://github.com/tile-ai/tilelang.git" +ARG TILELANG_GFX942_BRANCH="main" +ARG TILELANG_GFX942_COMMIT="2d8d3676eda18bd3d8e6fa783399ff96d3cd4ded" + ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git" ARG FHT_BRANCH="rocm" ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1" @@ -215,66 +219,82 @@ ENV LIBGL_ALWAYS_INDIRECT=1 RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment RUN /bin/bash -lc 'set -euo pipefail; \ - # Build TileLang only for gfx950 - if [ "${GPU_ARCH:-}" != "gfx950" ]; then \ + # Build TileLang for gfx950 and gfx942-rocm700 + if [ "${GPU_ARCH:-}" != "gfx950" ] && [ "${GPU_ARCH:-}" != "gfx942-rocm700" ]; then \ echo "[TileLang] Skipping (GPU_ARCH=${GPU_ARCH:-unset})"; \ exit 0; \ fi; \ echo "[TileLang] Building TileLang for ${GPU_ARCH}"; \ - \ - # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing) - apt-get update && apt-get install -y --no-install-recommends \ - build-essential git wget curl ca-certificates gnupg \ - libgtest-dev libgmock-dev \ - libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \ - python3 python3-dev python3-setuptools python3-pip \ - gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \ - cmake ninja-build pkg-config libstdc++6 \ - && rm -rf /var/lib/apt/lists/*; \ - \ - # Build GoogleTest static libs (Ubuntu package ships sources only) - cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \ - cmake --build /tmp/build-gtest -j"$(nproc)" && \ - cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \ - rm -rf /tmp/build-gtest; \ - \ - # Keep setuptools < 80 (compat with base image) - python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \ - python3 -m pip cache purge || true; \ - \ - # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing - LLVM_CONFIG_PATH=""; \ - for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \ - if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \ - done; \ - if [ -z "$LLVM_CONFIG_PATH" ]; then \ - echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \ - curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \ - chmod +x /tmp/llvm.sh; \ - /tmp/llvm.sh 18; \ - LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \ - if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \ - fi; \ - echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \ - export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \ - export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \ - \ - # Optional shim for tools that expect llvm-config-16 - mkdir -p /usr/local/bin && \ - printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \ - chmod +x /usr/local/bin/llvm-config-16; \ - \ - # TVM Python bits need Cython - python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \ - \ - # Clone + pin TileLang (bundled TVM), then build - git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \ - cd /opt/tilelang && \ - git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \ - git checkout -f "${TILELANG_COMMIT}" && \ - git submodule update --init --recursive && \ - export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \ - bash ./install_rocm.sh' + if [ "$GPU_ARCH" = "gfx950" ]; then \ + \ + # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing) + apt-get update && apt-get install -y --no-install-recommends \ + build-essential git wget curl ca-certificates gnupg \ + libgtest-dev libgmock-dev \ + libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \ + python3 python3-dev python3-setuptools python3-pip \ + gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \ + cmake ninja-build pkg-config libstdc++6 \ + && rm -rf /var/lib/apt/lists/*; \ + \ + # Build GoogleTest static libs (Ubuntu package ships sources only) + cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \ + cmake --build /tmp/build-gtest -j"$(nproc)" && \ + cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \ + rm -rf /tmp/build-gtest; \ + \ + # Keep setuptools < 80 (compat with base image) + python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \ + python3 -m pip cache purge || true; \ + \ + # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing + LLVM_CONFIG_PATH=""; \ + for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \ + if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \ + done; \ + if [ -z "$LLVM_CONFIG_PATH" ]; then \ + echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \ + curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \ + chmod +x /tmp/llvm.sh; \ + /tmp/llvm.sh 18; \ + LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \ + if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \ + fi; \ + echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \ + export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \ + export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \ + \ + # Optional shim for tools that expect llvm-config-16 + mkdir -p /usr/local/bin && \ + printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \ + chmod +x /usr/local/bin/llvm-config-16; \ + \ + # TVM Python bits need Cython + python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \ + \ + # Clone + pin TileLang (bundled TVM), then build + git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \ + cd /opt/tilelang && \ + git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \ + git checkout -f "${TILELANG_COMMIT}" && \ + git submodule update --init --recursive && \ + export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \ + bash ./install_rocm.sh; \ + else \ + # Build GoogleTest static libs (Ubuntu package ships sources only) + apt-get install -y libgtest-dev libgmock-dev && \ + cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \ + cmake --build /tmp/build-gtest -j && \ + cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \ + rm -rf /tmp/build-gtest; \ + # Build TileLang for gfx942-rocm700 + git clone --branch "${TILELANG_GFX942_BRANCH}" "${TILELANG_GFX942_REPO}" /opt/tilelang && \ + cd /opt/tilelang && \ + git checkout -f "${TILELANG_GFX942_COMMIT}" && \ + git submodule update --init --recursive && \ + sed -i "/^[[:space:]]*\"torch/d" pyproject.toml && \ + USE_ROCM=1 USE_CUDA=0 pip install -e . -v ; \ + fi' # ----------------------- # Hadamard-transform (HIP build) diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index 42dbc1589..1cdf65b91 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -4,8 +4,12 @@ import torch import triton import triton.language as tl +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.utils import is_hip +_is_hip = is_hip() +_is_fp8_fnuz = is_fp8_fnuz() + if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool @@ -349,21 +353,24 @@ def _set_k_and_s_triton( raise ValueError( f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}" ) - if is_hip(): + if _is_hip: assert buf_numel_per_page == 1 * (128 + 4) else: assert buf_numel_per_page == 64 * (128 + 4) assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__ assert index_head_dim == 128 assert scale_dim == 1 - if is_hip(): + if _is_hip: assert page_size == 1 else: assert page_size == 64 assert buf.dtype == torch.uint8 assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32 - assert index_k.dtype == torch.float8_e4m3fn + if _is_fp8_fnuz: + assert index_k.dtype == torch.float8_e4m3fnuz + else: + assert index_k.dtype == torch.float8_e4m3fn assert index_k_scale.dtype == torch.float32 assert buf.is_contiguous() @@ -371,7 +378,10 @@ def _set_k_and_s_triton( assert index_k.is_contiguous() assert index_k_scale.is_contiguous() - buf_fp8 = buf.view(torch.float8_e4m3fn) + if _is_fp8_fnuz: + buf_fp8 = buf.view(torch.float8_e4m3fnuz) + else: + buf_fp8 = buf.view(torch.float8_e4m3fn) buf_fp32 = buf.view(torch.float32) _set_k_and_s_triton_kernel[(num_tokens_to_write,)]( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 0a3c5016b..116244339 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -8,6 +8,7 @@ import torch from einops import rearrange from sglang.srt.layers.layernorm import LayerNorm +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu @@ -15,6 +16,7 @@ global _use_multi_stream _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() +_is_fp8_fnuz = is_fp8_fnuz() if _is_cuda: try: import deep_gemm @@ -504,8 +506,10 @@ class Indexer(MultiPlatformOp): ) k_fp8_list.append(k_fp8) k_scale_list.append(k_scale) - - k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) + if _is_fp8_fnuz: + k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fnuz) + else: + k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) kv_fp8 = (k_fp8, k_scale) ks, ke = metadata.get_indexer_kvcache_range() @@ -569,9 +573,9 @@ class Indexer(MultiPlatformOp): from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits kv, scale = kv_fp8 - logits = fp8_mqa_logits( + logits_chunk = fp8_mqa_logits( q_fp8[start:end], - kv_fp8, + kv, scale, weights[start:end], ks[start:end], @@ -1001,11 +1005,12 @@ class Indexer(MultiPlatformOp): .view(m, ng, group) .mul_(x_s.to(torch.float32).unsqueeze(-1)) .view(m, n) + .to(torch.bfloat16) ) else: - x_for_gate = x_q.to(torch.float32) + x_for_gate = x_q.to(torch.bfloat16) else: - x_for_gate = x_q.to(torch.float32) + x_for_gate = x_q.to(torch.bfloat16) else: x_for_gate = x diff --git a/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py b/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py index c32d74fb9..1088bd3d1 100644 --- a/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py +++ b/python/sglang/srt/layers/attention/nsa/tilelang_kernel.py @@ -4,21 +4,28 @@ import tilelang import tilelang.language as T import torch -from sglang.srt.utils import is_hip +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.utils import is_gfx95_supported, is_hip tilelang.set_log_level("WARNING") pass_configs = { tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_FAST_MATH: True, } - -BF16 = "bfloat16" -FP8 = "float8_e4m3" -FP32 = "float32" +# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang +if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True +elif hasattr(tilelang.PassConfigKey, "TL_ENABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = False _is_hip = is_hip() +_is_gfx95_supported = is_gfx95_supported() +_is_fp8_fnuz = is_fp8_fnuz() + +BF16 = "bfloat16" +FP8 = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3" +FP32 = "float32" def fast_log2_ceil(x): @@ -42,8 +49,8 @@ def act_quant_kernel( N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False ): M = T.symbolic("M") - fp8_min = -448.0 - fp8_max = 448.0 + fp8_min = -224.0 if _is_fp8_fnuz else -448.0 + fp8_max = 224.0 if _is_fp8_fnuz else 448.0 fp8_max_inv = 1 / fp8_max num_stages = 0 if round_scale else 2 blk_m = 32 @@ -108,7 +115,10 @@ def act_quant( x.size(-1) % block_size == 0 ), f"Last dimension size must be divisible by block_size (block_size={block_size})" N = x.size(-1) - y = torch.empty_like(x, dtype=torch.float8_e4m3fn) + if _is_fp8_fnuz: + y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) + else: + y = torch.empty_like(x, dtype=torch.float8_e4m3fn) s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32) kernel = act_quant_kernel(N, round_scale=scale_fmt is not None) kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) @@ -777,9 +787,21 @@ def tilelang_sparse_fwd( topk = indices.shape[-1] assert topk == 2048 if _is_hip: - kernel = sparse_attention_fwd_kernel_v1( - num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, num_stages=1 - ) + if _is_gfx95_supported: + kernel = sparse_attention_fwd_kernel_v1( + num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, num_stages=1 + ) + else: # reduce LDS usage on gfx942 target + kernel = sparse_attention_fwd_kernel_v1( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=32, + num_stages=1, + threads=128, + ) else: kernel = sparse_attention_fwd_kernel_v2( num_heads, d_v, tail_dim, topk, sm_scale=sm_scale