[AMD] Enable mamba-extra-buffer for Qwen3.5 on ROCm

.github/workflows/nightly-test-nvidia.yml

@@ -29,6 +29,11 @@ on:
           - 'nightly-test-diffusion'
           - 'nightly-test-kernel-8-gpu-h200'
           - 'nightly-test-precision-8-gpu-h200'
+      force_baseline_update:
+        description: 'precision job only: refresh the rolling baseline instead of comparing (sets SGLANG_PRECISION_FORCE_UPDATE=1). Dispatch once after an intentional forward-path dtype/precision change stales the baseline; later scheduled runs compare against it and go green.'
+        required: false
+        type: boolean
+        default: false
   workflow_call:
     inputs:
       ref:
@@ -718,6 +723,8 @@ jobs:
           SGLANG_PRECISION_HF_REVISION: ${{ vars.SGLANG_PRECISION_HF_REVISION || 'main' }}
           HF_TOKEN: ${{ secrets.HF_TOKEN_PRECISION_STORE }}
           SGLANG_PRECISION_COMMIT: ${{ github.sha }}
+          # '0' on scheduled runs (input unset); '1' only on a manual baseline-refresh dispatch.
+          SGLANG_PRECISION_FORCE_UPDATE: ${{ inputs.force_baseline_update && '1' || '0' }}
         run: |
           cd test
           python3 run_suite.py --hw cuda --suite nightly-precision-8-gpu-h200 --nightly --continue-on-error --timeout-per-file 3600

docker/Dockerfile

@@ -12,7 +12,7 @@ ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
 ARG BUILD_AND_DOWNLOAD_PARALLEL=8
 ARG SGL_KERNEL_VERSION=0.4.4
 ARG SGL_VERSION
-ARG SGL_DEEP_GEMM_VERSION=0.1.3
+ARG SGL_DEEP_GEMM_VERSION=0.1.4
 ARG USE_LATEST_SGLANG=0
 ARG GDRCOPY_VERSION=2.5.1
 ARG PIP_DEFAULT_INDEX
@@ -239,6 +239,11 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && echo '__version__ = "0.0.0"' > sglang/version.py \
     && touch README.md \
     && touch LICENSE \
+    && if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
+           sed -i 's/cuda-python>=13\.0/cuda-python>=12,<13/' pyproject.toml && \
+           sed -i 's/flashinfer_python\[cu13\]/flashinfer_python[cu12]/' pyproject.toml && \
+           sed -i 's/nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' pyproject.toml; \
+       fi \
     && python3 -m pip install --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} ".[${BUILD_TYPE}]" \
     && if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
            pip list --format=freeze | awk -F'==' '/-cu13(==|$)/ {print $1}' \
@@ -665,6 +670,11 @@ RUN if [ "$BRANCH_TYPE" = "local" ]; then \
 # Clean up __pycache__/tests/pyc in same RUN to avoid writing ~28k files to layer
 RUN --mount=type=cache,target=/root/.cache/pip \
     cd /sgl-workspace/sglang \
