Bump sglang to v0.5.16 and add GLM-5.2 support - #13
Open
TerrenceZhangX wants to merge 3 commits into
Open
Conversation
The pinned sglang v0.5.5 predates GLM-5.2, so the model could not be profiled at all. v0.5.16 registers GlmMoeDsaForCausalLM natively. sglang submodule: 0c006b8 (v0.5.5) -> fdebc93 (v0.5.16). - workload/framework/patches/hook-v0516.patch: port of hook.patch to v0.5.16, whose triton kernels moved into the new sglang/kernels/ops/ package. apply_auto_profile() patches a module attribute, and call sites use `from X import y` (early binding), so hooks must be registered on the *consumer* module rather than the definition site. That makes most of the original source edits unnecessary: only raw triton.JITFunction objects still need a launch_* wrapper extracted, since wrapping one directly breaks `kernel[grid](...)` syntax. Also fixes a KeyError in the launcher hook (env["..."] on an unset var). - dockerfiles/sglang-cu130.dockerfile: layer FlowSim on the official lmsysorg/sglang release image instead of building sglang from source. cuda12.6.dockerfile targets sm_90 and below, so it cannot run on Blackwell (GB200/B200, sm_100+), and building sglang + DeepGEMM from source on aarch64 is slow and fragile. The base image is a git checkout of the pinned commit with sglang installed editable, so the patch applies to it directly. Needs no network access, which GPU hosts frequently lack. - workload/models/configs/GLM-5.2: truncated config following the existing deepseek convention (num_hidden_layers 2 + first_k_dense_replace 1 => one dense layer and one MoE layer, full 256 experts). config_full.json keeps the upstream config so the truncation stays auditable. - utils/shape_merge: don't fail when a kernel's launch count differs between the CUDA-graph and eager passes. Disabling CUDA graphs can legitimately change the launch sequence -- GLM-5.2's DSA decode backend replays one fused _fused_dsa_decode_metadata_kernel under CUDA graph but issues arange/fill/cumsum/index/add separately in eager mode. Positional matching would attach the wrong shapes to those rows, so they now keep their timing-pass values and are reported; strict=True restores the old behaviour. Previously this aborted the whole merge over a handful of metadata kernels. - tests/integration/test_model_config.py: --batch is ambiguous in v0.5.16 (--batch-notify-size vs --batch-size); cap --context-length so GLM-5.2's 1M max_position_embeddings doesn't OOM the KV pool. Verified on GB200 (aarch64, CUDA 13.0, TP=1): --collect all completes for GLM-5.2 and produces merged EXTEND and DECODE kernel CSVs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
GLM-5.2 selects the DeepSeek sparse-attention backend, so the existing FlashInfer-MLA registrations were no-ops for it, and its attention, MoE and norm work is launched from FlashInfer / TensorRT-LLM / DeepGEMM custom ops that never reach aten dispatch. The profiler's own record_shapes cannot see those, so the kernels that matter most for simulation arrived in the CSVs with Dims=N/A: MLA attention (fmhaSm100, bmm), MoE routing/finalize, the DSA indexer (paged_mqa_logits, topk) and RMSNorm. Register the launchers that actually sit on GLM-5.2's path, chosen by walking the call stacks of the shape-less kernels in a real trace: DeepseekSparseAttnBackend (forward_extend/forward_decode plus the one-shot MHA, paged-MQA-schedule, topk-plan and cuda-graph-metadata helpers), Indexer.forward_cuda and its fused K prepare/store, DeepseekV2MoE.forward and the trtllm-gen bf16 runner, RMSNorm.forward_cuda, and the paged allocator decode path. Measured on GB200 (bs=1, input=2048, TP=1), share of kernels lacking Dims: prefill 19% -> 2% (by kernel time 36% -> 0%) decode 34% -> 6% (by kernel time 21% -> 2%) What remains is metadata kernels worth ~0.3% (prefill) and ~2% (decode) of kernel time, which FlowSim resolves by name via kernels.json. Total kernel time is unchanged (3.60 -> 3.61 ms extend, 2.95 -> 2.96 ms decode), so the added record_function ranges cost under 0.5%. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Profiling a MoE model with --load-format dummy gives an untrained router. On GLM-5.2 that router collapses onto a small subset of the 256 experts, so the grouped-GEMM launches stream far less expert weight than a real deployment would and the kernel timings understate MoE cost. Measured on GB200 at bs=1, input_len=2048: prefill GateUp took 612 us, but DRAM peak over 612 us moves only 4.59 GB against the 12.88 GB that 256 BF16 experts require, so at most 84 experts can have been active. Patch 0004 already added SGLANG_BALANCED_MOE, but it does not reach GLM-5.2: it rewrites topk_ids inside fused_moe_triton while GLM uses the TRT-LLM grouped GEMM, it hooks DeepseekV2MoE.forward while GLM's decoder layer builds Glm4MoeSparseMoeBlock, and with the flashinfer/trtllm backend TopK.forward_cuda returns a BypassedTopKOutput so topk_ids never materialises in Python at all -- routing happens inside moe::dev::routing from router_logits. Rewriting router_logits is therefore the only injection point that works. This hooks TopK itself rather than any one model's forward, because Glm4MoeSparseMoeBlock has three routing paths (forward_normal, forward_normal_dual_stream under CUDA-graph capture, and forward_deepep) and every one of them funnels into TopK; patching the model class covers only whichever path that deployment happens to take. GLM-5.2 has n_group=1 and topk_group=1, so grouped top-k degenerates to plain top-k over all experts and a large logit margin deterministically selects the forced set. The noaux_tc correction bias is zeroed too, since sigmoid scoring saturates at 1.0 and a dummy-initialised bias would otherwise outrank the forced choice. With this applied, the same prefill GateUp launch takes 1765 us at 99.9% of DRAM peak, and Down 911 us at 98.4% -- both the DRAM-bound answer for 2048 tokens spread evenly over 256 experts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.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
1 participant
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.
The pinned sglang v0.5.5 predates GLM-5.2, so the model could not be profiled at all. v0.5.16 registers GlmMoeDsaForCausalLM natively.
sglang submodule: 0c006b8 (v0.5.5) -> fdebc93 (v0.5.16).
workload/framework/patches/hook-v0516.patch: port of hook.patch to v0.5.16, whose triton kernels moved into the new sglang/kernels/ops/ package. apply_auto_profile() patches a module attribute, and call sites use
from X import y(early binding), so hooks must be registered on the consumer module rather than the definition site. That makes most of the original source edits unnecessary: only raw triton.JITFunction objects still need a launch_* wrapper extracted, since wrapping one directly breakskernel[grid](...)syntax. Also fixes a KeyError in the launcher hook (env["..."] on an unset var).dockerfiles/sglang-cu130.dockerfile: layer FlowSim on the official lmsysorg/sglang release image instead of building sglang from source. cuda12.6.dockerfile targets sm_90 and below, so it cannot run on Blackwell (GB200/B200, sm_100+), and building sglang + DeepGEMM from source on aarch64 is slow and fragile. The base image is a git checkout of the pinned commit with sglang installed editable, so the patch applies to it directly. Needs no network access, which GPU hosts frequently lack.
workload/models/configs/GLM-5.2: truncated config following the existing deepseek convention (num_hidden_layers 2 + first_k_dense_replace 1 => one dense layer and one MoE layer, full 256 experts). config_full.json keeps the upstream config so the truncation stays auditable.
utils/shape_merge: don't fail when a kernel's launch count differs between the CUDA-graph and eager passes. Disabling CUDA graphs can legitimately change the launch sequence -- GLM-5.2's DSA decode backend replays one fused _fused_dsa_decode_metadata_kernel under CUDA graph but issues arange/fill/cumsum/index/add separately in eager mode. Positional matching would attach the wrong shapes to those rows, so they now keep their timing-pass values and are reported; strict=True restores the old behaviour. Previously this aborted the whole merge over a handful of metadata kernels.
tests/integration/test_model_config.py: --batch is ambiguous in v0.5.16 (--batch-notify-size vs --batch-size); cap --context-length so GLM-5.2's 1M max_position_embeddings doesn't OOM the KV pool.
Verified on GB200 (aarch64, CUDA 13.0, TP=1): --collect all completes for GLM-5.2 and produces merged EXTEND and DECODE kernel CSVs.