ayh015/myLightningOPD
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: ayh015/myLightningOPD
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-08-27 23:50:14 +08:00
commit d4e0a1af66
368 changed files with 559583 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
slime/backends/fsdp_utils/kernels/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

384
slime/backends/fsdp_utils/kernels/fused_experts.py Normal file
View File

@@ -0,0 +1,384 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import torch
import triton.language as tl
from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
invoke_fused_moe_kernel,
moe_align_block_size,
moe_sum_reduce,
silu_and_mul,
)
from .fused_moe_triton_backward_kernels import invoke_fused_moe_backward_kernel
class GateUpProjFunction(torch.autograd.Function):
@staticmethod
def forward(
ctx,
hidden_states: torch.Tensor,
w1: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
):
num_tokens, _ = hidden_states.shape
E, N, _ = w1.shape
# We execute the fused_moe kernel in chunks to circumvent this issue:
# https://github.com/vllm-project/vllm/issues/5938
CHUNK_SIZE = 64 * 1024
# default deterministic config
config = {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
}
topk = topk_ids.shape[1]
intermediate_cache1 = torch.empty(
(num_tokens * topk, N),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
begin_chunk_idx, end_chunk_idx = (
chunk * CHUNK_SIZE,
min((chunk + 1) * CHUNK_SIZE, num_tokens),
)
curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
cur_intermediate_cache1 = intermediate_cache1[begin_chunk_idx * topk : end_chunk_idx * topk]
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
curr_topk_ids, config["BLOCK_SIZE_M"], E
)
invoke_fused_moe_kernel(
curr_hidden_states,
w1,
None,
cur_intermediate_cache1,
None,
None,
None,
curr_topk_weights,
curr_topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
False,
topk_ids.shape[1],
config,
compute_type=tl.bfloat16,
use_fp8_w8a8=False,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=None,
c_sorted=False,
filter_expert=True,
)
ctx.save_for_backward(hidden_states, w1, topk_weights, topk_ids)
ctx.config = config
ctx.num_tokens = num_tokens
ctx.topk = topk
return intermediate_cache1
@staticmethod
def backward(ctx, grad_output):
"""
Backward pass for GateUpProjFunction using Triton kernels.
Args:
grad_output: shape (num_tokens * topk, N)
Returns:
(grad_hidden_states, grad_w1, grad_topk_weights, None)
"""
hidden_states, w1, topk_weights, topk_ids = ctx.saved_tensors
config = ctx.config
num_tokens = ctx.num_tokens
topk = ctx.topk
E, N, D_in = w1.shape
CHUNK_SIZE = 64 * 1024
# Initialize gradient tensors
grad_hidden_states = torch.zeros_like(hidden_states)
grad_w1 = torch.zeros_like(w1)
# GateUpProj stage doesn't need topk_weights gradient
grad_topk_weights = torch.zeros_like(topk_weights)
# Process in chunks to match forward pass
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
begin_chunk_idx, end_chunk_idx = (
chunk * CHUNK_SIZE,
min((chunk + 1) * CHUNK_SIZE, num_tokens),
)
curr_num_tokens = end_chunk_idx - begin_chunk_idx
if curr_num_tokens == 0:
continue
curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
curr_grad_output = grad_output[begin_chunk_idx * topk : end_chunk_idx * topk]
# Get aligned metadata
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
curr_topk_ids, config["BLOCK_SIZE_M"], E
)
# Prepare gradient buffer for this chunk
curr_grad_hidden_states = torch.zeros_like(curr_hidden_states)
curr_grad_w1 = torch.zeros_like(w1)
# Call Triton backward kernel with MUL_ROUTED_WEIGHT=False
# Use chunk of hidden_states to match sorted_token_ids indices
invoke_fused_moe_backward_kernel(
grad_output=curr_grad_output,
input=curr_hidden_states, # Use chunk of hidden_states to match sorted_token_ids
weight=w1,
grad_input=curr_grad_hidden_states,
grad_weight=curr_grad_w1,
grad_topk_weights=None, # Not needed for GateUpProj
topk_weights=curr_topk_weights,
topk_ids=curr_topk_ids,
sorted_token_ids=sorted_token_ids,
expert_ids=expert_ids,
num_tokens_post_padded=num_tokens_post_padded,
mul_routed_weight=False,
top_k=topk,
config=config,
compute_type=tl.bfloat16,
)
# Accumulate gradients
grad_hidden_states[begin_chunk_idx:end_chunk_idx] += curr_grad_hidden_states
grad_w1 += curr_grad_w1
return grad_hidden_states, grad_w1, grad_topk_weights, None
class SiluAndMulFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, intermediate_cache1: torch.Tensor):
num_tokens, N = intermediate_cache1.shape
intermediate_cache2 = torch.empty(
(num_tokens, N // 2),
device=intermediate_cache1.device,
dtype=intermediate_cache1.dtype,
)
silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
ctx.save_for_backward(intermediate_cache1)
return intermediate_cache2
@staticmethod
def backward(ctx, grad_output):
(intermediate_cache1,) = ctx.saved_tensors
N = intermediate_cache1.shape[-1]
x1, x2 = intermediate_cache1.view(-1, N).chunk(2, dim=-1)
silu_x1 = torch.nn.functional.silu(x1)
sig = torch.sigmoid(x1)
dsilu_dx1 = sig + x1 * sig * (1 - sig)
grad_x1 = grad_output * x2 * dsilu_dx1
grad_x2 = grad_output * silu_x1
grad_input = torch.cat([grad_x1, grad_x2], dim=-1)
return grad_input.view_as(intermediate_cache1)
class DownProjFunction(torch.autograd.Function):
@staticmethod
def forward(
ctx,
intermediate_cache2: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
):
num_tokens, _ = intermediate_cache2.shape
topk = topk_ids.shape[1]
num_tokens //= topk
E, _, _ = w2.shape
# We execute the fused_moe kernel in chunks to circumvent this issue:
# https://github.com/vllm-project/vllm/issues/5938
CHUNK_SIZE = 64 * 1024
# default deterministic config
config = {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8,
}
intermediate_cache3 = torch.empty(
(num_tokens, topk, w2.shape[1]),
device=intermediate_cache2.device,
dtype=intermediate_cache2.dtype,
)
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
begin_chunk_idx, end_chunk_idx = (
chunk * CHUNK_SIZE,
min((chunk + 1) * CHUNK_SIZE, num_tokens),
)
cur_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
cur_intermediate_cache3 = intermediate_cache3[begin_chunk_idx:end_chunk_idx]
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
curr_topk_ids, config["BLOCK_SIZE_M"], E
)
invoke_fused_moe_kernel(
cur_intermediate_cache2,
w2,
None,
cur_intermediate_cache3,
None,
None,
None,
curr_topk_weights,
curr_topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
True,
1,
config,
compute_type=tl.bfloat16,
use_fp8_w8a8=False,
use_int8_w8a8=False,
use_int8_w8a16=False,
use_int4_w4a16=False,
per_channel_quant=False,
block_shape=None,
a_use_tma=False,
b_use_tma=False,
)
ctx.save_for_backward(intermediate_cache2, w2, topk_weights, topk_ids)
ctx.config = config
ctx.num_tokens = num_tokens
ctx.topk = topk
return intermediate_cache3
@staticmethod
def backward(ctx, grad_output):
"""
Backward pass for DownProjFunction using Triton kernels.
Args:
grad_output: shape (num_tokens, topk, hidden_size)
Returns:
(grad_intermediate_cache2, grad_w2, grad_topk_weights, None)
"""
intermediate_cache2, w2, topk_weights, topk_ids = ctx.saved_tensors
config = ctx.config
num_tokens = ctx.num_tokens
topk = ctx.topk
E, hidden_size, intermediate_size = w2.shape
CHUNK_SIZE = 64 * 1024
# Initialize gradient tensors
grad_intermediate_cache2 = torch.zeros_like(intermediate_cache2)
grad_w2 = torch.zeros_like(w2)
grad_topk_weights = torch.zeros_like(topk_weights)
# Process in chunks to match forward pass
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
begin_chunk_idx, end_chunk_idx = (
chunk * CHUNK_SIZE,
min((chunk + 1) * CHUNK_SIZE, num_tokens),
)
curr_num_tokens = end_chunk_idx - begin_chunk_idx
if curr_num_tokens == 0:
continue
curr_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
curr_grad_output = grad_output[begin_chunk_idx:end_chunk_idx]
# Get aligned metadata
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
curr_topk_ids, config["BLOCK_SIZE_M"], E
)
# Prepare gradient buffers for this chunk
curr_grad_intermediate_cache2 = torch.zeros_like(curr_intermediate_cache2)
curr_grad_w2 = torch.zeros_like(w2)
curr_grad_topk_weights = torch.zeros_like(curr_topk_weights)
# Call Triton backward kernel with MUL_ROUTED_WEIGHT=True
# Note: Use top_k=1 to match forward pass indexing
invoke_fused_moe_backward_kernel(
grad_output=curr_grad_output,
input=curr_intermediate_cache2,
weight=w2,
grad_input=curr_grad_intermediate_cache2,
grad_weight=curr_grad_w2,
grad_topk_weights=curr_grad_topk_weights,
topk_weights=curr_topk_weights,
topk_ids=curr_topk_ids,
sorted_token_ids=sorted_token_ids,
expert_ids=expert_ids,
num_tokens_post_padded=num_tokens_post_padded,
mul_routed_weight=True,
top_k=1,
config=config,
compute_type=tl.bfloat16,
)
# Accumulate gradients
grad_intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk] = curr_grad_intermediate_cache2
grad_w2 += curr_grad_w2
grad_topk_weights[begin_chunk_idx:end_chunk_idx] = curr_grad_topk_weights
return grad_intermediate_cache2, grad_w2, grad_topk_weights, None
class MoeSumReduceFunction(torch.autograd.Function):
@staticmethod
def forward(
ctx,
intermediate_cache3: torch.Tensor,
hidden_states_shape,
):
out_hidden_states = torch.empty(
hidden_states_shape, device=intermediate_cache3.device, dtype=intermediate_cache3.dtype
)
moe_sum_reduce(
intermediate_cache3,
out_hidden_states,
1.0,
)
ctx.save_for_backward(intermediate_cache3)
return out_hidden_states
@staticmethod
def backward(ctx, grad_output):
(intermediate_cache3,) = ctx.saved_tensors
return grad_output.unsqueeze(1).expand_as(intermediate_cache3), None

