[Kimi-K3][NPU] Support Kimi-K3 on NPU (#33465)

2026-08-12 21:20:14 +08:00
parent dc5f6c4883
commit 197832bcf5
24 changed files with 1294 additions and 129 deletions
@@ -558,7 +558,10 @@ class MambaPool:
                    )

                    conv_state = _init_npu_conv_state(
-                        conv_state[0], conv_state_shape, speculative_num_draft_tokens
+                        conv_state[0],
+                        conv_state_shape,
+                        speculative_num_draft_tokens,
+                        is_kda=cache_params.is_kda,
                    )

                if _is_cpu and _cpu_has_amx_support:
@@ -759,6 +762,14 @@ class MambaPool:
                    # Original dense layout (NPU/CPU, or EAGLE tree verify): one
                    # [dim, K-1] window per draft token.
                    # Shape: [num_layers, size+1, draft_tokens, dim, K-1]
+                    dense_conv_shapes = [
+                        (
+                            (conv_shape[1], conv_shape[0])
+                            if _is_npu and cache_params.is_kda
+                            else conv_shape
+                        )
+                        for conv_shape in conv_state_shape
+                    ]
                    intermediate_conv_window_cache = [
                        torch.zeros(
                            size=(
@@ -771,7 +782,7 @@ class MambaPool:
                            dtype=conv_dtype,
                            device="cuda",
                        )
-                        for conv_shape in conv_state_shape
+                        for conv_shape in dense_conv_shapes
                    ]
                    self._intermediate_conv_window_phys = intermediate_conv_window_cache
                self.mamba_cache = self.SpeculativeState(