fix: apply x_stride1 in the rmsnorm normalize pass

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

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:

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