diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py new file mode 100644 index 0000000000..101d316937 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py @@ -0,0 +1,232 @@ +import random +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from typing import Dict + + +@triton.jit +def _fwd_kernel_ep_scatter_1( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_EXPERT_NUM: tl.constexpr, +): + cur_expert = tl.program_id(0) + + offset_cumsum = tl.arange(0, BLOCK_EXPERT_NUM) + tokens_per_expert = tl.load(num_recv_tokens_per_expert + offset_cumsum, mask=offset_cumsum < num_experts, other=0) + cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert + tl.store(expert_start_loc + offset_cumsum, cumsum, mask=offset_cumsum < num_experts) + + cur_expert_start = tl.load(expert_start_loc + cur_expert) + cur_expert_token_num = tl.load(num_recv_tokens_per_expert + cur_expert) + + m_indices_start_ptr = m_indices + cur_expert_start + off_expert = tl.arange(0, BLOCK_E) + + for start_m in tl.range(0, cur_expert_token_num, BLOCK_E, num_stages=4): + tl.store( + m_indices_start_ptr + start_m + off_expert, + cur_expert, + ) + + +@triton.jit +def _fwd_kernel_ep_scatter_2( + total_token_num, + expert_start_loc, + recv_x, + recv_x_stride0, + recv_x_stride1, + recv_x_scale, + recv_x_scale_stride0, + recv_x_scale_stride1, + recv_topk, + recv_topk_stride0, + recv_topk_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + output_tensor_scale, + output_tensor_scale_stride0, + output_tensor_scale_stride1, + output_index, + output_index_stride0, + output_index_stride1, + topk_num: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + HIDDEN_SIZE_PAD: tl.constexpr, + SCALE_HIDDEN_SIZE: tl.constexpr, + SCALE_HIDDEN_SIZE_PAD: tl.constexpr, +): + start_token_id = tl.program_id(0) + grid_num = tl.num_programs(0) + + offset_in = tl.arange(0, HIDDEN_SIZE_PAD) + mask = offset_in < HIDDEN_SIZE + + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = offset_in_s < SCALE_HIDDEN_SIZE + for token_id in range(start_token_id, total_token_num, grid_num): + to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) + to_copy_s = tl.load(recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s) + + for topk_index in tl.range(0, topk_num, 1, num_stages=4): + expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) + if expert_id >= 0: + dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1) + dest_token_index = dest_token_index.to(tl.int64) + tl.store(output_index + token_id * output_index_stride0 + topk_index, dest_token_index) + output_tensor_ptr = output_tensor + dest_token_index * output_tensor_stride0 + output_tensor_scale_ptr = output_tensor_scale + dest_token_index * output_tensor_scale_stride0 + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s) + + +@torch.no_grad() +def ep_scatter( + recv_x: torch.Tensor, + recv_x_scale: torch.Tensor, + recv_topk: torch.Tensor, + num_recv_tokens_per_expert: torch.Tensor, + expert_start_loc: torch.Tensor, + output_tensor: torch.Tensor, + output_tensor_scale: torch.Tensor, + m_indices: torch.Tensor, + output_index: torch.Tensor, +): + BLOCK_E = 128 # token num of per expert is aligned to 128 + BLOCK_D = 128 # block size of quantization + num_warps = 8 + num_experts = num_recv_tokens_per_expert.shape[0] # 获取num_recv_tokens_per_expert的元素个数 + hidden_size = recv_x.shape[1] + # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts) + grid = num_experts + + assert m_indices.shape[0] % BLOCK_E == 0 + + _fwd_kernel_ep_scatter_1[(grid,)]( + num_recv_tokens_per_expert, + expert_start_loc, + m_indices, + num_experts=num_experts, + num_warps=num_warps, + BLOCK_E=BLOCK_E, + BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts), + ) + + grid = min(recv_topk.shape[0], 1024 * 8) + + _fwd_kernel_ep_scatter_2[(grid,)]( + recv_topk.shape[0], + expert_start_loc, + recv_x, + recv_x.stride(0), + recv_x.stride(1), + recv_x_scale, + recv_x_scale.stride(0), + recv_x_scale.stride(1), + recv_topk, + recv_topk.stride(0), + recv_topk.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + output_tensor_scale, + output_tensor_scale.stride(0), + output_tensor_scale.stride(1), + output_index, + output_index.stride(0), + output_index.stride(1), + topk_num=recv_topk.shape[1], + num_warps=num_warps, + HIDDEN_SIZE=hidden_size, + HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), + SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D, + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D), + ) + return + + +@triton.jit +def _fwd_kernel_ep_gather( + total_token_num, + input_tensor, + input_tensor_stride0, + input_tensor_stride1, + recv_topk_ids, + recv_topk_ids_stride0, + recv_topk_ids_stride1, + recv_topk_weight, + recv_topk_weight_stride0, + recv_topk_weight_stride1, + input_index, + input_index_stride0, + input_index_stride1, + output_tensor, + output_tensor_stride0, + output_tensor_stride1, + topk_num: tl.constexpr, + BLOCK_D: tl.constexpr, +): + cur_block = tl.program_id(0) + start_cur_token = tl.program_id(1) + grid_num = tl.num_programs(1) + + for cur_token in range(start_cur_token, total_token_num, grid_num): + off_d = tl.arange(0, BLOCK_D) + accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) + for topk_index in range(0, topk_num): + expert_id = tl.load(recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index) + if expert_id >= 0: + source_token_index = tl.load(input_index + cur_token * input_index_stride0 + topk_index) + acc_weight = tl.load(recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index) + tmp = tl.load(input_tensor + source_token_index * input_tensor_stride0 + cur_block * BLOCK_D + off_d) + accumulator += tmp.to(tl.float32) * acc_weight + + tl.store( + output_tensor + cur_token * output_tensor_stride0 + cur_block * BLOCK_D + off_d, + accumulator.to(output_tensor.dtype.element_ty), + ) + + +@torch.no_grad() +def ep_gather( + input_tensor: torch.Tensor, + recv_topk_ids: torch.Tensor, + recv_topk_weight: torch.Tensor, + input_index: torch.Tensor, + output_tensor: torch.Tensor, +): + BLOCK_D = 1024 # block size of quantization + num_warps = 2 + num_tokens = output_tensor.shape[0] + hidden_size = input_tensor.shape[1] + assert hidden_size % BLOCK_D == 0 + grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024)) + _fwd_kernel_ep_gather[grid]( + num_tokens, + input_tensor, + input_tensor.stride(0), + input_tensor.stride(1), + recv_topk_ids, + recv_topk_ids.stride(0), + recv_topk_ids.stride(1), + recv_topk_weight, + recv_topk_weight.stride(0), + recv_topk_weight.stride(1), + input_index, + input_index.stride(0), + input_index.stride(1), + output_tensor, + output_tensor.stride(0), + output_tensor.stride(1), + topk_num=recv_topk_ids.shape[1], + num_warps=num_warps, + BLOCK_D=BLOCK_D, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py b/lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py index 8dc8558922..f3a1c0e73d 100644 --- a/lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py +++ b/lightllm/common/basemodel/triton_kernel/norm/rmsnorm.py @@ -39,7 +39,7 @@ def _rms_norm_fwd_fused( mask = cols < N if HAS_WEIGHT: w = tl.load(W + cols, mask=mask).to(tl.float32) - x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + x = tl.load(X + cols * x_stride1, mask=mask, other=0.0).to(tl.float32) x_hat = x * rstd y = x_hat if HAS_WEIGHT: diff --git a/unit_tests/common/basemodel/triton_kernel/test_rmsnorm.py b/unit_tests/common/basemodel/triton_kernel/test_rmsnorm.py new file mode 100644 index 0000000000..66d59fb086 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_rmsnorm.py @@ -0,0 +1,49 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward, torch_rms_norm + + +@pytest.mark.parametrize("M,N", [(64, 256), (17, 1024), (128, 128)]) +@pytest.mark.parametrize("has_weight", [True, False]) +def test_rmsnorm_contiguous(M, N, has_weight): + """Ordinary contiguous input: the shape the model paths use today.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for rmsnorm test") + + torch.manual_seed(0) + x = torch.randn((M, N), device="cuda", dtype=torch.float32) + weight = torch.rand((N,), device="cuda", dtype=torch.float32) if has_weight else None + + out = rmsnorm_forward(x, weight, eps=1e-6) + ref = torch_rms_norm(x, weight if weight is not None else 1.0, 1e-6) + + assert (out - ref).abs().max().item() < 1e-5 + + +@pytest.mark.parametrize("M,N", [(64, 256), (17, 1024), (128, 128)]) +@pytest.mark.parametrize("has_weight", [True, False]) +def test_rmsnorm_last_dim_stride(M, N, has_weight): + """The kernel takes x_stride1, so an input whose last-dim stride is not 1 must work. + + `rmsnorm_forward` reaches the kernel with one without complaint: `x.view(-1, N)` is + shape-preserving for a 2-D input, and `torch.empty_like` gives the output the same + strides. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for rmsnorm test") + + torch.manual_seed(0) + x = torch.randn((N, M), device="cuda", dtype=torch.float32).t() + assert x.shape == (M, N) and x.stride(1) != 1 + weight = torch.rand((N,), device="cuda", dtype=torch.float32) if has_weight else None + + out = rmsnorm_forward(x, weight, eps=1e-6) + ref = torch_rms_norm(x, weight if weight is not None else 1.0, 1e-6) + + max_diff = (out - ref).abs().max().item() + assert max_diff < 1e-5, f"max diff too large: {max_diff}" + + # and it must agree with the same values laid out contiguously + contiguous_out = rmsnorm_forward(x.contiguous(), weight, eps=1e-6) + assert (out - contiguous_out).abs().max().item() < 1e-5