Add some fused elementwise kernels for grok-1 (#4398)

Co-authored-by: dhou-xai <dhou@x.ai>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>

python/sglang/srt/layers/elementwise.py

@@ -0,0 +1,411 @@
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+fused_softcap_autotune = triton.autotune(
+    configs=[
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 512}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 32768}, num_warps=32),
+    ],
+    key=["n_ele"],
+)
+
+
+@triton.jit
+def fused_softcap_kernel(
+    output_ptr,
+    input_ptr,
+    n_ele,
+    softcap_const: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_ele
+    x = tl.load(input_ptr + offsets, mask=mask)
+    fx = x.to(tl.float32)
+    fxs = fx / softcap_const
+    exped = tl.exp(2 * fxs)
+    top = exped - 1
+    bottom = exped + 1
+    output = top / bottom * softcap_const
+    tl.store(output_ptr + offsets, output, mask=mask)
+
+
+fused_softcap_kernel_autotuned = fused_softcap_autotune(fused_softcap_kernel)
+
+
+def fused_softcap(x, softcap_const, autotune=False):
+    output = torch.empty_like(x, dtype=torch.float32)
+    n_elements = output.numel()
+    if autotune:
+        grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+        fused_softcap_kernel_autotuned[grid](output, x, n_elements, softcap_const)
+    else:
+        fused_softcap_kernel[(triton.cdiv(n_elements, 128),)](
+            output, x, n_elements, softcap_const, BLOCK_SIZE=128, num_warps=8
+        )
+    return output
+
+
+# cast to float + softcap
+class Softcap:
+    def __init__(self, softcap_const: float):
+        self.softcap_const = softcap_const
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.is_cuda:
+            return self.forward_cuda(x)
+        else:
+            return self.forward_native(x)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.tanh(x.float() / self.softcap_const) * self.softcap_const
+
+    def forward_cuda(self, x: torch.Tensor, autotune=False) -> torch.Tensor:
+        return fused_softcap(x, self.softcap_const, autotune=autotune)
+
+
+rmsnorm_autotune = triton.autotune(
+    configs=[
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=4, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=8, num_stages=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 1024}, num_warps=16, num_stages=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 2048}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 4096}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 8192}, num_warps=32, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 16384}, num_warps=32, num_stages=4),
+    ],
+    key=["hidden_dim"],
+)
+
+
+@triton.jit
+def fused_dual_residual_rmsnorm_kernel(
+    output_ptr,
+    mid_ptr,
+    activ_ptr,
+    residual_ptr,
+    weight1_ptr,
+    weight2_ptr,
+    eps: tl.constexpr,
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    input_start = pid * hidden_dim
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < hidden_dim
+
+    a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0)
+    a = a_.to(tl.float32)
+    rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps)
+
+    r = tl.load(residual_ptr + input_start + offsets, mask=mask, other=0.0)
+    w1_ = tl.load(weight1_ptr + offsets, mask=mask, other=0.0)
+    w1 = w1_.to(tl.float32)
+
+    a2r = r + (a / rms * w1).to(r.dtype)
+    tl.store(
+        mid_ptr + input_start + offsets,
+        a2r,
+        mask=mask,
+    )
+
+    a2r = a2r.to(tl.float32)
+    rms2 = tl.sqrt(tl.sum(a2r * a2r, axis=0) / hidden_dim + eps)
+
+    w2_ = tl.load(weight2_ptr + offsets, mask=mask, other=0.0)
+    w2 = w2_.to(tl.float32)
+
+    tl.store(
+        output_ptr + input_start + offsets,
+        a2r / rms2 * w2,  # implicitly casts to output dtype here
+        mask=mask,
+    )
+
+
+fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune(
+    fused_dual_residual_rmsnorm_kernel
+)
+
+
+def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False):
+    assert len(x.shape) == 2
+    assert x.shape == residual.shape and x.dtype == residual.dtype
+    output, mid = torch.empty_like(x), torch.empty_like(x)
+    bs, hidden_dim = x.shape
+    if autotune:
+        fused_dual_residual_rmsnorm_kernel_autotune[(bs,)](
+            output, mid, x, residual, weight1, weight2, eps=eps, hidden_dim=hidden_dim
+        )
+    else:
+        config = {
+            "BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
+            "num_warps": max(
+                min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), 32), 4
+            ),
+        }
+
+        fused_dual_residual_rmsnorm_kernel[(bs,)](
+            output,
+            mid,
+            x,
+            residual,
+            weight1,
+            weight2,
+            eps=eps,
+            hidden_dim=hidden_dim,
+            **config,
+        )
+
+    return output, mid
+
+
+@triton.jit
+def fused_rmsnorm_kernel(
+    output_ptr,
+    activ_ptr,
+    weight_ptr,
+    eps: tl.constexpr,
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    input_start = pid * hidden_dim
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < hidden_dim
+
+    a_ = tl.load(activ_ptr + input_start + offsets, mask=mask, other=0.0)
+    a = a_.to(tl.float32)
+    rms = tl.sqrt(tl.sum(a * a, axis=0) / hidden_dim + eps)
+
+    w1_ = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    w1 = w1_.to(tl.float32)
+
+    a_rms = a / rms * w1
+
+    tl.store(
+        output_ptr + input_start + offsets,
+        a_rms,  # implicitly casts to output dtype here
+        mask=mask,
+    )
+
+
+def fused_rmsnorm(x, weight, eps, autotune=False, inplace=False):
+    assert len(x.shape) == 2
+    if inplace:
+        output = x
+    else:
+        output = torch.empty_like(x)
+    bs, hidden_dim = x.shape
+    config = {
+        "BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
+        "num_warps": max(
+            min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), 32), 4
+        ),
+    }
+
+    fused_rmsnorm_kernel[(bs,)](
+        output, x, weight, eps=eps, hidden_dim=hidden_dim, **config
+    )
+    return output
+
+
+class FusedDualResidualRMSNorm:
+    """
+    Fused implementation of
+    y = RMSNorm2(RMSNorm1(x) + residual))
+    """
+
+    def __init__(self, rmsnorm1, rmsnorm2) -> None:  # the one after rmsnorm1
+        self.rmsnorm1 = rmsnorm1
+        self.rmsnorm2 = rmsnorm2
+        self.variance_epsilon = self.rmsnorm1.variance_epsilon
+        assert self.rmsnorm1.variance_epsilon == self.rmsnorm2.variance_epsilon
+        assert self.rmsnorm1.weight.shape == self.rmsnorm2.weight.shape
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(
+        self, x: torch.Tensor, residual: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if x.is_cuda:
+            return self.forward_cuda(x, residual)
+        else:
+            return self.forward_flashinfer(x, residual)
+
+    def forward_cuda(
+        self, x: torch.Tensor, residual: torch.Tensor, autotune=False
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return fused_dual_residual_rmsnorm(
+            x,
+            residual,
+            self.rmsnorm1.weight,
+            self.rmsnorm2.weight,
+            self.variance_epsilon,
+            autotune=autotune,
+        )
+
+    def forward_flashinfer(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        normed1 = self.rmsnorm1(x)
+        residual = normed1 + residual
+        return self.rmsnorm2(residual), residual
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        normed1 = self.rmsnorm1.forward_native(x)
+        residual = normed1 + residual
+        return self.rmsnorm2.forward_native(residual), residual
+
+
+# gelu on first half of vector
+@triton.jit
+def gelu_and_mul_kernel(
+    out_hidden_states_ptr,  # (bs, hidden_dim)
+    out_scales_ptr,  # (bs,)
+    hidden_states_ptr,  # (bs, hidden_dim * 2)
+    quant_max: tl.constexpr,
+    static_scale: tl.constexpr,
+    hidden_dim: tl.constexpr,  # the output hidden_dim
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    input_start = pid * hidden_dim * 2
+    output_start = pid * hidden_dim
+
+    input1_offs = tl.arange(0, BLOCK_SIZE)
+    mask = tl.arange(0, BLOCK_SIZE) < hidden_dim  # shared for input1, input3, output
+    input3_offs = hidden_dim + tl.arange(0, BLOCK_SIZE)
+    output_offs = tl.arange(0, BLOCK_SIZE)
+
+    x1 = tl.load(
+        hidden_states_ptr + input_start + input1_offs, mask=mask, other=0.0
+    ).to(tl.float32)
+    x3 = tl.load(
+        hidden_states_ptr + input_start + input3_offs, mask=mask, other=0.0
+    ).to(tl.float32)
+
+    # gelu
+    # cast down before mul to better match training?
+    gelu_x1 = 0.5 * (1.0 + tl.erf(x1 * 0.7071067811865475)) * x1
+    out = x3 * gelu_x1.to(hidden_states_ptr.dtype.element_ty)
+
+    if quant_max is not None:
+        raise NotImplementedError()
+
+    tl.store(out_hidden_states_ptr + output_start + output_offs, out, mask=mask)
+
+
+def gelu_and_mul_triton(
+    hidden_states,
+    scales=None,
+    quantize=None,  # dtype to quantize to
+    out=None,
+):
+    bs, in_hidden_dim = hidden_states.shape
+    hidden_dim = in_hidden_dim // 2
+
+    if out is None:
+        out_hidden_states = torch.empty(
+            (bs, hidden_dim),
+            dtype=quantize or hidden_states.dtype,
+            device=hidden_states.device,
+        )
+    else:
+        assert out.shape == (bs, hidden_dim)
+        assert out.dtype == (quantize or hidden_states.dtype)
+        out_hidden_states = out
+    out_scales = None
+    static_scale = False
+    if quantize is not None:
+        if scales is None:
+            out_scales = torch.empty(
+                (bs,), dtype=torch.float32, device=hidden_states.device
+            )
+        else:
+            out_scales = scales
+            static_scale = True
+
+    config = {
+        # 8 ele per thread (not tuned)
+        "num_warps": max(
+            min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), 32), 4
+        ),
+    }
+
+    gelu_and_mul_kernel[(bs,)](
+        out_hidden_states,
+        out_scales,
+        hidden_states,
+        quant_max=torch.finfo(quantize).max if quantize is not None else None,
+        static_scale=static_scale,
+        hidden_dim=hidden_dim,
+        BLOCK_SIZE=triton.next_power_of_2(hidden_dim),
+        **config,
+    )
+
+    if quantize is not None:
+        return out_hidden_states, out_scales
+    else:
+        return out_hidden_states, None

python/sglang/srt/layers/moe/router.py

@@ -0,0 +1,342 @@
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.moe.topk import fused_topk
+
+
+@triton.jit
+def fused_moe_router_kernel(
+    input_ptr,  # input (bs, hidden_dim)
+    moe_router_weight_ptr,  # input (num_experts, hidden_dim)
+    topk_weights_ptr,  # output (bs, topk)
+    topk_ids_ptr,  # output (bs, topk)
+    num_experts: tl.constexpr,
+    topk: tl.constexpr,
+    moe_softcapping: tl.constexpr,
+    moe_renormalize: tl.constexpr,  # not supported
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < hidden_dim
+
+    # moe_router_weight is k major
+    expert_offsets = tl.arange(0, num_experts)[:, None]
+    router_mask = mask[None, :]
+    w_router = tl.load(
+        moe_router_weight_ptr + expert_offsets * hidden_dim + offsets[None, :],
+        mask=router_mask,
+        other=0.0,
+    )
+
+    x = tl.load(input_ptr + pid * hidden_dim + offsets, mask=mask, other=0.0)
+
+    # todo: tl.dot?
+    logits = tl.sum((w_router.to(tl.float32) * x[None, :].to(tl.float32)), axis=-1)
+
+    # logit softcap
+    logits_scaled = logits / moe_softcapping
+    exped = tl.exp(2 * logits_scaled)
+    top = exped - 1
+    bottom = exped + 1
+    logits_softcapped = top / bottom * moe_softcapping
+
+    # topk
+    # assert 1 <= topk <= num_experts
+
+    # 5.38 us
+
+    top1 = tl.argmax(logits_softcapped, axis=0)
+    tl.store(topk_ids_ptr + pid * topk + 0, top1)  # 5.63 us
+
+    top1_v = tl.max(logits_softcapped, axis=0)
+    invsumexp = 1.0 / tl.sum(tl.exp(logits_softcapped - top1_v), axis=0)
+
+    tl.store(
+        topk_weights_ptr + pid * topk + 0,
+        invsumexp,
+    )  # 5.73 us
+
+    if topk >= 2:
+        top2 = tl.argmax(
+            tl.where(
+                tl.arange(0, num_experts) != top1, logits_softcapped, float("-inf")
+            ),
+            axis=0,
+        )
+        tl.store(topk_ids_ptr + pid * topk + 1, top2)
+        top2_v = tl.sum(logits_softcapped * (tl.arange(0, num_experts) == top2), axis=0)
+        tl.store(
+            topk_weights_ptr + pid * topk + 1,
+            tl.exp(top2_v - top1_v) * invsumexp,
+        )  # 5.95us
+
+    # probably slow
+    if topk > 2:
+        topk_mask = tl.full(logits_softcapped.shape, 1.0, dtype=logits_softcapped.dtype)
+        topk_mask = tl.where(
+            tl.arange(0, num_experts) != top1, topk_mask, float("-inf")
+        )
+        topk_mask = tl.where(
+            tl.arange(0, num_experts) != top2, topk_mask, float("-inf")
+        )
+        for i in range(2, topk):
+            topi = tl.argmax(logits_softcapped + topk_mask, axis=0)
+            topk_mask = tl.where(
+                tl.arange(0, num_experts) != topi, topk_mask, float("-inf")
+            )
+            tl.store(topk_ids_ptr + pid * topk + i, topi)
+            topi_v = tl.sum(
+                logits_softcapped * (tl.arange(0, num_experts) == topi), axis=0
+            )
+            tl.store(
+                topk_weights_ptr + pid * topk + i,
+                tl.exp(topi_v - top1_v) * invsumexp,
+            )
+    # assert not moe_renormalize, "moe weight renormalization not implemented"
+
+
+def fused_moe_router_impl(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    topk: int,
+    moe_softcapping: float,
+):
+    assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1]
+    bs, hidden_dim = x.shape
+    num_experts = router_weight.shape[0]
+
+    # router_logits = torch.empty((bs, num_experts), dtype=torch.float32, device=x.device)
+    topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device)
+    topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device)
+
+    grid = lambda meta: (bs,)
+    config = {
+        "BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
+        "num_warps": max(
+            min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), 32), 4
+        ),
+    }
+
+    fused_moe_router_kernel[grid](
+        x,
+        router_weight,
+        topk_weights,
+        topk_ids,
+        num_experts=num_experts,
+        topk=topk,
+        moe_softcapping=moe_softcapping,
+        moe_renormalize=False,
+        hidden_dim=hidden_dim,
+        **config,
+    )
+
+    return topk_weights, topk_ids
+
+
+@triton.jit
+def fused_moe_router_large_bs_kernel(
+    a_ptr,  # input (bs, hidden_dim)
+    b_ptr,  # input (num_experts, hidden_dim)
+    topk_weights_ptr,  # output (bs, topk)
+    topk_ids_ptr,  # output (bs, topk)
+    bs,
+    num_experts: tl.constexpr,
+    topk: tl.constexpr,  # only support topk == 1
+    moe_softcapping: tl.constexpr,
+    moe_renormalize: tl.constexpr,  # not supported
+    K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    stride_am: tl.constexpr,
+    stride_bn: tl.constexpr,
+):
+
+    # 1. get block id
+    pid = tl.program_id(axis=0)
+
+    # 2. create pointers for the first block of A and B
+    # 2.1. setup a_ptrs with offsets in m and k
+    offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)[:, None]
+    bs_mask = offs_m < bs
+    offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :]
+    a_ptrs = a_ptr + (offs_m * stride_am + offs_k)
+
+    # 2.2. setup b_ptrs with offsets in k and n.
+    #      Note: b matrix is k-major.
+    offs_k = tl.arange(0, BLOCK_SIZE_K)[None, :]
+    offs_n = tl.arange(0, BLOCK_SIZE_N)[:, None]
+    expert_mask = offs_n < num_experts
+    b_ptrs = b_ptr + (offs_n * stride_bn + offs_k)
+
+    # 3. Create an accumulator of float32 of size [BLOCK_SIZE_M, BLOCK_SIZE_N]
+    #    3.1. iterate in K dimension
+    #    3.2. transpose tile B
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, K // BLOCK_SIZE_K):  # hidden_dim % BLOCK_SIZE_K == 0
+        a = tl.load(
+            a_ptrs,
+            mask=bs_mask,
+            other=0.0,
+        ).to(tl.float32)
+        b = tl.load(b_ptrs, mask=expert_mask, other=0.0).to(tl.float32).T
+        acc += tl.dot(a, b)
+
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K
+        b_ptrs += BLOCK_SIZE_K
+
+    # 4. logit softcap
+    logits_scaled = acc / moe_softcapping
+    exped = tl.exp(2 * logits_scaled)
+    logits_softcapped = (exped - 1) / (exped + 1) * moe_softcapping
+
+    # 5. top1
+    cond = tl.arange(0, BLOCK_SIZE_N)[None, :] < num_experts
+    top1 = tl.argmax(tl.where(cond, logits_softcapped, float("-inf")), axis=1)
+    top1_v = tl.max(
+        tl.where(cond, logits_softcapped, float("-inf")), axis=1, keep_dims=True
+    )
+    invsumexp = 1.0 / tl.sum(
+        tl.where(cond, tl.exp(logits_softcapped - top1_v), 0.0), axis=1
+    )
+
+    # 6. store to output
+    offs_topk = pid * topk * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    topk_mask = offs_topk < bs
+    tl.store(topk_ids_ptr + offs_topk, top1, mask=topk_mask)
+    tl.store(
+        topk_weights_ptr + offs_topk,
+        invsumexp,
+        mask=topk_mask,
+    )
+
+
+def fused_moe_router_large_bs_impl(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    topk: int,
+    moe_softcapping: float,
+    BLOCK_SIZE_M: int,
+    BLOCK_SIZE_N: int,
+    BLOCK_SIZE_K: int,
+):
+    assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1]
+    bs, hidden_dim = x.shape
+    num_experts = router_weight.shape[0]
+
+    assert num_experts <= BLOCK_SIZE_N
+    assert hidden_dim % BLOCK_SIZE_K == 0
+    assert topk == 1
+
+    topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device)
+    topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device)
+
+    grid = (triton.cdiv(bs, BLOCK_SIZE_M) * triton.cdiv(num_experts, BLOCK_SIZE_N),)
+
+    fused_moe_router_large_bs_kernel[grid](
+        a_ptr=x,
+        b_ptr=router_weight,
+        topk_weights_ptr=topk_weights,
+        topk_ids_ptr=topk_ids,
+        bs=bs,
+        num_experts=num_experts,
+        topk=topk,
+        moe_softcapping=moe_softcapping,
+        moe_renormalize=False,
+        K=hidden_dim,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        stride_am=hidden_dim,
+        stride_bn=hidden_dim,
+    )
+
+    return topk_weights, topk_ids
+
+
+def fused_moe_router_shim(
+    moe_softcapping,
+    hidden_states,
+    gating_output,
+    topk,
+    renormalize,
+):
+    assert not renormalize
+    assert (
+        len(hidden_states.shape) == 2
+        and hidden_states.shape[1] == gating_output.shape[1]
+    )
+    bs, hidden_dim = hidden_states.shape
+    num_experts = gating_output.shape[0]
+    BLOCK_SIZE_M = 32
+    BLOCK_SIZE_N = 16
+    BLOCK_SIZE_K = 256
+    if (
+        bs >= 512
+        and topk == 1
+        and num_experts <= BLOCK_SIZE_N
+        and hidden_dim % BLOCK_SIZE_K == 0
+    ):
+        return fused_moe_router_large_bs_impl(
+            x=hidden_states,
+            router_weight=gating_output,
+            topk=topk,
+            moe_softcapping=moe_softcapping,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+        )
+    else:
+        return fused_moe_router_impl(
+            x=hidden_states,
+            router_weight=gating_output,
+            topk=topk,
+            moe_softcapping=moe_softcapping,
+        )
+
+
+class FusedMoeRouter:
+    def __init__(self, router_linear, topk, moe_softcapping) -> None:
+        self.router_linear = router_linear
+        self.topk = topk
+        self.moe_softcapping = moe_softcapping
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+    def forward(
+        self, x: torch.Tensor, residual: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if x.is_cuda:
+            return self.forward_cuda(x, residual)
+        else:
+            return self.forward_vllm(x, residual)
+
+    def forward_cuda(
+        self, x: torch.Tensor, autotune=False
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return fused_moe_router_shim(
+            moe_softcapping=self.moe_softcapping,
+            hidden_states=x,
+            gating_output=self.router_linear.weight,
+            topk=self.topk,
+            renormalize=False,
+        )
+
+    def forward_vllm(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # g, _ = self.router_linear.forward(x)
+        g = x.float() @ self.router_linear.weight.T.float()
+
+        g = torch.tanh(g.float() / self.moe_softcapping) * self.moe_softcapping
+
+        return fused_topk(x, g, self.topk, False)

python/sglang/srt/models/grok.py

@@ -15,28 +15,36 @@
 # Adapted from
 # https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/mixtral.py#L1
 """Inference-only Grok1 model."""
+import functools
+import json
+import logging
+import math
+import os
+import warnings
+from typing import Iterable, Optional, Tuple
 
-from typing import Iterable, List, Optional, Tuple
+import numpy as np
-
 import torch
-import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
 )
-from sglang.srt.layers.activation import GeluAndMul
+from sglang.srt.layers.elementwise import fused_dual_residual_rmsnorm, fused_rmsnorm
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
-    MergedColumnParallelLinear,
     QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.moe.ep_moe.layer import EPMoE
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+from sglang.srt.layers.moe.router import fused_moe_router_shim
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
@@ -44,47 +52,17 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import add_prefix
+from sglang.srt.utils import dump_to_file
+
+logger = logging.getLogger(__name__)
 
 
-class Grok1MLP(nn.Module):
+debug_tensor_dump_output_folder = None
-    def __init__(
+debug_tensor_dump_inject = False
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        quant_config: Optional[QuantizationConfig] = None,
-        prefix: str = "",
-        reduce_results=True,
-        use_presharded_weights: bool = False,
-    ) -> None:
-        super().__init__()
-        self.gate_up_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [intermediate_size] * 2,
-            bias=False,
-            quant_config=quant_config,
-            prefix=add_prefix("gate_up_proj", prefix),
-            use_presharded_weights=use_presharded_weights,
-        )
-        self.down_proj = RowParallelLinear(
-            intermediate_size,
-            hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=add_prefix("down_proj", prefix),
-            reduce_results=reduce_results,
-            use_presharded_weights=use_presharded_weights,
-        )
-        self.act_fn = GeluAndMul(approximate="tanh")
-
-    def forward(self, x):
-        gate_up, _ = self.gate_up_proj(x)
-        x = self.act_fn(gate_up)
-        x, _ = self.down_proj(x)
-        return x
 
 
 class Grok1MoE(nn.Module):
@@ -108,51 +86,55 @@ class Grok1MoE(nn.Module):
         tp_size: Optional[int] = None,
         reduce_results=True,
         use_presharded_weights: bool = False,
-        prefix: str = "",
+        inplace: bool = True,
+        no_combine: bool = False,
     ):
         super().__init__()
         self.hidden_size = hidden_size
 
