diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 0c6f4f234dec..fcea8829442b 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -61,6 +61,9 @@ list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") +# Marlin is built as its own Hopper-only OBJECT library below. +list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") +list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -79,7 +82,26 @@ if(FAST_BUILD) STATUS "FAST_BUILD enabled for kernels: using -O1 for CUDA compilation") endif() -add_library(kernels_src STATIC ${SRC_CPP} ${SRC_CU}) +# Marlin NVFP4: Hopper-only OBJECT library. Pinned to sm_90 so the global +# CMAKE_CUDA_ARCHITECTURES doesn't propagate. +file(GLOB_RECURSE MARLIN_SRC "marlin/*.cu" "marlin/*.cpp") +if(MARLIN_SRC) + add_library(marlin_src OBJECT ${MARLIN_SRC}) + set_property(TARGET marlin_src PROPERTY POSITION_INDEPENDENT_CODE ON) + set_property(TARGET marlin_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) + target_include_directories( + marlin_src + PRIVATE + $ + ) + target_link_libraries(marlin_src PRIVATE trtllm_gen_fmha_interface) + set_cuda_architectures(marlin_src 90) +endif() + +add_library( + kernels_src STATIC + ${SRC_CPP} ${SRC_CU} + $<$:$>) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_include_directories( diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh new file mode 100644 index 000000000000..63864080b0f2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh @@ -0,0 +1,397 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Internal device-side header for the Marlin NVFP4 kernels. + +#pragma once + +#ifndef _marlin_cuh +#define _marlin_cuh + +#include +#include +#include +#include +#include + +#include + +#ifndef MARLIN_NAMESPACE_NAME +#define MARLIN_NAMESPACE_NAME marlin +#endif + +namespace MARLIN_NAMESPACE_NAME +{ + +static constexpr int default_threads = 256; +static constexpr int pipe_stages = 4; + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; +static constexpr int max_thread_n = 256; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +static constexpr int repack_stages = 8; +static constexpr int repack_threads = 256; + +static constexpr int tile_k_size = tile_size; +static constexpr int tile_n_size = tile_k_size * 4; + +template +struct Vec +{ + T elems[n]; + + __device__ T& operator[](int i) + { + return elems[i]; + } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) +{ + return (a + b - 1) / b; +} + +// cp.async wrappers (SM 7.x fallback / SM 8.x+ inline asm). + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + +__device__ inline void cp_async1_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + if (pred) + { + reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async2_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + if (pred) + { + reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + if (pred) + { + reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + if (pred) + { + reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4(void* smem_ptr, void const* glob_ptr) +{ + reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; +} + +__device__ inline void cp_async_fence() {} + +template +__device__ inline void cp_async_wait() +{ +} + +#else + +__device__ inline void cp_async1_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + int const BYTES = 4; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int) pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async2_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + int const BYTES = 8; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int) pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4_ca_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + int const BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int) pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4_pred(void* smem_ptr, void const* glob_ptr, bool pred = true) +{ + int const BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int) pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void* smem_ptr, void const* glob_ptr) +{ + int const BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() +{ + asm volatile("cp.async.commit_group;\n" ::); +} + +template +__device__ inline void cp_async_wait() +{ + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + +#endif + +// MarlinType traits + fragment aliases. +// MMA m16n8k16 fragment layouts: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type +template +struct MarlinType +{ +}; + +template <> +struct MarlinType +{ + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + using scalar_t4 = nv_bfloat162; + using scalar_32bit_t = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; + using FragZP = Vec; + +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) + { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) + { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, const nv_bfloat16 x2) + { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(float const x) + { + return __float2bfloat16(x); + } + + static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) + { + return __bfloat1622float2(x); + } +#endif +}; + +// Fast FP4 E2M1 -> BF16 and FP8 E4M3 -> BF16 dequantization. +// FP4->BF16 places the 3 FP4 bits into BF16's exponent/mantissa via bitwise +// ops; a subsequent multiply applies the exponent-bias correction (or +// ``skip_flop=true`` defers it to fuse with a scale multiply downstream). + +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750 + +// Lookup-table based 3-input logical operation; the compiler does not always +// recognize the pattern automatically. +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; +} + +// FP4 E2M1 -> BF16 +// +// skip_flop=true: just place bits, caller multiplies exponent bias later. +template +__device__ inline void dequant_fp4(int q, nv_bfloat162* frag_b) +{ + // Constants for FP4 (E2M1) -> BF16 (E8M7) + constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; + constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; + constexpr int MASK = 0x70007000; + + // Extract and shift FP4 values to BF16 format + int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 4; + int Out2 = (q & 0x80008000) | ((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); + + if constexpr (!skip_flop) + { + // Apply exponent bias correction + constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + + frag_b[1] = __hmul2(frag_b[1], bias_reg); + frag_b[0] = __hmul2(frag_b[0], bias_reg); + } +} + +// FP8 E4M3 scale -> BF16 +__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); +} + +#endif // __CUDA_ARCH__ >= 750 + +// m16n8k16 tensor-core MMA: BF16 inputs, FP32 accumulation. +template +__device__ inline void mma(const typename MarlinType::FragA& a_frag, + const typename MarlinType::FragB& frag_b, typename MarlinType::FragC& frag_c) +{ + uint32_t const* a = reinterpret_cast(&a_frag); + uint32_t const* b = reinterpret_cast(&frag_b); + + static_assert(std::is_same::value, "Only BF16 is supported for Marlin NVFP4 MMA"); + + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); +} + +// Transposed variant: column-major B weight loading. +template +__device__ inline void mma_trans(const typename MarlinType::FragA& a_frag, + const typename MarlinType::FragB& frag_b, const typename MarlinType::FragB& frag_b2, + typename MarlinType::FragC& frag_c) +{ + uint32_t const* a = reinterpret_cast(&a_frag); + uint32_t const* b = reinterpret_cast(&frag_b); + uint32_t const* b2 = reinterpret_cast(&frag_b2); + + static_assert(std::is_same::value, "Only BF16 is supported for Marlin NVFP4 MMA"); + + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); +} + +} // namespace MARLIN_NAMESPACE_NAME + +// Single-expert kernel forward decl. Opt in with +// ``#define MARLIN_DECLARE_SINGLE_EXPERT_KERNEL`` before including. The MoE +// TU does NOT define it (the MoE kernel has a different parameter list, in +// marlin_nvfp4_moe_template.h). +#ifdef MARLIN_DECLARE_SINGLE_EXPERT_KERNEL + +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ global_scale_ptr, \ + const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, \ + int prob_k, int lda, int *locks, bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem + +namespace MARLIN_NAMESPACE_NAME +{ + +// clang-format off +// NOTE: keep this template parameter list out of clang-format. East-const +// (QualifierAlignment: Right) miscompiles the *last* NTTP as `int X const`, +// which is invalid syntax. Non-type template parameters are implicitly +// const anyway, so we omit it on the last entry. +template +__global__ void Marlin(MARLIN_KERNEL_PARAMS); +// clang-format on + +} // namespace MARLIN_NAMESPACE_NAME + +#endif // MARLIN_DECLARE_SINGLE_EXPERT_KERNEL + +#endif // _marlin_cuh diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h new file mode 100644 index 000000000000..ad7d55471514 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h @@ -0,0 +1,73 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Host-side header for the Marlin NVFP4 W4A16 kernels. + +#pragma once + +#include +#include + +namespace marlin_nvfp4 +{ + +void dequantFp4Activations( + void const* act_fp4, void const* act_sf, float const* alpha, void* act_bf16, int m, int k, cudaStream_t stream); + +void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, void* C_tmp, void const* weight_sf, + void const* global_scale_bf16, int m, int n, int k, int* workspace, int num_groups, int group_size, + bool use_fp32_reduce, cudaStream_t stream); + +void marlinNvfp4MoeGemmDispatcher(void const* A, void const* B, void* C, void* C_tmp, void const* b_scales, + void const* global_scale, void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, + void const* topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, + void* workspace, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, cudaDataType_t outType, + cudaStream_t stream); + +void gptq_marlin_repack_dispatch(uint32_t const* b_q_weight_ptr, uint32_t const* perm_ptr, uint32_t* out_ptr, + int size_k, int size_n, int num_bits, bool has_perm, bool is_a_8bit, cudaStream_t stream); + +} // namespace marlin_nvfp4 + +namespace marlin_nvfp4_dispatch +{ + +struct thread_config_t +{ + int thread_k; + int thread_n; + int num_threads; +}; + +struct exec_config_t +{ + int blocks_per_sm; + thread_config_t tb_cfg; +}; + +extern thread_config_t const kSmallBatchConfigs[]; +extern thread_config_t const kLargeBatchConfigs[]; + +extern int const kSmallBatchConfigCount; +extern int const kLargeBatchConfigCount; + +int get_scales_cache_size( + thread_config_t const& th_config, int prob_n, int prob_k, int num_bits, int group_size, int stages); + +bool is_config_feasible(thread_config_t const& cfg, int prob_n, int prob_k); + +} // namespace marlin_nvfp4_dispatch diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp new file mode 100644 index 000000000000..9b3dac04ed25 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_dispatch_utils.cpp @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "marlin_nvfp4.h" + +namespace marlin_nvfp4_dispatch +{ + +thread_config_t const kSmallBatchConfigs[] = {{128, 128, 256}, {64, 128, 128}, {128, 64, 128}}; +thread_config_t const kLargeBatchConfigs[] = {{64, 256, 256}, {64, 128, 128}, {128, 64, 128}}; + +int const kSmallBatchConfigCount = sizeof(kSmallBatchConfigs) / sizeof(thread_config_t); +int const kLargeBatchConfigCount = sizeof(kLargeBatchConfigs) / sizeof(thread_config_t); + +int get_scales_cache_size( + thread_config_t const& th_config, int prob_n, int prob_k, int num_bits, int group_size, int stages) +{ + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + int tb_groups; + if (group_size == -1) + tb_groups = 1; + else if (group_size == 0) + tb_groups = (tb_k + 31) / 32; // div_ceil(tb_k, 32) + else + tb_groups = (tb_k + group_size - 1) / group_size; // div_ceil(tb_k, group_size) + return tb_groups * tb_n * 2 * stages; +} + +bool is_config_feasible(thread_config_t const& cfg, int prob_n, int prob_k) +{ + if (cfg.thread_k == -1 || cfg.thread_n == -1 || cfg.num_threads == -1) + return false; + if (prob_k % cfg.thread_k != 0 || prob_n % cfg.thread_n != 0) + return false; + // min_thread_n = 64, min_thread_k = 64 (from marlin.cuh constants) + if (cfg.thread_n < 64 || cfg.thread_k < 64) + return false; + if (cfg.num_threads < 128) + return false; + return true; +} + +} // namespace marlin_nvfp4_dispatch diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu new file mode 100644 index 000000000000..f9cbbe3d64ab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu @@ -0,0 +1,340 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MARLIN_NAMESPACE_NAME +#define MARLIN_NAMESPACE_NAME marlin +#endif + +#define MARLIN_DECLARE_SINGLE_EXPERT_KERNEL +#include "marlin_nvfp4.h" +#include "marlin_nvfp4_template.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include +#include + +namespace marlin +{ + +using namespace marlin_nvfp4_dispatch; + +__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; + +using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); + +// Single-expert shared-memory size (no block-meta overhead). +int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, + int group_size, int stages) +{ + int pack_factor = 32 / num_bits; + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + int tb_m = thread_m_blocks * 16; + int sh_a_size = stages * (tb_m * tb_k) * 2; + int sh_b_size = 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 = std::max(std::max(sh_b_size, sh_red_size), tmp_size); + int sh_s_size = get_scales_cache_size(th_config, prob_n, prob_k, num_bits, group_size, stages); + return tmp_size + sh_a_size + sh_s_size; +} + +bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, + int group_size, int stages, int max_shared_mem) +{ + if (!is_config_feasible(th_config, prob_n, prob_k)) + return false; + return get_kernel_cache_size(th_config, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages) + <= max_shared_mem; +} + +MarlinFuncPtr get_marlin_kernel(int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, + int group_blocks, int threads, int stages) +{ +#define MARLIN_KERNEL_MATCH(T, M, N, K, M8) \ + (threads == (T) && thread_m_blocks == (M) && thread_n_blocks == (N) && thread_k_blocks == (K) \ + && m_block_size_8 == (M8) && stages == 4 && group_blocks == 1) +#define MARLIN_KERNEL_IF(T, M, N, K, M8) \ + if (MARLIN_KERNEL_MATCH(T, M, N, K, M8)) \ + return Marlin; + + MARLIN_KERNEL_IF(256, 1, 8, 8, true) + MARLIN_KERNEL_IF(128, 1, 8, 4, true) + MARLIN_KERNEL_IF(128, 1, 4, 8, true) + MARLIN_KERNEL_IF(256, 1, 8, 8, false) + MARLIN_KERNEL_IF(128, 1, 8, 4, false) + MARLIN_KERNEL_IF(128, 1, 4, 8, false) + MARLIN_KERNEL_IF(256, 2, 16, 4, false) + MARLIN_KERNEL_IF(128, 2, 8, 4, false) + MARLIN_KERNEL_IF(128, 2, 4, 8, false) + MARLIN_KERNEL_IF(256, 3, 16, 4, false) + MARLIN_KERNEL_IF(128, 3, 8, 4, false) + MARLIN_KERNEL_IF(128, 3, 4, 8, false) + MARLIN_KERNEL_IF(256, 4, 16, 4, false) + MARLIN_KERNEL_IF(128, 4, 8, 4, false) + MARLIN_KERNEL_IF(128, 4, 4, 8, false) + +#undef MARLIN_KERNEL_MATCH +#undef MARLIN_KERNEL_IF + return MarlinDefault; +} + +exec_config_t determine_exec_config(int prob_m, int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, + int num_bits, int group_size, int stages, int max_shared_mem, int sms) +{ + exec_config_t exec_cfg{1, {-1, -1, -1}}; + thread_config_t const* cfgs = thread_m_blocks > 1 ? kLargeBatchConfigs : kSmallBatchConfigs; + int cfg_count = thread_m_blocks > 1 ? kLargeBatchConfigCount : kSmallBatchConfigCount; + + for (int i = 0; i < cfg_count; i++) + { + thread_config_t th = cfgs[i]; + if (!is_valid_config(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, max_shared_mem - 512)) + continue; + int group_blocks = group_size == -1 ? -1 : group_size / 16; + auto kernel = get_marlin_kernel( + thread_m_blocks, th.thread_n / 16, th.thread_k / 16, m_block_size_8, group_blocks, th.num_threads, stages); + if (kernel == MarlinDefault) + continue; + return {1, th}; + } + return exec_cfg; +} + +void marlin_mm_nvfp4(void const* A, void const* B, void* C, void* C_tmp, void const* b_s, void const* g_s, int prob_m, + int prob_n, int prob_k, int* locks, int num_groups, int group_size, bool use_fp32_reduce, int dev, + cudaStream_t stream) +{ + constexpr int num_bits = 4; + + int group_blocks = group_size == -1 ? -1 : group_size / 16; + + int4 const* A_ptr = (int4 const*) A; + int4 const* B_ptr = (int4 const*) B; + int4* C_ptr = (int4*) C; + int4* C_tmp_ptr = (int4*) C_tmp; + int4 const* b_s_ptr = (int4 const*) b_s; + uint16_t const* g_s_ptr = (uint16_t const*) g_s; + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + + int stages = 4; + int sms = -1; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + + int max_par_val = 16; + if (prob_n <= 4096) + max_par_val = 16 * 8; + int max_shared_mem_new = max_shared_mem; + int rest_m = prob_m; + int max_thread_m_blocks = 4; + int lda = prob_k; + + while (rest_m) + { + int par_count = std::min(rest_m / (max_thread_m_blocks * 16), max_par_val); + int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; + + int thread_m_blocks = std::min(div_ceil(prob_m_split, 16), max_thread_m_blocks); + bool m_block_size_8 = prob_m_split <= 8; + + exec_config_t exec_cfg = determine_exec_config(prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8, + num_bits, group_size, stages, max_shared_mem, sms); + thread_config_t thread_tfg = exec_cfg.tb_cfg; + + if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) + { + max_thread_m_blocks--; + continue; + } + + if (thread_tfg.thread_k == -1) + { + break; + } + + // Small wave optimization + if (thread_tfg.thread_n != -1) + { + if (prob_n / thread_tfg.thread_n * div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= sms) + { + if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, + max_shared_mem_new)) + { + thread_tfg = {128, 64, 128}; + exec_cfg = {1, thread_tfg}; + } + } + } + + int num_threads = thread_tfg.num_threads; + int thread_k = thread_tfg.thread_k; + int thread_n = thread_tfg.thread_n; + int blocks = sms * exec_cfg.blocks_per_sm; + if (exec_cfg.blocks_per_sm > 1) + max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + auto kernel = get_marlin_kernel( + thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); + + if (kernel == MarlinDefault) + { + break; + } + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new); + + // clang-format off + kernel<<>>( + A_ptr, B_ptr, C_ptr, C_tmp_ptr, + nullptr, // b_bias + nullptr, // a_scales + b_s_ptr, g_s_ptr, + nullptr, // zp + nullptr, // g_idx + num_groups, prob_m_split, prob_n, prob_k, lda, locks, + false, // has_bias + false, // use_atomic_add + use_fp32_reduce, max_shared_mem_new); + // clang-format on + + A_ptr += prob_m_split * (lda / 8); + C_ptr += prob_m_split * (prob_n / 8); + rest_m -= prob_m_split; + } +} + +// Explicit template instantiations for BF16 + NVFP4 Marlin kernels. +// clang-format off +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +// clang-format on + +// FP4 E2M1 -> BF16 activation dequant: out[i] = fp4_to_bf16(act[i]) * +// block_scale[i/16] * global_scale. Block scales are FP8 E4M3 (swizzled). +__global__ void dequant_fp4_act_kernel(uint8_t const* __restrict__ act_fp4, // [M, K/2] packed FP4 + uint8_t const* __restrict__ act_sf, // FP8 E4M3 block scales (swizzled) + float const* __restrict__ alpha, // global scale + nv_bfloat16* __restrict__ out, // [M, K] BF16 output + int M, int K) +{ + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total_pairs = M * (K / 2); + if (idx >= total_pairs) + return; + + float global_s = *alpha; + int row = idx / (K / 2); + int col_pair = idx % (K / 2); + + uint8_t packed = act_fp4[idx]; + + // Unpack two FP4 E2M1 values (low nibble first) + auto fp4_to_float = [](uint8_t nibble) -> float + { + // FP4 E2M1: 1 sign + 2 exponent + 1 mantissa + uint8_t sign = (nibble >> 3) & 1; + uint8_t exp = (nibble >> 1) & 0x3; + uint8_t mant = nibble & 1; + float val; + if (exp == 0) + { + // subnormal: (-1)^s * 0.mantissa * 2^(1-bias) = (-1)^s * mant * 0.5 + val = mant * 0.5f; + } + else + { + // normal: (-1)^s * 1.mantissa * 2^(exp-bias), bias=1 + val = (1.0f + mant * 0.5f) * (float) (1 << (exp - 1)); + } + return sign ? -val : val; + }; + + float v0 = fp4_to_float(packed & 0x0F); + float v1 = fp4_to_float((packed >> 4) & 0x0F); + + // Block scale: one FP8 E4M3 per 16 FP4 elements = per 8 bytes + // The scale layout is swizzled 128x4 — for now use linear indexing + // as a reasonable approximation. TODO: handle swizzled layout properly. + int elem0 = col_pair * 2; + int scale_idx = row * (K / 16) + elem0 / 16; + uint8_t sf_byte = act_sf[scale_idx]; + // FP8 E4M3 -> float: reinterpret as __nv_fp8_e4m3 + __nv_fp8_e4m3 sf_fp8 = *reinterpret_cast<__nv_fp8_e4m3 const*>(&sf_byte); + float sf = float(sf_fp8); + + float scale = sf * global_s; + int out_idx = row * K + col_pair * 2; + out[out_idx] = __float2bfloat16(v0 * scale); + out[out_idx + 1] = __float2bfloat16(v1 * scale); +} + +} // namespace marlin + +namespace marlin_nvfp4 +{ + +void dequantFp4Activations( + void const* act_fp4, void const* act_sf, float const* alpha, void* act_bf16, int m, int k, cudaStream_t stream) +{ + + int total_pairs = m * (k / 2); + int threads = 256; + int blocks = (total_pairs + threads - 1) / threads; + ::marlin::dequant_fp4_act_kernel<<>>( + (uint8_t const*) act_fp4, (uint8_t const*) act_sf, alpha, (nv_bfloat16*) act_bf16, m, k); +} + +void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, void* C_tmp, void const* weight_sf, + void const* global_scale_bf16, int m, int n, int k, int* workspace, int num_groups, int group_size, + bool use_fp32_reduce, cudaStream_t stream) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TLLM_CHECK_WITH_INFO( + sm >= 90 && sm < 100, "Marlin NVFP4 GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); + + int dev; + cudaGetDevice(&dev); + + ::marlin::marlin_mm_nvfp4(act_bf16, weight, output, C_tmp, weight_sf, global_scale_bf16, m, n, k, workspace, + num_groups, group_size, use_fp32_reduce, dev, stream); +} + +} // namespace marlin_nvfp4 diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu new file mode 100644 index 000000000000..35f28b309345 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu @@ -0,0 +1,269 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MARLIN_NAMESPACE_NAME +#define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#endif + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" + +#include "marlin_nvfp4.h" +#include "marlin_nvfp4_moe_template.h" + +#include +#include + +namespace marlin_moe_wna16 +{ + +using namespace marlin_nvfp4_dispatch; + +__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; + +using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); + +// MoE shared-memory size includes block-meta overhead for sorted_token_ids. +int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, + int group_size, int stages) +{ + int pack_factor = 32 / num_bits; + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + int tb_m = thread_m_blocks * 16; + int sh_block_meta_size = tb_m * 16; + int sh_a_size = stages * (tb_m * tb_k) * 2; + int sh_b_size = 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 = std::max(std::max(sh_b_size, sh_red_size), tmp_size); + int sh_s_size = get_scales_cache_size(th_config, prob_n, prob_k, num_bits, group_size, stages); + return tmp_size + sh_a_size + sh_s_size + sh_block_meta_size; +} + +bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks, int prob_n, int prob_k, int num_bits, + int group_size, int stages, int max_shared_mem) +{ + if (!is_config_feasible(th_config, prob_n, prob_k)) + return false; + return get_kernel_cache_size(th_config, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages) + <= max_shared_mem; +} + +MarlinFuncPtr get_marlin_kernel(int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, bool m_block_size_8, + int group_blocks, int threads, int stages) +{ + +#define MARLIN_KERNEL_MATCH(T, M, N, K, M8) \ + (threads == (T) && thread_m_blocks == (M) && thread_n_blocks == (N) && thread_k_blocks == (K) \ + && m_block_size_8 == (M8) && stages == 4 && group_blocks == 1) +#define MARLIN_KERNEL_IF(T, M, N, K, M8) \ + if (MARLIN_KERNEL_MATCH(T, M, N, K, M8)) \ + return Marlin; + + MARLIN_KERNEL_IF(256, 1, 8, 8, true) + MARLIN_KERNEL_IF(128, 1, 8, 4, true) + MARLIN_KERNEL_IF(128, 1, 4, 8, true) + MARLIN_KERNEL_IF(256, 1, 8, 8, false) + MARLIN_KERNEL_IF(128, 1, 8, 4, false) + MARLIN_KERNEL_IF(128, 1, 4, 8, false) + MARLIN_KERNEL_IF(256, 2, 16, 4, false) + MARLIN_KERNEL_IF(128, 2, 8, 4, false) + MARLIN_KERNEL_IF(128, 2, 4, 8, false) + MARLIN_KERNEL_IF(256, 3, 16, 4, false) + MARLIN_KERNEL_IF(128, 3, 8, 4, false) + MARLIN_KERNEL_IF(128, 3, 4, 8, false) + MARLIN_KERNEL_IF(256, 4, 16, 4, false) + MARLIN_KERNEL_IF(128, 4, 8, 4, false) + MARLIN_KERNEL_IF(128, 4, 4, 8, false) + +#undef MARLIN_KERNEL_MATCH +#undef MARLIN_KERNEL_IF + return MarlinDefault; +} + +// MoE config selection with occupancy-based multi-block logic. +exec_config_t determine_exec_config(int prob_m, int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks, + bool m_block_size_8, int num_bits, int group_size, int stages, int max_shared_mem, int sms) +{ + exec_config_t exec_cfg{1, {-1, -1, -1}}; + thread_config_t const* cfgs = thread_m_blocks > 1 ? kLargeBatchConfigs : kSmallBatchConfigs; + int cfg_count = thread_m_blocks > 1 ? kLargeBatchConfigCount : kSmallBatchConfigCount; + + int count = 0; + constexpr int device_max_reg_size = 255 * 1024; + int group_blocks = group_size == -1 ? -1 : (group_size / 16); + + for (int i = 0; i < cfg_count; i++) + { + thread_config_t th = cfgs[i]; + if (!is_valid_config(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages, max_shared_mem - 512)) + continue; + + int cache_size = get_kernel_cache_size(th, thread_m_blocks, prob_n, prob_k, num_bits, group_size, stages); + + auto kernel = get_marlin_kernel( + thread_m_blocks, th.thread_n / 16, th.thread_k / 16, m_block_size_8, group_blocks, th.num_threads, stages); + if (kernel == MarlinDefault) + continue; + + cudaFuncAttributes attr; + cudaFuncGetAttributes(&attr, kernel); + int reg_size = std::max(attr.numRegs, 1) * th.num_threads * 4; + int allow_count = std::min(device_max_reg_size / reg_size, max_shared_mem / (cache_size + 1536)); + if (thread_m_blocks == 1) + allow_count = std::max(std::min(allow_count, 4), 1); + else + allow_count = std::max(std::min(allow_count, 2), 1); + + if (prob_n / th.thread_n * prob_m * top_k * 4 < sms * allow_count) + allow_count = std::max(prob_n / th.thread_n * prob_m * top_k * 4 / sms, 1); + + if (allow_count > count) + { + count = allow_count; + exec_cfg = {count, th}; + }; + } + return exec_cfg; +} + +void marlin_mm_moe_nvfp4(void const* A, void const* B, void* C, void* C_tmp, void const* b_s, void const* g_s, + void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, void const* topk_weights, + int moe_block_size, int num_experts, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, + int* locks, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, int dev, cudaStream_t stream) +{ + + constexpr int num_bits = 4; + + int thread_m_blocks = div_ceil(moe_block_size, 16); + bool m_block_size_8 = moe_block_size == 8; + int group_blocks = group_size == -1 ? -1 : group_size / 16; + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + int stages = 4; + int sms = -1; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + + exec_config_t exec_cfg = determine_exec_config(prob_m, prob_n, prob_k, num_experts, top_k, thread_m_blocks, + m_block_size_8, num_bits, group_size, stages, max_shared_mem, sms); + thread_config_t thread_tfg = exec_cfg.tb_cfg; + + if (thread_tfg.thread_k == -1) + return; + + int num_threads = thread_tfg.num_threads; + int thread_k = thread_tfg.thread_k; + int thread_n = thread_tfg.thread_n; + int blocks = sms * exec_cfg.blocks_per_sm; + if (exec_cfg.blocks_per_sm > 1) + max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + auto kernel = get_marlin_kernel( + thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); + + if (kernel == MarlinDefault) + { + TLLM_LOG_ERROR( + "xuantengh debug error: kernel is MarlinDefault, cannot find corresponding instantiated kernel for threads " + "= %d, " + "thread_n_blocks = %d, thread_k_blocks = %d, m_block_size_8 = %d, group_blocks = %d, num_threads = %d, " + "stages = %d", + num_threads, thread_n_blocks, thread_k_blocks, m_block_size_8, group_blocks, num_threads, stages); + return; + } + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); + + int4 const* A_ptr = (int4 const*) A; + int4 const* B_ptr = (int4 const*) B; + int4* C_ptr = (int4*) C; + int4* C_tmp_ptr = (int4*) C_tmp; + int4 const* b_s_ptr = (int4 const*) b_s; + uint16_t const* g_s_ptr = (uint16_t const*) g_s; + int32_t const* sorted_token_ids_ptr = (int32_t const*) sorted_token_ids; + int32_t const* expert_ids_ptr = (int32_t const*) expert_ids; + int32_t const* num_tokens_past_padded_ptr = (int32_t const*) num_tokens_past_padded; + float const* topk_weights_ptr = (float const*) topk_weights; + + // clang-format off + kernel<<>>( + A_ptr, B_ptr, C_ptr, C_tmp_ptr, + nullptr, // b_bias + nullptr, // a_scales + b_s_ptr, g_s_ptr, + nullptr, // zp + nullptr, // g_idx + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, + topk_weights_ptr, top_k, mul_topk_weights, num_groups, + prob_m, prob_n, prob_k, locks, + false, // has_bias + use_atomic_add, use_fp32_reduce); + // clang-format on +} + +// Explicit template instantiations for BF16 + NVFP4 MoE Marlin kernels. +// clang-format off +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); +// clang-format on + +} // namespace marlin_moe_wna16 + +namespace marlin_nvfp4 +{ + +void marlinNvfp4MoeGemmDispatcher(void const* A, void const* B, void* C, void* C_tmp, void const* b_scales, + void const* global_scale, void const* sorted_token_ids, void const* expert_ids, void const* num_tokens_past_padded, + void const* topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, int prob_m, int prob_n, int prob_k, + void* workspace, int num_groups, int group_size, bool use_fp32_reduce, bool use_atomic_add, cudaDataType_t outType, + cudaStream_t stream) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TLLM_CHECK_WITH_INFO( + sm >= 90 && sm < 100, "Marlin NVFP4 MoE GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); + + int dev; + cudaGetDevice(&dev); + + int num_experts = 1; // Not used in kernel dispatch, only in config selection + + ::marlin_moe_wna16::marlin_mm_moe_nvfp4(A, B, C, C_tmp, b_scales, global_scale, sorted_token_ids, expert_ids, + num_tokens_past_padded, topk_weights, moe_block_size, num_experts, top_k, mul_topk_weights, prob_m, prob_n, + prob_k, (int*) workspace, num_groups, group_size, use_fp32_reduce, use_atomic_add, dev, stream); +} + +} // namespace marlin_nvfp4 diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h new file mode 100644 index 000000000000..9159e0feea61 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h @@ -0,0 +1,2175 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +#ifndef MARLIN_NAMESPACE_NAME +#define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#endif + +#include "marlin.cuh" + +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ global_scale_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, int num_groups, int prob_m, int prob_n, int prob_k, int *locks, bool has_bias, \ + bool use_atomic_add, bool use_fp32_reduce + +namespace MARLIN_NAMESPACE_NAME +{ + +template +__global__ void Marlin(MARLIN_KERNEL_PARAMS); + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +// Empty kernel stub for non-Hopper device passes; see marlin.cuh. +#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) + +template +__global__ void Marlin(MARLIN_KERNEL_PARAMS) +{ +} + +#else + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm(typename MarlinType::FragA& frag_a, void const* smem_ptr) +{ + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (count == 4) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } + else if constexpr (count == 2) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); + } + else if constexpr (count == 1) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); + } + else + { + static_assert(count == 1 || count == 2 || count == 4, "invalid count"); + } +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale( + typename MarlinType::FragB& frag_b, typename MarlinType::FragS& frag_s, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 s = MarlinType::num2num2(reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +template +__device__ inline void scale_and_sub(typename MarlinType::FragB& frag_b, scalar_t s, scalar_t zp) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 s2 = MarlinType::num2num2(s); + scalar_t2 zp2 = MarlinType::num2num2(zp); + frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); + frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); +} + +template +__device__ inline void sub_zp( + typename MarlinType::FragB& frag_b, typename MarlinType::scalar_t2& frag_zp, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 zp = MarlinType::num2num2(reinterpret_cast(&frag_zp)[i]); + frag_b[0] = __hsub2(frag_b[0], zp); + frag_b[1] = __hsub2(frag_b[1], zp); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4(typename MarlinType::FragB& frag_b, + typename MarlinType::FragS& frag_s_1, typename MarlinType::FragS& frag_s_2, + typename MarlinType::FragS& frag_s_3, typename MarlinType::FragS& frag_s_4, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float(float* c, typename MarlinType::FragS& s) +{ + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], MarlinType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], MarlinType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) +{ + if (threadIdx.x == 0) + { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) +{ + __syncthreads(); + if (threadIdx.x == 0) + { + if (reset) + { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); + } +} + +// Wait until value of lock to be negative, and then add 1 +__device__ inline void wait_negative_and_add(int* lock) +{ + if (threadIdx.x == 0) + { + int state = 0; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + while (state >= 0); + atomicAdd(lock, 1); + } + __syncthreads(); +} + +template +__global__ void Marlin(MARLIN_KERNEL_PARAMS) +{ + // 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 + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + + // NVFP4 kernel: BF16 activations only, no FP8/Turing arch guards needed. + static_assert(std::is_same::value, "NVFP4 kernel only supports BF16 compute type"); + + int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; + constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); + + constexpr bool use_fp16_accum = false; + using Dtype = MarlinType; + + using scalar_t2 = typename MarlinType::scalar_t2; + using scalar_32bit_t = typename MarlinType::scalar_32bit_t; + + using FragA = typename MarlinType::FragA; + using FragB = typename MarlinType::FragB; + using FragC = typename MarlinType::FragC; + using FragS = typename MarlinType::FragS; + using FragZP = typename MarlinType::FragZP; + + extern __shared__ int4 sh[]; + // NVFP4: b_type=FP4_E2M1, s_type=FP8_E4M3, a_type=c_type=BF16 + constexpr bool is_a_8bit = false; // BF16 activations are 16-bit + constexpr bool has_zp = false; // FP4 E2M1 has no zero-points + constexpr bool is_int_type = false; // FP4 E2M1 is not int type + constexpr bool dequant_skip_flop = true; // FP4 E2M1 + FP8 E4M3 scales + + scalar_t2 global_scale; + + constexpr bool has_act_order = group_blocks == 0; + + constexpr int pack_factor = 8; // 32 / 4 bits for FP4 E2M1 + static_assert(thread_m_blocks == 1 || !m_block_size_8); + int const group_size = (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; + int const scales_expert_stride = prob_n * prob_k / group_size / 16; + int const zp_expert_stride = 0; // No zero-points for NVFP4 + int const b_bias_expert_stride = prob_n / 8; + + // parallel: num valid moe blocks + int parallel = num_tokens_past_padded / moe_block_size; + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + + int global_mn_tiles = parallel * n_tiles; + int part2_mn_tiles = global_mn_tiles; + int part1_mn_iters = 0; + bool in_part2 = false; + + // we use DP + two-tile SK here + // part1: DP + // part2: two-tile SK + // see https://github.com/vllm-project/vllm/pull/24722 for more details + if (global_mn_tiles > gridDim.x) + { + part2_mn_tiles = global_mn_tiles % gridDim.x; + if (part2_mn_tiles * 3 <= gridDim.x) + part2_mn_tiles += gridDim.x; + part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; + } + + int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) + { + if (group_blocks >= thread_k_blocks) + { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = 0; + int slice_col_par = blockIdx.x; + int slice_col; + int slice_iters = k_tiles; // number of threadblock tiles in the current slice + // total number of active threadblocks in the current slice + int slice_count = 1; + // index of threadblock in current slice; numbered bottom to top + int slice_idx = 0; + + int par_id = 0; + int block_id = -1; + int64_t expert_id = 0; // use int64 to avoid computation result overflow + int old_expert_id = 0; + int64_t B_expert_off = 0; + + float* sh_a_s = reinterpret_cast(sh); + int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); + 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; + int32_t* sh_block_sorted_ids = reinterpret_cast(sh_block_sorted_ids_int4); + 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); + + int32_t block_num_valid_tokens = 0; + int32_t locks_off = 0; + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (part2_mn_tiles >= gridDim.x) + { + // when part2_mn_tiles >= sms + // then there are at most $sms$ conflict tile blocks + locks_off = blockIdx.x; + } + else + { + locks_off = (iters * blockIdx.x) / k_tiles - 1; + } + + int prob_m_top_k = prob_m * top_k; + // read moe block data given block_id + // block_sorted_ids / block_num_valid_tokens / block_topk_weights + auto read_moe_block_data = [&](int block_id) + { + block_num_valid_tokens = moe_block_size; + + cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x, + reinterpret_cast(sorted_token_ids_ptr) + (block_id * moe_block_size / 4 + threadIdx.x), + threadIdx.x < moe_block_size / 4); + + cp_async_fence(); + cp_async_wait<0>(); + + __syncthreads(); + + if (threadIdx.x >= threads - 32) + { + constexpr int size_per_thread = div_ceil(moe_block_size, 32); + int lane_id = threadIdx.x - (threads - 32); + + int local_count = 0; +#pragma unroll + for (int i = 0; i < size_per_thread; i++) + { + int j = lane_id * size_per_thread + i; + if (j < moe_block_size) + { + int idx = sh_block_sorted_ids[j]; + if (idx < prob_m_top_k) + local_count++; + } + } + + block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count); + + if (lane_id == 0) + reinterpret_cast(sh_new)[0] = block_num_valid_tokens; + } + + if (threadIdx.x < moe_block_size) + { + int idx = sh_block_sorted_ids[threadIdx.x]; + sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k; + + if (mul_topk_weights) + { + idx = idx < prob_m_top_k ? idx : 0; + scalar_t2 topk_weight_val = Dtype::num2num2(Dtype::float2num(topk_weights_ptr[idx])); + topk_weight_val = __hmul2(topk_weight_val, global_scale); + sh_block_topk_weights[threadIdx.x] = topk_weight_val; + } + } + + __syncthreads(); + + block_num_valid_tokens = reinterpret_cast(sh_new)[0]; + __syncthreads(); + }; + + // when move to next moe block, find the next block_id and expert_id + // and then read moe block data + auto update_next_moe_block_data = [&]() + { + if (par_id >= parallel) + return; + + old_expert_id = expert_id; + block_id = par_id; + expert_id = expert_ids_ptr[block_id]; + + { + uint16_t val = global_scale_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) + { + zp_ptr += (expert_id - old_expert_id) * zp_expert_stride; + } + 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); + }; + + // Compute all information about the current slice which is required for + // synchronization. + bool first_init = true; + auto init_part2_slice = [&]() + { + slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) + slice_iters = 0; + if (slice_iters == 0) + return; + if (slice_row + slice_iters > k_tiles) + slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) + { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) + slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else + { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) + slice_idx--; + } + } + if (part2_mn_tiles >= gridDim.x) + { + if (slice_count > 1 && slice_idx == slice_count - 1) + { + locks_off++; + } + } + else + { + locks_off++; + } + + if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) + { + constexpr int threads_per_m = 16 * thread_n_blocks / 8; + int m_per_thread = div_ceil(block_num_valid_tokens, threads / threads_per_m); + for (int i = 0; i < m_per_thread; i++) + { + int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; + if (row < block_num_valid_tokens) + { + int64_t sorted_row = sh_block_sorted_ids[row]; + int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; + C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0}; + } + } + // After write zero to output, write a negative value to lock. + // Every SM that processes the same slice would wait for + // the negative value, and then atomicAdd 1 to it. + // After all SMs are processed, the lock value would back to 0 again. + __syncthreads(); + if (threadIdx.x == 0) + locks[locks_off] = 1 - slice_count; + } + + if (slice_col == n_tiles) + { + slice_col = 0; + par_id++; + update_next_moe_block_data(); + } + if (is_a_8bit && (first_init || slice_col == 0)) + { + __syncthreads(); + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], + threadIdx.x < block_num_valid_tokens); + } + }; + + auto init_part1_slice = [&]() + { + if (part1_mn_iters) + { + part1_mn_iters--; + par_id = slice_col_par / n_tiles; + slice_col = slice_col_par % n_tiles; + slice_iters = k_tiles; + update_next_moe_block_data(); + if (is_a_8bit) + { + __syncthreads(); + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], + threadIdx.x < block_num_valid_tokens); + } + } + }; + + auto init_slice = [&]() + { + if (!in_part2 && !part1_mn_iters) + { + in_part2 = true; + slice_col_par = (iters * blockIdx.x) / k_tiles; + slice_row = (iters * blockIdx.x) % k_tiles; + slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; + par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; + update_next_moe_block_data(); + } + if (!in_part2) + { + init_part1_slice(); + } + else + { + init_part2_slice(); + first_init = false; + } + }; + + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8); + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); + constexpr int b_thread_vecs = 1; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / 16; + constexpr int s_sh_stride = 16 * thread_n_blocks / 16; + 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_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + 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; + + constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Zero-points sizes/strides + int zp_gl_stride = 0; // No zero-points for NVFP4 + constexpr int zp_sh_stride = 0; // No zero-points for NVFP4 + constexpr int zp_tb_groups = s_tb_groups; + constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; + int zp_gl_rd_delta = zp_gl_stride; + + // Global A read index of current thread. + 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. + int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); + a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; + + int b_gl_rd; + if (threads <= b_sh_stride) + { + b_gl_rd = threadIdx.x; + } + else + { + b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + } + + b_gl_rd += B_expert_off + b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + auto b_sh_rd = threadIdx.x * b_thread_vecs; + b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); + + // For act_order + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) + { + if constexpr (group_blocks == -1) + { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + s_gl_rd + = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; + } + else + { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + } + } + auto s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stage; + + // Zero-points + int zp_gl_rd; + if constexpr (has_zp) + { + if constexpr (group_blocks == -1) + { + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + + threadIdx.x; + } + else + { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + auto zp_sh_wr = threadIdx.x; + bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; + + // 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 (is_a_8bit) + { + s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); + } + else if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (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) % tb_n_warps) + (threadIdx.x % 32) / 8; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; + + int bias_sh_rd; + if constexpr (m_block_size_8) + { + bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + } + else + { + bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (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; + constexpr int num_row_threads = 4; + constexpr int num_ints_per_thread = 8 / pack_factor; + int zp_sh_rd; + if constexpr (has_zp) + { + if (is_a_8bit) + { + zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + else + { + zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + } + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // 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 % 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 + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + { +#pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + + // Shared memory storage for global fetch pipelines. + 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_a = sh_s + sh_s_size; + + // 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][is_a_8bit ? 2 : 4][2]; + FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + 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 + FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ + + if constexpr (is_a_8bit && group_blocks != -1) + { +#pragma unroll + for (int j = 0; j < 2; j++) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + frag_c_tmp[i][j][0][g] = 0.0f; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + + // Zero accumulators. + auto zero_accums = [&]() + { +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + + 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 > act_s_max_num_groups) + { + sh_num_groups = act_s_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) + { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) + { + for (int i = 0; i < sh_num_groups; i++) + { + if (threadIdx.x < s_sh_stride) + { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); + } + } + } + else + { + for (int i = 0; i < sh_num_groups; i++) + { + if (threadIdx.x < s_sh_stride) + { + sh_s[(i * s_sh_stride) + threadIdx.x] + = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) + { + if (pred) + { + int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + { + 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_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 * b_thread_vecs); i++) + { + constexpr int count = div_ceil(b_sh_stride, threads); + int b_gl_idx + = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); + + cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); + } + + b_gl_rd += b_gl_rd_delta_o; + + if constexpr (has_act_order) + { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) + { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) + { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } + else + { + if constexpr (group_blocks != -1) + { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta * s_tb_groups; + } + } + + if constexpr (has_zp && group_blocks != -1) + { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) + { + if (zp_sh_wr_pred) + { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + auto fetch_col_zp_to_shared = [&]() + { + if (zp_sh_wr_pred) + { + cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + }; + + auto fetch_col_scale_to_shared = [&]() + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() + { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // 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 + moe_block_size * a_sh_stride * pipe; +#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]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + +#pragma unroll + for (int i = 0; i < b_thread_vecs; i++) + { + frag_b_quant[k % 2][i] + = *reinterpret_cast(&sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) + { + if constexpr (!has_act_order) + { + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) + { + int pipe = full_pipe % stages; + using IT1 = typename std::conditional_t; + using IT0 = typename std::conditional_t; + constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + if constexpr (!has_act_order) + { + // No act-order case + if constexpr (group_blocks == -1) + { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 && dequant_skip_flop) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + else if constexpr (group_blocks != -1) + { + if constexpr (group_blocks >= thread_k_blocks) + { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) + { + if (k % b_sh_wr_iters == 0) + { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + 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 if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) + { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] + = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } + else if (group_blocks >= b_sh_wr_iters) + { + reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) + { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k % b_sh_wr_iters; + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + int warp_col = warp_id % tb_n_warps; + + cur_k += warp_row * 16 * b_sh_wr_iters; + + 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 = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) + { + if (k % 2 == 0) + { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) + = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; + } + else + { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) + = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) + { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) + = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread + +#pragma unroll + for (int i = 0; i < 4; i++) + { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + auto fetch_zp_to_registers = [&](int k, int full_pipe) + { + // This code does not handle group_blocks == 0, + // which signifies act_order. + // has_zp implies AWQ, which doesn't have act_order, + static_assert(!has_zp || group_blocks != 0); + + if constexpr (has_zp) + { + int pipe = full_pipe % stages; + + if constexpr (group_blocks == -1) + { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 || is_a_8bit) + { +#pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) + { + frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; + } + } + } + else if constexpr (group_blocks >= thread_k_blocks) + { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) + { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); +#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 + { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + 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]; + } + } + } + }; + + auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) + { + if constexpr (is_a_8bit && has_zp) + { + dequant_fp4(q, frag_b_ptr, zp); + } + else + { + dequant_fp4(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, int pipe) + { + if (is_a_8bit) + return; + int k2 = k % 2; + constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; + bool const is_new_zp = (group_blocks == 0) + || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) + || (group_blocks == -1 && is_first_matmul_in_slice); + if constexpr (has_zp) + { + if (is_new_zp) + { + if constexpr (group_blocks == -1) + is_first_matmul_in_slice = false; + int zp_quant_0 = frag_qzp[k2][0]; + int zp_quant_1 = zp_quant_0 >> 8; + + dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); + dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); + } + } + { + 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 + for (int j = 0; j < 4; j++) + { + FragB frag_b0; + FragB frag_b1; + int b_quant_1 = frag_b_quant[k2][0][j]; + int b_quant_0 = b_quant_1 << 8; + + 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_a_8bit) + { + 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 && !is_a_8bit) + { + static_assert(group_blocks != -1); + 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[k2][2][j], act_frag_s[k2][3][j], 1); + } + else if constexpr (!dequant_skip_flop && has_zp && group_blocks == -1 && !is_a_8bit) + { + int idx = (threadIdx.x / 4) % 2; + scalar_t2 s2 = Dtype::nums2num2(reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); + 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 (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) + { + if (is_new_zp) + frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); + 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 && !is_a_8bit) + { + scale(frag_b0, frag_s[k2][j], 0); + scale(frag_b1, frag_s[k2][j], 1); + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + if constexpr (m_block_size_8) + { + mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); + } + else + { + mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); + } + } + } + }; + + auto matmul_a8 = [&](int k) + { + int k2 = k % 2; +#pragma unroll + for (int j = 0; j < 2; j++) + { + FragB frag_b[2]; + + if (is_a_8bit && !has_zp) + { + dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b)); + dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2); + } + else if (is_a_8bit && has_zp) + { + int off = (threadIdx.x / 32) % 2 * 2 + j; + int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b), zp); + zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2, zp); + } + else + { + reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; + reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + mma(frag_a[k2][i], frag_b[0], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); + mma(frag_a[k2][i], frag_b[1], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); + } + + if constexpr (group_blocks != -1) + { + if (group_blocks == 2 || k == 1) + { + { + float2 s_vals[2]; + s_vals[0] = Dtype::num22float2(frag_s[k2][j * 2][0]); + s_vals[1] = Dtype::num22float2(frag_s[k2][j * 2 + 1][0]); + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&s_vals[0])[g % 2]; + frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&s_vals[1])[g % 2]; + frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() + { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) + { + auto red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 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); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + +#pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) + { +#pragma unroll + for (int i = red_off; i > 0; i /= 2) + { + if (i <= red_idx && red_idx < 2 * i) + { +#pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; j += (m_block_size_8 ? 2 : 1)) + { + int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) + { + float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); +#pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] + += c_rd[k] + c_wr[k]; + } + sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) + { +#pragma unroll + for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; i += (m_block_size_8 ? 2 : 1)) + { + float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); +#pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce_fp16 = [&](bool first = false, bool last = false) + { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + if (!is_th_active) + { + return; + } + + int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1); + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr; + if constexpr (m_block_size_8) + { + c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + } + else + { + c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); + } + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + if (!first) + { + +#pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) + { + int c_idx; + if constexpr (m_block_size_8) + c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; + else + c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); + if (c_idx / c_gl_stride < block_num_valid_tokens) + { + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; + if constexpr (is_a_8bit) + { + int2* sh_red_int2 = reinterpret_cast(sh_red); + int2* c_int2 = reinterpret_cast(C); + sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx]; + } + else + { + sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx]; + } + } + } + } + +#pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) + { + if (!first) + { + scalar_t* c_red_f16; + if constexpr (is_a_8bit) + { + int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } + else + { + int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } +#pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) + { + int delta = 0; + if constexpr (m_block_size_8) + { + delta = j % 2 == 1 ? -2 : 0; + } + reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] + += Dtype::num2float(c_red_f16[j]); + } + } + if (!last) + { + scalar_t c_f16[is_a_8bit ? 4 : 8]; +#pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) + { + int delta = 0; + if constexpr (m_block_size_8) + { + delta = j % 2 == 1 ? -2 : 0; + } + c_f16[j] = Dtype::float2num(reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); + } + + int c_idx; + if constexpr (m_block_size_8) + c_idx = c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; + else + c_idx = c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); + if (c_idx / c_gl_stride < block_num_valid_tokens) + { + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; + if constexpr (is_a_8bit) + { + int2* c_int2 = reinterpret_cast(C); + c_int2[true_idx] = *reinterpret_cast(c_f16); + } + else + { + C[true_idx] = *reinterpret_cast(c_f16); + } + } + } + } + }; + + // Globally reduce over threadblocks that compute the same column block. + // We use a tmp C buffer to reduce in full fp32 precision. + auto global_reduce_fp32 = [&](bool first = false, bool last = false) + { + constexpr int tb_m = thread_m_blocks * 16; + constexpr int tb_n = thread_n_blocks * 16; + + constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; + + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + + constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; + constexpr int th_size = num_floats * sizeof(float) / 16; + + int c_cur_offset = locks_off * c_size; + + if (!is_th_active) + { + return; + } + + if (!first) + { + float* frag_c_ptr = reinterpret_cast(&frag_c); +#pragma unroll + for (int k = 0; k < th_size; k++) + { + if constexpr (m_block_size_8) + { + if (k % 2) + continue; + } + else + { + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) + continue; + } + + sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; + + float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); +#pragma unroll + for (int f = 0; f < 4; f++) + { + frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; + } + } + } + + if (!last) + { + int4* frag_c_ptr = reinterpret_cast(&frag_c); +#pragma unroll + for (int k = 0; k < th_size; k++) + { + if constexpr (m_block_size_8) + { + if (k % 2) + continue; + } + else + { + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) + continue; + } + + C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; + } + } + }; + + // 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 = [&](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)); + constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr; + if constexpr (m_block_size_8) + { + c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; + c_sh_wr += 64 * (threadIdx.x / 32); + } + else + { + c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); + } + + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); + + // 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, 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 && !is_a_8bit + && (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 (!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) + { + ((scalar_t*) sh_red)[idx] = res.x; + ((scalar_t*) sh_red)[idx + 8 * c_sh_stride] = res.y; + } + else + { + ((scalar_t2*) sh_red)[idx] = res; + } + }; + + if (threadIdx.x / 32 < tb_n_warps) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 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], + 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], 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], 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], 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], frag_bias[j / 2][2 * (j % 2) + 1]); + } + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + +#pragma unroll + for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) + { + int row = c_gl_wr / c_gl_stride; + if (row < block_num_valid_tokens) + { + int64_t sorted_row = sh_block_sorted_ids[row]; + int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride; + scalar_t2 topk_weight_score; + if (mul_topk_weights) + topk_weight_score = sh_block_topk_weights[row]; + if (use_atomic_add && slice_count > 1 || mul_topk_weights) + { + scalar_t2* C_half2 = reinterpret_cast(&C[true_idx]); + scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); + if (mul_topk_weights) + { +#pragma unroll + for (int a = 0; a < 4; a++) + { + sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score); + } + } + + if (use_atomic_add && slice_count > 1) + { +#pragma unroll + for (int a = 0; a < 4; a++) + { + atomicAdd(&C_half2[a], sh_red_half2[a]); + } + } + else + { + C[true_idx] = *reinterpret_cast(sh_red_half2); + } + } + else + { + C[true_idx] = sh_red[c_sh_rd]; + } + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + __syncthreads(); + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() + { + +#pragma unroll + for (int i = 0; i < stages - 1; i++) + { + if (has_act_order && i == 0) + { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) + { + last_g_idx = prob_k - 1; + } + fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + + if constexpr (has_zp && group_blocks == -1) + { + if (i == 0) + { + fetch_col_zp_to_shared(); + if constexpr (!dequant_skip_flop) + { + fetch_col_scale_to_shared(); + } + } + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + fetch_zp_to_registers(0, 0); + 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(); + } + + // Main loop. + while (slice_iters) + { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + +#pragma unroll + 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); + } + + if constexpr (!is_a_8bit) + { + matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); + } + else + { + static_assert(group_blocks != 0 && group_blocks != 1); + matmul_a8(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(); + } + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) + { + // convert fp16 accum to fp32 for reduction + if constexpr (use_fp16_accum) + { +#pragma unroll + for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) + { + float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; + scalar_t* frag_c_part_half = reinterpret_cast(frag_c_part_float); + +#pragma unroll + for (int i = 3; i >= 0; i--) + { + frag_c_part_float[i] = Dtype::num2float(frag_c_part_half[i]); + } + } + } + + if constexpr (is_a_8bit) + { + float frag_a_s[2 * thread_m_blocks]; + + for (int i = 0; i < 2 * thread_m_blocks; i++) + frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; + +#pragma unroll + for (int j = 0; j < 2; j++) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float c_val = frag_c[i][j][0][g]; + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][0][g] = c_val * s_val; + } +#pragma unroll + for (int g = 0; g < 4; g++) + { + float c_val = frag_c[i][j][1][g]; + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][1][g] = c_val * s_val; + } + } + } + } + + cp_async_wait<0>(); + 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 && dequant_skip_flop || !has_zp)) + { + if ((last || use_atomic_add) || is_a_8bit) + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + + thread_block_reduce(); + + 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 || is_a_8bit)) + { + if constexpr (is_a_8bit) + { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + } + } + else if (last || use_atomic_add) + { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + if constexpr (m_block_size_8) + { + int idx = (threadIdx.x / 4) % 2; + scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); +#pragma unroll + for (int i = 0; i < 8; i++) + { + frag_s_half2[i] = Dtype::num2num2(reinterpret_cast(&frag_s_half2[i])[idx]); + } + } + } + } + } + + // 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 && is_a_8bit) + { +#pragma unroll + for (int j = 0; j < 2; j++) + { + float2 aa[2]; + aa[0] = Dtype::num22float2(frag_s[0][j * 2][0]); + aa[1] = Dtype::num22float2(frag_s[0][j * 2 + 1][0]); + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&aa[0])[g % 2]; + frag_c[i][j][0][g] *= scale; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&aa[1])[g % 2]; + frag_c[i][j][1][g] *= scale; + } + } + } + } + + if (slice_count > 1 && !use_atomic_add) + { + // only globally reduce if there is more than one block in a slice + barrier_acquire(&locks[locks_off], slice_idx); + if (use_fp32_reduce) + { + global_reduce_fp32(slice_idx == 0, last); + } + else + { + global_reduce_fp16(slice_idx == 0, last); + } + 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]; + if constexpr (!is_a_8bit) + 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(last); + slice_row = 0; + if (!in_part2) + { + slice_col_par += gridDim.x; + } + else + { + slice_col_par++; + slice_col++; + } + is_first_matmul_in_slice = true; + init_slice(); + + if (slice_iters) + { + a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; + b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; + + 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 + { + if constexpr (group_blocks == -1) + { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + + threadIdx.x; + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } + else + { + s_gl_rd + = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd + = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + start_pipes(); + } + } + } +} + +#endif + +} // namespace MARLIN_NAMESPACE_NAME diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h new file mode 100644 index 000000000000..fa02e8387430 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h @@ -0,0 +1,2073 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +#ifndef MARLIN_NAMESPACE_NAME +#define MARLIN_NAMESPACE_NAME marlin +#endif + +#include "marlin.cuh" + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +namespace MARLIN_NAMESPACE_NAME +{ + +// Empty kernel stub for non-Hopper device passes; see marlin.cuh. +#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) + +template shared + // fetch pipeline + bool const has_act_order, // whether act_order is enabled + int group_blocks // number of consecutive 16x16 blocks + // with a separate quantization scale + // (implicit const: trailing NTTP) + > +__global__ void Marlin(int4 const* __restrict__ A, // fp16 input matrix of shape mxk + int4 const* __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) + int4 const* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + int const* __restrict__ g_idx, // int32 group indices of shape k + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks, // extra global storage for barrier synchronization + bool use_fp32_reduce // whether to use fp32 global reduce +) +{ +} + +} // namespace marlin + +#else + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm(typename MarlinType::FragA& frag_a, void const* smem_ptr) +{ + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (count == 4) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } + else if constexpr (count == 2) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" : "=r"(a[0]), "=r"(a[1]) : "r"(smem)); + } + else if constexpr (count == 1) + { + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" : "=r"(a[0]) : "r"(smem)); + } + else + { + static_assert(count == 1 || count == 2 || count == 4, "invalid count"); + } +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale( + typename MarlinType::FragB& frag_b, typename MarlinType::FragS& frag_s, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 s = MarlinType::num2num2(reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +template +__device__ inline void scale_and_sub(typename MarlinType::FragB& frag_b, scalar_t s, scalar_t zp) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 s2 = MarlinType::num2num2(s); + scalar_t2 zp2 = MarlinType::num2num2(zp); + frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); + frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); +} + +template +__device__ inline void sub_zp( + typename MarlinType::FragB& frag_b, typename MarlinType::scalar_t2& frag_zp, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + scalar_t2 zp = MarlinType::num2num2(reinterpret_cast(&frag_zp)[i]); + frag_b[0] = __hsub2(frag_b[0], zp); + frag_b[1] = __hsub2(frag_b[1], zp); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4(typename MarlinType::FragB& frag_b, + typename MarlinType::FragS& frag_s_1, typename MarlinType::FragS& frag_s_2, + typename MarlinType::FragS& frag_s_3, typename MarlinType::FragS& frag_s_4, int i) +{ + using scalar_t2 = typename MarlinType::scalar_t2; + + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float(float* c, typename MarlinType::FragS& s) +{ + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], MarlinType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], MarlinType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) +{ + if (threadIdx.x == 0) + { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) +{ + __syncthreads(); + if (threadIdx.x == 0) + { + if (reset) + { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(lock), "r"(val)); + } +} + +// Wait until value of lock to be negative, and then add 1 +__device__ inline void wait_negative_and_add(int* lock) +{ + if (threadIdx.x == 0) + { + int state = 0; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + while (state >= 0); + atomicAdd(lock, 1); + } + __syncthreads(); +} + +template shared + // fetch pipeline + int group_blocks // number of consecutive 16x16 blocks + // with a separate quantization scale + // (implicit const: trailing NTTP) + > +__global__ void Marlin(int4 const* __restrict__ A0, // fp16 input matrix of shape mxk + int4 const* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C0, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + int4 const* __restrict__ b_bias_ptr, + // float scales of input matrix, only used when is_a_8bit == true. + // shape (m,) + float const* __restrict__ a_scales_ptr, + // fp16 quantization scales. shape (k/groupsize, n) + int4 const* __restrict__ scales_ptr, + // fp16 global scale (for nvfp4// only) + uint16_t const* __restrict__ global_scale_ptr, + // 4bit packed zero-points of shape + // (k/groupsize, n/pack_factor) + int4 const* __restrict__ zp_ptr, + // int32 group indices of shape k + int const* __restrict__ g_idx, + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int lda, // A.stride(0), equal to prob_k is A is contiguous + int* locks, // extra global storage for barrier synchronization + 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 + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + + constexpr bool use_fp16_accum = false; + + using scalar_t2 = typename MarlinType::scalar_t2; + using scalar_32bit_t = typename MarlinType::scalar_32bit_t; + + using c_scalar_t = scalar_t; + using c_scalar_t2 = scalar_t2; + + using FragA = typename MarlinType::FragA; + using FragB = typename MarlinType::FragB; + using FragC = typename MarlinType::FragC; + using FragS = typename MarlinType::FragS; + using FragZP = typename MarlinType::FragZP; + + int4 const* A = A0; + int4* C = C0; + + constexpr bool is_a_8bit = false; // BF16 is 16-bit + constexpr bool has_zp = false; // NVFP4 has no zero-points + // For NVFP4: dequant places bits, scale applied separately + constexpr bool dequant_skip_flop = true; + + c_scalar_t2 global_scale; + + { + uint16_t val = global_scale_ptr[0]; + global_scale = MarlinType::num2num2(*reinterpret_cast(&val)); + } + + constexpr bool has_act_order = group_blocks == 0; + constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); + + extern __shared__ int4 sh[]; + float* sh_a_s = reinterpret_cast(sh); + int4* sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); + constexpr int pack_factor = 8; // 32 / 4 (FP4) + static_assert(thread_m_blocks == 1 || !m_block_size_8); + + // For larger GEMMs we run multiple batchsize 64 versions in parallel for a + // better partitioning with less reductions + int parallel = 1; + if (prob_m > m_block_size) + { + parallel = prob_m / m_block_size; + prob_m = m_block_size; + } + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + + int global_mn_tiles = parallel * n_tiles; + int part2_mn_tiles = global_mn_tiles; + int part1_mn_iters = 0; + bool in_part2 = false; + + if (global_mn_tiles > gridDim.x) + { + part2_mn_tiles = global_mn_tiles % gridDim.x; + if (part2_mn_tiles * 3 <= gridDim.x) + part2_mn_tiles += gridDim.x; + part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; + } + + int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) + { + if (group_blocks >= thread_k_blocks) + { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = 0; + int slice_col_par = blockIdx.x; + int slice_col; + int slice_iters = k_tiles; // number of threadblock tiles in the current slice + // total number of active threadblocks in the current slice + int slice_count = 1; + // index of threadblock in current slice; numbered bottom to top + int slice_idx = 0; + + int par_id = 0; + int locks_off = 0; + + if (part2_mn_tiles >= gridDim.x) + { + // when part2_mn_tiles >= sms + // then there are at most $sms$ conflict tile blocks + locks_off = blockIdx.x; + } + else + { + locks_off = (iters * blockIdx.x) / k_tiles - 1; + } + + // Compute all information about the current slice which is required for + // synchronization. + bool first_init = true; + auto init_part2_slice = [&]() + { + slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) + slice_iters = 0; + if (slice_iters == 0) + return; + if (slice_row + slice_iters > k_tiles) + slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) + { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) + slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else + { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) + slice_idx--; + } + } + if (part2_mn_tiles >= gridDim.x) + { + if (slice_count > 1 && slice_idx == slice_count - 1) + { + locks_off++; + } + } + else + { + locks_off++; + } + + if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) + { + constexpr int threads_per_m = 16 * thread_n_blocks / 8; + int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); + if (m_block_size_8) + m_per_thread = div_ceil(8, threads / threads_per_m); + for (int i = 0; i < m_per_thread; i++) + { + int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; + if (row < prob_m) + { + int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; + C[row * prob_n / 8 + col] = {0, 0, 0, 0}; + } + } + // After write zero to output, write a negative value to lock. + // Every SM that processes the same slice would wait for + // the negative value, and then atomicAdd 1 to it. + // After all SMs are processed, the lock value would back to 0 again. + __syncthreads(); + if (threadIdx.x == 0) + locks[locks_off] = 1 - slice_count; + } + + if (slice_col == n_tiles) + { + A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); + C += 16 * thread_m_blocks * prob_n / 8; + slice_col = 0; + par_id++; + } + if (is_a_8bit && (first_init || slice_col == 0)) + { + __syncthreads(); + int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], threadIdx.x < prob_m); + } + }; + + auto init_part1_slice = [&]() + { + if (part1_mn_iters) + { + part1_mn_iters--; + par_id = slice_col_par / n_tiles; + slice_col = slice_col_par % n_tiles; + slice_iters = k_tiles; + A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; + C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; + if (is_a_8bit) + { + __syncthreads(); + int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; + cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], threadIdx.x < prob_m); + } + } + }; + + auto init_slice = [&]() + { + if (!in_part2 && !part1_mn_iters) + { + in_part2 = true; + slice_col_par = (iters * blockIdx.x) / k_tiles; + slice_row = (iters * blockIdx.x) % k_tiles; + slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; + par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; + A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; + C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; + } + if (!in_part2) + { + init_part1_slice(); + } + else + { + init_part2_slice(); + first_init = false; + } + }; + + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = lda / (is_a_8bit ? 16 : 8); + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * m_block_size; + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); + constexpr int b_thread_vecs = 1; // FP4: 1 vec per thread + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / (16 /* NVFP4 scale stride */); + constexpr int s_sh_stride = 16 * thread_n_blocks / (16 /* NVFP4 scale stride */); + 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_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + 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; + + constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Zero-points sizes/strides + int zp_gl_stride = (prob_n / pack_factor) / 4; + constexpr int zp_sh_stride = ((16 * thread_n_blocks) / pack_factor) / 4; + constexpr int zp_tb_groups = s_tb_groups; + constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; + 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; + // 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. + int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); + a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; + + int b_gl_rd; + if (threads <= b_sh_stride) + { + b_gl_rd = threadIdx.x; + } + else + { + b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + } + + b_gl_rd += b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + auto b_sh_rd = threadIdx.x * b_thread_vecs; + b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); + + // For act_order + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) + { + if constexpr (group_blocks == -1) + { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + s_gl_rd + = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; + } + else + { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + } + } + auto s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stage; + + // Zero-points + int zp_gl_rd; + if constexpr (has_zp) + { + if constexpr (group_blocks == -1) + { + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + + threadIdx.x; + } + else + { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + auto zp_sh_wr = threadIdx.x; + bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; + + // 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 (is_a_8bit) + { + s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); + } + else if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (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) % tb_n_warps) + (threadIdx.x % 32) / 8; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; + + int bias_sh_rd; + if constexpr (m_block_size_8) + { + bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + } + else + { + bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (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; + constexpr int num_row_threads = 4; + constexpr int num_ints_per_thread = 8 / pack_factor; + int zp_sh_rd; + if constexpr (has_zp) + { + if (is_a_8bit) + { + zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + else + { + zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + } + + // Precompute which thread should not read memory in which iterations; this is + // needed if there are more threads than required for a certain tilesize or + // when the batchsize is not a multiple of 16. + bool a_sh_wr_pred[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // 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 % 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 + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; +#pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) + { +#pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + + // Shared memory storage for global fetch pipelines. + 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_a = sh_s + sh_s_size; + + // 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][is_a_8bit ? 2 : 4][2]; + FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + 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 + FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ + + if constexpr (is_a_8bit) + { +#pragma unroll + for (int j = 0; j < 2; j++) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + frag_c_tmp[i][j][0][g] = 0.0f; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + + // Zero accumulators. + auto zero_accums = [&]() + { +#pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + + 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 > act_s_max_num_groups) + { + sh_num_groups = act_s_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) + { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) + { + for (int i = 0; i < sh_num_groups; i++) + { + if (threadIdx.x < s_sh_stride) + { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); + } + } + } + else + { + for (int i = 0; i < sh_num_groups; i++) + { + if (threadIdx.x < s_sh_stride) + { + sh_s[(i * s_sh_stride) + threadIdx.x] + = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) + { + if (pred) + { + int4* sh_a_stage = sh_a + a_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + { + cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], + &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], a_sh_wr_pred[i]); + } + int4* sh_b_stage = sh_b + b_sh_stage * pipe; +#pragma unroll + for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) + { + constexpr int count = div_ceil(b_sh_stride, threads); + int b_gl_idx + = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); + + cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); + } + + b_gl_rd += b_gl_rd_delta_o; + + if constexpr (has_act_order) + { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) + { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) + { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } + else + { + if constexpr (group_blocks != -1) + { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta * s_tb_groups; + } + } + + if constexpr (has_zp && group_blocks != -1) + { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) + { + if (zp_sh_wr_pred) + { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + auto fetch_col_zp_to_shared = [&]() + { + if (zp_sh_wr_pred) + { + cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + }; + + auto fetch_col_scale_to_shared = [&]() + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() + { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // 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; +#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]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + +#pragma unroll + for (int i = 0; i < b_thread_vecs; i++) + { + frag_b_quant[k % 2][i] + = *reinterpret_cast(&sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) + { + if constexpr (!has_act_order) + { + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) + { + int pipe = full_pipe % stages; + using IT1 = typename std::conditional_t; + using IT0 = typename std::conditional_t; + constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + if constexpr (!has_act_order) + { + // No act-order case + if constexpr (group_blocks == -1) + { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 && dequant_skip_flop) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } + else if constexpr (group_blocks != -1) + { + if constexpr (group_blocks >= thread_k_blocks) + { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) + { + if (k % b_sh_wr_iters == 0) + { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + 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 if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) + { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + reinterpret_cast(&frag_s[k % 2])[0] + = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } + else if (group_blocks >= b_sh_wr_iters) + { + reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) + { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k % b_sh_wr_iters; + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + int warp_col = warp_id % tb_n_warps; + + cur_k += warp_row * 16 * b_sh_wr_iters; + + 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 = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) + { + if (k % 2 == 0) + { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) + = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; + } + else + { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) + = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) + { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) + = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread + +#pragma unroll + for (int i = 0; i < 4; i++) + { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + auto fetch_zp_to_registers = [&](int k, int full_pipe) + { + // This code does not handle group_blocks == 0, + // which signifies act_order. + // has_zp implies AWQ, which doesn't have act_order, + static_assert(!has_zp || group_blocks != 0); + + if constexpr (has_zp) + { + int pipe = full_pipe % stages; + + if constexpr (group_blocks == -1) + { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 || is_a_8bit) + { +#pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) + { + frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; + } + } + } + else if constexpr (group_blocks >= thread_k_blocks) + { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) + { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); +#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 + { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + 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]; + } + } + } + }; + + auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) + { + // 16-bit != 4-bit, always dequant + dequant_fp4(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, int pipe) + { + if (is_a_8bit) + return; + int k2 = k % 2; + constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; + const bool is_new_zp = (group_blocks == 0) + || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) + || (group_blocks == -1 && is_first_matmul_in_slice); + if constexpr (has_zp) + { + if (is_new_zp) + { + if constexpr (group_blocks == -1) + is_first_matmul_in_slice = false; + int zp_quant_0, zp_quant_1; + + // FP4: 4-bit zero-points + zp_quant_0 = frag_qzp[k2][0]; + zp_quant_1 = zp_quant_0 >> 8; + + dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); + dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); + } + } + + // NVFP4: dequant FP8 scales to BF16 + { + 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 + for (int j = 0; j < 4; j++) + { + FragB frag_b0; + FragB frag_b1; + int b_quant_0, b_quant_1; + + // NVFP4 (FE2M1f): shift to extract two halves + b_quant_1 = frag_b_quant[k2][0][j]; + b_quant_0 = b_quant_1 << 8; + + 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_a_8bit) + { + 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 && !is_a_8bit) + { + static_assert(group_blocks != -1); + 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[k2][2][j], act_frag_s[k2][3][j], 1); + } + else if constexpr (!dequant_skip_flop && has_zp && group_blocks == -1 && !is_a_8bit) + { + int idx = (threadIdx.x / 4) % 2; + scalar_t2 s2 + = MarlinType::nums2num2(reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); + 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 (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) + { + if (is_new_zp) + frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); + 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 && !is_a_8bit) + { + scale(frag_b0, frag_s[k2][j], 0); + scale(frag_b1, frag_s[k2][j], 1); + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + if constexpr (m_block_size_8) + { + mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); + } + else + { + mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); + mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); + } + } + } + }; + + auto matmul_a8 = [&](int k) + { + int k2 = k % 2; +#pragma unroll + for (int j = 0; j < 2; j++) + { + FragB frag_b[2]; + + if (is_a_8bit && !has_zp) + { + dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b)); + dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2); + } + else if (is_a_8bit && has_zp) + { + int off = (threadIdx.x / 32) % 2 * 2 + j; + int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2], reinterpret_cast(&frag_b), zp); + zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2 + 1], reinterpret_cast(&frag_b) + 2, zp); + } + else + { + reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; + reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; + } + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { + mma(frag_a[k2][i], frag_b[0], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); + mma(frag_a[k2][i], frag_b[1], (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); + } + + if constexpr (group_blocks != -1) + { + if (group_blocks == 2 || k == 1) + { + { + float2 s_vals[2]; + s_vals[0] = MarlinType::num22float2(frag_s[k2][j * 2][0]); + s_vals[1] = MarlinType::num22float2(frag_s[k2][j * 2 + 1][0]); + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&s_vals[0])[g % 2]; + frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&s_vals[1])[g % 2]; + frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() + { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) + { + auto red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 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); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + +#pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) + { +#pragma unroll + for (int i = red_off; i > 0; i /= 2) + { + if (i <= red_idx && red_idx < 2 * i) + { +#pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; j += (m_block_size_8 ? 2 : 1)) + { + int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) + { + float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); +#pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] + += c_rd[k] + c_wr[k]; + } + sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) + { +#pragma unroll + for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; i += (m_block_size_8 ? 2 : 1)) + { + float* c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); +#pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast(frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce_fp16 = [&](bool first = false, bool last = false) + { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * tb_n_warps; + if (threadIdx.x < active_threads) + { + int c_gl_stride = prob_n / 8; + int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr; + if constexpr (m_block_size_8) + { + c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + } + else + { + c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + 4 * (threadIdx.x / 32) + + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); + } + constexpr int c_sh_wr_delta = active_threads; + auto c_sh_wr = threadIdx.x; + + int row = (threadIdx.x % 32) / 4; + + if (!first) + { +// Interestingly, doing direct global accesses here really seems to mess up +// the compiler and lead to slowdowns, hence we also use async-copies even +// though these fetches are not actually asynchronous. +#pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) + { + if constexpr (m_block_size_8) + { + cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], + (threadIdx.x % 4) * 2 + i < prob_m); + } + else if constexpr (is_a_8bit) + { + int2* sh_red_int2 = reinterpret_cast(sh_red); + int2* c_int2 = reinterpret_cast(C); + cp_async2_ca_pred(&sh_red_int2[c_sh_wr + c_sh_wr_delta * i], + &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + else + { + cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], + &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], + i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); + } + } + cp_async_fence(); + cp_async_wait<0>(); + } + +#pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) + { + bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) + || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); + if (mask) + { + if (!first) + { + c_scalar_t* c_red_f16; + if constexpr (is_a_8bit) + { + int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } + else + { + int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } +#pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) + { + int delta = 0; + if constexpr (m_block_size_8) + { + delta = j % 2 == 1 ? -2 : 0; + } + reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] + += MarlinType::num2float(c_red_f16[j]); + } + } + if (!last) + { + c_scalar_t c_f16[is_a_8bit ? 4 : 8]; +#pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) + { + int delta = 0; + if constexpr (m_block_size_8) + { + delta = j % 2 == 1 ? -2 : 0; + } + c_f16[j] = MarlinType::float2num(reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); + } + if constexpr (m_block_size_8) + { + C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] + = *reinterpret_cast(c_f16); + } + else if constexpr (is_a_8bit) + { + int2* c_int2 = reinterpret_cast(C); + c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] + = *reinterpret_cast(c_f16); + } + else + { + C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] + = *reinterpret_cast(c_f16); + } + } + } + } + } + }; + + // Globally reduce over threadblocks that compute the same column block. + // We use a tmp C buffer to reduce in full fp32 precision. + auto global_reduce_fp32 = [&](bool first = false, bool last = false) + { + constexpr int tb_m = thread_m_blocks * 16; + constexpr int tb_n = thread_n_blocks * 16; + + constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; + + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + + constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; + constexpr int th_size = num_floats * sizeof(float) / 16; + + int c_cur_offset = locks_off * c_size; + + if (!is_th_active) + { + return; + } + + if (!first) + { + float* frag_c_ptr = reinterpret_cast(&frag_c); +#pragma unroll + for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) + { + sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; + + float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); +#pragma unroll + for (int f = 0; f < 4; f++) + { + frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; + } + } + } + + if (!last) + { + int4* frag_c_ptr = reinterpret_cast(&frag_c); +#pragma unroll + for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) + { + C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; + } + } + }; + + // 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 = [&](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)); + constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr; + if constexpr (m_block_size_8) + { + c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; + c_sh_wr += 64 * (threadIdx.x / 32); + } + else + { + c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); + } + + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); + + int c_gl_wr_end = c_gl_stride * prob_m; + // 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, FragS& b_bias) + { + c_scalar_t2 res = MarlinType::nums2num2( + MarlinType::float2num(c0), MarlinType::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit + && (has_zp && dequant_skip_flop || !has_zp)) + { + c_scalar_t2 tmp_scale = s[0]; + if constexpr (m_block_size_8) + { + tmp_scale + = MarlinType::num2num2(reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); + } + res = __hmul2(res, tmp_scale); + } + + // NVFP4 with FP8 (E4M3) scales + res = __hmul2(res, global_scale); + if (has_bias && last) + { + c_scalar_t2 tmp_bias = b_bias[0]; + if constexpr (m_block_size_8) + { + tmp_bias = MarlinType::num2num2( + reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + } + res = __hadd2(res, tmp_bias); + } + + if constexpr (m_block_size_8) + { + ((c_scalar_t*) sh_red)[idx] = res.x; + ((c_scalar_t*) sh_red)[idx + 8 * c_sh_stride] = res.y; + } + else + { + ((c_scalar_t2*) sh_red)[idx] = res; + } + }; + + if (threadIdx.x / 32 < tb_n_warps) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 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], + 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], 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], 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], 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], frag_bias[j / 2][2 * (j % 2) + 1]); + } + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + +#pragma unroll + for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) + { + if (c_gl_wr < c_gl_wr_end) + { + if (use_atomic_add && slice_count > 1) + { + c_scalar_t2* C_half2 = reinterpret_cast(&C[c_gl_wr]); + c_scalar_t2* sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); +#pragma unroll + for (int a = 0; a < 4; a++) + { + atomicAdd(&C_half2[a], sh_red_half2[a]); + } + } + else + { + C[c_gl_wr] = sh_red[c_sh_rd]; + } + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + __syncthreads(); + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() + { + +#pragma unroll + for (int i = 0; i < stages - 1; i++) + { + if (has_act_order && i == 0) + { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) + { + last_g_idx = prob_k - 1; + } + fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); + } + + if constexpr (has_zp && group_blocks == -1) + { + if (i == 0) + { + fetch_col_zp_to_shared(); + if constexpr (!dequant_skip_flop) + { + fetch_col_scale_to_shared(); + } + } + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + 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); + if constexpr (has_act_order) + { + slice_k_start_shared_fetch += tb_k * (stages - 1); + } + }; + if (slice_iters) + { + start_pipes(); + } + + // Main loop. + while (slice_iters) + { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + +#pragma unroll + 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); + } + + if constexpr (!is_a_8bit) + { + matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); + } + else + { + static_assert(group_blocks != 0 && group_blocks != 1); + matmul_a8(k); + } + } + slice_iters--; + if (slice_iters == 0) + { + break; + } + } + + a_gl_rd += 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(); + } + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) + { + // convert fp16 accum to fp32 for reduction + if constexpr (use_fp16_accum) + { +#pragma unroll + for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) + { + float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; + scalar_t* frag_c_part_half = reinterpret_cast(frag_c_part_float); + +#pragma unroll + for (int i = 3; i >= 0; i--) + { + frag_c_part_float[i] = MarlinType::num2float(frag_c_part_half[i]); + } + } + } + + if constexpr (is_a_8bit) + { + float frag_a_s[2 * thread_m_blocks]; + + for (int i = 0; i < 2 * thread_m_blocks; i++) + frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; + +#pragma unroll + for (int j = 0; j < 2; j++) + { +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float c_val = frag_c[i][j][0][g]; + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][0][g] = c_val * s_val; + } +#pragma unroll + for (int g = 0; g < 4; g++) + { + float c_val = frag_c[i][j][1][g]; + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][1][g] = c_val * s_val; + } + } + } + } + + cp_async_wait<0>(); + 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 && dequant_skip_flop || !has_zp)) + { + if ((last || use_atomic_add) || is_a_8bit) + { + if (s_sh_wr_pred) + { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + + thread_block_reduce(); + + 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 || is_a_8bit)) + { + if constexpr (is_a_8bit) + { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + } + } + else if (last || use_atomic_add) + { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) + { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + if constexpr (m_block_size_8) + { + int idx = (threadIdx.x / 4) % 2; + c_scalar_t2* frag_s_half2 = reinterpret_cast(frag_s); +#pragma unroll + for (int i = 0; i < 8; i++) + { + frag_s_half2[i] = MarlinType::num2num2( + reinterpret_cast(&frag_s_half2[i])[idx]); + } + } + } + } + } + + // 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 && is_a_8bit) + { +#pragma unroll + for (int j = 0; j < 2; j++) + { + float2 aa[2]; + aa[0] = MarlinType::num22float2(frag_s[0][j * 2][0]); + aa[1] = MarlinType::num22float2(frag_s[0][j * 2 + 1][0]); + +#pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + { +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&aa[0])[g % 2]; + frag_c[i][j][0][g] *= scale; + } + +#pragma unroll + for (int g = 0; g < 4; g++) + { + float scale = reinterpret_cast(&aa[1])[g % 2]; + frag_c[i][j][1][g] *= scale; + } + } + } + } + + if (slice_count > 1 && !use_atomic_add) + { + // only globally reduce if there is more than one block in a slice + barrier_acquire(&locks[locks_off], slice_idx); + if (use_fp32_reduce) + { + global_reduce_fp32(slice_idx == 0, last); + } + else + { + global_reduce_fp16(slice_idx == 0, last); + } + 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]; + if constexpr (!is_a_8bit) + 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(last); + slice_row = 0; + if (!in_part2) + { + slice_col_par += gridDim.x; + } + else + { + slice_col_par++; + slice_col++; + } + is_first_matmul_in_slice = true; + init_slice(); + + 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 += a_gl_rd_delta_o * slice_row; + b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; + + 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 + { + if constexpr (group_blocks == -1) + { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } + else if constexpr (group_blocks >= thread_k_blocks) + { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + + threadIdx.x; + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } + else + { + s_gl_rd + = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd + = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + start_pipes(); + } + } + } +} + +} // namespace MARLIN_NAMESPACE_NAME + +#endif diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu new file mode 100644 index 000000000000..03996eebbf68 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_repack.cu @@ -0,0 +1,350 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "marlin.cuh" +#include "marlin_nvfp4.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +namespace marlin +{ + +template +__global__ void gptq_marlin_repack_kernel(uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) +{ + constexpr int pack_factor = 32 / num_bits; + + constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); + constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); + int k_tiles = size_k / target_tile_k_size; + int n_tiles = size_n / target_tile_n_size; + int block_k_tiles = div_ceil(k_tiles, gridDim.x); + + auto start_k_tile = blockIdx.x * block_k_tiles; + if (start_k_tile >= k_tiles) + { + return; + } + + int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() + { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + extern __shared__ int4 sh[]; + + constexpr int perm_size = target_tile_k_size / 4; + + int4* sh_perm_ptr = sh; + int4* sh_pipe_ptr = sh_perm_ptr; + if constexpr (has_perm) + { + sh_pipe_ptr += perm_size; + } + + constexpr int tile_ints = target_tile_k_size / pack_factor; + + constexpr int stage_n_threads = target_tile_n_size / 4; + constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; + constexpr int stage_size = stage_k_threads * stage_n_threads; + + auto load_perm_to_shared = [&](int k_tile_id) + { + int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; + + int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); + + if (threadIdx.x < perm_size) + { + sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; + } + __syncthreads(); + }; + + auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) + { + if (n_tile_id >= n_tiles) + { + cp_async_fence(); + return; + } + + int first_n = n_tile_id * target_tile_n_size; + + int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; + + if constexpr (has_perm) + { + if (threadIdx.x < stage_size) + { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; + + uint32_t const* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); + + int src_k = sh_perm_int_ptr[k_id]; + int src_k_packed = src_k / pack_factor; + + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); + } + } + else + { + if (threadIdx.x < stage_size) + { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; + + int first_k = k_tile_id * target_tile_k_size; + int first_k_packed = first_k / pack_factor; + + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); + } + } + + cp_async_fence(); + }; + + auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) + { + if (n_tile_id >= n_tiles) + { + return; + } + + auto warp_id = threadIdx.x / 32; + auto th_id = threadIdx.x % 32; + + if (warp_id >= 4) + { + return; + } + + int tc_col = th_id / 4; + int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); + + constexpr int tc_offsets[4] = {0, 1, 8, 9}; + + int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; + + constexpr int sh_stride = target_tile_n_size; + constexpr uint32_t mask = (1 << num_bits) - 1; + + int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; + uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); + + uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); + + uint32_t vals[8]; + + if constexpr (has_perm) + { + static_assert(!is_a_8bit); + for (int i = 0; i < 4; i++) + { + int k_idx = tc_row + tc_offsets[i]; + + uint32_t src_k = sh_perm_int_ptr[k_idx]; + uint32_t src_k_pos = src_k % pack_factor; + + uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; + uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; + + uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; + uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; + + vals[i] = b1_cur_val; + vals[4 + i] = b2_cur_val; + } + } + else + { + uint32_t b1_vals[tile_ints]; + uint32_t b2_vals[tile_ints]; + +#pragma unroll + for (int i = 0; i < tile_ints; i++) + { + if constexpr (is_a_8bit) + { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; + } + else + { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; + b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; + } + } + +#pragma unroll + for (int i = 0; i < 4; i++) + { + int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); + int cur_int = cur_elem / pack_factor; + int cur_pos = cur_elem % pack_factor; + + vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; + if constexpr (is_a_8bit) + vals[4 + i] = (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; + else + vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; + } + } + + constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; + int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; + + // Result of: + // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + if constexpr (!is_a_8bit && num_bits == 4) + { + int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + uint32_t res = 0; +#pragma unroll + for (int i = 0; i < 8; i++) + { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + } + else if constexpr (is_a_8bit && num_bits == 4) + { + int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; + + uint32_t res = 0; +#pragma unroll + for (int i = 0; i < 8; i++) + { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + } + else + { + constexpr int pack_idx[4] = {0, 2, 1, 3}; + + uint32_t res1 = 0; + uint32_t res2 = 0; +#pragma unroll + for (int i = 0; i < 4; i++) + { + const int ii = is_a_8bit ? i : pack_idx[i]; + res1 |= vals[ii] << (i * 8); + res2 |= vals[4 + ii] << (i * 8); + } + + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; + } + }; + + auto start_pipes = [&](int k_tile_id, int n_tile_id) + { +#pragma unroll + for (int pipe = 0; pipe < repack_stages - 1; pipe++) + { + fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); + } + + wait_for_stage(); + }; +#pragma unroll + for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) + { + int n_tile_id = 0; + + if constexpr (has_perm) + { + load_perm_to_shared(k_tile_id); + } + + start_pipes(k_tile_id, n_tile_id); + + while (n_tile_id < n_tiles) + { +#pragma unroll + for (int pipe = 0; pipe < repack_stages; pipe++) + { + fetch_to_shared( + (pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); + repack_tile(pipe, k_tile_id, n_tile_id + pipe); + wait_for_stage(); + } + n_tile_id += repack_stages; + } + } +} + +} // namespace marlin + +#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM && is_a_8bit == IS_A_8BIT) \ + { \ + cudaFuncSetAttribute(marlin::gptq_marlin_repack_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + marlin::gptq_marlin_repack_kernel \ + <<>>( \ + b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ + } + +namespace marlin_nvfp4 +{ + +void gptq_marlin_repack_dispatch(uint32_t const* b_q_weight_ptr, uint32_t const* perm_ptr, uint32_t* out_ptr, + int size_k, int size_n, int num_bits, bool has_perm, bool is_a_8bit, cudaStream_t stream) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TLLM_CHECK_WITH_INFO( + sm >= 90 && sm < 100, "Marlin NVFP4 repack is only supported on Hopper (SM 9.x); current SM = %d", sm); + + int blocks; + int dev; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + + if (false) + { + } + CALL_IF(4, false, false) + CALL_IF(4, true, false) + CALL_IF(8, false, false) + CALL_IF(8, true, false) + CALL_IF(4, false, true) + CALL_IF(8, false, true) +} + +} // namespace marlin_nvfp4 + +#undef CALL_IF diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 6f343a055eb4..5f3c8e24a134 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -50,6 +50,9 @@ add_library( cublasScaledMM.cpp cublasFp4ScaledMM.cpp cudaNvfp4MM.cpp + marlinNvfp4MM.cpp + marlinNvfp4MoeMM.cpp + marlinRepack.cpp cudaScaledMM.cpp dynamicDecodeOp.cpp fmhaPackMaskOp.cpp diff --git a/cpp/tensorrt_llm/thop/marlinNvfp4MM.cpp b/cpp/tensorrt_llm/thop/marlinNvfp4MM.cpp new file mode 100644 index 000000000000..40981f647e09 --- /dev/null +++ b/cpp/tensorrt_llm/thop/marlinNvfp4MM.cpp @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/marlin/marlin_nvfp4.h" +#include "tensorrt_llm/thop/outputTensor.h" +#include "tensorrt_llm/thop/thUtils.h" +#include + +using torch::Tensor; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Marlin NVFP4 GEMM: W4A(4|16) interface. +// +// Accepts either FP4 E2M1 packed activations or BF16 activations, plus Marlin-tiled FP4 weights. +// FP4 path: dequantizes activations to BF16 using scale_a/alpha, then runs W4A16 Marlin GEMM. +// BF16 path: skips dequant, runs W4A16 Marlin GEMM directly (scale_a/alpha ignored). +// +// mat_a: [M, K/2] FP4 E2M1 packed as uint8 (2 values per byte) OR [M, K] BF16 +// mat_b: Marlin-packed FP4 weights (int32), from gptq_marlin_repack +// scale_a: activation block scales (FP8 E4M3, swizzled, stored as uint8) +// scale_b: Marlin-processed weight block scales (from marlin_permute_scales + nvfp4_marlin_process_scales) +// alpha: float32 global scale (applied to activation dequant only) +// weight_global_scale: BF16 Marlin-processed weight global scale (from nvfp4_marlin_process_global_scale) +// bias: optional bias (not supported) +// out_dtype: output dtype (fp16 or bf16) +// size_n: output dimension N +// size_k: reduction dimension K +// output_buffer_kind: output allocation kind (0=default, 1=userbuffers, 2=nccl_window) +// group: communicator ranks, used when output_buffer_kind selects an NCCL window +Tensor marlin_nvfp4_gemm(Tensor const& mat_a, Tensor const& mat_b, std::optional const& scale_a, + Tensor const& scale_b, std::optional const& alpha, Tensor const& weight_global_scale, + std::optional const& bias, std::optional out_dtype, int64_t size_n, int64_t size_k, + int64_t output_buffer_kind = 0, c10::optional> group = c10::nullopt) +{ + CHECK_TH_CUDA(mat_a); + TORCH_CHECK(mat_a.scalar_type() == FLOAT4_E2M1X2 || mat_a.scalar_type() == at::ScalarType::BFloat16, + "mat_a must be FP4 E2M1X2 or BFloat16, got ", mat_a.scalar_type()); + CHECK_TH_CUDA(mat_b); + CHECK_TH_CUDA(scale_b); + CHECK_TH_CUDA(weight_global_scale); + if (mat_a.scalar_type() != at::ScalarType::BFloat16) + { + TORCH_CHECK(scale_a.has_value() and alpha.has_value(), "scale_a must be provided for FP4 activations"); + CHECK_INPUT(scale_a.value(), SF_DTYPE); // e4m3 + CHECK_INPUT(alpha.value(), at::ScalarType::Float); + } + + TORCH_CHECK(mat_a.dim() == 2, "A must be 2D tensor"); + TORCH_CHECK(!bias.has_value(), "bias is not supported yet"); + + auto const out_dtype_ = out_dtype.value_or(at::ScalarType::Half); + TORCH_CHECK( + out_dtype_ == at::ScalarType::Half || out_dtype_ == at::ScalarType::BFloat16, "Output must be fp16 or bf16"); + + int32_t m = mat_a.sizes()[0]; + int32_t n = static_cast(size_n); + int32_t k = static_cast(size_k); + + // Allocate output + auto [out, _] = torch_ext::allocate_output( + {m, n}, out_dtype_, mat_a.device(), static_cast(output_buffer_kind), group); + + if (m == 0) + return out; + + auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); + + Tensor act_bf16; + if (mat_a.scalar_type() == at::ScalarType::BFloat16) + { + // BF16 activations — skip FP4 dequant, use directly + act_bf16 = mat_a; + } + else + { + // FP4 activations — dequantize to BF16 + act_bf16 = at::empty({m, k}, mat_a.options().dtype(at::ScalarType::BFloat16)); + ::marlin_nvfp4::dequantFp4Activations(mat_a.data_ptr(), scale_a.value().data_ptr(), + reinterpret_cast(alpha.value().data_ptr()), act_bf16.data_ptr(), m, k, stream); + } + + int sms = 0; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, mat_a.get_device()); + auto workspace = at::zeros({sms}, mat_a.options().dtype(at::kInt)); + + int num_groups = k / 16; + int group_size = 16; + + // Step 2: Run Marlin W4A16 GEMM + // weight_global_scale is BF16, already Marlin-processed (includes 2^119 bias correction) + ::marlin_nvfp4::marlinNvfp4Gemm(act_bf16.data_ptr(), mat_b.data_ptr(), out.data_ptr(), + nullptr, // C_tmp + scale_b.data_ptr(), weight_global_scale.data_ptr(), m, n, k, reinterpret_cast(workspace.data_ptr()), + num_groups, group_size, + false, // use_fp32_reduce + stream); + + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "marlin_nvfp4_gemm(Tensor mat_a, Tensor mat_b, Tensor? scale_a, Tensor scale_b, Tensor? alpha," + " Tensor weight_global_scale, Tensor? bias, ScalarType? out_dtype," + " int size_n, int size_k, int output_buffer_kind=0, int[]? group=None) -> (Tensor out)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("marlin_nvfp4_gemm", &tensorrt_llm::torch_ext::marlin_nvfp4_gemm); +} diff --git a/cpp/tensorrt_llm/thop/marlinNvfp4MoeMM.cpp b/cpp/tensorrt_llm/thop/marlinNvfp4MoeMM.cpp new file mode 100644 index 000000000000..88924631bd83 --- /dev/null +++ b/cpp/tensorrt_llm/thop/marlinNvfp4MoeMM.cpp @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/marlin/marlin_nvfp4.h" +#include "tensorrt_llm/runtime/torchUtils.h" +#include "tensorrt_llm/thop/thUtils.h" +#include + +using torch::Tensor; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// W4A16 Fused MoE Marlin NVFP4 GEMM: BF16 activations + FP4 weights +// +// a: [M, K] BF16 activations +// b_q_weight: [num_experts, K/tile_size, N*tile_size/pack_factor] Marlin-packed FP4 +// b_scales: [num_experts, num_groups, N] FP8 E4M3 block scales +// global_scale: [num_experts] BF16 per-expert global scales +// workspace: int32 lock buffer +// sorted_token_ids: [max_num_tokens_padded] int32 +// expert_ids: [max_num_tokens_padded / block_size] int32 +// num_tokens_past_padded: [1] int32 device tensor +// topk_weights: [M, top_k] float32 router weights +// moe_block_size: MoE block size (typically 16) +// top_k: experts per token +// mul_topk_weights: whether to multiply topk weights in-kernel +// size_n: output dimension N +// size_k: reduction dimension K +// out_dtype: output data type +// use_fp32_reduce: use FP32 for intermediate reduction +Tensor marlin_nvfp4_moe_gemm(Tensor const& a, Tensor const& b_q_weight, Tensor const& b_scales, + Tensor const& global_scale, Tensor const& workspace, Tensor const& sorted_token_ids, Tensor const& expert_ids, + Tensor const& num_tokens_past_padded, Tensor const& topk_weights, int64_t moe_block_size, int64_t top_k, + bool mul_topk_weights, int64_t size_n, int64_t size_k, std::optional out_dtype, + bool use_fp32_reduce = false) +{ + CHECK_INPUT(a, at::kBFloat16); + CHECK_INPUT(b_q_weight, at::kLong); // 16x nvfp4 + CHECK_INPUT(b_scales, at::kInt); + CHECK_TH_CUDA(global_scale); + CHECK_TH_CUDA(workspace); + CHECK_INPUT(sorted_token_ids, at::kInt); + CHECK_INPUT(expert_ids, at::kInt); + CHECK_INPUT(num_tokens_past_padded, at::kInt); + CHECK_TH_CUDA(topk_weights); + + TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]"); + TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [num_experts, ...]"); + + int64_t size_m = a.size(0); + + auto const out_dtype_ = out_dtype.value_or(at::ScalarType::BFloat16); + TORCH_CHECK(out_dtype_ == at::ScalarType::BFloat16, "Output must be bf16"); + + // Output: [max_num_tokens_padded, N] (MoE output with sorted token ordering) + auto out = at::zeros({size_m * top_k, size_n}, a.options().dtype(out_dtype_)); + + if (size_m == 0) + return out; + + auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); + cudaDataType_t outType = convert_torch_dtype(out.scalar_type()); + + // Compute num_groups from b_scales shape [num_experts, num_groups, N] + int num_groups = b_scales.size(1); + int group_size = (num_groups > 1) ? (static_cast(size_k) / num_groups) : -1; + + // Allocate C_tmp for FP32 reduce + Tensor c_tmp; + if (use_fp32_reduce) + { + int sms = -1; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device()); + long max_c_tmp_size = std::min((long) size_n * sorted_token_ids.size(0), + (long) sms * 4 * moe_block_size * 256); // max_thread_n = 256 + if (moe_block_size == 8) + max_c_tmp_size *= 2; + c_tmp = at::empty({max_c_tmp_size}, a.options().dtype(at::kFloat)); + } + else + { + c_tmp = at::empty({0}, a.options().dtype(at::kFloat)); + } + + bool use_atomic_add = false; + + ::marlin_nvfp4::marlinNvfp4MoeGemmDispatcher(a.data_ptr(), b_q_weight.data_ptr(), out.data_ptr(), c_tmp.data_ptr(), + b_scales.data_ptr(), global_scale.data_ptr(), sorted_token_ids.data_ptr(), expert_ids.data_ptr(), + num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(), static_cast(moe_block_size), + static_cast(top_k), mul_topk_weights, static_cast(size_m), static_cast(size_n), + static_cast(size_k), const_cast(workspace.data_ptr()), num_groups, group_size, use_fp32_reduce, + use_atomic_add, outType, stream); + + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "marlin_nvfp4_moe_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, Tensor global_scale," + " Tensor workspace, Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_past_padded," + " Tensor topk_weights, int moe_block_size, int top_k, bool mul_topk_weights," + " int size_n, int size_k, ScalarType? out_dtype, bool use_fp32_reduce=False) -> (Tensor out)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("marlin_nvfp4_moe_gemm", &tensorrt_llm::torch_ext::marlin_nvfp4_moe_gemm); +} diff --git a/cpp/tensorrt_llm/thop/marlinRepack.cpp b/cpp/tensorrt_llm/thop/marlinRepack.cpp new file mode 100644 index 000000000000..5222b6a4ba5b --- /dev/null +++ b/cpp/tensorrt_llm/thop/marlinRepack.cpp @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/marlin/marlin_nvfp4.h" +#include "tensorrt_llm/thop/thUtils.h" +#include +#include +#include + +using torch::Tensor; + +// Marlin tile constants (must match marlin.cuh) +static constexpr int kMarlinTileSize = 16; +static constexpr int kMarlinTileKSize = kMarlinTileSize; +static constexpr int kMarlinTileNSize = kMarlinTileKSize * 4; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Repack quantized weights from row-major to Marlin tiled format. +// +// b_q_weight: [K/pack_factor, N] int32 row-major packed weights +// perm: [K] int32 permutation (empty for no-perm) +// size_k: reduction dimension K +// size_n: output dimension N +// num_bits: quantization bits (4 for FP4) +Tensor gptq_marlin_repack( + Tensor& b_q_weight, Tensor& perm, int64_t size_k, int64_t size_n, int64_t num_bits, bool is_a_8bit = false) +{ + TORCH_CHECK( + size_k % kMarlinTileKSize == 0, "size_k = ", size_k, " not divisible by tile_k_size = ", kMarlinTileKSize); + TORCH_CHECK( + size_n % kMarlinTileNSize == 0, "size_n = ", size_n, " not divisible by tile_n_size = ", kMarlinTileNSize); + TORCH_CHECK(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); + + int const pack_factor = 32 / static_cast(num_bits); + + TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), ", size_k = ", size_k, + ", pack_factor = ", pack_factor); + TORCH_CHECK(b_q_weight.size(1) == size_n, "b_q_weight.size(1) = ", b_q_weight.size(1), " != size_n = ", size_n); + + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); + auto options = torch::TensorOptions().dtype(b_q_weight.dtype()).device(b_q_weight.device()); + torch::Tensor out = torch::empty({size_k / kMarlinTileSize, size_n * kMarlinTileSize / pack_factor}, options); + + bool has_perm = perm.size(0) != 0; + + uint32_t const* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); + uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(b_q_weight.get_device()); + + ::marlin_nvfp4::gptq_marlin_repack_dispatch(b_q_weight_ptr, perm_ptr, out_ptr, static_cast(size_k), + static_cast(size_n), static_cast(num_bits), has_perm, is_a_8bit, stream); + + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, int size_k, int size_n," + " int num_bits, bool is_a_8bit=False) -> (Tensor out)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("gptq_marlin_repack", &tensorrt_llm::torch_ext::gptq_marlin_repack); +} diff --git a/examples/configs/curated/nemotron-super-marlin.yaml b/examples/configs/curated/nemotron-super-marlin.yaml new file mode 100644 index 000000000000..bba3f28dc128 --- /dev/null +++ b/examples/configs/curated/nemotron-super-marlin.yaml @@ -0,0 +1,12 @@ +max_batch_size: 8 +max_seq_len: 66048 + +tensor_parallel_size: 2 +pipeline_parallel_size: 1 +moe_expert_parallel_size: 2 +context_parallel_size: 1 + +nvfp4_gemm_config: + allowed_backends: [marlin] +moe_config: + backend: MARLIN diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 77fc6f71eeb8..b61ae74d23c0 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1125,6 +1125,61 @@ def _(mat_a: torch.Tensor, n = mat_b.shape[0] return mat_a.new_empty((m, n), dtype=out_dtype) + @torch.library.register_fake("trtllm::marlin_nvfp4_gemm") + def _(mat_a: torch.Tensor, + mat_b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + alpha: torch.Tensor, + weight_global_scale: torch.Tensor, + bias: Optional[torch.Tensor], + out_dtype: Optional[torch.dtype], + size_n: int, + size_k: int, + output_buffer_kind: int = 0, + group: Optional[List[int]] = None): + # mat_a: [M, K/2] FP4 packed (or BF16 when W4A16) + # mat_b: Marlin-packed weights + # Output: [M, size_n] with dtype=out_dtype + m = mat_a.shape[0] + return mat_a.new_empty((m, size_n), dtype=out_dtype) + + @torch.library.register_fake("trtllm::marlin_nvfp4_moe_gemm") + def _(a: torch.Tensor, + b_q_weight: torch.Tensor, + b_scales: torch.Tensor, + global_scale: torch.Tensor, + workspace: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_past_padded: torch.Tensor, + topk_weights: torch.Tensor, + moe_block_size: int, + top_k: int, + mul_topk_weights: bool, + size_n: int, + size_k: int, + out_dtype: Optional[torch.dtype], + use_fp32_reduce: bool = False): + # a: [M, K] BF16, b_q_weight: [num_experts, ...] Marlin-packed FP4 + # Output: [M * top_k, size_n] + m = a.shape[0] + dtype = out_dtype if out_dtype is not None else torch.bfloat16 + return a.new_empty((m * top_k, size_n), dtype=dtype) + + @torch.library.register_fake("trtllm::gptq_marlin_repack") + def _(b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False): + pack_factor = 32 // num_bits + tile_size = 16 + return b_q_weight.new_empty( + (size_k // tile_size, size_n * tile_size // pack_factor), + dtype=b_q_weight.dtype) + @torch.library.register_fake("trtllm::mla_rope_generation") def _( fused_q: torch.Tensor, diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 4ad69ecf21d2..fc7cc4b3781d 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -817,6 +817,119 @@ def forward( return result +class MarlinNVFP4Runner(TunableRunner): + """Marlin-based NVFP4 GEMM for SM90 (Hopper). + + Weights are eagerly repacked to Marlin tiled format during + ``get_valid_tactics`` (before CUDA graph capture) so that ``forward()`` does + not allocate any memory. + """ + + tuning_config = TuningConfig() # single tactic, no tuning + + MIN_SM_VERSION = 90 + MAX_SM_VERSION = 99 # SM90-series only (Hopper) + NVFP4_SCALE_VECTOR_SIZE = 16 + + def __init__(self, output_buffer_kind: int, output_dtype: torch.dtype): + super().__init__() + self.output_buffer_kind = int(output_buffer_kind) + self.output_dtype = output_dtype + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + if not torch.cuda.is_available(): + return [] + capability = torch.cuda.get_device_capability(torch.device('cuda:0')) + sm_version = capability[0] * 10 + capability[1] + if sm_version < self.MIN_SM_VERSION or sm_version > self.MAX_SM_VERSION: + return [] + + # Eagerly prepare Marlin weights so that forward() never allocates + # memory (safe for CUDA graph capture). + _, weight, _, weight_scale, _ = inputs + self._prepare_marlin_weights(weight, weight_scale) + + return [0] + + @classmethod + def _prepare_marlin_weights( + cls, weight: torch.Tensor, weight_scale: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Convert raw NVFP4 weights + scales to Marlin format.""" + from tensorrt_llm.math_utils import pad_up + from tensorrt_llm.quantization.utils import marlin_utils + + assert torch.iinfo(weight.dtype).bits == 8 + # weight: [N, K/2] FP4 packed as int8/uint8; view as int32 repacks to + # the int32-based Marlin tiled layout. + size_n = weight.shape[0] + size_k = weight.shape[1] * 2 + + qweight_int32 = weight.view(torch.int32) + qweight_int32 = qweight_int32.T.contiguous() + perm = torch.empty(0, dtype=torch.int32, device=weight.device) + marlin_weight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight_int32, + perm=perm, + size_k=size_k, + size_n=size_n, + num_bits=4, + is_a_8bit=False, + ) + + num_groups = size_k // cls.NVFP4_SCALE_VECTOR_SIZE + n_padded = pad_up(size_n, 128) + scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + weight_scale.view(n_padded, -1)) + scale_2d = scale_unswizzled[:size_n, :num_groups].view( + torch.float8_e4m3fn).T.contiguous() + marlin_scale = marlin_utils.marlin_permute_scales( + scale_2d.to(torch.half), + size_k, + size_n, + group_size=cls.NVFP4_SCALE_VECTOR_SIZE) + marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) + + marlin_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( + torch.tensor(1.0, dtype=torch.bfloat16, device=weight.device)) + + return marlin_weight, marlin_scale, marlin_global_scale, size_n, size_k + + def forward( + self, + /, + inputs: List[torch.Tensor], + tactic: int = -1, + do_preparation: bool = False, + **kwargs, + ) -> torch.Tensor: + act_fp4, weight, act_sf, weight_scale, alpha = inputs + + (marlin_weight, marlin_scale, marlin_global_scale, size_n, + size_k) = self._prepare_marlin_weights(weight, weight_scale) + + m = act_fp4.shape[0] + m_padded = (m + 128 - 1) // 128 * 128 + act_sf_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + act_sf.view(m_padded, -1)).flatten() + + result = torch.ops.trtllm.marlin_nvfp4_gemm( + act_fp4, + marlin_weight, + scale_a=act_sf_unswizzled, + scale_b=marlin_scale, + alpha=alpha, + weight_global_scale=marlin_global_scale, + bias=None, + out_dtype=self.output_dtype, + size_n=size_n, + size_k=size_k, + output_buffer_kind=self.output_buffer_kind, + ) + return result + + @torch.library.custom_op("trtllm::nvfp4_gemm_cublaslt", mutates_args=()) def nvfp4_gemm_cublaslt( act_fp4: torch.Tensor, @@ -977,6 +1090,22 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], tactics = [] act_fp4, weight, act_sf, weight_scale, alpha = inputs + # Add Marlin tactics (SM90 Hopper only) — users must opt-in explicitly + # by listing "marlin" in ``allowed_backends``. + if self._is_backend_allowed("marlin"): + marlin_runner = MarlinNVFP4Runner(self.output_buffer_kind, + self.output_dtype) + marlin_tactics = marlin_runner.get_valid_tactics(inputs, profile) + if marlin_tactics: + tactics.extend([("marlin", tactic) + for tactic in marlin_tactics]) + elif self._is_only_backend("marlin"): + sm_version = get_sm_version() + raise ValueError( + f"Marlin backend requires SM 90-99 (Hopper), but got SM " + f"{sm_version}. Please add other backends to " + "allowed_backends.") + # Add CUDA Core tactics if available if self._is_backend_allowed("cuda_core"): is_cuda_core_supported = False @@ -1086,12 +1215,17 @@ def forward( ) -> torch.Tensor: # Handle fallback tactic on cache miss if tactic == -1: - # Prefer cutlass as fallback if available, otherwise use first valid backend + # Prefer marlin on Hopper (SM90) when explicitly allowed, cutlass + # otherwise, falling back to whatever backend is available. assert len( self.allowed_backends) > 0, "No allowed backends available" - tactic = ("cutlass", - -1) if "cutlass" in self.allowed_backends else ( - self.allowed_backends[0], -1) + sm_version = get_sm_version() + if "marlin" in self.allowed_backends and 90 <= sm_version <= 99: + tactic = ("marlin", -1) + elif "cutlass" in self.allowed_backends: + tactic = ("cutlass", -1) + else: + tactic = (self.allowed_backends[0], -1) backend, sub_tactic = tactic if backend == "cuda_core": @@ -1119,6 +1253,10 @@ def forward( self.group)(inputs, tactic=sub_tactic, bias=bias) + elif backend == "marlin": + return MarlinNVFP4Runner(self.output_buffer_kind, + self.output_dtype)(inputs, + tactic=sub_tactic) else: raise ValueError(f"Invalid tactic: {tactic}") @@ -1143,6 +1281,7 @@ def nvfp4_gemm( - cuBLASLt: Heuristic-based algorithms from cuBLASLt library - CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid) - CUDA Core: CUDA Core implementation (requires SM >= 100 and M <= 8) + - Marlin: Hopper W4A16 NVFP4 implementation (requires SM 90-99) The AutoTuner profiles all available backends during the first run and caches the best choice for each input shape. Subsequent calls use the cached selection @@ -1168,7 +1307,9 @@ def nvfp4_gemm( ValueError: If backend is invalid/unavailable """ - valid_individual_backends = {'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'} + valid_individual_backends = { + 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin' + } # Parse comma-separated string to list backends_list = [ diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index 852d8cb4b5d3..3f8013564c1a 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -94,6 +94,11 @@ def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor: w = w.to(torch.float32) new_weights[key] = w elif "mixer.in_proj" in key: + # Restrict the mamba2 in_proj split to the actual weight tensor. + # NVFP4 checkpoints attach companion tensors (``input_scale``, + # ``weight_scale``, ``weight_scale_2``, …) under ``mixer.in_proj.*`` + # — those are scalars / 1-D scales and must not go through the + # Mamba2 split rearrangement. new_weights[key] = _split_mamba2_mixer_in_proj(weights[name]) elif "conv1d" in key: w = weights[name] diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index b10d90a3b759..3362b5e6c8c9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -804,12 +804,6 @@ def _force_moe_backend_for_w4a16_on_hopper( if model_config.moe_backend.upper() in ('CUTLASS', 'AUTO'): return - logger.warning( - f"Nemotron-H SM{get_sm_version()}: forcing moe_backend " - f"'{model_config.moe_backend}' -> 'CUTLASS' for W4A16 fallback") - model_config._frozen = False - model_config.moe_backend = 'CUTLASS' - model_config._frozen = True @contextmanager diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index cc0694728829..4b401c575c9c 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -28,8 +28,8 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - is_torch_compiling, maybe_compiled_cat, - maybe_compiled_copy_) + is_nvfp4_marlin_enabled, is_torch_compiling, + maybe_compiled_cat, maybe_compiled_copy_) from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from .multi_stream_utils import maybe_execute_in_parallel from .rms_norm import RMSNorm @@ -542,6 +542,8 @@ def __init__( attn_cls = get_attention_backend( self.attn_backend, sparse_attention_config=sparse_attn_cfg) + self.is_marlin_enabled: bool = is_nvfp4_marlin_enabled() + # These two modules are mutually exclusive - either splitted_qkv_lora or fused_qkv_lora will be used, # but never both at the same time. splitted_qkv_lora handles Q,K,V separately while fused_qkv_lora # handles them as a single fused operation. @@ -662,7 +664,7 @@ def _use_quantize_output(self): self.o_proj, 'pre_quant_scale') and self.o_proj.pre_quant_scale is not None - return self.has_quant_scale and not self.attn_output_gate and not has_awq_pre_quant_scale + return self.has_quant_scale and not self.attn_output_gate and not has_awq_pre_quant_scale and not self.is_marlin_enabled def create_output(self, q: torch.Tensor, attn_metadata: AttentionMetadata, mask_type: str): @@ -821,7 +823,8 @@ def forward_impl( use_custom_inplace_op = (self.register_to_config and (self.attn_backend == "TRTLLM" or self.attn_backend == "FLASHINFER") - and is_torch_compiling()) + and is_torch_compiling() + and not self.is_marlin_enabled) if use_custom_inplace_op: outputs = create_attn_outputs(q, attention_mask, self.layer_idx_str) diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index 556345b2029f..e8b9b7af88c1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -150,6 +150,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (auto-promoted when flashinfer is importable) | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | +| `fused_moe_marlin.py` | `MarlinFusedMoE` | SM90 only | W4A16 NVFP4 on Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); no dynamic EPLB | `EXTERNAL_COMM` | | `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | | `fused_moe_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecating — use ConfigurableMoE instead | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | @@ -197,19 +198,19 @@ is available. Each backend's `can_implement(quant_algo, dtype_activation, swiglu_gptoss_style, ...)` method declares supported quantizations. Source of truth: the `can_implement` classmethod in each backend file. -| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | WideEP | Vanilla | -|---|---|---|---|---|---|---|---|---|---|---| -| Unquantized (BF16/FP16) | Y (SM80+) | N | N | N | N | N | N | Y (SM90, BF16) | Y | Y | -| FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | Y | Y | -| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | Y (SM100/103) | N | N | N | Y | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y | Y | -| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | -| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | -| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | -| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | Y (SM100/103, requires `hidden_size % 512 == 0`) | N | N | N | N | -| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | N | N | N | -| W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | N | -| INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | N | Y | N | +| Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | WideEP | Vanilla | +|---|---|---|---|---|---|---|---|---|---|---|---| +| Unquantized (BF16/FP16) | Y (SM80+) | N | N | N | N | N | N | Y (SM90, BF16) | N | Y | Y | +| FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | Y | +| FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | Y (SM100/103) | N | N | N | N | Y | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM90, W4A16) | Y | Y | +| W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | N | +| W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | +| W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | +| W4A8 MXFP4 MXFP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | Y (SM100/103, requires `hidden_size % 512 == 0`) | N | N | N | N | N | +| W4A8 AWQ | Y (SM89/90) | N | N | N | N | N | N | N | N | N | N | +| W8A16 | Y (SM80+) | N | N | N | N | N | N | N | N | N | N | +| INT4 WoQ (W4AFP8) | N | N | N | N | N | N | N | N | N | Y | N | ### Scheduler / EPLB Constraints diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index 931ea93e30a4..f578feb247bf 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -3,6 +3,7 @@ from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from .fused_moe_cutlass import CutlassFusedMoE +from .fused_moe_marlin import MarlinFusedMoE from .fused_moe_triton import TritonFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE @@ -37,6 +38,7 @@ "FusedMoEQuantScalesFP8", "get_moe_cls", "Llama4RenormalizeMoeRoutingMethod", + "MarlinFusedMoE", "LoadBalancedMoeRoutingMethod", "moe_load_balancer_set_repeated_for_next_layer", "MoE", diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index affac8e417db..f7b713e3de61 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -17,6 +17,7 @@ from .fused_moe_cutlass import CutlassFusedMoE from .fused_moe_deepgemm import DeepGemmFusedMoE from .fused_moe_densegemm import DenseGEMMFusedMoE +from .fused_moe_marlin import MarlinFusedMoE from .fused_moe_triton import TritonFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .fused_moe_vanilla import VanillaMoE @@ -62,6 +63,12 @@ def get_moe_cls( if override_quant_config is not None: quant_config = override_quant_config layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" + if moe_backend.upper() == "MARLIN": + # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Require nvfp4 + # quantization explicitly so a misconfigured model fails fast. + if quant_config is None or not quant_config.quant_mode.has_nvfp4(): + raise ValueError("MarlinFusedMoE only supports NVFP4 quantization.") + return MarlinFusedMoE if moe_backend.upper() == "CUTLASS": return CutlassFusedMoE elif moe_backend.upper() == "VANILLA": @@ -359,7 +366,7 @@ def create_moe_backend( without_comm=without_comm, activation_type=activation_type, ) - elif moe_cls is CutlassFusedMoE: + elif moe_cls in (CutlassFusedMoE, MarlinFusedMoE): # CuteDslFusedMoE, DeepGemmFusedMoE, and CuteDslB12xFusedMoE # also subclass CutlassFusedMoE but have narrower constructors, so # they take their own branches below. @@ -593,7 +600,7 @@ def create_moe( CuteDslB12xFusedMoE): if moe_cls in (DeepGemmFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, CuteDslB12xFusedMoE, CutlassFusedMoE, DenseGEMMFusedMoE, - MegaMoEDeepGemm, MegaMoECuteDsl): + MegaMoEDeepGemm, MegaMoECuteDsl, MarlinFusedMoE): return ConfigurableMoE( routing_method=routing_method, num_experts=num_experts, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py new file mode 100644 index 000000000000..e6e9f3c7ea7c --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Marlin-based MoE backend for NVFP4 on SM90 (Hopper). + +Uses a fused ``marlin_nvfp4_moe_gemm`` CUDA kernel that processes ALL experts +in a single launch. W4A16 approach: BF16 activations + FP4 weights, +dequantized FP4→BF16 in registers, using BF16 m16n8k16 MMA. +No activation quantization overhead. In-kernel topk_weights multiplication. +""" + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.models.modeling_utils import QuantAlgo + +from ...utils import ActivationType, is_gated_activation, relu2 +from .fused_moe_cutlass import CutlassFusedMoE +from .interface import _warn_and_return +from .quantization import NVFP4MarlinFusedMoEMethod + +# Block size for moe_align_block_size — must match TILE_M in the kernel +_MOE_BLOCK_SIZE = 16 + + +def _has_fused_moe_kernel() -> bool: + """Check if the fused marlin_nvfp4_moe_gemm op is available.""" + return hasattr(torch.ops.trtllm, "marlin_nvfp4_moe_gemm") + + +class MarlinFusedMoE(CutlassFusedMoE): + """MoE backend using Marlin W4A16 NVFP4 GEMM for SM90 (Hopper). + + Uses ``marlin_nvfp4_moe_gemm`` with BF16 activations to process all experts + in a single kernel launch via sorted token dispatch. In-kernel topk_weights + multiplication eliminates separate scatter-weight step. CUDA-graph + compatible. Requires the fused kernel to be built (no fallback path). + """ + + _QUANT_SUPPORT_TABLE = { + QuantAlgo.NVFP4: { + "sm_constraint": ("in", {90}), + "dtypes": {torch.bfloat16}, + }, + } + + @classmethod + def can_implement( + cls, + quant_algo: Optional[QuantAlgo], + dtype_activation: torch.dtype = torch.bfloat16, + swiglu_gptoss_style: bool = False, + ) -> Tuple[bool, Optional[str]]: + sm_version = get_sm_version() + + if quant_algo != QuantAlgo.NVFP4: + return _warn_and_return( + f"MarlinFusedMoE only supports NVFP4 (got quant_algo={quant_algo})" + ) + + if sm_version != 90: + return _warn_and_return( + f"MarlinFusedMoE only supports SM90 (Hopper), got SM{sm_version}" + ) + + if swiglu_gptoss_style: + return _warn_and_return("MarlinFusedMoE does not support swiglu_gptoss_style") + + if dtype_activation != torch.bfloat16: + return _warn_and_return( + f"MarlinFusedMoE W4A16 requires bfloat16 activations, got {dtype_activation}" + ) + + return True, None + + def quantize_input( + self, x: torch.Tensor | Fp4QuantizedTensor, post_quant_comm: bool = True, **kwargs + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + return x, None + + def _get_quant_method(self): + if self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4(): + assert self.moe_backend == "MARLIN", ( + "MarlinFusedMoE only supports NVFP4, got {self.moe_backend}" + ) + return NVFP4MarlinFusedMoEMethod() + raise ValueError(f"MarlinFusedMoE only supports NVFP4, got {self.quant_config}") + + def _supports_load_balancer(self) -> bool: + return False + + def validate_configurable_moe(self, moe) -> None: + """Reject configs that require external-communication MoE dispatch. + + Marlin is W4A16 (``quantize_input`` produces no activation scale) and + routes internally inside ``run_moe``. The host-side all-to-all + dispatch/combine path that ConfigurableMoE uses for attention-DP + expert parallelism needs the routing decided *before* dispatch and a + per-token scale payload, so it is incompatible with Marlin and fails + inside ``moe_a2a_dispatch``. + + ConfigurableMoE only creates an external communication strategy when + ``enable_attention_dp and dp_size > 1`` (see CommunicationFactory); + ``moe.comm`` is not assigned yet when this hook runs, so check that + same condition via the mapping. Single-GPU, and TP/EP without + attention DP, are supported. + """ + if moe.use_dp and moe.mapping.dp_size > 1: + raise ValueError( + "MarlinFusedMoE does not support external-communication MoE " + "(attention data parallelism combined with expert " + "parallelism): its W4A16 layout and internal routing are " + "incompatible with the all-to-all dispatch path. Use Marlin " + "single-node, or with TP/EP without attention DP." + ) + + def _apply_activation(self, gemm1_out: torch.Tensor) -> torch.Tensor: + """Apply the activation function to the gemm1 output. + + For gated activations (SwiGLU, GeGLU), gemm1_out has shape + [tokens, 2 * intermediate_size] — split into gate + up. + For non-gated activations (Relu2, Relu), gemm1_out has shape + [tokens, intermediate_size] — apply element-wise. + + Returns [tokens, intermediate_size] in all cases. + """ + inter_size = self.intermediate_size_per_partition + if is_gated_activation(self.activation_type): + gate = gemm1_out[:, :inter_size] + up = gemm1_out[:, inter_size : 2 * inter_size] + if self.activation_type == ActivationType.Geglu: + return F.gelu(gate) * up + else: + return F.silu(gate) * up # SwiGLU + else: + if self.activation_type == ActivationType.Relu2: + return relu2(gemm1_out) + else: + return F.relu(gemm1_out) + + def _ensure_workspace(self, device: torch.device): + """Lazily allocate workspace tensor for Marlin kernel.""" + if not hasattr(self, "_marlin_workspace") or self._marlin_workspace is None: + props = torch.cuda.get_device_properties(device) + sms = props.multi_processor_count + max_blocks_per_sm = 4 + self._marlin_workspace = torch.zeros( + sms * max_blocks_per_sm, dtype=torch.int32, device=device + ) + return self._marlin_workspace + + # ==================================================================== + # Main entry point + # ==================================================================== + + def run_moe( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + x_sf: Optional[torch.Tensor] = None, + is_sf_swizzled: bool = True, + output_dtype: Optional[torch.dtype] = None, + tuner_num_tokens: Optional[int] = None, + tuner_top_k: Optional[int] = None, + moe_output: Optional[torch.Tensor] = None, + enable_alltoall: Optional[bool] = None, + router_logits: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + assert output_dtype is None or output_dtype == torch.bfloat16 + assert _has_fused_moe_kernel(), ( + "marlin_nvfp4_moe_gemm is not available. Rebuild TensorRT-LLM " + "with the fused Marlin MoE kernel for NVFP4." + ) + assert x.dtype == torch.bfloat16 + + output_dtype = torch.bfloat16 + + if token_selected_experts is None: + assert router_logits is not None, ( + "MarlinFusedMoE.run_moe needs token_selected_experts or router_logits" + ) + token_selected_experts, token_final_scales = self.routing_method.apply(router_logits) + + num_tokens = x.shape[0] + top_k = token_selected_experts.shape[1] + + local_n = self.expert_size_per_partition + if local_n != self.num_experts: + slot_start = self.slot_start + is_local = (token_selected_experts >= slot_start) & ( + token_selected_experts < slot_start + local_n + ) + token_selected_experts = (token_selected_experts - slot_start).clamp(0, local_n - 1) + if token_final_scales is None: + token_final_scales = torch.ones( + num_tokens, top_k, dtype=torch.float32, device=x.device + ) + token_final_scales = token_final_scales * is_local.to(token_final_scales.dtype) + num_experts = local_n + + workspace = self._ensure_workspace(x.device) # [num_sms * max_blocks_per_sm(4)] int32 + + # Step 1: Sort tokens by expert assignment + topk_ids = token_selected_experts.to(torch.int32).contiguous() + max_num_tokens_padded = num_tokens * top_k + num_experts * _MOE_BLOCK_SIZE + + sorted_token_ids = torch.empty(max_num_tokens_padded, dtype=torch.int32, device=x.device) + expert_ids_out = torch.empty( + (max_num_tokens_padded + _MOE_BLOCK_SIZE - 1) // _MOE_BLOCK_SIZE, + dtype=torch.int32, + device=x.device, + ) + num_tokens_post_pad = torch.empty(1, dtype=torch.int32, device=x.device) + + torch.ops.trtllm.moe_align_block_size( + topk_ids, + num_experts, + _MOE_BLOCK_SIZE, + sorted_token_ids, + expert_ids_out, + num_tokens_post_pad, + ) + + # Prepare topk_weights for in-kernel multiplication + if token_final_scales is not None: + topk_weights = token_final_scales.float().contiguous() + else: + topk_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=x.device) + + hidden_size = x.shape[1] + k1 = hidden_size + n1 = self.expand_intermediate_size_per_partition + + # Step 2: Fused gemm1 — ALL experts in ONE kernel launch (W4A16) + gemm1_out = torch.ops.trtllm.marlin_nvfp4_moe_gemm( + x.contiguous(), + self.w3_w1_weight, + b_scales=self.w3_w1_weight_scale, + global_scale=self.fc31_alpha, + workspace=workspace, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids_out, + num_tokens_past_padded=num_tokens_post_pad, + topk_weights=topk_weights, + moe_block_size=_MOE_BLOCK_SIZE, + top_k=top_k, + mul_topk_weights=False, # Don't multiply weights in gemm1 + size_n=n1, + size_k=k1, + out_dtype=output_dtype, + use_fp32_reduce=False, + ) # [num_tokens * top_k, expand_intermediate] + + # Step 3: Activation (element-wise) + hidden = self._apply_activation(gemm1_out) + + k2 = self.intermediate_size_per_partition + n2 = self.unpadded_hidden_size + + # Step 4: Fused gemm2 — ALL experts in ONE kernel launch (W4A16) + # hidden is [num_tokens * top_k, intermediate_size]. Each row is an + # independent token-expert pair. We re-sort with top_k=1 so the + # kernel maps each row to the correct expert without expanding again. + num_tokens_gemm2 = num_tokens * top_k + + # Build per-row expert assignment from the original topk_ids + gemm2_topk_ids = topk_ids.reshape(-1, 1)[:num_tokens_gemm2].contiguous() + + max_padded_g2 = num_tokens_gemm2 + num_experts * _MOE_BLOCK_SIZE + sorted_ids_g2 = torch.empty(max_padded_g2, dtype=torch.int32, device=x.device) + expert_ids_g2 = torch.empty( + (max_padded_g2 + _MOE_BLOCK_SIZE - 1) // _MOE_BLOCK_SIZE, + dtype=torch.int32, + device=x.device, + ) + num_post_pad_g2 = torch.empty(1, dtype=torch.int32, device=x.device) + + torch.ops.trtllm.moe_align_block_size( + gemm2_topk_ids, + num_experts, + _MOE_BLOCK_SIZE, + sorted_ids_g2, + expert_ids_g2, + num_post_pad_g2, + ) + + # topk_weights for gemm2: flatten to [num_tokens*top_k, 1] for top_k=1 + topk_weights_g2 = topk_weights.reshape(-1, 1)[:num_tokens_gemm2].contiguous() + + gemm2_out = torch.ops.trtllm.marlin_nvfp4_moe_gemm( + hidden.contiguous(), + self.w2_weight, + b_scales=self.w2_weight_scale, + global_scale=self.fc2_alpha, + workspace=workspace, + sorted_token_ids=sorted_ids_g2, + expert_ids=expert_ids_g2, + num_tokens_past_padded=num_post_pad_g2, + topk_weights=topk_weights_g2, + moe_block_size=_MOE_BLOCK_SIZE, + top_k=1, + mul_topk_weights=True, + size_n=n2, + size_k=k2, + out_dtype=output_dtype, + use_fp32_reduce=False, + ) # [num_tokens_gemm2, hidden_size] + + # Step 5: Scatter-reduce — sum weighted expert outputs. + # gemm2_out rows correspond to flattened (token_idx * top_k + k) pairs. + gemm2_out = gemm2_out[:num_tokens_gemm2, : self.unpadded_hidden_size] + + # Map each row back to its original token index + row_indices = torch.arange(num_tokens_gemm2, device=x.device) + orig_tokens = row_indices // top_k + + final_hidden_states = torch.zeros( + (num_tokens, self.unpadded_hidden_size), + dtype=output_dtype, + device=x.device, + ) + final_hidden_states.index_add_(0, orig_tokens, gemm2_out.to(output_dtype)) + + return final_hidden_states diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index b9ed4af37c01..5acfb2d5d66b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -57,6 +57,7 @@ from .fused_moe_cutlass import CutlassFusedMoE from .fused_moe_deepgemm import DeepGemmFusedMoE from .fused_moe_densegemm import DenseGEMMFusedMoE +from .fused_moe_marlin import MarlinFusedMoE from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .interface import MoESchedulerKind @@ -803,6 +804,9 @@ def _get_backend_kwargs( all_rank_num_tokens=all_rank_num_tokens, output_dtype=output_dtype ) + elif moe.backend.__class__ == MarlinFusedMoE: + kwargs["router_logits"] = router_logits + return kwargs diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 368c5b1a761c..d7167b16151f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -2822,7 +2822,7 @@ def load_expert_w3_w1_weight(self, if not hasattr(module, 'tmp_cutlass_w3_w1_weights'): module.tmp_cutlass_w3_w1_weights = {} assert expert_idx >= 0, "expert_idx must be provided for stable dict key" - dst_base = dst_w3_w1_weight.storage().data_ptr() + dst_base = dst_w3_w1_weight.untyped_storage().data_ptr() dict_key = (dst_base, expert_idx) expert_entry = module.tmp_cutlass_w3_w1_weights.setdefault(dict_key, {}) expert_entry['dst'] = dst_w3_w1_weight @@ -2911,6 +2911,113 @@ def process_weights_after_loading(self, module: torch.nn.Module): super().process_weights_after_loading(module) +class NVFP4MarlinFusedMoEMethod(NVFP4CutlassFusedMoEMethod): + """NVFP4 MoE quantization method for the Marlin backend. + + Inherits weight loading from the CUTLASS method, then transforms the loaded + weights to Marlin tiled format in ``post_load_weights``. + + The Marlin kernel is W4A16 (BF16 activations, no activation quantization), + so global_scale must be the raw ``weight_scale_2`` — not the CUTLASS alpha + which folds in ``input_scale``. We intercept alpha loading to save the + raw ``weight_scale_2`` values. + """ + + # Marlin's ``post_load_weights`` repacks weights into Marlin tiled format + # and rebuilds the module parameters, which is incompatible with dynamic + # EPLB weight migration. + eplb_support_status = EplbSupportStatus.NOT_SUPPORTED + + def load_expert_fc31_alpha_nvfp4(self, w1_weight_scale_2, w3_weight_scale_2, + final_fc31_input_scale, dst_fc31_alpha): + # Store raw weight_scale_2 for Marlin (W4A16: no input_scale needed). + w1_ws2 = w1_weight_scale_2[...].reshape([]) + dst_fc31_alpha.copy_(w1_ws2) + + def load_expert_fc2_alpha_nvfp4(self, w2_weight_scale_2, + final_fc2_input_scale, dst_w2_alpha): + w2_ws2 = w2_weight_scale_2[...].reshape([]) + dst_w2_alpha.copy_(w2_ws2) + + def post_load_weights(self, module): + """Transform CUTLASS-format NVFP4 weights to Marlin tiled format.""" + from tensorrt_llm.quantization.utils import marlin_utils + + # Standard CUTLASS loading (swizzles scales, computes alpha, etc.) + super().post_load_weights(module) + + num_experts = module.expert_size_per_partition + hidden_size = module.hidden_size + intermediate_size = module.intermediate_size_per_partition + is_act_and_mul = module.intermediate_size_expand_ratio == 2 + group_size = module.scaling_vector_size # 16 + + # Actual (unpadded) dimensions + N1 = intermediate_size * (2 if is_act_and_mul else 1) + K1 = hidden_size + N2 = hidden_size + K2 = intermediate_size + + def unswizzle_scales(scale_3d, N_actual, K_actual): + """Unswizzle packed int32 scales -> FP8 [num_experts, N, num_groups].""" + num_groups = K_actual // group_size + result = [] + for i in range(num_experts): + # [N_padded, K//64] int32 -> float4_sf_dtype -> reverse -> FP8 + s_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + scale_3d[i].view(float4_sf_dtype)) + s_fp8 = s_unswizzled.view( + torch.float8_e4m3fn)[:N_actual, :num_groups] + result.append(s_fp8.unsqueeze(0)) + return torch.cat(result, 0) + + w13_scale = unswizzle_scales(module.w3_w1_weight_scale, N1, K1) + w2_scale = unswizzle_scales(module.w2_weight_scale, N2, K2) + w13_weight = module.w3_w1_weight.view(torch.uint8)[:, :N1, :K1 // + 2].contiguous() + w2_weight = module.w2_weight.view(torch.uint8)[:, :N2, :K2 // + 2].contiguous() + + w13_gs = module.fc31_alpha.data.clone() + w2_gs = module.fc2_alpha.data.clone() + + (w13, w13_s, w13_gs, w2, w2_s, + w2_gs) = marlin_utils.prepare_nvfp4_moe_weights_for_marlin( + w13=w13_weight, + w13_scale=w13_scale, + w13_global_scale=w13_gs, + w2=w2_weight, + w2_scale=w2_scale, + w2_global_scale=w2_gs, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size, + num_experts=num_experts, + is_act_and_mul=is_act_and_mul, + param_dtype=torch.bfloat16, + ) + + for name in ( + "w3_w1_weight", + "w2_weight", + "w3_w1_weight_scale", + "w2_weight_scale", + "fc31_alpha", + "fc2_alpha", + ): + getattr(module, name).data.untyped_storage().resize_(0) + + module.w3_w1_weight = nn.Parameter(w13.view(torch.int64), + requires_grad=False) + module.w3_w1_weight_scale = nn.Parameter(w13_s.view(torch.int32), + requires_grad=False) + module.fc31_alpha = nn.Parameter(w13_gs, requires_grad=False) + module.w2_weight = nn.Parameter(w2.view(torch.int64), + requires_grad=False) + module.w2_weight_scale = nn.Parameter(w2_s.view(torch.int32), + requires_grad=False) + module.fc2_alpha = nn.Parameter(w2_gs, requires_grad=False) + + class W4A16NVFP4CutlassFusedMoEMethod(NVFP4CutlassFusedMoEMethod): """W4A16 dequant-on-the-fly variant of NVFP4 MoE for SM<100. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index aae3f3a65ab2..e15dc0cc7fd8 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -31,6 +31,7 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, + is_nvfp4_marlin_enabled, replace_parameter_and_save_metadata, unswizzle_sf) @@ -2792,6 +2793,124 @@ def apply(self, module: Linear, input: torch.Tensor, return output +class MarlinNVFP4LinearMethod(NVFP4LinearMethod): + """NVFP4 Linear method backed by the Marlin W4A16 kernel (Hopper only).""" + + def post_load_weights(self, module: Linear): + from tensorrt_llm.quantization.utils import marlin_utils + + weight = module.weight.data + weight_scale = module.weight_scale.data + size_n = module.out_features + size_k = module.in_features + group_size = module.scaling_vector_size # 16 + + assert size_k % group_size == 0, ( + f"size_k {size_k} must be divisible by group_size {group_size}") + + size_k_pad = fp4_utils.pad_up(size_k, 64) + size_n_pad = fp4_utils.pad_up(size_n, 128) + + num_groups = size_k // group_size + n_padded = size_n_pad + scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + weight_scale.view(n_padded, -1)) + # [size_n, num_groups] block scales; uint8 storage (reverse interleave), + # reinterpreted as E4M3 after any padding. Pad in uint8 since F.pad does + # not support float8. + scale_2d = scale_unswizzled[:size_n, :num_groups] + + if size_k_pad != size_k or size_n_pad != size_n: + num_groups_pad = size_k_pad // group_size + # weight: [N, K/2] uint8 -> [N_pad, K_pad/2] (FP4 zero == 0.0) + weight = F.pad(weight, + (0, + (size_k_pad - size_k) // 2, 0, size_n_pad - size_n)) + # scales: [N, num_groups] -> [N_pad, num_groups_pad]. + # The Marlin S0E5M3 fast-dequant is NOT zero-safe: a zero scale on + # a (zero-weight) padded K-group still corrupts that tile's output. + # Since K is the contraction dim, one bad group-scale poisons every + # output row, so padded K-groups must carry a valid non-zero fp8 + # scale -- use the smallest-normal e4m3 value (0x08), matching the + # quantizer's own zero-block scale. N-row padding is sliced off in + # ``apply`` and can stay zero. + scale_2d = F.pad(scale_2d, (0, num_groups_pad - num_groups), + value=0x08) + scale_2d = F.pad(scale_2d, (0, 0, 0, size_n_pad - size_n), value=0) + + qweight_int32 = weight.view( + torch.int32).T.contiguous() # [K_pad/4, N_pad] + perm = torch.empty(0, dtype=torch.int32, device=weight.device) + marlin_weight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight_int32, + perm=perm, + size_k=size_k_pad, + size_n=size_n_pad, + num_bits=4, + is_a_8bit=False, + ) + + scale_2d = scale_2d.view( + torch.float8_e4m3fn).T.contiguous() # [num_groups_pad, N_pad] + marlin_scale = marlin_utils.marlin_permute_scales(scale_2d.to( + torch.half), + size_k_pad, + size_n_pad, + group_size=group_size) + marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) + + ws2 = module.weight_scale_2.data + if ws2.numel() == 0 or not torch.isfinite(ws2).all() or ws2.item() == 0: + ws2 = torch.tensor([1.0], dtype=torch.float32, device=weight.device) + weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( + ws2.to(torch.bfloat16)) + + module.weight = Parameter(marlin_weight, requires_grad=False) + module.weight_scale = Parameter(marlin_scale, requires_grad=False) + module.weight_global_scale = Parameter(weight_global_scale, + requires_grad=False) + # Padded GEMM dims consumed by ``apply``; default to the real sizes. + module._marlin_size_k = size_k_pad + module._marlin_size_n = size_n_pad + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + assert is_nvfp4_marlin_enabled() + size_k = module.in_features + size_n = module.out_features + # Set by post_load_weights; equal to size_k/size_n when 64-aligned. + size_k_pad = getattr(module, "_marlin_size_k", size_k) + size_n_pad = getattr(module, "_marlin_size_n", size_n) + + x = input.bfloat16() + if size_k_pad != size_k: + x = F.pad(x, (0, size_k_pad - size_k)) + output = torch.ops.trtllm.marlin_nvfp4_gemm( + x, + module.weight, + scale_a=None, + scale_b=module.weight_scale, + alpha=None, + weight_global_scale=module.weight_global_scale, + bias=None, + out_dtype=module.dtype, + size_n=size_n_pad, + size_k=size_k_pad, + output_buffer_kind=int(BufferKind.DEFAULT), + ) + if size_n_pad != size_n: + output = output[..., :size_n].contiguous() + if bias is not None: + output = output + bias + return output + + def apply_linear_allreduce(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor], tp_rank: int, + tp_group: List[int]): + raise RuntimeError( + "MarlinNVFP4LinearMethod does not support apply_linear_allreduce") + + def get_quant_method(quant_config: Optional[QuantConfig] = None): if quant_config is None or not quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True): @@ -2805,6 +2924,8 @@ def get_quant_method(quant_config: Optional[QuantConfig] = None): if quant_config.layer_quant_mode.has_nvfp4(): if quant_config.quant_algo == QuantAlgo.NVFP4_ARC: return NVFP4ARCLinearMethod() + elif is_nvfp4_marlin_enabled(): + return MarlinNVFP4LinearMethod() else: return NVFP4LinearMethod() if quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8(): diff --git a/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py b/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py index 7bf3d7f33d5f..d968d8745e3d 100644 --- a/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py +++ b/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py @@ -20,6 +20,7 @@ import triton import triton.language as tl +from ...._utils import get_sm_version from ...utils import Fp4QuantizedTensor @@ -226,6 +227,7 @@ def forward( # NVFP4 quantized path - uses optimized fused CUDA kernel # Fuses: SiLU gating + Group RMSNorm + FP4 quantization if self.is_nvfp4 and z is not None and not self.norm_before_gate and \ + get_sm_version() >= 100 and \ fused_gated_rmsnorm_quant_shape_ok(self.hidden_size, self.group_size): if self.nvfp4_scale is None: raise ValueError( diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 03f9d8068679..b6687c2fe13f 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -75,7 +75,6 @@ def __init__( super().__init__() config = config or ModelConfig() - self.mapping = config.mapping if config.mapping.enable_attention_dp: self.mapping = Mapping( @@ -157,22 +156,25 @@ def __init__( # Choose between flashinfer and native implementation. (default to flashinfer) self._mamba_ssm_cache_dtype = config.quant_config.mamba_ssm_cache_dtype - # TODO: Update head_dims once flashinfer is updated. - # Nemotron-v2-Nano (mamba_head_dim=80) is not supported by flashinfer yet. - supported_head_dims = [64, 128] - self._use_flashinfer = head_dim in supported_head_dims self._stochastic_rounding_requested = ( config.quant_config.mamba_ssm_stochastic_rounding) self._philox_rounds = config.quant_config.mamba_ssm_philox_rounds - # SR needs fp16 cache. Replay and flashinfer each supply a Philox impl; - # custom_op does not. Only use_replay is resolved per-forward (from the - # cache manager), so precompute both gate values here. - sr_base = (self._stochastic_rounding_requested - and self._mamba_ssm_cache_dtype == torch.float16) - # Keep replay SSM-cache writes on the same stochastic-rounding policy - # as flashinfer; the replay kernel masks stale slots before using them. - self._stochastic_rounding_for_replay = sr_base - self._stochastic_rounding_for_flashinfer = sr_base and self._use_flashinfer + + # TODO: Update head_dims once flashinfer is updated. + # Nemotron-v2-Nano (mamba_head_dim=80) is not supported by flashinfer yet. + supported_head_dims = [64, 128] + supported_head_group_ratios = [1, 8, 16] + supported_d_states = [64, 128, 256] + head_group_ratio = (self.tp_nheads // + self.tp_ngroups if self.tp_ngroups > 0 else 0) + self._use_flashinfer = (head_dim in supported_head_dims and + head_group_ratio in supported_head_group_ratios + and d_state in supported_d_states) + + self._stochastic_rounding_for_replay = ( + self._stochastic_rounding_requested + and self._mamba_ssm_cache_dtype == torch.float16) + self._stochastic_rounding_for_flashinfer = self._stochastic_rounding_for_replay and self._use_flashinfer self._use_mtp_custom_op = os.environ.get( "TRTLLM_MAMBA2_MTP_USE_CUSTOM_OP", "0") == "1" diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index 6d910bf215eb..faf3036e40d1 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -4,6 +4,7 @@ import torch from torch import nn +from tensorrt_llm._utils import get_sm_version from tensorrt_llm.mapping import Mapping from ..model_config import ModelConfig @@ -101,8 +102,14 @@ def create_weights(self): has_kernel = hasattr(torch.ops.trtllm, 'fused_relu2_quantize') has_scale = hasattr(self.down_proj, 'input_scale') is_relu2 = self.activation is relu2 - - self._use_fused_relu2_quant = has_nvfp4 and has_kernel and has_scale and is_relu2 + # The fused relu2+fp4_quantize kernel body is guarded by + # ``__CUDA_ARCH__ >= 1000`` (see fusedActivationQuant.cu). On pre-SM100 + # GPUs the kernel is a no-op, so fall back to unfused relu2 → separate + # quantize in the downstream linear layer. + is_sm100_or_later = get_sm_version() >= 100 + + self._use_fused_relu2_quant = (has_nvfp4 and has_kernel and has_scale + and is_relu2 and is_sm100_or_later) def forward( self, diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py index 8c9a744e3bdd..8575f006be7b 100644 --- a/tensorrt_llm/_torch/modules/rms_norm.py +++ b/tensorrt_llm/_torch/modules/rms_norm.py @@ -23,7 +23,7 @@ from ..._utils import get_sm_version from ..cuda_tile_utils import IS_CUDA_TILE_AVAILABLE from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE -from ..utils import Fp4QuantizedTensor +from ..utils import Fp4QuantizedTensor, is_nvfp4_marlin_enabled class RMSNorm(nn.Module): @@ -82,7 +82,7 @@ def __init__( # the downstream linear layer handle FP4 quantization. if self.is_nvfp4: sm_version = get_sm_version() - if not (90 <= sm_version < 120): + if not (90 <= sm_version < 120) or is_nvfp4_marlin_enabled(): self.is_nvfp4 = False return_hp_output = False self.return_hp_output = return_hp_output @@ -186,8 +186,8 @@ def _ensure_contiguous_with_dtype(t: torch.Tensor, key: str): gather=True, use_gemma=self.use_gemma, ) - elif IS_FLASHINFER_AVAILABLE and hidden_states.dtype in ( - torch.float16, torch.bfloat16): + elif IS_FLASHINFER_AVAILABLE and not is_nvfp4_marlin_enabled( + ) and hidden_states.dtype in (torch.float16, torch.bfloat16): from ..custom_ops import (flashinfer_fused_add_rmsnorm, flashinfer_gemma_fused_add_rmsnorm, flashinfer_gemma_rmsnorm, diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 4e9c92c9ba76..e13eed6ce5e2 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -10,7 +10,7 @@ from torch.nn import functional as F from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, - torch_dtype_to_str) + get_sm_version, torch_dtype_to_str) from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import ceil_div, pad_up from tensorrt_llm.quantization.utils import fp4_utils @@ -104,6 +104,15 @@ def get_model_extra_attrs(): return getattr(_model_extra_attrs, 'attrs', None) +def is_nvfp4_marlin_enabled() -> bool: + is_hopper = get_sm_version() == 90 + has_marlin_kernel = hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") + attrs = get_model_extra_attrs() + is_marlin_specified = attrs is not None and "marlin" in attrs.get( + 'nvfp4_gemm_allowed_backends', []) + return is_hopper and has_marlin_kernel and is_marlin_specified + + @contextlib.contextmanager def model_extra_attrs(attrs: Dict): old_attrs = getattr(_model_extra_attrs, 'attrs', None) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 73bcc1936192..7c6b57eadd0a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -989,7 +989,7 @@ class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ "AUTO", "CUTLASS", "CUTEDSL", "WIDEEP", "TRTLLM", "DEEPGEMM", - "DENSEGEMM", "VANILLA", "TRITON"] = Field( + "DENSEGEMM", "VANILLA", "TRITON", "MARLIN"] = Field( default='AUTO', description="MoE backend to use. " "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." @@ -1019,7 +1019,7 @@ class MoeConfig(StrictBaseModel): ) -Nvfp4Backend = Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core'] +Nvfp4Backend = Literal['cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'] # Short aliases for built-in custom tokenizers. # Maps alias → full import path (module.ClassName). diff --git a/tensorrt_llm/quantization/utils/marlin_utils.py b/tensorrt_llm/quantization/utils/marlin_utils.py new file mode 100644 index 000000000000..be4763846b49 --- /dev/null +++ b/tensorrt_llm/quantization/utils/marlin_utils.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Marlin weight repacking and scale processing utilities. + +Ported from vLLM's marlin_utils.py and marlin_utils_fp4.py for use with +the vLLM-style Marlin NVFP4 kernels in TensorRT-LLM. +""" + +import torch + +GPTQ_MARLIN_TILE = 16 +MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128] +FP4_MARLIN_SUPPORTED_GROUP_SIZES = [16] +USE_FP32_REDUCE_DEFAULT = True + + +def get_scale_perms(): + """Get the permutation indices for Marlin scale layout.""" + scale_perm = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single = [] + for i in range(4): + scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]]) + return scale_perm, scale_perm_single + + +def marlin_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int, is_a_8bit: bool = False +) -> torch.Tensor: + """Permute scale tensor from [num_groups, N] to Marlin interleaved layout.""" + scale_perm, scale_perm_single = get_scale_perms() + if group_size < size_k and group_size != -1 and not is_a_8bit: + s = s.reshape((-1, len(scale_perm)))[:, scale_perm] + else: + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + s = s.reshape((-1, size_n)).contiguous() + return s + + +def marlin_permute_bias(s: torch.Tensor) -> torch.Tensor: + """Permute bias to match Marlin kernel layout.""" + origin_shape = s.shape + _, scale_perm_single = get_scale_perms() + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + return s.reshape(*origin_shape).contiguous() + + +def marlin_make_workspace(device: torch.device, max_blocks_per_sm: int = 1) -> torch.Tensor: + """Allocate int32 workspace tensor sized for grid parallelism.""" + props = torch.cuda.get_device_properties(device) + sms = props.multi_processor_count + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device, requires_grad=False) + + +def nvfp4_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor: + """Convert FP8-S1E4M3 scales to special FP8-S0E5M3 format for fast dequant. + + This assumes scales are non-negative. The conversion multiplies by 2^7 and + left-shifts by 1 to create a format where the top bit is always 1 when + scale > 0, allowing the kernel to use an exponent bias closer to zero. + """ + # Convert to half for manipulation + marlin_scales = marlin_scales.to(torch.half) + + # Fit the layout of fp8 dequantization + marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(marlin_scales.size(0), -1) + + # Convert to S0E5M3 format + marlin_scales = (marlin_scales * (2**7)).view(torch.int16) << 1 + marlin_scales = marlin_scales.view(torch.float8_e4m3fn) + marlin_scales = marlin_scales[:, 1::2].contiguous() + + return marlin_scales + + +def nvfp4_marlin_process_global_scale(global_scale: torch.Tensor) -> torch.Tensor: + """Adjust global scale with exponent bias for BF16/FP16 dequantization.""" + assert global_scale.dtype in [torch.half, torch.bfloat16] + fp4_exponent = 2 + if global_scale.dtype == torch.half: + target_exponent = 5 + elif global_scale.dtype == torch.bfloat16: + target_exponent = 8 + exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1) + return global_scale * (2.0 ** (exponent_bias - 7)) + + +def prepare_nvfp4_moe_weights_for_marlin( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w13_global_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + w2_global_scale: torch.Tensor, + hidden_size: int, + intermediate_size_per_partition: int, + num_experts: int, + is_act_and_mul: bool, + param_dtype: torch.dtype, +): + """Repack NVFP4 MoE weights and scales to Marlin tiled format. + + Returns: (w13, w13_scale, w13_global_scale, w2, w2_scale, w2_global_scale) + """ + GROUP_SIZE = 16 + K = hidden_size + N = intermediate_size_per_partition + device = w13.device + + perm = torch.empty(0, dtype=torch.int, device=device) + + def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: + tensor_list = [] + num_shards = 2 if is_act_and_mul else 1 + if "w13" in name: + size_n, size_k = N * num_shards, K + else: + size_n, size_k = K, N + + assert weight.shape == (num_experts, size_n, size_k // 2), ( + f"Expected {(num_experts, size_n, size_k // 2)}, got {weight.shape}" + ) + + for i in range(num_experts): + qweight = weight[i].view(torch.int32).T.contiguous() + marlin_qweight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight, + perm=perm, + size_k=size_k, + size_n=size_n, + num_bits=4, + is_a_8bit=False, + ) + tensor_list.append(marlin_qweight) + return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + + def permute_scales(scales: torch.Tensor, g_scales: torch.Tensor, name: str): + scales = scales.to(param_dtype) + g_scales = g_scales.to(param_dtype) + + tensor_list = [] + num_shards = 2 if is_act_and_mul else 1 + if "w13" in name: + size_n, size_k = N * num_shards, K + else: + size_n, size_k = K, N + + for i in range(num_experts): + scale = scales[i].T + marlin_scales = marlin_permute_scales( + s=scale, + size_k=size_k, + size_n=size_n, + group_size=GROUP_SIZE, + is_a_8bit=False, + ) + marlin_scales = nvfp4_marlin_process_scales(marlin_scales) + tensor_list.append(marlin_scales) + + scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) + g_scales = nvfp4_marlin_process_global_scale(g_scales) + return scales, g_scales + + w13 = repack_weight(w13, "w13") + w2 = repack_weight(w2, "w2") + + w13_scale, w13_global_scale = permute_scales(w13_scale, w13_global_scale, "w13") + w2_scale, w2_global_scale = permute_scales(w2_scale, w2_global_scale, "w2") + + return w13, w13_scale, w13_global_scale, w2, w2_scale, w2_global_scale diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 9666884d93e5..0bf29a72b199 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -42,8 +42,8 @@ get_device_memory, llm_models_root, parametrize_with_ids, skip_no_hopper, skip_no_mxfp4_swizzle, skip_post_blackwell, - skip_pre_ada, skip_pre_blackwell, skip_pre_hopper, - skip_ray) + skip_post_hopper, skip_pre_ada, skip_pre_blackwell, + skip_pre_hopper, skip_ray) from .accuracy_core import (GSM8K, MMLU, CnnDailymail, GPQADiamond, JsonModeEval, LlmapiAccuracyTestHarness, LongBenchV1, LongBenchV2) @@ -6547,6 +6547,32 @@ def test_fp8(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + @skip_post_hopper + @pytest.mark.skip_less_device_memory(80000) + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("tp_size", [1, 2, 4, 8]) + def test_nvfp4_marlin_multi_gpus(self, tp_size): + ep_size = tp_size + with LLM( + f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + mamba_ssm_cache_dtype="float16", + ), + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + max_batch_size=32, + moe_config=MoeConfig(backend="MARLIN"), + nvfp4_gemm_config={"allowed_backends": ["marlin"]}, + ) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + class TestNemotronV3Super(LlmapiAccuracyTestHarness): MODEL_NAME = "nvidia/Nemotron-Super-V3" @@ -6695,7 +6721,6 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend): layer_updates_per_iter=2) self._run_nvfp4_4gpus_eplb(moe_backend, eplb_config, model_path) - @skip_pre_hopper @skip_post_blackwell @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.skip_less_device_memory(80000) @@ -6724,6 +6749,28 @@ def test_nvfp4_4gpus_hopper_w4a16(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + @skip_post_hopper + @pytest.mark.skip_less_device_memory(80000) + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("tp_size", + [2, 4, 8]) # starting from TP=2 to avoid OOM + def test_nvfp4_marlin_multi_gpus(self, tp_size): + model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" + kv_cache_config = KvCacheConfig(enable_block_reuse=False) + ep_size = tp_size + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + moe_config=MoeConfig(backend="MARLIN"), + kv_cache_config=kv_cache_config, + max_batch_size=16, + nvfp4_gemm_config={"allowed_backends": ["marlin"]}) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.skip_less_device_memory(40000) @@ -7175,6 +7222,7 @@ def test_nvfp4_8gpus_mtp_custom_op(self, monkeypatch): class TestNemotronV3Ultra(LlmapiAccuracyTestHarness): MODEL_NAME = "nvidia/Nemotron-Ultra-V3" + MODEL_PATH = f"{llm_models_root()}/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4" # No thinking mode for now. EXTRA_EVALUATOR_KWARGS = dict(chat_template_kwargs=dict( enable_thinking=False)) @@ -7207,6 +7255,30 @@ def _run_nvfp4_4gpus_eplb(self, moe_backend, eplb_config, model_path): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper + @skip_post_hopper + @pytest.mark.skip_less_device_memory(80000) + @pytest.mark.skip_less_mpi_world_size(8) + def test_nvfp4_marlin_8gpus(self): + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + mamba_ssm_cache_dtype="float16", + ) + + with LLM(self.MODEL_PATH, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + moe_config=MoeConfig(backend="MARLIN"), + kv_cache_config=kv_cache_config, + max_batch_size=8, + nvfp4_gemm_config={"allowed_backends": ["marlin"]}) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.skip_less_device_memory(80000) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index af1cdf5277d9..59598f11c9de 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -1811,6 +1811,11 @@ def check_device_contain(keyword_list): reason="This test is not supported in pre-Hopper architecture", ) +skip_post_hopper = pytest.mark.skipif( + get_sm_version() > 90, + reason="This test is not supported in post-Hopper architecture", +) + skip_pre_blackwell = pytest.mark.skipif( get_sm_version() < 100, reason="This test is not supported in pre-Blackwell architecture", diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 56e1d4fbe60c..76928aec8485 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -713,6 +713,9 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attentio accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus[attention_dp_on-cutedsl] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_8gpus_mtp_custom_op +accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] +accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] +accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_marlin_8gpus accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TEP4_PP2] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TEP8_PP1] accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index 462551b1b555..b3c2b8ca8c0a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -23,6 +23,9 @@ l0_dgx_h200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=8] + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_marlin_multi_gpus[tp_size=8] + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_marlin_8gpus - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=False] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] diff --git a/tests/unittest/_torch/modules/moe/moe_test_utils.py b/tests/unittest/_torch/modules/moe/moe_test_utils.py index 3f6864c130d8..556535a360ce 100644 --- a/tests/unittest/_torch/modules/moe/moe_test_utils.py +++ b/tests/unittest/_torch/modules/moe/moe_test_utils.py @@ -42,6 +42,7 @@ from tensorrt_llm._torch.modules.fused_moe import ( CuteDslFusedMoE, CutlassFusedMoE, + MarlinFusedMoE, TRTLLMGenFusedMoE, ) from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE @@ -78,6 +79,7 @@ class MoeBackendType(str, Enum): MEGAMOE_DEEPGEMM = "MEGAMOE_DEEPGEMM" MEGAMOE_CUTEDSL = "MEGAMOE_CUTEDSL" CUTE_DSL_B12X = "CUTE_DSL_B12X" + MARLIN = "MARLIN" def get_backend_class(backend_type: MoeBackendType) -> Type[MoE]: @@ -91,6 +93,7 @@ def get_backend_class(backend_type: MoeBackendType) -> Type[MoE]: MoeBackendType.MEGAMOE_DEEPGEMM: MegaMoEDeepGemm, MoeBackendType.MEGAMOE_CUTEDSL: MegaMoECuteDsl, MoeBackendType.CUTE_DSL_B12X: CuteDslB12xFusedMoE, + MoeBackendType.MARLIN: MarlinFusedMoE, } return backend_class_map[backend_type] @@ -1097,7 +1100,7 @@ def supports_autotuner_capture( Returns: True if autotuner capture/replay is supported, False otherwise """ - # DEEPGEMM, both MegaMoE backends, and CUTE_DSL_B12X do not support + # DEEPGEMM, both MegaMoE backends, CUTE_DSL_B12X, and MARLIN do not support # autotuner capture (fused kernels own dispatch+combine, b12x has its own # dispatch/replay state). if backend_type in ( @@ -1105,6 +1108,7 @@ def supports_autotuner_capture( MoeBackendType.MEGAMOE_DEEPGEMM, MoeBackendType.MEGAMOE_CUTEDSL, MoeBackendType.CUTE_DSL_B12X, + MoeBackendType.MARLIN, ): return False @@ -1379,6 +1383,8 @@ def should_skip_to_accelerate_ci( Rules applied (in order): 0. Skip unquantized (quant=None) for most paths, but keep TRTLLM BF16 unquantized coverage enabled. + 0a. MARLIN backend: only NVFP4 on Hopper (SM90); skip all other + quant_algo / architecture combinations. 1. e256 model: only DeepSeekV3 routing, bfloat16, seq=1, non-gptoss 2. Multi-GPU: only DEP and TTP parallel modes 3. Routing: full 6 routing methods only on (CUTLASS or TRTLLM) with NVFP4; @@ -1413,6 +1419,16 @@ def should_skip_to_accelerate_ci( ): return "[CI accel] Skip unquantized (quant=None) in CI" + # --- Rule 0a: MARLIN backend only runs NVFP4 on Hopper (SM90) --- + if backend_type == MoeBackendType.MARLIN: + from tensorrt_llm._utils import get_sm_version + + if quant_algo != QuantAlgo.NVFP4: + return f"[CI accel] MARLIN only tests NVFP4 in CI (got {quant_algo})" + sm_version = get_sm_version() + if sm_version != 90: + return f"[CI accel] MARLIN only runs on Hopper (SM90) in CI (got SM{sm_version})" + # Any e256-class model_config triggers CI Rule-1 minimal coverage: # the full dtype x seq_len x swiglu x routing matrix on e256 models # otherwise blows the per-stage Slurm wall-clock budget (B200 stage diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/modules/moe/quantize_utils.py index 0cb3fd613beb..ce220c6d5a61 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/modules/moe/quantize_utils.py @@ -121,16 +121,21 @@ def get_test_quant_params(quant_algo, x, backend_type=None): x_scale = x_scale.float().squeeze() quant_kwargs["x_scale"] = x_scale elif quant_algo == QuantAlgo.NVFP4: - quantize_util_cls = NVFP4QuantizeUtil quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - x_sf_global = (448 * 6) / x.abs().max().float() - quant_kwargs["x_sf_global"] = x_sf_global - # MegaMoE CuteDSL runs the deepgemm graph (routing weight folded into the - # SwiGLU output before the fc1-output NVFP4 quant), so it needs a - # graph-matched reference; the generic transformers-graph NVFP4 reference - # mismatches systematically. See NVFP4RefMegaMoECuteDsl. - if _normalize_backend_name(backend_type) == "MEGAMOE_CUTEDSL": - quant_kwargs["ref_cls"] = NVFP4RefMegaMoECuteDsl + backend_name = _normalize_backend_name(backend_type) + if backend_name == "MARLIN": + # Marlin is W4A16 — no activation quantization required. + quantize_util_cls = MarlinNVFP4QuantizeUtil + else: + quantize_util_cls = NVFP4QuantizeUtil + x_sf_global = (448 * 6) / x.abs().max().float() + quant_kwargs["x_sf_global"] = x_sf_global + # MegaMoE CuteDSL runs the deepgemm graph (routing weight folded into + # the SwiGLU output before the fc1-output NVFP4 quant), so it needs a + # graph-matched reference; the generic transformers-graph NVFP4 + # reference mismatches systematically. See NVFP4RefMegaMoECuteDsl. + if backend_name == "MEGAMOE_CUTEDSL": + quant_kwargs["ref_cls"] = NVFP4RefMegaMoECuteDsl elif quant_algo == QuantAlgo.FP8_BLOCK_SCALES: quant_config = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) # Different backends have different numerical behaviors for FP8 block scaling: @@ -886,6 +891,245 @@ def create_ref_module(self, routing_method, ref_cls=NVFP4RefMLPFusedMoE) -> torc return ref_cls(**kwargs) +class MarlinNVFP4RefGatedMLPFusedMoE(nn.Module): + """Pure torch.matmul reference for NVFP4 MoE on Hopper (Marlin backend). + + Dequantizes FP4 weights to float32 via the CPU reference op and runs plain + matmul + SiLU activation. This avoids any quantized GEMM kernel so it works + on all architectures and serves as a kernel-independent ground truth. + """ + + SF_VEC_SIZE = 16 + + def __init__( + self, + num_experts: int, + routing_method, + hidden_size: int, + intermediate_size: int, + dtype=None, + model_config=None, + bias=False, + swiglu_gptoss_style: bool = False, + swiglu_alpha=None, + swiglu_beta=None, + swiglu_limit=None, + **kwargs, + ): + super().__init__() + self.num_experts = num_experts + self.routing_method = routing_method + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.dtype = dtype + self.bias = bias + self.swiglu_gptoss_style = swiglu_gptoss_style + self.swiglu_alpha_val = swiglu_alpha + self.swiglu_beta_val = swiglu_beta + self.swiglu_limit_val = swiglu_limit + self._expert_weights = {} + + @staticmethod + def _dequant_fp4_weight(weight_fp4, weight_scale_fp8, weight_scale_2): + """Dequantize FP4 weight to float32 using CPU reference op.""" + weight_fp4_cpu = weight_fp4.cpu() + scale_cpu = weight_scale_fp8.view(torch.uint8).flatten().cpu() + if isinstance(weight_scale_2, (int, float)): + global_scale = torch.tensor([weight_scale_2], dtype=torch.float32) + else: + global_scale = weight_scale_2.float().cpu().reshape(1) + return torch.ops.tensorrt_llm.e2m1_and_ufp8sf_scale_to_float_v2( + weight_fp4_cpu, + scale_cpu, + global_scale, + MarlinNVFP4RefGatedMLPFusedMoE.SF_VEC_SIZE, + 1, + False, + ) + + def load_weights(self, weights_list): + assert len(weights_list) == 1 + weights = weights_list[0] + for expert_id in range(self.num_experts): + ew = {} + for proj in ("w1", "w2", "w3"): + fp4 = weights[f"{expert_id}.{proj}.weight"] + scale = weights[f"{expert_id}.{proj}.weight_scale"] + gs = weights[f"{expert_id}.{proj}.weight_scale_2"] + ew[proj] = self._dequant_fp4_weight(fp4, scale, gs).cuda() + if self.bias and f"{expert_id}.{proj}.bias" in weights: + ew[f"{proj}_bias"] = weights[f"{expert_id}.{proj}.bias"].cuda() + self._expert_weights[expert_id] = ew + + def cuda(self, device=None): + super().cuda(device) + for ew in self._expert_weights.values(): + for name, tensor in ew.items(): + ew[name] = tensor.cuda(device) + return self + + def _activation(self, gate, value): + if self.swiglu_gptoss_style: + limit = self.swiglu_limit_val + if limit is not None and limit != float("inf"): + gate = gate.clamp(max=limit) + value = value.clamp(min=-limit, max=limit) + alpha = self.swiglu_alpha_val if self.swiglu_alpha_val is not None else 1.0 + gate_act = gate * torch.sigmoid(gate * alpha) + beta = self.swiglu_beta_val if self.swiglu_beta_val is not None else 0.0 + return gate_act * (value + beta) + return F.silu(gate) * value + + def forward(self, hidden_states, router_logits): + hidden_states = hidden_states.view(-1, self.hidden_size) + selected_experts, routing_weights = self.routing_method.apply(router_logits) + final_hidden_states = torch.zeros( + hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device + ) + for expert_id in range(self.num_experts): + if not torch.any(selected_experts == expert_id): + continue + batch_idx, nth_expert = torch.where(selected_experts == expert_id) + expert_inputs = hidden_states[batch_idx].float() + ew = self._expert_weights[expert_id] + gate = expert_inputs @ ew["w3"].T + up = expert_inputs @ ew["w1"].T + if self.bias: + if "w1_bias" in ew: + gate = gate + ew["w1_bias"] + if "w3_bias" in ew: + up = up + ew["w3_bias"] + hidden = self._activation(gate, up) + out = hidden @ ew["w2"].T + if self.bias and "w2_bias" in ew: + out = out + ew["w2_bias"] + final_hidden_states[batch_idx] += routing_weights[batch_idx, nth_expert, None] * out + return final_hidden_states + + def check_accuracy(self, output, ref_output): + if self.swiglu_gptoss_style: + check_accuracy(output, ref_output, rtol=0.1, atol=0.1, percent=0.95) + else: + check_accuracy(output, ref_output, rtol=1e-2, atol=0.15, percent=0.97) + + +class MarlinNVFP4QuantizeUtil(BaseQuantizeUtil): + """QuantizeUtil for Marlin NVFP4 MoE on Hopper (SM90). + + Uses the CPU-based ``float_to_e2m1_and_ufp8sf_scale`` op to create FP4 + weights with unswizzled FP8 E4M3 block scales. This avoids the SM100+ + ``fp4_quantize`` kernel so the tests can run on Hopper. + + Weight format matches modelopt checkpoint convention (unswizzled scales, + per-tensor ``weight_scale_2``). + """ + + SF_VEC_SIZE = 16 + + def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]: + assert self.quant_config is not None and self.quant_config.quant_algo == QuantAlgo.NVFP4, ( + "expect quant_algo to be NVFP4" + ) + + weights = {} + for expert_id in range(self.num_experts): + w1_float = ( + torch.randn( + (self.intermediate_size, self.hidden_size), + dtype=torch.float32, + ) + * 0.05 + ).cpu() + w2_float = ( + torch.randn( + (self.hidden_size, self.intermediate_size), + dtype=torch.float32, + ) + * 0.05 + ).cpu() + w3_float = ( + torch.randn( + (self.intermediate_size, self.hidden_size), + dtype=torch.float32, + ) + * 0.05 + ).cpu() + + w1_fp4, w1_sf, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( + w1_float, self.SF_VEC_SIZE, 1, False + ) + w2_fp4, w2_sf, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( + w2_float, self.SF_VEC_SIZE, 1, False + ) + w3_fp4, w3_sf, _ = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( + w3_float, self.SF_VEC_SIZE, 1, False + ) + + num_groups_w1 = self.hidden_size // self.SF_VEC_SIZE + num_groups_w2 = self.intermediate_size // self.SF_VEC_SIZE + w1_sf_2d = ( + w1_sf.view(self.intermediate_size, -1)[:, :num_groups_w1] + .contiguous() + .view(torch.float8_e4m3fn) + ) + w2_sf_2d = ( + w2_sf.view(self.hidden_size, -1)[:, :num_groups_w2] + .contiguous() + .view(torch.float8_e4m3fn) + ) + w3_sf_2d = ( + w3_sf.view(self.intermediate_size, -1)[:, :num_groups_w1] + .contiguous() + .view(torch.float8_e4m3fn) + ) + + weights[f"{expert_id}.w1.weight"] = w1_fp4.cuda() + weights[f"{expert_id}.w2.weight"] = w2_fp4.cuda() + weights[f"{expert_id}.w3.weight"] = w3_fp4.cuda() + weights[f"{expert_id}.w1.weight_scale"] = w1_sf_2d.cuda() + weights[f"{expert_id}.w2.weight_scale"] = w2_sf_2d.cuda() + weights[f"{expert_id}.w3.weight_scale"] = w3_sf_2d.cuda() + weights[f"{expert_id}.w1.weight_scale_2"] = torch.tensor(1.0, dtype=torch.float32) + weights[f"{expert_id}.w2.weight_scale_2"] = torch.tensor(1.0, dtype=torch.float32) + weights[f"{expert_id}.w3.weight_scale_2"] = torch.tensor(1.0, dtype=torch.float32) + weights[f"{expert_id}.w1.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + weights[f"{expert_id}.w2.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + weights[f"{expert_id}.w3.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + + if self.bias: + weights[f"{expert_id}.w1.bias"] = torch.randn( + self.intermediate_size, device="cuda", dtype=torch.float + ) + weights[f"{expert_id}.w2.bias"] = torch.randn( + self.hidden_size, device="cuda", dtype=torch.float + ) + weights[f"{expert_id}.w3.bias"] = torch.randn( + self.intermediate_size, device="cuda", dtype=torch.float + ) + + return weights + + def create_ref_module( + self, + routing_method, + ref_cls=MarlinNVFP4RefGatedMLPFusedMoE, + ) -> torch.nn.Module: + ref_fused_moe = ref_cls( + num_experts=self.num_experts, + routing_method=routing_method, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + dtype=self.dtype, + model_config=ModelConfig(quant_config=self.quant_config), + bias=self.bias, + swiglu_gptoss_style=self.swiglu_gptoss_style, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta, + swiglu_limit=self.swiglu_limit, + ) + return ref_fused_moe + + class FP8BlockScalesRefGatedMLPFusedMoE(RefMLPFusedMoE): """Reference implementation of FP8 block-wise quantization for correctness testing.""" diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 8006a8c890c8..ff974c11dd2b 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -316,6 +316,7 @@ def run_backend_moe( MoeBackendType.MEGAMOE_DEEPGEMM, MoeBackendType.MEGAMOE_CUTEDSL, MoeBackendType.CUTE_DSL_B12X, + MoeBackendType.MARLIN, ] # Data types to test diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index 544791860c34..eaa4052064d5 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -849,6 +849,7 @@ def init_worker(custom_paths, comm_method_type, master_port, moe_backend): MoeBackendType.MEGAMOE_DEEPGEMM, MoeBackendType.MEGAMOE_CUTEDSL, MoeBackendType.CUTE_DSL_B12X, + MoeBackendType.MARLIN, ] # Data types to test @@ -1599,7 +1600,9 @@ def test_trtllm_gen_fp32_routing_bias(routing_method_cls, moe_model_config, quan model_configs=MOE_MODEL_CONFIGS, seq_lens=[8] if IS_CI_MODE else SEQ_LENS, dtypes=DTYPES, - backend_types=BACKEND_TYPES, + backend_types=[ + b for b in BACKEND_TYPES if b != MoeBackendType.MARLIN + ], # Marlin doesn't support fused routing quant_algos=QUANT_ALGOS, routing_methods=MULTI_GPU_ROUTING_METHODS, ) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 339131905b63..18888cab9c3d 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -8,6 +8,7 @@ from tensorrt_llm._torch.autotuner import autotune from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.utils import model_extra_attrs from tensorrt_llm._utils import get_sm_version from tensorrt_llm.math_utils import pad_up from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -749,6 +750,91 @@ def test_fp4_linear_cuda_core(dtype, mnk): ) +@pytest.mark.skipif( + get_sm_version() < 90 or get_sm_version() >= 100, + reason="Marlin NVFP4 backend runs Hopper", +) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize( + "mnk", + [ + (1, 1024, 1024), + (8, 1024, 2048), + (128, 2048, 1024), + (1, 18560, 4096), + (128, 18560, 4096), + (1, 4096, 8192), + (128, 4096, 8192), + # Non-64-aligned N and/or K (e.g. from TP sharding) + (8, 2576, 672), + (128, 2576, 672), + (8, 2576, 4096), + (8, 4096, 2576), + (8, 2576, 544), + (1, 96, 80), + (3, 176, 144), + (128, 928, 1360), + ]) +def test_fp4_linear_marlin(dtype, mnk): + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + torch.manual_seed(0) + + w_float = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=torch.float32) + w_fp4, w_sf_swizzled, w_dequant = torch.ops.tensorrt_llm.float_to_e2m1_and_ufp8sf_scale( + w_float, + scaling_vector_size, + 1, # ufp8_type=1 (e4m3) + True, # is_sf_swizzled_layout=True (modelopt checkpoint native layout) + ) + assert torch.iinfo(w_sf_swizzled.dtype).bits == 8 # torch.uint8 + w_sf_2d = torch.ops.trtllm.block_scale_interleave_reverse( + w_sf_swizzled.view(pad_up(OUTPUT_SIZE, 128), + -1)).view(torch.float8_e4m3fn) + + with model_extra_attrs({'nvfp4_gemm_allowed_backends': ['marlin']}): + l_marlin = Linear( + in_features=HIDDEN_SIZE, + out_features=OUTPUT_SIZE, + bias=False, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + nvfp4_allowed_backends=['marlin'], # key + ) + + # ``float_to_e2m1_and_ufp8sf_scale`` returns ``w_dequant`` that already + # encodes the per-block FP8 scale. The Marlin BF16-activation path + # multiplies the kernel output by ``weight_global_scale`` (derived from + # ``weight_scale_2``); we want that scalar to be 1 so the kernel result + # matches the reference ``torch.mm(x, w_dequant.T)``. Mirrors the + # passing GEMM test (test_fp4_gemm.py:453-454, is_bf16_act=True branch). + l_marlin.load_weights([{ + 'weight': + w_fp4, + 'weight_scale': + w_sf_2d, + 'weight_scale_2': + torch.tensor(1.0, dtype=torch.float32), + }]) + l_marlin = l_marlin.cuda() + + l_marlin.post_load_weights() + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + x_sf_global = (448 * 6) / x.abs().max().float() + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( + x, x_sf_global, scaling_vector_size, False) + + with torch.inference_mode(): + output = l_marlin(x) + + w_dequant_bf16 = w_dequant.to(dtype).cuda() + with torch.inference_mode(): + ref_output = torch.mm(x, w_dequant_bf16.T) + + torch.cuda.synchronize() + torch.testing.assert_close(output, ref_output, atol=0.5, rtol=2e-2) + + if __name__ == "__main__": # m, n, k nvfp4_gemm_perf_test(torch.bfloat16, 128, 7168, 16384) diff --git a/tests/unittest/trt/functional/test_fp4_gemm.py b/tests/unittest/trt/functional/test_fp4_gemm.py index 92b80b0b2430..a4b24129e805 100644 --- a/tests/unittest/trt/functional/test_fp4_gemm.py +++ b/tests/unittest/trt/functional/test_fp4_gemm.py @@ -19,7 +19,8 @@ import tensorrt as trt import torch from parameterized import parameterized -from utils.util import skip_pre_blackwell_unittest, unittest_name_func +from utils.util import (skip_non_hopper_unittest, skip_pre_blackwell_unittest, + unittest_name_func) import tensorrt_llm from tensorrt_llm import Tensor @@ -388,3 +389,99 @@ def test_input_quant_and_fp4_gemm(self, input_dim, output_dim, batch_size, ref_output_fp32, atol=1e-3, rtol=1e-3) + + @parameterized.expand(list( + product([1024, 2048], [1024, 2048], [1, 8, 128], [16], [1.0, 2.0], + ['nvfp4', 'bf16'])), + name_func=unittest_name_func) + @skip_non_hopper_unittest + def test_nvfp4_marlin_gemm(self, input_dim, output_dim, batch_size, + sf_vec_size, alpha, act_dtype): + from tensorrt_llm.quantization.utils import marlin_utils + + torch.random.manual_seed(0) + + is_bf16_act = (act_dtype == 'bf16') + + if is_bf16_act: + # BF16 activations — no FP4 quantization needed + input_bf16 = torch.randn((batch_size, input_dim), + dtype=torch.bfloat16).cuda() + input_fp32 = input_bf16.float().cpu() + else: + # FP4 activations with swizzled scales (for act dequant kernel) + input_e2m1, input_e4m3_scale, input_fp32 = random_fp4_tensor_and_sf( + (batch_size, input_dim), sf_vec_size) + + # FP4 weights with UN-swizzled scales (for Marlin processing) + weights_e2m1, weights_e4m3_scale_raw, weights_fp32 = \ + float_tensor_to_e2m1_and_ufp8_scale( + torch.randn((output_dim, input_dim), dtype=torch.float32), + sf_vec_size, ufp8_type=1, is_sf_swizzled_layout=False) + + weights_fp32_transposed = torch.transpose(weights_fp32, 0, 1) + alpha_tensor = torch.FloatTensor([alpha]).cuda() + + ref_output_fp32 = torch.matmul(input_fp32, weights_fp32_transposed) + if not is_bf16_act: + ref_output_fp32 *= alpha + + # Step 1: Repack weights to Marlin tiled format + qweight_int32 = weights_e2m1.cuda().view(torch.int32).T.contiguous() + perm = torch.empty(0, dtype=torch.int32, device='cuda') + marlin_weight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight_int32, + perm=perm, + size_k=input_dim, + size_n=output_dim, + num_bits=4, + is_a_8bit=False, + ) + + # Step 2: Process weight scales for Marlin kernel + num_groups = input_dim // sf_vec_size + scale_fp8 = weights_e4m3_scale_raw.cuda().view(torch.float8_e4m3fn) + scale_2d = scale_fp8.reshape(output_dim, num_groups).T.contiguous() + marlin_scale = marlin_utils.marlin_permute_scales( + scale_2d.to(torch.half), + input_dim, + output_dim, + group_size=sf_vec_size) + marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) + + # Step 3: Process global scale (includes exponent bias correction) + weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( + torch.tensor(1.0, dtype=torch.bfloat16, device='cuda')) + + if is_bf16_act: + # BF16 path: pass BF16 activations directly, dummy scale_a/alpha + mat_a = input_bf16 + scale_a = torch.zeros(1, dtype=torch.uint8, device='cuda') + alpha_arg = torch.ones(1, dtype=torch.float32, device='cuda') + else: + # FP4 path: unswizzle activation scales for the dequant kernel + mat_a = input_e2m1.cuda() + m_padded = (batch_size + 128 - 1) // 128 * 128 + act_sf_gpu = input_e4m3_scale.cuda() + scale_a = torch.ops.trtllm.block_scale_interleave_reverse( + act_sf_gpu.view(m_padded, -1)).flatten() + alpha_arg = alpha_tensor + + output = torch.ops.trtllm.marlin_nvfp4_gemm( + mat_a, + marlin_weight, + scale_a=scale_a, + scale_b=marlin_scale, + alpha=alpha_arg, + weight_global_scale=weight_global_scale, + bias=None, + out_dtype=torch.bfloat16, + size_n=output_dim, + size_k=input_dim, + ) + + output_cpu_float = output.float().cpu() + assert torch.allclose(output_cpu_float, + ref_output_fp32, + atol=0.75, + rtol=0.02)