Skip to content

Enable native Kev inference with the Vulkan backend #23098

Description

@mergennachin

Enable the native Kev example on the Vulkan backend, extending the XNNPACK and MLX support added in #23023.

Kev runs prefill followed by pointer-head scoring, with an immutable prefix snapshot reused across question batches. The goal is to support this workflow through the existing C++ API on Vulkan-capable GPUs, without Python at inference.

Start with Kev's model/export code, the Vulkan export example, and the Vulkan backend. Qwen3.5 MoE's GatedDeltaNet implementation provides a reference for GDN math and state layouts; Vulkan will need its own lowering or shader support where coverage is missing.

Work to cover:

  • Identify operator and dynamic-shape gaps in Kev's prefill and score graphs. Add the required Vulkan support, including GDN, with focused backend regression tests. A small GDN parity test is a useful first step.
  • Add a Vulkan export option and CMake linkage so the existing kev_runner and kev_benchmark run the exported model through Module.
  • Preserve system_one, explicit prefill/evaluate, configurable token limits, variable question/option counts, and batching beyond eight questions. Repeated evaluations must leave the prefix unchanged.
  • Start with an unquantized FP32 path, retaining the fitted temperature and checkpoint metadata. Check logits, probabilities, and reused state against upstream Kev's PyTorch implementation and the existing FP32 path, with documented tolerances.
  • Verify the backbone, including GDN, executes on Vulkan and report any CPU fallback. Document a tested GPU, precision requirements, export/build/run commands, and measurements using kev_benchmark.

Keep the integration in the existing example and reuse the runner and benchmark. A Vulkan-capable GPU and the Vulkan SDK are needed for validation.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    good first issueGood for newcomersmodule: examplesIssues related to demos under examples/module: vulkanIssues related to the Vulkan delegate and code under backends/vulkan/

    Type

    No type

    Projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions