diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index f01342fa5..1344675b3 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -299,7 +299,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { */ m.def( "moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none," - "Tensor! b_q_weight, Tensor! b_scales, Tensor? b_zeros_or_none," + "Tensor! b_q_weight, Tensor? b_bias_or_none, Tensor! b_scales," + "Tensor? global_scale_or_none, Tensor? b_zeros_or_none," "Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace," "Tensor sorted_token_ids," "Tensor! expert_ids, Tensor! num_tokens_past_padded," diff --git a/sgl-kernel/csrc/gemm/marlin/dequant.h b/sgl-kernel/csrc/gemm/marlin/dequant.h index e91b8b9f4..37893f42e 100644 --- a/sgl-kernel/csrc/gemm/marlin/dequant.h +++ b/sgl-kernel/csrc/gemm/marlin/dequant.h @@ -449,6 +449,51 @@ __device__ inline void dequant_fp8_scales(int q, nv_bfloat162* fra q <<= 8; int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +}; + +// New version with s_type_id parameter for marlin_moe_wna16_v2 +template +__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b); + +template <> +__device__ inline void dequant_fp8_scales(int q, half2* frag_b) { + int Out1 = (q & 0xFF00FF00) >> 1; + ; + q <<= 8; + int Out2 = (q & 0xFF00FF00) >> 1; + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +}; + +template <> +__device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) { + constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; + constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; + constexpr int MASK = 0x7F007F00; + + // Extract and shift FP8 values to BF16 format + int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 8; + int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant_fp8_scales(int q, nv_bfloat162* frag_b) { + // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, + // but we assume that such a extreme value would not occur in real models. + int Out1 = (q & 0xFF00FF00) >> 1; + q <<= 7; + int Out2 = q & 0x7F807F80; + // Note: reverse indexing is intentional because weights are permuted frag_b[1] = *reinterpret_cast(&Out1); frag_b[0] = *reinterpret_cast(&Out2); diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/generate_kernels.py b/sgl-kernel/csrc/moe/marlin_moe_wna16/generate_kernels.py index b3ed863a3..dea951f7f 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/generate_kernels.py +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/generate_kernels.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob import itertools import os @@ -11,9 +12,6 @@ FILE_HEAD = """ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { """.strip() @@ -21,14 +19,13 @@ TEMPLATE = ( "template __global__ void Marlin<" "{{scalar_t}}, " "{{w_type_id}}, " + "{{s_type_id}}, " "{{threads}}, " "{{thread_m_blocks}}, " "{{thread_n_blocks}}, " "{{thread_k_blocks}}, " "{{'true' if m_block_size_8 else 'false'}}, " "{{stages}}, " - "{{'true' if has_act_order else 'false'}}, " - "{{'true' if has_zp else 'false'}}, " "{{group_blocks}}, " "{{'true' if is_zp_float else 'false'}}>" "( MARLIN_KERNEL_PARAMS );" @@ -47,10 +44,15 @@ KERNEL_FILE_NAME = "kernel_marlin.cuh" # int8 with zero point case (sglang::kU8) is also supported, # we don't add it to reduce wheel size. -SCALAR_TYPES = ["sglang::kU4", "sglang::kU4B8", "sglang::kU8B128"] -THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)] +# Only keep the most commonly used types to reduce compilation time +SCALAR_TYPES = [ + "sglang::kU4", + "sglang::kU4B8", + "sglang::kU8B128", +] +THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256)] -THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] +THREAD_M_BLOCKS = [0.5, 1, 2, 4] # group_blocks: # = 0 : act order case # = -1 : channelwise quantization @@ -67,39 +69,65 @@ def remove_old_kernels(): def generate_new_kernels(): kernel_files = set() for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES): - has_zp = "B" not in scalar_type all_template_str_list = [] for group_blocks, m_blocks, thread_configs in itertools.product( GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS ): - - has_act_order = group_blocks == 0 - if has_zp and has_act_order: + # act order case only support gptq-int4 and gptq-int8 + if group_blocks == 0 and scalar_type not in [ + "sglang::kU4B8", + "sglang::kU8B128", + ]: continue if thread_configs[2] == 256: + # for small batch (m_blocks == 1), we only need (128, 128, 256) + # for large batch (m_blocks > 1), we only need (64, 256, 256) if m_blocks <= 1 and thread_configs[0] != 128: continue if m_blocks > 1 and thread_configs[0] != 64: continue + # we only support channelwise quantization and group_size == 128 + # for fp8 + if scalar_type == "sglang::kFE4M3fn" and group_blocks not in [-1, 8]: + continue + # nvfp4 only supports group_size == 16 + # mxfp4 only supports group_size == 32 + if scalar_type == "sglang::kFE2M1f" and group_blocks not in [1, 2]: + continue + # other quantization methods don't support group_size = 16 + if scalar_type != "sglang::kFE2M1f" and group_blocks == 1: + continue + k_blocks = thread_configs[0] // 16 n_blocks = thread_configs[1] // 16 threads = thread_configs[2] c_dtype = "half" if dtype == "fp16" else "nv_bfloat16" + if scalar_type == "sglang::kFE2M1f" and group_blocks == 1: + s_type = "sglang::kFE4M3fn" + elif scalar_type == "sglang::kFE2M1f" and group_blocks == 2: + s_type = "sglang::kFE8M0fnu" + if dtype == "fp16": + # we cannot safely dequantize e8m0 to fp16, so skip this + continue + elif dtype == "fp16": + s_type = "sglang::kFloat16" + elif dtype == "bf16": + s_type = "sglang::kBFloat16" + template_str = jinja2.Template(TEMPLATE).render( scalar_t=c_dtype, w_type_id=scalar_type + ".id()", + s_type_id=s_type + ".id()", threads=threads, thread_m_blocks=max(m_blocks, 1), thread_n_blocks=n_blocks, thread_k_blocks=k_blocks, m_block_size_8=m_blocks == 0.5, stages="pipe_stages", - has_act_order=has_act_order, - has_zp=has_zp, group_blocks=group_blocks, is_zp_float=False, ) @@ -108,6 +136,7 @@ def generate_new_kernels(): file_content = FILE_HEAD + "\n\n" file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" + # Remove "sglang::" prefix (8 chars) from scalar_type for filename filename = f"kernel_{dtype}_{scalar_type[8:].lower()}.cuh" with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel.h b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel.h index afa7c377b..f3f0bf03c 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel.h +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel.h @@ -1,7 +1,6 @@ -#pragma once #ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#define MARLIN_NAMESPACE_NAME marlin_moe_wna16_v2 #endif #include "gemm/marlin/marlin.cuh" @@ -10,16 +9,18 @@ #define MARLIN_KERNEL_PARAMS \ const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ scales_ptr, const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ + const int4 *__restrict__ b_bias_ptr, const int4 *__restrict__ scales_ptr, \ + const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ const int32_t *__restrict__ sorted_token_ids_ptr, const int32_t *__restrict__ expert_ids_ptr, \ const int32_t *__restrict__ num_tokens_past_padded_ptr, const float *__restrict__ topk_weights_ptr, int top_k, \ bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, int prob_n, int prob_k, int *locks, \ - bool use_atomic_add, bool use_fp32_reduce + bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem namespace MARLIN_NAMESPACE_NAME { template < typename scalar_t, // compute dtype, half or nv_float16 const sglang::ScalarTypeId w_type_id, // weight ScalarType id + const sglang::ScalarTypeId s_type_id, // weight scale ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the @@ -30,8 +31,6 @@ template < // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const bool has_zp, // whether zero-points are enabled const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4.cuh index 7e83bed8f..51619bb5a 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4.cuh @@ -2,89 +2,38 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4b8.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4b8.cuh index 60e2dea31..e192eb56a 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4b8.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku4b8.cuh @@ -2,109 +2,46 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku8b128.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku8b128.cuh index 7eb6b18de..789d6c5f2 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku8b128.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_bf16_ku8b128.cuh @@ -2,109 +2,46 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4.cuh index ec41e018b..f69131038 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4.cuh @@ -2,89 +2,38 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4b8.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4b8.cuh index 7df28701b..b9611adc9 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4b8.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku4b8.cuh @@ -2,109 +2,46 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku8b128.cuh b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku8b128.cuh index 1150844e2..1fbe1cf4c 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku8b128.cuh +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/kernel_fp16_ku8b128.cuh @@ -2,109 +2,46 @@ // clang-format off #pragma once -#include "kernel.h" -#include "marlin_template.h" - namespace MARLIN_NAMESPACE_NAME { -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); - -template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/marlin_template.h b/sgl-kernel/csrc/moe/marlin_moe_wna16/marlin_template.h index ade562af6..1ba99787b 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/marlin_template.h +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/marlin_template.h @@ -18,12 +18,12 @@ /* * Adapted from https://github.com/IST-DASLab/marlin */ -#pragma once #ifndef MARLIN_NAMESPACE_NAME #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif +#include "gemm/marlin/dequant.h" #include "gemm/marlin/marlin.cuh" #include "gemm/marlin/marlin_dtypes.cuh" #include "scalar_type.hpp" @@ -50,8 +50,6 @@ template < // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const bool has_zp, // whether zero-points are enabled const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? @@ -79,8 +77,8 @@ __global__ void Marlin( int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce // whether to use fp32 global reduce -) {} + bool use_fp32_reduce, // whether to use fp32 global reduce + int max_shared_mem) {} } // namespace MARLIN_NAMESPACE_NAME @@ -177,130 +175,6 @@ __device__ inline void ldsm(typename ScalarType::FragA& frag_a, const } } -// Lookup-table based 3-input logical operation; explicitly used for -// dequantization as the compiler does not seem to automatically recognize it in -// all cases. -template -__device__ inline int lop3(int a, int b, int c) { - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -// Constructs destination register by taking bytes from 2 sources (based on -// mask) -template -__device__ inline uint32_t prmt(uint32_t a) { - uint32_t res; - asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(res) : "r"(a), "n"(start_byte), "n"(mask)); - return res; -} - -template -__device__ inline typename ScalarType::FragB dequant(int q, typename ScalarType::FragB& frag_b); - -// -// Efficiently dequantize 4bit values packed in an int32 value into a full -// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, -// with some small changes: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 -// -template <> -__device__ inline typename ScalarType::FragB dequant(int q, typename ScalarType::FragB& frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64086408; - const int MUL = 0x2c002c00; - const int ADD = 0xd480d480; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2( - *reinterpret_cast(&hi), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); - return frag_b; -} - -template <> -__device__ inline typename ScalarType::FragB -dequant(int q, typename ScalarType::FragB& frag_b) { - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t EX = 0x43004300; - - // Guarantee that the `(a & b) | c` operations are LOP3s. - - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - - static constexpr uint32_t MUL = 0x3F803F80; - static constexpr uint32_t ADD = 0xC308C308; - - frag_b[0] = __hfma2( - *reinterpret_cast(&lo), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - frag_b[1] = __hfma2( - *reinterpret_cast(&hi), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); - return frag_b; -} - -// -// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or -// bf16 Reference: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 -// -template <> -__device__ inline typename ScalarType::FragB dequant(int q, typename ScalarType::FragB& frag_b) { - static constexpr uint32_t mask_for_elt_01 = 0x5250; - static constexpr uint32_t mask_for_elt_23 = 0x5351; - static constexpr uint32_t start_byte_for_fp16 = 0x64646464; - - uint32_t lo = prmt(q); - uint32_t hi = prmt(q); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; - - frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(*reinterpret_cast(&hi), *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - return frag_b; -} - -template <> -__device__ inline typename ScalarType::FragB -dequant(int q, typename ScalarType::FragB& frag_b) { - float fp32_intermediates[4]; - uint32_t* fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388736.f; - fp32_intermediates[1] -= 8388736.f; - fp32_intermediates[2] -= 8388736.f; - fp32_intermediates[3] -= 8388736.f; - - uint32_t* bf16_result_ptr = reinterpret_cast(&frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632); - - return frag_b; -} - // Multiply dequantized values by the corresponding quantization scale; used // only for grouped quantization. template @@ -406,6 +280,7 @@ __device__ inline void wait_negative_and_add(int* lock) { template < typename scalar_t, // compute dtype, half or nv_float16 const sglang::ScalarTypeId w_type_id, // weight ScalarType id + const sglang::ScalarTypeId s_type_id, // weight scale ScalarType id const int threads, // number of threads in a threadblock const int thread_m_blocks, // number of 16x16 blocks in the m // dimension (batchsize) of the @@ -416,19 +291,20 @@ template < // only works when thread_m_blocks == 1 const int stages, // number of stages for the async global->shared // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const bool has_zp, // whether zero-points are enabled const int group_blocks, // number of consecutive 16x16 blocks // with a separate quantization scale const bool is_zp_float // is zero point of float16 type? > __global__ void Marlin( - const int4* __restrict__ A, // fp16 input matrix of shape mxk - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn - int4* __restrict__ C, // fp16 output buffer of shape mxn - int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ b_bias_ptr, const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape // (k/groupsize)xn + const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4 + // only) const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape // (k/groupsize)x(n/pack_factor) const int* __restrict__ g_idx, // int32 group indices of shape k @@ -444,9 +320,10 @@ __global__ void Marlin( int prob_n, // output dimension n int prob_k, // reduction dimension k int* locks, // extra global storage for barrier synchronization - bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce // whether to use fp32 global reduce -) { + bool has_bias, + bool use_atomic_add, // whether to use atomic add to reduce + bool use_fp32_reduce, // whether to use fp32 global reduce + int max_shared_mem) { // Each threadblock processes one "stripe" of the B matrix with (roughly) the // same size, which might involve multiple column "slices" (of width 16 * // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM @@ -468,14 +345,36 @@ __global__ void Marlin( extern __shared__ int4 sh[]; static constexpr auto w_type = sglang::ScalarType::from_id(w_type_id); + static constexpr auto s_type = sglang::ScalarType::from_id(s_type_id); + if constexpr (w_type == sglang::kFE2M1f) { + static_assert(s_type == sglang::kFE4M3fn && group_blocks == 1 || s_type == sglang::kFE8M0fnu && group_blocks == 2); + } else if constexpr (std::is_same::value) { + static_assert(s_type == sglang::kBFloat16); + } else if constexpr (std::is_same::value) { + static_assert(s_type == sglang::kFloat16); + } + + constexpr bool has_zp = w_type == sglang::kU4 || w_type == sglang::kU8; + constexpr bool is_int_type = + w_type == sglang::kU4 || w_type == sglang::kU8 || w_type == sglang::kU4B8 || w_type == sglang::kU8B128; + // see comments of dequant.h for more details + constexpr bool dequant_skip_flop = w_type == sglang::kFE4M3fn || + w_type == sglang::kFE2M1f && s_type == sglang::kFE4M3fn || + has_zp && !is_zp_float && !std::is_same::value || + has_zp && !is_zp_float && !(w_type == sglang::kU8); + + scalar_t2 global_scale; + + constexpr bool has_act_order = group_blocks == 0; constexpr int pack_factor = 32 / w_type.size_bits(); static_assert(thread_m_blocks == 1 || !m_block_size_8); constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); const int group_size = (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; - const int scales_expert_stride = prob_n * prob_k / group_size / 8; + const int scales_expert_stride = prob_n * prob_k / group_size / (w_type == sglang::kFE2M1f ? 16 : 8); const int zp_expert_stride = is_zp_float ? prob_n * prob_k / group_size / 8 : prob_n * prob_k / group_size / (pack_factor * 4); + const int b_bias_expert_stride = prob_n / 8; // parallel: num valid moe blocks int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; @@ -517,10 +416,14 @@ __global__ void Marlin( int64_t B_expert_off = 0; int4* sh_block_sorted_ids_int4 = sh; + int4* sh_rd_block_sorted_ids_int4 = sh_block_sorted_ids_int4 + moe_block_size / 4; + int4* sh_block_topk_weights_int4 = sh_rd_block_sorted_ids_int4 + moe_block_size / 4; + // sh_block_topk_weights_int4 only need (moe_block_size / 4); + // but we pad to align to 256 bytes + int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2 + moe_block_size; int32_t* sh_block_sorted_ids = reinterpret_cast(sh_block_sorted_ids_int4); - int4* sh_block_topk_weights_int4 = sh_block_sorted_ids_int4 + moe_block_size / 4; + int32_t* sh_rd_block_sorted_ids = reinterpret_cast(sh_rd_block_sorted_ids_int4); scalar_t2* sh_block_topk_weights = reinterpret_cast(sh_block_topk_weights_int4); - int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 4; int32_t block_num_valid_tokens = 0; int32_t locks_off = 0; @@ -564,11 +467,24 @@ __global__ void Marlin( sh_block_sorted_ids_int4[tid4] = reinterpret_cast(sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4]; +#pragma unroll + for (int i = 0; i < 4; i++) + sh_rd_block_sorted_ids[tid4 * 4 + i] = sh_block_sorted_ids[tid4 * 4 + i] / top_k; + if (mul_topk_weights) { #pragma unroll for (int i = 0; i < 4; i++) { - sh_block_topk_weights[tid4 * 4 + i] = - Dtype::num2num2(Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]])); + int idx = tid4 * 4 + i; + // idx = idx < block_num_valid_tokens ? idx : 0; + if (idx < block_num_valid_tokens) { + if constexpr (w_type == sglang::kFE2M1f && s_type == sglang::kFE4M3fn) { + sh_block_topk_weights[idx] = + __hmul2(global_scale, Dtype::num2num2(Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]))); + } else { + sh_block_topk_weights[idx] = + Dtype::num2num2(Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]])); + } + } } } } @@ -599,6 +515,11 @@ __global__ void Marlin( expert_id = expert_ids_ptr[block_id]; } + if constexpr (w_type == sglang::kFE2M1f && s_type == sglang::kFE4M3fn) { + uint16_t val = scale2_ptr[expert_id]; + global_scale = Dtype::num2num2(*reinterpret_cast(&val)); + } + B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4); scales_ptr += (expert_id - old_expert_id) * scales_expert_stride; if constexpr (has_zp) { @@ -607,6 +528,9 @@ __global__ void Marlin( if constexpr (has_act_order) { g_idx += (expert_id - old_expert_id) * prob_k; } + if (has_bias) { + b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride; + } read_moe_block_data(block_id); }; @@ -707,8 +631,9 @@ __global__ void Marlin( // Scale sizes/strides without act_order int s_gl_stride = prob_n / 8; constexpr int s_sh_stride = 16 * thread_n_blocks / 8; - constexpr int s_tb_groups = - !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks ? thread_k_blocks / group_blocks : 1; + constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks / (w_type == sglang::kFE2M1f ? 2 : 1) + : 1; constexpr int s_sh_stage = s_tb_groups * s_sh_stride; int s_gl_rd_delta = s_gl_stride; @@ -717,6 +642,7 @@ __global__ void Marlin( constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; // constexpr int act_s_row_stride = 1; // int act_s_col_stride = act_s_row_stride * num_groups; + constexpr int act_s_max_num_groups = 32; int act_s_col_stride = 1; int act_s_col_warp_stride = act_s_col_stride * 8; int tb_n_warps = thread_n_blocks / 4; @@ -730,8 +656,9 @@ __global__ void Marlin( int zp_gl_rd_delta = zp_gl_stride; // Global A read index of current thread. - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; + int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o; + int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; + // Shared write index of current thread. int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); // Shared read index. @@ -742,8 +669,8 @@ __global__ void Marlin( int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; b_gl_rd += b_sh_stride * slice_col; b_gl_rd += b_gl_rd_delta_o * slice_row; - int b_sh_wr = threadIdx.x * b_thread_vecs; - int b_sh_rd = threadIdx.x * b_thread_vecs; + auto b_sh_wr = threadIdx.x * b_thread_vecs; + auto b_sh_rd = threadIdx.x * b_thread_vecs; // For act_order constexpr int k_iter_size = tb_k / b_sh_wr_iters; @@ -758,10 +685,11 @@ __global__ void Marlin( if constexpr (group_blocks == -1) { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) / (w_type == sglang::kFE2M1f ? 2 : 1) + + s_sh_stride * slice_col + threadIdx.x; } } - int s_sh_wr = threadIdx.x; + auto s_sh_wr = threadIdx.x; bool s_sh_wr_pred = threadIdx.x < s_sh_stride; // Zero-points @@ -773,20 +701,38 @@ __global__ void Marlin( zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; } } - int zp_sh_wr = threadIdx.x; + auto zp_sh_wr = threadIdx.x; bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride; // We use a different scale layout for grouped and column-wise quantization as // we scale a `half2` tile in column-major layout in the former and in // row-major in the latter case. int s_sh_rd; - if constexpr (group_blocks != -1) + if constexpr (group_blocks != -1 && w_type == sglang::kFE2M1f) { + auto warp_id = threadIdx.x / 32; + int n_warps = thread_n_blocks / 4; + int warp_row = warp_id / n_warps; + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; - else if constexpr (group_blocks == -1 && (m_block_size_8 || has_zp)) + s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2; + + } else if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; + else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; else s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; + int bias_sh_rd; + if constexpr (m_block_size_8) { + bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; + } else { + bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; + } + + int bias_sh_wr = threadIdx.x; + int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; + // Zero-points have the same read layout as the scales // (without column-wise case) constexpr int num_col_threads = 8; @@ -812,7 +758,7 @@ __global__ void Marlin( // each warp must also write a consecutive memory segment? auto transform_a = [&](int i) { int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); }; // Since the computation of this remapping is non-trivial and, due to our main // loop unrolls, all shared memory accesses are static, we simply precompute @@ -839,18 +785,38 @@ __global__ void Marlin( B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; // Shared memory storage for global fetch pipelines. - int4* sh_a = sh_new; - int4* sh_b = sh_a + (stages * a_sh_stage); - int4* sh_g_idx = sh_b + (stages * b_sh_stage); + constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; + constexpr int sh_b_size = stages * b_sh_stage; + int4* sh_b = sh_new; + int4* sh_red = sh_new; + + constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); + constexpr int sh_b_red_bias_size = + sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) ? sh_size_b_red_max : (sh_size_b_red_min + sh_bias_size); + + int4* sh_bias = sh_new + sh_size_b_red_min; + int4* sh_g_idx = sh_new + sh_b_red_bias_size; int4* sh_zp = sh_g_idx + (stages * g_idx_stage); + constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); int4* sh_s = sh_zp + (stages * zp_sh_stage); - int4* sh_red = sh_b; + // shared memory reused by reduction should be smaller than + // shared memory used by weight. + static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <= stages * b_sh_stage); + int4* sh_a = sh_s + sh_s_size; + constexpr int shm_size_used = moe_block_size + stages * (g_idx_stage + zp_sh_stage) + sh_s_size + sh_b_red_bias_size; + + // all remaining shared memory is used to cache A (input) + // sh_a_max_row is at least ` stages * 16 * thread_m_blocks ` + int sh_a_max_row = ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2); // Register storage for double buffer of shared memory reads. FragA frag_a[2][thread_m_blocks]; I4 frag_b_quant[2][b_thread_vecs]; FragC frag_c[thread_m_blocks][4][2]; - FragS frag_s[2][4]; // No act-order + FragS frag_s[2][4]; // No act-order + FragS frag_bias[2][4]; FragS act_frag_s[2][4][4]; // For act-order int frag_qzp[2][num_ints_per_thread]; // Zero-points FragZP frag_zp; // Zero-points in fp16 @@ -865,14 +831,13 @@ __global__ void Marlin( int sh_first_group_id = -1; int sh_num_groups = -1; - constexpr int sh_max_num_groups = 32; auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) { sh_first_group_id = first_group_id; sh_num_groups = last_group_id - first_group_id + 1; - if (sh_num_groups < sh_max_num_groups) { - sh_num_groups = sh_max_num_groups; + if (sh_num_groups > act_s_max_num_groups) { + sh_num_groups = act_s_max_num_groups; } if (sh_first_group_id + sh_num_groups > num_groups) { @@ -898,24 +863,26 @@ __global__ void Marlin( } } }; + // Asynchronously fetch the next A, B and s tile from global to the next // shared memory pipeline location. - int a_remaining_load_count_in_slice = stages; - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + bool should_load_a = true; + int max_num_stage_groups = ((sh_a_max_row - moe_block_size) / moe_block_size + 1) / stages; + max_num_stage_groups = max(max_num_stage_groups, 1); + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true, int pipe_a = 0) { if (pred) { - int4* sh_a_stage = sh_a + a_sh_stage * pipe; - if (prob_k > thread_k_blocks * 16 * stages || slice_col == 0 || a_remaining_load_count_in_slice > 0) { - a_remaining_load_count_in_slice--; + if (should_load_a) { + int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a; #pragma unroll for (int i = 0; i < a_sh_wr_iters; i++) { - int a_idx = a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off; - int row = a_idx / a_gl_stride; + int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row; int64_t sorted_row = 0; - if (!m_block_size_8 || row < 8) sorted_row = sh_block_sorted_ids[row] / top_k; - int64_t true_idx = sorted_row * a_gl_stride + a_idx % a_gl_stride; + if (!m_block_size_8 || row < 8) sorted_row = sh_rd_block_sorted_ids[row]; + int64_t true_idx = sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off; cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx], row < block_num_valid_tokens); } } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) { @@ -1013,8 +980,8 @@ __global__ void Marlin( // Load the next sub-tile from the current location in the shared memory pipe // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) { - int4* sh_a_stage = sh_a + a_sh_stage * pipe; + auto fetch_to_registers = [&](int k, int pipe, int pipe_a = 0) { + int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a; #pragma unroll for (int i = 0; i < thread_m_blocks; i++) ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); @@ -1057,11 +1024,15 @@ __global__ void Marlin( } } else if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { - int4* sh_s_stage = - sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + if (k % b_sh_wr_iters == 0) { + int4* sh_s_stage = + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; + } } else { - int warp_id = threadIdx.x / 32; + auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; @@ -1070,11 +1041,19 @@ __global__ void Marlin( cur_k += k_iter_size * (k % b_sh_wr_iters); int k_blocks = cur_k / 16; - int cur_group_id = k_blocks / group_blocks; + int cur_group_id = k_blocks / (group_blocks * (w_type == sglang::kFE2M1f ? 2 : 1)); int4* sh_s_stage = sh_s + s_sh_stage * pipe; - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + if constexpr (w_type_id != sglang::kFE2M1f.id()) { + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) { + reinterpret_cast(&frag_s[k % 2])[0] = + reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } else { + reinterpret_cast(&frag_s[k % 2])[0] = + reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) + k % 2]; + } } } @@ -1098,7 +1077,7 @@ __global__ void Marlin( // Determine "position" inside the thread-block (based on warp and // thread-id) - int warp_id = threadIdx.x / 32; + auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N int warp_row = warp_id / n_warps; @@ -1106,7 +1085,7 @@ __global__ void Marlin( cur_k += warp_row * 16; - int th_id = threadIdx.x % 32; + auto th_id = threadIdx.x % 32; cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix int s_col_shift = @@ -1162,13 +1141,16 @@ __global__ void Marlin( } } else if constexpr (group_blocks >= thread_k_blocks) { - int4* sh_zp_stage = - sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + if (k % b_sh_wr_iters == 0) { + int4* sh_zp_stage = + sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); +#pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + } } } else { - int warp_id = threadIdx.x / 32; + auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; @@ -1189,6 +1171,7 @@ __global__ void Marlin( sh_zp_stage += cur_group_id * zp_sh_stride; +#pragma unroll for (int i = 0; i < num_ints_per_thread; i++) { frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; } @@ -1200,11 +1183,13 @@ __global__ void Marlin( if constexpr (group_blocks != -1) { if constexpr (group_blocks >= thread_k_blocks) { - int4* sh_zp_stage = - sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; + if (k % b_sh_wr_iters == 0) { + int4* sh_zp_stage = + sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); + reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; + } } else { - int warp_id = threadIdx.x / 32; + auto warp_id = threadIdx.x / 32; int n_warps = thread_n_blocks / 4; int warp_row = warp_id / n_warps; @@ -1227,6 +1212,10 @@ __global__ void Marlin( } }; + auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) { + dequant(q, frag_b_ptr); + }; + // Execute the actual tensor core matmul of a sub-tile. bool is_first_matmul_in_slice = true; auto matmul = [&](int k) { @@ -1236,8 +1225,6 @@ __global__ void Marlin( if constexpr (has_zp && !is_zp_float) { if (is_new_zp) { if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; - FragB frag_zp_0; - FragB frag_zp_1; int zp_quant_0, zp_quant_1; if constexpr (w_type.size_bits() == 4) { @@ -1249,15 +1236,28 @@ __global__ void Marlin( zp_quant_1 = frag_qzp[k2][1]; } - dequant(zp_quant_0, frag_zp_0); - dequant(zp_quant_1, frag_zp_1); - - frag_zp[0] = frag_zp_0[0]; - frag_zp[1] = frag_zp_0[1]; - frag_zp[2] = frag_zp_1[0]; - frag_zp[3] = frag_zp_1[1]; + dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); + dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); } } + if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { + if (is_new_zp) { + reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; + } + } + + // Commented out FP4/FP8 scale dequantization since we don't generate + // kFE2M1f kernels to reduce compilation time + // if constexpr (w_type == sglang::kFE2M1f) { + // int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; + // int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; + // + // dequant_fp8_scales( + // s_quant_0, reinterpret_cast(&frag_s[k2])); + // dequant_fp8_scales( + // s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); + // } + // We have the m dimension as the inner loop in order to encourage overlapping // dequantization and matmul operations. #pragma unroll @@ -1266,7 +1266,10 @@ __global__ void Marlin( FragB frag_b1; int b_quant_0, b_quant_1; - if constexpr (w_type.size_bits() == 4) { + if constexpr (w_type_id == sglang::kFE2M1f.id()) { + b_quant_1 = frag_b_quant[k2][0][j]; + b_quant_0 = b_quant_1 << 8; + } else if constexpr (w_type.size_bits() == 4) { b_quant_0 = frag_b_quant[k2][0][j]; b_quant_1 = b_quant_0 >> 8; } else { @@ -1276,8 +1279,13 @@ __global__ void Marlin( b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; } - dequant(b_quant_0, frag_b0); - dequant(b_quant_1, frag_b1); + dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); + dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); + + if constexpr (dequant_skip_flop && has_zp && !is_zp_float) { + sub_zp(frag_b0, frag_zp[j], 0); + sub_zp(frag_b1, frag_zp[j], 1); + } // Apply scale to frag_b0 if constexpr (has_act_order) { @@ -1285,9 +1293,8 @@ __global__ void Marlin( scale4( frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); scale4( - frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k][2][j], act_frag_s[k2][3][j], 1); - - } else if constexpr (has_zp && !is_zp_float && group_blocks == -1) { + frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); + } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1) { int idx = (threadIdx.x / 4) % 2; scalar_t2 s2 = Dtype::nums2num2( reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], @@ -1295,14 +1302,10 @@ __global__ void Marlin( if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); scale_and_sub(frag_b0, s2.x, frag_zp[j].x); scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } else if constexpr (has_zp && !is_zp_float && group_blocks != -1) { + } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) { if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k % 2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k % 2][j][0].y, frag_zp[j].y); - } else if constexpr (has_zp && is_zp_float && group_blocks != -1) { - if (is_new_zp) frag_zpf[k2][j] = __hmul2(frag_zpf[k2][j], *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j].x, frag_zpf[k2][j].x); - scale_and_sub(frag_b1, frag_s[k2][j].y, frag_zpf[k2][j].y); + scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); + scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); } else if constexpr (group_blocks != -1) { scale(frag_b0, frag_s[k2][j], 0); scale(frag_b1, frag_s[k2][j], 1); @@ -1327,7 +1330,7 @@ __global__ void Marlin( auto thread_block_reduce = [&]() { constexpr int red_off = threads / b_sh_stride_threads / 2; if (red_off >= 1) { - int red_idx = threadIdx.x / b_sh_stride_threads; + auto red_idx = threadIdx.x / b_sh_stride_threads; constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; constexpr int red_sh_delta = b_sh_stride_threads; int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); @@ -1513,7 +1516,7 @@ __global__ void Marlin( // Write out the reduce final result in the correct layout. We only actually // reshuffle matrix fragments in this step, the reduction above is performed // in fragment layout. - auto write_result = [&]() { + auto write_result = [&](bool last) { int c_gl_stride = prob_n / 8; constexpr int c_sh_stride = 2 * thread_n_blocks + 1; int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); @@ -1534,13 +1537,31 @@ __global__ void Marlin( // We first reorder in shared memory to guarantee the most efficient final // global write patterns - auto write = [&](int idx, float c0, float c1, FragS& s) { + auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) { scalar_t2 res = Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); // For per-column quantization we finally apply the scale here (only for // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && w_type.size_bits() == 4 && !has_zp) { - res = __hmul2(res, s[0]); + if constexpr ( + !has_act_order && group_blocks == -1 && w_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { + scalar_t2 tmp_scale = s[0]; + if constexpr (m_block_size_8) { + tmp_scale = Dtype::num2num2(reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); + } + res = __hmul2(res, tmp_scale); + } + + if constexpr (w_type == sglang::kFE2M1f && s_type == sglang::kFE4M3fn) { + if (!mul_topk_weights) { + res = __hmul2(res, global_scale); + } + } + if (has_bias && last) { + scalar_t2 tmp_bias = b_bias[0]; + if constexpr (m_block_size_8) { + tmp_bias = Dtype::num2num2(reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + } + res = __hadd2(res, tmp_bias); } if constexpr (m_block_size_8) { @@ -1558,18 +1579,44 @@ __global__ void Marlin( for (int j = 0; j < 4; j++) { if constexpr (m_block_size_8) { int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1]); + write( + wr, + frag_c[i][j][0][0], + frag_c[i][j][0][1], + frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write( + wr + 8, + frag_c[i][j][0][2], + frag_c[i][j][0][3], + frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); } else { int wr = c_sh_wr + 8 * j; write( - wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); + wr + (4 * c_sh_stride) * 0 + 0, + frag_c[i][j][0][0], + frag_c[i][j][0][1], + frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); write( - wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); + wr + (4 * c_sh_stride) * 8 + 0, + frag_c[i][j][0][2], + frag_c[i][j][0][3], + frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); write( - wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); + wr + (4 * c_sh_stride) * 0 + 4, + frag_c[i][j][1][0], + frag_c[i][j][1][1], + frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); write( - wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); + wr + (4 * c_sh_stride) * 8 + 4, + frag_c[i][j][1][2], + frag_c[i][j][1][3], + frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); } } c_sh_wr += 16 * (4 * c_sh_stride); @@ -1627,10 +1674,12 @@ __global__ void Marlin( if constexpr (has_zp && !is_zp_float && group_blocks == -1) { if (i == 0) { fetch_col_zp_to_shared(); - fetch_col_scale_to_shared(); + if constexpr (!dequant_skip_flop) { + fetch_col_scale_to_shared(); + } } } - fetch_to_shared(i, i, i < slice_iters); + fetch_to_shared(i, i, i < slice_iters, i); } zero_accums(); @@ -1639,8 +1688,10 @@ __global__ void Marlin( fetch_to_registers(0, 0); fetch_scales_to_registers(0, 0); fetch_zp_to_registers(0, 0); - a_gl_rd += a_gl_rd_delta_o * (stages - 1); - slice_k_start_shared_fetch += tb_k * (stages - 1); + a_gl_rd_col += a_gl_rd_delta_o * (stages - 1); + if constexpr (has_act_order) { + slice_k_start_shared_fetch += tb_k * (stages - 1); + } }; if (slice_iters) { start_pipes(); @@ -1653,44 +1704,56 @@ __global__ void Marlin( // have even length meaning that the next iteration will always start at // index 0. + for (int stage_group_id = 0; stage_group_id < max_num_stage_groups; stage_group_id++) { #pragma unroll - for (int pipe = 0; pipe < stages;) { + for (int pipe = 0; pipe < stages;) { #pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) { - fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); + for (int k = 0; k < b_sh_wr_iters; k++) { + int idx = (pipe >= stages && stage_group_id == max_num_stage_groups - 1) ? (pipe - stages) + : (pipe + stage_group_id * stages); + fetch_to_registers(k + 1, pipe % stages, idx); + fetch_scales_to_registers(k + 1, pipe); + fetch_zp_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + int idx = (pipe >= 1 && stage_group_id == max_num_stage_groups - 1) + ? (pipe - 1) + : (pipe + (stage_group_id + 1) * stages - 1); + fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages, idx); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + matmul(k); + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd_col += a_gl_rd_delta_o * stages; + + if constexpr (has_act_order) { + slice_k_start += tb_k * stages; + + if (slice_k_start < prob_k) { + slice_k_start_shared_fetch += tb_k * stages; + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); + __syncthreads(); + } } - matmul(k); } - slice_iters--; if (slice_iters == 0) { break; } } - a_remaining_load_count_in_slice = 0; - - a_gl_rd += a_gl_rd_delta_o * stages; - slice_k_start += tb_k * stages; - slice_k_start_shared_fetch += tb_k * stages; - - if constexpr (has_act_order) { - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); - __syncthreads(); - } - } // Process results and, if necessary, proceed to the next column slice. // While this pattern may not be the most readable, other ways of writing @@ -1700,7 +1763,7 @@ __global__ void Marlin( bool last = slice_idx == slice_count - 1; // For per-column scales, we only fetch them here in the final step before // write-out - if constexpr (!has_act_order && group_blocks == -1 && !has_zp) { + if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { if (s_sh_wr_pred) { cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); @@ -1710,7 +1773,14 @@ __global__ void Marlin( } thread_block_reduce(); - if constexpr (!has_act_order && group_blocks == -1 && !has_zp) { + + if (has_bias && last) { + __syncthreads(); + cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], threadIdx.x < 16 * thread_n_blocks / 8); + cp_async_fence(); + } + + if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { if (w_type.size_bits() == 8 || (last || use_atomic_add)) { cp_async_wait<0>(); __syncthreads(); @@ -1732,7 +1802,8 @@ __global__ void Marlin( // For 8-bit channelwise, we apply the scale before the global reduction // that converts the fp32 results to fp16 (so that we avoid possible // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && w_type.size_bits() == 8 && !has_zp) { + if constexpr ( + !has_act_order && group_blocks == -1 && w_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { if (threadIdx.x / 32 < thread_n_blocks / 4) { #pragma unroll for (int i = 0; i < thread_m_blocks; i++) { @@ -1761,18 +1832,42 @@ __global__ void Marlin( } barrier_release(&locks[locks_off], last); } + + if (has_bias && last) { + cp_async_wait<0>(); + __syncthreads(); + reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; + reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; + __syncthreads(); + } + if (use_atomic_add && slice_count > 1 && slice_idx != 0) wait_negative_and_add(&locks[locks_off]); if (last || use_atomic_add) // only the last block in a slice actually writes the result - write_result(); - if (slice_row) a_remaining_load_count_in_slice = stages; + write_result(last); + int old_slice_row = slice_row; slice_row = 0; slice_col_par++; slice_col++; is_first_matmul_in_slice = true; init_slice(); + + // Should we load A matrix in next slice? + // `slice_col == 0`: when move to a new moe block + // `old_slice_row > 0`: + // when the last slice is not starting from k_index == 0 + // (only happen when it is the first slice of a threadblock) + // `prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups`: + // when the required shared memory size is larger than + // the remaining shared memory + if (slice_col == 0 || old_slice_row || prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups) { + should_load_a = true; + } else { + should_load_a = false; + } + if (slice_iters) { - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); + a_gl_rd_col = (threadIdx.x % a_gl_rd_delta_o); #pragma unroll for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; @@ -1782,18 +1877,17 @@ __global__ void Marlin( B_ptr[i] -= b_gl_stride; } + bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; // Update slice k/n for scales loading if constexpr (has_act_order) { slice_k_start = tb_k * slice_row; slice_k_finish = slice_k_start + tb_k * slice_iters; slice_k_start_shared_fetch = slice_k_start; slice_n_offset = act_s_col_tb_stride * slice_col; - } else { s_gl_rd = s_sh_stride * slice_col + threadIdx.x; zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; } - start_pipes(); } } diff --git a/sgl-kernel/csrc/moe/marlin_moe_wna16/ops.cu b/sgl-kernel/csrc/moe/marlin_moe_wna16/ops.cu index b249f6415..57334663a 100644 --- a/sgl-kernel/csrc/moe/marlin_moe_wna16/ops.cu +++ b/sgl-kernel/csrc/moe/marlin_moe_wna16/ops.cu @@ -25,6 +25,7 @@ #include "kernel.h" #include "kernel_marlin.cuh" +#include "marlin_template.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert( \ @@ -50,13 +51,16 @@ __global__ void permute_cols_kernel( int size_m, int size_k, int top_k) {}; -} + +} // namespace marlin torch::Tensor moe_wna16_marlin_gemm( torch::Tensor& a, - std::optional const& c_or_none, + std::optional c_or_none, torch::Tensor& b_q_weight, + std::optional const& b_bias_or_none, torch::Tensor& b_scales, + std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, std::optional const& g_idx_or_none, std::optional const& perm_or_none, @@ -131,7 +135,7 @@ __global__ void permute_cols_kernel( int base_k = 0; for (int i = 0; i < iters; i++) { - int cur_k = base_k + threadIdx.x; + auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; @@ -141,7 +145,7 @@ __global__ void permute_cols_kernel( if (rest) { if (threadIdx.x < rest) { - int cur_k = base_k + threadIdx.x; + auto cur_k = base_k + threadIdx.x; int src_pos = perm_int_ptr[cur_k]; out_half[cur_k] = a_row_half[src_pos]; @@ -215,7 +219,6 @@ int get_scales_cache_size( int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K load_groups = max(load_groups, 32); // We load at least 32 scale groups return load_groups * tb_n * 2; - } else { int tb_scales = tb_groups * tb_n * 2; @@ -225,6 +228,7 @@ int get_scales_cache_size( int get_kernel_cache_size( thread_config_t const& th_config, + bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, @@ -242,11 +246,16 @@ int get_kernel_cache_size( int tb_n = th_config.thread_n; int tb_m = thread_m_blocks * 16; - // shm size for block_sorted_ids/block_topk_weights + // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32) - int sh_block_meta_size = tb_m * 4 * 2; + int sh_block_meta_size = tb_m * 4; int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; + int sh_red_size = tb_m * (tb_n + 8) * 2; + int sh_bias_size = tb_n * 2; + int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; + tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); + int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; @@ -260,13 +269,14 @@ int get_kernel_cache_size( sh_zp_size = sh_s_size / 2; } - int total_size = sh_a_size + sh_b_size + sh_s_size + sh_zp_size + sh_g_idx_size + sh_block_meta_size; + int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size + sh_block_meta_size; return total_size; } bool is_valid_config( thread_config_t const& th_config, + bool m_block_size_8, int thread_m_blocks, int prob_m, int prob_n, @@ -301,6 +311,7 @@ bool is_valid_config( // Check that pipeline fits into cache int cache_size = get_kernel_cache_size( th_config, + m_block_size_8, thread_m_blocks, prob_m, prob_n, @@ -311,105 +322,151 @@ bool is_valid_config( is_k_full, has_zp, is_zp_float); - return cache_size <= max_shared_mem; + return cache_size + 512 <= max_shared_mem; } -#define __GET_IF( \ - W_TYPE, \ - THREAD_M_BLOCKS, \ - THREAD_N_BLOCKS, \ - THREAD_K_BLOCKS, \ - M_BLOCK_SIZE_8, \ - HAS_ACT_ORDER, \ - HAS_ZP, \ - GROUP_BLOCKS, \ - NUM_THREADS, \ - IS_ZP_FLOAT) \ - else if ( \ - q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && \ - thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && has_act_order == HAS_ACT_ORDER && \ - has_zp == HAS_ZP && group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ - kernel = Marlin< \ - scalar_t, \ - W_TYPE.id(), \ - NUM_THREADS, \ - THREAD_M_BLOCKS, \ - THREAD_N_BLOCKS, \ - THREAD_K_BLOCKS, \ - M_BLOCK_SIZE_8, \ - pipe_stages, \ - HAS_ACT_ORDER, \ - HAS_ZP, \ - GROUP_BLOCKS, \ - IS_ZP_FLOAT>; \ +#define _GET_IF( \ + W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ + else if ( \ + q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && \ + thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && \ + num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ + constexpr auto S_TYPE = W_TYPE == sglang::kFE2M1f \ + ? (GROUP_BLOCKS == 1 ? sglang::kFE4M3fn : sglang::kFE8M0fnu) \ + : (std::is_same::value ? sglang::kFloat16 : sglang::kBFloat16); \ + kernel = Marlin< \ + scalar_t, \ + W_TYPE.id(), \ + S_TYPE.id(), \ + NUM_THREADS, \ + THREAD_M_BLOCKS, \ + THREAD_N_BLOCKS, \ + THREAD_K_BLOCKS, \ + M_BLOCK_SIZE_8, \ + pipe_stages, \ + GROUP_BLOCKS, \ + IS_ZP_FLOAT>; \ } -#define GPTQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, true, false, 0, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, false, 0, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 8, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 8, NUM_THREADS, false) +// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) +// this is the most common cases +// BIGGROUP: cases for big group size (group_blocks in [-1, 8]) +// FZP: cases for float-zero-point (is_zp_float = true) +// ACT: cases for act order case (group_blocks == 0) +// FP4: cases for nvfp4(e2m1) (group_blocks == 1) +#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) -#define GPTQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, false, 0, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, false, 0, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, false, 0, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 8, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 8, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 8, NUM_THREADS, false) +#define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ + \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ + \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) -#define AWQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 8, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 8, NUM_THREADS, false) +#define COMMON_GET_IF(W_TYPE) \ + COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ + COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ + COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ + COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) -#define AWQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 8, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 8, NUM_THREADS, false) \ - \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, -1, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 2, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, false) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 8, NUM_THREADS, false) +#define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) + +#define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) + +#define BIGGROUP_GET_IF(W_TYPE) \ + BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ + BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ + BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ + BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) + +#define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) + +#define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) + +#define NVFP4_GET_IF(W_TYPE) \ + NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ + NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ + NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ + NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128) + +#define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) + +#define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) + +#define MXFP4_GET_IF(W_TYPE) \ + MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ + MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ + MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ + MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128) // We currently have 4-bit models only with group_blocks == 4 -#define HQQ_GET_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, true) \ - __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, true) \ - __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, true) \ - __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, true) \ - __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, NUM_THREADS, true) +#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) + +#define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) + +#define FZP_GET_IF(W_TYPE) \ + FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ + FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ + FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ + FZP_GET_IF_M234(W_TYPE, 8, 4, 128) + +// We currently have 4-bit models only with group_blocks == 4 +#define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) + +#define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ + _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ + _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) + +#define ACT_GET_IF(W_TYPE) \ + ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ + ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ + ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ + ACT_GET_IF_M234(W_TYPE, 8, 4, 128) template MarlinFuncPtr get_marlin_kernel( @@ -427,23 +484,22 @@ MarlinFuncPtr get_marlin_kernel( auto kernel = MarlinDefault; if (false) { } - GPTQ_GET_IF_M1(sglang::kU4B8, 8, 8, 256) - GPTQ_GET_IF_M1(sglang::kU4B8, 8, 4, 128) - GPTQ_GET_IF_M234(sglang::kU4B8, 16, 4, 256) - GPTQ_GET_IF_M234(sglang::kU4B8, 8, 4, 128) + COMMON_GET_IF(sglang::kU4) + COMMON_GET_IF(sglang::kU4B8) + COMMON_GET_IF(sglang::kU8B128) - GPTQ_GET_IF_M1(sglang::kU8B128, 8, 8, 256) - GPTQ_GET_IF_M1(sglang::kU8B128, 8, 4, 128) + NVFP4_GET_IF(sglang::kFE2M1f) - GPTQ_GET_IF_M234(sglang::kU8B128, 16, 4, 256) - GPTQ_GET_IF_M234(sglang::kU8B128, 8, 4, 128) + BIGGROUP_GET_IF(sglang::kFE4M3fn) - AWQ_GET_IF_M1(sglang::kU4, 8, 8, 256) - AWQ_GET_IF_M1(sglang::kU4, 8, 4, 128) - - AWQ_GET_IF_M234(sglang::kU4, 16, 4, 256) - AWQ_GET_IF_M234(sglang::kU4, 8, 4, 128) + ACT_GET_IF(sglang::kU4B8) + ACT_GET_IF(sglang::kU8B128) + if (std::is_same::value) { + if (false) { + } + MXFP4_GET_IF(sglang::kFE2M1f) + } return kernel; } @@ -475,6 +531,7 @@ exec_config_t determine_exec_config( if (!is_valid_config( th_config, + m_block_size_8, thread_m_blocks, prob_m, prob_n, @@ -491,6 +548,7 @@ exec_config_t determine_exec_config( int cache_size = get_kernel_cache_size( th_config, + m_block_size_8, thread_m_blocks, prob_m, prob_n, @@ -504,7 +562,7 @@ exec_config_t determine_exec_config( int group_blocks = 0; if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; + group_blocks = group_size == -1 ? -1 : (group_size / 16); } auto kernel = get_marlin_kernel( @@ -546,7 +604,9 @@ void marlin_mm( const void* B, void* C, void* C_tmp, + void* b_bias, void* s, + void* s2, void* zp, void* g_idx, void* perm, @@ -564,6 +624,7 @@ void marlin_mm( int prob_k, void* workspace, sglang::ScalarType const& q_type, + bool has_bias, bool has_act_order, bool is_k_full, bool has_zp, @@ -587,8 +648,9 @@ void marlin_mm( q_type.str()); } else { TORCH_CHECK( - q_type == sglang::kU4B8 || q_type == sglang::kU8B128, - "q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ", + q_type == sglang::kU4B8 || q_type == sglang::kU8B128 || q_type == sglang::kFE4M3fn || q_type == sglang::kFE2M1f, + "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " + "has_zp = False. Got = ", q_type.str()); } @@ -620,7 +682,9 @@ void marlin_mm( const int4* B_ptr = (const int4*)B; int4* C_ptr = (int4*)C; int4* C_tmp_ptr = (int4*)C_tmp; + const int4* bias_ptr = (const int4*)b_bias; const int4* s_ptr = (const int4*)s; + const uint16_t* s2_ptr = (const uint16_t*)s2; const int4* zp_ptr = (const int4*)zp; const int* g_idx_ptr = (const int*)g_idx; const int* perm_ptr = (const int*)perm; @@ -705,6 +769,7 @@ void marlin_mm( TORCH_CHECK( is_valid_config( thread_tfg, + m_block_size_8, thread_m_blocks, prob_m, prob_n, @@ -787,10 +852,10 @@ void marlin_mm( // avoid ">>>" being formatted to "> > >" // clang-format off kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, + A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr, sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m, - prob_n, prob_k, locks, use_atomic_add, use_fp32_reduce); + prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce, max_shared_mem); // clang-format on } @@ -800,7 +865,9 @@ torch::Tensor moe_wna16_marlin_gemm( torch::Tensor& a, std::optional const& c_or_none, torch::Tensor& b_q_weight, + std::optional const& b_bias_or_none, torch::Tensor& b_scales, + std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, std::optional const& g_idx_or_none, std::optional const& perm_or_none, @@ -915,7 +982,6 @@ torch::Tensor moe_wna16_marlin_gemm( num_groups = b_scales.size(1); torch::Tensor g_idx, perm, a_tmp; - ; if (g_idx_or_none.has_value() && perm_or_none.has_value()) { g_idx = g_idx_or_none.value(); perm = perm_or_none.value(); @@ -962,6 +1028,29 @@ torch::Tensor moe_wna16_marlin_gemm( } } + torch::Tensor global_scale; + if (global_scale_or_none.has_value()) { + global_scale = global_scale_or_none.value(); + TORCH_CHECK(b_q_type == sglang::kFE2M1f && group_size == 16, "global_scale can only be used for nvfp4 format."); + } else { + global_scale = torch::empty({0}, options); + TORCH_CHECK( + !(b_q_type == sglang::kFE2M1f && group_size == 16), + "the global_scale parameter must be passed for nvfp4 format."); + } + + bool has_bias = b_bias_or_none.has_value(); + torch::Tensor b_bias; + if (has_bias) { + b_bias = b_bias_or_none.value(); + TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); + TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); + TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(0) != size_n"); + TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1"); + } else { + b_bias = torch::empty({0}, options); + } + torch::Tensor b_zeros; if (b_zeros_or_none.has_value()) { b_zeros = b_zeros_or_none.value(); @@ -971,13 +1060,18 @@ torch::Tensor moe_wna16_marlin_gemm( b_zeros = torch::empty({0}, options); } bool has_zp = b_zeros.size(-1) > 0; - if (has_zp) { - TORCH_CHECK(b_q_type == sglang::kU4, "b_q_type must be u4 when has_zp = True. Got = ", b_q_type.str()); + TORCH_CHECK( + b_q_type == sglang::kU4 || b_q_type == sglang::kU8, + "b_q_type must be u4 or u8 when has_zp = True. Got = ", + b_q_type.str()); } else { TORCH_CHECK( - b_q_type == sglang::kU4B8 || b_q_type == sglang::kU8B128, - "b_q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ", + b_q_type == sglang::kU4B8 || b_q_type == sglang::kU8B128 || b_q_type == sglang::kFE4M3fn || + b_q_type == sglang::kFE2M1f, + "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or " + "float4_e2m1f when " + "has_zp = False. Got = ", b_q_type.str()); } @@ -1028,12 +1122,26 @@ torch::Tensor moe_wna16_marlin_gemm( int dev = a.get_device(); if (a.scalar_type() == at::ScalarType::Half) { + void* scales_ptr; + if (b_q_type == sglang::kFE2M1f) { + if (group_size == 16) + scales_ptr = b_scales.data_ptr(); + else if (group_size == 32) + scales_ptr = b_scales.data_ptr(); + else + TORCH_CHECK(false, "float4_e2m1f only supports group_size == 16 (NVFP4) ", "and group_size == 32 (MXFP4)"); + } else { + scales_ptr = b_scales.data_ptr(); + } + MARLIN_NAMESPACE_NAME::marlin_mm( a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), - b_scales.data_ptr(), + b_bias.data_ptr(), + scales_ptr, + global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), @@ -1051,6 +1159,7 @@ torch::Tensor moe_wna16_marlin_gemm( size_k, workspace.data_ptr(), b_q_type, + has_bias, has_act_order, is_k_full, has_zp, @@ -1065,12 +1174,26 @@ torch::Tensor moe_wna16_marlin_gemm( use_fp32_reduce, is_zp_float); } else if (a.scalar_type() == at::ScalarType::BFloat16) { + void* scales_ptr; + if (b_q_type == sglang::kFE2M1f) { + if (group_size == 16) + scales_ptr = b_scales.data_ptr(); + else if (group_size == 32) + scales_ptr = b_scales.data_ptr(); + else + TORCH_CHECK(false, "float4_e2m1f only supports group_size == 16 (NVFP4) ", "and group_size == 32 (MXFP4)"); + } else { + scales_ptr = b_scales.data_ptr(); + } + MARLIN_NAMESPACE_NAME::marlin_mm( a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), - b_scales.data_ptr(), + b_bias.data_ptr(), + scales_ptr, + global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(), @@ -1088,6 +1211,7 @@ torch::Tensor moe_wna16_marlin_gemm( size_k, workspace.data_ptr(), b_q_type, + has_bias, has_act_order, is_k_full, has_zp, @@ -1109,3 +1233,5 @@ torch::Tensor moe_wna16_marlin_gemm( } #endif + +// Registration is done in common_extension.cc for v2 version diff --git a/sgl-kernel/include/scalar_type.hpp b/sgl-kernel/include/scalar_type.hpp index 9161a38b7..3ed1467bf 100644 --- a/sgl-kernel/include/scalar_type.hpp +++ b/sgl-kernel/include/scalar_type.hpp @@ -301,6 +301,7 @@ static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); +static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 312756c1c..552519fa4 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -459,7 +459,9 @@ torch::Tensor moe_wna16_marlin_gemm( torch::Tensor& a, std::optional const& c_or_none, torch::Tensor& b_q_weight, + std::optional const& b_bias_or_none, torch::Tensor& b_scales, + std::optional const& global_scale_or_none, std::optional const& b_zeros_or_none, std::optional const& g_idx_or_none, std::optional const& perm_or_none, diff --git a/sgl-kernel/python/sgl_kernel/fused_moe.py b/sgl-kernel/python/sgl_kernel/fused_moe.py index 8a7c3dcdf..15f3a2beb 100644 --- a/sgl-kernel/python/sgl_kernel/fused_moe.py +++ b/sgl-kernel/python/sgl_kernel/fused_moe.py @@ -7,7 +7,9 @@ def moe_wna16_marlin_gemm( a: torch.Tensor, c_or_none: Optional[torch.Tensor], b_q_weight: torch.Tensor, + b_bias_or_none: Optional[torch.Tensor], b_scales: torch.Tensor, + global_scale_or_none: Optional[torch.Tensor], b_zeros_or_none: Optional[torch.Tensor], g_idx_or_none: Optional[torch.Tensor], perm_or_none: Optional[torch.Tensor], @@ -33,7 +35,9 @@ def moe_wna16_marlin_gemm( a, c_or_none, b_q_weight, + b_bias_or_none, b_scales, + global_scale_or_none, b_zeros_or_none, g_idx_or_none, perm_or_none,