diff --git a/.dockerignore b/.dockerignore deleted file mode 120000 index 3e4e48b0b5fe..000000000000 --- a/.dockerignore +++ /dev/null @@ -1 +0,0 @@ -.gitignore \ No newline at end of file diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 000000000000..f9a8a460bd09 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,53 @@ +# Build-context allowlist. +# +# This file used to be a symlink to .gitignore. That let the entire repository +# into the build context, and `COPY . /sgl-workspace/sglang` then shipped it all +# inside a distributable image: the test suite, internal benchmarks, the docs +# site, CI workflows, the model gateway, and the agent skill directories under +# .claude -- including python/sglang/multimodal_gen/.claude, which sits inside +# the runtime tree and so survives any top-level filtering. +# +# None of that has a runtime role, and all of it is proprietary: anyone able to +# pull the image could read it. The context is therefore an explicit allowlist of +# what the image actually needs, rather than a denylist inherited from git. +# +# .gitignore and .dockerignore answer different questions -- "what should not be +# committed" versus "what should not leave the building" -- and sharing one file +# between them is what let this happen. + +# Exclude everything, then re-admit only what is needed. +* + +# The runtime package: what `pip install -e python` installs. +!python/ +# The Dockerfiles run cuda_pins.sh and apply_deepep_v2_patch.sh out of the copied +# tree, not only from /opt/qwen38, so the build needs this directory. +!docker/ +# Required at BUILD time, not run time: python/setup.py points setuptools-rust at +# ../rust (lines 3 and 178), so `pip install -e python` fails without it -- +# "no cargo workspace at rust/Cargo.toml". Excluding it would force +# SGLANG_BUILD_RUST_EXTS=none, i.e. an image with no Rust extensions at all, +# which is a capability change disguised as a size reduction. The Dockerfile +# still deletes rust/*/target afterwards, so no build output ships. +!rust/ +# Both are tiny. Dropping them on a guess that nothing reads them would risk +# breaking gRPC codegen or a vendored build for no meaningful reduction. +!proto/ +!3rdparty/ +# Apache-2.0 requires the licence to accompany distributions of a derived work, +# and this image is one. It is not runtime code, but it does have to ship. +!LICENSE + +# Re-exclude what must not travel even inside the allowlisted trees. These come +# after the allowlist entries on purpose: the last matching pattern wins. +**/.claude/ + +# Build artifacts. Previously covered by .gitignore through the symlink; restated +# so a `docker build` from a dirty checkout cannot smuggle them in either. +**/__pycache__/ +**/*.py[cod] +**/*.so +**/build/ +**/dist/ +**/*.egg-info/ +.git/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a53a6af395f5..45893c61c7d9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_stages: [pre-commit, pre-push, manual] -exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/kernels/ops/diffusion/render|python/sglang/kernels/ops/attention/flash_attn/cute) +exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/kernels/ops/diffusion/render|python/sglang/kernels/ops/attention/flash_attn/cute|python/sglang/srt/layers/flashinfer_fallback/comm/(mnnvl_cutedsl/|mnnvl_cutedsl_ar\.py)) repos: - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/docker/qwen38/apply_deepep_v2_patch.sh b/docker/qwen38/apply_deepep_v2_patch.sh new file mode 100755 index 000000000000..b5fef974270e --- /dev/null +++ b/docker/qwen38/apply_deepep_v2_patch.sh @@ -0,0 +1,120 @@ +#!/bin/bash +# Qwen38 DeepEP v2 patch + build -- replaces the base image's stock DeepEP. +# +# The base image ships DeepEP at the legacy pin, which does not serve GB300 on +# CUDA 13. This installs the v2 tree instead: +# - v2 moved to a CMake/JIT layout, so the cross-node timeout knobs now live +# in csrc/kernels/legacy/compiled.cuh behind LEGACY_-prefixed macros +# - the wheel links the nvidia-nvshmem/nccl wheels rather than a system NCCL +# - CUDA 13 relocated the cccl headers; the legacy setuptools path needs an +# extra include dir, the CMake path resolves them itself +# +# Both timeout bumps are asserted after the fact because `sed -i` exits 0 on no +# match: shipping the stock 100s CPU timeout is invisible at build time and +# surfaces only as GB300 multi-node init aborts in production. +# +# Env knobs (all optional): +# DEEPEP_V2_COMMIT pinned commit to build +# TORCH_CUDA_ARCH_LIST arches to compile cubins for +# MAX_JOBS nvcc build parallelism +# DEEPEP_DIR where the source tree lands in the image +set -euo pipefail + +: "${DEEPEP_V2_COMMIT:=01dc3aaac82068020353dce2c302e38153c0bfaa}" +: "${TORCH_CUDA_ARCH_LIST:=9.0;10.0;10.3}" +: "${MAX_JOBS:=8}" +: "${DEEPEP_DIR:=/sgl-workspace/DeepEP}" +: "${CUDA_HOME:=/usr/local/cuda}" + +BUILD_DIR=/build/DeepEP +TIMEOUT_HEADER=csrc/kernels/legacy/compiled.cuh + +rm -rf "${BUILD_DIR}" +git clone https://github.com/deepseek-ai/DeepEP.git "${BUILD_DIR}" +cd "${BUILD_DIR}" +git checkout "${DEEPEP_V2_COMMIT}" + +# --- Cross-node timeout headroom --- +sed -i \ + 's/#define LEGACY_NUM_CPU_TIMEOUT_SECS 100/#define LEGACY_NUM_CPU_TIMEOUT_SECS 1000/' \ + "${TIMEOUT_HEADER}" +sed -i \ + 's/#define LEGACY_NUM_TIMEOUT_CYCLES 200000000000ull/#define LEGACY_NUM_TIMEOUT_CYCLES 2000000000000ull/' \ + "${TIMEOUT_HEADER}" + +if ! grep -q '#define LEGACY_NUM_CPU_TIMEOUT_SECS 1000' "${TIMEOUT_HEADER}"; then + echo "ERROR: CPU timeout bump did not apply; the macro moved or was reformatted" >&2 + exit 1 +fi +if ! grep -q '#define LEGACY_NUM_TIMEOUT_CYCLES 2000000000000ull' "${TIMEOUT_HEADER}"; then + echo "ERROR: cycle timeout bump did not apply; the macro moved or was reformatted" >&2 + exit 1 +fi + +# --- CUDA 13 cccl include dir (legacy setuptools path only) --- +# A missing anchor is the expected outcome on a CMake tree, so report which +# path was taken rather than letting sed no-op silently. +if grep -q "^ include_dirs = \['csrc/'\]" setup.py; then + sed -i \ + "/^ include_dirs = \['csrc\/'\]/a\\ include_dirs.append('${CUDA_HOME}/include/cccl')" \ + setup.py + echo "Applied the CUDA 13 cccl include fix to setup.py" +else + echo "setup.py has no legacy include_dirs anchor; relying on the CMake build to resolve cccl" +fi + +# --- Build and swap in --- +TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" MAX_JOBS="${MAX_JOBS}" \ + python3 setup.py bdist_wheel -d /wheels + +# Remove whatever currently owns the deep_ep module, by asking the metadata +# rather than by name. The base image's provider has changed name: v0.5.17 +# shipped a distribution literally called `deep_ep`, while the dev images ship +# `sgl-deep-ep`, which installs the same deep_ep/ path (plus its own +# deep_ep_cpp*.so). `pip uninstall deep_ep` exits 0 with only a warning when no +# such distribution exists, so hardcoding either name leaves the other in place +# and two distributions then contend for deep_ep/ -- whichever wrote the files +# last wins, which is not a decision this build should leave to chance. +# +# packages_distributions() reads metadata and does not import deep_ep, so the +# packaged provider's driver check cannot fire here. +# +# Run from / rather than the build tree: `setup.py bdist_wheel` above leaves a +# deep_ep.egg-info in ${BUILD_DIR}, and `python3 -` puts the cwd first on +# sys.path, so importlib.metadata would report that build artifact as an +# installed distribution named deep_ep on top of the real one. +mapfile -t DEEP_EP_DISTS < <(cd / && python3 - <<'PY' +from importlib.metadata import packages_distributions + +for dist in sorted(set(packages_distributions().get("deep_ep", ()))): + print(dist) +PY +) +if [ "${#DEEP_EP_DISTS[@]}" -gt 0 ]; then + echo "Removing existing deep_ep provider(s): ${DEEP_EP_DISTS[*]}" + python3 -m pip uninstall -y "${DEEP_EP_DISTS[@]}" +else + echo "No installed distribution provides deep_ep; nothing to remove" +fi + +python3 -m pip install /wheels/*.whl + +# Exactly one provider must remain. Two would mean the removal above missed a +# name and the source build is now sharing deep_ep/ with a packaged copy. +# From / for the same reason as above -- the egg-info left in the build tree is a +# build artifact, not an installed distribution, and counting it here reported +# ['deep_ep', 'deep_ep'] and failed a build that was in fact correct. +( cd / && python3 - <<'PY' +from importlib.metadata import packages_distributions + +providers = sorted(set(packages_distributions().get("deep_ep", ()))) +if len(providers) != 1: + raise SystemExit( + f"expected exactly one distribution to provide deep_ep, found {providers}" + ) +print(f"deep_ep is provided by {providers[0]} alone") +PY +) + +rm -rf "${DEEPEP_DIR}" /wheels "${BUILD_DIR}/build" "${BUILD_DIR}/dist" /root/.cache/pip +mv "${BUILD_DIR}" "${DEEPEP_DIR}" diff --git a/docker/qwen38/cuda_pins.sh b/docker/qwen38/cuda_pins.sh new file mode 100755 index 000000000000..1930217db32c --- /dev/null +++ b/docker/qwen38/cuda_pins.sh @@ -0,0 +1,236 @@ +#!/usr/bin/env bash +# Keep the image's dependency reality in step with what this tree declares. +# +# The image installs SGLang with `--no-deps` on purpose: the base ships +# CUDA-tagged wheels (sglang-kernel X+cu129, sgl-deep-gemm Y+cu129) and a plain +# `pip install` would swap them for the untagged PyPI builds, which are compiled +# against a different libtorch and die at import with an undefined symbol. +# +# What `--no-deps` costs is the other direction: a pin this tree *raises* is +# silently dropped. The image keeps whatever the base had, builds green, gets +# pushed, gets pulled onto every node, and only fails when a server starts and +# srt/entrypoints/engine.py runs assert_pkg_version. These subcommands move each +# of those failures to build time, where one line of output is the whole answer. +# +# check-torch +# Fail if the base's torch is not the one this tree pins. Nothing below +# can be satisfied when this is wrong: the CUDA-tagged kernel wheels are +# built against a specific libtorch, and they declare no Requires-Dist, +# so pip cannot see the conflict and will happily install a broken combo. +# Depends only on the base image, so callers should run it as early as +# possible -- a wrong base is knowable in two seconds, and every expensive +# layer after it is wasted work. +# +# reconcile ... +# Install the + build of each pkg at the version +# pins, from SGLang's own wheel index, then import it -- an ABI mismatch +# shows up here rather than on the first forward pass. The import step is +# skipped (loudly) on build hosts with no NVIDIA driver; see below. +# +# verify ... +# Run the same assert_pkg_version call the server makes at launch. +# +# import-if-gpu ... +# Import each module, but only where a CUDA device is reachable. For +# modules that refuse to load without a driver this is the difference +# between a check and a guaranteed build failure -- see have_cuda_device. +set -euo pipefail + +# Point at the GitHub Pages origin directly. The docs.sglang.ai alias works but +# costs two cross-host redirects (docs.sglang.ai -> docs.sglang.io -> +# sgl-project.github.io) on every wheel fetch. +SGL_WHL_INDEX="${SGL_WHL_INDEX:-https://sgl-project.github.io/whl}" + +# Echo the version pins for a dependency, or fail if it is unpinned. +pinned_version() { + python3 - "$1" "$2" <<'PY' +import pathlib, sys, tomllib + +data = tomllib.loads(pathlib.Path(sys.argv[1]).read_text()) +want = sys.argv[2] +for spec in data["project"]["dependencies"]: + spec = spec.replace(" ", "") + if spec.startswith(want + "=="): + print(spec.split("==", 1)[1]) + break +else: + raise SystemExit(f"{want} is not pinned in {sys.argv[1]}") +PY +} + +# Distribution name -> import name. Deliberately an explicit table rather than a +# derivation: for these wheels the two names do not correspond, and a wrong +# guess surfaces as ModuleNotFoundError that reads exactly like the ABI failure +# this script exists to catch. Verified by listing the wheels' own contents: +# sglang-kernel ships sgl_kernel/ +# sgl-deep-gemm ships deep_gemm/ <- NOT sgl_deep_gemm +module_for() { + case "$1" in + sglang-kernel) echo sgl_kernel ;; + sgl-deep-gemm) echo deep_gemm ;; + *) + cat >&2 < 0 else 1) +except Exception: + sys.exit(1) +PY +} + +cmd_check_torch() { + local pyproject="$1" + local want got + want="$(pinned_version "$pyproject" "torch")" + got="$(python3 -c 'import torch; print(torch.__version__.split("+")[0])')" + if [ "$want" != "$got" ]; then + cat >&2 <=-on-release comparison assert_pkg_version applies. The local segment +# (+cu129) is deliberately ignored: it records which CUDA a wheel was built for, +# not a version ordering, and the base's own choice of build is the one to keep. +already_satisfied() { + PKG="$1" WANT="$2" python3 - <<'PY' +import os, sys +from importlib.metadata import PackageNotFoundError, version + +from packaging.version import Version + +try: + got = version(os.environ["PKG"]) +except PackageNotFoundError: + sys.exit(1) +sys.exit(0 if Version(got) >= Version(os.environ["WANT"]) else 1) +PY +} + +cmd_reconcile() { + local pyproject="$1" cuda_tag="$2" + shift 2 + local pkg ver mod + for pkg in "$@"; do + ver="$(pinned_version "$pyproject" "$pkg")" + mod="$(module_for "$pkg")" + # Reinstalling what the base already satisfies is not free: the base may + # have installed an untagged build on purpose (upstream's CUDA 13 path + # does exactly that), and replacing it with the + build swaps a + # known-good artifact for a different one. Only act when the pin is + # genuinely higher than what is installed. + if already_satisfied "$pkg" "$ver"; then + echo "[cuda_pins] ${pkg}: already satisfies the ${ver} pin" \ + "(installed $(python3 -c "from importlib.metadata import version; print(version('${pkg}'))"))," \ + "leaving the base's build in place" + continue + fi + echo "[cuda_pins] reconciling ${pkg}==${ver}+${cuda_tag}" + python3 -m pip install --no-deps "${pkg}==${ver}+${cuda_tag}" \ + --index-url "${SGL_WHL_INDEX}/${cuda_tag}" + # Import now: a wheel built against another libtorch installs cleanly + # and only fails when it is first loaded. + if ! have_cuda_device; then + echo "[cuda_pins] ${pkg}: SKIPPED the import check -- no reachable CUDA device on" \ + "this build host, so no CUDA extension can load here regardless of correctness." + echo "[cuda_pins] ${pkg}: 'import ${mod}' must be covered by a GPU smoke test." \ + "The build has verified the installed version only." + continue + fi + python3 -c "import ${mod}" || { + echo "ERROR: ${pkg}==${ver}+${cuda_tag} installed but 'import ${mod}' fails" >&2 + echo " against this base's torch -- the two were built apart." >&2 + exit 1 + } + done +} + +cmd_import_if_gpu() { + local mod + if ! have_cuda_device; then + echo "[cuda_pins] SKIPPED importing $*: no reachable CUDA device on this build host." + echo "[cuda_pins] These imports must be covered by a GPU smoke test." + return 0 + fi + for mod in "$@"; do + python3 -c "import ${mod}; print('[cuda_pins] import ${mod} OK')" + done +} + +cmd_verify() { + local pyproject="$1" + shift + python3 - "$pyproject" "$@" <<'PY' +import pathlib, sys, tomllib +from importlib.metadata import version + +from sglang.srt.utils.common import assert_pkg_version + +data = tomllib.loads(pathlib.Path(sys.argv[1]).read_text()) +pins = {} +for spec in data["project"]["dependencies"]: + spec = spec.replace(" ", "") + if "==" in spec: + name, ver = spec.split("==", 1) + pins[name] = ver + +for pkg in sys.argv[2:]: + if pkg not in pins: + raise SystemExit(f"{pkg} is not pinned in {sys.argv[1]}") + # The same call srt/entrypoints/engine.py makes when a server starts. Note + # that engine.py hardcodes its minimum rather than reading pyproject, so the + # two can drift apart; this asserts the tree's own pin, which is the stricter + # reading of what the source expects. + assert_pkg_version(pkg, pins[pkg], "image build did not reconcile this pin") + print(f"[cuda_pins] {pkg} OK: pinned {pins[pkg]}, installed {version(pkg)}") +PY +} + +case "${1:?usage: $0 check-torch|reconcile|verify|import-if-gpu ...}" in + check-torch) shift; cmd_check_torch "$@" ;; + reconcile) shift; cmd_reconcile "$@" ;; + verify) shift; cmd_verify "$@" ;; + import-if-gpu) shift; cmd_import_if_gpu "$@" ;; + *) echo "unknown subcommand: $1" >&2; exit 2 ;; +esac diff --git a/docker/qwen38/qwen38_cu12.Dockerfile b/docker/qwen38/qwen38_cu12.Dockerfile new file mode 100644 index 000000000000..72ab12d2474c --- /dev/null +++ b/docker/qwen38/qwen38_cu12.Dockerfile @@ -0,0 +1,237 @@ +# Qwen38 serving image (x86_64 / CUDA 12.9 / sm_90a + sm_100a). +# +# Base ships stock SGLang (editable at /sgl-workspace/sglang), DeepEP source +# (at /sgl-workspace/DeepEP), a released FlashInfer trio, the Rust toolchain at +# /root/.cargo, and the CUDA 12.9 toolchain. +# +# This image adds the two Qwen38-specific pieces that stock lacks: +# 1. FlashInfer, replacing the base's release: flashinfer-python from a pinned +# git commit (PR #4358), with the nightly's prebuilt flashinfer-cubin and +# flashinfer-jit-cache alongside it, since neither is published per commit. +# nvidia-cutlass-dsl is floored at 4.7.0 because the CuTe DSL kernels call +# PipelineTmaAsync.create(enable_multicast_signaling=..) +# 2. this repo's SGLang code, copied from the build context and editable-installed +# +# Unlike the CUDA 13 recipe this image keeps the base's DeepEP: the DeepEP v2 +# source build and its NCCL 2.30.7 preload exist for GB300 on CUDA 13 only. +# +# The build refuses to start unless the base's torch is the one +# python/pyproject.toml pins, reinstalls the CUDA-tagged wheels this tree pins +# above what the base ships, and asserts those pins before finishing. All three +# live in docker/qwen38/cuda_pins.sh, which explains why each is needed. +# +# Build (no GPU needed). The context must be the repo root -- the SGLang source +# is taken from it rather than cloned: +# docker build -f docker/qwen38/qwen38_cu12.Dockerfile -t qwen38-cu129 . + +# The nightly, not a release tag: v0.5.17-cu129 ships torch 2.11.0 / +# sglang-kernel 0.4.5, which the gate below rejects against this tree's pins. +# `dev-cu12` is the CUDA 12.9 nightly and carries torch 2.13.0+cu129, +# sglang-kernel 0.4.6.post1+cu129 and sgl-deep-gemm 0.1.5.post2+cu129 -- exactly +# what python/pyproject.toml asks for, already in the +cu129 builds. +# +# `dev-cu12` MOVES every night. Builds are therefore not reproducible from this +# tag alone; pin `dev-cu12@sha256:` when a build has to be repeatable. +FROM lmsysorg/sglang:dev-cu12 AS base + +# --- 0. Base/tree compatibility gate, FIRST --- +# Whether this base can host this tree is decided entirely by the base's torch, +# and it is knowable before any work happens. Run it here so a wrong base costs +# two seconds instead of a full FlashInfer nightly download and DeepEP link. +# +# pyproject.toml is copied to its own path purely to make that possible: the tree +# itself only lands at COPY below, far too late to gate anything. Both copies +# come from the same build context, so they cannot disagree. +COPY docker/qwen38/cuda_pins.sh /opt/qwen38/cuda_pins.sh +COPY python/pyproject.toml /opt/qwen38/pyproject.toml + +RUN bash /opt/qwen38/cuda_pins.sh check-torch /opt/qwen38/pyproject.toml + +# --- 1. FlashInfer: nightly cubin/jit-cache, python from a pinned commit --- +# Named apart from the base's ENV FLASHINFER_VERSION, which would otherwise +# shadow a same-named ARG and silently resolve to the base's 0.6.15.post1. +# +# The nightly version below no longer covers flashinfer-python: that comes from +# FLASHINFER_GIT_COMMIT instead. cubin and jit-cache stay on the nightly because +# they are prebuilt artifacts published per nightly date and per release only -- +# no wheel of either exists for an arbitrary commit. The nightly named here must +# therefore be the one whose main the pinned commit merges, so the prebuilt +# kernels correspond to the source they were compiled from. +ARG FLASHINFER_NIGHTLY_VERSION=0.6.18.dev20260807 +ARG FLASHINFER_JIT_CACHE_CUDA_TAG=cu129 +ARG CUTLASS_DSL_MIN_VERSION=4.7.0 + +# flashinfer-ai/flashinfer#4358, "feat(comm): add Blackwell MNNVL CuTe DSL +# all-reduce fusion backend". +# Pinned to the PR's merge commit on main, which is what keeps it fetchable. An +# earlier pin of a PR-branch commit (23922f9a) built fine until that branch was +# deleted, after which git refused to serve the object at all -- "upload-pack: +# not our ref", because no remaining ref reaches it. A commit on main cannot rot +# that way. It also carries the PR's two later review fixes, which the branch +# snapshot predated. +ARG FLASHINFER_GIT_REPO=https://github.com/flashinfer-ai/flashinfer.git +ARG FLASHINFER_GIT_COMMIT=906181e3f4cf4bcc81835fb480db4011bbd80b62 + +# Uninstall first: a mixed python/cubin/jit-cache installation fails at import, +# and pip would otherwise leave the base's jit-cache shadowing JIT compilation. +# Installed with dependency resolution so apache-tvm-ffi lands at whatever the +# jit-cache wheel's metadata requires -- an exact pin here could contradict it. +# +# The git clone needs its submodules, and not as an optimisation: the released +# wheel packages the cccl, cutlass and spdlog headers into flashinfer/data for +# runtime JIT compilation, so a non-recursive clone yields a package that +# imports fine and then cannot compile a kernel. Shallow, blob-filtered, and +# deleted in the same layer because cutlass and cccl are large. +RUN python3 -m pip uninstall -y \ + flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ + rm -rf /root/.cache/flashinfer && \ + python3 -m pip install \ + "flashinfer-cubin==${FLASHINFER_NIGHTLY_VERSION}" \ + "flashinfer-jit-cache==${FLASHINFER_NIGHTLY_VERSION}+${FLASHINFER_JIT_CACHE_CUDA_TAG}" \ + --extra-index-url https://flashinfer.ai/whl/nightly/ \ + --extra-index-url "https://flashinfer.ai/whl/nightly/${FLASHINFER_JIT_CACHE_CUDA_TAG}/" && \ + git clone --filter=blob:none "${FLASHINFER_GIT_REPO}" /tmp/flashinfer && \ + git -C /tmp/flashinfer checkout --detach "${FLASHINFER_GIT_COMMIT}" && \ + git -C /tmp/flashinfer submodule update --init --recursive --depth 1 && \ + python3 -m pip install --no-deps /tmp/flashinfer && \ + python3 -m pip install "nvidia-cutlass-dsl>=${CUTLASS_DSL_MIN_VERSION}" && \ + # Assert the prebuilt pair is the nightly this commit was matched against. + # flashinfer-python is deliberately NOT compared -- see the opt-out below. + FLASHINFER_EXPECTED="${FLASHINFER_NIGHTLY_VERSION}" \ + FLASHINFER_CUDA_TAG="${FLASHINFER_JIT_CACHE_CUDA_TAG}" \ + python3 -c 'import os; from importlib.metadata import version; e = os.environ["FLASHINFER_EXPECTED"]; tag = os.environ["FLASHINFER_CUDA_TAG"]; got = {p: version(p) for p in ("flashinfer-cubin", "flashinfer-jit-cache")}; assert got["flashinfer-cubin"].split("+")[0] == e, got; assert got["flashinfer-jit-cache"].startswith(e + "+" + tag), got' && \ + rm -rf /tmp/flashinfer /root/.cache/pip + +# flashinfer-python now reports the pinned commit's version while cubin and +# jit-cache report the nightly's. flashinfer/jit/env.py raises RuntimeError at +# import on exactly that mismatch, for both packages, and this variable is the +# opt-out its own error message tells you to use. The alternative -- installing +# no cubin and no jit-cache, which makes the check skip itself -- would send +# every existing kernel through runtime JIT on first use, so the mismatch is +# accepted knowingly instead. +# +# The cost is that the check is now off for good, including for a mismatch +# nobody intended, which is what the assertions above and below are for. +ENV FLASHINFER_DISABLE_VERSION_CHECK=1 + +# FLASHINFER_VERSION describes the prebuilt artifacts, which is what a consumer +# reading it wants to know; the python tree is recorded separately because the +# two genuinely differ in this image. +ENV FLASHINFER_VERSION=${FLASHINFER_NIGHTLY_VERSION} +ENV FLASHINFER_PYTHON_GIT_COMMIT=${FLASHINFER_GIT_COMMIT} + +LABEL ai.radixark.flashinfer.python_git_commit="${FLASHINFER_GIT_COMMIT}" \ + ai.radixark.flashinfer.prebuilt_nightly="${FLASHINFER_NIGHTLY_VERSION}" + +# This tree's BF16 Split-K GEMM loads FlashInfer PR #4266's standalone direct +# kernel by file path out of SGLANG_FLASHINFER_PR4266_SOURCE (see +# python/sglang/srt/layers/quantization/unquant.py). On SM100 that path is on by +# DEFAULT -- bf16_gemm_backend=auto resolves to cutedsl, and +# SGLANG_ENABLE_BF16_SPLITK_GEMM defaults to True -- so an unset variable makes +# the server raise at startup rather than degrade. (On Hopper the backend stays +# unoptimized and the variable is never read, but B200 runs this image too.) +# Point it at the installed FlashInfer through a stable symlink instead of a +# hardcoded dist-packages path, and fail the build now if the kernel is missing. +# +# `import flashinfer` here is load-bearing beyond resolving the path: it runs +# flashinfer/jit/env.py, so it is where a failed version-check opt-out would +# surface. The two file checks then confirm the pinned commit is the tree that +# actually got installed -- checked as files rather than imports because these +# modules pull in CuTe DSL, which cannot load on a CPU build host. +RUN FI_ROOT="$(python3 -c 'import pathlib, flashinfer; print(pathlib.Path(flashinfer.__file__).resolve().parent.parent)')" && \ + ln -sfn "${FI_ROOT}" /opt/flashinfer-src && \ + if [ ! -f /opt/flashinfer-src/flashinfer/gemm/kernels/dense_bf16_gemm_direct.py ]; then \ + echo "ERROR: flashinfer at ${FLASHINFER_GIT_COMMIT} does not ship flashinfer/gemm/kernels/dense_bf16_gemm_direct.py (PR #4266)." >&2; \ + echo " Pin a FlashInfer that carries it, or serve with SGLANG_ENABLE_BF16_SPLITK_GEMM=0." >&2; \ + exit 1; \ + fi && \ + if [ ! -f /opt/flashinfer-src/flashinfer/comm/mnnvl_cutedsl/__init__.py ]; then \ + echo "ERROR: flashinfer at ${FLASHINFER_GIT_COMMIT} does not ship flashinfer/comm/mnnvl_cutedsl (PR #4358)." >&2; \ + echo " That commit is the reason this image builds flashinfer-python from git;" >&2; \ + echo " if it is absent, the wrong ref was installed." >&2; \ + exit 1; \ + fi + +ENV SGLANG_FLASHINFER_PR4266_SOURCE=/opt/flashinfer-src + +# --- 2. Qwen38 SGLang code (replaces the base's stock sglang, editable) --- +# rm first: COPY merges into an existing directory, so files the stock release +# has and this tree does not would otherwise survive. +RUN rm -rf /sgl-workspace/sglang + +COPY . /sgl-workspace/sglang + +# .git is discarded, so setuptools-scm cannot derive a version and would fall +# back to 0.0.0.dev0; pass SGLANG_VERSION to label the build. +# Keep the installed extension modules, but discard Rust and pip build +# artifacts that are not used at runtime. +ARG SGLANG_VERSION=0.0.0.dev0 +# Which + build of the SGLang wheels to pull. Kept separate from the +# FlashInfer jit-cache tag even though both read "cu129" today: they index two +# unrelated wheel sets, and silently reusing one for the other would make a +# FlashInfer retag quietly change which kernel ABI gets installed. +ARG SGL_WHL_CUDA_TAG=cu129 +RUN cd /sgl-workspace/sglang && \ + rm -rf .git && \ + test ! -e .git && \ + SETUPTOOLS_SCM_PRETEND_VERSION="${SGLANG_VERSION}" \ + pip install -e python --no-deps && \ + # --no-deps protects the base's CUDA-tagged wheels from being replaced by + # untagged PyPI builds, but it also drops any pin this tree raised above + # what the base ships. Reinstall those from SGLang's index at the pinned + # version, read out of pyproject.toml so a future bump needs no edit here. + bash docker/qwen38/cuda_pins.sh reconcile \ + python/pyproject.toml "${SGL_WHL_CUDA_TAG}" \ + sglang-kernel sgl-deep-gemm && \ + kernels lock python && \ + ( success=0; \ + if [ "$(uname -m)" = "aarch64" ]; then \ + echo "Skipping sgl-flash-attn3 cubin download on aarch64; kernels will be JIT-compiled at runtime"; \ + success=1; \ + else \ + for i in 1 2 3; do \ + echo "Attempt $i/3: downloading sgl-kernel cubins..."; \ + if kernels download python; then success=1; break; fi; \ + [ "$i" = "3" ] || { echo "sgl-kernel cubin download failed, retrying in 30s..."; sleep 30; }; \ + done; \ + fi; \ + [ "$success" = "1" ] || \ + echo "WARNING: no matching sgl-flash-attn3 cubin variant; kernels will be JIT-compiled at runtime" ) && \ + mkdir -p /root/.cache/huggingface /root/.cache/sglang && \ + ( if [ -f python/kernels.lock ]; then mv python/kernels.lock /root/.cache/sglang/; fi ) && \ + rm -rf \ + rust/target \ + rust/sglang-grpc/target \ + rust/sglang-mm/target \ + rust/sglang-server/target \ + /root/.cargo/registry \ + /root/.cache/pip + +# --- 3. Verify the image can actually import what it ships --- +# This image keeps the base's DeepEP rather than rebuilding it, so the NCCL +# re-pin the CUDA 13 recipe needs does not apply here. What does apply is the +# check: the FlashInfer install above resolves dependencies, and torch declares +# a hard `nvidia-nccl-cu12==`, so the NCCL under DeepEP can move without +# any build step failing. The CUDA 13 image shipped green while deep_ep._C had an +# unresolved ncclGetLsaDevicePointer, so assert it instead of assuming it. +# +# deep_ep goes through import-if-gpu rather than a bare import. On the dev-cu12 +# base the module is provided by the packaged `sgl-deep-ep`, which checks for a +# usable CUDA device at import and raises "The NVIDIA driver does not expose a +# usable CUDA device" -- so a bare import fails on every CPU build host, which is +# all of them. The check is kept rather than dropped so it still runs when the +# image is built on a GPU host; a GPU smoke test has to cover the rest. +# +# `import sglang` proves nothing about the dependency versions: assert_pkg_version +# lives in srt/entrypoints/engine.py and only runs once a server starts. A tree +# that outgrew the base's sglang-kernel therefore builds green here, ships, and +# dies on every rank at launch -- exactly the shape of the deep_ep bug above. +# Run that same assertion at build time so the mismatch fails the build instead. +# Unlike the imports, verify reads installed metadata, so it works everywhere. +RUN cd /sgl-workspace/sglang && \ + bash docker/qwen38/cuda_pins.sh import-if-gpu deep_ep && \ + python3 -c 'import sglang; print("sglang", sglang.__version__)' && \ + bash docker/qwen38/cuda_pins.sh verify \ + python/pyproject.toml sglang-kernel sgl-deep-gemm + +WORKDIR /sgl-workspace/sglang diff --git a/docker/qwen38/qwen38_cu13.Dockerfile b/docker/qwen38/qwen38_cu13.Dockerfile new file mode 100644 index 000000000000..e31460924926 --- /dev/null +++ b/docker/qwen38/qwen38_cu13.Dockerfile @@ -0,0 +1,297 @@ +# Qwen38 serving image (CUDA 13 / GB300 aarch64: sm_90a + sm_100a + sm_103a). +# +# Base ships stock SGLang (editable at /sgl-workspace/sglang), a DeepEP the +# recipe below replaces, a released FlashInfer trio, the Rust toolchain at +# /root/.cargo, and the CUDA 13 toolchain (nvcc + +# /usr/local/cuda/include/cccl). How DeepEP arrives in the base is not assumed: +# v0.5.17 shipped a source tree at /sgl-workspace/DeepEP plus a `deep_ep` +# distribution, the dev images ship a packaged `sgl-deep-ep` instead, and +# apply_deepep_v2_patch.sh handles either. +# +# This image adds the four Qwen38-specific pieces that stock lacks: +# 1. DeepEP v2 (deepseek-ai@01dc3aa) patched and built from source; see +# apply_deepep_v2_patch.sh for what the patch does and why +# 2. NCCL pinned to the version the DeepEP v2 wheel requires, staged at +# /opt/nccl-2.30.7/lib so the srt-slurm configs can LD_PRELOAD it without +# a host library mount +# 3. FlashInfer, replacing the base's release: flashinfer-python from a pinned +# git commit (PR #4358, the Blackwell MNNVL CuTe DSL all-reduce backend), +# with the nightly's prebuilt flashinfer-cubin and flashinfer-jit-cache +# alongside it, since neither is published per commit. +# nvidia-cutlass-dsl is floored at 4.7.0 because the CuTe DSL kernels call +# PipelineTmaAsync.create(enable_multicast_signaling=..) +# 4. this repo's SGLang code, copied from the build context and editable-installed +# +# The build also refuses to start unless the base's torch is the one +# python/pyproject.toml pins, reinstalls the CUDA-tagged wheels this tree pins +# above what the base ships, and asserts those pins before finishing. All three +# live in docker/qwen38/cuda_pins.sh, which explains why each is needed. +# +# Build (no GPU needed; nvcc cross-compiles the DeepEP cubins). The context must +# be the repo root -- the SGLang source is taken from it rather than cloned: +# docker build -f docker/qwen38/qwen38_cu13.Dockerfile -t qwen38-cu130 . + +# The nightly, not a release tag: v0.5.17 ships torch 2.11.0 / sglang-kernel +# 0.4.5, which the gate below rejects against this tree's pins. `dev` is the +# CUDA 13 nightly and carries torch 2.13.0+cu130, sglang-kernel 0.4.6.post1 and +# sgl-deep-gemm 0.1.5.post2 -- exactly what python/pyproject.toml asks for. +# +# `dev` MOVES every night. Builds are therefore not reproducible from this tag +# alone; pin `dev@sha256:` when a build has to be repeatable. +FROM lmsysorg/sglang:dev AS base + +# --- 0. Base/tree compatibility gate, FIRST --- +# Whether this base can host this tree is decided entirely by the base's torch, +# and it is knowable before any work happens. That matters most here: the DeepEP +# v2 nvcc compile below is the longest step in either recipe, and running it +# against a base that cannot host this tree is pure waste. +# +# pyproject.toml is copied to its own path purely to make that possible: the tree +# itself only lands at COPY below, far too late to gate anything. Both copies +# come from the same build context, so they cannot disagree. +COPY docker/qwen38/cuda_pins.sh /opt/qwen38/cuda_pins.sh +COPY python/pyproject.toml /opt/qwen38/pyproject.toml + +RUN bash /opt/qwen38/cuda_pins.sh check-torch /opt/qwen38/pyproject.toml + +# --- 1. DeepEP v2: patch + build, replacing the base's stock tree and wheel --- +ARG DEEPEP_V2_COMMIT=01dc3aaac82068020353dce2c302e38153c0bfaa +ARG DEEPEP_CUDA_ARCH_LIST="9.0;10.0;10.3" +ARG BUILD_AND_DOWNLOAD_PARALLEL=8 + +COPY docker/qwen38/apply_deepep_v2_patch.sh /opt/qwen38/ + +# DeepEP v2's NCCL backend uses the GIN API (ncclCommProperties, +# ncclGinRequest_t, NCCL_GIN_*), which first shipped in NCCL 2.30 -- and the base +# is below that (2.29.7 on `dev` as of 2026-08-10), so its headers do not compile +# it. Install the pinned NCCL BEFORE the DeepEP build so the extension compiles +# and links against the same version that gets preloaded at run time. +ARG NCCL_PIN_VERSION=2.30.7 +RUN python3 -m pip install "nvidia-nccl-cu13==${NCCL_PIN_VERSION}" && \ + rm -rf /root/.cache/pip + +RUN DEEPEP_V2_COMMIT="${DEEPEP_V2_COMMIT}" \ + TORCH_CUDA_ARCH_LIST="${DEEPEP_CUDA_ARCH_LIST}" \ + MAX_JOBS="${BUILD_AND_DOWNLOAD_PARALLEL}" \ + bash /opt/qwen38/apply_deepep_v2_patch.sh + +# Assert the pin was still live for the DeepEP compile above. `pip install X==v` +# does not re-resolve the graph, so nothing should have moved it here -- but the +# headers DeepEP just compiled against are only correct if this holds. +RUN NCCL_EXPECTED="${NCCL_PIN_VERSION}" python3 -c 'import os; from importlib.metadata import version; e = os.environ["NCCL_EXPECTED"]; v = version("nvidia-nccl-cu13"); assert v == e, (v, e)' + +# --- 2. NCCL preload path: staged at the END of the build, see below --- + +# --- 3. FlashInfer: nightly cubin/jit-cache, python from a pinned commit --- +# Named apart from the base's ENV FLASHINFER_VERSION, which would otherwise +# shadow a same-named ARG and silently resolve to the base's 0.6.15.post1. +# +# The nightly version below no longer covers flashinfer-python: that comes from +# FLASHINFER_GIT_COMMIT instead. cubin and jit-cache stay on the nightly because +# they are prebuilt artifacts published per nightly date and per release only -- +# no wheel of either exists for an arbitrary commit. The nightly named here must +# therefore be the one whose main the pinned commit merges, so the prebuilt +# kernels correspond to the source they were compiled from. +ARG FLASHINFER_NIGHTLY_VERSION=0.6.18.dev20260807 +ARG FLASHINFER_JIT_CACHE_CUDA_TAG=cu130 +ARG CUTLASS_DSL_MIN_VERSION=4.7.0 + +# flashinfer-ai/flashinfer#4358, "feat(comm): add Blackwell MNNVL CuTe DSL +# all-reduce fusion backend" -- the MNNVL path this image's GB300 target needs. +# Pinned to the PR's merge commit on main, which is what keeps it fetchable. An +# earlier pin of a PR-branch commit (23922f9a) built fine until that branch was +# deleted, after which git refused to serve the object at all -- "upload-pack: +# not our ref", because no remaining ref reaches it. A commit on main cannot rot +# that way. It also carries the PR's two later review fixes, which the branch +# snapshot predated. +ARG FLASHINFER_GIT_REPO=https://github.com/flashinfer-ai/flashinfer.git +ARG FLASHINFER_GIT_COMMIT=906181e3f4cf4bcc81835fb480db4011bbd80b62 + +# Uninstall first: a mixed python/cubin/jit-cache installation fails at import, +# and pip would otherwise leave the base's jit-cache shadowing JIT compilation. +# Installed with dependency resolution so apache-tvm-ffi lands at whatever the +# jit-cache wheel's metadata requires -- an exact pin here could contradict it. +# +# NOTE: this resolution DOWNGRADES nvidia-nccl-cu13 back off the 2.30.7 pin, +# because torch declares its own hard `nvidia-nccl-cu13==` (2.29.7 on the `dev` +# base) and flashinfer pulls in nccl4py, dragging NCCL into the resolve. The +# exact version torch drags in tracks the base and is not worth hardcoding here; +# what matters is that it is below 2.30 and therefore lacks the GIN API. The pin +# is re-applied after all pip work completes -- see the final stage, which +# asserts the version rather than trusting it. +# +# The git clone needs its submodules, and not as an optimisation: the released +# wheel packages the cccl, cutlass and spdlog headers into flashinfer/data for +# runtime JIT compilation, so a non-recursive clone yields a package that +# imports fine and then cannot compile a kernel. Shallow, blob-filtered, and +# deleted in the same layer because cutlass and cccl are large. +RUN python3 -m pip uninstall -y \ + flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ + rm -rf /root/.cache/flashinfer && \ + python3 -m pip install \ + "flashinfer-cubin==${FLASHINFER_NIGHTLY_VERSION}" \ + "flashinfer-jit-cache==${FLASHINFER_NIGHTLY_VERSION}+${FLASHINFER_JIT_CACHE_CUDA_TAG}" \ + --extra-index-url https://flashinfer.ai/whl/nightly/ \ + --extra-index-url "https://flashinfer.ai/whl/nightly/${FLASHINFER_JIT_CACHE_CUDA_TAG}/" && \ + git clone --filter=blob:none "${FLASHINFER_GIT_REPO}" /tmp/flashinfer && \ + git -C /tmp/flashinfer checkout --detach "${FLASHINFER_GIT_COMMIT}" && \ + git -C /tmp/flashinfer submodule update --init --recursive --depth 1 && \ + python3 -m pip install --no-deps /tmp/flashinfer && \ + python3 -m pip install "nvidia-cutlass-dsl[cu13]>=${CUTLASS_DSL_MIN_VERSION}" && \ + # Assert the prebuilt pair is the nightly this commit was matched against. + # flashinfer-python is deliberately NOT compared -- see the opt-out below. + FLASHINFER_EXPECTED="${FLASHINFER_NIGHTLY_VERSION}" \ + FLASHINFER_CUDA_TAG="${FLASHINFER_JIT_CACHE_CUDA_TAG}" \ + python3 -c 'import os; from importlib.metadata import version; e = os.environ["FLASHINFER_EXPECTED"]; tag = os.environ["FLASHINFER_CUDA_TAG"]; got = {p: version(p) for p in ("flashinfer-cubin", "flashinfer-jit-cache")}; assert got["flashinfer-cubin"].split("+")[0] == e, got; assert got["flashinfer-jit-cache"].startswith(e + "+" + tag), got' && \ + rm -rf /tmp/flashinfer /root/.cache/pip + +# flashinfer-python now reports the pinned commit's version while cubin and +# jit-cache report the nightly's. flashinfer/jit/env.py raises RuntimeError at +# import on exactly that mismatch, for both packages, and this variable is the +# opt-out its own error message tells you to use. The alternative -- installing +# no cubin and no jit-cache, which makes the check skip itself -- would send +# every existing kernel through runtime JIT on first use, so the mismatch is +# accepted knowingly instead. +# +# The cost is that the check is now off for good, including for a mismatch +# nobody intended, which is what the assertions above and below are for. +ENV FLASHINFER_DISABLE_VERSION_CHECK=1 + +# FLASHINFER_VERSION describes the prebuilt artifacts, which is what a consumer +# reading it wants to know; the python tree is recorded separately because the +# two genuinely differ in this image. +ENV FLASHINFER_VERSION=${FLASHINFER_NIGHTLY_VERSION} +ENV FLASHINFER_PYTHON_GIT_COMMIT=${FLASHINFER_GIT_COMMIT} + +LABEL ai.radixark.flashinfer.python_git_commit="${FLASHINFER_GIT_COMMIT}" \ + ai.radixark.flashinfer.prebuilt_nightly="${FLASHINFER_NIGHTLY_VERSION}" + +# This tree's BF16 Split-K GEMM loads FlashInfer PR #4266's standalone direct +# kernel by file path out of SGLANG_FLASHINFER_PR4266_SOURCE (see +# python/sglang/srt/layers/quantization/unquant.py). On SM100/SM103 that path is +# on by DEFAULT -- bf16_gemm_backend=auto resolves to cutedsl, and +# SGLANG_ENABLE_BF16_SPLITK_GEMM defaults to True -- so an unset variable makes +# the server raise at startup rather than degrade. Point it at the installed +# FlashInfer through a stable symlink instead of a hardcoded dist-packages path, +# and fail the build now if the kernel is missing: a GB300 image that cannot +# start is worse than a build that stops here. +# +# `import flashinfer` here is load-bearing beyond resolving the path: it runs +# flashinfer/jit/env.py, so it is where a failed version-check opt-out would +# surface. The two file checks then confirm the pinned commit is the tree that +# actually got installed -- checked as files rather than imports because these +# modules pull in CuTe DSL, which cannot load on a CPU build host. +RUN FI_ROOT="$(python3 -c 'import pathlib, flashinfer; print(pathlib.Path(flashinfer.__file__).resolve().parent.parent)')" && \ + ln -sfn "${FI_ROOT}" /opt/flashinfer-src && \ + if [ ! -f /opt/flashinfer-src/flashinfer/gemm/kernels/dense_bf16_gemm_direct.py ]; then \ + echo "ERROR: flashinfer at ${FLASHINFER_GIT_COMMIT} does not ship flashinfer/gemm/kernels/dense_bf16_gemm_direct.py (PR #4266)." >&2; \ + echo " Pin a FlashInfer that carries it, or serve with SGLANG_ENABLE_BF16_SPLITK_GEMM=0." >&2; \ + exit 1; \ + fi && \ + if [ ! -f /opt/flashinfer-src/flashinfer/comm/mnnvl_cutedsl/__init__.py ]; then \ + echo "ERROR: flashinfer at ${FLASHINFER_GIT_COMMIT} does not ship flashinfer/comm/mnnvl_cutedsl (PR #4358)." >&2; \ + echo " That commit is the reason this image builds flashinfer-python from git;" >&2; \ + echo " if it is absent, the wrong ref was installed." >&2; \ + exit 1; \ + fi + +ENV SGLANG_FLASHINFER_PR4266_SOURCE=/opt/flashinfer-src + +# --- 4. Qwen38 SGLang code (replaces the base's stock sglang, editable) --- +# rm first: COPY merges into an existing directory, so files the stock release +# has and this tree does not would otherwise survive. +RUN rm -rf /sgl-workspace/sglang + +COPY . /sgl-workspace/sglang + +# .git is discarded, so setuptools-scm cannot derive a version and would fall +# back to 0.0.0.dev0; pass SGLANG_VERSION to label the build. +# Keep the installed extension modules, but discard Rust and pip build +# artifacts that are not used at runtime. +ARG SGLANG_VERSION=0.0.0.dev0 +# Which + build of the SGLang wheels to pull. Kept separate from the +# FlashInfer jit-cache tag even though both read "cu130" today: they index two +# unrelated wheel sets, and silently reusing one for the other would make a +# FlashInfer retag quietly change which kernel ABI gets installed. +ARG SGL_WHL_CUDA_TAG=cu130 +RUN cd /sgl-workspace/sglang && \ + rm -rf .git && \ + test ! -e .git && \ + SETUPTOOLS_SCM_PRETEND_VERSION="${SGLANG_VERSION}" \ + pip install -e python --no-deps && \ + # --no-deps protects the base's CUDA-tagged wheels from being replaced by + # untagged PyPI builds, but it also drops any pin this tree raised above + # what the base ships. Reinstall those from SGLang's index at the pinned + # version, read out of pyproject.toml so a future bump needs no edit here. + bash docker/qwen38/cuda_pins.sh reconcile \ + python/pyproject.toml "${SGL_WHL_CUDA_TAG}" \ + sglang-kernel sgl-deep-gemm && \ + kernels lock python && \ + ( success=0; \ + if [ "$(uname -m)" = "aarch64" ]; then \ + echo "Skipping sgl-flash-attn3 cubin download on aarch64; kernels will be JIT-compiled at runtime"; \ + success=1; \ + else \ + for i in 1 2 3; do \ + echo "Attempt $i/3: downloading sgl-kernel cubins..."; \ + if kernels download python; then success=1; break; fi; \ + [ "$i" = "3" ] || { echo "sgl-kernel cubin download failed, retrying in 30s..."; sleep 30; }; \ + done; \ + fi; \ + [ "$success" = "1" ] || \ + echo "WARNING: no matching sgl-flash-attn3 cubin variant; kernels will be JIT-compiled at runtime" ) && \ + mkdir -p /root/.cache/huggingface /root/.cache/sglang && \ + ( if [ -f python/kernels.lock ]; then mv python/kernels.lock /root/.cache/sglang/; fi ) && \ + rm -rf \ + rust/target \ + rust/sglang-grpc/target \ + rust/sglang-mm/target \ + rust/sglang-server/target \ + /root/.cargo/registry \ + /root/.cache/pip + +# --- 2 (deferred). Re-pin NCCL and stage the preload path, LAST --- +# torch declares `nvidia-nccl-cu13==2.28.9`, so every pip install above that +# resolves dependencies drags NCCL back down to 2.28.9 -- silently undoing the +# pin that DeepEP v2 was compiled against, and leaving deep_ep._C with an +# unresolved `ncclGetLsaDevicePointer`. Re-pinning here, after all pip work, is +# what makes the runtime match the compile. --no-deps so this install cannot +# itself re-resolve; pip will warn that torch's pin is now unsatisfied, which is +# the intended override. +# +# The link target is resolved from the installed package rather than hardcoded +# to /usr/local/lib/python3.12/dist-packages: a base that moves to another +# interpreter or to /usr/lib/python3 would otherwise leave a dangling link, and +# LD_PRELOAD of a dangling path fails silently at runtime. +# nvidia.nccl is a namespace package: __file__ is None, only __path__ is set. +RUN python3 -m pip install --no-deps --force-reinstall \ + "nvidia-nccl-cu13==${NCCL_PIN_VERSION}" && \ + rm -rf /root/.cache/pip && \ + NCCL_LIB="$(python3 -c 'import os, nvidia.nccl; print(os.path.join(list(nvidia.nccl.__path__)[0], "lib"))')" && \ + test -d "${NCCL_LIB}" && \ + mkdir -p /opt/nccl-${NCCL_PIN_VERSION} && \ + ln -sfn "${NCCL_LIB}" /opt/nccl-${NCCL_PIN_VERSION}/lib && \ + test -e /opt/nccl-${NCCL_PIN_VERSION}/lib/libnccl.so.2 && \ + # Assert the VERSION, not just the path: the old `test -e` guard passed + # happily against the downgraded 2.28.9 sitting at the same location. + NCCL_EXPECTED="${NCCL_PIN_VERSION}" python3 -c 'import os; from importlib.metadata import version; e = os.environ["NCCL_EXPECTED"]; v = version("nvidia-nccl-cu13"); assert v == e, (v, e)' && \ + # The check that actually matters, and the one whose absence let a broken + # image build green: DeepEP v2 must load against the NCCL now installed. + python3 -c 'import deep_ep; print("deep_ep import OK")' + +# --- 5. Assert the dependency pins the tree declares, LAST --- +# `import sglang` proves nothing about dependency versions: assert_pkg_version +# lives in srt/entrypoints/engine.py and only runs once a server starts. A tree +# that outgrew the base's sglang-kernel therefore builds green here, ships, and +# dies on every rank at launch -- the same shape as the deep_ep bug above, and it +# has already happened once: an image built on a base carrying sglang-kernel +# 0.4.5 shipped while engine.py asserts a 0.4.6.post1 floor. +# +# Runs after every pip step in the file, so nothing below can quietly move a +# version back down. +RUN cd /sgl-workspace/sglang && \ + python3 -c 'import sglang; print("sglang", sglang.__version__)' && \ + bash docker/qwen38/cuda_pins.sh verify \ + python/pyproject.toml sglang-kernel sgl-deep-gemm + +WORKDIR /sgl-workspace/sglang diff --git a/python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh b/python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh index b82f599955fa..fbecd7b7675b 100644 --- a/python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh +++ b/python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh @@ -151,11 +151,10 @@ __global__ __launch_bounds__(1024, 2) void // // Read this token's kTopK destinations once (fully unrolled). const auto* src2dst_row = params.src2dst + static_cast(token_id) * kTopK; - const auto* topk_ids_row = params.topk_ids + static_cast(token_id) * kTopK; int32_t dst_rows[kTopK]; #pragma unroll for (uint32_t i = 0; i < kTopK; ++i) { - dst_rows[i] = (topk_ids_row[i] >= 0) ? src2dst_row[i] : -1; + dst_rows[i] = src2dst_row[i]; } const uint32_t group_id = tid / kThreadsPerGroup; diff --git a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py index 076a1cc3f7be..6bfafaed9911 100644 --- a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py +++ b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py @@ -310,6 +310,321 @@ def fused_qkvzba_split_reshape_cat_contiguous( return mixed_qkv, z, b, a +# ============================================================================= +# Decode-only Qwen3.5 contiguous projection unpack + causal Conv1D update. +# +# This deliberately leaves the quantized projection GEMMs unchanged. The safe +# fusion boundary begins at their activation outputs: +# qkvz = [all_q | all_k | all_v | all_z] +# ba = [all_b | all_a] +# ============================================================================= + + +@triton.jit +def _fused_qkvzba_causal_conv1d_update_contiguous_kernel( + mixed_qkv, + z, + b, + a, + mixed_qkvz, + mixed_ba, + conv_state, + conv_weight, + conv_bias, + conv_state_indices, + stride_qkvz_batch: tl.constexpr, + stride_qkvz_dim: tl.constexpr, + stride_ba_batch: tl.constexpr, + stride_ba_dim: tl.constexpr, + stride_state_batch: tl.constexpr, + stride_state_dim: tl.constexpr, + stride_state_pos: tl.constexpr, + stride_weight_dim: tl.constexpr, + stride_weight_width: tl.constexpr, + stride_state_indices: tl.constexpr, + QKV_DIM: tl.constexpr, + V_DIM: tl.constexpr, + NUM_V_HEADS: tl.constexpr, + NUM_STATE_SLOTS: tl.constexpr, + STATE_LEN: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + dim_idx = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + qkv_mask = dim_idx < QKV_DIM + + x = tl.load( + mixed_qkvz + batch_idx * stride_qkvz_batch + dim_idx * stride_qkvz_dim, + mask=qkv_mask, + other=0.0, + ) + + state_slot = tl.load(conv_state_indices + batch_idx * stride_state_indices).to( + tl.int64 + ) + # Treat every out-of-range index as padding. Replay metadata normally uses + # exactly PAD_SLOT_ID, but this extra bound prevents malformed/stale graph + # metadata from turning an indexed state update into an OOB access. + valid_slot = ( + (state_slot != PAD_SLOT_ID) + & (state_slot >= 0) + & (state_slot < NUM_STATE_SLOTS) + ) + state_base = ( + conv_state + state_slot * stride_state_batch + dim_idx * stride_state_dim + ) + + acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + if HAS_BIAS: + acc += tl.load(conv_bias + dim_idx, mask=qkv_mask, other=0.0).to(tl.float32) + + # Match the deployed direct-Triton update exactly. Its effective decode + # state length is width-1 even when the physical cache tensor is wider. + for pos in tl.static_range(KERNEL_WIDTH - 1): + state_value = tl.load( + state_base + pos * stride_state_pos, + mask=qkv_mask & valid_slot, + other=0.0, + ) + weight_value = tl.load( + conv_weight + dim_idx * stride_weight_dim + pos * stride_weight_width, + mask=qkv_mask, + other=0.0, + ) + # Do not force an FP32 multiply here. This expression deliberately + # retains the operand types/order of causal_conv1d_triton.py. + acc += state_value * weight_value + + last_weight = tl.load( + conv_weight + + dim_idx * stride_weight_dim + + (KERNEL_WIDTH - 1) * stride_weight_width, + mask=qkv_mask, + other=0.0, + ) + acc += x * last_weight + if SILU_ACTIVATION: + conv_out = acc / (1.0 + tl.exp(-acc)) + else: + conv_out = acc + + # The legacy kernel leaves padded rows' input unchanged. + conv_out = tl.where(valid_slot, conv_out, x) + tl.store( + mixed_qkv + batch_idx * QKV_DIM + dim_idx, + conv_out, + mask=qkv_mask, + ) + + # The direct-Triton wrapper sets effective state_len=width-1 for decode. + for pos in tl.static_range(KERNEL_WIDTH - 2): + next_value = tl.load( + state_base + (pos + 1) * stride_state_pos, + mask=qkv_mask & valid_slot, + other=0.0, + ) + tl.store( + state_base + pos * stride_state_pos, + next_value, + mask=qkv_mask & valid_slot, + ) + tl.store( + state_base + (KERNEL_WIDTH - 2) * stride_state_pos, + x, + mask=qkv_mask & valid_slot, + ) + + # The first feature lanes also materialize the smaller downstream tensors. + z_mask = dim_idx < V_DIM + z_value = tl.load( + mixed_qkvz + + batch_idx * stride_qkvz_batch + + (QKV_DIM + dim_idx) * stride_qkvz_dim, + mask=z_mask, + other=0.0, + ) + tl.store(z + batch_idx * V_DIM + dim_idx, z_value, mask=z_mask) + + gate_mask = dim_idx < NUM_V_HEADS + b_value = tl.load( + mixed_ba + batch_idx * stride_ba_batch + dim_idx * stride_ba_dim, + mask=gate_mask, + other=0.0, + ) + a_value = tl.load( + mixed_ba + + batch_idx * stride_ba_batch + + (NUM_V_HEADS + dim_idx) * stride_ba_dim, + mask=gate_mask, + other=0.0, + ) + tl.store(b + batch_idx * NUM_V_HEADS + dim_idx, b_value, mask=gate_mask) + tl.store(a + batch_idx * NUM_V_HEADS + dim_idx, a_value, mask=gate_mask) + + +def can_use_fused_qkvzba_causal_conv1d_update_contiguous( + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + conv_state_indices: torch.Tensor, + *, + qkv_dim: int, + v_dim: int, + num_v_heads: int, + activation: str | None, +) -> tuple[bool, str]: + """Return an explicit eligibility decision for the decode fusion.""" + tensors = (mixed_qkvz, mixed_ba, conv_state, conv_weight, conv_state_indices) + if not all(isinstance(tensor, torch.Tensor) for tensor in tensors): + return False, "all inputs must be torch.Tensor instances" + if not all(tensor.is_cuda for tensor in tensors): + return False, "CUDA tensors are required" + if mixed_qkvz.ndim != 2 or mixed_ba.ndim != 2: + return False, "projection outputs must be rank-2" + if conv_state.ndim != 3 or conv_weight.ndim != 2: + return False, "Conv1D state/weight ranks must be 3/2" + if conv_state_indices.ndim != 1: + return False, "conv_state_indices must be rank-1" + batch = mixed_qkvz.shape[0] + if mixed_ba.shape[0] != batch or conv_state_indices.shape[0] != batch: + return False, "batch dimensions must match" + if qkv_dim <= 0 or v_dim <= 0 or num_v_heads <= 0: + return False, "TP-local dimensions must be positive" + if mixed_qkvz.shape[1] != qkv_dim + v_dim: + return False, "qkvz layout is not contiguous [Q|K|V|Z]" + if mixed_ba.shape[1] != 2 * num_v_heads: + return False, "ba layout is not contiguous [B|A]" + if conv_state.shape[1] != qkv_dim or conv_weight.shape[0] != qkv_dim: + return False, "Conv1D feature dimension does not match packed QKV" + width = conv_weight.shape[1] + if width < 2 or width > 4: + return False, "only Conv1D widths 2 through 4 are supported" + if conv_state.shape[2] < width - 1: + return False, "Conv1D state is shorter than width - 1" + supported_dtypes = (torch.float16, torch.bfloat16, torch.float32) + if mixed_qkvz.dtype not in supported_dtypes: + return False, "QKVZ activation dtype must be FP16, BF16, or FP32" + if conv_state.dtype != mixed_qkvz.dtype or conv_weight.dtype != mixed_qkvz.dtype: + return False, "QKVZ, Conv1D state, and weight dtypes must match" + if mixed_ba.dtype not in supported_dtypes: + return False, "BA activation dtype must be FP16, BF16, or FP32" + if conv_bias is not None: + if ( + not isinstance(conv_bias, torch.Tensor) + or not conv_bias.is_cuda + or conv_bias.ndim != 1 + or conv_bias.shape[0] != qkv_dim + or conv_bias.dtype != mixed_qkvz.dtype + ): + return False, "Conv1D bias contract is incompatible" + if activation not in (None, "silu", "swish"): + return False, "activation must be None, silu, or swish" + if mixed_qkvz.stride(1) != 1 or mixed_ba.stride(1) != 1: + return False, "projection feature dimensions must be contiguous" + if conv_weight.stride(1) != 1: + return False, "Conv1D weight width dimension must be contiguous" + if conv_state_indices.dtype not in (torch.int32, torch.int64): + return False, "conv_state_indices must be int32 or int64" + return True, "eligible" + + +def fused_qkvzba_causal_conv1d_update_contiguous( + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + conv_state: torch.Tensor, + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + conv_state_indices: torch.Tensor, + *, + qkv_dim: int, + v_dim: int, + num_v_heads: int, + head_v_dim: int, + activation: str | None, + pad_slot_id: int = -1, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Decode-only fused Qwen3.5 projection unpack and Conv1D state update.""" + eligible, reason = can_use_fused_qkvzba_causal_conv1d_update_contiguous( + mixed_qkvz, + mixed_ba, + conv_state, + conv_weight, + conv_bias, + conv_state_indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + activation=activation, + ) + if not eligible: + raise ValueError(f"Ineligible fused GDN decode projection/Conv1D: {reason}") + if v_dim != num_v_heads * head_v_dim: + raise ValueError( + "Ineligible fused GDN decode projection/Conv1D: " + "v_dim must equal num_v_heads * head_v_dim" + ) + + batch = mixed_qkvz.shape[0] + mixed_qkv = torch.empty( + (batch, qkv_dim), dtype=mixed_qkvz.dtype, device=mixed_qkvz.device + ) + z = torch.empty( + (batch, num_v_heads, head_v_dim), + dtype=mixed_qkvz.dtype, + device=mixed_qkvz.device, + ) + b = torch.empty( + (batch, num_v_heads), + dtype=mixed_ba.dtype, + device=mixed_ba.device, + ) + a = torch.empty_like(b) + + block_size = 256 + grid = (batch, triton.cdiv(qkv_dim, block_size)) + _fused_qkvzba_causal_conv1d_update_contiguous_kernel[grid]( + mixed_qkv, + z, + b, + a, + mixed_qkvz, + mixed_ba, + conv_state, + conv_weight, + conv_bias, + conv_state_indices, + mixed_qkvz.stride(0), + mixed_qkvz.stride(1), + mixed_ba.stride(0), + mixed_ba.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + conv_weight.stride(0), + conv_weight.stride(1), + conv_state_indices.stride(0), + QKV_DIM=qkv_dim, + V_DIM=v_dim, + NUM_V_HEADS=num_v_heads, + NUM_STATE_SLOTS=conv_state.shape[0], + STATE_LEN=conv_state.shape[2], + KERNEL_WIDTH=conv_weight.shape[1], + HAS_BIAS=conv_bias is not None, + SILU_ACTIVATION=activation in ("silu", "swish"), + PAD_SLOT_ID=pad_slot_id, + BLOCK_SIZE=block_size, + num_warps=8, + num_stages=2, + ) + return mixed_qkv, z, b, a + + @triton.jit def fused_qkv_split_gdn_prefill_kernel( q, diff --git a/python/sglang/kernels/ops/elementwise/elementwise.py b/python/sglang/kernels/ops/elementwise/elementwise.py index f0eccba29e08..a9aaeee6b9c1 100644 --- a/python/sglang/kernels/ops/elementwise/elementwise.py +++ b/python/sglang/kernels/ops/elementwise/elementwise.py @@ -465,9 +465,10 @@ def _fused_gate_sigmoid_mul_add_kernel( hidden_states_ptr, # [num_tokens, hidden_dim] gate_weight_ptr, # [hidden_dim] shared_output_ptr, # [num_tokens, hidden_dim] - final_hidden_states_ptr, # [num_tokens, hidden_dim] + output_ptr, # [num_tokens, hidden_dim], optionally also the addend hidden_dim: tl.constexpr, BLOCK_SIZE: tl.constexpr, + DO_ADD: tl.constexpr = True, USE_PDL: tl.constexpr = False, ): pid = tl.program_id(axis=0).to(tl.int64) @@ -487,41 +488,39 @@ def _fused_gate_sigmoid_mul_add_kernel( s = tl.load(shared_output_ptr + row_offset + offsets, mask=mask, other=0.0).to( tl.float32 ) - f = tl.load( - final_hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0 - ).to(tl.float32) + if DO_ADD: + f = tl.load(output_ptr + row_offset + offsets, mask=mask, other=0.0).to( + tl.float32 + ) if USE_PDL: tl.extra.cuda.gdc_launch_dependents() gate_val = tl.sigmoid(tl.sum(h * w, axis=0)) - result = f + gate_val * s + result = gate_val * s + if DO_ADD: + result += f - tl.store(final_hidden_states_ptr + row_offset + offsets, result, mask=mask) + tl.store(output_ptr + row_offset + offsets, result, mask=mask) -def fused_gate_sigmoid_mul_add( +def _launch_fused_gate_sigmoid_mul( hidden_states: torch.Tensor, gate_weight: torch.Tensor, shared_output: torch.Tensor, - final_hidden_states: torch.Tensor, + output: torch.Tensor, + *, + do_add: bool, ) -> None: - """ - Fused gate-sigmoid-mul-add for MoE shared expert gating. - - Equivalent to: - gate = hidden_states @ gate_weight - final_hidden_states += sigmoid(gate).unsqueeze(1) * shared_output - """ assert hidden_states.is_contiguous(), "hidden_states must be contiguous" assert gate_weight.is_contiguous(), "gate_weight must be contiguous" assert shared_output.is_contiguous(), "shared_output must be contiguous" - assert final_hidden_states.is_contiguous(), "final_hidden_states must be contiguous" + assert output.is_contiguous(), "output must be contiguous" num_tokens, hidden_dim = hidden_states.shape assert gate_weight.shape == (hidden_dim,) assert shared_output.shape == (num_tokens, hidden_dim) - assert final_hidden_states.shape == (num_tokens, hidden_dim) + assert output.shape == (num_tokens, hidden_dim) max_warps = 16 if _is_hip else 32 config = { @@ -540,8 +539,42 @@ def fused_gate_sigmoid_mul_add( hidden_states, gate_weight, shared_output, - final_hidden_states, + output, hidden_dim=hidden_dim, + DO_ADD=do_add, **config, **pdl_kwargs, ) + + +def fused_gate_sigmoid_mul( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + shared_output: torch.Tensor, +) -> torch.Tensor: + """Materialize the gated shared-expert contribution without an add/copy.""" + output = torch.empty_like(shared_output) + _launch_fused_gate_sigmoid_mul( + hidden_states, + gate_weight, + shared_output, + output, + do_add=False, + ) + return output + + +def fused_gate_sigmoid_mul_add( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + shared_output: torch.Tensor, + final_hidden_states: torch.Tensor, +) -> None: + """Add the gated shared-expert contribution to routed-expert output.""" + _launch_fused_gate_sigmoid_mul( + hidden_states, + gate_weight, + shared_output, + final_hidden_states, + do_add=True, + ) diff --git a/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py new file mode 100644 index 000000000000..ad405f3d2836 --- /dev/null +++ b/python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py @@ -0,0 +1,1041 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# Vendored from flashinfer-ai/flashinfer PR #4266 at 629147317d4149a12e53bcef27808bac380c283f. +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = 227 * 1024 + +#: K extent of one CTA tile. +_CTA_K = 128 + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 + +#: A/B pipeline stage bounds. +_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 + + +@dataclasses.dataclass(frozen=True, slots=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], " + f"got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across " + f"split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + base = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(_MIN_AB_STAGES, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) if k <= 4 * _CTA_K else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: cute.Pointer, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +) -> Int32: + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: cute.Pointer, + mbar_ptr: cute.Pointer, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank( + smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank( + mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ) -> None: + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = ( + (tactic.split_k - 1) * self.epilog_threads * values_per_thread + ) + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel(a, b, c, bias).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc( + tiled_mma, self.mma_tiler_mn, acc_stage=1 + ) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + cute.arch.alloc_tmem(num_tmem_cols, tmem_base_ptr, is_two_cta=False) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr( + self.epilog_threads * values_per_thread + ) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + + epi_tid * values_per_thread + + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + +import torch as _torch + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [ + leading_dim + ] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic( + _make_layout_tensor(shape, dtype, leading_dim), leading_dim + ) + for shape, leading_dim in zip( + ((batch, n, k), (batch, k, m), (batch, n, m)), + (a_leading, b_leading, c_leading), + strict=True, + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple( + _detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap) + ) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip( + (a_swap, b_swap, c_swap), leading_dims, strict=True + ) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided( + size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0) + ), + 1, + 2, + ), + leading_dims, + ) + + +# Tactic hashes all compile-time tile, split, and stage fields. +_SPLITK_COMPILE_CACHE: dict = {} + + +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + key = ( + dtype, + tactic, + use_pdl, + has_bias, + *leading_dims, + ) + cached = _SPLITK_COMPILE_CACHE.get(key) + if cached is not None: + return cached + + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError( + f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}" + ) + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + _SPLITK_COMPILE_CACHE[key] = compiled + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D tensors") + if a.device.type != "cuda" or any(tensor.device != a.device for tensor in tensors): + raise ValueError("all tensors must be on the same CUDA device") + if a.dtype not in _SUPPORTED_TORCH_DTYPES or any( + tensor.dtype != a.dtype for tensor in tensors + ): + raise ValueError("a, b, out, and bias must share BF16 or FP16 dtype") + if any( + not (tensor.is_contiguous() or tensor.t().is_contiguous()) + for tensor in (a, b, out) + ): + raise ValueError( + "a, b, and out must be dense row-major or column-major matrices" + ) + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError( + f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}" + ) + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and ( + bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous() + ): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 0cebb241b5fa..ee0278869cb6 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -911,9 +911,7 @@ def post_reorder_deepgemm_triton_kernel( BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr, ): - """`expert_id >= 0` includes the shared expert at num_experts (padding=-1); don't - switch to the cutlass `!= num_local_experts` gate. routed_scaling_factor is folded into the store. - """ + """Accumulate valid permuted rows; routed_scaling_factor is folded into the store.""" OutDtype = output_ptr.dtype.element_ty offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) @@ -935,9 +933,8 @@ def post_reorder_deepgemm_triton_kernel( sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for idx in range(topk): - expert_id = tl.load(token_topk_ids_ptr + idx) - if expert_id >= 0: - dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64) + dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64) + if dst_idx >= 0: weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32) load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32) @@ -1074,6 +1071,8 @@ def _fwd_kernel_ep_scatter_2( output_index, output_index_stride0, output_index_stride1, + expert_start, + num_experts, topk_num: tl.constexpr, HIDDEN_SIZE: tl.constexpr, HIDDEN_SIZE_PAD: tl.constexpr, @@ -1105,17 +1104,24 @@ def _fwd_kernel_ep_scatter_2( for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4): topk_index = topk_idx_int32.to(tl.int64) - expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) - if expert_id >= 0: + global_expert_id = tl.load( + recv_topk + token_id * recv_topk_stride0 + topk_index + ) + expert_id = global_expert_id - expert_start + output_index_ptr = ( + output_index + token_id * output_index_stride0 + topk_index + ) + valid = (expert_id >= 0) & (expert_id < num_experts) + # The post-permute path uses this index as the validity sentinel. + # Initialize non-local/padding lanes in this same kernel. + tl.store(output_index_ptr, -1) + if valid: dest_token_index_int32 = tl.atomic_add( expert_start_loc + expert_id, 1, sem=ATOMIC_ADD_SEM ) dest_token_index = dest_token_index_int32.to(tl.int64) - tl.store( - output_index + token_id * output_index_stride0 + topk_index, - dest_token_index_int32, - ) + tl.store(output_index_ptr, dest_token_index_int32) output_tensor_ptr = ( output_tensor + dest_token_index * output_tensor_stride0 ) @@ -1148,6 +1154,7 @@ def ep_scatter( output_index: torch.Tensor, scale_ue8m0: bool = False, quant_block_size: int = 128, + expert_start: int = 0, ): BLOCK_E = 128 # token num of per expert is aligned to 128 BLOCK_D = quant_block_size # block size of quantization @@ -1208,6 +1215,8 @@ def ep_scatter( output_index, output_index.stride(0), output_index.stride(1), + expert_start, + num_experts, topk_num=recv_topk.shape[1], num_warps=num_warps, HIDDEN_SIZE=hidden_size, @@ -1221,6 +1230,144 @@ def ep_scatter( return +@triton.jit +def _fwd_kernel_ep_scatter_psum_init( + psum_num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + BLOCK_E: tl.constexpr, +): + cur_expert = tl.program_id(0) + cur_end = tl.load(psum_num_recv_tokens_per_expert + cur_expert) + cur_start = tl.load( + psum_num_recv_tokens_per_expert + cur_expert - 1, + mask=cur_expert > 0, + other=0, + ) + cur_token_num = cur_end - cur_start + tl.store(expert_start_loc + cur_expert, cur_start) + + off_expert = tl.arange(0, BLOCK_E) + for start_m in tl.range(0, cur_token_num, BLOCK_E, num_stages=4): + # cur_token_num need not be a multiple of BLOCK_E; mask the tail block so + # the final partial iteration does not write past this expert's region + # (which is packed right up against the next expert) and corrupt it. + idx = cur_start + start_m + off_expert + tl.store(m_indices + idx, cur_expert, mask=idx < cur_end) + + +@torch.no_grad() +def ep_scatter_from_psum( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, + scale_ue8m0: bool = False, +): + BLOCK_E = 128 + BLOCK_D = 128 + num_warps = 8 + num_experts = psum_num_recv_tokens_per_expert.shape[0] + hidden_size = recv_x.shape[1] + scale_hidden_size = hidden_size // BLOCK_D + if scale_ue8m0: + scale_hidden_size = ceil_div(scale_hidden_size, 4) + + assert m_indices.shape[0] % BLOCK_E == 0 + is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 + if is_fp8: + assert recv_x_scale.dtype == output_tensor_scale.dtype + assert ( + recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size + ) + + _fwd_kernel_ep_scatter_psum_init[(num_experts,)]( + psum_num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_warps=num_warps, + BLOCK_E=BLOCK_E, + ) + + grid = min(recv_topk.shape[0], 1024 * 8) + _fwd_kernel_ep_scatter_2[(grid,)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0) if is_fp8 else 0, + recv_x_scale.stride(1) if is_fp8 else 0, + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0) if is_fp8 else 0, + output_tensor_scale.stride(1) if is_fp8 else 0, + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + num_warps=num_warps, + HIDDEN_SIZE=hidden_size, + HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), + SCALE_HIDDEN_SIZE=scale_hidden_size, + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), + ATOMIC_ADD_SEM=None if not _is_musa else "relaxed", + IS_FP8=is_fp8, + ) + return + + +@triton.jit +def _fwd_kernel_ep_expand_m_indices_init( + psum_num_recv_tokens_per_expert, + m_indices, + BLOCK_E: tl.constexpr, +): + cur_expert = tl.program_id(0) + cur_end = tl.load(psum_num_recv_tokens_per_expert + cur_expert) + prev_end = tl.load( + psum_num_recv_tokens_per_expert + cur_expert - 1, + mask=cur_expert > 0, + other=0, + ) + cur_start = ((prev_end + BLOCK_E - 1) // BLOCK_E) * BLOCK_E + aligned_end = ((cur_end + BLOCK_E - 1) // BLOCK_E) * BLOCK_E + + off_expert = tl.arange(0, BLOCK_E) + for start_m in tl.range(0, aligned_end - cur_start, BLOCK_E, num_stages=4): + idx = cur_start + start_m + off_expert + tl.store(m_indices + idx, cur_expert, mask=idx < aligned_end) + + +@torch.no_grad() +def ep_expand_init_m_indices_from_psum( + psum_num_recv_tokens_per_expert: torch.Tensor, + m_indices: torch.Tensor, +): + BLOCK_E = 128 + num_warps = 8 + num_experts = psum_num_recv_tokens_per_expert.shape[0] + assert m_indices.shape[0] % BLOCK_E == 0 + _fwd_kernel_ep_expand_m_indices_init[(num_experts,)]( + psum_num_recv_tokens_per_expert, + m_indices, + num_warps=num_warps, + BLOCK_E=BLOCK_E, + ) + return + + @triton.jit def _fwd_kernel_ep_gather( total_token_num, @@ -1404,12 +1551,13 @@ def tma_align_input_scale(input_scale: torch.Tensor): @triton.jit def fused_moe_dispatch_index_triton_kernel( - topk_ids_ptr, # flat (num_toks,) int32; -1 = padding (drives the `expert >= 0` gate) + topk_ids_ptr, # flat (num_toks,) int32; global or already-local expert IDs src2dst_ptr, masked_m_ptr, m_max, num_toks, num_experts, + expert_start, BLOCK_SIZE: tl.constexpr, ZERO_INIT: tl.constexpr, ): @@ -1427,19 +1575,24 @@ def fused_moe_dispatch_index_triton_kernel( tl.debug_barrier() offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offs < num_toks - expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) - valid = mask & (expert >= 0) + global_expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) + expert = global_expert - expert_start + valid = mask & (expert >= 0) & (expert < num_experts) # Clamp masked lanes to bin 0 so the masked atomic's pointer stays in-bounds. expert_safe = tl.where(valid, expert, 0) offset = tl.atomic_add(masked_m_ptr + expert_safe, 1, mask=valid) dst = expert_safe * m_max + offset - tl.store(src2dst_ptr + offs, dst, mask=valid) + # post_reorder checks src2dst rather than topk_ids. Explicitly initialize + # padding/non-local lanes to -1 in this same kernel so they cannot consume + # an uninitialized positive offset and add a bogus expert contribution. + tl.store(src2dst_ptr + offs, tl.where(valid, dst, -1), mask=mask) def fused_moe_dispatch_index( topk_ids: torch.Tensor, num_local_experts: int, m_max: int, + expert_start: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: num_toks = topk_ids.numel() src2dst = torch.empty(num_toks, device=topk_ids.device, dtype=torch.int32) @@ -1464,6 +1617,7 @@ def fused_moe_dispatch_index( m_max, num_toks, num_local_experts, + expert_start, BLOCK_SIZE=BLOCK_SIZE, ZERO_INIT=single_block, ) @@ -1499,9 +1653,8 @@ def fill_gateup_input_triton_kernel( vec = tl.arange(0, BLOCK_SIZE) for idx in range(topk): - expert_id = tl.load(topk_ids_ptr + idx) - if expert_id >= 0: - dst_idx_int32 = tl.load(src2dst_ptr + idx) + dst_idx_int32 = tl.load(src2dst_ptr + idx) + if dst_idx_int32 >= 0: dst_idx = dst_idx_int32.to(tl.int64) dst_ptr = gateup_input_ptr + dst_idx * hidden_size for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): @@ -1543,6 +1696,7 @@ def moe_ep_deepgemm_preprocess( block_shape, output_dtype: torch.dtype = torch.float8_e4m3fn, use_mxfp8: bool = False, + expert_start: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165 m_max = (hidden_states.size(0) // 256 + 1) * 256 @@ -1562,12 +1716,16 @@ def moe_ep_deepgemm_preprocess( # correctness is unconditional (m_cap >= max(masked_m) by # construction, and the final src2dst below is built with the same # capped stride). - masked_m_probe, _ = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max) + masked_m_probe, _ = fused_moe_dispatch_index( + topk_ids, num_local_experts, m_max, expert_start=expert_start + ) m_cap = (int(masked_m_probe.max().item()) + 255) // 256 * 256 m_max = min(m_max, max(m_cap, 256)) expected_m = (topk_ids.numel() - 1) // num_local_experts + 1 - masked_m, src2dst = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max) + masked_m, src2dst = fused_moe_dispatch_index( + topk_ids, num_local_experts, m_max, expert_start=expert_start + ) gateup_input = torch.empty( (num_local_experts, m_max, hidden_states.size(1)), @@ -2003,6 +2161,267 @@ def fp8_per_token_to_per_tensor_quant_triton( ) +# --------------------------------------------------------------------------- +# DeepEP v2 decode masked-GEMM bridge: repack the expanded expert-packed +# dispatch buffer into a regular [E_local, max_m, hidden] slab so DeepGEMM's +# *masked* grouped GEMM can bound compute by per-expert real counts (masked_m) +# instead of the dispatch capacity. All-GPU, static shapes -> cuda-graph safe. +# Expanded psum semantics (DeepEP v2): psum[e] = align(psum[e-1], ALIGN) + count_e, +# so expert e occupies recv rows [align(psum[e-1]) : psum[e]); count_e real tokens. +# Non-expand (contiguous) psum semantics differ: psum[e] is the inclusive prefix +# sum of alignment-PADDED counts, so every psum[e] is a multiple of ALIGN and +# psum[e-1] is expert e's aligned group start (consumed by ep_scatter_from_psum). +# --------------------------------------------------------------------------- + +_EPV2_REPACK_WORKERS_PER_EXPERT = 64 + + +@triton.jit +def _fwd_kernel_expand_to_masked_slab( + psum_ptr, + recv_x_ptr, + recv_x_stride0, + recv_x_scale_ptr, + recv_x_scale_stride0, + recv_x_scale_stride1, + slab_ptr, + slab_stride0, + slab_scale_ptr, + slab_scale_stride0, + masked_m_ptr, + overflow_ptr, + MAX_M: tl.constexpr, + ALIGN: tl.constexpr, + HIDDEN: tl.constexpr, + HIDDEN_PAD: tl.constexpr, + SCALE_HIDDEN: tl.constexpr, + SCALE_HIDDEN_PAD: tl.constexpr, + IS_FP8: tl.constexpr, + CHECK_OVERFLOW: tl.constexpr, + NUM_WORKERS: tl.constexpr, +): + # Keep a fixed worker pool per expert and let each worker walk only real rows. + # This avoids launching cdiv(MAX_M, BLOCK_M) programs for a conservative + # max_m when decode traffic contains only a few rows per expert. The grid is + # still static and therefore cuda-graph safe. + e = tl.program_id(0) + worker = tl.program_id(1) + prev_end = tl.load(psum_ptr + e - 1, mask=e > 0, other=0) + start = ((prev_end + ALIGN - 1) // ALIGN) * ALIGN + end = tl.load(psum_ptr + e) + raw_count = end - start + count = tl.minimum(raw_count, MAX_M) + if worker == 0: + tl.store(masked_m_ptr + e, count) + if CHECK_OVERFLOW: + # Eager execution reports an invalid bound instead of truncating. + # Graph replay uses the proven cap * ep_size upper bound and omits + # this host-observable flag and its per-layer reset kernel. + ovf = tl.arange(0, 1) + tl.store(overflow_ptr + ovf, 1, mask=raw_count > MAX_M) + off = tl.arange(0, HIDDEN_PAD) + mask = off < HIDDEN + off_s = tl.arange(0, SCALE_HIDDEN_PAD) + mask_s = off_s < SCALE_HIDDEN + for j in tl.range(worker, count, NUM_WORKERS): + src = (start + j).to(tl.int64) + dst = (e * MAX_M + j).to(tl.int64) + v = tl.load(recv_x_ptr + src * recv_x_stride0 + off, mask=mask) + tl.store(slab_ptr + dst * slab_stride0 + off, v, mask=mask) + if IS_FP8: + vs = tl.load( + recv_x_scale_ptr + + src * recv_x_scale_stride0 + + off_s * recv_x_scale_stride1, + mask=mask_s, + ) + # mn-major write: physical layout [E, SCALE_HIDDEN, MAX_M], element + # (e, s, j). Viewed as [E, MAX_M, SCALE_HIDDEN] this is the mn-major + # TMA-aligned layout deep_gemm wants, so the GEMM-side transpose + # (get_mn_major_tma_aligned_tensor) becomes a no-op. + tl.store( + slab_scale_ptr + e * SCALE_HIDDEN * MAX_M + off_s * MAX_M + j, + vs, + mask=mask_s, + ) + + +@torch.no_grad() +def expand_to_masked_slab( + recv_x: torch.Tensor, + recv_x_scale, + psum_num_recv_tokens_per_expert: torch.Tensor, + num_local_experts: int, + max_m: int, + expert_alignment: int, +): + """expanded [total, hidden] -> ([E_local, max_m, hidden], [E_local, max_m, sh] or None, masked_m[E_local]).""" + hidden = recv_x.shape[1] + is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 + slab = torch.empty( + (num_local_experts * max_m, hidden), device=recv_x.device, dtype=recv_x.dtype + ) + masked_m = torch.empty( + (num_local_experts,), device=recv_x.device, dtype=torch.int32 + ) + check_overflow = not torch.cuda.is_current_stream_capturing() + overflow = ( + torch.zeros((1,), device=recv_x.device, dtype=torch.int32) + if check_overflow + else masked_m + ) + if is_fp8: + sh = recv_x_scale.shape[1] + # mn-major slab_scale: store physically as [E, sh, max_m] (contiguous), + # return a [E, max_m, sh] view with mn-major stride. This matches + # deep_gemm's mn-major TMA-aligned scale layout, so the per-layer + # get_mn_major_tma_aligned_tensor call on the GEMM side is a no-op. + # (That call still runs and would transpose if the layout ever failed to + # match, so correctness does not depend on this optimization.) + slab_scale = torch.empty( + (num_local_experts * sh, max_m), + device=recv_x.device, + dtype=recv_x_scale.dtype, + ) + scale_arg = recv_x_scale + scale_s0 = recv_x_scale.stride(0) + scale_s1 = recv_x_scale.stride(1) + slab_scale_s0 = 0 # unused: scale write uses mn-major addressing + else: + sh = 1 + slab_scale = None + scale_arg = recv_x + scale_s0 = 0 + scale_s1 = 0 + slab_scale_s0 = 0 + num_workers = min(max_m, _EPV2_REPACK_WORKERS_PER_EXPERT) + _fwd_kernel_expand_to_masked_slab[(num_local_experts, num_workers)]( + psum_num_recv_tokens_per_expert, + recv_x, + recv_x.stride(0), + scale_arg, + scale_s0, + scale_s1, + slab, + slab.stride(0), + slab_scale if is_fp8 else scale_arg, + slab_scale_s0, + masked_m, + overflow, + MAX_M=max_m, + ALIGN=expert_alignment, + HIDDEN=hidden, + HIDDEN_PAD=triton.next_power_of_2(hidden), + SCALE_HIDDEN=sh, + SCALE_HIDDEN_PAD=triton.next_power_of_2(sh), + IS_FP8=is_fp8, + CHECK_OVERFLOW=check_overflow, + NUM_WORKERS=num_workers, + num_warps=4, + ) + # Outside cuda graph capture, fail fast on slab overflow rather than return a + # silently truncated result. During capture we skip the host read to keep the + # path graph-safe; the eager warmup forward validates representative shapes. + # Safety under graph replay therefore relies on the static upper bound + # max_m = cap * ep_group_size holding: each rank sends at most `cap` tokens + # (enforced by the dispatch-entry assert) and a token contributes at most once + # per local expert, so no expert can exceed max_m. If those invariants change, + # graph replay would NOT fail-fast on overflow — re-validate before relying on it. + if check_overflow and int(overflow.item()) != 0: + raise RuntimeError( + f"DeepEP v2 masked slab overflow: an expert received more than max_m=" + f"{max_m} tokens; increase " + f"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK." + ) + slab = slab.view(num_local_experts, max_m, hidden) + if is_fp8: + # physical [E, sh, max_m] -> [E, max_m, sh] view with mn-major stride (no copy) + slab_scale = slab_scale.view(num_local_experts, sh, max_m).transpose(1, 2) + return slab, slab_scale, masked_m + + +@triton.jit +def _fwd_kernel_masked_slab_to_expand( + psum_ptr, + slab_ptr, + slab_stride0, + out_ptr, + out_stride0, + weight_ptr, + MAX_M: tl.constexpr, + ALIGN: tl.constexpr, + HIDDEN: tl.constexpr, + HIDDEN_PAD: tl.constexpr, + HAS_W: tl.constexpr, + NUM_WORKERS: tl.constexpr, +): + # Fixed worker pool; see _fwd_kernel_expand_to_masked_slab. cuda-graph safe. + e = tl.program_id(0) + worker = tl.program_id(1) + prev_end = tl.load(psum_ptr + e - 1, mask=e > 0, other=0) + start = ((prev_end + ALIGN - 1) // ALIGN) * ALIGN + end = tl.load(psum_ptr + e) + count = end - start + count = tl.minimum(count, MAX_M) + off = tl.arange(0, HIDDEN_PAD) + mask = off < HIDDEN + for j in tl.range(worker, count, NUM_WORKERS): + src = (e * MAX_M + j).to(tl.int64) + dst = (start + j).to(tl.int64) + v = tl.load(slab_ptr + src * slab_stride0 + off, mask=mask) + if HAS_W: + w = tl.load(weight_ptr + dst) + v = (v.to(tl.float32) * w).to(v.dtype) + tl.store(out_ptr + dst * out_stride0 + off, v, mask=mask) + + +@torch.no_grad() +def masked_slab_to_expand( + slab: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + total_expanded_tokens: int, + expert_alignment: int, + topk_weights=None, +): + """[E_local, max_m, hidden] masked-GEMM output -> [total, hidden] expanded order. + + Only real rows are written; padding rows are uninitialized (the output is + torch.empty) and are never read -- combine consumes only real rows via handle + metadata. When topk_weights is given ([total_expanded], per expanded row), the + top-k weight is fused into the copy so the weighted-combine multiply happens + only on real rows (not the worst-case buffer). + """ + num_local_experts, max_m, hidden = slab.shape + # combine reads only real rows via handle metadata, so padding need not be + # zeroed -> use empty to skip the worst-case-buffer memset. + out = torch.empty( + (total_expanded_tokens, hidden), device=slab.device, dtype=slab.dtype + ) + slab2d = slab.view(num_local_experts * max_m, hidden) + has_w = topk_weights is not None + if has_w: + weight_arg = topk_weights.reshape(-1).to(torch.float32).contiguous() + else: + weight_arg = slab2d # dummy, unused + num_workers = min(max_m, _EPV2_REPACK_WORKERS_PER_EXPERT) + _fwd_kernel_masked_slab_to_expand[(num_local_experts, num_workers)]( + psum_num_recv_tokens_per_expert, + slab2d, + slab2d.stride(0), + out, + out.stride(0), + weight_arg, + MAX_M=max_m, + ALIGN=expert_alignment, + HIDDEN=hidden, + HIDDEN_PAD=triton.next_power_of_2(hidden), + HAS_W=has_w, + NUM_WORKERS=num_workers, + num_warps=4, + ) + return out + + def moe_permute( inputs: torch.Tensor, topk_ids: torch.Tensor, diff --git a/python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py b/python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py index d67786eca262..75d2966e99b1 100644 --- a/python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py +++ b/python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py @@ -89,8 +89,9 @@ def per_token_quant_fp8_ue8m0_scatter( then writes them to each of the token's ``topk`` destination rows: ``gateup_input`` fp8 ``[E, m_max, hidden]`` (row ``src2dst[token, i]``) ``gateup_input_scale`` int32 ``[E, hidden//group//4, m_max]`` (MN-major; byte-scattered) - Slots with ``topk_ids[token, i] < 0`` are skipped. Byte-identical to the - two-kernel path on every written row. + Slots with ``src2dst[token, i] < 0`` are skipped. Byte-identical to the + two-kernel path on every written row. ``topk_ids`` remains in the ABI for + compatibility with already-compiled JIT modules. """ assert x.is_cuda and x.dtype == torch.bfloat16 and x.dim() == 2 assert x.is_contiguous() diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 7ef95918167e..c65b8ee21657 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1926,6 +1926,21 @@ def _flashinfer_allreduce_fusion_auto_enable(view: Any) -> dict: moe_a2a_backend (after the DeepSeek CP and a2a declarations), exactly like the legacy tail block.""" model_arch = view.get_model_config().hf_config.architectures[0] + if envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() and model_arch in { + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + }: + # The Qwen-specific backend owns one workspace for both ordinary AR + # and MoE finalize patterns. Do not allocate the legacy FlashInfer + # TRTLLM/MNNVL workspace or let it become a graph-path fallback. + if view.flashinfer_allreduce_fusion_backend is not None: + logger.warning( + "SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION owns both Qwen3.5 " + "AllReduce fusion patterns; suppressing the separately configured " + "--flashinfer-allreduce-fusion-backend=%s", + view.flashinfer_allreduce_fusion_backend, + ) + return {"flashinfer_allreduce_fusion_backend": None} if ( view.flashinfer_allreduce_fusion_backend is None and model_arch in _FLASHINFER_ALLREDUCE_FUSION_ARCHS @@ -2455,6 +2470,11 @@ def _a2a_fusion_adjustments(view: Any) -> dict: "Flashinfer MoE A2A is enabled. --disable-shared-experts-fusion is automatically set." ) return {"disable_shared_experts_fusion": True} + if view.moe_a2a_backend == "deepep_v2": + logger.warning( + "DeepEP v2 MoE A2A is enabled. --disable-shared-experts-fusion is automatically set." + ) + return {"disable_shared_experts_fusion": True} return {} @@ -2483,6 +2503,7 @@ def _cutlass_moe_env_override(view: Any) -> dict: "flashinfer", "mori", "pplx", + "deepep_v2", } ) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index c25d4755a166..c939467efe98 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1369,6 +1369,7 @@ def _get_bootstrap_info_from_server( response = _get_bootstrap_session(self.bootstrap_addr).get(url, timeout=5) if response.status_code == 200: bootstrap_info = response.json() + bootstrap_info["pp_rank"] = int(target_pp_rank) return bootstrap_info else: logger.error( diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 823bd94605ae..dc50cee64af9 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -108,11 +108,15 @@ def register_wm_subscriber(self, receiver, session_id: str) -> None: self._wm_subscribers[key] = (receiver, session_id) def num_writers_for(self, receiver) -> int: - """Compute num_writers for a specific request based on its prefill TP.""" - prefill_tp = receiver.prefill_info.attn_tp_size + """Compute all TP and PP writers expected for a staging chunk.""" + prefill_info = receiver.prefill_info + prefill_tp = prefill_info.attn_tp_size if prefill_tp > self.decode_tp: - return prefill_tp // max(1, self.decode_tp) - return 1 + tp_writers = prefill_tp // max(1, self.decode_tp) + else: + tp_writers = 1 + pp_writers = prefill_info.pp_size // self.kv_manager.pp_size + return tp_writers * pp_writers @classmethod def create(cls, kv_manager, scheduler, tp_rank: int) -> DecodeStagingHandler: @@ -535,10 +539,14 @@ class StagingRegisterInfo: base_ptr: int = 0 total_size: int = 0 + # Staging slot order is [all K layers, all V layers], which differs from the + # kv_data_ptrs order once draft KV buffers are appended. Empty when the peer + # predates this field; callers then fall back to kv_layer_ids. + slot_layer_ids: List[int] = dataclasses.field(default_factory=list) @classmethod def from_zmq_fields( - cls, msg: list, msg_start_offset: int + cls, msg: list, msg_start_offset: int, slot_ids_index: Optional[int] = None ) -> Optional[StagingRegisterInfo]: i = msg_start_offset base_ptr = ( @@ -551,7 +559,17 @@ def from_zmq_fields( ) if base_ptr == 0 and total_size == 0: return None - return cls(base_ptr=base_ptr, total_size=total_size) + slot_layer_ids: List[int] = [] + if ( + slot_ids_index is not None + and len(msg) > slot_ids_index + and len(msg[slot_ids_index]) > 0 + ): + raw = msg[slot_ids_index] + slot_layer_ids = list(struct.unpack(f"{len(raw) // 8}Q", raw)) + return cls( + base_ptr=base_ptr, total_size=total_size, slot_layer_ids=slot_layer_ids + ) class PrefillStagingStrategy: @@ -643,7 +661,13 @@ def transfer( target_info.dst_tp_rank, target_info.dst_attn_tp_size, target_info.dst_kv_item_len, + target_info.dst_kv_layer_ids, staging_buffer=self.staging_buffer, + dst_slot_layer_ids=( + target_info.staging.slot_layer_ids + if target_info.staging is not None + else None + ), ) except Exception as e: raise RuntimeError( @@ -742,6 +766,7 @@ def handle_staging_req( chunk_idx = int(msg[2].decode("ascii")) chunk_num_pages = int(msg[3].decode("ascii")) session_id = msg[4].decode("ascii") + requester_pp_rank = int(msg[5].decode("ascii")) if len(msg) > 5 else None if staging_allocator is None: logger.warning( @@ -824,6 +849,8 @@ def handle_staging_req( bootstrap_infos = room_bootstrap.get(room) if bootstrap_infos: for bi in bootstrap_infos: + if requester_pp_rank is not None and bi["pp_rank"] != requester_pp_rank: + continue try: sock, lock = receiver._connect_to_bootstrap_server(bi) with lock: @@ -849,6 +876,7 @@ def prefetch_staging_reqs( chunked_prefill_size: int, staging_requested: set, prefetch_sockets: dict, + requester_pp_rank: Optional[int] = None, ) -> None: """Send STAGING_REQ for all chunks before the prefill forward starts. @@ -894,14 +922,15 @@ def prefetch_staging_reqs( sock.setsockopt(zmq.IPV6, 1) sock.connect(ep) prefetch_sockets[ep] = sock - prefetch_sockets[ep].send_multipart( - [ - b"STAGING_REQ", - str(room).encode("ascii"), - str(chunk_idx).encode("ascii"), - str(chunk_pages).encode("ascii"), - session_id.encode("ascii"), - ] - ) + request = [ + b"STAGING_REQ", + str(room).encode("ascii"), + str(chunk_idx).encode("ascii"), + str(chunk_pages).encode("ascii"), + session_id.encode("ascii"), + ] + if requester_pp_rank is not None: + request.append(str(requester_pp_rank).encode("ascii")) + prefetch_sockets[ep].send_multipart(request) except Exception: staging_requested.discard(stg_key) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 0bdd51034ad9..4f1561b3a923 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -50,6 +50,8 @@ ReqToMetadataIdxAllocator, TransferBackend, _is_fake_transfer, + build_kv_layer_ids, + build_staging_slot_metadata, get_dsv4_c128_state_indices, get_kv_class, is_dsv4_c128_online_enabled, @@ -455,6 +457,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_data_lens += device_kv_data_lens[c4_layer_num:] kv_item_lens += device_kv_item_lens[c4_layer_num:] kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) + num_draft_entries = 0 if self.draft_token_to_kv_pool is not None: # We should also transfer draft model kv cache. The indices are # always shared with a target model. @@ -465,15 +468,16 @@ def _init_kv_manager(self) -> CommonKVManager: kv_data_lens += draft_kv_data_lens kv_item_lens += draft_kv_item_lens kv_data_mem_kinds += ["VRAM"] * len(draft_kv_data_ptrs) + num_draft_entries = len(draft_kv_data_ptrs) kv_args.kv_data_ptrs = kv_data_ptrs kv_args.kv_data_lens = kv_data_lens kv_args.kv_item_lens = kv_item_lens - kv_args.kv_layer_ids = ( - self.token_to_kv_pool.get_kv_layer_ids() - if self.draft_token_to_kv_pool is None - and hasattr(self.token_to_kv_pool, "get_kv_layer_ids") - else [] + kv_args.kv_layer_ids = build_kv_layer_ids( + token_to_kv_pool=self.token_to_kv_pool, + draft_token_to_kv_pool=self.draft_token_to_kv_pool, + num_draft_entries=num_draft_entries, + num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) if self.transfer_backend == TransferBackend.NIXL: kv_args.kv_data_mem_kinds = kv_data_mem_kinds @@ -508,12 +512,24 @@ def _init_kv_manager(self) -> CommonKVManager: per_rank_kv_heads = getattr(kv_pool_for_heads, "head_num", 0) if per_rank_kv_heads > 0: kv_args.kv_head_num = per_rank_kv_heads - kv_args.total_kv_head_num = per_rank_kv_heads * attn_tp_size + kv_args.total_kv_head_num = ( + self.scheduler.model_config.get_total_num_kv_heads() + ) if hasattr(kv_manager, "set_kv_buffer_tensors"): kv_pool = kv_pool_for_heads - if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"): + staging_slots = build_staging_slot_metadata( + kv_layer_ids=kv_args.kv_layer_ids, + num_draft_entries=num_draft_entries, + kv_pool=kv_pool, + draft_kv_pool=self.draft_token_to_kv_pool, + ) + if staging_slots is not None: + k_buffers, v_buffers, slot_layer_ids = staging_slots kv_manager.set_kv_buffer_tensors( - kv_pool.k_buffer, kv_pool.v_buffer, kv_pool.page_size + k_buffers, + v_buffers, + kv_pool.page_size, + slot_layer_ids=slot_layer_ids, ) return kv_manager diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 353e50bbe0c2..d93f7b8b503b 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -184,7 +184,7 @@ def from_zmq(cls, msg: List[bytes]): int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0 ), # Note: always put the staging field at the final - staging=StagingRegisterInfo.from_zmq_fields(msg, 14), + staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18), ) @@ -320,11 +320,20 @@ def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): self._staging_ctx.room_bootstrap[room] = bootstrap_infos self._staging_ctx.room_receivers[room] = receiver - def set_kv_buffer_tensors(self, k_buffers: list, v_buffers: list, page_size: int): + def set_kv_buffer_tensors( + self, + k_buffers: list, + v_buffers: list, + page_size: int, + slot_layer_ids: Optional[List[int]] = None, + ): + # slot_layer_ids follows the staging slot order (every k_buffer, then + # every v_buffer), which is not kv_args.kv_layer_ids once a draft exists. self.kv_buffer_tensors = { "k_buffers": k_buffers, "v_buffers": v_buffers, "page_size": page_size, + "slot_layer_ids": list(slot_layer_ids or []), } def _init_staging_buffers(self, count: int): @@ -510,6 +519,7 @@ def _prefetch_staging_reqs(self, room: int): get_schedule().chunked_prefill_size, self._staging_ctx.prefetch_requested, self._staging_ctx.prefetch_sockets, + requester_pp_rank=self.pp_rank, ) def send_kvcache_staged( @@ -521,7 +531,9 @@ def send_kvcache_staged( dst_tp_rank: int, dst_attn_tp_size: int, dst_kv_item_len: int, + dst_layer_ids: List[int], staging_buffer=None, + dst_slot_layer_ids: Optional[List[int]] = None, ) -> int: """Transfer KV cache via staging buffers (gather -> bulk RDMA -> scatter on decode).""" from sglang.srt.disaggregation.common.staging_buffer import ( @@ -553,7 +565,27 @@ def send_kvcache_staged( num_tokens = len(prefill_kv_indices) * page_size per_layer_bytes = num_tokens * num_heads_to_send * head_dim * dtype_size - per_rank_bytes = per_layer_bytes * num_layers * 2 + local_bytes = per_layer_bytes * num_layers * 2 + + if self.pp_size > 1: + # Pair staging slots, not kv_data_ptrs entries: the gather lays out + # [every k_buffer, every v_buffer], which stops matching kv_layer_ids + # once draft KV buffers are appended. + src_slot_ids = ( + self.kv_buffer_tensors.get("slot_layer_ids") + or self.kv_args.kv_layer_ids + ) + dst_slot_ids = dst_slot_layer_ids or dst_layer_ids + pairs = build_transfer_entry_pairs( + src_slot_ids, + dst_slot_ids, + num_layers * 2, + len(dst_slot_ids), + ) + dst_num_layers = len(dst_slot_ids) // 2 + else: + pairs = None + dst_num_layers = num_layers num_writers, writer_rank_bytes, total_staging_needed = compute_staging_layout( self.attn_tp_size, @@ -562,21 +594,20 @@ def send_kvcache_staged( total_kv_heads, num_tokens, head_dim * dtype_size, - num_layers, + dst_num_layers, ) writer_idx = local_tp_rank % num_writers if num_writers > 1 else 0 rank_offset = sum(writer_rank_bytes[:writer_idx]) - if not staging_buffer.fits(per_rank_bytes): + if not staging_buffer.fits(local_bytes): logger.warning( - f"Prefill staging too small for {per_rank_bytes} bytes, falling back" + f"Prefill staging too small for {local_bytes} bytes, falling back" ) return -1 if dst_staging_size < total_staging_needed: logger.warning( f"Decode staging too small: need {total_staging_needed} bytes " - f"({num_writers if self.attn_tp_size > dst_attn_tp_size else 1} writers " - f"x {per_rank_bytes} bytes/rank), have {dst_staging_size}, falling back" + f"for {dst_num_layers} layers, have {dst_staging_size}, falling back" ) return -1 @@ -595,16 +626,29 @@ def send_kvcache_staged( self.kv_args.gpu_id, ) - dst_write_ptr = dst_staging_ptr + rank_offset - ret = self._transfer_data( - mooncake_session_id, - [(staging_buffer.get_ptr(), dst_write_ptr, per_rank_bytes)], - ) + if pairs is None: + transfer_blocks = [ + ( + staging_buffer.get_ptr(), + dst_staging_ptr + rank_offset, + local_bytes, + ) + ] + else: + transfer_blocks = [ + ( + staging_buffer.get_ptr() + src_idx * per_layer_bytes, + dst_staging_ptr + rank_offset + dst_idx * per_layer_bytes, + per_layer_bytes, + ) + for src_idx, dst_idx in pairs + ] + ret = self._transfer_data(mooncake_session_id, transfer_blocks) if ret != 0: raise RuntimeError( f"[Staging] Bulk RDMA transfer failed with ret={ret}. " f"src_ptr=0x{staging_buffer.get_ptr():x}, " - f"dst_ptr=0x{dst_write_ptr:x}, size={per_rank_bytes}. " + f"dst_ptr=0x{dst_staging_ptr + rank_offset:x}, size={local_bytes}. " f"The decode staging buffer may not be properly registered." ) return ret @@ -658,10 +702,19 @@ def _send_kvcache_generic( layers_params = None # Decode pp size should be equal to prefill pp size or 1 - if self.is_mla_backend or self.is_hybrid_mla_backend or force_flat: + # When both peers publish layer ids the pairing is exact, so prefer it + # over positional slicing regardless of backend; a plain-MHA model + # publishes no ids and is unaffected. + has_layer_ids = bool(src_layer_ids or dst_layer_ids) + if ( + self.is_mla_backend + or self.is_hybrid_mla_backend + or force_flat + or has_layer_ids + ): # Layer IDs map PP-local buffers to global decode entries. # Registrations without them retain the existing PP mapping. - if src_layer_ids or dst_layer_ids: + if has_layer_ids: pairs = build_transfer_entry_pairs( src_layer_ids, dst_layer_ids, @@ -916,6 +969,7 @@ def send_kvcache_slice( dst_attn_tp_size: int, dst_kv_item_len: int, executor: concurrent.futures.ThreadPoolExecutor, + dst_layer_ids: Optional[List[int]] = None, ): """ Sends KV cache slices from this Prefill rank to a target Decode rank, @@ -968,9 +1022,35 @@ def send_kvcache_slice( num_heads_to_send = dst_heads_per_rank dst_head_start_offset = 0 - src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = ( - self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs) - ) + src_data_ptrs = self.kv_args.kv_data_ptrs + src_layer_ids = self.kv_args.kv_layer_ids + if src_layer_ids or dst_layer_ids: + # Pair by layer id. Required once draft KV buffers are appended: the + # flat list is then no longer [K block, V block], so the half-split + # in get_mha_kv_ptrs_with_pp mislabels entries. + if any(l != src_kv_item_len for l in self.kv_args.kv_item_lens): + logger.error( + f"[{mooncake_session_id}] head-sliced transfer assumes one item " + f"length for every KV entry, got {set(self.kv_args.kv_item_lens)}" + ) + return -1 + layer_ptr_pairs = [ + (src_data_ptrs[i], dst_kv_ptrs[j]) + for i, j in build_transfer_entry_pairs( + src_layer_ids, + dst_layer_ids or [], + len(src_data_ptrs), + len(dst_kv_ptrs), + allow_positional_fallback=self.pp_size == 1, + ) + ] + else: + src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = ( + self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_kv_ptrs) + ) + layer_ptr_pairs = [ + (src_k_ptrs[i], dst_k_ptrs[i]) for i in range(layers_current_pp_stage) + ] + [(src_v_ptrs[i], dst_v_ptrs[i]) for i in range(layers_current_pp_stage)] # Calculate precise byte offset and length for the sub-slice within the token src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send @@ -1015,15 +1095,10 @@ def process_layer_tp_aware(src_layer_ptr, dst_layer_ptr): mooncake_session_id, src_addr_list, dst_addr_list, length_list ) - futures = [] - for i in range(layers_current_pp_stage): - futures.append( - executor.submit(process_layer_tp_aware, src_k_ptrs[i], dst_k_ptrs[i]) - ) - for i in range(layers_current_pp_stage): - futures.append( - executor.submit(process_layer_tp_aware, src_v_ptrs[i], dst_v_ptrs[i]) - ) + futures = [ + executor.submit(process_layer_tp_aware, src_layer_ptr, dst_layer_ptr) + for src_layer_ptr, dst_layer_ptr in layer_ptr_pairs + ] for future in concurrent.futures.as_completed(futures): status = future.result() @@ -1719,6 +1794,7 @@ def transfer_worker( target_rank_registration_info.dst_attn_tp_size, target_rank_registration_info.dst_kv_item_len, executor, + target_rank_registration_info.dst_kv_layer_ids, ) if ret != 0: with self.session_lock: @@ -2310,6 +2386,11 @@ def _register_kv_args(self) -> bool: else: packed_staging_base_ptr = b"" staging_total_size_str = b"" + staging_slots = getattr(self.kv_mgr, "kv_buffer_tensors", None) or {} + packed_staging_slot_layer_ids = b"".join( + struct.pack("Q", layer_id) + for layer_id in (staging_slots.get("slot_layer_ids") or []) + ) sock, lock = self._connect_to_bootstrap_server(bootstrap_info) try: @@ -2334,6 +2415,7 @@ def _register_kv_args(self) -> bool: staging_total_size_str, dst_dcp_size, dst_dcp_rank, + packed_staging_slot_layer_ids, ] ) except zmq.ZMQError: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 02db16862af3..5d1e17352929 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -43,6 +43,8 @@ MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, + build_kv_layer_ids, + build_staging_slot_metadata, get_dsv4_c128_state_indices, get_kv_class, is_aborted, @@ -68,6 +70,7 @@ release_kv_cache, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool from sglang.srt.observability.req_time_stats import set_schedule_time_batch from sglang.srt.runtime_context import get_disagg from sglang.srt.utils import is_npu @@ -95,6 +98,24 @@ def should_force_retry(req: Req) -> bool: return int.from_bytes(digest[:8], "big") < retry_prob * 2**64 +def _transfer_start_layer(*, pool, hf_text_config) -> int: + """Offset of this stage's first KV entry inside the peer's dense KV list. + + A hybrid-linear pool stores KV only for its full-attention layers, but its + ``start_layer`` is a global layer index that also counts linear layers. The + decode peer's pointer list is dense over full-attention layers, so the global + index over-shoots it. Translate to a full-attention-relative offset. + + The pool only knows this stage's own layer ids, so the count has to come from + the model-wide layer table. + """ + if not isinstance(pool, HybridLinearKVPool): + return pool.start_layer + return sum( + 1 for lid in hf_text_config.full_attention_layer_ids if lid < pool.start_layer + ) + + def maybe_release_metadata_buffer( req: Req, allocator: ReqToMetadataIdxAllocator ) -> None: @@ -169,10 +190,10 @@ def __init__( f"chunked_prefill_size that is a multiple of page_size " f"({page_size}); got {server_args.chunked_prefill_size}." ) - if self.pp_size > 1: - # Staging writer accounting has no pp dimension. + if self.pp_size > 1 and self.transfer_backend != TransferBackend.MOONCAKE: raise RuntimeError( - "SGLANG_DISAGG_STAGING_BUFFER does not support pp_size > 1." + "SGLANG_DISAGG_STAGING_BUFFER with pp_size > 1 is only " + "supported by Mooncake." ) if server_args.enable_prefill_context_parallel: # CP rewrites index_slice per rank, breaking the chunk grid. @@ -206,7 +227,10 @@ def _init_kv_manager(self) -> CommonKVManager: self.token_to_kv_pool.start_layer, ) if layer_shard_enabled - else self.token_to_kv_pool.start_layer + else _transfer_start_layer( + pool=self.token_to_kv_pool, + hf_text_config=self.scheduler.model_config.hf_text_config, + ) ) kv_args.mla_compression_ratios = None kv_data_ptrs, kv_data_lens, kv_item_lens = ( @@ -218,24 +242,27 @@ def _init_kv_manager(self) -> CommonKVManager: else getattr(self.token_to_kv_pool, "end_layer", None) ) - if self.draft_token_to_kv_pool is not None and transfer_draft_cache: + draft_kv_pool = self.draft_token_to_kv_pool if transfer_draft_cache else None + num_draft_entries = 0 + if draft_kv_pool is not None: # We should also transfer draft model kv cache. The indices are # always shared with a target model. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( - self.draft_token_to_kv_pool.get_contiguous_buf_infos() + draft_kv_pool.get_contiguous_buf_infos() ) kv_data_ptrs += draft_kv_data_ptrs kv_data_lens += draft_kv_data_lens kv_item_lens += draft_kv_item_lens + num_draft_entries = len(draft_kv_data_ptrs) kv_args.kv_data_ptrs = kv_data_ptrs kv_args.kv_data_lens = kv_data_lens kv_args.kv_item_lens = kv_item_lens - kv_args.kv_layer_ids = ( - self.token_to_kv_pool.get_kv_layer_ids() - if self.draft_token_to_kv_pool is None - and hasattr(self.token_to_kv_pool, "get_kv_layer_ids") - else [] + kv_args.kv_layer_ids = build_kv_layer_ids( + token_to_kv_pool=self.token_to_kv_pool, + draft_token_to_kv_pool=draft_kv_pool, + num_draft_entries=num_draft_entries, + num_hidden_layers=self.scheduler.model_config.num_hidden_layers, ) if not self.is_mla_backend: kv_args.kv_head_num = self.token_to_kv_pool.head_num @@ -282,11 +309,19 @@ def _init_kv_manager(self) -> CommonKVManager: kv_pool = self.token_to_kv_pool if hasattr(kv_pool, "full_kv_pool"): kv_pool = kv_pool.full_kv_pool - if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"): + staging_slots = build_staging_slot_metadata( + kv_layer_ids=kv_args.kv_layer_ids, + num_draft_entries=num_draft_entries, + kv_pool=kv_pool, + draft_kv_pool=draft_kv_pool, + ) + if staging_slots is not None: + k_buffers, v_buffers, slot_layer_ids = staging_slots kv_manager.set_kv_buffer_tensors( - kv_pool.k_buffer, - kv_pool.v_buffer, + k_buffers, + v_buffers, kv_pool.page_size, + slot_layer_ids=slot_layer_ids, ) return kv_manager @@ -684,6 +719,16 @@ def process_batch_result_disagg_prefill( logprob_pt = 0 assert batch.spec_info is result.next_draft_input draft_input = result.next_draft_input + draft_hidden_states_cpu = None + draft_dsa_topk_indices_cpu = None + if self.spec_algorithm.is_eagle() and draft_input is not None: + draft_hidden_states_cpu = draft_input.hidden_states.to( + "cpu", non_blocking=False + ) + if batch.spec_info.dsa_topk_indices is not None: + draft_dsa_topk_indices_cpu = batch.spec_info.dsa_topk_indices.to( + "cpu", non_blocking=False + ) # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue next_token_ids = result.next_token_ids.tolist() self.batch_result_processor.move_logprobs_to_cpu( @@ -718,12 +763,11 @@ def advance_logprob_pt(i: int, req: Req) -> None: if self.spec_algorithm.is_eagle() and draft_input is not None: req.output_topk_p = draft_input.topk_p[i] req.output_topk_index = draft_input.topk_index[i] - req.hidden_states_tensor = ( - draft_input.hidden_states[i].cpu().clone() - ) - dsa_topk_indices = batch.spec_info.dsa_topk_indices - if dsa_topk_indices is not None: - req.output_dsa_topk_indices = dsa_topk_indices[i].cpu().clone() + req.hidden_states_tensor = draft_hidden_states_cpu[i].clone() + if draft_dsa_topk_indices_cpu is not None: + req.output_dsa_topk_indices = draft_dsa_topk_indices_cpu[ + i + ].clone() else: req.output_dsa_topk_indices = None else: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index e6f35fac713e..7df7bfbda4ec 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -941,6 +941,65 @@ def build_transfer_entry_pairs( return [(i, i) for i in range(n_src)] +def build_kv_layer_ids( + *, + token_to_kv_pool, + draft_token_to_kv_pool, + num_draft_entries: int, + num_hidden_layers: int, +) -> List[int]: + """Global layer id for every entry in ``kv_args.kv_data_ptrs``. + + Draft KV buffers are appended after the target's, so they need ids of their + own: build_transfer_entry_pairs requires the id list to cover every entry, + and a target-only list would be rejected. The draft numbers its layers from + zero, which would collide with the target's, so its entries are remapped + into a reserved band above the target's layer range. Both PD peers run this + against the same draft config and so agree on the band. + + Returns [] for pools that cannot report ids, leaving the peers on positional + pairing. + """ + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + if not isinstance(token_to_kv_pool, HybridLinearKVPool): + return [] + layer_ids = token_to_kv_pool.get_kv_layer_ids() + if draft_token_to_kv_pool is None: + return layer_ids + + draft_ids = _draft_entry_layer_ids( + pool=draft_token_to_kv_pool, num_entries=num_draft_entries + ) + # Rank the draft's own ids by first appearance, so the band stays dense and + # contiguous whatever the draft config numbers its layers. + band_index = {lid: i for i, lid in enumerate(dict.fromkeys(draft_ids))} + return layer_ids + [num_hidden_layers + band_index[lid] for lid in draft_ids] + + +def _draft_entry_layer_ids(*, pool, num_entries: int) -> List[int]: + """One draft-local layer id per registered draft KV entry.""" + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + if isinstance(pool, HybridLinearKVPool): + ids = pool.get_kv_layer_ids() + else: + # Pools register k0..k(L-1) then v0..v(L-1), so ids repeat once per + # group; derive the group count rather than assuming MHA vs MLA. + if pool.layer_num <= 0 or num_entries % pool.layer_num != 0: + raise RuntimeError( + "Draft KV buffers must register a whole number of per-layer " + f"groups: entries={num_entries}, layers={pool.layer_num}" + ) + ids = list(range(pool.layer_num)) * (num_entries // pool.layer_num) + if len(ids) != num_entries: + raise RuntimeError( + "Draft KV layer ids must cover every registered entry: " + f"ids={len(ids)}, entries={num_entries}" + ) + return ids + + def resolve_dcp_dst_entry_indices( src_layer_ids: List[int], dst_layer_ids: List[int], @@ -966,6 +1025,52 @@ def resolve_dcp_dst_entry_indices( ] +def build_staging_slot_metadata( + *, + kv_layer_ids: List[int], + num_draft_entries: int, + kv_pool, + draft_kv_pool, +): + """Buffers and per-slot layer ids for the staging gather. + + The gather writes every k_buffer and then every v_buffer, while kv_layer_ids + follows kv_data_ptrs ([K target, V target, K draft, V draft]), so the two + orders diverge as soon as a draft pool is registered. + + Returns (k_buffers, v_buffers, slot_layer_ids), or None for a pool that has + no contiguous K/V tensors to stage. + """ + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool + + # A hybrid pool keeps its contiguous K/V tensors on the inner full-attention + # pool, and the draft pool is wrapped the same way. + if isinstance(kv_pool, HybridLinearKVPool): + kv_pool = kv_pool.full_kv_pool + if isinstance(draft_kv_pool, HybridLinearKVPool): + draft_kv_pool = draft_kv_pool.full_kv_pool + if not isinstance(kv_pool, MHATokenToKVPool): + return None + + ids = list(kv_layer_ids or []) + num_target = len(ids) - num_draft_entries + half = num_target // 2 + k_buffers, k_ids = list(kv_pool.k_buffer), ids[:half] + v_buffers, v_ids = list(kv_pool.v_buffer), ids[half:num_target] + + draft_half = num_draft_entries // 2 + if draft_half: + if not isinstance(draft_kv_pool, MHATokenToKVPool): + # An empty id list puts the sender back on kv_data_ptrs order, which + # is what staging did before draft KV existed. + return k_buffers, v_buffers, [] + k_buffers += list(draft_kv_pool.k_buffer) + v_buffers += list(draft_kv_pool.v_buffer) + k_ids += ids[num_target : num_target + draft_half] + v_ids += ids[num_target + draft_half :] + return k_buffers, v_buffers, k_ids + v_ids + + def append_state_component( kv_args: KVArgs, state_type: StateType, diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index e0c37793ddbd..e5be613d6459 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -105,10 +105,14 @@ def init_torch_distributed( tp_size=ps.tp_size, pp_size=ps.pp_size, moe_ep_size=ps.moe_ep_size ) + # The draft worker reuses the target's resolved memory_pool_config, so its own + # pre_model_load_memory feeds nothing cross-rank. Keep the WORLD reduction on the + # target only: a draft that exists on a subset of ranks (prefill-side PP builds it + # on the last stage) would otherwise enter a collective the other ranks never join. pre_model_load_memory = get_available_gpu_memory( device, ps.gpu_id, - distributed=get_world_group().world_size > 1, + distributed=get_world_group().world_size > 1 and not is_draft_worker, cpu_group=get_world_group().cpu_group, ) tp_group = get_tp_group() diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index aa25d7ddbbf1..db54f9411f71 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1946,6 +1946,12 @@ def get_moe_tp_group() -> GroupCoordinator: get_tensor_model_parallel_group = get_tp_group _PP: Optional[GroupCoordinator] = None +_SELF_PP: Optional[GroupCoordinator] = None + + +def get_self_pp_group() -> GroupCoordinator: + assert _SELF_PP is not None, "self pipeline group is not initialized" + return _SELF_PP def get_pp_group() -> GroupCoordinator: @@ -2550,6 +2556,19 @@ def initialize_model_parallel( max_world_size=max_world_size, ) + # A single-member pipeline group per rank. The speculative draft is one layer + # and never spans stages, so it is built and run against this group instead of + # the real one. new_group is collective, so every rank creates all of them. + global _SELF_PP + if _SELF_PP is None: + _SELF_PP = init_model_parallel_group( + [[r] for r in range(world_size)], + get_world_group().local_rank, + backend, + use_custom_allreduce=False, + group_name="self_pp", + ) + def create_custom_parallel_group( group_ranks: List[int], backend: str = "gloo" @@ -2646,6 +2665,28 @@ def model_parallel_is_initialized(): _TP_STATE_PATCHED = False +_PP_STATE_PATCHED = False + + +@contextmanager +def patch_pipeline_parallel_group(pp_group: GroupCoordinator): + """Patch the pp group temporarily until this function ends. + + This method is for draft workers of speculative decoding, whose model does not + span pipeline stages and must not read the target's pp topology. + """ + global _PP_STATE_PATCHED + assert not _PP_STATE_PATCHED, "Should not call when it's already patched" + + _PP_STATE_PATCHED = True + old_pp_group = get_pp_group() + global _PP + _PP = pp_group + try: + yield + finally: + _PP_STATE_PATCHED = False + _PP = old_pp_group @contextmanager diff --git a/python/sglang/srt/distributed/utils.py b/python/sglang/srt/distributed/utils.py index 61f2bb1f4554..81cfde07e098 100644 --- a/python/sglang/srt/distributed/utils.py +++ b/python/sglang/srt/distributed/utils.py @@ -101,6 +101,13 @@ def get_pp_indices( """ # partition_list_str can be set to None in sglang partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None) + if pp_size == 1: + # A single-stage pipeline owns every layer, so a partition list cannot + # apply to it. The env var is process-global, so a worker built with + # pp_size=1 inside a pipelined process -- the speculative draft, which + # never spans stages -- would otherwise read the target's list and + # reject it for having the wrong length. + partition_list_str = None if partition_list_str is not None: try: partitions = [int(layer) for layer in partition_list_str.split(",")] diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index a2f48d9b76c5..717a6a27528c 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -660,6 +660,15 @@ async def health_generate(request: Request) -> Response: if _global_state.tokenizer_manager.server_status == ServerStatus.Starting: return Response(status_code=503) + # Diagnostic only: allow an external E2E driver to establish a balanced + # DP16 batch instead of letting the router's singleton health request + # create the [1, 0, ..., 0] idle-rank case first. + if ( + os.getenv("SGLANG_DIAG_BYPASS_HEALTH_GENERATE", "0") == "1" + and request.url.path in ("/health", "/health_generate") + ): + return Response(status_code=200) + if ( not envs.SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION.get() and request.url.path == "/health" diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9d5987f25dc5..b8de9e46b230 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -744,6 +744,9 @@ class Envs: # DeepGemm SGLANG_ENABLE_JIT_DEEPGEMM = EnvBool(True) + # Enable the allowlisted low-M BF16 Split-K GEMM path on Blackwell. Shapes + # outside the measured allowlist continue to use CuTe DSL/cuBLAS. + SGLANG_ENABLE_BF16_SPLITK_GEMM = EnvBool(True) SGLANG_DEEPGEMM_STANDARD_LAYOUT = EnvStr("auto") SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION = EnvFloat(0.25) # Cap the DeepGEMM masked grouped-GEMM per-expert padded capacity at @@ -752,6 +755,13 @@ class Envs: # load imbalance (they otherwise OOM saturated --moe-runner-backend # deep_gemm serving). Costs one D2H sync per MoE layer. SGLANG_OPT_DG_MASKED_M_CAP = EnvBool(False) + # Use the compact standard-to-DeepGEMM layout outside CUDA graph capture. + # Wide-DP AG+RS prefill can expose every rank to hundreds of thousands of + # tokens. With skewed routing, the masked layout multiplies the hottest + # expert capacity by num_local_experts and can require tens of GiB per MoE + # intermediate. The compact layout scales with routed assignments instead. + # Decode CUDA graphs keep the faster masked layout. + SGLANG_OPT_DG_COMPACT_EAGER = EnvBool(False) # Drop dp-attention MAX_LEN pad rows from MoE dispatch (StandardDispatcher # post-translation topk_ids -> -1): pad rows otherwise run the router on # stale hidden values and burn expert compute whose outputs are discarded; @@ -777,6 +787,12 @@ class Envs: # DeepEP SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) # This argument is deprecated SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # DeepEP v2 per-rank communication buffer capacity. This is not a model + # semantic token limit; large prefill/chunked-prefill workloads may need a + # larger value. + SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + # 0 lets DeepEP v2 ElasticBuffer choose the communication SM count. + SGLANG_DEEPEP_V2_NUM_SMS = EnvInt(0) SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32) SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False) # Force dynamic Waterfill with runtime EP all-reduce instead of the default @@ -1338,6 +1354,13 @@ class Envs: SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang")) SGLANG_FLASHINFER_AUTOTUNE_CACHE = EnvBool(True) SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True) + # Qwen3.5 experimental integration for FlashInfer's MNNVL CuTe DSL + # AllReduce-fusion backend. One switch enables deferred MoE finalize and + # ordinary AR + residual + RMSNorm in decode and prefill. + SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION = EnvBool(False) + # Distinct workspace configurations allowed in one process. Production + # uses one model/configuration per rank, so fail closed on accidental reuse. + SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES = EnvInt(1) # Plugin system SGLANG_PLATFORM = EnvStr("") @@ -1436,6 +1459,8 @@ def _convert_SGL_to_SGLANG(): _print_deprecated_env("SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN") # sconv-family kernels always use the CUDA-JIT ports when supported; no toggle. _print_deprecated_env("SGLANG_OPT_USE_CUDA_SCONV") + # The PR #4266 direct dense BF16 GEMM kernel now ships in flashinfer itself. + _print_deprecated_env("SGLANG_FLASHINFER_PR4266_SOURCE") _print_deprecated_env("SGLANG_ENABLE_THINKING", "SGLANG_DEFAULT_THINKING") _print_deprecated_env("SGLANG_REASONING_EFFORT", "SGLANG_DSV4_REASONING_EFFORT") _print_deprecated_env( diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index b524a9afcb85..11f81721970e 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -281,6 +281,7 @@ def fast_prefill_plan( fixed_split_size if fixed_split_size is not None else -1, False, # disable_split_kv 0, # num_colocated_ctas + 0, # uniform_q_len ] self._plan_info = self._cached_module.plan(*args) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index f62abfe89c8f..d659a80318e4 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import os from typing import TYPE_CHECKING, Optional, Union import torch @@ -21,6 +22,9 @@ ForwardMetadata, Mamba2Metadata, ) +from sglang.srt.layers.attention.mamba.replay_state_indices_validator import ( + validate_replay_state_indices_cpu, +) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -33,6 +37,9 @@ from sglang.srt.layers.attention.verify_mask import VerifyMask logger = logging.getLogger(__name__) +_validate_mamba_replay_state_indices = ( + os.environ.get("SGLANG_VALIDATE_MAMBA_REPLAY_STATE_INDICES", "0") == "1" +) class MambaAttnBackendBase(AttentionBackend): @@ -579,6 +586,18 @@ def _replay_metadata( mamba_indices = self._translate_mamba_indices(mamba_indices) mamba_indices[bs - num_padding :] = -1 self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + if _validate_mamba_replay_state_indices and not in_capture: + # This function runs before graph replay. The diagnostic CPU copy + # deliberately synchronizes here so malformed live/padded indices + # fail before any captured indexed state update can consume them. + valid_bs = bs - int(num_padding) + validate_replay_state_indices_cpu( + mamba_indices.detach().cpu(), + valid_bs=valid_bs, + total_bs=bs, + num_state_slots=self.req_to_token_pool.mamba_pool.size + 1, + pad_slot_id=self.pad_slot_id, + ) # Refresh the static track-dest buffer in-place (translated); the captured # track-save reads it, leaving the handed-in InputBuffer slot read-only. # Hand out only the refreshed [:bs] prefix — Mamba2's track-save slices diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 7539cc1cdb67..a317518f4fca 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -1,3 +1,4 @@ +import os from typing import Optional, Tuple, Union import torch @@ -30,10 +31,23 @@ if is_cuda() or is_hip(): from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + can_use_fused_qkvzba_causal_conv1d_update_contiguous, fused_qkv_split_gdn_prefill, + fused_qkvzba_causal_conv1d_update_contiguous, + fused_qkvzba_split_reshape_cat_contiguous, ) MAX_FUSED_QKV_SPLIT_DIM = 8192 +_fused_decode_proj_conv_logged = False +_fused_decode_proj_conv_fallback_logged = False +_fused_decode_proj_conv_layers_logged: set[int] = set() +_fused_decode_real_tensor_verified_layers: set[int] = set() +_fused_decode_log_layer_hits = ( + os.environ.get("SGLANG_GDN_DECODE_FUSION_LOG_LAYER_HITS", "0") == "1" +) +_fused_decode_verify_real_tensors = ( + os.environ.get("SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS", "0") == "1" +) if is_cuda(): from sglang.srt.layers.attention.mamba.causal_conv1d import ( @@ -385,6 +399,11 @@ def forward_decode( b: torch.Tensor, **kwargs, ): + global _fused_decode_proj_conv_fallback_logged + global _fused_decode_proj_conv_logged + global _fused_decode_proj_conv_layers_logged + global _fused_decode_real_tensor_verified_layers + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) conv_states = layer_cache.conv[0] ssm_states = layer_cache.temporal @@ -402,15 +421,166 @@ def forward_decode( replayssm_k = layer_cache.replayssm_k replayssm_g = layer_cache.replayssm_g - assert isinstance(mixed_qkv, torch.Tensor) - mixed_qkv = causal_conv1d_update( - mixed_qkv, - conv_states, - layer.conv_weights, - layer.bias, - layer.activation, - conv_state_indices=cache_indices, - ) + return_z = False + conv_already_applied = False + if isinstance(mixed_qkv, tuple): + if len(mixed_qkv) != 2: + raise ValueError( + "Fused GDN decode projection input must be " + "(projected_qkvz, projected_ba)" + ) + projected_qkvz, projected_ba = mixed_qkv + eligible, eligibility_reason = ( + can_use_fused_qkvzba_causal_conv1d_update_contiguous( + projected_qkvz, + projected_ba, + conv_states, + layer.conv_weights, + layer.bias, + cache_indices, + qkv_dim=layer.q_dim + layer.k_dim + layer.v_dim, + v_dim=layer.v_dim, + num_v_heads=layer.num_v_heads, + activation=layer.activation, + ) + ) + if eligible: + qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim + fused_backend = "triton_direct_oracle_exact" + if not _fused_decode_proj_conv_logged: + rank0_log("Using fused GDN decode QKVZ/BA unpack + indexed Conv1D.") + _fused_decode_proj_conv_logged = True + if ( + _fused_decode_log_layer_hits or _fused_decode_verify_real_tensors + ) and layer.layer_id not in _fused_decode_proj_conv_layers_logged: + rank0_log( + "GDN_FUSED_DECODE_BACKEND " + f"layer_id={layer.layer_id} backend={fused_backend} " + f"batch={projected_qkvz.shape[0]} " + f"qkv_dim={qkv_dim} state_shape={tuple(conv_states.shape)} " + f"state_indices_dtype={cache_indices.dtype}" + ) + _fused_decode_proj_conv_layers_logged.add(layer.layer_id) + + # Diagnostic-only real-tensor oracle. It consumes the actual + # projection activations and selected pre-update cache rows, + # but runs the deployed direct-Triton update on a compact + # state copy so the live cache is mutated only by the + # candidate. Any mismatch aborts at the first GDN layer. + verify_real_tensors = ( + _fused_decode_verify_real_tensors + and layer.layer_id not in _fused_decode_real_tensor_verified_layers + ) + if verify_real_tensors: + if bool(torch.any(cache_indices < 0).item()): + raise AssertionError( + "Real-tensor GDN fusion verification requires " + "non-padding cache indices" + ) + ref_indices = torch.arange( + cache_indices.numel(), + device=cache_indices.device, + dtype=torch.int32, + ) + ref_state = torch.index_select( + conv_states, 0, cache_indices.to(torch.int64) + ) + ref_mixed_qkv, ref_z, ref_b, ref_a = ( + fused_qkvzba_split_reshape_cat_contiguous( + projected_qkvz, + projected_ba, + layer.num_q_heads, + layer.num_v_heads, + layer.head_q_dim, + layer.head_v_dim, + ) + ) + ref_mixed_qkv = causal_conv1d_update( + ref_mixed_qkv, + ref_state, + layer.conv_weights, + layer.bias, + layer.activation, + conv_state_indices=ref_indices, + ) + + mixed_qkv, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous( + projected_qkvz, + projected_ba, + conv_states, + layer.conv_weights, + layer.bias, + cache_indices, + qkv_dim=qkv_dim, + v_dim=layer.v_dim, + num_v_heads=layer.num_v_heads, + head_v_dim=layer.head_v_dim, + activation=layer.activation, + ) + if verify_real_tensors: + candidate_state = torch.index_select( + conv_states, 0, cache_indices.to(torch.int64) + ) + named_pairs = ( + ("qkv", mixed_qkv, ref_mixed_qkv), + ("z", z, ref_z), + ("b", b, ref_b), + ("a", a, ref_a), + ("state", candidate_state, ref_state), + ) + report = [] + mismatch = False + for tensor_name, candidate, reference in named_pairs: + diff = (candidate.float() - reference.float()).abs() + nonzero = int(torch.count_nonzero(diff).item()) + mismatch |= nonzero != 0 + report.append( + f"{tensor_name}_nonzero={nonzero}/" + f"{diff.numel()} {tensor_name}_max=" + f"{diff.max().item()}" + ) + rank0_log( + "GDN_FUSED_REAL_TENSOR_PARITY " + f"layer_id={layer.layer_id} backend={fused_backend} " + + " ".join(report) + ) + _fused_decode_real_tensor_verified_layers.add(layer.layer_id) + if mismatch: + raise AssertionError( + "GDN fused real-tensor parity failed at " + f"layer_id={layer.layer_id}; " + " ".join(report) + ) + conv_already_applied = True + else: + # Explicit correctness fallback for an unexpected runtime + # tensor/state contract. This still returns Z to the model. + if not _fused_decode_proj_conv_fallback_logged: + rank0_log( + "Falling back from fused GDN decode projection/Conv1D: " + f"{eligibility_reason}" + ) + _fused_decode_proj_conv_fallback_logged = True + mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous( + projected_qkvz, + projected_ba, + layer.num_q_heads, + layer.num_v_heads, + layer.head_q_dim, + layer.head_v_dim, + ) + return_z = True + else: + assert isinstance(mixed_qkv, torch.Tensor) + + if not conv_already_applied: + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + layer.conv_weights, + layer.bias, + layer.activation, + conv_state_indices=cache_indices, + ) # Skip split + reshape + separate gating kernel by consuming # the packed mixed_qkv directly in a single fused Triton kernel. @@ -435,7 +605,7 @@ def forward_decode( self._track_mamba_state_decode( forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) - return core_attn_out + return (core_attn_out, z) if return_z else core_attn_out query, key, value = torch.split( mixed_qkv, @@ -465,7 +635,7 @@ def forward_decode( forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id ) - return core_attn_out + return (core_attn_out, z) if return_z else core_attn_out def forward_extend( self, @@ -681,6 +851,7 @@ def forward_extend( state_checkpoint_every_n_tokens=( forward_metadata.state_checkpoint_every_n_tokens ), + output=kwargs.get("linear_attn_output"), ) if is_npu() and last_recurrent_state is not None: diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index f7f4d974d64c..38b057f70b2c 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -28,6 +28,30 @@ logger = logging.getLogger(__name__) +_FLASHINFER_GDN_ALIGNMENT = 32 + + +def _empty_aligned_like( + tensor: torch.Tensor, alignment: int = _FLASHINFER_GDN_ALIGNMENT +) -> torch.Tensor: + """Return an uninitialized contiguous tensor with an aligned data pointer.""" + element_size = tensor.dtype.itemsize + alignment_elements = max(1, alignment // element_size) + storage = torch.empty( + tensor.numel() + alignment_elements - 1, + dtype=tensor.dtype, + device=tensor.device, + ) + start_bytes = (-storage.data_ptr()) % alignment + if start_bytes % element_size: + raise RuntimeError( + f"Cannot align {tensor.dtype} storage at {storage.data_ptr()} " + f"to {alignment} bytes" + ) + start = start_bytes // element_size + return storage[start : start + tensor.numel()].view(tensor.shape) + + # --------------------------------------------------------------------------- # Lazy import for FlashInfer GDN kernels # --------------------------------------------------------------------------- @@ -162,6 +186,20 @@ def __init__(self): sm_major = torch.cuda.get_device_capability()[0] self.use_state_pool = sm_major >= 10 self.supports_target_verify = sm_major in (9, 10) + self._aligned_input_buffers: dict[tuple, torch.Tensor] = {} + self._aligned_parameter_cache: dict[ + tuple, tuple[torch.Tensor, torch.Tensor] + ] = {} + self._verify_intermediate_buffers: dict[tuple, torch.Tensor] = {} + self._alignment_fallback_warned = False + # Mutable state/workspace cannot be repaired with a temporary copy: + # FlashInfer writes through those pointers. Triton has no 32-byte ABI + # requirement and is therefore the semantics-preserving fallback. + from sglang.srt.layers.attention.linear.kernels.gdn_triton import ( + TritonGDNKernel, + ) + + self._alignment_fallback_kernel = TritonGDNKernel() if sm_major == 9 and self._prefill_fn is None: raise RuntimeError("FlashInfer GDN prefill kernel is unavailable.") @@ -204,6 +242,142 @@ def _mtp_bf16_adapted( logger.info("Using FlashInfer GDN kernels") + def _prepare_dynamic_input(self, name: str, tensor: torch.Tensor) -> torch.Tensor: + """Repair an under-aligned read-only input using reusable scratch. + + The normal producer path returns aligned tensors and takes the zero-cost + branch. The scratch path is a safety net for other models or view + layouts. A buffer is allocated once per argument/shape/stream and then + reused, avoiding allocator churn in eager decode and providing stable + addresses for CUDA Graph capture. + """ + if tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT == 0: + return tensor + + stream_key = ( + torch.cuda.current_stream(tensor.device).cuda_stream + if tensor.device.type == "cuda" + else None + ) + key = ( + name, + tensor.device, + tensor.dtype, + tuple(tensor.shape), + stream_key, + ) + aligned = self._aligned_input_buffers.get(key) + if aligned is None: + aligned = _empty_aligned_like(tensor) + self._aligned_input_buffers[key] = aligned + aligned.copy_(tensor) + return aligned + + def _prepare_parameter( + self, + name: str, + tensor: torch.Tensor, + *, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + """Return a stable aligned view/copy of an immutable kernel parameter.""" + key = ( + name, + id(tensor), + dtype, + ) + cached_entry = self._aligned_parameter_cache.get(key) + if cached_entry is not None and cached_entry[0] is tensor: + return cached_entry[1] + + prepared = tensor.detach().reshape(-1) + if dtype is not None: + prepared = prepared.to(dtype=dtype, copy=False) + if ( + not prepared.is_contiguous() + or prepared.data_ptr() % _FLASHINFER_GDN_ALIGNMENT + ): + aligned = _empty_aligned_like(prepared) + aligned.copy_(prepared) + prepared = aligned + # Retaining the source also prevents a recycled Python id from + # colliding with an older cache entry. + self._aligned_parameter_cache[key] = (tensor, prepared) + return prepared + + def _prepare_gate_parameters( + self, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + A_log_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + self._prepare_parameter("A_log", A_log, dtype=A_log_dtype), + self._prepare_parameter("dt_bias", dt_bias), + ) + + def _mutable_inputs_are_aligned( + self, *named_tensors: tuple[str, Optional[torch.Tensor]] + ) -> bool: + """Return whether mutable buffers satisfy FlashInfer's ABI. + + Copying a state or workspace into temporary storage would lose kernel + writeback. The caller falls back to Triton instead. + """ + for name, tensor in named_tensors: + if tensor is None or tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT == 0: + continue + if not self._alignment_fallback_warned: + logger.warning( + "FlashInfer GDN mutable buffer %r has data_ptr %d " + "(mod 32 = %d); falling back to Triton for this call.", + name, + tensor.data_ptr(), + tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT, + ) + self._alignment_fallback_warned = True + return False + return True + + def _prepare_verify_intermediate_buffer( + self, + intermediate_states_buffer: torch.Tensor, + batch_size: int, + ) -> tuple[torch.Tensor, bool]: + """Return FlashInfer's exact-batch verify workspace. + + SGLang owns a pool-scoped speculative-state buffer, while the SM100 + FlashInfer MTP kernel requires its workspace dim 0 to equal the + captured batch size exactly. CUDA Graph capture can pad the batch past + the speculative pool (for example, pool rows 7 and capture B=8). Keep + the normal path zero-copy and use a stable, stream-local aligned + scratch only for that padded capture tier. The caller copies the rows + owned by the pool back after FlashInfer writes them. + """ + direct = intermediate_states_buffer[:batch_size] + if direct.shape[0] == batch_size: + return direct, False + + stream_key = ( + torch.cuda.current_stream(intermediate_states_buffer.device).cuda_stream + if intermediate_states_buffer.device.type == "cuda" + else None + ) + shape = (batch_size, *intermediate_states_buffer.shape[1:]) + key = ( + intermediate_states_buffer.device, + intermediate_states_buffer.dtype, + shape, + stream_key, + ) + scratch = self._verify_intermediate_buffers.get(key) + if scratch is None: + template = intermediate_states_buffer.new_empty(shape) + scratch = _empty_aligned_like(template) + self._verify_intermediate_buffers[key] = scratch + return scratch, True + # ---- decode ---- def decode( @@ -221,6 +395,21 @@ def decode( query_start_loc: torch.Tensor, **kwargs, ) -> torch.Tensor: + if not self._mutable_inputs_are_aligned(("ssm_states", ssm_states)): + return self._alignment_fallback_kernel.decode( + q, + k, + v, + a, + b, + A_log=A_log, + dt_bias=dt_bias, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + **kwargs, + ) + batch_size = cache_indices.shape[0] num_heads = q.shape[2] head_k_dim = q.shape[3] @@ -232,20 +421,35 @@ def decode( value_fi = v.view(batch_size, 1, num_v_heads, head_v_dim) a_fi = a.view(batch_size, 1, num_v_heads) b_fi = b.view(batch_size, 1, num_v_heads) + query_fi = self._prepare_dynamic_input("decode_q", query_fi) + key_fi = self._prepare_dynamic_input("decode_k", key_fi) + value_fi = self._prepare_dynamic_input("decode_v", value_fi) + a_fi = self._prepare_dynamic_input("decode_a", a_fi) + b_fi = self._prepare_dynamic_input("decode_b", b_fi) + A_log_fi, dt_bias_fi = self._prepare_gate_parameters( + A_log, + dt_bias, + # Preserve the original backend contract: the SM100 state-pool + # kernel consumes float32 A_log, while SM90 uses the source dtype. + A_log_dtype=torch.float32 if self.use_state_pool else None, + ) if self.use_state_pool: + cache_indices_fi = self._prepare_dynamic_input( + "decode_cache_indices", cache_indices + ) output_fi, _ = self._decode_fn( q=query_fi, k=key_fi, v=value_fi, state=None, - A_log=A_log.detach().float(), + A_log=A_log_fi, a=a_fi, - dt_bias=dt_bias.detach(), + dt_bias=dt_bias_fi, b=b_fi, use_qk_l2norm=True, initial_state=ssm_states, - initial_state_indices=cache_indices, + initial_state_indices=cache_indices_fi, ) else: # TODO: Once FlashInfer PR#2521 is merged for SM90, gather/scatter @@ -256,9 +460,9 @@ def decode( k=key_fi, v=value_fi, state=state_batch, - A_log=A_log.detach(), + A_log=A_log_fi, a=a_fi, - dt_bias=dt_bias.detach(), + dt_bias=dt_bias_fi, b=b_fi, scale=None, output=None, @@ -296,6 +500,30 @@ def extend( k_fi = l2norm_fwd(k[0].contiguous()) v_fi = v[0].contiguous() + output = kwargs.get("output") + output_fi = None + if output is not None: + expected_output_shape = ( + 1, + total_seq_len, + num_v_heads, + head_v_dim, + ) + if tuple(output.shape) != expected_output_shape: + raise ValueError( + "FlashInfer GDN prefill output shape mismatch: " + f"expected {expected_output_shape}, got {tuple(output.shape)}" + ) + if output.dtype != v.dtype or output.device != v.device: + raise ValueError( + "FlashInfer GDN prefill output must match v dtype/device, " + f"got output=({output.dtype}, {output.device}) and " + f"v=({v.dtype}, {v.device})" + ) + output_fi = output[0] + if not output_fi.is_contiguous(): + raise ValueError("FlashInfer GDN prefill output must be contiguous") + # g (alpha) and beta: [1, seq, HV] -> [seq, HV], float32 for FlashInfer alpha_fi = torch.exp(g[0].to(torch.float32)) beta_fi = beta[0].to(torch.float32) @@ -338,6 +566,7 @@ def extend( output_final_state=True, cu_seqlens=cu_seqlens, use_qk_l2norm_in_kernel=False, + output=output_fi, output_state=output_state_fi, state_checkpoints=state_checkpoints, checkpoint_cu_starts=state_checkpoint_cu_starts, @@ -395,34 +624,80 @@ def target_verify( num_v_heads = v.shape[2] head_v_dim = v.shape[3] - query_mtp = q.view(batch_size, draft_token_num, num_heads, head_k_dim) - key_mtp = k.view(batch_size, draft_token_num, num_heads, head_k_dim) - value_mtp = v.view(batch_size, draft_token_num, num_v_heads, head_v_dim) - if a is None or b is None or A_log is None or dt_bias is None: raise RuntimeError( "FlashInfer GDN MTP kernel requires a, b, A_log, dt_bias." ) - a_mtp = a.view(batch_size, draft_token_num, num_v_heads) - b_mtp = b.view(batch_size, draft_token_num, num_v_heads) - intermediate_states_buffer_mtp = intermediate_states_buffer + copy_verify_intermediate_back = False if self.use_state_pool and intermediate_states_buffer is not None: # The SM100 bf16 MTP kernel indexes this scratch buffer by the # per-call batch id, while SGLang's speculative state cache is - # pool-scoped and may include an extra dummy slot. - intermediate_states_buffer_mtp = intermediate_states_buffer[:batch_size] + # pool-scoped. Graph capture can pad B beyond that pool, so use a + # stable exact-B scratch for the padded tier and copy owned rows + # back before post-verify commit reads the pool. + ( + intermediate_states_buffer_mtp, + copy_verify_intermediate_back, + ) = self._prepare_verify_intermediate_buffer( + intermediate_states_buffer, batch_size + ) + if not self._mutable_inputs_are_aligned( + ("ssm_states", ssm_states), + ("intermediate_states_buffer", intermediate_states_buffer_mtp), + ): + return self._alignment_fallback_kernel.target_verify( + A_log=A_log, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + a=a, + b=b, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + intermediate_states_buffer=intermediate_states_buffer, + intermediate_state_indices=intermediate_state_indices, + cache_steps=cache_steps, + retrieve_parent_token=retrieve_parent_token, + **kwargs, + ) + + query_mtp = self._prepare_dynamic_input( + "verify_q", + q.view(batch_size, draft_token_num, num_heads, head_k_dim), + ) + key_mtp = self._prepare_dynamic_input( + "verify_k", + k.view(batch_size, draft_token_num, num_heads, head_k_dim), + ) + value_mtp = self._prepare_dynamic_input( + "verify_v", + v.view(batch_size, draft_token_num, num_v_heads, head_v_dim), + ) + + a_mtp = self._prepare_dynamic_input( + "verify_a", a.view(batch_size, draft_token_num, num_v_heads) + ) + b_mtp = self._prepare_dynamic_input( + "verify_b", b.view(batch_size, draft_token_num, num_v_heads) + ) + A_log_fi, dt_bias_fi = self._prepare_gate_parameters(A_log, dt_bias) + cache_indices_fi = self._prepare_dynamic_input( + "verify_cache_indices", cache_indices + ) output_fi, _ = self._mtp_fn( q=query_mtp, k=key_mtp, v=value_mtp, initial_state=ssm_states, - initial_state_indices=cache_indices, - A_log=A_log.detach(), + initial_state_indices=cache_indices_fi, + A_log=A_log_fi, a=a_mtp, - dt_bias=dt_bias.detach(), + dt_bias=dt_bias_fi, b=b_mtp, scale=None, output=None, @@ -431,4 +706,9 @@ def target_verify( use_qk_l2norm=True, ) + if copy_verify_intermediate_back: + intermediate_states_buffer.copy_( + intermediate_states_buffer_mtp[: intermediate_states_buffer.shape[0]] + ) + return output_fi.view(1, seq_len, num_v_heads, head_v_dim) diff --git a/python/sglang/srt/layers/attention/mamba/replay_state_indices_validator.py b/python/sglang/srt/layers/attention/mamba/replay_state_indices_validator.py new file mode 100644 index 000000000000..fda320e3f41b --- /dev/null +++ b/python/sglang/srt/layers/attention/mamba/replay_state_indices_validator.py @@ -0,0 +1,82 @@ +"""Debug validation for Mamba CUDA-graph replay state indices. + +The validator is intentionally CPU-only. Replay metadata is prepared outside +the captured graph, so a diagnostic run may synchronize once here without +putting host reads or dynamic assertions into CUDA graph capture/replay. +""" + +from __future__ import annotations + +import torch + + +def validate_replay_state_indices_cpu( + state_indices: torch.Tensor, + *, + valid_bs: int, + total_bs: int, + num_state_slots: int, + pad_slot_id: int = -1, +) -> None: + """Validate live and padded rows of a replay state-index buffer. + + Live rows must own distinct in-range slots in ``[0, num_state_slots)``. + Slot zero is reserved for CUDA-graph dummy/idle traffic but is still a + valid storage row. All padded rows must carry exactly ``pad_slot_id`` so + indexed state kernels skip them. + """ + if state_indices.device.type != "cpu": + raise ValueError("state_indices must be copied to CPU before validation") + if state_indices.ndim != 1: + raise ValueError("state_indices must be rank-1") + if not 0 <= valid_bs <= total_bs <= state_indices.numel(): + raise ValueError( + "expected 0 <= valid_bs <= total_bs <= state_indices.numel(), got " + f"valid_bs={valid_bs} total_bs={total_bs} " + f"numel={state_indices.numel()}" + ) + if num_state_slots <= 1: + raise ValueError( + f"num_state_slots must include real slots, got {num_state_slots}" + ) + + indices = state_indices[:total_bs].to(dtype=torch.int64) + live = indices[:valid_bs] + padded = indices[valid_bs:] + errors: list[str] = [] + + live_in_range = (live >= 0) & (live < num_state_slots) + if not bool(torch.all(live_in_range)): + bad_rows = torch.nonzero(~live_in_range, as_tuple=False).flatten() + errors.append( + "live rows must contain in-range slots in " + f"[0, {num_state_slots}); bad_rows={bad_rows.tolist()} " + f"bad_values={live[bad_rows].tolist()}" + ) + + if live.numel() > 1: + unique_live, counts = torch.unique(live, sorted=True, return_counts=True) + duplicate_mask = counts > 1 + if bool(torch.any(duplicate_mask)): + errors.append( + "live rows must own unique slots; " + f"duplicate_slots={unique_live[duplicate_mask].tolist()} " + f"counts={counts[duplicate_mask].tolist()}" + ) + + bad_padding = padded != pad_slot_id + if bool(torch.any(bad_padding)): + bad_offsets = torch.nonzero(bad_padding, as_tuple=False).flatten() + errors.append( + f"padded rows must equal pad_slot_id={pad_slot_id}; " + f"bad_rows={(bad_offsets + valid_bs).tolist()} " + f"bad_values={padded[bad_offsets].tolist()}" + ) + + if errors: + raise AssertionError( + "Invalid Mamba replay state indices: " + + "; ".join(errors) + + f"; valid_bs={valid_bs} total_bs={total_bs} " + f"indices={indices.tolist()}" + ) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index eb8f4e7df29a..f406647c0f61 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -94,7 +94,17 @@ def get_dp_padding_mode( # Force MAX_LEN so all ranks are padded to equal token counts. from sglang.srt.layers.moe.utils import get_moe_a2a_backend - if get_moe_a2a_backend().is_pplx(): + moe_a2a_backend = get_moe_a2a_backend() + if moe_a2a_backend.is_pplx(): + return DpPaddingMode.MAX_LEN + + # Diagnostic gate for DeepEP v2 ElasticBuffer on ragged DP batches. + # Keeping this behind an environment variable makes the A/B change + # strictly limited to padding: communication mode, kernels, model and + # request payload remain identical. + if moe_a2a_backend.is_deepep_v2() and get_bool_env_var( + "SGLANG_DEEPEP_V2_FORCE_MAX_LEN" + ): return DpPaddingMode.MAX_LEN # When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding diff --git a/python/sglang/srt/layers/flashinfer_fallback/README.md b/python/sglang/srt/layers/flashinfer_fallback/README.md new file mode 100644 index 000000000000..5dbc92d12d20 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/README.md @@ -0,0 +1,15 @@ +# Temporary FlashInfer MNNVL CuTe DSL provider + +This package mirrors only the `flashinfer.comm` surface used by the SGLang +integration. It is removable once the serving image contains a FlashInfer +release with the MNNVL CuTe DSL all-reduce-fusion backend. + +`comm/mnnvl_cutedsl_ar.py` is copied from the in-development FlashInfer branch; +only imports of existing FlashInfer infrastructure are redirected to the +installed package. `comm/mnnvl_cutedsl/` is copied unchanged so later kernel +refreshes remain mechanical directory syncs. The current snapshot comes from +FlashInfer commit `b23d193d92d77227ecdf575eb651e8a69e78c720` plus its local +in-progress kernel changes. + +Model code must import `sglang.srt.layers.flashinfer_provider`, never this +package directly. diff --git a/python/sglang/srt/layers/flashinfer_fallback/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/__init__.py new file mode 100644 index 000000000000..7592ce85ab4c --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/__init__.py @@ -0,0 +1 @@ +"""Private temporary FlashInfer-compatible implementation details.""" diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/comm/__init__.py new file mode 100644 index 000000000000..423ee0657dc4 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/__init__.py @@ -0,0 +1,39 @@ +"""FlashInfer-compatible facade for the copied MNNVL CuTe DSL backend.""" + +from __future__ import annotations + +import flashinfer.comm as _upstream_comm + +from .mnnvl_cutedsl_ar import ( + MNNVLCuteDSLAllReduceFusionWorkspace, + _mnnvl_cutedsl_allreduce_fusion, +) + +AllReduceFusionPattern = _upstream_comm.AllReduceFusionPattern + + +def allreduce_fusion(*, input, workspace, pattern, **kwargs): + """Use the copied backend locally and preserve upstream dispatch otherwise.""" + if isinstance(workspace, MNNVLCuteDSLAllReduceFusionWorkspace): + # The unified API documents this as ignored by MNNVL backends and does + # not forward it to the backend-specific implementation. + kwargs.pop("trigger_completion_at_end", None) + return _mnnvl_cutedsl_allreduce_fusion( + input=input, + workspace=workspace, + pattern=pattern, + **kwargs, + ) + return _upstream_comm.allreduce_fusion( + input=input, + workspace=workspace, + pattern=pattern, + **kwargs, + ) + + +__all__ = [ + "AllReduceFusionPattern", + "MNNVLCuteDSLAllReduceFusionWorkspace", + "allreduce_fusion", +] diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/__init__.py new file mode 100644 index 000000000000..d7a782cd5e86 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MNNVL CuTe DSL AllReduce fusion backend internals.""" + +from importlib import import_module + +from .config import ( + KernelTarget, + MNNVLCuteDSLConfig, + MRangeDispatch, + ProtocolKind, + StaticProfile, +) + + +def __getattr__(name: str): + if name in { + "BT_ONLY_CONFIG", + "DEFAULT_CONFIG", + "HT_ONLY_CONFIG", + "LL_ONLY_CONFIG", + }: + presets = import_module(f"{__name__}.presets") + value = getattr(presets, name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "BT_ONLY_CONFIG", + "DEFAULT_CONFIG", + "HT_ONLY_CONFIG", + "LL_ONLY_CONFIG", + "KernelTarget", + "MNNVLCuteDSLConfig", + "MRangeDispatch", + "ProtocolKind", + "StaticProfile", +] diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/config.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/config.py new file mode 100644 index 000000000000..60bb7a9babaa --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/config.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration and token-count routing for the MNNVL CuTe DSL backend.""" + +from __future__ import annotations + +from bisect import bisect_left +from dataclasses import dataclass +from enum import Enum +from typing import Generic, TypeVar + +import torch + +__all__ = [ + "KernelTarget", + "MNNVLCuteDSLConfig", + "MRangeDispatch", + "ProtocolKind", + "StaticProfile", +] + + +class ProtocolKind(Enum): + LL = "ll" + BT = "bt" + HT = "ht" + + +PresetT = TypeVar("PresetT") + + +@dataclass(frozen=True, slots=True) +class KernelTarget(Generic[PresetT]): + protocol: ProtocolKind + preset: PresetT + + +TargetT = TypeVar("TargetT") + + +@dataclass(frozen=True, slots=True) +class MRangeDispatch(Generic[TargetT]): + """Map contiguous positive token-count ranges to kernel targets.""" + + upper_bounds: tuple[int | None, ...] + targets: tuple[TargetT, ...] + + def __post_init__(self) -> None: + if not self.upper_bounds: + raise ValueError("M range dispatch must contain at least one range") + if len(self.upper_bounds) != len(self.targets): + raise ValueError("M range upper bounds and targets must have equal length") + + previous = 0 + for index, upper_bound in enumerate(self.upper_bounds): + if upper_bound is None: + if index != len(self.upper_bounds) - 1: + raise ValueError("An unbounded M range must be the final range") + continue + if upper_bound <= previous: + raise ValueError("M range upper bounds must be strictly increasing") + previous = upper_bound + + @property + def is_unbounded(self) -> bool: + return self.upper_bounds[-1] is None + + @property + def finite_upper_bound(self) -> int | None: + return None if self.is_unbounded else self.upper_bounds[-1] + + def supports(self, m: int) -> bool: + if m <= 0: + return False + upper_bound = self.finite_upper_bound + return upper_bound is None or m <= upper_bound + + def select(self, m: int) -> TargetT: + if not self.supports(m): + raise ValueError(f"No kernel route supports M={m}") + + finite_bounds = tuple( + upper_bound for upper_bound in self.upper_bounds if upper_bound is not None + ) + index = bisect_left(finite_bounds, m) + return self.targets[index] + + def referenced_protocols(self) -> frozenset[ProtocolKind]: + protocols = { + target.protocol + for target in self.targets + if isinstance(target, KernelTarget) + } + return frozenset(protocols) + + def targets_for_capacity(self, capacity_m: int) -> tuple[TargetT, ...]: + if capacity_m <= 0: + return () + selected = [] + lower_bound = 1 + for upper_bound, target in zip(self.upper_bounds, self.targets, strict=True): + if lower_bound > capacity_m: + break + selected.append(target) + if upper_bound is None: + break + lower_bound = upper_bound + 1 + return tuple(selected) + + def max_m_for_protocol( + self, protocol: ProtocolKind, *, capacity_m: int + ) -> int | None: + lower_bound = 1 + maximum = None + for upper_bound, target in zip(self.upper_bounds, self.targets, strict=True): + effective_upper_bound = capacity_m if upper_bound is None else upper_bound + if ( + isinstance(target, KernelTarget) + and target.protocol is protocol + and lower_bound <= capacity_m + ): + maximum = min(effective_upper_bound, capacity_m) + lower_bound = effective_upper_bound + 1 + return maximum + + +@dataclass(frozen=True, slots=True) +class StaticProfile: + tp_size: int + hidden_size: int + top_k: int + dtype: torch.dtype + finalize_routes: MRangeDispatch[KernelTarget[object]] + all_reduce_routes: MRangeDispatch[KernelTarget[object]] + + def __post_init__(self) -> None: + if self.hidden_size <= 0 or self.hidden_size % 8: + raise ValueError("hidden_size must be a positive multiple of 8") + + def matches( + self, + *, + tp_size: int, + hidden_size: int, + top_k: int, + dtype: torch.dtype, + ) -> bool: + return ( + self.tp_size == tp_size + and self.hidden_size == hidden_size + and self.top_k == top_k + and self.dtype == dtype + ) + + def validate_capacity(self, capacity_m: int) -> None: + if capacity_m <= 0: + raise ValueError("capacity_m must be positive") + if not self.finalize_routes.supports(capacity_m): + raise ValueError( + "Finalize routes do not cover the requested workspace capacity" + ) + if not self.all_reduce_routes.supports(capacity_m): + raise ValueError( + "AllReduce routes do not cover the requested workspace capacity" + ) + + @property + def referenced_protocols(self) -> frozenset[ProtocolKind]: + return ( + self.finalize_routes.referenced_protocols() + | self.all_reduce_routes.referenced_protocols() + ) + + def protocol_capacity( + self, protocol: ProtocolKind, *, capacity_m: int + ) -> int | None: + maxima = ( + self.finalize_routes.max_m_for_protocol(protocol, capacity_m=capacity_m), + self.all_reduce_routes.max_m_for_protocol(protocol, capacity_m=capacity_m), + ) + present = tuple(value for value in maxima if value is not None) + return max(present) if present else None + + +@dataclass(frozen=True, slots=True) +class MNNVLCuteDSLConfig: + """Static profiles and routing policy for one backend configuration.""" + + profiles: tuple[StaticProfile, ...] + + def __post_init__(self) -> None: + keys = [ + (profile.tp_size, profile.hidden_size, profile.top_k, profile.dtype) + for profile in self.profiles + ] + if not keys: + raise ValueError("A backend config must contain at least one profile") + if len(keys) != len(set(keys)): + raise ValueError("Backend config profiles must have unique static shapes") + + def resolve( + self, + *, + tp_size: int, + hidden_size: int, + top_k: int, + dtype: torch.dtype, + capacity_m: int, + ) -> StaticProfile: + for profile in self.profiles: + if profile.matches( + tp_size=tp_size, + hidden_size=hidden_size, + top_k=top_k, + dtype=dtype, + ): + profile.validate_capacity(capacity_m) + return profile + raise ValueError("No MNNVL CuTe DSL profile supports this static shape") diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/cute_dsl_primitives.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/cute_dsl_primitives.py new file mode 100644 index 000000000000..1b10ee6fbc40 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/cute_dsl_primitives.py @@ -0,0 +1,876 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small standalone CuTe DSL and PTX primitives shared by Kernel backends.""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cutlass_dsl import T, dsl_user_op + +WARP_SIZE = 32 +VEC_BF16 = 8 +QUAD_BF16 = 4 +NEGATIVE_ZERO_BF16_BITS = 0x8000 +NEGATIVE_ZERO_BF16_PAIR = 0x80008000 +# CUTLASS cute::TMA::CacheHintSm100::EVICT_FIRST policy descriptor. +L2_EVICT_FIRST = 0x12F0000000000000 + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + address = pointer.toint(loc=loc, ip=ip) + if volatile: + opcode = "ld.volatile.global.v4.u32" + else: + opcode = "ld.global.v4.u32" + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=volatile, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def load_global_u32x4_predicated( + pointer: cute.Pointer, + predicate: Int32, + *, + loc=None, + ip=None, +): + address = pointer.toint(loc=loc, ip=ip) + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [ + address.ir_value(loc=loc, ip=ip), + Int32(predicate).ir_value(loc=loc, ip=ip), + ], + ( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.s32 p, $5, 0;\n\t" + "@!p mov.u32 $0, 0;\n\t" + "@!p mov.u32 $1, 0;\n\t" + "@!p mov.u32 $2, 0;\n\t" + "@!p mov.u32 $3, 0;\n\t" + "@p ld.global.v4.u32 {$0, $1, $2, $3}, [$4];\n\t" + "}" + ), + "=r,=r,=r,=r,l,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def load_global_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.u32 $0, [$1];", + "=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def load_global_u32_predicated( + pointer: cute.Pointer, + predicate: Int32, + *, + loc=None, + ip=None, +) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [ + address.ir_value(loc=loc, ip=ip), + Int32(predicate).ir_value(loc=loc, ip=ip), + ], + ( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.s32 p, $2, 0;\n\t" + "@!p mov.u32 $0, 0;\n\t" + "@p ld.global.u32 $0, [$1];\n\t" + "}" + ), + "=r,l,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): + address = pointer.toint(loc=loc, ip=ip) + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 2), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(2) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def load_global_u32x2_predicated( + pointer: cute.Pointer, + predicate: Int32, + *, + loc=None, + ip=None, +): + address = pointer.toint(loc=loc, ip=ip) + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 2), + [ + address.ir_value(loc=loc, ip=ip), + Int32(predicate).ir_value(loc=loc, ip=ip), + ], + ( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.s32 p, $3, 0;\n\t" + "@!p mov.u32 $0, 0;\n\t" + "@!p mov.u32 $1, 0;\n\t" + "@p ld.global.v2.u32 {$0, $1}, [$2];\n\t" + "}" + ), + "=r,=r,l,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(2) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None: + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32_address( + address: Int64, + value: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32x2(address: Int64, packed, *, loc=None, ip=None) -> None: + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u16_bits( + address: Int64, + value: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + ( + "{\n\t" + ".reg .b16 bits;\n\t" + "cvt.u16.u32 bits, $1;\n\t" + "st.global.u16 [$0], bits;\n\t" + "}" + ), + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_u32x4( + address: Int64, + *, + loc=None, + ip=None, +) -> None: + sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_global_bf16_as_f32( + address: Int64, + *, + loc=None, + ip=None, +) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [address.ir_value(loc=loc, ip=ip)], + ( + "{\n\t" + ".reg .b16 bits;\n\t" + "ld.global.b16 bits, [$1];\n\t" + "cvt.f32.bf16 $0, bits;\n\t" + "}" + ), + "=f,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def load_global_bf16_as_f32_predicated( + address: Int64, + predicate: Int32, + *, + loc=None, + ip=None, +) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [ + address.ir_value(loc=loc, ip=ip), + Int32(predicate).ir_value(loc=loc, ip=ip), + ], + ( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b16 bits;\n\t" + "setp.ne.s32 p, $2, 0;\n\t" + "@!p mov.b16 bits, 0;\n\t" + "@p ld.global.b16 bits, [$1];\n\t" + "cvt.f32.bf16 $0, bits;\n\t" + "}" + ), + "=f,l,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [value.ir_value(loc=loc, ip=ip)], + ( + "{\n\t" + ".reg .b16 bits;\n\t" + "cvt.rn.bf16.f32 bits, $1;\n\t" + "cvt.u32.u16 $0, bits;\n\t" + "}" + ), + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def shuffle_sync_idx_u32( + value: Uint32, + source_lane: Int32, + *, + loc=None, + ip=None, +) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + value.ir_value(loc=loc, ip=ip), + source_lane.ir_value(loc=loc, ip=ip), + ], + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", + "=r,r,r", + # Preserve full-warp execution across later divergent consumers. + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + llvm.inline_asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_global_u32( + pointer: cute.Pointer, + value: Uint32, + *, + loc=None, + ip=None, +) -> None: + address = pointer.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, BFloat16) + + +@dsl_user_op +def packed_u32_to_bf16x2(packed: Uint32, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([2], BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, 2, BFloat16) + + +@dsl_user_op +def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([QUAD_BF16], BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, QUAD_BF16, BFloat16) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([4], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def bf16x2_to_packed_u32(values, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.bitcast( + T.i32(), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([2], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@cute.jit +def sanitize_negative_zero_u32x4(packed): + sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for index in cutlass.range_constexpr(4): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def sanitize_negative_zero_u32(word: Uint32) -> Uint32: + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for index in cutlass.range_constexpr(2): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def fragment_has_negative_zero(packed): + dirty = False + for index in cutlass.range_constexpr(4): + word = packed[index] + dirty = ( + dirty + | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + ) + return dirty + + +@dsl_user_op +def map_shared_to_peer( + smem_pointer: cute.Pointer, + peer_rank: Int32, + *, + loc=None, + ip=None, +) -> Int32: + address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return Int32( + llvm.inline_asm( + T.i32(), + [address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, + value: Float32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [ + remote_address.ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_shared_u32x4(pointer: cute.Pointer, *, loc=None, ip=None): + address = pointer.toint(loc=loc, ip=ip) + # Prevent motion across the named-barrier pipeline protocol. + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [Int32(address).ir_value(loc=loc, ip=ip)], + "ld.shared.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def store_shared_u32x4( + pointer: cute.Pointer, + packed, + *, + loc=None, + ip=None, +) -> None: + address = pointer.toint(loc=loc, ip=ip) + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] + llvm.inline_asm( + None, + [Int32(address).ir_value(loc=loc, ip=ip), *words], + "st.shared.v4.u32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_global_u32x4_address( + address: Int64, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=volatile, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def packed_negative_zero_bf16x8(*, loc=None, ip=None): + word = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [word, word, word, word], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def cpasync_bulk_g2s( + gmem_ptr: cute.Pointer, + smem_ptr: cute.Pointer, + barrier_ptr: cute.Pointer, + size_bytes: Int32, + *, + loc=None, + ip=None, +) -> None: + operands = [ + gmem_ptr.toint(loc=loc, ip=ip).ir_value(), + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + barrier_ptr.toint(loc=loc, ip=ip).ir_value(), + size_bytes.ir_value(loc=loc, ip=ip), + Int64(L2_EVICT_FIRST).ir_value(), + ] + llvm.inline_asm( + None, + operands, + ( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes" + ".L2::cache_hint [$1], [$0], $3, [$2], $4;" + ), + "l,r,r,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def fence_proxy_async_shared_cta(*, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [], + "fence.proxy.async.shared::cta;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def remote_release_add1_u32(address: Int64, *, loc=None, ip=None) -> None: + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip)], + "red.release.sys.global.add.u32 [$0], 1;", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ldmc_bf16x8(address: Int64, *, loc=None, ip=None): + loaded = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.bf16x2 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def stmc_bf16x2( + address: Int64, + packed: Uint32, + *, + loc=None, + ip=None, +) -> None: + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)], + "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None: + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None: + words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)] + llvm.inline_asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/__init__.py new file mode 100644 index 000000000000..2e59041a61fb --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Balanced MNNVL protocol.""" + +from .protocol import ( + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1, + BTAllReduceTuning, + BTCollectiveTuning, + BTFinalizeTuning, +) + +__all__ = [ + "BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0", + "BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1", + "BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0", + "BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1", + "BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0", + "BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1", + "BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0", + "BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1", + "BTAllReduceTuning", + "BTCollectiveTuning", + "BTFinalizeTuning", +] diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py new file mode 100644 index 000000000000..50af0ebb1b8a --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py @@ -0,0 +1,1285 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Three-stage BF16 MoE finalize, TP reduction, and RMSNorm for SM100.""" + +from __future__ import annotations + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from ..cute_dsl_primitives import ( + NEGATIVE_ZERO_BF16_BITS, + VEC_BF16, + WARP_SIZE, + bf16x2_to_packed_u32, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + f32_to_bf16_bits, + fragment_has_negative_zero, + load_global_bf16_as_f32, + load_global_bf16_as_f32_predicated, + load_global_u32, + load_global_u32_predicated, + load_global_u32x2, + load_global_u32x2_predicated, + load_global_u32x4, + load_volatile_u32, + packed_u32_to_bf16x2, + packed_u32x2_to_bf16x4, + packed_u32x4_to_bf16x8, + sanitize_negative_zero_u32, + sanitize_negative_zero_u32x2, + sanitize_negative_zero_u32x4, + stmc_bf16x8, + store_global_u16_bits, + store_global_u32, + store_global_u32_address, + store_global_u32x2, + store_global_u32x4, + store_lamport_sentinel_u32x4, +) + +LAMPORT_GENERATIONS = 3 +NEXT_STAGE = 0 +ACTIVE_STAGE = 1 + + +@cute.jit +def _block_sum( + value: Float32, + warp_sums: cute.Tensor, + warps: cutlass.Constexpr[int], +) -> Float32: + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + value = cute.arch.warp_reduction_sum(value) + if lane == 0: + cute.arch.store((warp_sums + warp).llvm_ptr, value) + cute.arch.barrier() + + result = Float32(0.0) + if warp == 0: + if lane < Int32(warps): + result = cute.arch.load((warp_sums + lane).llvm_ptr, Float32) + result = cute.arch.warp_reduction_sum(result) + if lane == 0: + cute.arch.store(warp_sums.llvm_ptr, result) + cute.arch.barrier() + return cute.arch.load(warp_sums.llvm_ptr, Float32) + + +class _ScalarFinalizeUnicastDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + local_capacity: int, + threads: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.local_capacity = local_capacity + self.threads = threads + self.routed_scaling_factor = routed_scaling_factor + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + self.ctas_per_token = math.ceil(hidden_size / threads) + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_peer_addresses, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.top_k * 8, + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + hidden_index = (block % self.ctas_per_token) * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + metadata_offset = Int64(token) * self.top_k + Int64(metadata_index) + routed_index = cute.arch.load( + (permuted_indices.iterator + metadata_offset).llvm_ptr, + Int32, + ) + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + metadata_offset).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if routed_index == Int32(-1): + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + routed_index, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + shared_value = Float32(0.0) + if cutlass.const_expr( + self.include_shared_expert and self.load_shared_expert_before_pdl + ): + if hidden_index < Int32(self.hidden_size): + shared_offset = Int64(token) * self.hidden_size + Int64(hidden_index) + shared_value = load_global_bf16_as_f32( + Int64((shared_output.iterator + shared_offset).toint()) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if hidden_index < Int32(self.hidden_size): + shared_offset = Int64(token) * self.hidden_size + Int64(hidden_index) + shared_value = load_global_bf16_as_f32( + Int64((shared_output.iterator + shared_offset).toint()) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if hidden_index < Int32(self.hidden_size): + accumulator = Float32(0.0) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + source_offset = Int64(routed_index) * self.hidden_size + Int64( + hidden_index + ) + accumulator = ( + accumulator + + load_global_bf16_as_f32_predicated( + Int64((routed_output.iterator + source_offset).toint()), + Int32(routed_index != Int32(-1)), + ) + * weight + ) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group,)), + Float32, + ) + inputs.fill(Float32(0.0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + source_offset = Int64( + routed_index + ) * self.hidden_size + Int64(hidden_index) + inputs[item] = load_global_bf16_as_f32_predicated( + Int64((routed_output.iterator + source_offset).toint()), + Int32(routed_index != Int32(-1)), + ) + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + accumulator = accumulator + inputs[item] * cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + + if cutlass.const_expr(self.include_shared_expert): + accumulator = accumulator + shared_value + bits = f32_to_bf16_bits(accumulator) + if bits == Uint32(NEGATIVE_ZERO_BF16_BITS): + bits = Uint32(0) + + destination_rank = token % self.tp_size + local_token = token // self.tp_size + destination_base = cute.arch.load( + ( + contribution_mailbox_peer_addresses.iterator + destination_rank + ).llvm_ptr, + Int64, + ) + destination_offset = ( + (Int64(stage) * self.tp_size + self.rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(hidden_index) + store_global_u16_bits( + destination_base + destination_offset * 2, + bits, + ) + + if block == 0 and tidx == 0: + store_global_u32(stage_state.iterator + ACTIVE_STAGE, stage) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _NarrowVectorFinalizeUnicastDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + local_capacity: int, + threads: int, + elements_per_thread: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.local_capacity = local_capacity + self.threads = threads + self.elements_per_thread = elements_per_thread + self.routed_scaling_factor = routed_scaling_factor + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + self.words_per_fragment = elements_per_thread // 2 + self.fragments = hidden_size // elements_per_thread + self.ctas_per_token = math.ceil(self.fragments / threads) + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_peer_addresses, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.top_k * 8, + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + fragment = (block % self.ctas_per_token) * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + metadata_offset = Int64(token) * self.top_k + Int64(metadata_index) + routed_index = cute.arch.load( + (permuted_indices.iterator + metadata_offset).llvm_ptr, + Int32, + ) + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + metadata_offset).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if routed_index == Int32(-1): + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + routed_index, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + if cutlass.const_expr(self.include_shared_expert): + shared_values = cute.make_rmem_tensor( + cute.make_layout((self.elements_per_thread,)), + BFloat16, + ) + shared_values.fill(BFloat16(0.0)) + if cutlass.const_expr(self.load_shared_expert_before_pdl): + if fragment < self.fragments: + shared_offset = ( + Int64(token) * self.hidden_size + + Int64(fragment) * self.elements_per_thread + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=self.elements_per_thread * 2, + ) + if cutlass.const_expr(self.elements_per_thread == 2): + shared_values.store( + packed_u32_to_bf16x2(load_global_u32(shared_pointer)) + ) + else: + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if fragment < self.fragments: + shared_offset = ( + Int64(token) * self.hidden_size + + Int64(fragment) * self.elements_per_thread + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=self.elements_per_thread * 2, + ) + if cutlass.const_expr(self.elements_per_thread == 2): + shared_values.store( + packed_u32_to_bf16x2(load_global_u32(shared_pointer)) + ) + else: + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if fragment < self.fragments: + accumulator = cute.make_rmem_tensor( + cute.make_layout((self.elements_per_thread,)), + Float32, + ) + accumulator.fill(Float32(0.0)) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + source_offset = ( + Int64(routed_index) * self.hidden_size + + Int64(fragment) * self.elements_per_thread + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=self.elements_per_thread * 2, + ) + if cutlass.const_expr(self.elements_per_thread == 2): + source = packed_u32_to_bf16x2( + load_global_u32_predicated( + source_pointer, + Int32(routed_index != Int32(-1)), + ) + ).to(Float32) + else: + source = packed_u32x2_to_bf16x4( + load_global_u32x2_predicated( + source_pointer, + Int32(routed_index != Int32(-1)), + ) + ).to(Float32) + accumulator.store(accumulator.load() + source * weight) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group, self.words_per_fragment)), + Uint32, + ) + inputs.fill(Uint32(0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + source_offset = ( + Int64(routed_index) * self.hidden_size + + Int64(fragment) * self.elements_per_thread + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=self.elements_per_thread * 2, + ) + if cutlass.const_expr(self.elements_per_thread == 2): + inputs[item, 0] = load_global_u32_predicated( + source_pointer, + Int32(routed_index != Int32(-1)), + ) + else: + source = load_global_u32x2_predicated( + source_pointer, + Int32(routed_index != Int32(-1)), + ) + for word in cutlass.range_constexpr( + self.words_per_fragment + ): + inputs[item, word] = source[word] + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + if cutlass.const_expr(self.elements_per_thread == 2): + source_values = packed_u32_to_bf16x2( + inputs[item, 0] + ).to(Float32) + else: + source = cute.make_rmem_tensor( + cute.make_layout((self.words_per_fragment,)), + Uint32, + ) + for word in cutlass.range_constexpr( + self.words_per_fragment + ): + source[word] = inputs[item, word] + source_values = packed_u32x2_to_bf16x4( + source.load() + ).to(Float32) + accumulator.store( + accumulator.load() + + source_values + * cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + ) + + result = accumulator.load() + if cutlass.const_expr(self.include_shared_expert): + result = result + shared_values.load().to(Float32) + result = result.to(BFloat16) + destination_rank = token % self.tp_size + local_token = token // self.tp_size + destination_base = cute.arch.load( + ( + contribution_mailbox_peer_addresses.iterator + destination_rank + ).llvm_ptr, + Int64, + ) + destination_offset = ( + (Int64(stage) * self.tp_size + self.rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * self.elements_per_thread + if cutlass.const_expr(self.elements_per_thread == 2): + word = sanitize_negative_zero_u32(bf16x2_to_packed_u32(result)) + store_global_u32_address( + destination_base + destination_offset * 2, + word, + ) + else: + half = sanitize_negative_zero_u32x2(bf16x4_to_packed_u32x2(result)) + store_global_u32x2( + destination_base + destination_offset * 2, + half, + ) + + if block == 0 and tidx == 0: + store_global_u32(stage_state.iterator + ACTIVE_STAGE, stage) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _VectorFinalizeUnicastDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + local_capacity: int, + threads: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.local_capacity = local_capacity + self.threads = threads + self.routed_scaling_factor = routed_scaling_factor + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + self.words_per_fragment = VEC_BF16 // 2 + self.fragments = hidden_size // VEC_BF16 + self.ctas_per_token = math.ceil(self.fragments / threads) + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_peer_addresses, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.top_k * 8, + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + fragment = (block % self.ctas_per_token) * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + metadata_offset = Int64(token) * self.top_k + Int64(metadata_index) + routed_index = cute.arch.load( + (permuted_indices.iterator + metadata_offset).llvm_ptr, + Int32, + ) + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + metadata_offset).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if routed_index == Int32(-1): + # Vector loads are unpredicated; row zero is safe and its weight is zero. + routed_index = Int32(0) + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + routed_index, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + if cutlass.const_expr(self.include_shared_expert): + shared_values = cute.make_rmem_tensor( + cute.make_layout((VEC_BF16,)), + BFloat16, + ) + shared_values.fill(BFloat16(0.0)) + if cutlass.const_expr(self.load_shared_expert_before_pdl): + if fragment < self.fragments: + shared_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + shared_values.store( + packed_u32x4_to_bf16x8(load_global_u32x4(shared_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if fragment < self.fragments: + shared_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + shared_values.store( + packed_u32x4_to_bf16x8(load_global_u32x4(shared_pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if fragment < self.fragments: + accumulator = cute.make_rmem_tensor( + cute.make_layout((VEC_BF16,)), + Float32, + ) + accumulator.fill(Float32(0.0)) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + source_offset = ( + Int64(routed_index) * self.hidden_size + + Int64(fragment) * VEC_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + source_values = packed_u32x4_to_bf16x8( + load_global_u32x4(source_pointer) + ).to(Float32) + if weight != Float32(0.0): + accumulator.store(accumulator.load() + source_values * weight) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group, self.words_per_fragment)), + Uint32, + ) + inputs.fill(Uint32(0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + routed_index = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + source_offset = ( + Int64(routed_index) * self.hidden_size + + Int64(fragment) * VEC_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + source = load_global_u32x4(source_pointer) + for word in cutlass.range_constexpr( + self.words_per_fragment + ): + inputs[item, word] = source[word] + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + source = cute.make_rmem_tensor( + cute.make_layout((self.words_per_fragment,)), + Uint32, + ) + for word in cutlass.range_constexpr( + self.words_per_fragment + ): + source[word] = inputs[item, word] + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + if weight != Float32(0.0): + accumulator.store( + accumulator.load() + + packed_u32x4_to_bf16x8(source.load()).to(Float32) + * weight + ) + + result = accumulator.load() + if cutlass.const_expr(self.include_shared_expert): + result = result + shared_values.load().to(Float32) + result = result.to(BFloat16) + packed = sanitize_negative_zero_u32x4(bf16x8_to_packed_u32x4(result)) + destination_rank = token % self.tp_size + local_token = token // self.tp_size + destination_base = cute.arch.load( + ( + contribution_mailbox_peer_addresses.iterator + destination_rank + ).llvm_ptr, + Int64, + ) + destination_offset = ( + (Int64(stage) * self.tp_size + self.rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + store_global_u32x4( + destination_base + destination_offset * 2, + packed, + ) + + if block == 0 and tidx == 0: + store_global_u32(stage_state.iterator + ACTIVE_STAGE, stage) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _SharedOnlyPublishDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + tp_size: int, + rank: int, + local_capacity: int, + threads: int, + vectors_per_thread: int, + enable_pdl: bool, + ) -> None: + self.hidden_size = hidden_size + self.tp_size = tp_size + self.rank = rank + self.local_capacity = local_capacity + self.threads = threads + self.vectors_per_thread = vectors_per_thread + self.enable_pdl = enable_pdl + self.fragments = hidden_size // VEC_BF16 + self.fragments_per_cta = threads * vectors_per_thread + self.ctas_per_token = math.ceil(self.fragments / self.fragments_per_cta) + + @cute.jit + def __call__( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + shared_output, + stage_state, + contribution_mailbox_peer_addresses, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_peer_addresses: cute.Tensor, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + fragment_base = (block % self.ctas_per_token) * self.fragments_per_cta + tidx + + inputs = cute.make_rmem_tensor( + cute.make_layout( + (self.vectors_per_thread, 4), + stride=(4, 1), + ), + Uint32, + ) + inputs.fill(Uint32(0)) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + for trip in cutlass.range_constexpr(self.vectors_per_thread): + fragment = fragment_base + trip * self.threads + if fragment < self.fragments: + source_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source_pointer) + for word in cutlass.range_constexpr(4): + inputs[trip, word] = packed[word] + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + destination_rank = token % self.tp_size + local_token = token // self.tp_size + destination_base = cute.arch.load( + (contribution_mailbox_peer_addresses.iterator + destination_rank).llvm_ptr, + Int64, + ) + for trip in cutlass.range_constexpr(self.vectors_per_thread): + fragment = fragment_base + trip * self.threads + if fragment < self.fragments: + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = inputs[trip, word] + destination_offset = ( + (Int64(stage) * self.tp_size + self.rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + store_global_u32x4( + destination_base + destination_offset * 2, + sanitize_negative_zero_u32x4(packed.load()), + ) + + if block == 0 and tidx == 0: + store_global_u32(stage_state.iterator + ACTIVE_STAGE, stage) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _OwnerReduceMulticastDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + tp_size: int, + rank: int, + capacity_m: int, + local_capacity: int, + threads: int, + add_residual: bool, + enable_pdl: bool, + ) -> None: + self.hidden_size = hidden_size + self.tp_size = tp_size + self.rank = rank + self.capacity_m = capacity_m + self.local_capacity = local_capacity + self.threads = threads + self.add_residual = add_residual + self.enable_pdl = enable_pdl + self.fragments = hidden_size // VEC_BF16 + self.ctas_per_token = math.ceil(self.fragments / threads) + + @cute.jit + def __call__( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + stage_state: cute.Tensor, + prenorm_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + local_tokens = (m + Int32(self.tp_size - 1)) // Int32(self.tp_size) + self.kernel( + contribution_mailbox, + residual_source, + stage_state, + prenorm_mailbox_multicast_address, + m, + ).launch( + grid=(local_tokens * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + stage_state: cute.Tensor, + prenorm_mailbox_multicast_address: Int64, + m: Int32, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + local_token = block // self.ctas_per_token + fragment = (block % self.ctas_per_token) * self.threads + tidx + token = local_token * self.tp_size + self.rank + active = token < m and fragment < self.fragments + + if cutlass.const_expr(self.add_residual): + residual = cute.make_rmem_tensor( + cute.make_layout((VEC_BF16,)), + BFloat16, + ) + residual.fill(BFloat16(0.0)) + if active: + residual_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + residual_pointer = cute.make_ptr( + BFloat16, + (residual_source.iterator + residual_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + residual.store( + packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + rank_values = cute.make_rmem_tensor( + cute.make_layout((self.tp_size, 4)), + Uint32, + ) + rank_values.fill(Uint32(0)) + dirty = active + while dirty: + dirty = False + for source_rank in cutlass.range_constexpr(self.tp_size): + source_offset = ( + (Int64(stage) * self.tp_size + source_rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source_pointer, volatile=True) + dirty = dirty | fragment_has_negative_zero(packed) + for word in cutlass.range_constexpr(4): + rank_values[source_rank, word] = packed[word] + if active: + reduced = cute.make_rmem_tensor( + cute.make_layout((VEC_BF16,)), + Float32, + ) + reduced.fill(Float32(0.0)) + for source_rank in cutlass.range_constexpr(self.tp_size): + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = rank_values[source_rank, word] + reduced.store( + reduced.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32) + ) + + prenorm = reduced.load() + if cutlass.const_expr(self.add_residual): + prenorm = prenorm + residual.load().to(Float32) + prenorm_bf16 = prenorm.to(BFloat16) + output_offset = ( + Int64(stage) * self.capacity_m + Int64(token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + stmc_bf16x8( + prenorm_mailbox_multicast_address + output_offset * 2, + sanitize_negative_zero_u32x4(bf16x8_to_packed_u32x4(prenorm_bf16)), + ) + + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + NEXT_STAGE, + (stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + if active: + for source_rank in cutlass.range_constexpr(self.tp_size): + source_offset = ( + (Int64(stage) * self.tp_size + source_rank) * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((contribution_mailbox.iterator + source_offset).toint()) + ) + + +class _MaterializeRMSNormDeviceKernel: + def __init__( + self, + *, + hidden_size: int, + capacity_m: int, + threads: int, + rms_epsilon: float, + weight_bias: float, + write_residual_output: bool, + enable_pdl: bool, + ) -> None: + self.hidden_size = hidden_size + self.capacity_m = capacity_m + self.threads = threads + self.rms_epsilon = rms_epsilon + self.weight_bias = weight_bias + self.write_residual_output = write_residual_output + self.enable_pdl = enable_pdl + self.fragments = hidden_size // VEC_BF16 + self.trips = math.ceil(self.fragments / threads) + self.warps = threads // WARP_SIZE + + @cute.jit + def __call__( + self, + prenorm_mailbox: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + gamma: cute.Tensor, + stage_state: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + prenorm_mailbox, + residual_output, + norm_output, + gamma, + stage_state, + ).launch( + grid=(m, 1, 1), + block=(self.threads, 1, 1), + smem=self.warps * 4, + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + prenorm_mailbox: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + gamma: cute.Tensor, + stage_state: cute.Tensor, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + token, _, _ = cute.arch.block_idx() + + gamma_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + prenorm_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + gamma_fragments.fill(BFloat16(0.0)) + prenorm_fragments.fill(BFloat16(0.0)) + for trip in cutlass.range_constexpr(self.trips): + fragment = tidx + trip * self.threads + if fragment < self.fragments: + gamma_pointer = cute.make_ptr( + BFloat16, + (gamma.iterator + Int64(fragment) * VEC_BF16).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + thread_sum = Float32(0.0) + for trip in cutlass.range_constexpr(self.trips): + fragment = tidx + trip * self.threads + if fragment < self.fragments: + mailbox_offset = ( + Int64(stage) * self.capacity_m + Int64(token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + mailbox_pointer = cute.make_ptr( + BFloat16, + (prenorm_mailbox.iterator + mailbox_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(mailbox_pointer, volatile=True) + while fragment_has_negative_zero(packed): + packed = load_global_u32x4(mailbox_pointer, volatile=True) + prenorm = packed_u32x4_to_bf16x8(packed) + prenorm_fragments[trip, None].store(prenorm) + + if cutlass.const_expr(self.write_residual_output): + output_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + store_global_u32x4( + Int64((residual_output.iterator + output_offset).toint()), + packed, + ) + + values = prenorm.to(Float32) + thread_sum = thread_sum + (values * values).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_array(Float32, self.warps) + full_sum = _block_sum(thread_sum, warp_sums, self.warps) + inverse_rms = cute.math.rsqrt( + full_sum / Float32(self.hidden_size) + Float32(self.rms_epsilon), + fastmath=True, + ) + for trip in cutlass.range_constexpr(self.trips): + fragment = tidx + trip * self.threads + if fragment < self.fragments: + gamma_value = gamma_fragments[trip, None].load().to(Float32) + if cutlass.const_expr(self.weight_bias != 0.0): + gamma_value = gamma_value + Float32(self.weight_bias) + result = ( + prenorm_fragments[trip, None].load().to(Float32) + * inverse_rms + * gamma_value + ).to(BFloat16) + output_offset = ( + Int64(token) * self.hidden_size + Int64(fragment) * VEC_BF16 + ) + store_global_u32x4( + Int64((norm_output.iterator + output_offset).toint()), + bf16x8_to_packed_u32x4(result), + ) + + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + for trip in cutlass.range_constexpr(self.trips): + fragment = tidx + trip * self.threads + if fragment < self.fragments: + mailbox_offset = ( + Int64(stage) * self.capacity_m + Int64(token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((prenorm_mailbox.iterator + mailbox_offset).toint()) + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/protocol.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/protocol.py new file mode 100644 index 000000000000..373afd44f7fb --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_bt/protocol.py @@ -0,0 +1,511 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Balanced MNNVL protocol and its two operation paths.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, TypedDict, cast + +import cutlass.cute as cute +import torch +import torch.distributed as dist +from cutlass import BFloat16, Int32, Int64 +from cutlass.cute.runtime import make_fake_compact_tensor + +from ..cute_dsl_primitives import VEC_BF16 +from ..runtime import ( + current_cu_stream, + make_fake_dynamic_compact_tensor, + to_cute, + to_cute_dynamic, +) +from ..symmetric_buffer import SymmetricBuffer +from .device_kernels import ( + LAMPORT_GENERATIONS, + _MaterializeRMSNormDeviceKernel, + _NarrowVectorFinalizeUnicastDeviceKernel, + _OwnerReduceMulticastDeviceKernel, + _ScalarFinalizeUnicastDeviceKernel, + _SharedOnlyPublishDeviceKernel, + _VectorFinalizeUnicastDeviceKernel, +) + + +@dataclass(frozen=True, slots=True) +class BTCollectiveTuning: + reduction_threads: int = 128 + rms_threads: int = 1024 + enable_pdl: bool = True + + +@dataclass(frozen=True, slots=True) +class BTFinalizeTuning: + elements_per_thread: int = VEC_BF16 + threads: int = 128 + prefetch_group: int = 1 + load_shared_expert_before_pdl: bool = False + collective: BTCollectiveTuning = BTCollectiveTuning() + + +@dataclass(frozen=True, slots=True) +class BTAllReduceTuning: + publish_threads: int = 128 + publish_vectors_per_thread: int = 1 + collective: BTCollectiveTuning = BTCollectiveTuning(reduction_threads=32) + + +BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0 = BTFinalizeTuning( + elements_per_thread=2, threads=256 +) +BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1 = BTFinalizeTuning() +BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0 = BTFinalizeTuning( + elements_per_thread=2, threads=256 +) +BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1 = BTFinalizeTuning() +BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0 = BTAllReduceTuning() +BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1 = BTAllReduceTuning( + collective=BTCollectiveTuning(reduction_threads=320) +) +BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0 = BTAllReduceTuning() +BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1 = BTAllReduceTuning( + collective=BTCollectiveTuning(reduction_threads=320) +) + + +@dataclass(slots=True) +class BTProtocolState: + contribution_mailbox: SymmetricBuffer + prenorm_mailbox: SymmetricBuffer + stage_state: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _CompiledTail: + reduce: Any + rms_norm: Any + + +@dataclass(frozen=True, slots=True) +class _CompiledFinalize: + publish: Any + tail: _CompiledTail + + +@dataclass(frozen=True, slots=True) +class _CompiledAllReduce: + publish: Any + tail: _CompiledTail + + +class _PathKwargs(TypedDict): + hidden_size: int + top_k: int + capacity_m: int + write_residual_output: bool + + +class _BTPath: + def __init__( + self, + *, + hidden_size: int, + top_k: int, + capacity_m: int, + write_residual_output: bool, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.capacity_m = capacity_m + self.write_residual_output = write_residual_output + + def _outputs( + self, + m: int, + norm_output: torch.Tensor | None, + residual_output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + shape = (m, self.hidden_size) + device = torch.device("cuda", torch.cuda.current_device()) + if norm_output is None: + norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + if self.write_residual_output and residual_output is None: + residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + return norm_output, residual_output + + def _validate_state(self, state: BTProtocolState, m: int) -> None: + if not 1 <= m <= self.capacity_m: + raise ValueError(f"m must be in [1, {self.capacity_m}]") + if state.contribution_mailbox.peer_addresses is None: + raise ValueError("BT contribution mailbox requires peer addresses") + address = state.prenorm_mailbox.multicast_address + if address is None or address % 16: + raise ValueError( + "BT prenorm mailbox requires a 16-byte-aligned multicast address" + ) + + def _launch_tail( + self, + tail: _CompiledTail, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + state: BTProtocolState, + norm_output: torch.Tensor, + residual_output: torch.Tensor | None, + m: int, + ) -> None: + residual_arg = residual_source if residual_source is not None else norm_output + residual_output_arg = ( + residual_output if residual_output is not None else norm_output + ) + stream = current_cu_stream() + tail.reduce( + to_cute(state.contribution_mailbox.tensor.flatten(), 16), + to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(state.stage_state, 4), + Int64(cast(int, state.prenorm_mailbox.multicast_address)), + Int32(m), + stream, + ) + tail.rms_norm( + to_cute(state.prenorm_mailbox.tensor.flatten(), 16), + to_cute_dynamic( + residual_output_arg.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size), + to_cute(gamma, 16), + to_cute(state.stage_state, 4), + Int32(m), + stream, + ) + + +class FinalizeAllReduceRMSNormBTKernel(_BTPath): + def __init__(self, *, compiled: _CompiledFinalize, **kwargs) -> None: + super().__init__(**kwargs) + self._compiled = compiled + + def __call__( + self, + routed_output: torch.Tensor, + expert_weights: torch.Tensor, + permuted_indices: torch.Tensor, + shared_output: torch.Tensor | None, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: BTProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_state(state, m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + shared_arg = shared_output if shared_output is not None else norm_output + peers = cast(torch.Tensor, state.contribution_mailbox.peer_addresses) + self._compiled.publish( + to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size), + to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k), + to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k), + to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(state.stage_state, 4), + to_cute(peers, 8), + Int32(m), + current_cu_stream(), + ) + self._launch_tail( + self._compiled.tail, + residual_source, + gamma, + state, + norm_output, + residual_output, + m, + ) + return norm_output, residual_output + + +class AllReduceRMSNormBTKernel(_BTPath): + def __init__(self, *, compiled: _CompiledAllReduce, **kwargs) -> None: + super().__init__(**kwargs) + self._compiled = compiled + + def __call__( + self, + local_contribution: torch.Tensor, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: BTProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_state(state, m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + peers = cast(torch.Tensor, state.contribution_mailbox.peer_addresses) + self._compiled.publish( + to_cute_dynamic( + local_contribution.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute(state.stage_state, 4), + to_cute(peers, 8), + Int32(m), + current_cu_stream(), + ) + self._launch_tail( + self._compiled.tail, + residual_source, + gamma, + state, + norm_output, + residual_output, + m, + ) + return norm_output, residual_output + + +class BTProtocol: + """Own BT State and protocol-local compiled variants for both paths.""" + + def __init__( + self, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + capacity_m: int, + rms_epsilon: float, + routed_scaling_factor: float, + weight_bias: float, + *, + include_shared_expert: bool, + add_residual: bool, + write_residual_output: bool, + finalize_tunings: tuple[BTFinalizeTuning, ...], + all_reduce_tunings: tuple[BTAllReduceTuning, ...], + group: dist.ProcessGroup, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.capacity_m = capacity_m + self.local_capacity = math.ceil(capacity_m / tp_size) + self.rms_epsilon = rms_epsilon + self.routed_scaling_factor = routed_scaling_factor + self.weight_bias = weight_bias + self.include_shared_expert = include_shared_expert + self.add_residual = add_residual + self.write_residual_output = write_residual_output + + tail_cache = { + tuning: self._compile_tail(tuning) + for tuning in { + *(item.collective for item in finalize_tunings), + *(item.collective for item in all_reduce_tunings), + } + } + self.finalize_kernels = { + tuning: FinalizeAllReduceRMSNormBTKernel( + compiled=_CompiledFinalize( + publish=self._compile_finalize(tuning), + tail=tail_cache[tuning.collective], + ), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(finalize_tunings) + } + self.all_reduce_kernels = { + tuning: AllReduceRMSNormBTKernel( + compiled=_CompiledAllReduce( + publish=self._compile_all_reduce_publish(tuning), + tail=tail_cache[tuning.collective], + ), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(all_reduce_tunings) + } + self.state = self._create_state(group) + + def _path_kwargs(self) -> _PathKwargs: + return { + "hidden_size": self.hidden_size, + "top_k": self.top_k, + "capacity_m": self.capacity_m, + "write_residual_output": self.write_residual_output, + } + + def _compile_finalize(self, tuning: BTFinalizeTuning): + if tuning.elements_per_thread not in (1, 2, 4, VEC_BF16): + raise ValueError("BT finalize elements_per_thread must be 1, 2, 4, or 8") + kwargs: dict[str, Any] = { + "hidden_size": self.hidden_size, + "top_k": self.top_k, + "tp_size": self.tp_size, + "rank": self.rank, + "local_capacity": self.local_capacity, + "threads": tuning.threads, + "routed_scaling_factor": self.routed_scaling_factor, + "include_shared_expert": self.include_shared_expert, + "load_shared_expert_before_pdl": tuning.load_shared_expert_before_pdl, + "enable_pdl": tuning.collective.enable_pdl, + "prefetch_group": tuning.prefetch_group, + } + device_kernel: Any + if tuning.elements_per_thread == 1: + device_kernel = _ScalarFinalizeUnicastDeviceKernel(**kwargs) + elif tuning.elements_per_thread == VEC_BF16: + device_kernel = _VectorFinalizeUnicastDeviceKernel(**kwargs) + else: + device_kernel = _NarrowVectorFinalizeUnicastDeviceKernel( + **kwargs, elements_per_thread=tuning.elements_per_thread + ) + return cute.compile( + device_kernel, + *self._publish_compile_args(include_routed=True), + ) + + def _compile_all_reduce_publish(self, tuning: BTAllReduceTuning): + device_kernel = _SharedOnlyPublishDeviceKernel( + hidden_size=self.hidden_size, + tp_size=self.tp_size, + rank=self.rank, + local_capacity=self.local_capacity, + threads=tuning.publish_threads, + vectors_per_thread=tuning.publish_vectors_per_thread, + enable_pdl=tuning.collective.enable_pdl, + ) + return cute.compile( + device_kernel, + *self._publish_compile_args(include_routed=False), + ) + + def _publish_compile_args(self, *, include_routed: bool) -> tuple: + activation = make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ) + common = ( + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + make_fake_compact_tensor(Int64, (self.tp_size,), assumed_align=8), + Int32(self.capacity_m), + current_cu_stream(), + ) + if not include_routed: + return (activation, *common) + return ( + activation, + make_fake_dynamic_compact_tensor( + BFloat16, alignment=2, divisibility=self.top_k + ), + make_fake_dynamic_compact_tensor( + Int32, alignment=4, divisibility=self.top_k + ), + activation, + *common, + ) + + def _compile_tail(self, tuning: BTCollectiveTuning) -> _CompiledTail: + reduce_kernel = _OwnerReduceMulticastDeviceKernel( + hidden_size=self.hidden_size, + tp_size=self.tp_size, + rank=self.rank, + capacity_m=self.capacity_m, + local_capacity=self.local_capacity, + threads=tuning.reduction_threads, + add_residual=self.add_residual, + enable_pdl=tuning.enable_pdl, + ) + reduce_elements = ( + LAMPORT_GENERATIONS * self.tp_size * self.local_capacity * self.hidden_size + ) + reduce = cute.compile( + reduce_kernel, + make_fake_compact_tensor(BFloat16, (reduce_elements,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(self.capacity_m), + current_cu_stream(), + ) + rms_kernel = _MaterializeRMSNormDeviceKernel( + hidden_size=self.hidden_size, + capacity_m=self.capacity_m, + threads=tuning.rms_threads, + rms_epsilon=self.rms_epsilon, + weight_bias=self.weight_bias, + write_residual_output=self.write_residual_output, + enable_pdl=tuning.enable_pdl, + ) + prenorm_elements = LAMPORT_GENERATIONS * self.capacity_m * self.hidden_size + rms_norm = cute.compile( + rms_kernel, + make_fake_compact_tensor(BFloat16, (prenorm_elements,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int32(self.capacity_m), + current_cu_stream(), + ) + return _CompiledTail(reduce=reduce, rms_norm=rms_norm) + + def _create_state(self, group: dist.ProcessGroup) -> BTProtocolState: + if dist.get_world_size(group) != self.tp_size: + raise ValueError("ProcessGroup size does not match tp_size") + if dist.get_rank(group) != self.rank: + raise ValueError("ProcessGroup rank does not match rank") + device = torch.device("cuda", torch.cuda.current_device()) + contribution = SymmetricBuffer.allocate( + ( + LAMPORT_GENERATIONS, + self.tp_size, + self.local_capacity, + self.hidden_size, + ), + torch.bfloat16, + device, + group, + materialize_peer_addresses=True, + ) + contribution.tensor.view(torch.int16).fill_(-32768) + prenorm = SymmetricBuffer.allocate( + ( + LAMPORT_GENERATIONS, + self.capacity_m, + self.hidden_size, + ), + torch.bfloat16, + device, + group, + require_multicast=True, + ) + prenorm.tensor.view(torch.int16).fill_(-32768) + return BTProtocolState( + contribution_mailbox=contribution, + prenorm_mailbox=prenorm, + stage_state=torch.zeros((2,), dtype=torch.int32, device=device), + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/__init__.py new file mode 100644 index 000000000000..357f07916802 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""High-throughput MNNVL protocol.""" + +from .protocol import ( + HT_ALL_REDUCE_GB300_TP16_H8192, + HT_ALL_REDUCE_GB300_TP8_H8192, + HT_FINALIZE_GB300_TP16_H8192_K10, + HT_FINALIZE_GB300_TP8_H8192_K10, + HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049, + HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048, + HTAllReduceTuning, + HTFinalizeTuning, +) + +__all__ = [ + "HT_ALL_REDUCE_GB300_TP8_H8192", + "HT_ALL_REDUCE_GB300_TP16_H8192", + "HT_FINALIZE_GB300_TP8_H8192_K10", + "HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049", + "HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048", + "HT_FINALIZE_GB300_TP16_H8192_K10", + "HTAllReduceTuning", + "HTFinalizeTuning", +] diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py new file mode 100644 index 000000000000..5c8421590778 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py @@ -0,0 +1,978 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Persistent BF16 MoE finalize, TP reduction, and RMSNorm for SM100.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from ..cute_dsl_primitives import ( + VEC_BF16, + WARP_SIZE, + bf16x8_to_packed_u32x4, + cpasync_bulk_g2s, + fence_proxy_async_shared_cta, + fragment_has_negative_zero, + ldmc_bf16x8, + load_global_bf16_as_f32, + load_global_u32x4_address, + load_shared_u32x4, + packed_negative_zero_bf16x8, + packed_u32x4_to_bf16x8, + remote_release_add1_u32, + sanitize_negative_zero_u32x4, + stmc_bf16x8, + store_global_u32x4, + store_shared_u32x4, +) + +SMEM_ALIGNMENT = 1024 + + +class _MoeFinalizeAllReduceRMSNormHTDeviceKernel: + def __init__( + self, + *, + hidden: int, + top_k: int, + tp: int, + rank: int, + active_ctas: int, + stages: int, + consumer_threads: int, + vectors_per_thread: int, + reduction_warps: int, + reduction_cta_groups: int | None, + rms_token_groups: int, + rms_pipeline_stages: int, + rms_shard_major: bool, + rms_epsilon: float, + routed_scaling_factor: float, + weight_bias: float, + include_shared_expert: bool, + add_residual: bool, + write_residual_output: bool, + enable_pdl: bool, + ) -> None: + if tp not in (2, 4, 8, 16): + raise ValueError("tp must be 2, 4, 8, or 16") + if rank < 0 or rank >= tp: + raise ValueError("rank must be in [0, tp)") + if hidden <= 0 or hidden % VEC_BF16: + raise ValueError("hidden must be a positive multiple of 8") + if top_k < 0: + raise ValueError("top_k must be nonnegative") + if active_ctas <= 0 or active_ctas % tp: + raise ValueError("active_ctas must be positive and divisible by tp") + if stages < 2: + raise ValueError("stages must be at least 2") + if consumer_threads <= 0 or consumer_threads % WARP_SIZE: + raise ValueError("consumer_threads must be a positive warp multiple") + if vectors_per_thread <= 0: + raise ValueError("vectors_per_thread must be positive") + if reduction_warps not in (1, 2, 4, 8): + raise ValueError("reduction_warps must be 1, 2, 4, or 8") + if rms_token_groups not in (1, 2, 4): + raise ValueError("rms_token_groups must be 1, 2, or 4") + if consumer_threads % rms_token_groups: + raise ValueError("consumer threads must divide across RMS token groups") + if rms_pipeline_stages not in (1, 2, 3): + raise ValueError("rms_pipeline_stages must be 1, 2, or 3") + block_threads = consumer_threads + (2 + reduction_warps) * WARP_SIZE + if block_threads > 1024: + raise ValueError("warp roles exceed the CUDA block limit") + shard_elements = consumer_threads * VEC_BF16 * vectors_per_thread + if hidden <= 0 or hidden % shard_elements: + raise ValueError(f"hidden must be divisible by {shard_elements}") + cta_groups = active_ctas // tp + if reduction_cta_groups is None: + reduction_cta_groups = active_ctas // tp + if reduction_cta_groups <= 0 or reduction_cta_groups * tp > active_ctas: + raise ValueError("reduction CTA groups and shards must fit the grid") + contributions = top_k + int(include_shared_expert) + if contributions <= 0: + raise ValueError("at least one local contribution is required") + self.hidden = hidden + self.top_k = top_k + self.tp = tp + self.rank = rank + self.active_ctas = active_ctas + self.stages = stages + self.vectors_per_thread = vectors_per_thread + self.consumer_threads = consumer_threads + self.reduction_warps = reduction_warps + self.reduction_cta_groups = reduction_cta_groups + self.reduction_ctas = reduction_cta_groups * tp + self.rms_token_groups = rms_token_groups + self.rms_pipeline_stages = rms_pipeline_stages + self.rms_shard_major = rms_shard_major + self.rms_epsilon = rms_epsilon + self.routed_scaling_factor = routed_scaling_factor + self.weight_bias = weight_bias + self.include_shared_expert = include_shared_expert + self.add_residual = add_residual + self.write_residual_output = write_residual_output + self.enable_pdl = enable_pdl + self.metadata_chunks = (top_k + WARP_SIZE - 1) // WARP_SIZE + self.metadata_slots = max(top_k, 1) + self.consumer_warps = consumer_threads // WARP_SIZE + self.rms_threads_per_token = consumer_threads // rms_token_groups + self.rms_warps_per_token = self.rms_threads_per_token // WARP_SIZE + self.rms_stage_slots = rms_token_groups * rms_pipeline_stages + self.rms_warp_sum_slots = self.rms_stage_slots * self.rms_warps_per_token + self.publisher_warp = 1 + self.consumer_warps + self.reduction_warp_begin = self.publisher_warp + 1 + self.reduction_threads = reduction_warps * WARP_SIZE + self.block_threads = block_threads + self.shard_elements = shard_elements + self.shard_bytes = shard_elements * 2 + self.hidden_shards = hidden // shard_elements + self.contributions = contributions + if ( + rms_pipeline_stages > 1 + and self.rms_stage_slots * hidden > self.shard_elements * stages + ): + raise ValueError("finalize stage storage cannot hold the RMS pipeline") + self.cta_groups = cta_groups + self.packs_per_token = hidden // VEC_BF16 + if self.packs_per_token % tp: + raise ValueError("hidden vector count must be divisible by tp") + if self.packs_per_token % consumer_threads: + raise ValueError("token vectors must divide evenly across consumers") + self.clear_vectors_per_thread = self.packs_per_token // consumer_threads + self.copy_threads = self.rms_threads_per_token + if self.packs_per_token % self.copy_threads: + raise ValueError("token vectors must divide evenly across finalize threads") + self.rms_vectors_per_thread = self.packs_per_token // self.copy_threads + self.packs_per_reduction_shard = self.packs_per_token // tp + if rms_shard_major: + if tp < self.rms_warps_per_token or tp % self.rms_warps_per_token: + raise ValueError( + "shard-major RMS requires an integer number of reduction " + "shards per RMS warp" + ) + self.reduction_shards_per_rms_warp = tp // self.rms_warps_per_token + if ( + self.rms_vectors_per_thread * WARP_SIZE + != self.packs_per_reduction_shard * self.reduction_shards_per_rms_warp + ): + raise ValueError( + "shard-major RMS warp coverage must match its reduction shards" + ) + else: + self.reduction_shards_per_rms_warp = 0 + if self.packs_per_reduction_shard % self.reduction_threads: + raise ValueError("the reduction shard must divide evenly across threads") + self.reduction_vectors_per_thread = ( + self.packs_per_reduction_shard // self.reduction_threads + ) + + @cute.jit + def _rms_arrive_and_wait(self, rms_group: Int32) -> None: + barrier_0 = pipeline.NamedBarrier( + barrier_id=2, num_threads=self.rms_threads_per_token + ) + if cutlass.const_expr(self.rms_token_groups > 1): + barrier_1 = pipeline.NamedBarrier( + barrier_id=3, num_threads=self.rms_threads_per_token + ) + if cutlass.const_expr(self.rms_token_groups == 4): + barrier_2 = pipeline.NamedBarrier( + barrier_id=4, num_threads=self.rms_threads_per_token + ) + barrier_3 = pipeline.NamedBarrier( + barrier_id=5, num_threads=self.rms_threads_per_token + ) + if rms_group == 0: + barrier_0.arrive_and_wait() + elif cutlass.const_expr(self.rms_token_groups == 2): # noqa: SIM114 + barrier_1.arrive_and_wait() + elif rms_group == 1: + barrier_1.arrive_and_wait() + elif rms_group == 2: + barrier_2.arrive_and_wait() + else: + barrier_3.arrive_and_wait() + else: + barrier_0.arrive_and_wait() + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + local_contributions: cute.Tensor, + prenorm_mailbox: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + ready_counter_peer_addresses: cute.Tensor, + ready_counters: cute.Tensor, + processed_counters: cute.Tensor, + local_contributions_multicast_address: Int64, + prenorm_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + smem_layout = cute.make_layout( + (self.shard_elements * self.stages,), stride=(1,) + ) + + @cute.struct + class SharedStorage: + barriers: cute.struct.MemRange[Int64, 2 * self.stages] + stage_probs: cute.struct.MemRange[Float32, self.stages] + cached_rows: cute.struct.MemRange[Int32, self.metadata_slots] + cached_probs: cute.struct.MemRange[Float32, self.metadata_slots] + consumer_progress: cute.struct.MemRange[Int32, self.consumer_warps] + norm_warp_sums: cute.struct.MemRange[Float32, self.rms_warp_sum_slots] + norm_inv_rms: cute.struct.MemRange[Float32, self.rms_stage_slots] + rows: cute.struct.Align[ + cute.struct.MemRange[BFloat16, cute.cosize(smem_layout)], + SMEM_ALIGNMENT, + ] + + self.shared_storage: type[cute.struct.Struct] = SharedStorage + self.kernel( + routed_output, + shared_output, + residual_source, + gamma, + expert_weights, + permuted_indices, + local_contributions, + prenorm_mailbox, + residual_output, + norm_output, + ready_counter_peer_addresses, + ready_counters, + processed_counters, + local_contributions_multicast_address, + prenorm_mailbox_multicast_address, + m, + smem_layout, + ).launch( + grid=(self.active_ctas, 1, 1), + block=(self.block_threads, 1, 1), + min_blocks_per_mp=1, + use_pdl=self.enable_pdl, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + routed_source: cute.Tensor, + shared_source: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + local_contributions: cute.Tensor, + prenorm_mailbox: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + ready_counter_peer_addresses: cute.Tensor, + ready_counters: cute.Tensor, + processed_counters: cute.Tensor, + local_contributions_multicast_address: Int64, + prenorm_mailbox_multicast_address: Int64, + m: Int32, + smem_layout: cute.Layout, + ) -> None: + block = cute.arch.block_idx()[0] + tidx = cute.arch.thread_idx()[0] + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + cta_group = block // self.tp + cta_slot = block % self.tp + wave = Int64(cta_group) + token = wave * self.tp + cta_slot + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + rows = storage.rows.get_tensor(smem_layout) + barrier_storage = storage.barriers.data_ptr() + stage_probs = storage.stage_probs.data_ptr() + cached_rows = storage.cached_rows.data_ptr() + cached_probs = storage.cached_probs.data_ptr() + consumer_progress = storage.consumer_progress.data_ptr() + norm_warp_sums = storage.norm_warp_sums.data_ptr() + norm_inv_rms = storage.norm_inv_rms.data_ptr() + if tidx < self.consumer_warps: + cute.arch.store((consumer_progress + tidx).llvm_ptr, Int32(0)) + cute.arch.sync_threads() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + load_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=barrier_storage, + num_stages=self.stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.consumer_warps + ), + tx_count=self.shard_bytes, + ) + if warp == 0: + producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.stages + ) + peek_empty = cutlass.Boolean(1) + if token < Int64(m): + peek_empty = load_pipeline.producer_try_acquire(producer_state) + while token < Int64(m): + for metadata_chunk in cutlass.range_constexpr(self.metadata_chunks): + metadata_slot = metadata_chunk * WARP_SIZE + lane + if metadata_slot < self.top_k: + item = Int64(token) * self.top_k + metadata_slot + row = cute.arch.load( + (permuted_indices.iterator + item).llvm_ptr, Int32 + ) + prob = load_global_bf16_as_f32( + Int64((expert_weights.iterator + item).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + prob = prob * Float32(self.routed_scaling_factor) + if row == Int32(-1): + row = Int32(0) + prob = Float32(0.0) + cute.arch.store((cached_rows + metadata_slot).llvm_ptr, row) + cute.arch.store((cached_probs + metadata_slot).llvm_ptr, prob) + cute.arch.sync_warp() + for shard in cutlass.range_constexpr(self.hidden_shards): + for contribution in cutlass.range_constexpr(self.contributions): + load_pipeline.producer_acquire(producer_state, peek_empty) + if lane == 0: + if cutlass.const_expr(contribution < self.top_k): + prob = cute.arch.load( + (cached_probs + contribution).llvm_ptr, + Float32, + ) + row = cute.arch.load( + (cached_rows + contribution).llvm_ptr, + Int32, + ) + source_element = ( + Int64(row) * self.hidden + + shard * self.shard_elements + ) + source = routed_source.iterator + source_element + else: + prob = Float32(1.0) + source_element = ( + Int64(token) * self.hidden + + shard * self.shard_elements + ) + source = shared_source.iterator + source_element + cute.arch.store( + (stage_probs + producer_state.index).llvm_ptr, + prob, + ) + fence_proxy_async_shared_cta() + cpasync_bulk_g2s( + source, + rows.iterator + + producer_state.index * self.shard_elements, + load_pipeline.producer_get_barrier(producer_state), + Int32(self.shard_bytes), + ) + producer_state.advance() + peek_empty = load_pipeline.producer_try_acquire(producer_state) + wave += self.cta_groups + token = wave * self.tp + cta_slot + load_pipeline.producer_tail(producer_state) + elif warp > 0 and warp <= self.consumer_warps: + finalize_join = pipeline.NamedBarrier( + barrier_id=6, num_threads=self.consumer_threads + ) + consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.stages + ) + consumer_tid = tidx - WARP_SIZE + consumer_wave = Int64(cta_group) + token = consumer_wave * self.tp + cta_slot + consumer_token_progress = Int32(0) + while token < Int64(m): + for shard in cutlass.range_constexpr(self.hidden_shards): + if cutlass.const_expr(self.top_k == 0): + load_pipeline.consumer_wait(consumer_state) + for trip in cutlass.range_constexpr(self.vectors_per_thread): + output_element = ( + Int64(token) * self.hidden + + shard * self.shard_elements + + trip * self.consumer_threads * VEC_BF16 + + consumer_tid * VEC_BF16 + ) + store_global_u32x4( + Int64( + ( + local_contributions.iterator + output_element + ).toint() + ), + load_shared_u32x4( + rows.iterator + + consumer_state.index * self.shard_elements + + trip * self.consumer_threads * VEC_BF16 + + consumer_tid * VEC_BF16 + ), + ) + load_pipeline.consumer_release(consumer_state) + consumer_state.advance() + accum = cute.make_rmem_tensor( + cute.make_layout( + (self.vectors_per_thread, VEC_BF16), + stride=(VEC_BF16, 1), + ), + Float32, + ) + accum.fill(Float32(0.0)) + for _ in cutlass.range_constexpr( + self.contributions if self.top_k > 0 else 0 + ): + load_pipeline.consumer_wait(consumer_state) + prob = cute.arch.load( + (stage_probs + consumer_state.index).llvm_ptr, + Float32, + ) + if prob != Float32(0.0): + for trip in cutlass.range_constexpr( + self.vectors_per_thread + ): + stage_ptr = ( + rows.iterator + + consumer_state.index * self.shard_elements + + trip * self.consumer_threads * VEC_BF16 + + consumer_tid * VEC_BF16 + ) + values = packed_u32x4_to_bf16x8( + load_shared_u32x4(stage_ptr) + ).to(Float32) + accum[trip, None].store( + accum[trip, None].load() + values * prob + ) + load_pipeline.consumer_release(consumer_state) + consumer_state.advance() + for trip in cutlass.range_constexpr( + self.vectors_per_thread if self.top_k > 0 else 0 + ): + output_element = ( + Int64(token) * self.hidden + + shard * self.shard_elements + + trip * self.consumer_threads * VEC_BF16 + + consumer_tid * VEC_BF16 + ) + store_global_u32x4( + Int64( + (local_contributions.iterator + output_element).toint() + ), + bf16x8_to_packed_u32x4( + accum[trip, None].load().to(BFloat16) + ), + ) + clear_value = packed_negative_zero_bf16x8() + token_pack = token * self.packs_per_token + for clear_item in cutlass.range_constexpr( + self.clear_vectors_per_thread + ): + clear_pack = consumer_tid + clear_item * self.consumer_threads + clear_element = (token_pack + clear_pack) * VEC_BF16 + store_global_u32x4( + Int64((prenorm_mailbox.iterator + clear_element).toint()), + clear_value, + ) + cute.arch.sync_warp() + consumer_token_progress += 1 + if lane == 0: + cute.arch.store( + (consumer_progress + warp - 1).llvm_ptr, + consumer_token_progress, + sem="release", + scope="cta", + ) + consumer_wave += self.cta_groups + token = consumer_wave * self.tp + cta_slot + finalize_join.arrive_and_wait() + if cutlass.const_expr(self.rms_token_groups > 1): + rms_group = (warp - 1) // self.rms_warps_per_token + rms_group_warp = warp - 1 - rms_group * self.rms_warps_per_token + copy_tid = rms_group_warp * WARP_SIZE + lane + else: + rms_group = Int32(0) + rms_group_warp = warp - 1 + copy_tid = consumer_tid + rms_pack_base = copy_tid + rms_pack_stride = self.copy_threads + if cutlass.const_expr(self.rms_shard_major): + rms_pack_base = ( + rms_group_warp + * self.reduction_shards_per_rms_warp + * self.packs_per_reduction_shard + + lane + ) + rms_pack_stride = WARP_SIZE + copy_wave = Int64(cta_group) + Int64(rms_group) * self.cta_groups + copy_token = copy_wave * self.tp + cta_slot + if cutlass.const_expr(self.rms_pipeline_stages > 1): + rms_wave_stride = self.cta_groups * self.rms_token_groups + while copy_token < Int64(m): + for rms_stage in cutlass.range_constexpr(self.rms_pipeline_stages): + stage_wave = copy_wave + rms_stage * rms_wave_stride + stage_token = stage_wave * self.tp + cta_slot + if stage_token < Int64(m): + stage_slot = ( + rms_group * self.rms_pipeline_stages + rms_stage + ) + token_pack = stage_token * self.packs_per_token + copy_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.rms_vectors_per_thread, 4), + stride=(4, 1), + ), + Uint32, + ) + all_ready = cutlass.Boolean(0) + while not all_ready: + all_ready = cutlass.Boolean(1) + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + packed = load_global_u32x4_address( + Int64( + ( + prenorm_mailbox.iterator + + linear_pack * VEC_BF16 + ).toint() + ), + volatile=True, + ) + copy_fragments[item, None].store(packed) + all_ready = all_ready and ( + not fragment_has_negative_zero(packed) + ) + thread_sum = Float32(0.0) + if cutlass.const_expr(self.write_residual_output): + prenorm_packed = [] + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + pack = rms_pack_base + item * rms_pack_stride + prenorm = packed_u32x4_to_bf16x8( + copy_fragments[item, None].load() + ) + packed_prenorm = bf16x8_to_packed_u32x4(prenorm) + if cutlass.const_expr(self.write_residual_output): + prenorm_packed.append(packed_prenorm) + # Finalize has drained `rows`; __init__ verifies the RMS layout fits. + store_shared_u32x4( + rows.iterator + + stage_slot * self.hidden + + pack * VEC_BF16, + packed_prenorm, + ) + prenorm_f32 = prenorm.to(Float32) + thread_sum = thread_sum + ( + prenorm_f32 * prenorm_f32 + ).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + if cutlass.const_expr(self.write_residual_output): + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + residual_address = Int64( + ( + residual_output.iterator + + linear_pack * VEC_BF16 + ).toint() + ) + store_global_u32x4( + residual_address, prenorm_packed[item] + ) + warp_sum = cute.arch.warp_reduction_sum(thread_sum) + if lane == 0: + cute.arch.store( + ( + norm_warp_sums + + stage_slot * self.rms_warps_per_token + + rms_group_warp + ).llvm_ptr, + warp_sum, + ) + self._rms_arrive_and_wait(rms_group) + if warp == 1 + rms_group * self.rms_warps_per_token: + for rms_stage in cutlass.range_constexpr( + self.rms_pipeline_stages + ): + stage_wave = copy_wave + rms_stage * rms_wave_stride + stage_token = stage_wave * self.tp + cta_slot + if stage_token < Int64(m): + stage_slot = ( + rms_group * self.rms_pipeline_stages + rms_stage + ) + cta_sum = Float32(0.0) + if lane < self.rms_warps_per_token: + cta_sum = cute.arch.load( + ( + norm_warp_sums + + stage_slot * self.rms_warps_per_token + + lane + ).llvm_ptr, + Float32, + ) + cta_sum = cute.arch.warp_reduction_sum(cta_sum) + if lane == 0: + inv_rms = cute.math.rsqrt( + cta_sum / Float32(self.hidden) + + Float32(self.rms_epsilon), + fastmath=True, + ) + cute.arch.store( + (norm_inv_rms + stage_slot).llvm_ptr, + inv_rms, + ) + self._rms_arrive_and_wait(rms_group) + gamma_values = [] + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + gamma_value = packed_u32x4_to_bf16x8( + load_global_u32x4_address( + Int64((gamma.iterator + pack * VEC_BF16).toint()) + ) + ).to(Float32) + if cutlass.const_expr(self.weight_bias != 0.0): + gamma_value = gamma_value + Float32(self.weight_bias) + gamma_values.append(gamma_value) + for rms_stage in cutlass.range_constexpr(self.rms_pipeline_stages): + stage_wave = copy_wave + rms_stage * rms_wave_stride + stage_token = stage_wave * self.tp + cta_slot + if stage_token < Int64(m): + stage_slot = ( + rms_group * self.rms_pipeline_stages + rms_stage + ) + token_pack = stage_token * self.packs_per_token + inv_rms = cute.arch.load( + (norm_inv_rms + stage_slot).llvm_ptr, Float32 + ) + norm_packed = [] + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + pack = rms_pack_base + item * rms_pack_stride + prenorm = packed_u32x4_to_bf16x8( + load_shared_u32x4( + rows.iterator + + stage_slot * self.hidden + + pack * VEC_BF16 + ) + ).to(Float32) + result = (prenorm * inv_rms * gamma_values[item]).to( + BFloat16 + ) + norm_packed.append(bf16x8_to_packed_u32x4(result)) + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + store_global_u32x4( + Int64( + ( + norm_output.iterator + + linear_pack * VEC_BF16 + ).toint() + ), + norm_packed[item], + ) + copy_wave += rms_wave_stride * self.rms_pipeline_stages + copy_token = copy_wave * self.tp + cta_slot + if cutlass.const_expr(self.rms_pipeline_stages == 1): + while copy_token < Int64(m): + token_pack = copy_token * self.packs_per_token + copy_values = [] + copy_sources = [] + if cutlass.const_expr(self.write_residual_output): + copy_destinations = [] + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + source_address = Int64( + (prenorm_mailbox.iterator + linear_pack * VEC_BF16).toint() + ) + copy_sources.append(source_address) + if cutlass.const_expr(self.write_residual_output): + copy_destinations.append( + Int64( + ( + residual_output.iterator + + linear_pack * VEC_BF16 + ).toint() + ) + ) + copy_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.rms_vectors_per_thread, 4), + stride=(4, 1), + ), + Uint32, + ) + all_ready = cutlass.Boolean(0) + while not all_ready: + all_ready = cutlass.Boolean(1) + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + packed = load_global_u32x4_address( + copy_sources[item], + volatile=True, + ) + copy_fragments[item, None].store(packed) + all_ready = all_ready and ( + not fragment_has_negative_zero(packed) + ) + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + copy_values.append(copy_fragments[item, None].load()) + prenorm_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.rms_vectors_per_thread, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + thread_sum = Float32(0.0) + if cutlass.const_expr(self.write_residual_output): + prenorm_packed = [] + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + prenorm = packed_u32x4_to_bf16x8(copy_values[item]) + prenorm_fragments[item, None].store(prenorm) + packed_prenorm = bf16x8_to_packed_u32x4(prenorm) + if cutlass.const_expr(self.write_residual_output): + prenorm_packed.append(packed_prenorm) + prenorm_f32 = prenorm.to(Float32) + thread_sum = thread_sum + (prenorm_f32 * prenorm_f32).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + if cutlass.const_expr(self.write_residual_output): + for item in cutlass.range_constexpr( + self.rms_vectors_per_thread + ): + store_global_u32x4( + copy_destinations[item], prenorm_packed[item] + ) + warp_sum = cute.arch.warp_reduction_sum(thread_sum) + if lane == 0: + cute.arch.store((norm_warp_sums + warp - 1).llvm_ptr, warp_sum) + self._rms_arrive_and_wait(rms_group) + if warp == 1 + rms_group * self.rms_warps_per_token: + cta_sum = Float32(0.0) + if lane < self.rms_warps_per_token: + cta_sum = cute.arch.load( + ( + norm_warp_sums + + rms_group * self.rms_warps_per_token + + lane + ).llvm_ptr, + Float32, + ) + cta_sum = cute.arch.warp_reduction_sum(cta_sum) + if lane == 0: + inv_rms = cute.math.rsqrt( + cta_sum / Float32(self.hidden) + + Float32(self.rms_epsilon), + fastmath=True, + ) + cute.arch.store( + (norm_inv_rms + rms_group).llvm_ptr, inv_rms + ) + self._rms_arrive_and_wait(rms_group) + inv_rms = cute.arch.load( + (norm_inv_rms + rms_group).llvm_ptr, Float32 + ) + gamma_values = [] + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + gamma_value = packed_u32x4_to_bf16x8( + load_global_u32x4_address( + Int64((gamma.iterator + pack * VEC_BF16).toint()) + ) + ).to(Float32) + if cutlass.const_expr(self.weight_bias != 0.0): + gamma_value = gamma_value + Float32(self.weight_bias) + gamma_values.append(gamma_value) + norm_packed = [] + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + prenorm_for_norm = ( + prenorm_fragments[item, None].load().to(Float32) + ) + result = (prenorm_for_norm * inv_rms * gamma_values[item]).to( + BFloat16 + ) + norm_packed.append(bf16x8_to_packed_u32x4(result)) + for item in cutlass.range_constexpr(self.rms_vectors_per_thread): + pack = rms_pack_base + item * rms_pack_stride + linear_pack = token_pack + pack + store_global_u32x4( + Int64( + (norm_output.iterator + linear_pack * VEC_BF16).toint() + ), + norm_packed[item], + ) + copy_wave += self.cta_groups * self.rms_token_groups + copy_token = copy_wave * self.tp + cta_slot + elif warp == self.publisher_warp: + owner_ready_address = cute.arch.load( + (ready_counter_peer_addresses.iterator + cta_slot).llvm_ptr, + Int64, + ) + first_token = Int64(cta_group) * self.tp + cta_slot + token_count = Int32(0) + if first_token < Int64(m): + token_count = Int32( + (Int64(m) + self.active_ctas - 1 - first_token) // self.active_ctas + ) + published = Int32(0) + while published < token_count: + observed = token_count + if lane < self.consumer_warps: + observed = cute.arch.load( + (consumer_progress + lane).llvm_ptr, + Int32, + sem="relaxed", + scope="cta", + ) + frontier = cute.arch.warp_reduction( + observed, lambda x, y: cutlass.min(x, y) + ) + if frontier > published: + acquired = token_count + if lane < self.consumer_warps: + acquired = cute.arch.load( + (consumer_progress + lane).llvm_ptr, + Int32, + sem="acquire", + scope="cta", + ) + frontier = cute.arch.warp_reduction( + acquired, lambda x, y: cutlass.min(x, y) + ) + cute.arch.sync_warp() + batch = cutlass.min(frontier - published, Int32(WARP_SIZE)) + if lane < batch: + sequence = Int64(published + lane) + publish_token = ( + Int64(cta_group) + sequence * self.cta_groups + ) * self.tp + cta_slot + owner_token = publish_token // self.tp + remote_release_add1_u32(owner_ready_address + owner_token * 4) + published += batch + elif ( + block < self.reduction_ctas + and warp >= self.reduction_warp_begin + and (warp < self.reduction_warp_begin + self.reduction_warps) + ): + reduction_warp = warp - self.reduction_warp_begin + reduction_tid = reduction_warp * WARP_SIZE + lane + reduction_barrier = pipeline.NamedBarrier( + barrier_id=1, num_threads=self.reduction_threads + ) + reduction_shard = block % self.tp + local_token = Int64(block // self.tp) + token = local_token * self.tp + self.rank + while token < Int64(m): + processed_index = local_token * self.tp + reduction_shard + target = Uint32(0) + if reduction_tid == 0: + ready_counter_address = ( + ready_counters.iterator + local_token + ).llvm_ptr + processed_counter_address = ( + processed_counters.iterator + processed_index + ).llvm_ptr + target = cute.arch.load(processed_counter_address, Uint32) + Uint32( + self.tp + ) + observed = Uint32(0) + while observed != target: + observed = cute.arch.load( + ready_counter_address, + Uint32, + sem="relaxed", + scope="sys", + ) + cute.arch.load( + ready_counter_address, + Uint32, + sem="acquire", + scope="sys", + ) + reduction_barrier.arrive_and_wait() + values = [] + addresses = [] + token_pack = token * self.packs_per_token + shard_pack = reduction_shard * self.packs_per_reduction_shard + for item in cutlass.range_constexpr(self.reduction_vectors_per_thread): + pack = shard_pack + reduction_tid + item * self.reduction_threads + input_address = ( + local_contributions_multicast_address + (token_pack + pack) * 16 + ) + output_address = ( + prenorm_mailbox_multicast_address + (token_pack + pack) * 16 + ) + reduced_packed = ldmc_bf16x8(input_address) + reduced_values = packed_u32x4_to_bf16x8(reduced_packed).to(Float32) + if cutlass.const_expr(self.add_residual): + residual_values = packed_u32x4_to_bf16x8( + load_global_u32x4_address( + Int64( + ( + residual_source.iterator + + (token_pack + pack) * VEC_BF16 + ).toint() + ) + ) + ).to(Float32) + reduced_values = reduced_values + residual_values + reduced_packed = bf16x8_to_packed_u32x4(reduced_values.to(BFloat16)) + values.append(sanitize_negative_zero_u32x4(reduced_packed)) + addresses.append(output_address) + for item in cutlass.range_constexpr(self.reduction_vectors_per_thread): + stmc_bf16x8(addresses[item], values[item]) + if reduction_tid == 0: + cute.arch.store( + (processed_counters.iterator + processed_index).llvm_ptr, + target, + ) + local_token += self.reduction_cta_groups + token = local_token * self.tp + self.rank + cute.arch.sync_threads() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/protocol.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/protocol.py new file mode 100644 index 000000000000..b024387b9d7d --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ht/protocol.py @@ -0,0 +1,469 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""High-throughput MNNVL protocol and its two operation paths.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypedDict, cast + +import cutlass.cute as cute +import torch +import torch.distributed as dist +from cutlass import BFloat16, Int32, Int64, Uint32 +from cutlass.cute.runtime import make_fake_compact_tensor + +from ..runtime import ( + current_cu_stream, + make_fake_dynamic_compact_tensor, + to_cute, + to_cute_dynamic, +) +from ..symmetric_buffer import SymmetricBuffer +from .device_kernel import _MoeFinalizeAllReduceRMSNormHTDeviceKernel + + +@dataclass(frozen=True, slots=True) +class HTFinalizeTuning: + persistent_ctas: int | None = None + consumer_threads: int = 512 + vectors_per_thread: int = 2 + stages: int = 6 + reduction_warps: int = 1 + reduction_cta_groups: int | None = None + rms_token_groups: int = 2 + rms_pipeline_stages: int = 2 + rms_shard_major: bool = False + enable_pdl: bool = True + + +@dataclass(frozen=True, slots=True) +class HTAllReduceTuning: + persistent_ctas: int | None = None + consumer_threads: int = 512 + vectors_per_thread: int = 2 + stages: int = 2 + reduction_warps: int = 2 + reduction_cta_groups: int | None = None + rms_token_groups: int = 2 + rms_pipeline_stages: int = 1 + rms_shard_major: bool = False + enable_pdl: bool = True + + +HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048 = HTFinalizeTuning( + stages=7, + reduction_warps=2, + rms_pipeline_stages=3, + rms_shard_major=True, +) +HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049 = HTFinalizeTuning() +HT_FINALIZE_GB300_TP8_H8192_K10 = HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049 +HT_FINALIZE_GB300_TP16_H8192_K10 = HTFinalizeTuning( + stages=7, + reduction_warps=2, + rms_pipeline_stages=3, + rms_shard_major=True, +) +HT_ALL_REDUCE_GB300_TP8_H8192 = HTAllReduceTuning() +HT_ALL_REDUCE_GB300_TP16_H8192 = HTAllReduceTuning() + + +@dataclass(slots=True) +class HTProtocolState: + local_contributions: SymmetricBuffer + prenorm_mailbox: SymmetricBuffer + routed_ready_counters: SymmetricBuffer + routed_processed_counters: torch.Tensor + all_reduce_ready_counters: SymmetricBuffer + all_reduce_processed_counters: torch.Tensor + + +class _PathKwargs(TypedDict): + hidden_size: int + top_k: int + capacity_m: int + write_residual_output: bool + + +class _HTPath: + def __init__( + self, + *, + compiled: Any, + hidden_size: int, + top_k: int, + capacity_m: int, + write_residual_output: bool, + ) -> None: + self._compiled = compiled + self.hidden_size = hidden_size + self.top_k = top_k + self.capacity_m = capacity_m + self.write_residual_output = write_residual_output + + def _outputs( + self, + m: int, + norm_output: torch.Tensor | None, + residual_output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + shape = (m, self.hidden_size) + device = torch.device("cuda", torch.cuda.current_device()) + if norm_output is None: + norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + if self.write_residual_output and residual_output is None: + residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + return norm_output, residual_output + + def _state_buffers( + self, state: HTProtocolState + ) -> tuple[SymmetricBuffer, SymmetricBuffer]: + return state.local_contributions, state.prenorm_mailbox + + def _validate_m(self, m: int) -> None: + if not 1 <= m <= self.capacity_m: + raise ValueError(f"m must be in [1, {self.capacity_m}]") + + +class FinalizeAllReduceRMSNormHTKernel(_HTPath): + def __call__( + self, + routed_output: torch.Tensor, + expert_weights: torch.Tensor, + permuted_indices: torch.Tensor, + shared_output: torch.Tensor | None, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: HTProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_m(m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + local, prenorm = self._state_buffers(state) + peers = cast(torch.Tensor, state.routed_ready_counters.peer_addresses) + shared_arg = shared_output if shared_output is not None else norm_output + residual_arg = residual_source if residual_source is not None else norm_output + residual_output_arg = ( + residual_output if residual_output is not None else norm_output + ) + self._compiled( + to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size), + to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k), + to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k), + to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(gamma, 16), + to_cute(local.tensor.flatten(), 16), + to_cute(prenorm.tensor.flatten(), 16), + to_cute_dynamic( + residual_output_arg.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size), + to_cute(peers, 8), + to_cute(state.routed_ready_counters.tensor, 4), + to_cute(state.routed_processed_counters.flatten(), 4), + Int64(cast(int, local.multicast_address)), + Int64(cast(int, prenorm.multicast_address)), + Int32(m), + current_cu_stream(), + ) + return norm_output, residual_output + + +class AllReduceRMSNormHTKernel(_HTPath): + def __call__( + self, + local_contribution: torch.Tensor, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: HTProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_m(m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + local, prenorm = self._state_buffers(state) + peers = cast(torch.Tensor, state.all_reduce_ready_counters.peer_addresses) + residual_arg = residual_source if residual_source is not None else norm_output + residual_output_arg = ( + residual_output if residual_output is not None else norm_output + ) + index_arg = state.all_reduce_processed_counters.view(torch.int32) + # top_k=0 disables metadata reads, so the aliased placeholders stay unused. + self._compiled( + to_cute_dynamic( + local_contribution.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(local_contribution.flatten(), 2, divisibility=1), + to_cute_dynamic(index_arg.flatten(), 4, divisibility=1), + to_cute_dynamic( + local_contribution.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(gamma, 16), + to_cute(local.tensor.flatten(), 16), + to_cute(prenorm.tensor.flatten(), 16), + to_cute_dynamic( + residual_output_arg.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size), + to_cute(peers, 8), + to_cute(state.all_reduce_ready_counters.tensor, 4), + to_cute(state.all_reduce_processed_counters.flatten(), 4), + Int64(cast(int, local.multicast_address)), + Int64(cast(int, prenorm.multicast_address)), + Int32(m), + current_cu_stream(), + ) + return norm_output, residual_output + + +class HTProtocol: + """Own tuning-independent HT State and both persistent path variants.""" + + def __init__( + self, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + capacity_m: int, + rms_epsilon: float, + routed_scaling_factor: float, + weight_bias: float, + *, + include_shared_expert: bool, + add_residual: bool, + write_residual_output: bool, + finalize_tunings: tuple[HTFinalizeTuning, ...], + all_reduce_tunings: tuple[HTAllReduceTuning, ...], + group: dist.ProcessGroup, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.capacity_m = capacity_m + self.rms_epsilon = rms_epsilon + self.routed_scaling_factor = routed_scaling_factor + self.weight_bias = weight_bias + self.include_shared_expert = include_shared_expert + self.add_residual = add_residual + self.write_residual_output = write_residual_output + + self.finalize_kernels = { + tuning: FinalizeAllReduceRMSNormHTKernel( + compiled=self._compile_finalize(tuning), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(finalize_tunings) + } + self.all_reduce_kernels = { + tuning: AllReduceRMSNormHTKernel( + compiled=self._compile_all_reduce(tuning), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(all_reduce_tunings) + } + self.state = self._create_state(group) + + def _path_kwargs(self) -> _PathKwargs: + return { + "hidden_size": self.hidden_size, + "top_k": self.top_k, + "capacity_m": self.capacity_m, + "write_residual_output": self.write_residual_output, + } + + def _resolve_ctas(self, persistent_ctas: int | None) -> int: + sm_count = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + # min_blocks_per_mp=1 guarantees one resident CTA per SM for this kernel. + resident_ctas = (sm_count // self.tp_size) * self.tp_size + if resident_ctas == 0: + raise ValueError("tp_size exceeds the available SM count") + if persistent_ctas is None: + return resident_ctas + if persistent_ctas <= 0 or persistent_ctas % self.tp_size: + raise ValueError( + "persistent_ctas must be positive and divisible by tp_size" + ) + return min(persistent_ctas, resident_ctas) + + def _compile_finalize(self, tuning: HTFinalizeTuning): + active_ctas = self._resolve_ctas(tuning.persistent_ctas) + groups = tuning.reduction_cta_groups or active_ctas // self.tp_size + kernel = _MoeFinalizeAllReduceRMSNormHTDeviceKernel( + hidden=self.hidden_size, + top_k=self.top_k, + tp=self.tp_size, + rank=self.rank, + active_ctas=active_ctas, + stages=tuning.stages, + consumer_threads=tuning.consumer_threads, + vectors_per_thread=tuning.vectors_per_thread, + reduction_warps=tuning.reduction_warps, + reduction_cta_groups=groups, + rms_token_groups=tuning.rms_token_groups, + rms_pipeline_stages=tuning.rms_pipeline_stages, + rms_shard_major=tuning.rms_shard_major, + rms_epsilon=self.rms_epsilon, + routed_scaling_factor=self.routed_scaling_factor, + weight_bias=self.weight_bias, + include_shared_expert=self.include_shared_expert, + add_residual=self.add_residual, + write_residual_output=self.write_residual_output, + enable_pdl=tuning.enable_pdl, + ) + return self._compile(kernel, top_k=self.top_k) + + def _compile_all_reduce(self, tuning: HTAllReduceTuning): + active_ctas = self._resolve_ctas(tuning.persistent_ctas) + groups = tuning.reduction_cta_groups or active_ctas // self.tp_size + kernel = _MoeFinalizeAllReduceRMSNormHTDeviceKernel( + hidden=self.hidden_size, + top_k=0, + tp=self.tp_size, + rank=self.rank, + active_ctas=active_ctas, + stages=tuning.stages, + consumer_threads=tuning.consumer_threads, + vectors_per_thread=tuning.vectors_per_thread, + reduction_warps=tuning.reduction_warps, + reduction_cta_groups=groups, + rms_token_groups=tuning.rms_token_groups, + rms_pipeline_stages=tuning.rms_pipeline_stages, + rms_shard_major=tuning.rms_shard_major, + rms_epsilon=self.rms_epsilon, + routed_scaling_factor=1.0, + weight_bias=self.weight_bias, + include_shared_expert=True, + add_residual=self.add_residual, + write_residual_output=self.write_residual_output, + enable_pdl=tuning.enable_pdl, + ) + return self._compile(kernel, top_k=0) + + def _compile( + self, + kernel: _MoeFinalizeAllReduceRMSNormHTDeviceKernel, + *, + top_k: int, + ): + activation = self.capacity_m * self.hidden_size + token_slots = (self.capacity_m + self.tp_size - 1) // self.tp_size + args = ( + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=2, divisibility=max(top_k, 1) + ), + make_fake_dynamic_compact_tensor( + Int32, alignment=4, divisibility=max(top_k, 1) + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16), + make_fake_compact_tensor(BFloat16, (activation,), assumed_align=16), + make_fake_compact_tensor(BFloat16, (activation,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(Int64, (self.tp_size,), assumed_align=8), + make_fake_compact_tensor(Uint32, (token_slots,), assumed_align=4), + make_fake_compact_tensor( + Uint32, (token_slots * self.tp_size,), assumed_align=4 + ), + Int64(0), + Int64(0), + Int32(self.capacity_m), + current_cu_stream(), + ) + return cute.compile(kernel, *args) + + def _allocate_large_buffers( + self, group: dist.ProcessGroup + ) -> tuple[SymmetricBuffer, SymmetricBuffer]: + shape = (self.capacity_m, self.hidden_size) + device = torch.device("cuda", torch.cuda.current_device()) + return ( + SymmetricBuffer.allocate( + shape, + torch.bfloat16, + device, + group, + require_multicast=True, + ), + SymmetricBuffer.allocate( + shape, + torch.bfloat16, + device, + group, + require_multicast=True, + ), + ) + + def _create_state(self, group: dist.ProcessGroup) -> HTProtocolState: + device = torch.device("cuda", torch.cuda.current_device()) + token_slots = (self.capacity_m + self.tp_size - 1) // self.tp_size + + def counters() -> tuple[SymmetricBuffer, torch.Tensor]: + ready = SymmetricBuffer.allocate( + (token_slots,), + torch.uint32, + device, + group, + materialize_peer_addresses=True, + ) + ready.tensor.zero_() + processed = torch.zeros( + (token_slots, self.tp_size), dtype=torch.uint32, device=device + ) + return ready, processed + + routed_ready, routed_processed = counters() + all_reduce_ready, all_reduce_processed = counters() + local, prenorm = self._allocate_large_buffers(group) + return HTProtocolState( + local_contributions=local, + prenorm_mailbox=prenorm, + routed_ready_counters=routed_ready, + routed_processed_counters=routed_processed, + all_reduce_ready_counters=all_reduce_ready, + all_reduce_processed_counters=all_reduce_processed, + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/__init__.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/__init__.py new file mode 100644 index 000000000000..18d69a296c65 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/__init__.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Low-latency MNNVL protocol.""" + +from .protocol import ( + LL_ALL_REDUCE_GB300_TP16_H8192, + LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17, + LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18, + LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10, + LL_ALL_REDUCE_GB300_TP8_H8192, + LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5, + LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4, + LL_FINALIZE_GB300_TP16_H8192_K10, + LL_FINALIZE_GB300_TP8_H8192_K10, + LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20, + LLAllReduceTuning, + LLCollectiveTuning, + LLFinalizeTuning, +) + +__all__ = [ + "LL_ALL_REDUCE_GB300_TP8_H8192", + "LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5", + "LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4", + "LL_ALL_REDUCE_GB300_TP16_H8192", + "LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17", + "LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18", + "LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10", + "LL_FINALIZE_GB300_TP8_H8192_K10", + "LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20", + "LL_FINALIZE_GB300_TP16_H8192_K10", + "LLAllReduceTuning", + "LLCollectiveTuning", + "LLFinalizeTuning", +] diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py new file mode 100644 index 000000000000..92b6b1755af6 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py @@ -0,0 +1,1003 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Split low-latency BF16 MoE finalize, TP reduction, and RMSNorm.""" + +from __future__ import annotations + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Int64, Uint32 + +from ..cute_dsl_primitives import ( + NEGATIVE_ZERO_BF16_BITS, + QUAD_BF16, + VEC_BF16, + WARP_SIZE, + bf16x4_to_packed_u32x2, + bf16x8_to_packed_u32x4, + f32_to_bf16_bits, + fragment_has_negative_zero, + load_global_bf16_as_f32, + load_global_bf16_as_f32_predicated, + load_global_u32x2, + load_global_u32x4, + load_volatile_u32, + map_shared_to_peer, + packed_u32x2_to_bf16x4, + packed_u32x4_to_bf16x8, + sanitize_negative_zero_u32x2, + sanitize_negative_zero_u32x4, + shuffle_sync_idx_u32, + stmc_bf16x2, + stmc_bf16x4, + stmc_bf16x8, + store_global_u32, + store_global_u32x4, + store_lamport_sentinel_u32x4, + store_shared_cluster_f32, +) + +LAMPORT_GENERATIONS = 3 +NEXT_STAGE = 0 +ACTIVE_STAGE = 1 + + +@cute.jit +def _group_leader_block_sum( + value: Float32, + warp_sums: cute.Tensor, + warps: cutlass.Constexpr[int], + leader_stride: cutlass.Constexpr[int], +) -> Float32: + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + for offset in cutlass.range_constexpr(1, WARP_SIZE): + if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + if lane == 0: + cute.arch.store((warp_sums + warp).llvm_ptr, value) + cute.arch.barrier() + + result = Float32(0.0) + if warp == 0: + if lane < Int32(warps): + result = cute.arch.load( + (warp_sums + lane).llvm_ptr, + Float32, + ) + result = cute.arch.warp_reduction_sum(result) + if lane == 0: + cute.arch.store(warp_sums.llvm_ptr, result) + cute.arch.barrier() + return cute.arch.load(warp_sums.llvm_ptr, Float32) + + +class _ScalarFinalizePublishDeviceKernel: + def __init__( + self, + *, + hidden: int, + top_k: int, + tp: int, + rank: int, + capacity_m: int, + threads: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + ) -> None: + if hidden <= 0 or hidden % 2: + raise ValueError("hidden must be a positive multiple of 2") + self.hidden = hidden + self.top_k = top_k + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.threads = threads + self.routed_scaling_factor = routed_scaling_factor + self.ctas_per_token = math.ceil(hidden / threads) + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + + def smem_size_in_bytes(self) -> int: + return self.top_k * 8 + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_multicast_address, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + cta_in_token = block % self.ctas_per_token + hidden_index = cta_in_token * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + element = Int64(token) * self.top_k + Int64(metadata_index) + row = cute.arch.load( + (permuted_indices.iterator + element).llvm_ptr, + Int32, + ) + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + element).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if row == Int32(-1): + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + row, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + shared_value = Float32(0.0) + if cutlass.const_expr( + self.include_shared_expert and self.load_shared_expert_before_pdl + ): + if hidden_index < Int32(self.hidden): + shared_element = Int64(token) * self.hidden + Int64(hidden_index) + shared_value = load_global_bf16_as_f32( + Int64((shared_output.iterator + shared_element).toint()) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if hidden_index < Int32(self.hidden): + shared_element = Int64(token) * self.hidden + Int64(hidden_index) + shared_value = load_global_bf16_as_f32( + Int64((shared_output.iterator + shared_element).toint()) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + bits = Uint32(0) + if hidden_index < Int32(self.hidden): + accumulator = Float32(0.0) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + row = cute.arch.load((staged_indices + k).llvm_ptr, Int32) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + source_element = Int64(row) * self.hidden + Int64(hidden_index) + accumulator = ( + accumulator + + load_global_bf16_as_f32_predicated( + Int64((routed_output.iterator + source_element).toint()), + Int32(row != Int32(-1)), + ) + * weight + ) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group,)), + Float32, + ) + inputs.fill(Float32(0.0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + row = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + source_element = Int64(row) * self.hidden + Int64( + hidden_index + ) + inputs[item] = load_global_bf16_as_f32_predicated( + Int64( + (routed_output.iterator + source_element).toint() + ), + Int32(row != Int32(-1)), + ) + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + accumulator = accumulator + inputs[item] * cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + if cutlass.const_expr(self.include_shared_expert): + accumulator = accumulator + shared_value + bits = f32_to_bf16_bits(accumulator) + if bits == Uint32(NEGATIVE_ZERO_BF16_BITS): + bits = Uint32(0) + + # Unlike the quad path, this publishes first; Lamport sentinels gate consumers. + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + lane = cute.arch.lane_idx() + partner_bits = shuffle_sync_idx_u32( + bits, + Int32(lane | Int32(1)), + ) + if (lane & Int32(1)) == Int32(0) and hidden_index < Int32(self.hidden): + packed = bits | (partner_bits << Uint32(16)) + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token) + ) * self.hidden + Int64(hidden_index) + stmc_bf16x2( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed, + ) + + +class _QuadFinalizePublishDeviceKernel: + def __init__( + self, + *, + hidden: int, + top_k: int, + tp: int, + rank: int, + capacity_m: int, + threads: int, + routed_scaling_factor: float, + include_shared_expert: bool, + load_shared_expert_before_pdl: bool, + enable_pdl: bool, + prefetch_group: int, + ) -> None: + if hidden <= 0 or hidden % QUAD_BF16: + raise ValueError("hidden must be a positive multiple of 4") + self.hidden = hidden + self.top_k = top_k + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.threads = threads + self.routed_scaling_factor = routed_scaling_factor + self.fragments = hidden // QUAD_BF16 + self.ctas_per_token = math.ceil(self.fragments / threads) + self.include_shared_expert = include_shared_expert + self.load_shared_expert_before_pdl = load_shared_expert_before_pdl + self.enable_pdl = enable_pdl + self.prefetch_group = prefetch_group + self.prefetch_groups = (top_k + prefetch_group - 1) // prefetch_group + + def smem_size_in_bytes(self) -> int: + return self.top_k * 8 + + @cute.jit + def __call__( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + stage_state, + contribution_mailbox_multicast_address, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + routed_output: cute.Tensor, + expert_weights: cute.Tensor, + permuted_indices: cute.Tensor, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + cta_in_token = block % self.ctas_per_token + fragment = cta_in_token * self.threads + tidx + + smem = cutlass.utils.SmemAllocator() + staged_indices = smem.allocate_array(Int32, self.top_k) + staged_weights = smem.allocate_array(Float32, self.top_k) + metadata_index = Int32(tidx) + while metadata_index < Int32(self.top_k): + element = Int64(token) * self.top_k + Int64(metadata_index) + row = cute.arch.load( + (permuted_indices.iterator + element).llvm_ptr, + Int32, + ) + weight = load_global_bf16_as_f32( + Int64((expert_weights.iterator + element).toint()) + ) + if cutlass.const_expr(self.routed_scaling_factor != 1.0): + weight = weight * Float32(self.routed_scaling_factor) + if row == Int32(-1): + weight = Float32(0.0) + cute.arch.store( + (staged_indices + metadata_index).llvm_ptr, + row, + ) + cute.arch.store( + (staged_weights + metadata_index).llvm_ptr, + weight, + ) + metadata_index = metadata_index + self.threads + cute.arch.barrier() + + if cutlass.const_expr(self.include_shared_expert): + shared_values = cute.make_rmem_tensor( + cute.make_layout((QUAD_BF16,)), BFloat16 + ) + shared_values.fill(BFloat16(0.0)) + if cutlass.const_expr(self.load_shared_expert_before_pdl): + if fragment < self.fragments: + shared_element = ( + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr( + self.include_shared_expert and not self.load_shared_expert_before_pdl + ): + if fragment < self.fragments: + shared_element = ( + Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + shared_pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + shared_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + shared_values.store( + packed_u32x2_to_bf16x4(load_global_u32x2(shared_pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if fragment < self.fragments: + accumulator = cute.make_rmem_tensor(cute.make_layout((QUAD_BF16,)), Float32) + accumulator.fill(Float32(0.0)) + if cutlass.const_expr(self.prefetch_group == 1): + for k in cutlass.range_constexpr(self.top_k): + row = cute.arch.load((staged_indices + k).llvm_ptr, Int32) + weight = cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + if row != Int32(-1): + source_element = ( + Int64(row) * self.hidden + Int64(fragment) * QUAD_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + accumulator.store( + accumulator.load() + + packed_u32x2_to_bf16x4( + load_global_u32x2(source_pointer) + ).to(Float32) + * weight + ) + else: + inputs = cute.make_rmem_tensor( + cute.make_layout((self.prefetch_group, 2)), + Uint32, + ) + inputs.fill(Uint32(0)) + for group in cutlass.range_constexpr(self.prefetch_groups): + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + row = cute.arch.load( + (staged_indices + k).llvm_ptr, + Int32, + ) + for word in cutlass.range_constexpr(2): + inputs[item, word] = Uint32(0) + if row != Int32(-1): + source_element = ( + Int64(row) * self.hidden + + Int64(fragment) * QUAD_BF16 + ) + source_pointer = cute.make_ptr( + BFloat16, + (routed_output.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + source = load_global_u32x2(source_pointer) + for word in cutlass.range_constexpr(2): + inputs[item, word] = source[word] + for item in cutlass.range_constexpr(self.prefetch_group): + k = group * self.prefetch_group + item + if cutlass.const_expr(k < self.top_k): + source = cute.make_rmem_tensor( + cute.make_layout((2,)), + Uint32, + ) + for word in cutlass.range_constexpr(2): + source[word] = inputs[item, word] + accumulator.store( + accumulator.load() + + packed_u32x2_to_bf16x4(source.load()).to(Float32) + * cute.arch.load( + (staged_weights + k).llvm_ptr, + Float32, + ) + ) + result = accumulator.load() + if cutlass.const_expr(self.include_shared_expert): + result = result + shared_values.load().to(Float32) + result_packed = bf16x4_to_packed_u32x2(result.to(BFloat16)) + packed = sanitize_negative_zero_u32x2(result_packed) + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + Int64(token) + ) * self.hidden + Int64(fragment) * QUAD_BF16 + stmc_bf16x4( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed, + ) + + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _SharedOnlyPublishDeviceKernel: + def __init__( + self, + *, + hidden: int, + tp: int, + rank: int, + capacity_m: int, + elements_per_thread: int, + threads: int, + release_before_store: bool, + enable_pdl: bool, + ) -> None: + if elements_per_thread not in (1, QUAD_BF16, VEC_BF16): + raise ValueError("elements_per_thread must be 1, 4, or 8") + if hidden <= 0 or hidden % elements_per_thread: + raise ValueError("hidden must divide evenly across thread fragments") + self.hidden = hidden + self.tp = tp + self.rank = rank + self.capacity_m = capacity_m + self.elements_per_thread = elements_per_thread + self.threads = threads + self.fragments = hidden // elements_per_thread + self.ctas_per_token = math.ceil(self.fragments / threads) + self.release_before_store = release_before_store + self.enable_pdl = enable_pdl + + @cute.jit + def __call__( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + shared_output, + stage_state, + contribution_mailbox_multicast_address, + ).launch( + grid=(m * self.ctas_per_token, 1, 1), + block=(self.threads, 1, 1), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + shared_output: cute.Tensor, + stage_state: cute.Tensor, + contribution_mailbox_multicast_address: Int64, + ) -> None: + block, _, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + token = block // self.ctas_per_token + cta_in_token = block % self.ctas_per_token + fragment = cta_in_token * self.threads + tidx + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + if cutlass.const_expr(self.elements_per_thread == 1): + bits = Uint32(0) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) + bits = f32_to_bf16_bits( + load_global_bf16_as_f32( + Int64((shared_output.iterator + element).toint()) + ) + ) + if bits == Uint32(NEGATIVE_ZERO_BF16_BITS): + bits = Uint32(0) + partner_bits = shuffle_sync_idx_u32( + bits, + Int32(cute.arch.lane_idx() | Int32(1)), + ) + packed_word = bits | (partner_bits << Uint32(16)) + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16): + packed_words = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + packed_words.fill(Uint32(0)) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) * QUAD_BF16 + pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=8, + ) + packed_words.store( + sanitize_negative_zero_u32x2(load_global_u32x2(pointer)) + ) + else: + packed_words = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + packed_words.fill(Uint32(0)) + if fragment < self.fragments: + element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + pointer = cute.make_ptr( + BFloat16, + (shared_output.iterator + element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed_words.store( + sanitize_negative_zero_u32x4(load_global_u32x4(pointer)) + ) + + stage = load_volatile_u32(stage_state.iterator + NEXT_STAGE) + if cutlass.const_expr(self.release_before_store): + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.elements_per_thread == 1): + if (cute.arch.lane_idx() & Int32(1)) == Int32( + 0 + ) and fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) + stmc_bf16x2( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_word, + ) + elif cutlass.const_expr(self.elements_per_thread == QUAD_BF16): + if fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * QUAD_BF16 + stmc_bf16x4( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_words, + ) + else: + if fragment < self.fragments: + mailbox_element = ( + (Int64(stage) * self.tp + self.rank) * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + stmc_bf16x8( + contribution_mailbox_multicast_address + mailbox_element * 2, + packed_words, + ) + + if cutlass.const_expr(not self.release_before_store): + if block == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + ACTIVE_STAGE, + stage, + ) + cute.arch.barrier() + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +class _LamportResidualRMSNormDeviceKernel: + def __init__( + self, + *, + hidden: int, + tp: int, + capacity_m: int, + cluster_size: int, + rank_lanes: int, + threads: int, + rms_epsilon: float, + weight_bias: float, + add_residual: bool, + write_residual_output: bool, + enable_pdl: bool, + ) -> None: + if rank_lanes not in (1, 2, 4, 8): + raise ValueError("rank_lanes must be 1, 2, 4, or 8") + if tp % rank_lanes: + raise ValueError("tp must be divisible by rank_lanes") + if threads <= 0 or threads % WARP_SIZE or threads % rank_lanes: + raise ValueError("threads must be a positive warp and rank-lane multiple") + if hidden <= 0 or hidden % VEC_BF16: + raise ValueError("hidden must be a positive multiple of 8") + self.hidden = hidden + self.tp = tp + self.capacity_m = capacity_m + self.cluster_size = cluster_size + self.rank_lanes = rank_lanes + self.threads = threads + self.rms_epsilon = rms_epsilon + self.weight_bias = weight_bias + self.add_residual = add_residual + self.write_residual_output = write_residual_output + self.enable_pdl = enable_pdl + self.fragments = hidden // VEC_BF16 + self.groups_per_cta = threads // rank_lanes + self.fragment_stride = cluster_size * self.groups_per_cta + self.trips = math.ceil(self.fragments / self.fragment_stride) + self.warps = threads // WARP_SIZE + self.rank_waves = tp // rank_lanes + + def smem_size_in_bytes(self) -> int: + return (self.warps + self.cluster_size) * 4 + + @cute.jit + def __call__( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + stage_state: cute.Tensor, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + contribution_mailbox, + residual_source, + gamma, + residual_output, + norm_output, + stage_state, + ).launch( + grid=(m, self.cluster_size, 1), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_size, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + norm_output: cute.Tensor, + stage_state: cute.Tensor, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + token, _, _ = cute.arch.block_idx() + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + rank_lane = tidx % self.rank_lanes + group = tidx // self.rank_lanes + base_fragment = cluster_rank * self.groups_per_cta + group + + prenorm_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + prenorm_fragments.fill(BFloat16(0.0)) + gamma_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + gamma_fragments.fill(BFloat16(0.0)) + if cutlass.const_expr(self.add_residual): + residual_fragments = cute.make_rmem_tensor( + cute.make_layout( + (self.trips, VEC_BF16), + stride=(VEC_BF16, 1), + ), + BFloat16, + ) + residual_fragments.fill(BFloat16(0.0)) + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + gamma_element = Int64(fragment) * VEC_BF16 + gamma_pointer = cute.make_ptr( + BFloat16, + (gamma.iterator + gamma_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer)) + ) + if cutlass.const_expr(self.add_residual): + residual_element = Int64(token) * self.hidden + gamma_element + residual_pointer = cute.make_ptr( + BFloat16, + (residual_source.iterator + residual_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + residual_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + lane_packed = cute.make_rmem_tensor( + cute.make_layout((self.rank_waves, 4)), + Uint32, + ) + lane_packed.fill(Uint32(0)) + dirty = fragment < self.fragments + while dirty: + dirty = False + for wave in cutlass.range_constexpr(self.rank_waves): + source_rank = wave * self.rank_lanes + rank_lane + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4( + source_pointer, + volatile=True, + ) + dirty = dirty | fragment_has_negative_zero(packed) + for word in cutlass.range_constexpr(4): + lane_packed[wave, word] = packed[word] + + lane_sum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + lane_sum.fill(Float32(0.0)) + for wave in cutlass.range_constexpr(self.rank_waves): + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = lane_packed[wave, word] + lane_sum.store( + lane_sum.load() + packed_u32x4_to_bf16x8(packed.load()).to(Float32) + ) + for offset in cutlass.range_constexpr(1, 5): + if cutlass.const_expr(offset < self.rank_lanes and offset in (1, 2, 4)): + for element in cutlass.range_constexpr(VEC_BF16): + lane_sum[element] = lane_sum[ + element + ] + cute.arch.shuffle_sync_bfly( + lane_sum[element], + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + + if fragment < self.fragments and rank_lane == 0: + prenorm = lane_sum.load() + if cutlass.const_expr(self.add_residual): + prenorm = prenorm + residual_fragments[trip, None].load().to( + Float32 + ) + prenorm_bf16 = prenorm.to(BFloat16) + prenorm_fragments[trip, None].store(prenorm_bf16) + if cutlass.const_expr(self.write_residual_output): + output_element = ( + Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + ) + store_global_u32x4( + Int64((residual_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(prenorm_bf16), + ) + + if token == 0 and cluster_rank == 0 and tidx == 0: + store_global_u32( + stage_state.iterator + NEXT_STAGE, + (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + for wave in cutlass.range_constexpr(self.rank_waves): + source_rank = wave * self.rank_lanes + rank_lane + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((contribution_mailbox.iterator + source_element).toint()) + ) + + thread_sum = Float32(0.0) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + values = prenorm_fragments[trip, None].load().to(Float32) + thread_sum = thread_sum + (values * values).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_array(Float32, self.warps) + cluster_sums = smem.allocate_array(Float32, self.cluster_size) + cta_sum = _group_leader_block_sum( + thread_sum, + warp_sums, + self.warps, + self.rank_lanes, + ) + if tidx < self.cluster_size: + local_slot = cluster_sums + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, cta_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_size): + full_sum = full_sum + cute.arch.load( + (cluster_sums + peer).llvm_ptr, + Float32, + ) + inv_rms = cute.math.rsqrt( + full_sum / Float32(self.hidden) + Float32(self.rms_epsilon), + fastmath=True, + ) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments and rank_lane == 0: + gamma_values = gamma_fragments[trip, None].load().to(Float32) + if cutlass.const_expr(self.weight_bias != 0.0): + gamma_values = gamma_values + Float32(self.weight_bias) + result = ( + prenorm_fragments[trip, None].load().to(Float32) + * inv_rms + * gamma_values + ).to(BFloat16) + output_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + store_global_u32x4( + Int64((norm_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(result), + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/protocol.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/protocol.py new file mode 100644 index 000000000000..3a26bb485079 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/kernel_ll/protocol.py @@ -0,0 +1,461 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Low-latency MNNVL protocol and its two operation paths.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypedDict, cast + +import cutlass.cute as cute +import torch +import torch.distributed as dist +from cutlass import BFloat16, Int32, Int64 +from cutlass.cute.runtime import make_fake_compact_tensor + +from ..cute_dsl_primitives import QUAD_BF16 +from ..runtime import ( + current_cu_stream, + make_fake_dynamic_compact_tensor, + to_cute, + to_cute_dynamic, +) +from ..symmetric_buffer import SymmetricBuffer +from .device_kernels import ( + LAMPORT_GENERATIONS, + _LamportResidualRMSNormDeviceKernel, + _QuadFinalizePublishDeviceKernel, + _ScalarFinalizePublishDeviceKernel, + _SharedOnlyPublishDeviceKernel, +) + + +@dataclass(frozen=True, slots=True) +class LLCollectiveTuning: + cluster_size: int = 8 + rank_lanes: int = 1 + threads: int = 128 + enable_pdl: bool = True + + +@dataclass(frozen=True, slots=True) +class LLFinalizeTuning: + elements_per_thread: int = 4 + threads: int = 128 + prefetch_group: int = 10 + load_shared_expert_before_pdl: bool = False + collective: LLCollectiveTuning = LLCollectiveTuning() + + +@dataclass(frozen=True, slots=True) +class LLAllReduceTuning: + publish_elements_per_thread: int = 8 + publish_threads: int = 128 + publish_release_before_store: bool = False + collective: LLCollectiveTuning = LLCollectiveTuning() + + +LL_FINALIZE_GB300_TP8_H8192_K10 = LLFinalizeTuning() +LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20 = LLFinalizeTuning( + collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2) +) +LL_FINALIZE_GB300_TP16_H8192_K10 = LLFinalizeTuning( + collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2) +) +LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4 = LLAllReduceTuning( + collective=LLCollectiveTuning(cluster_size=16, threads=64) +) +LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5 = LLAllReduceTuning() +LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10 = LLAllReduceTuning( + collective=LLCollectiveTuning(cluster_size=16, threads=64) +) +LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17 = LLAllReduceTuning() +LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18 = LLAllReduceTuning( + collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2) +) +LL_ALL_REDUCE_GB300_TP8_H8192 = LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5 +LL_ALL_REDUCE_GB300_TP16_H8192 = LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10 + + +@dataclass(slots=True) +class LLProtocolState: + contribution_mailbox: SymmetricBuffer + stage_state: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _CompiledFinalize: + publish: Any + collective: Any + + +@dataclass(frozen=True, slots=True) +class _CompiledAllReduce: + publish: Any + collective: Any + + +class _PathKwargs(TypedDict): + hidden_size: int + top_k: int + capacity_m: int + write_residual_output: bool + + +class _LLPath: + def __init__( + self, + *, + hidden_size: int, + top_k: int, + capacity_m: int, + write_residual_output: bool, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.capacity_m = capacity_m + self.write_residual_output = write_residual_output + + def _outputs( + self, + m: int, + norm_output: torch.Tensor | None, + residual_output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + shape = (m, self.hidden_size) + device = torch.device("cuda", torch.cuda.current_device()) + if norm_output is None: + norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + if self.write_residual_output and residual_output is None: + residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device) + return norm_output, residual_output + + def _validate_state(self, state: LLProtocolState, m: int) -> None: + if not 1 <= m <= self.capacity_m: + raise ValueError(f"m must be in [1, {self.capacity_m}]") + address = state.contribution_mailbox.multicast_address + if address is None or address % 16: + raise ValueError( + "LL contribution mailbox requires a 16-byte-aligned multicast address" + ) + + def _launch_collective( + self, + collective, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + state: LLProtocolState, + norm_output: torch.Tensor, + residual_output: torch.Tensor | None, + m: int, + ) -> None: + residual_arg = residual_source if residual_source is not None else norm_output + residual_output_arg = ( + residual_output if residual_output is not None else norm_output + ) + collective( + to_cute(state.contribution_mailbox.tensor.flatten(), 16), + to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(gamma, 16), + to_cute_dynamic( + residual_output_arg.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size), + to_cute(state.stage_state, 4), + Int32(m), + current_cu_stream(), + ) + + +class FinalizeAllReduceRMSNormLLKernel(_LLPath): + def __init__(self, *, compiled: _CompiledFinalize, **kwargs) -> None: + super().__init__(**kwargs) + self._compiled = compiled + + def __call__( + self, + routed_output: torch.Tensor, + expert_weights: torch.Tensor, + permuted_indices: torch.Tensor, + shared_output: torch.Tensor | None, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: LLProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_state(state, m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + shared_arg = shared_output if shared_output is not None else norm_output + self._compiled.publish( + to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size), + to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k), + to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k), + to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size), + to_cute(state.stage_state, 4), + Int64(cast(int, state.contribution_mailbox.multicast_address)), + Int32(m), + current_cu_stream(), + ) + self._launch_collective( + self._compiled.collective, + residual_source, + gamma, + state, + norm_output, + residual_output, + m, + ) + return norm_output, residual_output + + +class AllReduceRMSNormLLKernel(_LLPath): + def __init__(self, *, compiled: _CompiledAllReduce, **kwargs) -> None: + super().__init__(**kwargs) + self._compiled = compiled + + def __call__( + self, + local_contribution: torch.Tensor, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + state: LLProtocolState, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + self._validate_state(state, m) + norm_output, residual_output = self._outputs(m, norm_output, residual_output) + self._compiled.publish( + to_cute_dynamic( + local_contribution.flatten(), + 16, + divisibility=self.hidden_size, + ), + to_cute(state.stage_state, 4), + Int64(cast(int, state.contribution_mailbox.multicast_address)), + Int32(m), + current_cu_stream(), + ) + self._launch_collective( + self._compiled.collective, + residual_source, + gamma, + state, + norm_output, + residual_output, + m, + ) + return norm_output, residual_output + + +class LLProtocol: + """Own LL State and protocol-local compiled variants for both paths.""" + + def __init__( + self, + hidden_size: int, + top_k: int, + tp_size: int, + rank: int, + capacity_m: int, + rms_epsilon: float, + routed_scaling_factor: float, + weight_bias: float, + *, + include_shared_expert: bool, + add_residual: bool, + write_residual_output: bool, + finalize_tunings: tuple[LLFinalizeTuning, ...], + all_reduce_tunings: tuple[LLAllReduceTuning, ...], + group: dist.ProcessGroup, + ) -> None: + self.hidden_size = hidden_size + self.top_k = top_k + self.tp_size = tp_size + self.rank = rank + self.capacity_m = capacity_m + self.rms_epsilon = rms_epsilon + self.routed_scaling_factor = routed_scaling_factor + self.weight_bias = weight_bias + self.include_shared_expert = include_shared_expert + self.add_residual = add_residual + self.write_residual_output = write_residual_output + + collective_cache = { + tuning: self._compile_collective(tuning) + for tuning in { + *(item.collective for item in finalize_tunings), + *(item.collective for item in all_reduce_tunings), + } + } + self.finalize_kernels = { + tuning: FinalizeAllReduceRMSNormLLKernel( + compiled=_CompiledFinalize( + publish=self._compile_finalize(tuning), + collective=collective_cache[tuning.collective], + ), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(finalize_tunings) + } + self.all_reduce_kernels = { + tuning: AllReduceRMSNormLLKernel( + compiled=_CompiledAllReduce( + publish=self._compile_all_reduce_publish(tuning), + collective=collective_cache[tuning.collective], + ), + **self._path_kwargs(), + ) + for tuning in dict.fromkeys(all_reduce_tunings) + } + self.state = self._create_state(group) + + def _path_kwargs(self) -> _PathKwargs: + return { + "hidden_size": self.hidden_size, + "top_k": self.top_k, + "capacity_m": self.capacity_m, + "write_residual_output": self.write_residual_output, + } + + def _compile_finalize(self, tuning: LLFinalizeTuning): + if tuning.elements_per_thread not in (1, QUAD_BF16): + raise ValueError("LL finalize elements_per_thread must be 1 or 4") + kwargs: dict[str, Any] = { + "hidden": self.hidden_size, + "top_k": self.top_k, + "tp": self.tp_size, + "rank": self.rank, + "capacity_m": self.capacity_m, + "threads": tuning.threads, + "routed_scaling_factor": self.routed_scaling_factor, + "include_shared_expert": self.include_shared_expert, + "load_shared_expert_before_pdl": tuning.load_shared_expert_before_pdl, + "enable_pdl": tuning.collective.enable_pdl, + "prefetch_group": tuning.prefetch_group, + } + device_kernel = ( + _ScalarFinalizePublishDeviceKernel(**kwargs) + if tuning.elements_per_thread == 1 + else _QuadFinalizePublishDeviceKernel(**kwargs) + ) + args = ( + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=2, divisibility=self.top_k + ), + make_fake_dynamic_compact_tensor( + Int32, alignment=4, divisibility=self.top_k + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(self.capacity_m), + current_cu_stream(), + ) + return cute.compile(device_kernel, *args) + + def _compile_all_reduce_publish(self, tuning: LLAllReduceTuning): + device_kernel = _SharedOnlyPublishDeviceKernel( + hidden=self.hidden_size, + tp=self.tp_size, + rank=self.rank, + capacity_m=self.capacity_m, + elements_per_thread=tuning.publish_elements_per_thread, + threads=tuning.publish_threads, + release_before_store=tuning.publish_release_before_store, + enable_pdl=tuning.collective.enable_pdl, + ) + args = ( + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(self.capacity_m), + current_cu_stream(), + ) + return cute.compile(device_kernel, *args) + + def _compile_collective(self, tuning: LLCollectiveTuning): + device_kernel = _LamportResidualRMSNormDeviceKernel( + hidden=self.hidden_size, + tp=self.tp_size, + capacity_m=self.capacity_m, + cluster_size=tuning.cluster_size, + rank_lanes=tuning.rank_lanes, + threads=tuning.threads, + rms_epsilon=self.rms_epsilon, + weight_bias=self.weight_bias, + add_residual=self.add_residual, + write_residual_output=self.write_residual_output, + enable_pdl=tuning.enable_pdl, + ) + activation = self.capacity_m * self.hidden_size + args = ( + make_fake_compact_tensor( + BFloat16, + (LAMPORT_GENERATIONS * self.tp_size * activation,), + assumed_align=16, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=self.hidden_size + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int32(self.capacity_m), + current_cu_stream(), + ) + return cute.compile(device_kernel, *args) + + def _create_state(self, group: dist.ProcessGroup) -> LLProtocolState: + if dist.get_world_size(group) != self.tp_size: + raise ValueError("ProcessGroup size does not match tp_size") + if dist.get_rank(group) != self.rank: + raise ValueError("ProcessGroup rank does not match rank") + device = torch.device("cuda", torch.cuda.current_device()) + mailbox = SymmetricBuffer.allocate( + ( + LAMPORT_GENERATIONS, + self.tp_size, + self.capacity_m, + self.hidden_size, + ), + torch.bfloat16, + device, + group, + require_multicast=True, + ) + mailbox.tensor.view(torch.int16).fill_(-32768) + return LLProtocolState( + contribution_mailbox=mailbox, + stage_state=torch.zeros((2,), dtype=torch.int32, device=device), + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/presets.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/presets.py new file mode 100644 index 000000000000..a389d39e280c --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/presets.py @@ -0,0 +1,338 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Built-in routing configurations for the MNNVL CuTe DSL backend.""" + +import torch + +from .config import ( + KernelTarget, + MNNVLCuteDSLConfig, + MRangeDispatch, + ProtocolKind, + StaticProfile, +) +from .kernel_bt import ( + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1, +) +from .kernel_ll import ( + LL_ALL_REDUCE_GB300_TP16_H8192, + LL_ALL_REDUCE_GB300_TP8_H8192, + LL_FINALIZE_GB300_TP16_H8192_K10, + LL_FINALIZE_GB300_TP8_H8192_K10, +) +from .kernel_ht import ( + HT_ALL_REDUCE_GB300_TP16_H8192, + HT_ALL_REDUCE_GB300_TP8_H8192, + HT_FINALIZE_GB300_TP16_H8192_K10, + HT_FINALIZE_GB300_TP8_H8192_K10, +) + +__all__ = [ + "BT_ONLY_CONFIG", + "DEFAULT_CONFIG", + "HT_ONLY_CONFIG", + "LL_ONLY_CONFIG", +] + + +def _target(protocol: ProtocolKind, preset: object) -> KernelTarget[object]: + return KernelTarget(protocol=protocol, preset=preset) + + +LL_ONLY_CONFIG = MNNVLCuteDSLConfig( + profiles=( + StaticProfile( + tp_size=8, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.LL, + LL_FINALIZE_GB300_TP8_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.LL, + LL_ALL_REDUCE_GB300_TP8_H8192, + ), + ), + ), + ), + StaticProfile( + tp_size=16, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.LL, + LL_FINALIZE_GB300_TP16_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.LL, + LL_ALL_REDUCE_GB300_TP16_H8192, + ), + ), + ), + ), + ) +) + + +BT_ONLY_CONFIG = MNNVLCuteDSLConfig( + profiles=( + StaticProfile( + tp_size=8, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(48, 1024), + targets=( + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(256, 1024), + targets=( + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1, + ), + ), + ), + ), + StaticProfile( + tp_size=16, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(52, 1024), + targets=( + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(512, 1024), + targets=( + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1, + ), + ), + ), + ), + ) +) + + +HT_ONLY_CONFIG = MNNVLCuteDSLConfig( + profiles=( + StaticProfile( + tp_size=8, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.HT, + HT_FINALIZE_GB300_TP8_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.HT, + HT_ALL_REDUCE_GB300_TP8_H8192, + ), + ), + ), + ), + StaticProfile( + tp_size=16, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.HT, + HT_FINALIZE_GB300_TP16_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(None,), + targets=( + _target( + ProtocolKind.HT, + HT_ALL_REDUCE_GB300_TP16_H8192, + ), + ), + ), + ), + ) +) + + +DEFAULT_CONFIG = MNNVLCuteDSLConfig( + profiles=( + StaticProfile( + tp_size=8, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(23, 48, 703, None), + targets=( + _target( + ProtocolKind.LL, + LL_FINALIZE_GB300_TP8_H8192_K10, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1, + ), + _target( + ProtocolKind.HT, + HT_FINALIZE_GB300_TP8_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(15, 256, 1024, None), + targets=( + _target( + ProtocolKind.LL, + LL_ALL_REDUCE_GB300_TP8_H8192, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1, + ), + _target( + ProtocolKind.HT, + HT_ALL_REDUCE_GB300_TP8_H8192, + ), + ), + ), + ), + StaticProfile( + tp_size=16, + hidden_size=8192, + top_k=10, + dtype=torch.bfloat16, + finalize_routes=MRangeDispatch( + upper_bounds=(7, 52, 703, None), + targets=( + _target( + ProtocolKind.LL, + LL_FINALIZE_GB300_TP16_H8192_K10, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1, + ), + _target( + ProtocolKind.HT, + HT_FINALIZE_GB300_TP16_H8192_K10, + ), + ), + ), + all_reduce_routes=MRangeDispatch( + upper_bounds=(5, 512, 959, None), + targets=( + _target( + ProtocolKind.LL, + LL_ALL_REDUCE_GB300_TP16_H8192, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0, + ), + _target( + ProtocolKind.BT, + BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1, + ), + _target( + ProtocolKind.HT, + HT_ALL_REDUCE_GB300_TP16_H8192, + ), + ), + ), + ), + ) +) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/runtime.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/runtime.py new file mode 100644 index 000000000000..a706f82cea2a --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/runtime.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Framework-side facilities shared by production Kernel wrappers.""" + +import cuda.bindings.driver as cuda +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor + + +class _GraphSafeDLPack: + __slots__ = ("tensor",) + + def __init__(self, tensor: torch.Tensor) -> None: + self.tensor = tensor + + def __dlpack__(self, stream=None): + # stream=-1 skips producer sync; CuTe launches on the current captured stream. + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor: + return from_dlpack( + _GraphSafeDLPack(tensor.detach()), + assumed_align=alignment, + ) + + +def to_cute_dynamic( + tensor: torch.Tensor, + alignment: int, + *, + divisibility: int, +) -> cute.Tensor: + return to_cute(tensor, alignment).mark_compact_shape_dynamic( + mode=0, + divisibility=divisibility, + ) + + +def make_fake_dynamic_compact_tensor( + dtype, + *, + alignment: int, + divisibility: int, +) -> cute.Tensor: + return make_fake_compact_tensor( + dtype, + (cute.sym_int32(divisibility=divisibility),), + assumed_align=alignment, + ) + + +def current_cu_stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/symmetric_buffer.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/symmetric_buffer.py new file mode 100644 index 000000000000..ded27bacbe16 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl/symmetric_buffer.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Typed ownership for one rendezvoused symmetric Tensor.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + + +def _enable_symm_mem_for_group(group_name: str) -> None: + # Local copy of flashinfer.comm.torch_symmetric_memory's helper so the + # fallback package never imports a private symbol from an installed + # flashinfer that may predate it. PyTorch 2.11+ enables groups lazily. + torch_version = tuple(int(x) for x in torch.__version__.split(".")[:2]) + if torch_version >= (2, 11): + return + symm_mem.enable_symm_mem_for_group(group_name) + + +@dataclass(frozen=True, slots=True) +class SymmetricBuffer: + """A symmetric Tensor and the mapping resources derived at rendezvous.""" + + tensor: torch.Tensor + # Keep the rendezvous mapping alive without exposing the backend handle as + # part of a Kernel State's public surface. + _handle: object = field(repr=False) + multicast_address: int | None = field(default=None, repr=False) + peer_addresses: torch.Tensor | None = field(default=None, repr=False) + + @classmethod + def allocate( + cls, + shape: Sequence[int], + dtype: torch.dtype, + device: torch.device, + group: dist.ProcessGroup, + *, + require_multicast: bool = False, + materialize_peer_addresses: bool = False, + ) -> SymmetricBuffer: + """Allocate with the current SymmMem backend and verify requested mappings.""" + if symm_mem.get_backend(device) is None: + raise RuntimeError( + "PyTorch Symmetric Memory has no backend for the current device" + ) + _enable_symm_mem_for_group(group.group_name) + return cls.rendezvous( + symm_mem.empty(shape, dtype=dtype, device=device), + group, + require_multicast=require_multicast, + materialize_peer_addresses=materialize_peer_addresses, + ) + + @classmethod + def rendezvous( + cls, + tensor: torch.Tensor, + group: dist.ProcessGroup, + *, + require_multicast: bool = False, + materialize_peer_addresses: bool = False, + ) -> SymmetricBuffer: + _enable_symm_mem_for_group(group.group_name) + handle = symm_mem.rendezvous(tensor, group) + multicast_address = None + if require_multicast: + multicast_address = int(handle.multicast_ptr or 0) + if not multicast_address: + raise RuntimeError("NVLink multicast mapping is unavailable") + + peer_addresses = None + if materialize_peer_addresses: + # Preserve the rendezvous offset for SymmMem Pool suballocations. + addresses = [ + handle.get_remote_tensor( + peer, + tensor.shape, + tensor.dtype, + ).data_ptr() + for peer in range(dist.get_world_size(group)) + ] + if any(not address for address in addresses): + raise RuntimeError("Symmetric peer mapping is unavailable") + peer_addresses = torch.tensor( + addresses, + dtype=torch.int64, + device=tensor.device, + ) + + return cls( + tensor=tensor, + _handle=handle, + multicast_address=multicast_address, + peer_addresses=peer_addresses, + ) diff --git a/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl_ar.py b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl_ar.py new file mode 100644 index 000000000000..6808ccd7fd92 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_fallback/comm/mnnvl_cutedsl_ar.py @@ -0,0 +1,568 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MNNVL AllReduce fusion backend implemented with CuTe DSL.""" + +from __future__ import annotations + +import logging +from typing import Any, Optional, cast + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from torch.distributed import ProcessGroup + +from .mnnvl_cutedsl import DEFAULT_CONFIG, MNNVLCuteDSLConfig, ProtocolKind +from .mnnvl_cutedsl.config import StaticProfile +from .mnnvl_cutedsl.kernel_bt import BTAllReduceTuning, BTFinalizeTuning +from .mnnvl_cutedsl.kernel_bt.protocol import BTProtocol +from .mnnvl_cutedsl.kernel_ht import HTAllReduceTuning, HTFinalizeTuning +from .mnnvl_cutedsl.kernel_ht.protocol import HTProtocol +from .mnnvl_cutedsl.kernel_ll import LLAllReduceTuning, LLFinalizeTuning +from .mnnvl_cutedsl.kernel_ll.protocol import LLProtocol +# Keep the copied backend and kernel package self-contained while reusing the +# stable communication infrastructure already supplied by the serving image. +from flashinfer.comm.mnnvl import is_multicast_supported +from flashinfer.comm.trtllm_ar import AllReduceFusionPattern +from flashinfer.comm.workspace_base import AllReduceFusionWorkspace + +logger = logging.getLogger(__name__) + +__all__ = ["MNNVLCuteDSLAllReduceFusionWorkspace"] + + +def _check_tensor( + tensor: torch.Tensor, + name: str, + *, + shape: tuple[int | None, ...], + dtype: torch.dtype, + device: torch.device, + alignment: int, +) -> None: + if tensor.device != device: + raise ValueError(f"{name} must be on {device}") + if tensor.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}") + if tensor.ndim != len(shape) or any( + expected is not None and actual != expected + for actual, expected in zip(tensor.shape, shape, strict=True) + ): + raise ValueError(f"{name} has an unsupported shape") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.data_ptr() % alignment: + raise ValueError(f"{name} must be {alignment}-byte aligned") + + +def _warn_pdl_mismatch( + workspace: "MNNVLCuteDSLAllReduceFusionWorkspace", + pattern: int, + m: int, + launch_with_pdl: bool, +) -> None: + preset_pdl = workspace._uses_pdl(pattern, m) + if launch_with_pdl != preset_pdl: + logger.warning( + "launch_with_pdl does not match the selected MNNVL CuTe DSL " + "preset; using enable_pdl=%s", + preset_pdl, + ) + + +class MNNVLCuteDSLAllReduceFusionWorkspace(AllReduceFusionWorkspace): + """Compiled LL, BT, and HT protocols for one static problem shape. + + Workspace construction compiles the selected kernels and must finish before + the first invocation. Calls using the same workspace must not overlap. + Feature-disabled tensor slots use internal placeholders that are not read. + """ + + _destroyed: bool + + def __init__( + self, + tp_size: int, + tp_rank: int, + max_token_num: int, + hidden_dim: int, + dtype: torch.dtype, + *, + group: Optional[ProcessGroup] = None, + top_k: int = 10, + rms_eps: float = 1e-6, + routed_scaling_factor: float = 1.0, + weight_bias: float = 0.0, + include_shared_expert: bool = True, + add_residual: bool = True, + write_residual_output: bool = True, + config: MNNVLCuteDSLConfig = DEFAULT_CONFIG, + ) -> None: + if tp_size not in (2, 4, 8, 16): + raise ValueError("tp_size must be 2, 4, 8, or 16") + if not 0 <= tp_rank < tp_size: + raise ValueError("tp_rank must be in [0, tp_size)") + if max_token_num <= 0: + raise ValueError("max_token_num must be positive") + if dtype != torch.bfloat16: + raise ValueError("MNNVL CuTe DSL kernels only support torch.bfloat16") + if not torch.cuda.is_available(): + raise RuntimeError("MNNVL CuTe DSL kernels require CUDA") + device = torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.get_device_capability(device)[0] < 10: + raise RuntimeError("MNNVL CuTe DSL kernels require a Blackwell GPU") + if symm_mem.get_backend(device) is None: + raise RuntimeError("PyTorch Symmetric Memory is unavailable") + if not is_multicast_supported(device.index): + raise RuntimeError("NVLink multicast is unavailable") + if group is None: + if not dist.is_initialized(): + raise ValueError("A ProcessGroup is required before initialization") + group = dist.group.WORLD + if dist.get_world_size(group) != tp_size: + raise ValueError("ProcessGroup size does not match tp_size") + if dist.get_rank(group) != tp_rank: + raise ValueError("ProcessGroup rank does not match tp_rank") + + super().__init__(tp_size, tp_rank) + self._protocols: dict[ProtocolKind, LLProtocol | BTProtocol | HTProtocol] = {} + self.max_token_num = max_token_num + self.hidden_dim = hidden_dim + self.top_k = top_k + self.dtype = dtype + self.group = group + self.rms_eps = rms_eps + self.routed_scaling_factor = routed_scaling_factor + self.weight_bias = weight_bias + self.include_shared_expert = include_shared_expert + self.add_residual = add_residual + self.write_residual_output = write_residual_output + self.config = config + self.profile = config.resolve( + tp_size=tp_size, + hidden_size=hidden_dim, + top_k=top_k, + dtype=dtype, + capacity_m=max_token_num, + ) + + for protocol in (ProtocolKind.LL, ProtocolKind.BT, ProtocolKind.HT): + capacity = self.profile.protocol_capacity( + protocol, capacity_m=max_token_num + ) + if capacity is None: + continue + finalize_tunings = self._tunings( + self.profile, protocol, finalize=True, capacity_m=capacity + ) + all_reduce_tunings = self._tunings( + self.profile, protocol, finalize=False, capacity_m=capacity + ) + common = dict( + hidden_size=hidden_dim, + top_k=top_k, + tp_size=tp_size, + rank=tp_rank, + capacity_m=capacity, + rms_epsilon=rms_eps, + routed_scaling_factor=routed_scaling_factor, + weight_bias=weight_bias, + include_shared_expert=include_shared_expert, + add_residual=add_residual, + write_residual_output=write_residual_output, + group=group, + ) + instance: LLProtocol | BTProtocol | HTProtocol + if protocol is ProtocolKind.LL: + instance = LLProtocol( + **common, + finalize_tunings=finalize_tunings, + all_reduce_tunings=all_reduce_tunings, + ) + elif protocol is ProtocolKind.BT: + instance = BTProtocol( + **common, + finalize_tunings=finalize_tunings, + all_reduce_tunings=all_reduce_tunings, + ) + else: + instance = HTProtocol( + **common, + finalize_tunings=finalize_tunings, + all_reduce_tunings=all_reduce_tunings, + ) + self._protocols[protocol] = instance + + torch.cuda.synchronize(device) + dist.barrier(group=group) + + @staticmethod + def _tunings( + profile: StaticProfile, + protocol: ProtocolKind, + *, + finalize: bool, + capacity_m: int, + ) -> tuple: + routes = profile.finalize_routes if finalize else profile.all_reduce_routes + tunings = tuple( + dict.fromkeys( + target.preset + for target in routes.targets_for_capacity(capacity_m) + if target.protocol is protocol + ) + ) + expected_type = { + (ProtocolKind.LL, True): LLFinalizeTuning, + (ProtocolKind.LL, False): LLAllReduceTuning, + (ProtocolKind.BT, True): BTFinalizeTuning, + (ProtocolKind.BT, False): BTAllReduceTuning, + (ProtocolKind.HT, True): HTFinalizeTuning, + (ProtocolKind.HT, False): HTAllReduceTuning, + }[(protocol, finalize)] + if not all(isinstance(tuning, expected_type) for tuning in tunings): + path = "finalize" if finalize else "all-reduce" + raise TypeError(f"Invalid {protocol.value} {path} preset") + return tunings + + def _uses_pdl(self, pattern: int, m: int) -> bool: + if pattern == AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm: + target = self.profile.finalize_routes.select(m) + elif pattern == AllReduceFusionPattern.kARResidualRMSNorm: + target = self.profile.all_reduce_routes.select(m) + else: + raise NotImplementedError("Unsupported MNNVL CuTe DSL fusion pattern") + preset = cast(Any, target.preset) + enabled = getattr(preset, "enable_pdl", None) + if enabled is None: + enabled = preset.collective.enable_pdl + return bool(enabled) + + @property + def backend(self) -> str: + return "mnnvl-cutedsl" + + def is_buffer_size_sufficient( + self, + tp_size: int, + num_tokens: int, + hidden_dim: int, + dtype: torch.dtype, + use_oneshot=None, + ) -> bool: + del use_oneshot + return ( + tp_size == self.world_size + and num_tokens <= self.max_token_num + and hidden_dim == self.hidden_dim + and dtype == self.dtype + ) + + def _finalize_all_reduce_rms_norm( + self, + routed_output: torch.Tensor, + expert_weights: torch.Tensor, + permuted_indices: torch.Tensor, + shared_output: torch.Tensor | None, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + norm_output: torch.Tensor | None, + residual_output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + target = self.profile.finalize_routes.select(m) + protocol = cast(Any, self._protocols[target.protocol]) + kernel = protocol.finalize_kernels[target.preset] + return kernel( + routed_output, + expert_weights, + permuted_indices, + shared_output, + residual_source, + gamma, + m, + state=protocol.state, + norm_output=norm_output, + residual_output=residual_output, + ) + + def _all_reduce_rms_norm( + self, + local_contribution: torch.Tensor, + residual_source: torch.Tensor | None, + gamma: torch.Tensor, + m: int, + *, + norm_output: torch.Tensor | None, + residual_output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + target = self.profile.all_reduce_routes.select(m) + protocol = cast(Any, self._protocols[target.protocol]) + kernel = protocol.all_reduce_kernels[target.preset] + return kernel( + local_contribution, + residual_source, + gamma, + m, + state=protocol.state, + norm_output=norm_output, + residual_output=residual_output, + ) + + def destroy(self) -> None: + if self._destroyed: + return + self._protocols.clear() + self._destroyed = True + + +def _mnnvl_cutedsl_allreduce_fusion( + input: torch.Tensor, + workspace: MNNVLCuteDSLAllReduceFusionWorkspace, + pattern: int, + *, + launch_with_pdl: bool, + output: Optional[torch.Tensor] = None, + residual_in: Optional[torch.Tensor] = None, + residual_out: Optional[torch.Tensor] = None, + norm_out: Optional[torch.Tensor] = None, + quant_out: Optional[torch.Tensor] = None, + scale_out: Optional[torch.Tensor] = None, + rms_gamma: Optional[torch.Tensor] = None, + rms_eps: float = 1e-6, + scale_factor: Optional[torch.Tensor | float] = None, + layout_code: Optional[int] = None, + use_oneshot: Optional[bool] = None, + fp32_acc: bool = False, + moe_reduction_device_num_experts: Optional[int] = None, + moe_reduction_scale_input: Optional[torch.Tensor] = None, + moe_reduction_active_experts_token_input: Optional[torch.Tensor] = None, + moe_reduction_token_input: Optional[torch.Tensor] = None, + weight_bias: float = 0.0, + expanded_idx_to_permuted_idx: Optional[torch.Tensor] = None, + expert_scale_factor: Optional[torch.Tensor] = None, + shared_expert_output: Optional[torch.Tensor] = None, + block_quant_group_size: Optional[int] = None, +) -> torch.Tensor: + if workspace._destroyed: + raise RuntimeError( + "The MNNVLCuteDSLAllReduceFusionWorkspace has been destroyed" + ) + if pattern not in ( + AllReduceFusionPattern.kARResidualRMSNorm, + AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm, + ): + raise NotImplementedError("Unsupported MNNVL CuTe DSL fusion pattern") + unsupported = [ + name + for name, value in ( + ("output", output), + ("quant_out", quant_out), + ("scale_out", scale_out), + ("scale_factor", scale_factor), + ("layout_code", layout_code), + ("use_oneshot", use_oneshot), + ("block_quant_group_size", block_quant_group_size), + ("moe_reduction_scale_input", moe_reduction_scale_input), + ( + "moe_reduction_active_experts_token_input", + moe_reduction_active_experts_token_input, + ), + ("moe_reduction_token_input", moe_reduction_token_input), + ) + if value is not None + ] + if fp32_acc: + unsupported.append("fp32_acc") + if moe_reduction_device_num_experts is not None: + unsupported.append("moe_reduction_device_num_experts") + if unsupported: + raise ValueError("MNNVL CuTe DSL does not support: " + ", ".join(unsupported)) + + if rms_eps != workspace.rms_eps: + raise ValueError("rms_eps does not match the compiled workspace") + if weight_bias != workspace.weight_bias: + raise ValueError("weight_bias does not match the compiled workspace") + if rms_gamma is None: + raise ValueError("rms_gamma is required") + if workspace.add_residual and residual_in is None: + raise ValueError("residual_in is required by the compiled workspace") + if not workspace.add_residual and residual_in is not None: + raise ValueError("residual_in must be None for this compiled workspace") + if not workspace.write_residual_output and residual_out is not None: + raise ValueError("residual_out must be None for this compiled workspace") + + device = torch.device("cuda", torch.cuda.current_device()) + hidden = workspace.hidden_dim + _check_tensor( + rms_gamma, + "rms_gamma", + shape=(hidden,), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + + if pattern == AllReduceFusionPattern.kARResidualRMSNorm: + if any( + value is not None + for value in ( + expanded_idx_to_permuted_idx, + expert_scale_factor, + shared_expert_output, + ) + ): + raise ValueError("MoE finalize operands require the finalize pattern") + _check_tensor( + input, + "input", + shape=(None, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + m = input.shape[0] + if not 1 <= m <= workspace.max_token_num: + raise ValueError("input token count exceeds workspace capacity") + _warn_pdl_mismatch(workspace, pattern, m, launch_with_pdl) + if residual_in is not None: + _check_tensor( + residual_in, + "residual_in", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + if norm_out is not None: + _check_tensor( + norm_out, + "norm_out", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + if residual_out is not None: + _check_tensor( + residual_out, + "residual_out", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + norm_out, _ = workspace._all_reduce_rms_norm( + input, + residual_in, + rms_gamma, + input.shape[0], + norm_output=norm_out, + residual_output=residual_out, + ) + return norm_out + + if pattern == AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm: + if expanded_idx_to_permuted_idx is None: + raise ValueError("expanded_idx_to_permuted_idx is required") + if expert_scale_factor is None: + raise ValueError("expert_scale_factor is required") + if workspace.include_shared_expert and shared_expert_output is None: + raise ValueError( + "shared_expert_output is required by the compiled workspace" + ) + if not workspace.include_shared_expert and shared_expert_output is not None: + raise ValueError( + "shared_expert_output must be None for this compiled workspace" + ) + m = expanded_idx_to_permuted_idx.shape[0] + if not 1 <= m <= workspace.max_token_num: + raise ValueError("input token count exceeds workspace capacity") + _warn_pdl_mismatch(workspace, pattern, m, launch_with_pdl) + _check_tensor( + input, + "input", + shape=(None, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + _check_tensor( + expert_scale_factor, + "expert_scale_factor", + shape=(m, workspace.top_k), + dtype=torch.bfloat16, + device=device, + alignment=2, + ) + _check_tensor( + expanded_idx_to_permuted_idx, + "expanded_idx_to_permuted_idx", + shape=(m, workspace.top_k), + dtype=torch.int32, + device=device, + alignment=4, + ) + if shared_expert_output is not None: + _check_tensor( + shared_expert_output, + "shared_expert_output", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + if residual_in is not None: + _check_tensor( + residual_in, + "residual_in", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + if norm_out is not None: + _check_tensor( + norm_out, + "norm_out", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + if residual_out is not None: + _check_tensor( + residual_out, + "residual_out", + shape=(m, hidden), + dtype=torch.bfloat16, + device=device, + alignment=16, + ) + norm_out, _ = workspace._finalize_all_reduce_rms_norm( + input, + expert_scale_factor, + expanded_idx_to_permuted_idx, + shared_expert_output, + residual_in, + rms_gamma, + m, + norm_output=norm_out, + residual_output=residual_out, + ) + return norm_out + + raise AssertionError("unreachable") diff --git a/python/sglang/srt/layers/flashinfer_mnnvl_cutedsl.py b/python/sglang/srt/layers/flashinfer_mnnvl_cutedsl.py new file mode 100644 index 000000000000..c237f76bf91d --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_mnnvl_cutedsl.py @@ -0,0 +1,366 @@ +"""Process-local access to FlashInfer's MNNVL CuTe DSL fusion workspace.""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass, replace +from functools import lru_cache +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +from sglang.srt.layers.flashinfer_provider import get_flashinfer_comm_provider + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + +logger = logging.getLogger(__name__) + + +def _with_early_finalize_shared_load(config): + """Copy a provider config and opt its LL/BT finalize presets into overlap.""" + profiles = [] + updated_presets = 0 + for profile in config.profiles: + targets = [] + for target in profile.finalize_routes.targets: + preset = target.preset + if hasattr(preset, "load_shared_expert_before_pdl"): + preset = replace(preset, load_shared_expert_before_pdl=True) + target = replace(target, preset=preset) + updated_presets += 1 + targets.append(target) + profiles.append( + replace( + profile, + finalize_routes=replace( + profile.finalize_routes, + targets=tuple(targets), + ), + ) + ) + + if updated_presets == 0: + raise RuntimeError( + "FlashInfer MNNVL config does not expose the finalize shared-load " + "PDL ordering option" + ) + return replace(config, profiles=tuple(profiles)) + + +@dataclass(frozen=True, slots=True) +class _WorkspaceSignature: + hidden_size: int + top_k: int + rms_epsilon: float + weight_bias: float + max_m: int + device_index: int + process_group_identity: int + + +class FlashInferMNNVLCuteDSLARFusion: + """One graph-stable workspace serving both supported fusion patterns.""" + + def __init__( + self, + *, + hidden_size: int, + top_k: int, + max_m: int, + rms_epsilon: float, + weight_bias: float, + process_group: ProcessGroup, + device: torch.device, + ) -> None: + if hidden_size <= 0 or top_k <= 0 or max_m <= 0: + raise ValueError("hidden_size, top_k, and max_m must be positive") + if device.type != "cuda": + raise ValueError(f"MNNVL CuTe DSL fusion requires CUDA, got {device}") + + self.hidden_size = int(hidden_size) + self.top_k = int(top_k) + self.max_m = int(max_m) + self.rms_epsilon = float(rms_epsilon) + self.weight_bias = float(weight_bias) + self.process_group = process_group + self.device = torch.device(device) + self._destroyed = False + + with torch.cuda.device(self.device): + self.device = torch.device("cuda", torch.cuda.current_device()) + # The CuTe DSL provider obtains its NVLS workspace from PyTorch + # symmetric memory. This selects PyTorch's allocation/rendezvous + # backend, not the FlashInfer fusion provider implemented below. + # It is process-local setup and must precede workspace construction. + import torch.distributed._symmetric_memory as symm_mem + + symmetric_memory_backend = symm_mem.get_backend(self.device) + if symmetric_memory_backend is None: + symm_mem.set_backend("NCCL") + symmetric_memory_backend = symm_mem.get_backend(self.device) + if symmetric_memory_backend is None: + raise RuntimeError( + "PyTorch symmetric memory has no backend for the current device" + ) + logger.info( + "Using PyTorch symmetric-memory backend %s for %s", + symmetric_memory_backend, + self.device, + ) + + self.provider = get_flashinfer_comm_provider() + # Qwen's shared-expert handoff is complete before this fused finalize + # launch. Opt only finalize kernels into FlashInfer's faster early + # shared load; standalone AllReduce kernels retain the safe ordering. + from sglang.srt.runtime_context import get_spec + + if get_spec().speculative_algorithm is None: + self.workspace_config = _with_early_finalize_shared_load( + self.provider.default_config + ) + else: + # The early shared load lets the fused finalize PDL-preload + # the shared-expert buffer before its predecessor completes. + # That is only safe for the single looping decode graph; + # speculative decoding alternates draft/verify graph replays + # and intermittently reads a not-yet-ready buffer. + logger.info( + "Speculative decoding active: keeping the FlashInfer MNNVL " + "CuTe DSL finalize presets on the safe (non-early-load) " + "ordering." + ) + self.workspace_config = self.provider.default_config + self.workspace = self.provider.workspace_type( + tp_size=dist.get_world_size(process_group), + tp_rank=dist.get_rank(process_group), + max_token_num=self.max_m, + hidden_dim=self.hidden_size, + dtype=torch.bfloat16, + group=process_group, + top_k=self.top_k, + rms_eps=self.rms_epsilon, + routed_scaling_factor=1.0, + weight_bias=self.weight_bias, + include_shared_expert=True, + add_residual=True, + write_residual_output=True, + config=self.workspace_config, + ) + + # Defensive publish barrier: current workspace classes already + # synchronize + barrier at the end of __init__, but an older + # upstream FlashInfer workspace may not, and publishing an + # incompletely initialized mailbox desynchronizes the Lamport + # stages permanently. + torch.cuda.synchronize(self.device) + dist.barrier(group=process_group) + + def supports(self, m: int) -> bool: + if self._destroyed or not 1 <= int(m) <= self.max_m: + return False + return self.workspace.is_buffer_size_sufficient( + tp_size=dist.get_world_size(self.process_group), + num_tokens=int(m), + hidden_dim=self.hidden_size, + dtype=torch.bfloat16, + ) + + def moe_finalize_all_reduce_rms_norm( + self, + *, + routed_output: torch.Tensor, + expert_weights: torch.Tensor, + permuted_indices: torch.Tensor, + gated_shared_output: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + m = int(permuted_indices.shape[0]) + if not self.supports(m): + raise ValueError(f"workspace does not support M={m}") + shape = (m, self.hidden_size) + if norm_output is None: + norm_output = torch.empty(shape, dtype=torch.bfloat16, device=self.device) + if residual_output is None: + residual_output = torch.empty( + shape, dtype=torch.bfloat16, device=self.device + ) + + pattern = self.provider.patterns.kMoEFinalizeARResidualRMSNorm + self.provider.allreduce_fusion( + input=routed_output, + workspace=self.workspace, + pattern=pattern, + # The public API carries the caller's PDL intent. The backend's + # routing profile owns the compiled choice and validates it. + launch_with_pdl=True, + residual_in=residual, + residual_out=residual_output, + norm_out=norm_output, + rms_gamma=gamma, + rms_eps=self.rms_epsilon, + weight_bias=self.weight_bias, + expanded_idx_to_permuted_idx=permuted_indices, + expert_scale_factor=expert_weights, + shared_expert_output=gated_shared_output, + ) + return norm_output, residual_output + + def all_reduce_residual_rms_norm( + self, + *, + local_contribution: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + norm_output: torch.Tensor | None = None, + residual_output: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + m = int(local_contribution.shape[0]) + if not self.supports(m): + raise ValueError(f"workspace does not support M={m}") + if norm_output is None: + norm_output = torch.empty_like(local_contribution) + if residual_output is None: + residual_output = torch.empty_like(local_contribution) + + pattern = self.provider.patterns.kARResidualRMSNorm + self.provider.allreduce_fusion( + input=local_contribution, + workspace=self.workspace, + pattern=pattern, + launch_with_pdl=True, + residual_in=residual, + residual_out=residual_output, + norm_out=norm_output, + rms_gamma=gamma, + rms_eps=self.rms_epsilon, + weight_bias=self.weight_bias, + ) + return norm_output, residual_output + + def destroy(self) -> None: + if self._destroyed: + return + self.workspace.destroy() + self._destroyed = True + + +_WORKSPACES: dict[_WorkspaceSignature, FlashInferMNNVLCuteDSLARFusion] = {} +_WORKSPACES_LOCK = threading.RLock() + + +@lru_cache(maxsize=1) +def _max_workspace_instances() -> int: + from sglang.srt.environ import envs + + value = int(envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES.get()) + if value < 1: + raise ValueError( + "SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES must be positive" + ) + return value + + +def get_flashinfer_mnnvl_cutedsl_ar_fusion( + *, + hidden_size: int | None = None, + top_k: int | None = None, + max_m: int | None = None, + rms_epsilon: float | None = None, + weight_bias: float | None = None, +) -> FlashInferMNNVLCuteDSLARFusion: + """Lookup, or before graph capture create, the process-local workspace.""" + supplied = (hidden_size, top_k, max_m, rms_epsilon, weight_bias) + if all(value is None for value in supplied): + with _WORKSPACES_LOCK: + if len(_WORKSPACES) == 1: + return next(iter(_WORKSPACES.values())) + if not _WORKSPACES: + raise RuntimeError( + "MNNVL CuTe DSL fusion workspace was not initialized before use" + ) + raise RuntimeError( + "multiple MNNVL CuTe DSL fusion workspaces exist; configuration " + "arguments are required" + ) + if any(value is None for value in supplied): + raise TypeError( + "hidden_size, top_k, max_m, rms_epsilon, and weight_bias must be " + "supplied together" + ) + if not torch.cuda.is_available(): + raise RuntimeError("MNNVL CuTe DSL fusion requires CUDA") + + assert hidden_size is not None + assert top_k is not None + assert max_m is not None + assert rms_epsilon is not None + assert weight_bias is not None + from sglang.srt.distributed.parallel_state import get_tp_group + + device = torch.device("cuda", torch.cuda.current_device()) + process_group = get_tp_group().device_group + domain = ( + int(hidden_size), + int(top_k), + float(rms_epsilon), + float(weight_bias), + int(device.index), + id(process_group), + ) + + with _WORKSPACES_LOCK: + compatible = [ + (signature.max_m, instance) + for signature, instance in _WORKSPACES.items() + if ( + signature.hidden_size, + signature.top_k, + signature.rms_epsilon, + signature.weight_bias, + signature.device_index, + signature.process_group_identity, + ) + == domain + and signature.max_m >= int(max_m) + ] + if compatible: + return min(compatible, key=lambda item: item[0])[1] + + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "creating an MNNVL CuTe DSL fusion workspace during CUDA Graph " + "capture is forbidden" + ) + if len(_WORKSPACES) >= _max_workspace_instances(): + raise RuntimeError( + "MNNVL CuTe DSL fusion workspace instance limit exceeded; " + "increase SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES " + "only when multiple model configurations intentionally coexist" + ) + + signature = _WorkspaceSignature( + hidden_size=int(hidden_size), + top_k=int(top_k), + rms_epsilon=float(rms_epsilon), + weight_bias=float(weight_bias), + max_m=int(max_m), + device_index=int(device.index), + process_group_identity=id(process_group), + ) + instance = FlashInferMNNVLCuteDSLARFusion( + hidden_size=hidden_size, + top_k=top_k, + max_m=max_m, + rms_epsilon=rms_epsilon, + weight_bias=weight_bias, + process_group=process_group, + device=device, + ) + _WORKSPACES[signature] = instance + return instance diff --git a/python/sglang/srt/layers/flashinfer_provider.py b/python/sglang/srt/layers/flashinfer_provider.py new file mode 100644 index 000000000000..0cbab32068b7 --- /dev/null +++ b/python/sglang/srt/layers/flashinfer_provider.py @@ -0,0 +1,177 @@ +"""Resolve the FlashInfer MNNVL CuTe DSL communication provider. + +The stable FlashInfer surface is intentionally small: construct the backend- +specific workspace, then invoke the unified ``flashinfer.comm.allreduce_fusion`` +function. Until the serving image contains that workspace implementation, +SGLang supplies the same two-part surface from a temporary copied provider. +""" + +from __future__ import annotations + +import inspect +import logging +from functools import lru_cache +from types import ModuleType +from typing import Any, Callable + +import msgspec + +logger = logging.getLogger(__name__) + + +_REQUIRED_WORKSPACE_PARAMETERS = { + "tp_size", + "tp_rank", + "max_token_num", + "hidden_dim", + "dtype", + "group", + "top_k", + "rms_eps", + "routed_scaling_factor", + "weight_bias", + "include_shared_expert", + "add_residual", + "write_residual_output", + "config", +} +_REQUIRED_ALLREDUCE_PARAMETERS = { + "input", + "workspace", + "pattern", + "launch_with_pdl", + "residual_in", + "residual_out", + "norm_out", + "rms_gamma", + "rms_eps", + "weight_bias", + "expanded_idx_to_permuted_idx", + "expert_scale_factor", + "shared_expert_output", +} +_REQUIRED_SUPPORTS_PARAMETERS = { + "tp_size", + "num_tokens", + "hidden_dim", + "dtype", +} + + +class FlashInferMNNVLCuteDSLProvider(msgspec.Struct, frozen=True): + workspace_type: type + allreduce_fusion: Callable[..., Any] + patterns: type + default_config: Any + + +def _accepts_required_parameters(callable_object, required_parameters) -> bool: + try: + parameters = inspect.signature(callable_object).parameters + except (TypeError, ValueError): + return False + return bool( + required_parameters <= parameters.keys() + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ) + ) + + +def _make_provider( + comm: ModuleType, workspace_type: type, *, default_config: Any +) -> FlashInferMNNVLCuteDSLProvider | None: + if not all( + hasattr(workspace_type, method) + for method in ("is_buffer_size_sufficient", "destroy") + ): + return None + if not _accepts_required_parameters(workspace_type, _REQUIRED_WORKSPACE_PARAMETERS): + return None + if not _accepts_required_parameters( + workspace_type.is_buffer_size_sufficient, _REQUIRED_SUPPORTS_PARAMETERS + ): + return None + allreduce_fusion = getattr(comm, "allreduce_fusion", None) + if allreduce_fusion is None: + return None + if not _accepts_required_parameters( + allreduce_fusion, _REQUIRED_ALLREDUCE_PARAMETERS + ): + return None + patterns = getattr(comm, "AllReduceFusionPattern", None) + if patterns is None or not all( + hasattr(patterns, name) + for name in ( + "kARResidualRMSNorm", + "kMoEFinalizeARResidualRMSNorm", + ) + ): + return None + return FlashInferMNNVLCuteDSLProvider( + workspace_type=workspace_type, + allreduce_fusion=allreduce_fusion, + patterns=patterns, + default_config=default_config, + ) + + +@lru_cache(maxsize=1) +def get_flashinfer_comm_provider() -> FlashInferMNNVLCuteDSLProvider: + """Return the upstream provider, or the API-compatible copied fallback.""" + import flashinfer.comm as upstream_comm + + try: + from flashinfer.comm.mnnvl_cutedsl_ar import ( + MNNVLCuteDSLAllReduceFusionWorkspace as upstream_workspace_type, + ) + from flashinfer.comm.mnnvl_cutedsl import ( + DEFAULT_CONFIG as upstream_default_config, + ) + except ImportError as error: + logger.debug("Upstream FlashInfer MNNVL CuTe DSL import failed: %s", error) + else: + provider = _make_provider( + upstream_comm, + upstream_workspace_type, + default_config=upstream_default_config, + ) + if provider is not None: + logger.info("Using upstream FlashInfer MNNVL CuTe DSL fusion backend") + return provider + logger.debug( + "Installed FlashInfer contains an incompatible MNNVL CuTe DSL API; " + "using SGLang's copied provider" + ) + + try: + from sglang.srt.layers.flashinfer_fallback import comm as fallback_comm + from sglang.srt.layers.flashinfer_fallback.comm.mnnvl_cutedsl_ar import ( + MNNVLCuteDSLAllReduceFusionWorkspace as fallback_workspace_type, + ) + from sglang.srt.layers.flashinfer_fallback.comm.mnnvl_cutedsl import ( + DEFAULT_CONFIG as fallback_default_config, + ) + except ImportError as error: + raise RuntimeError( + "MNNVL CuTe DSL fusion requires either a FlashInfer release with " + "the backend or SGLang's copied provider dependencies, including " + "nvidia-cutlass-dsl and cuda-python" + ) from error + + provider = _make_provider( + fallback_comm, + fallback_workspace_type, + default_config=fallback_default_config, + ) + if provider is None: + raise RuntimeError( + "SGLang's copied FlashInfer MNNVL CuTe DSL provider has an " + "incompatible API" + ) + logger.warning( + "Installed FlashInfer does not provide the stable MNNVL CuTe DSL API; " + "using SGLang's copied provider" + ) + return provider diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 480ed81a5e30..ba25e8de1a21 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -15,6 +15,7 @@ import dataclasses import logging +import os from contextlib import contextmanager from typing import Any, Dict, List, Optional, Tuple, Union @@ -71,6 +72,22 @@ } +def _trace_e2e_logits(stage: str, **fields) -> None: + """Opt-in stage trace for diagnosing post-MoE DP-attention stalls.""" + if os.getenv("SGLANG_TRACE_LOGITS_E2E", "0") != "1": + return + try: + parallel = get_parallel() + rank = ( + f"dp={parallel.attn_dp_rank} " + f"tp={parallel.tp_rank}" + ) + except Exception: + rank = "rank=unknown" + details = " ".join(f"{key}={value}" for key, value in fields.items()) + print(f"SGLANG_TRACE_LOGITS_E2E {rank} stage={stage} {details}", flush=True) + + def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool: return all(hasattr(lm_head, attr_name) for attr_name in attr_names) @@ -718,24 +735,50 @@ def _get_logits( last position (e.g., extend without input logprobs). The caller should guarantee the given hidden_states follow this constraint. """ + _trace_e2e_logits( + "get_logits_enter", + hidden_shape=tuple(hidden_states.shape), + dp_gather=self.do_tensor_parallel_all_gather_dp_attn, + tp_gather=self.do_tensor_parallel_all_gather, + ) hidden_states, local_hidden_states = self._gather_dp_attn_hidden_states( hidden_states, logits_metadata ) + _trace_e2e_logits( + "dp_hidden_gather_returned", + global_shape=tuple(hidden_states.shape), + local_shape=tuple(local_hidden_states.shape), + ) + + if os.getenv("SGLANG_TRACE_LOGITS_E2E_SYNC", "0") == "1": + _trace_e2e_logits("pre_lm_head_sync_enter") + torch.cuda.synchronize() + _trace_e2e_logits("pre_lm_head_sync_returned") + _trace_e2e_logits("lm_head_enter", hidden_shape=tuple(hidden_states.shape)) logits = self._compute_lm_head(hidden_states, lm_head, embedding_bias) + _trace_e2e_logits("lm_head_returned", logits_shape=tuple(logits.shape)) + if os.getenv("SGLANG_TRACE_LOGITS_E2E_SYNC", "0") == "1": + _trace_e2e_logits("post_lm_head_sync_enter") + torch.cuda.synchronize() + _trace_e2e_logits("post_lm_head_sync_returned") if self.logit_scale is not None: logits.mul_(self.logit_scale) if self.do_tensor_parallel_all_gather: + _trace_e2e_logits("tp_logits_gather_enter", logits_shape=tuple(logits.shape)) if self.use_attn_tp_group: logits = self._gather_attn_tp_logits(logits) else: logits = self._logits_gatherer(logits) + _trace_e2e_logits("tp_logits_gather_returned", logits_shape=tuple(logits.shape)) + _trace_e2e_logits("dp_logits_scatter_enter", logits_shape=tuple(logits.shape)) logits = self._scatter_dp_attn_logits( logits, local_hidden_states, logits_metadata ) + _trace_e2e_logits("dp_logits_scatter_returned", logits_shape=tuple(logits.shape)) logits = self._copy_logits_to_buffer( logits, logits_metadata, use_buffer=use_logits_buffer @@ -803,10 +846,27 @@ def _gather_dp_attn_hidden_states( self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata ) -> Tuple[torch.Tensor, torch.Tensor]: if self.do_tensor_parallel_all_gather_dp_attn: + _trace_e2e_logits( + "dp_metadata_enter", + local_shape=tuple(hidden_states.shape), + global_counts_cpu=logits_metadata.global_num_tokens_for_logprob_cpu, + ) logits_metadata.compute_dp_attention_metadata() + _trace_e2e_logits( + "dp_metadata_returned", + buffer_shape=tuple(logits_metadata.gathered_buffer.shape), + local_start=logits_metadata.dp_local_start_pos, + local_tokens=logits_metadata.dp_local_num_tokens, + ) local_hidden_states = hidden_states hidden_states = logits_metadata.gathered_buffer + _trace_e2e_logits( + "dp_hidden_gather_enter", + global_shape=tuple(hidden_states.shape), + local_shape=tuple(local_hidden_states.shape), + ) dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata) + _trace_e2e_logits("dp_hidden_gather_collective_returned") return hidden_states, local_hidden_states return hidden_states, hidden_states diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 513ba0aaa667..77a1eee4a9f2 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -103,7 +103,12 @@ def __init__( and quant_config is not None and quant_config.get_name() == "humming" ) - if is_humming: + if get_moe_a2a_backend().is_deepep_v2(): + # deepep_v2 runs on the base FusedMoE forward via its own + # DeepEPv2Dispatcher, so always delegate (never use DeepEPMoE's + # v1-specific dispatch/run_moe_core path). + self.deprecate_flag = True + elif is_humming: self.deprecate_flag = True elif _use_aiter: self.deprecate_flag = True @@ -354,6 +359,7 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): if ( get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_pplx() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index b3f1291f484f..0b451f57e85c 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -38,6 +38,7 @@ AscendTPDispatcher, ) from sglang.srt.layers.moe.token_dispatcher.base import BaseDispatcher +from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2Dispatcher from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardDispatcher, @@ -166,6 +167,15 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: async_finish=True, return_recv_hook=True, ) + elif a2a_backend.is_deepep_v2(): + return DeepEPv2Dispatcher( + group=get_tp_group().device_group, + router_topk=moe_runner_config.top_k, + num_experts=moe_runner_config.num_experts, + num_local_experts=moe_runner_config.num_local_experts, + hidden_size=moe_runner_config.hidden_size, + params_dtype=moe_runner_config.params_dtype, + ) elif a2a_backend.is_flashinfer(): return FlashinferDispatcher( group=get_tp_group().device_group, @@ -173,6 +183,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: num_experts=moe_runner_config.num_experts, num_local_experts=moe_runner_config.num_local_experts, hidden_size=moe_runner_config.hidden_size, + moe_runner_config=moe_runner_config, ) else: raise NotImplementedError(f"Unsupported a2a backend: {a2a_backend}") @@ -222,6 +233,8 @@ class FusedMoE(torch.nn.Module): reduce_results: Whether to apply all_reduce on the output of the layer quant_config: Quantization configuration. inplace: suggestion to compute inplace (modify input activation). + enable_qwen35_fp8_deferred_finalize: Whether this concrete Qwen3.5 + layer may expose FlashInfer's block-FP8 deferred MoE output. """ # True on shared-expert FusedMoE subclasses (e.g. Inkling's sink); lets @@ -256,6 +269,7 @@ def __init__( routing_method_type: Optional[RoutingMethodType] = None, is_gated: bool = True, gate_up_interleaved: bool = True, + enable_qwen35_fp8_deferred_finalize: bool = False, ): super().__init__() if params_dtype is None: @@ -379,10 +393,17 @@ def __init__( self.use_deep_gemm, ) _validate_hpc_ops_quant_method(self.quant_method) + nvfp4_deferred = envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get() and isinstance( + self.quant_method, ModelOptNvFp4FusedMoEMethod + ) + qwen35_fp8_deferred = ( + enable_qwen35_fp8_deferred_finalize + and isinstance(self.quant_method, Fp8MoEMethod) + and self.quant_method.block_quant + ) self.supports_deferred_finalize = ( - envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get() - and get_moe_runner_backend().is_flashinfer_trtllm() - and isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) + get_moe_runner_backend().is_flashinfer_trtllm() + and (nvfp4_deferred or qwen35_fp8_deferred) ) global _deferred_finalize_info_logged if not _deferred_finalize_info_logged: @@ -435,6 +456,15 @@ def __init__( and ( get_moe_runner_backend().is_cutlass() or get_moe_runner_backend().is_flashinfer_trtllm_routed() + # FlashInfer A2A materializes routing before TRT-LLM Gen + # MoE. The regular backend name therefore enters the same + # packed routed kernel as flashinfer_trtllm_routed and must + # use the same top-k scaling contract. Keep the original + # fused-routing behavior unchanged when A2A is disabled. + or ( + get_moe_a2a_backend().is_flashinfer() + and get_moe_runner_backend().is_flashinfer_trtllm() + ) ) ) or ( diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 67bdb5ea8b6f..05a58c39a7c2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import os from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional, Tuple @@ -31,7 +32,7 @@ register_pre_permute, ) from sglang.srt.layers.moe.utils import MoeRunnerBackend -from sglang.srt.runtime_context import get_exec +from sglang.srt.runtime_context import get_exec, get_flags from sglang.srt.utils import ( ceil_div, dispose_tensor, @@ -50,6 +51,13 @@ DeepEPNormalCombineInput, DeepEPNormalDispatchOutput, ) + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import ( + DeepEPv2CombineInput, + DeepEPv2DispatchOutput, + ) + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferDispatchOutput, + ) from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardCombineInput, StandardDispatchOutput, @@ -158,6 +166,14 @@ def _should_use_masked_standard_layout( quant_info: DeepGemmMoeQuantInfo, hidden_states: torch.Tensor, ) -> bool: + # Preserve the Oakhaven WideEP escape hatch while adopting upstream's + # memory-budget-based auto policy. CUDA graph capture remains masked. + if ( + envs.SGLANG_OPT_DG_COMPACT_EAGER.get() + and not get_flags().capture.disable_dispose_tensor + ): + return False + mode = envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() if mode not in ("auto", "masked", "compact"): raise ValueError( @@ -207,6 +223,7 @@ class DeepGemmRunnerInput(RunnerInput): masked_m: Optional[torch.Tensor] = None expected_m: Optional[int] = None m_indices: Optional[torch.Tensor] = None + hidden_states_scale_tma_aligned: bool = False @property def runner_backend(self) -> MoeRunnerBackend: @@ -306,11 +323,24 @@ def _run_contiguous_gemm( hidden_states_dtype = running_state["hidden_states_dtype"] hidden_states_shape = running_state["hidden_states_shape"] m_indices = runner_input.m_indices + trace_deepep_v2_contig = ( + os.environ.get("SGLANG_DEEPEP_V2_TRACE_CONTIG") == "1" + and running_state.get("deepep_v2_expanded", False) + ) N = quant_info.w13_weight.size(1) K = hidden_states_shape[1] scale_block_size = 128 + if all_tokens == 0: + if trace_deepep_v2_contig: + logger.warning("DeepEP v2 expanded contig runner empty return") + dispose_tensor(hidden_states) + dispose_tensor(hidden_states_scale) + return torch.empty( + (0, K), device=hidden_states_device, dtype=torch.bfloat16 + ) + recipe_a, recipe_b = ( ((1, 128), (1, 32)) if quant_info.is_fp4_experts else (None, None) ) @@ -326,7 +356,22 @@ def _run_contiguous_gemm( device=hidden_states_device, dtype=torch.bfloat16, ) - if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: + if trace_deepep_v2_contig: + logger.warning( + "DeepEP v2 expanded contig runner enter: hidden=%s scale=%s " + "m_indices=%s", + tuple(hidden_states.shape), + None + if hidden_states_scale is None + else tuple(hidden_states_scale.shape), + m_indices.detach().cpu().tolist(), + ) + torch.cuda.synchronize() + logger.warning("DeepEP v2 expanded contig runner pre-sync returned") + if ( + deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES + and not runner_input.hidden_states_scale_tma_aligned + ): hidden_states_scale = tma_align_input_scale(hidden_states_scale) deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( @@ -337,6 +382,9 @@ def _run_contiguous_gemm( recipe_a=recipe_a, recipe_b=recipe_b, ) + if trace_deepep_v2_contig: + torch.cuda.synchronize() + logger.warning("DeepEP v2 expanded contig gateup GEMM returned") dispose_tensor(hidden_states) dispose_tensor(hidden_states_scale) @@ -397,6 +445,9 @@ def _run_contiguous_gemm( del down_input elif envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get(): swiglu_limit_arg: Optional[float] = self.swiglu_limit + use_contig_swizzle = self.use_swizzle and not running_state.get( + "deepep_v2_disable_contig_swizzle", False + ) down_input_fp8 = torch.empty( (all_tokens, N // 2), @@ -419,7 +470,7 @@ def _run_contiguous_gemm( scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, transposed=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, swiglu_limit=swiglu_limit_arg, - swizzle=self.use_swizzle, + swizzle=use_contig_swizzle, ) del gateup_output else: @@ -454,6 +505,9 @@ def _run_contiguous_gemm( scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ) del down_input + if trace_deepep_v2_contig: + torch.cuda.synchronize() + logger.warning("DeepEP v2 expanded contig activation returned") # Allocate the MoE output in the NCCL symmetric memory pool when symmetric # allocation is required, so the downstream all-reduce takes the low-latency @@ -478,6 +532,9 @@ def _run_contiguous_gemm( recipe_a=recipe_a, recipe_b=recipe_b, ) + if trace_deepep_v2_contig: + torch.cuda.synchronize() + logger.warning("DeepEP v2 expanded contig down GEMM returned") return down_output @@ -568,6 +625,22 @@ def _run_masked_gemm( w2_scale = quant_info.w2_scale hidden_states_device = running_state["hidden_states_device"] + trace_deepep_v2_masked = ( + os.environ.get("SGLANG_DEEPEP_V2_TRACE_MASKED") == "1" + ) + if trace_deepep_v2_masked: + logger.warning( + "DeepEP v2 masked runner enter: hidden=%s hidden_stride=%s " + "scale=%s scale_stride=%s masked_m=%s expected_m=%s", + tuple(hidden_states.shape), + hidden_states.stride(), + None + if hidden_states_scale is None + else tuple(hidden_states_scale.shape), + None if hidden_states_scale is None else hidden_states_scale.stride(), + masked_m.detach().cpu().tolist(), + expected_m, + ) use_mxfp8 = quant_info.use_mxfp8 scale_block_size = quant_info.block_shape[1] if quant_info.block_shape else 128 @@ -632,6 +705,9 @@ def _run_masked_gemm( recipe_a=recipe_a, recipe_b=recipe_b, ) + if trace_deepep_v2_masked: + torch.cuda.synchronize() + logger.warning("DeepEP v2 masked runner gateup GEMM returned") dispose_tensor(hidden_states) dispose_tensor(hidden_states_scale) @@ -690,6 +766,9 @@ def _run_masked_gemm( gemm1_clamp_limit=self.config.gemm1_clamp_limit, num_real_tokens=num_real_tokens, ) + if trace_deepep_v2_masked: + torch.cuda.synchronize() + logger.warning("DeepEP v2 masked runner activation returned") del gateup_output # Down activation is quantised locally at scale_block_size (never DeepEP-LL), @@ -748,6 +827,9 @@ def _run_masked_gemm( recipe_b=recipe_b, **gemm_overlap_args_dict, ) + if trace_deepep_v2_masked: + torch.cuda.synchronize() + logger.warning("DeepEP v2 masked runner down GEMM returned") meta_overlap_args = running_state.get("meta_overlap_args", None) # Returns (block_m, threshold) only with down-gemm overlap, else None; # meta_overlap_args may be set without overlap, so guard the unpack. @@ -836,6 +918,7 @@ def pre_permute_standard_to_deep_gemm( quant_info: DeepGemmMoeQuantInfo, runner_config: MoeRunnerConfig, running_state: dict, + expert_start: int = 0, ) -> DeepGemmRunnerInput: from sglang.kernels.ops.moe.ep_moe_kernels import ( ep_scatter, @@ -871,6 +954,7 @@ def pre_permute_standard_to_deep_gemm( quant_info.block_shape, output_dtype=output_dtype, use_mxfp8=quant_info.use_mxfp8, + expert_start=expert_start, ) ) # Use the global expert count because expected_m is a tuning hint, not @@ -912,7 +996,7 @@ def pre_permute_standard_to_deep_gemm( all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e) tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index( - topk_ids, num_experts, 1 + topk_ids, num_experts, 1, expert_start=expert_start ) dispose_tensor(unused_masked_dst) valid_tokens_per_expert = tokens_per_expert @@ -992,6 +1076,7 @@ def pre_permute_standard_to_deep_gemm( src2dst, scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128), + expert_start=expert_start, ) if packed_input_source is not hidden_states: dispose_tensor(packed_input_source) @@ -1021,6 +1106,49 @@ def pre_permute_standard_to_deep_gemm( ) +@register_pre_permute("flashinfer", "deep_gemm") +def pre_permute_flashinfer_to_deep_gemm( + dispatch_output: FlashinferDispatchOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepGemmRunnerInput: + """Feed one-sided A2A output into DeepGEMM with fused expert remapping.""" + + from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput + from sglang.srt.runtime_context import get_parallel + + if dispatch_output.hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer A2A + DeepGEMM requires a BF16 dispatch payload, got " + f"{dispatch_output.hidden_states.dtype}." + ) + if dispatch_output.hidden_states_scale is not None: + raise ValueError( + "FlashInfer A2A + DeepGEMM expects unquantized BF16 dispatch; " + "hidden_states_scale must be None." + ) + if dispatch_output.topk_output.topk_ids.dtype != torch.int32: + raise TypeError( + "FlashInfer A2A expert IDs must be int32 before DeepGEMM, got " + f"{dispatch_output.topk_output.topk_ids.dtype}." + ) + + standard_output = StandardDispatchOutput( + hidden_states=dispatch_output.hidden_states, + hidden_states_scale=None, + topk_output=dispatch_output.topk_output, + ) + expert_start = get_parallel().moe_ep_rank * runner_config.num_local_experts + return pre_permute_standard_to_deep_gemm( + standard_output, + quant_info, + runner_config, + running_state, + expert_start=expert_start, + ) + + @register_post_permute("deep_gemm", "standard") def post_permute_deep_gemm_to_standard( runner_output: DeepGemmRunnerOutput, @@ -1065,6 +1193,30 @@ def post_permute_deep_gemm_to_standard( ) +@register_post_permute("deep_gemm", "flashinfer") +def post_permute_deep_gemm_to_flashinfer( + runner_output: DeepGemmRunnerOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +): + """Reuse DeepGEMM's weighted post-permute and hand BF16 to A2A combine.""" + + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferCombineInput, + ) + + standard_input = post_permute_deep_gemm_to_standard( + runner_output, quant_info, runner_config, running_state + ) + if standard_input.hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer A2A + DeepGEMM combine payload must be BF16, got " + f"{standard_input.hidden_states.dtype}." + ) + return FlashinferCombineInput(hidden_states=standard_input.hidden_states) + + @register_pre_permute("deepep_ll", "deep_gemm") def pre_permute_deepep_ll_to_deep_gemm( dispatch_output: DeepEPLLDispatchOutput, @@ -1448,3 +1600,290 @@ def _apply_swiglu_limit( out = torch.cat([gate, up], dim=-1) assert out.shape == (num_tokens, hidden_size_x2) return out + + +@register_pre_permute("deepep_v2", "deep_gemm") +def pre_permute_deepep_v2_to_deep_gemm( + dispatch_output: DeepEPv2DispatchOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepGemmRunnerInput: + from sglang.kernels.ops.moe.ep_moe_kernels import ( + ep_expand_init_m_indices_from_psum, + ep_scatter, + ep_scatter_from_psum, + ) + + hidden_states = dispatch_output.hidden_states + hidden_states_scale = dispatch_output.hidden_states_scale + topk_ids = dispatch_output.topk_ids + topk_weights = dispatch_output.topk_weights + num_recv_tokens_per_expert = dispatch_output.num_recv_tokens_per_expert + psum_num_recv_tokens_per_expert = dispatch_output.psum_num_recv_tokens_per_expert + is_expanded = dispatch_output.is_expanded + hidden_states_scale_tma_aligned = dispatch_output.hidden_states_scale_tma_aligned + deepep_v2_use_masked = dispatch_output.use_masked_gemm + deepep_v2_expected_m = dispatch_output.expected_m + deepep_v2_masked_max_m = dispatch_output.masked_max_m + deepep_v2_total_expanded = dispatch_output.total_expanded + deepep_v2_expert_alignment = dispatch_output.expert_alignment + if hidden_states_scale is None: + raise RuntimeError( + "DeepEP v2 -> DeepGEMM requires FP8 dispatch output with activation scales. " + "Use --deepep-v2-dispatcher-output-dtype fp8 or select a BF16 runner such as triton." + ) + if envs.SGLANG_OPT_FIX_MEGA_MOE_MEMORY.get(): + # The MegaMoE memory optimization enables a swizzled activation kernel + # for its gran=8 interleaved gate/up layout. DeepEP v2's contiguous adapter + # is validated with the non-swizzled activation layout; using the + # swizzled reader here mixes gate/up pairs and breaks generation. + running_state["deepep_v2_disable_contig_swizzle"] = True + assert runner_config.activation == "silu" + + if is_expanded: + if psum_num_recv_tokens_per_expert is None: + raise RuntimeError( + "DeepEP v2 expanded layout requires native expert prefix sums." + ) + all_tokens = hidden_states.shape[0] + running_state["all_tokens"] = all_tokens + running_state["hidden_states_shape"] = hidden_states.shape + running_state["hidden_states_device"] = hidden_states.device + running_state["hidden_states_dtype"] = hidden_states.dtype + running_state["topk_ids"] = None + running_state["topk_weights"] = topk_weights + running_state["deepep_v2_expanded"] = True + + if deepep_v2_use_masked: + # Masked-GEMM bridge: repack the expanded expert-packed buffer into a + # regular [E_local, max_m, hidden] slab so DeepGEMM's masked grouped + # GEMM bounds compute by per-expert real counts (masked_m), decoupled + # from the dispatch capacity. Static shapes -> cuda-graph safe. + from sglang.kernels.ops.moe.ep_moe_kernels import expand_to_masked_slab + + num_local_experts = psum_num_recv_tokens_per_expert.shape[0] + trace_deepep_v2_masked = ( + os.environ.get("SGLANG_DEEPEP_V2_TRACE_MASKED") == "1" + ) + if trace_deepep_v2_masked: + logger.warning( + "DeepEP v2 masked repack enter: hidden=%s hidden_stride=%s " + "scale=%s scale_stride=%s psum=%s max_m=%s align=%s", + tuple(hidden_states.shape), + hidden_states.stride(), + None + if hidden_states_scale is None + else tuple(hidden_states_scale.shape), + None + if hidden_states_scale is None + else hidden_states_scale.stride(), + psum_num_recv_tokens_per_expert.detach().cpu().tolist(), + deepep_v2_masked_max_m, + deepep_v2_expert_alignment, + ) + torch.cuda.synchronize() + logger.warning("DeepEP v2 masked repack pre-sync returned") + slab, slab_scale, masked_m = expand_to_masked_slab( + hidden_states, + hidden_states_scale, + psum_num_recv_tokens_per_expert, + num_local_experts, + deepep_v2_masked_max_m, + deepep_v2_expert_alignment, + ) + if trace_deepep_v2_masked: + torch.cuda.synchronize() + logger.warning( + "DeepEP v2 masked repack returned: slab=%s scale=%s masked_m=%s", + tuple(slab.shape), + None if slab_scale is None else tuple(slab_scale.shape), + masked_m.detach().cpu().tolist(), + ) + running_state["deepep_v2_masked"] = True + running_state["deepep_v2_psum"] = psum_num_recv_tokens_per_expert + running_state["deepep_v2_total_expanded"] = deepep_v2_total_expanded + running_state["deepep_v2_expert_alignment"] = deepep_v2_expert_alignment + return DeepGemmRunnerInput( + hidden_states=slab, + hidden_states_scale=slab_scale, + use_masked_gemm=True, + masked_m=masked_m, + expected_m=deepep_v2_expected_m, + hidden_states_scale_tma_aligned=hidden_states_scale_tma_aligned, + ) + + # do_cpu_sync=False -> recv buffer is worst-case sized; ep_expand_init only + # writes real-token slots, so pre-fill the tail with -1 to skip padding rows. + m_indices = torch.full( + (all_tokens,), -1, device=hidden_states.device, dtype=torch.int32 + ) + ep_expand_init_m_indices_from_psum(psum_num_recv_tokens_per_expert, m_indices) + if os.environ.get("SGLANG_DEEPEP_V2_TRACE_CONTIG") == "1": + torch.cuda.synchronize() + logger.warning( + "DeepEP v2 expanded contig m_indices ready: all_tokens=%s " + "m_indices=%s", + all_tokens, + m_indices.detach().cpu().tolist(), + ) + return DeepGemmRunnerInput( + hidden_states=hidden_states, + hidden_states_scale=hidden_states_scale, + use_masked_gemm=False, + m_indices=m_indices, + hidden_states_scale_tma_aligned=hidden_states_scale_tma_aligned, + ) + + if psum_num_recv_tokens_per_expert is not None: + all_tokens = int(psum_num_recv_tokens_per_expert[-1].item()) + num_recv_tokens_per_expert_gpu = None + else: + num_recv_tokens_per_expert = [ + ceil_div(x, 128) * 128 for x in num_recv_tokens_per_expert + ] + all_tokens = sum(num_recv_tokens_per_expert) + num_recv_tokens_per_expert_gpu = torch.tensor( + num_recv_tokens_per_expert, dtype=torch.int32, pin_memory=True, device="cpu" + ).cuda(non_blocking=True) + K = hidden_states.shape[1] + running_state["all_tokens"] = all_tokens + running_state["hidden_states_shape"] = hidden_states.shape + running_state["hidden_states_device"] = hidden_states.device + running_state["hidden_states_dtype"] = hidden_states.dtype + running_state["topk_ids"] = topk_ids + running_state["topk_weights"] = topk_weights + + # Match the legacy deepep_normal adapter (same ep_scatter + grouped GEMM): the + # scatter writes only real-token rows and the post-permute ep_gather reads them + # back via output_index, so the per-expert alignment padding rows are never + # consumed and the activation buffer needs no zero-init. The ue8m0 packed-scale + # layout keeps zeros (its in-int32 padding lanes must be zero). + input_tensor = torch.empty( + (all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype + ) + if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + input_tensor_scale = torch.zeros( + (ceil_div(K // 128, 4), all_tokens), + device=hidden_states.device, + dtype=torch.int, + ).transpose(0, 1) + else: + input_tensor_scale = torch.empty( + (all_tokens, K // 128), device=hidden_states.device, dtype=torch.float32 + ) + m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32) + output_index = torch.empty_like(topk_ids) + if psum_num_recv_tokens_per_expert is not None: + # Contiguous-path alignment contract: this psum comes from ElasticBuffer + # dispatch(do_expand=False, expert_alignment=capability.expert_alignment), + # and DeepEP documents the non-expand psum as the inclusive prefix sum of + # alignment-PADDED per-expert counts (deep_ep/buffers/elastic.py). The + # deep_gemm capability pins expert_alignment=128 == + # get_m_alignment_for_contiguous_layout(), so psum[e-1] is a valid + # 128-aligned group start for the contiguous grouped GEMM. Do NOT re-align + # here: an align_up would silently mask an upstream contract break. + expert_start_loc = torch.empty_like(psum_num_recv_tokens_per_expert) + ep_scatter_from_psum( + hidden_states, + hidden_states_scale, + topk_ids, + psum_num_recv_tokens_per_expert, + expert_start_loc, + input_tensor, + input_tensor_scale, + m_indices, + output_index, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + else: + expert_start_loc = torch.empty_like(num_recv_tokens_per_expert_gpu) + ep_scatter( + hidden_states, + hidden_states_scale, + topk_ids, + num_recv_tokens_per_expert_gpu, + expert_start_loc, + input_tensor, + input_tensor_scale, + m_indices, + output_index, + scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + ) + dispose_tensor(hidden_states) + dispose_tensor(hidden_states_scale) + running_state["output_index"] = output_index + + return DeepGemmRunnerInput( + hidden_states=input_tensor, + hidden_states_scale=input_tensor_scale, + use_masked_gemm=False, + m_indices=m_indices, + ) + + +@register_post_permute("deep_gemm", "deepep_v2") +def post_permute_deep_gemm_to_deepep_v2( + runner_output: DeepGemmRunnerOutput, + quant_info: DeepGemmMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> DeepEPv2CombineInput: + from sglang.kernels.ops.moe.ep_moe_kernels import ep_gather + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2CombineInput + + if running_state.get("deepep_v2_expanded", False): + hidden_states = runner_output.hidden_states + topk_weights = running_state["topk_weights"] + trace_deepep_v2_contig = ( + os.environ.get("SGLANG_DEEPEP_V2_TRACE_CONTIG") == "1" + and not running_state.get("deepep_v2_masked", False) + ) + if trace_deepep_v2_contig: + torch.cuda.synchronize() + logger.warning( + "DeepEP v2 expanded contig post-permute enter: hidden=%s " + "topk_weights=%s", + tuple(hidden_states.shape), + None if topk_weights is None else tuple(topk_weights.shape), + ) + if running_state.get("deepep_v2_masked", False): + # Masked path: GEMM output is the [E_local, max_m, hidden] slab. Repack + # it back to expanded row order (padding rows zeroed) before combine. + from sglang.kernels.ops.moe.ep_moe_kernels import masked_slab_to_expand + + hidden_states = masked_slab_to_expand( + hidden_states, + running_state["deepep_v2_psum"], + running_state["deepep_v2_total_expanded"], + running_state["deepep_v2_expert_alignment"], + topk_weights=topk_weights, + ) + return DeepEPv2CombineInput(hidden_states, None, None) + if topk_weights is not None: + # Expanded combine does not consume top-k weights, so apply them to + # each expert slot before combine. Keep this out-of-place until the + # runner/communication buffer reuse contract is explicitly audited. + hidden_states = hidden_states * topk_weights.to( + hidden_states.dtype + ).unsqueeze(-1) + if trace_deepep_v2_contig: + torch.cuda.synchronize() + logger.warning("DeepEP v2 expanded contig post-permute returned") + return DeepEPv2CombineInput(hidden_states, None, None) + + hidden_states = runner_output.hidden_states + topk_ids = running_state["topk_ids"] + topk_weights = running_state["topk_weights"] + output_index = running_state["output_index"] + gather_out = torch.empty( + running_state["hidden_states_shape"], + device=running_state["hidden_states_device"], + dtype=torch.bfloat16, + ) + ep_gather(hidden_states, topk_ids, topk_weights, output_index, gather_out) + return DeepEPv2CombineInput( + hidden_states=gather_out, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index f821d095c474..4e530c02925e 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -87,6 +87,41 @@ def finalize_flashinfer_trtllm_deferred_output( ) +def _make_deferred_finalize_output( + result, + *, + top_k: int, +) -> FlashInferTrtllmDeferredFinalizeOutput: + """Validate and adapt FlashInfer's ``do_finalize=False`` output ABI.""" + gemm2_out, expert_weights, expanded_idx_to_permuted_idx = result[:3] + # Some FlashInfer versions size this buffer from routing_logits dtype while + # writing BF16 weights into it. Reinterpret only the live BF16 prefix. + if expert_weights.dtype == torch.float32: + n, k = expert_weights.shape + expert_weights = expert_weights.view(torch.bfloat16).view(-1, k)[:n] + if expert_weights.dtype != torch.bfloat16: + raise RuntimeError( + "FlashInfer deferred finalize must return BF16 expert weights, got " + f"{expert_weights.dtype}" + ) + if gemm2_out.dtype != torch.bfloat16: + raise RuntimeError( + "FlashInfer deferred finalize must return BF16 GEMM2 output, got " + f"{gemm2_out.dtype}" + ) + if expanded_idx_to_permuted_idx.dtype != torch.int32: + raise RuntimeError( + "FlashInfer deferred finalize must return Int32 permuted indices, got " + f"{expanded_idx_to_permuted_idx.dtype}" + ) + return FlashInferTrtllmDeferredFinalizeOutput( + gemm2_out=gemm2_out, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + top_k=top_k, + ) + + def round_up_to_multiple(x: int, m: int) -> int: """Round up *x* to the nearest multiple of *m*.""" return (x + m - 1) // m * m @@ -695,6 +730,16 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( else Fp8QuantizationType.DeepSeekFp8 ) use_shuffled_weight = quant_info.use_mxfp8 + defer_finalize = _deferred_finalize_enabled.get() + if defer_finalize and ( + not quant_info.block_quant + or use_routed_topk + or not TopKOutputChecker.format_is_bypassed(topk_output) + ): + raise RuntimeError( + "FP8 deferred finalize requires block quantization, the logits-based " + "FlashInfer TRTLLM backend, and bypassed TopK" + ) if quant_info.block_quant: assert quant_info.weight_block_k is not None @@ -717,16 +762,19 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( ) a_sf_t = a_sf.t() - # Allocate output inside symmetric memory context - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - symm_output = torch.empty( - hidden_states.shape[0], - hidden_states.shape[1], - dtype=hidden_states.dtype, - device=hidden_states.device, - ) + symm_output = None + if not defer_finalize: + # The deferred path returns FlashInfer's permuted/padded GEMM2 + # materialization and must not allocate the ordinary final output. + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + symm_output = torch.empty( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) # Move kernel call outside context manager to avoid graph breaks # during torch.compile for piecewise cuda graph. @@ -769,10 +817,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( fp8_quantization_type=int(fp8_quantization_type), activation_type=quant_info.activation_type, ) + output = cast(torch.Tensor, symm_output) else: assert TopKOutputChecker.format_is_bypassed(topk_output) - - trtllm_fp8_block_scale_moe_out_wrapper( + common_kwargs = dict( routing_logits=router_logits, routing_bias=correction_bias, hidden_states=a_q, @@ -781,7 +829,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( gemm1_weights_scale=quant_info.w13_weight_scale_inv, gemm2_weights=quant_info.w2_weight, gemm2_weights_scale=quant_info.w2_weight_scale_inv, - output=symm_output, num_experts=quant_info.global_num_experts, top_k=topk_config.top_k, n_group=topk_config.num_expert_group, @@ -797,10 +844,34 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( routing_method_type=routing_method_type, use_shuffled_weight=use_shuffled_weight, tune_max_num_tokens=next_power_of_2(a_q.shape[0]), - fp8_quantization_type=int(fp8_quantization_type), - activation_type=quant_info.activation_type, ) - output = symm_output + if defer_finalize: + from flashinfer.fused_moe import trtllm_fp8_block_scale_moe + + deferred_kwargs = dict( + **common_kwargs, + do_finalize=False, + fp8_quantization_type=fp8_quantization_type, + enable_pdl=a_q.shape[0] <= _TRTLLM_MOE_PDL_MAX_TOKENS, + ) + if quant_info.activation_type is not None: + from flashinfer.fused_moe.core import ActivationType + + deferred_kwargs["activation_type"] = ActivationType( + quant_info.activation_type + ) + output = _make_deferred_finalize_output( + trtllm_fp8_block_scale_moe(**deferred_kwargs), + top_k=topk_config.top_k, + ) + else: + trtllm_fp8_block_scale_moe_out_wrapper( + **common_kwargs, + output=cast(torch.Tensor, symm_output), + fp8_quantization_type=int(fp8_quantization_type), + activation_type=quant_info.activation_type, + ) + output = cast(torch.Tensor, symm_output) else: assert TopKOutputChecker.format_is_bypassed(topk_output) assert quant_info.w13_input_scale is not None @@ -1319,23 +1390,26 @@ def fused_experts_none_to_flashinfer_trtllm_routed( ) +@register_fused_func("flashinfer", "flashinfer_trtllm") @register_fused_func("flashinfer", "flashinfer_trtllm_routed") -def fused_experts_flashinfer_to_flashinfer_trtllm_routed( - dispatch_output: FlashinferDispatchOutput, +def fused_experts_flashinfer_to_flashinfer_trtllm( + dispatch_output: FlashinferDispatchOutput | StandardDispatchOutput, quant_info: MoeQuantInfo, runner_config: MoeRunnerConfig, -) -> FlashinferCombineInput: - """Fused function for flashinfer A2A + flashinfer_trtllm_routed runner. - - FlashinferDispatchOutput and StandardDispatchOutput share the same field - layout (hidden_states, hidden_states_scale, topk_output), so the existing - FP8/FP4/BF16 implementations work unchanged. We wrap the returned - StandardCombineInput into a FlashinferCombineInput for the FlashinferDispatcher - combine path. +) -> FlashinferCombineInput | StandardCombineInput: + """Fused function for FlashInfer A2A + TRT-LLM Gen MoE. + + Both one-sided decode and AG+RS prefill materialize routing IDs and weights, + so the regular and explicitly-routed backend names enter TRT-LLM's routed + kernel. The dispatch formats share the fields consumed by the implementation; + only the combine wrapper differs. """ from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( FlashinferCombineInput, ) + from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardDispatchOutput, + ) if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo): result = fused_experts_none_to_flashinfer_trtllm_fp4( @@ -1345,6 +1419,16 @@ def fused_experts_flashinfer_to_flashinfer_trtllm_routed( use_routed_topk=True, ) elif isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo): + if dispatch_output.hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer A2A + TRT-LLM Gen FP8 MoE requires a BF16 " + f"dispatch payload, got {dispatch_output.hidden_states.dtype}." + ) + if dispatch_output.hidden_states_scale is not None: + raise ValueError( + "FlashInfer A2A + TRT-LLM Gen FP8 MoE quantizes locally; " + "the BF16 dispatch payload must not carry activation scales." + ) result = fused_experts_none_to_flashinfer_trtllm_fp8( dispatch_output, quant_info, @@ -1360,8 +1444,18 @@ def fused_experts_flashinfer_to_flashinfer_trtllm_routed( ) else: raise TypeError( - f"Unexpected quant_info type for flashinfer a2a + flashinfer_trtllm_routed: {type(quant_info)}" + f"Unexpected quant_info type for flashinfer a2a + flashinfer_trtllm: {type(quant_info)}" + ) + if ( + isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo) + and result.hidden_states.dtype != torch.bfloat16 + ): + raise TypeError( + "FlashInfer A2A + TRT-LLM Gen FP8 MoE must return a BF16 combine " + f"payload, got {result.hidden_states.dtype}." ) + if isinstance(dispatch_output, StandardDispatchOutput): + return result return FlashinferCombineInput(hidden_states=result.hidden_states) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 7548086dda13..1ace6e9db1a7 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -19,6 +19,7 @@ from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2DispatchOutput from sglang.srt.layers.moe.token_dispatcher.standard import ( StandardCombineInput, StandardDispatchOutput, @@ -157,7 +158,9 @@ def run( no_combine=self.config.no_combine, inplace=self.config.inplace, apply_router_weight_on_input=self.config.apply_router_weight_on_input, - routed_scaling_factor=self.config.routed_scaling_factor, + routed_scaling_factor=running_state.get( + "deepep_v2_routed_scaling_factor", self.config.routed_scaling_factor + ), gemm1_alpha=self.config.gemm1_alpha, gemm1_limit=self.config.gemm1_clamp_limit, filter_expert=filter_expert, @@ -325,3 +328,105 @@ def post_permute_triton_to_standard( return StandardCombineInput( hidden_states=runner_output.hidden_states, ) + + +def _prepare_triton_runner_input( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + quant_info: TritonMoeQuantInfo, + running_state: dict, +) -> TritonRunnerInput: + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( + _prepare_fused_moe_run, + ) + + ( + config, + down_config, + down_moe_use_tma, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + ) = _prepare_fused_moe_run( + hidden_states, + quant_info.w13_weight, + quant_info.w2_weight, + topk_ids, + use_fp8_w8a8=quant_info.use_fp8_w8a8, + use_int8_w8a8=quant_info.use_int8_w8a8, + use_int8_w8a16=quant_info.use_int8_w8a16, + use_int4_w4a16=quant_info.use_int4_w4a16, + per_channel_quant=quant_info.per_channel_quant, + block_shape=quant_info.block_shape, + ) + running_state["config"] = config + running_state["down_config"] = down_config + running_state["down_moe_use_tma"] = down_moe_use_tma + return TritonRunnerInput( + hidden_states=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + ) + + +@register_pre_permute("deepep_v2", "triton") +def pre_permute_deepep_v2_to_triton( + dispatch_output: DeepEPv2DispatchOutput, + quant_info: TritonMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +) -> TritonRunnerInput: + hidden_states = dispatch_output.hidden_states + hidden_states_scale = dispatch_output.hidden_states_scale + # ElasticBuffer's dispatch API uses int64 expert indices, while the + # non-TMA down-MoE JIT activation consumes the flattened topk_ids directly + # and requires int32. The standard Triton path normally starts with int32, + # so normalize at this format boundary. + topk_ids = dispatch_output.topk_ids.to(torch.int32) + topk_weights = dispatch_output.topk_weights + if hidden_states_scale is not None or hidden_states.dtype != torch.bfloat16: + raise RuntimeError( + "DeepEP v2 -> Triton expects BF16 dispatch output without activation scales. " + "Use --deepep-v2-dispatcher-output-dtype bf16." + ) + # A2A EP combine inputs are kept unscaled. The model-level MoE forward + # applies the routed scaling factor once after combine. + running_state["deepep_v2_routed_scaling_factor"] = None + valid_rows = (topk_ids >= 0).any(dim=1) + running_state["deepep_v2_output_shape"] = hidden_states.shape + running_state["deepep_v2_valid_rows"] = valid_rows + running_state["deepep_v2_topk_ids"] = topk_ids + running_state["deepep_v2_topk_weights"] = topk_weights + hidden_states = hidden_states[valid_rows].contiguous() + topk_ids = topk_ids[valid_rows].contiguous() + topk_weights = topk_weights[valid_rows].contiguous() + return _prepare_triton_runner_input( + hidden_states, topk_ids, topk_weights, quant_info, running_state + ) + + +@register_post_permute("triton", "deepep_v2") +def post_permute_triton_to_deepep_v2( + runner_output: TritonRunnerOutput, + quant_info: TritonMoeQuantInfo, + runner_config: MoeRunnerConfig, + running_state: dict, +): + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import DeepEPv2CombineInput + + valid_rows = running_state["deepep_v2_valid_rows"] + output = torch.zeros( + running_state["deepep_v2_output_shape"], + device=runner_output.hidden_states.device, + dtype=runner_output.hidden_states.dtype, + ) + output[valid_rows] = runner_output.hidden_states + return DeepEPv2CombineInput( + hidden_states=output, + topk_ids=running_state["deepep_v2_topk_ids"], + topk_weights=running_state["deepep_v2_topk_weights"], + ) diff --git a/python/sglang/srt/layers/moe/qwen35_flashinfer_fusion.py b/python/sglang/srt/layers/moe/qwen35_flashinfer_fusion.py new file mode 100644 index 000000000000..b14be44f8aaf --- /dev/null +++ b/python/sglang/srt/layers/moe/qwen35_flashinfer_fusion.py @@ -0,0 +1,346 @@ +"""Qwen3.5 integration for FlashInfer MNNVL CuTe DSL AllReduce fusion.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.srt.layers.communicator import ( + CommunicateWithAllReduceAndLayerNormFn, + LayerCommunicator, + ScatterMode, + get_attn_tp_context, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.moe import get_moe_a2a_backend +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_exec, get_parallel + +logger = logging.getLogger(__name__) + + +def is_supported_forward_mode(forward_mode: ForwardMode) -> bool: + return forward_mode in ( + ForwardMode.DECODE, + ForwardMode.EXTEND, + ForwardMode.TARGET_VERIFY, + ) + + +def resolve_max_m(model_runner) -> int: + """Use framework token bounds as the workspace-capacity source of truth.""" + server_args = model_runner.server_args + decode_config = server_args.cuda_graph_config.decode + prefill_config = server_args.cuda_graph_config.prefill + candidates = [ + server_args.cutedsl_moe_max_num_tokens(), + model_runner.max_running_requests, + decode_config.max_bs, + prefill_config.max_bs, + *(decode_config.bs or []), + *(prefill_config.bs or []), + ] + positive = [ + int(value) for value in candidates if value is not None and int(value) > 0 + ] + if not positive: + raise RuntimeError("framework reported no positive fusion workspace M bound") + return max(positive) + + +@dataclass(frozen=True) +class Qwen35MoeFinalizeHandoff: + """Unfinalized routed output plus the separately gated shared contribution.""" + + routed_output: torch.Tensor + expert_weights: torch.Tensor + permuted_indices: torch.Tensor + gated_shared_output: torch.Tensor + m: int + + @classmethod + def from_flashinfer( + cls, + deferred_output, + *, + gated_shared_output: torch.Tensor, + m: int, + ) -> Qwen35MoeFinalizeHandoff: + top_k = int(deferred_output.top_k) + return cls( + routed_output=deferred_output.gemm2_out.view( + -1, deferred_output.gemm2_out.shape[-1] + ), + expert_weights=deferred_output.expert_weights.view(-1, top_k)[:m], + permuted_indices=deferred_output.expanded_idx_to_permuted_idx.view( + -1, top_k + )[:m], + gated_shared_output=gated_shared_output, + m=int(m), + ) + + +class Qwen35FlashInferFusionService: + """A lightweight model handle for the process-local FlashInfer workspace.""" + + def __init__( + self, + *, + hidden_size: int, + top_k: int, + rms_epsilon: float, + ) -> None: + self.hidden_size = int(hidden_size) + self.top_k = int(top_k) + self.rms_epsilon = float(rms_epsilon) + self.max_m: int | None = None + self._workspace = None + + @property + def is_prepared(self) -> bool: + return self._workspace is not None + + def prepare(self, *, max_m: int) -> None: + if self._workspace is not None: + assert self.max_m is not None + if int(max_m) > self.max_m: + raise RuntimeError( + f"fusion workspace is already prepared for M_max={self.max_m}; " + f"refusing M_max={max_m}" + ) + return + from sglang.srt.layers.flashinfer_mnnvl_cutedsl import ( + get_flashinfer_mnnvl_cutedsl_ar_fusion, + ) + + workspace = get_flashinfer_mnnvl_cutedsl_ar_fusion( + hidden_size=self.hidden_size, + top_k=self.top_k, + max_m=int(max_m), + rms_epsilon=self.rms_epsilon, + # GemmaRMSNorm.gemma_weight is already checkpoint weight + 1. + weight_bias=0.0, + ) + self._workspace = workspace + self.max_m = workspace.max_m + + def supports(self, m: int) -> bool: + if self._workspace is None or self.max_m is None: + return False + return 1 <= int(m) <= self.max_m and self._workspace.supports(m) + + def finalize( + self, + handoff: Qwen35MoeFinalizeHandoff, + residual: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_finalize(handoff, residual, gamma) + assert self._workspace is not None + return self._workspace.moe_finalize_all_reduce_rms_norm( + routed_output=handoff.routed_output, + expert_weights=handoff.expert_weights, + permuted_indices=handoff.permuted_indices, + gated_shared_output=handoff.gated_shared_output, + residual=residual, + gamma=gamma, + ) + + def all_reduce_residual_rms_norm( + self, + local_contribution: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_matrix(local_contribution, "local_contribution") + self._validate_matrix(residual, "residual", m=local_contribution.shape[0]) + self._validate_gamma(gamma) + if not self.supports(local_contribution.shape[0]): + raise ValueError(f"unsupported M={local_contribution.shape[0]}") + assert self._workspace is not None + return self._workspace.all_reduce_residual_rms_norm( + local_contribution=local_contribution, + residual=residual, + gamma=gamma, + ) + + def _validate_finalize( + self, + handoff: Qwen35MoeFinalizeHandoff, + residual: torch.Tensor, + gamma: torch.Tensor, + ) -> None: + if not self.supports(handoff.m): + raise ValueError(f"unsupported M={handoff.m}") + self._validate_matrix(handoff.routed_output, "routed_output", exact_m=False) + expected_metadata = (handoff.m, self.top_k) + if tuple(handoff.expert_weights.shape) != expected_metadata: + raise ValueError("expert_weights must have shape [M, top_k]") + if tuple(handoff.permuted_indices.shape) != expected_metadata: + raise ValueError("permuted_indices must have shape [M, top_k]") + if handoff.expert_weights.dtype != torch.bfloat16: + raise ValueError("expert_weights must be BF16") + if handoff.permuted_indices.dtype != torch.int32: + raise ValueError("permuted_indices must be Int32") + self._validate_matrix( + handoff.gated_shared_output, "gated_shared_output", m=handoff.m + ) + self._validate_matrix(residual, "residual", m=handoff.m) + self._validate_gamma(gamma) + + def _validate_matrix( + self, + tensor: torch.Tensor, + name: str, + *, + m: int | None = None, + exact_m: bool = True, + ) -> None: + if tensor.ndim != 2 or tensor.shape[1] != self.hidden_size: + raise ValueError(f"{name} must have shape [M, hidden_size]") + if m is not None and exact_m and tensor.shape[0] != int(m): + raise ValueError(f"{name} has the wrong M dimension") + if tensor.dtype != torch.bfloat16 or not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous BF16") + + def _validate_gamma(self, gamma: torch.Tensor) -> None: + if ( + tuple(gamma.shape) != (self.hidden_size,) + or gamma.dtype != torch.bfloat16 + or not gamma.is_contiguous() + ): + raise ValueError("gamma must be contiguous BF16 [hidden_size]") + + +class Qwen35FlashInferLayerCommunicator(LayerCommunicator): + """Qwen-only hooks; generic LayerCommunicator remains backend agnostic.""" + + fusion_service: Qwen35FlashInferFusionService | None = None + + def prepare_attn( + self, + hidden_states, + residual, + forward_batch, + quant_format: str = "", + post_residual_addition=None, + ): + if isinstance(hidden_states, Qwen35MoeFinalizeHandoff): + if not self.should_use_finalize(forward_batch, hidden_states.m): + raise RuntimeError("received deferred MoE output on an ineligible path") + if residual is None: + raise RuntimeError("deferred MoE finalize requires residual input") + if not hasattr(self.input_layernorm, "gemma_weight"): + raise RuntimeError("deferred Qwen finalize requires GemmaRMSNorm") + if post_residual_addition is not None: + residual = residual + post_residual_addition + assert self.fusion_service is not None + return self.fusion_service.finalize( + hidden_states, residual, self.input_layernorm.gemma_weight + ) + return super().prepare_attn( + hidden_states, + residual, + forward_batch, + quant_format=quant_format, + post_residual_addition=post_residual_addition, + ) + + def prepare_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + cache=None, + ): + if cache is not None: + self._context.cache = cache + if self.should_use_all_reduce_rms_norm( + forward_batch, int(hidden_states.shape[0]), residual + ): + assert self.fusion_service is not None and residual is not None + return self.fusion_service.all_reduce_residual_rms_norm( + hidden_states, + residual, + self.post_attention_layernorm.gemma_weight, + ) + return super().prepare_mlp(hidden_states, residual, forward_batch, cache=cache) + + def should_use_all_reduce_rms_norm( + self, + forward_batch: ForwardBatch, + m: int, + residual: Optional[torch.Tensor], + ) -> bool: + communicate_fn = self._communicate_with_all_reduce_and_layer_norm_fn + norm_fn = getattr(communicate_fn, "func", communicate_fn) + residual_input_mode = getattr(communicate_fn, "keywords", {}).get( + "residual_input_mode" + ) + parallel = get_parallel() + return ( + self._common_eligible(forward_batch, m) + and residual is not None + and hasattr(self.post_attention_layernorm, "gemma_weight") + and norm_fn + is CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual + and residual_input_mode is ScatterMode.TP_ATTN_FULL + and self._context.attn_dp_size == 1 + and parallel.attn_tp_size == parallel.tp_size + and not get_exec().comm.enable_quant_communications + ) + + def should_use_finalize(self, forward_batch: ForwardBatch, m: int) -> bool: + parallel = get_parallel() + return ( + self._common_eligible(forward_batch, m) + and self.layer_scatter_modes.mlp_mode is not ScatterMode.SCATTERED + and parallel.moe_ep_size == 1 + ) + + def _common_eligible(self, forward_batch: ForwardBatch, m: int) -> bool: + parallel = get_parallel() + return bool( + self.fusion_service is not None + and self.fusion_service.is_prepared + and is_supported_forward_mode(forward_batch.forward_mode) + and self.fusion_service.supports(m) + and not is_dp_attention_enabled() + and parallel.attn_cp_size == 1 + and not get_attn_tp_context().input_scattered + and get_moe_a2a_backend().is_none() + and self._context.tp_size > 1 + ) + + def should_fuse_mlp_allreduce_with_next_layer( + self, forward_batch: ForwardBatch + ) -> bool: + m = ( + int(forward_batch.input_ids.shape[0]) + if getattr(forward_batch, "input_ids", None) is not None + else 0 + ) + if self.should_use_finalize(forward_batch, m): + # The Qwen model consumes the final layer's handoff with its final + # GemmaRMSNorm, so this is intentionally also true for that layer. + return True + return super().should_fuse_mlp_allreduce_with_next_layer(forward_batch) + + +def prepare_qwen35_flashinfer_fusion(model, model_runner) -> None: + service = getattr(model, "flashinfer_mnnvl_cutedsl_fusion", None) + if service is None: + return + if model_runner.server_args.enable_pdmux: + raise RuntimeError( + "FlashInfer MNNVL CuTe DSL fusion does not support concurrent PDMux " + "streams sharing one mutable workspace" + ) + service.prepare(max_m=resolve_max_m(model_runner)) + logger.info( + "Prepared Qwen3.5 FlashInfer MNNVL CuTe DSL fusion workspace for M_max=%d", + service.max_m, + ) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py index 7f2c0942f95c..806dabf7d98d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/__init__.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/__init__.py @@ -21,6 +21,11 @@ DeepEPNormalCombineInput, DeepEPNormalDispatchOutput, ) +from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import ( + DeepEPv2CombineInput, + DeepEPv2Dispatcher, + DeepEPv2DispatchOutput, +) from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( FlashinferDispatcher, FlashinferDispatchOutput, @@ -72,6 +77,9 @@ "MoriEPLLDispatchOutput", "MoriEPLLCombineInput", "MoriEPDispatcher", + "DeepEPv2Dispatcher", + "DeepEPv2DispatchOutput", + "DeepEPv2CombineInput", "NixlEPCombineInput", "NixlEPDispatchOutput", "NixlEPDispatcher", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index 1ff2beb5bb33..e718c00857ef 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -27,6 +27,8 @@ DeepEPLLDispatchOutput, DeepEPNormalCombineInput, DeepEPNormalDispatchOutput, + DeepEPv2CombineInput, + DeepEPv2DispatchOutput, FlashinferCombineInput, FlashinferDispatchOutput, StandardCombineInput, @@ -165,6 +167,12 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferDispatchOutput]: return dispatch_output.format.is_flashinfer() + @staticmethod + def format_is_deepep_v2( + dispatch_output: DispatchOutput, + ) -> TypeGuard[DeepEPv2DispatchOutput]: + return dispatch_output.format.is_deepep_v2() + class DispatchOutputFormat(Enum): @@ -172,6 +180,7 @@ class DispatchOutputFormat(Enum): DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" FLASHINFER = "flashinfer" + DEEPEP_V2 = "deepep_v2" ASCEND_TP = "ascend_tp" def is_standard(self) -> bool: @@ -195,6 +204,9 @@ def is_deepep(self) -> bool: def is_flashinfer(self) -> bool: return self == DispatchOutputFormat.FLASHINFER + def is_deepep_v2(self) -> bool: + return self == DispatchOutputFormat.DEEPEP_V2 + @runtime_checkable class DispatchOutput(Protocol): @@ -249,12 +261,19 @@ def format_is_flashinfer( ) -> TypeGuard[FlashinferCombineInput]: return combine_input.format == CombineInputFormat.FLASHINFER + @staticmethod + def format_is_deepep_v2( + combine_input: CombineInput, + ) -> TypeGuard[DeepEPv2CombineInput]: + return combine_input.format == CombineInputFormat.DEEPEP_V2 + class CombineInputFormat(Enum): STANDARD = "standard" DEEPEP_NORMAL = "deepep_normal" DEEPEP_LL = "deepep_ll" FLASHINFER = "flashinfer" + DEEPEP_V2 = "deepep_v2" ASCEND_TP = "ascend_tp" diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py new file mode 100644 index 000000000000..58b21a17ff76 --- /dev/null +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py @@ -0,0 +1,845 @@ +from __future__ import annotations + +import logging +import os +from typing import List, NamedTuple, Optional, Tuple + +import torch +import torch.distributed as dist + +from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import get_is_extend_in_batch +from sglang.srt.layers.moe.token_dispatcher.base import ( + BaseDispatcher, + CombineInput, + CombineInputFormat, + DispatchOutput, + DispatchOutputFormat, +) +from sglang.srt.layers.moe.topk import TopKOutput +from sglang.srt.layers.moe.utils import ( + DeepEPv2OutputDtype, + DeepEPv2RunnerCapability, + get_deepep_v2_runner_capability, +) + +logger = logging.getLogger(__name__) + +_SCALE_BLOCK_SIZE = 128 +_deepep_v2_import_error: Optional[BaseException] = None +_fp8_quant_import_error: Optional[BaseException] = None +sglang_per_token_group_quant_fp8 = None + +try: + from deep_ep import ElasticBuffer + + use_deepep_v2 = True +except (ImportError, OSError) as exc: + use_deepep_v2 = False + _deepep_v2_import_error = exc + +if use_deepep_v2: + try: + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + except (ImportError, OSError) as exc: + _fp8_quant_import_error = exc + + +class DeepEPv2DispatchOutput(NamedTuple): + hidden_states: torch.Tensor + hidden_states_scale: Optional[torch.Tensor] + topk_ids: Optional[torch.Tensor] + topk_weights: torch.Tensor + num_recv_tokens_per_expert: List[int] + psum_num_recv_tokens_per_expert: Optional[torch.Tensor] = None + is_expanded: bool = False + hidden_states_scale_tma_aligned: bool = False + use_masked_gemm: bool = False + expected_m: int = 0 + masked_max_m: int = 0 + total_expanded: int = 0 + expert_alignment: int = 128 + + @property + def format(self) -> DispatchOutputFormat: + return DispatchOutputFormat.DEEPEP_V2 + + +class DeepEPv2CombineInput(NamedTuple): + hidden_states: torch.Tensor + topk_ids: Optional[torch.Tensor] + topk_weights: Optional[torch.Tensor] + + @property + def format(self) -> CombineInputFormat: + return CombineInputFormat.DEEPEP_V2 + + +assert isinstance(DeepEPv2DispatchOutput, DispatchOutput) +assert isinstance(DeepEPv2CombineInput, CombineInput) + + +def _raise_deepep_v2_import_error() -> None: + detail = ( + f" Original import error: {_deepep_v2_import_error}" + if _deepep_v2_import_error is not None + else "" + ) + raise ImportError( + "DeepEP v2 (ElasticBuffer) is not available. Install DeepEP v2 from " + "https://github.com/deepseek-ai/DeepEP." + detail + ) + + +def _ensure_deepep_v2_available() -> None: + if not use_deepep_v2: + _raise_deepep_v2_import_error() + + +def _ensure_fp8_quant_available() -> None: + _ensure_deepep_v2_available() + if sglang_per_token_group_quant_fp8 is None: + detail = ( + f" Original import error: {_fp8_quant_import_error}" + if _fp8_quant_import_error is not None + else "" + ) + raise ImportError( + "DeepEP v2 FP8 dispatch requires the SGLang FP8 quantization kernel." + + detail + ) + + +def _get_allow_hybrid_mode() -> bool: + # direct/hybrid is a communication-topology knob resolved from ServerArgs. + # Callers without a running server (synthetic/unit tests) must pass + # allow_hybrid_mode explicitly instead (get_server_args() raises when the + # process-wide ServerArgs is not set). + from sglang.srt.runtime_context import get_server_args + + return get_server_args().deepep_v2_mode == "hybrid" + + +def _quantize_for_deepep_v2_dispatch( + hidden_states: torch.Tensor, capability: DeepEPv2RunnerCapability +): + _ensure_fp8_quant_available() + return sglang_per_token_group_quant_fp8( + hidden_states, + _SCALE_BLOCK_SIZE, + column_major_scales=capability.fp8_scale_tma_aligned, + scale_tma_aligned=capability.fp8_scale_tma_aligned, + scale_ue8m0=capability.fp8_scale_ue8m0, + ) + + +class DeepEPv2Buffer: + _buffer: Optional[ElasticBuffer] = None + _buffer_key: Optional[Tuple] = None + + @classmethod + def get_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + router_topk: int, + num_max_dispatch_tokens_per_rank: int, + use_fp8_dispatch: bool, + allow_hybrid_mode: Optional[bool] = None, + ) -> ElasticBuffer: + _ensure_deepep_v2_available() + + if allow_hybrid_mode is None: + allow_hybrid_mode = _get_allow_hybrid_mode() + key = ( + id(group), + hidden_size, + router_topk, + num_max_dispatch_tokens_per_rank, + use_fp8_dispatch, + allow_hybrid_mode, + dist.get_world_size(group), + ) + if cls._buffer is not None and cls._buffer_key == key: + return cls._buffer + + if cls._buffer is not None: + cls.destroy() + + # DeepEP reuses the torch process group's internal NCCL communicator + # when EP_REUSE_NCCL_COMM=1 (its default). That path requires the group + # to be device-bound at init_process_group time (eager comm init), + # which SGLang's shared init does not do -- reusing then reads an + # uninitialized communicator and ElasticBuffer sizing segfaults in + # ncclTeamWorld. Default to letting DeepEP create its own communicator + # (it binds to the already-set current device); setdefault keeps any + # explicit user override. + os.environ.setdefault("EP_REUSE_NCCL_COMM", "0") + cls._buffer = ElasticBuffer( + group, + num_max_tokens_per_rank=num_max_dispatch_tokens_per_rank, + hidden=hidden_size, + num_topk=router_topk, + use_fp8_dispatch=use_fp8_dispatch, + allow_hybrid_mode=allow_hybrid_mode, + sl_idx=0, + prefer_overlap_with_compute=False, + ) + cls._buffer_key = key + logger.info( + "Initialized DeepEP v2 ElasticBuffer: world_size=%s hidden_size=%s " + "num_topk=%s max_dispatch_tokens_per_rank=%s use_fp8_dispatch=%s " + "allow_hybrid_mode=%s num_bytes=%s", + dist.get_world_size(group), + hidden_size, + router_topk, + num_max_dispatch_tokens_per_rank, + use_fp8_dispatch, + allow_hybrid_mode, + cls._buffer.num_bytes, + ) + return cls._buffer + + @classmethod + def destroy(cls) -> None: + cls._buffer = None + cls._buffer_key = None + + +class _DeepEPv2Impl: + def __init__( + self, + group: dist.ProcessGroup, + router_topk: int, + num_experts: int, + num_local_experts: int, + hidden_size: int, + capability: DeepEPv2RunnerCapability, + num_max_dispatch_tokens_per_rank: int, + allow_hybrid_mode: Optional[bool] = None, + ): + self.group = group + self.router_topk = router_topk + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.capability = capability + self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + # Prefill and decode have different static-shape requirements. A large + # one-pass prefill needs a correspondingly large ElasticBuffer, but the + # decode masked-GEMM slab only needs to cover the largest decode batch. + # Reusing the prefill cap for that slab can allocate multiple GiB per + # expert during CUDA graph capture. Keep the communication buffer cap + # unchanged and allow a smaller, fixed, cross-rank decode slab cap. + masked_cap = int( + os.environ.get( + "SGLANG_DEEPEP_V2_MASKED_NUM_MAX_DISPATCH_TOKENS_PER_RANK", + str(num_max_dispatch_tokens_per_rank), + ) + ) + if masked_cap < 1 or masked_cap > num_max_dispatch_tokens_per_rank: + raise ValueError( + "SGLANG_DEEPEP_V2_MASKED_NUM_MAX_DISPATCH_TOKENS_PER_RANK " + f"must be in [1, {num_max_dispatch_tokens_per_rank}], got " + f"{masked_cap}" + ) + self.masked_num_max_dispatch_tokens_per_rank = masked_cap + self.allow_hybrid_mode = allow_hybrid_mode + self.rank = dist.get_rank(group) + self._handle = None + self._pad_empty_combine = False + self._dispatch_seq = 0 + + def set_runner_capability(self, capability: DeepEPv2RunnerCapability) -> None: + if self.capability != capability: + self._destroy_handle() + self.capability = capability + + def _uses_fp8_dispatch_output(self) -> bool: + return self.capability.output_dtype == DeepEPv2OutputDtype.FP8 + + def _destroy_handle(self) -> None: + self._handle = None + + def _get_buffer(self) -> ElasticBuffer: + return DeepEPv2Buffer.get_buffer( + self.group, + self.hidden_size, + self.router_topk, + self.num_max_dispatch_tokens_per_rank, + self._uses_fp8_dispatch_output(), + allow_hybrid_mode=self.allow_hybrid_mode, + ) + + def _resolve_num_sms_qps(self, buffer: ElasticBuffer) -> Tuple[int, int]: + # num_sms/num_qps are NOT auto-resolved by ElasticBuffer when left at 0; + # 0 means "0 SMs / 0 QPs". Multi-node RDMA dispatch needs real QPs, so + # resolve them from the theoretical helpers (matches the DeepEP elastic + # test harness). Single-node NVLink works with 0 QPs. + # get_theoretical_num_sms is @weak_lru-cached in DeepEP with fixed inputs + # here, and its first (modeling) call happens during eager warmup -- so on + # the CUDA-graph decode path this is a cache lookup: pure host work, no + # device sync, capture-safe. + num_sms = envs.SGLANG_DEEPEP_V2_NUM_SMS.get() + if num_sms == 0: + num_sms = buffer.get_theoretical_num_sms(self.num_experts, self.router_topk) + num_qps = buffer.get_theoretical_num_qps(num_sms) + return num_sms, num_qps + + def _validate_common( + self, hidden_states: torch.Tensor, topk_ids: torch.Tensor + ) -> None: + if hidden_states.shape[0] > self.num_max_dispatch_tokens_per_rank: + raise ValueError( + f"DeepEP v2 dispatch input exceeds the per-rank buffer capacity " + f"{self.num_max_dispatch_tokens_per_rank}, got {hidden_states.shape[0]}. " + "Increase SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK." + ) + if hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"DeepEP v2 hidden size mismatch: expected {self.hidden_size}, " + f"got {hidden_states.shape[1]}" + ) + if ( + self._uses_fp8_dispatch_output() + and self.hidden_size % _SCALE_BLOCK_SIZE != 0 + ): + raise ValueError( + "DeepEP v2 FP8 dispatch requires hidden_size multiple of " + f"{_SCALE_BLOCK_SIZE}, got {self.hidden_size}" + ) + if topk_ids.shape[1] != self.router_topk: + raise ValueError( + f"DeepEP v2 topk mismatch: expected {self.router_topk}, " + f"got {topk_ids.shape[1]}" + ) + + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> DeepEPv2DispatchOutput: + # Handle lifecycle: dispatch produces exactly one handle that the next + # combine() consumes. Guard-first (before the import check) so misuse is + # reportable without DeepEP installed. + if self._handle is not None: + raise RuntimeError( + "DeepEP v2 dispatch called while the previous dispatch handle is " + "still unconsumed (missing combine)" + ) + _ensure_deepep_v2_available() + topk_weights = topk_output.topk_weights + topk_ids = topk_output.topk_ids.to(torch.int64) + self._validate_common(hidden_states, topk_ids) + # DeepEP v2's native expanded layout is profitable for decode-like DeepGEMM + # FP8 workloads but regresses prefill-like ones, so layout is chosen by + # inference PHASE, independently of the comm mode (direct/hybrid is a topology + # knob fixed at server init): decode (non-extend) -> native expanded layout; + # prefill/extend -> non-expanded contiguous layout. This decouples the + # masked-GEMM + CUDA-graph decode fast path from the comm mode, so it is + # available under multi-node `hybrid` too. + force_expand_prefill = ( + os.environ.get("SGLANG_DEEPEP_V2_EXPAND_PREFILL") == "1" + ) + use_expand_layout = self.capability.use_expanded_layout and ( + force_expand_prefill or not get_is_extend_in_batch() + ) + # Decode uses the graph-oriented expanded -> masked-GEMM bridge. Qwen + # prefill uses DeepEP's expanded communication layout but consumes it via + # the existing m_indices/contiguous-GEMM adapter; this is the split that + # existed before layout selection was changed to decode-vs-prefill in + # 941b17a9d. Treating "expanded prefill" as "masked decode" allocates a + # cap-sized slab and is not the Qwen prefill contract. + use_masked = use_expand_layout and not get_is_extend_in_batch() + if ( + use_masked + and hidden_states.shape[0] + > self.masked_num_max_dispatch_tokens_per_rank + ): + raise ValueError( + "DeepEP v2 masked decode input exceeds the per-rank slab " + f"capacity {self.masked_num_max_dispatch_tokens_per_rank}, got " + f"{hidden_states.shape[0]}. Increase " + "SGLANG_DEEPEP_V2_MASKED_NUM_MAX_DISPATCH_TOKENS_PER_RANK." + ) + + self._dispatch_seq += 1 + trace_dispatch = os.environ.get("SGLANG_DEEPEP_V2_TRACE_DISPATCH") == "1" + if trace_dispatch: + logger.warning( + "DeepEP v2 dispatch enter: ep_rank=%s seq=%s tokens=%s " + "is_extend_in_batch=%s use_expand_layout=%s use_masked=%s", + self.rank, + self._dispatch_seq, + hidden_states.shape[0], + get_is_extend_in_batch(), + use_expand_layout, + use_masked, + ) + + # ElasticBuffer requires >=1 token per rank on the non-masked (contiguous / + # extend) path: DeepEP's own ElasticBuffer test pads every rank to + # `max(1, num_tokens)` (tests/elastic/test_ep.py). An idle DP rank with 0 + # tokens never fires the dispatch notify / scale-up-reduction warps, so no + # rank's recv count becomes "ready" and the do_cpu_sync CPU readback times + # out ("Dispatch CPU wait", buffer.hpp:1032). Pad an empty local batch to a + # single dummy token (routed to local expert 0); the contiguous slice in + # dispatch_b yields 0 real rows and combine_b drops it back to an empty + # output. The masked decode path tolerates empty (do_cpu_sync=False), so it + # is left untouched. + # ElasticBuffer's expanded layout handles an empty sender even with + # do_cpu_sync=True (validated on DEP16). Padding is required only for + # the non-expanded exact-count path. Keying this off `not use_masked` + # incorrectly pads expanded-prefill ranks and routes every dummy to + # global experts 0..topk-1, creating a large artificial hotspot on EP0. + self._pad_empty_combine = ( + not use_expand_layout and hidden_states.shape[0] == 0 + ) + if self._pad_empty_combine: + empty_pad_tokens = int( + os.environ.get("SGLANG_DEEPEP_V2_EMPTY_PAD_TOKENS", "1") + ) + if empty_pad_tokens < 1: + raise ValueError( + "SGLANG_DEEPEP_V2_EMPTY_PAD_TOKENS must be at least 1" + ) + hidden_states = hidden_states.new_zeros( + (empty_pad_tokens, hidden_states.shape[-1]) + ) + # A token's top-k experts must be DISTINCT valid ids: duplicates (e.g. + # all-zero -> expert 0 repeated) fault the dispatch kernel. Route the + # dummy to experts [0, 1, ..., topk-1] with zero weights so it + # contributes nothing even before combine_b slices it off. + topk_ids = torch.arange( + topk_ids.shape[-1], dtype=topk_ids.dtype, device=topk_ids.device + ).unsqueeze(0).expand(empty_pad_tokens, -1).contiguous() + topk_weights = topk_weights.new_zeros( + (empty_pad_tokens, topk_weights.shape[-1]) + ) + if trace_dispatch: + logger.warning( + "DeepEP v2 padded empty dispatch: ep_rank=%s seq=%s " + "dispatch_tokens=%s", + self.rank, + self._dispatch_seq, + empty_pad_tokens, + ) + + # Deterministic mixed-route reproducer. Some production warmups leave + # a subset of DP ranks idle; the adapter pads those ranks with zero + # tokens routed to experts [0..topk). Scheduler timing makes the idle + # subset nondeterministic, so this diagnostic can reproduce the same + # payload shape on a fixed suffix of ranks while keeping all collective + # arguments and tensor shapes identical. + dummy_rank_from = int( + os.environ.get("SGLANG_DEEPEP_V2_DUMMY_RANK_FROM", "-1") + ) + if ( + self._dispatch_seq == 1 + and dummy_rank_from >= 0 + and self.rank >= dummy_rank_from + ): + hidden_states = torch.zeros_like(hidden_states) + topk_ids = ( + torch.arange( + topk_ids.shape[-1], + dtype=topk_ids.dtype, + device=topk_ids.device, + ) + .unsqueeze(0) + .expand(hidden_states.shape[0], -1) + .contiguous() + ) + topk_weights = torch.zeros_like(topk_weights) + if trace_dispatch: + logger.warning( + "DeepEP v2 forced dummy route: ep_rank=%s seq=%s tokens=%s", + self.rank, + self._dispatch_seq, + hidden_states.shape[0], + ) + elif ( + self._dispatch_seq == 1 + and dummy_rank_from >= 0 + and self.rank < dummy_rank_from + and self._pad_empty_combine + ): + # If scheduler timing happened to leave a designated "real" rank + # idle, replace its already-padded dummy payload with deterministic + # nonzero data and distributed valid routes. combine() will still + # slice the synthetic local output back to zero rows. + hidden_states = torch.ones_like(hidden_states) + token_offsets = torch.arange( + hidden_states.shape[0], device=topk_ids.device, dtype=topk_ids.dtype + ).unsqueeze(1) + expert_offsets = torch.arange( + topk_ids.shape[-1], device=topk_ids.device, dtype=topk_ids.dtype + ).unsqueeze(0) + topk_ids = ( + self.rank * self.router_topk + + token_offsets * self.router_topk + + expert_offsets + ) % self.num_experts + topk_weights = torch.full_like( + topk_weights, 1.0 / self.router_topk + ) + if trace_dispatch: + logger.warning( + "DeepEP v2 synthesized real route: ep_rank=%s seq=%s tokens=%s", + self.rank, + self._dispatch_seq, + hidden_states.shape[0], + ) + + if self._uses_fp8_dispatch_output(): + _ensure_fp8_quant_available() + if use_masked: + # Follow the hardware scale format (DEEPGEMM_SCALE_UE8M0 via + # capability.fp8_scale_ue8m0). Hopper (False): plain row-major + # fp32 scale, and _run_masked_gemm does its own e8m0/tma-major + # alignment. Blackwell (True): pre-quantize the activation + # against a col-major UE8M0 scale so it already matches the + # layout the masked GEMM consumes. + _ue8m0 = self.capability.fp8_scale_ue8m0 + dispatch_x = sglang_per_token_group_quant_fp8( + hidden_states, + _SCALE_BLOCK_SIZE, + column_major_scales=_ue8m0, + scale_tma_aligned=_ue8m0, + scale_ue8m0=_ue8m0, + ) + use_tma_aligned_col_major_sf = _ue8m0 + else: + # Diagnostic A/B: DeepEP's elastic unit test exercises the + # non-expanded FP8 path with ordinary row-major scales, whereas + # the GB300 SGLang path normally supplies packed UE8M0, + # TMA-aligned column-major scales. Keep overrides opt-in so + # every other dispatcher input remains identical. + force_rowmajor_fp8 = ( + os.environ.get("SGLANG_DEEPEP_V2_FORCE_ROWMAJOR_FP8") == "1" + ) + if force_rowmajor_fp8: + dispatch_x = sglang_per_token_group_quant_fp8( + hidden_states, + _SCALE_BLOCK_SIZE, + column_major_scales=False, + scale_tma_aligned=False, + scale_ue8m0=False, + ) + use_tma_aligned_col_major_sf = False + else: + dispatch_x = _quantize_for_deepep_v2_dispatch( + hidden_states, self.capability + ) + layout_ab = os.environ.get( + "SGLANG_DEEPEP_V2_FP8_LAYOUT_AB", "col_in_col_out" + ) + if layout_ab not in { + "col_in_col_out", + "col_in_row_out", + "row_in_col_out", + "row_in_row_out", + }: + raise ValueError( + "Invalid SGLANG_DEEPEP_V2_FP8_LAYOUT_AB=" + f"{layout_ab}" + ) + if layout_ab.startswith("row_in_"): + dispatch_x = (dispatch_x[0], dispatch_x[1].contiguous()) + use_tma_aligned_col_major_sf = layout_ab.endswith("_col_out") + if trace_dispatch: + logger.warning( + "DeepEP v2 FP8 layout A/B: ep_rank=%s seq=%s " + "mode=%s sf_shape=%s sf_stride=%s sf_dtype=%s " + "col_major_output=%s", + self.rank, + self._dispatch_seq, + layout_ab, + tuple(dispatch_x[1].shape), + dispatch_x[1].stride(), + dispatch_x[1].dtype, + use_tma_aligned_col_major_sf, + ) + else: + dispatch_x = hidden_states + use_tma_aligned_col_major_sf = False + + # num_max_tokens_per_rank is a COLLECTIVE dispatch arg (ElasticBuffer + # requires the same value on all ranks). Keep it at the fixed buffer cap + # (class-level, cross-rank-consistent), matching DeepEP LL which uses a + # fixed _num_max_dispatch_tokens_per_rank rather than a per-forward token + # count. Do NOT derive it from the local hidden_states.shape[0]: under + # ragged DP load (or TP attention) the ranks would disagree on this + # collective arg. (The masked slab max_m below is likewise fixed at + # cap * ep_group_size for the same cross-rank / overflow safety; only + # expected_m, a per-rank-local GEMM schedule hint, uses the actual batch.) + num_max_tokens = self.num_max_dispatch_tokens_per_rank + # Non-masked (hybrid / direct-extend) path reads exact per-expert recv + # counts on the CPU, so it must wait for the GPU to finish writing them + # (matches the DeepEP elastic test which passes do_cpu_sync=1). Leaving + # it None lets the CPU read zeros on multi-node (scaleup) dispatch. Only + # the masked decode path keeps do_cpu_sync=False for graph capturability. + do_cpu_sync_val = True + if use_masked: + do_cpu_sync_val = False + # Diagnostic for the Qwen prefill-expanded contract. The upstream + # ElasticBuffer test exercises do_expand with CPU synchronization by + # default, while the initial SGLang integration tied expanded layout to + # the graph-oriented decode path and therefore forced async metadata. + # Keep decode unchanged, but allow an eager/BCG-breakable prefill to use + # the synchronized expanded path so we can distinguish an unfinished + # dispatch epilogue from the downstream masked-slab kernel. + if ( + force_expand_prefill + and get_is_extend_in_batch() + and os.environ.get("SGLANG_DEEPEP_V2_EXPAND_PREFILL_CPU_SYNC") == "1" + ): + do_cpu_sync_val = True + + buffer = self._get_buffer() + _num_sms, _num_qps = self._resolve_num_sms_qps(buffer) + recv_x, recv_topk_idx, recv_topk_weights, handle, event = buffer.dispatch( + dispatch_x, + topk_idx=topk_ids, + topk_weights=topk_weights, + num_experts=self.num_experts, + num_max_tokens_per_rank=num_max_tokens, + expert_alignment=self.capability.expert_alignment, + num_sms=_num_sms, + num_qps=_num_qps, + use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf, + do_cpu_sync=do_cpu_sync_val, + do_expand=use_expand_layout, + ) + if trace_dispatch: + logger.warning( + "DeepEP v2 dispatch returned: ep_rank=%s seq=%s " + "use_expand_layout=%s do_cpu_sync=%s", + self.rank, + self._dispatch_seq, + use_expand_layout, + do_cpu_sync_val, + ) + self._handle = handle + local_tokens = hidden_states.shape[0] + # event.current_stream_wait() is a GPU stream dependency (not a CPU + # sync); the do_cpu_sync=False masked decode path stays CUDA-graph + # capturable. + if event.event is not None: + event.current_stream_wait() + + if os.environ.get("SGLANG_DEEPEP_V2_TRACE_PSUM") == "1": + trace_recv = recv_x[0] if isinstance(recv_x, tuple) else recv_x + trace_scale = recv_x[1] if isinstance(recv_x, tuple) else None + logger.warning( + "DeepEP v2 dispatch metadata: ep_rank=%s seq=%s recv_x_shape=%s " + "recv_x_stride=%s scale_shape=%s scale_stride=%s scale_dtype=%s " + "psum=%s", + self.rank, + self._dispatch_seq, + tuple(trace_recv.shape), + trace_recv.stride(), + None if trace_scale is None else tuple(trace_scale.shape), + None if trace_scale is None else trace_scale.stride(), + None if trace_scale is None else trace_scale.dtype, + handle.psum_num_recv_tokens_per_expert.detach().cpu().tolist(), + ) + + if isinstance(recv_x, tuple): + recv_hidden_states, recv_hidden_states_scale = recv_x + else: + recv_hidden_states = recv_x + recv_hidden_states_scale = None + + if use_expand_layout: + # Expanded layout already has one row per local expert slot. There is + # no recv_topk_idx tensor in this native layout; combine uses handle + # metadata and expects top-k weights to be applied before combine. + # Avoid exact-count CPU reads that are only needed by non-expanded + # slicing/scatter paths. + local_topk_ids = None + num_recv_tokens_per_expert = [] + else: + num_recv_tokens = int( + handle.psum_num_recv_tokens_per_scaleup_rank[-1].item() + ) + recv_topk_idx = recv_topk_idx[:num_recv_tokens] + recv_topk_weights = recv_topk_weights[:num_recv_tokens] + recv_hidden_states = recv_hidden_states[:num_recv_tokens] + if recv_hidden_states_scale is not None: + recv_hidden_states_scale = recv_hidden_states_scale[:num_recv_tokens] + + # Elastic dispatch epilogue already converts global expert ids to local + # expert ids and marks non-local choices as -1. Keep it on-GPU and avoid + # an unnecessary max().item() synchronization in the decode path. + local_topk_ids = recv_topk_idx + num_recv_tokens_per_expert = list(handle.num_recv_tokens_per_expert_list) + + expected_m = 0 + masked_max_m = 0 + total_expanded = 0 + if use_masked: + # expected_m: average tokens-per-expert across the EP group, a + # per-rank-local schedule hint for the masked GEMM (NOT a hard bound; + # the real per-expert bound is masked_m on the GPU). Derive it from + # the actual local batch * EP group size, matching DeepEP LL + # (deepep.py dispatch_a uses hidden_states.shape[0]). Per-rank-local, + # so the actual batch is safe here even under ragged DP. group size + # == ep world size == num_experts // num_local_experts. + ep_group_size = max(1, self.num_experts // self.num_local_experts) + expected_m = max( + 1, + (local_tokens * ep_group_size * self.router_topk + self.num_experts) + // self.num_experts, + ) + # Size the masked slab to the FIXED worst case cap * ep_group_size, + # matching DeepEP LL's fixed buffer. A local expert receives the sum + # over all ranks of the tokens routed to it; each rank sends at most + # `cap` tokens (enforced by the dispatch-entry assert), so the count + # is bounded by cap * ep_group_size regardless of DP padding mode + # (MAX_LEN / SUM_LEN / skewed). Using the local batch for the slab + # would be unsafe: under skewed SUM_LEN decode another rank's larger + # batch could overflow this rank's slab. + masked_max_m = ( + self.masked_num_max_dispatch_tokens_per_rank * ep_group_size + ) + total_expanded = recv_hidden_states.shape[0] + + return DeepEPv2DispatchOutput( + recv_hidden_states, + recv_hidden_states_scale, + local_topk_ids, + recv_topk_weights, + num_recv_tokens_per_expert, + handle.psum_num_recv_tokens_per_expert, + use_expand_layout, + use_tma_aligned_col_major_sf, + use_masked, + expected_m, + masked_max_m, + total_expanded, + self.capability.expert_alignment, + ) + + def combine(self, combine_input: DeepEPv2CombineInput) -> torch.Tensor: + # Guard-first (before any DeepEP work) so misuse is reportable without + # DeepEP installed. + if self._handle is None: + raise RuntimeError( + "DeepEP v2 combine called without a valid dispatch handle" + ) + # The handle is single-use: release it whether combine succeeds or + # raises, so a failed step cannot poison the next dispatch. + try: + buffer = self._get_buffer() + _num_sms, _num_qps = self._resolve_num_sms_qps(buffer) + trace_contig = os.environ.get("SGLANG_DEEPEP_V2_TRACE_CONTIG") == "1" + if trace_contig: + torch.cuda.synchronize() + logger.warning( + "DeepEP v2 combine enter: ep_rank=%s seq=%s hidden=%s", + self.rank, + self._dispatch_seq, + tuple(combine_input.hidden_states.shape), + ) + combined_x, _, event = buffer.combine( + combine_input.hidden_states, + handle=self._handle, + topk_weights=combine_input.topk_weights, + num_sms=_num_sms, + num_qps=_num_qps, + ) + # Stream dependency, not a CPU sync (graph-safe). + if event.event is not None: + event.current_stream_wait() + if trace_contig: + torch.cuda.synchronize() + logger.warning( + "DeepEP v2 combine returned: ep_rank=%s seq=%s output=%s " + "stride=%s dtype=%s contiguous=%s", + self.rank, + self._dispatch_seq, + tuple(combined_x.shape), + combined_x.stride(), + combined_x.dtype, + combined_x.is_contiguous(), + ) + if os.environ.get("SGLANG_DEEPEP_V2_CLONE_COMBINE_OUTPUT") == "1": + if trace_contig: + logger.warning( + "DeepEP v2 combine clone enter: ep_rank=%s seq=%s", + self.rank, + self._dispatch_seq, + ) + combined_x = combined_x.clone() + torch.cuda.synchronize() + if trace_contig: + logger.warning( + "DeepEP v2 combine clone returned: ep_rank=%s seq=%s", + self.rank, + self._dispatch_seq, + ) + if self._pad_empty_combine: + # Drop the dummy token padded onto an empty local batch in + # dispatch so this idle rank's combined output is empty again. + combined_x = combined_x[:0] + return combined_x + finally: + self._pad_empty_combine = False + self._destroy_handle() + + +class DeepEPv2Dispatcher(BaseDispatcher): + def __init__( + self, + group: dist.ProcessGroup, + router_topk: int, + num_experts: int, + num_local_experts: int, + hidden_size: int, + params_dtype: torch.dtype, + allow_hybrid_mode: Optional[bool] = None, + ): + super().__init__() + if params_dtype != torch.bfloat16: + raise NotImplementedError( + "DeepEP v2 dispatch adapter currently expects BF16 model activations, " + f"got {params_dtype}" + ) + capability = get_deepep_v2_runner_capability(self) + self.output_dtype = capability.output_dtype + self.num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + self._impl = _DeepEPv2Impl( + group=group, + router_topk=router_topk, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + capability=capability, + num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, + allow_hybrid_mode=allow_hybrid_mode, + ) + + def set_quant_config(self, quant_config: dict) -> None: + self.quant_config = quant_config + capability = get_deepep_v2_runner_capability(self) + self.output_dtype = capability.output_dtype + self._impl.set_runner_capability(capability) + + # This backend intentionally exposes only single-shot dispatch()/combine(): + # TBO/SBO are rejected at server start, and our overlap PoC showed the naive + # two-phase split cannot overlap anyway (ElasticBuffer.dispatch is + # host-blocking); a split API will land together with real TBO support. + def dispatch( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> DispatchOutput: + return self._impl.dispatch(hidden_states, topk_output) + + def combine(self, combine_input: CombineInput) -> torch.Tensor: + if combine_input.format != CombineInputFormat.DEEPEP_V2: + raise TypeError( + f"Expected DeepEP v2 combine input, got {combine_input.format}" + ) + return self._impl.combine(combine_input) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 35759b9199b3..048a3dda0bc8 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -9,6 +9,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( get_dp_global_num_tokens, + get_is_extend_in_batch, is_dp_attention_enabled, ) from sglang.srt.layers.moe.token_dispatcher import ( @@ -21,13 +22,18 @@ from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import ( TorchDistributedCommBackend, ) +from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardCombineInput, + StandardDispatcher, + StandardDispatchOutput, +) from sglang.srt.layers.moe.topk import ( StandardTopKOutput, TopKOutput, TopKOutputChecker, ) from sglang.srt.layers.moe.utils import get_moe_runner_backend -from sglang.srt.runtime_context import get_schedule, get_spec +from sglang.srt.runtime_context import get_flags, get_parallel, get_schedule, get_spec from sglang.srt.speculative.spec_info import SpeculativeAlgorithm try: @@ -46,6 +52,45 @@ MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get() +# FlashInfer caches MoeAlltoAll's MNNVL allocation by workspace size. A tiny, +# aligned tail padding gives concurrently live target/draft and decode/prefill +# paths distinct cache keys without changing the usable token geometry. MNNVL +# rounds allocations to its own mapping granularity, so this does not add any +# dispatch/combine work; it only leases a separate persistent workspace. +_WORKSPACE_NAMESPACE_ALIGNMENT = 128 + + +def _max_tokens_per_scattered_source( + dp_global_num_tokens: list[int], attn_tp_size: int +) -> int: + """Return the largest token shard owned by one physical EP source rank.""" + + assert attn_tp_size > 0 + max_dp_tokens = max(dp_global_num_tokens) + return (max_dp_tokens + attn_tp_size - 1) // attn_tp_size + + +def _scattered_source_token_counts( + dp_global_num_tokens: list[int], attn_tp_size: int +) -> list[int]: + """Expand DP token counts into the physical source-rank tensor splits.""" + + assert attn_tp_size > 0 + counts = [] + for num_tokens in dp_global_num_tokens: + base, remainder = divmod(num_tokens, attn_tp_size) + counts.extend( + base + int(attn_tp_rank < remainder) for attn_tp_rank in range(attn_tp_size) + ) + return counts + + +def _workspace_size_for_namespace(workspace_size: int, *, speculative: bool) -> int: + """Return distinct FlashInfer cache keys for target and draft decode.""" + + slot = int(speculative) + return workspace_size + slot * _WORKSPACE_NAMESPACE_ALIGNMENT + class FlashinferDispatchOutput(NamedTuple): """Flashinfer EP dispatch output.""" @@ -88,6 +133,7 @@ def __init__( num_local_experts: int = None, # Unused hidden_size: int = None, params_dtype: torch.dtype = None, # Unused + moe_runner_config=None, ): super().__init__() if not use_flashinfer: @@ -102,13 +148,28 @@ def __init__( self.hidden_size = hidden_size self.num_experts = num_experts self.num_local_experts = num_local_experts + runner_backend = get_moe_runner_backend() self.invalid_token_expert_id = ( -1 - if get_moe_runner_backend().is_flashinfer_trtllm_routed() + if ( + runner_backend.is_deep_gemm() + or runner_backend.is_flashinfer_trtllm() + or runner_backend.is_flashinfer_trtllm_routed() + ) else self.num_experts ) # TODO: Can other moe runners use payload_in_workspace too? self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass() + if moe_runner_config is None: + from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig + + moe_runner_config = MoeRunnerConfig( + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + top_k=router_topk, + ) + self.prefill_dispatcher = StandardDispatcher(moe_runner_config) # FlashInfer sizes the workspace from the maximum dispatched tokens per # EP rank. See FlashInfer's moe_a2a_get_workspace_size_per_rank(), @@ -167,19 +228,117 @@ def __init__( pp_size=1, cp_size=1, ) - self.moe_a2a = MoeAlltoAll( - mapping=self.mapping, - max_num_tokens=self.max_num_tokens, - top_k=self.router_topk, - num_experts=self.num_experts, - workspace_size_per_rank=self.workspace_size, - mnnvl_config=MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)), + mnnvl_config = MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)) + is_speculative_model = get_flags().moe.speculative_context + + def make_moe_a2a() -> MoeAlltoAll: + # Target and draft decode graphs can coexist. Prefill and mixed + # extend use AG+RS below and therefore do not lease an MNNVL A2A + # workspace at all. + workspace_size = _workspace_size_for_namespace( + self.workspace_size, + speculative=is_speculative_model, + ) + return MoeAlltoAll( + mapping=self.mapping, + max_num_tokens=self.max_num_tokens, + top_k=self.router_topk, + num_experts=self.num_experts, + workspace_size_per_rank=workspace_size, + mnnvl_config=mnnvl_config, + ) + + self.moe_a2a = make_moe_a2a() + + def set_quant_config(self, quant_config: dict) -> None: + super().set_quant_config(quant_config) + self.prefill_dispatcher.set_quant_config(quant_config) + + def _dispatch_prefill_allgather( + self, hidden_states: torch.Tensor, topk_output: TopKOutput + ) -> StandardDispatchOutput: + """Use exact-size BF16 AG for prefill and mixed extend. + + FlashInfer one-sided A2A owns the pure-decode CUDA Graph fast path. + Eager extend can overlap a preceding graph on another stream, and the + one-sided transport's signal state is not safe to reuse across those + streams. A conventional all-gatherv here avoids that communicator + race without adding synchronization to decode. + """ + + if hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer WideEP prefill AG requires BF16 hidden states, got " + f"{hidden_states.dtype}." + ) + if TopKOutputChecker.format_is_bypassed(topk_output): + topk_output = topk_output.to_standard() + if not TopKOutputChecker.format_is_standard(topk_output): + raise TypeError( + "FlashInfer WideEP prefill AG requires materialized top-k " + f"routing, got {type(topk_output).__name__}." + ) + + dp_global = get_dp_global_num_tokens() + if dp_global is None: + source_sizes = [hidden_states.shape[0]] * self.ep_size + else: + source_sizes = _scattered_source_token_counts( + dp_global, get_parallel().attn_tp_size + ) + if len(source_sizes) != self.ep_size: + raise RuntimeError( + "FlashInfer WideEP prefill AG source geometry does not match " + f"EP: len(source_sizes)={len(source_sizes)}, ep_size={self.ep_size}." + ) + if source_sizes[self.ep_rank] != hidden_states.shape[0]: + raise RuntimeError( + "FlashInfer WideEP prefill AG local source geometry mismatch: " + f"source_sizes[{self.ep_rank}]={source_sizes[self.ep_rank]} != " + f"hidden_states.shape[0]={hidden_states.shape[0]}." + ) + + topk_ids = topk_output.topk_ids.to(torch.int32) + hidden_states, topk_ids, topk_weights = get_parallel().tp_group.all_gatherv( + [hidden_states, topk_ids, topk_output.topk_weights], + sizes=source_sizes, + ) + self.prefill_source_sizes = source_sizes + return self.prefill_dispatcher.dispatch( + hidden_states, + StandardTopKOutput(topk_weights, topk_ids, topk_output.router_logits), ) @debug_kernel_api def dispatch( self, hidden_states: torch.Tensor, topk_output: TopKOutput - ) -> FlashinferDispatchOutput: + ) -> FlashinferDispatchOutput | StandardDispatchOutput: + if get_is_extend_in_batch(): + return self._dispatch_prefill_allgather(hidden_states, topk_output) + self.active_moe_a2a = self.moe_a2a + # Block-wise FP8 expert runners quantize activations locally immediately + # before GEMM. Keep the one-sided dispatch wire format in BF16 so both + # TRT-LLM Gen MoE and DeepGEMM consume the same lossless payload and the + # combine reduction also stays BF16. FP4 dispatch is intentionally + # excluded: it advertises input_global_scale and uses its existing + # packed payload path below. + runner_backend = get_moe_runner_backend() + weight_dtype = self.quant_config.get("weight_dtype") + uses_bf16_fp8_payload = weight_dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) and ( + runner_backend.is_deep_gemm() + or runner_backend.is_flashinfer_trtllm() + or runner_backend.is_flashinfer_trtllm_routed() + ) + if uses_bf16_fp8_payload and hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer A2A with an FP8 DeepGEMM/TRT-LLM Gen MoE runner " + "requires BF16 dispatch and combine payloads, but received " + f"{hidden_states.dtype}." + ) + output_dtype = hidden_states.dtype x = hidden_states x_sf = None @@ -218,16 +377,18 @@ def dispatch( # CUDA-graph *capture*; on *replay* dispatch() is not re-executed and the # value baked at capture is reused. Two cases, both rank-invariant: # - # Case 1 — max(dp_global): DP attention feeding EP. The scheduler - # all-gathers per-DP-rank token counts into dp_global (length dp_size, - # identical on every rank), which differ across ranks, so we must take - # the max. FlashInfer A2A forces require_mlp_tp_gather=True (see - # require_mlp_tp_gather()), so: eager reads the live list; capture sees - # [num_tokens] * dp_size (uniform capture bs) and bakes max() == the - # bucket; replay reuses that baked value and every rank replays the same - # bucket because the decode graph runner sizes it from the cross-rank - # max. Without this, per-rank buckets could diverge -> geometry mismatch - # -> illegal memory access (issue #30242). + # Case 1 — DP attention feeding EP. The scheduler all-gathers the token + # count for every DP replica into dp_global (identical on every rank). + # Before MoE, LayerCommunicator token-scatters each replica across its + # attention-TP group, so a physical EP source rank owns at most + # ceil(max(dp_global) / attn_tp_size) tokens. Using max(dp_global) + # directly is safe but can create attn_tp_size times too much padding + # and MoE work for DPxTP layouts. FlashInfer A2A forces + # require_mlp_tp_gather=True (see require_mlp_tp_gather()), so eager + # reads the live list; capture sees a uniform bucket list; replay uses + # the same baked geometry on all ranks. Without the cross-rank max, + # ranks could replay different-sized graphs -> geometry mismatch -> + # illegal memory access (issue #30242). # # Case 2 — x.shape[0]: no per-rank DP list (dp_global absent or scalar). # This is SP attention feeding EP (tokens are sequence-parallel scattered @@ -237,7 +398,10 @@ def dispatch( dp_global = get_dp_global_num_tokens() if dp_global is not None and len(dp_global) > 1: # Case 1 - self.runtime_max_tokens_per_rank = max(dp_global) + attn_tp_size = get_parallel().attn_tp_size + self.runtime_max_tokens_per_rank = _max_tokens_per_scattered_source( + dp_global, attn_tp_size + ) else: # Case 2. Guard against the #30242 failure mode: DP attention must # never land here with ep_size > 1, because there x.shape[0] differs @@ -252,9 +416,22 @@ def dispatch( ) self.runtime_max_tokens_per_rank = x.shape[0] + # MoeAlltoAll's workspace is allocated once from max_num_tokens. Passing + # a larger runtime geometry does not resize it and can otherwise turn a + # mixed prefill/speculative batch into a delayed illegal memory access. + # Fail at the dispatch boundary with the exact required capacity. + assert self.runtime_max_tokens_per_rank <= self.max_num_tokens, ( + "FlashInfer A2A runtime token geometry exceeds its fixed workspace: " + f"runtime_max_tokens_per_rank={self.runtime_max_tokens_per_rank} > " + f"max_num_tokens={self.max_num_tokens}. Increase " + "SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK to cover the " + "largest mixed prefill and speculative-verify batch." + ) + # The recv buffer reserves runtime_max_tokens_per_rank slots for THIS # rank, so it must cover this rank's own tokens. This holds in both cases - # (Case 1: max(dp_global) >= the local count; Case 2: exactly x.shape[0]), + # (Case 1: ceil(max(dp_global) / attn_tp_size) covers every token-scatter + # shard; Case 2: exactly x.shape[0]), # so a violation signals a sizing/plumbing bug (e.g. an un-adjusted spec # count) rather than a benign case. assert self.runtime_max_tokens_per_rank >= x.shape[0], ( @@ -268,7 +445,7 @@ def dispatch( # padding slots whose expert_id would otherwise route to a real expert # and waste downstream MoE compute. Sanitizing the padding to a # sentinel id is structural, not optional. - recv_tensors = self.moe_a2a.dispatch( + recv_tensors = self.active_moe_a2a.dispatch( topk_ids, payloads, self.runtime_max_tokens_per_rank, @@ -290,7 +467,7 @@ def dispatch( # Provide an output tensor to fused_moe so it writes directly to our buffer moe_output = None if self.payload_in_workspace: - moe_output = self.moe_a2a.get_combine_payload_tensor_in_workspace( + moe_output = self.active_moe_a2a.get_combine_payload_tensor_in_workspace( self.runtime_max_tokens_per_rank, self.hidden_size, output_dtype ).view(-1, self.hidden_size) return FlashinferDispatchOutput( @@ -301,10 +478,40 @@ def dispatch( ) @debug_kernel_api - def combine(self, combine_input: FlashinferCombineInput) -> torch.Tensor: + def combine( + self, combine_input: FlashinferCombineInput | StandardCombineInput + ) -> torch.Tensor: hidden_states = combine_input.hidden_states + if combine_input.format == CombineInputFormat.STANDARD: + if hidden_states.dtype != torch.bfloat16: + raise TypeError( + "FlashInfer WideEP prefill RS requires BF16 expert output, " + f"got {hidden_states.dtype}." + ) + source_sizes = self.prefill_source_sizes + hidden_states = get_parallel().tp_group.reduce_scatterv( + hidden_states, sizes=source_sizes + ) + del self.prefill_source_sizes + return hidden_states + + weight_dtype = self.quant_config.get("weight_dtype") + runner_backend = get_moe_runner_backend() + if ( + weight_dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + and ( + runner_backend.is_deep_gemm() + or runner_backend.is_flashinfer_trtllm() + or runner_backend.is_flashinfer_trtllm_routed() + ) + and hidden_states.dtype != torch.bfloat16 + ): + raise TypeError( + "FlashInfer A2A FP8 MoE combine payload must be BF16, but " + f"received {hidden_states.dtype}." + ) output_hidden_size = hidden_states.shape[-1] - hidden_states = self.moe_a2a.combine( + hidden_states = self.active_moe_a2a.combine( hidden_states.view( self.ep_size, self.runtime_max_tokens_per_rank, output_hidden_size ), @@ -313,4 +520,5 @@ def combine(self, combine_input: FlashinferCombineInput) -> torch.Tensor: ) del self.runtime_max_tokens_per_rank + del self.active_moe_a2a return hidden_states diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 9ecaec671404..eb1bde93e026 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -4,7 +4,7 @@ import os from contextlib import contextmanager from enum import Enum, IntEnum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import torch @@ -12,7 +12,13 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) -from sglang.srt.runtime_context import get_exec, get_flags, get_forward, get_parallel +from sglang.srt.runtime_context import ( + get_exec, + get_flags, + get_forward, + get_parallel, + get_server_args, +) from sglang.srt.utils import is_cuda, is_npu _is_npu = is_npu() @@ -20,8 +26,6 @@ if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs -from sglang.srt.runtime_context import get_server_args - logger = logging.getLogger(__name__) @@ -37,6 +41,7 @@ class MoeA2ABackend(Enum): FLASHINFER = "flashinfer" MEGAMOE = "megamoe" PPLX = "pplx" + DEEPEP_V2 = "deepep_v2" CUSTOMIZED = "customized" @classmethod @@ -78,6 +83,9 @@ def is_megamoe(self): def is_pplx(self): return self == MoeA2ABackend.PPLX + def is_deepep_v2(self): + return self == MoeA2ABackend.DEEPEP_V2 + def is_customized(self): return self == MoeA2ABackend.CUSTOMIZED @@ -174,6 +182,35 @@ def is_aiter(self): return self == MoeRunnerBackend.AITER +class DeepEPv2OutputDtype(Enum): + """ + Describes the dispatch output data type for DeepEP v2. + + - BF16: dispatch hidden states in bf16, without activation scales. + - FP8: dispatch hidden states in fp8, with activation scales. + """ + + BF16 = "bf16" + FP8 = "fp8" + + +class DeepEPv2RunnerCapability(NamedTuple): + """ + Describes the DeepEP v2 dispatcher contract required by the active MoE runner. + + This capability is resolved once (in get_deepep_v2_runner_capability, which reads + runner-side flags such as DeepGEMM JIT TMA/UE8M0 settings) and then consumed + by the dispatcher. The dispatcher depends only on this resolved contract and + does not peek at runner implementation details itself. + """ + + output_dtype: DeepEPv2OutputDtype + expert_alignment: int + fp8_scale_tma_aligned: bool = False + fp8_scale_ue8m0: bool = False + use_expanded_layout: bool = False + + class DeepEPMode(Enum): NORMAL = "normal" @@ -295,6 +332,74 @@ def get_ascend_dispatcher_output_dtype(dispatcher): return DispatcherOutputDtype.BF16 +def get_deepep_v2_output_dtype(self) -> DeepEPv2OutputDtype: + """ + Automatically choose the dispatch output dtype for DeepEP v2. + + The decision follows several checks in priority order: + 0. Parse server argument. + 1. Parse quant config. + 2. DeepGEMM expects FP8 activation + scales. + 3. Triton consumes BF16 activation without dispatcher-provided scales. + """ + + server_args = get_server_args() + if server_args and server_args.deepep_v2_dispatcher_output_dtype != "auto": + return DeepEPv2OutputDtype(server_args.deepep_v2_dispatcher_output_dtype) + + if self.quant_config is not None: + dispatcher_output_dtype = self.quant_config.get("dispatcher_output_dtype", None) + if dispatcher_output_dtype is not None: + return DeepEPv2OutputDtype(dispatcher_output_dtype) + + runner_backend = get_moe_runner_backend() + if runner_backend.is_deep_gemm(): + return DeepEPv2OutputDtype.FP8 + if runner_backend.is_triton(): + return DeepEPv2OutputDtype.BF16 + + raise ValueError( + "DeepEP v2 auto dispatcher output dtype only supports deep_gemm and triton " + f"runner backends for now, got {runner_backend.value}. Set " + "--deepep-v2-dispatcher-output-dtype explicitly only after adding a matching " + "DeepEP v2 runner adapter." + ) + + +def get_deepep_v2_runner_capability(self) -> DeepEPv2RunnerCapability: + output_dtype = get_deepep_v2_output_dtype(self) + runner_backend = get_moe_runner_backend() + if output_dtype == DeepEPv2OutputDtype.FP8: + if not runner_backend.is_deep_gemm(): + raise ValueError( + "DeepEP v2 FP8 dispatch output currently requires " + "--moe-runner-backend deep_gemm because the adapter must consume " + f"activation scales. Got {runner_backend.value}." + ) + from sglang.srt.layers import deep_gemm_wrapper + + # DeepGEMM consumes expert-major grouped activations. Use DeepEP v2's + # native expanded layout so the dispatcher copy epilogue writes one + # row per local expert slot, avoiding an extra SGLang scatter/gather + # adapter round-trip on the decode path. + return DeepEPv2RunnerCapability( + output_dtype=output_dtype, + expert_alignment=128, + fp8_scale_tma_aligned=( + deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES + or deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + ), + fp8_scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, + use_expanded_layout=True, + ) + if not runner_backend.is_triton(): + raise ValueError( + "DeepEP v2 BF16 dispatch output currently requires " + f"--moe-runner-backend triton. Got {runner_backend.value}." + ) + return DeepEPv2RunnerCapability(output_dtype=output_dtype, expert_alignment=1) + + def initialize_moe_config(server_args: ServerArgs): moe = get_flags().moe moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend) @@ -545,14 +650,17 @@ def speculative_moe_a2a_backend_context(): moe = get_flags().moe original_backend = moe.a2a_backend original_disable_fp4_allgather = moe.disable_fp4_allgather + original_speculative_context = moe.speculative_context try: moe.a2a_backend = get_speculative_moe_a2a_backend() # Disable FP4 allgather for spec decode since MTP layers are unquantized moe.disable_fp4_allgather = True + moe.speculative_context = True yield finally: moe.a2a_backend = original_backend moe.disable_fp4_allgather = original_disable_fp4_allgather + moe.speculative_context = original_speculative_context # The type of method in top-K routing, for use in torch custom op diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index fe78b4fccdd7..f2cf5f28e8c5 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -68,6 +68,7 @@ class Bf16GemmBackend(Enum): AUTO = "auto" CUTEDSL = "cutedsl" + FLASHINFER_PR4266 = "flashinfer_pr4266" TORCH = "torch" def is_auto(self) -> bool: @@ -76,14 +77,85 @@ def is_auto(self) -> bool: def is_cutedsl(self) -> bool: return self == Bf16GemmBackend.CUTEDSL + def is_flashinfer_pr4266(self) -> bool: + return self == Bf16GemmBackend.FLASHINFER_PR4266 + + def is_optimized(self) -> bool: + return self.is_cutedsl() or self.is_flashinfer_pr4266() + _BF16_GEMM_BACKEND: Optional[Bf16GemmBackend] = None _cutedsl_bf16_gemm = None _use_cutedsl_bf16_gemm = None +_flashinfer_pr4266_splitk_tactic = None +_flashinfer_pr4266_run_splitk_dense = None +_flashinfer_pr4266_direct_default_tactic = None +_flashinfer_pr4266_prefer_direct = None +_flashinfer_pr4266_run_direct_dense = None +_enable_bf16_splitk_gemm = False + +_FLASHINFER_SPLITK_GEMM_HINT = ( + "The BF16 Split-K GEMM path needs the direct dense kernel added by " + "FlashInfer PR #4266. Reinstall a newer FlashInfer, or set " + "SGLANG_ENABLE_BF16_SPLITK_GEMM=0 to disable this path." +) + + +# Oakhaven-Max TP16 tactics measured on GB300 under CUDA graph replay with PDL +# and L2-defeating weight rotation. Every entry passed the strict correctness +# gate and beat SGLang's existing dispatch by at least 1.26x. Unlisted M/N/K, +# including M=64, retain SGLang's existing TGV/cuBLAS path. +_FLASHINFER_PR4266_TUNED_TACTICS = { + (1, 256, 8192): (64, 8, 4, 11), + (2, 256, 8192): (64, 8, 4, 11), + (4, 256, 8192): (64, 8, 4, 11), + (8, 256, 8192): (64, 8, 4, 11), + (16, 256, 8192): (64, 8, 4, 10), + (24, 256, 8192): (64, 8, 4, 11), + (32, 256, 8192): (64, 8, 4, 12), + (1, 512, 8192): (64, 8, 4, 11), + (2, 512, 8192): (64, 8, 4, 12), + (4, 512, 8192): (64, 8, 4, 10), + (8, 512, 8192): (64, 8, 4, 12), + (16, 512, 8192): (64, 8, 4, 12), + (24, 512, 8192): (64, 8, 4, 12), + (32, 512, 8192): (64, 16, 4, 9), + (1, 2304, 8192): (128, 8, 4, 6), + (2, 2304, 8192): (64, 8, 2, 12), + (4, 2304, 8192): (128, 8, 4, 6), + (8, 2304, 8192): (64, 8, 4, 10), + (16, 2304, 8192): (64, 16, 4, 9), + (24, 2304, 8192): (64, 32, 2, 9), + (32, 2304, 8192): (64, 32, 2, 9), + (1, 2560, 8192): (64, 8, 2, 10), + (2, 2560, 8192): (64, 8, 2, 10), + (4, 2560, 8192): (64, 8, 2, 10), + (8, 2560, 8192): (64, 8, 2, 10), + (16, 2560, 8192): (64, 16, 2, 11), + (24, 2560, 8192): (64, 32, 2, 9), + (32, 2560, 8192): (64, 32, 2, 9), +} + + +def use_flashinfer_pr4266_bf16_gemm(m: int, n: int, k: int) -> bool: + """Return whether the PR #4266 low-M kernel is selected for this shape.""" + return (m, n, k) in _FLASHINFER_PR4266_TUNED_TACTICS + + +def should_enable_bf16_splitk_gemm(backend: Bf16GemmBackend) -> bool: + """Return whether the optional Split-K path should be initialized.""" + return backend.is_optimized() and envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.get() def initialize_bf16_gemm_config(server_args: ServerArgs) -> None: - global _BF16_GEMM_BACKEND, _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm + global _BF16_GEMM_BACKEND + global _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm + global _flashinfer_pr4266_splitk_tactic + global _flashinfer_pr4266_run_splitk_dense + global _flashinfer_pr4266_direct_default_tactic + global _flashinfer_pr4266_prefer_direct + global _flashinfer_pr4266_run_direct_dense + global _enable_bf16_splitk_gemm from sglang.srt.utils import is_sm100_supported @@ -93,9 +165,12 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None: backend = Bf16GemmBackend(backend_str) - if backend.is_cutedsl(): + if backend.is_optimized(): if not is_sm100_supported(): - raise ValueError("--bf16-gemm-backend cutedsl requires an SM10x GPU") + raise ValueError( + f"--bf16-gemm-backend {backend.value} requires " + "SM100/SM103 (Blackwell)" + ) from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import ( cutedsl_bf16_gemm, @@ -105,6 +180,29 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None: _cutedsl_bf16_gemm = cutedsl_bf16_gemm _use_cutedsl_bf16_gemm = use_cutedsl_bf16_gemm + _enable_bf16_splitk_gemm = False + if should_enable_bf16_splitk_gemm(backend): + from sglang.kernels.ops.gemm.flashinfer_pr4266_dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + try: + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + default_tactic, + prefer_direct_bf16_gemm_sm100, + run_direct_dense, + ) + except ImportError as exc: + raise ImportError(_FLASHINFER_SPLITK_GEMM_HINT) from exc + + _flashinfer_pr4266_splitk_tactic = SplitKTactic + _flashinfer_pr4266_run_splitk_dense = run_splitk_dense + _flashinfer_pr4266_direct_default_tactic = default_tactic + _flashinfer_pr4266_prefer_direct = prefer_direct_bf16_gemm_sm100 + _flashinfer_pr4266_run_direct_dense = run_direct_dense + _enable_bf16_splitk_gemm = True + _BF16_GEMM_BACKEND = backend @@ -114,12 +212,40 @@ def _bf16_gemm_dispatch_fake( return x.new_empty((*x.shape[:-1], weight.shape[0])) -@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake) -def bf16_gemm_dispatch( +def _flashinfer_pr4266_bf16_gemm( x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] ) -> torch.Tensor: + x_2d = x.view(-1, x.shape[-1]) + out = torch.empty((x_2d.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device) + m, n, k = x_2d.shape[0], weight.shape[0], weight.shape[1] + if bias is None and _flashinfer_pr4266_prefer_direct(m, n, k): + tactic = _flashinfer_pr4266_direct_default_tactic(m, n, k) + _flashinfer_pr4266_run_direct_dense(x_2d, weight.T, out, True, tactic) + else: + tactic = _flashinfer_pr4266_splitk_tactic( + *_FLASHINFER_PR4266_TUNED_TACTICS[(m, n, k)] + ) + _flashinfer_pr4266_run_splitk_dense( + x_2d, + weight.T, + bias, + out, + True, + tactic, + ) + return out.view(*x.shape[:-1], weight.shape[0]) + + +def _bf16_gemm_dispatch_impl( + x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] +) -> torch.Tensor: + m = x.numel() // x.shape[-1] + if _enable_bf16_splitk_gemm and use_flashinfer_pr4266_bf16_gemm( + m, weight.shape[0], weight.shape[1] + ): + return _flashinfer_pr4266_bf16_gemm(x, weight, bias) if _use_cutedsl_bf16_gemm is not None and _use_cutedsl_bf16_gemm( - x.numel() // x.shape[-1], weight.shape[0], weight.shape[1] + m, weight.shape[0], weight.shape[1] ): return _cutedsl_bf16_gemm(x.view(-1, x.shape[-1]), weight, bias).view( *x.shape[:-1], -1 @@ -127,6 +253,13 @@ def bf16_gemm_dispatch( return F.linear(x, weight, bias) +@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake) +def bf16_gemm_dispatch( + x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] +) -> torch.Tensor: + return _bf16_gemm_dispatch_impl(x, weight, bias) + + def get_bf16_gemm_backend() -> Bf16GemmBackend: global _BF16_GEMM_BACKEND if _BF16_GEMM_BACKEND is None: @@ -225,7 +358,7 @@ def apply( return tgemm.mm(x, layer.weight, bias, otype=x.dtype) elif ( - get_bf16_gemm_backend().is_cutedsl() + get_bf16_gemm_backend().is_optimized() and x.is_cuda and x.dtype == torch.bfloat16 and layer.weight.dtype == torch.bfloat16 @@ -239,17 +372,7 @@ def apply( # opaque op resolves it at runtime with concrete shapes, # keeping the per-shape kernel choice. return bf16_gemm_dispatch(x, layer.weight, bias) - if _use_cutedsl_bf16_gemm( - x.numel() // x.shape[-1], - layer.weight.shape[0], - layer.weight.shape[1], - ): - x_shapes = x.shape - output = _cutedsl_bf16_gemm( - x.view(-1, x_shapes[-1]), layer.weight, bias - ) - return output.view(*x_shapes[:-1], -1) - return F.linear(x, layer.weight, bias) + return _bf16_gemm_dispatch_impl(x, layer.weight, bias) return F.linear(x, layer.weight, bias) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 2dfc00d7915f..594f6c5c5c67 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -83,10 +83,8 @@ def forward( a: torch.Tensor, b: torch.Tensor, ) -> torch.Tensor: - if ( - forward_batch.forward_mode.is_extend() - and get_tc_piecewise_forward_context() is not None - ): + is_extend = forward_batch.forward_mode.is_extend() + if is_extend and get_tc_piecewise_forward_context() is not None: # Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim) seq_len = mixed_qkv.shape[0] output = torch.empty( @@ -111,52 +109,127 @@ def forward( self.layer_id, ) return output - else: - return get_attn_backend().forward( - layer=self, - forward_batch=forward_batch, + + # Target verify rebuilds query_start_loc from the physical padded input, + # unlike ordinary extend where it retains the logical sequence ends. + should_trim_padded_extend = ( + is_extend and not forward_batch.forward_mode.is_target_verify() + ) + real_num_tokens = ( + getattr(forward_batch, "num_token_non_padded_cpu", None) + if should_trim_padded_extend + else None + ) + if real_num_tokens is not None and real_num_tokens < mixed_qkv.shape[0]: + # Eager DP/attention-TP synchronization may append physical token rows + # while extend_seq_lens/query_start_loc keep the logical sequence + # lengths. Varlen linear-attention kernels require their packed input + # length to agree with those logical lengths. Compute only the real + # prefix, then restore the physical shape expected by the following + # residual, MLP, and collective operations. + output = torch.empty( + (1, mixed_qkv.shape[0], self.num_v_heads, self.head_v_dim), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + _linear_attention_with_output_impl( mixed_qkv=mixed_qkv, a=a, b=b, + output=output, + attention_layer=self, + forward_batch=forward_batch, ) + return output + return get_attn_backend().forward( + layer=self, + forward_batch=forward_batch, + mixed_qkv=mixed_qkv, + a=a, + b=b, + ) -@register_custom_op(mutates_args=["output"]) -@register_split_op() -def unified_linear_attention_with_output( + +def _linear_attention_with_output_impl( mixed_qkv: torch.Tensor, a: torch.Tensor, b: torch.Tensor, output: torch.Tensor, - layer_id: int, + attention_layer: RadixLinearAttention, + forward_batch: ForwardBatch, ) -> None: - """ - Custom op wrapper for linear attention computation only. - """ - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layers = context.attention_layers - attention_layer = attention_layers[layer_id] - real_num_tokens = forward_batch.num_token_non_padded_cpu + """Run linear attention on the real prefix and initialize physical padding.""" + real_num_tokens = min(forward_batch.num_token_non_padded_cpu, mixed_qkv.shape[0]) original_out_cache_loc = forward_batch.out_cache_loc # Keep the original ForwardBatch object and only narrow cache locations for # this backend call so model/backend state is still written to the same batch. forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens] - - ret = get_attn_backend().forward( - layer=attention_layer, + logical_output = output[:, :real_num_tokens] + try: + ret = get_attn_backend().forward( + layer=attention_layer, + forward_batch=forward_batch, + mixed_qkv=mixed_qkv[:real_num_tokens], + a=a[:real_num_tokens], + b=b[:real_num_tokens], + linear_attn_output=logical_output, + ) + finally: + forward_batch.out_cache_loc = original_out_cache_loc + + # FlashInfer GDN can write directly into the physical output's logical + # prefix. Other backends return their own tensor and keep the copy fallback. + if ret.data_ptr() != logical_output.data_ptr(): + logical_output.copy_(ret) + # Physical padding participates in following residual, router, expert/MoE, + # and collective operations. Keep those inputs finite and deterministic. + output[:, real_num_tokens:].zero_() + + +def _unified_linear_attention_with_output_impl( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + output: torch.Tensor, + layer_id: int, +) -> None: + """Eager implementation kept separate for backend-independent tests.""" + context = get_tc_piecewise_forward_context() + forward_batch = context.forward_batch + attention_layers = context.attention_layers + attention_layer = attention_layers[layer_id] + _linear_attention_with_output_impl( + mixed_qkv=mixed_qkv, + a=a, + b=b, + output=output, + attention_layer=attention_layer, forward_batch=forward_batch, - mixed_qkv=mixed_qkv[:real_num_tokens], - a=a[:real_num_tokens], - b=b[:real_num_tokens], ) - forward_batch.out_cache_loc = original_out_cache_loc - - output[:, :real_num_tokens].copy_(ret) return +@register_custom_op(mutates_args=["output"]) +@register_split_op() +def unified_linear_attention_with_output( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + output: torch.Tensor, + layer_id: int, +) -> None: + """Custom op wrapper for linear attention computation only.""" + _unified_linear_attention_with_output_impl( + mixed_qkv=mixed_qkv, + a=a, + b=b, + output=output, + layer_id=layer_id, + ) + + bcg_unified_linear_attention_with_output = eager_on_graph(True)( unified_linear_attention_with_output ) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index dc9e5128d312..f5becd73e0f3 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -1,4 +1,5 @@ import logging +import os from functools import partial from typing import Callable, Dict, List, Optional, Tuple @@ -67,6 +68,19 @@ _BUILT_IN_SAMPLING_BACKENDS = {"flashinfer", "pytorch", "ascend"} +def _trace_e2e_sampler(stage: str, **fields) -> None: + """Opt-in sampler stage trace for DP-attention E2E diagnostics.""" + if os.getenv("SGLANG_TRACE_SAMPLER_E2E", "0") != "1": + return + try: + parallel = get_parallel() + rank = f"dp={parallel.attn_dp_rank} tp={parallel.tp_rank}" + except Exception: + rank = "rank=unknown" + details = " ".join(f"{key}={value}" for key, value in fields.items()) + print(f"SGLANG_TRACE_SAMPLER_E2E {rank} stage={stage} {details}", flush=True) + + class Sampler(nn.Module): def __init__(self): super().__init__() @@ -118,12 +132,20 @@ def forward( to get the unique seed for each position. """ logits = logits_output.next_token_logits + _trace_e2e_sampler( + "forward_enter", + logits_shape=tuple(logits.shape), + all_greedy=sampling_info.is_all_greedy, + ) # Preprocess logits (custom processors and NaN handling) + _trace_e2e_sampler("preprocess_enter") logits = self._preprocess_logits(logits, sampling_info) + _trace_e2e_sampler("preprocess_returned") return_sampling_mask = any(sampling_info.return_sampling_masks or []) if sampling_info.is_all_greedy: + _trace_e2e_sampler("greedy_enter") if _use_aiter and not _disable_aiter_greedy_sample: batch_next_token_ids = torch.empty( logits.shape[0], device=logits.device, dtype=torch.int32 @@ -131,6 +153,9 @@ def forward( _aiter_greedy_sample(batch_next_token_ids, logits) else: batch_next_token_ids = torch.argmax(logits, -1) + _trace_e2e_sampler( + "greedy_returned", output_shape=tuple(batch_next_token_ids.shape) + ) if return_sampling_mask: self._attach_greedy_sampling_mask_to_output( logits_output, sampling_info, batch_next_token_ids @@ -239,8 +264,11 @@ def forward( ) logprob_result.write_output_to(logits_output) + _trace_e2e_sampler("token_sync_enter") self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) + _trace_e2e_sampler("token_sync_returned") + _trace_e2e_sampler("forward_returned") return batch_next_token_ids def _sample_from_probs( diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 9e2f191aa2ff..89333c25e71d 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -302,7 +302,13 @@ def _lazy_init_forward_buf(self, payload: RelayPayload): # Spec extras are gated by spec_algo, not by the payload's shape, so a # non-spec stash allocates no extra bufs (only output_tokens_buf). - self.need_topk = self.spec_algo.is_some() and self.spec_algo.need_topk() + # Mirror need_hidden_states: a relay that carries no draft proposal (the PP + # ring rebuilds results without one) must not allocate the topk buffers. + self.need_topk = ( + self.spec_algo.is_some() + and self.spec_algo.need_topk() + and payload.topk_p is not None + ) self.need_hidden_states = ( self.spec_algo.is_some() and spec_need_hidden_states() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c8bdceb570ba..9950e8849707 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -283,7 +283,9 @@ from sglang.srt.session.session_controller import SessionController from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.dflash_utils import validate_dflash_request -from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec +from sglang.srt.speculative.eagle_utils import ( + get_draft_recurrent_hidden_state_spec_from_config, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( DynamicGradMode, @@ -1286,11 +1288,20 @@ def init_disaggregation(self): ) if self.spec_algorithm.carries_draft_hidden_states(): - # `draft_runner` aliases `draft_runner_list[0]` in the multi-layer - # worker, so a single accessor covers both shapes. - draft_runner = self.draft_worker.draft_worker.draft_runner + # Derive from the draft config, not the draft runner: the runner does + # not exist on ranks that do not host the draft (prefill-side PP builds + # it only on the last stage), and the PD metadata wire schema has to be + # identical on every rank. + draft_model_config = ModelConfig.from_server_args( + self.server_args, + model_path=self.server_args.speculative_draft_model_path, + model_revision=self.server_args.speculative_draft_model_revision, + is_draft_model=True, + ) disagg_hidden_size, disagg_hidden_states_dtype = ( - get_draft_recurrent_hidden_state_spec(draft_runner) + get_draft_recurrent_hidden_state_spec_from_config( + draft_model_config, self.spec_algorithm + ) ) else: disagg_hidden_size = 16 # minimal padding size for RDMA @@ -3671,7 +3682,9 @@ def run_batch( # future_map relay / on_publish). resolve_forward_inputs(batch, self.future_map) with self._forward_isolation(batch, overlap=False): - batch_result = self.model_worker.forward_batch_generation(batch) + batch_result = self.model_worker.forward_batch_generation( + batch, pp_proxy_tensors=pp_proxy_tensors + ) # The isolation restore reverted the worker's in-forward SB edits; # re-apply what must carry to the next iter. batch.spec_info = batch_result.next_draft_input @@ -3682,12 +3695,18 @@ def run_batch( batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) batch.input_ids = None # rebuilt next iter from draft_token self.update_cache_from_scheduler(batch, batch_result) - # Sync D2H so the result processor can read CPU tensors. + # Sync D2H so the result processor can read CPU tensors. A non-last + # PP rank produced only proxy tensors, so there is nothing to copy. + # Under PP this result is not the one that gets processed -- every + # rank consumes the copy rebuilt from the output ring -- and the + # ring carries device tensors, so copying here would only move + # next_token_ids to the host behind the ring's back. batch_result.copy_done = self.device_module.Event() - batch_result.copy_to_cpu( - return_logprob=batch.return_logprob, - return_hidden_states=batch.return_hidden_states, - ) + if batch_result.has_sampled_token_ids and self.ps.pp_size == 1: + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, + return_hidden_states=batch.return_hidden_states, + ) else: kwargs = ( {"pp_proxy_tensors": pp_proxy_tensors} diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 715e20ebf34f..bad96ff59986 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -278,6 +278,8 @@ def event_loop_pp_disagg_prefill(self: Scheduler): ) self._pp_commit_comm_work(self.send_proxy_work) if cur_batch: + if self.enable_staging: + self.maybe_prefetch_staging_for_batch(cur_batch) result, self.launch_event = self._pp_launch_batch( mb_id, cur_batch, @@ -1018,6 +1020,15 @@ def _pp_prepare_tensor_dict( "next_token_ids": result.next_token_ids, } + # The draft extend only runs on the last stage, but every rank needs its + # output: process_batch_result_disagg_prefill reads it off the relayed + # result to fill the PD aux buffers. + draft_input = result.next_draft_input + if draft_input is not None and draft_input.topk_p is not None: + tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous() + tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous() + tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous() + if batch.return_logprob: logprob_dict = get_logprob_dict_from_result(result) tensor_dict = { @@ -1149,17 +1160,46 @@ def _pp_prep_batch_result( extend_logprob_start_len_per_req, ) = get_logprob_from_pp_outputs(pp_outputs) next_token_ids = pp_outputs["next_token_ids"].to(torch.int64) + + # Rebuild the draft proposal the last stage put on the ring. Rebinding + # batch.spec_info to the same object keeps the identity check in + # process_batch_result_disagg_prefill true on every rank. + next_draft_input = None + if "draft_topk_p" in pp_outputs.tensors: + from sglang.srt.speculative.eagle_info import EagleDraftInput + + next_draft_input = EagleDraftInput( + topk_p=pp_outputs["draft_topk_p"], + topk_index=pp_outputs["draft_topk_index"], + hidden_states=pp_outputs["draft_hidden_states"], + bonus_tokens=next_token_ids, + num_tokens_per_req=1, + num_tokens_for_logprob_per_req=1, + ) + batch.spec_info = next_draft_input + # PP rank 0 also relays into output_tokens_buf so the next iter's # resolve_forward_inputs finds these tokens for the decode portion # of mixed-chunk batches (which gather via mix_running_indices). self.future_map.stash( - batch.req_pool_indices, RelayPayload(bonus_tokens=next_token_ids) + batch.req_pool_indices, + RelayPayload( + bonus_tokens=next_token_ids, + topk_p=None if next_draft_input is None else next_draft_input.topk_p, + topk_index=( + None if next_draft_input is None else next_draft_input.topk_index + ), + hidden_states=( + None if next_draft_input is None else next_draft_input.hidden_states + ), + ), ) batch.input_ids = None output_result = GenerationBatchResult( logits_output=logits_output, pp_hidden_states_proxy_tensors=None, next_token_ids=pp_outputs["next_token_ids"], + next_draft_input=next_draft_input, extend_input_len_per_req=extend_input_len_per_req, extend_logprob_start_len_per_req=extend_logprob_start_len_per_req, can_run_cuda_graph=mb_metadata.can_run_cuda_graph, diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f6eec4b37156..e06593ced0c3 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -311,6 +311,7 @@ def __init__( is_multi_layer_eagle: bool = False, context_length: Optional[int] = None, draft_attention_backend: Optional[str] = None, + random_seed: Optional[int] = None, ): # Parse args self.server_args = server_args @@ -370,8 +371,13 @@ def __init__( self.world_group = get_world_group() # Sync random seed across TP workers. - # Elastic joiners cannot enter the launch-time WORLD broadcast. - if server_args.is_ep_joiner: + # Elastic joiners cannot enter the launch-time WORLD broadcast. Neither can a + # draft worker that exists on a subset of ranks (prefill-side PP builds it on + # the last stage only); it takes the seed the target already broadcast, which + # keeps the value identical to today's re-broadcast on every other setup. + if random_seed is not None: + self.random_seed = random_seed + elif server_args.is_ep_joiner: self.random_seed = server_args.random_seed else: self.random_seed = broadcast_pyobj( diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index cdcc2b5c76dd..cdc99316f3db 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -51,6 +51,30 @@ class KVCacheBuildResult: from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +def get_draft_kv_pool( + *, + draft_worker: BaseTpWorker, + spec_algorithm: SpeculativeAlgorithm, + server_args: ServerArgs, +): + """Return the draft token-to-KV pool for the current draft worker, + or None when no draft KV pool is available.""" + if draft_worker is None or spec_algorithm.is_ngram(): + return None + + # V2 workers nest the draft runner under `.draft_worker`. That inner worker is + # None on ranks that do not host the draft (prefill-side PP builds it only on + # the last stage), and those ranks own no draft KV pool. + if draft_worker.draft_worker is None: + return None + + if server_args.enable_multi_layer_eagle: + draft_runner = draft_worker.draft_worker.draft_runner_list[0] + else: + draft_runner = draft_worker.draft_worker.draft_runner + return draft_runner.token_to_kv_pool + + def maybe_register_hicache_draft( *, tree_cache, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index de9efbf20227..a52d435c1f9d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -122,6 +122,32 @@ def _should_enable_lazy_compaction() -> bool: MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1 MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_BUFFER = 1 + +def _pp_local_per_request_bytes( + total_bytes: int, + layer_ids: list[int], + start_layer: int, + end_layer: int, +) -> int: + """Scale a layer-linear state cost to the current PP stage. + + ``BaseLinearStateParams`` reports bytes for every linear-attention layer in + the model config, while the PP memory pools below allocate only layers in + ``[start_layer, end_layer)``. Budgeting the global value makes the error + grow with PP size and can reject configurations whose real local pools fit. + """ + if not layer_ids: + return 0 + if total_bytes % len(layer_ids) != 0: + raise ValueError( + "Linear-state bytes must be uniform per layer: " + f"total_bytes={total_bytes}, num_layers={len(layer_ids)}" + ) + local_layer_count = sum( + start_layer <= layer_id < end_layer for layer_id in layer_ids + ) + return total_bytes // len(layer_ids) * local_layer_count + if TYPE_CHECKING: from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.mem_cache.unified_memory_pool import ( @@ -791,7 +817,14 @@ def _build_hybrid_req_pool( ), enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), enable_mamba_extra_buffer_lazy=self.server_args.enable_mamba_extra_buffer_lazy(), - speculative_num_draft_tokens=self.server_args.max_speculative_num_draft_tokens, + # A PD prefill server never runs TARGET_VERIFY, so skip the + # verify-only per-draft-token state snapshots (see the draft-head + # case above: None => the pool skips SpeculativeState). + speculative_num_draft_tokens=( + None + if get_disagg().disaggregation_mode == "prefill" + else self.server_args.max_speculative_num_draft_tokens + ), speculative_eagle_topk=get_spec().speculative_eagle_topk, enable_overlap_schedule=not get_schedule().disable_overlap_schedule, start_layer=self.layer_info.start_layer, diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 5e60753b4023..a90a732c252f 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -938,6 +938,37 @@ def _prefill_num_token_non_padded_post_fill(buf, fb, ctx): "prefill registry; cannot adopt." ) reg.register_slot(slot, bind=bind) + + # Pipeline-parallel stage inputs are token-axis tensors carried outside + # ForwardBatch. Adopt the runner-owned backing buffers so capture and + # replay use the same addresses, copy the live stage input through + # FillContext, and clear bucket padding because prefill graphs execute + # every padded token. + if source is not None: + pp = getattr(source, "pp_proxy_tensors", None) + if pp is not None: + + def _pp_source(key): + def _fn(_fb, ctx): + ppx = ctx.pp_proxy_tensors + return None if ppx is None else ppx.tensors[key] + + return _fn + + for _key, _backing in pp.items(): + reg.register_slot( + GraphSlot( + name=f"pp_proxy_tensors.{_key}", + shape_fn=lambda _bs, mt, _tail=tuple( + _backing.shape[1:] + ): (mt, *_tail), + dtype=_backing.dtype, + axis="tokens", + padding_policy=PaddingPolicy.ZERO, + source_fn=_pp_source(_key), + ), + bind=_backing, + ) return reg diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 2bd30e03d9ad..0559160a9fd6 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -3,7 +3,7 @@ import logging import time from collections import defaultdict -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import msgspec @@ -69,6 +69,38 @@ def should_skip_auto_prefill_cuda_graph_for_memory( ) +def has_standard_gqa_for_all_local_layers( + *, attention_layer_count: int, start_layer: int, end_layer: int +) -> bool: + """Check the layers materialized on this pipeline rank, not the full model.""" + return attention_layer_count >= end_layer - start_layer + + +def index_attention_layers_by_global_id( + attention_layers: list[Any], mha_companion_layers: list[Any] +) -> tuple[list[Any], list[Any]]: + """Pad PP-local attention metadata so global layer_id remains a valid index.""" + if len(attention_layers) != len(mha_companion_layers): + raise ValueError("attention and MHA companion metadata must be parallel") + populated = [layer for layer in attention_layers if layer is not None] + if not populated or any(not hasattr(layer, "layer_id") for layer in populated): + return attention_layers, mha_companion_layers + max_layer_id = max(int(layer.layer_id) for layer in populated) + indexed_attention = [None] * (max_layer_id + 1) + indexed_companions = [None] * (max_layer_id + 1) + for attention, companion in zip(attention_layers, mha_companion_layers): + if attention is None: + if companion is not None: + raise ValueError("MHA companion has no primary attention layer") + continue + layer_id = int(attention.layer_id) + if layer_id < 0 or indexed_attention[layer_id] is not None: + raise ValueError(f"invalid or duplicate attention layer_id: {layer_id}") + indexed_attention[layer_id] = attention + indexed_companions[layer_id] = companion + return indexed_attention, indexed_companions + + class GraphCapture(msgspec.Struct, frozen=True, kw_only=True): runner: Optional[BaseRunner] memory_phase: str @@ -385,8 +417,20 @@ def result( model_runner.dsa_indexers, model_runner.mha_companion_layers, ) = compute_attention_and_moe_layers(layer_model) + ( + model_runner.attention_layers, + model_runner.mha_companion_layers, + ) = index_attention_layers_by_global_id( + model_runner.attention_layers, model_runner.mha_companion_layers + ) - if len(model_runner.attention_layers) < model_runner.model_config.num_hidden_layers: + if not has_standard_gqa_for_all_local_layers( + attention_layer_count=sum( + layer is not None for layer in model_runner.attention_layers + ), + start_layer=model_runner.layer_info.start_layer, + end_layer=model_runner.layer_info.end_layer, + ): # TODO(yuwei): support Non-Standard GQA log_info_on_rank0( logger, @@ -446,6 +490,14 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> GraphCapture: capture_time=0, ) + # A PD prefill server never replays the target-verify graph, and its pool + # is built without the spec-verify scratch the capture would need. + if ( + model_runner.spec_algorithm.is_speculative() + and not model_runner.is_draft_worker + and model_runner.server_args.disaggregation_mode == "prefill" + ): + return no_capture if not model_runner.is_generation: # TODO: Currently, cuda graph only captures decode steps, which only exists for generation models return no_capture diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 288b1e274be1..9a73d11279aa 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -113,7 +113,7 @@ def _allocate_decode_buffers( is_mhc = hc_hidden_size is not None hs = hc_hidden_size if is_mhc else hidden_size pp_proxy_tensors = { - "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), + "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype), } if not is_mhc: # Only Kimi K3 supplies num_blocks: its PP bank is token-major @@ -121,7 +121,7 @@ def _allocate_decode_buffers( residual_shape = ( (max_num_token, pp_proxy_residual_num_blocks, hidden_size) if pp_proxy_residual_num_blocks is not None - else (max_bs, hidden_size) + else (max_num_token, hidden_size) ) pp_proxy_tensors["residual"] = torch.zeros(residual_shape, dtype=dtype) if pp_proxy_topk_size is not None: @@ -231,6 +231,14 @@ def warmup(self) -> None: self._pre_initialize_flashinfer_allreduce_workspace() self._pre_initialize_fi_a2a_workspace() + # Model-owned communication resources may depend on the resolved + # request pool and must be compiled/allocated before graph capture. + prepare_model_resources = getattr( + mr.model, "prepare_before_cuda_graph_capture", None + ) + if prepare_model_resources is not None: + prepare_model_resources(mr) + if should_run_flashinfer_autotune(self.model_runner): buffers, batch_size = self._autotune_buffers() assert ( @@ -385,7 +393,13 @@ def _dummy_run( else get_server_return_hidden_states_mode(mr.server_args) ) num_tokens_per_req = 1 - if mr.spec_algorithm.is_speculative(): + # A PD prefill target worker's pool has no SpeculativeState, so a + # TARGET_VERIFY dummy forward would trip the linear-attn backend's + # pool-type assert. Warm up in plain DECODE instead. + _is_pd_prefill_target = ( + mr.server_args.disaggregation_mode == "prefill" and not mr.is_draft_worker + ) + if mr.spec_algorithm.is_speculative() and not _is_pd_prefill_target: if mr.is_draft_worker: assert ( mr.spec_algorithm.supports_target_verify_for_draft() diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 2e7a0bf536e5..7637a69c00e4 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -38,6 +38,7 @@ from __future__ import annotations import copy +import dataclasses import inspect import logging from contextlib import contextmanager @@ -307,6 +308,11 @@ def __init__(self, model_runner: ModelRunner): hidden_size=self.model_runner.model_config.hidden_size, dtype=self.model_runner.dtype, enable_mamba_track=self.mamba_track_enabled, + pp_size=self.model_runner.pp_group.world_size, + hc_hidden_size=getattr( + self.model_runner.model_config, "hc_hidden_size", None + ), + pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(), ) self.buffers.share_buffers() # Token-axis FB-shared slot registry adopting PrefillInputBuffers @@ -608,6 +614,14 @@ def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tenso return forward_batch.positions + def _static_pp_proxy_tensors(self, num_tokens: int) -> Optional[PPProxyTensors]: + buffers = self.buffers.pp_proxy_tensors + if buffers is None: + return None + return PPProxyTensors( + {key: value[:num_tokens] for key, value in buffers.items()} + ) + @contextmanager def _prefill_forward_context( self, @@ -664,6 +678,9 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int): set_is_extend_in_batch(False) with self._prefill_forward_context(forward_batch): + pp_kwargs = self.model_runner._pp_kwargs( + self._static_pp_proxy_tensors(num_tokens) + ) if self._uses_eager_prefill_tail(): # BCG / Full: capture the transformer body only. positions = self._get_layer_model_positions(forward_batch) @@ -672,12 +689,14 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int): positions, forward_batch, forward_batch.input_embeds, + **pp_kwargs, ) # tc_piecewise: compile/capture the outer model.forward path. return self.model_runner.model.forward( forward_batch.input_ids, forward_batch.positions, forward_batch, + **pp_kwargs, ) def _run_dummy_forward(self, num_tokens: int) -> None: @@ -1429,6 +1448,7 @@ def load_batch(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch: padded_bs=bs, raw_num_tokens=num_tokens, padded_num_tokens=static_num_tokens, + pp_proxy_tensors=kwargs.get("pp_proxy_tensors"), ) registry = self.buffer_registry @@ -1481,6 +1501,21 @@ def _slot(name): else forward_batch.global_forward_mode ) + # The eager tail of the draft model concatenates spec hidden states + # with embeddings of the padded input_ids, so the batch must carry the + # bucket-sized static buffer view rather than the raw-length live + # tensor. The buffer prefix is refreshed from the live tensor below. + padded_spec_info = forward_batch.spec_info + if ( + self.static_draft_hidden_states is not None + and padded_spec_info is not None + and getattr(padded_spec_info, "hidden_states", None) is not None + ): + padded_spec_info = dataclasses.replace( + padded_spec_info, + hidden_states=self.static_draft_hidden_states[:static_num_tokens], + ) + static_forward_batch = ForwardBatch( forward_mode=pcg_forward_mode, batch_size=bs, @@ -1524,7 +1559,7 @@ def _slot(name): global_dp_buffer_len=forward_batch.global_dp_buffer_len, mrope_positions=mrope_positions, spec_algorithm=forward_batch.spec_algorithm, - spec_info=forward_batch.spec_info, + spec_info=padded_spec_info, capture_hidden_mode=forward_batch.capture_hidden_mode, num_token_non_padded=num_token_non_padded, num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu, @@ -1717,8 +1752,8 @@ def _finalize_execute_output( if isinstance(output, EmbeddingPoolerOutput): return output assert isinstance(output, PPProxyTensors) - raise NotImplementedError( - "PPProxyTensors is not supported in PrefillCudaGraphRunner yet." + return PPProxyTensors( + {key: value[: self.raw_num_tokens] for key, value in output.tensors.items()} ) def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None: diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 97c7702c7288..e68e6364ba99 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -133,7 +133,7 @@ def create( is_mhc = hc_hidden_size is not None hs = hc_hidden_size if is_mhc else hidden_size pp_proxy_tensors = { - "hidden_states": torch.zeros((max_bs, hs), dtype=dtype), + "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype), } if not is_mhc: # Only Kimi K3 supplies num_blocks: its PP bank is token-major @@ -141,7 +141,7 @@ def create( residual_shape = ( (max_num_token, pp_proxy_residual_num_blocks, hidden_size) if pp_proxy_residual_num_blocks is not None - else (max_bs, hidden_size) + else (max_num_token, hidden_size) ) pp_proxy_tensors["residual"] = torch.zeros( residual_shape, dtype=dtype @@ -342,6 +342,7 @@ class PrefillInputBuffers(ForwardInputBuffers): positions: torch.Tensor input_embeds: Optional[torch.Tensor] mrope_positions: Optional[torch.Tensor] + pp_proxy_tensors: Optional[Dict[str, torch.Tensor]] @classmethod def create( @@ -355,6 +356,9 @@ def create( hidden_size: int, dtype: torch.dtype, enable_mamba_track: bool, + pp_size: int = 1, + hc_hidden_size: Optional[int] = None, + pp_proxy_topk_size: Optional[int] = None, ) -> PrefillInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64) @@ -382,6 +386,25 @@ def create( input_embeds = None mrope_positions = None + if pp_size > 1: + is_mhc = hc_hidden_size is not None + proxy_hidden_size = hc_hidden_size if is_mhc else hidden_size + pp_proxy_tensors = { + "hidden_states": torch.zeros( + (max_num_tokens, proxy_hidden_size), dtype=dtype + ), + } + if not is_mhc: + pp_proxy_tensors["residual"] = torch.zeros( + (max_num_tokens, hidden_size), dtype=dtype + ) + if pp_proxy_topk_size is not None: + pp_proxy_tensors["topk_indices"] = torch.zeros( + (max_num_tokens, pp_proxy_topk_size), dtype=torch.int32 + ) + else: + pp_proxy_tensors = None + return cls( input_ids=input_ids, out_cache_loc=out_cache_loc, @@ -392,6 +415,7 @@ def create( positions=positions, input_embeds=input_embeds, mrope_positions=mrope_positions, + pp_proxy_tensors=pp_proxy_tensors, ) def populate_from_forward_batch( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e8c8a5d17891..3ca2a1b65a23 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -731,6 +731,7 @@ def __init__( or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() or get_moe_a2a_backend().is_megamoe() + or get_moe_a2a_backend().is_deepep_v2() or should_use_flashinfer_cutlass_moe_fp4_allgather() or envs.SGLANG_SHARED_EXPERT_TP1.get() ) @@ -811,6 +812,7 @@ def __init__( or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() + or get_moe_a2a_backend().is_deepep_v2() ): # TODO: we will support tp < ep in the future self.ep_size = get_parallel().moe_ep_size @@ -833,6 +835,7 @@ def __init__( or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() + or get_moe_a2a_backend().is_deepep_v2() ) self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() # SGLANG_OPT_MOE_QUANT_ONCE eligibility, resolved lazily on first @@ -2669,7 +2672,12 @@ def __init__( for i in range(len(self.layers)): if isinstance(self.layers[i].mlp, DeepseekV2MoE): # tp_size = get_parallel().tp_size - is_a2a_moe = is_deepep_class_backend() + # Keep the original deepep-class scope here and only add DeepEP v2, + # so unrelated backends' allocator sizing is unchanged. + is_a2a_moe = ( + is_deepep_class_backend() + or get_moe_a2a_backend().is_deepep_v2() + ) tp_size = 1 if is_a2a_moe else get_parallel().tp_size intermediate_size = ( config.moe_intermediate_size * config.n_shared_experts @@ -2689,10 +2697,11 @@ def __init__( ) ) self.layers_to_capture = [] - if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): - self.enable_a2a_moe = True - else: - self.enable_a2a_moe = False + self.enable_a2a_moe = ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_deepep_v2() + ) # llama_4_scaling: for supporting Mistral-Large-3 model self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 87594757a171..3349fb01ba64 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -19,6 +19,7 @@ """Inference-only Qwen2MoE model compatible with HuggingFace weights.""" import logging +import os from contextlib import nullcontext from typing import Any, Dict, Iterable, List, Optional, Tuple, Union @@ -27,7 +28,10 @@ from torch import nn from transformers import PretrainedConfig -from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add +from sglang.kernels.ops.elementwise.elementwise import ( + fused_gate_sigmoid_mul, + fused_gate_sigmoid_mul_add, +) from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_pp_group, @@ -151,6 +155,7 @@ def can_fuse_shared_expert( or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size or get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() ): return False @@ -299,6 +304,10 @@ def __init__( routing_method_type=RoutingMethodType.RenormalizeNaive, num_fused_shared_experts=self.num_fused_shared_experts, inplace=not _needs_hidden_after_experts, + enable_qwen35_fp8_deferred_finalize=( + config.model_type == "qwen3_5_moe_text" + and envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() + ), ) self.gate = ReplicatedLinear( @@ -326,6 +335,7 @@ def __init__( dict(tp_rank=0, tp_size=1) if ( get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() or get_moe_a2a_backend().is_flashinfer() ) else {} @@ -344,7 +354,10 @@ def __init__( else: self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) - if get_moe_a2a_backend().is_deepep(): + if ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() + ): # TODO: we will support tp < ep in the future self.ep_size = get_parallel().moe_ep_size self.num_experts = ( @@ -459,6 +472,23 @@ def _forward_shared_experts( return shared_output def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + trace_e2e = os.getenv("SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E", "0") == "1" + + def trace_sync(stage: str): + if not trace_e2e: + return + print( + "SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E " + f"stage={stage}_sync_enter tokens={hidden_states.shape[0]}", + flush=True, + ) + torch.cuda.synchronize() + print( + "SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E " + f"stage={stage}_sync_returned tokens={hidden_states.shape[0]}", + flush=True, + ) + enable_dual_stream = ( is_npu() and envs.SGLANG_NPU_USE_MULTI_STREAM.get() @@ -474,6 +504,7 @@ def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBat ) else: shared_output = self._forward_shared_experts(hidden_states) + trace_sync("shared_expert") topk_output = self.topk( hidden_states, router_logits, @@ -488,34 +519,95 @@ def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBat ) else: topk_output = self.topk.empty_topk_output(hidden_states.device) + trace_sync("pre_experts") final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) + trace_sync("post_experts") if enable_dual_stream: wait_share_stream() if shared_output is not None: + trace_sync("pre_shared_add") final_hidden_states.add_(shared_output) + trace_sync("post_shared_add") return final_hidden_states - def _forward_router_experts(self, hidden_states: torch.Tensor): + @property + def supports_deferred_finalize(self) -> bool: + return bool( + self.experts.supports_deferred_finalize and self.shared_expert is not None + ) + + def _forward_router_experts( + self, + hidden_states: torch.Tensor, + *, + defer_finalize: bool = False, + ): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) + if defer_finalize: + if not self.supports_deferred_finalize: + raise RuntimeError( + "Qwen deferred finalize requires a compatible FlashInfer " + "TRTLLM MoE producer and a separate shared expert" + ) + if not TopKOutputChecker.format_is_bypassed(topk_output): + raise RuntimeError( + "Qwen deferred finalize requires logits-based bypassed TopK" + ) + return self.experts.forward_deferred_finalize(hidden_states, topk_output) if self.enable_shared_expert_fusion and TopKOutputChecker.format_is_standard( topk_output ): topk_output = self._append_shared_to_topk_output(topk_output, hidden_states) return self.experts(hidden_states, topk_output) + def _gate_shared_output_out_of_place( + self, + hidden_states: torch.Tensor, + shared_output: torch.Tensor, + ) -> torch.Tensor: + if self.shared_expert_gate is None: + return shared_output + return fused_gate_sigmoid_mul( + hidden_states, + self.shared_expert_gate.weight.squeeze(0), + shared_output, + ) + def forward_normal_dual_stream( self, hidden_states: torch.Tensor, use_fused_gate: bool = False, + defer_finalize: bool = False, ) -> torch.Tensor: current_stream = torch.cuda.current_stream() + + if defer_finalize: + # Keep routed FC2 on the current stream so the following finalize + # kernel can be its PDL dependent. The shared branch only reads the + # same live input and writes a fresh gated output; no D2D clone is + # introduced on this path. + self.alt_stream.wait_stream(current_stream) + router_output = self._forward_router_experts( + hidden_states, defer_finalize=True + ) + with torch.cuda.stream(self.alt_stream): + shared_output = self._forward_shared_experts( + hidden_states, apply_gate=False + ) + if shared_output is not None: + shared_output = self._gate_shared_output_out_of_place( + hidden_states, shared_output + ) + current_stream.wait_stream(self.alt_stream) + return router_output, shared_output + self.alt_stream.wait_stream(current_stream) shared_output = ( self._forward_shared_experts( @@ -556,11 +648,17 @@ def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, + defer_finalize: bool = False, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) + if defer_finalize and num_tokens == 0: + raise RuntimeError("Qwen deferred finalize does not support M=0") - if get_moe_a2a_backend().is_deepep(): + if ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_deepep_v2() + ): return self._forward_deepep(hidden_states, forward_batch) use_fused_gate = ( @@ -582,13 +680,34 @@ def forward( and not torch.compiler.is_compiling() ): final_hidden_states, shared_output = self.forward_normal_dual_stream( - hidden_states, use_fused_gate=use_fused_gate + hidden_states, + use_fused_gate=use_fused_gate, + defer_finalize=defer_finalize, ) else: shared_output = self._forward_shared_experts( - hidden_states, apply_gate=not use_fused_gate + hidden_states, apply_gate=not use_fused_gate and not defer_finalize + ) + if defer_finalize and shared_output is not None: + shared_output = self._gate_shared_output_out_of_place( + hidden_states, shared_output + ) + final_hidden_states = self._forward_router_experts( + hidden_states, defer_finalize=defer_finalize + ) + + if defer_finalize: + if shared_output is None: + raise RuntimeError("Qwen deferred finalize requires shared output") + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35MoeFinalizeHandoff, + ) + + return Qwen35MoeFinalizeHandoff.from_flashinfer( + final_hidden_states, + gated_shared_output=shared_output, + m=num_tokens, ) - final_hidden_states = self._forward_router_experts(hidden_states) if shared_output is not None: if use_fused_gate: diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 08c47b9942b5..d11cae1ee2de 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -15,6 +15,7 @@ """Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" import logging +import os from functools import lru_cache from typing import Iterable, Optional, Set, Tuple, Union @@ -38,6 +39,7 @@ # Distributed from sglang.srt.distributed import get_pp_group +from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader @@ -131,6 +133,9 @@ _qknorm_use_alt_stream = _is_cuda or ( get_bool_env_var("SGLANG_QK_NORM_ALT_STREAM", "False") and _hip_use_alt_stream ) +_gdn_decode_fused_proj_conv = _is_cuda and get_bool_env_var( + "SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV", "True" +) _is_amx_available = cpu_has_amx_support() cached_get_processor = lru_cache(get_processor) @@ -139,7 +144,30 @@ def _disable_shared_experts_fusion() -> bool: # Resolved lazily: the global server args is not set at module import time # (e.g. when this module is imported by unit tests). - return get_exec().moe.disable_shared_experts_fusion + # The deferred-finalize ABI needs the shared expert as a separate, gated + # local contribution; it cannot consume a shared slot fused into routed MoE. + return bool( + envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() + or get_exec().moe.disable_shared_experts_fusion + ) + + +def _use_mnnvl_cutedsl_fusion(config: Qwen3_5TextConfig, is_nextn: bool) -> bool: + return bool( + not is_nextn + and config.model_type == "qwen3_5_moe_text" + and envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() + ) + + +def _layer_communicator_class(config: Qwen3_5TextConfig, is_nextn: bool): + if _use_mnnvl_cutedsl_fusion(config, is_nextn): + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35FlashInferLayerCommunicator, + ) + + return Qwen35FlashInferLayerCommunicator + return LayerCommunicator if _is_cuda: @@ -207,6 +235,28 @@ def _select_fused_ar_input_for_linear(hidden_states, linear: nn.Module): ) +def _finish_mlp_output(hidden_states, *, expect_deferred: bool): + """Preserve the ordinary marker or validate the real deferred handoff.""" + if not expect_deferred: + if not isinstance(hidden_states, torch.Tensor): + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35MoeFinalizeHandoff, + ) + + if isinstance(hidden_states, Qwen35MoeFinalizeHandoff): + raise RuntimeError("unexpected deferred-finalize handoff") + hidden_states._sglang_needs_allreduce_fusion = True + return hidden_states + + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35MoeFinalizeHandoff, + ) + + if not isinstance(hidden_states, Qwen35MoeFinalizeHandoff): + raise RuntimeError("Qwen3.5 expected a FlashInfer deferred-finalize handoff") + return hidden_states + + if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( split_qkvgate_gemma_rmsnorm_rope, @@ -632,7 +682,23 @@ def forward( hidden_states ) - if self.num_v_heads // self.num_k_heads in [1, 2, 4] and not _is_npu: + use_fused_decode_proj_conv = ( + _gdn_decode_fused_proj_conv + and forward_batch.forward_mode.is_decode() + and isinstance(projected_states_qkvz, torch.Tensor) + and isinstance(projected_states_ba, torch.Tensor) + ) + value_to_key_head_ratio = self.num_v_heads // self.num_k_heads + use_fused_contiguous_unpack = value_to_key_head_ratio in [1, 2, 4, 8] + if use_fused_decode_proj_conv: + # The GDN backend owns the indexed Conv1D state and therefore owns + # the safe unpack+Conv fusion boundary. B/A are passed as temporary + # placeholders and replaced by the backend before recurrent GDN. + mixed_qkv = (projected_states_qkvz, projected_states_ba) + z = None + b = projected_states_ba + a = projected_states_ba + elif use_fused_contiguous_unpack and not _is_npu: if _is_cpu: num_k_heads_tp = self.num_k_heads // self.attn_tp_size num_v_heads_tp = self.num_v_heads // self.attn_tp_size @@ -659,12 +725,22 @@ def forward( ) mixed_qkv = torch.cat((query, key, value), dim=-1) - core_attn_out = self.attn( + attn_result = self.attn( forward_batch, mixed_qkv=mixed_qkv, a=a, b=b, ) + if use_fused_decode_proj_conv: + if not isinstance(attn_result, tuple) or len(attn_result) != 2: + raise RuntimeError( + "Fused GDN decode projection/Conv1D backend must return " + "(core_attn_out, z)" + ) + core_attn_out, z = attn_result + else: + core_attn_out = attn_result + assert z is not None z_shape_og = z.shape # reshape input data into 2D tensor @@ -757,7 +833,7 @@ def __init__( _enable_qwen35_fused_ar_quant() and _linear_accepts_fp8_tuple(self.linear_attn.in_proj_qkvz) ) - self.layer_communicator = LayerCommunicator( + self.layer_communicator = _layer_communicator_class(config, is_nextn)( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, @@ -806,6 +882,24 @@ def forward( forward_batch ) ) + defer_moe_finalize = ( + fuse_mlp_allreduce + and isinstance(hidden_states, torch.Tensor) + and isinstance(self.mlp, Qwen2MoeSparseMoeBlock) + and hasattr(self.layer_communicator, "should_use_finalize") + and self.layer_communicator.should_use_finalize( + forward_batch, int(hidden_states.shape[0]) + ) + ) + if ( + fuse_mlp_allreduce + and self.layer_communicator.is_last_layer + and not defer_moe_finalize + ): + # The last layer has no following prepare_attn to consume the + # ordinary deferred-AllReduce marker. Fall back before the MLP so + # postprocess_layer performs the collective instead of dropping it. + fuse_mlp_allreduce = False with get_forward().scoped( fuse_mlp_allreduce=fuse_mlp_allreduce, mlp_reduce_scatter=mlp_reduce_scatter, @@ -814,11 +908,14 @@ def forward( hidden_states = self.mlp( hidden_states, forward_batch, + defer_finalize=defer_moe_finalize, ) else: hidden_states = self.mlp(hidden_states) if fuse_mlp_allreduce: - hidden_states._sglang_needs_allreduce_fusion = True + hidden_states = _finish_mlp_output( + hidden_states, expect_deferred=defer_moe_finalize + ) else: hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch @@ -968,7 +1065,7 @@ def __init__( enable_fused_ar_quant = ( _enable_qwen35_fused_ar_quant() and _linear_accepts_fp8_tuple(self.qkv_proj) ) - self.layer_communicator = LayerCommunicator( + self.layer_communicator = _layer_communicator_class(config, is_nextn)( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, @@ -1201,6 +1298,21 @@ def forward( forward_batch ) ) + defer_moe_finalize = ( + fuse_mlp_allreduce + and isinstance(hidden_states, torch.Tensor) + and isinstance(self.mlp, Qwen2MoeSparseMoeBlock) + and hasattr(self.layer_communicator, "should_use_finalize") + and self.layer_communicator.should_use_finalize( + forward_batch, int(hidden_states.shape[0]) + ) + ) + if ( + fuse_mlp_allreduce + and self.layer_communicator.is_last_layer + and not defer_moe_finalize + ): + fuse_mlp_allreduce = False with get_forward().scoped( fuse_mlp_allreduce=fuse_mlp_allreduce, mlp_reduce_scatter=mlp_reduce_scatter, @@ -1209,11 +1321,14 @@ def forward( hidden_states = self.mlp( hidden_states, forward_batch, + defer_finalize=defer_moe_finalize, ) else: hidden_states = self.mlp(hidden_states) if fuse_mlp_allreduce: - hidden_states._sglang_needs_allreduce_fusion = True + hidden_states = _finish_mlp_output( + hidden_states, expect_deferred=defer_moe_finalize + ) else: hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch @@ -1380,6 +1495,48 @@ def get_layer(idx: int, prefix: str): prefix=f"{prefix}.layers", ) + self.flashinfer_mnnvl_cutedsl_fusion = None + if _use_mnnvl_cutedsl_fusion(config, is_nextn): + if self.pp_group.world_size != 1: + raise RuntimeError( + "Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently requires PP=1" + ) + unsupported_layers = [ + layer.layer_id + for layer in self.layers + if not isinstance(layer.mlp, Qwen2MoeSparseMoeBlock) + or not layer.mlp.supports_deferred_finalize + ] + if unsupported_layers: + raise RuntimeError( + "Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently " + "requires block-FP8 MoE weights with FlashInfer TRTLLM " + "deferred-finalize support on every layer; unsupported " + "layers: " + f"{unsupported_layers}" + ) + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35FlashInferFusionService, + Qwen35FlashInferLayerCommunicator, + ) + + self.flashinfer_mnnvl_cutedsl_fusion = Qwen35FlashInferFusionService( + hidden_size=config.hidden_size, + top_k=config.num_experts_per_tok, + rms_epsilon=config.rms_norm_eps, + ) + for layer in self.layers: + communicator = layer.layer_communicator + if not isinstance(communicator, Qwen35FlashInferLayerCommunicator): + raise RuntimeError( + "Qwen3.5 fusion-enabled layer has the wrong communicator" + ) + communicator.fusion_service = self.flashinfer_mnnvl_cutedsl_fusion + logger.info( + "Installed one Qwen3.5 FlashInfer fusion handle for %d layers", + len(self.layers), + ) + # Final normalization if self.pp_group.is_last_rank: self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -1391,6 +1548,15 @@ def get_layer(idx: int, prefix: str): def get_input_embeddings(self): return self.embed_tokens + def prepare_before_cuda_graph_capture(self, model_runner) -> None: + if self.flashinfer_mnnvl_cutedsl_fusion is None: + return + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + prepare_qwen35_flashinfer_fusion, + ) + + prepare_qwen35_flashinfer_fusion(self, model_runner) + def set_dflash_layers_to_capture(self, layers_to_capture: list[int]): self.layers_to_capture = layers_to_capture for layer_id in self.layers_to_capture: @@ -1414,6 +1580,16 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, input_deepstack_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, PPProxyTensors]: + if ( + self.flashinfer_mnnvl_cutedsl_fusion is not None + and input_deepstack_embeds is not None + and input_deepstack_embeds.numel() > 0 + ): + raise RuntimeError( + "Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently supports " + "the text-only path, not deepstack visual inputs" + ) + # Initialize hidden states if self.pp_group.is_first_rank: if input_embeds is None: @@ -1465,12 +1641,67 @@ def forward( } ) - # Apply final normalization - if hidden_states.shape[0] != 0: + # Apply final normalization. The final decoder layer has no following + # layer to consume its deferred MoE tail, so consume it here with the + # model's own GemmaRMSNorm gamma. Preserve the sam/dev native-final-norm + # diagnostic path for the ordinary (non-deferred) case. + trace_final_norm = os.getenv("SGLANG_TRACE_QWEN35_FINAL_NORM", "0") == "1" + use_native_final_norm = ( + os.getenv("SGLANG_QWEN35_NATIVE_FINAL_NORM", "0") == "1" + ) + is_deferred_finalize = False + if self.flashinfer_mnnvl_cutedsl_fusion is not None: + from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35MoeFinalizeHandoff, + ) + + is_deferred_finalize = isinstance(hidden_states, Qwen35MoeFinalizeHandoff) + + if is_deferred_finalize: + if residual is None or self.flashinfer_mnnvl_cutedsl_fusion is None: + raise RuntimeError("invalid final deferred MoE handoff") + hidden_states, _ = self.flashinfer_mnnvl_cutedsl_fusion.finalize( + hidden_states, residual, self.norm.gemma_weight + ) + elif hidden_states.shape[0] != 0: + if trace_final_norm: + print( + "SGLANG_TRACE_QWEN35_FINAL_NORM " + f"stage=pre_sync_enter hidden={tuple(hidden_states.shape)} " + f"hidden_stride={hidden_states.stride()} " + f"hidden_dtype={hidden_states.dtype} " + f"hidden_contiguous={hidden_states.is_contiguous()} " + f"residual={None if residual is None else tuple(residual.shape)} " + f"native={use_native_final_norm}", + flush=True, + ) + torch.cuda.synchronize() + print( + "SGLANG_TRACE_QWEN35_FINAL_NORM stage=pre_sync_returned", + flush=True, + ) if residual is None: - hidden_states = self.norm(hidden_states) + hidden_states = ( + self.norm.forward_native(hidden_states) + if use_native_final_norm + else self.norm(hidden_states) + ) else: - hidden_states, _ = self.norm(hidden_states, residual) + hidden_states, _ = ( + self.norm.forward_native(hidden_states, residual) + if use_native_final_norm + else self.norm(hidden_states, residual) + ) + if trace_final_norm: + print( + "SGLANG_TRACE_QWEN35_FINAL_NORM stage=post_sync_enter", + flush=True, + ) + torch.cuda.synchronize() + print( + "SGLANG_TRACE_QWEN35_FINAL_NORM stage=post_sync_returned", + flush=True, + ) if len(aux_hidden_states) == 0: return hidden_states @@ -1956,6 +2187,11 @@ def __init__( def get_hidden_dim(self, module_name: str, layer_idx: int): return self.model.get_hidden_dim(module_name, layer_idx) + def prepare_before_cuda_graph_capture(self, model_runner) -> None: + prepare = getattr(self.model, "prepare_before_cuda_graph_capture", None) + if prepare is not None: + prepare(model_runner) + def should_apply_lora(self, module_name: str) -> bool: # Accept all language model layer modules (attention, linear_attn, mlp). return module_name.startswith("model.layers.") diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index fd146e1a67ed..b135ac1377bd 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -134,12 +134,16 @@ def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight def set_embed_and_head(self, embed, head): - del self.model.embed_tokens.weight - if not self.config.tie_word_embeddings: + # Under prefill-side pipeline parallelism the target's embed lives on the + # first stage and its lm_head on the last, so only one of them reaches a + # draft that sits on the last stage. Keep whatever the draft loaded itself + # for the half the target cannot share. + if embed is not None: + del self.model.embed_tokens.weight + self.model.embed_tokens.weight = embed + if head is not None and not self.config.tie_word_embeddings: del self.lm_head.weight - - self.model.embed_tokens.weight = embed - self.lm_head.weight = head + self.lm_head.weight = head torch.cuda.empty_cache() torch.cuda.synchronize() @@ -194,6 +198,18 @@ def forward( if not forward_batch.forward_mode.is_idle(): input_embeds = self.pre_fc_norm_embedding(input_embeds) hidden_states = self.pre_fc_norm_hidden(hidden_states) + # A captured prefill graph hands the model its static token slot, so + # input_embeds is the padded height while the target's hidden states + # arrive at this chunk's real height. Place the real rows into a slot + # of the same height; the padding rows are never read downstream. + if hidden_states.shape[0] != input_embeds.shape[0]: + rows = min(hidden_states.shape[0], input_embeds.shape[0]) + slot = hidden_states.new_zeros( + (input_embeds.shape[0], hidden_states.shape[1]) + ) + slot[:rows] = hidden_states[:rows] + hidden_states = slot + hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) hidden_states = self.fc(hidden_states) diff --git a/python/sglang/srt/models/qwen3_5_text.py b/python/sglang/srt/models/qwen3_5_text.py index 7366f9f88d12..93e29b606e27 100644 --- a/python/sglang/srt/models/qwen3_5_text.py +++ b/python/sglang/srt/models/qwen3_5_text.py @@ -100,8 +100,37 @@ def end_layer(self) -> int: def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens + def prepare_before_cuda_graph_capture(self, model_runner) -> None: + """Forward model-owned warmup to the text backbone.""" + self.model.prepare_before_cuda_graph_capture(model_runner) + + def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None: + if self.pp_group.world_size > 1: + raise NotImplementedError("DFLASH/DSPARK aux hidden capture requires PP=1.") + num_layers = len(self.model.layers) + if sorted(set(layer_ids)) != list(layer_ids) or not all( + 0 <= layer_id < num_layers - 1 for layer_id in layer_ids + ): + raise ValueError( + "target_layer_ids must be unique, strictly increasing, and in " + f"[0, {num_layers - 1}); got {layer_ids}" + ) + self.capture_aux_hidden_states = True + self.model.set_dflash_layers_to_capture( + [layer_id + 1 for layer_id in layer_ids] + ) + def get_embed_and_head(self): - return self.model.embed_tokens.weight, self.lm_head.weight + # Under PP the embedding lives on the first stage and the lm_head on the + # last, so a stage holds at most one of them; the draft keeps its own + # copy for the half it does not receive. + embed = ( + None + if isinstance(self.model.embed_tokens, PPMissingLayer) + else self.model.embed_tokens.weight + ) + head = None if isinstance(self.lm_head, PPMissingLayer) else self.lm_head.weight + return embed, head def set_embed_and_head(self, embed, head): del self.model.embed_tokens.weight @@ -147,21 +176,30 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() - body_weights = [] - for name, loaded_weight in weights: - if name.startswith(_MODEL_PREFIX): - body_weights.append((name[len(_MODEL_PREFIX) :], loaded_weight)) - elif name == "lm_head.weight": - if self.config.tie_word_embeddings: - continue - if "lm_head.weight" not in params_dict: - continue - param = params_dict["lm_head.weight"] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add("lm_head.weight") - - body_loaded = self.model.load_weights(body_weights) + # Keep the upstream checkpoint iterator lazy. Materializing all body + # tensors in a list retains every mmap-backed CPU tensor until the + # complete Oakhaven checkpoint has been scanned. With four model + # processes per Grace node that transient host-RSS peak exceeds the + # Slurm memory allocation before the tensors can be copied to GPU. + # The body loader already consumes an Iterable one tensor at a time, + # so prefix stripping can remain a streaming generator. + def body_weights(): + for name, loaded_weight in weights: + if name.startswith(_MODEL_PREFIX): + yield name[len(_MODEL_PREFIX) :], loaded_weight + elif name == "lm_head.weight": + if self.config.tie_word_embeddings: + continue + if "lm_head.weight" not in params_dict: + continue + param = params_dict["lm_head.weight"] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add("lm_head.weight") + + body_loaded = self.model.load_weights(body_weights()) loaded_params.update(f"{_MODEL_PREFIX}{n}" for n in body_loaded) if self.config.tie_word_embeddings and self.pp_group.is_last_rank: diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 9c742a864532..48c48e939911 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -376,6 +376,10 @@ class MoeFlags(_FlagGroupBase): tbo_token_distribution_threshold: float | None = None disable_fp4_allgather: bool | None = None quantization: str | None = None + # True only while constructing/running the speculative draft model. A2A + # dispatchers use it to keep draft CUDA graphs off the target model's + # one-sided communication workspace. + speculative_context: bool = False @dataclasses.dataclass diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c166bb21c335..6122971df754 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -274,6 +274,7 @@ "flashinfer", "megamoe", "pplx", + "deepep_v2", "ascend_tp", ] @@ -304,7 +305,12 @@ "marlin", ] -BF16_GEMM_BACKEND_CHOICES = ["auto", "cutedsl", "torch"] +BF16_GEMM_BACKEND_CHOICES = [ + "auto", + "cutedsl", + "flashinfer_pr4266", + "torch", +] RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"] RETRACTION_POLICY_CHOICES = ["length", "priority"] @@ -1729,7 +1735,7 @@ class ServerArgs: bf16_gemm_backend: A[ str, Arg( - help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the CuTe DSL kernel and cuBLAS), 'torch' (always uses cuBLAS via torch.nn.functional.linear).", + help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM100/SM103 (Blackwell), otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10X; dispatches between the allowlisted low-M Split-K kernel, the CuTe DSL kernel, and cuBLAS; set SGLANG_ENABLE_BF16_SPLITK_GEMM=0 to disable Split-K), 'flashinfer_pr4266' (legacy compatibility alias for the optimized CuTe DSL path), 'torch' (always uses cuBLAS via torch.nn.functional.linear, even on SM100/SM103).", cli_name="--bf16-gemm-backend", choices=BF16_GEMM_BACKEND_CHOICES, ), @@ -2291,6 +2297,8 @@ class ServerArgs: "flashinfer", "megamoe", "pplx", + "deepep_v2", + "ascend_tp", ], Arg( help="Choose the backend for MoE A2A.", @@ -2299,6 +2307,21 @@ class ServerArgs: ), NS("exec.moe"), ] = "none" + deepep_v2_mode: A[ + Literal["direct", "hybrid"], + "DeepEP v2 ElasticBuffer communication topology, fixed at server init: " + "`direct` (single-node NVLink) or `hybrid` (multi-node scale-out). " + "Layout/grouped-GEMM and the decode CUDA graph are chosen per batch by " + "inference phase, independent of this knob; not equivalent to DeepEP v1 " + "normal/low_latency.", + NS("exec.moe"), + ] = "direct" + deepep_v2_dispatcher_output_dtype: A[ + Literal["auto", "bf16", "fp8"], + "DeepEP v2 dispatcher output dtype. `auto`: fp8 for the DeepGEMM runner, bf16 for " + "Triton.", + NS("exec.moe"), + ] = "auto" moe_runner_backend: A[ str, Arg( @@ -6710,6 +6733,119 @@ def _handle_a2a_moe(self): self.cuda_graph_config.decode.backend = Backend.DISABLED self.cuda_graph_config.prefill.backend = Backend.DISABLED + if a2a_backend == "deepep_v2": + if self.moe_runner_backend == "auto": + # The generic auto -> runner resolution above only fires for + # moe_a2a_backend "none", so deepep_v2 would otherwise reach the + # check below with the default "auto" and fail. deep_gemm is the + # production FP8 path, triton the BF16 functional path. Key off the + # dispatcher output dtype and default to the deep_gemm FP8 path: + # self.quantization is not reliably resolved at server-args time + # (FP8 is detected from the checkpoint later), so it cannot drive + # this; a genuinely BF16 run should set + # --deepep-v2-dispatcher-output-dtype bf16 (or --moe-runner-backend + # triton). + self.moe_runner_backend = ( + "triton" + if self.deepep_v2_dispatcher_output_dtype == "bf16" + else "deep_gemm" + ) + logger.warning( + "DeepEP v2 MoE: resolved --moe-runner-backend auto -> %s " + "(--deepep-v2-dispatcher-output-dtype=%s).", + self.moe_runner_backend, + self.deepep_v2_dispatcher_output_dtype, + ) + if self.moe_runner_backend not in ["deep_gemm", "triton"]: + raise ValueError( + "DeepEP v2 MoE currently supports only " + "--moe-runner-backend deep_gemm or triton. " + f"Got {self.moe_runner_backend!r}. Add a runner adapter before " + "enabling DeepEP v2 with other MoE runners." + ) + if self.enable_two_batch_overlap or self.enable_single_batch_overlap: + raise ValueError( + "DeepEP v2 MoE has not implemented the TBO/SBO overlap hooks yet. " + "Disable --enable-two-batch-overlap and " + "--enable-single-batch-overlap when using --moe-a2a-backend deepep_v2." + ) + if self.enforce_shared_experts_fusion: + raise ValueError( + "DeepEP v2 MoE has not validated fused shared experts yet. " + "Remove --enforce-shared-experts-fusion when using " + "--moe-a2a-backend deepep_v2." + ) + deepep_v2_cap = envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + # Prefill capacity pre-check: the ElasticBuffer capacity is per + # physical MoE sender rank. At this point chunked_prefill_size has + # already been divided by attention DP. Before an A2A dispatcher, + # LayerCommunicator uses SCATTERED MLP mode and reduce-scatters that + # per-DP chunk over attention TP, so divide by attention TP as well. + # Without this second division DP4 x TP4 would incorrectly compare + # its 2048-token DP-worker chunk to the buffer even though each + # DeepEP sender receives only 512 tokens. + if ( + self.chunked_prefill_size + and self.chunked_prefill_size > 0 + and (self.disaggregation_mode != "decode") + ): + attn_dp_size = ( + self.dp_size if resolved_view(self).enable_dp_attention else 1 + ) + attn_tp_size = max(1, self.tp_size // attn_dp_size // self.attn_cp_size) + prefill_tokens_per_sender = math.ceil( + self.chunked_prefill_size / attn_tp_size + ) + if prefill_tokens_per_sender > deepep_v2_cap: + raise ValueError( + "DeepEP v2 MoE: the per-sender prefill dispatch budget " + f"({prefill_tokens_per_sender} tokens, from a " + f"{self.chunked_prefill_size}-token DP-worker chunk " + f"reduce-scattered over attention TP={attn_tp_size}) " + "exceeds the per-rank dispatch buffer " + "capacity SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_" + f"RANK={deepep_v2_cap}. Raise the env (it sizes the " + "communication buffer) or lower --chunked-prefill-size." + ) + # CUDA graph is safe on the DeepEP v2 decode masked-GEMM path under ANY + # comm mode (direct or hybrid): the masked layout is chosen per-batch by + # inference phase (decode), not by the comm mode, giving static shapes + # with no host readback. deep_gemm runner + fp8 dispatch are required; + # every other combination (triton/bf16, or the prefill/extend contiguous + # path) needs a host readback / cpu_sync and is not capturable, so the + # decode graph is disabled there (the prefill graph is always disabled). + deepep_v2_fp8 = self.deepep_v2_dispatcher_output_dtype == "fp8" or ( + self.deepep_v2_dispatcher_output_dtype == "auto" + and self.moe_runner_backend == "deep_gemm" + ) + deepep_v2_graph_ok = ( + self.moe_runner_backend == "deep_gemm" and deepep_v2_fp8 + ) + if not deepep_v2_graph_ok: + self.cuda_graph_config.decode.backend = Backend.DISABLED + self.cuda_graph_config.prefill.backend = Backend.DISABLED + else: + # The decode masked-GEMM path is capture-safe under any comm mode + # (static shapes, no host readback). The prefill/extend path goes + # through the non-masked contiguous layout with a host readback and + # is not capturable, so keep the decode graph but always disable the + # prefill graph under DeepEP v2. + self.cuda_graph_config.prefill.backend = Backend.DISABLED + logger.warning( + f"DeepEP v2 MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." + ) + logger.warning( + "DeepEP v2 MoE is using deepep_v2_mode=%s. This controls " + "ElasticBuffer direct/hybrid mode and is independent from " + "--deepep-mode normal/low_latency. DeepEP v2 MoE enables the " + "decode CUDA graph on the deep_gemm + fp8 masked decode path " + "(any comm mode) and disables shared expert fusion. " + "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK is a " + "per-rank communication buffer capacity, not a model limit; " + "increase it for large prefill/chunked-prefill workloads.", + self.deepep_v2_mode, + ) + if ( self.moe_a2a_backend == "none" and is_npu() ) or self.moe_a2a_backend == "ascend_tp": @@ -6718,8 +6854,16 @@ def _handle_a2a_moe(self): if self.moe_a2a_backend == "flashinfer": assert ( - resolved_view(self).enable_dp_attention and self.dp_size == self.tp_size - ), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + resolved_view(self).enable_dp_attention + and self.dp_size > 1 + and self.tp_size % self.dp_size == 0 + ), ( + "FlashInfer MoE A2A requires --enable-dp-attention and a " + "data-parallel size that divides the TP/EP world." + ) + logger.warning( + f"Flashinfer MoE A2A is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." + ) if self.deepep_mode != "auto": logger.warning("--deepep-mode is ignored for Flashinfer MoE A2A") if not envs.SGLANG_MOE_NVFP4_DISPATCH.is_set() and ( @@ -6733,8 +6877,14 @@ def _handle_a2a_moe(self): assert resolved_view(self).moe_runner_backend in [ "flashinfer_cutlass", "flashinfer_cutedsl", + "flashinfer_trtllm", "flashinfer_trtllm_routed", - ], "Flashinfer MoE A2A is only supported with flashinfer_cutlass, flashinfer_cutedsl or flashinfer_trtllm_routed moe runner backend" + "deep_gemm", + ], ( + "FlashInfer MoE A2A is supported with flashinfer_cutlass, " + "flashinfer_cutedsl, flashinfer_trtllm, " + "flashinfer_trtllm_routed, or deep_gemm." + ) if a2a_backend == "mori": if self.deepep_mode == "auto": @@ -8763,8 +8913,16 @@ def check_server_args(self): if self.pp_size > 1: assert ( - self.disable_overlap_schedule and self.speculative_algorithm is None - ), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + self.disable_overlap_schedule + ), "Pipeline parallelism is not compatible with overlap schedule" + # A PD prefill engine runs speculative decoding as a single extend step + # (target forward + one draft extend); there is no accept length and no + # per-step hidden-state feedback, so it composes with the pipeline. The + # decode side still owns the draft loop and stays unsupported. + assert ( + self.speculative_algorithm is None + or self.disaggregation_mode == "prefill" + ), "Pipeline parallelism is only compatible with speculative decoding on a PD prefill engine" assert self.min_free_slots_delay is None, ( "--min-free-slots-delay is not supported with pipeline " "parallelism: allocatable slots per microbatch are bounded by " diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 4f3c100b8bd5..45237bbd279b 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -195,7 +195,7 @@ def __init__( bundle = build_draft_tp_worker( server_args=server_args, gpu_id=gpu_id, - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, target_model_config=target_worker.model_runner.model_config, algo_label="DFLASH", diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index ed64e21f37c8..1a2c6886149d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -111,7 +111,7 @@ def __init__( bundle = build_draft_tp_worker( server_args=server_args, gpu_id=gpu_id, - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, target_model_config=target_worker.model_runner.model_config, algo_label="DSPARK", diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index e87ecd91ef9e..d1b1ed3016ab 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -478,13 +478,27 @@ def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int: return target_hidden * num_aux +def get_draft_recurrent_hidden_state_spec_from_config( + model_config, spec_algorithm +) -> tuple[Optional[int], Optional[torch.dtype]]: + """Return hidden_states width/dtype carried between draft decode steps. + + Config-only so callers without a draft runner can reach it: prefill-side PP + builds the draft on the last stage alone, but the PD metadata wire schema it + feeds has to come out identical on every rank. + """ + if spec_algorithm.is_standalone(): + return None, None + return model_config.spec_hidden_size, model_config.dtype + + def get_draft_recurrent_hidden_state_spec( model_runner: ModelRunner, ) -> tuple[Optional[int], Optional[torch.dtype]]: """Return hidden_states width/dtype carried between draft decode steps.""" - if model_runner.spec_algorithm.is_standalone(): - return None, None - return model_runner.model_config.spec_hidden_size, model_runner.model_config.dtype + return get_draft_recurrent_hidden_state_spec_from_config( + model_runner.model_config, model_runner.spec_algorithm + ) def eagle_prepare_for_verify( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 682c60937ec8..4882a1b8e869 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -7,6 +7,7 @@ import torch from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess +from sglang.srt.distributed import get_pp_group from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import ( @@ -84,6 +85,7 @@ ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( + draft_pp_context, draft_tp_context, fast_sample, get_plan_stream, @@ -162,16 +164,17 @@ def __init__( ctx = empty_context() with ( ctx - ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + ), draft_pp_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): self.draft_worker = TpModelWorker( server_args=server_args, gpu_id=gpu_id, # spec workers don't support pipeline parallelism - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, is_draft_worker=True, # The draft runs at absolute target positions. context_length=target_worker.model_runner.model_config.context_len, + random_seed=target_worker.random_seed, ) # Alias for better readability @@ -304,7 +307,7 @@ def maybe_share_target_lm_head(): ) else: - if self.hot_token_id is not None: + if self.hot_token_id is not None and head is not None: head = head.clone() self.hot_token_id = self.hot_token_id.to(head.device) head.data = head.data[self.hot_token_id] @@ -1024,17 +1027,26 @@ def __init__( server_args.speculative_algorithm ) - self._draft_worker = EagleDraftWorker( - server_args, - gpu_id, - ps, - nccl_port, - target_worker, + # The draft runs where the target's last-layer hidden states and sampled + # tokens exist, i.e. only on the last pipeline stage. Other stages keep an + # EAGLEWorkerV2 that forwards the target and returns proxy tensors, so the + # scheduler's dispatch and run_batch branching stay rank-uniform. + self._hosts_draft = get_pp_group().is_last_rank + self._draft_worker = ( + EagleDraftWorker( + server_args, + gpu_id, + ps, + nccl_port, + target_worker, + ) + if self._hosts_draft + else None ) # Adaptive speculative self.adaptive_controller: Optional[AdaptiveController] = None - if server_args.speculative_adaptive: + if server_args.speculative_adaptive and self._hosts_draft: self.adaptive_controller = AdaptiveController( self, config_path=server_args.speculative_adaptive_config, @@ -1097,7 +1109,11 @@ def init_cuda_graphs(self): ) def forward_batch_generation( - self, batch: ScheduleBatch, on_publish=None, grammar_barrier=None + self, + batch: ScheduleBatch, + on_publish=None, + grammar_barrier=None, + pp_proxy_tensors=None, ): if batch.forward_mode.is_extend() or batch.is_extend_in_batch: # Target prefill @@ -1107,7 +1123,9 @@ def forward_batch_generation( else CaptureHiddenMode.FULL ) batch_output = self.target_worker.forward_batch_generation( - batch, capture_hidden_mode=target_capture_mode + batch, + pp_proxy_tensors=pp_proxy_tensors, + capture_hidden_mode=target_capture_mode, ) # Spec_v2 convention: batch.seq_lens = length BEFORE this iter's tokens. @@ -1117,6 +1135,11 @@ def forward_batch_generation( if on_publish is not None: on_publish(batch_output.new_seq_lens) + # A rank that does not host the draft (prefill-side PP builds it only on + # the last stage) forwards the target's proxy tensors and stops here. + if self._draft_worker is None: + return batch_output + # Draft prefill with ( self.draft_worker.draft_tp_context( diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index c32402333bcc..a13dd4706b87 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -135,7 +135,7 @@ def __init__( server_args=server_args, gpu_id=gpu_id, # spec workers don't support pipeline parallelism - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, is_draft_worker=True, # The draft runs at absolute target positions. diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index d1b347ce880d..d62d0a9aecf9 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -155,7 +155,7 @@ def __init__( server_args=server_args, gpu_id=gpu_id, # spec workers don't support pipeline parallelism - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, is_draft_worker=True, is_multi_layer_eagle=True, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 2845a2d6c5a2..75d123710184 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -35,6 +35,8 @@ from sglang.srt.constrained.base_grammar_backend import GrammarMask from sglang.srt.distributed.parallel_state import ( GroupCoordinator, + get_self_pp_group, + patch_pipeline_parallel_group, patch_tensor_parallel_group, ) from sglang.srt.environ import envs @@ -669,6 +671,14 @@ def load_token_map(token_map_path: str) -> List[int]: return torch.tensor(hot_token_id, dtype=torch.int64) +@contextmanager +def draft_pp_context(): + # The draft model is one layer and never spans pipeline stages; give it a + # single-member pp group so it initializes as if pp were off. + with patch_pipeline_parallel_group(get_self_pp_group()): + yield + + @contextmanager def draft_tp_context(tp_group: GroupCoordinator): # Draft model doesn't use dp and has its own tp group. diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index d830b984883b..3e6a73fb1fbe 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -72,7 +72,7 @@ def __init__( server_args=server_args, gpu_id=gpu_id, # spec workers don't support pipeline parallelism - ps=replace(ps, pp_rank=0), + ps=replace(ps, pp_rank=0, pp_size=1), nccl_port=nccl_port, is_draft_worker=True, # The draft runs at absolute target positions. diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index 1207b13f9510..cd156054be3e 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -21,6 +21,17 @@ from sglang.srt.state_capturer.base import BaseTopkCapturer +def _is_scattered_a2a_backend() -> bool: + """DeepEP-class a2a dispatchers hand the MoE layer only this attention + rank's DP-local tokens (dispatch happens after top-k), so the capturer + must attn-TP-gather at capture time and read back from the buffer head. + DeepEP v2 shares this token topology with legacy DeepEP; misclassifying + it as a TP-MoE backend would make dp_rank > 0 read unwritten buffer rows. + """ + backend = get_moe_a2a_backend() + return backend.is_deepep() or backend.is_deepep_v2() + + class RoutedExpertsCapturer(BaseTopkCapturer): """Capturer for routed experts with host buffer. @@ -89,11 +100,12 @@ def __init__( device_topk_size=topk_size + num_fused_shared_experts, ) - # DeepEP a2a path: each attn-TP rank only sees its scattered slice of + # DeepEP-class a2a path (see _is_scattered_a2a_backend): each attn-TP + # rank only sees its scattered slice of # topk_ids. All-gather across attn-TP at capture time so device_cache # holds the full batch and the existing _get_local_slice / D2H sync # paths work unchanged. Pre-allocate the gather target. - if get_moe_a2a_backend().is_deepep(): + if _is_scattered_a2a_backend(): attn_tp_size = ( get_parallel().attn_tp_size if is_dp_attention_enabled() else 1 ) @@ -107,7 +119,7 @@ def __init__( ) def capture(self, layer_id: int, topk_indices: torch.Tensor): - if get_moe_a2a_backend().is_deepep(): + if _is_scattered_a2a_backend(): local_topk = topk_indices topk_indices = self.gather_buffer[ : local_topk.size(0) * get_parallel().attn_tp_size @@ -121,10 +133,11 @@ def _get_local_slice( can_run_graph: bool, cuda_graph_batch: Optional[int], ) -> torch.Tensor: - # Under DeepEP, capture() already attn_tp_all_gathered into the head of + # Under DeepEP-class backends, capture() already attn_tp_all_gathered + # into the head of # the per-rank buffer, so the local DP rank's data lives at [0:N_local] # rather than at the global [start_pos:end_pos] offset. - if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep(): + if is_dp_attention_enabled() and not _is_scattered_a2a_backend(): # GPU->CPU sync would break overlap; operate on CPU directly. local_start_pos, local_num_tokens = get_dp_local_slice_cpu( forward_batch, can_run_graph, cuda_graph_batch diff --git a/test/registered/attention/unittests/mamba/test_replay_state_indices_validator.py b/test/registered/attention/unittests/mamba/test_replay_state_indices_validator.py new file mode 100644 index 000000000000..75ed2be1cd31 --- /dev/null +++ b/test/registered/attention/unittests/mamba/test_replay_state_indices_validator.py @@ -0,0 +1,70 @@ +import sys +import unittest +from pathlib import Path + +import torch + +# Keep this CPU-only contract test lightweight: importing the public sglang +# package initializes unrelated frontend/runtime dependencies. +sys.path.insert( + 0, + str( + Path(__file__).resolve().parents[5] + / "python/sglang/srt/layers/attention/mamba" + ), +) +from replay_state_indices_validator import validate_replay_state_indices_cpu + + +class TestReplayStateIndicesValidator(unittest.TestCase): + def test_valid_unique_live_slots_and_padding(self): + validate_replay_state_indices_cpu( + torch.tensor([0, 2, 9, -1, -1], dtype=torch.int32), + valid_bs=3, + total_bs=5, + num_state_slots=10, + ) + + def test_rejects_duplicate_live_slot(self): + with self.assertRaisesRegex(AssertionError, r"duplicate_slots=\[7\]"): + validate_replay_state_indices_cpu( + torch.tensor([7, 2, 7, -1], dtype=torch.int32), + valid_bs=3, + total_bs=4, + num_state_slots=10, + ) + + def test_rejects_out_of_range_live_slots(self): + for bad_slot in (-1, -2, 10): + with self.subTest(bad_slot=bad_slot): + with self.assertRaisesRegex(AssertionError, "live rows"): + validate_replay_state_indices_cpu( + torch.tensor([3, bad_slot, -1], dtype=torch.int64), + valid_bs=2, + total_bs=3, + num_state_slots=10, + ) + + def test_rejects_non_sentinel_padding(self): + with self.assertRaisesRegex(AssertionError, "padded rows"): + validate_replay_state_indices_cpu( + torch.tensor([3, 5, 5], dtype=torch.int32), + valid_bs=2, + total_bs=3, + num_state_slots=10, + ) + + def test_requires_cpu_tensor(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is unavailable") + with self.assertRaisesRegex(ValueError, "copied to CPU"): + validate_replay_state_indices_cpu( + torch.tensor([1], dtype=torch.int32, device="cuda"), + valid_bs=1, + total_bs=1, + num_state_slots=2, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/ep/test_routed_experts_dp_readback.py b/test/registered/ep/test_routed_experts_dp_readback.py new file mode 100644 index 000000000000..e84750c0228e --- /dev/null +++ b/test/registered/ep/test_routed_experts_dp_readback.py @@ -0,0 +1,209 @@ +"""DP>1 readback of routed experts over DeepEP-class a2a backends. + +With DP attention + a DeepEP-class a2a backend, the MoE layer sees only the +attention rank's DP-local tokens, so RoutedExpertsCapturer must gather at +capture time and read back from the buffer head. If the backend is not +recognized, requests owned by dp_rank > 0 read unwritten buffer rows and +silently return garbage expert ids (dp_rank 0 sits at offset 0 and looks +correct, which is why a DP>1 test is required). + +Oracle: solo-vs-concurrent consistency. A request served alone is correct +even on a misclassifying tree (with the other rank empty, the global offset +degenerates to 0), so its per-token expert sets form a valid baseline. The +same prompts served concurrently must reproduce those sets; a misclassified +backend instead reads whatever the offset region holds (often well-formed +rows belonging to other tokens or graph warmup, which per-row validity +checks cannot catch). Radix cache is disabled so the concurrent phase cannot +serve cached prefix rows written by the solo phase. + +Uses a dummy-weight single-layer 24-expert DeepSeek-V3 so each server boots +in seconds (same pattern as test_deepseek_v3_cutedsl_4gpu.py); generation +quality is irrelevant — only the capture/readback plumbing is under test. +""" + +import concurrent.futures +import json +import os +import unittest + +import numpy as np +import pybase64 +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=900, stage="base-c", runner_config="deepep-8-gpu-h200") + +_MODEL = os.environ.get("SGLANG_ROUTED_EXPERTS_TEST_MODEL", "deepseek-ai/DeepSeek-V3") +_NUM_EXPERTS = 24 +_NUM_LAYERS = 1 +_TOPK = 8 # DeepSeek-V3 num_experts_per_tok + +_DUMMY_WEIGHT_ENV = { + # Dummy random weights legitimately produce NaN logits; sanitize instead + # of crashing (same rationale as test_deepseek_v3_cutedsl_4gpu.py). + "SGLANG_ENABLE_ASYNC_ASSERT": "0", + "SGLANG_SANITIZE_NAN_LOGITS": "1", + "SGLANG_CUDA_COREDUMP": "0", + "CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "0", + "SGLANG_CUDA_COREDUMP_BEFORE_CRASH": "0", +} + + +def _deep_ep_has(attr: str) -> bool: + try: + import deep_ep # noqa: F401 + except ImportError: + return False + return hasattr(deep_ep, attr) + + +class _ReadbackMixin: + backend_args: list + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--load-format", + "dummy", + "--json-model-override-args", + json.dumps( + { + "num_hidden_layers": _NUM_LAYERS, + "first_k_dense_replace": 0, + "n_routed_experts": _NUM_EXPERTS, + } + ), + "--tp", + "2", + "--dp", + "2", + "--ep", + "2", + "--enable-dp-attention", + "--enable-return-routed-experts", + "--disable-cuda-graph", + "--disable-radix-cache", + "--mem-fraction-static", + "0.5", + *cls.backend_args, + ] + cls.process = popen_launch_server( + _MODEL, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + env={ + **os.environ, + **_DUMMY_WEIGHT_ENV, + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", + "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def _one_request(self, i: int): + resp = requests.post( + self.base_url + "/generate", + json={ + "text": f"{self._WORDS[i]} is item number {i}. Describe it in detail.", + "sampling_params": {"max_new_tokens": 24, "temperature": 0}, + "return_routed_experts": True, + }, + timeout=300, + ) + self.assertEqual(resp.status_code, 200) + meta = resp.json()["meta_info"] + self.assertIn("routed_experts", meta) + arr = np.frombuffer(pybase64.b64decode(meta["routed_experts"]), dtype=np.int32) + self.assertEqual( + arr.size % (_NUM_LAYERS * _TOPK), + 0, + f"req{i}: payload size {arr.size} not a multiple of layers*topk", + ) + rows = arr.reshape(-1, _NUM_LAYERS, _TOPK) + self.assertGreater(rows.shape[0], 0) + self.assertTrue( + bool(((rows >= 0) & (rows < _NUM_EXPERTS)).all()), + f"req{i}: expert id out of range [{rows.min()}, {rows.max()}]", + ) + return rows + + _WORDS = ["Alpha", "Bravo", "Charlie", "Delta", "Echo", "Foxtrot"] + _N_REQ = 6 + + def test_dp2_readback(self): + # Phase 1 — solo baselines: sequential requests leave the other DP + # rank empty, the global offset degenerates to 0, and the readback is + # correct even when the backend is misclassified. + solo = [self._one_request(i) for i in range(self._N_REQ)] + + # Phase 2 — the same prompts concurrently: joint forward batches give + # dp_rank > 0 requests a non-zero global offset, which is exactly the + # path a misclassified backend gets wrong. + with concurrent.futures.ThreadPoolExecutor(max_workers=self._N_REQ) as ex: + conc = list(ex.map(self._one_request, range(self._N_REQ))) + + for i in range(self._N_REQ): + a, b = solo[i], conc[i] + n = min(a.shape[0], b.shape[0]) + total = match = 0 + for t in range(n): + for layer in range(_NUM_LAYERS): + total += 1 + if set(a[t, layer].tolist()) == set(b[t, layer].tolist()): + match += 1 + frac = match / max(1, total) + self.assertGreaterEqual( + frac, + 0.9, + f"req{i}: only {frac:.1%} of per-token expert sets match the " + "solo baseline — the capturer is reading rows that belong to " + "other tokens (DeepEP-class backend misclassification)", + ) + + +@unittest.skipUnless(_deep_ep_has("Buffer"), "DeepEP (v1 Buffer) not installed") +class TestRoutedExpertsReadbackDeepEP(_ReadbackMixin, CustomTestCase): + backend_args = [ + "--moe-a2a-backend", + "deepep", + "--deepep-mode", + "low_latency", + "--deepep-dispatcher-output-dtype", + "fp8", + "--moe-runner-backend", + "deep_gemm", + ] + + +@unittest.skipUnless( + _deep_ep_has("ElasticBuffer"), "DeepEP v2 (ElasticBuffer) not installed" +) +class TestRoutedExpertsReadbackDeepEPv2(_ReadbackMixin, CustomTestCase): + backend_args = [ + "--moe-a2a-backend", + "deepep_v2", + "--deepep-v2-mode", + "direct", + "--deepep-v2-dispatcher-output-dtype", + "fp8", + "--moe-runner-backend", + "deep_gemm", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/attention/test_gdn_decode_fused_proj_conv.py b/test/registered/kernels/ops/attention/test_gdn_decode_fused_proj_conv.py new file mode 100644 index 000000000000..874b3c3b8a37 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_gdn_decode_fused_proj_conv.py @@ -0,0 +1,395 @@ +import unittest + +import torch + +from sglang.kernels.ops.attention.triton_gdn_fused_proj import ( + can_use_fused_qkvzba_causal_conv1d_update_contiguous, + fused_qkvzba_causal_conv1d_update_contiguous, + fused_qkvzba_split_reshape_cat_contiguous, +) +# This is also the update implementation imported directly by GDNBackend on +# CUDA; the presence of the optional sgl_kernel AOT extension does not reroute +# GDN decode through srt.layers.attention.mamba.causal_conv1d. +from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-large") + + +def _reference( + qkvz, + ba, + state, + weight, + bias, + indices, + *, + qkv_dim, + v_dim, + num_v_heads, + head_v_dim, + activation, +): + qkv = qkvz[:, :qkv_dim] + out = torch.empty_like(qkv) + state_out = state.clone() + width = weight.shape[1] + for row, slot_tensor in enumerate(indices.cpu()): + slot = int(slot_tensor) + if slot < 0 or slot >= state.shape[0]: + out[row].copy_(qkv[row]) + continue + # The deployed direct-Triton decode wrapper uses an effective + # state_len=width-1 even if the physical cache envelope is wider. + history = state[slot, :, : width - 1].float() + values = torch.cat((history, qkv[row, :, None].float()), dim=-1) + acc = (values * weight.float()).sum(dim=-1) + if bias is not None: + acc = acc + bias.float() + if activation in ("silu", "swish"): + acc = torch.nn.functional.silu(acc) + out[row].copy_(acc.to(qkv.dtype)) + if width > 2: + state_out[slot, :, : width - 2].copy_(state[slot, :, 1 : width - 1]) + state_out[slot, :, width - 2].copy_(qkv[row]) + + z = qkvz[:, qkv_dim:].reshape(-1, num_v_heads, head_v_dim).contiguous() + b, a = ba.split([num_v_heads, num_v_heads], dim=-1) + return out, z, b.contiguous(), a.contiguous(), state_out + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") +class TestGDNDecodeFusedProjectionConv1D(unittest.TestCase): + def test_contiguous_unpack_ratio8_microbenchmark_baseline(self): + batch = 9 + num_qk_heads = 1 + num_v_heads = 8 + head_dim = 128 + qkv_dim = (2 * num_qk_heads + num_v_heads) * head_dim + v_dim = num_v_heads * head_dim + qkvz = torch.randn( + batch, + qkv_dim + v_dim, + device="cuda", + dtype=torch.bfloat16, + ) + ba = torch.randn( + batch, + 2 * num_v_heads, + device="cuda", + dtype=torch.bfloat16, + ) + mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous( + qkvz, + ba, + num_qk_heads, + num_v_heads, + head_dim, + head_dim, + ) + torch.testing.assert_close(mixed_qkv, qkvz[:, :qkv_dim]) + torch.testing.assert_close( + z, qkvz[:, qkv_dim:].reshape(batch, num_v_heads, head_dim) + ) + torch.testing.assert_close(b, ba[:, :num_v_heads]) + torch.testing.assert_close(a, ba[:, num_v_heads:]) + + def _run_case( + self, + *, + batch, + q_dim, + k_dim, + v_dim, + num_v_heads, + head_v_dim, + width, + state_len, + dtype, + with_bias, + activation, + strided_state=False, + with_padding=False, + ): + torch.manual_seed(17) + device = "cuda" + qkv_dim = q_dim + k_dim + v_dim + qkvz = torch.randn(batch, qkv_dim + v_dim, device=device, dtype=dtype) + ba = torch.randn(batch, 2 * num_v_heads, device=device, dtype=dtype) + weight = torch.randn(qkv_dim, width, device=device, dtype=dtype) * 0.1 + bias = ( + torch.randn(qkv_dim, device=device, dtype=dtype) * 0.1 + if with_bias + else None + ) + slots = batch + 3 + if strided_state: + backing = torch.randn( + slots, + state_len, + qkv_dim * 2, + device=device, + dtype=dtype, + ) + state = backing[:, :, ::2].transpose(1, 2) + self.assertFalse(state.is_contiguous()) + else: + state = torch.randn(slots, qkv_dim, state_len, device=device, dtype=dtype) + indices = torch.randperm(slots, device=device, dtype=torch.int64)[:batch] + if with_padding: + indices[-1] = -1 + indices = indices.to(torch.int32) + + ref = _reference( + qkvz, + ba, + state, + weight, + bias, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation=activation, + ) + state_test = state.clone(memory_format=torch.preserve_format) + out, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous( + qkvz, + ba, + state_test, + weight, + bias, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation=activation, + ) + torch.cuda.synchronize() + + atol = 2e-2 if dtype == torch.bfloat16 else 3e-3 + output_max_diff = (out.float() - ref[0].float()).abs().max().item() + state_max_diff = (state_test.float() - ref[4].float()).abs().max().item() + print( + "case " + f"B={batch} QKV={qkv_dim} V={v_dim} W={width} " + f"dtype={dtype} output_max_diff={output_max_diff:.8g} " + f"state_max_diff={state_max_diff:.8g}", + flush=True, + ) + torch.testing.assert_close(out, ref[0], rtol=0, atol=atol) + torch.testing.assert_close(z, ref[1], rtol=0, atol=0) + torch.testing.assert_close(b, ref[2], rtol=0, atol=0) + torch.testing.assert_close(a, ref[3], rtol=0, atol=0) + torch.testing.assert_close(state_test, ref[4], rtol=0, atol=0) + self.assertEqual(b.data_ptr() % 32, 0) + self.assertEqual(a.data_ptr() % 32, 0) + + def test_random_shapes_widths_dtypes_and_state_updates(self): + cases = ( + # Small boundary shapes. + dict( + batch=1, + q_dim=16, + k_dim=16, + v_dim=32, + num_v_heads=2, + head_v_dim=16, + width=2, + state_len=1, + dtype=torch.float16, + with_bias=False, + activation=None, + ), + dict( + batch=17, + q_dim=32, + k_dim=32, + v_dim=64, + num_v_heads=4, + head_v_dim=16, + width=3, + state_len=5, + dtype=torch.bfloat16, + with_bias=True, + activation="silu", + strided_state=True, + with_padding=True, + ), + # Qwen3.5-35B TP16 local GDN dimensions. + dict( + batch=32, + q_dim=128, + k_dim=128, + v_dim=256, + num_v_heads=2, + head_v_dim=128, + width=4, + state_len=3, + dtype=torch.bfloat16, + with_bias=False, + activation="silu", + ), + # Large TP-local GDN dimensions with an 8:1 value/key head ratio. + dict( + batch=8, + q_dim=128, + k_dim=128, + v_dim=1024, + num_v_heads=8, + head_v_dim=128, + width=4, + state_len=3, + dtype=torch.bfloat16, + with_bias=False, + activation="silu", + ), + ) + for case in cases: + with self.subTest(case=case): + self._run_case(**case) + + def test_cuda_graph_replay(self): + batch = 4 + qkv_dim, v_dim, num_v_heads, head_v_dim = 128, 64, 2, 32 + qkvz = torch.randn(batch, qkv_dim + v_dim, device="cuda", dtype=torch.bfloat16) + ba = torch.randn(batch, 2 * num_v_heads, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(qkv_dim, 4, device="cuda", dtype=torch.bfloat16) + state = torch.randn(batch + 1, qkv_dim, 3, device="cuda", dtype=torch.bfloat16) + initial_state = state.clone() + indices = torch.arange(batch, device="cuda", dtype=torch.int32) + # Compile before capture; Triton compilation itself is not graph-safe. + fused_qkvzba_causal_conv1d_update_contiguous( + qkvz, + ba, + state, + weight, + None, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation="silu", + ) + state.copy_(initial_state) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = fused_qkvzba_causal_conv1d_update_contiguous( + qkvz, + ba, + state, + weight, + None, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation="silu", + ) + state.copy_(initial_state) + graph.replay() + ref_state = initial_state.clone() + ref_qkv, ref_z, ref_b, ref_a = fused_qkvzba_split_reshape_cat_contiguous( + qkvz, + ba, + 1, + num_v_heads, + 32, + head_v_dim, + ) + ref_qkv = causal_conv1d_update( + ref_qkv, + ref_state, + weight, + None, + "silu", + conv_state_indices=indices, + ) + torch.testing.assert_close(captured[0], ref_qkv, rtol=0, atol=0) + torch.testing.assert_close(captured[1], ref_z, rtol=0, atol=0) + torch.testing.assert_close(captured[2], ref_b, rtol=0, atol=0) + torch.testing.assert_close(captured[3], ref_a, rtol=0, atol=0) + torch.testing.assert_close(state, ref_state, rtol=0, atol=0) + + def test_out_of_range_state_slots_are_safely_masked(self): + torch.manual_seed(29) + batch = 4 + qkv_dim, v_dim, num_v_heads, head_v_dim = 128, 64, 2, 32 + qkvz = torch.randn( + batch, qkv_dim + v_dim, device="cuda", dtype=torch.bfloat16 + ) + ba = torch.randn( + batch, 2 * num_v_heads, device="cuda", dtype=torch.bfloat16 + ) + weight = torch.randn(qkv_dim, 4, device="cuda", dtype=torch.bfloat16) + state = torch.randn(3, qkv_dim, 3, device="cuda", dtype=torch.bfloat16) + # -1 is the expected padding sentinel; -2 and len(state) exercise the + # hard lower/upper bounds. Slot 1 remains a normal live update. + indices = torch.tensor([-2, -1, state.shape[0], 1], device="cuda") + indices = indices.to(torch.int32) + + ref = _reference( + qkvz, + ba, + state, + weight, + None, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation="silu", + ) + state_test = state.clone() + out, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous( + qkvz, + ba, + state_test, + weight, + None, + indices, + qkv_dim=qkv_dim, + v_dim=v_dim, + num_v_heads=num_v_heads, + head_v_dim=head_v_dim, + activation="silu", + ) + torch.cuda.synchronize() + + torch.testing.assert_close(out, ref[0], rtol=0, atol=2e-2) + torch.testing.assert_close(z, ref[1], rtol=0, atol=0) + torch.testing.assert_close(b, ref[2], rtol=0, atol=0) + torch.testing.assert_close(a, ref[3], rtol=0, atol=0) + torch.testing.assert_close(state_test, ref[4], rtol=0, atol=0) + + def test_fp8_activation_is_an_explicit_fallback(self): + if not hasattr(torch, "float8_e4m3fn"): + self.skipTest("PyTorch has no FP8 dtype") + qkvz = torch.empty(1, 96, device="cuda", dtype=torch.float8_e4m3fn) + ba = torch.empty(1, 4, device="cuda", dtype=torch.bfloat16) + state = torch.empty(2, 64, 3, device="cuda", dtype=torch.bfloat16) + weight = torch.empty(64, 4, device="cuda", dtype=torch.bfloat16) + indices = torch.zeros(1, device="cuda", dtype=torch.int32) + eligible, reason = can_use_fused_qkvzba_causal_conv1d_update_contiguous( + qkvz, + ba, + state, + weight, + None, + indices, + qkv_dim=64, + v_dim=32, + num_v_heads=2, + activation="silu", + ) + self.assertFalse(eligible) + self.assertIn("FP16, BF16, or FP32", reason) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py index 99ecf1a0b229..7007fd9897fc 100644 --- a/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py +++ b/test/registered/kernels/ops/moe/test_minimax_quant_scatter.py @@ -1,6 +1,7 @@ import random import sys from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch @@ -159,6 +160,42 @@ def test_compact_all_tokens_uses_tight_routing_independent_bound( ) +def test_compact_eager_keeps_masked_layout_for_cuda_graph(monkeypatch): + config = MoeRunnerConfig( + num_experts=128, + num_local_experts=16, + hidden_size=2048, + intermediate_size_per_partition=4096, + top_k=4, + activation="silu", + is_gated=True, + inplace=False, + ) + monkeypatch.setattr( + deep_gemm_runner.envs.SGLANG_OPT_DG_COMPACT_EAGER, "get", lambda: True + ) + capture = SimpleNamespace(disable_dispose_tensor=False) + monkeypatch.setattr( + deep_gemm_runner, "get_flags", lambda: SimpleNamespace(capture=capture) + ) + hidden_states = torch.empty((128, 2048), device="meta") + quant_info = DeepGemmMoeQuantInfo( + w13_weight=torch.empty((1, 4096, 1), dtype=torch.float8_e4m3fn), + w2_weight=torch.empty((1, 2048, 1), dtype=torch.float8_e4m3fn), + use_fp8=True, + block_shape=[128, 128], + ) + with envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.override("masked"): + assert not deep_gemm_runner._should_use_masked_standard_layout( + config, quant_info, hidden_states + ) + + capture.disable_dispose_tensor = True + assert deep_gemm_runner._should_use_masked_standard_layout( + config, quant_info, hidden_states + ) + + def test_standard_layout_auto_memory_policy(monkeypatch): config = MoeRunnerConfig( num_experts=512, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 1fe8f0a47b02..236715749b72 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -1,11 +1,15 @@ +import threading import unittest from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import numpy as np import torch from sglang.srt.disaggregation.base.conn import KVArgs, StateType +from sglang.srt.disaggregation.common.staging_handler import ( + handle_staging_req, +) from sglang.srt.disaggregation.common.utils import ( group_concurrent_contiguous, pack_int_lists, @@ -13,6 +17,7 @@ unpack_int_lists, unpack_list_of_buffers, ) +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager from sglang.srt.disaggregation.utils import ( MetadataBuffers, get_dsv4_c128_state_indices, @@ -101,6 +106,84 @@ def test_mismatched_nonempty_lengths_raise(self): group_concurrent_contiguous(self._arr([1, 2, 3]), self._arr([1, 2])) +class TestMooncakePPStaging(unittest.TestCase): + def test_staging_response_targets_requesting_pp_rank(self): + sock = Mock() + receiver = SimpleNamespace( + chunk_staging_infos=[], + _connect_to_bootstrap_server=Mock(return_value=(sock, threading.Lock())), + ) + allocator = SimpleNamespace( + assign=Mock(return_value=(3, 128, 0)), total_size=1 << 20 + ) + kv_args = SimpleNamespace( + page_size=64, + kv_item_lens=[4096, 4096], + total_kv_head_num=4, + engine_rank=0, + ) + target = {"pp_rank": 3} + + handle_staging_req( + [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"], + allocator, + kv_args, + attn_tp_size=16, + prefill_attn_tp_size=1, + kv_buffer_tensors=None, + room_receivers={7: receiver}, + room_bootstrap={7: [{"pp_rank": 2}, target]}, + ) + + receiver._connect_to_bootstrap_server.assert_called_once_with(target) + sock.send_multipart.assert_called_once() + + @patch( + "sglang.srt.disaggregation.common.staging_buffer.gather_all_layers_to_staging" + ) + def test_pp_stage_writes_its_global_layer_slots(self, gather): + manager = object.__new__(MooncakeKVManager) + tensor = SimpleNamespace(shape=(1, 1, 8), element_size=lambda: 2) + manager.kv_buffer_tensors = { + "k_buffers": [tensor], + "v_buffers": [tensor], + "page_size": 2, + } + manager.attn_tp_size = 1 + manager.pp_size = 16 + manager.kv_args = SimpleNamespace( + engine_rank=0, + gpu_id=0, + total_kv_head_num=4, + kv_head_num=4, + kv_layer_ids=[7, 7], + ) + manager._transfer_data = Mock(return_value=0) + staging = SimpleNamespace(fits=lambda size: True, get_ptr=lambda: 0x9000) + + ret = manager.send_kvcache_staged( + "peer", + np.array([1, 2], dtype=np.int32), + dst_staging_ptr=0x100000, + dst_staging_size=1 << 20, + dst_tp_rank=0, + dst_attn_tp_size=16, + dst_kv_item_len=128, + dst_layer_ids=[3, 7, 11, 3, 7, 11], + staging_buffer=staging, + ) + + self.assertEqual(ret, 0) + gather.assert_called_once() + manager._transfer_data.assert_called_once_with( + "peer", + [ + (0x9000, 0x100000 + 64, 64), + (0x9000 + 64, 0x100000 + 4 * 64, 64), + ], + ) + + class TestEagleDsaSeedTransfer(unittest.TestCase): @staticmethod def _make_req(seed, metadata_buffer_index=0): diff --git a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py new file mode 100644 index 000000000000..e3a8bc7f4cdb --- /dev/null +++ b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py @@ -0,0 +1,284 @@ +"""Unit tests for full-attention KV transfer with prefill pp_size > 1 on +hybrid-linear models (HybridLinearKVPool).""" + +import unittest +from types import SimpleNamespace + +import numpy as np + +from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager +from sglang.srt.disaggregation.prefill import _transfer_start_layer +from sglang.srt.disaggregation.utils import ( + build_kv_layer_ids, + build_transfer_entry_pairs, +) +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def _full_attention_ids(*, num_layers: int, interval: int) -> list: + return [i for i in range(num_layers) if i % interval == interval - 1] + + +def _hybrid_pool(*, start_layer: int) -> HybridLinearKVPool: + pool = HybridLinearKVPool.__new__(HybridLinearKVPool) + pool.start_layer = start_layer + return pool + + +class TestTransferStartLayer(CustomTestCase): + """Bug regression: with prefill pp_size=2 on a 60-layer hybrid-linear model + (full_attention_interval=4), stage 1's pool.start_layer is 30 — a global + layer index counting linear layers. The decode peer's KV pointer list is + dense over the 15 full-attention layers only, so slicing dst[30:38] yielded + [] and an IndexError in mooncake send_kvcache_slice. The transfer offset + must be the count of full-attention layers before the stage boundary.""" + + def test_hybrid_stage1_translates_to_full_attention_offset(self): + cfg = SimpleNamespace( + full_attention_layer_ids=_full_attention_ids(num_layers=60, interval=4) + ) + self.assertEqual( + _transfer_start_layer( + pool=_hybrid_pool(start_layer=30), hf_text_config=cfg + ), + 7, + ) + + def test_hybrid_stage0_is_zero(self): + cfg = SimpleNamespace( + full_attention_layer_ids=_full_attention_ids(num_layers=60, interval=4) + ) + self.assertEqual( + _transfer_start_layer(pool=_hybrid_pool(start_layer=0), hf_text_config=cfg), + 0, + ) + + def test_non_hybrid_pool_keeps_global_start_layer(self): + cfg = SimpleNamespace(full_attention_layer_ids=[]) + self.assertEqual( + _transfer_start_layer( + pool=SimpleNamespace(start_layer=30), hf_text_config=cfg + ), + 30, + ) + + +class _RecordingKVManager: + get_mha_kv_ptrs_with_pp = CommonKVManager.get_mha_kv_ptrs_with_pp + + def __init__(self, *, prefill_start_layer: int, pp_size: int): + self.is_mla_backend = False + self.is_hybrid_mla_backend = False + self.enable_custom_mem_pool = False + self.pp_size = pp_size + self.kv_args = SimpleNamespace(prefill_start_layer=prefill_start_layer) + self.blocks = [] + + def _transfer_data(self, mooncake_session_id, transfer_blocks): + self.blocks.extend(transfer_blocks) + return 0 + + +class TestHybridSendUsesLayerIdPairing(CustomTestCase): + """Bug regression: a hybrid-linear (non-MLA-flagged) backend fell into the + positional MHA slicing path of _send_kvcache_generic even when both peers + published layer ids. For a stage with F full-attention layers against a + decode peer with N (F < N, F not dividing N), the draft-KV modulo heuristic + silently placed the V block at F * (N // F) instead of N — wrong layers + transferred, no error. With layer ids published on both sides the pairing + must be exact.""" + + def _run_case( + self, *, model_full_ids: list, stage_full_ids: list, start_offset: int + ): + num_stage = len(stage_full_ids) + num_model = len(model_full_ids) + src_ptrs = [1000 + i for i in range(2 * num_stage)] + dst_ptrs = [2000 + i for i in range(2 * num_model)] + item_lens = [10 + i for i in range(2 * num_stage)] + manager = _RecordingKVManager(prefill_start_layer=start_offset, pp_size=2) + rc = MooncakeKVManager._send_kvcache_generic( + manager, + mooncake_session_id="session", + src_data_ptrs=src_ptrs, + dst_data_ptrs=dst_ptrs, + item_lens=item_lens, + prefill_data_indices=np.array([0], dtype=np.int32), + dst_data_indices=np.array([0], dtype=np.int32), + executor=None, + src_layer_ids=stage_full_ids * 2, + dst_layer_ids=model_full_ids * 2, + ) + self.assertEqual(rc, 0) + expected = [ + (src_ptrs[i], dst_ptrs[start_offset + i], item_lens[i]) + for i in range(num_stage) + ] + [ + ( + src_ptrs[num_stage + i], + dst_ptrs[num_model + start_offset + i], + item_lens[num_stage + i], + ) + for i in range(num_stage) + ] + self.assertEqual(manager.blocks, expected) + + def test_stage1_f8_of_n15(self): + ids = _full_attention_ids(num_layers=60, interval=4) + self._run_case(model_full_ids=ids, stage_full_ids=ids[7:], start_offset=7) + + def test_stage0_f7_of_n15(self): + ids = _full_attention_ids(num_layers=60, interval=4) + self._run_case(model_full_ids=ids, stage_full_ids=ids[:7], start_offset=0) + + def test_f5_of_n12(self): + ids = _full_attention_ids(num_layers=48, interval=4) + self._run_case(model_full_ids=ids, stage_full_ids=ids[:5], start_offset=0) + + +class TestGetMhaKvPtrsWithPp(CustomTestCase): + """Derived property: the modulo heuristic in get_mha_kv_ptrs_with_pp exists + for the decode-has-draft-KV layout [K_main, V_main, draft_K, draft_V]. Pin + that geometry (15 main + 1 draft layer) so a future rewrite of the + heuristic (e.g. to fix the plain-MHA pp>1 F-not-dividing-N case) keeps the + draft case intact.""" + + def test_draft_kv_geometry_selects_main_v_block(self): + manager = SimpleNamespace(kv_args=SimpleNamespace(prefill_start_layer=0)) + src_kv_ptrs = list(range(30)) + dst_kv_ptrs = list(range(100, 132)) + src_k, src_v, dst_k, dst_v, num_layers = ( + CommonKVManager.get_mha_kv_ptrs_with_pp(manager, src_kv_ptrs, dst_kv_ptrs) + ) + self.assertEqual(src_k, src_kv_ptrs[:15]) + self.assertEqual(src_v, src_kv_ptrs[15:]) + self.assertEqual(dst_k, dst_kv_ptrs[:15]) + self.assertEqual(dst_v, dst_kv_ptrs[15:30]) + self.assertEqual(num_layers, 15) + + +class TestBuildTransferEntryPairsDuplicateIds(CustomTestCase): + """Derived property: layer ids repeat across the K and V tensor groups, so + pairing must consume dst occurrences in order (K with K, V with V) rather + than by plain id lookup.""" + + def test_k_then_v_occurrence_ordering(self): + pairs = build_transfer_entry_pairs( + src_layer_ids=[3, 7, 3, 7], + dst_layer_ids=[3, 7, 11, 3, 7, 11], + n_src=4, + n_dst=6, + allow_positional_fallback=False, + ) + self.assertEqual(pairs, [(0, 0), (1, 1), (2, 3), (3, 4)]) + + +def _hybrid_pool_with_ids(*, layer_ids: list) -> HybridLinearKVPool: + pool = HybridLinearKVPool.__new__(HybridLinearKVPool) + pool.full_attention_layer_id_mapping = layer_ids + pool.use_mla = False + return pool + + +class TestBuildKvLayerIds(CustomTestCase): + """Bug regression: enabling EAGLE appended draft KV buffers to kv_data_ptrs + while kv_layer_ids described only the target's entries, so the ids were + suppressed entirely and the transfer fell back to positional slicing. Under + prefill pp_size > 1 that slices the wrong layers -- prefill pp=2 + EAGLE + produced garbled decode output while pp=1 + EAGLE did not.""" + + def _stage1_ids(self) -> list: + full = _full_attention_ids(num_layers=60, interval=4) + return [lid for lid in full if lid >= 30] + + def test_draft_entries_get_a_reserved_band_above_the_target_range(self): + """A draft pool that only reports a layer count, not ids.""" + ids = build_kv_layer_ids( + token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()), + draft_token_to_kv_pool=SimpleNamespace(layer_num=1), + num_draft_entries=2, + num_hidden_layers=60, + ) + stage1 = self._stage1_ids() + # k0..k(L-1) then v0..v(L-1) per pool, and the pools are concatenated -- + # so the band repeats per group after the target's ids, not interleaved. + self.assertEqual(ids, stage1 + stage1 + [60, 60]) + + def test_hybrid_draft_pool_is_remapped_out_of_the_target_range(self): + """The EAGLE draft pool for a hybrid-linear model is itself a + HybridLinearKVPool that numbers its single MTP layer from zero, so its + raw ids collide with target layer 0 and must be remapped into the band.""" + ids = build_kv_layer_ids( + token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()), + draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]), + num_draft_entries=2, + num_hidden_layers=60, + ) + stage1 = self._stage1_ids() + self.assertEqual(ids, stage1 + stage1 + [60, 60]) + + def test_non_hybrid_pool_publishes_nothing(self): + self.assertEqual( + build_kv_layer_ids( + token_to_kv_pool=SimpleNamespace(), + draft_token_to_kv_pool=None, + num_draft_entries=0, + num_hidden_layers=60, + ), + [], + ) + + def test_ragged_draft_registration_is_rejected(self): + with self.assertRaises(RuntimeError): + build_kv_layer_ids( + token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()), + draft_token_to_kv_pool=SimpleNamespace(layer_num=2), + num_draft_entries=3, + num_hidden_layers=60, + ) + + +class TestDraftBandPairsAcrossPipelineStages(CustomTestCase): + """Derived property: a pp=2 prefill stage and a pp=1 decode peer, both with + an EAGLE draft pool, must pair on layer id -- the stage's 8 full-attention + layers land on the decode peer's matching K and V entries, and the draft + band lands on the decode peer's draft entries rather than on layer 0.""" + + def test_stage1_pairs_onto_the_decode_layout(self): + full = _full_attention_ids(num_layers=60, interval=4) + stage1 = [lid for lid in full if lid >= 30] + src = build_kv_layer_ids( + token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=stage1), + draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]), + num_draft_entries=2, + num_hidden_layers=60, + ) + dst = build_kv_layer_ids( + token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=full), + draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]), + num_draft_entries=2, + num_hidden_layers=60, + ) + pairs = build_transfer_entry_pairs( + src, dst, len(src), len(dst), allow_positional_fallback=False + ) + k_offset = len(full) - len(stage1) + self.assertEqual( + pairs, + # K block, then V block, then the two draft entries at the tail. + [(i, k_offset + i) for i in range(len(stage1))] + + [(len(stage1) + i, len(full) + k_offset + i) for i in range(len(stage1))] + + [ + (2 * len(stage1), 2 * len(full)), + (2 * len(stage1) + 1, 2 * len(full) + 1), + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_staging_draft_kv_slots.py b/test/registered/unit/disaggregation/test_staging_draft_kv_slots.py new file mode 100644 index 000000000000..8350b5d54cd0 --- /dev/null +++ b/test/registered/unit/disaggregation/test_staging_draft_kv_slots.py @@ -0,0 +1,133 @@ +"""Staging slot ids stay aligned once a draft KV pool is registered. + +The staging gather writes every k_buffer and then every v_buffer, while +kv_data_ptrs (and therefore kv_layer_ids) is ordered +[K target, V target, K draft, V draft]. Labelling slots with kv_layer_ids +silently pairs a layer's KV with another layer's staging slot as soon as a +draft pool exists. +""" + +import unittest + +from sglang.srt.disaggregation.utils import ( + build_staging_slot_metadata, + build_transfer_entry_pairs, +) +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class _Pool(MHATokenToKVPool): + def __init__(self, tag, layer_ids): + self.k_buffer = [f"{tag}K{i}" for i in layer_ids] + self.v_buffer = [f"{tag}V{i}" for i in layer_ids] + + +class _Wrapper(HybridLinearKVPool): + def __init__(self, inner): + self.full_kv_pool = inner + + +def _kv_layer_ids(target_ids, draft_ids): + """kv_data_ptrs order: K target, V target, K draft, V draft.""" + return list(target_ids) + list(target_ids) + list(draft_ids) + list(draft_ids) + + +class TestStagingDraftKvSlots(CustomTestCase): + def test_draft_slots_follow_gather_order(self): + target, draft = [87, 91], [92] + k_buffers, v_buffers, slot_ids = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids(target, draft), + num_draft_entries=2, + kv_pool=_Pool("t", target), + draft_kv_pool=_Pool("d", draft), + ) + self.assertEqual(k_buffers, ["tK87", "tK91", "dK92"]) + self.assertEqual(v_buffers, ["tV87", "tV91", "dV92"]) + self.assertEqual(slot_ids, [87, 91, 92, 87, 91, 92]) + self.assertNotEqual(slot_ids, _kv_layer_ids(target, draft)) + + def test_without_draft_matches_kv_layer_ids(self): + # The two orders coincide with no draft pool, so every deployment that + # predates draft KV must keep its exact slot labelling. + target = [3, 7] + _, _, slot_ids = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids(target, []), + num_draft_entries=0, + kv_pool=_Pool("t", target), + draft_kv_pool=None, + ) + self.assertEqual(slot_ids, _kv_layer_ids(target, [])) + + def test_pp_stage_pairs_against_full_decode(self): + # A prefill stage holds a slice of the layers while decode holds them + # all, so the ids -- not the positions -- have to drive the pairing. + src = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids([87, 91], [92]), + num_draft_entries=2, + kv_pool=_Pool("t", [87, 91]), + draft_kv_pool=_Pool("d", [92]), + )[2] + decode_target = [3, 7, 11, 87, 91] + dst = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids(decode_target, [92]), + num_draft_entries=2, + kv_pool=_Pool("t", decode_target), + draft_kv_pool=_Pool("d", [92]), + )[2] + pairs = build_transfer_entry_pairs(src, dst, len(src), len(dst)) + self.assertEqual(len(pairs), len(src)) + for i, j in pairs: + self.assertEqual(src[i], dst[j]) + self.assertEqual(len({j for _, j in pairs}), len(pairs)) + + def test_hybrid_wrapper_pools_are_unwrapped(self): + # A hybrid draft pool that is left wrapped looks exactly like a draft + # pool with no buffers, which drops draft KV out of staging. + target, draft = [87, 91], [92] + k_buffers, _, slot_ids = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids(target, draft), + num_draft_entries=2, + kv_pool=_Wrapper(_Pool("t", target)), + draft_kv_pool=_Wrapper(_Pool("d", draft)), + ) + self.assertEqual(k_buffers, ["tK87", "tK91", "dK92"]) + self.assertEqual(slot_ids, [87, 91, 92, 87, 91, 92]) + + def test_undescribable_draft_still_yields_target_buffers(self): + # Returning nothing here left the caller skipping set_kv_buffer_tensors + # entirely, and staging then came up with no buffers at all. + class _NoBuffers: + pass + + k_buffers, v_buffers, slot_ids = build_staging_slot_metadata( + kv_layer_ids=_kv_layer_ids([87], [92]), + num_draft_entries=2, + kv_pool=_Pool("t", [87]), + draft_kv_pool=_NoBuffers(), + ) + self.assertEqual(k_buffers, ["tK87"]) + self.assertEqual(v_buffers, ["tV87"]) + self.assertEqual(slot_ids, []) + + def test_pool_without_contiguous_tensors_is_declined(self): + # MLA pools have no k_buffer/v_buffer to stage; the caller relies on None + # to skip the registration rather than register empty lists. + class _NoBuffers: + pass + + self.assertIsNone( + build_staging_slot_metadata( + kv_layer_ids=[], + num_draft_entries=0, + kv_pool=_NoBuffers(), + draft_kv_pool=None, + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py b/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py new file mode 100644 index 000000000000..302ed72f198c --- /dev/null +++ b/test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py @@ -0,0 +1,411 @@ +import unittest +from unittest import mock + +import torch + +from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( + FlashInferGDNKernel, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _view_with_pointer_mod( + shape: tuple[int, ...], dtype: torch.dtype, pointer_mod: int +) -> torch.Tensor: + numel = 1 + for dim in shape: + numel *= dim + element_size = dtype.itemsize + base = torch.empty(numel + 32 // element_size, dtype=dtype) + for offset in range(32 // element_size): + view = base[offset : offset + numel] + if view.data_ptr() % 32 == pointer_mod: + return view.view(shape) + raise AssertionError(f"Could not construct a pointer with mod32={pointer_mod}") + + +def _make_kernel_without_flashinfer() -> FlashInferGDNKernel: + kernel = object.__new__(FlashInferGDNKernel) + kernel._aligned_input_buffers = {} + kernel._aligned_parameter_cache = {} + kernel._verify_intermediate_buffers = {} + kernel._alignment_fallback_warned = False + return kernel + + +class TestFlashInferGDNAlignment(unittest.TestCase): + def test_extend_writes_directly_to_preallocated_output(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + captured = {} + + def fake_prefill(**kwargs): + captured.update(kwargs) + kwargs["output"].fill_(7.0) + kwargs["output_state"].copy_(kwargs["initial_state"]) + return kwargs["output"], kwargs["output_state"] + + kernel._prefill_fn = fake_prefill + q = torch.ones((1, 3, 1, 4), dtype=torch.bfloat16) + k = torch.ones_like(q) + v = torch.ones((1, 3, 2, 4), dtype=torch.bfloat16) + g = torch.zeros((1, 3, 2), dtype=torch.bfloat16) + beta = torch.ones_like(g) + ssm_states = torch.zeros((3, 2, 4, 4), dtype=torch.bfloat16) + physical_output = torch.empty((1, 5, 2, 4), dtype=v.dtype) + preallocated_output = physical_output[:, :3] + + with mock.patch( + "sglang.kernels.ops.attention.fla.l2norm.l2norm_fwd", + side_effect=lambda tensor: tensor, + ): + result, _, checkpoints = kernel.extend( + q, + k, + v, + g, + beta, + ssm_states=ssm_states, + cache_indices=torch.tensor([1], dtype=torch.int32), + query_start_loc=torch.tensor([0, 3], dtype=torch.int32), + output=preallocated_output, + ) + + self.assertEqual(captured["output"].data_ptr(), preallocated_output.data_ptr()) + self.assertEqual(result.data_ptr(), preallocated_output.data_ptr()) + torch.testing.assert_close(result, torch.full_like(result, 7.0)) + self.assertIsNone(checkpoints) + + def test_ratio8_bs1_split_view_reproduces_under_alignment(self): + # A TP-sharded ratio-8 projection packs [b_local(8) | a_local(8)] in + # BF16. The a view begins 16 bytes after the storage base. At BS=1 + # PyTorch considers the strided view contiguous, so contiguous() is a + # no-op and cannot satisfy FlashInfer's stricter 32-byte ABI. + projected_ba = torch.empty((2, 16), dtype=torch.bfloat16) + _, a = projected_ba.split((8, 8), dim=-1) + + a_bs1 = a[:1] + self.assertEqual(a_bs1.stride(), (16, 1)) + self.assertTrue(a_bs1.is_contiguous()) + self.assertEqual(a_bs1.data_ptr() % 32, 16) + self.assertEqual(a_bs1.contiguous().data_ptr(), a_bs1.data_ptr()) + + # BS>1 exposes the row gap, so contiguous() does allocate a rebased, + # allocator-aligned tensor. This explains why only BS=1 failed. + a_bs2 = a[:2] + self.assertFalse(a_bs2.is_contiguous()) + repaired = a_bs2.contiguous() + self.assertNotEqual(repaired.data_ptr(), a_bs2.data_ptr()) + self.assertEqual(repaired.data_ptr() % 32, 0) + + def test_dynamic_repair_buffer_is_reused_without_allocator_churn(self): + kernel = _make_kernel_without_flashinfer() + source = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16) + source.fill_(1) + + first = kernel._prepare_dynamic_input("decode_a", source) + first_ptr = first.data_ptr() + self.assertEqual(first_ptr % 32, 0) + torch.testing.assert_close(first, source) + self.assertEqual(len(kernel._aligned_input_buffers), 1) + + source.fill_(2) + second = kernel._prepare_dynamic_input("decode_a", source) + self.assertIs(second, first) + self.assertEqual(second.data_ptr(), first_ptr) + torch.testing.assert_close(second, source) + self.assertEqual(len(kernel._aligned_input_buffers), 1) + + # Distinct kernel arguments cannot alias because both are live at the + # FlashInfer call boundary. + other = kernel._prepare_dynamic_input("decode_b", source) + self.assertNotEqual(other.data_ptr(), first_ptr) + self.assertEqual(len(kernel._aligned_input_buffers), 2) + + def test_decode_repairs_read_only_arguments_before_flashinfer(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + captured = {} + + def fake_decode(**kwargs): + captured.update(kwargs) + v = kwargs["v"] + return ( + torch.zeros( + v.shape[0], + 1, + v.shape[2], + v.shape[3], + dtype=v.dtype, + ), + None, + ) + + kernel._decode_fn = fake_decode + + q = torch.empty(1, 1, 1, 128, dtype=torch.bfloat16) + k = torch.empty_like(q) + v = torch.empty(1, 1, 8, 128, dtype=torch.bfloat16) + a = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16) + b = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16) + A_log = _view_with_pointer_mod((8,), torch.float32, 4) + dt_bias = _view_with_pointer_mod((8,), torch.bfloat16, 2) + state = torch.zeros(2, 8, 128, 128, dtype=torch.bfloat16) + cache_indices = _view_with_pointer_mod((1,), torch.int32, 4) + + result = kernel.decode( + q, + k, + v, + a, + b, + A_log=A_log, + dt_bias=dt_bias, + ssm_states=state, + cache_indices=cache_indices, + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + ) + + self.assertEqual(result.shape, (1, 1, 8, 128)) + for name in ( + "q", + "k", + "v", + "A_log", + "a", + "dt_bias", + "b", + "initial_state", + "initial_state_indices", + ): + with self.subTest(name=name): + self.assertEqual(captured[name].data_ptr() % 32, 0) + torch.testing.assert_close(captured["a"], a) + torch.testing.assert_close(captured["b"], b) + + def test_gate_parameter_cache_preserves_backend_dtype_contract(self): + kernel = _make_kernel_without_flashinfer() + A_log = torch.empty(8, dtype=torch.bfloat16) + dt_bias = torch.empty(8, dtype=torch.bfloat16) + + A_log_sm90, _ = kernel._prepare_gate_parameters(A_log, dt_bias) + A_log_sm100, _ = kernel._prepare_gate_parameters( + A_log, dt_bias, A_log_dtype=torch.float32 + ) + + self.assertEqual(A_log_sm90.dtype, torch.bfloat16) + self.assertEqual(A_log_sm100.dtype, torch.float32) + self.assertEqual(A_log_sm90.data_ptr() % 32, 0) + self.assertEqual(A_log_sm100.data_ptr() % 32, 0) + self.assertIs( + kernel._prepare_gate_parameters(A_log, dt_bias)[0], + A_log_sm90, + ) + + def test_mutable_state_falls_back_without_losing_writeback(self): + kernel = _make_kernel_without_flashinfer() + captured = {} + expected = torch.empty(1) + + class FakeFallback: + def decode(self, *args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return expected + + kernel._alignment_fallback_kernel = FakeFallback() + state = _view_with_pointer_mod((2, 8, 4, 4), torch.bfloat16, 16) + q = torch.empty(1, 1, 1, 4, dtype=torch.bfloat16) + k = torch.empty_like(q) + v = torch.empty(1, 1, 8, 4, dtype=torch.bfloat16) + a = torch.empty(1, 1, 8, dtype=torch.bfloat16) + b = torch.empty_like(a) + cache_indices = torch.zeros(1, dtype=torch.int32) + query_start_loc = torch.tensor([0, 1], dtype=torch.int32) + + result = kernel.decode( + q, + k, + v, + a, + b, + A_log=torch.zeros(8), + dt_bias=torch.zeros(8, dtype=torch.bfloat16), + ssm_states=state, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + ) + + self.assertIs(result, expected) + self.assertIs(captured["kwargs"]["ssm_states"], state) + self.assertEqual(len(kernel._aligned_input_buffers), 0) + + def test_mutable_mtp_workspace_falls_back_without_copying(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + captured = {} + expected = torch.empty(1) + + class FakeFallback: + def target_verify(self, **kwargs): + captured.update(kwargs) + return expected + + kernel._alignment_fallback_kernel = FakeFallback() + q = torch.empty(1, 2, 1, 4, dtype=torch.bfloat16) + k = torch.empty_like(q) + v = torch.empty(1, 2, 8, 4, dtype=torch.bfloat16) + a = torch.empty(1, 2, 8, dtype=torch.bfloat16) + b = torch.empty_like(a) + state = torch.empty(2, 8, 4, 4, dtype=torch.bfloat16) + workspace = _view_with_pointer_mod((2, 2, 8, 4, 4), torch.bfloat16, 16) + + result = kernel.target_verify( + torch.zeros(8), + torch.zeros(8, dtype=torch.bfloat16), + q, + k, + v, + a, + b, + ssm_states=state, + cache_indices=torch.zeros(1, dtype=torch.int32), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32), + intermediate_states_buffer=workspace, + intermediate_state_indices=torch.zeros(1, 2, dtype=torch.int32), + cache_steps=2, + retrieve_parent_token=None, + ) + + self.assertIs(result, expected) + self.assertIs(captured["intermediate_states_buffer"], workspace) + self.assertEqual(len(kernel._aligned_input_buffers), 0) + + def test_mtp_padded_capture_uses_stable_exact_batch_workspace_and_copies_back(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + captured_ptrs = [] + + def fake_mtp(**kwargs): + workspace = kwargs["intermediate_states_buffer"] + captured_ptrs.append(workspace.data_ptr()) + self.assertEqual(workspace.shape[0], 8) + for row in range(workspace.shape[0]): + workspace[row].fill_(row + 1) + return torch.zeros_like(kwargs["v"]), None + + kernel._mtp_fn = fake_mtp + workspace = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16) + + def run_once(): + return kernel.target_verify( + torch.zeros(8), + torch.zeros(8, dtype=torch.bfloat16), + torch.empty(1, 16, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 8, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 8, dtype=torch.bfloat16), + torch.empty(1, 16, 8, dtype=torch.bfloat16), + ssm_states=torch.zeros(8, 8, 4, 4, dtype=torch.bfloat16), + cache_indices=torch.zeros(8, dtype=torch.int32), + query_start_loc=torch.arange(0, 18, 2, dtype=torch.int32), + intermediate_states_buffer=workspace, + intermediate_state_indices=torch.arange(8, dtype=torch.int32), + cache_steps=2, + retrieve_parent_token=None, + ) + + self.assertEqual(run_once().shape, (1, 16, 8, 4)) + for row in range(workspace.shape[0]): + torch.testing.assert_close( + workspace[row], torch.full_like(workspace[row], row + 1) + ) + self.assertEqual(len(kernel._verify_intermediate_buffers), 1) + + workspace.zero_() + run_once() + self.assertEqual(captured_ptrs[0], captured_ptrs[1]) + self.assertEqual(len(kernel._verify_intermediate_buffers), 1) + + def test_mtp_pool_sized_batch_keeps_zero_copy_fast_path(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + workspace = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16) + captured = {} + + def fake_mtp(**kwargs): + captured.update(kwargs) + return torch.zeros_like(kwargs["v"]), None + + kernel._mtp_fn = fake_mtp + result = kernel.target_verify( + torch.zeros(8), + torch.zeros(8, dtype=torch.bfloat16), + torch.empty(1, 4, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 4, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 4, 8, 4, dtype=torch.bfloat16), + torch.empty(1, 4, 8, dtype=torch.bfloat16), + torch.empty(1, 4, 8, dtype=torch.bfloat16), + ssm_states=torch.zeros(7, 8, 4, 4, dtype=torch.bfloat16), + cache_indices=torch.zeros(2, dtype=torch.int32), + query_start_loc=torch.arange(0, 6, 2, dtype=torch.int32), + intermediate_states_buffer=workspace, + intermediate_state_indices=torch.arange(7, dtype=torch.int32), + cache_steps=2, + retrieve_parent_token=None, + ) + + self.assertEqual(result.shape, (1, 4, 8, 4)) + self.assertEqual( + captured["intermediate_states_buffer"].data_ptr(), workspace.data_ptr() + ) + self.assertEqual(len(kernel._verify_intermediate_buffers), 0) + + def test_mtp_padded_workspace_is_reused_across_sequential_layer_pools(self): + kernel = _make_kernel_without_flashinfer() + kernel.use_state_pool = True + captured_ptrs = [] + call_value = 0 + + def fake_mtp(**kwargs): + nonlocal call_value + call_value += 1 + scratch = kwargs["intermediate_states_buffer"] + captured_ptrs.append(scratch.data_ptr()) + scratch.fill_(call_value) + return torch.zeros_like(kwargs["v"]), None + + kernel._mtp_fn = fake_mtp + + def run(pool): + kernel.target_verify( + torch.zeros(8), + torch.zeros(8, dtype=torch.bfloat16), + torch.empty(1, 16, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 1, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 8, 4, dtype=torch.bfloat16), + torch.empty(1, 16, 8, dtype=torch.bfloat16), + torch.empty(1, 16, 8, dtype=torch.bfloat16), + ssm_states=torch.zeros(8, 8, 4, 4, dtype=torch.bfloat16), + cache_indices=torch.zeros(8, dtype=torch.int32), + query_start_loc=torch.arange(0, 18, 2, dtype=torch.int32), + intermediate_states_buffer=pool, + intermediate_state_indices=torch.arange(8, dtype=torch.int32), + cache_steps=2, + retrieve_parent_token=None, + ) + + first_pool = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16) + second_pool = torch.zeros_like(first_pool) + run(first_pool) + run(second_pool) + + self.assertEqual(captured_ptrs[0], captured_ptrs[1]) + torch.testing.assert_close(first_pool, torch.ones_like(first_pool)) + torch.testing.assert_close(second_pool, torch.full_like(second_pool, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py b/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py new file mode 100644 index 000000000000..5ab8af045070 --- /dev/null +++ b/test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py @@ -0,0 +1,192 @@ +"""Unit tests for the DeepEP v2 masked-slab repack Triton kernels. + +Covers the corner cases flagged in review: empty expert, single hot expert, +per-expert count near / over max_m (overflow -> fail-fast, not silent truncation), +top-k weight fusion on real rows only, expanded<->slab round-trip layout, and the +fp8 activation+scale path. +""" + +import unittest + +import torch + +from sglang.kernels.ops.moe.ep_moe_kernels import ( + expand_to_masked_slab, + masked_slab_to_expand, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large") + +DEVICE = "cuda" + + +def _build_layout(counts, align, hidden, dtype, with_scale=False, scale_hidden=4): + """Build (recv_x, recv_x_scale, psum, starts, total) for given per-expert counts. + + Mirrors the DeepEP v2 expanded layout: expert e occupies rows + [align(psum[e-1]), psum[e]) with psum[-1] == 0. + """ + starts, psum = [], [] + prev_end = 0 + for c in counts: + start = ((prev_end + align - 1) // align) * align + end = start + c + starts.append(start) + psum.append(end) + prev_end = end + total = max(((prev_end + align - 1) // align) * align, 1) + + # unique, exactly-representable values per real row (row index, kept small) + base = torch.zeros((total, hidden), dtype=torch.float32, device=DEVICE) + for e, (s, c) in enumerate(zip(starts, counts)): + for j in range(c): + base[s + j] = float((s + j) % 200 + 1) + recv_x = base.to(dtype) + + scale = None + if with_scale: + scale = torch.zeros((total, scale_hidden), dtype=torch.float32, device=DEVICE) + for e, (s, c) in enumerate(zip(starts, counts)): + for j in range(c): + scale[s + j] = float((s + j) % 50 + 1) * 0.5 + + psum_t = torch.tensor(psum, dtype=torch.int32, device=DEVICE) + return recv_x, scale, psum_t, starts, total + + +def _real_rows(starts, counts): + rows = [] + for s, c in zip(starts, counts): + rows.extend(range(s, s + c)) + return rows + + +class TestDeepEPv2MaskedSlab(CustomTestCase): + ALIGN = 16 + HIDDEN = 8 + MAX_M = 32 + + def _check_expand_roundtrip(self, counts, dtype, with_scale, topk=False): + recv_x, scale, psum, starts, total = _build_layout( + counts, self.ALIGN, self.HIDDEN, dtype, with_scale=with_scale + ) + E = len(counts) + slab, slab_scale, masked_m = expand_to_masked_slab( + recv_x, scale, psum, E, self.MAX_M, self.ALIGN + ) + + # masked_m == real per-expert count + self.assertEqual(masked_m.tolist(), list(counts)) + self.assertEqual(tuple(slab.shape), (E, self.MAX_M, self.HIDDEN)) + + # slab real rows == source expanded rows + for e, (s, c) in enumerate(zip(starts, counts)): + for j in range(c): + torch.testing.assert_close(slab[e, j].float(), recv_x[s + j].float()) + if with_scale: + torch.testing.assert_close( + slab_scale[e, j].float(), scale[s + j].float() + ) + + # round-trip back to expanded order + weights = None + if topk: + weights = torch.zeros(total, dtype=torch.float32, device=DEVICE) + for r in _real_rows(starts, counts): + weights[r] = 0.25 + (r % 7) * 0.1 + out = masked_slab_to_expand(slab, psum, total, self.ALIGN, topk_weights=weights) + self.assertEqual(tuple(out.shape), (total, self.HIDDEN)) + for e, (s, c) in enumerate(zip(starts, counts)): + for j in range(c): + expected = slab[e, j].float() + if topk: + expected = (expected * weights[s + j]).to(slab.dtype).float() + torch.testing.assert_close(out[s + j].float(), expected) + + def test_roundtrip_bf16(self): + self._check_expand_roundtrip([3, 0, 5, 1], torch.bfloat16, with_scale=False) + + def test_roundtrip_bf16_with_topk_weight(self): + self._check_expand_roundtrip( + [2, 4, 0, 7], torch.bfloat16, with_scale=False, topk=True + ) + + def test_roundtrip_fp8_with_scale(self): + self._check_expand_roundtrip([3, 1, 6, 2], torch.float8_e4m3fn, with_scale=True) + + def test_empty_experts(self): + # all experts empty + self._check_expand_roundtrip([0, 0, 0, 0], torch.bfloat16, with_scale=False) + + def test_single_hot_expert(self): + # one expert holds many tokens, the rest empty + self._check_expand_roundtrip( + [0, self.MAX_M, 0, 0], torch.bfloat16, with_scale=False, topk=True + ) + + def test_count_at_max_m_boundary(self): + # exactly max_m must be kept (no overflow, no truncation) + self._check_expand_roundtrip( + [self.MAX_M, 1, self.MAX_M], torch.bfloat16, with_scale=False + ) + + def test_overflow_fails_fast(self): + # one expert exceeds max_m -> must raise, not silently truncate + counts = [self.MAX_M + 1, 2] + recv_x, scale, psum, starts, total = _build_layout( + counts, self.ALIGN, self.HIDDEN, torch.bfloat16 + ) + with self.assertRaises(RuntimeError): + expand_to_masked_slab( + recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN + ) + + +class TestDeepEPv2HandleLifecycle(CustomTestCase): + """CPU-only guards of the dispatch/combine handle lifecycle. + + The guards are ordered before any DeepEP work, so misuse is testable + without deep_ep installed and without a GPU. The positive dispatch -> + combine path needs real ElasticBuffer communication and is covered by the + GPU accuracy runs instead. + """ + + @staticmethod + def _bare_impl(): + from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import _DeepEPv2Impl + + impl = object.__new__(_DeepEPv2Impl) + impl._handle = None + impl._pad_empty_combine = False + return impl + + def test_combine_without_dispatch_raises(self): + impl = self._bare_impl() + with self.assertRaisesRegex(RuntimeError, "without a valid dispatch handle"): + impl.combine(None) + + def test_dispatch_with_unconsumed_handle_raises(self): + impl = self._bare_impl() + impl._handle = object() + with self.assertRaisesRegex(RuntimeError, "unconsumed"): + impl.dispatch(None, None) + + def test_handle_cleared_when_combine_fails(self): + impl = self._bare_impl() + impl._handle = object() + impl._pad_empty_combine = True + + def _boom(): + raise RuntimeError("boom") + + impl._get_buffer = _boom + with self.assertRaisesRegex(RuntimeError, "boom"): + impl.combine(None) + self.assertIsNone(impl._handle) + self.assertFalse(impl._pad_empty_combine) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py new file mode 100644 index 000000000000..16197d843271 --- /dev/null +++ b/test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py @@ -0,0 +1,87 @@ +import unittest + +import torch + +from sglang.kernels.ops.moe.ep_moe_kernels import fused_moe_dispatch_index +from sglang.srt.layers.moe.moe_runner.base import ( + FusedOpPool, + PermuteMethodPool, +) +from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + _max_tokens_per_scattered_source, + _scattered_source_token_counts, + _workspace_size_for_namespace, +) +from sglang.srt.layers.quantization import fp8 # noqa: F401 +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +class TestFlashinferA2AWideEPPlumbing(CustomTestCase): + def test_runner_paths_are_registered(self): + self.assertIn(("flashinfer", "flashinfer_trtllm"), FusedOpPool._fused_funcs) + self.assertIn( + ("flashinfer", "flashinfer_trtllm_routed"), FusedOpPool._fused_funcs + ) + self.assertIn( + ("flashinfer", "deep_gemm"), PermuteMethodPool._pre_permute_methods + ) + self.assertIn( + ("deep_gemm", "flashinfer"), PermuteMethodPool._post_permute_methods + ) + + def test_dp4_tp4_uses_physical_source_rank_geometry(self): + self.assertEqual(_max_tokens_per_scattered_source([2048] * 4, 4), 512) + self.assertEqual(_max_tokens_per_scattered_source([1, 0, 0, 0], 4), 1) + self.assertEqual(_max_tokens_per_scattered_source([7, 3, 2, 1], 4), 2) + self.assertEqual(_max_tokens_per_scattered_source([512] * 16, 1), 512) + + def test_target_and_draft_decode_use_distinct_workspaces(self): + sizes = { + _workspace_size_for_namespace(4096, speculative=speculative) + for speculative in (False, True) + } + self.assertEqual(sizes, {4096, 4224}) + + def test_prefill_ag_expands_dp_counts_to_physical_source_ranks(self): + self.assertEqual( + _scattered_source_token_counts([7, 3], 4), + [2, 2, 2, 1, 1, 1, 1, 0], + ) + self.assertEqual(_scattered_source_token_counts([4] * 16, 1), [4] * 16) + + def test_deepgemm_dispatch_marks_empty_expert_lanes_invalid(self): + topk_ids = torch.tensor([-1, 0, -1, 1], dtype=torch.int32, device="cuda") + masked_m, src2dst = fused_moe_dispatch_index( + topk_ids, num_local_experts=2, m_max=4 + ) + + torch.testing.assert_close( + masked_m, torch.tensor([1, 1], dtype=torch.int32, device="cuda") + ) + torch.testing.assert_close( + src2dst, + torch.tensor([-1, 0, -1, 4], dtype=torch.int32, device="cuda"), + ) + + def test_global_expert_mapping_is_fused_into_dispatch_index(self): + global_ids = torch.tensor( + [-1, 15, 16, 17, 31, 32], dtype=torch.int32, device="cuda" + ) + masked_m, src2dst = fused_moe_dispatch_index( + global_ids, num_local_experts=2, m_max=4, expert_start=16 + ) + + torch.testing.assert_close( + masked_m, torch.tensor([1, 1], dtype=torch.int32, device="cuda") + ) + torch.testing.assert_close( + src2dst, + torch.tensor([-1, -1, 0, 4, -1, -1], dtype=torch.int32, device="cuda"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py b/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py new file mode 100644 index 000000000000..1d1ab54022ae --- /dev/null +++ b/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py @@ -0,0 +1,395 @@ +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.layers.flashinfer_mnnvl_cutedsl import ( + FlashInferMNNVLCuteDSLARFusion, + _with_early_finalize_shared_load, +) +from sglang.srt.layers.flashinfer_provider import _make_provider +from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( + Qwen35MoeFinalizeHandoff, + is_supported_forward_mode, + resolve_max_m, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.models.qwen3_5_text import Qwen3_5ForCausalLM +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="base-c-test-cpu") + + +@dataclass(frozen=True) +class _TestPreset: + load_shared_expert_before_pdl: bool = False + + +@dataclass(frozen=True) +class _TestTarget: + preset: object + + +@dataclass(frozen=True) +class _TestRoutes: + targets: tuple[_TestTarget, ...] + + +@dataclass(frozen=True) +class _TestProfile: + finalize_routes: _TestRoutes + + +@dataclass(frozen=True) +class _TestConfig: + profiles: tuple[_TestProfile, ...] + + +_TEST_DEFAULT_CONFIG = _TestConfig( + profiles=( + _TestProfile( + finalize_routes=_TestRoutes(targets=(_TestTarget(_TestPreset()),)) + ), + ) +) + + +@pytest.mark.parametrize( + ("forward_mode", "expected"), + [ + (ForwardMode.DECODE, True), + (ForwardMode.EXTEND, True), + (ForwardMode.IDLE, False), + (ForwardMode.TARGET_VERIFY, False), + ], +) +def test_supported_forward_modes(forward_mode, expected): + assert is_supported_forward_mode(forward_mode) is expected + + +def test_framework_capacity_is_maximum_of_all_sources(): + graph = SimpleNamespace( + decode=SimpleNamespace(max_bs=512, bs=[1, 64, 256]), + prefill=SimpleNamespace(max_bs=4096, bs=[1024, 2048, 4096]), + ) + server_args = SimpleNamespace( + cuda_graph_config=graph, + cutedsl_moe_max_num_tokens=lambda: 8192, + ) + runner = SimpleNamespace(server_args=server_args, max_running_requests=2048) + + assert resolve_max_m(runner) == 8192 + + +def test_deferred_handoff_reuses_producer_storage(): + m, top_k, hidden_size = 3, 10, 16 + gemm2_out = torch.empty(m * top_k + 4, hidden_size, dtype=torch.bfloat16) + expert_weights = torch.empty(m, top_k, dtype=torch.bfloat16) + permuted_indices = torch.empty(m, top_k, dtype=torch.int32) + gated_shared_output = torch.empty(m, hidden_size, dtype=torch.bfloat16) + deferred = SimpleNamespace( + gemm2_out=gemm2_out, + expert_weights=expert_weights, + expanded_idx_to_permuted_idx=permuted_indices, + top_k=top_k, + ) + + handoff = Qwen35MoeFinalizeHandoff.from_flashinfer( + deferred, + gated_shared_output=gated_shared_output, + m=m, + ) + + assert handoff.routed_output.data_ptr() == gemm2_out.data_ptr() + assert handoff.expert_weights.data_ptr() == expert_weights.data_ptr() + assert handoff.permuted_indices.data_ptr() == permuted_indices.data_ptr() + assert handoff.gated_shared_output is gated_shared_output + + +class _CompleteWorkspace: + def __init__( + self, + tp_size, + tp_rank, + max_token_num, + hidden_dim, + dtype, + *, + group, + top_k, + rms_eps, + routed_scaling_factor, + weight_bias, + include_shared_expert, + add_residual, + write_residual_output, + config=_TEST_DEFAULT_CONFIG, + ): + pass + + def is_buffer_size_sufficient( + self, tp_size, num_tokens, hidden_dim, dtype, use_oneshot=None + ): + pass + + def destroy(self): + pass + + +def _complete_allreduce( + input, + workspace, + pattern, + launch_with_pdl, + residual_in, + residual_out, + norm_out, + rms_gamma, + rms_eps, + weight_bias, + expanded_idx_to_permuted_idx, + expert_scale_factor, + shared_expert_output, +): + pass + + +def _complete_comm(allreduce_fusion=_complete_allreduce): + return SimpleNamespace( + AllReduceFusionPattern=SimpleNamespace( + kARResidualRMSNorm=1, + kMoEFinalizeARResidualRMSNorm=7, + ), + allreduce_fusion=allreduce_fusion, + ) + + +def test_provider_requires_the_stable_backend_specific_abi(): + provider = _make_provider( + _complete_comm(), + _CompleteWorkspace, + default_config=_TEST_DEFAULT_CONFIG, + ) + assert provider is not None + assert provider.workspace_type is _CompleteWorkspace + assert provider.default_config is _TEST_DEFAULT_CONFIG + + class WorkspaceWithoutRoutedScale: + def __init__( + self, + tp_size, + tp_rank, + max_token_num, + hidden_dim, + dtype, + *, + group, + top_k, + rms_eps, + weight_bias, + include_shared_expert, + add_residual, + write_residual_output, + config=_TEST_DEFAULT_CONFIG, + ): + pass + + def is_buffer_size_sufficient( + self, tp_size, num_tokens, hidden_dim, dtype, use_oneshot=None + ): + pass + + def destroy(self): + pass + + assert ( + _make_provider( + _complete_comm(), + WorkspaceWithoutRoutedScale, + default_config=_TEST_DEFAULT_CONFIG, + ) + is None + ) + + +def test_provider_accepts_forward_compatible_kwargs_abi(): + class Workspace: + def __init__(self, **kwargs): + pass + + def is_buffer_size_sufficient(self, **kwargs): + pass + + def destroy(self): + pass + + def allreduce_fusion(**kwargs): + pass + + assert ( + _make_provider( + _complete_comm(allreduce_fusion), + Workspace, + default_config=_TEST_DEFAULT_CONFIG, + ) + is not None + ) + + +def test_provider_rejects_incomplete_unified_api(): + assert ( + _make_provider( + SimpleNamespace( + AllReduceFusionPattern=_complete_comm().AllReduceFusionPattern + ), + _CompleteWorkspace, + default_config=_TEST_DEFAULT_CONFIG, + ) + is None + ) + + def allreduce_without_pdl( + input, + workspace, + pattern, + residual_in, + residual_out, + norm_out, + rms_gamma, + rms_eps, + weight_bias, + expanded_idx_to_permuted_idx, + expert_scale_factor, + shared_expert_output, + ): + pass + + assert ( + _make_provider( + _complete_comm(allreduce_without_pdl), + _CompleteWorkspace, + default_config=_TEST_DEFAULT_CONFIG, + ) + is None + ) + assert ( + _make_provider( + SimpleNamespace( + AllReduceFusionPattern=SimpleNamespace(kARResidualRMSNorm=1), + allreduce_fusion=_complete_allreduce, + ), + _CompleteWorkspace, + default_config=_TEST_DEFAULT_CONFIG, + ) + is None + ) + + +def test_qwen_workspace_config_enables_only_supported_finalize_presets(): + untouched_preset = object() + default_config = _TestConfig( + profiles=( + _TestProfile( + finalize_routes=_TestRoutes( + targets=( + _TestTarget(_TestPreset()), + _TestTarget(untouched_preset), + ) + ) + ), + ) + ) + + qwen_config = _with_early_finalize_shared_load(default_config) + + assert qwen_config is not default_config + assert ( + default_config.profiles[0] + .finalize_routes.targets[0] + .preset.load_shared_expert_before_pdl + is False + ) + assert ( + qwen_config.profiles[0] + .finalize_routes.targets[0] + .preset.load_shared_expert_before_pdl + is True + ) + assert qwen_config.profiles[0].finalize_routes.targets[1].preset is untouched_preset + + +def test_wrapper_calls_only_the_stable_unified_api(): + calls = [] + wrapper = object.__new__(FlashInferMNNVLCuteDSLARFusion) + wrapper.hidden_size = 8 + wrapper.top_k = 2 + wrapper.max_m = 4 + wrapper.rms_epsilon = 1e-5 + wrapper.weight_bias = 0.0 + wrapper.device = torch.device("cpu") + wrapper.workspace = object() + wrapper.supports = lambda m: True + wrapper.provider = SimpleNamespace( + patterns=SimpleNamespace( + kARResidualRMSNorm=1, + kMoEFinalizeARResidualRMSNorm=7, + ), + allreduce_fusion=lambda **kwargs: calls.append(kwargs), + ) + + routed_output = torch.empty(8, 8, dtype=torch.bfloat16) + expert_weights = torch.empty(4, 2, dtype=torch.bfloat16) + permuted_indices = torch.empty(4, 2, dtype=torch.int32) + gated_shared_output = torch.empty(4, 8, dtype=torch.bfloat16) + residual = torch.empty(4, 8, dtype=torch.bfloat16) + gamma = torch.empty(8, dtype=torch.bfloat16) + norm_output = torch.empty_like(residual) + residual_output = torch.empty_like(residual) + + wrapper.moe_finalize_all_reduce_rms_norm( + routed_output=routed_output, + expert_weights=expert_weights, + permuted_indices=permuted_indices, + gated_shared_output=gated_shared_output, + residual=residual, + gamma=gamma, + norm_output=norm_output, + residual_output=residual_output, + ) + + assert calls[0]["launch_with_pdl"] is True + assert "routed_scaling_factor" not in calls[0] + + wrapper.all_reduce_residual_rms_norm( + local_contribution=residual, + residual=residual, + gamma=gamma, + norm_output=norm_output, + residual_output=residual_output, + ) + + assert calls[1]["pattern"] == 1 + assert calls[1]["launch_with_pdl"] is True + assert "routed_scaling_factor" not in calls[1] + assert "expanded_idx_to_permuted_idx" not in calls[1] + + +def test_text_entry_wrapper_delegates_pre_capture_prepare(): + calls = [] + runner = object() + wrapper = SimpleNamespace( + model=SimpleNamespace( + prepare_before_cuda_graph_capture=lambda value: calls.append(value) + ) + ) + + Qwen3_5ForCausalLM.prepare_before_cuda_graph_capture(wrapper, runner) + + assert calls == [runner] + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py b/test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py new file mode 100644 index 000000000000..cf00d9da3500 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py @@ -0,0 +1,54 @@ +import pytest + +from sglang.srt.layers.quantization.unquant import ( + _FLASHINFER_PR4266_TUNED_TACTICS, + Bf16GemmBackend, + should_enable_bf16_splitk_gemm, + use_flashinfer_pr4266_bf16_gemm, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@pytest.mark.parametrize("m,n,k", _FLASHINFER_PR4266_TUNED_TACTICS) +def test_flashinfer_pr4266_selects_tuned_oakhaven_shape(m: int, n: int, k: int): + assert use_flashinfer_pr4266_bf16_gemm(m, n, k) + + +@pytest.mark.parametrize("m", [0, 33, 64]) +@pytest.mark.parametrize("n,k", [(256, 8192), (512, 8192), (2304, 8192), (2560, 8192)]) +def test_flashinfer_pr4266_keeps_large_m_on_existing_path(m: int, n: int, k: int): + assert not use_flashinfer_pr4266_bf16_gemm(m, n, k) + + +@pytest.mark.parametrize( + "shape", + [ + (1, 1024, 2048), + (3, 256, 8192), + (16, 8192, 4096), + (32, 4096, 8192), + ], +) +def test_flashinfer_pr4266_rejects_unmeasured_shapes(shape: tuple[int, int, int]): + assert not use_flashinfer_pr4266_bf16_gemm(*shape) + + +def test_flashinfer_pr4266_backend_is_explicit(): + assert Bf16GemmBackend.FLASHINFER_PR4266.value == "flashinfer_pr4266" + + +def test_bf16_splitk_is_enabled_by_default(monkeypatch): + monkeypatch.delenv("SGLANG_ENABLE_BF16_SPLITK_GEMM", raising=False) + assert should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL) + + +def test_bf16_splitk_env_kill_switch(monkeypatch): + monkeypatch.setenv("SGLANG_ENABLE_BF16_SPLITK_GEMM", "0") + assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL) + + +def test_bf16_splitk_does_not_override_torch_backend(monkeypatch): + monkeypatch.setenv("SGLANG_ENABLE_BF16_SPLITK_GEMM", "1") + assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.TORCH) diff --git a/test/registered/unit/layers/test_radix_linear_attention.py b/test/registered/unit/layers/test_radix_linear_attention.py new file mode 100644 index 000000000000..e05d185e09e3 --- /dev/null +++ b/test/registered/unit/layers/test_radix_linear_attention.py @@ -0,0 +1,232 @@ +"""CPU regression coverage for padded linear-attention inputs and outputs.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang.srt.layers.radix_linear_attention as radix_linear_attention +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeAttentionBackend: + def forward( + self, + *, + layer, + forward_batch, + mixed_qkv, + a, + b, + linear_attn_output=None, + ): + del layer + torch.testing.assert_close(forward_batch.out_cache_loc, torch.arange(3)) + assert mixed_qkv.shape[0] == 3 + assert a.shape[0] == 3 + assert b.shape[0] == 3 + if linear_attn_output is None: + return torch.full((1, 3, 2, 4), 5.0) + linear_attn_output.fill_(5.0) + return linear_attn_output + + +class _FailingAttentionBackend: + def forward(self, **kwargs): + del kwargs + raise RuntimeError("backend failure") + + +class _ExtendMode: + def is_extend(self): + return True + + def is_target_verify(self): + return False + + +class _TargetVerifyMode: + def is_extend(self): + return True + + def is_target_verify(self): + return True + + +class _PhysicalAttentionBackend: + def forward(self, *, layer, forward_batch, mixed_qkv, a, b): + del layer, forward_batch + assert mixed_qkv.shape[0] == 5 + assert a.shape[0] == 5 + assert b.shape[0] == 5 + return torch.full((1, 5, 2, 4), 9.0) + + +class TestRadixLinearAttentionPadding(CustomTestCase): + def test_eager_padded_input_is_sliced_and_output_shape_is_restored(self): + layer = radix_linear_attention.RadixLinearAttention( + layer_id=0, + num_q_heads=1, + num_k_heads=1, + num_v_heads=2, + head_q_dim=4, + head_k_dim=4, + head_v_dim=4, + ) + original_out_cache_loc = torch.arange(5) + forward_batch = SimpleNamespace( + forward_mode=_ExtendMode(), + num_token_non_padded_cpu=3, + out_cache_loc=original_out_cache_loc, + ) + + with ( + patch.object( + radix_linear_attention, + "get_tc_piecewise_forward_context", + return_value=None, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FakeAttentionBackend(), + ), + ): + output = layer.forward( + forward_batch=forward_batch, + mixed_qkv=torch.zeros((5, 8)), + a=torch.zeros((5, 2)), + b=torch.zeros((5, 2)), + ) + + torch.testing.assert_close(output[:, :3], torch.full((1, 3, 2, 4), 5.0)) + torch.testing.assert_close(output[:, 3:], torch.zeros((1, 2, 2, 4))) + self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) + + def test_target_verify_keeps_physical_rows_matching_its_metadata(self): + layer = radix_linear_attention.RadixLinearAttention( + layer_id=0, + num_q_heads=1, + num_k_heads=1, + num_v_heads=2, + head_q_dim=4, + head_k_dim=4, + head_v_dim=4, + ) + original_out_cache_loc = torch.arange(5) + forward_batch = SimpleNamespace( + forward_mode=_TargetVerifyMode(), + num_token_non_padded_cpu=3, + out_cache_loc=original_out_cache_loc, + ) + + with ( + patch.object( + radix_linear_attention, + "get_tc_piecewise_forward_context", + return_value=None, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_PhysicalAttentionBackend(), + ), + ): + output = layer.forward( + forward_batch=forward_batch, + mixed_qkv=torch.zeros((5, 8)), + a=torch.zeros((5, 2)), + b=torch.zeros((5, 2)), + ) + + torch.testing.assert_close(output, torch.full((1, 5, 2, 4), 9.0)) + self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) + + def test_eager_backend_failure_restores_out_cache_loc(self): + layer = radix_linear_attention.RadixLinearAttention( + layer_id=0, + num_q_heads=1, + num_k_heads=1, + num_v_heads=2, + head_q_dim=4, + head_k_dim=4, + head_v_dim=4, + ) + original_out_cache_loc = torch.arange(5) + forward_batch = SimpleNamespace( + forward_mode=_ExtendMode(), + num_token_non_padded_cpu=3, + out_cache_loc=original_out_cache_loc, + ) + + with ( + patch.object( + radix_linear_attention, + "get_tc_piecewise_forward_context", + return_value=None, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FailingAttentionBackend(), + ), + self.assertRaisesRegex(RuntimeError, "backend failure"), + ): + layer.forward( + forward_batch=forward_batch, + mixed_qkv=torch.zeros((5, 8)), + a=torch.zeros((5, 2)), + b=torch.zeros((5, 2)), + ) + + self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) + + def test_padded_output_tail_is_initialized(self): + for padded_num_tokens in (3, 5): + with self.subTest(padded_num_tokens=padded_num_tokens): + original_out_cache_loc = torch.arange(padded_num_tokens) + forward_batch = SimpleNamespace( + num_token_non_padded_cpu=3, + out_cache_loc=original_out_cache_loc, + ) + context = SimpleNamespace( + forward_batch=forward_batch, + attention_layers=[object()], + ) + output = torch.full((1, padded_num_tokens, 2, 4), float("nan")) + + with ( + patch.object( + radix_linear_attention, + "get_tc_piecewise_forward_context", + return_value=context, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FakeAttentionBackend(), + ), + ): + radix_linear_attention._unified_linear_attention_with_output_impl( + mixed_qkv=torch.zeros((padded_num_tokens, 8)), + a=torch.zeros((padded_num_tokens, 2)), + b=torch.zeros((padded_num_tokens, 2)), + output=output, + layer_id=0, + ) + + torch.testing.assert_close(output[:, :3], torch.full((1, 3, 2, 4), 5.0)) + torch.testing.assert_close( + output[:, 3:], + torch.zeros((1, padded_num_tokens - 3, 2, 4)), + ) + self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_decision_batch_params.py b/test/registered/unit/managers/test_scheduler_decision_batch_params.py index dc30ceacea98..30c888fefa51 100644 --- a/test/registered/unit/managers/test_scheduler_decision_batch_params.py +++ b/test/registered/unit/managers/test_scheduler_decision_batch_params.py @@ -1,5 +1,7 @@ import inspect import unittest +from types import SimpleNamespace +from unittest.mock import patch from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import maybe_stub_sgl_kernel @@ -48,5 +50,39 @@ def test_decision_methods_take_batches_as_params_not_self(self): ) +class TestMtpPhaseBoundaryOverlap(unittest.TestCase): + @staticmethod + def _batch(*, is_extend: bool, is_speculative: bool = True): + return SimpleNamespace( + is_extend_in_batch=is_extend, + forward_mode=SimpleNamespace( + is_extend=lambda: is_extend, + is_decode=lambda: not is_extend, + ), + spec_algorithm=SimpleNamespace(is_none=lambda: not is_speculative), + grammar_needs_sync=lambda: False, + ) + + def _scheduler(self, *, require_mlp_sync: bool): + scheduler = object.__new__(Scheduler) + scheduler.require_mlp_sync = require_mlp_sync + scheduler.result_queue = [object()] + return scheduler + + @patch( + "sglang.srt.managers.scheduler.envs." + "SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get", + return_value=False, + ) + def test_mtp_phase_crossing_keeps_overlap(self, _disable_consecutive_prefill): + extend = self._batch(is_extend=True) + decode = self._batch(is_extend=False) + + for require_mlp_sync in (False, True): + scheduler = self._scheduler(require_mlp_sync=require_mlp_sync) + self.assertFalse(scheduler.is_disable_overlap_for_batch(decode, extend)) + self.assertFalse(scheduler.is_disable_overlap_for_batch(extend, decode)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py index 6bcf78c4b8c9..578bc0b1fff9 100644 --- a/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py +++ b/test/registered/unit/mem_cache/test_replayssm_ring_accounting.py @@ -19,6 +19,7 @@ Mamba2StateDType, Mamba2StateShape, ) +from sglang.srt.mem_cache.kv_cache_configurator import _pp_local_per_request_bytes from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -81,6 +82,19 @@ def test_zero_len_ring(self): self.assertEqual(_gdn_params().replayssm_ring_bytes_per_req(record_len=0), 0) self.assertEqual(_kda_params().replayssm_ring_bytes_per_req(record_len=0), 0) + def test_pp_local_state_budget(self): + # Four equal-cost linear layers globally, two owned by this PP stage. + self.assertEqual( + _pp_local_per_request_bytes(4096, [0, 1, 3, 4], 1, 4), + 2048, + ) + + def test_pp_local_state_budget_empty_stage(self): + self.assertEqual( + _pp_local_per_request_bytes(4096, [0, 1, 3, 4], 5, 8), + 0, + ) + if __name__ == "__main__": import sys diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index 0632fe9c3341..0101c15df9fa 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py @@ -6,6 +6,8 @@ from sglang.srt.model_executor.cuda_graph_config import Phase from sglang.srt.model_executor.model_runner_components import cuda_graph_setup from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( + has_standard_gqa_for_all_local_layers, + index_attention_layers_by_global_id, capture_decode_graph, should_skip_auto_prefill_cuda_graph_for_memory, ) @@ -25,6 +27,39 @@ def test_explicit_prefill_backend_bypasses_memory_gate(): ) +def test_standard_gqa_gate_uses_pipeline_local_layer_range(): + # PP rank owns layers [23, 46), while the full model has 92 layers. + assert has_standard_gqa_for_all_local_layers( + attention_layer_count=23, start_layer=23, end_layer=46 + ) + assert not has_standard_gqa_for_all_local_layers( + attention_layer_count=22, start_layer=23, end_layer=46 + ) + + +def test_standard_gqa_gate_is_unchanged_without_pipeline_parallelism(): + assert has_standard_gqa_for_all_local_layers( + attention_layer_count=92, start_layer=0, end_layer=92 + ) + + +def test_pipeline_attention_metadata_is_indexed_by_global_layer_id(): + layer23 = SimpleNamespace(layer_id=23) + layer24 = SimpleNamespace(layer_id=24) + companion24 = object() + + attention, companions = index_attention_layers_by_global_id( + [layer23, layer24], [None, companion24] + ) + + assert len(attention) == 25 + assert all(layer is None for layer in attention[:23]) + assert attention[23] is layer23 + assert attention[24] is layer24 + assert companions[23] is None + assert companions[24] is companion24 + + def test_model_runner_can_override_decode_graph_runner(monkeypatch): class CustomGraphRunner: def __init__(self, model_runner): diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 3ef7f5a10720..fe2543d66315 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -1175,6 +1175,78 @@ def test_mamba_bs_axis_copy(self): reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=3, padded_num_tokens=8) self.assertTrue(torch.equal(idx, torch.tensor([3, 4], dtype=torch.int64))) + def test_pp_proxy_token_slots_copy_head_and_zero_bucket_tail(self): + from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + build_prefill_registry, + ) + from sglang.srt.model_executor.forward_batch_info import PPProxyTensors + + hidden = torch.full((16, 4), 7.0) + residual = torch.full((16, 4), 7.0) + src = self._src( + pp_proxy_tensors={ + "hidden_states": hidden, + "residual": residual, + } + ) + reg = build_prefill_registry( + device=torch.device("cpu"), + max_bs=1, + max_num_token=16, + cache_loc_dtype=torch.int64, + source=src, + ) + self.assertTrue(reg.has_slot("pp_proxy_tensors.hidden_states")) + fb = _MiniForwardBatch( + input_ids=torch.zeros(3, dtype=torch.int64), + positions=torch.zeros(3, dtype=torch.int64), + out_cache_loc=torch.zeros(3, dtype=torch.int64), + ) + pp_proxy = PPProxyTensors( + { + "hidden_states": torch.ones((3, 4)), + "residual": torch.full((3, 4), 2.0), + } + ) + reg.fill_from( + fb, + raw_bs=1, + padded_bs=1, + raw_num_tokens=3, + padded_num_tokens=8, + pp_proxy_tensors=pp_proxy, + ) + self.assertTrue(torch.all(hidden[:3] == 1.0)) + self.assertTrue(torch.all(residual[:3] == 2.0)) + self.assertTrue(torch.all(hidden[3:8] == 0.0)) + self.assertTrue(torch.all(residual[3:8] == 0.0)) + self.assertTrue(torch.all(hidden[8:] == 7.0)) + + def test_prefill_input_buffers_allocate_pp_proxy_by_token(self): + from sglang.srt.model_executor.runner_utils.buffers import ( + PrefillInputBuffers, + ) + + buffers = PrefillInputBuffers.create( + device=torch.device("cpu"), + max_bs=4, + max_num_tokens=16, + cache_loc_dtype=torch.int64, + is_multimodal=False, + hidden_size=8, + dtype=torch.bfloat16, + enable_mamba_track=False, + pp_size=2, + pp_proxy_topk_size=3, + ) + self.assertEqual( + tuple(buffers.pp_proxy_tensors["hidden_states"].shape), (16, 8) + ) + self.assertEqual(tuple(buffers.pp_proxy_tensors["residual"].shape), (16, 8)) + self.assertEqual( + tuple(buffers.pp_proxy_tensors["topk_indices"].shape), (16, 3) + ) + def test_source_none_owns_allocated_buffers(self): # source=None -> the registry allocates (owns) every slot. from sglang.srt.model_executor.cuda_graph_buffer_registry import ( diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index e52b08e0c213..4e8db011e502 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py @@ -9,7 +9,10 @@ import sglang.srt.model_executor.model_runner_components.cuda_graph_setup as graph_setup import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module from sglang.srt.model_executor.cuda_graph_config import Backend -from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + PPProxyTensors, +) from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( capture_prefill_graph, ) @@ -84,6 +87,22 @@ def test_eagle_target_tc_piecewise_skips_last_mode_capture(self): self.assertIs(capture.runner, eager_runner) + def test_pp_proxy_output_is_trimmed_to_raw_prefill_tokens(self): + runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) + runner.raw_num_tokens = 3 + output = PPProxyTensors( + { + "hidden_states": torch.arange(32).view(8, 4), + "residual": torch.arange(32, 64).view(8, 4), + } + ) + + trimmed = runner._finalize_execute_output(output) + + self.assertIsInstance(trimmed, PPProxyTensors) + self.assertEqual(tuple(trimmed["hidden_states"].shape), (3, 4)) + self.assertEqual(tuple(trimmed["residual"].shape), (3, 4)) + def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self): model_runner = SimpleNamespace( server_args=SimpleNamespace( diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 6fbdbeee9cba..3ae440f02b76 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1681,6 +1681,122 @@ def test_token_oracle_accepted_when_env_enabled(self): self.assertEqual(parsed.sampling_backend, "token_oracle") +class TestDeepEPv2Args(CustomTestCase): + """DeepEP v2 server-args resolution + validation. The dummy-model path + short-circuits __post_init__, so _handle_a2a_moe() is invoked directly.""" + + def _args(self, **overrides): + server_args = ServerArgs(model_path="dummy", moe_a2a_backend="deepep_v2") + # The deepep_v2 branch mutates cuda_graph_config.{decode,prefill}.backend, + # so it must exist (the dummy path leaves it unset otherwise). + server_args.cuda_graph_config = CudaGraphConfig( + decode=PhaseConfig(backend=Backend.FULL, max_bs=512), + prefill=PhaseConfig(backend=Backend.FULL, max_bs=512), + ) + for key, value in overrides.items(): + setattr(server_args, key, value) + return server_args + + def test_auto_runner_resolves_to_deep_gemm_for_fp8(self): + args = self._args( + moe_runner_backend="auto", deepep_v2_dispatcher_output_dtype="fp8" + ) + args._handle_a2a_moe() + self.assertEqual(args.moe_runner_backend, "deep_gemm") + + def test_auto_runner_resolves_to_triton_for_bf16(self): + args = self._args( + moe_runner_backend="auto", deepep_v2_dispatcher_output_dtype="bf16" + ) + args._handle_a2a_moe() + self.assertEqual(args.moe_runner_backend, "triton") + + def test_auto_runner_defaults_to_deep_gemm(self): + args = self._args( + moe_runner_backend="auto", deepep_v2_dispatcher_output_dtype="auto" + ) + args._handle_a2a_moe() + self.assertEqual(args.moe_runner_backend, "deep_gemm") + + def test_unsupported_runner_rejected(self): + args = self._args(moe_runner_backend="flashinfer_trtllm") + with self.assertRaises(ValueError): + args._handle_a2a_moe() + + def test_two_batch_overlap_rejected(self): + args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True) + with self.assertRaises(ValueError): + args._handle_a2a_moe() + + # --- prefill capacity pre-check (per-sender chunk vs dispatch buffer cap) --- + _CAP_ENV = "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK" + + def test_prefill_chunk_exceeding_cap_rejected(self): + args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=2048) + with patch.dict(os.environ, {self._CAP_ENV: "1024"}): + with self.assertRaisesRegex(ValueError, "NUM_MAX_DISPATCH_TOKENS_PER_RANK"): + args._handle_a2a_moe() + + def test_prefill_chunk_at_cap_boundary_accepted(self): + # chunk == cap is the documented (and currently benchmarked) edge; the + # guard must be strict-greater-than. + args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024) + with patch.dict(os.environ, {self._CAP_ENV: "1024"}): + args._handle_a2a_moe() + self.assertEqual(args.moe_runner_backend, "deep_gemm") + + def test_prefill_dp4_tp4_scatter_reduces_sender_budget(self): + # _handle_data_parallelism has already converted the global 8192-token + # chunk to 2048 per DP worker. The SCATTERED MoE path then reduce-scatters + # it over attention TP4, so each DeepEP sender sees only 512 tokens. + args = self._args( + moe_runner_backend="deep_gemm", + chunked_prefill_size=2048, + tp_size=16, + dp_size=4, + attn_cp_size=1, + enable_dp_attention=True, + ) + with patch.dict(os.environ, {self._CAP_ENV: "512"}): + args._handle_a2a_moe() + self.assertEqual(args.moe_runner_backend, "deep_gemm") + + def test_prefill_dp4_tp4_scatter_uses_ceiling_sender_budget(self): + args = self._args( + moe_runner_backend="deep_gemm", + chunked_prefill_size=2049, + tp_size=16, + dp_size=4, + attn_cp_size=1, + enable_dp_attention=True, + ) + with patch.dict(os.environ, {self._CAP_ENV: "512"}): + with self.assertRaisesRegex(ValueError, "per-sender prefill"): + args._handle_a2a_moe() + + def test_prefill_chunk_rejected_under_default_cap(self): + # Default cap is 128: a typical 1024-token per-rank chunk must be + # rejected at boot instead of at the first full prefill chunk. + args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024) + with self.assertRaisesRegex(ValueError, "prefill dispatch budget"): + args._handle_a2a_moe() + + def test_prefill_chunk_check_skipped_for_decode_disaggregation(self): + args = self._args( + moe_runner_backend="deep_gemm", + chunked_prefill_size=4096, + disaggregation_mode="decode", + ) + args._handle_a2a_moe() + + def test_prefill_chunk_check_skipped_when_chunking_disabled(self): + for disabled in (None, 0, -1): + args = self._args( + moe_runner_backend="deep_gemm", chunked_prefill_size=disabled + ) + args._handle_a2a_moe() + + class TestHandleCrashDumpEnv(CustomTestCase): _COREDUMP_ENV_KEYS = ( "CUDA_ENABLE_COREDUMP_ON_EXCEPTION", diff --git a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py new file mode 100644 index 000000000000..a7f509c6ae42 --- /dev/null +++ b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py @@ -0,0 +1,98 @@ +"""DeepEP-class backend recognition in RoutedExpertsCapturer. + +The capturer keys its buffer layout on the a2a backend: DeepEP-class +dispatchers hand the MoE layer only the attention rank's DP-local tokens, so +``capture()`` must attn-TP-gather and ``_get_local_slice()`` must read the +buffer head instead of the global DP offset. These tests pin that DeepEP v2 +is classified like DeepEP (it shares that token topology); a miss makes +dp_rank > 0 read unwritten rows (silent wrong data), see the DP>1 readback +test in test/registered/ep/test_routed_experts_dp_readback.py. +""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.layers.moe.utils import MoeA2ABackend +from sglang.srt.state_capturer import routed_experts as re_mod +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large") + + +class TestScatteredA2ABackendHelper(CustomTestCase): + def test_classification(self): + # deepep_v2 shares DeepEP's scattered token topology. Other backends + # keep their existing classification (mooncake/mori are deliberately + # not reclassified here). + expected = { + "deepep": True, + "deepep_v2": True, + "none": False, + "mooncake": False, + } + for value, exp in expected.items(): + with mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(value) + ): + self.assertEqual( + re_mod._is_scattered_a2a_backend(), exp, f"backend={value}" + ) + + +class TestGetLocalSliceBackendBranch(CustomTestCase): + T, L, K = 16, 3, 4 # buffer tokens, layers, top-k + + def _capturer(self): + cap = object.__new__(re_mod.RoutedExpertsCapturer) + buf = torch.arange(self.T * self.L * self.K, dtype=torch.int32).reshape( + self.T, self.L, self.K + ) + cap.device_cache = SimpleNamespace(buffer=buf) + cap.topk_size = self.K + return cap, buf + + def _slice(self, cap, n_local): + fb = SimpleNamespace(out_cache_loc=torch.empty(n_local)) + return cap._get_local_slice(fb, can_run_graph=False, cuda_graph_batch=None) + + def test_deepep_v2_reads_buffer_head(self): + cap, buf = self._capturer() + with mock.patch.object( + re_mod, "is_dp_attention_enabled", return_value=True + ), mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("deepep_v2") + ): + out = self._slice(cap, n_local=5) + self.assertTrue(torch.equal(out, buf[0:5, :, : self.K])) + + def test_deepep_v2_matches_deepep(self): + cap, _ = self._capturer() + outs = [] + for backend in ("deepep", "deepep_v2"): + with mock.patch.object( + re_mod, "is_dp_attention_enabled", return_value=True + ), mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(backend) + ): + outs.append(self._slice(cap, n_local=7)) + self.assertTrue(torch.equal(outs[0], outs[1])) + + def test_tp_moe_reads_global_offset(self): + cap, buf = self._capturer() + with mock.patch.object( + re_mod, "is_dp_attention_enabled", return_value=True + ), mock.patch.object( + re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("none") + ), mock.patch.object( + re_mod, "get_dp_local_slice_cpu", return_value=(6, 4) + ): + out = self._slice(cap, n_local=999) + self.assertTrue(torch.equal(out, buf[6:10, :, : self.K])) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 6079f9614fb1..2059909b1be3 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -448,9 +448,11 @@ def test_speculative_swap_and_restore(self): self.assertTrue(get_moe_a2a_backend().is_none()) # MTP layers are unquantized: fp4 allgather is forced off self.assertTrue(get_flags().moe.disable_fp4_allgather) + self.assertTrue(get_flags().moe.speculative_context) self.assertEqual(get_moe_runner_backend().name, "TRITON") self.assertTrue(get_moe_a2a_backend().is_deepep()) self.assertFalse(get_flags().moe.disable_fp4_allgather) + self.assertFalse(get_flags().moe.speculative_context) def test_swap_restores_on_exception(self): from sglang.srt.layers.moe.utils import (