Support Kimi K3 (#507)

kimi-k3 1/N

kimi-k3 2/N

kimi-k3 3/N

kimi-k3: fix MoE routing/pre-sort and compressed-tensors scheme probing on P800



fused_moe_ct_w4a16: moe_sigmoid_group_topk_norm and moe_pre_sorted only
support expert_num <= 512 on P800 (the former returns ret=1 and leaves its
outputs untouched, so the uninitialized topk ids made the pre-sort kernels
read out of bounds). Add torch fallbacks for >512 experts, verified
bit-exact against the kernels at expert_num=512.
compressed_tensors_moe: probe the expert quant scheme with the checkpoint’s
projection names (Kimi-K3 uses w1/w3/w2) as well, otherwise the scheme
lookup misses and the MoE silently falls back to UnquantizedFusedMoEMethod.
kunlun_communicator: stop forwarding use_all2all to DeviceCommunicatorBase,
which derives it itself and rejects the kwarg.

Co-authored-by: Kiro

Revert “kimi-k3: fix MoE routing/pre-sort and compressed-tensors scheme probing on P800”

This reverts commit 00db81aca5d530da8dd6be5ab3f4e2e8c71c8cdc.

kimi-k3 4/N moe_sigmoid_group_topk_norm & moe_pre_sorted  do not support more than 512 experts

add KDA

reorder_batch_threshold 512 → 1

fix recurrent_state oom

for vllm main branch

fix mamba_cache_mode align

fix fused_kimi_k3_mla_decode_q_concat_kv_cache_insert

use xspeedgate_ops.attn_res

use torch.ops.xspeedgate_ops.fused_recurrent_kda_packed_decode

fix cuda graph oom and precision issues

use xspeedgate_ops.l2norm_fwd and xspeedgate_ops.gather_and_maybe_dequant_cache

SituAndMul patched with situ_and_mul_kernel

use xspeedgate_ops.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert

enable chunked prefill

split chunk KDA to match vllm stages

Revert “split chunk KDA to match vllm stages”


This reverts commit 560809acae2ff9fa9737d555ed4290f22ca2fcfb.

remove cocopod

kda: use xspeedgate_ops fused_kda_gate_chunk_cumsum and layer_norm_gated_fwd

kda: route layer_norm_gated_fwd to xspeedgate_ops (sigmoid gate for o_norm)

torch native triton_merge_attn_states

mla prefill use fused_kimi_k3_mla_key_concat_kv_cache_insert

split chunk KDA to match vllm stages and use kunlun_ops.chunk_gla_fwd_o_gk

补回chunk_kda_with_fused_gate_fwd

kda:prefill use recompute_w_u_fwd_k3

kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context

Revert “kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context”


This reverts commit 87eb39a6928cf636cc4ee0e2b008636906baf565.

kimi-k3 use flashkda

use moe_sigmoid_group_topk_norm op

kda: use fused kunlun_ops causal_conv1d kernels


Replace the temporary torch implementations of causal_conv1d_fn and
causal_conv1d_update with kunlun_ops.causal_conv1d_fwd /
causal_conv1d_update.
Both kernels write in place, so the prefill path hands them a private
contiguous buffer (callers pass strided views into the packed QKV tensor and
still expect their input to survive) and the decode path runs on the caller’s
out buffer.
K3 stores conv1d weights in fp32 (ColumnParallelLinear(params_dtype=
torch.float32)) against fp16 activations. The 20260818 kunlun_ops build
accepts that combination and accumulates in fp32, which matches an fp64
reference to 5.99e-04 on prefill and 9.21e-08 on decode – an order of
magnitude tighter than casting the weight down to fp16 first.
Padded cuda-graph lanes carry conv_state_indices == -1; pad_slot_id=-1 makes
the kernel skip them instead of writing a real conv-state row, which is what
the previous clamp-to-slot-0 workaround had to do.
_conv1d_query_start_loc_cpu prefers the CPU mirror the GDN metadata already
keeps over .tolist(), to avoid a device sync on every layer.

kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context

fix mla and merge_attn_states

fix causal_conv1d_update used

fix some issues

kimi-k3: use moe_pre_sorted

kimi-k3 fix cuda grpah d2h copy

kimi-k3 fix cu_seqlens d2h copy

kda: route FusedRMSNormGated to the xspeedgate layer_norm_gated_fwd

kimi-k3: support expert parallel for W4A16 MoE

kimi-k3: support data parallel and mooncake PD

mla: use xspeedgate_ops concat_k_nope_k_pe for the chunked-context key concat

deepep: support the high-throughput all2all backend

kimi-k3 w4a8_int8 moe

deepep: support the low_latency backend

kda: gather/scatter causal_conv1d_update state on contiguous page rows

remove cpu() once


Signed-off-by: chengxiaokang chengxiaokang@baidu.com

mla: remove strided_slice_copy2d before gemm

perf: cache block_sizes device tensor to avoid per-step H2D sync


Signed-off-by: chengxiaokang chengxiaokang@baidu.com

support vllm 0.28.0

remove conv_1d_update input memory copy

revert conv_1d_update input memory copy, fix accuracy


Signed-off-by: baoqian426 1354987947@qq.com

Support MRV2

Signed-off-by: chengxiaokang 1783281729@qq.com

MRV2: support prefix cache

Signed-off-by: chengxiaokang 1783281729@qq.com

MRV2: add 3 kernels

Signed-off-by: chengxiaokang 1783281729@qq.com

use _prepare_prefill_inputs_torch

decode support dcp

remove conv_1d_update input memory copy

Revert “remove conv_1d_update input memory copy”


This reverts commit 8f25c9fc1282f732b17c8ec308bcc6f32d2c9b92.

Adapt native Mamba ops and support vLLM v0.28

Fix Kimi K3 decode convolution for token-strided input

Guard Scheduler._update_from_kv_xfer_finished against stale req_ids


vllm Scheduler._update_from_kv_xfer_finished asserts req_id in
self.requests while consuming finished_recving/finished_sending. When a
client batch-abort and a late KV-transfer-finished event land in the same
scheduler step, the req is already removed from self.requests and the
assert kills EngineCore (EngineDeadError) in PD-disaggregation.
Add a post-import hook that wraps the method to filter out stale req_ids
before delegating to the original, so the asserts never fire. Present
req_ids are untouched; stale transfers are reclaimed by the Mooncake
timeout GC.
Verified on the 6-node Kimi-K3 PD service: cancel-wave repro hit the race
33 times with 0 crashes (previously a single hit killed the engine).

kimi-k3 1/N

kimi-k3 2/N

kimi-k3 3/N

add KDA

kimi-k3: fix MoE routing/pre-sort and compressed-tensors scheme probing on P800



fused_moe_ct_w4a16: moe_sigmoid_group_topk_norm and moe_pre_sorted only
support expert_num <= 512 on P800 (the former returns ret=1 and leaves its
outputs untouched, so the uninitialized topk ids made the pre-sort kernels
read out of bounds). Add torch fallbacks for >512 experts, verified
bit-exact against the kernels at expert_num=512.
compressed_tensors_moe: probe the expert quant scheme with the checkpoint’s
projection names (Kimi-K3 uses w1/w3/w2) as well, otherwise the scheme
lookup misses and the MoE silently falls back to UnquantizedFusedMoEMethod.
kunlun_communicator: stop forwarding use_all2all to DeviceCommunicatorBase,
which derives it itself and rejects the kwarg.

Co-authored-by: Kiro

Revert “kimi-k3: fix MoE routing/pre-sort and compressed-tensors scheme probing on P800”

This reverts commit 00db81aca5d530da8dd6be5ab3f4e2e8c71c8cdc.

kimi-k3 4/N moe_sigmoid_group_topk_norm & moe_pre_sorted  do not support more than 512 experts

reorder_batch_threshold 512 → 1

fix recurrent_state oom

for vllm main branch

fix mamba_cache_mode align

fix fused_kimi_k3_mla_decode_q_concat_kv_cache_insert

use xspeedgate_ops.attn_res

use torch.ops.xspeedgate_ops.fused_recurrent_kda_packed_decode

fix cuda graph oom and precision issues

use xspeedgate_ops.l2norm_fwd and xspeedgate_ops.gather_and_maybe_dequant_cache

SituAndMul patched with situ_and_mul_kernel

use xspeedgate_ops.fused_kimi_k3_mla_decode_q_concat_kv_cache_insert

enable chunked prefill

split chunk KDA to match vllm stages

kda: use xspeedgate_ops fused_kda_gate_chunk_cumsum and layer_norm_gated_fwd

Revert “split chunk KDA to match vllm stages”


This reverts commit 560809acae2ff9fa9737d555ed4290f22ca2fcfb.

remove cocopod

kda: route layer_norm_gated_fwd to xspeedgate_ops (sigmoid gate for o_norm)

torch native triton_merge_attn_states

mla prefill use fused_kimi_k3_mla_key_concat_kv_cache_insert

补回chunk_kda_with_fused_gate_fwd

split chunk KDA to match vllm stages and use kunlun_ops.chunk_gla_fwd_o_gk

kda:prefill use recompute_w_u_fwd_k3

kda: use fused kunlun_ops causal_conv1d kernels


Replace the temporary torch implementations of causal_conv1d_fn and
causal_conv1d_update with kunlun_ops.causal_conv1d_fwd /
causal_conv1d_update.
Both kernels write in place, so the prefill path hands them a private
contiguous buffer (callers pass strided views into the packed QKV tensor and
still expect their input to survive) and the decode path runs on the caller’s
out buffer.
K3 stores conv1d weights in fp32 (ColumnParallelLinear(params_dtype=
torch.float32)) against fp16 activations. The 20260818 kunlun_ops build
accepts that combination and accumulates in fp32, which matches an fp64
reference to 5.99e-04 on prefill and 9.21e-08 on decode – an order of
magnitude tighter than casting the weight down to fp16 first.
Padded cuda-graph lanes carry conv_state_indices == -1; pad_slot_id=-1 makes
the kernel skip them instead of writing a real conv-state row, which is what
the previous clamp-to-slot-0 workaround had to do.
_conv1d_query_start_loc_cpu prefers the CPU mirror the GDN metadata already
keeps over .tolist(), to avoid a device sync on every layer.

kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context

Revert “kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context”


This reverts commit 87eb39a6928cf636cc4ee0e2b008636906baf565.

kimi-k3 use flashkda

use moe_sigmoid_group_topk_norm op

kimi-k3 use xspeedgate_ops gather_initial_states and mask_empty_context

fix mla and merge_attn_states

fix causal_conv1d_update used

fix some issues

kimi-k3: use moe_pre_sorted

kimi-k3 fix cuda grpah d2h copy

kimi-k3 fix cu_seqlens d2h copy

kda: route FusedRMSNormGated to the xspeedgate layer_norm_gated_fwd

kimi-k3: support expert parallel for W4A16 MoE

kimi-k3: support data parallel and mooncake PD

mla: use xspeedgate_ops concat_k_nope_k_pe for the chunked-context key concat

deepep: support the high-throughput all2all backend

kimi-k3 w4a8_int8 moe

deepep: support the low_latency backend

kda: gather/scatter causal_conv1d_update state on contiguous page rows

remove cpu() once


Signed-off-by: chengxiaokang chengxiaokang@baidu.com

mla: remove strided_slice_copy2d before gemm

perf: cache block_sizes device tensor to avoid per-step H2D sync


Signed-off-by: chengxiaokang chengxiaokang@baidu.com

support vllm 0.28.0

remove conv_1d_update input memory copy

revert conv_1d_update input memory copy, fix accuracy


Signed-off-by: baoqian426 1354987947@qq.com

Support MRV2

Signed-off-by: chengxiaokang 1783281729@qq.com

MRV2: support prefix cache

Signed-off-by: chengxiaokang 1783281729@qq.com

MRV2: add 3 kernels

Signed-off-by: chengxiaokang 1783281729@qq.com

use _prepare_prefill_inputs_torch

decode support dcp

remove conv_1d_update input memory copy

Revert “remove conv_1d_update input memory copy”


This reverts commit 8f25c9fc1282f732b17c8ec308bcc6f32d2c9b92.

Adapt native Mamba ops and support vLLM v0.28

Fix Kimi K3 decode convolution for token-strided input

Guard Scheduler._update_from_kv_xfer_finished against stale req_ids


vllm Scheduler._update_from_kv_xfer_finished asserts req_id in
self.requests while consuming finished_recving/finished_sending. When a
client batch-abort and a late KV-transfer-finished event land in the same
scheduler step, the req is already removed from self.requests and the
assert kills EngineCore (EngineDeadError) in PD-disaggregation.
Add a post-import hook that wraps the method to filter out stale req_ids
before delegating to the original, so the asserts never fire. Present
req_ids are untouched; stale transfers are reclaimed by the Mooncake
timeout GC.
Verified on the 6-node Kimi-K3 PD service: cancel-wave repro hit the race
33 times with 0 crashes (previously a single hit killed the engine).

execute pre-commit

Signed-off-by: chengxiaokang chengxiaokang@baidu.com

Signed-off-by: chengxiaokang chengxiaokang@baidu.com
Signed-off-by: baoqian426 1354987947@qq.com
Signed-off-by: chengxiaokang 1783281729@qq.com
Co-authored-by: baoqian426 1354987947@qq.com
Co-authored-by: wangjiasheng 1833390160@qq.com
Co-authored-by: youzeyu youzeyu@baidu.com
Co-authored-by: yangmingyu02 yangmingyu02@baidu.com
Co-authored-by: kiro <kiro@localhost>
Co-authored-by: chengxiaokang chengxiaokang@baidu.com
Co-authored-by: longxin 1351172767@qq.com