-        # Gate always runs at half / full precision for now.
+        # Gate always runs at full precision for stability (see https://arxiv.org/pdf/2101.03961)
         self.gate = ReplicatedLinear(
             hidden_size,
             num_experts,
             bias=False,
-            params_dtype=params_dtype,
+            params_dtype=torch.float32,
             quant_config=None,
-            prefix=add_prefix("gate", prefix),
         )
 
         self.router_logit_softcapping = getattr(
             config, "router_logit_softcapping", 30.0
         )
-        self.experts = FusedMoE(
+        custom_routing_function = functools.partial(
+            fused_moe_router_shim, self.router_logit_softcapping
+        )
+
+        kwargs = {}
+        if global_server_args_dict["enable_ep_moe"]:
+            MoEImpl = EPMoE
+        else:
+            MoEImpl = FusedMoE
+            kwargs["reduce_results"] = reduce_results
+            kwargs["use_presharded_weights"] = use_presharded_weights
+            kwargs["inplace"] = inplace
+            kwargs["no_combine"] = no_combine
+
+        self.experts = MoEImpl(
             num_experts=num_experts,
             top_k=top_k,
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             params_dtype=params_dtype,
-            reduce_results=reduce_results,
             renormalize=False,
             quant_config=quant_config,
             tp_size=tp_size,
+            custom_routing_function=custom_routing_function,
             activation="gelu",
-            use_presharded_weights=use_presharded_weights,
+            **kwargs,
-            prefix=add_prefix("experts", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # NOTE: hidden_states can have either 1D or 2D shape.
-        orig_shape = hidden_states.shape
-        hidden_states = hidden_states.view(-1, self.hidden_size)
-
-        # router_logits: (num_tokens, n_experts)
-        router_logits, _ = self.gate(hidden_states)
-        router_logits = 30.0 * F.tanh(router_logits / 30.0)
-
         # need to assert self.gate.quant_method is unquantized
-        final_hidden_states = self.experts(hidden_states, router_logits)
+        return self.experts(hidden_states, self.gate.weight)
-        return final_hidden_states.view(orig_shape)
 
 
 class Grok1Attention(nn.Module):
@@ -167,31 +149,33 @@ class Grok1Attention(nn.Module):
         rope_theta: float = 10000,
         quant_config: Optional[QuantizationConfig] = None,
         reduce_results: bool = True,
-        prefix: str = "",
+        load_presharded_attn: bool = False,
     ) -> None:
         super().__init__()
         self.config = config
         self.layer_id = layer_id
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        attn_tp_rank = get_tensor_model_parallel_rank()
+        attn_tp_size = get_tensor_model_parallel_world_size()
         self.total_num_heads = num_heads
-        assert self.total_num_heads % tp_size == 0
+        assert self.total_num_heads % attn_tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
+        self.num_heads = self.total_num_heads // attn_tp_size
         self.total_num_kv_heads = num_kv_heads
-        if self.total_num_kv_heads >= tp_size:
+        if self.total_num_kv_heads >= attn_tp_size:
             # Number of KV heads is greater than TP size, so we partition
             # the KV heads across multiple tensor parallel GPUs.
-            assert self.total_num_kv_heads % tp_size == 0
+            assert self.total_num_kv_heads % attn_tp_size == 0
         else:
             # Number of KV heads is less than TP size, so we replicate
             # the KV heads across multiple tensor parallel GPUs.
-            assert tp_size % self.total_num_kv_heads == 0
+            assert attn_tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size)
         self.head_dim = getattr(config, "head_dim", 128)
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self.rope_theta = rope_theta
+        self.load_presharded_attn = load_presharded_attn
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -200,7 +184,9 @@ class Grok1Attention(nn.Module):
             self.total_num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix=add_prefix("qkv_proj", prefix),
+            tp_rank=attn_tp_rank,
+            tp_size=attn_tp_size,
+            load_presharded_attn=self.load_presharded_attn,
         )
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
@@ -208,7 +194,9 @@ class Grok1Attention(nn.Module):
             bias=False,
             quant_config=quant_config,
             reduce_results=reduce_results,
