Repository navigation
[DSV4.1] Keep Engram extend batch shape out of Triton constexprs (fix for sgl-project/sglang#42159) - #15
Merged
Oasis-Git merged 1 commit intoOct 12, 2026
Conversation
engram_hash_extend_and_commit passed the batch size (EXTEND_BS, SEARCH_STEPS) and the token count (NUM_TOKENS) as tl.constexpr, so every new prefill shape compiled and device-loaded a fresh specialization of _engram_hash_kernel and _engram_commit_extend_history_kernel while serving. With little free device memory those late cuModuleLoadData calls can fail with CUDA OOM and kill the scheduler. Pass them as runtime arguments instead (the row search loop takes a runtime trip count), keeping a single ROWS_FROM_STARTS constexpr. Hash and history outputs are unchanged and the eager timing is the same (~70 us on a 4090 for 8192 tokens at bs 1/64/256). Add a test that new batch shapes reuse the compiled kernels; it fails with 4 compiled variants before this change. Co-Authored-By: Claude Opus 5.5 <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
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.
Follow-up fix for sgl-project#42159, targeting its branch.
Problem
engram_hash_extend_and_commitpasses the batch size (EXTEND_BS,SEARCH_STEPS) and the token count (NUM_TOKENS) astl.constexpr. So every new prefill shape compiles and device-loads a new specialization of_engram_hash_kernel/_engram_commit_extend_history_kernelwhile the server is running.In production with
SGLANG_OPT_ENGRAM_EXTEND_HASH=1, the warningTriton kernel '_engram_commit_extend_history_kernel' device-loaded after serving started (free device mem: 0.37–0.62 GiB)was logged hundreds of times. The scheduler then died with CUDA OOM and the pod was killed by the liveness probe. Because the number of shapes is unbounded, warming up at startup can't cover them.Fix
extend_bs,search_stepsandnum_tokensbecome runtime arguments. The row binary search now uses a runtime-boundedrange.ROWS_FROM_STARTSconstexpr chooses between the row map and the search, so the number of specializations stays fixed.Verification (RTX 4090, triton 3.8.0, torch 2.14.1)
test/manual/quantization/test_engram_extend_hash.py: all existing bit-exact parity subtests pass (28).test_batch_shapes_reuse_compiled_kernels: on the current branch head it fails, with each kernel compiled 4 times for 4 batch shapes. With this change it passes, with 1 compiled kernel each.🤖 Generated with Claude Code