diff --git a/vllm_musa/jit_kernel/csrc/norm.py b/vllm_musa/jit_kernel/csrc/norm.py index 2fc1b8621651..f924a0f163cf 100644 --- a/vllm_musa/jit_kernel/csrc/norm.py +++ b/vllm_musa/jit_kernel/csrc/norm.py @@ -150,3 +150,138 @@ def _fused_add_rmsnorm_custom_fake( mutates_args=["input", "residual"], fake_impl=_fused_add_rmsnorm_custom_fake, ) + + +@cache_once +def _qk_mrope_module(): + import tilelang + + tilelang_dir = Path(tilelang.__file__).resolve().parent + return load_musa_jit( + "vllm_musa_norm_qk_mrope", + ("norm/qk_mrope.mu",), + extra_musa_cflags=( + f"-I{(tilelang_dir / 'src').resolve()}", + f"-I{(tilelang_dir / '3rdparty' / 'mutlass' / 'include').resolve()}", + "-Wno-error=address-of-temporary", + "-fmusa-flush-denormals-to-zero", + "-fno-signed-zeros", + "-D__MUSA_ARCH_LIST__=310", + "-mllvm", + "-mtgpu-opt-level=1", + "-mllvm", + "-mtgpu-load-store-opt=1", + "-mllvm", + "-mtgpu-fold-global-ldst=1", + "-mllvm", + "-mtgpu-load-cluster-mutation=1", + "-mllvm", + "-mtgpu-store-cluster-mutation=1", + "-mllvm", + "-mtgpu-memory-sched-mutation=1", + "-mllvm", + "-mtgpu-alloc-shared-memory-from-zero=1", + ), + ) + + +def fused_qk_rmsnorm_mrope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + is_neox: bool, + mrope_section_t: int, + mrope_section_h: int, + mrope_section_w: int, + is_interleaved: bool, + eps: float = 1e-6, + gemma: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """QK-RMSNorm + MRoPE in one kernel. q/k are (tokens, heads, head_dim).""" + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + torch.ops.vllm.musa_csrc_fused_qk_rmsnorm_mrope( + q, + k, + q_weight, + k_weight, + positions, + cos_sin_cache, + q_out, + k_out, + bool(is_neox), + int(mrope_section_t), + int(mrope_section_h), + int(mrope_section_w), + bool(is_interleaved), + float(eps), + bool(gemma), + ) + return q_out, k_out + + +def _fused_qk_rmsnorm_mrope_custom( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + is_neox: bool, + mrope_section_t: int, + mrope_section_h: int, + mrope_section_w: int, + is_interleaved: bool, + eps: float, + gemma: bool, +) -> None: + _qk_mrope_module().sgl_musa_fused_qk_rmsnorm_mrope( + q, + k, + q_weight, + k_weight, + positions, + cos_sin_cache, + q_out, + k_out, + bool(is_neox), + int(mrope_section_t), + int(mrope_section_h), + int(mrope_section_w), + bool(is_interleaved), + float(eps), + bool(gemma), + ) + + +def _fused_qk_rmsnorm_mrope_custom_fake( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + is_neox: bool, + mrope_section_t: int, + mrope_section_h: int, + mrope_section_w: int, + is_interleaved: bool, + eps: float, + gemma: bool, +) -> None: + return + + +direct_register_custom_op( + op_name="musa_csrc_fused_qk_rmsnorm_mrope", + op_func=_fused_qk_rmsnorm_mrope_custom, + mutates_args=["q_out", "k_out"], + fake_impl=_fused_qk_rmsnorm_mrope_custom_fake, +) diff --git a/vllm_musa/jit_kernel/csrc/norm/common.mu b/vllm_musa/jit_kernel/csrc/norm/common.mu new file mode 100644 index 000000000000..8345c991799b --- /dev/null +++ b/vllm_musa/jit_kernel/csrc/norm/common.mu @@ -0,0 +1,158 @@ +template +struct __align__(16) Vec8Storage { + T elem[8]; +}; + +struct __align__(32) Float8Storage { + float elem[8]; +}; + +template +struct __align__(16) Vec8 { + union { + Vec8Storage storage; + T elem[8]; + } val; + + __device__ __forceinline__ Vec8() {} + + template + static __device__ __forceinline__ Vec8 load(const T* ptr, Offset idx) { + return *(const Vec8*)(ptr + idx); + } + + template + static __device__ __forceinline__ Vec8 load_byp_slc(const T* ptr, Offset idx) { +#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310)) + Vec8 dst; + const T* addr = ptr + idx; + asm volatile( + "LSU.LD.B128 %0, %1, _, 16, 1, 1, inner_persist=0, outer_persist=2, " + "chrnt=l2_l3, slc=byp, persist=0, stride_add_first=0" + : "=R"(dst) + : "R"(addr)); + return dst; +#else + return *(const Vec8*)(ptr + idx); +#endif + } +}; + +struct __align__(32) Float8 { + union { + Float8Storage storage; + float elem[8]; + } val; + + __device__ __forceinline__ Float8() {} +}; + +__device__ __forceinline__ int mrope_24_20_20_interleaved_axis(int rot_offset) { + constexpr unsigned long long axis1_mask = 0x492492492492492ULL; + constexpr unsigned long long axis2_mask = 0x924924924924924ULL; + const unsigned long long bit = 1ULL << rot_offset; + return ((axis1_mask & bit) != 0ULL) + (((axis2_mask & bit) != 0ULL) << 1); +} + +__device__ __forceinline__ int mrope_11_11_10_interleaved_axis(int rot_offset) { + constexpr unsigned int axis1_mask = 0x92492492U; + constexpr unsigned int axis2_mask = 0x24924924U; + const unsigned int bit = 1U << rot_offset; + return ((axis1_mask & bit) != 0U) + (((axis2_mask & bit) != 0U) << 1); +} + +__device__ __forceinline__ float fast_rsqrt(float value) { +#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310)) + const float half_value = 0.5f * value; + float y = __frsqrt_rn(value); + y = y * (1.5f - half_value * y * y); + return y; +#else + return rsqrtf(value); +#endif +} + +__device__ __forceinline__ float block_sum(float value, float* warp_sums) { + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int num_warps = ((int)blockDim.x + 31) >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[warp] = value; + } + __syncthreads_lm(); + + value = tid < num_warps ? warp_sums[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[0] = value; + } + } + __syncthreads_lm(); + return warp_sums[0]; +} + +__device__ __forceinline__ float block_sum_8warps(float value, float* warp_sums) { + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[warp] = value; + } + __syncthreads_lm(); + + value = lane < 8 ? warp_sums[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[0] = value; + } + } + __syncthreads_lm(); + return warp_sums[0]; +} + +__device__ __forceinline__ float block_sum_4warps(float value, float* warp_sums) { + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[warp] = value; + } + __syncthreads_lm(); + + value = lane < 4 ? warp_sums[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset, 32); + } + if (lane == 0) { + warp_sums[0] = value; + } + } + __syncthreads_lm(); + return warp_sums[0]; +} diff --git a/vllm_musa/jit_kernel/csrc/norm/qk_mrope.mu b/vllm_musa/jit_kernel/csrc/norm/qk_mrope.mu new file mode 100644 index 000000000000..68ef141d9433 --- /dev/null +++ b/vllm_musa/jit_kernel/csrc/norm/qk_mrope.mu @@ -0,0 +1,2941 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../device_utils.h" +#include "common.mu" + +void check_fused_qk_mrope_inputs(ffi::TensorView q, ffi::TensorView k, + ffi::TensorView q_weight, + ffi::TensorView k_weight, + ffi::TensorView positions, + ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out) { + CHECK_MUSA(q); + CHECK_MUSA(k); + CHECK_MUSA(q_weight); + CHECK_MUSA(k_weight); + CHECK_MUSA(positions); + CHECK_MUSA(cos_sin_cache); + CHECK_MUSA(q_out); + CHECK_MUSA(k_out); + TVM_FFI_ICHECK_EQ(q.ndim(), 3); + TVM_FFI_ICHECK_EQ(k.ndim(), 3); + TVM_FFI_ICHECK_EQ(q_out.ndim(), 3); + TVM_FFI_ICHECK_EQ(k_out.ndim(), 3); + TVM_FFI_ICHECK_EQ(positions.ndim(), 2); + TVM_FFI_ICHECK_EQ(positions.size(0), 3); + TVM_FFI_ICHECK_EQ(positions.size(1), q.size(0)); + TVM_FFI_ICHECK_EQ(cos_sin_cache.ndim(), 2); + TVM_FFI_ICHECK_EQ(q.size(2), k.size(2)); + TVM_FFI_ICHECK_EQ(q.size(2) % 2, 0); + TVM_FFI_ICHECK_EQ(q_weight.size(0), q.size(2)); + TVM_FFI_ICHECK_EQ(k_weight.size(0), k.size(2)); + TVM_FFI_ICHECK_GT(cos_sin_cache.size(1), 0); + TVM_FFI_ICHECK_LE(cos_sin_cache.size(1), q.size(2)); + TVM_FFI_ICHECK_EQ(cos_sin_cache.size(1) % 2, 0); + TVM_FFI_ICHECK_EQ(q_out.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(q_out.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(q_out.size(2), q.size(2)); + TVM_FFI_ICHECK_EQ(k_out.size(0), k.size(0)); + TVM_FFI_ICHECK_EQ(k_out.size(1), k.size(1)); + TVM_FFI_ICHECK_EQ(k_out.size(2), k.size(2)); + TVM_FFI_ICHECK_EQ(q.stride(2), 1); + TVM_FFI_ICHECK_EQ(k.stride(2), 1); + TVM_FFI_ICHECK_EQ(q_out.stride(2), 1); + TVM_FFI_ICHECK_EQ(k_out.stride(2), 1); + TVM_FFI_ICHECK_EQ(cos_sin_cache.stride(1), 1); + TVM_FFI_ICHECK_EQ(q.device().device_id, k.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, q_weight.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, k_weight.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, positions.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, cos_sin_cache.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, q_out.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, k_out.device().device_id); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), dl_bfloat16)); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), q_weight.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k_weight.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), cos_sin_cache.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), q_out.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k_out.dtype())); + TVM_FFI_ICHECK(dtype_equal(positions.dtype(), dl_int64)); +} + +void check_fused_qk_mrope_cache_inputs( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices) { + CHECK_MUSA(q); + CHECK_MUSA(k); + CHECK_MUSA(v); + CHECK_MUSA(q_weight); + CHECK_MUSA(k_weight); + CHECK_MUSA(positions); + CHECK_MUSA(cos_sin_cache); + CHECK_MUSA(q_out); + CHECK_MUSA(k_cache); + CHECK_MUSA(v_cache); + CHECK_MUSA(indices); + TVM_FFI_ICHECK_EQ(q.ndim(), 3); + TVM_FFI_ICHECK_EQ(k.ndim(), 3); + TVM_FFI_ICHECK_EQ(v.ndim(), 3); + TVM_FFI_ICHECK_EQ(q_out.ndim(), 3); + TVM_FFI_ICHECK_EQ(k_cache.ndim(), 2); + TVM_FFI_ICHECK_EQ(v_cache.ndim(), 2); + TVM_FFI_ICHECK_EQ(positions.ndim(), 2); + TVM_FFI_ICHECK_EQ(indices.ndim(), 1); + TVM_FFI_ICHECK_EQ(positions.size(0), 3); + TVM_FFI_ICHECK_EQ(positions.size(1), q.size(0)); + TVM_FFI_ICHECK_EQ(indices.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(cos_sin_cache.ndim(), 2); + TVM_FFI_ICHECK_EQ(k.size(1), v.size(1)); + TVM_FFI_ICHECK_EQ(q.size(2), k.size(2)); + TVM_FFI_ICHECK_EQ(k.size(2), v.size(2)); + TVM_FFI_ICHECK_EQ(q.size(2) % 2, 0); + TVM_FFI_ICHECK_EQ(q_weight.size(0), q.size(2)); + TVM_FFI_ICHECK_EQ(k_weight.size(0), k.size(2)); + TVM_FFI_ICHECK_GT(cos_sin_cache.size(1), 0); + TVM_FFI_ICHECK_LE(cos_sin_cache.size(1), q.size(2)); + TVM_FFI_ICHECK_EQ(cos_sin_cache.size(1) % 2, 0); + TVM_FFI_ICHECK_EQ(k_cache.size(1), k.size(1) * k.size(2)); + TVM_FFI_ICHECK_EQ(v_cache.size(1), v.size(1) * v.size(2)); + TVM_FFI_ICHECK_EQ(q_out.size(0), q.size(0)); + TVM_FFI_ICHECK_EQ(q_out.size(1), q.size(1)); + TVM_FFI_ICHECK_EQ(q_out.size(2), q.size(2)); + TVM_FFI_ICHECK_EQ(q.stride(2), 1); + TVM_FFI_ICHECK_EQ(k.stride(2), 1); + TVM_FFI_ICHECK_EQ(v.stride(2), 1); + TVM_FFI_ICHECK_EQ(q_out.stride(2), 1); + TVM_FFI_ICHECK_EQ(k_cache.stride(1), 1); + TVM_FFI_ICHECK_EQ(v_cache.stride(1), 1); + TVM_FFI_ICHECK_EQ(cos_sin_cache.stride(1), 1); + TVM_FFI_ICHECK_EQ(q.device().device_id, k.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, v.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, q_weight.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, k_weight.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, positions.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, cos_sin_cache.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, q_out.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, k_cache.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, v_cache.device().device_id); + TVM_FFI_ICHECK_EQ(q.device().device_id, indices.device().device_id); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), dl_bfloat16)); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), v.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), q_weight.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k_weight.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), cos_sin_cache.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), q_out.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), k_cache.dtype())); + TVM_FFI_ICHECK(dtype_equal(q.dtype(), v_cache.dtype())); + TVM_FFI_ICHECK(dtype_equal(positions.dtype(), dl_int64)); + TVM_FFI_ICHECK(dtype_equal(indices.dtype(), dl_int32) || + dtype_equal(indices.dtype(), dl_int64)); +} + + +template +__global__ void fused_qk_rmsnorm_mrope_32q4kv_h128_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + int batch, int64_t position_stride, int64_t q_batch_stride, + int64_t q_head_stride, int64_t k_batch_stride, int64_t k_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, float eps) { + constexpr int q_heads = 32; + constexpr int k_heads = 4; + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = 8; + constexpr int groups_per_token = 5; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= q_heads + k_heads) { + return; + } + + const bool is_q = global_head < q_heads; + const int head = is_q ? global_head : global_head - q_heads; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + __mt_bfloat16 *__restrict__ out = is_q ? q_out : k_out; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int64_t out_base = + is_q ? ((int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride) + : ((int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + float sum = + data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + data_3 * data_3; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int rot_offset0 = lane; + const int axis0 = INTERLEAVED + ? mrope_24_20_20_interleaved_axis(rot_offset0) + : (rot_offset0 < 24 ? 0 : (rot_offset0 < 44 ? 1 : 2)); + const int64_t pos0 = positions[(int64_t)axis0 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr0 = cos_sin_cache + pos0 * rot_dim; + const float cos_v0 = __bfloat162float(cache_ptr0[rot_offset0]); + const float sin_v0 = __bfloat162float(cache_ptr0[embed_dim + rot_offset0]); + const int x_index0 = rot_offset0; + const int y_index0 = embed_dim + rot_offset0; + const float x_weight0 = + __bfloat162float(weight[x_index0]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight0 = + __bfloat162float(weight[y_index0]) + (GEMMA ? 1.0f : 0.0f); + const float x0 = data_0 * scale * x_weight0; + const float y0 = data_2 * scale * y_weight0; + out[out_base + x_index0] = __float2bfloat16_rn(x0 * cos_v0 - y0 * sin_v0); + out[out_base + y_index0] = __float2bfloat16_rn(y0 * cos_v0 + x0 * sin_v0); + + const int rot_offset1 = lane + 32; + const int axis1 = INTERLEAVED + ? mrope_24_20_20_interleaved_axis(rot_offset1) + : (rot_offset1 < 44 ? 1 : 2); + const int64_t pos1 = positions[(int64_t)axis1 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr1 = cos_sin_cache + pos1 * rot_dim; + const float cos_v1 = __bfloat162float(cache_ptr1[rot_offset1]); + const float sin_v1 = __bfloat162float(cache_ptr1[embed_dim + rot_offset1]); + const int x_index1 = rot_offset1; + const int y_index1 = embed_dim + rot_offset1; + const float x_weight1 = + __bfloat162float(weight[x_index1]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight1 = + __bfloat162float(weight[y_index1]) + (GEMMA ? 1.0f : 0.0f); + const float x1 = data_1 * scale * x_weight1; + const float y1 = data_3 * scale * y_weight1; + out[out_base + x_index1] = __float2bfloat16_rn(x1 * cos_v1 - y1 * sin_v1); + out[out_base + y_index1] = __float2bfloat16_rn(y1 * cos_v1 + x1 * sin_v1); +} + +template +__global__ void fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int64_t position_stride, + int64_t q_batch_stride, int64_t q_head_stride, int64_t k_batch_stride, + int64_t k_head_stride, int64_t v_batch_stride, int64_t v_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, + int64_t k_cache_row_stride, int64_t v_cache_row_stride, + int64_t indices_stride, float eps) { + constexpr int q_heads = 32; + constexpr int k_heads = 4; + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = 8; + constexpr int groups_per_token = 5; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int64_t cache_idx = + static_cast(indices[(int64_t)token * indices_stride]); + + if (group == 0) { + constexpr int kVec = 8; + constexpr int row_dim = k_heads * hidden; + constexpr int vec_count = row_dim / kVec; + const int64_t v_token_base = (int64_t)token * v_batch_stride; + const int64_t v_out_base = cache_idx * v_cache_row_stride; + for (int vec_idx = tid; vec_idx < vec_count; vec_idx += (int)blockDim.x) { + const int head = vec_idx >> 4; + const int col = (vec_idx & 15) * kVec; + Vec8<__mt_bfloat16> v_vec = Vec8<__mt_bfloat16>::load( + v + v_token_base + (int64_t)head * v_head_stride, col); + *(Vec8<__mt_bfloat16> *)(v_cache + v_out_base + (int64_t)head * hidden + + col) = v_vec; + } + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= q_heads + k_heads) { + return; + } + + const bool is_q = global_head < q_heads; + const int head = is_q ? global_head : global_head - q_heads; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int64_t q_out_base = + (int64_t)token * q_out_batch_stride + (int64_t)head * q_out_head_stride; + const int64_t k_out_base = + (int64_t)token * k_out_batch_stride + (int64_t)head * k_out_head_stride; + const int64_t k_cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + float sum = + data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + data_3 * data_3; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int rot_offset0 = lane; + const int axis0 = INTERLEAVED + ? mrope_24_20_20_interleaved_axis(rot_offset0) + : (rot_offset0 < 24 ? 0 : (rot_offset0 < 44 ? 1 : 2)); + const int64_t pos0 = positions[(int64_t)axis0 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr0 = cos_sin_cache + pos0 * rot_dim; + const float cos_v0 = __bfloat162float(cache_ptr0[rot_offset0]); + const float sin_v0 = __bfloat162float(cache_ptr0[embed_dim + rot_offset0]); + const int x_index0 = rot_offset0; + const int y_index0 = embed_dim + rot_offset0; + const float x_weight0 = + __bfloat162float(weight[x_index0]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight0 = + __bfloat162float(weight[y_index0]) + (GEMMA ? 1.0f : 0.0f); + const float x0 = data_0 * scale * x_weight0; + const float y0 = data_2 * scale * y_weight0; + const __mt_bfloat16 x_rot0 = __float2bfloat16_rn(x0 * cos_v0 - y0 * sin_v0); + const __mt_bfloat16 y_rot0 = __float2bfloat16_rn(y0 * cos_v0 + x0 * sin_v0); + if (is_q) { + q_out[q_out_base + x_index0] = x_rot0; + q_out[q_out_base + y_index0] = y_rot0; + } else { + k_cache[k_cache_base + x_index0] = x_rot0; + k_cache[k_cache_base + y_index0] = y_rot0; + if constexpr (STORE_K_OUT) { + k_out[k_out_base + x_index0] = x_rot0; + k_out[k_out_base + y_index0] = y_rot0; + } + } + + const int rot_offset1 = lane + 32; + const int axis1 = INTERLEAVED + ? mrope_24_20_20_interleaved_axis(rot_offset1) + : (rot_offset1 < 44 ? 1 : 2); + const int64_t pos1 = positions[(int64_t)axis1 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr1 = cos_sin_cache + pos1 * rot_dim; + const float cos_v1 = __bfloat162float(cache_ptr1[rot_offset1]); + const float sin_v1 = __bfloat162float(cache_ptr1[embed_dim + rot_offset1]); + const int x_index1 = rot_offset1; + const int y_index1 = embed_dim + rot_offset1; + const float x_weight1 = + __bfloat162float(weight[x_index1]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight1 = + __bfloat162float(weight[y_index1]) + (GEMMA ? 1.0f : 0.0f); + const float x1 = data_1 * scale * x_weight1; + const float y1 = data_3 * scale * y_weight1; + const __mt_bfloat16 x_rot1 = __float2bfloat16_rn(x1 * cos_v1 - y1 * sin_v1); + const __mt_bfloat16 y_rot1 = __float2bfloat16_rn(y1 * cos_v1 + x1 * sin_v1); + if (is_q) { + q_out[q_out_base + x_index1] = x_rot1; + q_out[q_out_base + y_index1] = y_rot1; + } else { + k_cache[k_cache_base + x_index1] = x_rot1; + k_cache[k_cache_base + y_index1] = y_rot1; + if constexpr (STORE_K_OUT) { + k_out[k_out_base + x_index1] = x_rot1; + k_out[k_out_base + y_index1] = y_rot1; + } + } +} + +template +__global__ void fused_qk_rmsnorm_mrope_cache_h128_full_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int64_t position_stride, + int64_t q_batch_stride, int64_t q_head_stride, int64_t k_batch_stride, + int64_t k_head_stride, int64_t v_batch_stride, int64_t v_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, + int64_t k_cache_row_stride, int64_t v_cache_row_stride, + int64_t indices_stride, float eps) { + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int64_t cache_idx = + static_cast(indices[(int64_t)token * indices_stride]); + + if (group == 0) { + constexpr int kVec = 8; + constexpr int row_dim = K_HEADS * hidden; + constexpr int vec_count = row_dim / kVec; + const int64_t v_token_base = (int64_t)token * v_batch_stride; + const int64_t v_out_base = cache_idx * v_cache_row_stride; + for (int vec_idx = tid; vec_idx < vec_count; vec_idx += (int)blockDim.x) { + const int head = vec_idx >> 4; + const int col = (vec_idx & 15) * kVec; + Vec8<__mt_bfloat16> v_vec = Vec8<__mt_bfloat16>::load( + v + v_token_base + (int64_t)head * v_head_stride, col); + *(Vec8<__mt_bfloat16> *)(v_cache + v_out_base + (int64_t)head * hidden + + col) = v_vec; + } + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + float sum = + data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + data_3 * data_3; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int rot_offset0 = lane; + int64_t pos0; + int64_t pos1; + if constexpr (SAME_POSITION) { + pos0 = positions[token]; + pos1 = pos0; + } else { + constexpr unsigned int axis0_mask1 = 0x92492492U; + constexpr unsigned int axis0_mask2 = 0x24924924U; + constexpr unsigned int axis1_mask1 = 0x04924924U; + constexpr unsigned int axis1_mask2 = 0x09249249U; + const unsigned int lane_bit = 1U << lane; + const int axis0 = + ((axis0_mask1 & lane_bit) != 0U) + + (((axis0_mask2 & lane_bit) != 0U) << 1); + pos0 = positions[(int64_t)axis0 * position_stride + token]; + const int axis1 = + ((axis1_mask1 & lane_bit) != 0U) + + (((axis1_mask2 & lane_bit) != 0U) << 1); + pos1 = positions[(int64_t)axis1 * position_stride + token]; + } + const __mt_bfloat16 *__restrict__ cache_ptr0 = cos_sin_cache + pos0 * rot_dim; + const float cos_v0 = __bfloat162float(cache_ptr0[rot_offset0]); + const float sin_v0 = __bfloat162float(cache_ptr0[embed_dim + rot_offset0]); + const float x_weight0 = + __bfloat162float(weight[rot_offset0]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight0 = + __bfloat162float(weight[embed_dim + rot_offset0]) + + (GEMMA ? 1.0f : 0.0f); + const float x0 = data_0 * scale * x_weight0; + const float y0 = data_2 * scale * y_weight0; + const __mt_bfloat16 x_rot0 = __float2bfloat16_rn(x0 * cos_v0 - y0 * sin_v0); + const __mt_bfloat16 y_rot0 = __float2bfloat16_rn(y0 * cos_v0 + x0 * sin_v0); + + const int rot_offset1 = lane + 32; + const __mt_bfloat16 *__restrict__ cache_ptr1 = cos_sin_cache + pos1 * rot_dim; + const float cos_v1 = __bfloat162float(cache_ptr1[rot_offset1]); + const float sin_v1 = __bfloat162float(cache_ptr1[embed_dim + rot_offset1]); + const float x_weight1 = + __bfloat162float(weight[rot_offset1]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight1 = + __bfloat162float(weight[embed_dim + rot_offset1]) + + (GEMMA ? 1.0f : 0.0f); + const float x1 = data_1 * scale * x_weight1; + const float y1 = data_3 * scale * y_weight1; + const __mt_bfloat16 x_rot1 = __float2bfloat16_rn(x1 * cos_v1 - y1 * sin_v1); + const __mt_bfloat16 y_rot1 = __float2bfloat16_rn(y1 * cos_v1 + x1 * sin_v1); + + if (is_q) { + const int64_t q_out_base = + (int64_t)token * q_out_batch_stride + (int64_t)head * q_out_head_stride; + q_out[q_out_base + rot_offset0] = x_rot0; + q_out[q_out_base + embed_dim + rot_offset0] = y_rot0; + q_out[q_out_base + rot_offset1] = x_rot1; + q_out[q_out_base + embed_dim + rot_offset1] = y_rot1; + } else { + const int64_t k_cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + k_cache[k_cache_base + rot_offset0] = x_rot0; + k_cache[k_cache_base + embed_dim + rot_offset0] = y_rot0; + k_cache[k_cache_base + rot_offset1] = x_rot1; + k_cache[k_cache_base + embed_dim + rot_offset1] = y_rot1; + if constexpr (STORE_K_OUT) { + const int64_t k_out_base = (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[k_out_base + rot_offset0] = x_rot0; + k_out[k_out_base + embed_dim + rot_offset0] = y_rot0; + k_out[k_out_base + rot_offset1] = x_rot1; + k_out[k_out_base + embed_dim + rot_offset1] = y_rot1; + } + } +} + +template +__global__ void fused_qk_rmsnorm_rope_cache_h128_full_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int64_t position_stride, + int64_t q_batch_stride, int64_t q_head_stride, int64_t k_batch_stride, + int64_t k_head_stride, int64_t v_batch_stride, int64_t v_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, + int64_t k_cache_row_stride, int64_t v_cache_row_stride, + int64_t indices_stride, float eps) { + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int64_t cache_idx = + static_cast(indices[(int64_t)token * indices_stride]); + + if (group == 0) { + constexpr int kVec = 8; + constexpr int row_dim = K_HEADS * hidden; + constexpr int vec_count = row_dim / kVec; + const int64_t v_token_base = (int64_t)token * v_batch_stride; + const int64_t v_out_base = cache_idx * v_cache_row_stride; + for (int vec_idx = tid; vec_idx < vec_count; vec_idx += (int)blockDim.x) { + const int head = vec_idx >> 4; + const int col = (vec_idx & 15) * kVec; + Vec8<__mt_bfloat16> v_vec = Vec8<__mt_bfloat16>::load( + v + v_token_base + (int64_t)head * v_head_stride, col); + *(Vec8<__mt_bfloat16> *)(v_cache + v_out_base + (int64_t)head * hidden + + col) = v_vec; + } + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + float sum = + data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + data_3 * data_3; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int64_t pos = positions[(int64_t)0 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; + + const float cos_v0 = __bfloat162float(cache_ptr[lane]); + const float sin_v0 = __bfloat162float(cache_ptr[embed_dim + lane]); + const float x_weight0 = + __bfloat162float(weight[lane]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight0 = + __bfloat162float(weight[embed_dim + lane]) + (GEMMA ? 1.0f : 0.0f); + const float x0 = data_0 * scale * x_weight0; + const float y0 = data_2 * scale * y_weight0; + const __mt_bfloat16 x_rot0 = __float2bfloat16_rn(x0 * cos_v0 - y0 * sin_v0); + const __mt_bfloat16 y_rot0 = __float2bfloat16_rn(y0 * cos_v0 + x0 * sin_v0); + + const int rot_offset1 = lane + 32; + const float cos_v1 = __bfloat162float(cache_ptr[rot_offset1]); + const float sin_v1 = __bfloat162float(cache_ptr[embed_dim + rot_offset1]); + const float x_weight1 = + __bfloat162float(weight[rot_offset1]) + (GEMMA ? 1.0f : 0.0f); + const float y_weight1 = + __bfloat162float(weight[embed_dim + rot_offset1]) + + (GEMMA ? 1.0f : 0.0f); + const float x1 = data_1 * scale * x_weight1; + const float y1 = data_3 * scale * y_weight1; + const __mt_bfloat16 x_rot1 = __float2bfloat16_rn(x1 * cos_v1 - y1 * sin_v1); + const __mt_bfloat16 y_rot1 = __float2bfloat16_rn(y1 * cos_v1 + x1 * sin_v1); + + if (is_q) { + const int64_t q_out_base = + (int64_t)token * q_out_batch_stride + (int64_t)head * q_out_head_stride; + q_out[q_out_base + lane] = x_rot0; + q_out[q_out_base + embed_dim + lane] = y_rot0; + q_out[q_out_base + rot_offset1] = x_rot1; + q_out[q_out_base + embed_dim + rot_offset1] = y_rot1; + } else { + const int64_t k_cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + k_cache[k_cache_base + lane] = x_rot0; + k_cache[k_cache_base + embed_dim + lane] = y_rot0; + k_cache[k_cache_base + rot_offset1] = x_rot1; + k_cache[k_cache_base + embed_dim + rot_offset1] = y_rot1; + if constexpr (STORE_K_OUT) { + const int64_t k_out_base = (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[k_out_base + lane] = x_rot0; + k_out[k_out_base + embed_dim + lane] = y_rot0; + k_out[k_out_base + rot_offset1] = x_rot1; + k_out[k_out_base + embed_dim + rot_offset1] = y_rot1; + } + } +} + +template +__global__ void fused_qk_rmsnorm_rope_cache_h128r64_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int64_t position_stride, + int64_t q_batch_stride, int64_t q_head_stride, int64_t k_batch_stride, + int64_t k_head_stride, int64_t v_batch_stride, int64_t v_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, + int64_t k_cache_row_stride, int64_t v_cache_row_stride, + int64_t indices_stride, float eps) { + constexpr int hidden = 128; + constexpr int embed_dim = 32; + constexpr int rot_dim = 64; + constexpr int warps_per_block = HEADS_PER_BLOCK; + constexpr int heads_per_block = warps_per_block * 2; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int half = lane >> 4; + const int sublane = lane & 15; + const int pair_col = sublane * 2; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int64_t cache_idx = + static_cast(indices[(int64_t)token * indices_stride]); + + if (group == 0) { + constexpr int kVec = 8; + constexpr int row_dim = K_HEADS * hidden; + constexpr int vec_count = row_dim / kVec; + const int64_t v_token_base = (int64_t)token * v_batch_stride; + const int64_t v_out_base = cache_idx * v_cache_row_stride; + for (int vec_idx = tid; vec_idx < vec_count; vec_idx += (int)blockDim.x) { + const int head = vec_idx >> 4; + const int col = (vec_idx & 15) * kVec; + Vec8<__mt_bfloat16> v_vec = Vec8<__mt_bfloat16>::load( + v + v_token_base + (int64_t)head * v_head_stride, col); + *(Vec8<__mt_bfloat16> *)(v_cache + v_out_base + (int64_t)head * hidden + + col) = v_vec; + } + } + + const int global_head = group * heads_per_block + warp * 2 + half; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + + const float2 data_0 = __bfloat1622float2( + *reinterpret_cast(data + base + pair_col)); + const float2 data_1 = __bfloat1622float2(*reinterpret_cast( + data + base + embed_dim + pair_col)); + const float2 data_2 = __bfloat1622float2(*reinterpret_cast( + data + base + 64 + pair_col)); + const float2 data_3 = __bfloat1622float2(*reinterpret_cast( + data + base + 96 + pair_col)); + float sum = data_0.x * data_0.x + data_0.y * data_0.y + + data_1.x * data_1.x + data_1.y * data_1.y + + data_2.x * data_2.x + data_2.y * data_2.y + + data_3.x * data_3.x + data_3.y * data_3.y; +#pragma unroll + for (int mask = 8; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int64_t pos = positions[(int64_t)0 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; + + const float2 cos_v = __bfloat1622float2( + *reinterpret_cast(cache_ptr + pair_col)); + const float2 sin_v = __bfloat1622float2(*reinterpret_cast( + cache_ptr + embed_dim + pair_col)); + float2 x_weight = __bfloat1622float2( + *reinterpret_cast(weight + pair_col)); + float2 y_weight = __bfloat1622float2(*reinterpret_cast( + weight + embed_dim + pair_col)); + float2 pass_weight0 = __bfloat1622float2(*reinterpret_cast( + weight + 64 + pair_col)); + float2 pass_weight1 = __bfloat1622float2(*reinterpret_cast( + weight + 96 + pair_col)); + if constexpr (GEMMA) { + x_weight.x += 1.0f; + x_weight.y += 1.0f; + y_weight.x += 1.0f; + y_weight.y += 1.0f; + pass_weight0.x += 1.0f; + pass_weight0.y += 1.0f; + pass_weight1.x += 1.0f; + pass_weight1.y += 1.0f; + } + const float2 x = {data_0.x * scale * x_weight.x, + data_0.y * scale * x_weight.y}; + const float2 y = {data_1.x * scale * y_weight.x, + data_1.y * scale * y_weight.y}; + const float2 x_rot_f = {x.x * cos_v.x - y.x * sin_v.x, + x.y * cos_v.y - y.y * sin_v.y}; + const float2 y_rot_f = {y.x * cos_v.x + x.x * sin_v.x, + y.y * cos_v.y + x.y * sin_v.y}; + const float2 pass0_f = {data_2.x * scale * pass_weight0.x, + data_2.y * scale * pass_weight0.y}; + const float2 pass1_f = {data_3.x * scale * pass_weight1.x, + data_3.y * scale * pass_weight1.y}; + const __mt_bfloat162 x_rot = __float22bfloat162_rn(x_rot_f); + const __mt_bfloat162 y_rot = __float22bfloat162_rn(y_rot_f); + const __mt_bfloat162 pass0 = __float22bfloat162_rn(pass0_f); + const __mt_bfloat162 pass1 = __float22bfloat162_rn(pass1_f); + + if (is_q) { + const int64_t q_out_base = + (int64_t)token * q_out_batch_stride + (int64_t)head * q_out_head_stride; + *reinterpret_cast<__mt_bfloat162 *>(q_out + q_out_base + pair_col) = x_rot; + *reinterpret_cast<__mt_bfloat162 *>(q_out + q_out_base + embed_dim + pair_col) = + y_rot; + *reinterpret_cast<__mt_bfloat162 *>(q_out + q_out_base + 64 + pair_col) = + pass0; + *reinterpret_cast<__mt_bfloat162 *>(q_out + q_out_base + 96 + pair_col) = + pass1; + } else { + const int64_t k_cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + *reinterpret_cast<__mt_bfloat162 *>(k_cache + k_cache_base + pair_col) = + x_rot; + *reinterpret_cast<__mt_bfloat162 *>(k_cache + k_cache_base + embed_dim + + pair_col) = y_rot; + *reinterpret_cast<__mt_bfloat162 *>(k_cache + k_cache_base + 64 + pair_col) = + pass0; + *reinterpret_cast<__mt_bfloat162 *>(k_cache + k_cache_base + 96 + pair_col) = + pass1; + if constexpr (STORE_K_OUT) { + const int64_t k_out_base = (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + *reinterpret_cast<__mt_bfloat162 *>(k_out + k_out_base + pair_col) = + x_rot; + *reinterpret_cast<__mt_bfloat162 *>(k_out + k_out_base + embed_dim + + pair_col) = y_rot; + *reinterpret_cast<__mt_bfloat162 *>(k_out + k_out_base + 64 + pair_col) = + pass0; + *reinterpret_cast<__mt_bfloat162 *>(k_out + k_out_base + 96 + pair_col) = + pass1; + } + } +} + +template +__global__ void fused_qk_rmsnorm_mrope_cache_h256r64_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int64_t position_stride, + int64_t q_batch_stride, int64_t q_head_stride, int64_t k_batch_stride, + int64_t k_head_stride, int64_t v_batch_stride, int64_t v_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, + int64_t k_cache_row_stride, int64_t v_cache_row_stride, + int64_t indices_stride, float eps) { + constexpr int hidden = 256; + constexpr int embed_dim = 32; + constexpr int rot_dim = 64; + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int64_t cache_idx = + static_cast(indices[(int64_t)token * indices_stride]); + + if (group == 0) { + constexpr int kVec = 8; + constexpr int row_dim = K_HEADS * hidden; + constexpr int vec_count = row_dim / kVec; + const int64_t v_token_base = (int64_t)token * v_batch_stride; + const int64_t v_out_base = cache_idx * v_cache_row_stride; + for (int vec_idx = tid; vec_idx < vec_count; vec_idx += (int)blockDim.x) { + const int head = vec_idx >> 5; + const int col = (vec_idx & 31) * kVec; + Vec8<__mt_bfloat16> v_vec = Vec8<__mt_bfloat16>::load( + v + v_token_base + (int64_t)head * v_head_stride, col); + *(Vec8<__mt_bfloat16> *)(v_cache + v_out_base + (int64_t)head * hidden + + col) = v_vec; + } + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + const float data_4 = __bfloat162float(data[base + lane + 128]); + const float data_5 = __bfloat162float(data[base + lane + 160]); + const float data_6 = __bfloat162float(data[base + lane + 192]); + const float data_7 = __bfloat162float(data[base + lane + 224]); + float sum = data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + + data_3 * data_3 + data_4 * data_4 + data_5 * data_5 + + data_6 * data_6 + data_7 * data_7; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 256.0f) + eps); + + const int rot_offset = lane; + const int axis = mrope_11_11_10_interleaved_axis(rot_offset); + const int64_t pos = positions[(int64_t)axis * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; + const float cos_v = __bfloat162float(cache_ptr[rot_offset]); + const float sin_v = __bfloat162float(cache_ptr[embed_dim + rot_offset]); + const float x_weight = __bfloat162float(weight[rot_offset]) + 1.0f; + const float y_weight = __bfloat162float(weight[embed_dim + rot_offset]) + 1.0f; + const float x = data_0 * scale * x_weight; + const float y = data_1 * scale * y_weight; + const __mt_bfloat16 x_rot = __float2bfloat16_rn(x * cos_v - y * sin_v); + const __mt_bfloat16 y_rot = __float2bfloat16_rn(y * cos_v + x * sin_v); + + const __mt_bfloat16 out_2 = __float2bfloat16_rn( + data_2 * scale * (__bfloat162float(weight[lane + 64]) + 1.0f)); + const __mt_bfloat16 out_3 = __float2bfloat16_rn( + data_3 * scale * (__bfloat162float(weight[lane + 96]) + 1.0f)); + const __mt_bfloat16 out_4 = __float2bfloat16_rn( + data_4 * scale * (__bfloat162float(weight[lane + 128]) + 1.0f)); + const __mt_bfloat16 out_5 = __float2bfloat16_rn( + data_5 * scale * (__bfloat162float(weight[lane + 160]) + 1.0f)); + const __mt_bfloat16 out_6 = __float2bfloat16_rn( + data_6 * scale * (__bfloat162float(weight[lane + 192]) + 1.0f)); + const __mt_bfloat16 out_7 = __float2bfloat16_rn( + data_7 * scale * (__bfloat162float(weight[lane + 224]) + 1.0f)); + + if (is_q) { + const int64_t q_out_base = + (int64_t)token * q_out_batch_stride + (int64_t)head * q_out_head_stride; + q_out[q_out_base + lane] = x_rot; + q_out[q_out_base + lane + 32] = y_rot; + q_out[q_out_base + lane + 64] = out_2; + q_out[q_out_base + lane + 96] = out_3; + q_out[q_out_base + lane + 128] = out_4; + q_out[q_out_base + lane + 160] = out_5; + q_out[q_out_base + lane + 192] = out_6; + q_out[q_out_base + lane + 224] = out_7; + } else { + const int64_t k_cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + k_cache[k_cache_base + lane] = x_rot; + k_cache[k_cache_base + lane + 32] = y_rot; + k_cache[k_cache_base + lane + 64] = out_2; + k_cache[k_cache_base + lane + 96] = out_3; + k_cache[k_cache_base + lane + 128] = out_4; + k_cache[k_cache_base + lane + 160] = out_5; + k_cache[k_cache_base + lane + 192] = out_6; + k_cache[k_cache_base + lane + 224] = out_7; + if constexpr (STORE_K_OUT) { + const int64_t k_out_base = (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[k_out_base + lane] = x_rot; + k_out[k_out_base + lane + 32] = y_rot; + k_out[k_out_base + lane + 64] = out_2; + k_out[k_out_base + lane + 96] = out_3; + k_out[k_out_base + lane + 128] = out_4; + k_out[k_out_base + lane + 160] = out_5; + k_out[k_out_base + lane + 192] = out_6; + k_out[k_out_base + lane + 224] = out_7; + } + } +} + +template +__global__ void fused_qk_rmsnorm_mrope_h256r64_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + int batch, int64_t position_stride, int64_t q_batch_stride, + int64_t q_head_stride, int64_t k_batch_stride, int64_t k_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, float eps) { + constexpr int hidden = 256; + constexpr int embed_dim = 32; + constexpr int rot_dim = 64; + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + __mt_bfloat16 *__restrict__ out = is_q ? q_out : k_out; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int64_t out_base = + is_q ? ((int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride) + : ((int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + const float data_4 = __bfloat162float(data[base + lane + 128]); + const float data_5 = __bfloat162float(data[base + lane + 160]); + const float data_6 = __bfloat162float(data[base + lane + 192]); + const float data_7 = __bfloat162float(data[base + lane + 224]); + float sum = data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + + data_3 * data_3 + data_4 * data_4 + data_5 * data_5 + + data_6 * data_6 + data_7 * data_7; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 256.0f) + eps); + + const int rot_offset = lane; + const int axis = mrope_11_11_10_interleaved_axis(rot_offset); + const int64_t pos = positions[(int64_t)axis * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; + const float cos_v = __bfloat162float(cache_ptr[rot_offset]); + const float sin_v = __bfloat162float(cache_ptr[embed_dim + rot_offset]); + const float x_weight = __bfloat162float(weight[rot_offset]) + 1.0f; + const float y_weight = __bfloat162float(weight[embed_dim + rot_offset]) + 1.0f; + const float x = data_0 * scale * x_weight; + const float y = data_1 * scale * y_weight; + out[out_base + lane] = __float2bfloat16_rn(x * cos_v - y * sin_v); + out[out_base + lane + 32] = __float2bfloat16_rn(y * cos_v + x * sin_v); + out[out_base + lane + 64] = __float2bfloat16_rn( + data_2 * scale * (__bfloat162float(weight[lane + 64]) + 1.0f)); + out[out_base + lane + 96] = __float2bfloat16_rn( + data_3 * scale * (__bfloat162float(weight[lane + 96]) + 1.0f)); + out[out_base + lane + 128] = __float2bfloat16_rn( + data_4 * scale * (__bfloat162float(weight[lane + 128]) + 1.0f)); + out[out_base + lane + 160] = __float2bfloat16_rn( + data_5 * scale * (__bfloat162float(weight[lane + 160]) + 1.0f)); + out[out_base + lane + 192] = __float2bfloat16_rn( + data_6 * scale * (__bfloat162float(weight[lane + 192]) + 1.0f)); + out[out_base + lane + 224] = __float2bfloat16_rn( + data_7 * scale * (__bfloat162float(weight[lane + 224]) + 1.0f)); +} + +template +__global__ void fused_qk_rmsnorm_mrope_h128_full_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + int batch, int64_t position_stride, int64_t q_batch_stride, + int64_t q_head_stride, int64_t k_batch_stride, int64_t k_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, float eps) { + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + __mt_bfloat16 *__restrict__ out = is_q ? q_out : k_out; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int64_t out_base = + is_q ? ((int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride) + : ((int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride); + + const float data_0 = __bfloat162float(data[base + lane]); + const float data_1 = __bfloat162float(data[base + lane + 32]); + const float data_2 = __bfloat162float(data[base + lane + 64]); + const float data_3 = __bfloat162float(data[base + lane + 96]); + float sum = + data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + data_3 * data_3; +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + + const int rot_offset0 = lane; + const int axis0 = mrope_24_20_20_interleaved_axis(rot_offset0); + const int64_t pos0 = positions[(int64_t)axis0 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr0 = cos_sin_cache + pos0 * rot_dim; + const float cos_v0 = __bfloat162float(cache_ptr0[rot_offset0]); + const float sin_v0 = __bfloat162float(cache_ptr0[embed_dim + rot_offset0]); + const float x_weight0 = __bfloat162float(weight[rot_offset0]) + 1.0f; + const float y_weight0 = + __bfloat162float(weight[embed_dim + rot_offset0]) + 1.0f; + const float x0 = data_0 * scale * x_weight0; + const float y0 = data_2 * scale * y_weight0; + out[out_base + rot_offset0] = + __float2bfloat16_rn(x0 * cos_v0 - y0 * sin_v0); + out[out_base + embed_dim + rot_offset0] = + __float2bfloat16_rn(y0 * cos_v0 + x0 * sin_v0); + + const int rot_offset1 = lane + 32; + const int axis1 = mrope_24_20_20_interleaved_axis(rot_offset1); + const int64_t pos1 = positions[(int64_t)axis1 * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr1 = cos_sin_cache + pos1 * rot_dim; + const float cos_v1 = __bfloat162float(cache_ptr1[rot_offset1]); + const float sin_v1 = __bfloat162float(cache_ptr1[embed_dim + rot_offset1]); + const float x_weight1 = __bfloat162float(weight[rot_offset1]) + 1.0f; + const float y_weight1 = + __bfloat162float(weight[embed_dim + rot_offset1]) + 1.0f; + const float x1 = data_1 * scale * x_weight1; + const float y1 = data_3 * scale * y_weight1; + out[out_base + rot_offset1] = + __float2bfloat16_rn(x1 * cos_v1 - y1 * sin_v1); + out[out_base + embed_dim + rot_offset1] = + __float2bfloat16_rn(y1 * cos_v1 + x1 * sin_v1); +} + +template +__global__ void fused_qk_rmsnorm_mrope_h128_full_halfwarp_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + int batch, int64_t position_stride, int64_t q_batch_stride, + int64_t q_head_stride, int64_t k_batch_stride, int64_t k_head_stride, + int64_t q_out_batch_stride, int64_t q_out_head_stride, + int64_t k_out_batch_stride, int64_t k_out_head_stride, float eps) { + constexpr int hidden = 128; + constexpr int embed_dim = 64; + constexpr int rot_dim = 128; + constexpr int heads_per_block = 16; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int half = lane >> 4; + const int sublane = lane & 15; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + const int global_head = group * heads_per_block + warp * 2 + half; + if (global_head >= Q_HEADS + K_HEADS) { + return; + } + + const bool is_q = global_head < Q_HEADS; + const int head = is_q ? global_head : global_head - Q_HEADS; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + __mt_bfloat16 *__restrict__ out = is_q ? q_out : k_out; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int64_t out_base = + is_q ? ((int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride) + : ((int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride); + + const int o0 = sublane; + const int o1 = sublane + 16; + const int o2 = sublane + 32; + const int o3 = sublane + 48; + const float data_0 = __bfloat162float(data[base + o0]); + const float data_1 = __bfloat162float(data[base + o1]); + const float data_2 = __bfloat162float(data[base + o2]); + const float data_3 = __bfloat162float(data[base + o3]); + const float data_4 = __bfloat162float(data[base + o0 + embed_dim]); + const float data_5 = __bfloat162float(data[base + o1 + embed_dim]); + const float data_6 = __bfloat162float(data[base + o2 + embed_dim]); + const float data_7 = __bfloat162float(data[base + o3 + embed_dim]); + float sum = data_0 * data_0 + data_1 * data_1 + data_2 * data_2 + + data_3 * data_3 + data_4 * data_4 + data_5 * data_5 + + data_6 * data_6 + data_7 * data_7; +#pragma unroll + for (int mask = 8; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum * (1.0f / 128.0f) + eps); + +#define SGLANG_H128_FULL_HALF_ROTATE(ROT_OFFSET, DATA_X, DATA_Y) \ + do { \ + const int axis = mrope_24_20_20_interleaved_axis(ROT_OFFSET); \ + const int64_t pos = positions[(int64_t)axis * position_stride + token]; \ + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; \ + const float cos_v = __bfloat162float(cache_ptr[ROT_OFFSET]); \ + const float sin_v = __bfloat162float(cache_ptr[embed_dim + ROT_OFFSET]); \ + const float x_weight = __bfloat162float(weight[ROT_OFFSET]) + 1.0f; \ + const float y_weight = \ + __bfloat162float(weight[embed_dim + ROT_OFFSET]) + 1.0f; \ + const float x = (DATA_X)*scale * x_weight; \ + const float y = (DATA_Y)*scale * y_weight; \ + out[out_base + ROT_OFFSET] = __float2bfloat16_rn(x * cos_v - y * sin_v); \ + out[out_base + embed_dim + ROT_OFFSET] = \ + __float2bfloat16_rn(y * cos_v + x * sin_v); \ + } while (0) + + SGLANG_H128_FULL_HALF_ROTATE(o0, data_0, data_4); + SGLANG_H128_FULL_HALF_ROTATE(o1, data_1, data_5); + SGLANG_H128_FULL_HALF_ROTATE(o2, data_2, data_6); + SGLANG_H128_FULL_HALF_ROTATE(o3, data_3, data_7); +#undef SGLANG_H128_FULL_HALF_ROTATE +} + +template +__global__ void fused_qk_rmsnorm_mrope_generic_bf16_kernel( + const __mt_bfloat16 *__restrict__ q, const __mt_bfloat16 *__restrict__ k, + const __mt_bfloat16 *__restrict__ v, + const __mt_bfloat16 *__restrict__ q_weight, + const __mt_bfloat16 *__restrict__ k_weight, + const int64_t *__restrict__ positions, + const __mt_bfloat16 *__restrict__ cos_sin_cache, + __mt_bfloat16 *__restrict__ q_out, __mt_bfloat16 *__restrict__ k_out, + __mt_bfloat16 *__restrict__ k_cache, __mt_bfloat16 *__restrict__ v_cache, + const index_t *__restrict__ indices, int batch, int q_heads, int k_heads, + int hidden, int rot_dim, int64_t position_stride, int64_t q_batch_stride, + int64_t q_head_stride, int64_t k_batch_stride, int64_t k_head_stride, + int64_t v_batch_stride, int64_t v_head_stride, int64_t q_out_batch_stride, + int64_t q_out_head_stride, int64_t k_out_batch_stride, + int64_t k_out_head_stride, int64_t k_cache_row_stride, + int64_t v_cache_row_stride, int64_t indices_stride, int mrope_section_t, + int mrope_section_h, int mrope_section_w, bool is_interleaved, float eps) { + constexpr int heads_per_block = 8; + const int tid = (int)threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int total_heads = q_heads + k_heads; + const int groups_per_token = + (total_heads + heads_per_block - 1) / heads_per_block; + const int token = (int)blockIdx.x / groups_per_token; + const int group = (int)blockIdx.x - token * groups_per_token; + if (token >= batch) { + return; + } + + int64_t cache_idx = 0; + if constexpr (WITH_CACHE) { + cache_idx = static_cast(indices[(int64_t)token * indices_stride]); + if (group == 0) { + const int row_dim = k_heads * hidden; + const int64_t v_base = (int64_t)token * v_batch_stride; + const int64_t v_dst = cache_idx * v_cache_row_stride; + for (int col = tid; col < row_dim; col += (int)blockDim.x) { + const int head = col / hidden; + const int head_col = col - head * hidden; + v_cache[v_dst + col] = + v[v_base + (int64_t)head * v_head_stride + head_col]; + } + } + } + + const int global_head = group * heads_per_block + warp; + if (global_head >= total_heads) { + return; + } + const bool is_q = global_head < q_heads; + const int head = is_q ? global_head : global_head - q_heads; + const __mt_bfloat16 *__restrict__ data = is_q ? q : k; + const __mt_bfloat16 *__restrict__ weight = is_q ? q_weight : k_weight; + const int64_t base = + is_q ? ((int64_t)token * q_batch_stride + (int64_t)head * q_head_stride) + : ((int64_t)token * k_batch_stride + (int64_t)head * k_head_stride); + const int embed_dim = rot_dim >> 1; + + float sum = 0.0f; + for (int col = lane; col < hidden; col += 32) { + const float value = __bfloat162float(data[base + col]); + sum += value * value; + } +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, mask); + } + const float scale = fast_rsqrt(sum / static_cast(hidden) + eps); + + for (int rot_offset = lane; rot_offset < embed_dim; rot_offset += 32) { + int axis; + if (is_interleaved) { + const bool use_h = + (rot_offset % 3 == 1) && (rot_offset <= 3 * mrope_section_h); + const bool use_w = + (rot_offset % 3 == 2) && (rot_offset <= 3 * mrope_section_w); + axis = use_h ? 1 : (use_w ? 2 : 0); + } else { + axis = rot_offset < mrope_section_t + ? 0 + : (rot_offset < mrope_section_t + mrope_section_h ? 1 : 2); + } + const int64_t pos = positions[(int64_t)axis * position_stride + token]; + const __mt_bfloat16 *__restrict__ cache_ptr = cos_sin_cache + pos * rot_dim; + const int x_index = rot_offset; + const int y_index = embed_dim + rot_offset; + const float cos_v = __bfloat162float(cache_ptr[x_index]); + const float sin_v = __bfloat162float(cache_ptr[y_index]); + const float x_weight = __bfloat162float(weight[x_index]) + + (GEMMA ? 1.0f : 0.0f); + const float y_weight = __bfloat162float(weight[y_index]) + + (GEMMA ? 1.0f : 0.0f); + const float x = __bfloat162float(data[base + x_index]) * scale * x_weight; + const float y = __bfloat162float(data[base + y_index]) * scale * y_weight; + const __mt_bfloat16 x_rot = __float2bfloat16_rn(x * cos_v - y * sin_v); + const __mt_bfloat16 y_rot = __float2bfloat16_rn(y * cos_v + x * sin_v); + if (is_q) { + const int64_t out_base = + (int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride; + q_out[out_base + x_index] = x_rot; + q_out[out_base + y_index] = y_rot; + } else if constexpr (WITH_CACHE) { + const int64_t cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + k_cache[cache_base + x_index] = x_rot; + k_cache[cache_base + y_index] = y_rot; + if (k_out != nullptr) { + const int64_t out_base = + (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[out_base + x_index] = x_rot; + k_out[out_base + y_index] = y_rot; + } + } else { + const int64_t out_base = + (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[out_base + x_index] = x_rot; + k_out[out_base + y_index] = y_rot; + } + } + + for (int col = rot_dim + lane; col < hidden; col += 32) { + const float weight_value = + __bfloat162float(weight[col]) + (GEMMA ? 1.0f : 0.0f); + const __mt_bfloat16 out_value = + __float2bfloat16_rn(__bfloat162float(data[base + col]) * scale * weight_value); + if (is_q) { + const int64_t out_base = + (int64_t)token * q_out_batch_stride + + (int64_t)head * q_out_head_stride; + q_out[out_base + col] = out_value; + } else if constexpr (WITH_CACHE) { + const int64_t cache_base = + cache_idx * k_cache_row_stride + (int64_t)head * hidden; + k_cache[cache_base + col] = out_value; + if (k_out != nullptr) { + const int64_t out_base = + (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[out_base + col] = out_value; + } + } else { + const int64_t out_base = + (int64_t)token * k_out_batch_stride + + (int64_t)head * k_out_head_stride; + k_out[out_base + col] = out_value; + } + } +} + +void launch_qk_mrope_kernel(ffi::TensorView q, ffi::TensorView k, + ffi::TensorView q_weight, + ffi::TensorView k_weight, + ffi::TensorView positions, + ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, + int batch, bool gemma, bool is_interleaved, float eps, + musaStream_t stream) { + constexpr int threads = 256; + constexpr int groups_per_token = 5; + if (gemma && is_interleaved) { + fused_qk_rmsnorm_mrope_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); + return; + } + if (gemma) { + fused_qk_rmsnorm_mrope_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); + return; + } + if (is_interleaved) { + fused_qk_rmsnorm_mrope_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); + return; + } + fused_qk_rmsnorm_mrope_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); +} + +template +void launch_h256r64_mrope_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView q_weight, + ffi::TensorView k_weight, ffi::TensorView positions, + ffi::TensorView cos_sin_cache, ffi::TensorView q_out, + ffi::TensorView k_out, int batch, float eps, musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_mrope_h256r64_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); +} + +template +void launch_h128_full_mrope_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView q_weight, + ffi::TensorView k_weight, ffi::TensorView positions, + ffi::TensorView cos_sin_cache, ffi::TensorView q_out, + ffi::TensorView k_out, int batch, float eps, musaStream_t stream) { + constexpr int warps_per_block = HEADS_PER_BLOCK; + constexpr int heads_per_block = warps_per_block * 2; + constexpr int threads = 32 * warps_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_mrope_h128_full_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); +} + +template +void launch_h128_full_halfwarp_mrope_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView q_weight, + ffi::TensorView k_weight, ffi::TensorView positions, + ffi::TensorView cos_sin_cache, ffi::TensorView q_out, + ffi::TensorView k_out, int batch, float eps, musaStream_t stream) { + constexpr int threads = 256; + constexpr int heads_per_block = 16; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_mrope_h128_full_halfwarp_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), eps); +} + +void launch_fused_qk_rmsnorm_mrope_bf16( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView q_weight, + ffi::TensorView k_weight, ffi::TensorView positions, + ffi::TensorView cos_sin_cache, ffi::TensorView q_out, ffi::TensorView k_out, + int mrope_section_t, int mrope_section_h, int mrope_section_w, + bool is_interleaved, bool gemma, float eps) { + const int batch = static_cast(q.size(0)); + if (batch == 0) { + return; + } + musaStream_t stream = get_stream(q.device()); + const int q_heads = static_cast(q.size(1)); + const int k_heads = static_cast(k.size(1)); + const int hidden = static_cast(q.size(2)); + const int rot_dim = static_cast(cos_sin_cache.size(1)); + const bool use_32q4kv_h128_mrope_specialization = + q_heads == 32 && k_heads == 4 && hidden == 128 && rot_dim == 128 && + mrope_section_t == 24 && mrope_section_h == 20 && + mrope_section_w == 20; + const bool use_h256r64_mrope_specialization = + gemma && is_interleaved && hidden == 256 && rot_dim == 64 && + mrope_section_t == 11 && mrope_section_h == 11 && + mrope_section_w == 10; + if (use_h256r64_mrope_specialization) { + if (q_heads == 4 && k_heads == 1) { + launch_h256r64_mrope_kernel<4, 1>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 8 && k_heads == 1) { + launch_h256r64_mrope_kernel<8, 1, 16>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 1) { + launch_h256r64_mrope_kernel<16, 1>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 2) { + launch_h256r64_mrope_kernel<16, 2>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 4) { + launch_h256r64_mrope_kernel<16, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 24 && k_heads == 4) { + launch_h256r64_mrope_kernel<24, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 2) { + launch_h256r64_mrope_kernel<32, 2>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 4) { + launch_h256r64_mrope_kernel<32, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 4) { + launch_h256r64_mrope_kernel<64, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 8) { + launch_h256r64_mrope_kernel<64, 8>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + } + if (use_32q4kv_h128_mrope_specialization && batch < 128) { + launch_qk_mrope_kernel(q, k, q_weight, k_weight, positions, cos_sin_cache, + q_out, k_out, batch, gemma, is_interleaved, eps, + stream); + return; + } + const bool use_h128_full_mrope_specialization = + gemma && is_interleaved && hidden == 128 && rot_dim == 128 && + mrope_section_t == 24 && mrope_section_h == 20 && + mrope_section_w == 20; + if (use_h128_full_mrope_specialization) { + if (q_heads == 2 && k_heads == 2) { + launch_h128_full_mrope_kernel<2, 2>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 4 && k_heads == 4) { + launch_h128_full_halfwarp_mrope_kernel<4, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 8 && k_heads == 8) { + launch_h128_full_halfwarp_mrope_kernel<8, 8>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 32) { + launch_h128_full_halfwarp_mrope_kernel<32, 32>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 16) { + launch_h128_full_halfwarp_mrope_kernel<16, 16>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 4) { + launch_h128_full_halfwarp_mrope_kernel<32, 4>( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + batch, eps, stream); + return; + } + } + if (!use_32q4kv_h128_mrope_specialization) { + constexpr int threads = 256; + constexpr int heads_per_block = 8; + const int groups_per_token = + (q_heads + k_heads + heads_per_block - 1) / heads_per_block; + if (gemma) { + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(nullptr), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast(nullptr), batch, q_heads, k_heads, + hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), static_cast(0), + static_cast(0), static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), static_cast(0), + static_cast(0), static_cast(0), + mrope_section_t, mrope_section_h, mrope_section_w, + is_interleaved, eps); + return; + } + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(nullptr), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast(nullptr), batch, q_heads, k_heads, + hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(0), static_cast(0), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), static_cast(0), + static_cast(0), static_cast(0), mrope_section_t, + mrope_section_h, mrope_section_w, is_interleaved, eps); + return; + } + launch_qk_mrope_kernel(q, k, q_weight, k_weight, positions, cos_sin_cache, + q_out, k_out, batch, gemma, is_interleaved, eps, + stream); +} + +template +void launch_qk_mrope_cache_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int batch, bool gemma, bool is_interleaved, float eps, + musaStream_t stream) { + constexpr int threads = 256; + constexpr int groups_per_token = 5; + if (gemma && is_interleaved) { + fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); + return; + } + if (gemma) { + fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); + return; + } + if (is_interleaved) { + fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); + return; + } + fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h128_full_mrope_cache_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int batch, float eps, musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; +#define LAUNCH_H128_FULL_MROPE_CACHE(SAME_POSITION_VALUE) \ + fused_qk_rmsnorm_mrope_cache_h128_full_bf16_kernel< \ + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, false, \ + SAME_POSITION_VALUE> \ + <<>>( \ + static_cast(q.data_ptr()), \ + static_cast(k.data_ptr()), \ + static_cast(v.data_ptr()), \ + static_cast(q_weight.data_ptr()), \ + static_cast(k_weight.data_ptr()), \ + static_cast(positions.data_ptr()), \ + static_cast(cos_sin_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, \ + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), \ + static_cast(indices.data_ptr()), batch, \ + static_cast(positions.stride(0)), \ + static_cast(q.stride(0)), \ + static_cast(q.stride(1)), \ + static_cast(k.stride(0)), \ + static_cast(k.stride(1)), \ + static_cast(v.stride(0)), \ + static_cast(v.stride(1)), \ + static_cast(q_out.stride(0)), \ + static_cast(q_out.stride(1)), 0, 0, \ + static_cast(k_cache.stride(0)), \ + static_cast(v_cache.stride(0)), \ + static_cast(indices.stride(0)), eps) + if (positions.stride(0) == 0) { + LAUNCH_H128_FULL_MROPE_CACHE(true); + } else { + LAUNCH_H128_FULL_MROPE_CACHE(false); + } +#undef LAUNCH_H128_FULL_MROPE_CACHE +} + +template +void launch_h128_full_mrope_cache_out_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int batch, float eps, + musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; +#define LAUNCH_H128_FULL_MROPE_CACHE_OUT(SAME_POSITION_VALUE) \ + fused_qk_rmsnorm_mrope_cache_h128_full_bf16_kernel< \ + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, true, \ + SAME_POSITION_VALUE> \ + <<>>( \ + static_cast(q.data_ptr()), \ + static_cast(k.data_ptr()), \ + static_cast(v.data_ptr()), \ + static_cast(q_weight.data_ptr()), \ + static_cast(k_weight.data_ptr()), \ + static_cast(positions.data_ptr()), \ + static_cast(cos_sin_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), \ + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), \ + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), \ + static_cast(indices.data_ptr()), batch, \ + static_cast(positions.stride(0)), \ + static_cast(q.stride(0)), \ + static_cast(q.stride(1)), \ + static_cast(k.stride(0)), \ + static_cast(k.stride(1)), \ + static_cast(v.stride(0)), \ + static_cast(v.stride(1)), \ + static_cast(q_out.stride(0)), \ + static_cast(q_out.stride(1)), \ + static_cast(k_out.stride(0)), \ + static_cast(k_out.stride(1)), \ + static_cast(k_cache.stride(0)), \ + static_cast(v_cache.stride(0)), \ + static_cast(indices.stride(0)), eps) + if (positions.stride(0) == 0) { + LAUNCH_H128_FULL_MROPE_CACHE_OUT(true); + } else { + LAUNCH_H128_FULL_MROPE_CACHE_OUT(false); + } +#undef LAUNCH_H128_FULL_MROPE_CACHE_OUT +} + +template +void launch_h128_full_rope_cache_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int batch, float eps, musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_rope_cache_h128_full_bf16_kernel< + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, false> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h128_full_rope_cache_out_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int batch, float eps, + musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_rope_cache_h128_full_bf16_kernel< + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, true> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h128r64_rope_cache_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int batch, float eps, musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_rope_cache_h128r64_bf16_kernel< + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, false> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h128r64_rope_cache_out_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int batch, float eps, + musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_rope_cache_h128r64_bf16_kernel< + index_t, Q_HEADS, K_HEADS, GEMMA, HEADS_PER_BLOCK, true> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h256r64_mrope_cache_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int batch, float eps, musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_mrope_cache_h256r64_bf16_kernel< + index_t, Q_HEADS, K_HEADS, HEADS_PER_BLOCK, false> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), nullptr, + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), 0, 0, + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_h256r64_mrope_cache_out_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int batch, float eps, + musaStream_t stream) { + constexpr int heads_per_block = HEADS_PER_BLOCK; + constexpr int threads = 32 * heads_per_block; + constexpr int groups_per_token = + (Q_HEADS + K_HEADS + heads_per_block - 1) / heads_per_block; + fused_qk_rmsnorm_mrope_cache_h256r64_bf16_kernel< + index_t, Q_HEADS, K_HEADS, HEADS_PER_BLOCK> + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, + static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), eps); +} + +template +void launch_fused_qk_rmsnorm_mrope_cache_bf16( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, int mrope_section_t, int mrope_section_h, + int mrope_section_w, bool is_interleaved, bool gemma, float eps) { + const int batch = static_cast(q.size(0)); + if (batch == 0) { + return; + } + musaStream_t stream = get_stream(q.device()); + const int q_heads = static_cast(q.size(1)); + const int k_heads = static_cast(k.size(1)); + const int hidden = static_cast(q.size(2)); + const int rot_dim = static_cast(cos_sin_cache.size(1)); + const bool use_32q4kv_h128_mrope_specialization = + q_heads == 32 && k_heads == 4 && hidden == 128 && rot_dim == 128 && + mrope_section_t == 24 && mrope_section_h == 20 && + mrope_section_w == 20; + const bool use_h256r64_mrope_specialization = + gemma && is_interleaved && hidden == 256 && rot_dim == 64 && + mrope_section_t == 11 && mrope_section_h == 11 && + mrope_section_w == 10; + if (use_h256r64_mrope_specialization) { + if (q_heads == 4 && k_heads == 1) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 8 && k_heads == 1) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 1) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 2) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 4) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 24 && k_heads == 4) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 2) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 4) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 4) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 8) { + launch_h256r64_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, eps, stream); + return; + } + } + const bool use_h128_full_mrope_specialization = + is_interleaved && hidden == 128 && rot_dim == 128 && mrope_section_t == 24 && + mrope_section_h == 20 && mrope_section_w == 20; + if (use_h128_full_mrope_specialization) { +#define LAUNCH_H128_MROPE_CACHE(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128_full_mrope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128_full_mrope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128_MROPE_CACHE(2, 1); + LAUNCH_H128_MROPE_CACHE(4, 1); + LAUNCH_H128_MROPE_CACHE(4, 2); + LAUNCH_H128_MROPE_CACHE(8, 2); + LAUNCH_H128_MROPE_CACHE(8, 4); + LAUNCH_H128_MROPE_CACHE(16, 4); + LAUNCH_H128_MROPE_CACHE(16, 8); + LAUNCH_H128_MROPE_CACHE(32, 8); + LAUNCH_H128_MROPE_CACHE(2, 2); + LAUNCH_H128_MROPE_CACHE(4, 4); + LAUNCH_H128_MROPE_CACHE(8, 8); + LAUNCH_H128_MROPE_CACHE(16, 16); + LAUNCH_H128_MROPE_CACHE(32, 32); +#undef LAUNCH_H128_MROPE_CACHE + } + const bool use_h128_full_rope_specialization = + hidden == 128 && rot_dim == 128 && mrope_section_t == 64 && + mrope_section_h == 0 && mrope_section_w == 0 && !is_interleaved; + if (use_h128_full_rope_specialization) { +#define LAUNCH_H128_ROPE_CACHE(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128_full_rope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128_full_rope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128_ROPE_CACHE(8, 8); + LAUNCH_H128_ROPE_CACHE(16, 8); + LAUNCH_H128_ROPE_CACHE(24, 8); + LAUNCH_H128_ROPE_CACHE(32, 8); + LAUNCH_H128_ROPE_CACHE(40, 8); + LAUNCH_H128_ROPE_CACHE(48, 8); + LAUNCH_H128_ROPE_CACHE(64, 8); + LAUNCH_H128_ROPE_CACHE(5, 1); + LAUNCH_H128_ROPE_CACHE(8, 1); + LAUNCH_H128_ROPE_CACHE(10, 2); + LAUNCH_H128_ROPE_CACHE(16, 1); + LAUNCH_H128_ROPE_CACHE(16, 2); + LAUNCH_H128_ROPE_CACHE(16, 4); + LAUNCH_H128_ROPE_CACHE(20, 4); + LAUNCH_H128_ROPE_CACHE(24, 4); + LAUNCH_H128_ROPE_CACHE(32, 2); + LAUNCH_H128_ROPE_CACHE(32, 4); + LAUNCH_H128_ROPE_CACHE(40, 4); + LAUNCH_H128_ROPE_CACHE(4, 1); + LAUNCH_H128_ROPE_CACHE(8, 2); +#undef LAUNCH_H128_ROPE_CACHE + } + const bool use_h128r64_rope_specialization = + hidden == 128 && rot_dim == 64 && mrope_section_t == 32 && + mrope_section_h == 0 && mrope_section_w == 0 && !is_interleaved; + if (use_h128r64_rope_specialization) { +#define LAUNCH_H128R64_ROPE_CACHE(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128r64_rope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128r64_rope_cache_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128R64_ROPE_CACHE(6, 1); + LAUNCH_H128R64_ROPE_CACHE(8, 1); + LAUNCH_H128R64_ROPE_CACHE(12, 1); + LAUNCH_H128R64_ROPE_CACHE(12, 2); + LAUNCH_H128R64_ROPE_CACHE(16, 2); + LAUNCH_H128R64_ROPE_CACHE(24, 2); + LAUNCH_H128R64_ROPE_CACHE(24, 4); + LAUNCH_H128R64_ROPE_CACHE(32, 4); + LAUNCH_H128R64_ROPE_CACHE(48, 4); + LAUNCH_H128R64_ROPE_CACHE(48, 8); + LAUNCH_H128R64_ROPE_CACHE(64, 8); + LAUNCH_H128R64_ROPE_CACHE(96, 8); +#undef LAUNCH_H128R64_ROPE_CACHE + } + if (!use_32q4kv_h128_mrope_specialization) { + constexpr int threads = 256; + constexpr int heads_per_block = 8; + const int groups_per_token = + (q_heads + k_heads + heads_per_block - 1) / heads_per_block; + if (gemma) { + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, q_heads, + k_heads, hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), static_cast(0), + static_cast(0), static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), mrope_section_t, + mrope_section_h, mrope_section_w, is_interleaved, eps); + return; + } + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(nullptr), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, q_heads, + k_heads, hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), static_cast(0), + static_cast(0), static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), mrope_section_t, + mrope_section_h, mrope_section_w, is_interleaved, eps); + return; + } + launch_qk_mrope_cache_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, batch, gemma, is_interleaved, eps, stream); +} + +template +void launch_qk_mrope_cache_out_kernel( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int batch, bool gemma, + bool is_interleaved, float eps, musaStream_t stream) { + constexpr int threads = 256; + constexpr int groups_per_token = 5; +#define LAUNCH_CACHE_OUT_32Q4(GEMMA_VALUE, INTERLEAVED_VALUE) \ + fused_qk_rmsnorm_mrope_cache_32q4kv_h128_bf16_kernel< \ + index_t, GEMMA_VALUE, INTERLEAVED_VALUE> \ + <<>>( \ + static_cast(q.data_ptr()), \ + static_cast(k.data_ptr()), \ + static_cast(v.data_ptr()), \ + static_cast(q_weight.data_ptr()), \ + static_cast(k_weight.data_ptr()), \ + static_cast(positions.data_ptr()), \ + static_cast(cos_sin_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), \ + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), \ + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), \ + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), \ + static_cast(indices.data_ptr()), batch, \ + static_cast(positions.stride(0)), \ + static_cast(q.stride(0)), \ + static_cast(q.stride(1)), \ + static_cast(k.stride(0)), \ + static_cast(k.stride(1)), \ + static_cast(v.stride(0)), \ + static_cast(v.stride(1)), \ + static_cast(q_out.stride(0)), \ + static_cast(q_out.stride(1)), \ + static_cast(k_out.stride(0)), \ + static_cast(k_out.stride(1)), \ + static_cast(k_cache.stride(0)), \ + static_cast(v_cache.stride(0)), \ + static_cast(indices.stride(0)), eps) + + if (gemma && is_interleaved) { + LAUNCH_CACHE_OUT_32Q4(true, true); + return; + } + if (gemma) { + LAUNCH_CACHE_OUT_32Q4(true, false); + return; + } + if (is_interleaved) { + LAUNCH_CACHE_OUT_32Q4(false, true); + return; + } + LAUNCH_CACHE_OUT_32Q4(false, false); +#undef LAUNCH_CACHE_OUT_32Q4 +} + +void sgl_musa_fused_qk_rmsnorm_mrope( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView q_weight, + ffi::TensorView k_weight, ffi::TensorView positions, + ffi::TensorView cos_sin_cache, ffi::TensorView q_out, ffi::TensorView k_out, + bool is_neox, int mrope_section_t, int mrope_section_h, + int mrope_section_w, bool is_interleaved, double eps, bool gemma) { + check_fused_qk_mrope_inputs(q, k, q_weight, k_weight, positions, + cos_sin_cache, q_out, k_out); + TVM_FFI_ICHECK(is_neox) << "MUSA fused MRoPE path requires NeoX style"; + TVM_FFI_ICHECK_EQ(mrope_section_t + mrope_section_h + mrope_section_w, + cos_sin_cache.size(1) / 2); + ffi::MUSADeviceGuard device_guard(q.device().device_id); + launch_fused_qk_rmsnorm_mrope_bf16( + q, k, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + mrope_section_t, mrope_section_h, mrope_section_w, is_interleaved, gemma, + static_cast(eps)); + const musaError_t err = musaGetLastError(); + TVM_FFI_ICHECK_EQ(err, musaSuccess) + << "MUSA fused_qk_rmsnorm_mrope kernel failed: " + << musaGetErrorString(err); +} + +void sgl_musa_fused_qk_rmsnorm_mrope_cache( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_cache, ffi::TensorView v_cache, + ffi::TensorView indices, bool is_neox, int mrope_section_t, + int mrope_section_h, int mrope_section_w, bool is_interleaved, double eps, + bool gemma) { + check_fused_qk_mrope_cache_inputs(q, k, v, q_weight, k_weight, positions, + cos_sin_cache, q_out, k_cache, v_cache, + indices); + TVM_FFI_ICHECK(is_neox) + << "MUSA fused MRoPE cache path requires NeoX style"; + TVM_FFI_ICHECK_EQ(mrope_section_t + mrope_section_h + mrope_section_w, + cos_sin_cache.size(1) / 2); + ffi::MUSADeviceGuard device_guard(q.device().device_id); + if (dtype_equal(indices.dtype(), dl_int32)) { + launch_fused_qk_rmsnorm_mrope_cache_bf16( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, mrope_section_t, mrope_section_h, mrope_section_w, + is_interleaved, gemma, static_cast(eps)); + } else if (dtype_equal(indices.dtype(), dl_int64)) { + launch_fused_qk_rmsnorm_mrope_cache_bf16( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_cache, + v_cache, indices, mrope_section_t, mrope_section_h, mrope_section_w, + is_interleaved, gemma, static_cast(eps)); + } else { + TVM_FFI_THROW(ValueError) << "indices must be int32 or int64"; + } + const musaError_t err = musaGetLastError(); + TVM_FFI_ICHECK_EQ(err, musaSuccess) + << "MUSA fused_qk_rmsnorm_mrope_cache kernel failed: " + << musaGetErrorString(err); +} + +template +void launch_fused_qk_rmsnorm_mrope_cache_out_bf16( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, int mrope_section_t, + int mrope_section_h, int mrope_section_w, bool is_interleaved, bool gemma, + float eps) { + const int batch = static_cast(q.size(0)); + if (batch == 0) { + return; + } + musaStream_t stream = get_stream(q.device()); + const int q_heads = static_cast(q.size(1)); + const int k_heads = static_cast(k.size(1)); + const int hidden = static_cast(q.size(2)); + const int rot_dim = static_cast(cos_sin_cache.size(1)); + const bool use_32q4kv_h128_mrope_specialization = + q_heads == 32 && k_heads == 4 && hidden == 128 && rot_dim == 128 && + mrope_section_t == 24 && mrope_section_h == 20 && + mrope_section_w == 20; + const bool use_h256r64_mrope_specialization = + gemma && is_interleaved && hidden == 256 && rot_dim == 64 && + mrope_section_t == 11 && mrope_section_h == 11 && + mrope_section_w == 10; + if (use_h256r64_mrope_specialization) { + if (q_heads == 4 && k_heads == 1) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 8 && k_heads == 1) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 1) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 2) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 16 && k_heads == 4) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 24 && k_heads == 4) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 2) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 32 && k_heads == 4) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 4) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + if (q_heads == 64 && k_heads == 8) { + launch_h256r64_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, eps, stream); + return; + } + } + const bool use_h128_full_mrope_specialization = + is_interleaved && hidden == 128 && rot_dim == 128 && mrope_section_t == 24 && + mrope_section_h == 20 && mrope_section_w == 20; + if (use_h128_full_mrope_specialization) { +#define LAUNCH_H128_MROPE_CACHE_OUT(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128_full_mrope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128_full_mrope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128_MROPE_CACHE_OUT(2, 1); + LAUNCH_H128_MROPE_CACHE_OUT(4, 1); + LAUNCH_H128_MROPE_CACHE_OUT(4, 2); + LAUNCH_H128_MROPE_CACHE_OUT(8, 2); + LAUNCH_H128_MROPE_CACHE_OUT(8, 4); + LAUNCH_H128_MROPE_CACHE_OUT(16, 4); + LAUNCH_H128_MROPE_CACHE_OUT(16, 8); + LAUNCH_H128_MROPE_CACHE_OUT(32, 8); + LAUNCH_H128_MROPE_CACHE_OUT(2, 2); + LAUNCH_H128_MROPE_CACHE_OUT(4, 4); + LAUNCH_H128_MROPE_CACHE_OUT(8, 8); + LAUNCH_H128_MROPE_CACHE_OUT(16, 16); + LAUNCH_H128_MROPE_CACHE_OUT(32, 32); +#undef LAUNCH_H128_MROPE_CACHE_OUT + } + const bool use_h128_full_rope_specialization = + hidden == 128 && rot_dim == 128 && mrope_section_t == 64 && + mrope_section_h == 0 && mrope_section_w == 0 && !is_interleaved; + if (use_h128_full_rope_specialization) { +#define LAUNCH_H128_ROPE_CACHE_OUT(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128_full_rope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128_full_rope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128_ROPE_CACHE_OUT(8, 8); + LAUNCH_H128_ROPE_CACHE_OUT(16, 8); + LAUNCH_H128_ROPE_CACHE_OUT(24, 8); + LAUNCH_H128_ROPE_CACHE_OUT(32, 8); + LAUNCH_H128_ROPE_CACHE_OUT(40, 8); + LAUNCH_H128_ROPE_CACHE_OUT(48, 8); + LAUNCH_H128_ROPE_CACHE_OUT(64, 8); + LAUNCH_H128_ROPE_CACHE_OUT(5, 1); + LAUNCH_H128_ROPE_CACHE_OUT(8, 1); + LAUNCH_H128_ROPE_CACHE_OUT(10, 2); + LAUNCH_H128_ROPE_CACHE_OUT(16, 1); + LAUNCH_H128_ROPE_CACHE_OUT(16, 2); + LAUNCH_H128_ROPE_CACHE_OUT(16, 4); + LAUNCH_H128_ROPE_CACHE_OUT(20, 4); + LAUNCH_H128_ROPE_CACHE_OUT(24, 4); + LAUNCH_H128_ROPE_CACHE_OUT(32, 2); + LAUNCH_H128_ROPE_CACHE_OUT(32, 4); + LAUNCH_H128_ROPE_CACHE_OUT(40, 4); + LAUNCH_H128_ROPE_CACHE_OUT(4, 1); + LAUNCH_H128_ROPE_CACHE_OUT(8, 2); +#undef LAUNCH_H128_ROPE_CACHE_OUT + } + const bool use_h128r64_rope_specialization = + hidden == 128 && rot_dim == 64 && mrope_section_t == 32 && + mrope_section_h == 0 && mrope_section_w == 0 && !is_interleaved; + if (use_h128r64_rope_specialization) { +#define LAUNCH_H128R64_ROPE_CACHE_OUT(QH, KH) \ + do { \ + if (q_heads == (QH) && k_heads == (KH)) { \ + if (gemma) { \ + launch_h128r64_rope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } else { \ + launch_h128r64_rope_cache_out_kernel( \ + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, \ + k_out, k_cache, v_cache, indices, batch, eps, stream); \ + } \ + return; \ + } \ + } while (0) + LAUNCH_H128R64_ROPE_CACHE_OUT(6, 1); + LAUNCH_H128R64_ROPE_CACHE_OUT(8, 1); + LAUNCH_H128R64_ROPE_CACHE_OUT(12, 1); + LAUNCH_H128R64_ROPE_CACHE_OUT(12, 2); + LAUNCH_H128R64_ROPE_CACHE_OUT(16, 2); + LAUNCH_H128R64_ROPE_CACHE_OUT(24, 2); + LAUNCH_H128R64_ROPE_CACHE_OUT(24, 4); + LAUNCH_H128R64_ROPE_CACHE_OUT(32, 4); + LAUNCH_H128R64_ROPE_CACHE_OUT(48, 4); + LAUNCH_H128R64_ROPE_CACHE_OUT(48, 8); + LAUNCH_H128R64_ROPE_CACHE_OUT(64, 8); + LAUNCH_H128R64_ROPE_CACHE_OUT(96, 8); +#undef LAUNCH_H128R64_ROPE_CACHE_OUT + } + if (use_32q4kv_h128_mrope_specialization) { + launch_qk_mrope_cache_out_kernel( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, batch, gemma, is_interleaved, eps, stream); + return; + } + constexpr int threads = 256; + constexpr int heads_per_block = 8; + const int groups_per_token = + (q_heads + k_heads + heads_per_block - 1) / heads_per_block; + if (gemma) { + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, q_heads, + k_heads, hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), + static_cast(q.stride(1)), + static_cast(k.stride(0)), + static_cast(k.stride(1)), + static_cast(v.stride(0)), + static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), mrope_section_t, + mrope_section_h, mrope_section_w, is_interleaved, eps); + return; + } + fused_qk_rmsnorm_mrope_generic_bf16_kernel + <<>>( + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(v.data_ptr()), + static_cast(q_weight.data_ptr()), + static_cast(k_weight.data_ptr()), + static_cast(positions.data_ptr()), + static_cast(cos_sin_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(q_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_out.data_ptr()), + static_cast<__mt_bfloat16 *>(k_cache.data_ptr()), + static_cast<__mt_bfloat16 *>(v_cache.data_ptr()), + static_cast(indices.data_ptr()), batch, q_heads, + k_heads, hidden, rot_dim, static_cast(positions.stride(0)), + static_cast(q.stride(0)), static_cast(q.stride(1)), + static_cast(k.stride(0)), static_cast(k.stride(1)), + static_cast(v.stride(0)), static_cast(v.stride(1)), + static_cast(q_out.stride(0)), + static_cast(q_out.stride(1)), + static_cast(k_out.stride(0)), + static_cast(k_out.stride(1)), + static_cast(k_cache.stride(0)), + static_cast(v_cache.stride(0)), + static_cast(indices.stride(0)), mrope_section_t, + mrope_section_h, mrope_section_w, is_interleaved, eps); +} + +void sgl_musa_fused_qk_rmsnorm_mrope_cache_out( + ffi::TensorView q, ffi::TensorView k, ffi::TensorView v, + ffi::TensorView q_weight, ffi::TensorView k_weight, + ffi::TensorView positions, ffi::TensorView cos_sin_cache, + ffi::TensorView q_out, ffi::TensorView k_out, ffi::TensorView k_cache, + ffi::TensorView v_cache, ffi::TensorView indices, bool is_neox, + int mrope_section_t, int mrope_section_h, int mrope_section_w, + bool is_interleaved, double eps, bool gemma) { + check_fused_qk_mrope_inputs(q, k, q_weight, k_weight, positions, + cos_sin_cache, q_out, k_out); + check_fused_qk_mrope_cache_inputs(q, k, v, q_weight, k_weight, positions, + cos_sin_cache, q_out, k_cache, v_cache, + indices); + TVM_FFI_ICHECK(is_neox) + << "MUSA fused MRoPE cache-out path requires NeoX style"; + TVM_FFI_ICHECK_EQ(mrope_section_t + mrope_section_h + mrope_section_w, + cos_sin_cache.size(1) / 2); + ffi::MUSADeviceGuard device_guard(q.device().device_id); + if (dtype_equal(indices.dtype(), dl_int32)) { + launch_fused_qk_rmsnorm_mrope_cache_out_bf16( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, mrope_section_t, mrope_section_h, + mrope_section_w, is_interleaved, gemma, static_cast(eps)); + } else if (dtype_equal(indices.dtype(), dl_int64)) { + launch_fused_qk_rmsnorm_mrope_cache_out_bf16( + q, k, v, q_weight, k_weight, positions, cos_sin_cache, q_out, k_out, + k_cache, v_cache, indices, mrope_section_t, mrope_section_h, + mrope_section_w, is_interleaved, gemma, static_cast(eps)); + } else { + TVM_FFI_THROW(ValueError) << "indices must be int32 or int64"; + } + const musaError_t err = musaGetLastError(); + TVM_FFI_ICHECK_EQ(err, musaSuccess) + << "MUSA fused_qk_rmsnorm_mrope_cache_out kernel failed: " + << musaGetErrorString(err); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_musa_fused_qk_rmsnorm_mrope, + sgl_musa_fused_qk_rmsnorm_mrope); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_musa_fused_qk_rmsnorm_mrope_cache, + sgl_musa_fused_qk_rmsnorm_mrope_cache); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_musa_fused_qk_rmsnorm_mrope_cache_out, + sgl_musa_fused_qk_rmsnorm_mrope_cache_out); diff --git a/vllm_musa/jit_kernel/csrc/quant.py b/vllm_musa/jit_kernel/csrc/quant.py new file mode 100644 index 000000000000..6e75d334e147 --- /dev/null +++ b/vllm_musa/jit_kernel/csrc/quant.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import torch +from vllm.utils.torch_utils import direct_register_custom_op + +from vllm_musa.jit_kernel.csrc.jit import load_musa_jit +from vllm_musa.jit_kernel.utils import cache_once + +_SILU_ACTIVATION_TYPE = 0 + + +@cache_once +def _quant_v2_module(): + return load_musa_jit( + "vllm_musa_quant_v2", + ("quant/per_token_group_quant_8bit_v2.mu",), + extra_musa_cflags=( + "-fmusa-flush-denormals-to-zero", + "-fno-signed-zeros", + "-mllvm", + "-mtgpu-opt-level=1", + "-mllvm", + "-mtgpu-load-store-opt=1", + "-mllvm", + "-mtgpu-fold-global-ldst=1", + ), + ) + + +def per_token_group_quant_8bit( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, + group_size: int, + eps: float, + min_8bit: float, + max_8bit: float, + scale_ue8m0: bool = False, + fuse_silu_and_mul: bool = False, +) -> None: + """Per-token-group 8-bit quant, optionally fusing silu+mul over a gated input. + + With fuse_silu_and_mul, `input` is (tokens, 2 * hidden) and `output_q` is + (tokens, hidden): the gate/up halves are combined before quantizing. + """ + torch.ops.vllm.musa_csrc_per_token_group_quant_8bit( + input, + output_q, + output_s, + int(group_size), + float(eps), + float(min_8bit), + float(max_8bit), + bool(scale_ue8m0), + bool(fuse_silu_and_mul), + ) + + +_DUMMY_MASKED_M: dict[torch.device, torch.Tensor] = {} + + +def _dummy_masked_m(device: torch.device) -> torch.Tensor: + t = _DUMMY_MASKED_M.get(device) + if t is None: + t = torch.empty((1,), device=device, dtype=torch.int32) + _DUMMY_MASKED_M[device] = t + return t + + +def _per_token_group_quant_8bit_custom( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, + group_size: int, + eps: float, + min_8bit: float, + max_8bit: float, + scale_ue8m0: bool, + fuse_silu_and_mul: bool, +) -> None: + # The kernel takes a masked_m tensor for the masked-layout callers; the contiguous + # layout used here has no mask, so pass a cached dummy and flag it off. + masked_m = _dummy_masked_m(input.device) + _quant_v2_module().sgl_per_token_group_quant_8bit_v2( + input, + output_q, + output_s, + int(group_size), + float(eps), + float(min_8bit), + float(max_8bit), + bool(scale_ue8m0), + bool(fuse_silu_and_mul), + _SILU_ACTIVATION_TYPE, + masked_m, + False, + ) + + +def _per_token_group_quant_8bit_custom_fake( + input: torch.Tensor, + output_q: torch.Tensor, + output_s: torch.Tensor, + group_size: int, + eps: float, + min_8bit: float, + max_8bit: float, + scale_ue8m0: bool, + fuse_silu_and_mul: bool, +) -> None: + return + + +direct_register_custom_op( + op_name="musa_csrc_per_token_group_quant_8bit", + op_func=_per_token_group_quant_8bit_custom, + mutates_args=["output_q", "output_s"], + fake_impl=_per_token_group_quant_8bit_custom_fake, +) diff --git a/vllm_musa/jit_kernel/csrc/quant/per_token_group_quant_8bit_v2.mu b/vllm_musa/jit_kernel/csrc/quant/per_token_group_quant_8bit_v2.mu new file mode 100644 index 000000000000..fbb6b1e44f8d --- /dev/null +++ b/vllm_musa/jit_kernel/csrc/quant/per_token_group_quant_8bit_v2.mu @@ -0,0 +1,1398 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../device_utils.h" + +// Copied and modified from DeepEP +__forceinline__ __device__ float fast_pow2(int x) { + // We can ensure `-126 <= x and x <= 127` + uint32_t bits_x = (x + 127) << 23; + return *reinterpret_cast(&bits_x); +} + +// Copied and modified from DeepEP +__forceinline__ __device__ int fast_log2_ceil(float x) { + auto bits_x = *reinterpret_cast(&x); + auto exp_x = (bits_x >> 23) & 0xff; + auto man_bits = bits_x & ((1 << 23) - 1); + return exp_x - 127 + (man_bits != 0); +} + +// Copied and modified from DeepEP +template +__forceinline__ __device__ void calculate_fp8_scales(float amax, float& scale, float& scale_inv) { + constexpr float MAX_8BIT_INV = 1.0f / dtype_info::MAX; + if constexpr (ROUND_SCALE) { + auto exp_scale_inv = fast_log2_ceil(amax * MAX_8BIT_INV); + scale = fast_pow2(-exp_scale_inv); + scale_inv = fast_pow2(exp_scale_inv); + } else { + scale_inv = amax * MAX_8BIT_INV; + scale = dtype_info::MAX / amax; + } +} + +// Copied and modified from DeepEP +template > +__forceinline__ __device__ OUT_DTYPE_T extract_required_scale_format(float value) { + if constexpr (SCALE_UE8M0) { + return static_cast((*reinterpret_cast(&value)) >> 23); + } else { + return value; + } +} + +template +struct DtypeInfo; + +template <> +struct DtypeInfo { + static constexpr float MIN = -128; + static constexpr float MAX = 127; +}; + +template <> +struct DtypeInfo<__mt_fp8_e4m3> { + static constexpr float MIN = -448; + static constexpr float MAX = 448; +}; + +template +__device__ __forceinline__ int compute_input_group_start_offset( + int expert_idx, + int token_idx, + int hidden_dim_group_idx, + int hidden_size, + int num_tokens_per_expert, + int group_size) { + return expert_idx * num_tokens_per_expert * hidden_size * (FUSE_SILU_AND_MUL ? 2 : 1) + + token_idx * hidden_size * (FUSE_SILU_AND_MUL ? 2 : 1) + hidden_dim_group_idx * group_size; +} + +constexpr float LOCAL_ABSMAX_ABS = 1e-10; +constexpr uint32_t INPUT_PRIMARY_VEC_NUM_BYTES = 32; +constexpr uint32_t FP8_FAST_VEC_ELEMS = 4; +constexpr int BF16_FP8_G128_THREADS_PER_GROUP = 16; +constexpr int kSiluActivation = 0; +constexpr int kGeluActivation = 1; +constexpr int kGeluTanhActivation = 2; + +__device__ __forceinline__ float fast_erf_musa(float x) { + constexpr float kP = 0.3275911f; + constexpr float kA1 = 0.254829592f; + constexpr float kA2 = -0.284496736f; + constexpr float kA3 = 1.421413741f; + constexpr float kA4 = -1.453152027f; + constexpr float kA5 = 1.061405429f; + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fabsf(x); + const float t = 1.0f / (1.0f + kP * ax); + const float poly = + (((((kA5 * t + kA4) * t + kA3) * t + kA2) * t + kA1) * t); + return sign * (1.0f - poly * expf(-ax * ax)); +} + +__device__ __forceinline__ float gelu_exact_musa(float x) { + return 0.5f * x * (1.0f + fast_erf_musa(x * 0.7071067811865475f)); +} + +__device__ __forceinline__ float gelu_tanh_musa(float x) { + constexpr float kSqrt2OverPi = 0.7978845608028654f; + constexpr float kCoeff = 0.044715f; + return 0.5f * x * + (1.0f + tanhf(kSqrt2OverPi * (x + kCoeff * x * x * x))); +} + +__device__ __forceinline__ float fast_silu_musa(float x) { + const float half_x = 0.5f * x; + return half_x * (1.0f + tanhf(half_x)); +} + +__device__ __forceinline__ float apply_fused_activation( + float x, + int activation_type) { + if (activation_type == kGeluActivation) { + return gelu_exact_musa(x); + } + if (activation_type == kGeluTanhActivation) { + return gelu_tanh_musa(x); + } + return fast_silu_musa(x); +} + +struct RowScheduler { + static void compute_exec_config( + int threads_per_subwarp, + int num_local_experts, + int hidden_dim_num_groups, + int num_groups, + int& subwarps_per_block, + dim3& grid, + dim3& block) { + subwarps_per_block = ([=]() -> int { + if (hidden_dim_num_groups % 16 == 0) { + return 16; + } else if (hidden_dim_num_groups % 8 == 0) { + return 8; + } else if (hidden_dim_num_groups % 4 == 0) { + return 4; + } else if (hidden_dim_num_groups % 2 == 0) { + return 2; + } + return 1; + })(); + const int num_tokens_per_expert = num_groups / hidden_dim_num_groups; + grid = dim3(hidden_dim_num_groups / subwarps_per_block, num_tokens_per_expert); + block = dim3(subwarps_per_block * threads_per_subwarp); + } + + template + __device__ __forceinline__ static void execute( + const int subwarps_per_block, + const int hidden_dim_num_groups, + const int32_t* masked_m, + const int num_tokens_per_expert, + FUNC fn) { + constexpr int expert_idx = 0; + + const int64_t subwarp_id = threadIdx.x / THREADS_PER_SUBWARP; + const int lane_id = threadIdx.x % THREADS_PER_SUBWARP; + + const int token_idx = blockIdx.y; + const int hidden_dim_group_idx = blockIdx.x * subwarps_per_block + subwarp_id; + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int64_t input_group_start_offset = compute_input_group_start_offset( + expert_idx, token_idx, hidden_dim_group_idx, hidden_size, num_tokens_per_expert, GROUP_SIZE); + + fn(expert_idx, token_idx, hidden_dim_group_idx, lane_id, input_group_start_offset); + } +}; + +struct MaskedLayoutScheduler { + // TODO can be dynamically determined (which may be good when num rank is small) + static constexpr int TOKEN_DIM_BLOCK_NUM_PER_EXPERT = 1024; + + static void compute_exec_config( + int threads_per_subwarp, + int num_local_experts, + int hidden_dim_num_groups, + int num_groups, + int& subwarps_per_block, + dim3& grid, + dim3& block) { + subwarps_per_block = ([=]() -> int { + if (hidden_dim_num_groups % 16 == 0) { + return 16; + } else if (hidden_dim_num_groups % 8 == 0) { + return 8; + } else if (hidden_dim_num_groups % 4 == 0) { + return 4; + } else if (hidden_dim_num_groups % 2 == 0) { + return 2; + } + return 1; + })(); + TORCH_CHECK(hidden_dim_num_groups % subwarps_per_block == 0); + grid = dim3(hidden_dim_num_groups / subwarps_per_block, TOKEN_DIM_BLOCK_NUM_PER_EXPERT, num_local_experts); + block = dim3(subwarps_per_block * threads_per_subwarp); + } + + template + __device__ __forceinline__ static void execute( + const int subwarps_per_block, + const int hidden_dim_num_groups, + const int32_t* masked_m, + const int num_tokens_per_expert, + FUNC fn) { + const int64_t subwarp_id = threadIdx.x / THREADS_PER_SUBWARP; + const int lane_id = threadIdx.x % THREADS_PER_SUBWARP; + + const int expert_idx = blockIdx.z; + const int token_idx_start = blockIdx.y; + + const int64_t hidden_dim_group_idx = blockIdx.x * subwarps_per_block + subwarp_id; + + const int curr_expert_token_num = masked_m[expert_idx]; + + for (int token_idx = token_idx_start; token_idx < curr_expert_token_num; + token_idx += TOKEN_DIM_BLOCK_NUM_PER_EXPERT) { + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int64_t input_group_start_offset = compute_input_group_start_offset( + expert_idx, token_idx, hidden_dim_group_idx, hidden_size, num_tokens_per_expert, GROUP_SIZE); + fn(expert_idx, token_idx, hidden_dim_group_idx, lane_id, input_group_start_offset); + } + } +}; + +template < + typename SCHEDULER, + int GROUP_SIZE, + int THREADS_PER_SUBWARP, + typename T, + typename DST_DTYPE, + bool IS_COLUMN_MAJOR = false, + bool SCALE_UE8M0 = false, + bool FUSE_SILU_AND_MUL = false, + typename scale_packed_t = std::conditional_t> +__global__ void per_token_group_quant_8bit_kernel( + const T* __restrict__ input, + DST_DTYPE* __restrict__ output_q, + scale_packed_t* __restrict__ output_s, + const int32_t* __restrict__ masked_m, + const int subwarps_per_block, + const int hidden_dim_num_groups, + // TODO can this be removed? + const int scale_expert_stride, + const int scale_hidden_stride, + const int num_tokens_per_expert, + const int fused_activation_type) { + using dst_dtype_info = DtypeInfo; + using scale_element_t = std::conditional_t; + + SCHEDULER::template execute( + subwarps_per_block, + hidden_dim_num_groups, + masked_m, + num_tokens_per_expert, + [&](const int expert_idx, + const int token_idx, + const int hidden_dim_group_idx, + const int lane_id, + const int input_group_start_offset) { + constexpr uint32_t INPUT_PRIMARY_VEC_SIZE = INPUT_PRIMARY_VEC_NUM_BYTES / sizeof(T); + constexpr uint32_t INPUT_PRIMARY_INT4_SIZE = INPUT_PRIMARY_VEC_NUM_BYTES / sizeof(int4); + + const int offset_num_groups = expert_idx * num_tokens_per_expert * hidden_dim_num_groups + + token_idx * hidden_dim_num_groups + hidden_dim_group_idx; + + int4 input_primary_int4[INPUT_PRIMARY_INT4_SIZE]; + T* input_primary_vec = reinterpret_cast(input_primary_int4); + + int4 input_secondary_int4[INPUT_PRIMARY_INT4_SIZE]; + T* input_secondary_vec = reinterpret_cast(input_secondary_int4); + +#pragma unroll + for (uint32_t j = 0; j < INPUT_PRIMARY_INT4_SIZE; ++j) { + input_primary_int4[j] = ld_global_nc( + reinterpret_cast(input + input_group_start_offset + lane_id * INPUT_PRIMARY_VEC_SIZE) + j); + } + if constexpr (FUSE_SILU_AND_MUL) { + const int secondary_offset = hidden_dim_num_groups * GROUP_SIZE; +#pragma unroll + for (uint32_t j = 0; j < INPUT_PRIMARY_INT4_SIZE; ++j) { + input_secondary_int4[j] = ld_global_nc( + reinterpret_cast( + input + input_group_start_offset + lane_id * INPUT_PRIMARY_VEC_SIZE + secondary_offset) + + j); + } + } + + constexpr int num_elems_per_pack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); + scale_element_t* scale_output; + if constexpr (IS_COLUMN_MAJOR) { + constexpr int scale_token_stride = 1; + + const int hidden_idx_packed = hidden_dim_group_idx / num_elems_per_pack; + const int pack_idx = hidden_dim_group_idx % num_elems_per_pack; + scale_output = reinterpret_cast(output_s) + + (expert_idx * scale_expert_stride * num_elems_per_pack + + hidden_idx_packed * scale_hidden_stride * num_elems_per_pack + + token_idx * scale_token_stride * num_elems_per_pack + pack_idx); + } else { + scale_output = output_s + offset_num_groups; + } + + // can speed up if too slow + if constexpr (IS_COLUMN_MAJOR and SCALE_UE8M0) { + const int remainder_num_groups = hidden_dim_num_groups % num_elems_per_pack; + if ((remainder_num_groups != 0) and (hidden_dim_group_idx == hidden_dim_num_groups - 1) and + (lane_id < num_elems_per_pack - remainder_num_groups)) { + const int shift = 1 + lane_id; + *(scale_output + shift) = 0; + } + } + + float local_absmax = LOCAL_ABSMAX_ABS; + +#pragma unroll + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { + float val; + if constexpr (FUSE_SILU_AND_MUL) { + // TODO maybe vectorize + T val_lowprec = static_cast(apply_fused_activation( + static_cast(input_primary_vec[j]), + fused_activation_type)) * + input_secondary_vec[j]; + val = static_cast(val_lowprec); + input_primary_vec[j] = val_lowprec; + } else { + val = static_cast(input_primary_vec[j]); + } + + float abs_val = fabsf(val); + local_absmax = fmaxf(local_absmax, abs_val); + } + + local_absmax = GroupReduceMax(local_absmax, lane_id); + + float y_scale, y_scale_inv; + calculate_fp8_scales(local_absmax, y_scale, y_scale_inv); + if (lane_id == 0) { + *scale_output = extract_required_scale_format(y_scale_inv); + } + + int4 output_buf; + + if constexpr (std::is_same_v) { + const auto output_buf_ptr = reinterpret_cast<__mt_fp8x4_storage_t*>(&output_buf); + +#pragma unroll + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 4) { + float4 outputx4 = { + static_cast(input_primary_vec[j]) * y_scale, + static_cast(input_primary_vec[j + 1]) * y_scale, + static_cast(input_primary_vec[j + 2]) * y_scale, + static_cast(input_primary_vec[j + 3]) * y_scale}; + output_buf_ptr[j / 4] = __musa_cvt_float4_to_fp8x4(outputx4, __MT_SATFINITE, __MT_E4M3); + } + } else { + const auto output_buf_ptr = reinterpret_cast(&output_buf); + +#pragma unroll + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { + float val = static_cast(input_primary_vec[j]); + float q_val = fminf(fmaxf(val * y_scale, dst_dtype_info::MIN), dst_dtype_info::MAX); + output_buf_ptr[j] = DST_DTYPE(q_val); + } + } + + st_global( + reinterpret_cast(output_q + offset_num_groups * GROUP_SIZE + lane_id * INPUT_PRIMARY_VEC_SIZE), + output_buf); + }); +} + +template +__device__ __forceinline__ uint32_t pack_quant4(const float values[FP8_FAST_VEC_ELEMS], const float y_scale) { + if constexpr (std::is_same_v) { + float4 outputx4 = {values[0] * y_scale, values[1] * y_scale, values[2] * y_scale, values[3] * y_scale}; + return __musa_cvt_float4_to_fp8x4(outputx4, __MT_SATFINITE, __MT_E4M3); + } else { + DST_DTYPE packed[FP8_FAST_VEC_ELEMS]; +#pragma unroll + for (uint32_t j = 0; j < FP8_FAST_VEC_ELEMS; ++j) { + const float q_val = fminf(fmaxf(values[j] * y_scale, DtypeInfo::MIN), DtypeInfo::MAX); + packed[j] = DST_DTYPE(q_val); + } + return *reinterpret_cast(packed); + } +} + +template < + int GROUP_SIZE, + int THREADS_PER_GROUP, + typename DST_DTYPE, + bool IS_COLUMN_MAJOR = false, + bool FUSE_SILU_AND_MUL = false> +__global__ void per_token_group_quant_8bit_fast_kernel( + const half* __restrict__ input, + DST_DTYPE* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int scale_hidden_stride, + const int num_tokens_per_expert, + const int fused_activation_type) { + const int subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int lane_id = threadIdx.x - subwarp_id * THREADS_PER_GROUP; + const int subwarps_per_block = blockDim.x / THREADS_PER_GROUP; + + const int token_idx = blockIdx.y; + const int hidden_dim_group_idx = blockIdx.x * subwarps_per_block + subwarp_id; + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int input_group_start_offset = compute_input_group_start_offset( + 0, token_idx, hidden_dim_group_idx, hidden_size, num_tokens_per_expert, GROUP_SIZE); + const int elem_offset = lane_id * FP8_FAST_VEC_ELEMS; + + uint64_t input_primary_u64 = *reinterpret_cast(input + input_group_start_offset + elem_offset); + half* input_primary_vec = reinterpret_cast(&input_primary_u64); + + uint64_t input_secondary_u64; + half* input_secondary_vec = reinterpret_cast(&input_secondary_u64); + if constexpr (FUSE_SILU_AND_MUL) { + input_secondary_u64 = *reinterpret_cast(input + input_group_start_offset + hidden_size + elem_offset); + } + + float values[FP8_FAST_VEC_ELEMS]; + float local_absmax = LOCAL_ABSMAX_ABS; + +#pragma unroll + for (uint32_t j = 0; j < FP8_FAST_VEC_ELEMS; ++j) { + float val; + if constexpr (FUSE_SILU_AND_MUL) { + half val_lowprec = static_cast(fast_silu_musa(static_cast(input_primary_vec[j]))) * + input_secondary_vec[j]; + val = static_cast(val_lowprec); + values[j] = val; + } else { + val = static_cast(input_primary_vec[j]); + values[j] = val; + } + local_absmax = fmaxf(local_absmax, fabsf(val)); + } + + local_absmax = GroupReduceMax(local_absmax, lane_id); + const float y_scale_inv = local_absmax * (1.0f / DtypeInfo::MAX); + const float y_scale = DtypeInfo::MAX / local_absmax; + + if (lane_id == 0) { + if constexpr (IS_COLUMN_MAJOR) { + output_s[token_idx + hidden_dim_group_idx * scale_hidden_stride] = y_scale_inv; + } else { + output_s[token_idx * hidden_dim_num_groups + hidden_dim_group_idx] = y_scale_inv; + } + } + + const int output_offset = (token_idx * hidden_dim_num_groups + hidden_dim_group_idx) * GROUP_SIZE + elem_offset; + const uint32_t output_buf = pack_quant4(values, y_scale); + *reinterpret_cast(output_q + output_offset) = output_buf; +} + +template +__global__ void per_token_group_quant_8bit_i8_fused_tiled_kernel( + const half* __restrict__ input, + int8_t* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int num_tokens_per_expert) { + const int subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int lane_id = threadIdx.x - subwarp_id * THREADS_PER_GROUP; + const int token_tile_idx = subwarp_id / HIDDEN_GROUPS_PER_BLOCK; + const int hidden_tile_idx = subwarp_id - token_tile_idx * HIDDEN_GROUPS_PER_BLOCK; + + const int token_idx = blockIdx.y * TOKEN_GROUPS_PER_BLOCK + token_tile_idx; + const int hidden_dim_group_idx = blockIdx.x * HIDDEN_GROUPS_PER_BLOCK + hidden_tile_idx; + if (token_idx >= num_tokens_per_expert || hidden_dim_group_idx >= hidden_dim_num_groups) { + return; + } + + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int input_group_start_offset = token_idx * hidden_size * 2 + hidden_dim_group_idx * GROUP_SIZE; + const int elem_offset = lane_id * FP8_FAST_VEC_ELEMS; + + uint64_t input_primary_u64 = *reinterpret_cast(input + input_group_start_offset + elem_offset); + half* input_primary_vec = reinterpret_cast(&input_primary_u64); + uint64_t input_secondary_u64 = + *reinterpret_cast(input + input_group_start_offset + hidden_size + elem_offset); + half* input_secondary_vec = reinterpret_cast(&input_secondary_u64); + + float values[FP8_FAST_VEC_ELEMS]; + float local_absmax = LOCAL_ABSMAX_ABS; + +#pragma unroll + for (uint32_t j = 0; j < FP8_FAST_VEC_ELEMS; ++j) { + half val_lowprec = static_cast(fast_silu_musa(static_cast(input_primary_vec[j]))) * + input_secondary_vec[j]; + const float val = static_cast(val_lowprec); + values[j] = val; + local_absmax = fmaxf(local_absmax, fabsf(val)); + } + + local_absmax = GroupReduceMax(local_absmax, lane_id); + const float y_scale_inv = local_absmax * (1.0f / DtypeInfo::MAX); + const float y_scale = DtypeInfo::MAX / local_absmax; + + if (lane_id == 0) { + output_s[token_idx * hidden_dim_num_groups + hidden_dim_group_idx] = y_scale_inv; + } + + const int output_offset = (token_idx * hidden_dim_num_groups + hidden_dim_group_idx) * GROUP_SIZE + elem_offset; + const uint32_t output_buf = pack_quant4(values, y_scale); + *reinterpret_cast(output_q + output_offset) = output_buf; +} + +template +__global__ void per_token_group_quant_8bit_fp8_col_tiled_kernel( + const half* __restrict__ input, + __mt_fp8_e4m3* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int scale_hidden_stride, + const int num_tokens_per_expert) { + const int subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int lane_id = threadIdx.x - subwarp_id * THREADS_PER_GROUP; + const int token_tile_idx = subwarp_id / HIDDEN_GROUPS_PER_BLOCK; + const int hidden_tile_idx = subwarp_id - token_tile_idx * HIDDEN_GROUPS_PER_BLOCK; + + const int token_idx = blockIdx.y * TOKEN_GROUPS_PER_BLOCK + token_tile_idx; + const int hidden_dim_group_idx = blockIdx.x * HIDDEN_GROUPS_PER_BLOCK + hidden_tile_idx; + if (token_idx >= num_tokens_per_expert || hidden_dim_group_idx >= hidden_dim_num_groups) { + return; + } + + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int input_group_start_offset = + token_idx * hidden_size + hidden_dim_group_idx * GROUP_SIZE + lane_id * FP8_FAST_VEC_ELEMS; + + uint64_t input_primary_u64 = *reinterpret_cast(input + input_group_start_offset); + half* input_primary_vec = reinterpret_cast(&input_primary_u64); + + float values[FP8_FAST_VEC_ELEMS]; + float local_absmax = LOCAL_ABSMAX_ABS; + +#pragma unroll + for (uint32_t j = 0; j < FP8_FAST_VEC_ELEMS; ++j) { + const float val = static_cast(input_primary_vec[j]); + values[j] = val; + local_absmax = fmaxf(local_absmax, fabsf(val)); + } + + local_absmax = GroupReduceMax(local_absmax, lane_id); + const float y_scale_inv = local_absmax * (1.0f / DtypeInfo<__mt_fp8_e4m3>::MAX); + const float y_scale = DtypeInfo<__mt_fp8_e4m3>::MAX / local_absmax; + + if (lane_id == 0) { + output_s[token_idx + hidden_dim_group_idx * scale_hidden_stride] = y_scale_inv; + } + + const int output_offset = + (token_idx * hidden_dim_num_groups + hidden_dim_group_idx) * GROUP_SIZE + lane_id * FP8_FAST_VEC_ELEMS; + uint32_t output_buf; + auto output_buf_ptr = reinterpret_cast<__mt_fp8x4_storage_t*>(&output_buf); + float4 outputx4 = {values[0] * y_scale, values[1] * y_scale, values[2] * y_scale, values[3] * y_scale}; + output_buf_ptr[0] = __musa_cvt_float4_to_fp8x4(outputx4, __MT_SATFINITE, __MT_E4M3); + *reinterpret_cast(output_q + output_offset) = output_buf; +} + +template +__device__ __forceinline__ void quantize_bf16_fp8_g128_group( + const __mt_bfloat16* __restrict__ input, + __mt_fp8_e4m3* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int token_idx, + const int hidden_dim_group_idx, + const int fused_activation_type) { + constexpr int GROUP_SIZE = 128; + constexpr int ELEMS_PER_THREAD = GROUP_SIZE / THREADS_PER_GROUP; + + const int hidden_size = hidden_dim_num_groups * GROUP_SIZE; + const int lane_id = threadIdx.x % THREADS_PER_GROUP; + const int elem_offset = lane_id * ELEMS_PER_THREAD; + const int input_offset = + token_idx * hidden_size * (FUSE_SILU_AND_MUL ? 2 : 1) + hidden_dim_group_idx * GROUP_SIZE + elem_offset; + + float values[ELEMS_PER_THREAD]; + if constexpr (ELEMS_PER_THREAD >= 8) { + constexpr int INPUT_INT4_SIZE = ELEMS_PER_THREAD * sizeof(__mt_bfloat16) / sizeof(int4); + int4 input_int4[INPUT_INT4_SIZE]; + int4 input_secondary_int4[INPUT_INT4_SIZE]; + auto input_vec = reinterpret_cast<__mt_bfloat16*>(input_int4); + auto input_secondary_vec = reinterpret_cast<__mt_bfloat16*>(input_secondary_int4); +#pragma unroll + for (int j = 0; j < INPUT_INT4_SIZE; ++j) { + input_int4[j] = ld_global_nc(reinterpret_cast(input + input_offset) + j); + if constexpr (FUSE_SILU_AND_MUL) { + input_secondary_int4[j] = + ld_global_nc(reinterpret_cast(input + input_offset + hidden_size) + j); + } + } +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; ++j) { + if constexpr (FUSE_SILU_AND_MUL) { + __mt_bfloat16 val_lowprec = static_cast<__mt_bfloat16>( + fast_silu_musa(static_cast(input_vec[j]))) * + input_secondary_vec[j]; + values[j] = static_cast(val_lowprec); + } else { + values[j] = static_cast(input_vec[j]); + } + } + } else { + uint64_t input_u64 = *reinterpret_cast(input + input_offset); + auto input_vec = reinterpret_cast<__mt_bfloat16*>(&input_u64); + uint64_t input_secondary_u64; + auto input_secondary_vec = reinterpret_cast<__mt_bfloat16*>(&input_secondary_u64); + if constexpr (FUSE_SILU_AND_MUL) { + input_secondary_u64 = *reinterpret_cast(input + input_offset + hidden_size); + } +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; ++j) { + if constexpr (FUSE_SILU_AND_MUL) { + __mt_bfloat16 val_lowprec = static_cast<__mt_bfloat16>( + fast_silu_musa(static_cast(input_vec[j]))) * + input_secondary_vec[j]; + values[j] = static_cast(val_lowprec); + } else { + values[j] = static_cast(input_vec[j]); + } + } + } + + float local_absmax = LOCAL_ABSMAX_ABS; +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; ++j) { + const float val = values[j]; + local_absmax = fmaxf(local_absmax, fabsf(val)); + } + + local_absmax = GroupReduceMax(local_absmax, lane_id); + const float y_scale_inv = local_absmax * (1.0f / DtypeInfo<__mt_fp8_e4m3>::MAX); + const float y_scale = 1.0f / y_scale_inv; + + if (lane_id == 0) { + output_s[token_idx * hidden_dim_num_groups + hidden_dim_group_idx] = y_scale_inv; + } + + const int output_offset = (token_idx * hidden_dim_num_groups + hidden_dim_group_idx) * GROUP_SIZE + elem_offset; + if constexpr (ELEMS_PER_THREAD == 16) { + float4 outputx4_0 = {values[0] * y_scale, values[1] * y_scale, values[2] * y_scale, values[3] * y_scale}; + float4 outputx4_1 = {values[4] * y_scale, values[5] * y_scale, values[6] * y_scale, values[7] * y_scale}; + float4 outputx4_2 = {values[8] * y_scale, values[9] * y_scale, values[10] * y_scale, values[11] * y_scale}; + float4 outputx4_3 = {values[12] * y_scale, values[13] * y_scale, values[14] * y_scale, values[15] * y_scale}; + auto output_u32 = reinterpret_cast(output_q + output_offset); + output_u32[0] = __musa_cvt_float4_to_fp8x4(outputx4_0, __MT_SATFINITE, __MT_E4M3); + output_u32[1] = __musa_cvt_float4_to_fp8x4(outputx4_1, __MT_SATFINITE, __MT_E4M3); + output_u32[2] = __musa_cvt_float4_to_fp8x4(outputx4_2, __MT_SATFINITE, __MT_E4M3); + output_u32[3] = __musa_cvt_float4_to_fp8x4(outputx4_3, __MT_SATFINITE, __MT_E4M3); + } else { + const float4 outputx4_0 = {values[0] * y_scale, values[1] * y_scale, values[2] * y_scale, values[3] * y_scale}; + const uint32_t output0 = __musa_cvt_float4_to_fp8x4(outputx4_0, __MT_SATFINITE, __MT_E4M3); + if constexpr (ELEMS_PER_THREAD == 8) { + const float4 outputx4_1 = {values[4] * y_scale, values[5] * y_scale, values[6] * y_scale, values[7] * y_scale}; + const uint64_t output_buf = + static_cast(output0) | + (static_cast(__musa_cvt_float4_to_fp8x4(outputx4_1, __MT_SATFINITE, __MT_E4M3)) << 32); + st_global_b64_slc_new(reinterpret_cast(output_q + output_offset), output_buf); + } else { + *reinterpret_cast(output_q + output_offset) = output0; + } + } +} + +template +__global__ void per_token_group_quant_8bit_bf16_fp8_g128_row_kernel( + const __mt_bfloat16* __restrict__ input, + __mt_fp8_e4m3* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int fused_activation_type) { + const int subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int subwarps_per_block = blockDim.x / THREADS_PER_GROUP; + const int token_idx = blockIdx.y; + const int hidden_dim_group_idx = blockIdx.x * subwarps_per_block + subwarp_id; + quantize_bf16_fp8_g128_group( + input, output_q, output_s, hidden_dim_num_groups, token_idx, hidden_dim_group_idx, fused_activation_type); +} + +template +__global__ void per_token_group_quant_8bit_bf16_fp8_g128_row_tiled_kernel( + const __mt_bfloat16* __restrict__ input, + __mt_fp8_e4m3* __restrict__ output_q, + float* __restrict__ output_s, + const int hidden_dim_num_groups, + const int num_tokens_per_expert, + const int fused_activation_type) { + const int subwarp_id = threadIdx.x / THREADS_PER_GROUP; + const int token_tile_idx = subwarp_id / HIDDEN_GROUPS_PER_BLOCK; + const int hidden_tile_idx = subwarp_id - token_tile_idx * HIDDEN_GROUPS_PER_BLOCK; + + const int token_idx = blockIdx.y * TOKEN_GROUPS_PER_BLOCK + token_tile_idx; + const int hidden_dim_group_idx = blockIdx.x * HIDDEN_GROUPS_PER_BLOCK + hidden_tile_idx; + if (token_idx >= num_tokens_per_expert || hidden_dim_group_idx >= hidden_dim_num_groups) { + return; + } + + quantize_bf16_fp8_g128_group( + input, output_q, output_s, hidden_dim_num_groups, token_idx, hidden_dim_group_idx, fused_activation_type); +} + +inline int choose_subwarps_per_block(const int hidden_dim_num_groups) { + if (hidden_dim_num_groups % 16 == 0) { + return 16; + } else if (hidden_dim_num_groups % 8 == 0) { + return 8; + } else if (hidden_dim_num_groups % 4 == 0) { + return 4; + } else if (hidden_dim_num_groups % 2 == 0) { + return 2; + } + return 1; +} + +template +struct FastQuantConfig; + +template <> +struct FastQuantConfig<16> { + static constexpr int THREADS_PER_GROUP = 4; +}; + +template <> +struct FastQuantConfig<32> { + static constexpr int THREADS_PER_GROUP = 8; +}; + +template <> +struct FastQuantConfig<64> { + static constexpr int THREADS_PER_GROUP = 16; +}; + +template <> +struct FastQuantConfig<128> { + static constexpr int THREADS_PER_GROUP = 32; +}; + +template +void launch_fp8_fast_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int scale_hidden_stride, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + constexpr int THREADS_PER_GROUP = FastQuantConfig::THREADS_PER_GROUP; + const int subwarps_per_block = choose_subwarps_per_block(hidden_dim_num_groups); + dim3 grid(hidden_dim_num_groups / subwarps_per_block, num_tokens_per_expert); + dim3 block(subwarps_per_block * THREADS_PER_GROUP); + per_token_group_quant_8bit_fast_kernel< + GROUP_SIZE, + THREADS_PER_GROUP, + __mt_fp8_e4m3, + IS_COLUMN_MAJOR, + FUSE_SILU_AND_MUL><<>>( + static_cast(input.data_ptr()), + static_cast<__mt_fp8_e4m3*>(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert, + fused_activation_type); +} + +template +void dispatch_fp8_fast_kernel( + int64_t group_size, + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int scale_hidden_stride, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + switch (group_size) { + case 16: + launch_fp8_fast_kernel<16, IS_COLUMN_MAJOR, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, scale_hidden_stride, + num_tokens_per_expert, fused_activation_type, stream); + break; + case 32: + launch_fp8_fast_kernel<32, IS_COLUMN_MAJOR, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, scale_hidden_stride, + num_tokens_per_expert, fused_activation_type, stream); + break; + case 64: + launch_fp8_fast_kernel<64, IS_COLUMN_MAJOR, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, scale_hidden_stride, + num_tokens_per_expert, fused_activation_type, stream); + break; + case 128: + launch_fp8_fast_kernel<128, IS_COLUMN_MAJOR, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, scale_hidden_stride, + num_tokens_per_expert, fused_activation_type, stream); + break; + default: + TORCH_CHECK(false, "Unsupported group_size"); + } +} + +template +void launch_fp8_col_tiled_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int scale_hidden_stride, + int num_tokens_per_expert, + musaStream_t stream) { + constexpr int THREADS_PER_GROUP = FastQuantConfig::THREADS_PER_GROUP; + dim3 grid( + (hidden_dim_num_groups + HIDDEN_GROUPS_PER_BLOCK - 1) / HIDDEN_GROUPS_PER_BLOCK, + (num_tokens_per_expert + TOKEN_GROUPS_PER_BLOCK - 1) / TOKEN_GROUPS_PER_BLOCK); + dim3 block(THREADS_PER_GROUP * HIDDEN_GROUPS_PER_BLOCK * TOKEN_GROUPS_PER_BLOCK); + per_token_group_quant_8bit_fp8_col_tiled_kernel< + GROUP_SIZE, + THREADS_PER_GROUP, + HIDDEN_GROUPS_PER_BLOCK, + TOKEN_GROUPS_PER_BLOCK><<>>( + static_cast(input.data_ptr()), + static_cast<__mt_fp8_e4m3*>(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert); +} + +void dispatch_fp8_col_nonfused_kernel( + int64_t group_size, + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int scale_hidden_stride, + int num_tokens_per_expert, + musaStream_t stream) { + if (group_size == 32) { + launch_fp8_col_tiled_kernel<32, 4, 4>( + input, output_q, output_s, hidden_dim_num_groups, scale_hidden_stride, num_tokens_per_expert, stream); + return; + } + + dispatch_fp8_fast_kernel( + group_size, input, output_q, output_s, hidden_dim_num_groups, + scale_hidden_stride, num_tokens_per_expert, kSiluActivation, stream); +} + +template +void launch_bf16_fp8_g128_row_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + int subwarps_per_block = choose_subwarps_per_block(hidden_dim_num_groups); + if constexpr (THREADS_PER_GROUP == 32) { + subwarps_per_block = subwarps_per_block > 8 ? 8 : subwarps_per_block; + } + dim3 grid(hidden_dim_num_groups / subwarps_per_block, num_tokens_per_expert); + dim3 block(subwarps_per_block * THREADS_PER_GROUP); + per_token_group_quant_8bit_bf16_fp8_g128_row_kernel + <<>>( + static_cast<__mt_bfloat16*>(input.data_ptr()), + static_cast<__mt_fp8_e4m3*>(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + fused_activation_type); +} + +template +void launch_bf16_fp8_g128_row_tiled_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + dim3 grid( + (hidden_dim_num_groups + HIDDEN_GROUPS_PER_BLOCK - 1) / HIDDEN_GROUPS_PER_BLOCK, + (num_tokens_per_expert + TOKEN_GROUPS_PER_BLOCK - 1) / TOKEN_GROUPS_PER_BLOCK); + dim3 block(THREADS_PER_GROUP * HIDDEN_GROUPS_PER_BLOCK * TOKEN_GROUPS_PER_BLOCK); + per_token_group_quant_8bit_bf16_fp8_g128_row_tiled_kernel< + THREADS_PER_GROUP, + FUSE_SILU_AND_MUL, + HIDDEN_GROUPS_PER_BLOCK, + TOKEN_GROUPS_PER_BLOCK><<>>( + static_cast<__mt_bfloat16*>(input.data_ptr()), + static_cast<__mt_fp8_e4m3*>(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + num_tokens_per_expert, + fused_activation_type); +} + +template +bool try_launch_bf16_fp8_g128_row_tiled_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + if (hidden_dim_num_groups > 32) { + return false; + } + if (num_tokens_per_expert < 16) { + if constexpr (FUSE_SILU_AND_MUL) { + return false; + } else { + if (num_tokens_per_expert < 8 || hidden_dim_num_groups < 2 || hidden_dim_num_groups > 8) { + return false; + } + } + } + + if constexpr (FUSE_SILU_AND_MUL) { + if (hidden_dim_num_groups <= 4) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 8) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 16) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } + return true; + } + + if (hidden_dim_num_groups <= 1) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 2) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 4) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 8) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups <= 16) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else if (hidden_dim_num_groups % 4 == 0) { + return false; + } else if (hidden_dim_num_groups <= 24) { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } else { + launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); + } + return true; +} + +template +void dispatch_bf16_fp8_g128_row_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + constexpr int THREADS_PER_GROUP = BF16_FP8_G128_THREADS_PER_GROUP; + if (try_launch_bf16_fp8_g128_row_tiled_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream)) { + return; + } + launch_bf16_fp8_g128_row_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_activation_type, stream); +} + +template +void launch_i8_fast_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + constexpr int THREADS_PER_GROUP = FastQuantConfig::THREADS_PER_GROUP; + const int subwarps_per_block = choose_subwarps_per_block(hidden_dim_num_groups); + dim3 grid(hidden_dim_num_groups / subwarps_per_block, num_tokens_per_expert); + dim3 block(subwarps_per_block * THREADS_PER_GROUP); + per_token_group_quant_8bit_fast_kernel + <<>>( + static_cast(input.data_ptr()), + static_cast(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + 0, + num_tokens_per_expert, + fused_activation_type); +} + +template +void launch_i8_fused_tiled_kernel( + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + musaStream_t stream) { + constexpr int THREADS_PER_GROUP = FastQuantConfig::THREADS_PER_GROUP; + dim3 grid( + (hidden_dim_num_groups + HIDDEN_GROUPS_PER_BLOCK - 1) / HIDDEN_GROUPS_PER_BLOCK, + (num_tokens_per_expert + TOKEN_GROUPS_PER_BLOCK - 1) / TOKEN_GROUPS_PER_BLOCK); + dim3 block(THREADS_PER_GROUP * HIDDEN_GROUPS_PER_BLOCK * TOKEN_GROUPS_PER_BLOCK); + per_token_group_quant_8bit_i8_fused_tiled_kernel< + GROUP_SIZE, + THREADS_PER_GROUP, + HIDDEN_GROUPS_PER_BLOCK, + TOKEN_GROUPS_PER_BLOCK><<>>( + static_cast(input.data_ptr()), + static_cast(output_q.data_ptr()), + static_cast(output_s.data_ptr()), + hidden_dim_num_groups, + num_tokens_per_expert); +} + +bool try_launch_i8_fused_tiled_kernel( + int64_t group_size, + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + musaStream_t stream) { + if (num_tokens_per_expert >= 512) { + return false; + } + + if (group_size == 64) { + if (hidden_dim_num_groups <= 2 && num_tokens_per_expert >= 128) { + launch_i8_fused_tiled_kernel<64, 2, 8>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream); + } else if (hidden_dim_num_groups <= 4 && num_tokens_per_expert >= 256) { + launch_i8_fused_tiled_kernel<64, 4, 4>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream); + } else { + return false; + } + return true; + } + + if (group_size == 128) { + if (hidden_dim_num_groups <= 1 && num_tokens_per_expert >= 256) { + launch_i8_fused_tiled_kernel<128, 1, 8>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream); + } else if (hidden_dim_num_groups <= 2 && num_tokens_per_expert >= 256) { + launch_i8_fused_tiled_kernel<128, 2, 4>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream); + } else if (hidden_dim_num_groups <= 4 && num_tokens_per_expert >= 128 && num_tokens_per_expert <= 384) { + launch_i8_fused_tiled_kernel<128, 4, 2>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream); + } else { + return false; + } + return true; + } + + return false; +} + +template +void dispatch_i8_fast_kernel( + int64_t group_size, + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int hidden_dim_num_groups, + int num_tokens_per_expert, + int fused_activation_type, + musaStream_t stream) { + switch (group_size) { + case 16: + launch_i8_fast_kernel<16, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, + fused_activation_type, stream); + break; + case 32: + launch_i8_fast_kernel<32, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, + fused_activation_type, stream); + break; + case 64: + launch_i8_fast_kernel<64, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, + fused_activation_type, stream); + break; + case 128: + launch_i8_fast_kernel<128, FUSE_SILU_AND_MUL>( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, + fused_activation_type, stream); + break; + default: + TORCH_CHECK(false, "Unsupported group_size"); + } +} + +void sgl_per_token_group_quant_8bit_v2( + // vanilla: (num_tokens, hidden_size) + // fuse_silu_and_mul: (num_tokens, hidden_size * 2) + // fuse_silu_and_mul + masked_layout: (num_experts, num_tokens-with-padding, hidden_size * 2) + ffi::TensorView input, + ffi::TensorView output_q, + ffi::TensorView output_s, + int64_t group_size, + double eps, + double min_8bit, + double max_8bit, + bool scale_ue8m0, + bool fuse_silu_and_mul, + int64_t fused_activation_type, + ffi::TensorView masked_m, + bool has_masked_m) { + CHECK_INPUT(input); + CHECK_INPUT(output_q); + TVM_FFI_ICHECK_EQ(input.device().device_id, output_q.device().device_id); + TVM_FFI_ICHECK_EQ(input.device().device_id, output_s.device().device_id); + TVM_FFI_ICHECK(tensor_numel(input) > 0); + + TORCH_CHECK(std::abs(LOCAL_ABSMAX_ABS - eps) < 1e-13); + + CHECK_EQ(tensor_numel(input) % group_size, 0); + const int num_groups = static_cast(tensor_numel(input)) / group_size / (fuse_silu_and_mul ? 2 : 1); + + const bool masked_layout = has_masked_m; + TORCH_CHECK(output_s.ndim() == (masked_layout ? 3 : 2)); + if (masked_layout) { + TVM_FFI_ICHECK_EQ(masked_m.device().device_type, kDLExtDev); + TVM_FFI_ICHECK_EQ(masked_m.device().device_id, input.device().device_id); + TVM_FFI_ICHECK(masked_m.IsContiguous()); + TVM_FFI_ICHECK(dtype_equal(masked_m.dtype(), DLDataType{kDLInt, 32, 1})); + } + + const int num_local_experts = masked_layout ? input.size(0) : 1; + + ffi::MUSADeviceGuard device_guard(input.device().device_id); + musaStream_t stream = get_stream(input.device()); + + const bool is_column_major = output_s.stride(-2) < output_s.stride(-1); + const int hidden_dim_num_groups = static_cast(output_q.size(-1)) / group_size; + const int num_tokens_per_expert = static_cast(output_q.size(-2)); + const int scale_expert_stride = masked_layout ? static_cast(output_s.stride(0)) : 0; + const int scale_hidden_stride = static_cast(output_s.stride(-1)); + const int fused_act = static_cast(fused_activation_type); + TORCH_CHECK( + fused_act == kSiluActivation || fused_act == kGeluActivation || + fused_act == kGeluTanhActivation, + "Unsupported fused activation type"); + +#define LAUNCH_KERNEL_INNER(SCHEDULER, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, output_s_dtype, ...) \ + do { \ + int subwarps_per_block; \ + dim3 grid, block; \ + SCHEDULER::compute_exec_config( \ + THREADS_PER_SUBWARP, num_local_experts, hidden_dim_num_groups, num_groups, subwarps_per_block, grid, block); \ + \ + per_token_group_quant_8bit_kernel \ + <<>>( \ + static_cast(input.data_ptr()), \ + static_cast(output_q.data_ptr()), \ + static_cast(output_s.data_ptr()), \ + static_cast(masked_layout ? masked_m.data_ptr() : nullptr), \ + subwarps_per_block, \ + hidden_dim_num_groups, \ + scale_expert_stride, \ + scale_hidden_stride, \ + num_tokens_per_expert, \ + fused_act); \ + } while (0) + +#define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \ + do { \ + constexpr int THREADS_PER_SUBWARP = GROUP_SIZE / 16; \ + TORCH_CHECK(THREADS_PER_SUBWARP * INPUT_PRIMARY_VEC_NUM_BYTES == group_size * sizeof(T)); \ + \ + using dst_dtype_info = DtypeInfo; \ + CHECK_EQ(dst_dtype_info::MIN, min_8bit); \ + CHECK_EQ(dst_dtype_info::MAX, max_8bit); \ + \ + if (is_column_major) { \ + if (scale_ue8m0) { \ + if (fuse_silu_and_mul) { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER( \ + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, true); \ + } else { \ + LAUNCH_KERNEL_INNER( \ + RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, true); \ + } \ + } else { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER( \ + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true); \ + } else { \ + LAUNCH_KERNEL_INNER(RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true); \ + } \ + } \ + } else { \ + if (fuse_silu_and_mul) { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER( \ + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ + } else { \ + LAUNCH_KERNEL_INNER( \ + RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ + } \ + } else { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER( \ + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true); \ + } else { \ + LAUNCH_KERNEL_INNER(RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true); \ + } \ + } \ + } \ + } else { \ + if (fuse_silu_and_mul) { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER( \ + MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ + } else { \ + LAUNCH_KERNEL_INNER( \ + RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ + } \ + } else { \ + if (masked_layout) { \ + LAUNCH_KERNEL_INNER(MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false); \ + } else { \ + LAUNCH_KERNEL_INNER(RowScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false); \ + } \ + } \ + } \ + } while (0) + +#define LAUNCH_KERNEL_OUTER(...) \ + switch (group_size) { \ + case 16: \ + LAUNCH_KERNEL(16, __VA_ARGS__); \ + break; \ + case 32: \ + LAUNCH_KERNEL(32, __VA_ARGS__); \ + break; \ + case 64: \ + LAUNCH_KERNEL(64, __VA_ARGS__); \ + break; \ + case 128: \ + LAUNCH_KERNEL(128, __VA_ARGS__); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported group_size"); \ + } \ + while (0) + + if (dtype_equal(input.dtype(), dl_float16)) { + using scalar_t = half; + if (dtype_equal(output_q.dtype(), dl_int8)) { + if (!is_column_major && !scale_ue8m0 && !masked_layout) { + const bool prefer_generic_for_large_hidden = + (group_size == 64 && + ((hidden_dim_num_groups >= 32 && num_tokens_per_expert >= 2048) || + (hidden_dim_num_groups >= 24 && num_tokens_per_expert >= 4096))) || + (group_size == 128 && hidden_dim_num_groups >= 12 && num_tokens_per_expert >= 2048); + if (fuse_silu_and_mul && group_size <= 128 && !prefer_generic_for_large_hidden) { + if (!try_launch_i8_fused_tiled_kernel( + group_size, input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, stream)) { + dispatch_i8_fast_kernel( + group_size, input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_act, stream); + } + } else if (!fuse_silu_and_mul && group_size == 16) { + dispatch_i8_fast_kernel( + group_size, input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_act, stream); + } else { + LAUNCH_KERNEL_OUTER(scalar_t, int8_t); + } + } else { + LAUNCH_KERNEL_OUTER(scalar_t, int8_t); + } + } else if (dtype_equal(output_q.dtype(), dl_float8_e4m3fn)) { + if (!scale_ue8m0 && !masked_layout) { + if (is_column_major) { + if (fuse_silu_and_mul) { + dispatch_fp8_fast_kernel( + group_size, + input, + output_q, + output_s, + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert, + fused_act, + stream); + } else { + dispatch_fp8_col_nonfused_kernel( + group_size, + input, + output_q, + output_s, + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert, + stream); + } + } else { + if (fuse_silu_and_mul) { + dispatch_fp8_fast_kernel( + group_size, + input, + output_q, + output_s, + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert, + fused_act, + stream); + } else { + dispatch_fp8_fast_kernel( + group_size, + input, + output_q, + output_s, + hidden_dim_num_groups, + scale_hidden_stride, + num_tokens_per_expert, + fused_act, + stream); + } + } + } else { + LAUNCH_KERNEL_OUTER(scalar_t, __mt_fp8_e4m3); + } + } else { + TVM_FFI_THROW(ValueError) << "Unsupported output_q dtype"; + } + } else if (dtype_equal(input.dtype(), dl_bfloat16)) { + using scalar_t = __mt_bfloat16; + if (dtype_equal(output_q.dtype(), dl_int8)) { + LAUNCH_KERNEL_OUTER(scalar_t, int8_t); + } else if (dtype_equal(output_q.dtype(), dl_float8_e4m3fn)) { + if (group_size == 128 && !is_column_major && !scale_ue8m0 && !masked_layout) { + if (fuse_silu_and_mul) { + dispatch_bf16_fp8_g128_row_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_act, stream); + } else { + dispatch_bf16_fp8_g128_row_kernel( + input, output_q, output_s, hidden_dim_num_groups, num_tokens_per_expert, fused_act, stream); + } + } else { + LAUNCH_KERNEL_OUTER(scalar_t, __mt_fp8_e4m3); + } + } else { + TVM_FFI_THROW(ValueError) << "Unsupported output_q dtype"; + } + } else { + TVM_FFI_THROW(ValueError) << "Unsupported input dtype"; + } + +#undef LAUNCH_KERNEL +#undef LAUNCH_KERNEL_INNER +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(sgl_per_token_group_quant_8bit_v2, sgl_per_token_group_quant_8bit_v2); diff --git a/vllm_musa/jit_kernel/extend_topk_shared.py b/vllm_musa/jit_kernel/extend_topk_shared.py new file mode 100644 index 000000000000..2fabddfe1b4f --- /dev/null +++ b/vllm_musa/jit_kernel/extend_topk_shared.py @@ -0,0 +1,71 @@ +"""Append a folded shared expert's routing column in one kernel. + +The routed topk output is (tokens, top_k); a folded shared expert rides the grouped GEMM as one +extra always-selected slot, so the routing needs one more column carrying +`sigmoid(shared_gate(x))` and the shared expert's id. Done with tensor ops that is a +sigmoid + zeros_like + add + 2 cats per MoE layer; this writes all of it in a single pass. + +The shared weight stays out of any renormalization (the routed weights are already normalized +among themselves when the caller runs this), matching the reference implementation. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _extend_topk_with_shared_kernel( + topk_weights_ptr, + topk_ids_ptr, + shared_logits_ptr, + out_weights_ptr, + out_ids_ptr, + shared_expert_id, + top_k: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, BLOCK) + mask = offs < top_k + + w = tl.load(topk_weights_ptr + token * top_k + offs, mask=mask, other=0.0) + i = tl.load(topk_ids_ptr + token * top_k + offs, mask=mask, other=0) + + out_w = out_weights_ptr + token * (top_k + 1) + out_i = out_ids_ptr + token * (top_k + 1) + tl.store(out_w + offs, w, mask=mask) + tl.store(out_i + offs, i, mask=mask) + + logit = tl.load(shared_logits_ptr + token).to(tl.float32) + shared_w = 1.0 / (1.0 + tl.exp(-logit)) + tl.store(out_w + top_k, shared_w.to(out_w.dtype.element_ty)) + tl.store(out_i + top_k, shared_expert_id) + + +def extend_topk_with_shared( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_logits: torch.Tensor, + shared_expert_id: int, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, top_k = topk_weights.shape + out_w = torch.empty( + (tokens, top_k + 1), device=topk_weights.device, dtype=topk_weights.dtype + ) + out_i = torch.empty( + (tokens, top_k + 1), device=topk_ids.device, dtype=topk_ids.dtype + ) + if tokens == 0: + return out_w, out_i + _extend_topk_with_shared_kernel[(tokens,)]( + topk_weights, + topk_ids, + shared_logits, + out_w, + out_i, + int(shared_expert_id), + top_k=top_k, + BLOCK=triton.next_power_of_2(top_k), + ) + return out_w, out_i diff --git a/vllm_musa/jit_kernel/fused_gdn_gating.py b/vllm_musa/jit_kernel/fused_gdn_gating.py new file mode 100644 index 000000000000..8acb2d80c3d9 --- /dev/null +++ b/vllm_musa/jit_kernel/fused_gdn_gating.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Single fused GDN gating kernel for the strided-QKV MATE prefill path. + +Computes the two log-space GDN gating tensors in one Triton launch: + + g = -exp(A_log) * softplus(a + dt_bias) (log-space, no final exp) + beta = sigmoid(b) + +replacing the ~9 elementwise kernels (exp / mul / softplus / sigmoid / cast) +the strided path would otherwise issue per GDN layer. Ported from SGLang's +``fused_gdn_gating``; the wrapper returns 2D ``[L, HV]`` fp32 tensors matching +what the MATE ``chunk_gated_delta_rule`` prefill call expects. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_gdn_gating_kernel( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + NUM_HEADS: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_HEADS: tl.constexpr, +): + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + mask = head_off < NUM_HEADS + blk_A_log = tl.load(A_log + head_off, mask=mask) + blk_a = tl.load(a + off, mask=mask) + blk_b = tl.load(b + off, mask=mask) + blk_bias = tl.load(dt_bias + head_off, mask=mask) + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) + tl.store( + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + ) + + +def fused_gdn_gating( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused log-space GDN gating. + + Args: + A_log: [HV] log decay parameter (per v-head). + a: [L, HV] gating input, contiguous. + b: [L, HV] gating input, contiguous. + dt_bias: [HV] dt bias (per v-head). + + Returns: + g: [L, HV] float32, log-space ``-exp(A_log) * softplus(a + dt_bias)``. + beta: [L, HV] float32, ``sigmoid(b)``. + """ + assert a.shape == b.shape and a.dim() == 2 + L, num_heads = a.shape + g = torch.empty(L, num_heads, dtype=torch.float32, device=a.device) + beta_output = torch.empty(L, num_heads, dtype=torch.float32, device=b.device) + if L == 0: + return g, beta_output + a = a.contiguous() + b = b.contiguous() + grid = (L, 1, triton.cdiv(num_heads, 8)) + _fused_gdn_gating_kernel[grid]( + g, + beta_output, + A_log, + a, + b, + dt_bias, + 1, # seq_len: each token treated independently + num_heads, + beta, + threshold, + 8, + num_warps=1, + ) + return g, beta_output diff --git a/vllm_musa/jit_kernel/post_reorder.py b/vllm_musa/jit_kernel/post_reorder.py new file mode 100644 index 000000000000..fbc000afab72 --- /dev/null +++ b/vllm_musa/jit_kernel/post_reorder.py @@ -0,0 +1,44 @@ +"""Fused un-permute + weighted top-k reduce for the DeepGEMM grouped-MoE prefill path. + +Gathers each source token's top-k expert-output rows from the grouped-GEMM +output (indexed through ``src2dst``), scales by the router weights, and reduces +into the final hidden state — one launch replacing an unfused unpermute+reduce. +""" + +import triton +import triton.language as tl + + +@triton.jit +def post_reorder_triton_kernel( + down_output_ptr, + output_ptr, + src2dst_ptr, + topk_ids_ptr, + topk_weights_ptr, + topk, + hidden_size, + BLOCK_SIZE: tl.constexpr, +): + InDtype = down_output_ptr.dtype.element_ty + src_idx = tl.program_id(0).to(tl.int64) + src2dst_ptr = src2dst_ptr + src_idx * topk + topk_weights_ptr = topk_weights_ptr + src_idx * topk + store_ptr = output_ptr + src_idx * hidden_size + vec = tl.arange(0, BLOCK_SIZE) + for start in tl.range(0, hidden_size, BLOCK_SIZE): + offset = start + vec + mask = offset < hidden_size + acc = tl.zeros([BLOCK_SIZE], dtype=InDtype) + for idx in range(topk): + dst = tl.load(src2dst_ptr + idx) + if dst >= 0: + weight = tl.load(topk_weights_ptr + idx).to(InDtype) + acc += ( + tl.load( + down_output_ptr + dst.to(tl.int64) * hidden_size + offset, + mask=mask, + ) + * weight + ) + tl.store(store_ptr + offset, acc, mask=mask) diff --git a/vllm_musa/jit_kernel/tilelang/_atomic_helper.h b/vllm_musa/jit_kernel/tilelang/_atomic_helper.h new file mode 100644 index 000000000000..f4eb390f6a56 --- /dev/null +++ b/vllm_musa/jit_kernel/tilelang/_atomic_helper.h @@ -0,0 +1,36 @@ +#pragma once + +#include + +#include + +__device__ __forceinline__ int sgl_tl_atomic_add_offset(int *base, int offset, + int val) { + return atomicAdd(base + offset, val); +} + +__device__ __forceinline__ void sgl_tl_store_fp8e4m3x4(fp8_e4_t *base, + int64_t offset, float x0, + float x1, float x2, + float x3) { + const float4 values = {x0, x1, x2, x3}; + const fp8_e4_4_t packed = tl::cvt_float_to_fp8e4m3_x4(values); + *reinterpret_cast(base + offset) = + *reinterpret_cast(&packed); +} + +__device__ __forceinline__ void sgl_tl_copy_bf16x8(bfloat16_t *dst, + const bfloat16_t *src, + int64_t dst_offset, + int64_t src_offset) { + const int4 value = *reinterpret_cast(src + src_offset); + *reinterpret_cast(dst + dst_offset) = value; +} + +__device__ __forceinline__ void sgl_tl_copy_fp8x16(fp8_e4_t *dst, + const fp8_e4_t *src, + int64_t dst_offset, + int64_t src_offset) { + const int4 value = *reinterpret_cast(src + src_offset); + *reinterpret_cast(dst + dst_offset) = value; +} diff --git a/vllm_musa/jit_kernel/tilelang/causal_conv1d.py b/vllm_musa/jit_kernel/tilelang/causal_conv1d.py new file mode 100644 index 000000000000..16ce6b7b4926 --- /dev/null +++ b/vllm_musa/jit_kernel/tilelang/causal_conv1d.py @@ -0,0 +1,1646 @@ +"""MUSA TileLang causal conv1d forward kernel.""" + +import functools +from typing import List, Optional, Union + +import tilelang +import tilelang.language as T +import torch + +from vllm_musa.jit_kernel.tilelang.utils import ( + MUSA_COMMON_PASS_CONFIGS, + MUSA_COMPILE_FLAGS, + storage_window, + tilelang_dtype, +) + +PAD_SLOT_ID = -1 # MUSA: match vllm mamba causal_conv1d PAD_SLOT_ID + + +def register_custom_op(fn=None, **_kw): + def _wrap(f): + return f + + return _wrap if fn is None else _wrap(fn) + + +_LOG2E = 1.4426950408889634 +_ENABLE_WIDTH4_PREFILL_SPLIT = False + +_CAUSAL_CONV1D_PASS_CONFIGS = dict(MUSA_COMMON_PASS_CONFIGS) +for _key, _value in ( + ("TL_ENABLE_LOWER_LDGSTG", True), + ("TL_ENABLE_LOWER_LDGSTG_PREDICATED", True), + ("TL_DISABLE_SAFE_COPY_PREDICATION", True), + ("TL_DISABLE_SAFE_ROBUST_COPY_PREDICATION", True), + ("TL_CONFIG_INDEX_BITWIDTH", 32), +): + if hasattr(tilelang.PassConfigKey, _key): + _CAUSAL_CONV1D_PASS_CONFIGS[getattr(tilelang.PassConfigKey, _key)] = _value + + +def _next_power_of_2(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _tilelang_index_dtype(dtype: torch.dtype) -> str: + if dtype is torch.int32: + return "int32" + if dtype is torch.int64: + return "int64" + raise TypeError(f"Unsupported index dtype for TileLang MUSA kernel: {dtype}") + + +@functools.lru_cache(maxsize=64) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_CAUSAL_CONV1D_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _causal_conv1d_fwd_kernel( + dtype: str, + cache_indices_dtype: str, + width: int, + x_stride_dim: int, + x_stride_token: int, + w_stride_dim: int, + w_stride_width: int, + state_stride_seq: int, + state_stride_dim: int, + state_stride_token: int, + o_stride_dim: int, + o_stride_token: int, + has_bias: bool, + has_conv_states: bool, + has_cache_indices: bool, + has_cache_index_mapping: bool, + has_initial_states: bool, + use_pad_slot: bool, + silu_activation: bool, + block_m: int, + block_n: int, +): + x_numel = T.dynamic("x_numel") + w_numel = T.dynamic("w_numel") + bias_numel = T.dynamic("bias_numel") + state_numel = T.dynamic("state_numel") + cache_numel = T.dynamic("cache_numel") + mapping_numel = T.dynamic("mapping_numel") + init_numel = T.dynamic("init_numel") + query_numel = T.dynamic("query_numel") + out_numel = T.dynamic("out_numel") + state_len = width - 1 + + @T.prim_func + def musa_causal_conv1d_fwd( + x: T.Tensor((x_numel,), dtype), + weight: T.Tensor((w_numel,), dtype), + bias: T.Tensor((bias_numel,), dtype), + conv_states: T.Tensor((state_numel,), dtype), + cache_indices: T.Tensor((cache_numel,), cache_indices_dtype), + cache_index_mapping: T.Tensor((mapping_numel,), "int32"), + has_initial_state: T.Tensor((init_numel,), "bool"), + query_start_loc: T.Tensor((query_numel,), "int32"), + out: T.Tensor((out_numel,), dtype), + max_seq_len: T.int32, + dim: T.int32, + num_cache_lines: T.int32, + pad_slot_id: T.int32, + ): + with T.Kernel( + query_numel - 1, + T.ceildiv(max_seq_len, block_m), + T.ceildiv(dim, block_n), + threads=block_n, + ) as (seq_idx, chunk_idx, dim_block): + tid = T.get_thread_binding() + feat = dim_block * block_n + tid + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + seq_len = T.alloc_var("int32") + token_offset = T.alloc_var("int32") + segment_len = T.alloc_var("int32") + cache_idx = T.alloc_var("int32") + load_init = T.alloc_var("bool") + state_base = T.alloc_var("int32") + x_base = T.alloc_var("int32") + w_base = T.alloc_var("int32") + out_base = T.alloc_var("int32") + valid_seq = T.alloc_var("bool") + col0 = T.alloc_var("float32") + col1 = T.alloc_var("float32") + col2 = T.alloc_var("float32") + col3 = T.alloc_var("float32") + w0 = T.alloc_var("float32") + w1 = T.alloc_var("float32") + w2 = T.alloc_var("float32") + w3 = T.alloc_var("float32") + w4 = T.alloc_var("float32") + x_cur = T.alloc_var("float32") + acc = T.alloc_var("float32") + state_src = T.alloc_var("int32") + state_cut = T.alloc_var("int32") + + seq_start = query_start_loc[seq_idx] + seq_end = query_start_loc[seq_idx + 1] + seq_len = seq_end - seq_start + token_offset = chunk_idx * block_m + segment_len = seq_len - token_offset + if segment_len > block_m: + segment_len = block_m + + cache_idx = seq_idx + if has_cache_indices: + cache_idx = cache_indices[seq_idx] + if has_cache_index_mapping: + cache_idx = cache_index_mapping[cache_idx] + valid_seq = segment_len > 0 + if use_pad_slot and cache_idx == pad_slot_id: + valid_seq = False + if has_conv_states and cache_idx >= num_cache_lines: + valid_seq = False + + if valid_seq and feat < dim: + x_base = seq_start * x_stride_token + feat * x_stride_dim + w_base = feat * w_stride_dim + out_base = seq_start * o_stride_token + feat * o_stride_dim + state_base = cache_idx * state_stride_seq + feat * state_stride_dim + + col0 = 0.0 + col1 = 0.0 + col2 = 0.0 + col3 = 0.0 + load_init = False + if has_initial_states: + load_init = has_initial_state[seq_idx] + + if chunk_idx == 0: + if has_conv_states and load_init: + if width >= 2: + col0 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 1) * state_stride_token + ], + ) + if width >= 3: + col1 = col0 + col0 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 2) * state_stride_token + ], + ) + if width >= 4: + col2 = col1 + col1 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 2) * state_stride_token + ], + ) + col0 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 3) * state_stride_token + ], + ) + if width >= 5: + col3 = col2 + col2 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 2) * state_stride_token + ], + ) + col1 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 3) * state_stride_token + ], + ) + col0 = T.Cast( + "float32", + conv_states[ + state_base + (state_len - 4) * state_stride_token + ], + ) + + if has_conv_states: + state_cut = state_len - seq_len + for state_i in T.serial(state_len): + if state_len <= seq_len: + conv_states[ + state_base + state_i * state_stride_token + ] = x[ + x_base + + (seq_len - state_len + state_i) * x_stride_token + ] + else: + if load_init and state_i < state_cut: + state_src = state_i + seq_len + conv_states[ + state_base + state_i * state_stride_token + ] = conv_states[ + state_base + state_src * state_stride_token + ] + elif state_i >= state_cut: + conv_states[ + state_base + state_i * state_stride_token + ] = x[ + x_base + (state_i - state_cut) * x_stride_token + ] + else: + conv_states[ + state_base + state_i * state_stride_token + ] = T.Cast(dtype, 0.0) + else: + if width >= 2: + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 1) * x_stride_token], + ) + if width >= 3: + col1 = col0 + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 2) * x_stride_token], + ) + if width >= 4: + col2 = col1 + col1 = T.Cast( + "float32", + x[x_base + (token_offset - 2) * x_stride_token], + ) + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 3) * x_stride_token], + ) + if width >= 5: + col3 = col2 + col2 = T.Cast( + "float32", + x[x_base + (token_offset - 2) * x_stride_token], + ) + col1 = T.Cast( + "float32", + x[x_base + (token_offset - 3) * x_stride_token], + ) + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 4) * x_stride_token], + ) + + w0 = T.Cast("float32", weight[w_base]) + w1 = T.Cast("float32", weight[w_base + w_stride_width]) + if width >= 3: + w2 = T.Cast("float32", weight[w_base + 2 * w_stride_width]) + if width >= 4: + w3 = T.Cast("float32", weight[w_base + 3 * w_stride_width]) + if width >= 5: + w4 = T.Cast("float32", weight[w_base + 4 * w_stride_width]) + + for token_i in T.serial(block_m): + if token_i < segment_len: + x_cur = T.Cast( + "float32", + x[x_base + (token_offset + token_i) * x_stride_token], + ) + acc = 0.0 + if has_bias: + acc = T.Cast("float32", bias[feat]) + + if width == 2: + acc += col0 * w0 + x_cur * w1 + col0 = x_cur + elif width == 3: + acc += col0 * w0 + col1 * w1 + x_cur * w2 + col0 = col1 + col1 = x_cur + elif width == 4: + acc += col0 * w0 + col1 * w1 + col2 * w2 + x_cur * w3 + col0 = col1 + col1 = col2 + col2 = x_cur + else: + acc += ( + col0 * w0 + + col1 * w1 + + col2 * w2 + + col3 * w3 + + x_cur * w4 + ) + col0 = col1 + col1 = col2 + col2 = col3 + col3 = x_cur + + if silu_activation: + acc = acc / (1.0 + T.exp2(-acc * _LOG2E)) + out[out_base + (token_offset + token_i) * o_stride_token] = ( + T.Cast(dtype, acc) + ) + + return musa_causal_conv1d_fwd + + +_causal_conv1d_fwd_kernel.mode = "lazy" + + +@functools.lru_cache(maxsize=64) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_CAUSAL_CONV1D_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _causal_conv1d_fwd_width4_vec_kernel( + dtype: str, + cache_indices_dtype: str, + x_stride_token: int, + w_stride_dim: int, + state_stride_seq: int, + state_stride_dim: int, + state_stride_token: int, + o_stride_token: int, + has_bias: bool, + has_conv_states: bool, + has_cache_indices: bool, + has_cache_index_mapping: bool, + has_initial_states: bool, + use_pad_slot: bool, + silu_activation: bool, + block_m: int, + block_feats: int, + vec_elems: int, +): + x_numel = T.dynamic("x_numel") + w_numel = T.dynamic("w_numel") + bias_numel = T.dynamic("bias_numel") + state_numel = T.dynamic("state_numel") + cache_numel = T.dynamic("cache_numel") + mapping_numel = T.dynamic("mapping_numel") + init_numel = T.dynamic("init_numel") + query_numel = T.dynamic("query_numel") + out_numel = T.dynamic("out_numel") + num_threads = block_feats // vec_elems + state_len = 3 + + @T.prim_func + def musa_causal_conv1d_fwd_width4_vec( + x: T.Tensor((x_numel,), dtype), + weight: T.Tensor((w_numel,), dtype), + bias: T.Tensor((bias_numel,), dtype), + conv_states: T.Tensor((state_numel,), dtype), + cache_indices: T.Tensor((cache_numel,), cache_indices_dtype), + cache_index_mapping: T.Tensor((mapping_numel,), "int32"), + has_initial_state: T.Tensor((init_numel,), "bool"), + query_start_loc: T.Tensor((query_numel,), "int32"), + out: T.Tensor((out_numel,), dtype), + max_seq_len: T.int32, + dim: T.int32, + num_cache_lines: T.int32, + pad_slot_id: T.int32, + ): + with T.Kernel( + query_numel - 1, + T.ceildiv(max_seq_len, block_m), + T.ceildiv(dim, block_feats), + threads=num_threads, + ) as (seq_idx, chunk_idx, dim_block): + tid = T.get_thread_binding() + feat_base = dim_block * block_feats + tid * vec_elems + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + seq_len = T.alloc_var("int32") + token_offset = T.alloc_var("int32") + segment_len = T.alloc_var("int32") + cache_idx = T.alloc_var("int32") + load_init = T.alloc_var("bool") + valid_seq = T.alloc_var("bool") + state_cut = T.alloc_var("int32") + state_base = T.alloc_local((vec_elems,), "int32") + x_base = T.alloc_local((vec_elems,), "int32") + w_base = T.alloc_local((vec_elems,), "int32") + out_base = T.alloc_local((vec_elems,), "int32") + col0 = T.alloc_local((vec_elems,), "float32") + col1 = T.alloc_local((vec_elems,), "float32") + col2 = T.alloc_local((vec_elems,), "float32") + w0 = T.alloc_local((vec_elems,), "float32") + w1 = T.alloc_local((vec_elems,), "float32") + w2 = T.alloc_local((vec_elems,), "float32") + w3 = T.alloc_local((vec_elems,), "float32") + x_cur = T.alloc_local((vec_elems,), "float32") + acc = T.alloc_local((vec_elems,), "float32") + + seq_start = query_start_loc[seq_idx] + seq_end = query_start_loc[seq_idx + 1] + seq_len = seq_end - seq_start + token_offset = chunk_idx * block_m + segment_len = seq_len - token_offset + if segment_len > block_m: + segment_len = block_m + + cache_idx = seq_idx + if has_cache_indices: + cache_idx = cache_indices[seq_idx] + if has_cache_index_mapping: + cache_idx = cache_index_mapping[cache_idx] + valid_seq = segment_len > 0 + if use_pad_slot and cache_idx == pad_slot_id: + valid_seq = False + if has_conv_states and cache_idx >= num_cache_lines: + valid_seq = False + load_init = False + if has_initial_states: + load_init = has_initial_state[seq_idx] + + if valid_seq: + for v in T.vectorized(vec_elems): + feat = feat_base + v + col0[v] = 0.0 + col1[v] = 0.0 + col2[v] = 0.0 + x_base[v] = seq_start * x_stride_token + feat + w_base[v] = feat * w_stride_dim + out_base[v] = seq_start * o_stride_token + feat + state_base[v] = ( + cache_idx * state_stride_seq + feat * state_stride_dim + ) + if feat < dim: + if chunk_idx == 0: + if has_conv_states and load_init: + col2[v] = T.Cast( + "float32", + conv_states[state_base[v] + 2 * state_stride_token], + ) + col1[v] = T.Cast( + "float32", + conv_states[state_base[v] + state_stride_token], + ) + col0[v] = T.Cast("float32", conv_states[state_base[v]]) + else: + col2[v] = T.Cast( + "float32", + x[x_base[v] + (token_offset - 1) * x_stride_token], + ) + col1[v] = T.Cast( + "float32", + x[x_base[v] + (token_offset - 2) * x_stride_token], + ) + col0[v] = T.Cast( + "float32", + x[x_base[v] + (token_offset - 3) * x_stride_token], + ) + + w0[v] = T.Cast("float32", weight[w_base[v]]) + w1[v] = T.Cast("float32", weight[w_base[v] + 1]) + w2[v] = T.Cast("float32", weight[w_base[v] + 2]) + w3[v] = T.Cast("float32", weight[w_base[v] + 3]) + + if chunk_idx == 0 and has_conv_states: + state_cut = state_len - seq_len + for state_i in T.serial(state_len): + for v in T.vectorized(vec_elems): + feat = feat_base + v + if feat < dim: + if seq_len >= state_len: + conv_states[ + state_base[v] + state_i * state_stride_token + ] = x[ + x_base[v] + + (seq_len - state_len + state_i) + * x_stride_token + ] + else: + if load_init and state_i < state_cut: + conv_states[ + state_base[v] + state_i * state_stride_token + ] = conv_states[ + state_base[v] + + (state_i + seq_len) * state_stride_token + ] + elif state_i >= state_cut: + conv_states[ + state_base[v] + state_i * state_stride_token + ] = x[ + x_base[v] + + (state_i - state_cut) * x_stride_token + ] + else: + conv_states[ + state_base[v] + state_i * state_stride_token + ] = T.Cast(dtype, 0.0) + + for token_i in T.serial(block_m): + if token_i < segment_len: + for v in T.vectorized(vec_elems): + feat = feat_base + v + if feat < dim: + x_cur[v] = T.Cast( + "float32", + x[ + x_base[v] + + (token_offset + token_i) * x_stride_token + ], + ) + acc[v] = 0.0 + if has_bias: + acc[v] = T.Cast("float32", bias[feat]) + acc[v] += ( + col0[v] * w0[v] + + col1[v] * w1[v] + + col2[v] * w2[v] + + x_cur[v] * w3[v] + ) + col0[v] = col1[v] + col1[v] = col2[v] + col2[v] = x_cur[v] + if silu_activation: + acc[v] = acc[v] / (1.0 + T.exp2(-acc[v] * _LOG2E)) + out[ + out_base[v] + + (token_offset + token_i) * o_stride_token + ] = T.Cast(dtype, acc[v]) + + return musa_causal_conv1d_fwd_width4_vec + + +_causal_conv1d_fwd_width4_vec_kernel.mode = "lazy" + + +@functools.lru_cache(maxsize=64) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_CAUSAL_CONV1D_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _causal_conv1d_prefill_width4_kernel( + dtype: str, + cache_indices_dtype: str, + x_stride_token: int, + w_stride_dim: int, + state_stride_seq: int, + state_stride_dim: int, + state_stride_token: int, + o_stride_token: int, + has_bias: bool, + has_cache_indices: bool, + has_cache_index_mapping: bool, + has_initial_states: bool, + use_pad_slot: bool, + silu_activation: bool, + block_m: int, + block_n: int, +): + x_numel = T.dynamic("x_numel") + w_numel = T.dynamic("w_numel") + bias_numel = T.dynamic("bias_numel") + state_numel = T.dynamic("state_numel") + cache_numel = T.dynamic("cache_numel") + mapping_numel = T.dynamic("mapping_numel") + init_numel = T.dynamic("init_numel") + query_numel = T.dynamic("query_numel") + out_numel = T.dynamic("out_numel") + + @T.prim_func + def musa_causal_conv1d_prefill_width4( + x: T.Tensor((x_numel,), dtype), + weight: T.Tensor((w_numel,), dtype), + bias: T.Tensor((bias_numel,), dtype), + conv_states: T.Tensor((state_numel,), dtype), + cache_indices: T.Tensor((cache_numel,), cache_indices_dtype), + cache_index_mapping: T.Tensor((mapping_numel,), "int32"), + has_initial_state: T.Tensor((init_numel,), "bool"), + query_start_loc: T.Tensor((query_numel,), "int32"), + out: T.Tensor((out_numel,), dtype), + max_seq_len: T.int32, + dim: T.int32, + num_cache_lines: T.int32, + pad_slot_id: T.int32, + ): + with T.Kernel( + query_numel - 1, + T.ceildiv(max_seq_len, block_m), + T.ceildiv(dim, block_n), + threads=block_n, + ) as (seq_idx, chunk_idx, dim_block): + tid = T.get_thread_binding() + feat = dim_block * block_n + tid + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + seq_len = T.alloc_var("int32") + token_offset = T.alloc_var("int32") + segment_len = T.alloc_var("int32") + cache_idx = T.alloc_var("int32") + valid_seq = T.alloc_var("bool") + load_init = T.alloc_var("bool") + x_base = T.alloc_var("int32") + w_base = T.alloc_var("int32") + state_base = T.alloc_var("int32") + out_base = T.alloc_var("int32") + state_cut = T.alloc_var("int32") + col0 = T.alloc_var("float32") + col1 = T.alloc_var("float32") + col2 = T.alloc_var("float32") + w0 = T.alloc_var("float32") + w1 = T.alloc_var("float32") + w2 = T.alloc_var("float32") + w3 = T.alloc_var("float32") + x_cur = T.alloc_var("float32") + acc = T.alloc_var("float32") + + seq_start = query_start_loc[seq_idx] + seq_end = query_start_loc[seq_idx + 1] + seq_len = seq_end - seq_start + token_offset = chunk_idx * block_m + segment_len = seq_len - token_offset + if segment_len > block_m: + segment_len = block_m + + cache_idx = seq_idx + if has_cache_indices: + cache_idx = cache_indices[seq_idx] + if has_cache_index_mapping: + cache_idx = cache_index_mapping[cache_idx] + valid_seq = segment_len > 0 and feat < dim + if use_pad_slot and cache_idx == pad_slot_id: + valid_seq = False + if cache_idx >= num_cache_lines: + valid_seq = False + + if valid_seq: + x_base = seq_start * x_stride_token + feat + w_base = feat * w_stride_dim + out_base = seq_start * o_stride_token + feat + state_base = cache_idx * state_stride_seq + feat * state_stride_dim + + col0 = 0.0 + col1 = 0.0 + col2 = 0.0 + load_init = False + if has_initial_states: + load_init = has_initial_state[seq_idx] + + if chunk_idx == 0: + if load_init: + col0 = T.Cast("float32", conv_states[state_base]) + col1 = T.Cast( + "float32", + conv_states[state_base + state_stride_token], + ) + col2 = T.Cast( + "float32", + conv_states[state_base + 2 * state_stride_token], + ) + + if seq_len >= 3: + conv_states[state_base] = x[ + x_base + (seq_len - 3) * x_stride_token + ] + conv_states[state_base + state_stride_token] = x[ + x_base + (seq_len - 2) * x_stride_token + ] + conv_states[state_base + 2 * state_stride_token] = x[ + x_base + (seq_len - 1) * x_stride_token + ] + else: + state_cut = 3 - seq_len + if seq_len == 1: + if load_init: + conv_states[state_base] = conv_states[ + state_base + state_stride_token + ] + conv_states[state_base + state_stride_token] = ( + conv_states[state_base + 2 * state_stride_token] + ) + else: + conv_states[state_base] = T.Cast(dtype, 0.0) + conv_states[state_base + state_stride_token] = T.Cast( + dtype, 0.0 + ) + conv_states[state_base + 2 * state_stride_token] = x[x_base] + else: + if load_init: + conv_states[state_base] = conv_states[ + state_base + 2 * state_stride_token + ] + else: + conv_states[state_base] = T.Cast(dtype, 0.0) + conv_states[state_base + state_stride_token] = x[x_base] + conv_states[state_base + 2 * state_stride_token] = x[ + x_base + x_stride_token + ] + else: + col2 = T.Cast( + "float32", + x[x_base + (token_offset - 1) * x_stride_token], + ) + col1 = T.Cast( + "float32", + x[x_base + (token_offset - 2) * x_stride_token], + ) + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 3) * x_stride_token], + ) + + w0 = T.Cast("float32", weight[w_base]) + w1 = T.Cast("float32", weight[w_base + 1]) + w2 = T.Cast("float32", weight[w_base + 2]) + w3 = T.Cast("float32", weight[w_base + 3]) + + for token_i in T.serial(block_m): + if token_i < segment_len: + x_cur = T.Cast( + "float32", + x[x_base + (token_offset + token_i) * x_stride_token], + ) + acc = 0.0 + if has_bias: + acc = T.Cast("float32", bias[feat]) + acc += col0 * w0 + col1 * w1 + col2 * w2 + x_cur * w3 + col0 = col1 + col1 = col2 + col2 = x_cur + if silu_activation: + acc = acc / (1.0 + T.exp2(-acc * _LOG2E)) + out[out_base + (token_offset + token_i) * o_stride_token] = ( + T.Cast(dtype, acc) + ) + + return musa_causal_conv1d_prefill_width4 + + +_causal_conv1d_prefill_width4_kernel.mode = "lazy" + + +@functools.lru_cache(maxsize=64) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_CAUSAL_CONV1D_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _causal_conv1d_prefill_width4_body_kernel( + dtype: str, + cache_indices_dtype: str, + x_stride_token: int, + w_stride_dim: int, + o_stride_token: int, + has_bias: bool, + has_cache_indices: bool, + has_cache_index_mapping: bool, + use_pad_slot: bool, + silu_activation: bool, + block_m: int, + block_n: int, +): + x_numel = T.dynamic("x_numel") + w_numel = T.dynamic("w_numel") + bias_numel = T.dynamic("bias_numel") + cache_numel = T.dynamic("cache_numel") + mapping_numel = T.dynamic("mapping_numel") + query_numel = T.dynamic("query_numel") + out_numel = T.dynamic("out_numel") + + @T.prim_func + def musa_causal_conv1d_prefill_width4_body( + x: T.Tensor((x_numel,), dtype), + weight: T.Tensor((w_numel,), dtype), + bias: T.Tensor((bias_numel,), dtype), + cache_indices: T.Tensor((cache_numel,), cache_indices_dtype), + cache_index_mapping: T.Tensor((mapping_numel,), "int32"), + query_start_loc: T.Tensor((query_numel,), "int32"), + out: T.Tensor((out_numel,), dtype), + max_seq_len: T.int32, + dim: T.int32, + num_cache_lines: T.int32, + pad_slot_id: T.int32, + ): + with T.Kernel( + query_numel - 1, + T.ceildiv(max_seq_len - block_m, block_m), + T.ceildiv(dim, block_n), + threads=block_n, + ) as (seq_idx, body_chunk_idx, dim_block): + tid = T.get_thread_binding() + feat = dim_block * block_n + tid + seq_start = T.alloc_var("int32") + seq_end = T.alloc_var("int32") + seq_len = T.alloc_var("int32") + token_offset = T.alloc_var("int32") + segment_len = T.alloc_var("int32") + cache_idx = T.alloc_var("int32") + valid_seq = T.alloc_var("bool") + x_base = T.alloc_var("int32") + w_base = T.alloc_var("int32") + out_base = T.alloc_var("int32") + col0 = T.alloc_var("float32") + col1 = T.alloc_var("float32") + col2 = T.alloc_var("float32") + w0 = T.alloc_var("float32") + w1 = T.alloc_var("float32") + w2 = T.alloc_var("float32") + w3 = T.alloc_var("float32") + x_cur = T.alloc_var("float32") + acc = T.alloc_var("float32") + + seq_start = query_start_loc[seq_idx] + seq_end = query_start_loc[seq_idx + 1] + seq_len = seq_end - seq_start + token_offset = (body_chunk_idx + 1) * block_m + segment_len = seq_len - token_offset + if segment_len > block_m: + segment_len = block_m + + cache_idx = seq_idx + if has_cache_indices: + cache_idx = cache_indices[seq_idx] + if has_cache_index_mapping: + cache_idx = cache_index_mapping[cache_idx] + valid_seq = segment_len > 0 and feat < dim + if use_pad_slot and cache_idx == pad_slot_id: + valid_seq = False + if cache_idx >= num_cache_lines: + valid_seq = False + + if valid_seq: + x_base = seq_start * x_stride_token + feat + w_base = feat * w_stride_dim + out_base = seq_start * o_stride_token + feat + + col2 = T.Cast( + "float32", + x[x_base + (token_offset - 1) * x_stride_token], + ) + col1 = T.Cast( + "float32", + x[x_base + (token_offset - 2) * x_stride_token], + ) + col0 = T.Cast( + "float32", + x[x_base + (token_offset - 3) * x_stride_token], + ) + w0 = T.Cast("float32", weight[w_base]) + w1 = T.Cast("float32", weight[w_base + 1]) + w2 = T.Cast("float32", weight[w_base + 2]) + w3 = T.Cast("float32", weight[w_base + 3]) + + if segment_len == block_m: + for token_i in T.serial(block_m): + x_cur = T.Cast( + "float32", + x[x_base + (token_offset + token_i) * x_stride_token], + ) + acc = 0.0 + if has_bias: + acc = T.Cast("float32", bias[feat]) + acc += col0 * w0 + col1 * w1 + col2 * w2 + x_cur * w3 + col0 = col1 + col1 = col2 + col2 = x_cur + if silu_activation: + acc = acc / (1.0 + T.exp2(-acc * _LOG2E)) + out[out_base + (token_offset + token_i) * o_stride_token] = ( + T.Cast(dtype, acc) + ) + else: + for token_i in T.serial(block_m): + if token_i < segment_len: + x_cur = T.Cast( + "float32", + x[x_base + (token_offset + token_i) * x_stride_token], + ) + acc = 0.0 + if has_bias: + acc = T.Cast("float32", bias[feat]) + acc += col0 * w0 + col1 * w1 + col2 * w2 + x_cur * w3 + col0 = col1 + col1 = col2 + col2 = x_cur + if silu_activation: + acc = acc / (1.0 + T.exp2(-acc * _LOG2E)) + out[ + out_base + (token_offset + token_i) * o_stride_token + ] = T.Cast(dtype, acc) + + return musa_causal_conv1d_prefill_width4_body + + +_causal_conv1d_prefill_width4_body_kernel.mode = "lazy" + + +@functools.lru_cache(maxsize=64) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_CAUSAL_CONV1D_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _causal_conv1d_decode_width4_batched_kernel( + dtype: str, + cache_indices_dtype: str, + x_stride_token: int, + w_stride_dim: int, + state_stride_seq: int, + state_stride_dim: int, + state_stride_token: int, + o_stride_token: int, + has_bias: bool, + has_cache_indices: bool, + has_cache_index_mapping: bool, + has_initial_states: bool, + use_pad_slot: bool, + silu_activation: bool, + block_feats: int, + batch_per_block: int, +): + x_numel = T.dynamic("x_numel") + w_numel = T.dynamic("w_numel") + bias_numel = T.dynamic("bias_numel") + state_numel = T.dynamic("state_numel") + cache_numel = T.dynamic("cache_numel") + mapping_numel = T.dynamic("mapping_numel") + init_numel = T.dynamic("init_numel") + out_numel = T.dynamic("out_numel") + num_threads = block_feats * batch_per_block + + @T.prim_func + def musa_causal_conv1d_decode_width4_batched( + x: T.Tensor((x_numel,), dtype), + weight: T.Tensor((w_numel,), dtype), + bias: T.Tensor((bias_numel,), dtype), + conv_states: T.Tensor((state_numel,), dtype), + cache_indices: T.Tensor((cache_numel,), cache_indices_dtype), + cache_index_mapping: T.Tensor((mapping_numel,), "int32"), + has_initial_state: T.Tensor((init_numel,), "bool"), + out: T.Tensor((out_numel,), dtype), + batch: T.int32, + dim: T.int32, + num_cache_lines: T.int32, + pad_slot_id: T.int32, + ): + with T.Kernel( + T.ceildiv(batch, batch_per_block), + T.ceildiv(dim, block_feats), + threads=num_threads, + ) as (batch_block, dim_block): + tid = T.get_thread_binding() + batch_lane = tid // block_feats + feat_lane = tid - batch_lane * block_feats + seq_idx = batch_block * batch_per_block + batch_lane + feat = dim_block * block_feats + feat_lane + cache_idx = T.alloc_var("int32") + valid = T.alloc_var("bool") + load_init = T.alloc_var("bool") + x_base = T.alloc_var("int32") + w_base = T.alloc_var("int32") + state_base = T.alloc_var("int32") + col0 = T.alloc_var("float32") + col1 = T.alloc_var("float32") + col2 = T.alloc_var("float32") + x_cur = T.alloc_var("float32") + acc = T.alloc_var("float32") + + cache_idx = seq_idx + if has_cache_indices and seq_idx < batch: + cache_idx = cache_indices[seq_idx] + if has_cache_index_mapping: + cache_idx = cache_index_mapping[cache_idx] + valid = seq_idx < batch and feat < dim + if use_pad_slot and cache_idx == pad_slot_id: + valid = False + if cache_idx >= num_cache_lines: + valid = False + + if valid: + load_init = False + if has_initial_states: + load_init = has_initial_state[seq_idx] + + x_base = seq_idx * x_stride_token + feat + w_base = feat * w_stride_dim + state_base = cache_idx * state_stride_seq + feat * state_stride_dim + x_cur = T.Cast("float32", x[x_base]) + col0 = 0.0 + col1 = 0.0 + col2 = 0.0 + if load_init: + col0 = T.Cast("float32", conv_states[state_base]) + col1 = T.Cast( + "float32", conv_states[state_base + state_stride_token] + ) + col2 = T.Cast( + "float32", conv_states[state_base + 2 * state_stride_token] + ) + + acc = ( + col0 * T.Cast("float32", weight[w_base]) + + col1 * T.Cast("float32", weight[w_base + 1]) + + col2 * T.Cast("float32", weight[w_base + 2]) + + x_cur * T.Cast("float32", weight[w_base + 3]) + ) + if has_bias: + acc += T.Cast("float32", bias[feat]) + if silu_activation: + acc = acc / (1.0 + T.exp2(-acc * _LOG2E)) + + out[seq_idx * o_stride_token + feat] = T.Cast(dtype, acc) + if load_init: + conv_states[state_base] = conv_states[ + state_base + state_stride_token + ] + conv_states[state_base + state_stride_token] = conv_states[ + state_base + 2 * state_stride_token + ] + else: + conv_states[state_base] = T.Cast(dtype, 0.0) + conv_states[state_base + state_stride_token] = T.Cast(dtype, 0.0) + conv_states[state_base + 2 * state_stride_token] = T.Cast(dtype, x_cur) + + return musa_causal_conv1d_decode_width4_batched + + +_causal_conv1d_decode_width4_batched_kernel.mode = "lazy" + + +def _check_inputs( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + cache_indices: Optional[torch.Tensor], + has_initial_state: Optional[torch.Tensor], + cache_index_mapping: Optional[torch.Tensor], + activation: Optional[Union[str, bool]], + seq_lens_cpu: List[int], +) -> tuple[int, int, int, int, str]: + if isinstance(activation, bool) and activation: + activation = "silu" + if activation not in (None, "silu", "swish"): + raise NotImplementedError("activation must be None, silu, or swish") + if x.dim() != 2: + raise ValueError( + "TileLang causal_conv1d_fwd expects varlen x with shape (dim, total_tokens)" + ) + if weight.dim() != 2: + raise ValueError("weight must be a 2D tensor") + if query_start_loc is None or query_start_loc.dim() != 1: + raise ValueError("query_start_loc must be a 1D tensor") + if query_start_loc.dtype is not torch.int32: + raise TypeError("query_start_loc must be int32") + dim, total_tokens = x.shape + weight_dim, width = weight.shape + if weight_dim != dim: + raise ValueError("weight first dimension must match x dim") + if width < 2 or width > 5: + raise ValueError("TileLang causal_conv1d_fwd supports width in [2, 5]") + if bias is not None and (bias.dim() != 1 or bias.numel() != dim): + raise ValueError("bias must have shape (dim,)") + if cache_indices is not None: + if cache_indices.dim() != 1 or cache_indices.dtype not in ( + torch.int32, + torch.int64, + ): + raise TypeError("cache_indices must be a 1D int32 or int64 tensor") + if cache_indices.numel() != query_start_loc.numel() - 1: + raise ValueError("cache_indices length must match batch size") + if cache_index_mapping is not None: + if ( + cache_index_mapping.dim() != 1 + or cache_index_mapping.dtype is not torch.int32 + ): + raise TypeError("cache_index_mapping must be a 1D int32 tensor") + if cache_indices is None: + raise ValueError("cache_index_mapping requires cache_indices") + if has_initial_state is not None: + if has_initial_state.dim() != 1 or has_initial_state.dtype is not torch.bool: + raise TypeError("has_initial_state must be a 1D bool tensor") + if has_initial_state.numel() != query_start_loc.numel() - 1: + raise ValueError("has_initial_state length must match batch size") + if conv_states is None: + raise ValueError("has_initial_state requires conv_states") + if conv_states is not None: + if conv_states.dim() != 3: + raise ValueError( + "conv_states must have shape (num_cache_lines, dim, state_len)" + ) + if conv_states.size(1) != dim or conv_states.size(2) < width - 1: + raise ValueError("conv_states shape is incompatible with x/weight") + seq_lens_sum = sum(seq_lens_cpu) + max_seq_len = max(seq_lens_cpu, default=0) + if seq_lens_sum != total_tokens: + # CUDA graph / DP padding can leave seq_lens_cpu describing the padded + # batch while x and query_start_loc describe the actual packed tokens. + # The kernel gets true per-sequence bounds from query_start_loc, so use + # a conservative launch bound instead of rejecting a valid packed input. + max_seq_len = total_tokens + return ( + dim, + total_tokens, + width, + max_seq_len, + tilelang_dtype(x.dtype), + ) + + +def _causal_conv1d_fwd_impl( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + seq_lens_cpu: List[int], + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + cache_index_mapping: Optional[torch.Tensor] = None, + activation: Optional[Union[str, bool]] = "silu", + pad_slot_id: int = PAD_SLOT_ID, +) -> torch.Tensor: + if isinstance(activation, bool) and activation: + activation = "silu" + dim, _total_tokens, width, max_seq_len, dtype = _check_inputs( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + cache_index_mapping, + activation, + seq_lens_cpu, + ) + + out = torch.empty_like(x) + x_arg = storage_window(x) + weight_arg = storage_window(weight) + out_arg = storage_window(out) + bias_arg = storage_window(bias) if bias is not None else x_arg + conv_states_arg = storage_window(conv_states) if conv_states is not None else x_arg + cache_indices_arg = cache_indices if cache_indices is not None else query_start_loc + cache_index_mapping_arg = ( + cache_index_mapping if cache_index_mapping is not None else cache_indices_arg + ) + has_initial_state_arg = ( + has_initial_state + if has_initial_state is not None + else torch.empty((1,), dtype=torch.bool, device=x.device) + ) + + block_n = 256 if dim >= 256 else max(32, _next_power_of_2(dim)) + if max_seq_len == 1 and query_start_loc.numel() > 128 and dim >= 256: + block_n = 128 + block_m = 8 + if width == 4 and max_seq_len >= 128: + if max_seq_len <= 512: + block_m = 4 + block_n = 256 + elif max_seq_len < 4096: + if query_start_loc.numel() > 4: + block_m = 28 + block_n = 256 + else: + block_m = 12 + block_n = 128 if query_start_loc.numel() > 2 else 256 + else: + block_m = 28 + block_n = 256 + num_cache_lines = conv_states.size(0) if conv_states is not None else 0 + state_stride_seq = conv_states.stride(0) if conv_states is not None else 0 + state_stride_dim = conv_states.stride(1) if conv_states is not None else 0 + state_stride_token = conv_states.stride(2) if conv_states is not None else 0 + cache_indices_dtype = ( + _tilelang_index_dtype(cache_indices.dtype) + if cache_indices is not None + else "int32" + ) + + if ( + width == 4 + and max_seq_len == 1 + and query_start_loc.numel() > 256 + and conv_states is not None + and dim >= 4096 + and x.stride(0) == 1 + and out.stride(0) == 1 + and weight.stride(1) == 1 + ): + _causal_conv1d_decode_width4_batched_kernel( + dtype, + cache_indices_dtype, + int(x.stride(1)), + int(weight.stride(0)), + int(state_stride_seq), + int(state_stride_dim), + int(state_stride_token), + int(out.stride(1)), + bias is not None, + cache_indices is not None, + cache_index_mapping is not None, + has_initial_state is not None, + pad_slot_id is not None, + activation in ("silu", "swish"), + 256, + 1, + )( + x_arg, + weight_arg, + bias_arg, + conv_states_arg, + cache_indices_arg, + cache_index_mapping_arg, + has_initial_state_arg, + out_arg, + int(query_start_loc.numel() - 1), + int(dim), + int(num_cache_lines), + int(pad_slot_id if pad_slot_id is not None else PAD_SLOT_ID), + ) + return out + + if ( + _ENABLE_WIDTH4_PREFILL_SPLIT + and width == 4 + and max_seq_len >= 128 + and query_start_loc.numel() > 2 + and conv_states is not None + and cache_indices is not None + and x.stride(0) == 1 + and out.stride(0) == 1 + and weight.stride(1) == 1 + ): + _causal_conv1d_prefill_width4_kernel( + dtype, + cache_indices_dtype, + int(x.stride(1)), + int(weight.stride(0)), + int(state_stride_seq), + int(state_stride_dim), + int(state_stride_token), + int(out.stride(1)), + bias is not None, + cache_indices is not None, + cache_index_mapping is not None, + has_initial_state is not None, + pad_slot_id is not None, + activation in ("silu", "swish"), + int(block_m), + int(block_n), + )( + x_arg, + weight_arg, + bias_arg, + conv_states_arg, + cache_indices_arg, + cache_index_mapping_arg, + has_initial_state_arg, + query_start_loc, + out_arg, + int(block_m), + int(dim), + int(num_cache_lines), + int(pad_slot_id if pad_slot_id is not None else PAD_SLOT_ID), + ) + if max_seq_len > block_m: + _causal_conv1d_prefill_width4_body_kernel( + dtype, + cache_indices_dtype, + int(x.stride(1)), + int(weight.stride(0)), + int(out.stride(1)), + bias is not None, + cache_indices is not None, + cache_index_mapping is not None, + pad_slot_id is not None, + activation in ("silu", "swish"), + int(block_m), + int(block_n), + )( + x_arg, + weight_arg, + bias_arg, + cache_indices_arg, + cache_index_mapping_arg, + query_start_loc, + out_arg, + int(max_seq_len), + int(dim), + int(num_cache_lines), + int(pad_slot_id if pad_slot_id is not None else PAD_SLOT_ID), + ) + return out + + if ( + width == 4 + and dim >= 4096 + and max_seq_len == 1 + and query_start_loc.numel() == 2 + and x.stride(0) == 1 + and out.stride(0) == 1 + and weight.stride(1) == 1 + ): + block_feats = 256 + vec_elems = 1 + _causal_conv1d_fwd_width4_vec_kernel( + dtype, + cache_indices_dtype, + int(x.stride(1)), + int(weight.stride(0)), + int(state_stride_seq), + int(state_stride_dim), + int(state_stride_token), + int(out.stride(1)), + bias is not None, + conv_states is not None, + cache_indices is not None, + cache_index_mapping is not None, + has_initial_state is not None, + pad_slot_id is not None, + activation in ("silu", "swish"), + int(block_m), + int(block_feats), + int(vec_elems), + )( + x_arg, + weight_arg, + bias_arg, + conv_states_arg, + cache_indices_arg, + cache_index_mapping_arg, + has_initial_state_arg, + query_start_loc, + out_arg, + int(max_seq_len), + int(dim), + int(num_cache_lines), + int(pad_slot_id if pad_slot_id is not None else PAD_SLOT_ID), + ) + return out + + _causal_conv1d_fwd_kernel( + dtype, + cache_indices_dtype, + int(width), + int(x.stride(0)), + int(x.stride(1)), + int(weight.stride(0)), + int(weight.stride(1)), + int(state_stride_seq), + int(state_stride_dim), + int(state_stride_token), + int(out.stride(0)), + int(out.stride(1)), + bias is not None, + conv_states is not None, + cache_indices is not None, + cache_index_mapping is not None, + has_initial_state is not None, + pad_slot_id is not None, + activation in ("silu", "swish"), + int(block_m), + int(block_n), + )( + x_arg, + weight_arg, + bias_arg, + conv_states_arg, + cache_indices_arg, + cache_index_mapping_arg, + has_initial_state_arg, + query_start_loc, + out_arg, + int(max_seq_len), + int(dim), + int(num_cache_lines), + int(pad_slot_id if pad_slot_id is not None else PAD_SLOT_ID), + ) + return out + + +@register_custom_op( + op_name="musa_causal_conv1d_fwd", + mutates_args=["conv_states"], +) +def _causal_conv1d_fwd_custom( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + seq_lens_cpu: List[int], + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + activation: Optional[str] = "silu", + pad_slot_id: int = PAD_SLOT_ID, + cache_index_mapping: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return _causal_conv1d_fwd_impl( + x, + weight, + bias, + conv_states, + query_start_loc, + seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + cache_index_mapping=cache_index_mapping, + activation=activation, + pad_slot_id=pad_slot_id, + ) + + +def causal_conv1d_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + seq_lens_cpu: List[int], + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + activation: Optional[Union[str, bool]] = "silu", + pad_slot_id: int = PAD_SLOT_ID, + cache_index_mapping: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if isinstance(seq_lens_cpu, torch.Tensor) and seq_lens_cpu.device.type != "cpu": + # Backward-compatible positional form used by the generic mamba wrapper: + # causal_conv1d_fwd(..., query_start_loc, cache_indices, + # has_initial_state, activation, pad_slot_id) + old_cache_indices = seq_lens_cpu + old_has_initial_state = cache_indices + old_activation = has_initial_state + old_pad_slot_id = activation + seq_lens_cpu = query_start_loc.diff().detach().cpu().tolist() + cache_indices = old_cache_indices + has_initial_state = old_has_initial_state + activation = old_activation + pad_slot_id = old_pad_slot_id + elif isinstance(seq_lens_cpu, torch.Tensor): + seq_lens_cpu = seq_lens_cpu.detach().cpu().tolist() + if isinstance(activation, bool): + activation = "silu" if activation else None + return _causal_conv1d_fwd_impl( + x, + weight, + bias, + conv_states, + query_start_loc, + seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + cache_index_mapping=cache_index_mapping, + activation=activation, + pad_slot_id=pad_slot_id, + ) + + +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + conv_states: Optional[torch.Tensor], + query_start_loc: torch.Tensor, + seq_lens_cpu: List[int], + cache_indices: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + activation: Optional[Union[str, bool]] = "silu", + pad_slot_id: int = PAD_SLOT_ID, + cache_index_mapping: Optional[torch.Tensor] = None, + **_: object, +) -> torch.Tensor: + return causal_conv1d_fwd( + x, + weight, + bias, + conv_states, + query_start_loc, + seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + cache_index_mapping=cache_index_mapping, + activation=activation, + pad_slot_id=pad_slot_id, + ) + + +def musa_tilelang_causal_conv1d_fn( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices=None, + has_initial_state=None, + activation="silu", + pad_slot_id=PAD_SLOT_ID, + cache_index_mapping=None, + metadata=None, + **_, +): + """Drop-in for vllm Triton causal_conv1d_fn (prefill). Synthesizes + seq_lens_cpu from query_start_loc and matches conv_states dtype.""" + qsl_cpu = getattr(metadata, "non_spec_query_start_loc_cpu", None) + if qsl_cpu is not None: + seq_lens_cpu = qsl_cpu.diff().tolist() + else: + seq_lens_cpu = query_start_loc.diff().cpu().tolist() + orig_dtype = x.dtype + if conv_states is not None and x.dtype != conv_states.dtype: + x = x.to(conv_states.dtype) + out = causal_conv1d_fn( + x, + weight, + bias, + conv_states, + query_start_loc, + seq_lens_cpu, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + cache_index_mapping=cache_index_mapping, + ) + return out.to(orig_dtype) + + +# --- MUSA: single-token decode causal_conv1d (width-4) via the batched TileLang +# decode kernel. Drop-in for vllm Triton causal_conv1d_update; returns None when +# the fast path does not apply so the caller keeps the Triton path. +_DECODE_HAS_INIT_BUF = {} + + +def _decode_has_init_buffer(batch, device): + key = str(device) + buf = _DECODE_HAS_INIT_BUF.get(key) + if buf is None or buf.numel() < batch: + buf = torch.ones(max(int(batch), 2048), device=device, dtype=torch.bool) + _DECODE_HAS_INIT_BUF[key] = buf + return buf[:batch] + + +def musa_tilelang_causal_conv1d_update( + x, + conv_state, + weight, + bias=None, + activation=None, + conv_state_indices=None, + **_, +): + """Width-4 single-token decode conv. x:[batch,dim], conv_state:[lines,dim,>=3], + weight:[dim,width]. Updates conv_state in place. None => use Triton fallback.""" + if conv_state is None or weight.dim() != 2 or weight.shape[1] != 4: + return None + if conv_state.dim() != 3 or conv_state.size(2) < 3 or x.dim() != 2: + return None + batch, dim = x.shape + orig_dtype = x.dtype + if x.dtype != conv_state.dtype: + x = x.to(conv_state.dtype) + xt = x.transpose(0, 1) # [dim, batch] VIEW, stride(0)==1 + out_bd = torch.empty_like(x) # [batch, dim] + outt = out_bd.transpose(0, 1) # [dim, batch] VIEW, stride(0)==1 + if conv_state_indices is None: + idx = torch.arange(batch, device=x.device, dtype=torch.int32) + else: + idx = conv_state_indices + has_init = _decode_has_init_buffer(batch, x.device) + silu = activation in ("silu", "swish", True) + ker = _causal_conv1d_decode_width4_batched_kernel( + tilelang_dtype(xt.dtype), + _tilelang_index_dtype(idx.dtype), + int(xt.stride(1)), + int(weight.stride(0)), + int(conv_state.stride(0)), + int(conv_state.stride(1)), + int(conv_state.stride(2)), + int(outt.stride(1)), + bias is not None, + True, + False, + True, + True, + silu, + 256, + 1, + ) + ker( + storage_window(xt), + storage_window(weight), + storage_window(bias) if bias is not None else storage_window(xt), + storage_window(conv_state), + idx, + idx, + has_init, + storage_window(outt), + int(batch), + int(dim), + int(conv_state.size(0)), + int(PAD_SLOT_ID), + ) + return out_bd.to(orig_dtype) diff --git a/vllm_musa/jit_kernel/tilelang/deep_gemm_contig_preprocess.py b/vllm_musa/jit_kernel/tilelang/deep_gemm_contig_preprocess.py new file mode 100644 index 000000000000..798d8203e497 --- /dev/null +++ b/vllm_musa/jit_kernel/tilelang/deep_gemm_contig_preprocess.py @@ -0,0 +1,917 @@ +"""TileLang MUSA DeepGEMM contiguous preprocess fast paths.""" + +import functools +from pathlib import Path + +import tilelang +import tilelang.language as T +import torch + + +# MUSA: local passthrough for register_custom_op (SGLang default eager=True returns fn). +def register_custom_op(fn=None, **_kw): + def _wrap(f): + return f + + return _wrap if fn is None else _wrap(fn) + + +_ATOMIC_HELPER_H = str((Path(__file__).resolve().parent / "_atomic_helper.h").resolve()) + +_PASS_CONFIGS = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} +for _key, _value in ( + ("TL_ENABLE_FAST_MATH", True), + ("TL_DISABLE_THREAD_STORAGE_SYNC", True), + ("TL_ENABLE_MUSA_BURST", True), + ("TL_ENABLE_REDUCE_BURST", True), + ("TL_DISABLE_SAFE_MEMORY_ACCESS", True), + ("TL_DISABLE_INDEX_TYPE_PROMOTION", True), +): + if hasattr(tilelang.PassConfigKey, _key): + _PASS_CONFIGS[getattr(tilelang.PassConfigKey, _key)] = _value + +_COMPILE_FLAGS = [ + "-Od3", + "-fno-signed-zeros", + "-fmusa-flush-denormals-to-zero", + "-mllvm", + "-misched=mtgpu-max-ilp", + "-mllvm", + "-mtgpu-if-convert=1", + "-mllvm", + "-mtgpu-tiny-offset-hint=1", + "-mllvm", + "-mtgpu-enable-postra-sched=0", + "-mllvm", + "-misched-recompute-slotindex=1", + "-mllvm", + "-mtgpu-combine-fop-instr=1", +] + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _clear_i32_kernel(): + n = T.dynamic("n") + + @T.prim_func + def deep_gemm_contig_preprocess_clear_i32_kernel( + ptr: T.Tensor((n,), "int32"), blocks: T.int32, total: T.int32 + ): + with T.Kernel(blocks, threads=256) as (bid,): + tid = T.get_thread_binding() + idx = bid * 256 + tid + if idx < total: + ptr[idx] = T.int32(0) + + return deep_gemm_contig_preprocess_clear_i32_kernel + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _fill_i32_kernel(): + n = T.dynamic("n") + + @T.prim_func + def deep_gemm_contig_preprocess_fill_i32_kernel( + ptr: T.Tensor((n,), "int32"), blocks: T.int32, total: T.int32, value: T.int32 + ): + with T.Kernel(blocks, threads=256) as (bid,): + tid = T.get_thread_binding() + idx = bid * 256 + tid + if idx < total: + ptr[idx] = value + + return deep_gemm_contig_preprocess_fill_i32_kernel + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _count_topk_block_hist_kernel(max_experts: int): + n = T.dynamic("n") + e = T.dynamic("e") + slots_per_block = 1024 + + @T.prim_func + def deep_gemm_contig_preprocess_count_topk_block_hist_kernel( + topk_ids: T.Tensor((n,), "int32"), + counts: T.Tensor((e,), "int32"), + blocks: T.int32, + num_slots: T.int32, + num_local_experts: T.int32, + ): + with T.Kernel(blocks, threads=256) as (bid,): + tid = T.get_thread_binding() + local_counts = T.alloc_shared((max_experts,), "int32") + + for i in T.serial(T.ceildiv(max_experts, 256)): + expert = i * 256 + tid + if expert < num_local_experts: + local_counts[expert] = T.int32(0) + T.sync_threads() + + block_start = bid * slots_per_block + for i in T.serial(T.ceildiv(slots_per_block, 256)): + offset = i * 256 + tid + slot = block_start + offset + if slot < num_slots: + expert = topk_ids[slot] + if expert >= 0 and expert < num_local_experts: + T.atomic_add(local_counts[expert], T.int32(1)) + T.sync_threads() + + for i in T.serial(T.ceildiv(max_experts, 256)): + expert = i * 256 + tid + if expert < num_local_experts: + cnt = local_counts[expert] + if cnt != 0: + T.atomic_add(counts[expert], cnt) + + return deep_gemm_contig_preprocess_count_topk_block_hist_kernel + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _count_topk_single_block_kernel(max_experts: int): + n = T.dynamic("n") + e = T.dynamic("e") + slots_per_block = 1024 + + @T.prim_func + def deep_gemm_contig_preprocess_count_topk_single_block_kernel( + topk_ids: T.Tensor((n,), "int32"), + counts: T.Tensor((e,), "int32"), + num_slots: T.int32, + num_local_experts: T.int32, + ): + with T.Kernel(1, threads=256) as (_bid,): + tid = T.get_thread_binding() + local_counts = T.alloc_shared((max_experts,), "int32") + + for i in T.serial(T.ceildiv(max_experts, 256)): + expert = i * 256 + tid + if expert < num_local_experts: + local_counts[expert] = T.int32(0) + T.sync_threads() + + for i in T.serial(T.ceildiv(slots_per_block, 256)): + slot = i * 256 + tid + if slot < num_slots: + expert = topk_ids[slot] + if expert >= 0 and expert < num_local_experts: + T.atomic_add(local_counts[expert], T.int32(1)) + T.sync_threads() + + for i in T.serial(T.ceildiv(max_experts, 256)): + expert = i * 256 + tid + if expert < num_local_experts: + counts[expert] = local_counts[expert] + + return deep_gemm_contig_preprocess_count_topk_single_block_kernel + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _count_prefix_topk_single_block_kernel(max_experts: int, max_block_m: int): + n = T.dynamic("n") + e = T.dynamic("e") + m = T.dynamic("m") + slots_per_block = 1024 + scan_size = 1 << (max_experts - 1).bit_length() + threads = max(256, scan_size) + + @T.prim_func + def deep_gemm_contig_preprocess_count_prefix_topk_single_block_kernel( + topk_ids: T.Tensor((n,), "int32"), + counts: T.Tensor((e,), "int32"), + cursor: T.Tensor((e,), "int32"), + m_indices: T.Tensor((m,), "int32"), + num_slots: T.int32, + num_local_experts: T.int32, + block_m: T.int32, + ): + with T.Kernel(1, threads=threads) as (_bid,): + tid = T.get_thread_binding() + local_counts = T.alloc_shared((max_experts,), "int32") + prefix = T.alloc_shared((scan_size,), "int32") + expert = T.alloc_var("int32") + cnt = T.alloc_var("int32") + aligned_cnt = T.alloc_var("int32") + start = T.alloc_var("int32") + addend = T.alloc_var("int32") + + for i in T.serial(T.ceildiv(max_experts, threads)): + expert = i * threads + tid + if expert < num_local_experts: + local_counts[expert] = T.int32(0) + T.sync_threads() + + for i in T.serial(T.ceildiv(slots_per_block, threads)): + slot = i * threads + tid + if slot < num_slots: + expert = topk_ids[slot] + if expert >= 0 and expert < num_local_experts: + T.atomic_add(local_counts[expert], T.int32(1)) + T.sync_threads() + + if tid < scan_size: + if tid < num_local_experts: + cnt = local_counts[tid] + counts[tid] = cnt + prefix[tid] = T.ceildiv(cnt, block_m) * block_m + else: + prefix[tid] = T.int32(0) + T.sync_threads() + + for offset in T.serial((scan_size.bit_length() - 1)): + step = 1 << offset + addend = T.int32(0) + if tid < scan_size and tid >= step: + addend = prefix[tid - step] + T.sync_threads() + if tid < scan_size and tid >= step: + prefix[tid] = prefix[tid] + addend + T.sync_threads() + + if tid < num_local_experts: + cnt = local_counts[tid] + aligned_cnt = T.ceildiv(cnt, block_m) * block_m + start = prefix[tid] - aligned_cnt + cursor[tid] = start + for pad in T.serial(max_block_m): + if pad < aligned_cnt - cnt: + m_indices[start + cnt + pad] = tid + + return deep_gemm_contig_preprocess_count_prefix_topk_single_block_kernel + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _prefix_counts_kernel(max_experts: int, max_block_m: int): + e = T.dynamic("e") + m = T.dynamic("m") + + @T.prim_func + def deep_gemm_contig_preprocess_prefix_counts( + counts: T.Tensor((e,), "int32"), + cursor: T.Tensor((e,), "int32"), + m_indices: T.Tensor((m,), "int32"), + num_local_experts: T.int32, + block_m: T.int32, + ): + with T.Kernel(1, threads=256) as (_bid,): + tid = T.get_thread_binding() + acc = T.alloc_var("int32") + cnt = T.alloc_var("int32") + aligned_cnt = T.alloc_var("int32") + start = T.alloc_var("int32") + + if tid == 0: + acc = T.int32(0) + for expert in T.serial(max_experts): + if expert < num_local_experts: + cursor[expert] = acc + cnt = counts[expert] + aligned_cnt = T.ceildiv(cnt, block_m) * block_m + acc = acc + aligned_cnt + T.sync_threads() + + for i in T.serial(T.ceildiv(max_experts, 256)): + expert = i * 256 + tid + if expert < num_local_experts: + cnt = counts[expert] + aligned_cnt = T.ceildiv(cnt, block_m) * block_m + start = cursor[expert] + for pad in T.serial(max_block_m): + if pad < aligned_cnt - cnt: + m_indices[start + cnt + pad] = expert + + return deep_gemm_contig_preprocess_prefix_counts + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _prefix_counts_scan_kernel(max_experts: int, max_block_m: int): + e = T.dynamic("e") + m = T.dynamic("m") + scan_size = 1 << (max_experts - 1).bit_length() + threads = max(256, scan_size) + + @T.prim_func + def deep_gemm_contig_preprocess_prefix_counts_scan( + counts: T.Tensor((e,), "int32"), + cursor: T.Tensor((e,), "int32"), + m_indices: T.Tensor((m,), "int32"), + num_local_experts: T.int32, + block_m: T.int32, + ): + with T.Kernel(1, threads=threads) as (_bid,): + tid = T.get_thread_binding() + prefix = T.alloc_shared((scan_size,), "int32") + cnt = T.alloc_var("int32") + aligned_cnt = T.alloc_var("int32") + start = T.alloc_var("int32") + + if tid < scan_size: + if tid < num_local_experts: + cnt = counts[tid] + prefix[tid] = T.ceildiv(cnt, block_m) * block_m + else: + prefix[tid] = T.int32(0) + T.sync_threads() + + T.cumsum(prefix, prefix, dim=0) + T.sync_threads() + + if tid < num_local_experts: + cnt = counts[tid] + aligned_cnt = T.ceildiv(cnt, block_m) * block_m + start = prefix[tid] - aligned_cnt + cursor[tid] = start + for pad in T.serial(max_block_m): + if pad < aligned_cnt - cnt: + m_indices[start + cnt + pad] = tid + + return deep_gemm_contig_preprocess_prefix_counts_scan + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], target="musa", pass_configs=_PASS_CONFIGS, compile_flags=_COMPILE_FLAGS +) +def _prefix_counts_tree_kernel(max_experts: int, max_block_m: int): + e = T.dynamic("e") + m = T.dynamic("m") + scan_size = 1 << (max_experts - 1).bit_length() + threads = max(256, scan_size) + + @T.prim_func + def deep_gemm_contig_preprocess_prefix_counts_tree( + counts: T.Tensor((e,), "int32"), + cursor: T.Tensor((e,), "int32"), + m_indices: T.Tensor((m,), "int32"), + num_local_experts: T.int32, + block_m: T.int32, + ): + with T.Kernel(1, threads=threads) as (_bid,): + tid = T.get_thread_binding() + prefix = T.alloc_shared((scan_size,), "int32") + cnt = T.alloc_var("int32") + aligned_cnt = T.alloc_var("int32") + start = T.alloc_var("int32") + addend = T.alloc_var("int32") + + if tid < scan_size: + if tid < num_local_experts: + cnt = counts[tid] + prefix[tid] = T.ceildiv(cnt, block_m) * block_m + else: + prefix[tid] = T.int32(0) + T.sync_threads() + + for offset in T.serial((scan_size.bit_length() - 1)): + step = 1 << offset + addend = T.int32(0) + if tid < scan_size and tid >= step: + addend = prefix[tid - step] + T.sync_threads() + if tid < scan_size and tid >= step: + prefix[tid] = prefix[tid] + addend + T.sync_threads() + + if tid < num_local_experts: + cnt = counts[tid] + aligned_cnt = T.ceildiv(cnt, block_m) * block_m + start = prefix[tid] - aligned_cnt + cursor[tid] = start + for pad in T.serial(max_block_m): + if pad < aligned_cnt - cnt: + m_indices[start + cnt + pad] = tid + + return deep_gemm_contig_preprocess_prefix_counts_tree + + +def _prefix_counts_no_pad_kernel(max_experts: int): + e = T.dynamic("e") + + @T.prim_func + def deep_gemm_contig_preprocess_prefix_counts_no_pad( + counts: T.Tensor((e,), "int32"), + cursor: T.Tensor((e,), "int32"), + num_local_experts: T.int32, + ): + with T.Kernel(1, threads=256) as (_bid,): + tid = T.get_thread_binding() + acc = T.alloc_var("int32") + + if tid == 0: + acc = T.int32(0) + for expert in T.serial(max_experts): + if expert < num_local_experts: + cursor[expert] = acc + acc = acc + counts[expert] + + return deep_gemm_contig_preprocess_prefix_counts_no_pad + + +@functools.lru_cache(maxsize=8) +@tilelang.jit( + out_idx=[], + target="musa", + pass_configs=_PASS_CONFIGS, + compile_flags=_COMPILE_FLAGS + ["-include", _ATOMIC_HELPER_H], +) +def _fp8_assign_compact_kernel( + input_dtype, + output_dtype, + hidden_groups: int, + groups_per_block: int, + vec_elems: int, + topk: int, +): + input_numel = T.dynamic("input_numel") + output_numel = T.dynamic("output_numel") + scale_numel = T.dynamic("scale_numel") + topk_ids_numel = T.dynamic("topk_ids_numel") + src2dst_numel = T.dynamic("src2dst_numel") + m_indices_numel = T.dynamic("m_indices_numel") + expert_numel = T.dynamic("expert_numel") + group_size = 128 + hidden_size = hidden_groups * group_size + threads_per_group = group_size // vec_elems + num_threads = threads_per_group * groups_per_block + + @T.prim_func + def deep_gemm_contig_preprocess_fp8_assign_compact( + hidden: T.Tensor((input_numel,), input_dtype), + topk_ids: T.Tensor((topk_ids_numel,), "int32"), + topk_ids_for_combine: T.Tensor((topk_ids_numel,), "int32"), + cursor: T.Tensor((expert_numel,), "int32"), + src2dst: T.Tensor((src2dst_numel,), "int32"), + m_indices: T.Tensor((m_indices_numel,), "int32"), + output_q: T.Tensor((output_numel,), output_dtype), + output_s: T.Tensor((scale_numel,), "float32"), + num_tokens: T.int32, + num_local_experts: T.int32, + eps: T.float32, + max_8bit: T.float32, + ): + with T.Kernel(num_tokens, threads=num_threads) as (bt,): + tid = T.get_thread_binding() + subgroup = tid // threads_per_group + lane = tid % threads_per_group + elem_base = T.alloc_var("int32") + hidden_group = T.alloc_var("int32") + input_base = T.alloc_var("int64") + local_absmax = T.alloc_local((1,), "float32") + scale = T.alloc_local((1,), "float32") + scale_inv = T.alloc_local((1,), "float32") + values = T.alloc_local((vec_elems,), "float32") + dst = T.alloc_var("int32") + out_base = T.alloc_var("int64") + dst_shared = T.alloc_shared((topk,), "int32") + + if tid < topk: + slot = bt * topk + tid + expert = topk_ids[slot] + if expert >= 0 and expert < num_local_experts: + topk_ids_for_combine[slot] = expert + dst = T.call_extern( + "int32", + "sgl_tl_atomic_add_offset", + T.address_of(cursor[0]), + expert, + T.int32(1), + ) + if dst < m_indices_numel: + src2dst[slot] = dst + m_indices[dst] = expert + dst_shared[tid] = dst + else: + topk_ids_for_combine[slot] = num_local_experts + src2dst[slot] = T.int32(-1) + dst_shared[tid] = T.int32(-1) + else: + topk_ids_for_combine[slot] = num_local_experts + src2dst[slot] = T.int32(-1) + dst_shared[tid] = T.int32(-1) + T.sync_threads() + + elem_base = lane * vec_elems + for tile in T.serial(T.ceildiv(hidden_groups, groups_per_block)): + hidden_group = tile * groups_per_block + subgroup + if hidden_group < hidden_groups: + input_base = ( + T.Cast("int64", bt) * hidden_size + + hidden_group * group_size + + elem_base + ) + local_absmax[0] = eps + for i in T.vectorized(vec_elems): + values[i] = T.Cast("float32", hidden[input_base + i]) + local_absmax[0] = T.max(local_absmax[0], T.abs(values[i])) + + if threads_per_group >= 32: + local_absmax[0] = T.max( + local_absmax[0], T.shfl_xor(local_absmax[0], 16) + ) + if threads_per_group >= 16: + local_absmax[0] = T.max( + local_absmax[0], T.shfl_xor(local_absmax[0], 8) + ) + if threads_per_group >= 8: + local_absmax[0] = T.max( + local_absmax[0], T.shfl_xor(local_absmax[0], 4) + ) + if threads_per_group >= 4: + local_absmax[0] = T.max( + local_absmax[0], T.shfl_xor(local_absmax[0], 2) + ) + if threads_per_group >= 2: + local_absmax[0] = T.max( + local_absmax[0], T.shfl_xor(local_absmax[0], 1) + ) + + scale_inv[0] = local_absmax[0] / max_8bit + scale[0] = max_8bit / local_absmax[0] + + for topk_idx in T.serial(topk): + dst = dst_shared[topk_idx] + if dst >= 0: + out_base = ( + T.Cast("int64", dst) * hidden_size + + hidden_group * group_size + + elem_base + ) + if vec_elems == 4: + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base, + T.min( + T.max(values[0] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[1] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[2] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[3] * scale[0], -max_8bit), max_8bit + ), + ) + elif vec_elems == 8: + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base, + T.min( + T.max(values[0] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[1] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[2] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[3] * scale[0], -max_8bit), max_8bit + ), + ) + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base + 4, + T.min( + T.max(values[4] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[5] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[6] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[7] * scale[0], -max_8bit), max_8bit + ), + ) + else: + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base, + T.min( + T.max(values[0] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[1] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[2] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[3] * scale[0], -max_8bit), max_8bit + ), + ) + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base + 4, + T.min( + T.max(values[4] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[5] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[6] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[7] * scale[0], -max_8bit), max_8bit + ), + ) + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base + 8, + T.min( + T.max(values[8] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[9] * scale[0], -max_8bit), max_8bit + ), + T.min( + T.max(values[10] * scale[0], -max_8bit), + max_8bit, + ), + T.min( + T.max(values[11] * scale[0], -max_8bit), + max_8bit, + ), + ) + T.call_extern( + "handle", + "sgl_tl_store_fp8e4m3x4", + T.address_of(output_q[0]), + out_base + 12, + T.min( + T.max(values[12] * scale[0], -max_8bit), + max_8bit, + ), + T.min( + T.max(values[13] * scale[0], -max_8bit), + max_8bit, + ), + T.min( + T.max(values[14] * scale[0], -max_8bit), + max_8bit, + ), + T.min( + T.max(values[15] * scale[0], -max_8bit), + max_8bit, + ), + ) + if lane == 0: + for topk_idx in T.serial(topk): + dst = dst_shared[topk_idx] + if dst >= 0: + output_s[dst * hidden_groups + hidden_group] = ( + scale_inv[0] + ) + + return deep_gemm_contig_preprocess_fp8_assign_compact + + +def _fp8_config(hidden_size: int) -> tuple[int, int, int] | None: + if hidden_size <= 0 or hidden_size % 128 != 0: + return None + hidden_groups = hidden_size // 128 + if hidden_groups == 22: + return hidden_groups, hidden_groups, 8 + if hidden_groups == 24: + return hidden_groups, hidden_groups, 8 + if hidden_groups <= 32: + return hidden_groups, hidden_groups, 4 + if hidden_groups <= 64: + return hidden_groups, hidden_groups, 8 + if hidden_groups <= 128: + return hidden_groups, hidden_groups, 16 + return None + + +def _impl_fp8( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + m_indices: torch.Tensor, + src2dst: torch.Tensor, + topk_ids_for_combine: torch.Tensor, + counts: torch.Tensor, + cursor: torch.Tensor, + num_local_experts: int, + block_m: int, +) -> None: + num_slots = topk_ids.numel() + num_tokens = hidden_states.shape[0] + topk = topk_ids.shape[-1] + config = _fp8_config(hidden_states.shape[1]) + if config is None: + raise RuntimeError( + f"unsupported TileLang DeepGEMM preprocess hidden={hidden_states.shape[1]}" + ) + if topk > 16: + raise RuntimeError(f"unsupported TileLang DeepGEMM fp8 preprocess topk={topk}") + hidden_groups, groups_per_block, vec_elems = config + num_local_experts = int(num_local_experts) + if num_local_experts <= 0: + raise ValueError(f"num_local_experts must be positive, got {num_local_experts}") + clear = _clear_i32_kernel() + fill = _fill_i32_kernel() + single_block_count = _count_topk_single_block_kernel(num_local_experts) + count = _count_topk_block_hist_kernel(num_local_experts) + block_m = int(block_m) + if block_m <= 0: + raise ValueError(f"block_m must be positive, got {block_m}") + single_block_count_prefix = _count_prefix_topk_single_block_kernel( + num_local_experts, block_m + ) + use_single_block_count = num_slots <= 1024 + use_single_block_count_prefix = block_m != 1 and use_single_block_count + prefix = ( + _prefix_counts_no_pad_kernel(num_local_experts) + if block_m == 1 + else ( + _prefix_counts_tree_kernel(num_local_experts, block_m) + if num_tokens < 512 and num_local_experts <= 1024 + else ( + _prefix_counts_scan_kernel(num_local_experts, block_m) + if num_local_experts <= 1024 + else _prefix_counts_kernel(num_local_experts, block_m) + ) + ) + ) + compact = _fp8_assign_compact_kernel( + hidden_states.dtype, + output.dtype, + hidden_groups, + groups_per_block, + vec_elems, + topk, + ) + if not use_single_block_count and not use_single_block_count_prefix: + clear(counts, tilelang.cdiv(num_local_experts, 256), num_local_experts) + count( + topk_ids.reshape(-1), + counts, + tilelang.cdiv(num_slots, 1024), + num_slots, + num_local_experts, + ) + if block_m == 1: + if use_single_block_count: + single_block_count( + topk_ids.reshape(-1), counts, num_slots, num_local_experts + ) + prefix(counts, cursor, num_local_experts) + else: + fill( + m_indices, + tilelang.cdiv(m_indices.numel(), 256), + m_indices.numel(), + num_local_experts - 1, + ) + if use_single_block_count_prefix: + single_block_count_prefix( + topk_ids.reshape(-1), + counts, + cursor, + m_indices, + num_slots, + num_local_experts, + block_m, + ) + else: + prefix(counts, cursor, m_indices, num_local_experts, block_m) + compact( + hidden_states.reshape(-1), + topk_ids.reshape(-1), + topk_ids_for_combine.reshape(-1), + cursor, + src2dst, + m_indices, + output.reshape(-1), + output_scale.reshape(-1), + num_tokens, + num_local_experts, + 1.0e-10, + 448.0, + ) + + +@register_custom_op( + op_name="musa_deep_gemm_contig_preprocess_fp8_tilelang", + mutates_args=[ + "output", + "output_scale", + "m_indices", + "src2dst", + "topk_ids_for_combine", + "counts", + "cursor", + ], +) +def _custom_fp8( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + m_indices: torch.Tensor, + src2dst: torch.Tensor, + topk_ids_for_combine: torch.Tensor, + counts: torch.Tensor, + cursor: torch.Tensor, + num_local_experts: int, + block_m: int, +) -> None: + _impl_fp8( + hidden_states, + topk_ids, + output, + output_scale, + m_indices, + src2dst, + topk_ids_for_combine, + counts, + cursor, + int(num_local_experts), + int(block_m), + ) + + +def deep_gemm_contig_preprocess_fp8_tilelang( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + m_indices: torch.Tensor, + src2dst: torch.Tensor, + topk_ids_for_combine: torch.Tensor, + counts: torch.Tensor, + cursor: torch.Tensor, + num_local_experts: int, + block_m: int, +) -> None: + _custom_fp8( + hidden_states, + topk_ids, + output, + output_scale, + m_indices, + src2dst, + topk_ids_for_combine, + counts, + cursor, + int(num_local_experts), + int(block_m), + ) + + +def can_use_fp8_tilelang( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + num_local_experts: int, + use_fp8_quant: bool, +) -> bool: + return ( + use_fp8_quant + and hidden_states.dim() == 2 + and _fp8_config(hidden_states.shape[1]) is not None + and topk_ids.dim() == 2 + and topk_ids.shape[-1] <= 16 + and num_local_experts > 0 + ) diff --git a/vllm_musa/jit_kernel/tilelang/gdn_fused_proj.py b/vllm_musa/jit_kernel/tilelang/gdn_fused_proj.py new file mode 100644 index 000000000000..69dd05ed27b5 --- /dev/null +++ b/vllm_musa/jit_kernel/tilelang/gdn_fused_proj.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused z/b/a split for the strided-QKV GDN prefill path. + +One TileLang kernel writes contiguous ``z`` (gate), ``b`` and ``a`` out of the +``mixed_qkvz`` / ``mixed_ba`` projections, replacing the three separate copies +the strided path would otherwise issue per GDN layer: + - the strided-``z`` ``CopyLastContiguous`` in the output projection, + - ``b.contiguous()`` and ``a.contiguous()``. + +``mixed_qkv`` is intentionally left as a strided view of ``mixed_qkvz`` (the +strided path passes it to conv/MATE without materializing it), so unlike +SGLang's ``fused_qkvzba`` this kernel does NOT copy the large qkv block. + +Adapted from SGLang's ``fused_qkvzba_split_reshape_cat_contiguous`` (row kernel) +with the qkv writes dropped for the MUSA strided path. +""" + +import functools + +import tilelang +import tilelang.language as T +import torch + +from vllm_musa.jit_kernel.tilelang.utils import ( + MUSA_COMMON_PASS_CONFIGS, + MUSA_COMPILE_FLAGS, + tilelang_dtype, +) + +_PASS_CONFIGS = dict(MUSA_COMMON_PASS_CONFIGS) +for _key, _value in ( + ("TL_DISABLE_SAFE_COPY_PREDICATION", True), + ("TL_DISABLE_SAFE_ROBUST_COPY_PREDICATION", True), + ("TL_CONFIG_INDEX_BITWIDTH", 32), +): + if hasattr(tilelang.PassConfigKey, _key): + _PASS_CONFIGS[getattr(tilelang.PassConfigKey, _key)] = _value + +__all__ = ["fused_zba"] + + +@functools.lru_cache(maxsize=32) +@tilelang.jit( + target="musa", + pass_configs=_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _fused_zba_kernel( + num_heads_qk: int, + num_heads_v: int, + head_qk: int, + head_v: int, + input_dtype: str, + ba_dtype: str, + block_elems: int, +): + m = T.dynamic("m") + total_q = num_heads_qk * head_qk + total_v = num_heads_v * head_v + qkv_dim = total_q * 2 + total_v + total_qkvz = qkv_dim + total_v + total_ba = num_heads_v * 2 + v_blocks = T.ceildiv(total_v, block_elems) + ba_blocks = T.ceildiv(total_ba, block_elems) + + @T.prim_func + def zba_contiguous( + z: T.Tensor((m, num_heads_v, head_v), input_dtype), + b: T.Tensor((m, num_heads_v), ba_dtype), + a: T.Tensor((m, num_heads_v), ba_dtype), + mixed_qkvz: T.Tensor((m, total_qkvz), input_dtype), + mixed_ba: T.Tensor((m, total_ba), ba_dtype), + ): + with T.Kernel(m, threads=block_elems) as row: + for block in T.serial(v_blocks): + for i in T.Parallel(block_elems): + offset = block * block_elems + i + if offset < total_v: + z[row, offset // head_v, offset % head_v] = mixed_qkvz[ + row, qkv_dim + offset + ] + + for block in T.serial(ba_blocks): + for i in T.Parallel(block_elems): + offset = block * block_elems + i + if offset < num_heads_v: + b[row, offset] = mixed_ba[row, offset] + ba_offset = offset + num_heads_v + if ba_offset < total_ba: + a[row, offset] = mixed_ba[row, ba_offset] + + return zba_contiguous + + +_fused_zba_kernel.mode = "lazy" + + +def _block_elems(total_v: int) -> int: + be = 1 + while be < total_v and be < 1024: + be <<= 1 + return max(be, 32) + + +def fused_zba( + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + num_heads_qk: int, + num_heads_v: int, + head_qk: int, + head_v: int, +) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": + """Split contiguous z/b/a out of mixed_qkvz / mixed_ba in one kernel. + + Layout (Qwen3.5): mixed_qkvz = [q, k, v, z]; mixed_ba = [b, a] (chunk(2)). + Returns contiguous z [m, num_heads_v, head_v], b/a [m, num_heads_v]. + """ + m = mixed_qkvz.shape[0] + z = torch.empty( + (m, num_heads_v, head_v), dtype=mixed_qkvz.dtype, device=mixed_qkvz.device + ) + b = torch.empty((m, num_heads_v), dtype=mixed_ba.dtype, device=mixed_ba.device) + a = torch.empty_like(b) + if m == 0: + return z, b, a + kernel = _fused_zba_kernel( + num_heads_qk, + num_heads_v, + head_qk, + head_v, + tilelang_dtype(mixed_qkvz.dtype), + tilelang_dtype(mixed_ba.dtype), + _block_elems(num_heads_v * head_v), + ) + kernel(z, b, a, mixed_qkvz.contiguous(), mixed_ba.contiguous()) + return z, b, a diff --git a/vllm_musa/jit_kernel/tilelang/layernorm_gated.py b/vllm_musa/jit_kernel/tilelang/layernorm_gated.py new file mode 100644 index 000000000000..ec9713254552 --- /dev/null +++ b/vllm_musa/jit_kernel/tilelang/layernorm_gated.py @@ -0,0 +1,570 @@ +"""MUSA TileLang RMSNorm + gate kernels.""" + +import functools + +import tilelang +import tilelang.language as T +import torch + +from vllm_musa.jit_kernel.tilelang.utils import ( + MUSA_COMMON_PASS_CONFIGS, + MUSA_COMPILE_FLAGS, + tilelang_dtype, +) + + +def register_custom_op(fn=None, **_kw): + def _wrap(f): + return f + + return _wrap if fn is None else _wrap(fn) + + +__all__ = [ + "RMSNorm", + "_layer_norm_fwd", + "layernorm_fn", + "rms_norm_gated", +] + +_LOG2E = 1.4426950408889634 + +_RMSNORM_PASS_CONFIGS = dict(MUSA_COMMON_PASS_CONFIGS) +for _key, _value in ( + ("TL_DISABLE_DATA_RACE_CHECK", True), + ("TL_DISABLE_SAFE_COPY_PREDICATION", True), + ("TL_DISABLE_SAFE_ROBUST_COPY_PREDICATION", True), + ("TL_CONFIG_INDEX_BITWIDTH", 32), +): + if hasattr(tilelang.PassConfigKey, _key): + _RMSNORM_PASS_CONFIGS[getattr(tilelang.PassConfigKey, _key)] = _value + + +@functools.lru_cache(maxsize=32) +@tilelang.jit( + target="musa", + pass_configs=_RMSNORM_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _rms_norm_gated_kernel( + dtype: str, + hidden_size: int, + rows_per_block: int, + lanes_per_row: int, +): + m = T.dynamic("m") + x_stride_row = T.dynamic("x_stride_row") + y_stride_row = T.dynamic("y_stride_row") + z_stride_row = T.dynamic("z_stride_row") + vec_size = tilelang.cdiv(hidden_size, lanes_per_row) + num_shuffles = lanes_per_row.bit_length() - 1 + threads = rows_per_block * lanes_per_row + + @T.prim_func + def musa_rms_norm_gated( + x: T.StridedTensor((m, hidden_size), (x_stride_row, 1), dtype), + y: T.StridedTensor((m, hidden_size), (y_stride_row, 1), dtype), + w: T.Tensor((hidden_size,), dtype), + z: T.StridedTensor((m, hidden_size), (z_stride_row, 1), dtype), + rstd: T.Tensor((m,), "float32"), + eps: T.float32, + ): + with T.Kernel(T.ceildiv(m, rows_per_block), threads=threads) as (bid,): + tid = T.get_thread_binding() + lane = tid % lanes_per_row + row_in_block = tid // lanes_per_row + row = bid * rows_per_block + row_in_block + offset = T.alloc_var("int32") + x_local = T.alloc_local((vec_size,), "float32") + z_local = T.alloc_local((vec_size,), "float32") + sum_local = T.alloc_local((1,), "float32") + inv_rms = T.alloc_local((1,), "float32") + + if row < m: + sum_local[0] = 0.0 + for i in T.vectorized(vec_size): + offset = lane * vec_size + i + if offset < hidden_size: + x_local[i] = T.cast(x[row, offset], "float32") + sum_local[0] += x_local[i] * x_local[i] + + for i in T.unroll(num_shuffles): + sum_local[0] += T.shfl_xor(sum_local[0], (lanes_per_row // 2) >> i) + + inv_rms[0] = T.rsqrt(sum_local[0] / hidden_size + eps) + if lane == 0: + rstd[row] = inv_rms[0] + + for i in T.vectorized(vec_size): + offset = lane * vec_size + i + if offset < hidden_size: + z_local[i] = T.cast(z[row, offset], "float32") + y[row, offset] = T.cast( + x_local[i] + * inv_rms[0] + * T.cast(w[offset], "float32") + * z_local[i] + / (1.0 + T.exp2(-z_local[i] * _LOG2E)), + dtype, + ) + + return musa_rms_norm_gated + + +_rms_norm_gated_kernel.mode = "lazy" + + +@functools.lru_cache(maxsize=32) +@tilelang.jit( + target="musa", + pass_configs=_RMSNORM_PASS_CONFIGS, + compile_flags=MUSA_COMPILE_FLAGS, +) +def _rms_norm_gated_kernel_cta( + dtype: str, + hidden_size: int, + threads: int, +): + m = T.dynamic("m") + x_stride_row = T.dynamic("x_stride_row") + y_stride_row = T.dynamic("y_stride_row") + z_stride_row = T.dynamic("z_stride_row") + warps_per_cta = threads // 32 + + @T.prim_func + def musa_rms_norm_gated_cta( + x: T.StridedTensor((m, hidden_size), (x_stride_row, 1), dtype), + y: T.StridedTensor((m, hidden_size), (y_stride_row, 1), dtype), + w: T.Tensor((hidden_size,), dtype), + z: T.StridedTensor((m, hidden_size), (z_stride_row, 1), dtype), + rstd: T.Tensor((m,), "float32"), + eps: T.float32, + ): + with T.Kernel(m, threads=threads) as (row,): + tid = T.get_thread_binding() + lane = tid % 32 + warp = tid // 32 + sum_local = T.alloc_var("float32") + inv_rms = T.alloc_var("float32") + warp_sum = T.alloc_shared((warps_per_cta,), "float32") + inv_rms_shared = T.alloc_shared((1,), "float32") + + sum_local = 0.0 + for offset_base in T.serial(0, hidden_size, threads): + offset = offset_base + tid + if offset < hidden_size: + x_val = T.cast(x[row, offset], "float32") + sum_local += x_val * x_val + + sum_local = T.warp_reduce_sum(sum_local) + if lane == 0: + warp_sum[warp] = sum_local + T.sync_threads() + + sum_local = T.if_then_else(tid < warps_per_cta, warp_sum[tid], 0.0) + if warp == 0: + sum_local = T.warp_reduce_sum(sum_local) + if lane == 0: + inv_rms = T.rsqrt(sum_local / hidden_size + eps) + inv_rms_shared[0] = inv_rms + rstd[row] = inv_rms + T.sync_threads() + + for offset_base in T.serial(0, hidden_size, threads): + offset = offset_base + tid + if offset < hidden_size: + x_val = T.cast(x[row, offset], "float32") + z_val = T.cast(z[row, offset], "float32") + y[row, offset] = T.cast( + x_val + * inv_rms_shared[0] + * T.cast(w[offset], "float32") + * z_val + / (1.0 + T.exp2(-z_val * _LOG2E)), + dtype, + ) + + return musa_rms_norm_gated_cta + + +_rms_norm_gated_kernel_cta.mode = "lazy" + + +def _launch_rms_norm_gated( + x: torch.Tensor, + out: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor, + rstd: torch.Tensor, + eps: float, + rows_per_block: int = 8, + lanes_per_row: int = 16, + cta_threads: int | None = None, +) -> None: + if cta_threads is not None: + kernel = _rms_norm_gated_kernel_cta( + tilelang_dtype(x.dtype), + x.shape[-1], + cta_threads, + ) + kernel(x, out, weight, z, rstd, float(eps)) + return + + kernel = _rms_norm_gated_kernel( + tilelang_dtype(x.dtype), + x.shape[-1], + rows_per_block, + lanes_per_row, + ) + kernel(x, out, weight, z, rstd, float(eps)) + + +@register_custom_op( + op_name="musa_rms_norm_gated", + mutates_args=["out", "rstd"], +) +def _rms_norm_gated_custom( + x: torch.Tensor, + out: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor, + rstd: torch.Tensor, + eps: float, + rows_per_block: int = 8, + lanes_per_row: int = 16, + cta_threads: int | None = None, +) -> None: + # Keep the TileLang executable launch opaque to Dynamo. + _launch_rms_norm_gated( + x, + out, + weight, + z, + rstd, + eps, + rows_per_block, + lanes_per_row, + cta_threads, + ) + + +def _rms_norm_gated_impl( + x: torch.Tensor, + weight: torch.Tensor, + z: torch.Tensor, + eps: float, + out: torch.Tensor | None = None, + rows_per_block: int | None = None, + lanes_per_row: int | None = None, + cta_threads: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if x.dim() != 2 or z.dim() != 2: + raise RuntimeError("rms_norm_gated expects x/z with shape [M, N].") + if x.shape != z.shape: + raise RuntimeError("rms_norm_gated expects x and z to have the same shape.") + if weight.shape != (x.shape[-1],): + raise RuntimeError("rms_norm_gated weight shape mismatch.") + if x.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("rms_norm_gated expects fp16 or bf16 input.") + if z.dtype != x.dtype or weight.dtype != x.dtype: + raise RuntimeError("rms_norm_gated expects x/z/weight to have the same dtype.") + if x.stride(-1) != 1 or z.stride(-1) != 1: + raise RuntimeError("rms_norm_gated requires x/z contiguous in the last dim.") + if out is None: + out = torch.empty_like(x) + if out.shape != x.shape or out.stride(-1) != 1: + raise RuntimeError( + "rms_norm_gated output must match x and be last-dim contiguous." + ) + + hidden_size = x.shape[-1] + if cta_threads is None: + if 2048 <= hidden_size < 3072: + cta_threads = 512 if hidden_size == 2048 and x.shape[0] < 128 else 256 + elif hidden_size >= 3072: + cta_threads = 512 + if cta_threads is not None: + if cta_threads not in (128, 256, 512): + raise RuntimeError("cta_threads must be one of 128, 256, 512.") + if cta_threads > 1024: + raise RuntimeError("invalid cta_threads.") + + if lanes_per_row is None: + lanes_per_row = 32 if hidden_size % 32 == 0 else 16 + if rows_per_block is None: + rows_per_block = 4 if hidden_size >= 8192 and lanes_per_row == 32 else 8 + if hidden_size % lanes_per_row != 0: + raise RuntimeError("hidden_size must be divisible by lanes_per_row.") + if lanes_per_row not in (4, 8, 16, 32): + raise RuntimeError("lanes_per_row must be one of 4, 8, 16, 32.") + if rows_per_block < 1 or rows_per_block * lanes_per_row > 1024: + raise RuntimeError("invalid rows_per_block/lanes_per_row combination.") + + rstd = torch.empty((x.shape[0],), dtype=torch.float32, device=x.device) + _rms_norm_gated_custom( + x, + out, + weight, + z, + rstd, + eps, + rows_per_block, + lanes_per_row, + cta_threads, + ) + return out, rstd + + +def _pytorch_gated_norm( + x, weight, bias, eps, z, group_size, norm_before_gate, is_rms_norm, activation +): + """MUSA: pure-PyTorch general-case gated (layer/rms) norm fallback, used when + the fused TileLang fast path does not apply. Self-contained (no external dep).""" + import torch.nn.functional as _F + + dtype = x.dtype + xf = x.float() + + def _gate(t): + if z is None: + return t + zf = z.float() + if activation in ("swish", "silu"): + g = _F.silu(zf) + elif activation == "sigmoid": + g = torch.sigmoid(zf) + else: + g = zf + return t * g + + if not norm_before_gate: + xf = _gate(xf) + + def _norm(v): + if group_size is None or group_size == v.shape[-1]: + if is_rms_norm: + return v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + eps) + m = v.mean(-1, keepdim=True) + return (v - m) * torch.rsqrt((v - m).pow(2).mean(-1, keepdim=True) + eps) + M, N = v.shape + vg = v.reshape(M, N // group_size, group_size) + if is_rms_norm: + vn = vg * torch.rsqrt(vg.pow(2).mean(-1, keepdim=True) + eps) + else: + m = vg.mean(-1, keepdim=True) + vn = (vg - m) * torch.rsqrt((vg - m).pow(2).mean(-1, keepdim=True) + eps) + return vn.reshape(M, N) + + out = _norm(xf) * weight.float() + if bias is not None: + out = out + bias.float() + if norm_before_gate: + out = _gate(out) + return out.to(dtype) + + +def _layer_norm_fwd( + x, + weight, + bias, + eps, + z=None, + out=None, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, + activation: str = "swish", +): + if ( + z is not None + and bias is None + and group_size in (None, x.shape[-1]) + and norm_before_gate + and is_rms_norm + and activation in ("swish", "silu") + and x.dtype in (torch.float16, torch.bfloat16) + and x.dim() == 2 + and z.dim() == 2 + and x.stride(-1) == 1 + and z.stride(-1) == 1 + and weight.stride(-1) == 1 + ): + out, rstd = _rms_norm_gated_impl(x, weight, z, eps, out=out) + return out, None, rstd + + y = _pytorch_gated_norm( + x, weight, bias, eps, z, group_size, norm_before_gate, is_rms_norm, activation + ) + if out is not None: + out.copy_(y) + y = out + return y, None, None + + +def _fallback_layernorm_fn( + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, + activation: str = "swish", +): + return _pytorch_gated_norm( + x, weight, bias, eps, z, group_size, norm_before_gate, is_rms_norm, activation + ) + + +def layernorm_fn( + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, + activation: str = "swish", +): + x_shape_og = x.shape + x_2d = x.reshape(-1, x.shape[-1]) + z_2d = None + if z is not None: + if z.shape != x_shape_og: + raise RuntimeError("z shape must match x shape") + z_2d = z.reshape(-1, z.shape[-1]) + + if ( + z_2d is not None + and bias is None + and group_size in (None, x_2d.shape[-1]) + and norm_before_gate + and is_rms_norm + and activation in ("swish", "silu") + and x_2d.dtype in (torch.float16, torch.bfloat16) + and x_2d.stride(-1) == 1 + and z_2d.stride(-1) == 1 + and weight.stride(-1) == 1 + ): + y, _, _ = _layer_norm_fwd( + x_2d, + weight, + bias, + eps, + z=z_2d, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + activation=activation, + ) + return y.reshape(x_shape_og) + + return _fallback_layernorm_fn( + x, + weight, + bias, + z=z, + eps=eps, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + activation=activation, + ) + + +def rms_norm_gated( + *, + x, + weight, + bias, + z=None, + eps=1e-6, + group_size=None, + norm_before_gate=True, + is_rms_norm=False, + activation: str = "swish", +): + x_shape_og = x.shape + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + if z.shape != x_shape_og: + raise RuntimeError("z shape must match x shape") + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y, _, _ = _layer_norm_fwd( + x, + weight, + bias, + eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + activation=activation, + ) + return y.reshape(x_shape_og) + + +class RMSNorm(torch.nn.Module): + def __init__( + self, + hidden_size, + eps=1e-5, + group_size=None, + norm_before_gate=True, + device=None, + dtype=None, + activation: str = "swish", + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.activation = activation + self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x, z=None): + if ( + z is not None + and self.bias is None + and self.group_size in (None, x.shape[-1]) + and self.norm_before_gate + and self.activation in ("swish", "silu") + and x.dtype in (torch.float16, torch.bfloat16) + ): + x_shape_og = x.shape + x_2d = x.reshape(-1, x.shape[-1]) + z_2d = z.reshape(-1, z.shape[-1]) + if x_2d.stride(-1) == 1 and z_2d.stride(-1) == 1: + y, _ = _rms_norm_gated_impl( + x_2d, + self.weight, + z_2d, + self.eps, + ) + return y.reshape(x_shape_og) + return layernorm_fn( + x, + self.weight, + self.bias, + z=z, + eps=self.eps, + group_size=self.group_size, + norm_before_gate=self.norm_before_gate, + is_rms_norm=True, + activation=self.activation, + ) diff --git a/vllm_musa/model_executor/layers/fused_moe/fused_moe.py b/vllm_musa/model_executor/layers/fused_moe/fused_moe.py index e7179dc9503a..796dc9a99661 100644 --- a/vllm_musa/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm_musa/model_executor/layers/fused_moe/fused_moe.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect import json import os import time @@ -49,10 +50,8 @@ def disable_inplace() -> bool: ) _MOE_SHAPE_INVENTORY_RECORDS = 0 _MOE_SHAPE_INVENTORY_WARNED = False -_DEEPGEMM_PREFILL_ENV = "VLLM_MUSA_DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL" -_DEEPGEMM_PREFILL_MIN_TOKENS_ENV = ( - "VLLM_MUSA_DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL_MIN_TOKENS" -) +_DEEPGEMM_PREFILL_ENV = "VLLM_MUSA_MOE_DEEPGEMM_PREFILL" +_DEEPGEMM_PREFILL_MIN_TOKENS_ENV = "VLLM_MUSA_MOE_DEEPGEMM_PREFILL_MIN_TOKENS" _DEEPGEMM_PREFILL_WARNED = False @@ -61,6 +60,11 @@ def _env_flag_enabled(name: str) -> bool: return value.lower() not in {"", "0", "false", "no", "off"} +def _env_flag_disabled(name: str) -> bool: + """True only when explicitly set to a falsy value; unset means enabled.""" + return os.environ.get(name, "").strip().lower() in {"0", "false", "no", "off"} + + def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, default)) @@ -213,7 +217,7 @@ def _maybe_record_deepseek_v4_moe_shape_inventory( _MOE_SHAPE_INVENTORY_WARNED = True -def _can_use_deepseek_v4_moe_deepgemm_prefill( +def _can_use_moe_deepgemm_prefill( *, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -236,10 +240,10 @@ def _can_use_deepseek_v4_moe_deepgemm_prefill( w1_bias: torch.Tensor | None, w2_bias: torch.Tensor | None, ) -> bool: - if not _env_flag_enabled(_DEEPGEMM_PREFILL_ENV): + if _env_flag_disabled(_DEEPGEMM_PREFILL_ENV): return False - min_tokens = _env_int(_DEEPGEMM_PREFILL_MIN_TOKENS_ENV, 4096) + min_tokens = _env_int(_DEEPGEMM_PREFILL_MIN_TOKENS_ENV, 2500) if hidden_states.size(0) < min_tokens: return False @@ -266,7 +270,7 @@ def _can_use_deepseek_v4_moe_deepgemm_prefill( if block_shape != [128, 128]: return False - if topk_ids.size(1) != 6: + if topk_ids.size(1) > 16: return False if hidden_states.dtype != torch.bfloat16: @@ -288,8 +292,9 @@ def _can_use_deepseek_v4_moe_deepgemm_prefill( E, N, K = w1.shape return ( - E == 256 - and N == 512 + E in (256, 257) + and N % 256 == 0 + and K % 128 == 0 and K == hidden_states.size(1) and w2.shape == (E, K, N // 2) and w1_scale.shape == (E, N // 128, K // 128) @@ -312,20 +317,148 @@ def _silu_mul_per_token_group_fp8_quant_musa_large( device=input_tensor.device, dtype=torch.float32, ) + from vllm_musa.jit_kernel.csrc.quant import per_token_group_quant_8bit + fp8_min, fp8_max = get_fp8_min_max() - torch.ops._C_musa_ops.silu_and_mul_per_token_group_fp8_quant( + per_token_group_quant_8bit( input_tensor, output, output_s, - group_size, - 1e-10, - fp8_min, - fp8_max, + group_size=group_size, + eps=1e-10, + min_8bit=fp8_min, + max_8bit=fp8_max, + fuse_silu_and_mul=True, ) return output, output_s -def _deepseek_v4_moe_deepgemm_prefill_impl( +def _moe_deepgemm_prefill_impl( + *, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + inplace: bool, +) -> torch.Tensor: + # MUSA: SGLang-style fused-glue DeepGEMM MoE prefill. One fused preprocess + # kernel does quant + permute + contiguous group-layout + m_indices/src2dst + # (replacing moe_kernel_quantize_input + deepgemm_moe_permute), the two + # grouped FP8 GEMMs are unchanged, and one post_reorder kernel does the + # weighted un-permute (replacing deepgemm_unpermute_and_reduce). Falls back + # to the unfused path when the fused preprocess cannot handle the shape. + E = w1.shape[0] + try: + from vllm_musa.jit_kernel.post_reorder import post_reorder_triton_kernel + from vllm_musa.jit_kernel.tilelang.deep_gemm_contig_preprocess import ( + can_use_fp8_tilelang, + deep_gemm_contig_preprocess_fp8_tilelang, + ) + + use_fused = can_use_fp8_tilelang(hidden_states, topk_ids, E, True) + except Exception: + use_fused = False + + if not use_fused: + return _moe_deepgemm_prefill_impl_unfused( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + w1_scale=w1_scale, + w2_scale=w2_scale, + inplace=inplace, + ) + + from vllm.utils.deep_gemm import ( + get_mk_alignment_for_contiguous_layout, + m_grouped_fp8_gemm_nt_contiguous, + mk_alignment_scope, + ) + + logger.info_once("Using the MUSA fused-glue grouped DeepGEMM MoE prefill path.") + + _, N, K = w1.shape + M, top_k = topk_ids.shape + block_m = int(get_mk_alignment_for_contiguous_layout()[0]) + dev = hidden_states.device + + src2dst_numel = M * top_k + all_tokens = ( + (src2dst_numel + E * (block_m - 1) + block_m - 1) // block_m + ) * block_m + + qhidden_perm = torch.empty((all_tokens, K), device=dev, dtype=torch.float8_e4m3fn) + qhidden_scale_perm = torch.empty( + (all_tokens, K // 128), device=dev, dtype=torch.float32 + ) + m_indices = torch.empty(all_tokens, device=dev, dtype=torch.int32) + src2dst = torch.empty(src2dst_numel, device=dev, dtype=torch.int32) + topk_ids_for_combine = torch.empty_like(topk_ids) + counts = torch.empty(E, device=dev, dtype=torch.int32) + cursor = torch.empty(E, device=dev, dtype=torch.int32) + + deep_gemm_contig_preprocess_fp8_tilelang( + hidden_states, + topk_ids, + qhidden_perm, + qhidden_scale_perm, + m_indices, + src2dst, + topk_ids_for_combine, + counts, + cursor, + E, + block_m, + ) + + with mk_alignment_scope(block_m): + mm1_out = torch.empty((all_tokens, N), device=dev, dtype=hidden_states.dtype) + m_grouped_fp8_gemm_nt_contiguous( + (qhidden_perm, qhidden_scale_perm.contiguous()), + (w1, w1_scale.contiguous()), + mm1_out, + m_indices, + ) + + a2q = torch.empty((all_tokens, N // 2), device=dev, dtype=torch.float8_e4m3fn) + a2q, a2q_scale = _silu_mul_per_token_group_fp8_quant_musa_large( + mm1_out.view(-1, N), + a2q, + group_size=128, + ) + + mm2_out = torch.empty((all_tokens, K), device=dev, dtype=hidden_states.dtype) + m_grouped_fp8_gemm_nt_contiguous( + (a2q, a2q_scale.contiguous()), + (w2, w2_scale.contiguous()), + mm2_out, + m_indices, + ) + + if inplace and not disable_inplace(): + output = hidden_states + else: + output = torch.empty_like(hidden_states) + + post_reorder_triton_kernel[(M,)]( + mm2_out, + output, + src2dst, + topk_ids, + topk_weights, + top_k, + K, + BLOCK_SIZE=1024, + ) + return output + + +def _moe_deepgemm_prefill_impl_unfused( *, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -345,11 +478,7 @@ def _deepseek_v4_moe_deepgemm_prefill_impl( mk_alignment_scope, ) - logger.info_once( - "Using DeepSeek-V4 MUSA grouped DeepGEMM MoE prefill path " - "(set %s=0 to disable).", - _DEEPGEMM_PREFILL_ENV, - ) + logger.info_once("Using the MUSA grouped DeepGEMM MoE prefill path.") qhidden, a1_scale = moe_kernel_quantize_input( A=hidden_states, @@ -427,7 +556,7 @@ def _deepseek_v4_moe_deepgemm_prefill_impl( return output -def _maybe_deepseek_v4_moe_deepgemm_prefill( +def _maybe_moe_deepgemm_prefill( *, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -454,7 +583,7 @@ def _maybe_deepseek_v4_moe_deepgemm_prefill( ) -> torch.Tensor | None: global _DEEPGEMM_PREFILL_WARNED - if not _can_use_deepseek_v4_moe_deepgemm_prefill( + if not _can_use_moe_deepgemm_prefill( hidden_states=hidden_states, w1=w1, w2=w2, @@ -481,7 +610,7 @@ def _maybe_deepseek_v4_moe_deepgemm_prefill( try: assert w1_scale is not None assert w2_scale is not None - return _deepseek_v4_moe_deepgemm_prefill_impl( + return _moe_deepgemm_prefill_impl( hidden_states=hidden_states, w1=w1, w2=w2, @@ -494,8 +623,8 @@ def _maybe_deepseek_v4_moe_deepgemm_prefill( except Exception as exc: if not _DEEPGEMM_PREFILL_WARNED: logger.warning( - "DeepSeek-V4 grouped DeepGEMM MoE prefill path failed; " - "falling back to native GEMV path: %s", + "Grouped DeepGEMM MoE prefill path failed; " + "falling back to the upstream Triton path: %s", exc, ) _DEEPGEMM_PREFILL_WARNED = True @@ -695,7 +824,7 @@ def fused_experts_impl( global_num_experts=global_num_experts, ) - deepgemm_prefill_output = _maybe_deepseek_v4_moe_deepgemm_prefill( + deepgemm_prefill_output = _maybe_moe_deepgemm_prefill( hidden_states=hidden_states, w1=w1, w2=w2, @@ -837,6 +966,43 @@ def fused_experts_impl( ) +def _try_deepgemm_prefill_dispatch(args, kwargs) -> torch.Tensor | None: + try: + bound = inspect.signature(fused_experts_impl).bind(*args, **kwargs) + except TypeError: + return None + bound.apply_defaults() + a = bound.arguments + if not a.get("use_fp8_w8a8"): + return None + w1 = a["w1"] + w2 = a["w2"] + return _maybe_moe_deepgemm_prefill( + hidden_states=a["hidden_states"], + w1=w1, + w2=w2, + topk_weights=a["topk_weights"], + topk_ids=a["topk_ids"], + inplace=a.get("inplace", False), + activation=a.get("activation", "silu"), + apply_router_weight_on_input=a.get("apply_router_weight_on_input", False), + use_fp8_w8a8=a.get("use_fp8_w8a8", False), + use_int8_w8a8=a.get("use_int8_w8a8", False), + use_int8_w8a16=a.get("use_int8_w8a16", False), + use_int4_w4a16=a.get("use_int4_w4a16", False), + ocp_mx_scheme=a.get("ocp_mx_scheme"), + per_channel_quant=a.get("per_channel_quant", False), + expert_map=a.get("expert_map"), + w1_scale=_maybe_expand_fp8_moe_per_tensor_scale(a.get("w1_scale"), w1), + w2_scale=_maybe_expand_fp8_moe_per_tensor_scale(a.get("w2_scale"), w2), + a1_scale=a.get("a1_scale"), + a2_scale=a.get("a2_scale"), + block_shape=a.get("block_shape"), + w1_bias=a.get("w1_bias"), + w2_bias=a.get("w2_bias"), + ) + + def _musa_fused_experts_impl_dispatch(*args, **kwargs) -> torch.Tensor: # The native GEMV MoE path (fused_experts_impl) is opt-in: it can hurt some # prefill workloads, but DeepSeek-V4-Flash-Base TP8 graph+MTP decode depends @@ -844,6 +1010,11 @@ def _musa_fused_experts_impl_dispatch(*args, **kwargs) -> torch.Tensor: # other models; platform.py opts DeepSeek-V4 in per serving process. if os.environ.get("VLLM_MUSA_DEEPSEEK_V4_FUSED_MOE_GEMV", "0") == "1": return fused_experts_impl(*args, **kwargs) + # Grouped DeepGEMM for large-M FP8 block-128 MoE prefill; the shape gate declines + # decode and every non-matching shape, which then fall through to upstream Triton. + dg_out = _try_deepgemm_prefill_dispatch(args, kwargs) + if dg_out is not None: + return dg_out return _upstream_fused_moe._musa_original_fused_experts_impl(*args, **kwargs) diff --git a/vllm_musa/model_executor/layers/layernorm.py b/vllm_musa/model_executor/layers/layernorm.py index bffde8eb9164..3d23c0d66273 100644 --- a/vllm_musa/model_executor/layers/layernorm.py +++ b/vllm_musa/model_executor/layers/layernorm.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm, RMSNormGated from vllm_musa.jit_kernel.csrc import norm as musa_jit_norm from vllm_musa.utils.environ import envs @@ -69,13 +69,66 @@ def forward_oot( return self.forward_native(x, residual) if residual is not None: - # Residual calls flow through vLLM IR so its maybe_inplace contract - # validates donation. The MUSA provider owns only capability and - # measured profitability. + weight = self.weight.data + if ( + _can_use_musa_jit_rmsnorm(x, weight) + and residual.shape == x.shape + and residual.dtype == x.dtype + and residual.is_contiguous() + ): + # MUSA: fused residual-add + Gemma RMSNorm in one JIT kernel; + # weight is the raw zero-centered param, gemma=True applies +1. + return musa_jit_norm.fused_add_rmsnorm( + x, residual, weight, self.variance_epsilon, gemma=True + ) return self.forward_native(x, residual) weight = self.weight.data + if x.dim() != 2: + # MUSA: q/k GemmaRMSNorm gets a 3D [tokens, heads, head_dim] tensor; + # the fused rmsnorm gate needs 2D, so reshape to 2D, run the fused + # kernel, and reshape back (mirrors SGLang RMSNorm.forward_cuda). + x2 = x.contiguous().reshape(-1, x.shape[-1]) + if x2.shape[0] > 0 and _can_use_musa_jit_rmsnorm(x2, weight): + y = musa_jit_norm.gemma_rmsnorm(x2, weight, self.variance_epsilon) + return y.reshape(x.shape) + return self.forward_native(x, residual) + if _can_use_musa_jit_rmsnorm(x, weight): return musa_jit_norm.gemma_rmsnorm(x, weight, self.variance_epsilon) return self.forward_native(x, residual) + + +@RMSNormGated.register_oot +class MusaRMSNormGated(RMSNormGated): + def forward_oot(self, x, z=None): + if ( + z is not None + and self.bias is None + and (self.group_size is None or self.group_size == x.shape[-1]) + and self.norm_before_gate + and self.activation in ("silu", "swish") + and x.dtype in (torch.float16, torch.bfloat16) + and x.is_contiguous() + and z.is_contiguous() + ): + try: + from vllm_musa.jit_kernel.tilelang.layernorm_gated import ( + rms_norm_gated, + ) + + return rms_norm_gated( + x=x, + weight=self.weight.data, + bias=None, + z=z, + eps=self.eps, + group_size=None, + norm_before_gate=True, + is_rms_norm=True, + activation="silu", + ) + except Exception: + pass + return self.forward_native(x, z) diff --git a/vllm_musa/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm_musa/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index cc5aa0acfb5e..8aba0ece6250 100644 --- a/vllm_musa/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm_musa/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -6,6 +6,7 @@ import inspect import torch +from mate.gdn_decode import gated_delta_rule_decode from mate.gdn_prefill import chunk_gated_delta_rule from vllm.logger import init_logger from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -14,6 +15,9 @@ logger = init_logger(__name__) +_MATE_GDN_PREFILL_HAS_OUTPUT = ( + "output" in inspect.signature(chunk_gated_delta_rule).parameters +) _MATE_GDN_PREFILL_HAS_IS_LOG_SPACE = ( "is_log_space" in inspect.signature(chunk_gated_delta_rule).parameters ) @@ -58,6 +62,63 @@ def _forward_core( attn_metadata, ) + def forward_cuda( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + ) -> None: + # MUSA: Qwen3.5 GDN forward with a single fused z/b/a split kernel + # (contiguous z/b/a in one launch) replacing the strided-z output-proj + # copy + b/a contiguous copies. mixed_qkv stays a strided view (conv/MATE + # accept it) so the large qkv block is never materialized. Qwen3-Next's + # interleaved layout and the replicated-ba TP path keep the upstream flow. + if self.gqa_interleaved_layout: + return super().forward_cuda(hidden_states, output) + + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + _encode_layer_name, + ) + + num_tokens = hidden_states.size(0) + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) + + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + mixed_qkv = mixed_qkvz[:, :qkv_size] + + if getattr(self, "disable_tp_for_ba_proj", False) and self.tp_size > 1: + z = mixed_qkvz[:, qkv_size:].reshape(num_tokens, -1, self.head_v_dim) + b, a = self.split_ba(ba) + b = b.contiguous() + a = a.contiguous() + else: + from vllm_musa.jit_kernel.tilelang.gdn_fused_proj import fused_zba + + z, b, a = fused_zba( + mixed_qkvz, + ba, + self.num_k_heads // self.tp_size, + self.num_v_heads // self.tp_size, + self.head_k_dim, + self.head_v_dim, + ) + + core_attn_out = torch.zeros( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + torch.ops.vllm.qwen_gdn_attention_core( + mixed_qkv, + b, + a, + core_attn_out, + layer_name=_encode_layer_name(self.prefix), + ) + + self._output_projection(core_attn_out, z, output, num_tokens) + def _get_gdn_attention_metadata(self, mixed_qkv: torch.Tensor): from vllm.forward_context import get_forward_context from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata @@ -131,12 +192,78 @@ def _try_mate_decode( validate_data=False, ) - query, key, value = self.rearrange_mixed_qkv(mixed_qkv) + # MUSA: mate's decode kernel reads q/k/v from strided views; split the + # packed qkv without materializing contiguous copies. + _qk_dim = self.key_dim // self.tp_size + _v_dim = self.value_dim // self.tp_size + _q, _k, _v = torch.split(mixed_qkv, [_qk_dim, _qk_dim, _v_dim], dim=-1) + query = _q.unflatten(-1, (-1, self.head_k_dim)).unsqueeze(1) + key = _k.unflatten(-1, (-1, self.head_k_dim)).unsqueeze(1) + value = _v.unflatten(-1, (-1, self.head_v_dim)).unsqueeze(1) + + # MUSA: the mamba pool is page-aligned (it must not be laid out contiguously, + # or its bytes would alias the attention KV in the shared hybrid tensor), so + # ssm_state is not contiguous. Gather the active per-sequence states into a + # small contiguous buffer for mate's fp32 VK decode (identity mapping), then + # scatter the updated states back. This keeps the fast mate kernel without a + # whole-pool contiguity copy. + if ssm_state.dtype == torch.float32: + try: + import os as _os + + _musa_sep = _os.environ.get("VLLM_MUSA_MAMBA_SEPARATE_POOL", "1") == "1" + if _musa_sep and ssm_state.is_contiguous(): + # MUSA: separate contiguous mamba pool -> mate decodes in + # place (block b at b*state_numel); no gather/scatter copy. + output, _ = gated_delta_rule_decode( + q=query, + k=key, + v=value, + state=ssm_state, + state_layout="VK", + state_indices=state_indices, + scale=self.head_k_dim**-0.5, + A_log=self.A_log.detach().float(), + a=a.view(num_decode_tokens, 1, -1), + dt_bias=self.dt_bias.detach().float(), + b=b.view(num_decode_tokens, 1, -1), + disable_state_update=False, + use_qk_l2norm=True, + ) + _log_once( + "info", + "MUSA GDN mate in-place decode active (separate pool)", + ) + else: + active_state = ssm_state[state_indices] + output, updated_state = gated_delta_rule_decode( + q=query, + k=key, + v=value, + state=active_state, + state_layout="VK", + scale=self.head_k_dim**-0.5, + A_log=self.A_log.detach().float(), + a=a.view(num_decode_tokens, 1, -1), + dt_bias=self.dt_bias.detach().float(), + b=b.view(num_decode_tokens, 1, -1), + disable_state_update=False, + use_qk_l2norm=True, + ) + ssm_state[state_indices] = updated_state + core_attn_out[:num_decode_tokens] = output.view( + num_decode_tokens, + self.num_v_heads // self.tp_size, + self.head_v_dim, + ) + return True + except Exception as e: + _log_once( + "warning", + "MATE GDN decode failed; using recurrent fallback: %s", + e, + ) - # MUSA: the mate fp32-VK decode kernel is unavailable in this toolchain, - # so pure decode runs the fused recurrent update directly on the paged - # state pool (leaner than the upstream general prefill/decode path). Fall - # back to the general decode path if this kernel ever fails at runtime. try: core_attn_out_non_spec, _ = fused_sigmoid_gating_delta_rule_update( A_log=self.A_log, @@ -172,32 +299,35 @@ def _try_mate_prefill( non_spec_state_indices_tensor: torch.Tensor, non_spec_query_start_loc: torch.Tensor, has_initial_state: torch.Tensor | None, + out: torch.Tensor | None = None, ): - from vllm.model_executor.layers.fla.ops import fused_post_conv_prep - try: - q, k, v, g, beta = fused_post_conv_prep( - conv_output=mixed_qkv_non_spec, - a=a_non_spec, - b=b_non_spec, - A_log=self.A_log, - dt_bias=self.dt_bias, - num_k_heads=self.num_k_heads // self.tp_size, - head_k_dim=self.head_k_dim, - head_v_dim=self.head_v_dim, - apply_l2norm=False, - output_g_exp=not _MATE_GDN_PREFILL_HAS_IS_LOG_SPACE, - ) + # MUSA: feed mate no-copy strided q/k/v views (skip the qkv + # contiguous copy fused_post_conv_prep does) + one fused kernel for + # the log-space gating (g, beta). mate l2-norms q/k internally. + hk = self.num_k_heads // self.tp_size + hv = self.num_v_heads // self.tp_size + _kd = hk * self.head_k_dim + _vd = hv * self.head_v_dim + _qf, _kf, _vf = torch.split(mixed_qkv_non_spec, [_kd, _kd, _vd], dim=-1) + q = _qf.view(_qf.shape[0], hk, self.head_k_dim) + k = _kf.view(_kf.shape[0], hk, self.head_k_dim) + v = _vf.view(_vf.shape[0], hv, self.head_v_dim) + from vllm_musa.jit_kernel.fused_gdn_gating import fused_gdn_gating + + g, beta = fused_gdn_gating(self.A_log, a_non_spec, b_non_spec, self.dt_bias) if not _MATE_GDN_PREFILL_HAS_IS_LOG_SPACE: - # Clamp exp-space g away from exact zero to avoid MATE NaNs on - # long real-model prefills with very negative log-space gates. - g = g.clamp_min(1e-30) + # Exp-space fallback for pre-is_log_space mate; clamp away from + # exact zero to avoid MATE NaNs on long prefills. + g = torch.exp(g).clamp_min(1e-30) state_indices = non_spec_state_indices_tensor.to(torch.int64) initial_state = ssm_state[state_indices].to(torch.float32) if has_initial_state is not None: initial_state[~has_initial_state, ...] = 0 - cu_seqlens = non_spec_query_start_loc.to(torch.int64) + # mate compiles the GDN kernel against this dtype; int32 indexing is + # cheaper and the offsets are bounded by the per-forward token cap. + cu_seqlens = non_spec_query_start_loc.to(torch.int32) mate_kwargs = { "q": q, @@ -214,6 +344,11 @@ def _try_mate_prefill( if _MATE_GDN_PREFILL_HAS_IS_LOG_SPACE: mate_kwargs["is_log_space"] = True + # Let mate write straight into the caller's buffer instead of + # producing a temporary that is then copied into it. + if out is not None and _MATE_GDN_PREFILL_HAS_OUTPUT: + mate_kwargs["output"] = out + output, final_state = chunk_gated_delta_rule(**mate_kwargs) ssm_state.index_copy_(0, state_indices, final_state.to(ssm_state.dtype)) return output.unsqueeze(0) @@ -244,6 +379,13 @@ def _forward_core_mate_prefill( causal_conv1d_update, ) + try: + from vllm_musa.jit_kernel.tilelang.causal_conv1d import ( + musa_tilelang_causal_conv1d_fn as causal_conv1d_fn, + ) + except Exception: + pass # MUSA: fall back to Triton causal_conv1d_fn on import failure + has_initial_state = attn_metadata.has_initial_state spec_query_start_loc = attn_metadata.spec_query_start_loc non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc @@ -342,6 +484,13 @@ def _forward_core_mate_prefill( core_attn_out_spec = None assert non_spec_state_indices_tensor is not None + # With no spec-decode branch the mate output IS the destination, so hand the + # buffer down and let the kernel write it. + mate_out = ( + core_attn_out[:num_actual_tokens] + if (core_attn_out_spec is None and _MATE_GDN_PREFILL_HAS_OUTPUT) + else None + ) core_attn_out_non_spec = self._try_mate_prefill( mixed_qkv_non_spec, a_non_spec, @@ -350,7 +499,9 @@ def _forward_core_mate_prefill( non_spec_state_indices_tensor, non_spec_query_start_loc, has_initial_state, + out=mate_out, ) + wrote_in_place = mate_out is not None and core_attn_out_non_spec is not None if core_attn_out_non_spec is None: if has_initial_state is not None: zero_mask = ~has_initial_state @@ -384,7 +535,7 @@ def _forward_core_mate_prefill( merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec) merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) core_attn_out[:num_actual_tokens] = merged_out.squeeze(0) - else: + elif not wrote_in_place: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) diff --git a/vllm_musa/model_executor/layers/quantization/fp8.py b/vllm_musa/model_executor/layers/quantization/fp8.py index cda126a9f092..1b9ddcbc40ea 100644 --- a/vllm_musa/model_executor/layers/quantization/fp8.py +++ b/vllm_musa/model_executor/layers/quantization/fp8.py @@ -20,6 +20,7 @@ def _zero_fp8_weight(weight: torch.Tensor) -> None: _ORIGINAL_FP8_MOE_MAYBE_ROUNDUP_SIZES = vllm_fp8.Fp8MoEMethod.maybe_roundup_sizes _ORIGINAL_FP8_MOE_CREATE_WEIGHTS = vllm_fp8.Fp8MoEMethod.create_weights +_ORIGINAL_FP8_MOE_PROCESS_WEIGHTS = vllm_fp8.Fp8MoEMethod.process_weights_after_loading def maybe_roundup_sizes( @@ -109,6 +110,106 @@ def create_weights( ) +def _release_cached_blocks() -> None: + empty_cache = getattr(current_platform, "empty_cache", None) + if callable(empty_cache): + empty_cache() + elif hasattr(torch, "musa"): + torch.musa.empty_cache() + + +def _fold_shared_expert_weights(self, layer: FusedMoE) -> None: + """Append the model's shared expert to the routed FP8 expert stack. + + The MoE block stashes its shared MLP and gate on the routed-experts layer + when the two carry the same weight format. The shared gate_up/down + projections and their block scales are concatenated as one extra expert, so + a single grouped GEMM covers routed + shared work instead of running the + shared expert as separate dense GEMMs on every MoE layer. Runs once after + load, before any graph capture. + """ + mlp = getattr(layer, "_musa_shared_mlp", None) + if mlp is None or getattr(layer, "_musa_shared_folded", False): + return + + scale_name = self.weight_scale_name + parts = ( + ("w13_weight", f"w13_{scale_name}", mlp.gate_up_proj), + ("w2_weight", f"w2_{scale_name}", mlp.down_proj), + ) + folded: dict[str, torch.Tensor] = {} + for weight_attr, scale_attr, proj in parts: + routed_w = getattr(layer, weight_attr).data + routed_s = getattr(layer, scale_attr).data + shared_w = proj.weight.data + shared_s = getattr(proj, scale_name).data + if shared_w.shape != routed_w.shape[1:] or shared_w.dtype != routed_w.dtype: + raise RuntimeError( + f"MUSA shared-expert fold: {weight_attr} expects " + f"{tuple(routed_w.shape[1:])}/{routed_w.dtype}, shared expert has " + f"{tuple(shared_w.shape)}/{shared_w.dtype}" + ) + if shared_s.shape != routed_s.shape[1:] or shared_s.dtype != routed_s.dtype: + raise RuntimeError( + f"MUSA shared-expert fold: {scale_attr} expects " + f"{tuple(routed_s.shape[1:])}/{routed_s.dtype}, shared expert has " + f"{tuple(shared_s.shape)}/{shared_s.dtype}" + ) + folded[weight_attr] = torch.cat( + [routed_w, shared_w.unsqueeze(0)], dim=0 + ).contiguous() + folded[scale_attr] = torch.cat( + [routed_s, shared_s.unsqueeze(0)], dim=0 + ).contiguous() + + num_routed = int(layer.w13_weight.shape[0]) + # Rebuild the quant config / kernel so they bind the folded tensors. + self._setup_kernel( + layer, + folded["w13_weight"], + folded["w2_weight"], + folded[f"w13_{scale_name}"], + folded[f"w2_{scale_name}"], + layer.w13_input_scale, + layer.w2_input_scale, + ) + del folded + # Each expert stack is concatenated while the unfolded one is still alive, so the + # freed blocks leave expert-sized holes in the caching allocator. Left there, they + # are reserved-but-unusable and the KV cache later fails to find contiguous space. + _release_cached_blocks() + layer._musa_shared_expert_id = num_routed + layer._musa_shared_folded = True + logger.info_once( + "MUSA shared-expert fold active: shared expert folded into the routed " + "FP8 grouped GEMM as slot %d (%d routed + 1 shared).", + num_routed, + num_routed, + ) + + +def process_weights_after_loading(self, layer: FusedMoE) -> None: + _ORIGINAL_FP8_MOE_PROCESS_WEIGHTS(self, layer) + _fold_shared_expert_weights(self, layer) + + +def _extend_routing_with_shared_expert( + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + from vllm_musa.jit_kernel.extend_topk_shared import extend_topk_with_shared + + shared_logits, _ = layer._musa_shared_gate(x) + return extend_topk_with_shared( + topk_weights, + topk_ids, + shared_logits.view(-1), + layer._musa_shared_expert_id, + ) + + def apply( self, layer: FusedMoE, @@ -121,6 +222,21 @@ def apply( ep_size = getattr(layer, "ep_size", None) if ep_size is None: ep_size = layer.moe_config.ep_size + + # A folded shared expert rides the routed grouped GEMM as one extra + # always-selected slot, so extend the routing by a single column here. + folded_shared = getattr(layer, "_musa_shared_folded", False) + if folded_shared: + assert ( + shared_experts is None and shared_experts_input is None + ), "folded shared expert expects shared_experts=None" + topk_weights, topk_ids = _extend_routing_with_shared_expert( + layer, x, topk_weights, topk_ids + ) + global_num_experts = ( + layer._musa_shared_expert_id + 1 if folded_shared else layer.global_num_experts + ) + if ep_size != None and ep_size <= 1: # the legacy fused_experts() path only computes routed # experts. For the no-overlap path used by DeepSeek-V2/V3 on MUSA, the @@ -139,7 +255,7 @@ def apply( topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, - global_num_experts=layer.global_num_experts, + global_num_experts=global_num_experts, apply_router_weight_on_input=layer.apply_router_weight_on_input, expert_map=layer.expert_map, quant_config=self.moe_quant_config, @@ -157,7 +273,7 @@ def apply( topk_weights, topk_ids, activation=layer.activation, - global_num_experts=layer.global_num_experts, + global_num_experts=global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, shared_experts=shared_experts, @@ -167,4 +283,5 @@ def apply( vllm_fp8.Fp8MoEMethod.maybe_roundup_sizes = maybe_roundup_sizes vllm_fp8.Fp8MoEMethod.create_weights = create_weights +vllm_fp8.Fp8MoEMethod.process_weights_after_loading = process_weights_after_loading vllm_fp8.Fp8MoEMethod.apply = apply diff --git a/vllm_musa/patches/series/0060-MUSA-csrc-drop-unboxable-stable-ops-guard-non-core-a.patch b/vllm_musa/patches/series/0060-MUSA-csrc-drop-unboxable-stable-ops-guard-non-core-a.patch index 1dc4f9c89881..59ec7e5c89d2 100644 --- a/vllm_musa/patches/series/0060-MUSA-csrc-drop-unboxable-stable-ops-guard-non-core-a.patch +++ b/vllm_musa/patches/series/0060-MUSA-csrc-drop-unboxable-stable-ops-guard-non-core-a.patch @@ -323,8 +323,8 @@ index 3133243c4..1ca216c0d 100644 ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); // w8a8 per-token-group quant (MUSA-0/Option-3 consolidation): the upstream // stable kernels replace the native csrc/musa non-fused re-impl. The MUSA-only -@@ -982,4 +1054,5 @@ STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1, ops) { - ops.impl("per_token_group_quant_int8", TORCH_BOX(&per_token_group_quant_int8)); +@@ -984,4 +1056,5 @@ STABLE_TORCH_LIBRARY_IMPL(_C, PrivateUse1, ops) { + ops.impl("apply_repetition_penalties_", TORCH_BOX(&apply_repetition_penalties_)); } #endif + diff --git a/vllm_musa/patches/series/0076-MUSA-model-fold-the-Qwen3.5-shared-expert-into-fused.patch b/vllm_musa/patches/series/0076-MUSA-model-fold-the-Qwen3.5-shared-expert-into-fused.patch index 1a2c3f1673f9..9f567beac8fc 100644 --- a/vllm_musa/patches/series/0076-MUSA-model-fold-the-Qwen3.5-shared-expert-into-fused.patch +++ b/vllm_musa/patches/series/0076-MUSA-model-fold-the-Qwen3.5-shared-expert-into-fused.patch @@ -5,31 +5,93 @@ Subject: [PATCH] MUSA(model): fold the Qwen3.5 shared expert into fused_moe On MUSA the Qwen3.5/Qwen3.6 shared expert runs as a separate serial MLP because the shared-expert overlap stream is unavailable, adding per-layer kernels to the -decode path. When unquantized and the shared and routed intermediate sizes -match, pass no shared expert to FusedMoE and stash the shared MLP and its gate -on the routed-experts layer; the MUSA unquantized MoE method then folds the -shared expert into the routed grouped GEMM as one extra always-selected slot, +decode path. When the shared and routed intermediate sizes match and the shared +MLP carries a weight format the routed stack can absorb (unquantized, or FP8 +with [128, 128] block scales), pass no shared expert to FusedMoE and stash the +shared MLP and its gate on the routed-experts layer; the MUSA MoE method then +folds it into the routed grouped GEMM as one extra always-selected slot, producing numerically equivalent output from a single kernel. + +The fold appends one expert to the routed stack, so it stays off under expert +parallelism and EPLB, whose expert_map is built for the unfolded count. +Set VLLM_MUSA_MOE_SHARED_EXPERT_FUSION=0 to disable. --- - vllm/model_executor/models/qwen3_next.py | 17 ++++++++++++++++- - 1 file changed, 16 insertions(+), 1 deletion(-) + vllm/model_executor/models/qwen3_next.py | 60 +++--- + 1 file changed, 59 insertions(+), 1 deletions(-) diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py -index acce2a767..b41bdccce 100644 +index acce2a767..de163b0e6 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py -@@ -159,8 +159,18 @@ class Qwen3NextSparseMoeBlock(nn.Module): +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + """Inference-only Qwen3Next model.""" + ++import os + from collections.abc import Iterable + from itertools import islice + +@@ -88,6 +89,42 @@ + KVCache = tuple[torch.Tensor, torch.Tensor] + + ++MUSA_SHARED_EXPERT_FUSION_ENV = "VLLM_MUSA_MOE_SHARED_EXPERT_FUSION" ++ ++ ++def _musa_linear_is_block_fp8(linear: nn.Module) -> bool: ++ method = getattr(linear, "quant_method", None) ++ if method is None or not getattr(method, "block_quant", False): ++ return False ++ block_size = getattr(method, "weight_block_size", None) ++ if block_size is None: ++ block_size = getattr( ++ getattr(method, "quant_config", None), "weight_block_size", None ++ ) ++ return list(block_size or []) == [128, 128] ++ ++ ++def _musa_shared_expert_foldable(shared_expert: nn.Module | None) -> bool: ++ """Whether the shared expert's weights can join the routed expert stack. ++ ++ Foldable when the shared MLP's projections carry the same weight format as ++ the routed experts: unquantized, or FP8 with the [128, 128] block scales the ++ grouped GEMM expects. ++ """ ++ if shared_expert is None: ++ return False ++ gate_up = getattr(shared_expert, "gate_up_proj", None) ++ down = getattr(shared_expert, "down_proj", None) ++ if gate_up is None or down is None: ++ return False ++ unquantized = [ ++ type(getattr(proj, "quant_method", None)).__name__ for proj in (gate_up, down) ++ ] == ["UnquantizedLinearMethod"] * 2 ++ return unquantized or ( ++ _musa_linear_is_block_fp8(gate_up) and _musa_linear_is_block_fp8(down) ++ ) ++ ++ + class Qwen3NextSparseMoeBlock(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() +@@ -159,8 +196,24 @@ prefix=f"{prefix}.shared_expert", ) + # MUSA: fold the shared expert into the routed grouped GEMM as one extra + # always-selected slot instead of running it as a separate serial MLP. -+ # Only when unquantized and the shared/routed intermediate sizes match. ++ # Needs matching shared/routed intermediate sizes and a shared expert whose ++ # weight format the routed expert stack can absorb. The fold appends one ++ # expert to the routed stack, which an expert-parallel expert_map built for ++ # the unfolded count would not know about. + self._musa_shared_fold = ( -+ self.shared_expert is not None -+ and quant_config is None ++ getattr(torch.version, "musa", None) is not None ++ and os.environ.get(MUSA_SHARED_EXPERT_FUSION_ENV, "1").strip().lower() ++ not in ("0", "false", "off", "no") ++ and self.ep_size <= 1 ++ and not self.enable_eplb + and config.shared_expert_intermediate_size == config.moe_intermediate_size -+ and getattr(torch.version, "musa", None) is not None ++ and _musa_shared_expert_foldable(self.shared_expert) + ) + self.experts = FusedMoE( @@ -38,7 +100,7 @@ index acce2a767..b41bdccce 100644 gate=self.gate, num_experts=self.n_routed_experts, top_k=config.num_experts_per_tok, -@@ -178,6 +188,11 @@ class Qwen3NextSparseMoeBlock(nn.Module): +@@ -178,6 +231,11 @@ else None, ) diff --git a/vllm_musa/patches/series/0081-MUSA-vllm.v1.attention.backends.gdn_attn.patch b/vllm_musa/patches/series/0081-MUSA-vllm.v1.attention.backends.gdn_attn.patch new file mode 100644 index 000000000000..c4f37777632a --- /dev/null +++ b/vllm_musa/patches/series/0081-MUSA-vllm.v1.attention.backends.gdn_attn.patch @@ -0,0 +1,33 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Sun, 13 Jul 2026 00:00:00 +0800 +Subject: [PATCH] MUSA: vllm.v1.attention.backends.gdn_attn conv1d cpu seq_lens + +Plumb a CPU copy of non_spec_query_start_loc through GDNAttentionMetadata so +the tilelang causal_conv1d prefill wrapper derives seq_lens on CPU without a +per-call device->host sync. +--- + vllm/v1/attention/backends/gdn_attn.py | 4 +++- + 1 file changed, 3 insertions(+), 1 deletion(-) + +diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py +--- a/vllm/v1/attention/backends/gdn_attn.py ++++ b/vllm/v1/attention/backends/gdn_attn.py +@@ -54,6 +54,9 @@ + non_spec_query_start_loc: torch.Tensor | None = ( + None # shape: [batch - num_spec_decodes + 1,] + ) ++ # MUSA: CPU copy of non_spec_query_start_loc so the tilelang conv1d wrapper ++ # derives seq_lens without a per-call device->host sync. ++ non_spec_query_start_loc_cpu: torch.Tensor | None = None + + spec_state_indices_tensor: torch.Tensor | None = None # shape: [batch, num_spec] + non_spec_state_indices_tensor: torch.Tensor | None = ( +@@ -503,6 +506,7 @@ + prefill_has_initial_state=prefill_has_initial_state, + spec_query_start_loc=spec_query_start_loc, + non_spec_query_start_loc=non_spec_query_start_loc, ++ non_spec_query_start_loc_cpu=non_spec_query_start_loc_cpu, + spec_state_indices_tensor=spec_state_indices_tensor, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, diff --git a/vllm_musa/patches/series/0082-MUSA-vllm.v1.kv_cache_interface.patch b/vllm_musa/patches/series/0082-MUSA-vllm.v1.kv_cache_interface.patch new file mode 100644 index 000000000000..38931fbfa3e9 --- /dev/null +++ b/vllm_musa/patches/series/0082-MUSA-vllm.v1.kv_cache_interface.patch @@ -0,0 +1,27 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Subject: [PATCH] MUSA: vllm.v1.kv_cache_interface separate mamba pool fields + +--- +--- a/vllm/v1/kv_cache_interface.py ++++ b/vllm/v1/kv_cache_interface.py +@@ -859,6 +859,7 @@ + shared_by: list[str] # layer names that share the same KV cache tensor + offset: int = 0 # byte offset of this layer within a contiguous block + block_stride: int = 0 # total bytes per block in a packed layout (0 = not packed) ++ musa_mamba_pool: bool = False # MUSA: backs the separate mamba state pool + + + @dataclass +@@ -895,6 +896,11 @@ + see `_get_kv_cache_config_uniform_page_size` for more details. + """ + ++ # MUSA: block count of the separate mamba state pool when ++ # VLLM_MUSA_MAMBA_SEPARATE_POOL is on. None = mamba shares the ++ # attention pool (upstream behavior). ++ musa_mamba_num_blocks: int | None = None ++ + @property + def has_mamba_layers(self) -> bool: + return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups) diff --git a/vllm_musa/patches/series/0083-MUSA-vllm.v1.core.kv_cache_utils.patch b/vllm_musa/patches/series/0083-MUSA-vllm.v1.core.kv_cache_utils.patch new file mode 100644 index 000000000000..5d94ae489df6 --- /dev/null +++ b/vllm_musa/patches/series/0083-MUSA-vllm.v1.core.kv_cache_utils.patch @@ -0,0 +1,116 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Mon, 14 Jul 2026 12:00:00 +0800 +Subject: [PATCH] MUSA: separate mamba state BlockPool + +Give mamba state its own contiguous BlockPool, sized from max_num_seqs so it +never binds, and keep it out of the attention pool's proportional shrink. The +mamba groups then own a small contiguous block-id space and mate's GDN decode +updates the recurrent state in place instead of gathering and scattering it. + +On by default; set VLLM_MUSA_MAMBA_SEPARATE_POOL=0 to fall back to the shared pool. +--- + vllm/v1/core/kv_cache_utils.py | 76 +++--- + 1 file changed, 76 insertions(+), 0 deletions(-) + +diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py +index 3ea89d5af..c8425baa5 100644 +--- a/vllm/v1/core/kv_cache_utils.py ++++ b/vllm/v1/core/kv_cache_utils.py +@@ -1357,6 +1357,63 @@ + _get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed + + ++def musa_mamba_separate_pool_enabled() -> bool: ++ """MUSA: give mamba state its own contiguous BlockPool so mate's GDN ++ decode runs in place (no gather/scatter). On by default.""" ++ return os.environ.get("VLLM_MUSA_MAMBA_SEPARATE_POOL", "1") == "1" ++ ++ ++def _musa_has_mamba_and_attention(kv_cache_groups) -> bool: ++ has_mamba = any(isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups) ++ has_attn = any(not isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups) ++ return has_mamba and has_attn ++ ++ ++def _musa_separated_mamba_attn_tensors(vllm_config, kv_cache_groups, available_memory): ++ """Split KV memory into a small contiguous mamba pool and the attention pool. ++ ++ Mamba groups get one dedicated tensor per layer (segregated-contiguous, ++ indexed by their own block-id space); the mamba pool is sized so it never ++ binds (attention admission caps running requests at <= max_num_seqs, each ++ needing 1 block per mamba group). ++ Returns (kv_cache_tensors, attn_num_blocks, mamba_num_blocks). ++ """ ++ mamba_groups = [g for g in kv_cache_groups if isinstance(g.kv_cache_spec, MambaSpec)] ++ attn_groups = [ ++ g for g in kv_cache_groups if not isinstance(g.kv_cache_spec, MambaSpec) ++ ] ++ max_num_seqs = vllm_config.scheduler_config.max_num_seqs ++ # Peak live blocks = num_mamba_groups * max_num_seqs (1 block/seq/group); ++ # plus the null block and slack. ++ mamba_num_blocks = len(mamba_groups) * (max_num_seqs + 8) + 1 ++ mamba_page = get_uniform_page_size([g.kv_cache_spec for g in mamba_groups]) ++ mamba_layers = [ln for g in mamba_groups for ln in g.layer_names] ++ mamba_bytes = mamba_page * mamba_num_blocks * len(mamba_layers) ++ if mamba_bytes >= available_memory: ++ raise ValueError("MUSA separate mamba pool needs more memory than available") ++ remaining = available_memory - mamba_bytes ++ attn_group_size = max(len(g.layer_names) for g in attn_groups) ++ attn_page = get_uniform_page_size([g.kv_cache_spec for g in attn_groups]) ++ attn_num_blocks = get_num_blocks( ++ vllm_config, attn_group_size, remaining, attn_page ++ ) ++ kv_cache_tensors = [] ++ for i in range(attn_group_size): ++ shared_by = [g.layer_names[i] for g in attn_groups if i < len(g.layer_names)] ++ kv_cache_tensors.append( ++ KVCacheTensor(size=attn_page * attn_num_blocks, shared_by=shared_by) ++ ) ++ for ln in mamba_layers: ++ kv_cache_tensors.append( ++ KVCacheTensor( ++ size=mamba_page * mamba_num_blocks, ++ shared_by=[ln], ++ musa_mamba_pool=True, ++ ) ++ ) ++ return kv_cache_tensors, attn_num_blocks, mamba_num_blocks ++ ++ + def get_kv_cache_config_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +@@ -1407,6 +1464,21 @@ + num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( + vllm_config, kv_cache_groups, available_memory + ) ++ elif musa_mamba_separate_pool_enabled() and _musa_has_mamba_and_attention( ++ kv_cache_groups ++ ): ++ # MUSA: mamba state on its own contiguous pool (own block-id space). ++ kv_cache_tensors, num_blocks, musa_mamba_num_blocks = ( ++ _musa_separated_mamba_attn_tensors( ++ vllm_config, kv_cache_groups, available_memory ++ ) ++ ) ++ return KVCacheConfig( ++ num_blocks=num_blocks, ++ kv_cache_tensors=kv_cache_tensors, ++ kv_cache_groups=kv_cache_groups, ++ musa_mamba_num_blocks=musa_mamba_num_blocks, ++ ) + else: + # General case: + # We will have group_size memory pools, each is shared by one layer from +@@ -2182,6 +2254,10 @@ + + # Shrink tensor size proportionally + for tensor in kv_cache_config.kv_cache_tensors: ++ if getattr(tensor, "musa_mamba_pool", False): ++ # MUSA: mamba pool size is rank-invariant (from max_num_seqs), ++ # not shrunk with the attention pool's num_blocks. ++ continue + assert tensor.size % num_blocks_old == 0 + tensor.size = tensor.size // num_blocks_old * min_num_blocks + diff --git a/vllm_musa/patches/series/0084-MUSA-vllm.v1.core.kv_cache_coordinator.patch b/vllm_musa/patches/series/0084-MUSA-vllm.v1.core.kv_cache_coordinator.patch new file mode 100644 index 000000000000..3f6c52bcc581 --- /dev/null +++ b/vllm_musa/patches/series/0084-MUSA-vllm.v1.core.kv_cache_coordinator.patch @@ -0,0 +1,73 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Subject: [PATCH] MUSA: vllm.v1.core.kv_cache_coordinator separate mamba pool + +--- +--- a/vllm/v1/core/kv_cache_coordinator.py ++++ b/vllm/v1/core/kv_cache_coordinator.py +@@ -12,9 +12,11 @@ + BlockHashList, + BlockHashListWithBlockSize, + KVCacheBlock, ++ musa_mamba_separate_pool_enabled, + ) + from vllm.v1.core.single_type_kv_cache_manager import ( + CrossAttentionManager, ++ MambaManager, + SingleTypeKVCacheManager, + get_manager_for_kv_cache_spec, + ) +@@ -96,6 +98,22 @@ + metrics_collector=metrics_collector, + ) + ++ # MUSA: dedicated non-binding BlockPool for mamba state, so mamba ++ # groups own a small contiguous block-id space and the mate GDN ++ # decode runs in place. Sized in the config to never bind. ++ self.musa_mamba_block_pool = None ++ _musa_mamba_nb = getattr( ++ kv_cache_config, "musa_mamba_num_blocks", None ++ ) ++ if musa_mamba_separate_pool_enabled() and _musa_mamba_nb: ++ self.musa_mamba_block_pool = BlockPool( ++ num_gpu_blocks=_musa_mamba_nb, ++ enable_caching=False, ++ hash_block_size=hash_block_size, ++ enable_kv_cache_events=False, ++ metrics_collector=None, ++ ) ++ + # KV cache group indices that get the EAGLE last-block drop. + self.eagle_group_ids: set[int] = { + i for i, g in enumerate(kv_cache_config.kv_cache_groups) if g.is_eagle_group +@@ -109,7 +127,16 @@ + kv_cache_spec=kv_cache_group.kv_cache_spec, + max_num_batched_tokens=max_num_batched_tokens, + max_model_len=max_model_len, +- block_pool=self.block_pool, ++ block_pool=( ++ self.musa_mamba_block_pool ++ if ( ++ self.musa_mamba_block_pool is not None ++ and isinstance( ++ kv_cache_group.kv_cache_spec, MambaSpec ++ ) ++ ) ++ else self.block_pool ++ ), + enable_caching=enable_caching, + kv_cache_group_id=i, + dcp_world_size=dcp_world_size, +@@ -162,6 +189,12 @@ + """ + num_blocks_to_allocate = 0 + for i, manager in enumerate(self.single_type_managers): ++ if self.musa_mamba_block_pool is not None and isinstance( ++ manager, MambaManager ++ ): ++ # MUSA: mamba is on its own non-binding pool; do not count ++ # its blocks against the attention pool admission check. ++ continue + if isinstance(manager, CrossAttentionManager): + # For cross-attention, we issue a single static allocation + # of blocks based on the number of encoder input tokens. diff --git a/vllm_musa/patches/series/0085-MUSA-vllm.model_executor.layers.quantization.utils.fp8_utils.patch b/vllm_musa/patches/series/0085-MUSA-vllm.model_executor.layers.quantization.utils.fp8_utils.patch new file mode 100644 index 000000000000..3eb27ccdfeb5 --- /dev/null +++ b/vllm_musa/patches/series/0085-MUSA-vllm.model_executor.layers.quantization.utils.fp8_utils.patch @@ -0,0 +1,58 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Mon, 14 Jul 2026 12:00:00 +0800 +Subject: [PATCH] MUSA: vllm.model_executor.layers.quantization.utils.fp8_utils vec group quant + +Route group-128 bf16/fp16 activation quant to the MUSA register-resident vec +kernel, which keeps the group in registers instead of round-tripping the +stable per-token-group path. Falls through to the existing kernel for any +shape, dtype, alignment or scale layout it does not cover. +--- + vllm/model_executor/layers/quantization/utils/fp8_utils.py | 27 +++--- + 1 file changed, 27 insertions(+), 0 deletions(-) + +diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py +index 74db17f55..799e51a58 100644 +--- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py ++++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py +@@ -563,6 +563,19 @@ + tl.store(y_s_ptr, y_s) + + ++_MUSA_VEC_QUANT = None ++ ++ ++def _musa_vec_quant(): ++ """Cached handle to the MUSA register-resident group-quant kernel.""" ++ global _MUSA_VEC_QUANT ++ if _MUSA_VEC_QUANT is None: ++ import vllm_musa._C # noqa: F401 ++ ++ _MUSA_VEC_QUANT = torch.ops._C_musa_ops.per_token_group_quant_8bit_vec ++ return _MUSA_VEC_QUANT ++ ++ + def per_token_group_quant_fp8( + x: torch.Tensor, + group_size: int, +@@ -631,6 +644,20 @@ + + # prefer CUDA/XPU kernel if available + # TODO(bnell): this causes some fp8 moe test to fail. ++ # MUSA: register-resident vec kernel for group-128 bf16/fp16 activation quant. ++ if ( ++ current_platform.is_musa() ++ and x.is_contiguous() ++ and group_size == 128 ++ and not use_ue8m0 ++ and not column_major_scales ++ and not tma_aligned_scales ++ and x.dtype in (torch.bfloat16, torch.float16) ++ and x.data_ptr() % 16 == 0 ++ ): ++ _musa_vec_quant()(x, x_q, x_s, group_size, eps, fp8_min, fp8_max) ++ return x_q, x_s ++ + if ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() diff --git a/vllm_musa/patches/series/0087-MUSA-fuse-QK-RMSNorm-and-MRoPE-for-interleaved-MRoPE.patch b/vllm_musa/patches/series/0087-MUSA-fuse-QK-RMSNorm-and-MRoPE-for-interleaved-MRoPE.patch new file mode 100644 index 000000000000..805e31a10cdd --- /dev/null +++ b/vllm_musa/patches/series/0087-MUSA-fuse-QK-RMSNorm-and-MRoPE-for-interleaved-MRoPE.patch @@ -0,0 +1,118 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: musa +Date: Mon, 14 Jul 2026 12:00:00 +0800 +Subject: [PATCH] MUSA(model): fuse QK-RMSNorm and MRoPE for interleaved MRoPE + +Fuse QK-RMSNorm and MRoPE into one MUSA kernel for interleaved-MRoPE gated +attention. The text-only NeoX fusion rejects these models and the IR provider +keeps routed MoE on native, so they otherwise run the eager rotary path and +issue a long tail of elementwise kernels. + +The IR operation stays as the fallback for every shape this kernel declines. +Set VLLM_MUSA_FUSED_QK_MROPE=0 to disable. +--- + vllm/model_executor/models/qwen3_next.py | 77 +++--- + 1 file changed, 77 insertions(+), 0 deletions(-) + +diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py +index f7907f3c4..d73dfa9e4 100644 +--- a/vllm/model_executor/models/qwen3_next.py ++++ b/vllm/model_executor/models/qwen3_next.py +@@ -267,6 +267,30 @@ + return final_hidden_states.view(orig_shape) + + ++MUSA_FUSED_QK_MROPE_ENV = "VLLM_MUSA_FUSED_QK_MROPE" ++ ++_MUSA_MROPE_COS_SIN_CACHE: dict[tuple[int, torch.dtype], torch.Tensor] = {} ++ ++ ++def _musa_mrope_cos_sin_cache( ++ cos_sin_cache: torch.Tensor, dtype: torch.dtype ++) -> torch.Tensor: ++ """The fused kernel reads the rotary cache in the activation dtype. ++ ++ get_rope() hands every layer the same rotary module, so convert once and ++ share: this cache is num_positions x rotary_dim and reaches hundreds of MB ++ on long-context models. ++ """ ++ if cos_sin_cache.dtype == dtype: ++ return cos_sin_cache ++ key = (cos_sin_cache.data_ptr(), dtype) ++ converted = _MUSA_MROPE_COS_SIN_CACHE.get(key) ++ if converted is None: ++ converted = cos_sin_cache.to(dtype) ++ _MUSA_MROPE_COS_SIN_CACHE[key] = converted ++ return converted ++ ++ + class Qwen3NextAttention(nn.Module): + def __init__( + self, +@@ -379,6 +403,26 @@ + self.mrope_interleaved = bool( + getattr(self.rotary_emb, "mrope_interleaved", False) + ) ++ # MUSA: interleaved-MRoPE models reject the text-only NeoX fusion above ++ # and the IR provider keeps routed MoE on native, leaving the eager ++ # rotary path. Fuse QK-RMSNorm + MRoPE for them in one kernel. ++ self._musa_fused_qk_mrope = ( ++ self.attn_output_gate ++ and current_platform.is_musa() ++ and os.environ.get(MUSA_FUSED_QK_MROPE_ENV, "1").strip().lower() ++ not in ("0", "false", "off", "no") ++ and len(self.mrope_section) == 3 ++ and self.mrope_interleaved ++ ) ++ self._musa_mrope_cos_sin_cache = None ++ if self._musa_fused_qk_mrope: ++ dtype = model_config.dtype if model_config else torch.bfloat16 ++ self._musa_fused_qk_mrope = dtype == torch.bfloat16 ++ if self._musa_fused_qk_mrope: ++ self._musa_mrope_cos_sin_cache = _musa_mrope_cos_sin_cache( ++ self.rotary_emb.cos_sin_cache, dtype ++ ) ++ self._musa_mrope_section = tuple(int(x) for x in self.mrope_section) + + def _project_qkv_gate( + self, +@@ -414,6 +458,39 @@ + ) + return q, k, v, gate + ++ if self._musa_fused_qk_mrope and positions.ndim == 2: ++ from vllm_musa.jit_kernel.csrc.norm import fused_qk_rmsnorm_mrope ++ ++ q_gate, k, v = qkv.split( ++ [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 ++ ) ++ orig_shape = q_gate.shape[:-1] ++ q_gate = q_gate.view(*orig_shape, self.num_heads, -1) ++ q, gate = torch.chunk(q_gate, 2, dim=-1) ++ gate = gate.reshape(*orig_shape, -1) ++ t, h, w = self._musa_mrope_section ++ q, k = fused_qk_rmsnorm_mrope( ++ q.contiguous(), ++ k.view(-1, self.num_kv_heads, self.head_dim), ++ self.q_norm.weight, ++ self.k_norm.weight, ++ positions, ++ self._musa_mrope_cos_sin_cache, ++ is_neox=bool(self.rotary_emb.is_neox_style), ++ mrope_section_t=t, ++ mrope_section_h=h, ++ mrope_section_w=w, ++ is_interleaved=True, ++ eps=self.q_norm.variance_epsilon, ++ gemma=True, ++ ) ++ return ( ++ q.view(-1, self.num_heads * self.head_dim), ++ k.view(-1, self.num_kv_heads * self.head_dim), ++ v, ++ gate, ++ ) ++ + if self.use_gated_qkv_rms_norm_rope: + # Gate and V remain views of the packed projection. In particular, + # keeping gate outside the functional IR op lets Inductor fuse its diff --git a/vllm_musa/platform.py b/vllm_musa/platform.py index afb204e5fbe6..93c61fdbb68c 100644 --- a/vllm_musa/platform.py +++ b/vllm_musa/platform.py @@ -72,7 +72,6 @@ "VLLM_MUSA_DEEPSEEK_V4_INDEXER_TOPK_PREFILL_MATERIALIZED_DIRECT" ) _DEEPSEEK_V4_QNORM_KV_FUSED_ENV = "VLLM_MUSA_DEEPSEEK_V4_QNORM_ROPE_KV_INSERT_FUSED" -_DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL_ENV = "VLLM_MUSA_DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL" _DEEPSEEK_V4_TP8_BALANCED_LONG_PREFILL_DEFAULTS = { _DEEPSEEK_V4_GEMV_MOE_BLOCK_ENV: "16x8", "VLLM_MUSA_FUSED_ADD_RMSNORM_BLOCK_X": "256", @@ -95,7 +94,6 @@ _DEEPSEEK_V4_INDEXER_MATERIALIZED_TOPK_SORTED_ENV: "0", _DEEPSEEK_V4_INDEXER_MATERIALIZED_DIRECT_ENV: "1", _DEEPSEEK_V4_QNORM_KV_FUSED_ENV: "1", - _DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL_ENV: "1", } _DEEPSEEK_V4_TP8_PROFILE_DEFAULTS = { _DEEPSEEK_V4_TP8_BALANCED_LONG_PREFILL_PROFILE: (