-            prefix=add_prefix("o_proj", prefix),
+            tp_rank=attn_tp_rank,
+            tp_size=attn_tp_size,
+            use_presharded_weights=self.load_presharded_attn,
         )
         self.rotary_emb = get_rope(
             self.head_dim,
@@ -227,7 +215,6 @@ class Grok1Attention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
             logit_cap=logit_cap,
-            prefix=add_prefix("attn", prefix),
         )
 
     def forward(
@@ -236,10 +223,73 @@ class Grok1Attention(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
+        if hidden_states.shape[0] == 0:
+            assert (
+                not self.o_proj.reduce_results
+            ), "short-circuiting allreduce will lead to hangs"
+            return hidden_states
+        if debug_tensor_dump_output_folder:
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                f"attn_input_{self.layer_id}",
+                hidden_states,
+            )
+
+            if debug_tensor_dump_inject:
+                name = os.path.join(
+                    debug_tensor_dump_output_folder,
+                    f"jax_dump_attn_input_{self.layer_id}.npy",
+                )
+                logger.info(f"Load {name} from jax.")
+                x = np.load(name)
+                hidden_states = torch.tensor(x[0, : hidden_states.shape[0]]).to(
+                    hidden_states
+                )
+
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self.rotary_emb(positions, q, k)
+
+        if debug_tensor_dump_output_folder:
+            num_tokens = q.shape[0]
+            num_heads_q = self.num_heads
+            head_dim = self.head_dim
+            num_heads_kv = k.numel() // (num_tokens * head_dim)
+
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                f"q_{self.layer_id}",
+                tensor_model_parallel_all_gather(
+                    q.reshape(num_tokens, num_heads_q, head_dim).contiguous(), dim=1
+                ).contiguous(),
+            )
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                f"k_{self.layer_id}",
+                tensor_model_parallel_all_gather(
+                    k.reshape(num_tokens, num_heads_kv, head_dim).contiguous(), dim=1
+                ).contiguous(),
+            )
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                f"v_{self.layer_id}",
+                tensor_model_parallel_all_gather(
+                    v.reshape(num_tokens, num_heads_kv, head_dim).contiguous(), dim=1
+                ).contiguous(),
+            )
+
         attn_output = self.attn(q, k, v, forward_batch)
