Migrate KimiVLImagePixelInputs to TensorSchema

[bbeckca]

Purpose

This PR migrates KimiVLImagePixelInputs from a TypedDict-based definition to a structured TensorSchema model with runtime shape validation. This brings it in line with recent changes to Phi3VImagePixelInputs, and is part of a broader effort to improve input contract enforcement and debug-ability across multi-modal models.

Test Plan

Confirm validation works via standalone tests in tests/standalone_test/test_tensor_schema.py and rely on CI to check integration.

Test Result

(venv) benjibeck@Benjis-MBP vllm % python3 -m pytest tests/standalone_tests/test_tensor_schema.py -v --log-cli-level=DEBUG
=========================================== test session starts ============================================
platform darwin -- Python 3.9.6, pytest-8.4.1, pluggy-1.6.0 -- /Users/benjibeck/Projects/vllm/venv/bin/python3
cachedir: .pytest_cache
rootdir: /Users/benjibeck/Projects/vllm
configfile: pyproject.toml
plugins: anyio-4.9.0
collected 14 items                                                                                         

tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_valid_tensor PASSED                 [  7%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_optional_fields PASSED              [ 14%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_constant_dim_failure PASSED         [ 21%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_symbolic_dim_mismatch PASSED        [ 28%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_list_tensor_valid PASSED            [ 35%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_variable_patch_counts_valid PASSED  [ 42%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_tuple_tensor_valid PASSED           [ 50%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_inconsistent_shapes_in_list PASSED  [ 57%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_empty_list PASSED                   [ 64%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_validation_disabled_skips_shape_check PASSED [ 71%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_with_valid_resolve_binding_dims PASSED [ 78%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_with_invalid_resolve_binding_dims PASSED [ 85%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_with_list_of_symbolic_dim PASSED    [ 92%]
tests/standalone_tests/test_tensor_schema.py::test_tensor_schema_with_list_of_symbolic_dim_mismatch_in_length PASSED [100%]

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request is a great improvement. Migrating KimiVLImagePixelInputs from a TypedDict to a TensorSchema enhances type safety and introduces runtime shape validation, which improves the robustness and debuggability of the multi-modal input processing. The implementation is clean, and removing the manual assertion in favor of the schema-based validation is a solid refactoring step. The code changes look correct and align well with the stated objectives. Well done!

[bbeckca]

@DarkLight1337 @Isotr0py Heads up (not related to this diff), but I came across some prior discussions around Torchscript compatibility of Annotated modules with compilation. Seems there's some workarounds, if we encounter such but proactively flagging in case this is worth investigating further. patrick-kidger/torchtyping#13

[DarkLight1337]

Let's wait for #21754 to be merged first so we can test the correctness without worrying about OOM

[DarkLight1337]

but I came across some prior discussions around Torchscript compatibility of Annotated modules with compilation

torch.compile is currently not applied to the multi-modal part of the model so it shouldn't be a problem for now