[Test] Enhance JIT kvcache store kernel test coverage

[xingsy97]

Motivation

The existing test_store_cache only tests bf16 + int64 indices. This leaves several CUDA template paths untested (#17035).

Modifications

Add 4 tests to test_store_cache.py:

• test_store_cache_dtypes: fp16/bf16/fp32
• test_store_cache_int32_indices: int32 index path (would have caught Fix wrong kernel selection for int32/int64 indices #16912)
• test_store_cache_num_split: explicit num_split=1/2/4
• test_can_use_store_cache: guard function

Original tests unchanged.

Accuracy Tests

N/A — this PR only adds tests, no changes to kernel or model code.

Benchmarking and Profiling

N/A — no inference speed impact, test-only change.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the test coverage for the JIT KV cache store kernel. The changes introduce several new test functions designed to validate the kernel's behavior under a wider range of data types, index tensor types, and num_split configurations. This ensures the kernel's stability and correctness across various operational scenarios, improving the overall reliability of the system.

Highlights

• Enhanced JIT KV Cache Test Coverage: Expanded the test suite for the JIT KV cache store kernel to cover more CUDA template code paths and configurations, addressing a known gap in existing tests.
• New Test Cases for Data Types: Introduced test_store_cache_dtypes to validate the kernel's behavior across fp16, bf16, and fp32 data types, ensuring correct handling of different kAlignment branches.
• Validation for int32 Indices: Added test_store_cache_int32_indices to specifically test the int32_t template instantiation for index tensors, which was previously untested and could have caught issues like Fix wrong kernel selection for int32/int64 indices #16912.
• Comprehensive num_split Testing: Implemented test_store_cache_num_split to explicitly verify each valid num_split kernel path (1, 2, 4), ensuring robustness across different splitting strategies.
• Guard Function Verification: Included test_can_use_store_cache to test the can_use_store_cache guard function, ensuring it correctly identifies valid and invalid row_bytes values for the JIT kernel.
• Optimized Test Execution: The new tests utilize a smaller cache size (4096 rows) and representative parameter subsets to ensure faster execution without compromising coverage.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/jit_kernel/tests/test_store_cache.py
  • Imported can_use_store_cache from sglang.jit_kernel.kvcache.
  • Added REPR_BS, REPR_DIMS, and SMALL_CACHE constants for new test cases.
  • Introduced test_store_cache_dtypes to test different floating-point data types.
  • Added test_store_cache_int32_indices to verify functionality with int32 index tensors.
  • Implemented _valid_num_splits helper function to determine valid num_split values.
  • Created test_store_cache_num_split to test various num_split configurations.
  • Added test_can_use_store_cache to validate the can_use_store_cache guard function.

Activity

• The pull request author, xingsy97, created this PR to address the "Enhance tests on JIT kernel" item from [Roadmap] JIT kernel development #17035.
• The PR description indicates that no accuracy tests or benchmarking were performed as the changes are test-only and do not impact inference speed or kernel/model code.
• The author has provided a checklist for code formatting, unit tests, documentation, and benchmarking, with some items marked as N/A or unchecked.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request significantly enhances the test coverage for the JIT kvcache store kernel by adding tests for various data types, index types, and num_split configurations. The new tests are well-structured and use smaller, representative datasets for efficiency. I've identified a key area for improvement in the num_split test to ensure it covers all data types, and a minor refinement for code clarity and efficiency. Overall, these are valuable additions that increase confidence in the kernel's correctness.

[HydraQYH]

LGTM

[DarkSharpness]

my only concern is do we need to assert not supported? we might add support for corner cases later

[xingsy97]

removed assert not cases. could you help trigger CI tests for this PR?

[DarkSharpness]

/tag-and-rerun-ci

[xingsy97]

/rerun-failed-ci

[DarkSharpness]

/rerun-failed-ci

[DarkSharpness]

/rerun-stage jit-kernel-unit-test

[github-actions]

❌ Stage jit-kernel-unit-test doesn't support isolated runs yet.

NVIDIA stages:

• stage-a-test-1
• stage-a-cpu-only
• stage-b-test-small-1-gpu
• stage-b-test-large-1-gpu
• stage-b-test-large-2-gpu
• stage-b-test-4-gpu-b200
• stage-c-test-4-gpu-h100
• stage-c-test-8-gpu-h200
• stage-c-test-8-gpu-h20
• stage-c-test-4-gpu-b200
• stage-c-test-4-gpu-gb200
• stage-c-test-deepep-4-gpu
• stage-c-test-deepep-8-gpu-h200
• multimodal-gen-test-1-gpu
• multimodal-gen-test-2-gpu

AMD stages:

• sgl-kernel-unit-test-amd
• sgl-kernel-unit-test-2-gpu-amd
• stage-a-test-1-amd
• stage-b-test-small-1-gpu-amd
• stage-b-test-small-1-gpu-amd-nondeterministic
• stage-b-test-small-1-gpu-amd-mi35x
• stage-b-test-large-1-gpu-amd
• stage-b-test-large-2-gpu-amd
• multimodal-gen-test-1-gpu-amd
• multimodal-gen-test-2-gpu-amd
• stage-c-test-large-8-gpu-amd
• stage-c-test-large-8-gpu-amd-mi35x

Other stages will be added soon. For now, use /rerun-failed-ci for those stages.

[DarkSharpness]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[xingsy97]

Hi @DarkSharpness @HydraQYH I noticed this PR has been blocked by unrelated CI failures. Even multiple CI reruns didn't help. Could you help take a look at the CI or merge it directly if acceptable? thanks.