+
+        if debug_tensor_dump_output_folder:
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                f"attn_output_{self.layer_id}",
+                tensor_model_parallel_all_gather(
+                    attn_output.reshape(num_tokens, num_heads_q, head_dim).contiguous(),
+                    dim=1,
+                ).contiguous(),
+            )
+
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -250,8 +300,9 @@ class Grok1DecoderLayer(nn.Module):
         config: PretrainedConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
-        use_presharded_weights: bool = False,
+        load_presharded_moe: bool = False,
-        prefix: str = "",
+        load_presharded_attn: bool = False,
+        load_presharded_mlp: bool = False,
     ) -> None:
         super().__init__()
         self.num_experts = config.num_local_experts
@@ -268,7 +319,8 @@ class Grok1DecoderLayer(nn.Module):
             layer_id=layer_id,
             rope_theta=rope_theta,
             quant_config=quant_config,
-            prefix=add_prefix("attn", prefix),
+            reduce_results=False,
+            load_presharded_attn=load_presharded_attn,
         )
         self.block_sparse_moe = Grok1MoE(
             config=config,
@@ -282,38 +334,68 @@ class Grok1DecoderLayer(nn.Module):
             ),
             quant_config=quant_config,
             reduce_results=True,
-            use_presharded_weights=use_presharded_weights,
+            use_presharded_weights=load_presharded_moe,
-            prefix=add_prefix("block_sparse_moe", prefix),
+            inplace=True,
+            no_combine=False,  # just a suggestion to not combine topk
         )
