Skip to content

perf(kimi3): fuse KDA decode core #1086

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

Open
qedawkins wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(kimi3): fuse KDA decode core #1086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
80 changes: 57 additions & 23 deletions python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from typing import TYPE_CHECKING

import torch
from tokenspeed_kernel.ops.activation.triton import rmsnorm_gated_sigmoid
from tokenspeed_kernel.ops.attention import (
kda_paged_decode,
kda_paged_prefill,
Expand Down Expand Up @@ -134,29 +135,50 @@ def _decode(
attn_tp_size: int,
head_v_dim: int,
lower_bound: float | None,
output_gate: torch.Tensor | None,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> torch.Tensor | None:

if output_gate is not None and (norm_weight is None or norm_eps is None):
raise ValueError(
"norm_weight and norm_eps are required with a KDA output gate"
)
if f_a_out is None:
return None
else:
num_value_heads = value_dim // attn_tp_size // head_v_dim
return try_kda_fused_paged_decode(
mixed_qkv,
conv_weights,
conv_states,
f_a_out,
f_b_weight,
beta_raw,
A_log,
dt_bias,
state_pool=ssm_states,
read_indices=read_indices,
write_indices=write_indices,
num_heads=num_value_heads,
head_dim=head_v_dim,
cu_seqlens=self.forward_metadata.query_start_loc,
lower_bound=lower_bound,
)

num_value_heads = value_dim // attn_tp_size // head_v_dim
result = try_kda_fused_paged_decode(
mixed_qkv,
conv_weights,
conv_states,
f_a_out,
f_b_weight,
beta_raw,
A_log,
dt_bias,
state_pool=ssm_states,
read_indices=read_indices,
write_indices=write_indices,
num_heads=num_value_heads,
head_dim=head_v_dim,
cu_seqlens=self.forward_metadata.query_start_loc,
lower_bound=lower_bound,
output_gate=output_gate,
norm_weight=norm_weight,
norm_eps=norm_eps,
)
if result is None:
return None
if result.output_norm_applied or output_gate is None:
return result.out
return rmsnorm_gated_sigmoid(
result.out.reshape(-1, num_value_heads * head_v_dim).contiguous(),
output_gate.contiguous(),
norm_weight,
norm_eps,
num_value_heads,
head_v_dim,
).view(1, -1, num_value_heads, head_v_dim)

@override
def _decode_scan(
Expand All @@ -177,8 +199,10 @@ def _decode_scan(
f_b_weight: torch.Tensor | None,
beta_raw: torch.Tensor | None,
lower_bound: float | None,
output_gate: torch.Tensor | None,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> torch.Tensor:

seq_len = query.shape[0]
num_heads = query.shape[2]
head_k_dim = query.shape[3]
Expand All @@ -193,7 +217,7 @@ def _decode_scan(
g_kda = g_raw.view(1, seq_len, num_value_heads, head_k_dim)
beta_kda = beta_raw.view(1, seq_len, num_value_heads)

return kda_paged_decode(
core_attn_out = kda_paged_decode(
query,
key,
value,
Expand All @@ -206,7 +230,17 @@ def _decode_scan(
write_indices=write_indices,
cu_seqlens=query_start_loc,
lower_bound=lower_bound,
).squeeze(0)
)
if output_gate is not None:
core_attn_out = rmsnorm_gated_sigmoid(
core_attn_out.reshape(-1, num_value_heads * head_v_dim).contiguous(),
output_gate.contiguous(),
norm_weight,
norm_eps,
num_value_heads,
head_v_dim,
).view(1, -1, num_value_heads, head_v_dim)
return core_attn_out.squeeze(0)

@override
def _verify(
Expand Down
21 changes: 21 additions & 0 deletions python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1292,6 +1292,9 @@ def forward_decode(
f_b_weight = kwargs.get("f_b_weight")
g_raw = kwargs.get("g_raw")
beta_raw = kwargs.get("beta_raw")
output_gate = kwargs.get("output_gate")
norm_weight = kwargs.get("norm_weight")
norm_eps = kwargs.get("norm_eps")
gate_lower_bound = kwargs.get("lower_bound")
A_log = kwargs["A_log"]
dt_bias = kwargs["dt_bias"]
Expand Down Expand Up @@ -1320,6 +1323,9 @@ def forward_decode(
attn_tp_size=attn_tp_size,
head_v_dim=head_v_dim,
lower_bound=gate_lower_bound,
output_gate=output_gate,
norm_weight=norm_weight,
norm_eps=norm_eps,
)
if fused_out is not None:
return fused_out
Expand Down Expand Up @@ -1368,6 +1374,9 @@ def forward_decode(
f_b_weight=f_b_weight,
beta_raw=beta_raw,
lower_bound=gate_lower_bound,
output_gate=output_gate,
norm_weight=norm_weight,
norm_eps=norm_eps,
)

def _decode(
Expand All @@ -1388,6 +1397,9 @@ def _decode(
attn_tp_size: int,
head_v_dim: int,
lower_bound: float | None,
output_gate: torch.Tensor | None,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> torch.Tensor | None:
"""Whole-step decode attempt; ``None`` falls through to the shared flow.

Expand Down Expand Up @@ -1415,6 +1427,9 @@ def _decode(
attn_tp_size: Attention tensor-parallel size.
head_v_dim: Value head dimension.
lower_bound: KDA decay clamp.
output_gate: Optional KDA gated-norm logits.
norm_weight: Optional KDA output RMSNorm weight.
norm_eps: Optional KDA output RMSNorm epsilon.

Returns:
The layer output when a fused kernel ran, else None.
Expand All @@ -1439,6 +1454,9 @@ def _decode_scan(
f_b_weight: torch.Tensor | None,
beta_raw: torch.Tensor | None,
lower_bound: float | None,
output_gate: torch.Tensor | None,
norm_weight: torch.Tensor | None,
norm_eps: float | None,
) -> torch.Tensor:
"""Single-token recurrent scan over the split, conv'd projections.

Expand All @@ -1464,6 +1482,9 @@ def _decode_scan(
f_b_weight: KDA second gate projection.
beta_raw: KDA raw per-head beta logits.
lower_bound: KDA decay clamp.
output_gate: Optional KDA gated-norm logits.
norm_weight: Optional KDA output RMSNorm weight.
norm_eps: Optional KDA output RMSNorm epsilon.

Returns:
``[1, B, Hv, V]`` layer output.
Expand Down
25 changes: 16 additions & 9 deletions python/tokenspeed/runtime/models/kimi_k3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -918,6 +918,7 @@ def forward(
# plain GEMV on the prefill path.
# Fused [3*proj, k] conv kernel bank, built once in post_load_weights.
conv_weights = self.conv_weights
fuse_decode_output_norm = ctx.forward_mode.is_decode() and num_tokens == ctx.bs

core_out = ctx.attn_backend.forward(
q=None,
Expand All @@ -943,19 +944,25 @@ def forward(
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
output_gate=out_gate if fuse_decode_output_norm else None,
norm_weight=self.o_norm.weight if fuse_decode_output_norm else None,
norm_eps=self.o_norm.variance_epsilon if fuse_decode_output_norm else None,
layer_id=self.layer_id,
seq_len=num_tokens,
)

# Per-head gated RMSNorm + sigmoid output gate in one kernel.
core_out = rmsnorm_gated_sigmoid(
core_out.reshape(num_tokens, hn * hd).contiguous(),
out_gate.contiguous(),
self.o_norm.weight,
self.o_norm.variance_epsilon,
hn,
hd,
)
core_out = core_out.reshape(num_tokens, hn * hd)
if not fuse_decode_output_norm:
# Decode kernels may fuse this epilogue; prefill retains the shared
# per-head norm implementation.
core_out = rmsnorm_gated_sigmoid(
core_out.contiguous(),
out_gate.contiguous(),
self.o_norm.weight,
self.o_norm.variance_epsilon,
hn,
hd,
)
output, _ = self.o_proj(core_out)
return output

Expand Down
Loading
Loading