[SM120] Avoid per-token-count JIT recompilation in swapped GEMM paths

build_sgl_deep_gemm.sh

@@ -0,0 +1,157 @@
+#!/usr/bin/env bash
+#
+# Build a wheel for the `sgl-deep-gemm` distribution.
+#
+# Distribution name: sgl-deep-gemm. Top-level import name: `deep_gemm`
+# (so existing call sites like `import deep_gemm` in sglang keep working).
+#
+# Build flow:
+#   1. Initialises submodules (cutlass, fmt) — same prerequisite as `bash build.sh`.
+#   2. Stages the package layout under build/deep_gemm/ with the Python
+#      sub-modules pulled from the source deep_gemm/ tree (utils, testing,
+#      legacy, mega).
+#   3. Reads the version string from sgl_deep_gemm/VERSION.
+#   4. Pre-compiles the tvm-ffi `_C.so` extension and bundles it into the wheel.
+#   5. Invokes `python -m build` to produce dist/*.whl.
+
+set -euo pipefail
+
+PYTHON_EXE=$(which python3 || which python)
+ROOT_DIR=$(realpath "$(dirname "$0")")
+BUILD_DIR="${ROOT_DIR}/build"
+PKG_DIR="${BUILD_DIR}/deep_gemm"
+DIST_DIR="${ROOT_DIR}/dist"
+
+cd "$ROOT_DIR"
+
+if [[ ! -f "setup.py" || ! -d "sgl_deep_gemm" || ! -d "deep_gemm" || ! -d "csrc" ]]; then
+    echo "Error: Run from the DeepGEMM project root." >&2
+    exit 1
+fi
+
+echo "--- Initialising submodules ---"
+git submodule update --init --recursive
+
+echo "--- Linking CUTLASS headers into deep_gemm/include ---"
+ln -sfn "${ROOT_DIR}/third-party/cutlass/include/cutlass" "${ROOT_DIR}/deep_gemm/include/cutlass"
+ln -sfn "${ROOT_DIR}/third-party/cutlass/include/cute" "${ROOT_DIR}/deep_gemm/include/cute"
+
+echo "--- Preparing build directory ---"
+rm -rf "$BUILD_DIR"
+mkdir -p "$PKG_DIR"
+
+cp sgl_deep_gemm/LICENSE sgl_deep_gemm/README.md sgl_deep_gemm/pyproject.toml "$BUILD_DIR/"
+cp sgl_deep_gemm/__init__.py sgl_deep_gemm/cuda_helpers.py "$PKG_DIR/"
+
+# `__init__.py` imports `.utils`, `.testing`, `.legacy`, `.mega` — pulled from
+# the existing deep_gemm/ tree.
+for sub in utils testing legacy mega; do
+    cp -r "deep_gemm/${sub}" "$PKG_DIR/"
+done
+
+# Headers required by the runtime JIT (same set the deep_gemm wheel ships).
+mkdir -p "$PKG_DIR/include"
+cp -r "${ROOT_DIR}/deep_gemm/include/deep_gemm" "$PKG_DIR/include/deep_gemm"
+cp -r "${ROOT_DIR}/third-party/cutlass/include/cute" "$PKG_DIR/include/cute"
+cp -r "${ROOT_DIR}/third-party/cutlass/include/cutlass" "$PKG_DIR/include/cutlass"
+
+echo "--- Reading version from sgl_deep_gemm/VERSION ---"
+if [[ ! -f "sgl_deep_gemm/VERSION" ]]; then
+    echo "Error: sgl_deep_gemm/VERSION is missing — create it with the desired version (e.g. 0.0.1)." >&2
+    exit 1
+fi
+# Strip surrounding whitespace; the file is the single source of truth.
+tr -d '[:space:]' < sgl_deep_gemm/VERSION > "$PKG_DIR/VERSION"
+echo "Version: $(cat "$PKG_DIR/VERSION")"
+
+echo "--- Compiling _C.so ---"
+ROOT_DIR="$ROOT_DIR" PKG_DIR="$PKG_DIR" "$PYTHON_EXE" -u - <<'PY'
+import importlib.util
+import os, shutil, subprocess, sys, sysconfig
+root_dir = os.environ['ROOT_DIR']
+pkg_dir = os.environ['PKG_DIR']
+sys.path.insert(0, root_dir)
+
+# tvm_ffi.cpp.build runs ninja with capture_output=True, hiding compile logs
+# until a failure. Patch subprocess.run so the ninja invocation streams to the
+# terminal — leaves other internal calls (nvidia-smi, nvcc --version) alone.
+_orig_run = subprocess.run
+def _streamed_run(*args, **kwargs):
+    cmd = kwargs.get('args') if 'args' in kwargs else (args[0] if args else None)
+    is_ninja = isinstance(cmd, (list, tuple)) and cmd and 'ninja' in str(cmd[0])
+    if is_ninja:
+        kwargs.pop('capture_output', None)
+        kwargs['stdout'] = None
+        kwargs['stderr'] = None
+    return _orig_run(*args, **kwargs)
+subprocess.run = _streamed_run
+
+import torch
+import tvm_ffi.cpp
+
+helpers_path = os.path.join(root_dir, 'sgl_deep_gemm', 'cuda_helpers.py')
+helpers_spec = importlib.util.spec_from_file_location('sgl_deep_gemm_cuda_helpers', helpers_path)
+cuda_helpers = importlib.util.module_from_spec(helpers_spec)
+assert helpers_spec.loader is not None
+helpers_spec.loader.exec_module(cuda_helpers)
+
+cuda_home = cuda_helpers.find_cuda_home()
+os.environ.setdefault('TVM_FFI_CUDA_ARCH_LIST', cuda_helpers.get_cuda_arch())
+
+cxx_abi = int(torch.compiled_with_cxx11_abi())
+extra_cflags = [
+    '-std=c++17', '-O3', '-fPIC',
+    '-Wno-psabi', '-Wno-deprecated-declarations',
+    f'-D_GLIBCXX_USE_CXX11_ABI={cxx_abi}',
+]
+if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')):
+    extra_cflags.append('-DDG_JIT_USE_RUNTIME_API')
+
+torch_dir = os.path.dirname(torch.__file__)
+extra_include_paths = [
+    f'{cuda_home}/include',
+    sysconfig.get_path('include'),
+    os.path.join(torch_dir, 'include'),
+    os.path.join(torch_dir, 'include', 'torch', 'csrc', 'api', 'include'),
+    os.path.join(root_dir, 'deep_gemm', 'include'),
+    os.path.join(root_dir, 'third-party', 'cutlass', 'include'),
+    os.path.join(root_dir, 'third-party', 'fmt', 'include'),
+]
+cccl = f'{cuda_home}/include/cccl'
+if os.path.exists(cccl):
+    extra_include_paths.append(cccl)
+
+extra_ldflags = [
+    f'-L{cuda_home}/lib64',
+    f'-L{os.path.join(torch_dir, "lib")}',
+    '-lcudart', '-lnvrtc', '-lcublasLt', '-lcublas',
+    '-ltorch', '-ltorch_cpu', '-lc10', '-lc10_cuda', '-ltorch_cuda',
+]
+
+build_subdir = os.path.join(pkg_dir, '_C_build')
+os.makedirs(build_subdir, exist_ok=True)
+lib_path = tvm_ffi.cpp.build(
+    name='_C',
+    cpp_files=[os.path.join(root_dir, 'csrc', 'tvm_ffi_api.cpp')],
+    extra_cflags=extra_cflags,
+    extra_ldflags=extra_ldflags,
+    extra_include_paths=extra_include_paths,
+    build_directory=build_subdir,
+)
+target = os.path.join(pkg_dir, '_C.so')
+if os.path.exists(target):
+    os.remove(target)
+shutil.copy2(lib_path, target)
+shutil.rmtree(build_subdir, ignore_errors=True)
+print(f"Built {target}")
+PY
+
+echo "--- Installing build frontend ---"
+"$PYTHON_EXE" -m pip install --quiet --upgrade build
+
+echo "--- Building wheel ---"
+mkdir -p "$DIST_DIR"
+"$PYTHON_EXE" -m build --wheel "$BUILD_DIR" --outdir "$DIST_DIR"
+
+echo "--- Done ---"
+ls -lh "$DIST_DIR"/sgl_deep_gemm-*.whl 2>/dev/null || ls -lh "$DIST_DIR"/sgl-deep-gemm-*.whl 2>/dev/null || ls -lh "$DIST_DIR"/sgl_deep_gemm*.whl

