Uploaded using `kernel-builder`.

2fe9f62 verified about 1 month ago

6.42 kB

	# ********************************************************************************
	# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
	# ********************************************************************************

	import cuda.bindings.driver as cuda
	import cutlass.cute as cute
	import torch
	import triton
	import triton.language as tl
	from cutlass.cute.runtime import from_dlpack
	from ..quack.cute_dsl_utils import torch2cute_dtype_map
	from ..quack.gemm_interface import gemm, gemm_gated

	from .._ops_compat import add_op_namespace_prefix
	from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
	from .topk import Softmax_Over_TopK, TopK_Over_Softmax


	@torch.library.custom_op(add_op_namespace_prefix("_topk_fwd"), mutates_args={"values", "indices"})
	def _topk_fwd(
	x: torch.Tensor,
	k: int,
	values: torch.Tensor,
	indices: torch.Tensor,
	is_softmax_over_topk: bool,
	norm_topk_probs: bool,
	) -> None:
	"""Top-k forward pass.
	Args:
	x: Input tensor of shape (M, N)
	k: Number of top elements to return
	Returns:
	Tuple of (values tensor of shape (M, k), indices tensor of shape (M, k))
	"""
	N = x.size(1)

	input_dtype = torch2cute_dtype_map[x.dtype]
	output_dtype = torch2cute_dtype_map[values.dtype]
	convert_from_dlpack = lambda tensor: (
	from_dlpack(tensor.detach(), assumed_align=16).mark_compact_shape_dynamic(mode=0, stride_order=(0, 1))
	)

	x_tensor, values_tensor, indices_tensor = [convert_from_dlpack(tensor) for tensor in (x, values, indices)]
	current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
	if is_softmax_over_topk:
	compile_key = (input_dtype, output_dtype, N, k, True)
	else:
	compile_key = (input_dtype, output_dtype, N, k, False, norm_topk_probs)

	if compile_key not in _topk_fwd.compile_cache:
	if is_softmax_over_topk:
	topk_op = Softmax_Over_TopK(input_dtype, output_dtype, N, k)
	else:
	topk_op = TopK_Over_Softmax(input_dtype, output_dtype, N, k, norm_topk_probs)

	_topk_fwd.compile_cache[compile_key] = cute.compile(
	topk_op, x_tensor, values_tensor, indices_tensor, current_stream
	)
	_topk_fwd.compile_cache[compile_key](x_tensor, values_tensor, indices_tensor, current_stream)


	_topk_fwd.compile_cache = {}


	@torch.library.custom_op(add_op_namespace_prefix("_up_projection_forward"), mutates_args={"h", "a"})
	def _up_projection_forward(
	x: torch.Tensor,
	w1: torch.Tensor,
	h: torch.Tensor,
	a: torch.Tensor,
	b1: torch.Tensor \| None,
	expert_frequency_offset: torch.Tensor,
	x_gather_idx: torch.Tensor,
	activation_type: str,
	is_inference_mode_enabled: bool = False,
	concat_layout: bool = False,
	) -> None:
	assert activation_type in (
	"swiglu",
	"geglu",
	), f"QuACK gemm_gated only supports glu activations, got {activation_type}"
	gemm_gated(
	x,
	w1.permute(2, 1, 0),
	activation=activation_type,
	cu_seqlens_m=expert_frequency_offset,
	A_idx=x_gather_idx,
	preact_out=h,
	postact_out=a,
	store_preact=(not is_inference_mode_enabled),
	bias=b1,
	concat_layout=(("B", "bias") if b1 is not None else ("B",)) if concat_layout else None,
	)


	_up_projection_forward.compile_cache = {}


	@torch.library.custom_op(add_op_namespace_prefix("_down_projection_forward"), mutates_args={"y"})
	def _down_projection_forward(
	w2: torch.Tensor,
	a: torch.Tensor,
	y: torch.Tensor,
	b2: torch.Tensor \| None,
	expert_frequency_offset: torch.Tensor,
	) -> None:
	gemm(a, w2.permute(2, 1, 0), out=y, cu_seqlens_m=expert_frequency_offset, bias=b2)


	_down_projection_forward.compile_cache = {}


	@torch.library.custom_op(add_op_namespace_prefix("_router_forward"), mutates_args={"o"})
	def _router_forward(
	y: torch.Tensor,
	o: torch.Tensor,
	topk_scores: torch.Tensor,
	s_reverse_scatter_idx: torch.Tensor,
	num_activated_expert_per_token_offset: torch.Tensor,
	varlen_K_max: int,
	H: int,
	is_varlen_K: bool,
	) -> None:
	token_gather_and_sum_varlen_K_triton(
	y,
	topk_scores,
	o,
	s_reverse_scatter_idx,
	num_activated_expert_per_token_offset,
	o.size(0),
	varlen_K_max,
	H,
	is_varlen_K,
	)


	@triton.jit
	def _softmax_fwd_small_kernel(
	logits_ptr, stride_lm: tl.constexpr, stride_ln: tl.constexpr, K: tl.constexpr, BLOCK_K: tl.constexpr
	):
	row = tl.program_id(axis=0)

	# tl.assume(K <= BLOCK_K)
	k_offs = tl.arange(0, BLOCK_K)
	k_mask = k_offs < K

	# load full row (all columns) in one go (N is small)
	x = tl.load(logits_ptr + row * stride_lm + k_offs * stride_ln, mask=k_mask, other=-float("inf")).to(tl.float32)
	x = x - tl.max(x, axis=0)
	ex = tl.exp(x)
	y = ex / tl.sum(ex, axis=0)

	tl.store(logits_ptr + row * stride_lm + k_offs * stride_ln, y, mask=k_mask)


	@torch.library.custom_op(
	add_op_namespace_prefix("_softmax_topk_fwd"), mutates_args={"topk_router_score", "topk_router_indices"}
	)
	def _topk_softmax_fwd(
	router_logits: torch.Tensor,
	topk_router_score: torch.Tensor,
	topk_router_indices: torch.Tensor,
	E: int,
	K: int,
	is_softmax_over_topk: bool,
	norm_topk_probs: bool,
	) -> None:
	if E <= 4096 and K <= 16 and E % 8 == 0:
	_topk_fwd(
	router_logits,
	K,
	topk_router_score,
	topk_router_indices,
	is_softmax_over_topk=is_softmax_over_topk,
	norm_topk_probs=norm_topk_probs,
	)
	else:
	if is_softmax_over_topk:
	topk_results = router_logits.topk(K, dim=-1)
	vals = topk_results.values.softmax(dim=-1, dtype=torch.float32)
	topk_router_score.copy_(vals.to(topk_router_score.dtype))
	topk_router_indices.copy_(topk_results.indices.to(topk_router_indices.dtype))
	else:
	probs = router_logits.softmax(dim=-1, dtype=torch.float32)
	topk_results = probs.topk(K, dim=-1)
	vals = topk_results.values
	if norm_topk_probs:
	vals = vals / vals.sum(dim=-1, keepdim=True)
	topk_router_score.copy_(vals.to(topk_router_score.dtype))
	topk_router_indices.copy_(topk_results.indices.to(topk_router_indices.dtype))