Skip to content

feat(k3): Kimi-K3 vLLM plugin vision + DSpark draft support #1910

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
sajandhy wants to merge 5 commits into
base: main
Choose a base branch
from
Open

feat(k3): Kimi-K3 vLLM plugin vision + DSpark draft support #1910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b699f12
feat(k3): reapply remaining K3 DSpark/vision delta on main
whx-sjtu Aug 16, 2026
a4a4dfc
Added changes required due to other changes on main.
sajandhy Aug 17, 2026
916074a
Fix Black and Ruff warnings.
sajandhy Aug 18, 2026
af2b8da
fix(k3): restore vision imports dropped during rebase
sajandhy Aug 19, 2026
1b937ea
Removed unnecessary files.
sajandhy Sep 6, 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
1 change: 1 addition & 0 deletions atom/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -675,6 +675,7 @@ def glm5_kpool_block_size(index_kpool: int) -> int:
# multimodal models fully supported by plugin mode
_PLUGIN_SUPPORTED_MULTIMODAL_MODELS: set[str] = {
"kimi_k25",
"kimi_k3",
"qwen3_5",
"qwen3_5_moe",
}
Expand Down
2 changes: 1 addition & 1 deletion atom/model_ops/fla_ops/chunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ def chunk_gated_delta_rule(
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float = None,
scale: float | None = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
cu_seqlens: torch.LongTensor | None = None,
Expand Down
2 changes: 1 addition & 1 deletion atom/model_ops/fla_ops/chunk_vk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ def chunk_gated_delta_rule_vk(
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float = None,
scale: float | None = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
cu_seqlens: torch.Tensor | None = None,
Expand Down
2 changes: 1 addition & 1 deletion atom/model_ops/fla_ops/fused_sigmoid_gating.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def fused_sigmoid_gating_delta_rule_update(
o: torch.Tensor | None = None,
beta: float = 1.0,
threshold: float = 20.0,
scale: float = None,
scale: float | None = None,
initial_state: torch.Tensor = None,
inplace_final_state: bool = True,
cu_seqlens: torch.LongTensor | None = None,
Expand Down
2 changes: 1 addition & 1 deletion atom/model_ops/fla_ops/layernorm_guard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ def layer_norm_fwd(
eps: float,
z: torch.Tensor = None,
out: torch.Tensor = None,
group_size: int = None,
group_size: int | None = None,
norm_before_gate: bool = True,
is_rms_norm: bool = False,
):
Expand Down
2 changes: 1 addition & 1 deletion atom/model_ops/paged_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def __init__(
head_dim,
scale,
num_kv_heads,
alibi_slopes: list[float] = None,
alibi_slopes: list[float] | None = None,
kv_cache_dtype="bf16",
layer_num=0,
use_mla: bool = False,
Expand Down
15 changes: 10 additions & 5 deletions atom/models/kimi_k3_dspark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -545,19 +545,24 @@ def __init__(
prefix=f"{prefix}.mlp.gate_up_proj",
)

def write_context_kv(self, ctx_hidden, positions) -> None:
def write_context_kv(self, ctx_hidden, positions, slot_mapping=None) -> None:
"""Populate this layer's context rows.

The context source does NOT go through ``input_layernorm``: the reference
feeds every layer the same ``context_norm(context_proj(aux))`` tensor
straight into the KV projection, while ``input_layernorm`` applies only
to the residual stream carrying the draft block.

``slot_mapping`` is optional: the native path reads it from

@valarLip valarLip Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why can't plugin follow the same way?

@sajandhy sajandhy Aug 19, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Native ATOM can use the default path because DSparkProposer calls write_context_kv immediately after the target forward, while forward_context().attn_metadata still describes the target step — so slot_mapping[:N] is the verified-token slots we need.

In the vLLM plugin, context KV is written from the speculator via precompute_and_store_context_kv(..., slot_mappings=...), not from ATOM's proposer/forward context. vLLM builds those slots from block tables, and they can differ per layer when the draft spans multiple KV cache groups — so the plugin passes slot_mappings explicitly rather than reading a single global attn_metadata.slot_mapping.

Both paths end up in the same write_context_kv → write_context_kv_latent store; only the source of the slot indices differs. During the later draft backbone pass, attn_metadata.slot_mapping would describe the draft block, not the context rows — so we can't rely on forward context for this write in either design.

``forward_context``, while the vLLM plugin passes per-layer slots
explicitly (groups may differ across layers).
"""
from atom.utils.forward_context import get_forward_context
if slot_mapping is None:
from atom.utils.forward_context import get_forward_context

slot_mapping = get_forward_context().attn_metadata.slot_mapping[
: ctx_hidden.shape[0]
]
slot_mapping = get_forward_context().attn_metadata.slot_mapping[
: ctx_hidden.shape[0]
]
self.self_attn.write_context_kv(ctx_hidden, positions, slot_mapping)

def forward(
Expand Down
30 changes: 14 additions & 16 deletions atom/plugin/vllm/attention/layer_mha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING

import aiter
import torch
Expand Down Expand Up @@ -89,17 +89,17 @@ def __init__(
head_dim,
scale,
num_kv_heads,
alibi_slopes: list[float] = None,
alibi_slopes: list[float] | None = None,
kv_cache_dtype="bf16",
layer_num=0,
use_mla: bool = False,
mla_modules: Optional[MLAModules] = None,
sinks: Optional[nn.Parameter] = None,
per_layer_sliding_window: Optional[int] = None,
rotary_emb: Optional[torch.nn.Module] = None,
prefix: Optional[str] = None,
q_norm: Optional[torch.nn.Module] = None,
k_norm: Optional[torch.nn.Module] = None,
mla_modules: MLAModules | None = None,
sinks: nn.Parameter | None = None,
per_layer_sliding_window: int | None = None,
rotary_emb: torch.nn.Module | None = None,
prefix: str | None = None,
q_norm: torch.nn.Module | None = None,
k_norm: torch.nn.Module | None = None,
**kwargs,
):
from vllm.v1.attention.backend import AttentionType
Expand Down Expand Up @@ -193,7 +193,7 @@ def forward(
key: torch.Tensor,
value: torch.Tensor,
positions: torch.Tensor = None,
q_scale: Optional[torch.Tensor] = None,
q_scale: torch.Tensor | None = None,
qkv: torch.Tensor = None,
**kwargs,
):
Expand Down Expand Up @@ -527,8 +527,6 @@ def paged_attention_asm(
high_precision=0,
)

return

def extend_for_sliding_window(
self,
attn_metadata: "AiterMhaMetadataForVllm",
Expand All @@ -539,8 +537,8 @@ def extend_for_sliding_window(
cu_seqlens_q: torch.Tensor,
max_seqlen_q: int,
block_table: torch.Tensor,
k_scale: Optional[torch.Tensor],
v_scale: Optional[torch.Tensor],
k_scale: torch.Tensor | None,
v_scale: torch.Tensor | None,
):
assert attn_metadata.extend_metadata is not None
assert attn_metadata.extend_metadata.chunk_context_metadata is not None
Expand Down Expand Up @@ -613,8 +611,8 @@ def extend_forward(
min_seqlen_q: int,
block_table: torch.Tensor,
slot_mapping: torch.Tensor,
k_scale: Optional[torch.Tensor],
v_scale: Optional[torch.Tensor],
k_scale: torch.Tensor | None,
v_scale: torch.Tensor | None,
):
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states

Expand Down
18 changes: 13 additions & 5 deletions atom/plugin/vllm/attention/layer_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,12 @@ def forward_impl(
q = q[:num_actual_toks, ...]
k_c_normed = k_c_normed[:num_actual_toks, ...]
k_pe = k_pe[:num_actual_toks, ...]
# Model runner V2 sizes slot_mapping to the CUDA-graph token width and
# marks the tail with PAD_SLOT_ID, but leaves num_actual_tokens unpadded
# on the piecewise path; V1 never pads it. The cache kernels launch one
# block per slot and index q/kv by that block, so slot_mapping must not
# outrun the tensors sliced above. Same trim layer_mha already applies.
slot_mapping = attn_metadata.slot_mapping[:num_actual_toks]

decode_q = q[:num_decode_tokens]
prefill_q = q[num_decode_tokens:]
Expand Down Expand Up @@ -1056,7 +1062,7 @@ def forward_impl(
k_c_normed,
self.rotary_emb_cos_sin_cache,
self.rotary_emb.is_neox_style,
attn_metadata.slot_mapping,
slot_mapping,
kv_cache,
self.kv_cache_dtype,
self._k_scale,
Expand All @@ -1068,7 +1074,7 @@ def forward_impl(
k_c_normed,
k_pe.squeeze(1),
kv_cache,
attn_metadata.slot_mapping.flatten(),
slot_mapping.flatten(),
kv_cache_dtype=self.kv_cache_dtype,
scale=self._k_scale,
)
Expand Down Expand Up @@ -1166,7 +1172,7 @@ def forward_impl(
self.kv_lora_rank + self.qk_rope_head_dim,
),
decode_q[:, : self.num_heads] if fused_q_head_pad else decode_q,
attn_metadata.slot_mapping,
slot_mapping,
self._k_scale,
self._q_scale,
positions,
Expand Down Expand Up @@ -1338,6 +1344,7 @@ def forward_impl_sparse(
q = q[:num_actual_toks, ...]
k_c_normed = k_c_normed[:num_actual_toks, ...]
k_pe = k_pe[:num_actual_toks, ...].unsqueeze(1)
slot_mapping = sparse_meta.slot_mapping[:num_actual_toks]

positions = None
if self._is_vllm_forward_context_available():
Expand Down Expand Up @@ -1416,7 +1423,7 @@ def forward_impl_sparse(
kv_cache.shape[0], -1, self.kv_lora_rank + self.qk_rope_head_dim
),
q_out,
sparse_meta.slot_mapping,
slot_mapping,
self._k_scale,
self._q_scale,
positions,
Expand Down Expand Up @@ -1477,8 +1484,9 @@ def forward(
def get_kv_cache_spec(self, vllm_config):
from vllm.v1.kv_cache_interface import MLAAttentionSpec

block_size = vllm_config.cache_config.block_size
return MLAAttentionSpec(
block_size=vllm_config.cache_config.block_size,
block_size=block_size,
num_kv_heads=1,
head_size=self.head_size,
dtype=self.kv_cache_torch_dtype,
Expand Down
50 changes: 24 additions & 26 deletions atom/plugin/vllm/attention/minimax_m3_attnetion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@
MiniMax-M3's q/k norm + RoPE transform.
"""

from typing import Optional

import aiter
import torch
from aiter import dtypes
from torch import nn
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

from atom.config import get_current_atom_config
from atom.model_ops.minimax_m3.sparse_attn import (
Expand All @@ -31,8 +31,6 @@
_register_vllm_static_forward_context,
)
from atom.utils import mark_spliting_op
from vllm.forward_context import get_forward_context
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

_MINIMAX_M3_TOPK_CACHE_STATE: dict = {}

Expand Down Expand Up @@ -75,8 +73,8 @@ def __init__(
head_dim: int,
kv_cache_dtype: str,
) -> None:
from vllm.v1.attention.backend import AttentionType
from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
from vllm.v1.attention.backend import AttentionType

super().__init__()
atom_config = get_current_atom_config()
Expand Down Expand Up @@ -140,19 +138,19 @@ def __init__(
head_dim: int,
scale: float,
num_kv_heads: int,
alibi_slopes: Optional[list[float]] = None,
alibi_slopes: list[float] | None = None,
kv_cache_dtype: str = "bf16",
layer_num: int = 0,
use_mla: bool = False,
rotary_emb: Optional[nn.Module] = None,
prefix: Optional[str] = None,
q_norm: Optional[nn.Module] = None,
k_norm: Optional[nn.Module] = None,
rotary_emb: nn.Module | None = None,
prefix: str | None = None,
q_norm: nn.Module | None = None,
k_norm: nn.Module | None = None,
cache_config=None,
quant_config=None,
index_q_norm: Optional[nn.Module] = None,
index_k_norm: Optional[nn.Module] = None,
index_rotary_emb: Optional[nn.Module] = None,
index_q_norm: nn.Module | None = None,
index_k_norm: nn.Module | None = None,
index_rotary_emb: nn.Module | None = None,
index_q_size: int = 0,
index_head_dim: int = 0,
topk: int = 0,
Expand Down Expand Up @@ -670,8 +668,8 @@ def _run_sparse_attention(
def _forward_with_output(
self,
qkv: torch.Tensor,
positions: Optional[torch.Tensor] = None,
output: Optional[torch.Tensor] = None,
positions: torch.Tensor | None = None,
output: torch.Tensor | None = None,
) -> torch.Tensor:
main_metadata, index_metadata = self._metadata_for_layer()
num_tokens = qkv.shape[0]
Expand Down Expand Up @@ -714,9 +712,9 @@ def forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
positions: Optional[torch.Tensor] = None,
q_scale: Optional[torch.Tensor] = None,
qkv: Optional[torch.Tensor] = None,
positions: torch.Tensor | None = None,
q_scale: torch.Tensor | None = None,
qkv: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
del query, key, value, q_scale, kwargs
Expand All @@ -741,14 +739,14 @@ def __init__(
head_dim: int,
scale: float,
num_kv_heads: int,
alibi_slopes: Optional[list[float]] = None,
alibi_slopes: list[float] | None = None,
kv_cache_dtype: str = "bf16",
layer_num: int = 0,
use_mla: bool = False,
rotary_emb: Optional[nn.Module] = None,
prefix: Optional[str] = None,
q_norm: Optional[nn.Module] = None,
k_norm: Optional[nn.Module] = None,
rotary_emb: nn.Module | None = None,
prefix: str | None = None,
q_norm: nn.Module | None = None,
k_norm: nn.Module | None = None,
cache_config=None,
quant_config=None,
**kwargs,
Expand Down Expand Up @@ -878,9 +876,9 @@ def forward(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
positions: Optional[torch.Tensor] = None,
q_scale: Optional[torch.Tensor] = None,
qkv: Optional[torch.Tensor] = None,
positions: torch.Tensor | None = None,
q_scale: torch.Tensor | None = None,
qkv: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
del query, key, value, q_scale, kwargs
Expand Down
22 changes: 20 additions & 2 deletions atom/plugin/vllm/kda_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,29 @@ def build( # type: ignore[override]
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
fast_build=fast_build,
)
self._stage_packed_decode_query_start_loc(common_attn_metadata, metadata)
self._adapt_full_graph_decode_metadata(common_attn_metadata, metadata)
return metadata

def _stage_packed_decode_query_start_loc(
self,
common_attn_metadata: CommonAttentionMetadata,
metadata: KimiK3KDAMetadata,
) -> None:
if metadata.num_decodes <= 0 or metadata.non_spec_query_start_loc is not None:
return

batch_size = int(common_attn_metadata.num_reqs)
query_start_loc_buf = self.non_spec_query_start_loc[: batch_size + 1]
torch.arange(
metadata.num_decodes + 1,
dtype=torch.int32,
device=query_start_loc_buf.device,
out=query_start_loc_buf[: metadata.num_decodes + 1],
)
query_start_loc_buf[metadata.num_decodes + 1 :].fill_(metadata.num_decodes)
metadata.non_spec_query_start_loc = query_start_loc_buf

def _adapt_full_graph_decode_metadata(
self,
common_attn_metadata: CommonAttentionMetadata,
Expand All @@ -81,8 +101,6 @@ def _adapt_full_graph_decode_metadata(
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
real_decode_mask_cpu = query_lens_cpu > 0
real_num_decodes = int(real_decode_mask_cpu.sum().item())
if real_num_decodes == metadata.num_decodes:
return

batch_size = int(common_attn_metadata.num_reqs)
if batch_size > self.decode_cudagraph_max_bs:
Expand Down
2 changes: 1 addition & 1 deletion atom/plugin/vllm/model_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ def _maybe_set_v4_expert_dtype(atom_config, vllm_config) -> None:
"Qwen3_5ForConditionalGeneration": "atom.plugin.vllm.models.qwen3_5:Qwen3_5ForConditionalGeneration_",
"KimiK25ForConditionalGeneration": "atom.plugin.vllm.models.kimi_k25:KimiK25ForConditionalGeneration_",
"KimiK3ForConditionalGeneration": (
"atom.plugin.vllm.models.kimi_k3:KimiK3ForCausalLM"
"atom.plugin.vllm.models.kimi_k3:KimiK3ForConditionalGeneration_"
),
"MiniMaxM2ForCausalLM": "atom.models.minimax_m2:MiniMaxM2ForCausalLM",
"DeepseekV4ForCausalLM": "atom.plugin.vllm.models.deepseek_v4:DeepseekV4ForCausalLM",
Expand Down
Loading