[AMD] Support ds3.2 on gfx942 platform (#17504)

Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>

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
+  # Build TileLang for gfx950 and gfx942-rocm700
-  if [ "${GPU_ARCH:-}" != "gfx950" ]; then \
+  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}"; \
-  \
+  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 \
+    # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing)
-      build-essential git wget curl ca-certificates gnupg \
+    apt-get update && apt-get install -y --no-install-recommends \
-      libgtest-dev libgmock-dev \
+        build-essential git wget curl ca-certificates gnupg \
-      libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
+        libgtest-dev libgmock-dev \
-      python3 python3-dev python3-setuptools python3-pip \
+        libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
-      gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
+        python3 python3-dev python3-setuptools python3-pip \
-      cmake ninja-build pkg-config libstdc++6 \
+        gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
-  && rm -rf /var/lib/apt/lists/*; \
+        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 && \
+    # Build GoogleTest static libs (Ubuntu package ships sources only)
-  cmake --build /tmp/build-gtest -j"$(nproc)" && \
+    cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
-  cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
+    cmake --build /tmp/build-gtest -j"$(nproc)" && \
-  rm -rf /tmp/build-gtest; \
+    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 && \
+    # Keep setuptools < 80 (compat with base image)
-  python3 -m pip cache purge || true; \
+    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=""; \
+    # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing
-  for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
+    LLVM_CONFIG_PATH=""; \
-    if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
+    for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
-  done; \
+      if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
-  if [ -z "$LLVM_CONFIG_PATH" ]; then \
+    done; \
-    echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
+    if [ -z "$LLVM_CONFIG_PATH" ]; then \
-    curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \
+      echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
-    chmod +x /tmp/llvm.sh; \
+      curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \
-    /tmp/llvm.sh 18; \
+      chmod +x /tmp/llvm.sh; \
-    LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
+      /tmp/llvm.sh 18; \
-    if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
+      LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
-  fi; \
+      if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
-  echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
+    fi; \
-  export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \
+    echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
-  export LLVM_CONFIG="$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 && \
+    # Optional shim for tools that expect llvm-config-16
-  printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \
+    mkdir -p /usr/local/bin && \
-  chmod +x /usr/local/bin/llvm-config-16; \
+    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"; \
+    # 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 && \
+    # Clone + pin TileLang (bundled TVM), then build
-  cd /opt/tilelang && \
+    git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \
-  git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
+    cd /opt/tilelang && \
-  git checkout -f "${TILELANG_COMMIT}" && \
+    git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
-  git submodule update --init --recursive && \
+    git checkout -f "${TILELANG_COMMIT}" && \
-  export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \
+    git submodule update --init --recursive && \
-  bash ./install_rocm.sh'
+    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)

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,)](

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)
-
+        if _is_fp8_fnuz:
-        k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
+            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
 

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,
 }
-
+# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang
-BF16 = "bfloat16"
+if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"):
-FP8 = "float8_e4m3"
+    pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True
-FP32 = "float32"
+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_min = -224.0 if _is_fp8_fnuz else -448.0
-    fp8_max = 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(
+        if _is_gfx95_supported:
-            num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, num_stages=1
+            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