[LoRA][I] Add MOE LoRA JIT alignment kernel and tests

[yushengsu-thu]

Split this PR #14105 into 3 parts - Part I

Add JIT-compiled CUDA kernels for MOE LoRA block size alignment:

• moe_lora_align.py: JIT wrapper for moe_lora_align_block_size
• moe_lora_align_kernel.cu: CUDA kernels for token alignment, sorting, and expert counting
• test_moe_lora_align_block_size.py: Unit tests for the alignment kernel

Made-with: Cursor

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

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]

Warning

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

[yushengsu-thu]

Co-authored-by: Jonah Bernard jb2528@cornell.edu
Co-authored-by: cursor[bot] noreply@cursor.sh

[Copilot]

Pull request overview

Adds a JIT-compiled CUDA implementation for MoE+LoRA token alignment (block-size padding + per-expert sorting), along with a Python wrapper and a CUDA CI unit test, as part of the larger MoE LoRA enablement work split out from #14105.

Changes:

• Introduce moe_lora_align_block_size Python wrapper that JIT-loads a new CUDA kernel.
• Add CUDA kernels to build a per-LoRA token mask, align counts to block_size, and sort tokens by expert.
• Add a CUDA-registered pytest validating expert assignment and LoRA ownership of sorted blocks.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.

File	Description
python/sglang/jit_kernel/moe_lora_align.py	JIT loader + Python entrypoint for the new MOE LoRA alignment kernel.
python/sglang/jit_kernel/csrc/lora/moe_lora_align_kernel.cu	CUDA implementation for token masking, expert counting/padding, and sorting for MoE LoRA alignment.
python/sglang/jit_kernel/tests/test_moe_lora_align_block_size.py	CUDA CI test validating the alignment/sorting results.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yushengsu-thu]

/tag-and-rerun-ci

[yushengsu-thu]

/tag-and-rerun-ci