Repository navigation
[GLM5.3] Serve GLM-5.3-Flash on SM80 (A100): k-pool indexer boot + fp8 KV + sparse MLA split-K + CP threshold - #87
Merged
Conversation
What: replace the SM90-only deep_gemm calls of IndexerKPool with Triton tensor-core kernels (kpool_indexer_sm80.py: ragged-prefill MQA logits, paged decode MQA logits, act_quant), read the SM count from device properties, return no paged-MQA schedule, and route every k-pool fp8 store through a software e4m3fn codec (fp8_e4m3_sw.py), because Triton rejects fp8e4nv below SM89. Why: GLM-5.3-Flash (index_kpool > 1) cannot boot on A100: deep_gemm.fp8_mqa_logits, fp8_paged_mqa_logits, get_num_sms and every fp8 tl.store of the k-pool write kernels fail on compute capability 8.0. Enable: automatic on compute capability 8.0 to 8.8 (is_sm80_family). SM89 and newer, and ROCm, keep the original deep_gemm / hardware-fp8 path; SW_FP8 is a constexpr False there, so the generated kernels are unchanged. Evidence (A100-SXM4-80GB): the indexer kernels match a float64 reference (codec values bit-exact, top-512 sets equal, CUDA graph replay equals eager) and the k-pool write kernels match it byte-wise over 6 seeds; prefill logits for 8192 queries over 35k keys take 4.47 ms (37% of bf16 tensor-core peak). Upstream: none (not submitted yet)
What: build the deep_gemm paged-MQA schedule metadata only when index_kpool <= 1 (needs_paged_mqa_schedule). With a k-pool the indexer reads only the pooled schedule, so the per-token schedule is dead work. Why: the per-token schedule calls get_paged_mqa_logits_metadata, which needs deep_gemm and aborts metadata initialisation on SM8x. The same idea as upstream sgl-project#44378. Enable: automatic. This is not limited to SM8x: models with index_kpool > 1 skip the dead schedule on every GPU; index_kpool <= 1 is unchanged. Evidence: boot gate on A100, without it init_forward_metadata fails before the first token; no kernel timing was measured. Upstream: sgl-project#44378 (same idea, written independently)
What: two env-gated levers for the TileLang sparse MLA kernels on SM8x.
SGLANG_SM80_SPARSE_MLA_V1_CFG ("block_I,num_stages,threads") overrides the tile
config of the V1 prefill kernel. SGLANG_SM80_SPARSE_MLA_DECODE_SPLIT=1 runs
decode, target-verify and draft-extend steps with few query rows as a split-K:
the existing partial kernel writes into a reusable per-stream workspace and the
existing combine kernel merges it, when query tokens x padded heads <= 768.
Prefill and plain extend keep the V1 kernel (the caller passes
allow_decode_split), because under the breakable prefill graph their attention
runs as an eager break on bucket-padded q.
Why: at the default tile the A100 sparse MLA kernel is the largest term of a
verify step, and a few query rows leave most of the 108 SMs idle.
The settings are parsed strictly (new EnvField.get_strict, which raises a
ValueError naming the variable where get() would warn and fall back) and checked
at boot on SM8x with TileLang.
Enable: SGLANG_SM80_SPARSE_MLA_DECODE_SPLIT=1 and
SGLANG_SM80_SPARSE_MLA_V1_CFG=32,1,128. Both are off / unset by default and only
apply on compute capability 8.0 to 8.8.
Evidence (A100): split vs single kernel by token-heads: 1.19x faster at 48,
1.02x at 56, 0.91x at 64, which sets the 768 gate; 32,1,128 won the tile sweep
and held in the end-to-end runs.
Upstream: the split follows sgl-project#43917
What: A100 has no fp8 MMA, so keep the MLA KV pool in fp8_e4m3 and convert it to bf16 while the TileLang V1 and decode-partial kernels write their shared-memory tile: the raw bytes are gathered first, then a packed 8-wide extern converter (tl_fp8x8_to_bf16x8) fills the bf16 tile. On compute capability 8.0 to 8.8 the TileLang DSA constraint check accepts --kv-cache-dtype fp8_e4m3 instead of raising "requires SM89 or newer"; SM89 and newer use the upstream fp8 kernels. The converter writes shared memory outside the pipeline's sync, so the V1 tile must use num_stages == 1; this is checked at boot, together with the format of SGLANG_SM80_SPARSE_MLA_V1_CFG, and again in the kernel factories. Why: the KV bytes per token and layer drop from 1024 to 512, so the pool grows 1.94x (3,760,640 tokens at TP8, mem-fraction 0.80). That is what lets 1M-token contexts and a high prefix-cache hit rate fit on 8 x 80 GB. Enable: --kv-cache-dtype fp8_e4m3 --dsa-prefill-backend tilelang --dsa-decode-backend tilelang. The V1 tile comes from SGLANG_SM80_SPARSE_MLA_V1_CFG and is 32,1,128 on this path when it is unset. There is no codec switch: the packed converter is the only path. The existing CPU validation test is updated: compute capability (7, 5) still fails (now with "SM80") and (8, 0) passes. Evidence (A100): kernel time vs bf16 KV is 1.21-1.27x for prefill and 1.04-1.15x for decode; the accuracy gate passed and long-context retrieval is 27/27 on the fp8 KV build. Upstream: none (not submitted yet)
…graph What: SGLANG_CP_MIN_EXTEND_TOKENS=N makes interleave prefill CP apply only to extend batches with at least N new tokens (InterleaveCPStrategy.can_apply and can_dsa_prefill_cp_interleave). Smaller extends run the non-CP body and may replay the breakable prefill CUDA graph. With N set, the three places that forced the prefill graph off for any CP server are skipped (the CP-TP group sharing hook, the "context parallel" rule in cuda_graph_hook, the DSA + prefill CP branch in model_hook), the captured buckets are limited to sizes below N, the runner gate admits non-CP plain extends, and the DSA indexer decides per batch (attn_cp_metadata is not None) instead of per server. N > 0 with a CP strategy other than interleave fails at boot. --enable-cp-decode-attn-tp (DeepSeek-V4, GLM-5) does not slice the attention weights inside such a graph: its break functions run attention with the unsliced head count, so the slicing would fail at the first replay; decode and verify full graphs are unchanged. Why: with CP on, every prefill is eager, so even a 4-token request pays the launch cost of all 45 layers (a ~225 ms floor on a 30k cached prefix). Enable: SGLANG_CP_MIN_EXTEND_TOKENS=4096 together with --enable-prefill-cp --attn-cp-size 8 --cp-strategy interleave and --cuda-graph-bs-prefill buckets below N. Unset or 0 keeps the upstream behavior. Evidence (A100, TP8, CP8): bs1 TTFT p50 on a 30k cached prefix with a 4-token tail is 224.7 ms with CP forcing eager prefill and 63.6 ms once the small batch replays the graph; 0.75 rps TPOT p50 stays at about 21 ms. Upstream: related to sgl-project#44293 (CP + breakable graph, which excludes speculative decoding); this takes the threshold route instead.
What: register Glm5NextForConditionalGeneration (and its NextN) in CP_DECODE_ATTN_TP_SUPPORTED_ARCHS and wrap the DSA-MLA self_attn call of Glm5NextDecoderLayer in maybe_use_decode_attn_tp. KDA layers keep their TP partition (nullcontext); the MTP layer shares the decoder layer class. Why: under CP-TP group sharing the DSA MLA attention is replicated at attn_tp=1, so decode and verify run 64 heads per rank (q_b_proj, o_proj, w_kc / w_vc and the sparse MLA kernel are 8x too large). The flag slices them to 8 heads per rank for decode, target-verify and draft-extend and all-reduces at o_proj. Without the whitelist entry model_hook raises ValueError. Enable: --enable-prefill-cp --attn-cp-size 8 --cp-strategy interleave --enable-cp-decode-attn-tp. The context logs "Enable CP decode attention TP". Evidence (A100, TP8, CP8): 0.75 rps open-loop TPOT p50 36.5 -> 21.0 ms against the same CP configuration without the flag. Upstream: none (not submitted yet)
What: build four tiny index / length tensors with torch.tensor(..., pin_memory=True).to(device, non_blocking=True) instead of torch.tensor(..., device=cuda): interleave_kpool.local_query_inputs, InterleaveCPStrategy.moe_num_token_non_padded, the extend_seq_lens of DeepseekSparseAttnBackend.init_forward_metadata and the seq_len_t of the IndexerKPool ragged top-k. Why: a pageable host-to-device copy synchronizes the stream. Profiling five CP extends showed 249 ms of cudaStreamSynchronize (about 50 ms per extend) on the overlap scheduler. The copied values are unchanged and these sites run eagerly, never inside CUDA graph capture. Enable: always on. This is not gated by GPU or model: every path that reaches these four sites uses the pinned copy. Evidence (A100, TP8, CP8): bs1 TTFT p50 at 2048 / 4096 new tokens 264.5 / 307.0 ms, 3.8% / 3.0% lower than without; cold 35k floor 2.22 -> 2.17 s. Upstream: none (not submitted yet)
…he fp8 encoder What: CPU tests for the strict env parsers (SGLANG_CP_MIN_EXTEND_TOKENS, SGLANG_SM80_SPARSE_MLA_V1_CFG, SGLANG_SM80_SPARSE_MLA_DECODE_SPLIT: unset and valid values pass; malformed, negative and wrong-arity values raise with the variable name) and for the fp8 KV boot check (num_stages must be 1). A CUDA test compares the software e4m3fn encoder with torch.float8_e4m3fn, including random values, subnormals, the round-to-nearest-even ties 1.0625, 1.1875 and 15.5, and the exponent carry at 15.5; fixed bytes check the saturation to +-448 and NaN. Why: these settings fail the boot instead of falling back to a default, so the error paths need tests that run without a GPU; the encoder replaces a hardware convert and has to match it bit for bit. Enable: the settings tests join base-a-test-cpu and need no tilelang or triton; the encoder test is registered for the base-b-kernel-unit 1-gpu-large runner. Upstream: none (not submitted yet)
Author
|
머지하지 않고 열어 둡니다(#83과 같은 방식). 결과물은 브랜치 |
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.
GLM-5.3-Flash를 A100(SM80)에서 서빙하기 위한 패치 8개 (v0.5.22 base)
v0.5.22(baa9befa38,lmsysorg/sglang:v0.5.22이미지 소스와 동일 커밋)0.5.22-glm5.3flash-sm80@13d0665ec1(8커밋, 28파일 +1249 −86)quay.io/vessl-ai/sglang:0.5.22-glm5.3flash-sm80-13d0665@sha256:0aba1f18cd215fe1fadbd61b5171fe773755c89036e7e24a42ed1c42c14faa7f (고정 태그0.5.22-glm5.3flash-sm80동일 digest)canada-quant/GLM-5.3-Flash-W4A16-MTP(W4A16 Marlin MoE + 내장 MTP).pr-review-toolkit리뷰 루프 5회(코드·커널 효율·silent failure·주석·테스트). 마지막 라운드 Critical 0 / Important 0. 남은 Minor는 §7.1. 패치 지도 — forward pass 어디에 붙는가
flowchart LR BOOT["C1 SM8x 부팅<br/>k-pool indexer → Triton TC 커널<br/>fp8 store → SW e4m3fn 코덱"]:::boot SCHED["C2 k-pool 모델<br/>per-token paged-MQA 스케줄 생략"]:::boot subgraph PREFILL["prefill"] P1["C5 CP 토큰 임계<br/>N 미만 extend = prefill graph replay"]:::pre P2["C7 pinned non_blocking H2D<br/>CP·DSA 메타데이터"]:::pre P3["C3 V1 타일 cfg 32,1,128"]:::pre end subgraph KV["KV 풀"] K1["C4 fp8_e4m3 KV<br/>smem 타일에서 fp8x8→bf16x8 변환"]:::kv end subgraph DECODE["decode · MTP verify"] D1["C3 sparse MLA split-K<br/>q토큰×헤드 ≤ 768"]:::dec D2["C6 Glm5Next CP decode attn TP<br/>64 → 8 heads/rank"]:::dec end BOOT --> SCHED --> PREFILL --> KV --> DECODE classDef boot fill:#e8eefc,stroke:#4f46e5,color:#111 classDef pre fill:#e9f7ef,stroke:#15803d,color:#111 classDef kv fill:#f3e8ff,stroke:#7e22ce,color:#111 classDef dec fill:#fff4e5,stroke:#c2410c,color:#111e20f17bk-pool indexer SM8x 부팅fp8_mqa_logits·fp8_paged_mqa_logits·act_quant을 Triton 텐서코어 커널(kpool_indexer_sm80.py)로, k-pool fp8 store를 SW e4m3fn 코덱으로33ce759죽은 스케줄 생략index_kpool > 1이면 per-token paged-MQA 스케줄 메타데이터를 만들지 않음init_forward_metadata가 첫 토큰 전에 실패6eb4bd9sparse MLA split-K + 타일SGLANG_SM80_SPARSE_MLA_DECODE_SPLIT=1,SGLANG_SM80_SPARSE_MLA_V1_CFG=32,1,128fc5e76aSM8x fp8 KVfp8_e4m3허용--kv-cache-dtype fp8_e4m3+ tilelang DSA 백엔드1.27×, decode 1.041.15×. 장문 검색 27/275a8678dCP 토큰 임계SGLANG_CP_MIN_EXTEND_TOKENS=4096+ 버킷 < N15f7893Glm5Next CP decode attn TPGlm5NextForConditionalGeneration화이트리스트 + DSA-MLA를maybe_use_decode_attn_tp로 감쌈--enable-cp-decode-attn-tp2fe02a0pinned H2Dnon_blocking복사로cudaStreamSynchronize제거. 콜드 35k 2.22 → 2.17 s13d0665유닛 테스트torch.float8_e4m3fn(CUDA)base-a-test-cpu/base-b-kernel-unit1-gpu-large새 env
SGLANG_CP_MIN_EXTEND_TOKENSValueError(변수명 포함). interleave 외 CP 전략과 함께 N>0이면 부팅 실패SGLANG_SM80_SPARSE_MLA_DECODE_SPLITValueErrorSGLANG_SM80_SPARSE_MLA_V1_CFGValueError. fp8 KV이면num_stages == 1강제파싱은 새
EnvField.get_strict()로 합니다.get()은 형식 오류 때 경고 후 기본값으로 떨어지는데, 성능 레버가 조용히 꺼진 채 측정되는 것을 막으려고 부팅을 실패시킵니다.2. 누적 효과 — 저-hit trace 0.75 rps 창 (서버 p50)
같은 trace·시드 쌍으로 레버를 하나씩 더한 순서입니다. config 레버(hrrn 스케줄링, chunked-prefill 16384)는 이 PR의 코드가 아니지만 순서를 보이려고 함께 적었습니다.
xychart-beta title "TPOT p50 (ms), 0.75 rps 저-hit — 레버 누적" x-axis ["hrrn 적용 후", "+V1 타일(C3)", "+CP·decode attn TP(C6)", "+CP 임계(C5)", "+fp8 KV(C4)·chunk16k", "최종 정본"] y-axis "TPOT p50 (ms)" 0 --> 70 bar [62, 36.5, 21.0, 18.2, 28.7, 20.8] line [20, 20, 20, 20, 20, 20]막대가 TPOT p50, 선이 SLA 20 ms입니다. fp8 KV 단계는 hit가 낮은 다른 창(TPOT 108 ms 기준선)에서 잰 값이라 앞 단계와 직접 이어지지 않습니다. 같은 창에서 108 → 28.7 ms였습니다.
3. 처리 한계 (knee) — 실 trace open-loop, 서버 p50
xychart-beta title "작업점 서버 TPOT p50 (ms)" x-axis ["고-hit 0.75 (3창 중앙)", "고-hit 0.8", "저-hit 0.6", "저-hit 0.75 (3창 중앙)"] y-axis "TPOT p50 (ms)" 0 --> 25 bar [15.0, 16.4, 12.7, 22.1] line [20, 20, 20, 20]knee는 고-hit (0.8, 0.9], 저-hit (0.6, 0.75] 구간에 있습니다. 작업점 창은 창당 1회 측정입니다.
4. 정확도
5. 쓰는 법
정본 런처 전문은 inference-playbook
tpm-measurements/nvidia-a100/glm-5.3-flash/canon/run.sh입니다.6. 리뷰 포인트
_runs_with_prefill_cp == dsa_enable_prefill_cp,can_apply의 추가 조건seq_len >= 0은 항상 참,in_prefill_graph는 False입니다. graph·model·cp_tp 훅 결정,tilelang_sparse_fwd, k-pool 커널(SW_FP8=Falseconstexpr)은 upstream 경로 그대로입니다. SM89 이상은 SM8x 분기에 들어가지 않습니다(DeepSeek-V3.2/GLM-5의 SM89 tail_dim assert 경로 보존).index_kpool > 1모델은 모든 GPU에서 죽은 per-token 스케줄 생략,index_kpool <= 1은 불변)와 C7(pinned H2D, 값 동일·eager 전용 사이트). C7 사이트가 나중에 CUDA graph capture 안으로 들어가면 조용히 캡처될 수 있어 각 사이트 주석에 "eager 전용"을 적었습니다.num_stages == 1을 부팅과 커널 팩토리 양쪽에서 강제합니다.glm5_next.py에 compressed-tensors 분기가 없어--disable-shared-experts-fusion이 필요합니다. 이 PR은 고치지 않았습니다.7. 남은 Minor (follow-up)
parallel_hook.py:150-154의 info 로그 "has no effect without --enable-prefill-cp"는attn_cp_size > 1이면서 prefill CP가 꺼진 니치 구성에서 부정확합니다(cuda_graph_hook.py:175,cp_decode_attn_tp.py:89가 env를 소비). 수정안은 env 쪽 조건에cfg.enable_prefill_cp를 거는 것입니다.EnvField.get_strict()는_set_to_none과 deprecated alias를 처리하지 않습니다. 현재 호출처 3곳은 해당 없습니다.8. 검증
py_compile통과, 각 커밋이 쓰는 심볼은 그 커밋 또는 이전 커밋에 정의test_dsa_tilelang_fp8_validation.py갱신: (7,5) 실패 메시지 "SM80", (8,0) 통과;test_dsa_metadata_replay픽스처 호환)lmsysorg/sglang:v0.5.22@sha256:e3a4f060…+ 이 브랜치python/sglang덮어쓰기, 5,571파일 md5 대조 PASS(이미지 전용 파일 = Rust 확장.so5개 +_version.py)