csrc/apis/attention.hpp

@@ -6,7 +6,10 @@
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
 #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp"
+#include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp"
 #include "../jit_kernels/impls/sm100_mqa_logits.hpp"
+#include "../jit_kernels/impls/sm120_mqa_logits.hpp"
+#include "../jit_kernels/impls/sm120_paged_mqa_logits.hpp"
 #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp"
 #include "../jit_kernels/impls/smxx_clean_logits.hpp"
 #endif
@@ -68,6 +71,9 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
         // NOTES: Only granularity 128 and FP8 are exposed in the API
         sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k,
                                 128, 128, major_a, major_b, compiled_dims, epilogue_type);
+    } else if (arch_major == 12 and sfa.scalar_type() == torch::kInt) {
+        sm120_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k,
+                                128, 128, major_a, major_b, compiled_dims, epilogue_type);
     } else {
         DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
     }
@@ -91,9 +97,10 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt
         // Check FP4 Q
         std::tie(seq_len, num_heads, head_dim) = get_shape<3>(q_fp);
         head_dim *= 2;
-        DG_HOST_ASSERT(arch_major == 10);
+        DG_HOST_ASSERT(arch_major == 10 or arch_major == 12);
         DG_HOST_ASSERT(num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64);
         DG_HOST_ASSERT(head_dim == 64 or head_dim == 128);