+
         self.pre_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.pre_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_moe_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
+        self.ffn = self.block_sparse_moe
+
     def forward(
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
-    ) -> torch.Tensor:
+        residual: Optional[torch.Tensor] = None,
+        deferred_norm: Optional[RMSNorm] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, RMSNorm]:
         # Self Attention
-        hidden_states = (
+        if deferred_norm is not None:
-            self.post_attn_norm(
+            assert residual is not None
-                self.self_attn(
+            # here hidden_states is output of ffn, residual is residual from after previous attn layer
-                    positions=positions,
+            hidden_states, residual = fused_dual_residual_rmsnorm(
-                    hidden_states=self.pre_attn_norm(hidden_states),
+                hidden_states,
-                    forward_batch=forward_batch,
+                residual,
-                )
+                deferred_norm.weight,
+                self.pre_attn_norm.weight,
+                deferred_norm.variance_epsilon,
             )
-            + hidden_states
+        else:
+            # here hidden_states is the residual
+            hidden_states, residual = (
+                fused_rmsnorm(
+                    hidden_states,
+                    self.pre_attn_norm.weight,
+                    self.pre_attn_norm.variance_epsilon,
+                ),
+                hidden_states,
+            )
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            forward_batch=forward_batch,
+        )
+
+        if get_tensor_model_parallel_world_size() > 1:
+            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+        hidden_states, residual = fused_dual_residual_rmsnorm(
+            hidden_states,
+            residual,
+            self.post_attn_norm.weight,
+            self.pre_moe_norm.weight,
+            self.post_attn_norm.variance_epsilon,
         )
 
         # Fully Connected
-        hidden_states = (
+        hidden_states = self.ffn(hidden_states)
-            self.post_moe_norm(self.block_sparse_moe(self.pre_moe_norm(hidden_states)))
+        return hidden_states, residual, self.post_moe_norm  # defer layernorm
-            + hidden_states
-        )
-        return hidden_states
 
 
 class Grok1Model(nn.Module):
@@ -321,8 +403,10 @@ class Grok1Model(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
-        use_presharded_weights: bool = False,
+        load_presharded_moe: bool = False,
-        prefix: str = "",
+        load_presharded_embedding: bool = False,
+        load_presharded_attn: bool = False,
+        load_presharded_mlp: bool = False,
     ) -> None:
         super().__init__()
         self.config = config
@@ -332,7 +416,7 @@ class Grok1Model(nn.Module):
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
             config.hidden_size,
-            prefix=add_prefix("embed_tokens", prefix),
+            use_presharded_weights=load_presharded_embedding,
         )
         self.layers = nn.ModuleList(
             [
@@ -340,8 +424,9 @@ class Grok1Model(nn.Module):
                     config,
                     i,
                     quant_config=quant_config,
-                    use_presharded_weights=use_presharded_weights,
+                    load_presharded_moe=load_presharded_moe,
-                    prefix=add_prefix(f"layers.{i}", prefix),
+                    load_presharded_attn=load_presharded_attn,
+                    load_presharded_mlp=load_presharded_mlp,
                 )
                 for i in range(config.num_hidden_layers)
             ]
