[KDA] Add FlashInfer SM100 KDA decode + MTP (target_verify) backend (#30113)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-07-15 15:04:20 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent 1afab30577
commit a649b5a9db
6 changed files with 1077 additions and 39 deletions
+7 -9
View File
@@ -102,21 +102,19 @@ _use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip
def conv_window_dedup_enabled(
is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int]
is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int], is_kda: bool
) -> bool:
"""Whether the deduplicated sliding-window conv-intermediate layout is safe.
It is only correct for a *linear* draft chain (``speculative_eagle_topk <= 1``,
i.e. NEXTN / MTP): consecutive draft tokens then form a true sliding window, so
the overlapping physical columns hold identical values. Under EAGLE *tree*
verify (``topk > 1``) the conv kernel walks per-token tree ancestors, so aliased
columns can need different values from different parent chains -> fall back to
the dense layout. NPU/CPU also keep the dense layout (their kernels assume
contiguous per-step windows). See ``MambaPool.__init__``.
It is safe for CUDA linear draft chains whose kernels consume the window raw.
Tree verify, NPU/CPU, and KDA keep dense windows: tree ancestors need independent
windows, platform kernels expect contiguous steps, and KDA transposes the window
before conv so the overlapping ``as_strided`` layout would corrupt stores.
"""
return (
not is_npu
and not is_cpu
and not is_kda
and (speculative_eagle_topk is None or speculative_eagle_topk <= 1)
)
@@ -576,7 +574,7 @@ class MambaPool:
# `fused_conv_window_scatter_with_mask` scatter is layout-agnostic,
# so the dense fallback reads correctly through the same code path.
dedup_conv_window = conv_window_dedup_enabled(
_is_npu, _is_cpu, speculative_eagle_topk
_is_npu, _is_cpu, speculative_eagle_topk, cache_params.is_kda
)
self._intermediate_conv_window_phys = []
if dedup_conv_window: