Conversation
…rtup Add precompile_kernels_after_loading to Glm5NextForConditionalGeneration. When the vision attention backend resolves to triton_attn, the hook issues one context_attention_fwd call with the tower's specialization inputs (head_dim = hidden_size // num_heads, the tower dtype, kv_group_num=1, is_causal=False) so the Triton _fwd_kernel is compiled and device-loaded while ModelRunner.load_model runs, before the KV pool and CUDA graphs claim device memory. Without it the kernel loads on the first image request, which can arrive hours into serving with almost no free device memory. Only the first pipeline rank runs the hook, matching where general_mm_embed_routine embeds images. The hook synchronizes the device inside its try so asynchronously reported kernel errors are logged at WARNING together with synchronous ones; failures never abort loading. The model runner already calls the hook through maybe_precompile_model_kernels_after_loading; language-only models, non-Triton vision backends and later PP ranks return early. Tests: a CPU test covers the argument contract and the skip conditions; a GPU test runs the hook against the real kernel, then a representative multi-sequence strided vision call, and asserts through Triton's kernel cache and kernel_load_start_hook that no second compile or device load happens.
…ore pools exist swiglu_clamped is torch.compile'd with dynamo's default shape policy: the first token count compiles a static kernel and the next distinct count recompiles a dynamic one, each running inductor's Triton autotune with benchmark buffers on the device. On a fresh inductor cache the first image whose token count differs from the warmup image's autotunes during serving; at mem_fraction_static 0.99 that failed with CUDA OOM on the first 3840x2160 request after a 64x64 warmup image. Run the block MLP and the patch merger at two token counts in the precompile hook, before the memory pools are allocated.
ormandj
changed the title
Open
2 tasks
ormandj
changed the title
2 tasks
Contributor
Author
|
Refiled as #38214 against |
This was referenced Sep 6, 2026
Merged
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
The first image request can run out of GPU memory while compiling or loading vision kernels, after the KV cache and CUDA graphs have consumed the available headroom. A small warmup image also does not necessarily compile the dynamic MLP activation used by larger images.
This PR precompiles GLM-5-Next vision attention and MLP activations after loading the model but before allocating the KV pool. The request-time forward implementation is unchanged.
Modifications
Glm5NextForConditionalGeneration.Accuracy Tests
CPU recording subclasses verify the MLP and merger calls, token counts, dtype/device, skip conditions and warning behavior. The attention test substitutes a recording kernel to check its arguments. Author-reported CPU result: 8 passed.
The GPU regression runs the real attention precompile, then a strided three-sequence vision-attention call. It compares outputs with SDPA and checks Triton's cache and device-load hook to require no new specialization or device load on the second call. Author-reported pass on one RTX PRO 6000 Blackwell, Triton 3.7.1; that run preceded the expanded MLP CPU cases.
Author-observed serving failure: a first 3840x2160 image on a fresh compilation cache triggered MLP autotuning and OOM at memory fraction 0.99. No exact-head full-model reproduction log is attached, and the attention GPU regression does not itself validate the MLP autotune path.
Speed Tests and Profiling
The hook adds startup compilation and small forward calls before pool allocation. No isolated startup-duration or serving-speed benchmark for this revision.
Checklist
Developed with AI assistance.
CI States
Latest PR Test (Base): ❌ Run #33595057897
Latest PR Test (Extra): ❌ Run #33595057749
Latest PR Test (AMD ROCm 7.2): ❌ Run #33595058218