Repository navigation
[Fix] glm53_flash: Read the KDA layer cache from the linear-attention backend - #94
Merged
jhinpan merged 1 commit intoOct 30, 2026
Conversation
…backend sgl-project/sglang#46980 removes HybridReqToTokenPool.mamba2_layer_cache. On SGLang with that change, the GLM-5.3-Flash monokernel fails at CUDA-graph capture with an AttributeError in _bind_kda, and SGLANG_AMD_MONO_CHECK fails the same way. Read the cache from the backend's state_layer_cache(layer_id) instead, and update SGLANG_READS.
4 of 5 tasks
This was referenced Nov 1, 2026
jhinpan
pushed a commit
that referenced
this pull request
Nov 5, 2026
…#71) ## Motivation GLM-5.3-Flash is served with MTP, but the glm53_flash spec declares no features, so SGLang refuses `--monokernel-backend amd` once `--speculative-algorithm` is set: a server gets MTP or the monokernel, not both. This PR runs MTP verify steps on the monokernel when a step has at most 8 rows (4 at TP2): MTP3 at 1–2 requests, MTP5 at 1. The NextN draft, prefill and wider steps keep SGLang's path. AgentX P90 interactivity, same node: +27% at concurrency 1 (MTP5) and +35% at concurrency 4 (MTP3). SGLang: upstream `5d8e90a0f9` (has sgl-project/sglang#46036 and #46980), no other SGLang PRs. Rebased on #94. AgentX contexts pass 262,144 tokens, so the speed runs include #97 (now on `main`), which fixes a decode fault there. ## Modifications - **Which steps run the monokernel.** The spec declares `Feature.SPECULATIVE_DECODING` and builds every launch for `verify` = draft tokens rows a request. A step runs the monokernel only if each request has exactly `verify` tokens and the step fits the width; anything else runs SGLang's forward. - **Refused at startup.** A draft other than one chain (topk 1, draft tokens = steps + 1), ReplaySSM spec, `SGLANG_ENABLE_PP_SPEC`, and SGLang's fused KDA accept. With a FlashInfer KDA verify backend, SGLang skips committing the KDA state that a verify launch reads. FlashInfer KDA is CUDA-only, so MI355X never hits this. - **KDA.** One task per (request, head) runs the request's rows in order from the committed state and leaves it unchanged. After each draft token it writes the state and conv window to SGLang's verify scratch, which SGLang commits for the accepted token. - **DSA.** A request's rows write the kpool ring in order, and each row waits for the earlier rows of its chain before it attends their keys. - **Tests.** A verify launch equals its draft tokens' decode launches bitwise (kernel test, MTP2/3/5, TP2/4/8). `TestMonokernelMtp` serves GSM8K with MTP3, 2 requests at a time, so every verify step runs the kernels. Unit tests cover the startup refusals and the fallback. Every new SGLang name is in `SGLANG_READS`. `SGLANG_AMD_MONO_CHECK` still checks decode steps only. ## Accuracy Tests Setup: 8 × MI355X, `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261028` (SGLang `5d8e90a0f9`, no other PRs), `amd/GLM-5.3-Flash-Quark-MXFP4`, TP4, bf16 KV, tilelang DSA. Off and on differ only in `--monokernel-backend amd`, and ran at the same time on one node (GPUs 0–3 and 4–7). | GSM8K run | MTP | Accuracy, off → on | Accept length, off → on | | --- | --- | --- | --- | | `few_shot_gsm8k`, 1,319 questions, parallel 2 (every verify step fits 8 rows) | 3 | 0.905 → 0.923 | 3.31 → 3.32 | | `few_shot_gsm8k`, 1,319 questions, parallel 1 | 5 | 0.915 → 0.914 | 4.01 → 4.02 | | sgl-eval, 200 questions, 2 at a time (every verify step on the kernels), thinking (`TestMonokernelMtp`) | 3 | 0.985 → 0.970 | 3.10 → 3.24 | ## Speed Tests and Profiling **AgentX**, the setup above with InferenceX's MI355X MTP recipe (golden accept length), 3,600-s windows as InferenceX runs. Off and on ran at the same time on one node. | Concurrency | MTP | P90 interactivity, 1 / ITL p90 (tok/s/user) | P90 E2E-normalized interactivity (tok/s/user) | P90 TTFT (s) | Total tok/s/GPU | Requests completed | | --- | --- | --- | --- | --- | --- | --- | | 1 | 5 | 305.5 → 389.3 (**+27.4%**) | 193.1 → 222.7 (+15.3%) | 0.393 → 0.427 (+8.6%) | 5,346 → 6,327 (+18.3%) | 275 → 300 | | 4 | 3 | 257.1 → 345.8 (**+34.5%**) | 170.3 → 199.9 (+17.4%) | 0.395 → 0.391 (−1.0%) | 7,854 → 9,126 (+16.2%) | 612 → 697 | | 16 (control) | 3 | 153.9 → 153.2 (−0.5%) | 110.9 → 110.6 (−0.3%) | 0.723 → 0.700 (−3.2%) | 27,099 → 27,476 (+1.4%) | 2,533 → 2,566 | Off → on. E2E-normalized interactivity is 1 / the P90 of E2E latency per output token, so it includes TTFT. Tok/s/GPU counts prompt tokens, so it follows the requests completed. At concurrency 16 the slow steps pass 8 rows (SGLang's path on both sides); 57% of the on side's logged decode batches had 1–2 requests. By turn: the on arm completed every turn of the off arm (all but 2 at concurrency 16) and more of the same conversations' later turns (keyed by `conversation_id` and `turn_index`). P90 over the turns that both arms completed, and over the turns that only the on arm completed: | Concurrency | Turns, both / on only | P90 interactivity, off → on, both (tok/s/user) | On only | P90 E2E-normalized interactivity, off → on, both (tok/s/user) | On only | | --- | --- | --- | --- | --- | --- | | 1 | 275 / 25 | 305.5 → 381.1 (+24.7%) | 392.8 | 193.1 → 221.3 (+14.6%) | 262.3 | | 4 | 612 / 85 | 257.1 → 320.7 (+24.7%) | 369.7 | 170.3 → 204.2 (+19.9%) | 155.6 | | 16 (control) | 2,530 / 35 | 154.1 → 152.3 (−1.1%) | 220.0 | 111.2 → 110.8 (−0.4%) | 127.5 | The turns that only the on arm completed have a lower ITL, so they raise the on arm's P90 interactivity over all turns (+34.5% at concurrency 4) above that over the turns both arms completed (+24.7%). At concurrency 4 they have longer contexts and a higher TTFT (P90 0.58 s against 0.39 s), so their E2E-normalized interactivity is lower. Server, concurrency 4 (MTP3). Off drops `--monokernel-backend amd`: ```bash export SGLANG_SIMULATE_ACC_LEN=2.99 SGLANG_SIMULATE_ACC_METHOD=match-expected SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token \ SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION=0.04 SGLANG_TIMEOUT_KEEP_ALIVE=900 python3 -m sglang.launch_server --model-path amd/GLM-5.3-Flash-Quark-MXFP4 --served-model-name amd/GLM-5.3-Flash-Quark-MXFP4 \ --host 0.0.0.0 --port 29904 --trust-remote-code --tp 4 --ep-size 1 --kv-cache-dtype bfloat16 \ --dsa-prefill-backend tilelang --dsa-decode-backend tilelang --moe-runner-backend aiter \ --mamba-radix-cache-strategy extra_buffer --tool-call-parser glm47 --reasoning-parser glm45 \ --chunked-prefill-size 16384 --mem-fraction-static 0.85 --max-running-requests 8 --cuda-graph-max-bs-decode 8 \ --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 32}' --mamba-full-memory-ratio 0.9 \ --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --watchdog-timeout 1800 --enable-metrics --monokernel-backend amd ``` Client: aiperf 0.12.0 (SemiAnalysisAI/aiperf `754356e9`, in InferenceX `bee3405ca`): ```bash export AIPERF_HTTP_TCP_USER_TIMEOUT=900000 aiperf profile --scenario inferencex-agentx-mvp --url http://localhost:29904 --endpoint /v1/chat/completions \ --endpoint-type chat --streaming --model amd/GLM-5.3-Flash-Quark-MXFP4 --tokenizer amd/GLM-5.3-Flash-Quark-MXFP4 \ --tokenizer-trust-remote-code --concurrency 4 --benchmark-duration 3600 --warmup-grace-period 1800 \ --warmup-requests-per-lane 10 --stats-interval 30 --random-seed 42 --failed-request-threshold 0.10 \ --trajectory-start-min-ratio 0.25 --trajectory-start-max-ratio 0.75 --trace-idle-gap-cap-seconds 300 \ --use-server-token-count --no-gpu-telemetry --num-dataset-entries 393 --slice-duration 1.0 \ --public-dataset semianalysis_cc_traces_weka_062126 --server-metrics http://localhost:29904/metrics \ --output-artifact-dir aiperf_artifacts ``` | Point | `SGLANG_SIMULATE_ACC_LEN` | `--speculative-num-steps` / `--speculative-num-draft-tokens` | `--max-running-requests` / `--cuda-graph-max-bs-decode` | `--concurrency` | | --- | --- | --- | --- | --- | | Concurrency 1, MTP5 | 3.56 | 5 / 6 | 8 / 8 | 1 | | Concurrency 4, MTP3 | 2.99 | 3 / 4 | 8 / 8 | 4 | | Concurrency 16, MTP3 | 2.99 | 3 / 4 | 32 / 32 | 16 | ## Not covered - Verify steps of more than 8 rows (MTP3 at 3+ requests, MTP5 at 2+) run SGLang's path. - `SGLANG_AMD_MONO_CHECK` does not check verify steps. - TP2 and TP8 were tested at kernel level only. ## Checklist - [x] `pre-commit run --all-files` passes. (isort, ruff and ruff-format at the config's versions, on the changed files.) - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.)
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.
Motivation
sgl-project/sglang#46980 removes
HybridReqToTokenPool.mamba2_layer_cache. glm53_flash calls it in_bind_kdaon every routed step, andSGLANG_AMD_MONO_CHECKcalls it in_regions. On SGLang with that change, for examplelmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261028(SGLang5d8e90a0f9), a GLM-5.3-Flash server with the monokernel exits at CUDA-graph capture:test_contract.pydoes not catch it: the name is still inmemory_pool.py, onMambaPool.Modifications
spec.py:_bind_kdatakes the linear-attention backend and reads itsstate_layer_cache(layer_id), as SGLang's KDA backend does since #46980.check.py:_regionsreads the same cache throughget_attn_backend().linear_attn_backend.SGLANG_READS:state_layer_cacheinhybrid_linear_attn_backend.py;mamba2_layer_cacheandget_req_to_token_poolgo.This needs SGLang with sgl-project/sglang#46980.
Accuracy Tests
MI355X,
lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261028(SGLang5d8e90a0f9),amd/GLM-5.3-Flash-Quark-MXFP4, TP4,test/e2e/test_glm53_flash_server.py:mainTestMonokernelDecode(2 tests)TestMonokernelCheckTestMonokernelHiddenStatesTestMonokernelFaultmainis eb0f20d; glm53_flash has not changed since.TestMonokernelFaultalso fails onmainon SGLang5044779f, which is before #46980, so this PR does not cause it.Speed Tests and Profiling
None: the launches and their arguments do not change.
Checklist
pre-commit run --all-filespasses.test/cover the change. (test_glm53_flash_server.pyfails before this PR and passes after, on SGLang with #46980.).github/CODEOWNERSthat names its owner. (No new model.)