[Feature Request] - Supporting MegaKernel for DeepSeek Models

Feature Description

All unique layer types needed for the DeepSeek V3 MPK demo, regardless of where/how many times they appear.

Layer	Status	Signature	Data Types	Notes
embed_layer	✅	token_ids, embed_table -> embeddings	INT64 -> BF16
rmsnorm_layer	✅	activations, norm_weight -> normalized	BF16 -> BF16
linear_layer	✅	input, weight -> output	BF16 -> BF16	Only used for routing gate (BF16) and lm_head
silu_mul_layer	✅	input [B, 2K] -> output [B, K]	BF16 -> BF16	Shared expert and dense MLP activation
moe_topk_sigmoid_routing_layer	✅	logits + bias -> (topk_weight, routing_indices, masks)	BF16 -> FP32/INT32	Group-aware sigmoid routing
moe_w13_fp8_layer	✅	input_fp8, input_scale, weight_fp8, weight_scale, routing_indices, masks -> output	FP8+FP32 -> BF16	Sparse expert gate+up proj
moe_silu_mul_layer	✅	input [B, K, 2M] -> output [B, K, M]	BF16 -> BF16	Sparse expert activation
moe_w2_fp8_layer	✅	input_fp8, input_scale, weight_fp8, weight_scale, routing_indices, masks -> output	FP8+FP32 -> BF16	Sparse expert down proj
allreduce_layer	✅	input, buffer -> output	BF16 -> BF16	TP8 all-reduce
argmax_partial_layer	✅	logits -> (partial_values, partial_indices)	BF16 -> BF16/INT64
argmax_reduce_layer	✅	(partial_values, partial_indices) -> token_id	BF16/INT64 -> INT64
moe_local_reduction	❌	expert_outputs, topk_weights, residual -> output	BF16 -> BF16	Weighted sum + residual
bf16_to_fp8_quantize	❌	input BF16 -> output FP8 E4M3 + scale ue8m0	BF16 -> FP8+FP32	External branch. Block-wise (block_size=128) quantize before every FP8 GEMM. ~6x per MoE layer, ~2x per dense layer.
fp8_linear	❌	input_fp8, input_scale, weight_fp8, weight_scale -> output	FP8+FP32 -> BF16	External branch. Dense FP8 GEMM for MLA projections (q_a, kv_a, q_b, kv_b, o_proj), shared expert, dense MLP.
fp8_linear_with_residual	❌	input_fp8, input_scale, weight_fp8, weight_scale, residual -> output	FP8+FP32+BF16 -> BF16	External branch. FP8 GEMM + residual add for o_proj and dense/shared down_proj.
mla_layer	❌	q_nope, q_rope, latent_cache, rope_cache, page_table -> attn_output	BF16 -> BF16	C++ kernel mla_sm100_2sm.cuh exists, needs Python wrapper + task registration.

---

Layer Tensor Details

Concrete tensor shapes for each layer. B = max_num_batched_tokens. Assumes TP8 (8-way tensor parallelism), EP2 within TP8 for sparse experts (128 local experts per GPU).

embed_layer

Tensor	Role	Shape	Dtype
input	input	[B, 1]	INT64
weight	input	[vocab_size, hidden_size] = [129280, 7168]	BF16
output	output	[B, hidden_size] = [B, 7168]	BF16

rmsnorm_layer

Tensor	Role	Shape	Dtype
input	input	[B, D]	BF16
weight	input	[D]	BF16
output	output	[B, D]	BF16

Used with D = 7168 (input/post-attn layernorm), 1536 (q_a_layernorm), 512 (kv_a_layernorm).

linear_layer

Tensor	Role	Shape	Dtype
input	input	[B, K]	BF16
weight	input	[N, K]	BF16
output	output	[B, N]	BF16

Used for: routing gate (K=7168, N=256), lm_head (K=7168, N=vocab_padded).

silu_mul_layer

Tensor	Role	Shape	Dtype
input	input	[B, 2*intermediate_size]	BF16
output	output	[B, intermediate_size]	BF16

Used for: shared expert (intermediate_size = 2048/TP = 256 per GPU), dense MLP layers 0-2 (intermediate_size = 18432/TP = 2304 per GPU).

moe_topk_sigmoid_routing_layer

Tensor	Role	Shape	Dtype
input	input	[B, n_routed_experts] = [B, 256]	BF16
bias	input	[n_routed_experts] = [256]	BF16
topk_weight	output	[B, num_experts_per_tok] = [B, 8]	FP32
routing_indices	output	[num_experts_per_tok, B] = [8, B]	INT32
masks	output	[n_routed_experts + 1] = [257]	INT32

Params: num_groups=8, topk_group=4, routed_scaling_factor=2.5.

moe_w13_fp8_layer

Tensor	Role	Shape	Dtype
input_fp8	input	[B, hidden_size] = [B, 7168]	FP8 E4M3
input_scale	input	[B, hidden_size/128] = [B, 56]	FP32
weight_fp8	input	[n_local_experts, 2*moe_inter, hidden_size] = [128, 4096, 7168]	FP8 E4M3
weight_scale	input	[n_local_experts, 2*moe_inter, hidden_size/128] = [128, 4096, 56]	FP32
routing_indices	input	[num_experts_per_tok, B] = [8, B]	INT32
masks	input	[n_routed_experts + 1] = [257]	INT32
output	output	[B, num_experts_per_tok, 2*moe_inter] = [B, 8, 4096]	BF16