@@ -361,10 +446,48 @@ class Grok1Model(nn.Module):
         else:
             hidden_states = input_embeds
 
+        residual, deferred_norm = None, None
         for i in range(len(self.layers)):
-            hidden_states = self.layers[i](positions, hidden_states, forward_batch)
+            hidden_states, residual, deferred_norm = self.layers[i](
-        hidden_states = self.norm(hidden_states)
+                positions, hidden_states, forward_batch, residual, deferred_norm
-        hidden_states.mul_(self.config.output_multiplier_scale)
+            )
+
+        if debug_tensor_dump_output_folder:
+            hidden_states = (
+                fused_rmsnorm(
+                    hidden_states,
+                    deferred_norm.weight,
+                    deferred_norm.variance_epsilon,
+                )
+                + residual
+            )
+
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                "last_hidden_before_norm",
+                hidden_states,
+            )
+
+            hidden_states = fused_rmsnorm(
+                hidden_states,
+                self.norm.weight,
+                self.norm.variance_epsilon,
+            )
+
+            dump_to_file(
+                debug_tensor_dump_output_folder,
+                "last_hidden_after_norm",
+                hidden_states,
+            )
+        else:
+            hidden_states, _ = fused_dual_residual_rmsnorm(
+                hidden_states,
+                residual,
+                deferred_norm.weight,
+                self.norm.weight,
+                deferred_norm.variance_epsilon,
+            )
+
         return hidden_states
 
 
