import logging from typing import Tuple import torch import triton from sglang.srt.utils import ceil_div, is_cuda, is_musa logger = logging.getLogger(__name__) _is_cuda = is_cuda() _is_musa = is_musa() if _is_cuda or _is_musa: from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8 as per_token_group_quant_fp8, ) import triton.language as tl def _get_launch_config_1d(device, numel): MAX_THREADS_PER_BLOCK = 1024 MIN_THREADS_PER_BLOCK = 512 MAX_WAVES = 8 # empirical numbers props = torch.cuda.get_device_properties(device) sm_count = props.multi_processor_count max_threads_per_sm = props.max_threads_per_multi_processor max_num_blocks = sm_count * max_threads_per_sm // MAX_THREADS_PER_BLOCK block_dim = MAX_THREADS_PER_BLOCK def get_num_blocks(block_dim): return triton.cdiv(numel, block_dim) while ( block_dim > MIN_THREADS_PER_BLOCK and get_num_blocks(block_dim // 2) <= max_num_blocks ): block_dim = block_dim // 2 num_blocks = get_num_blocks(block_dim) grid_dim = min(num_blocks, max_num_blocks * MAX_WAVES) return (grid_dim,), block_dim def _get_launch_config_2d(device, m, n): MAX_THREADS_PER_BLOCK = 1024 MIN_THREADS_PER_BLOCK = 512 MAX_WAVES = 8 # empirical numbers props = torch.cuda.get_device_properties(device) sm_count = props.multi_processor_count max_threads_per_sm = props.max_threads_per_multi_processor max_num_blocks = sm_count * max_threads_per_sm // MAX_THREADS_PER_BLOCK block_dim = MAX_THREADS_PER_BLOCK def get_num_blocks(block_dim): return m * triton.cdiv(n, block_dim) while ( block_dim > MIN_THREADS_PER_BLOCK and get_num_blocks(block_dim // 2) <= max_num_blocks ): block_dim = block_dim // 2 grid_dim_x = triton.cdiv(n, block_dim) grid_dim_y = max(min(m, max_num_blocks * MAX_WAVES // grid_dim_x), 1) return (grid_dim_y, grid_dim_x), block_dim @triton.jit def deepep_permute_triton_kernel( input_ptr, gateup_input_ptr, src2dst_ptr, topk_ids_ptr, a1_scales_ptr, topk, hidden_size, BLOCK_SIZE: tl.constexpr, ): OutDtype = gateup_input_ptr.dtype.element_ty src_idx = tl.program_id(0) src2dst_ptr = src2dst_ptr + src_idx * topk topk_ids_ptr = topk_ids_ptr + src_idx * topk src_ptr = input_ptr + src_idx * hidden_size for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): offset = start_offset + tl.arange(0, BLOCK_SIZE) mask = offset < hidden_size in_data = tl.load(src_ptr + offset, mask=mask).to(OutDtype) for idx in range(topk): dst_idx = tl.load(src2dst_ptr + idx) if dst_idx >= 0: dst_ptr = gateup_input_ptr + dst_idx * hidden_size tl.store(dst_ptr + offset, in_data, mask=mask) @triton.jit def deepep_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) src2dst_ptr = src2dst_ptr + src_idx * topk topk_ids_ptr = topk_ids_ptr + src_idx * topk topk_weights_ptr = topk_weights_ptr + src_idx * topk store_ptr = output_ptr + src_idx * hidden_size for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): offset = start_offset + tl.arange(0, BLOCK_SIZE) mask = offset < hidden_size sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype) for idx in range(topk): dst_idx = tl.load(src2dst_ptr + idx) if dst_idx >= 0: weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype) load_ptr = down_output_ptr + dst_idx * hidden_size in_data = tl.load(load_ptr + offset, mask=mask) sum_vec += in_data * weigh_scale tl.store(store_ptr + offset, sum_vec, mask=mask) @triton.jit def compute_src2dst_triton_kernel( reorder_ids, src2dst, num_toks, BLOCK_SIZE: tl.constexpr ): pid = tl.program_id(axis=0) dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = dst_id < num_toks src_id = tl.load(reorder_ids + dst_id, mask=mask) tl.store(src2dst + src_id, dst_id, mask=mask) @triton.jit def deepep_compute_src2dst_triton_kernel( reorder_ids, src2dst, num_toks, num_minus_one, BLOCK_SIZE: tl.constexpr ): pid = tl.program_id(axis=0) dst_id = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = dst_id < num_toks src_id = tl.load(reorder_ids + dst_id, mask=mask) num_invalid = tl.load(num_minus_one) tl.store(src2dst + src_id, dst_id - num_invalid, mask=mask) def deepep_run_moe_deep_preprocess(topk_ids: torch.Tensor, num_experts: int): reorder_topk_ids, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) seg_indptr = torch.empty(num_experts + 1, device=topk_ids.device, dtype=torch.int64) src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int64) # Find offset expert_ids = torch.arange( num_experts + 1, device=topk_ids.device, dtype=reorder_topk_ids.dtype ) torch.searchsorted(reorder_topk_ids, expert_ids, out=seg_indptr) num_minus_one = seg_indptr[0] seg_indptr = seg_indptr - num_minus_one BLOCK_SIZE = 512 grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) deepep_compute_src2dst_triton_kernel[grid]( reorder_ids, src2dst, topk_ids.numel(), num_minus_one, BLOCK_SIZE ) reorder_topk_ids = reorder_topk_ids[num_minus_one:] return reorder_topk_ids, src2dst, seg_indptr def cutlass_w4_run_moe_ep_preproess(topk_ids: torch.Tensor): _, reorder_ids = torch.sort(topk_ids.view(-1), stable=True) BLOCK_SIZE = 512 grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) src2dst = torch.empty(topk_ids.numel(), device=topk_ids.device, dtype=torch.int32) compute_src2dst_triton_kernel[grid]( reorder_ids, src2dst, topk_ids.numel(), BLOCK_SIZE ) return src2dst @triton.jit def pre_reorder_triton_kernel_for_cutlass_moe( input_ptr, gateup_input_ptr, src2dst_ptr, topk_ids_ptr, a1_scales_ptr, num_local_experts, topk, num_tokens, hidden_size, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr, ): OutDtype = gateup_input_ptr.dtype.element_ty if a1_scales_ptr is not None: a1_scale = 1.0 / tl.load(a1_scales_ptr) else: a1_scale = 1.0 offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) mask = offset < hidden_size start_src_idx = tl.program_id(0) step = tl.num_programs(0) for src_idx_int32 in tl.range( start_src_idx, num_tokens, step, num_stages=NUM_STAGES ): src_idx = src_idx_int32.to(tl.int64) token_src2dst_ptr = src2dst_ptr + src_idx * topk token_topk_ids_ptr = topk_ids_ptr + src_idx * topk src_ptr_offs = input_ptr + src_idx * hidden_size + offset dst_ptr_offs = gateup_input_ptr + offset in_data = tl.load(src_ptr_offs, mask=mask).to(tl.float32) out_data = (in_data * a1_scale).to(OutDtype) for idx in range(topk): expert_id = tl.load(token_topk_ids_ptr + idx) if expert_id != num_local_experts: dst_idx = tl.load(token_src2dst_ptr + idx) tl.store(dst_ptr_offs + dst_idx * hidden_size, out_data, mask=mask) def pre_reorder_for_cutlass_moe( input, gateup_input, src2dst, topk_ids, a1_scales, num_local_experts, topk, num_tokens, hidden_size, ): grid, block_dim = _get_launch_config_2d(input.device, num_tokens, hidden_size) pre_reorder_triton_kernel_for_cutlass_moe[grid]( input_ptr=input, gateup_input_ptr=gateup_input, src2dst_ptr=src2dst, topk_ids_ptr=topk_ids, a1_scales_ptr=a1_scales, num_local_experts=num_local_experts, topk=topk, num_tokens=num_tokens, hidden_size=hidden_size, BLOCK_SIZE=block_dim, NUM_STAGES=3, ) # copy from https://github.com/ModelTC/lightllm/blob/a000ab69098654df4731f5b12587dd4e7f0a4f41/lightllm/common/fused_moe/moe_silu_and_mul_mix_quant_ep.py @triton.jit def _silu_and_mul_post_quant_kernel( input_ptr, stride_input_0, stride_input_1, stride_input_2, output_ptr, stride_output_0, stride_output_1, stride_output_2, output_scale_ptr, stride_output_scale_0, stride_output_scale_1, stride_output_scale_2, masked_m_ptr, size_n, fp8_max, fp8_min, QUANT_GROUP_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, SCALE_UE8M0: tl.constexpr, GEMM1_ALPHA: tl.constexpr, GEMM1_CLAMP_LIMIT: tl.constexpr, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) hidden_dim_block_index = tl.program_id(0) block_num_per_expert = tl.num_programs(1) token_num_cur_expert = tl.load(masked_m_ptr + expert_id) stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64) stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64) stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d scale_base = ( output_scale_ptr + expert_id * stride_output_scale_0 + hidden_dim_block_index * N_GROUPS * stride_output_scale_2 ) for token_index in tl.range( token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE ): gate = tl.load( input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0, ).to(tl.float32) up = tl.load( input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0, ) if GEMM1_ALPHA > 0: gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT) up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT) gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1) else: gate = gate / (1 + tl.exp(-gate)) gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE)) group_absmax = tl.max(tl.abs(gate_up_2d), axis=1) group_absmax = tl.maximum(group_absmax, 1e-10) output_s = group_absmax / fp8_max if SCALE_UE8M0: output_s = tl.exp2(tl.ceil(tl.log2(output_s))) inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1)) output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max) output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty) tl.store( output_ptr_offs + token_index * stride_output_1, output_q, mask=offs_in_d < size_n, ) scale_offs = scale_base + token_index * stride_output_scale_1 tl.store( scale_offs + tl.arange(0, N_GROUPS) * stride_output_scale_2, output_s, ) def silu_and_mul_masked_post_quant_fwd( input: torch.Tensor, output: torch.Tensor, output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, scale_ue8m0: bool = False, gemm1_alpha: float = 0.0, gemm1_clamp_limit: float = 0.0, ): """ input shape [expert_num, token_num_padded, hidden_dim] output shape [expert_num, token_num_padded, hidden_dim // 2], dtype fp8 output_scale [expert_num token_num_paddded, hidden_dim // 2 // 128] dtype float32 quant_group_size int, masked_m shape [expert_num], """ assert input.is_contiguous() assert output.dtype == torch.float8_e4m3fn assert output.is_contiguous() assert len(input.shape) == 3 assert input.shape[0] == masked_m.shape[0] assert input.shape[-1] % 2 == 0 size_n = input.shape[-1] // 2 assert size_n % quant_group_size == 0 expert_num = len(masked_m) if expert_num < 4: BLOCK_NUM_PER_EXPERT = 64 else: BLOCK_NUM_PER_EXPERT = 32 groups_total = size_n // quant_group_size gpb = 4 while gpb > 1: block_n = quant_group_size * gpb if (block_n & (block_n - 1) == 0) and (groups_total % gpb == 0): break gpb //= 2 BLOCK_N = quant_group_size * gpb num_warps = 1 NUM_STAGES = 6 hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) grid = ( hidden_dim_split_block_num, BLOCK_NUM_PER_EXPERT, expert_num, ) finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max fp8_min = -fp8_max _silu_and_mul_post_quant_kernel[grid]( input, *input.stride(), output, *output.stride(), output_scale, *output_scale.stride(), masked_m, size_n, fp8_max, fp8_min, QUANT_GROUP_SIZE=quant_group_size, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, num_warps=num_warps, SCALE_UE8M0=scale_ue8m0, GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0, GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0, ) return @triton.jit def _silu_and_mul_post_quant_packed_kernel( input_ptr, stride_input_0, stride_input_1, stride_input_2, output_ptr, stride_output_0, stride_output_1, stride_output_2, scale_ptr, # int32 [E, G//4, m_max] (MN-major: packed scale stored token-minor) stride_scale_e, stride_scale_g4, stride_scale_m, masked_m_ptr, num_experts, size_n, fp8_max, fp8_min, QUANT_GROUP_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, # == 4 * QUANT_GROUP_SIZE (one packed int32 per block) GEMM1_ALPHA: tl.constexpr, GEMM1_CLAMP_LIMIT: tl.constexpr, E_PADDED: tl.constexpr, ): # Flat work dim rides axis 0 (x): num_real_tokens*topk can exceed the 65535 grid.y/z limit. work_id = tl.program_id(0) hidden_dim_block_index = tl.program_id(1) e_off = tl.arange(0, E_PADDED) mm = tl.load(masked_m_ptr + e_off, mask=e_off < num_experts, other=0) incl = tl.cumsum(mm) total = tl.sum(mm) if work_id >= total: return excl = incl - mm # first global slot of each expert owner = (excl <= work_id) & (work_id < incl) expert_id = tl.sum(tl.where(owner, e_off, 0)) token_index = work_id - tl.sum(tl.where(owner, excl, 0)) stride_input_0 = tl.cast(stride_input_0, tl.int64) stride_input_1 = tl.cast(stride_input_1, tl.int64) stride_output_0 = tl.cast(stride_output_0, tl.int64) stride_output_1 = tl.cast(stride_output_1, tl.int64) N_GROUPS: tl.constexpr = BLOCK_N // QUANT_GROUP_SIZE offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) mask_d = offs_in_d < size_n in_base = input_ptr + expert_id * stride_input_0 + token_index * stride_input_1 out_base = output_ptr + expert_id * stride_output_0 + token_index * stride_output_1 gate = tl.load(in_base + offs_in_d, mask=mask_d, other=0.0).to(tl.float32) up = tl.load(in_base + offs_in_d + size_n, mask=mask_d, other=0.0) if GEMM1_ALPHA > 0: gate = tl.minimum(gate, GEMM1_CLAMP_LIMIT) up = tl.clamp(up, -GEMM1_CLAMP_LIMIT, GEMM1_CLAMP_LIMIT) gate_up = gate * tl.sigmoid(gate * GEMM1_ALPHA) * (up + 1) else: gate = gate / (1 + tl.exp(-gate)) gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate gate_up_2d = tl.reshape(gate_up, (N_GROUPS, QUANT_GROUP_SIZE)) group_absmax = tl.max(tl.abs(gate_up_2d), axis=1) group_absmax = tl.maximum(group_absmax, 1e-10) output_s = group_absmax / fp8_max # UE8M0: round to a power of two so the (>>23)&0xFF exponent extraction below is valid. output_s = tl.exp2(tl.ceil(tl.log2(output_s))) inv_s = tl.reshape(1.0 / output_s, (N_GROUPS, 1)) output_q_2d = tl.clamp(gate_up_2d * inv_s, fp8_min, fp8_max) output_q = tl.reshape(output_q_2d, (BLOCK_N,)).to(output_ptr.dtype.element_ty) tl.store(out_base + offs_in_d, output_q, mask=mask_d) # Pack 4 UE8M0 exponent bytes little-endian into one int32, matching deep_gemm's # get_mn_major_tma_aligned_packed_ue8m0_tensor (fp32>>23 -> uint8 -> 4-group int32 view). s_bits = output_s.to(tl.int32, bitcast=True) expo = (s_bits >> 23) & 0xFF shifts = tl.arange(0, N_GROUPS) * 8 packed = tl.sum(expo << shifts) scale_off = ( expert_id * stride_scale_e + hidden_dim_block_index * stride_scale_g4 + token_index * stride_scale_m ) tl.store(scale_ptr + scale_off, packed) def silu_and_mul_masked_post_quant_packed_fwd( input: torch.Tensor, output: torch.Tensor, output_scale_packed: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, num_real_tokens: int, topk: int, gemm1_alpha: float = 0.0, gemm1_clamp_limit: float = 0.0, ): assert input.is_contiguous() assert output.dtype == torch.float8_e4m3fn assert output.is_contiguous() assert input.dim() == 3 and input.shape[-1] % 2 == 0 E, m_max, _ = input.shape size_n = input.shape[-1] // 2 assert size_n % quant_group_size == 0 G = size_n // quant_group_size assert G % 4 == 0, "packed UE8M0 path requires num_groups % 4 == 0" BLOCK_N = quant_group_size * 4 assert size_n % BLOCK_N == 0, "packed UE8M0 path requires size_n % (4*group) == 0" hidden_dim_split = size_n // BLOCK_N assert tuple(output_scale_packed.shape) == (E, hidden_dim_split, m_max) assert output_scale_packed.dtype == torch.int32 finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max grid = (num_real_tokens * topk, hidden_dim_split) _silu_and_mul_post_quant_packed_kernel[grid]( input, *input.stride(), output, *output.stride(), output_scale_packed, *output_scale_packed.stride(), masked_m, E, size_n, fp8_max, -fp8_max, QUANT_GROUP_SIZE=quant_group_size, BLOCK_N=BLOCK_N, GEMM1_ALPHA=gemm1_alpha if gemm1_alpha is not None else 0.0, GEMM1_CLAMP_LIMIT=gemm1_clamp_limit if gemm1_clamp_limit is not None else 0.0, E_PADDED=triton.next_power_of_2(E), num_warps=1, ) return @triton.jit def _silu_and_mul_kernel( input_ptr, stride_input_0, stride_input_1, stride_input_2, output_ptr, stride_output_0, stride_output_1, stride_output_2, masked_m_ptr, size_n, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) hidden_dim_block_index = tl.program_id(0) block_num_per_expert = tl.num_programs(1) token_num_cur_expert = tl.load(masked_m_ptr + expert_id) stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64) stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64) stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d for token_index in tl.range( token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE ): gate = tl.load( input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0, ).to(tl.float32) up = tl.load( input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0, ).to(tl.float32) gate = gate / (1 + tl.exp(-gate)) gate_up = up * gate # Compute SiLU in fp32 for better precision, then cast back to the # input dtype. gate_up = gate_up.to(input_ptr.dtype.element_ty) tl.store( output_ptr_offs + token_index * stride_output_1, gate_up, mask=offs_in_d < size_n, ) def silu_and_mul_masked_fwd( input: torch.Tensor, output: torch.Tensor, masked_m: torch.Tensor, ): """ input shape [expert_num, token_num_padded, hidden_dim], dtype bf16 output shape [expert_num, token_num_padded, hidden_dim // 2], dtype bf16 masked_m shape [expert_num] """ assert input.is_contiguous() assert output.dtype == torch.bfloat16 assert input.dtype == torch.bfloat16 assert output.is_contiguous() assert len(input.shape) == 3 assert input.shape[0] == masked_m.shape[0] assert input.shape[-1] % 2 == 0 size_n = input.shape[-1] // 2 expert_num = len(masked_m) if expert_num < 4: BLOCK_NUM_PER_EXPERT = 64 else: BLOCK_NUM_PER_EXPERT = 32 BLOCK_N = 128 num_warps = 4 NUM_STAGES = 4 hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) grid = ( hidden_dim_split_block_num, BLOCK_NUM_PER_EXPERT, expert_num, ) _silu_and_mul_kernel[grid]( input, *input.stride(), output, *output.stride(), masked_m, size_n, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, num_warps=num_warps, ) return output @triton.jit def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe( input_ptr, output_ptr, scale_ptr, num_tokens_tensor_ptr, intermediate_size, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr, ): OutDtype = output_ptr.dtype.element_ty num_tokens = tl.load(num_tokens_tensor_ptr) numel = num_tokens * intermediate_size gate_ptr = input_ptr up_ptr = input_ptr + intermediate_size scale = 1.0 / tl.load(scale_ptr) start_idx = tl.program_id(0) * BLOCK_SIZE step = tl.num_programs(0) * BLOCK_SIZE for id in tl.range(start_idx, numel, step, num_stages=NUM_STAGES): ids = id + tl.arange(0, BLOCK_SIZE) token_ids = ids // intermediate_size mask = ids < numel offs = ids + token_ids * intermediate_size gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) output = gate / (1 + tl.exp(-gate)) * up * scale tl.store(output_ptr + ids, output.to(OutDtype), mask=mask) def silu_mul_static_tensorwise_quant_for_cutlass_moe( input: torch.Tensor, output: torch.Tensor, scale: torch.Tensor, num_tokens_tensor: torch.Tensor, expected_num_tokens: int, intermediate_size: int, ): grid, block_dim = _get_launch_config_1d( input.device, expected_num_tokens * intermediate_size ) silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe[grid]( input_ptr=input, output_ptr=output, scale_ptr=scale, num_tokens_tensor_ptr=num_tokens_tensor, intermediate_size=intermediate_size, BLOCK_SIZE=block_dim, NUM_STAGES=3, ) @triton.jit def post_reorder_triton_kernel_for_cutlass_moe( down_output_ptr, output_ptr, src2dst_ptr, topk_ids_ptr, topk_weights_ptr, num_local_experts, topk, num_tokens, hidden_size, routed_scaling_factor: float, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr, ): OutDtype = output_ptr.dtype.element_ty offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) mask = offset < hidden_size down_output_ptr_offs = down_output_ptr + offset output_ptr_offs = output_ptr + offset start_src_idx = tl.program_id(0) step = tl.num_programs(0) for src_idx_int32 in tl.range( start_src_idx, num_tokens, step, num_stages=NUM_STAGES ): src_idx = src_idx_int32.to(tl.int64) token_src2dst_ptr = src2dst_ptr + src_idx * topk token_topk_ids_ptr = topk_ids_ptr + src_idx * topk token_topk_weights_ptr = topk_weights_ptr + src_idx * topk sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for idx in range(topk): expert_id = tl.load(token_topk_ids_ptr + idx) if expert_id != num_local_experts: dst_idx_int32 = tl.load(token_src2dst_ptr + idx) dst_idx = dst_idx_int32.to(tl.int64) dst_idx = dst_idx weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32) load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32) sum_vec += in_data * weight_scale sum_vec *= routed_scaling_factor store_ptr_offs = output_ptr_offs + src_idx * hidden_size tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask) def post_reorder_for_cutlass_moe( down_output, output, src2dst, topk_ids, topk_weights, num_local_experts, topk, num_tokens, hidden_size, routed_scaling_factor: float, ): grid, block_dim = _get_launch_config_2d(down_output.device, num_tokens, hidden_size) post_reorder_triton_kernel_for_cutlass_moe[grid]( down_output_ptr=down_output, output_ptr=output, src2dst_ptr=src2dst, topk_ids_ptr=topk_ids, topk_weights_ptr=topk_weights, num_local_experts=num_local_experts, topk=topk, num_tokens=num_tokens, hidden_size=hidden_size, routed_scaling_factor=routed_scaling_factor, BLOCK_SIZE=block_dim, NUM_STAGES=3, ) @triton.jit def post_reorder_deepgemm_triton_kernel( down_output_ptr, output_ptr, src2dst_ptr, topk_ids_ptr, topk_weights_ptr, topk, num_tokens, hidden_size, routed_scaling_factor: float, BLOCK_SIZE: tl.constexpr, NUM_STAGES: tl.constexpr, ): """`expert_id >= 0` includes the shared expert at num_experts (padding=-1); don't switch to the cutlass `!= num_local_experts` gate. routed_scaling_factor is folded into the store. """ OutDtype = output_ptr.dtype.element_ty offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE) mask = offset < hidden_size down_output_ptr_offs = down_output_ptr + offset output_ptr_offs = output_ptr + offset start_src_idx = tl.program_id(0) step = tl.num_programs(0) for src_idx_int32 in tl.range( start_src_idx, num_tokens, step, num_stages=NUM_STAGES ): src_idx = src_idx_int32.to(tl.int64) token_src2dst_ptr = src2dst_ptr + src_idx * topk token_topk_ids_ptr = topk_ids_ptr + src_idx * topk token_topk_weights_ptr = topk_weights_ptr + src_idx * topk sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for idx in range(topk): expert_id = tl.load(token_topk_ids_ptr + idx) if expert_id >= 0: dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64) weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32) load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32) sum_vec += in_data * weight_scale sum_vec *= routed_scaling_factor store_ptr_offs = output_ptr_offs + src_idx * hidden_size tl.store(store_ptr_offs, sum_vec.to(OutDtype), mask=mask) def post_reorder_deepgemm( down_output, output, src2dst, topk_ids, topk_weights, topk, num_tokens, hidden_size, routed_scaling_factor: float, ): grid, block_dim = _get_launch_config_2d(down_output.device, num_tokens, hidden_size) post_reorder_deepgemm_triton_kernel[grid]( down_output, output, src2dst, topk_ids, topk_weights, topk, num_tokens, hidden_size, float(routed_scaling_factor), BLOCK_SIZE=block_dim, NUM_STAGES=3, ) @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_int32 = tl.program_id(0) src_idx = src_idx_int32.to(tl.int64) src2dst_ptr = src2dst_ptr + src_idx * topk topk_ids_ptr = topk_ids_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_offset in tl.range(0, hidden_size, BLOCK_SIZE): offset = start_offset + vec mask = offset < hidden_size sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for idx in range(topk): expert_id = tl.load(topk_ids_ptr + idx) if expert_id >= 0: dst_idx_int32 = tl.load(src2dst_ptr + idx) dst_idx = dst_idx_int32.to(tl.int64) weigh_scale = tl.load(topk_weights_ptr + idx).to(tl.float32) load_ptr = down_output_ptr + dst_idx * hidden_size # accumulate expert outputs in fp32 for better precision # before casting to the final output dtype. in_data = tl.load(load_ptr + offset, mask=mask).to(tl.float32) sum_vec += in_data * weigh_scale tl.store(store_ptr + offset, sum_vec.to(InDtype), mask=mask) @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, # Platform-specific semaphore for atomic_add performance tuning ATOMIC_ADD_SEM: tl.constexpr, IS_FP8: 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 index_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) mask_s = index_in_s < SCALE_HIDDEN_SIZE for token_id_int32 in range(start_token_id, total_token_num, grid_num): token_id = token_id_int32.to(tl.int64) to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) if IS_FP8: to_copy_s = tl.load( recv_x_scale + token_id * recv_x_scale_stride0 + index_in_s * recv_x_scale_stride1, mask=mask_s, ) for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4): topk_index = topk_idx_int32.to(tl.int64) expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) if expert_id >= 0: dest_token_index_int32 = tl.atomic_add( expert_start_loc + expert_id, 1, sem=ATOMIC_ADD_SEM ) dest_token_index = dest_token_index_int32.to(tl.int64) tl.store( output_index + token_id * output_index_stride0 + topk_index, dest_token_index_int32, ) output_tensor_ptr = ( output_tensor + dest_token_index * output_tensor_stride0 ) tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) if IS_FP8: output_tensor_scale_ptr = ( output_tensor_scale + dest_token_index * output_tensor_scale_stride0 ) tl.store( output_tensor_scale_ptr + index_in_s * output_tensor_scale_stride1, to_copy_s, mask=mask_s, ) # copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/deepep_scatter_gather.py @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, scale_ue8m0: bool = False, quant_block_size: int = 128, ): BLOCK_E = 128 # token num of per expert is aligned to 128 BLOCK_D = quant_block_size # block size of quantization num_warps = 8 num_experts = num_recv_tokens_per_expert.shape[0] hidden_size = recv_x.shape[1] # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts) grid = num_experts scale_hidden_size = hidden_size // BLOCK_D if scale_ue8m0: # ue8m0 scales are packed here (4 scales per int32), # hence the effective size of this dimension is divided by 4. scale_hidden_size = ceil_div(scale_hidden_size, 4) assert m_indices.shape[0] % BLOCK_E == 0 is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 if is_fp8: assert ( recv_x_scale.dtype == output_tensor_scale.dtype ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" assert ( recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size ) _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) if is_fp8 else 0, recv_x_scale.stride(1) if is_fp8 else 0, 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) if is_fp8 else 0, output_tensor_scale.stride(1) if is_fp8 else 0, 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=scale_hidden_size, SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), # XXX (MUSA): Atomic add with "relaxed" semaphore on musa backend for better performance ATOMIC_ADD_SEM=None if not _is_musa else "relaxed", IS_FP8=is_fp8, ) 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_int32 = tl.program_id(0) cur_block = cur_block_int32.to(tl.int64) start_cur_token_int32 = tl.program_id(1) grid_num = tl.num_programs(1) for cur_token_int32 in range(start_cur_token_int32, total_token_num, grid_num): cur_token = cur_token_int32.to(tl.int64) off_d = tl.arange(0, BLOCK_D) accumulator = tl.zeros([BLOCK_D], dtype=tl.float32) for topk_index_int32 in range(0, topk_num): topk_index = topk_index_int32.to(tl.int64) expert_id = tl.load( recv_topk_ids + cur_token * recv_topk_ids_stride0 + topk_index ) if expert_id >= 0: source_token_index_int32 = tl.load( input_index + cur_token * input_index_stride0 + topk_index ) source_token_index = source_token_index_int32.to(tl.int64) 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, ): num_warps = 2 num_tokens = output_tensor.shape[0] hidden_size = input_tensor.shape[1] BLOCK_D = 128 if hidden_size % 1024 != 0 else 1024 # block size of quantization 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 # copy from # https://github.com/deepseek-ai/DeepGEMM/blob/bd2a77552886b98c205af12f8d7d2d61247c4b27/deep_gemm/jit_kernels/utils.py#L58 def get_tma_aligned_size(x: int, element_size: int) -> int: """ Global memory address of TMA must be 16-byte aligned. Since we use column-major layout for the LHS scaling tensor, the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes. Arguments: x: original M-axis shape of the LHS scaling tensor. element_size: element size of the LHS scaling tensor. Returns: M-axis shape of the LHS scaling tensor after padding. """ tma_alignment_bytes = 16 assert tma_alignment_bytes % element_size == 0 alignment = tma_alignment_bytes // element_size return ceil_div(x, alignment) * alignment @triton.jit def _tma_align_input_scale_kernel( input_scale_ptr, output_ptr, m, k_div_block_size, input_scale_stride_m, input_scale_stride_k, output_stride_m, output_stride_k, BLOCK_SIZE_K: tl.constexpr, ): pid_m = tl.program_id(axis=0) grid_m = tl.num_programs(0) k_offsets = tl.arange(0, BLOCK_SIZE_K) for m_base in range(pid_m, m, grid_m): input_offset = ( input_scale_ptr + m_base * input_scale_stride_m + k_offsets * input_scale_stride_k ) input_data = tl.load(input_offset, mask=k_offsets < k_div_block_size) output_offset = ( output_ptr + k_offsets * output_stride_k + m_base * output_stride_m ) tl.store(output_offset, input_data, mask=k_offsets < k_div_block_size) # copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/quantization/triton_quant/fp8/fp8act_quant_kernel.py def tma_align_input_scale(input_scale: torch.Tensor): assert input_scale.dim() == 2 m, k_div_block_size = input_scale.shape padd_m = get_tma_aligned_size(m, input_scale.element_size()) output = torch.empty( (k_div_block_size, padd_m), dtype=input_scale.dtype, device=input_scale.device ) grid_m = min(m, 8192) BLOCK_SIZE_K = triton.next_power_of_2(k_div_block_size) _tma_align_input_scale_kernel[(grid_m,)]( input_scale_ptr=input_scale, output_ptr=output, m=m, k_div_block_size=k_div_block_size, input_scale_stride_m=input_scale.stride(0), input_scale_stride_k=input_scale.stride(1), output_stride_m=output.stride(1), # Note: these are swapped output_stride_k=output.stride(0), # for column-major BLOCK_SIZE_K=BLOCK_SIZE_K, ) return output.t()[:m] @triton.jit def fused_moe_dispatch_index_triton_kernel( topk_ids_ptr, # flat (num_toks,) int32; -1 = padding (drives the `expert >= 0` gate) src2dst_ptr, masked_m_ptr, m_max, num_toks, num_experts, BLOCK_SIZE: tl.constexpr, ZERO_INIT: tl.constexpr, ): # Each token picks top_k distinct experts, so per-expert count <= num_tokens < m_max; # dst = expert*m_max + offset never spills into the next expert's region. pid = tl.program_id(0) if ZERO_INIT: # Zero the cursor in-kernel and barrier before any atomic_add (single-block path). e_off = tl.arange(0, BLOCK_SIZE) tl.store( masked_m_ptr + e_off, tl.zeros((BLOCK_SIZE,), dtype=tl.int32), mask=e_off < num_experts, ) tl.debug_barrier() offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offs < num_toks expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) valid = mask & (expert >= 0) # Clamp masked lanes to bin 0 so the masked atomic's pointer stays in-bounds. expert_safe = tl.where(valid, expert, 0) offset = tl.atomic_add(masked_m_ptr + expert_safe, 1, mask=valid) dst = expert_safe * m_max + offset tl.store(src2dst_ptr + offs, dst, mask=valid) def fused_moe_dispatch_index( topk_ids: torch.Tensor, num_local_experts: int, m_max: int, ) -> Tuple[torch.Tensor, torch.Tensor]: num_toks = topk_ids.numel() src2dst = torch.empty(num_toks, device=topk_ids.device, dtype=torch.int32) # masked_m doubles as the atomic cursor and the final per-expert count; must be zeroed before any atomic_add. single_block = max(num_toks, num_local_experts) <= 1024 if single_block: BLOCK_SIZE = triton.next_power_of_2(max(num_toks, num_local_experts)) masked_m = torch.empty( num_local_experts, device=topk_ids.device, dtype=torch.int32 ) grid = (1,) else: BLOCK_SIZE = 256 masked_m = torch.zeros( num_local_experts, device=topk_ids.device, dtype=torch.int32 ) grid = (triton.cdiv(num_toks, BLOCK_SIZE),) fused_moe_dispatch_index_triton_kernel[grid]( topk_ids.view(-1), src2dst, masked_m, m_max, num_toks, num_local_experts, BLOCK_SIZE=BLOCK_SIZE, ZERO_INIT=single_block, ) return masked_m, src2dst @triton.jit def fill_gateup_input_triton_kernel( input_ptr, scale_ptr, gateup_input_ptr, gateup_input_scale_ptr, src2dst_ptr, topk_ids_ptr, topk, hidden_size, scale_size, m_max, scale_row_stride, scale_col_stride, BLOCK_SIZE: tl.constexpr, IS_FP8: tl.constexpr = True, SCALE_MN_MAJOR: tl.constexpr = False, ): src_idx_int32 = tl.program_id(0) src_idx = src_idx_int32.to(tl.int64) src2dst_ptr = src2dst_ptr + src_idx * topk topk_ids_ptr = topk_ids_ptr + src_idx * topk src_ptr = input_ptr + src_idx * hidden_size if IS_FP8: scale_src_ptr = scale_ptr + src_idx * scale_row_stride vec = tl.arange(0, BLOCK_SIZE) for idx in range(topk): expert_id = tl.load(topk_ids_ptr + idx) if expert_id >= 0: dst_idx_int32 = tl.load(src2dst_ptr + idx) dst_idx = dst_idx_int32.to(tl.int64) dst_ptr = gateup_input_ptr + dst_idx * hidden_size for start_offset in tl.range(0, hidden_size, BLOCK_SIZE): offset = start_offset + vec mask = offset < hidden_size in_data = tl.load(src_ptr + offset, mask=mask) tl.store(dst_ptr + offset, in_data, mask=mask) if IS_FP8: if SCALE_MN_MAJOR: expert = dst_idx // m_max m = dst_idx % m_max scale_dst_ptr = ( gateup_input_scale_ptr + expert * scale_size * m_max + m ) for start_offset in tl.range(0, scale_size, BLOCK_SIZE): offset = start_offset + vec mask = offset < scale_size in_scale = tl.load( scale_src_ptr + offset * scale_col_stride, mask=mask ) tl.store(scale_dst_ptr + offset * m_max, in_scale, mask=mask) else: scale_dst_ptr = gateup_input_scale_ptr + dst_idx * scale_size for start_offset in tl.range(0, scale_size, BLOCK_SIZE): offset = start_offset + vec mask = offset < scale_size in_scale = tl.load( scale_src_ptr + offset * scale_col_stride, mask=mask ) tl.store(scale_dst_ptr + offset, in_scale, mask=mask) def moe_ep_deepgemm_preprocess( topk_ids: torch.Tensor, num_local_experts: int, hidden_states: torch.Tensor, top_k: int, block_shape, output_dtype: torch.dtype = torch.float8_e4m3fn, use_mxfp8: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165 m_max = (hidden_states.size(0) // 256 + 1) * 256 expected_m = (topk_ids.numel() - 1) // num_local_experts + 1 masked_m, src2dst = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max) gateup_input = torch.empty( (num_local_experts, m_max, hidden_states.size(1)), device=hidden_states.device, dtype=output_dtype, ) if block_shape is None: block_shape = [128, 128] assert len(block_shape) == 2 block_n, block_k = block_shape[0], block_shape[1] is_fp8 = output_dtype == torch.float8_e4m3fn if is_fp8 and use_mxfp8: from sglang.jit_kernel.minimax_quant_ue8m0 import ( per_token_quant_fp8_ue8m0_scatter, ) num_groups = hidden_states.size(1) // block_k gateup_input_scale = torch.empty( (gateup_input.size(0), num_groups // 4, m_max), device=hidden_states.device, dtype=torch.int32, ) per_token_quant_fp8_ue8m0_scatter( hidden_states, gateup_input, gateup_input_scale, src2dst, topk_ids, top_k, m_max, group_size=block_k, ) gateup_input_scale = gateup_input_scale.transpose(1, 2) elif is_fp8: hidden_states, scale = per_token_group_quant_fp8(hidden_states, block_k) gateup_input_scale = torch.empty( (gateup_input.size(0), gateup_input.size(1), scale.size(1)), device=hidden_states.device, dtype=scale.dtype, ) fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( hidden_states, scale, gateup_input, gateup_input_scale, src2dst, topk_ids, top_k, hidden_states.size(1), scale.size(1), m_max, scale.stride(0), scale.stride(1), BLOCK_SIZE=1024, IS_FP8=True, SCALE_MN_MAJOR=False, ) else: scale = None gateup_input_scale = None fill_gateup_input_triton_kernel[(hidden_states.shape[0],)]( hidden_states, scale, gateup_input, gateup_input_scale, src2dst, topk_ids, top_k, hidden_states.size(1), 0, m_max, 0, 0, BLOCK_SIZE=1024, IS_FP8=False, SCALE_MN_MAJOR=False, ) return ( masked_m, expected_m, src2dst, gateup_input, gateup_input_scale, ) @triton.jit def compute_identity_kernel( top_k, hidden_states_ptr, expert_scales_ptr, num_tokens, output_ptr, hidden_dim, scales_stride, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) batch_id = pid // (hidden_dim // BLOCK_SIZE) dim_offset = pid % (hidden_dim // BLOCK_SIZE) * BLOCK_SIZE if batch_id >= num_tokens or dim_offset >= hidden_dim: return h = tl.load( hidden_states_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, ) result = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for i in range(top_k): scale = tl.load(expert_scales_ptr + batch_id * scales_stride + i) result += h * scale tl.store( output_ptr + batch_id * hidden_dim + dim_offset + tl.arange(0, BLOCK_SIZE), result, mask=(dim_offset + tl.arange(0, BLOCK_SIZE)) < hidden_dim, ) def zero_experts_compute_triton( expert_indices, expert_scales, num_experts, zero_expert_type, hidden_states ): N = expert_indices.numel() top_k = expert_indices.size(-1) grid = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),) if zero_expert_type == "identity": zero_expert_mask = expert_indices < num_experts zero_expert_scales = expert_scales.clone() zero_expert_scales[zero_expert_mask] = 0.0 normal_expert_mask = expert_indices >= num_experts # Keep a valid routed-expert id for MoE kernels that do not accept negative # ids. The zero scale below still removes the routed-expert contribution. expert_indices[normal_expert_mask] = 0 expert_scales[normal_expert_mask] = 0.0 output = torch.zeros_like(hidden_states).to(hidden_states.device) hidden_dim = hidden_states.size(-1) num_tokens = hidden_states.size(0) grid = lambda meta: (num_tokens * (hidden_dim // meta["BLOCK_SIZE"]),) compute_identity_kernel[grid]( top_k, hidden_states, zero_expert_scales, num_tokens, output, hidden_dim, zero_expert_scales.stride(0), BLOCK_SIZE=256, ) return output @triton.jit def compute_problem_sizes_w4a8_kernel( masked_m_ptr, problem_sizes1_ptr, problem_sizes2_ptr, n, k, num_experts, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = pid < num_experts final_occurrences = tl.load(masked_m_ptr + pid, mask=mask, other=0) ps1_idx_0 = pid * 3 ps1_idx_1 = ps1_idx_0 + 1 ps1_idx_2 = ps1_idx_0 + 2 ps2_idx_0 = pid * 3 ps2_idx_1 = ps2_idx_0 + 1 ps2_idx_2 = ps2_idx_0 + 2 ps1_mask_0 = ps1_idx_0 < num_experts * 3 ps1_mask_1 = ps1_idx_1 < num_experts * 3 ps1_mask_2 = ps1_idx_2 < num_experts * 3 ps2_mask_0 = ps2_idx_0 < num_experts * 3 ps2_mask_1 = ps2_idx_1 < num_experts * 3 ps2_mask_2 = ps2_idx_2 < num_experts * 3 tl.store(problem_sizes1_ptr + ps1_idx_0, 2 * n, mask=ps1_mask_0) tl.store(problem_sizes1_ptr + ps1_idx_1, final_occurrences, mask=ps1_mask_1) tl.store(problem_sizes1_ptr + ps1_idx_2, k, mask=ps1_mask_2) tl.store(problem_sizes2_ptr + ps2_idx_0, k, mask=ps2_mask_0) tl.store(problem_sizes2_ptr + ps2_idx_1, final_occurrences, mask=ps2_mask_1) tl.store(problem_sizes2_ptr + ps2_idx_2, n, mask=ps2_mask_2) def compute_problem_sizes_w4a8( masked_m, problem_sizes1, problem_sizes2, n, k, num_experts ): BLOCK_SIZE = 256 grid = lambda meta: (triton.cdiv(num_experts, meta["BLOCK_SIZE"]),) compute_problem_sizes_w4a8_kernel[grid]( masked_m, problem_sizes1, problem_sizes2, n, k, num_experts, BLOCK_SIZE=BLOCK_SIZE, ) return problem_sizes1, problem_sizes2 def deepep_ll_get_cutlass_w4a8_moe_mm_data( masked_m, problem_sizes1, problem_sizes2, num_experts, n, k, ): problem_sizes1, problem_sizes2 = compute_problem_sizes_w4a8( masked_m, problem_sizes1, problem_sizes2, n, k, num_experts ) return ( problem_sizes1.to(torch.int32), problem_sizes2.to(torch.int32), ) @triton.jit def _silu_and_mul_post_per_tensor_quant_kernel( input_ptr, stride_input_expert, stride_input_token, stride_input_dim, output_ptr, stride_output_expert, stride_output_token, stride_output_dim, scale_ptr, masked_m_ptr, inner_dim, fp8_max, fp8_min, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, ): """ Triton kernel: fused SiLU(gate) * up + per-tensor FP8 quantization. Shape: input: [E, T_padded, 2*D] -> gate: [:,:,D], up: [:,:,D] output: [E, T_padded, D], dtype=float8_e4m3fn """ expert_id = tl.program_id(2) block_id_token = tl.program_id(1) block_id_dim = tl.program_id(0) num_token_blocks = tl.num_programs(1) token_num_cur_expert = tl.load(masked_m_ptr + expert_id) scale = 1.0 / tl.load(scale_ptr).to(tl.float32) stride_input_expert = tl.cast(stride_input_expert, tl.int32) stride_output_expert = tl.cast(stride_output_expert, tl.int32) stride_input_token = tl.cast(stride_input_token, tl.int32) stride_output_token = tl.cast(stride_output_token, tl.int32) offset_d = block_id_dim * BLOCK_N + tl.arange(0, BLOCK_N) mask_d = offset_d < inner_dim # base pointers for current expert and dim block input_base_offs = input_ptr + expert_id * stride_input_expert + offset_d output_base_offs = output_ptr + expert_id * stride_output_expert + offset_d for token_idx in tl.range( block_id_token, token_num_cur_expert, num_token_blocks, num_stages=NUM_STAGE ): gate_ptr = input_base_offs + token_idx * stride_input_token up_ptr = gate_ptr + inner_dim gate = tl.load(gate_ptr, mask=mask_d, other=0.0).to(tl.float32) up = tl.load(up_ptr, mask=mask_d, other=0.0).to(tl.float32) # SiLU: x * sigmoid(x) gate = gate / (1 + tl.exp(-gate)) gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate scaled = gate_up * scale output_q = tl.clamp(scaled, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) out_ptr = output_base_offs + token_idx * stride_output_token tl.store(out_ptr, output_q, mask=mask_d) def silu_and_mul_masked_post_per_tensor_quant_fwd( input: torch.Tensor, output: torch.Tensor, masked_m: torch.Tensor, scale: torch.Tensor, ) -> torch.Tensor: """ Fused SiLU + Mul + Per-Tensor Quantization to FP8. Args: input: [expert_num, token_num_padded, 2 * inner_dim] output: [expert_num, token_num_padded, inner_dim], dtype=torch.float8_e4m3fn masked_m: [expert_num], actual token count for each expert scale: [1] or [expert_num], quantization scale (per-tensor or per-expert) Returns: output tensor """ assert input.is_contiguous() assert output.is_contiguous() assert output.dtype == torch.float8_e4m3fn assert input.ndim == 3 assert input.shape[0] == masked_m.shape[0] assert input.shape[-1] % 2 == 0 assert scale.numel() == 1 or scale.shape[0] == input.shape[0] expert_num = input.shape[0] # 3584 inner_dim = input.shape[-1] // 2 BLOCK_N = 256 BLOCK_M = 64 if expert_num < 4 else 32 NUM_STAGES = 3 hidden_dim_split_block_num = triton.cdiv(inner_dim, BLOCK_N) grid = (hidden_dim_split_block_num, BLOCK_M, expert_num) finfo = torch.finfo(torch.float8_e4m3fn) fp8_max = finfo.max fp8_min = -fp8_max _silu_and_mul_post_per_tensor_quant_kernel[grid]( input, *input.stride(), output, *output.stride(), scale, masked_m, inner_dim, fp8_max, fp8_min, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, ) return output @triton.jit def _fp8_per_token_quant_to_per_tensor_quant_kernel( x_ptr, x_scale_ptr, x_scale_stride0, x_scale_stride1, x_scale_stride2, masked_m_ptr, output_scale_ptr, output_ptr, m, k, K_SCALE_BLOCK_SIZE: tl.constexpr, K_BLOCK_SIZE: tl.constexpr, ): pid_k, pid_m, pid_e = ( tl.program_id(axis=0), tl.program_id(axis=1), tl.program_id(axis=2), ) pid_m_dim = tl.num_programs(1) token_id = pid_m last_effective_id = tl.load(masked_m_ptr + pid_e) if token_id >= last_effective_id: return output_scale_val_inv = 1.0 / tl.load(output_scale_ptr).to(tl.float32) k_offsets = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE) scale_offsets = (k_offsets // K_SCALE_BLOCK_SIZE) * x_scale_stride2 x_ptrs = x_ptr + pid_e * m * k + k_offsets output_ptrs = output_ptr + pid_e * m * k + k_offsets x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim): hidden = tl.load(x_ptrs + tok_idx * k).to(tl.float32) scale_fp32 = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1).to(tl.float32) hidden = hidden * scale_fp32 * output_scale_val_inv tl.store(output_ptrs + tok_idx * k, hidden.to(output_ptr.dtype.element_ty)) def fp8_per_token_to_per_tensor_quant_triton( x: torch.Tensor, x_scale: torch.Tensor, masked_m: torch.Tensor, output_scale: torch.Tensor, output: torch.Tensor, ): K_SCALE_BLOCK_SIZE = 128 assert len(x.shape) == 3 and x.size(2) % K_SCALE_BLOCK_SIZE == 0 assert x.is_contiguous() assert x_scale.size(2) == x.size(2) // K_SCALE_BLOCK_SIZE assert output_scale.numel() == 1 K_BLOCK_SIZE = 1024 assert x.size(2) % K_BLOCK_SIZE == 0 grid = (x.size(2) // K_BLOCK_SIZE, 32, x.size(0)) _fp8_per_token_quant_to_per_tensor_quant_kernel[grid]( x, x_scale, *x_scale.stride(), masked_m, output_scale, output, x.size(1), x.size(2), K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE, K_BLOCK_SIZE=K_BLOCK_SIZE, num_warps=8, )