Skip to content

[jit kernel] support dtype as a cpp template parameter#16452

Merged
BBuf merged 8 commits intosgl-project:mainfrom
yingluosanqian:dtype_wrapper
Jan 8, 2026
Merged

[jit kernel] support dtype as a cpp template parameter#16452
BBuf merged 8 commits intosgl-project:mainfrom
yingluosanqian:dtype_wrapper

Conversation

@yingluosanqian
Copy link
Copy Markdown
Collaborator

Motivation

In the current jit kernel implementation, dtype can not be passed as a cpp template parameters, which in some cases leads to larger binary size and longer compilation time. For example, when a kernel supports multiple dtype we have to generate separate kernels for each dtype instead of generating a single templated kernel.

This PR support dtype as a cpp template parameter.

Modifications

This PR modify make_cpp_args to deal with dtype template parameters. Then optimize the norm kernel to validate this modification.

Accuracy Tests

run test_qknorm and passed.

python python/sglang/jit_kernel/tests/test_qknorm.py

Benchmarking and Profiling

None.

Checklist

None.

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 (/tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci) or contact authorized users to do so.
  4. After green CI and required approvals, ask Merge Oncalls to merge.

@gemini-code-assist
Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@yingluosanqian yingluosanqian marked this pull request as ready for review January 5, 2026 05:45
@github-actions github-actions bot added Multi-modal multi-modal language model diffusion SGLang Diffusion labels Jan 5, 2026
@yingluosanqian
Copy link
Copy Markdown
Collaborator Author

/tag-run-ci-label

@DarkSharpness
Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions bot added the run-ci label Jan 5, 2026
@yhyang201
Copy link
Copy Markdown
Collaborator

/rerun-failed-ci

@DarkSharpness DarkSharpness added high priority ready-to-merge The PR is ready to merge after the CI is green. labels Jan 6, 2026
@yhyang201
Copy link
Copy Markdown
Collaborator

/rerun-failed-ci

@BBuf BBuf self-assigned this Jan 8, 2026
@BBuf BBuf merged commit 48b8dcd into sgl-project:main Jan 8, 2026
214 of 229 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

diffusion SGLang Diffusion high priority Multi-modal multi-modal language model ready-to-merge The PR is ready to merge after the CI is green. run-ci

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants