You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
perf(glm): reuse Triton kernels across query and pool lengths - #1684
Long-context dimensions were still part of several Triton compilation keys after the pool gather optimization. This made changing query counts, chunk bounds, or pool capacities compile new binaries even when the kernel algorithm and tensor layout were unchanged.
This PR:
keeps runtime dimensions dynamic for the FP8 prefill row-table, KV gather, top-k compaction, online merge, and selected-prefill kernels;
keeps TOKENS and MAX_POOLS dynamic in the GLM pool-range kernel;
adds GPU regression coverage that checks output correctness and stable Triton kernel hashes across changing lengths.
The pool gather/page reuse implementation from #1667 is retained; this PR only changes compilation specialization and does not change pool contents or page ownership.
Validation
On NVIDIA H200 with PyTorch 2.11.0+cu130, Triton 3.6.0, vLLM 0.22.1, and sglang-kernel 0.4.5:
240 GLM GPU tests passed, including the new compile-cache tests, FP8 KV, K-pool, indexer, and MTP regressions.
Black 21.12b0, Flake8 6.1.0, Python syntax checks, and git diff --check passed.
Summary by CodeRabbit
Performance Improvements
Improved GPU kernel compilation reuse across varying context lengths, query counts, token counts, and chunk boundaries.
Tests
Added GPU tests that verify output correctness and consistent compilation behavior across different input sizes and modes.
Review was skipped as selected files did not have any reviewable changes.
Run configuration
Source mode: Auto
Repository: Unavailable
Branch: Unavailable
File: No configuration file from this repository applied
Commit: Unavailable
Layers: defaults
Review profile: CHILL
Plan: Advanced
Run ID: 207262ad-8be6-4c0c-a2a9-36bcce21aa7b
Commits
Reviewing files that changed from the base of the PR and between 1f9e5d9 and 4803608.
You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.
Use the checkbox below for a quick retry:
Trigger review
No actionable comments were generated in the recent review.
Walkthrough
Selected GLM5 Triton kernels now exclude specified runtime arguments from specialization. CUDA-gated tests check outputs and compiled-kernel hash reuse across varied FP8 attention, prefill, and pool-range inputs.
Changes: GLM5 Triton kernel compile-cache reuse
Layer / File(s)
Summary
FP8 kernel specialization and validation lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py, unit_tests/models/glm5_next/test_fp8_compile_cache.py
Selected dimensions and strides are excluded from specialization in FP8 prefill and attention kernels. CUDA tests check outputs and hash reuse across varied context sizes, query counts, token counts, and chunk ranges.
Pool-range specialization and validation lightllm/models/glm5_next/triton_kernel/kpool.py, unit_tests/models/glm5_next/test_kpool_compile_cache.py
The pool-range kernel excludes token count and pool limit from specialization. CUDA tests check returned ranges and hash reuse across varied inputs.
The change aims to reuse compiled kernels across runtime sizes. No merge-blocking issue is established; normal GPU checks remain appropriate.
Pre-merge checks: 4 · 1
Failed checks: 1 warning
Check name
Status
Explanation
Resolution
Docstring Coverage
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed checks: 4 passed
Check name
Status
Explanation
Description Check
Check skipped - CodeRabbit’s high-level summary is enabled.
Title check
The title "perf(glm): reuse Triton kernels across query and pool lengths" matches the main change. The PR stops specializing runtime dimensions in several GLM Triton kernels, so query counts, chunk bo…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Full details: Title check
Explanation
The title "perf(glm): reuse Triton kernels across query and pool lengths" matches the main change. The PR stops specializing runtime dimensions in several GLM Triton kernels, so query counts, chunk bounds, and pool capacities no longer trigger new compilations. The raw summary lists the affected kernels, including the do_not_specialize markings in the FP8 kernels and the kpool get_pool_ranges kernel. The added compile-cache tests check that kernel hashes are reused across these varying lengths. The title is specific and concise, and it follows the conventional commit style.
Full details: Docstring Coverage
Explanation
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 5 functions across 4 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches
Generate docstrings
Commit to this branch
Create a new PR
Generate unit tests (beta)
Commit to this branch
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Long-context dimensions were still part of several Triton compilation keys after the pool gather optimization. This made changing query counts, chunk bounds, or pool capacities compile new binaries even when the kernel algorithm and tensor layout were unchanged.
This PR:
TOKENSandMAX_POOLSdynamic in the GLM pool-range kernel;The pool gather/page reuse implementation from #1667 is retained; this PR only changes compilation specialization and does not change pool contents or page ownership.
Validation
On NVIDIA H200 with PyTorch 2.11.0+cu130, Triton 3.6.0, vLLM 0.22.1, and sglang-kernel 0.4.5:
git diff --checkpassed.Summary by CodeRabbit
Performance Improvements
Tests