543
slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py Normal file
View File

@@ -0,0 +1,543 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from typing import Any
import torch
import triton
import triton.language as tl
@triton.jit
def fused_moe_backward_input_kernel(
# Pointers to matrices
grad_output_ptr,
weight_ptr,
grad_input_ptr,
grad_topk_weights_ptr,
topk_weights_ptr,
sorted_token_ids_ptr,
expert_ids_ptr,
num_tokens_post_padded_ptr,
# Matrix dimensions
N,
K,
EM,
num_valid_tokens,
# Strides
stride_gom,
stride_gon,
stride_we,
stride_wn,
stride_wk,
stride_gim,
stride_gik,
# Meta-parameters
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
MUL_ROUTED_WEIGHT: tl.constexpr,
top_k: tl.constexpr,
compute_type: tl.constexpr,
):
"""
Backward kernel for computing grad_input.
Forward: output = input @ weight.T (optionally multiplied by topk_weights)
Backward: grad_input = grad_output @ weight (optionally multiplied by topk_weights)
This kernel computes: grad_input[token] = sum_over_N(grad_output[token, n] * weight[expert, n, :])
If MUL_ROUTED_WEIGHT: grad_input[token] *= topk_weights[token]
Parallelization: Similar to forward, parallel over M and N dimensions, loop over K.
"""
# Map program ids to blocks (parallel over M and N, similar to forward)
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
# Check bounds
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
# Only process if this block is valid
if pid_m * BLOCK_SIZE_M < num_tokens_post_padded:
# Load token information
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
offs_token = offs_token.to(tl.int64)
token_mask = offs_token < num_valid_tokens
# Get expert ID for this block
off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
# Only process if expert is valid
if off_experts != -1:
# Initialize offsets for N dimension (current block)
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
offs_k = tl.arange(0, BLOCK_SIZE_K)
# Load grad_output block: shape (BLOCK_SIZE_M, BLOCK_SIZE_N)
grad_output_ptrs = grad_output_ptr + (offs_token[:, None] * stride_gom + offs_n[None, :] * stride_gon)
grad_out = tl.load(
grad_output_ptrs,
mask=token_mask[:, None] & (offs_n[None, :] < N),
other=0.0,
)
# Apply topk_weights to grad_output if needed
if MUL_ROUTED_WEIGHT:
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
grad_out = grad_out * moe_weight[:, None]
# Iterate over K dimension
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
# Current K offsets
curr_offs_k = k * BLOCK_SIZE_K + offs_k
# Load weight block: shape (BLOCK_SIZE_N, BLOCK_SIZE_K)
# weight: shape (E, N, K)
weight_ptrs = (
weight_ptr
+ off_experts * stride_we
+ offs_n[:, None] * stride_wn
+ curr_offs_k[None, :] * stride_wk
)
w = tl.load(
weight_ptrs,
mask=(offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
other=0.0,
)
# Compute contribution: grad_out @ weight
# grad_out: (BLOCK_SIZE_M, BLOCK_SIZE_N)
# w: (BLOCK_SIZE_N, BLOCK_SIZE_K)
# result: (BLOCK_SIZE_M, BLOCK_SIZE_K)
contribution = tl.dot(grad_out, w)
# Atomic add to grad_input because different N blocks contribute to same K
grad_input_ptrs = grad_input_ptr + (
(offs_token[:, None] // top_k) * stride_gim + curr_offs_k[None, :] * stride_gik
)
grad_input_mask = token_mask[:, None] & (curr_offs_k[None, :] < K)
tl.atomic_add(grad_input_ptrs, contribution.to(compute_type), mask=grad_input_mask)
@triton.jit
def fused_moe_backward_weight_kernel(
# Pointers to matrices
grad_output_ptr,
input_ptr,
grad_weight_ptr,
topk_weights_ptr,
sorted_token_ids_ptr,
expert_ids_ptr,
num_tokens_post_padded_ptr,
# Matrix dimensions
N,
K,
EM,
num_valid_tokens,
# Strides
stride_gom,
stride_gon,
stride_im,
stride_ik,
stride_gwe,
stride_gwn,
stride_gwk,
# Meta-parameters
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
MUL_ROUTED_WEIGHT: tl.constexpr,
top_k: tl.constexpr,
compute_type: tl.constexpr,
):
"""
Backward kernel for computing grad_weight.
Forward: output = input @ weight.T (optionally multiplied by topk_weights)
Backward: grad_weight = input.T @ grad_output (optionally multiplied by topk_weights)
This kernel computes: grad_weight[expert, n, k] = sum_over_tokens(input[token, k] * grad_output[token, n])
If MUL_ROUTED_WEIGHT: the accumulation is weighted by topk_weights[token]
Parallelization: Parallel over M and N dimensions with grouping, loop over K.
"""
# Map program ids to blocks (parallel over M and N with grouping, similar to forward and backward_input)
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
# Check bounds
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
# Only process if this block is valid
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
return
# Get expert ID for this M block
expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
# Only process if expert is valid
if expert_id == -1:
return
# Load token information for this M block
offs_m = tl.arange(0, BLOCK_SIZE_M)
offs_token_id = pid_m * BLOCK_SIZE_M + offs_m.to(tl.int64)
offs_token = tl.load(
sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
)
offs_token = offs_token.to(tl.int64)
token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
# Clamp offs_token to valid range
offs_token_clamped = tl.where(token_mask, offs_token, 0)
# Determine input token indices based on MUL_ROUTED_WEIGHT
if MUL_ROUTED_WEIGHT:
input_token_idx = offs_token_clamped
input_mask = token_mask
else:
input_token_idx = offs_token_clamped // top_k
num_input_tokens = num_valid_tokens // top_k
input_mask = token_mask & (input_token_idx < num_input_tokens)
# Load topk_weights if needed
if MUL_ROUTED_WEIGHT:
moe_weight = tl.load(topk_weights_ptr + offs_token_clamped, mask=token_mask, other=0.0)
# Current N offset for this program
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
# Load grad_output for this N block: shape (M, BLOCK_SIZE_N)
# grad_output is always indexed by sorted_token_ids (offs_token_clamped)
# because it has shape (num_tokens * topk, N)
grad_output_ptrs = grad_output_ptr + (offs_token_clamped[:, None] * stride_gom + offs_n[None, :] * stride_gon)
grad_out = tl.load(
grad_output_ptrs,
mask=token_mask[:, None] & (offs_n[None, :] < N),
other=0.0,
)
# Apply topk_weights if needed
if MUL_ROUTED_WEIGHT:
grad_out = grad_out * moe_weight[:, None]
# Zero out padding tokens
token_mask_col = token_mask[:, None]
grad_out = grad_out * token_mask_col
# Iterate over K blocks and accumulate
for k_block in range(tl.cdiv(K, BLOCK_SIZE_K)):
offs_k = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
# Load input for this K block
input_ptrs = input_ptr + (input_token_idx[:, None] * stride_im + offs_k[None, :] * stride_ik)
inp = tl.load(
input_ptrs,
mask=input_mask[:, None] & (offs_k[None, :] < K),
other=0.0,
)
# Zero out padding tokens - use input_mask for input, token_mask for grad_output
input_mask_col = input_mask[:, None]
inp = inp * input_mask_col
# Compute grad_weight contribution: grad_out.T @ inp
grad_w_contribution = tl.dot(grad_out.T, inp)
# Write back using atomic add
grad_weight_ptrs = (
grad_weight_ptr + expert_id * stride_gwe + offs_n[:, None] * stride_gwn + offs_k[None, :] * stride_gwk
)
grad_weight_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
tl.atomic_add(grad_weight_ptrs, grad_w_contribution.to(compute_type), mask=grad_weight_mask)
@triton.jit
def fused_moe_backward_topk_weights_kernel(
# Pointers to matrices
grad_output_ptr,
input_ptr,
weight_ptr,
grad_topk_weights_ptr,
sorted_token_ids_ptr,
expert_ids_ptr,
num_tokens_post_padded_ptr,
# Matrix dimensions
N,
K,
EM,
num_valid_tokens,
# Strides
stride_gom,
stride_gon,
stride_im,
stride_ik,
stride_we,
stride_wn,
stride_wk,
# Meta-parameters
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
top_k: tl.constexpr,
compute_type: tl.constexpr,
):
"""
Backward kernel for computing grad_topk_weights.
Forward: output = topk_weights * (input @ weight.T)
Backward: grad_topk_weights = sum(grad_output * (input @ weight.T))
This kernel computes the gradient of topk_weights by computing the dot product
of grad_output with the forward output before weight multiplication.
"""
# Map program id to token block
pid = tl.program_id(axis=0)
# Check bounds
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
# Only process if this block is valid
if pid * BLOCK_SIZE_M < num_tokens_post_padded:
# Load token information
offs_token_id = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
offs_token = tl.load(
sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
)
offs_token = offs_token.to(tl.int64)
token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
# Clamp offs_token to valid range for safe pointer arithmetic
offs_token_clamped = tl.where(token_mask, offs_token, 0)
# Get expert ID for this block
off_experts = tl.load(expert_ids_ptr + pid).to(tl.int64)
# Only process if expert is valid
if off_experts != -1:
# Initialize offsets
offs_n = tl.arange(0, BLOCK_SIZE_N)
offs_k = tl.arange(0, BLOCK_SIZE_K)
# Accumulator for grad_topk_weights
accumulator = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
# Iterate over N and K dimensions to compute forward output and gradient
for n in range(0, tl.cdiv(N, BLOCK_SIZE_N)):
# Current N offset
curr_offs_n = n * BLOCK_SIZE_N + offs_n
# Load grad_output block: (M, N)
grad_output_ptrs = grad_output_ptr + (
offs_token_clamped[:, None] * stride_gom + curr_offs_n[None, :] * stride_gon
)
grad_out = tl.load(
grad_output_ptrs,
mask=token_mask[:, None] & (curr_offs_n[None, :] < N),
other=0.0,
)
# Compute forward output for this N block: input @ weight[:, n, :].T
forward_output_n = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
# Current K offset
curr_offs_k = k * BLOCK_SIZE_K + offs_k
# Load input block: (M, K)
input_ptrs = input_ptr + (
(offs_token_clamped[:, None] // top_k) * stride_im + curr_offs_k[None, :] * stride_ik
)
inp = tl.load(
input_ptrs,
mask=token_mask[:, None] & (curr_offs_k[None, :] < K),
other=0.0,
)
# Load weight block: (N, K)
weight_ptrs = (
weight_ptr
+ off_experts * stride_we
+ curr_offs_n[:, None] * stride_wn
+ curr_offs_k[None, :] * stride_wk
)
w = tl.load(
weight_ptrs,
mask=(curr_offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
other=0.0,
)
# Accumulate forward output: input @ weight.T
# inp: (M, K), w.T: (K, N) -> (M, N)
forward_output_n += tl.dot(inp, w.T)
# Compute contribution to grad_topk_weights: sum(grad_out * forward_output)
# Sum over N dimension
accumulator += tl.sum(grad_out * forward_output_n, axis=1)
# Write back grad_topk_weights using atomic add with clamped token indices
tl.atomic_add(grad_topk_weights_ptr + offs_token_clamped, accumulator.to(compute_type), mask=token_mask)
def invoke_fused_moe_backward_kernel(
grad_output: torch.Tensor,
input: torch.Tensor,
weight: torch.Tensor,
grad_input: torch.Tensor,
grad_weight: torch.Tensor,
grad_topk_weights: torch.Tensor | None,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
sorted_token_ids: torch.Tensor,
expert_ids: torch.Tensor,
num_tokens_post_padded: torch.Tensor,
mul_routed_weight: bool,
top_k: int,
config: dict[str, Any],
compute_type: tl.dtype,
) -> None:
"""
Invoke the fused MOE backward kernels to compute gradients.
Args:
grad_output: Gradient of output, shape (num_tokens * topk, N) or (num_tokens, topk, N)
input: Input tensor, shape (num_tokens, K)
weight: Weight tensor, shape (E, N, K)
grad_input: Output gradient for input, shape (num_tokens, K)
grad_weight: Output gradient for weight, shape (E, N, K)
grad_topk_weights: Output gradient for topk_weights, shape (num_tokens, topk) or None
topk_weights: Top-K routing weights, shape (num_tokens, topk)
topk_ids: Top-K expert IDs, shape (num_tokens, topk)
sorted_token_ids: Sorted token IDs
expert_ids: Expert IDs for each block
num_tokens_post_padded: Number of tokens after padding
mul_routed_weight: Whether to multiply by routing weights
top_k: Number of experts per token
config: Kernel configuration
compute_type: Computation data type
"""
assert topk_weights.stride(1) == 1
assert sorted_token_ids.stride(0) == 1
# Flatten grad_output if needed
# Before: (num_tokens, topk, hidden_size)
# After: (num_tokens * topk, hidden_size)
if grad_output.ndim == 3:
grad_output = grad_output.reshape(-1, grad_output.shape[-1])
E, N, K = weight.shape
# ===================== Compute grad_input =====================
def grid_input(META):
return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
fused_moe_backward_input_kernel[grid_input](
grad_output,
weight,
grad_input,
grad_topk_weights if grad_topk_weights is not None else grad_input, # dummy pointer
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
N,
K,
sorted_token_ids.shape[0],
grad_output.shape[0],
grad_output.stride(0),
grad_output.stride(1),
weight.stride(0),
weight.stride(1),
weight.stride(2),
grad_input.stride(0),
grad_input.stride(1),
MUL_ROUTED_WEIGHT=mul_routed_weight,
top_k=top_k,
compute_type=compute_type,
**config,
)
# ===================== Compute grad_weight =====================
# Initialize grad_weight to zero
grad_weight.zero_()
# Use same grid configuration as forward kernel: encode both M and N dimensions
def grid_weight(META):
return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
fused_moe_backward_weight_kernel[grid_weight](
grad_output,
input,
grad_weight,
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
N,
K,
sorted_token_ids.shape[0],
grad_output.shape[0],
grad_output.stride(0),
grad_output.stride(1),
input.stride(0),
input.stride(1),
grad_weight.stride(0),
grad_weight.stride(1),
grad_weight.stride(2),
MUL_ROUTED_WEIGHT=mul_routed_weight,
top_k=top_k,
compute_type=compute_type,
**config,
)
# ===================== Compute grad_topk_weights (if needed) =====================
if mul_routed_weight and grad_topk_weights is not None:
def grid_topk(META):
return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]),)
fused_moe_backward_topk_weights_kernel[grid_topk](
grad_output,
input,
weight,
grad_topk_weights.view(-1),
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
N,
K,
sorted_token_ids.shape[0],
grad_output.shape[0],
grad_output.stride(0),
grad_output.stride(1),
input.stride(0),
input.stride(1),
weight.stride(0),
weight.stride(1),
weight.stride(2),
top_k=top_k,
compute_type=compute_type,
BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
)