[FEAT] Support GGUF format#2215
Merged
merrymercy merged 14 commits intosgl-project:mainfrom Nov 30, 2024
Merged
Conversation
Contributor
Author
|
lm_head.weight is directly used in many places, however, vllm changes it to be |
Contributor
|
Thanks for the contributions. Can you fix the CI errors? |
5c616a5 to
2cffa70
Compare
Contributor
Author
How to trigger the CI? |
Contributor
Author
Pass lm_head to LogitsProcessor and check the weight inside |
Contributor
|
@zhengy001 CI won't be triggered for you automatically because you are a first-time contributor. You can send a random typo fix PR and I can merge that for you so your future commits can trigger CI automatically. |
Contributor
|
@zhengy001 Can you fix the CI errors? |
Contributor
Author
@merrymercy Sure, working on it. |
merrymercy
requested changes
Nov 30, 2024
Contributor
|
#2269 adds you as a new contributor so your future commits will trigger CI automatically |
Contributor
Author
@merrymercy :) |
zhengy001
commented
Nov 30, 2024
Contributor
Author
There was a problem hiding this comment.
There won't be "lm_head.weight" if self.config.tie_word_embeddings is True
zhengy001
commented
Nov 30, 2024
Contributor
Author
There was a problem hiding this comment.
Compared the result with vllm's. Pls suggest if there is a better way.
merrymercy
approved these changes
Nov 30, 2024
merrymercy
reviewed
Nov 30, 2024
merrymercy
added a commit
that referenced
this pull request
Dec 1, 2024
This reverts commit 883c955.
timethink
pushed a commit
to timethink/sglang
that referenced
this pull request
Mar 9, 2025
Co-authored-by: Yang Zheng(SW)(Alex) <you@example.com>
11 tasks
hammersam
added a commit
to hammersam/sglang
that referenced
this pull request
Mar 8, 2026
Port the multi-CTA radix-based top-k kernel from flashinfer PR sgl-project#2215 (flashinfer-ai/flashinfer#2215) into sglang as a JIT-compiled kernel. This replaces the existing AOT single-CTA top-k implementation for NSA attention, providing better performance on long sequences (32K+) where the multi-CTA path activates. Key changes: - Add `python/sglang/jit_kernel/topk.py`: Python API exposing three JIT top-k variants (basic, page-table transform, ragged transform) with workspace management and lazy compilation via `cache_once`. - Add `python/sglang/jit_kernel/csrc/elementwise/topk.cuh`: CUDA wrapper providing TVM FFI entry points that dispatch to the flashinfer adaptive top-k kernels (TopKDispatch, TopKPageTableTransformDispatch, TopKRaggedTransformDispatch). - Add `python/sglang/jit_kernel/include/sgl_kernel/topk_fi.cuh`: Core CUDA implementation adapted from flashinfer, featuring: - 8-bit radix selection algorithm with multi-CTA support for large sequences (threshold configurable, default 32K) - Support for float32, float16, and bfloat16 input types - row_starts parameter for ragged input score layouts (sglang-specific) - Three output modes: indices-only, page-table lookup, and ragged offset addition - Update `python/sglang/srt/layers/attention/nsa_backend.py`: Switch NSA indexer to import from JIT kernel instead of AOT sgl_kernel. - Update `sgl-kernel/python/sgl_kernel/top_k.py`: Add JIT fallback path controlled by SGLANG_USE_JIT_TOPK env var (default enabled). When JIT is available, fast_topk_v2 / fast_topk_transform_fused / fast_topk_transform_ragged_fused transparently delegate to JIT kernels. - Add `sgl-kernel/tests/test_topk_jit.py`: Correctness tests covering basic, page-table, ragged, and trivial (length <= topk) cases across various batch sizes and sequence lengths up to 131K. - Add `sgl-kernel/benchmarks/bench_topk_jit.py`: Latency benchmark comparing JIT multi-CTA vs AOT single-CTA kernels. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Motivation
#1616
Modifications
Support GGUF format
Checklist