+        DG_HOST_ASSERT(arch_major != 12 or head_dim == 128);  // SM120 FP4 MQA logits: head_dim=128 only
         DG_HOST_ASSERT(q_fp.is_contiguous());
         DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4);
 
@@ -120,6 +127,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt
         // Check FP8 Q
         std::tie(seq_len, num_heads, head_dim) = get_shape<3>(q_fp);
         DG_HOST_ASSERT((arch_major == 10 and (num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64)) or
+                       (arch_major == 12 and (num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64)) or
                        (arch_major == 9 and (num_heads == 32 or num_heads == 64)));
         DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
         DG_HOST_ASSERT(q_fp.is_contiguous());
@@ -158,7 +166,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt
 
     // Allocate output
     constexpr int block_qh = 128;
-    constexpr int block_kv = 256;
+    const int block_kv = (arch_major == 12) ? 128 : 256;  // SM120 dense MQA uses block_kv=128
     const int block_q = block_qh / num_heads;
     DG_HOST_ASSERT(block_qh % num_heads == 0);
 
@@ -181,6 +189,9 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt
     if (arch_major == 10) {
         sm100_mqa_logits(is_fp4, q_fp, q_sf, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype,
                          seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv);
+    } else if (arch_major == 12) {
+        sm120_mqa_logits(is_fp4, q_fp, q_sf, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype,
+                         seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv);
     } else if (arch_major == 9 and not is_fp4) {
         DG_HOST_ASSERT(weights.scalar_type() == torch::kFloat);
         sm90_fp8_mqa_logits(q_fp, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype,
@@ -220,6 +231,11 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_
     } else if (arch_major == 10) {
         DG_HOST_ASSERT(block_kv == 64 or block_kv == 32);
         sm100_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, batch_size * next_n, next_n, num_sms, is_context_lens_2d, false, nullptr);
+    } else if (arch_major == 12) {
+        DG_HOST_ASSERT(block_kv == 64);
+        const int next_n_atom = (next_n >= 2) ? 2 : 1;
+        const int num_next_n_atoms = (next_n + next_n_atom - 1) / next_n_atom;
+        sm120_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, num_next_n_atoms, false, nullptr);
     } else if (arch_major == 9) {
         DG_HOST_ASSERT(block_kv == 64);
         sm90_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, false, nullptr);
@@ -256,9 +272,10 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
         std::tie(batch_size, next_n, num_heads, head_dim) = get_shape<4>(q_fp);
         head_dim *= 2;
         DG_HOST_ASSERT(next_n >= 1);
-        DG_HOST_ASSERT(arch_major == 10);
+        DG_HOST_ASSERT(arch_major == 10 or arch_major == 12);
         DG_HOST_ASSERT(num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64);
         DG_HOST_ASSERT(head_dim == 64 or head_dim == 128);
