Skip to content

(WIP) perf(kimi-k3): tune small-M MoE decode #1140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
panditsa wants to merge 4 commits into
base: sanket/wip-k3-a8w4-select
Choose a base branch
from
Draft

(WIP) perf(kimi-k3): tune small-M MoE decode #1140

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2b1d76a
perf(moe): tune small-batch router waves
panditsa Aug 14, 2026
c6facd8
perf(kimi-k3): fuse the MoE path through four tokens
panditsa Aug 19, 2026
205030c
perf(moe): rebalance multi-token decode work
panditsa Aug 19, 2026
377b289
perf(amd): tune TP-local A16W4 decode
panditsa Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/tokenspeed/runtime/models/kimi_k3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1820,7 +1820,7 @@ def forward(
max_num_tokens_per_gpu = num_global_tokens

if self.native_latent_moe is not None:
if self._use_fused_decode_pipeline and hidden_states.shape[0] == 1:
if self._use_fused_decode_pipeline and hidden_states.shape[0] <= 4:
output = self._forward_fused_decode_pipeline(hidden_states, prefix_sum)
else:
output = self.native_latent_moe(
Expand Down
32 changes: 24 additions & 8 deletions ...nspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/gemm/fp16/rmsnorm_linear_add.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
import torch
from tokenspeed_kernel_amd._triton import gl, gluon

_HIDDEN = gl.constexpr(7168)
_LATENT = gl.constexpr(3584)
_BLOCK_N = gl.constexpr(32)
_BLOCK_K = gl.constexpr(512)
Expand All @@ -41,9 +40,14 @@ def _rmsnorm_linear_add_kernel(
residual_ptr,
shared_ptr,
output_ptr,
latent_stride_m,
residual_stride_m,
shared_stride_m,
output_stride_m,
eps,
):
pid_n = gl.program_id(0)
pid_m = gl.program_id(0)
pid_n = gl.program_id(1)
layout: gl.constexpr = gl.BlockedLayout(
[1, _BLOCK_K // _LANES],
[1, _LANES],
Expand All @@ -60,7 +64,7 @@ def _rmsnorm_linear_add_kernel(
k_mask = offs_k < _LATENT
latent = gl.amd.cdna4.buffer_load(
latent_ptr,
offs_k.to(gl.int32),
(pid_m * latent_stride_m + offs_k).to(gl.int32),
mask=k_mask,
other=0.0,
).to(gl.float32)
Expand All @@ -73,7 +77,7 @@ def _rmsnorm_linear_add_kernel(
k_mask = offs_k < _LATENT
latent = gl.amd.cdna4.buffer_load(
latent_ptr,
offs_k.to(gl.int32),
(pid_m * latent_stride_m + offs_k).to(gl.int32),
mask=k_mask,
other=0.0,
).to(gl.float32)
Expand Down Expand Up @@ -101,9 +105,16 @@ def _rmsnorm_linear_add_kernel(

# Match the materialized BF16 projection before the residual additions.
acc = acc.to(gl.bfloat16).to(gl.float32)
acc += gl.amd.cdna4.buffer_load(residual_ptr, offs_n.to(gl.int32)).to(gl.float32)
acc += gl.amd.cdna4.buffer_load(shared_ptr, offs_n.to(gl.int32)).to(gl.float32)
gl.store(output_ptr + offs_n, acc)
residual_offset = pid_m * residual_stride_m + offs_n
shared_offset = pid_m * shared_stride_m + offs_n
output_offset = pid_m * output_stride_m + offs_n
acc += gl.amd.cdna4.buffer_load(residual_ptr, residual_offset.to(gl.int32)).to(
gl.float32
)
acc += gl.amd.cdna4.buffer_load(shared_ptr, shared_offset.to(gl.int32)).to(
gl.float32
)
gl.store(output_ptr + output_offset, acc)


def gluon_rmsnorm_linear_add_gfx950(
Expand All @@ -118,13 +129,18 @@ def gluon_rmsnorm_linear_add_gfx950(
) -> torch.Tensor:
if out is None:
out = torch.empty_like(residual)
_rmsnorm_linear_add_kernel[(_HIDDEN // _BLOCK_N,)](
m = latent.shape[0]
_rmsnorm_linear_add_kernel[(m, projection_weight.shape[0] // _BLOCK_N)](
latent,
norm_weight,
projection_weight,
residual,
shared,
out,
latent.stride(0),
residual.stride(0),
shared.stride(0),
out.stride(0),
float(eps),
num_warps=8,
num_stages=1,
Expand Down
4 changes: 2 additions & 2 deletions tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/moe/mxfp4/routing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,7 @@ def invoke_softmax_topk_route_gluon(
(M, topk), dtype=torch.float32, device=router_logits.device
)
bias = correction_bias if correction_bias is not None else topk_weights
nw = 1 if M <= 2 else 4
nw = min(_next_pow2(M), 4)
_softmax_topk_route_gluon_kernel[(1,)](
router_logits,
bias,
Expand Down Expand Up @@ -358,7 +358,7 @@ def _launch_sigmoid_bias_topk_route_gluon(
topk_weights = torch.empty(
(M, topk), dtype=torch.float32, device=router_logits.device
)
nw = 1 if M <= 2 else 4
nw = 1 if M == 1 else min(max(_next_pow2(M), 4), 8)
_sigmoid_bias_topk_route_gluon_kernel[(1,)](
router_logits,
correction_bias,
Expand Down
Loading
Loading