[AMD] Dsv4/pr1 fix run time issue (#25898)
Co-authored-by: wunhuang <wunhuang@amd.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: amd-danli103 <danli103@amd.com> Co-authored-by: Lin, Soga <soga.lin@amd.com> Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com> Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Co-authored-by: yichiche@amd.com <jacky.cheng> Co-authored-by: yctseng0211 <yctseng@amd.com> Co-authored-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
kk
co-authored by
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
parent
982f67d9a6
commit
af8f66940e
@@ -258,6 +258,7 @@ set(SOURCES
|
||||
"csrc/elementwise/activation.cu"
|
||||
"csrc/elementwise/concat_mla.cu"
|
||||
"csrc/elementwise/copy.cu"
|
||||
"csrc/elementwise/dsv4_norm_rope.cu"
|
||||
"csrc/elementwise/fused_add_rms_norm_kernel.cu"
|
||||
"csrc/elementwise/pos_enc.cu"
|
||||
"csrc/elementwise/topk.cu"
|
||||
|
||||
@@ -47,6 +47,25 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
|
||||
"topk_indices_offset, Tensor ? row_starts) -> ()");
|
||||
m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface);
|
||||
|
||||
m.def(
|
||||
"deepseek_v4_topk_transform_512(Tensor scores, Tensor seq_lens, Tensor page_table, Tensor! "
|
||||
"page_indices, int page_size, Tensor!? raw_indices) -> ()");
|
||||
m.impl("deepseek_v4_topk_transform_512", torch::kCUDA, &deepseek_v4_topk_transform_512);
|
||||
|
||||
m.def(
|
||||
"dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()");
|
||||
m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope);
|
||||
|
||||
m.def(
|
||||
"dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, "
|
||||
"Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()");
|
||||
m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla);
|
||||
|
||||
m.def(
|
||||
"dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, "
|
||||
"Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()");
|
||||
m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant);
|
||||
|
||||
/*
|
||||
* From csrc/allreduce
|
||||
*/
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
/* Copyright 2025 SGLang Team. All Rights Reserved.
|
||||
|
||||
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 <ATen/core/TensorBase.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <c10/macros/Macros.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <optional>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr uint32_t kMaxTopK = 1024;
|
||||
constexpr uint32_t kBlockSize = 512;
|
||||
|
||||
#ifdef SGL_TOPK_DYNAMIC_SMEM_BYTES
|
||||
constexpr size_t kSMEM = static_cast<size_t>(SGL_TOPK_DYNAMIC_SMEM_BYTES);
|
||||
#else
|
||||
constexpr size_t kSMEM = 48 * 1024; // bytes
|
||||
#endif
|
||||
static_assert(kSMEM % (2 * sizeof(int32_t)) == 0, "kSMEM must be a multiple of 8 bytes.");
|
||||
|
||||
struct TopKParams {
|
||||
const float* __restrict__ scores;
|
||||
const int32_t* __restrict__ seq_lens;
|
||||
const int32_t* __restrict__ page_table;
|
||||
int32_t* __restrict__ page_indices;
|
||||
int32_t* __restrict__ raw_indices;
|
||||
int64_t score_stride;
|
||||
int64_t page_table_stride;
|
||||
uint32_t page_bits;
|
||||
uint32_t topk;
|
||||
int64_t output_stride;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ uint8_t convert_to_uint8(float x) {
|
||||
__half h = __float2half_rn(x);
|
||||
uint16_t bits = __half_as_ushort(h);
|
||||
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
|
||||
return static_cast<uint8_t>(key >> 8);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t convert_to_uint32(float x) {
|
||||
uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int32_t
|
||||
page_to_slot(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
|
||||
const uint32_t mask = (1u << page_bits) - 1u;
|
||||
return (page_table[i >> page_bits] << page_bits) | static_cast<int32_t>(i & mask);
|
||||
}
|
||||
|
||||
__device__ void naive_paged_transform(
|
||||
int32_t length,
|
||||
uint32_t topk,
|
||||
uint32_t page_bits,
|
||||
const int32_t* __restrict__ page_table,
|
||||
int32_t* __restrict__ page_indices_out,
|
||||
int32_t* __restrict__ raw_indices_out) {
|
||||
for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) {
|
||||
if (i < static_cast<uint32_t>(length)) {
|
||||
page_indices_out[i] = page_to_slot(page_table, i, page_bits);
|
||||
if (raw_indices_out != nullptr) {
|
||||
raw_indices_out[i] = static_cast<int32_t>(i);
|
||||
}
|
||||
} else {
|
||||
page_indices_out[i] = -1;
|
||||
if (raw_indices_out != nullptr) {
|
||||
raw_indices_out[i] = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__device__ void
|
||||
radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, uint32_t length, uint32_t topk) {
|
||||
constexpr uint32_t RADIX = 256;
|
||||
constexpr uint32_t BLOCK_SIZE = kBlockSize;
|
||||
constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t));
|
||||
|
||||
alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32];
|
||||
alignas(128) __shared__ uint32_t s_counter;
|
||||
alignas(128) __shared__ uint32_t s_threshold_bin_id;
|
||||
alignas(128) __shared__ uint32_t s_num_input[2];
|
||||
alignas(128) __shared__ int32_t s_last_remain;
|
||||
|
||||
extern __shared__ uint32_t s_input_idx[][SMEM_INPUT_SIZE];
|
||||
|
||||
const uint32_t tx = threadIdx.x;
|
||||
uint32_t remain_topk = topk;
|
||||
auto& s_histogram = _s_histogram_buf[0];
|
||||
|
||||
const auto run_cumsum = [&] {
|
||||
#pragma unroll 8
|
||||
for (int32_t i = 0; i < 8; ++i) {
|
||||
static_assert(1 << 8 == RADIX);
|
||||
if (tx < RADIX) {
|
||||
const auto j = 1 << i;
|
||||
const auto k = i & 1;
|
||||
auto value = _s_histogram_buf[k][tx];
|
||||
if (tx + j < RADIX) {
|
||||
value += _s_histogram_buf[k][tx + j];
|
||||
}
|
||||
_s_histogram_buf[k ^ 1][tx] = value;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
};
|
||||
|
||||
// stage 1: 8bit coarse histogram
|
||||
if (tx < RADIX + 1) s_histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const auto bin = convert_to_uint8(input[idx]);
|
||||
::atomicAdd(&s_histogram[bin], 1);
|
||||
}
|
||||
__syncthreads();
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[0] = 0;
|
||||
s_counter = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
{
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
remain_topk -= s_histogram[threshold_bin + 1];
|
||||
if (remain_topk == 0) {
|
||||
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const uint32_t bin = convert_to_uint8(input[idx]);
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
output[pos] = static_cast<int32_t>(idx);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
return;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) s_histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
|
||||
for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
|
||||
const float raw_input = input[idx];
|
||||
const uint32_t bin = convert_to_uint8(raw_input);
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
output[pos] = static_cast<int32_t>(idx);
|
||||
} else if (bin == threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_num_input[0], 1);
|
||||
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
|
||||
s_input_idx[0][pos] = idx;
|
||||
const auto bin32 = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin32 >> 24) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// stage 2: refine with 8bit radix passes
|
||||
#pragma unroll 4
|
||||
for (int round = 0; round < 4; ++round) {
|
||||
const auto r_idx = round % 2;
|
||||
|
||||
const auto raw_num_input = s_num_input[r_idx];
|
||||
const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE;
|
||||
|
||||
run_cumsum();
|
||||
if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
|
||||
s_threshold_bin_id = tx;
|
||||
s_num_input[r_idx ^ 1] = 0;
|
||||
s_last_remain = static_cast<int32_t>(remain_topk - s_histogram[tx + 1]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto threshold_bin = s_threshold_bin_id;
|
||||
remain_topk -= s_histogram[threshold_bin + 1];
|
||||
|
||||
if (remain_topk == 0) {
|
||||
for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
output[pos] = static_cast<int32_t>(idx);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
break;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx < RADIX + 1) s_histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
|
||||
const auto idx = s_input_idx[r_idx][i];
|
||||
const auto raw_input = input[idx];
|
||||
const auto offset = 24 - round * 8;
|
||||
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
|
||||
if (bin > threshold_bin) {
|
||||
const auto pos = ::atomicAdd(&s_counter, 1);
|
||||
output[pos] = static_cast<int32_t>(idx);
|
||||
} else if (bin == threshold_bin) {
|
||||
if (round == 3) {
|
||||
const auto pos = ::atomicAdd(&s_last_remain, -1);
|
||||
if (pos > 0) {
|
||||
output[topk - pos] = static_cast<int32_t>(idx);
|
||||
}
|
||||
} else {
|
||||
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
|
||||
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
|
||||
s_input_idx[r_idx ^ 1][pos] = idx;
|
||||
const auto bin32 = convert_to_uint32(raw_input);
|
||||
const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF;
|
||||
::atomicAdd(&s_histogram[sub_bin], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(kBlockSize) void deepseek_v4_topk_transform_kernel(const TopKParams params) {
|
||||
const auto bid = blockIdx.x;
|
||||
const auto seq_len = params.seq_lens[bid];
|
||||
const auto topk = params.topk;
|
||||
const auto score_ptr = params.scores + bid * params.score_stride;
|
||||
const auto page_ptr = params.page_table + bid * params.page_table_stride;
|
||||
const auto indices_ptr = params.page_indices + bid * params.output_stride;
|
||||
const auto raw_indices_ptr =
|
||||
params.raw_indices != nullptr ? params.raw_indices + bid * params.output_stride : nullptr;
|
||||
|
||||
if (seq_len <= static_cast<int32_t>(topk)) {
|
||||
naive_paged_transform(seq_len, topk, params.page_bits, page_ptr, indices_ptr, raw_indices_ptr);
|
||||
return;
|
||||
}
|
||||
|
||||
__shared__ int32_t s_topk_indices[kMaxTopK];
|
||||
radix_topk(score_ptr, s_topk_indices, static_cast<uint32_t>(seq_len), topk);
|
||||
|
||||
__syncthreads();
|
||||
for (uint32_t i = threadIdx.x; i < topk; i += kBlockSize) {
|
||||
const auto raw = s_topk_indices[i];
|
||||
indices_ptr[i] = page_to_slot(page_ptr, static_cast<uint32_t>(raw), params.page_bits);
|
||||
if (raw_indices_ptr != nullptr) {
|
||||
raw_indices_ptr[i] = raw;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <auto* f, size_t kMaxDynamicSMEM>
|
||||
void setup_kernel_smem_once() {
|
||||
[[maybe_unused]]
|
||||
static const auto result = [] {
|
||||
#ifdef USE_ROCM
|
||||
return ::cudaFuncSetAttribute(
|
||||
reinterpret_cast<const void*>(f), ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
|
||||
#else
|
||||
return ::cudaFuncSetAttribute(f, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
|
||||
#endif
|
||||
}();
|
||||
TORCH_CHECK(
|
||||
result == cudaSuccess, "deepseek_v4_topk_transform: cudaFuncSetAttribute failed: ", ::cudaGetErrorString(result));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
|
||||
void deepseek_v4_topk_transform_512(
|
||||
const at::Tensor& scores,
|
||||
const at::Tensor& seq_lens,
|
||||
const at::Tensor& page_table,
|
||||
at::Tensor& page_indices,
|
||||
int64_t page_size,
|
||||
std::optional<at::Tensor> raw_indices_opt) {
|
||||
CHECK_CUDA(scores);
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(page_table);
|
||||
CHECK_CUDA(page_indices);
|
||||
if (raw_indices_opt.has_value()) {
|
||||
CHECK_CUDA(raw_indices_opt.value());
|
||||
}
|
||||
|
||||
TORCH_CHECK(
|
||||
scores.dim() == 2 && scores.scalar_type() == at::kFloat, "scores must be float32 with shape [B, max_seq_len]");
|
||||
TORCH_CHECK(scores.stride(1) == 1, "scores must be contiguous along the last dim");
|
||||
|
||||
TORCH_CHECK(
|
||||
seq_lens.dim() == 1 && seq_lens.is_contiguous() && seq_lens.scalar_type() == at::kInt,
|
||||
"seq_lens must be int32 with shape [B], contiguous");
|
||||
|
||||
TORCH_CHECK(
|
||||
page_table.dim() == 2 && page_table.scalar_type() == at::kInt,
|
||||
"page_table must be int32 with shape [B, num_pages]");
|
||||
TORCH_CHECK(page_table.stride(1) == 1, "page_table must be contiguous along the last dim");
|
||||
|
||||
const auto topk = page_indices.size(1);
|
||||
TORCH_CHECK(
|
||||
page_indices.dim() == 2 && page_indices.is_contiguous() && page_indices.scalar_type() == at::kInt,
|
||||
"page_indices must be int32 with shape [B, topk], contiguous");
|
||||
TORCH_CHECK(
|
||||
topk > 0 && topk <= static_cast<int64_t>(kMaxTopK),
|
||||
"page_indices last dim must be in [1, ",
|
||||
kMaxTopK,
|
||||
"], got ",
|
||||
topk);
|
||||
|
||||
const auto B = scores.size(0);
|
||||
TORCH_CHECK(
|
||||
seq_lens.size(0) == B && page_table.size(0) == B && page_indices.size(0) == B,
|
||||
"batch sizes must match across scores, seq_lens, page_table, page_indices");
|
||||
|
||||
TORCH_CHECK(
|
||||
page_size > 0 && (page_size & (page_size - 1)) == 0, "page_size must be a positive power of 2, got ", page_size);
|
||||
const auto page_bits = static_cast<uint32_t>(__builtin_ctzll(static_cast<unsigned long long>(page_size)));
|
||||
|
||||
int32_t* raw_ptr = nullptr;
|
||||
if (raw_indices_opt.has_value()) {
|
||||
auto& raw = raw_indices_opt.value();
|
||||
TORCH_CHECK(
|
||||
raw.dim() == 2 && raw.is_contiguous() && raw.scalar_type() == at::kInt,
|
||||
"raw_indices must be int32 with shape [B, topk], contiguous");
|
||||
TORCH_CHECK(raw.size(0) == B && raw.size(1) == topk, "raw_indices shape must match page_indices [B, ", topk, "]");
|
||||
raw_ptr = raw.data_ptr<int32_t>();
|
||||
}
|
||||
|
||||
const TopKParams params{
|
||||
.scores = scores.data_ptr<float>(),
|
||||
.seq_lens = seq_lens.data_ptr<int32_t>(),
|
||||
.page_table = page_table.data_ptr<int32_t>(),
|
||||
.page_indices = page_indices.data_ptr<int32_t>(),
|
||||
.raw_indices = raw_ptr,
|
||||
.score_stride = scores.stride(0),
|
||||
.page_table_stride = page_table.stride(0),
|
||||
.page_bits = page_bits,
|
||||
.topk = static_cast<uint32_t>(topk),
|
||||
.output_stride = topk,
|
||||
};
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
const dim3 grid(static_cast<uint32_t>(B));
|
||||
const dim3 block(kBlockSize);
|
||||
|
||||
setup_kernel_smem_once<deepseek_v4_topk_transform_kernel, kSMEM>();
|
||||
deepseek_v4_topk_transform_kernel<<<grid, block, kSMEM, stream>>>(params);
|
||||
|
||||
const auto err = cudaGetLastError();
|
||||
TORCH_CHECK(err == cudaSuccess, "deepseek_v4_topk_transform kernel launch failed: ", ::cudaGetErrorString(err));
|
||||
}
|
||||
@@ -0,0 +1,700 @@
|
||||
/* Copyright 2025 SGLang Team. All Rights Reserved.
|
||||
|
||||
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.
|
||||
==============================================================================*/
|
||||
|
||||
// DeepSeek-V4 fused norm + RoPE kernels, ported from JIT kernel
|
||||
// python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh
|
||||
// to sgl-kernel AOT compilation with CUDA + HIP (ROCm) support.
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#else
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
#include <hip/hip_runtime.h>
|
||||
#endif
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
// ============================================================================
|
||||
// Platform-compatible type aliases
|
||||
// ============================================================================
|
||||
#ifndef USE_ROCM
|
||||
using bf16_t = __nv_bfloat16;
|
||||
using bf16x2_t = __nv_bfloat162;
|
||||
using fp8x2_e4m3_t = __nv_fp8x2_e4m3;
|
||||
#else
|
||||
using bf16_t = __hip_bfloat16;
|
||||
using bf16x2_t = __hip_bfloat162;
|
||||
using fp8x2_e4m3_t = uint16_t;
|
||||
#ifndef __grid_constant__
|
||||
#define __grid_constant__
|
||||
#endif
|
||||
#endif
|
||||
|
||||
// ============================================================================
|
||||
// Utility helpers (inlined, no external header dependency)
|
||||
// ============================================================================
|
||||
|
||||
static constexpr uint32_t kWarpSize = 32;
|
||||
|
||||
template <uint32_t kNumThreads = kWarpSize>
|
||||
__device__ __forceinline__ float warp_reduce_sum(float val) {
|
||||
#pragma unroll
|
||||
for (uint32_t mask = kNumThreads / 2; mask > 0; mask >>= 1)
|
||||
val += SGLANG_SHFL_XOR_SYNC(FULL_MASK, val, mask);
|
||||
return val;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float warp_reduce_max(float val) {
|
||||
#pragma unroll
|
||||
for (uint32_t mask = kWarpSize / 2; mask > 0; mask >>= 1)
|
||||
val = fmaxf(val, SGLANG_SHFL_XOR_SYNC(FULL_MASK, val, mask));
|
||||
return val;
|
||||
}
|
||||
|
||||
// Aligned vector for coalesced memory access.
|
||||
template <typename T, int N>
|
||||
struct alignas(sizeof(T) * N) AlignedVec {
|
||||
T data[N];
|
||||
__device__ __forceinline__ T& operator[](int i) {
|
||||
return data[i];
|
||||
}
|
||||
__device__ __forceinline__ T operator[](int i) const {
|
||||
return data[i];
|
||||
}
|
||||
__device__ __forceinline__ void load(const void* ptr, int64_t offset = 0) {
|
||||
*this = reinterpret_cast<const AlignedVec*>(ptr)[offset];
|
||||
}
|
||||
__device__ __forceinline__ void store(void* ptr, int64_t offset = 0) const {
|
||||
reinterpret_cast<AlignedVec*>(ptr)[offset] = *this;
|
||||
}
|
||||
};
|
||||
|
||||
__device__ __forceinline__ float bf16_to_float(bf16_t v) {
|
||||
return __bfloat162float(v);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ bf16_t float_to_bf16(float v) {
|
||||
#ifndef USE_ROCM
|
||||
return __float2bfloat16_rn(v);
|
||||
#else
|
||||
return __float2bfloat16(v);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// FP8 E4M3 helpers (portable CUDA + HIP)
|
||||
// ============================================================================
|
||||
|
||||
// UE8M0 scale: round a positive float to the nearest power-of-two
|
||||
// representable in UE8M0 (unsigned 8-bit exponent, no mantissa).
|
||||
__device__ __forceinline__ int32_t cast_to_ue8m0(float x) {
|
||||
uint32_t u = __float_as_uint(x);
|
||||
int32_t exp = static_cast<int32_t>((u >> 23) & 0xFFu);
|
||||
uint32_t mant = u & 0x7FFFFFu;
|
||||
return exp + (mant != 0);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float inv_scale_ue8m0(int32_t exp) {
|
||||
return __uint_as_float(static_cast<uint32_t>((127 + 127 - exp) << 23));
|
||||
}
|
||||
|
||||
static constexpr float kFP8Max = 448.0f;
|
||||
|
||||
#ifndef USE_ROCM
|
||||
__device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) {
|
||||
x = fmaxf(fminf(x, kFP8Max), -kFP8Max);
|
||||
y = fmaxf(fminf(y, kFP8Max), -kFP8Max);
|
||||
return __nv_fp8x2_e4m3(float2{x, y});
|
||||
}
|
||||
#else
|
||||
// Software float -> FP8 E4M3 conversion for ROCm
|
||||
__device__ __forceinline__ uint8_t cvt_float_to_fp8_e4m3(float val) {
|
||||
constexpr float kMax = kFP8Max;
|
||||
val = fmaxf(fminf(val, kMax), -kMax);
|
||||
if (val == 0.0f) return 0;
|
||||
|
||||
uint32_t f32 = __float_as_uint(val);
|
||||
uint8_t sign = static_cast<uint8_t>((f32 >> 24) & 0x80u);
|
||||
f32 &= 0x7FFFFFFFu;
|
||||
|
||||
int32_t exp32 = static_cast<int32_t>((f32 >> 23) & 0xFFu);
|
||||
uint32_t mant32 = f32 & 0x7FFFFFu;
|
||||
|
||||
// FP8 E4M3 bias=7, FP32 bias=127, offset=120
|
||||
int32_t exp8 = exp32 - 120;
|
||||
|
||||
if (exp8 <= 0) {
|
||||
mant32 |= 0x800000u;
|
||||
int32_t shift = 1 - exp8;
|
||||
if (shift > 24) return sign;
|
||||
uint32_t shifted = mant32 >> (20 + shift);
|
||||
uint32_t rbit = (shift <= 23) ? ((mant32 >> (19 + shift)) & 1u) : 0u;
|
||||
uint32_t sbit = (shift <= 23) ? ((mant32 & ((1u << (19 + shift)) - 1u)) != 0) : 0u;
|
||||
shifted += (rbit && (sbit || (shifted & 1u)));
|
||||
return sign | static_cast<uint8_t>(shifted & 0x7u);
|
||||
}
|
||||
if (exp8 >= 15) return sign | 0x7Eu;
|
||||
|
||||
uint32_t mant3 = (mant32 >> 20) & 0x7u;
|
||||
uint32_t rbit = (mant32 >> 19) & 1u;
|
||||
uint32_t sbit = (mant32 & 0x7FFFFu) != 0;
|
||||
mant3 += (rbit && (sbit || (mant3 & 1u)));
|
||||
if (mant3 > 7) {
|
||||
mant3 = 0;
|
||||
exp8++;
|
||||
if (exp8 >= 15) return sign | 0x7Eu;
|
||||
}
|
||||
return sign | (static_cast<uint8_t>(exp8) << 3) | static_cast<uint8_t>(mant3);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ fp8x2_e4m3_t pack_fp8(float x, float y) {
|
||||
uint8_t x8 = cvt_float_to_fp8_e4m3(x);
|
||||
uint8_t y8 = cvt_float_to_fp8_e4m3(y);
|
||||
return static_cast<uint16_t>(x8) | (static_cast<uint16_t>(y8) << 8);
|
||||
}
|
||||
#endif
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 1: Fused Q Norm + RoPE
|
||||
// warp-per-(token, head), rmsnorm-self (no weight) + RoPE + write to q_out.
|
||||
// ============================================================================
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr uint32_t kFusedQBlockSize = 128;
|
||||
constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / kWarpSize;
|
||||
|
||||
constexpr uint32_t kFusedKBlockSize = 256;
|
||||
constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / kWarpSize;
|
||||
|
||||
struct FusedQNormRopeParams {
|
||||
const void* __restrict__ q_input;
|
||||
void* __restrict__ q_output;
|
||||
const float* __restrict__ freqs_cis;
|
||||
const int32_t* __restrict__ positions;
|
||||
int64_t q_input_stride_batch;
|
||||
int64_t q_output_stride_batch;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_q_heads;
|
||||
float eps;
|
||||
};
|
||||
|
||||
// Compute the largest power-of-2 vec size that divides both kHeadDim and
|
||||
// fits in 16 bytes, while also dividing kRopeDim.
|
||||
template <int64_t kHeadDim, int64_t kRopeDim>
|
||||
struct QKernelTraits {
|
||||
static constexpr int64_t kMaxVecSize = 16 / sizeof(bf16_t); // 8
|
||||
// Use kRopeDim/kWarpSize (=2 for kRopeDim=64) as the vec size.
|
||||
// This guarantees kRopeDim % kVecSize == 0 and works for all head dims
|
||||
// that are multiples of kWarpSize*kVecSize.
|
||||
static constexpr int64_t kVecSize = kRopeDim / kWarpSize; // 2
|
||||
static constexpr int64_t kLocalSize = kHeadDim / (kWarpSize * kVecSize);
|
||||
static constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
|
||||
static_assert(kHeadDim % (kWarpSize * kVecSize) == 0);
|
||||
static_assert(kRopeDim % kVecSize == 0);
|
||||
static_assert(kRopeDim == kWarpSize * 2, "1 (real, imag) pair per lane");
|
||||
};
|
||||
|
||||
template <int64_t kHeadDim, int64_t kRopeDim>
|
||||
__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_norm_rope_kernel(
|
||||
const __grid_constant__ FusedQNormRopeParams params) {
|
||||
using Traits = QKernelTraits<kHeadDim, kRopeDim>;
|
||||
constexpr int64_t kVecSize = Traits::kVecSize;
|
||||
constexpr int64_t kLocalSize = Traits::kLocalSize;
|
||||
constexpr uint32_t kRopeSize = Traits::kRopeSize;
|
||||
|
||||
using Storage = AlignedVec<bf16_t, kVecSize>;
|
||||
using Float2 = AlignedVec<float, 2>;
|
||||
|
||||
const auto warp_id = threadIdx.x / kWarpSize;
|
||||
const auto lane_id = threadIdx.x % kWarpSize;
|
||||
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
|
||||
|
||||
const uint32_t total_works = params.batch_size * params.num_q_heads;
|
||||
if (work_id >= total_works) return;
|
||||
|
||||
const uint32_t batch_id = work_id / params.num_q_heads;
|
||||
const uint32_t head_id = work_id % params.num_q_heads;
|
||||
const auto input_ptr =
|
||||
static_cast<const bf16_t*>(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim;
|
||||
const auto output_ptr =
|
||||
static_cast<bf16_t*>(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim;
|
||||
const auto position = params.positions[batch_id];
|
||||
|
||||
__shared__ Storage s_rope[kFusedQNumWarps][kRopeSize];
|
||||
|
||||
// Prefetch freq pair.
|
||||
Float2 freq;
|
||||
freq.load(params.freqs_cis + position * kRopeDim, lane_id);
|
||||
|
||||
// Part 1: rmsnorm-self (no weight).
|
||||
Storage input_vec[kLocalSize];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
input_vec[i].load(input_ptr, lane_id + i * kWarpSize);
|
||||
}
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kVecSize; ++j) {
|
||||
float x = bf16_to_float(input_vec[i][j]);
|
||||
sum_of_squares += x * x;
|
||||
}
|
||||
}
|
||||
sum_of_squares = warp_reduce_sum(sum_of_squares);
|
||||
const float norm_factor = rsqrtf(sum_of_squares / static_cast<float>(kHeadDim) + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kVecSize; ++j) {
|
||||
float x = bf16_to_float(input_vec[i][j]);
|
||||
input_vec[i][j] = float_to_bf16(x * norm_factor);
|
||||
}
|
||||
}
|
||||
|
||||
// Stash rope tail into shared memory; write nope tiles to gmem.
|
||||
const bool is_rope_lane = lane_id >= kWarpSize - kRopeSize;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLocalSize; ++i) {
|
||||
if (i == kLocalSize - 1 && is_rope_lane) {
|
||||
const auto rope_id = lane_id - (kWarpSize - kRopeSize);
|
||||
s_rope[warp_id][rope_id] = input_vec[i];
|
||||
} else {
|
||||
input_vec[i].store(output_ptr, lane_id + i * kWarpSize);
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// Part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane.
|
||||
auto elem_ptr = reinterpret_cast<bf16x2_t*>(&s_rope[warp_id][0]);
|
||||
bf16x2_t elem = elem_ptr[lane_id];
|
||||
#ifndef USE_ROCM
|
||||
float2 elem_f = __bfloat1622float2(elem);
|
||||
float x_real = elem_f.x, x_imag = elem_f.y;
|
||||
#else
|
||||
float x_real = __bfloat162float(elem.x), x_imag = __bfloat162float(elem.y);
|
||||
#endif
|
||||
float freq_real = freq[0], freq_imag = freq[1];
|
||||
float rot_real = x_real * freq_real - x_imag * freq_imag;
|
||||
float rot_imag = x_real * freq_imag + x_imag * freq_real;
|
||||
bf16x2_t rotated = __float22bfloat162_rn(make_float2(rot_real, rot_imag));
|
||||
auto out_elem = reinterpret_cast<bf16x2_t*>(output_ptr + (kHeadDim - kRopeDim));
|
||||
out_elem[lane_id] = rotated;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 2: Fused K Norm + RoPE + FlashMLA Store
|
||||
// block-per-token, rmsnorm (with kv_weight) + RoPE + FP8 quantized store.
|
||||
// ============================================================================
|
||||
|
||||
struct FusedKNormRopeFlashMLAParams {
|
||||
const void* __restrict__ kv;
|
||||
const void* __restrict__ kv_weight;
|
||||
const float* __restrict__ freqs_cis;
|
||||
const int32_t* __restrict__ positions;
|
||||
const int32_t* __restrict__ out_loc;
|
||||
uint8_t* __restrict__ kvcache;
|
||||
int64_t kv_stride_batch;
|
||||
uint32_t batch_size;
|
||||
float eps;
|
||||
};
|
||||
|
||||
template <int64_t kHeadDim, int64_t kRopeDim, int32_t kPageBits>
|
||||
__global__ __launch_bounds__(kFusedKBlockSize, 8) void fused_k_norm_rope_flashmla_kernel(
|
||||
const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
|
||||
constexpr int64_t kVecSize = 2;
|
||||
constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1;
|
||||
constexpr int64_t kPageBytes = ((584ll << kPageBits) + 575) / 576 * 576;
|
||||
static_assert(kHeadDim == kFusedKBlockSize * kVecSize);
|
||||
static_assert(kRopeDim == kWarpSize * kVecSize);
|
||||
|
||||
using Storage = AlignedVec<bf16_t, kVecSize>;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
const auto warp_id = tx / kWarpSize;
|
||||
const auto lane_id = tx % kWarpSize;
|
||||
const auto work_id = blockIdx.x;
|
||||
if (work_id >= params.batch_size) return;
|
||||
|
||||
const auto input_ptr = static_cast<const bf16_t*>(params.kv) + work_id * params.kv_stride_batch;
|
||||
const auto position = params.positions[work_id];
|
||||
const auto out_loc = params.out_loc[work_id];
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
AlignedVec<float, kVecSize> data, freq;
|
||||
|
||||
// Part 1: norm with block-wide reduction.
|
||||
{
|
||||
__shared__ float partial_sums[kFusedKNumWarps];
|
||||
|
||||
Storage input_vec, weight_vec;
|
||||
input_vec.load(input_ptr, tx);
|
||||
weight_vec.load(params.kv_weight, tx);
|
||||
if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
|
||||
|
||||
float sum_of_squares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
float x = bf16_to_float(input_vec[i]);
|
||||
sum_of_squares += x * x;
|
||||
}
|
||||
const float warp_sum = warp_reduce_sum(sum_of_squares);
|
||||
if (lane_id == 0) partial_sums[warp_id] = warp_sum;
|
||||
__syncthreads();
|
||||
sum_of_squares = warp_reduce_sum<kFusedKNumWarps>(partial_sums[lane_id % kFusedKNumWarps]);
|
||||
const float norm_factor = rsqrtf(sum_of_squares / static_cast<float>(kHeadDim) + params.eps);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
float x = bf16_to_float(input_vec[i]);
|
||||
float w = bf16_to_float(weight_vec[i]);
|
||||
data[i] = x * norm_factor * w;
|
||||
}
|
||||
}
|
||||
|
||||
const int32_t page = out_loc >> kPageBits;
|
||||
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
|
||||
const auto page_ptr = params.kvcache + page * kPageBytes;
|
||||
const auto value_ptr = page_ptr + offset * 576;
|
||||
|
||||
// Part 2: rope on last warp (BF16 store), per-warp UE8M0 quant + store on others.
|
||||
if (warp_id == kRopeWarp) {
|
||||
float x_real = data[0], x_imag = data[1];
|
||||
float freq_real = freq[0], freq_imag = freq[1];
|
||||
float rot_real = x_real * freq_real - x_imag * freq_imag;
|
||||
float rot_imag = x_real * freq_imag + x_imag * freq_real;
|
||||
bf16x2_t result = __float22bfloat162_rn(make_float2(rot_real, rot_imag));
|
||||
auto rope_ptr = value_ptr + 448;
|
||||
reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
|
||||
} else {
|
||||
float x = data[0], y = data[1];
|
||||
float abs_max = warp_reduce_max(fmaxf(fabsf(x), fabsf(y)));
|
||||
float scale_raw = fmaxf(1e-4f, abs_max) / kFP8Max;
|
||||
int32_t scale_ue8m0 = cast_to_ue8m0(scale_raw);
|
||||
float inv_scale = inv_scale_ue8m0(scale_ue8m0);
|
||||
fp8x2_e4m3_t result = pack_fp8(x * inv_scale, y * inv_scale);
|
||||
auto scale_ptr = page_ptr + (576ll << kPageBits) + offset * 8;
|
||||
reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
|
||||
if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = static_cast<uint8_t>(scale_ue8m0);
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 3: Fused Q Indexer RoPE + Hadamard + FP8 Quantization
|
||||
// warp-per-(token, head), no norm, RoPE + Hadamard + fp8 act-quant.
|
||||
// ============================================================================
|
||||
|
||||
struct FusedQIndexerRopeHadamardQuantParams {
|
||||
const void* __restrict__ q_input;
|
||||
void* __restrict__ q_fp8;
|
||||
const void* __restrict__ weight;
|
||||
float* __restrict__ weights_out;
|
||||
float weight_scale;
|
||||
const float* __restrict__ freqs_cis;
|
||||
const int32_t* __restrict__ positions;
|
||||
uint32_t batch_size;
|
||||
uint32_t num_heads;
|
||||
};
|
||||
|
||||
__global__ __launch_bounds__(kFusedQBlockSize, 16) void fused_q_indexer_rope_hadamard_quant_kernel(
|
||||
const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) {
|
||||
constexpr int64_t kHeadDim = 128;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
constexpr int64_t kVecSize = 4;
|
||||
constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
|
||||
static_assert(kHeadDim == kWarpSize * kVecSize);
|
||||
|
||||
using Storage = AlignedVec<bf16_t, kVecSize>;
|
||||
using Float4 = AlignedVec<float, kVecSize>;
|
||||
using OutStorage = AlignedVec<fp8x2_e4m3_t, 2>;
|
||||
|
||||
const auto warp_id = threadIdx.x / kWarpSize;
|
||||
const auto lane_id = threadIdx.x % kWarpSize;
|
||||
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
|
||||
const bool is_rope_lane = lane_id >= kWarpSize - kRopeSize;
|
||||
|
||||
const uint32_t total_works = params.batch_size * params.num_heads;
|
||||
if (work_id >= total_works) return;
|
||||
|
||||
const uint32_t batch_id = work_id / params.num_heads;
|
||||
const auto input_ptr = static_cast<const bf16_t*>(params.q_input) + work_id * kHeadDim;
|
||||
const auto position = params.positions[batch_id];
|
||||
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
|
||||
|
||||
Float4 data, freq;
|
||||
const float weight_val = bf16_to_float(static_cast<const bf16_t*>(params.weight)[work_id]);
|
||||
|
||||
// Part 1: load (no norm).
|
||||
{
|
||||
Storage input_vec;
|
||||
input_vec.load(input_ptr, lane_id);
|
||||
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpSize - kRopeSize));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i)
|
||||
data[i] = bf16_to_float(input_vec[i]);
|
||||
}
|
||||
|
||||
// Part 2: rope on rope lanes.
|
||||
if (is_rope_lane) {
|
||||
float x_r = data[0], x_i = data[1], y_r = data[2], y_i = data[3];
|
||||
float fxr = freq[0], fxi = freq[1], fyr = freq[2], fyi = freq[3];
|
||||
data[0] = x_r * fxr - x_i * fxi;
|
||||
data[1] = x_r * fxi + x_i * fxr;
|
||||
data[2] = y_r * fyr - y_i * fyi;
|
||||
data[3] = y_r * fyi + y_i * fyr;
|
||||
}
|
||||
|
||||
// Part 3: 128-point Hadamard (2 local + 5 cross-lane stages).
|
||||
{
|
||||
{
|
||||
float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a1;
|
||||
data[1] = a0 - a1;
|
||||
data[2] = a2 + a3;
|
||||
data[3] = a2 - a3;
|
||||
}
|
||||
{
|
||||
float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
|
||||
data[0] = a0 + a2;
|
||||
data[1] = a1 + a3;
|
||||
data[2] = a0 - a2;
|
||||
data[3] = a1 - a3;
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 1; mask < kWarpSize; mask <<= 1) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i) {
|
||||
float other = SGLANG_SHFL_XOR_SYNC_WIDTH(FULL_MASK, data[i], mask, kWarpSize);
|
||||
data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
|
||||
}
|
||||
}
|
||||
const float kHadamardScale = rsqrtf(static_cast<float>(kHeadDim));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kVecSize; ++i)
|
||||
data[i] *= kHadamardScale;
|
||||
}
|
||||
|
||||
// Part 4: per-warp FP8 quant + store.
|
||||
{
|
||||
float local_max = fabsf(data[0]);
|
||||
#pragma unroll
|
||||
for (int i = 1; i < kVecSize; ++i)
|
||||
local_max = fmaxf(local_max, fabsf(data[i]));
|
||||
float abs_max = warp_reduce_max(local_max);
|
||||
float scale = fmaxf(1e-4f, abs_max) / kFP8Max;
|
||||
float inv_scale = 1.0f / scale;
|
||||
|
||||
OutStorage result;
|
||||
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
|
||||
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
|
||||
|
||||
auto out_row = static_cast<uint8_t*>(params.q_fp8) + work_id * kHeadDim;
|
||||
result.store(out_row, lane_id);
|
||||
params.weights_out[work_id] = weight_val * params.weight_scale * scale;
|
||||
}
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// ============================================================================
|
||||
// Host-side launchers (PyTorch C++ extension API)
|
||||
// ============================================================================
|
||||
|
||||
void dsv4_fused_q_norm_rope(
|
||||
const at::Tensor& q_input,
|
||||
at::Tensor& q_output,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions,
|
||||
double eps) {
|
||||
TORCH_CHECK(q_input.is_cuda(), "q_input must be a CUDA tensor");
|
||||
TORCH_CHECK(q_output.is_cuda(), "q_output must be a CUDA tensor");
|
||||
TORCH_CHECK(q_input.scalar_type() == at::ScalarType::BFloat16, "q_input must be bfloat16");
|
||||
TORCH_CHECK(q_output.scalar_type() == at::ScalarType::BFloat16, "q_output must be bfloat16");
|
||||
TORCH_CHECK(q_input.dim() == 3, "q_input must be 3D: (B, H, D)");
|
||||
TORCH_CHECK(q_output.dim() == 3, "q_output must be 3D: (B, H, D)");
|
||||
TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32");
|
||||
|
||||
const int64_t B = q_input.size(0);
|
||||
const int64_t H = q_input.size(1);
|
||||
const int64_t D = q_input.size(2);
|
||||
TORCH_CHECK(
|
||||
q_output.size(0) == B && q_output.size(1) == H && q_output.size(2) == D, "q_output shape must match q_input");
|
||||
TORCH_CHECK(q_input.stride(2) == 1 && q_output.stride(2) == 1, "last dim must be contiguous");
|
||||
TORCH_CHECK(q_input.stride(1) == D && q_output.stride(1) == D, "head dim must be contiguous");
|
||||
|
||||
if (B == 0) return;
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream(q_input.get_device());
|
||||
const auto params = FusedQNormRopeParams{
|
||||
.q_input = q_input.data_ptr(),
|
||||
.q_output = q_output.data_ptr(),
|
||||
.freqs_cis = freqs_cis.data_ptr<float>(),
|
||||
.positions = positions.data_ptr<int32_t>(),
|
||||
.q_input_stride_batch = q_input.stride(0),
|
||||
.q_output_stride_batch = q_output.stride(0),
|
||||
.batch_size = static_cast<uint32_t>(B),
|
||||
.num_q_heads = static_cast<uint32_t>(H),
|
||||
.eps = static_cast<float>(eps),
|
||||
};
|
||||
const uint32_t total_works = static_cast<uint32_t>(B * H);
|
||||
const uint32_t num_blocks = CEILDIV(total_works, kFusedQNumWarps);
|
||||
|
||||
// Dispatch on head_dim. DeepSeek V4 uses D=192 with kRopeDim=64.
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
switch (D) {
|
||||
case 128:
|
||||
fused_q_norm_rope_kernel<128, kRopeDim><<<num_blocks, kFusedQBlockSize, 0, stream>>>(params);
|
||||
break;
|
||||
case 192:
|
||||
fused_q_norm_rope_kernel<192, kRopeDim><<<num_blocks, kFusedQBlockSize, 0, stream>>>(params);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unsupported head_dim for dsv4_fused_q_norm_rope: ", D);
|
||||
}
|
||||
}
|
||||
|
||||
void dsv4_fused_k_norm_rope_flashmla(
|
||||
const at::Tensor& kv,
|
||||
const at::Tensor& kv_weight,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions,
|
||||
const at::Tensor& out_loc,
|
||||
at::Tensor& kvcache,
|
||||
double eps,
|
||||
int64_t page_size) {
|
||||
TORCH_CHECK(kv.is_cuda(), "kv must be a CUDA tensor");
|
||||
TORCH_CHECK(kv.scalar_type() == at::ScalarType::BFloat16, "kv must be bfloat16");
|
||||
TORCH_CHECK(kv.dim() == 2, "kv must be 2D: (B, D)");
|
||||
TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32");
|
||||
TORCH_CHECK(out_loc.scalar_type() == at::ScalarType::Int, "out_loc must be int32");
|
||||
|
||||
const int64_t B = kv.size(0);
|
||||
const int64_t D = kv.size(1);
|
||||
TORCH_CHECK(D == 512, "kv head_dim must be 512 for FlashMLA");
|
||||
TORCH_CHECK(kv_weight.size(0) == D, "kv_weight size must match head_dim");
|
||||
|
||||
if (B == 0) return;
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream(kv.get_device());
|
||||
const auto params = FusedKNormRopeFlashMLAParams{
|
||||
.kv = kv.data_ptr(),
|
||||
.kv_weight = kv_weight.data_ptr(),
|
||||
.freqs_cis = freqs_cis.data_ptr<float>(),
|
||||
.positions = positions.data_ptr<int32_t>(),
|
||||
.out_loc = out_loc.data_ptr<int32_t>(),
|
||||
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
|
||||
.kv_stride_batch = kv.stride(0),
|
||||
.batch_size = static_cast<uint32_t>(B),
|
||||
.eps = static_cast<float>(eps),
|
||||
};
|
||||
|
||||
constexpr int64_t kHeadDim = 512;
|
||||
constexpr int64_t kRopeDim = 64;
|
||||
|
||||
// Dispatch on page_size (must be power of 2).
|
||||
TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0, "page_size must be a power of 2");
|
||||
|
||||
#define LAUNCH_K_KERNEL(PAGE_BITS) \
|
||||
fused_k_norm_rope_flashmla_kernel<kHeadDim, kRopeDim, PAGE_BITS> \
|
||||
<<<static_cast<uint32_t>(B), kFusedKBlockSize, 0, stream>>>(params)
|
||||
|
||||
switch (page_size) {
|
||||
case 1:
|
||||
LAUNCH_K_KERNEL(0);
|
||||
break;
|
||||
case 2:
|
||||
LAUNCH_K_KERNEL(1);
|
||||
break;
|
||||
case 4:
|
||||
LAUNCH_K_KERNEL(2);
|
||||
break;
|
||||
case 8:
|
||||
LAUNCH_K_KERNEL(3);
|
||||
break;
|
||||
case 16:
|
||||
LAUNCH_K_KERNEL(4);
|
||||
break;
|
||||
case 32:
|
||||
LAUNCH_K_KERNEL(5);
|
||||
break;
|
||||
case 64:
|
||||
LAUNCH_K_KERNEL(6);
|
||||
break;
|
||||
case 128:
|
||||
LAUNCH_K_KERNEL(7);
|
||||
break;
|
||||
case 256:
|
||||
LAUNCH_K_KERNEL(8);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unsupported page_size: ", page_size);
|
||||
}
|
||||
#undef LAUNCH_K_KERNEL
|
||||
}
|
||||
|
||||
void dsv4_fused_q_indexer_rope_hadamard_quant(
|
||||
const at::Tensor& q_input,
|
||||
at::Tensor& q_fp8,
|
||||
const at::Tensor& weight,
|
||||
at::Tensor& weights_out,
|
||||
double weight_scale,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions) {
|
||||
TORCH_CHECK(q_input.is_cuda(), "q_input must be a CUDA tensor");
|
||||
TORCH_CHECK(q_input.scalar_type() == at::ScalarType::BFloat16, "q_input must be bfloat16");
|
||||
TORCH_CHECK(q_input.dim() == 3, "q_input must be 3D: (B, H, D)");
|
||||
|
||||
const int64_t B = q_input.size(0);
|
||||
const int64_t H = q_input.size(1);
|
||||
constexpr int64_t kHeadDim = 128;
|
||||
TORCH_CHECK(q_input.size(2) == kHeadDim, "q_input head_dim must be 128 for indexer");
|
||||
TORCH_CHECK(
|
||||
q_input.stride(2) == 1 && q_input.stride(1) == kHeadDim, "q_input must be contiguous in (head, elem) dims");
|
||||
TORCH_CHECK(q_input.stride(0) == H * kHeadDim, "q_input must be contiguous (B, H, D)");
|
||||
TORCH_CHECK(q_fp8.stride(0) == H * kHeadDim, "q_fp8 must be contiguous (B, H, D)");
|
||||
TORCH_CHECK(positions.scalar_type() == at::ScalarType::Int, "positions must be int32");
|
||||
|
||||
if (B == 0) return;
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream(q_input.get_device());
|
||||
const auto params = FusedQIndexerRopeHadamardQuantParams{
|
||||
.q_input = q_input.data_ptr(),
|
||||
.q_fp8 = q_fp8.data_ptr(),
|
||||
.weight = weight.data_ptr(),
|
||||
.weights_out = weights_out.data_ptr<float>(),
|
||||
.weight_scale = static_cast<float>(weight_scale),
|
||||
.freqs_cis = freqs_cis.data_ptr<float>(),
|
||||
.positions = positions.data_ptr<int32_t>(),
|
||||
.batch_size = static_cast<uint32_t>(B),
|
||||
.num_heads = static_cast<uint32_t>(H),
|
||||
};
|
||||
const uint32_t total_works = static_cast<uint32_t>(B * H);
|
||||
const uint32_t num_blocks = CEILDIV(total_works, kFusedQNumWarps);
|
||||
|
||||
fused_q_indexer_rope_hadamard_quant_kernel<<<num_blocks, kFusedQBlockSize, 0, stream>>>(params);
|
||||
}
|
||||
@@ -172,8 +172,45 @@ void fast_topk_transform_ragged_interface(
|
||||
|
||||
#ifdef USE_ROCM
|
||||
void gelu_quick(at::Tensor& out, const at::Tensor& input);
|
||||
|
||||
void deepseek_v4_topk_transform_512(
|
||||
const at::Tensor& scores,
|
||||
const at::Tensor& seq_lens,
|
||||
const at::Tensor& page_table,
|
||||
at::Tensor& page_indices,
|
||||
int64_t page_size,
|
||||
std::optional<at::Tensor> raw_indices_opt = std::nullopt);
|
||||
#endif
|
||||
|
||||
/*
|
||||
* From csrc/elementwise (DeepSeek-V4 norm + rope)
|
||||
*/
|
||||
void dsv4_fused_q_norm_rope(
|
||||
const at::Tensor& q_input,
|
||||
at::Tensor& q_output,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions,
|
||||
double eps);
|
||||
|
||||
void dsv4_fused_k_norm_rope_flashmla(
|
||||
const at::Tensor& kv,
|
||||
const at::Tensor& kv_weight,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions,
|
||||
const at::Tensor& out_loc,
|
||||
at::Tensor& kvcache,
|
||||
double eps,
|
||||
int64_t page_size);
|
||||
|
||||
void dsv4_fused_q_indexer_rope_hadamard_quant(
|
||||
const at::Tensor& q_input,
|
||||
at::Tensor& q_fp8,
|
||||
const at::Tensor& weight,
|
||||
at::Tensor& weights_out,
|
||||
double weight_scale,
|
||||
const at::Tensor& freqs_cis,
|
||||
const at::Tensor& positions);
|
||||
|
||||
/*
|
||||
* From csrc/gemm
|
||||
*/
|
||||
|
||||
@@ -324,6 +324,85 @@ if torch.version.hip is not None:
|
||||
return out
|
||||
|
||||
|
||||
def dsv4_fused_q_norm_rope(
|
||||
q_input: torch.Tensor,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
q_output: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""DeepSeek-V4 fused Q RMSNorm (no weight) + RoPE.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
q_input : (B, num_q_heads, head_dim) bfloat16
|
||||
freqs_cis: (max_pos, rope_dim) float32, re/im interleaved
|
||||
positions: (B,) int32
|
||||
eps : RMSNorm epsilon
|
||||
q_output : optional pre-allocated output tensor
|
||||
"""
|
||||
if q_output is None:
|
||||
q_output = torch.empty_like(q_input)
|
||||
torch.ops.sgl_kernel.dsv4_fused_q_norm_rope.default(
|
||||
q_input, q_output, freqs_cis, positions, eps
|
||||
)
|
||||
return q_output
|
||||
|
||||
|
||||
def dsv4_fused_k_norm_rope_flashmla(
|
||||
kv: torch.Tensor,
|
||||
kv_weight: torch.Tensor,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
kvcache: torch.Tensor,
|
||||
eps: float = 1e-6,
|
||||
page_size: int = 1,
|
||||
) -> None:
|
||||
"""DeepSeek-V4 fused K RMSNorm + RoPE + FlashMLA FP8 store.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kv : (B, 512) bfloat16
|
||||
kv_weight: (512,) bfloat16
|
||||
freqs_cis: (max_pos, 64) float32
|
||||
positions: (B,) int32
|
||||
out_loc : (B,) int32 cache slot ids
|
||||
kvcache : (npages, page_bytes) uint8
|
||||
eps : RMSNorm epsilon
|
||||
page_size: page size (power of 2)
|
||||
"""
|
||||
torch.ops.sgl_kernel.dsv4_fused_k_norm_rope_flashmla.default(
|
||||
kv, kv_weight, freqs_cis, positions, out_loc, kvcache, eps, page_size
|
||||
)
|
||||
|
||||
|
||||
def dsv4_fused_q_indexer_rope_hadamard_quant(
|
||||
q_input: torch.Tensor,
|
||||
q_fp8: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weights_out: torch.Tensor,
|
||||
weight_scale: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> None:
|
||||
"""DeepSeek-V4 fused Q indexer: RoPE + Hadamard + FP8 quant.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
q_input : (B, num_heads, 128) bfloat16
|
||||
q_fp8 : (B, num_heads, 128) fp8_e4m3 output
|
||||
weight : (B, num_heads) bfloat16
|
||||
weights_out: (B, num_heads, 1) float32 output
|
||||
weight_scale: scalar
|
||||
freqs_cis : (max_pos, 64) float32
|
||||
positions : (B,) int32
|
||||
"""
|
||||
torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant.default(
|
||||
q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions
|
||||
)
|
||||
|
||||
|
||||
def rotary_embedding(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
|
||||
@@ -80,6 +80,38 @@ def fast_topk_transform_fused(
|
||||
return dst_page_table
|
||||
|
||||
|
||||
def deepseek_v4_topk_transform_512(
|
||||
scores: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
page_table: torch.Tensor,
|
||||
page_indices: torch.Tensor,
|
||||
page_size: int,
|
||||
raw_indices: Optional[torch.Tensor] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Performs the DeepSeek-V4 indexer top-k selection and writes the paged
|
||||
physical slot indices into ``page_indices``. Supports topk up to 1024.
|
||||
Optionally also writes the row-relative raw token positions into
|
||||
``raw_indices`` for hisparse capture.
|
||||
|
||||
Args:
|
||||
scores: float32 ``[B, max_seq_len]`` indexer logits, contiguous on dim 1.
|
||||
seq_lens: int32 ``[B]``, true KV length per batch row.
|
||||
page_table: int32 ``[B, num_pages]``, logical->physical page table,
|
||||
contiguous on dim 1.
|
||||
page_indices: int32 ``[B, topk]``, output buffer, contiguous. Filled
|
||||
with paged physical slots; -1 for padding entries.
|
||||
page_size: power-of-2 page size.
|
||||
raw_indices: optional int32 ``[B, topk]``, contiguous. If provided,
|
||||
filled with raw token positions within each row.
|
||||
"""
|
||||
if raw_indices is not None:
|
||||
assert raw_indices.dim() == 2
|
||||
torch.ops.sgl_kernel.deepseek_v4_topk_transform_512(
|
||||
scores, seq_lens, page_table, page_indices, page_size, raw_indices
|
||||
)
|
||||
|
||||
|
||||
def fast_topk_transform_ragged_fused(
|
||||
score: torch.Tensor,
|
||||
lengths: torch.Tensor,
|
||||
|
||||
@@ -46,6 +46,8 @@ sources = [
|
||||
"csrc/allreduce/quick_all_reduce.cu",
|
||||
"csrc/common_extension_rocm.cc",
|
||||
"csrc/elementwise/activation.cu",
|
||||
"csrc/elementwise/deepseek_v4_topk.cu",
|
||||
"csrc/elementwise/dsv4_norm_rope.cu",
|
||||
"csrc/elementwise/topk.cu",
|
||||
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
|
||||
"csrc/moe/moe_align_kernel.cu",
|
||||
|
||||
Reference in New Issue
Block a user