Skip to content

[TRTLLM-15498][refactor] consolidate Kimi KDA production frontend #17822

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

Merged
jiaganc merged 3 commits into from
Aug 23, 2026
Merged

[TRTLLM-15498][refactor] consolidate Kimi KDA production frontend #17822

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c86e2f7
[None][refactor] consolidate Kimi KDA production frontend
jiaganc Aug 20, 2026
09253a7
[fix] update KDA layer-wise benchmark marker
jiaganc Aug 21, 2026
0d5f955
[fix] clean up stale Kimi KDA runtime test
jiaganc Aug 22, 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,011 changes: 28 additions & 983 deletions tensorrt_llm/_torch/models/modeling_kimi_linear.py
View file
Open in desktop

Large diffs are not rendered by default.

7 changes: 2 additions & 5 deletions tensorrt_llm/_torch/modules/kimi_kda/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,6 @@
interface.
"""

from .kimi_kda_mixer import KimiKDAKernelPath, KimiKDALinearAttention
from .kimi_kda_mixer import KimiKDALinearAttention

__all__ = [
"KimiKDAKernelPath",
"KimiKDALinearAttention",
]
__all__ = ["KimiKDALinearAttention"]
17 changes: 1 addition & 16 deletions tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,10 +101,6 @@ def is_intree_prefill_available() -> bool:
return False


def _load_fla_chunk_kda() -> ModuleType:
return importlib.import_module("fla.ops.kda")


# ---------------------------------------------------------------------------
# In-tree KDA multi-token verify op (CuTe DSL, trtllm::kda_mtp_decode).
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -209,16 +205,6 @@ def __init__(
f"verify={self.verify_kernel_path}"
)

def get_prefill_source(self) -> str:
if self.prefill_kernel_path == "optimized":
return _load_prefill_module().__file__ or "<custom_ops.cute_dsl_kimi_k3>"
return _load_fla_chunk_kda().__file__ or "<fla.ops.kda>"

def get_decode_source(self) -> str:
if self.decode_kernel_path == "optimized":
return _kda_decode.__file__ or "<kimi_kda._kda_decode>"
return _load_fla_chunk_kda().__file__ or "<fla.ops.kda>"

def mtp_verify(self, **kwargs) -> torch.Tensor:
"""Run the fused KDA multi-token verify kernel.

Expand Down Expand Up @@ -354,7 +340,7 @@ def prefill_chunk_kda(

from fla.ops.kda import chunk_kda

o, final_state = chunk_kda(
return chunk_kda(
q=q,
k=k,
v=v,
Expand All @@ -373,7 +359,6 @@ def prefill_chunk_kda(
state_v_first=True,
cu_seqlens=cu_seqlens,
)
return o, final_state

def decode_kda(self, **kwargs) -> torch.Tensor:
"""Run the fused KDA single-token decode kernel.
Expand Down
Loading
Loading