moe_silu_mul_layer

Tensor	Role	Shape	Dtype
input	input	[B, num_experts_per_tok, 2*moe_inter] = [B, 8, 4096]	BF16
output	output	[B, num_experts_per_tok, moe_inter] = [B, 8, 2048]	BF16

moe_w2_fp8_layer

Tensor	Role	Shape	Dtype
input_fp8	input	[B, num_experts_per_tok, moe_inter] = [B, 8, 2048]	FP8 E4M3
input_scale	input	[B, num_experts_per_tok, moe_inter/128] = [B, 8, 16]	FP32
weight_fp8	input	[n_local_experts, hidden_size, moe_inter] = [128, 7168, 2048]	FP8 E4M3
weight_scale	input	[n_local_experts, hidden_size, moe_inter/128] = [128, 7168, 16]	FP32
routing_indices	input	[num_experts_per_tok, B] = [8, B]	INT32
masks	input	[n_routed_experts + 1] = [257]	INT32
output	output	[B, num_experts_per_tok, hidden_size] = [B, 8, 7168]	BF16

allreduce_layer

Tensor	Role	Shape	Dtype
input	input	[B, hidden_size] = [B, 7168]	BF16
buffer	input	[world_size, B, hidden_size] = [8, B, 7168]	BF16
output	output	[B, hidden_size] = [B, 7168]	BF16

argmax_partial_layer

Tensor	Role	Shape	Dtype
input	input	[B, vocab_padded]	BF16
output_value	output	[B, num_workers]	BF16
output_index	output	[B, num_workers]	INT64

argmax_reduce_layer

Tensor	Role	Shape	Dtype
input_value	input	[B, num_workers]	BF16
input_index	input	[B, num_workers]	INT64
output	output	[B, 1]	INT64

moe_local_reduction ❌

Tensor	Role	Shape	Dtype
input	input	[B, num_experts_per_tok, hidden_size] = [B, 8, 7168]	BF16
weight	input	[B, num_experts_per_tok] = [B, 8]	FP32
residual	input	[B, hidden_size] = [B, 7168]	BF16
output	output	[B, hidden_size] = [B, 7168]	BF16

Computes output = sum_k(input[:, k, :] * weight[:, k]) + residual.

bf16_to_fp8_quantize ❌

Tensor	Role	Shape	Dtype
input	input	[..., K] (2D or 3D)	BF16
output_fp8	output	[..., K] (same shape as input)	FP8 E4M3
output_scale	output	[..., K/128]	ue8m0

Block-wise quantize with block_size=128. Used before every FP8 GEMM. Example shapes: [B, 7168], [B, 1536], [B, 512], [B, 8, 2048].

fp8_linear ❌

Tensor	Role	Shape	Dtype
input_fp8	input	[B, K]	FP8 E4M3
input_scale	input	[B, K/32]	ue8m0
weight_fp8	input	[N, K]	FP8 E4M3
weight_scale	input	[N, K/128]	FP32
output	output	[B, N]	BF16

Used for all dense (non-MoE) FP8 GEMMs. Example (K, N) pairs: fused q_a+kv_a (7168, 2112), q_b (1536, 3072), kv_b (512, 4096), o_proj (2048, 7168), shared gate_up (7168, 512), shared down (256, 7168), dense gate_up (7168, 4608), dense down (2304, 7168).

fp8_linear_with_residual ❌

Tensor	Role	Shape	Dtype
input_fp8	input	[B, K]	FP8 E4M3
input_scale	input	[B, K/128]	FP32
weight_fp8	input	[N, K]	FP8 E4M3
weight_scale	input	[N, K/128]	FP32
residual	input	[B, N]	BF16
output	output	[B, N]	BF16

Computes output = matmul(input_fp8, weight_fp8.T) + residual. Used for o_proj and down_proj where residual connection follows.

mla_layer ❌

Tensor	Role	Shape	Dtype
q_nope	input	[B, num_local_heads, qk_nope_head_dim] = [B, 16, 128]	BF16
q_rope	input	[B, num_local_heads, qk_rope_head_dim] = [B, 16, 64]	BF16
latent_cache	input	[num_pages, page_size, kv_lora_rank] = [P, S, 512]	BF16
rope_cache	input	[num_pages, page_size, qk_rope_head_dim] = [P, S, 64]	BF16
output	output	[B, num_local_heads * v_head_dim] = [B, 2048]	BF16

C++ kernel mla_sm100_2sm.cuh exists. Two separate paged KV caches: latent (compressed, shared across heads) and rope (positional, shared across heads). Params: [num_q_heads, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, max_seq_len, page_size].

This task tracks the todos for supporting DeepSeek models in MK

• Layers

  • RMS Norm BF16
  • GEMM (MLA down proj & o proj FP8e4m3) @JudySun233
  • MLA @xinhaoc
  • Quantization from BF16 to FP8e4m3
  • MoE BMM FP8e4m3 (MoE w13 w2) @ZephyrZhao247
  • Expert parallelism We switch to Megatron-Style TP for now.

• Features

  • Quantization input
  • Weight loader to auto partition weights during launch, instead of AOT partitioning. @EH67
  • Deepseek Demo @ZephyrZhao247

• MTP

  • Eagle style speculative decoding, sequential attention