+    && if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
+           sed -i 's/cuda-python>=13\.0/cuda-python>=12,<13/' python/pyproject.toml && \
+           sed -i 's/flashinfer_python\[cu13\]/flashinfer_python[cu12]/' python/pyproject.toml && \
+           sed -i 's/nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' python/pyproject.toml; \
+       fi \
     && python3 -m pip install --no-deps -e "python[${BUILD_TYPE}]" \
     && kernels lock python \
     && ( success=0; \

python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh

@@ -0,0 +1,169 @@
+// DeepSeek-V4 wo_a activation quantization for DeepGEMM fp8_einsum.
+//
+// This is intentionally narrower than the generic per_token_group_quant_8bit_v2
+// kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is
+// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0
+// power-of-two values, and output_s is a logical [T, G, D/128] view backed by
+// group-major [G, T, D/128] storage.
+//
+// The generic kernel cannot read the strided DSV4 view while producing
+// contiguous [T, G, D] codes and group-major scales without an extra full-tensor
+// copy.
+#include <sgl_kernel/tensor.h>  // TensorMatcher, SymbolicSize/Device
+#include <sgl_kernel/utils.h>   // RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // fp8 aliases, PDL helpers
+#include <sgl_kernel/warp.cuh>   // warp::reduce_max
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>  // UE8M0 and FP8 helpers
+
+#include <tvm/ffi/container/tensor.h>  // tvm::ffi::TensorView
+
+#include <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
+constexpr uint32_t GROUP_SIZE = 128;
+constexpr uint32_t THREADS_PER_GROUP = 8;
+constexpr uint32_t SUBWARPS_PER_BLOCK = 16;
+constexpr uint32_t INPUT_VEC_NUM_BYTES = 32;
+constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4);
+
+template <int THREADS_PER_SUBWARP>
+SGL_DEVICE float GroupReduceMax(float val) {
+  static_assert(
+      (THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1,
+      "THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16");
+  // Tail subwarps can be inactive at the bounds check, so reduce with only the
+  // current subgroup's lanes rather than a full-warp mask.
+  constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1;
+  const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP));
+  return device::warp::reduce_max<THREADS_PER_SUBWARP>(val, mask);
+}
+
+template <typename T, bool kUsePDL>
+__global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
+    const T* __restrict__ input,
+    fp8_e4m3_t* __restrict__ output_q,
+    float* __restrict__ output_s,
+    int64_t total_scale_groups,
+    int64_t num_tokens,
+    int hidden_dim_groups,
+    int num_outer_groups,
+    int64_t input_stride_t) {
+  device::PDLWaitPrimary<kUsePDL>();
+
+  const int64_t subwarp_id = threadIdx.x / THREADS_PER_GROUP;
+  const int lane_id = threadIdx.x % THREADS_PER_GROUP;
+  const int64_t group_id = static_cast<int64_t>(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id;
+  if (group_id < total_scale_groups) {
+    const int hidden_group = group_id % hidden_dim_groups;
+    const int64_t token_outer = group_id / hidden_dim_groups;
+    const int outer_idx = token_outer % num_outer_groups;
+    const int64_t token_idx = token_outer / num_outer_groups;
+
+    constexpr uint32_t INPUT_VEC_SIZE = INPUT_VEC_NUM_BYTES / sizeof(T);
+    static_assert(INPUT_VEC_SIZE * THREADS_PER_GROUP == GROUP_SIZE);
+
+    const int64_t input_group_start_offset =
+        token_idx * input_stride_t + outer_idx * GROUP_SIZE * hidden_dim_groups + hidden_group * GROUP_SIZE;
+    const int64_t output_group_start_offset = group_id * GROUP_SIZE;
+
+    int4 input_int4[INPUT_INT4_SIZE];
+    T* input_vec = reinterpret_cast<T*>(input_int4);
+
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_INT4_SIZE; ++j) {
+      input_int4[j] = reinterpret_cast<const int4*>(input + input_group_start_offset + lane_id * INPUT_VEC_SIZE)[j];
+    }
+
+    float local_absmax = LOCAL_ABSMAX_ABS;
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_VEC_SIZE; ++j) {
+      const float val = static_cast<float>(input_vec[j]);
+      local_absmax = fmaxf(local_absmax, fabsf(val));
+    }
+
+    local_absmax = GroupReduceMax<THREADS_PER_GROUP>(local_absmax);
+
+    constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max;
+    const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv);
+    const float y_scale = inv_scale_ue8m0(scale_ue8m0);
+    const float y_scale_inv = __uint_as_float(static_cast<uint32_t>(scale_ue8m0) << 23);
+
+    int4 output_buf;
+    auto* output_buf_ptr = reinterpret_cast<fp8x2_e4m3_t*>(&output_buf);
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_VEC_SIZE; j += 2) {
+      output_buf_ptr[j / 2] =
+          pack_fp8(static_cast<float>(input_vec[j]) * y_scale, static_cast<float>(input_vec[j + 1]) * y_scale);
+    }
+
+    *reinterpret_cast<int4*>(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf;
+
+    if (lane_id == 0) {
+      output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv;
+    }
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename T, bool kUsePDL>
+struct FP8WoAGroupMajorQuantUE8M0Kernel {
+  static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    auto TSize = SymbolicSize{"num_tokens"};
+    auto GSize = SymbolicSize{"num_outer_groups"};
+    auto DSize = SymbolicSize{"hidden_dim"};
+    auto SSize = SymbolicSize{"hidden_dim_groups"};
+
+    TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype<T>().with_device(device).verify(input);
+    TensorMatcher({TSize, GSize, DSize}).with_dtype<fp8_e4m3_t>().with_device(device).verify(output_q);
+    TensorMatcher({GSize, TSize, SSize}).with_dtype<float>().with_device(device).verify(output_s);
+
+    const auto num_tokens = TSize.unwrap();
+    const auto num_outer_groups = GSize.unwrap();
+    const auto hidden_dim = DSize.unwrap();
+    const auto hidden_dim_groups = SSize.unwrap();
+    const auto input_stride_t = input.stride(0);
+    constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T);
+
+    RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128");
+    RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch");
+    RuntimeCheck(
+        reinterpret_cast<uintptr_t>(input.data_ptr()) % sizeof(int4) == 0,
+        "input base pointer must be 16-byte aligned");
+    RuntimeCheck(
+        num_tokens <= 1 || input_stride_t % kInputAlignElements == 0,
+        "input token stride must preserve 16-byte vector-load alignment");
+
+    const int64_t total_scale_groups = num_tokens * num_outer_groups * hidden_dim_groups;
+    if (total_scale_groups == 0) return;
+
+    const auto grid = dim3((total_scale_groups + SUBWARPS_PER_BLOCK - 1) / SUBWARPS_PER_BLOCK);
+    const auto block = dim3(SUBWARPS_PER_BLOCK * THREADS_PER_GROUP);
+    host::LaunchKernel(grid, block, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            fp8_wo_a_group_major_quant_ue8m0_kernel<T, kUsePDL>,
+            static_cast<const T*>(input.data_ptr()),
+            static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
+            static_cast<float*>(output_s.data_ptr()),
+            total_scale_groups,
+            static_cast<int64_t>(num_tokens),
+            static_cast<int>(hidden_dim_groups),
+            static_cast<int>(num_outer_groups),
+            static_cast<int64_t>(input_stride_t));
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu

@@ -43,7 +43,19 @@ inline uint32_t next_pow2(uint32_t x) noexcept {
 
 namespace moe {
 
-__device__ __forceinline__ int warp_exclusive_scan(int v, unsigned mask = 0xffffffffu) {
+#ifdef USE_ROCM
+// ROCm's templated __shfl_*_sync builtins require an 8-byte mask
+// (static_assert(sizeof(MaskT) == 8)) because AMD wavefront width varies by
+// arch (32 on RDNA, 64 on CDNA).
+using warp_mask_t = uint64_t;
+#define WARP_FULL_MASK 0xffffffffffffffffull
+#else
+// CUDA: 32-lane warps take a plain 32-bit mask
+using warp_mask_t = unsigned;
+#define WARP_FULL_MASK 0xffffffffu
+#endif
+
+__device__ __forceinline__ int warp_exclusive_scan(int v, warp_mask_t mask = WARP_FULL_MASK) {
   int original = v;
 #pragma unroll
   for (int offset = 1; offset < WARP_SIZE; offset <<= 1) {

python/sglang/jit_kernel/dsa/__init__.py

@@ -1,4 +1,13 @@
-from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
+from sglang.srt.utils import is_cuda
+
+_is_cuda = is_cuda()
+
+if _is_cuda:
+    from .cutedsl_paged_mqa_logits import CuteDSLPagedMQALogitsRunner, pick_dsl_expand
+else:
+    CuteDSLPagedMQALogitsRunner = None
+    pick_dsl_expand = None
+
 from .paged_mqa_logits import (
     aiter_paged_mqa_logits,
     cutedsl_paged_mqa_logits,

python/sglang/jit_kernel/dsv4/__init__.py

@@ -19,6 +19,7 @@
     fused_q_norm_rope,
     fused_rope_inplace,
 )
+from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
 from .gemm import linear_bf16_fp32
 from .moe import (
     hash_topk,
@@ -45,6 +46,7 @@
     "fused_q_indexer_rope_hadamard_fp4_quant",
     "fused_q_indexer_rope_hadamard_quant",
     "fused_k_norm_rope_flashmla",
+    "sglang_per_token_group_quant_fp8_dsv4_wo_a",
     "make_name",
     "linear_bf16_fp32",
     "get_paged_mqa_logits_metadata",

python/sglang/jit_kernel/dsv4/fp8_wo_a.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.srt.utils.custom_op import register_custom_op
+
+from .utils import make_name
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_GROUP_SIZE = 128
+
+
+@cache_once
+def _jit_module(in_dtype: torch.dtype, use_pdl: bool) -> Module:
+    args = make_cpp_args(in_dtype, use_pdl)
+    return load_jit(
+        make_name("fp8_wo_a_group_major_quant_ue8m0"),
+        *args,
+        cuda_files=["deepseek_v4/fp8_wo_a_group_major_quant.cuh"],
+        cuda_wrappers=[
+            (
+                "fp8_wo_a_group_major_quant_ue8m0",
+                f"FP8WoAGroupMajorQuantUE8M0Kernel<{args}>::run",
+            )
+        ],
+        # Match the AOT/JIT v2 quant path's fast-math build so FP8 rounding stays
+        # bit-identical for the DSV4 wo_a replacement.
+        extra_cuda_cflags=["--use_fast_math"],
+    )
+
+
+@register_custom_op(
+    op_name="fp8_wo_a_group_major_quant_ue8m0",
+    mutates_args=["output_q", "output_s"],
+)
+def _fp8_wo_a_group_major_quant_ue8m0_custom_op(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    """Opaque custom-op boundary for the DeepSeek-V4 wo_a quant JIT kernel."""
+    assert input.dtype in (torch.bfloat16, torch.float16)
+
+    module = _jit_module(input.dtype, is_arch_support_pdl())
+    module.fp8_wo_a_group_major_quant_ue8m0(input, output_q, output_s)
+
+
+@debug_kernel_api
+def fp8_wo_a_group_major_quant_ue8m0(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    _fp8_wo_a_group_major_quant_ue8m0_custom_op(input, output_q, output_s)
+
+
+def sglang_per_token_group_quant_fp8_dsv4_wo_a(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize DSV4 wo_a activations for DeepGEMM fp8_einsum.
+
+    The input is a [T, G, D] bf16/fp16 tensor whose hidden dimension is
+    contiguous. The output codes are contiguous [T, G, D] fp8 values. The scale
+    tensor is returned as logical [T, G, D/128] fp32 UE8M0 values backed by
+    contiguous [G, T, D/128] storage, so each group/head [T, S] panel is
+    contiguous for the DeepGEMM recipe=(1, 1, 128) consumer. Group size is fixed
+    to 128 and the absmax floor is fixed to 1e-10.
+    """
+    num_tokens, num_groups, hidden = x.shape
+    hidden_groups = hidden // _GROUP_SIZE
+    x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn)
+    x_s_storage = torch.empty(
+        (num_groups, num_tokens, hidden_groups),
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    if x.numel() > 0:
+        fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage)
+
+    # DeepGEMM fp8_einsum consumes each group/head [T, S] scale panel contiguously.
+    return x_q, x_s_storage.transpose(0, 1)