[Feature] Introduce mypy in a gradual rollout way

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Introduce mypy into codebase would help catch low-level bugs early, like previously done in vllm, jax, etc. Follow a gradual rollout plan to achieve this.

Rollout plan

Phase 1

• python/sgl_jax/srt/models
• python/sgl_jax/srt/layers
• python/sgl_jax/srt/utils
• python/sgl_jax/srt/model_loader
• python/sgl_jax/srt/model_executor

Suggested Workflow

1. Select a file/folder to work on (e.g., python/sgl_jax/srt/models/grok.py)
2. Add the file path to files = [ ... ] under [tool.mypy] in python/pyproject.toml
3. Run pre-commit run --all-files to see mypy errors
4. Fix the errors, typically by adjusting type annotations until pre-commit run --all-files passes
5. Validate locally by running mypy in your local environment with all dependencies:
   • Why? Because pre-commit runs mypy in an isolated virtualenv without your dependencies. You may need to ensure additional_dependencies in .pre-commit-config.yaml are properly configured for accurate type checking.
   • How? Rewrite the settings in [tool.mypy] section from python/pyproject.toml as a mypy.ini file and run mypy --config-file mypy.ini [filename]
   • If results differ between local runs and pre-commit, adjust additional_dependencies: in .pre-commit-config.yaml (mypy section) to include only necessary libraries until results match
6. Continue with other files, or submit a PR for completed changes

Reference: See #474 as a working example of this workflow.

Side Notes

• Work file-by-file within each folder (e.g., within models/, finish one model before moving to the next)
• Update related files as needed—type changes in one file may require updates in other modules or tests
• Consider type implications—decisions like optional (None) types can affect the entire pipeline. Avoid allowing None everywhere; if uncertain, discuss with the team
• Ensure all tests pass (pre-commit & CI) for each refactored component before proceeding

Related resources

https://docs.jax.dev/en/latest/jep/12049-type-annotations.html https://github.com/vllm-project/vllm/pull/25313/files#diff-6b980d504c08a48078accc6b5fb392b7b5d4c65dfe293a3956a7109200f184e2