+        DG_HOST_ASSERT(arch_major != 12 or head_dim == 128);  // SM120 FP4 paged: head_dim=128 only
         DG_HOST_ASSERT(q_fp.is_contiguous());
         DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4);
 
@@ -272,6 +289,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
         int num_heads_kv, fp4_with_sf_bytes;
         std::tie(num_kv_blocks, block_kv, num_heads_kv, fp4_with_sf_bytes) = get_shape<4>(fused_kv_cache);
         DG_HOST_ASSERT((arch_major == 10 and (block_kv == 32 or block_kv == 64)) or
+                       (arch_major == 12 and block_kv == 64) or
                        (arch_major == 9 and block_kv == 64));
         DG_HOST_ASSERT(num_heads_kv == 1 and fp4_with_sf_bytes == head_dim / 2 + static_cast<int>(sizeof(int)));
         DG_HOST_ASSERT(fused_kv_cache.stride(1) == fp4_with_sf_bytes and fused_kv_cache.stride(3) == 1);
@@ -297,6 +315,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
         std::tie(batch_size, next_n, num_heads, head_dim) = get_shape<4>(q_fp);
         DG_HOST_ASSERT(next_n >= 1);
         DG_HOST_ASSERT((arch_major == 10 and (num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64)) or
+                       (arch_major == 12 and (num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64)) or
                        (arch_major == 9 and (num_heads == 32 or num_heads == 64)));
         DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
         DG_HOST_ASSERT(q_fp.is_contiguous());
@@ -306,6 +325,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
         int num_heads_kv, head_dim_with_sf;
         std::tie(num_kv_blocks, block_kv, num_heads_kv, head_dim_with_sf) = get_shape<4>(fused_kv_cache);
         DG_HOST_ASSERT((arch_major == 10 and (block_kv == 32 or block_kv == 64)) or
+                       (arch_major == 12 and block_kv == 64) or
                        (arch_major == 9 and block_kv == 64));
         DG_HOST_ASSERT(num_heads_kv == 1 and head_dim_with_sf == head_dim + static_cast<int>(sizeof(float)));
         DG_HOST_ASSERT(fused_kv_cache.stride(1) == head_dim_with_sf and fused_kv_cache.stride(3) == 1);
@@ -371,7 +391,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
 
     // Allocate output
     DG_HOST_ASSERT(logits_dtype == torch::kFloat32 or logits_dtype == torch::kBFloat16);
-    constexpr int split_kv = 256;
+    const int split_kv = (arch_major == 12) ? 128 : 256;  // SM120: 2 groups × 64
     // Logits row stride must be 1024-byte aligned
     const int stride_logits_alignment = 1024 / static_cast<int>(c10::elementSize(logits_dtype));
     const auto aligned_max_context_len = align(align(max_context_len, split_kv), stride_logits_alignment);
@@ -384,6 +404,11 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
         sm100_paged_mqa_logits(is_fp4, q_fp, q_sf, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, schedule_meta,
                                logits_dtype, batch_size, batch_size * next_n, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
                                is_varlen, aligned_max_context_len, block_table_stride, num_sms, split_kv, splits_per_chunk);
+    } else if (arch_major == 12) {
+        constexpr int splits_per_chunk = 16;  // unused by SM120 atom scheduler
+        sm120_paged_mqa_logits(is_fp4, q_fp, q_sf, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, schedule_meta,
+                               logits_dtype, batch_size, batch_size * next_n, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
+                               is_varlen, aligned_max_context_len, block_table_stride, num_sms, split_kv, splits_per_chunk);
     } else if (arch_major == 9 and not is_fp4) {
         DG_HOST_ASSERT(weights.scalar_type() == torch::kFloat);
         sm90_fp8_paged_mqa_logits(q_fp, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, schedule_meta,
@@ -430,6 +455,8 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
 }
 #endif
 
+#if 0
+
 static void register_apis(pybind11::module_& m) {
 #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
     m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
@@ -467,4 +494,6 @@ static void register_apis(pybind11::module_& m) {
 #endif
 }
 
+#endif
+
 } // namespace deep_gemm::attention