@@ -373,31 +496,77 @@ class Grok1ForCausalLM(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
-        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
         self.quant_config = quant_config
 
-        if (
+        # Get presharded weights.
+        self.load_presharded_mlp = getattr(config, "load_presharded_mlp", False)
+        self.load_presharded_moe = (
             self.config.num_local_experts > 0
             and get_tensor_model_parallel_world_size() > 1
-        ):
+        )
-            self.use_presharded_weights = True
+        self.load_presharded_attn = getattr(config, "load_presharded_attn", False)
+        self.load_presharded_embedding = getattr(
+            config, "load_presharded_embedding", False
+        )
+
+        self.is_weights_presharded = (
+            self.load_presharded_mlp
+            or self.load_presharded_moe
+            or self.load_presharded_attn
+            or self.load_presharded_embedding
+        )
+
+        if self.is_weights_presharded:
             setattr(DefaultModelLoader, "_prepare_weights", _prepare_presharded_weights)
-        else:
+
-            self.use_presharded_weights = False
+        default_replicate_lm_head = False
+        self.replicate_lm_head = getattr(
+            config, "replicate_lm_head", default_replicate_lm_head
+        )
 
         self.model = Grok1Model(
             config,
             quant_config=quant_config,
-            use_presharded_weights=self.use_presharded_weights,
+            load_presharded_moe=self.load_presharded_moe,
-            prefix=add_prefix("model", prefix),
+            load_presharded_embedding=self.load_presharded_embedding,
+            load_presharded_attn=self.load_presharded_attn,
+            load_presharded_mlp=self.load_presharded_mlp,
         )
-        self.lm_head = ParallelLMHead(
+
-            config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix)
+        lm_head_params_dtype = None
-        )
+        if self.replicate_lm_head:
-        self.logits_processor = LogitsProcessor(config)
+            self.lm_head = ReplicatedLinear(
+                config.hidden_size,
+                config.vocab_size,
+                bias=False,
+                params_dtype=lm_head_params_dtype,
+            )
+            self.logits_processor = LogitsProcessor(config, skip_all_gather=True)
+        else:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                use_presharded_weights=self.load_presharded_embedding,
+                params_dtype=lm_head_params_dtype,
+            )
+            self.logits_processor = LogitsProcessor(config)
+
+        # Dump tensors for debugging
+        global debug_tensor_dump_output_folder, debug_tensor_dump_inject
+        debug_tensor_dump_output_folder = global_server_args_dict[
+            "debug_tensor_dump_output_folder"
+        ]
+        debug_tensor_dump_inject = global_server_args_dict["debug_tensor_dump_inject"]
+        warnings.filterwarnings("ignore", category=FutureWarning)
+
+        if get_tensor_model_parallel_rank() == 0:
+            logger.info(
+                f"#parameters (analytical): {self.get_num_params_analytical() / 1e9:.2f} B, "
+                f"#parameters (actual): {self.get_num_params_torch() / 1e9:.2f} B"
+            )
 
     def forward(
         self,
@@ -406,6 +575,9 @@ class Grok1ForCausalLM(nn.Module):
         forward_batch: ForwardBatch,
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
+        if debug_tensor_dump_output_folder:
+            dump_to_file(debug_tensor_dump_output_folder, "input_ids", input_ids)
+
         hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
         return self.logits_processor(
             input_ids, hidden_states, self.lm_head, forward_batch
@@ -414,21 +586,28 @@ class Grok1ForCausalLM(nn.Module):
     def load_weights(
         self,
         weights: Iterable[Tuple[str, torch.Tensor]],
-    ):
+        num_experts: Optional[int] = None,
-        num_experts = self.config.num_local_experts
+        ignore_parent_name: bool = False,
-
+    ) -> dict[str, torch.Tensor]:
-        stacked_params_mapping = [
+        if num_experts is None:
+            num_experts = self.config.num_local_experts
+        stacked_params_mapping = []
+        stacked_params_mapping += [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
+        ]
+        stacked_params_mapping += [
+            # (param_name, shard_name, shard_id)
             ("gate_up_proj", "gate_proj", 0),
             ("gate_up_proj", "up_proj", 1),
         ]
 
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
-        expert_params_mapping = FusedMoE.make_expert_params_mapping(
+        MoEImpl = EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE
+        expert_params_mapping = MoEImpl.make_expert_params_mapping(
             ckpt_gate_proj_name="w1",
             ckpt_down_proj_name="w2",
             ckpt_up_proj_name="w3",
@@ -439,14 +618,25 @@ class Grok1ForCausalLM(nn.Module):
         all_names = set(params_dict.keys())
         hit_names = set()
 
-        def load_weight_wrapper(name, loaded_weight, *args, **kwargs):
+        def load_weight_wrapper(
+            name: str, loaded_weight: torch.Tensor, *args, **kwargs
+        ):
+            if ignore_parent_name:
+                name = name.split(".")[-1]
+
             if name not in params_dict:
                 return
 
+            # Fuse constant multipliers into the weights
+            if "lm_head" in name:
+                loaded_weight = (
+                    loaded_weight.to(torch.float32)
+                    * self.config.output_multiplier_scale
+                )
+
             param = params_dict[name]
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
             weight_loader(param, loaded_weight, *args, **kwargs)
-
             hit_names.add(name)
 
         for name, loaded_weight in weights:
@@ -460,7 +650,6 @@ class Grok1ForCausalLM(nn.Module):
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue
-
                 load_weight_wrapper(name, loaded_weight, shard_id)
                 break
             else:
@@ -487,13 +676,79 @@ class Grok1ForCausalLM(nn.Module):
 
                     load_weight_wrapper(name=name, loaded_weight=loaded_weight)
 
+        if len(hit_names) > 5:
+            missing = all_names - hit_names
+            missing_exclude_scales = {x for x in missing if "scale" not in x}
+            logger.info(
+                f"#all_names: {len(all_names)}, #hit_names: {len(hit_names)}, #missing_exclude_scales: {len(missing_exclude_scales)}",
+            )
+            if len(missing_exclude_scales) > 0:
+                raise ValueError(
+                    f"load_weights failed because some weights are missing: {missing_exclude_scales=}."
+                )
+
+        elif len(hit_names) == 0:
+            raise ValueError("load_weights failed because it did not hit any names.")
+
+        return hit_names
+
+    def get_num_params_analytical(self):
+        cfg = self.config
+        moe_intermediate_size = getattr(
+            cfg,
+            "moe_intermediate_size",
+            getattr(cfg, "intermediate_size", None),
+        )
+        num_experts = cfg.num_local_experts
+
+        wq = (
+            cfg.num_hidden_layers
+            * cfg.hidden_size
+            * cfg.num_attention_heads
+            * cfg.head_dim
+        )
+        wkv = (
+            cfg.num_hidden_layers
+            * cfg.hidden_size
+            * cfg.num_key_value_heads
+            * cfg.head_dim
+            * 2
+        )
+        out = (
+            cfg.num_hidden_layers
+            * cfg.hidden_size
+            * cfg.num_attention_heads
+            * cfg.head_dim
+        )
+        ffn1 = (
+            cfg.num_hidden_layers
+            * num_experts
+            * cfg.hidden_size
+            * moe_intermediate_size
+            * 2
+        )
+        ffn2 = (
+            cfg.num_hidden_layers
+            * num_experts
+            * cfg.hidden_size
+            * moe_intermediate_size
+        )
+        embed = cfg.hidden_size * cfg.vocab_size * 2
+        return wq + wkv + out + ffn1 + ffn2 + embed
+
+    def get_num_params_torch(self):
+        return (
+            sum(p.numel() for p in self.parameters())
+            * get_tensor_model_parallel_world_size()
+        )
+
 
 old_prepare_weights = getattr(DefaultModelLoader, "_prepare_weights")
 
 
 def _prepare_presharded_weights(
     self, model_name_or_path: str, revision: Optional[str], fall_back_to_pt: bool
-) -> Tuple[str, List[str], bool]:
+) -> Tuple[str, list[str], bool]:
     import glob
     import os
 
@@ -522,7 +777,7 @@ def _prepare_presharded_weights(
     # The new format
     allow_patterns += [f"*-TP-{tp_rank:03d}.safetensors", "*-TP-common.safetensors"]
 
-    hf_weights_files: List[str] = []
+    hf_weights_files = []
     for pattern in allow_patterns:
         hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))