From 97c6978369ac1e04c91fcc01c98acc25129a6000 Mon Sep 17 00:00:00 2001 From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> Date: Sun, 6 Sep 2026 17:27:59 +0800 Subject: [PATCH] GLM-5.3-Flash support (#36507) Co-authored-by: zRzRzRzRzRzRzR Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com> Co-authored-by: zanes-ops Co-authored-by: Baizhou Zhang Co-authored-by: Jian Chen Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com> Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com> Co-authored-by: Ehsan Akhgari Co-authored-by: kpham-sgl Co-authored-by: BBuf <1182563586@qq.com> Co-authored-by: Raiden Makoto <81530826+Raiden-Makoto@users.noreply.github.com> --- .../ops/attention/dsa/quant_k_cache.py | 74 +- .../ops/attention/dsa/tilelang_kernel.py | 54 +- .../kernels/ops/attention/dsa_metadata.py | 16 +- .../sglang/kernels/ops/attention/fla/kda.py | 9 +- .../sglang/srt/arg_groups/attention_hook.py | 1 + .../sglang/srt/arg_groups/cuda_graph_hook.py | 9 +- python/sglang/srt/arg_groups/model_hook.py | 1 + .../arg_groups/model_overrides/deepseek_v2.py | 3 +- python/sglang/srt/arg_groups/overrides.py | 3 + .../srt/arg_groups/pd_disaggregation_hook.py | 3 +- python/sglang/srt/configs/__init__.py | 3 + python/sglang/srt/configs/glm5_next.py | 341 ++++ python/sglang/srt/configs/hybrid_arch.py | 11 + python/sglang/srt/configs/model_config.py | 74 +- python/sglang/srt/configs/qwen3_asr.py | 4 +- python/sglang/srt/disaggregation/base/conn.py | 3 + .../sglang/srt/disaggregation/common/conn.py | 2 +- python/sglang/srt/disaggregation/decode.py | 9 + .../disaggregation/encoder/preprocessor.py | 215 ++- .../srt/disaggregation/mooncake/conn.py | 57 +- python/sglang/srt/disaggregation/nixl/conn.py | 100 +- python/sglang/srt/disaggregation/prefill.py | 10 + python/sglang/srt/disaggregation/utils.py | 234 ++- .../hardware_backend/npu/memory_pool_npu.py | 12 +- .../layers/attention/attention_registry.py | 3 + .../layers/attention/dsa/dsa_backend_kpool.py | 304 +++ .../dsa/dsa_backend_mtp_precompute.py | 19 +- .../srt/layers/attention/dsa/dsa_indexer.py | 63 +- .../layers/attention/dsa/dsa_indexer_kpool.py | 1520 +++++++++++++++ .../srt/layers/attention/dsa/kpool_plan.py | 24 +- .../sglang/srt/layers/attention/dsa/utils.py | 14 +- .../srt/layers/attention/dsa_backend.py | 344 +++- .../attention/hybrid_linear_attn_backend.py | 35 +- .../layers/attention/linear/kda_backend.py | 104 +- .../srt/layers/attention/triton_backend.py | 4 + python/sglang/srt/layers/communicator.py | 61 +- python/sglang/srt/layers/communicator_mhc.py | 557 ++++++ python/sglang/srt/layers/dp_attention.py | 19 + .../srt/layers/moe/moe_runner/deep_gemm.py | 13 +- python/sglang/srt/managers/schedule_batch.py | 5 +- python/sglang/srt/managers/scheduler.py | 34 + python/sglang/srt/mem_cache/allocator/base.py | 6 +- .../sglang/srt/mem_cache/allocator/paged.py | 17 +- python/sglang/srt/mem_cache/allocator/swa.py | 17 +- .../sglang/srt/mem_cache/allocator/token.py | 17 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 6 +- .../srt/mem_cache/dsa_cache_layer_split.py | 15 +- .../srt/mem_cache/hisparse_memory_pool.py | 16 +- .../hybrid_cache/hybrid_pool_assembler.py | 5 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 2 + .../srt/mem_cache/kv_cache_configurator.py | 105 +- python/sglang/srt/mem_cache/memory_pool.py | 530 +++++- .../sglang/srt/mem_cache/swa_memory_pool.py | 14 +- .../srt/mem_cache/unified_memory_pool.py | 14 +- .../cuda_graph_buffer_registry.py | 6 +- .../srt/model_executor/forward_batch_info.py | 21 + .../sglang/srt/model_executor/model_runner.py | 2 +- .../cuda_graph_setup.py | 121 +- .../spec_aux_hidden_state.py | 68 +- .../srt/model_executor/pool_configurator.py | 51 +- .../runner/prefill_cuda_graph_runner.py | 19 +- .../breakable_cuda_graph.py | 12 +- python/sglang/srt/model_loader/loader.py | 6 +- .../sglang/srt/model_loader/weight_utils.py | 11 +- .../attention_backend_handler.py | 6 +- .../attention_forward_methods/forward_mha.py | 15 +- .../attention_forward_methods/forward_mla.py | 2 + .../deepseek_common/deepseek_weight_loader.py | 5 +- python/sglang/srt/models/deepseek_nextn.py | 32 +- python/sglang/srt/models/deepseek_v2.py | 19 +- python/sglang/srt/models/glm5_next.py | 1630 +++++++++++++++++ python/sglang/srt/models/glm5_next_nextn.py | 80 + python/sglang/srt/models/glm_ocr.py | 12 +- python/sglang/srt/multimodal/mm_utils.py | 64 + .../multimodal/processors/base_processor.py | 3 + .../sglang/srt/multimodal/processors/glm4v.py | 218 ++- .../srt/speculative/adaptive_spec_params.py | 2 +- .../eagle_draft_cuda_graph_runner.py | 2 +- .../eagle_draft_extend_cuda_graph_runner.py | 22 +- python/sglang/srt/speculative/eagle_utils.py | 32 +- .../srt/speculative/eagle_worker_common.py | 4 + .../sglang/srt/speculative/eagle_worker_v2.py | 24 +- python/sglang/srt/speculative/spec_info.py | 3 + .../srt/speculative/standalone_worker_v2.py | 1 + python/sglang/srt/utils/common.py | 5 +- .../srt/utils/hf_transformers/common.py | 4 + python/sglang/srt/utils/video_decoder.py | 42 +- .../kernels/ops/attention/test_dsa_indexer.py | 2 + .../attention/test_fused_verify_triton_gdn.py | 37 +- .../registered/kernels/test_lean_attention.py | 24 +- .../models_e2e/test_glm53_flash_b200.py | 138 ++ .../models_e2e/test_glm53_flash_h200.py | 119 ++ .../test_prefill_abort_result_cleanup.py | 1 + .../mlx/test_attention_patching.py | 1 + .../test_mla_decode_forced_splits.py | 14 +- .../test_generation_auxiliary_output.py | 2 + .../mem_cache/test_retraction_mamba_backup.py | 6 +- .../model_executor/test_pool_configurator.py | 2 + .../models/test_deepseek_nextn_mm_embed.py | 210 +++ .../models/test_glm5_next_dflash_capture.py | 40 + .../unit/parser/test_jinja_template_utils.py | 1 - .../unit/spec/test_adaptive_spec_params.py | 23 +- .../test_eagle_worker_v2_topk1_fastpath.py | 1 + 103 files changed, 7741 insertions(+), 559 deletions(-) create mode 100644 python/sglang/srt/configs/glm5_next.py create mode 100644 python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py create mode 100644 python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py create mode 100644 python/sglang/srt/layers/communicator_mhc.py create mode 100644 python/sglang/srt/models/glm5_next.py create mode 100644 python/sglang/srt/models/glm5_next_nextn.py create mode 100644 test/registered/models_e2e/test_glm53_flash_b200.py create mode 100644 test/registered/models_e2e/test_glm53_flash_h200.py create mode 100644 test/registered/unit/models/test_deepseek_nextn_mm_embed.py create mode 100644 test/registered/unit/models/test_glm5_next_dflash_capture.py diff --git a/python/sglang/kernels/ops/attention/dsa/quant_k_cache.py b/python/sglang/kernels/ops/attention/dsa/quant_k_cache.py index adaa948e1..d7b1e5a50 100644 --- a/python/sglang/kernels/ops/attention/dsa/quant_k_cache.py +++ b/python/sglang/kernels/ops/attention/dsa/quant_k_cache.py @@ -3,6 +3,57 @@ import triton import triton.language as tl +def gather_dsa_kv_scales( + scale_src, + scale_dst, + kv_indices, + kv_indptr, + kv_indptr_idx, +): + _gather_dsa_kv_scales[(32,)]( + scale_src, + scale_dst, + kv_indices, + kv_indptr, + scale_src.stride(0), + KV_INDPTR_IDX=kv_indptr_idx, + NUM_TILES=scale_src.shape[-1], + BLOCK=256, + ) + + +@triton.jit +def _gather_dsa_kv_scales( + scale_src, + scale_dst, + kv_indices, + kv_indptr, + scale_src_stride, + KV_INDPTR_IDX: tl.constexpr, + NUM_TILES: tl.constexpr, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + num_programs = tl.num_programs(0) + active = tl.load(kv_indptr + KV_INDPTR_IDX) + block_start = pid * BLOCK + tiles = tl.arange(0, NUM_TILES) + while block_start < active: + offsets = block_start + tl.arange(0, BLOCK) + mask = offsets < active + rows = tl.load(kv_indices + offsets, mask=mask, other=0) + values = tl.load( + scale_src + rows[:, None] * scale_src_stride + tiles[None, :], + mask=mask[:, None], + ) + tl.store( + scale_dst + rows[:, None] * NUM_TILES + tiles[None, :], + values, + mask=mask[:, None], + ) + block_start += num_programs * BLOCK + + def quantize_k_cache(cache_k): return _quantize_k_cache_fast_wrapped(cache_k) @@ -22,19 +73,27 @@ def quantize_k_cache_separate( k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope) Must have dim_nope=512 for FP8 MLA quantization k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope) - Must have dim_rope=64 for FP8 MLA quantization + Must have dim_rope=64 for FP8 MLA quantization, or dim_rope=0 + for no-PE MLA (empty rope); None is treated + the same as an empty rope. tile_size: quantization tile size (default 128) Returns: Tuple of (nope_part, rope_part) where: - nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)] - rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)] + (empty, (num_tokens, 1, 0), when dim_rope=0) These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part) """ # Squeeze middle dimension if present k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope - k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope + if k_rope is None or k_rope.numel() == 0: + k_rope_2d = torch.empty( + (k_nope_2d.shape[0], 0), dtype=k_nope_2d.dtype, device=k_nope_2d.device + ) + else: + k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope num_tokens = k_nope_2d.shape[0] dim_nope = k_nope_2d.shape[1] @@ -43,8 +102,8 @@ def quantize_k_cache_separate( # Validate dimensions for FP8 MLA if dim_nope != 512: raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}") - if dim_rope != 64: - raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}") + if dim_rope not in (0, 64): + raise ValueError(f"Expected dim_rope=64 (or 0 for no-PE MLA), got {dim_rope}") if k_rope_2d.shape[0] != num_tokens: raise ValueError( f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}" @@ -234,7 +293,12 @@ def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128): # Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)] nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn) nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32) - rope_view = rope_part_u8.view(torch.bfloat16) + if dim_rope > 0: + rope_view = rope_part_u8.view(torch.bfloat16) + else: + rope_view = torch.empty( + (num_tokens, 0), dtype=torch.bfloat16, device=k_rope.device + ) # Kernel launch parameters num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size) diff --git a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py index 23d5e3025..58dcd749a 100644 --- a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py +++ b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py @@ -272,12 +272,13 @@ def sparse_attention_fwd_kernel_v1( num_stages=2, threads=256, ): - assert dim == tilelang.math.next_power_of_2(dim), ( - f"haven't check padding correctness yet, dim={dim}" + assert dim == tilelang.math.next_power_of_2(dim) or dim % 64 == 0, ( + f"dim={dim} must be a power of 2 or a multiple of 64" ) - assert tail_dim == tilelang.math.next_power_of_2(tail_dim), ( - f"haven't check padding correctness yet, dim={tail_dim}" + assert tail_dim == 0 or tail_dim == tilelang.math.next_power_of_2(tail_dim), ( + f"tail_dim={tail_dim} must be 0 or a power of 2" ) + has_tail = tail_dim > 0 assert is_causal == True, "non-casual is not supported" assert topk % block_I == 0, ( "otherwise will load some index=0 thus causing wrong kv to be loaded" @@ -330,9 +331,11 @@ def sparse_attention_fwd_kernel_v1( bz, ): Q_shared = T.alloc_shared([H_per_block, D], dtype) - Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + if has_tail: + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) KV_shared = T.alloc_shared([BI, D], dtype) - K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + if has_tail: + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) O_shared = T.alloc_shared([H_per_block, D], dtype) mask = T.alloc_fragment([BI], "bool") @@ -358,7 +361,8 @@ def sparse_attention_fwd_kernel_v1( H1 = H0 + H_per_block T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + if has_tail: + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) for i_i in T.Pipelined(NI, num_stages=num_stages): for bi_i in T.Parallel(BI): @@ -368,10 +372,14 @@ def sparse_attention_fwd_kernel_v1( KV_shared[bi_i, d_i] = KV[ b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i ] - for bi_i, d_i in T.Parallel(BI, D_tail): - K_tail_shared[bi_i, d_i] = KV[ - b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i - ] + if has_tail: + for bi_i, d_i in T.Parallel(BI, D_tail): + K_tail_shared[bi_i, d_i] = KV[ + b_i, + Indices[b_i, s_i, g_i, i_i * BI + bi_i], + g_i, + D + d_i, + ] for h_i, bi_i in T.Parallel(H_per_block, BI): acc_s[h_i, bi_i] = T.if_then_else( @@ -384,13 +392,14 @@ def sparse_attention_fwd_kernel_v1( transpose_B=True, policy=T.GemmWarpPolicy.FullCol, ) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) + if has_tail: + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) T.copy(m_i, m_i_prev) T.reduce_max(acc_s, m_i, dim=1, clear=False) for h_i in T.Parallel(H_per_block): @@ -1325,7 +1334,7 @@ def tilelang_sparse_fwd( dim = q.shape[2] tail_dim = dim - d_v topk = indices.shape[-1] - assert topk == 2048 + assert topk % 64 == 0, "topk must be padded to a multiple of 64" if _is_hip: is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) @@ -1379,9 +1388,12 @@ def tilelang_sparse_fwd( ) out = kernel_combine(partial_o_batched, partial_lse_batched) else: - kernel = sparse_attention_fwd_kernel_v2( - num_heads, d_v, tail_dim, topk, sm_scale=sm_scale + kernel_factory = ( + sparse_attention_fwd_kernel_v1 + if tail_dim == 0 + else sparse_attention_fwd_kernel_v2 ) + kernel = kernel_factory(num_heads, d_v, tail_dim, topk, sm_scale=sm_scale) out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0)) # type: ignore return out diff --git a/python/sglang/kernels/ops/attention/dsa_metadata.py b/python/sglang/kernels/ops/attention/dsa_metadata.py index 21a1fedb4..df430081a 100644 --- a/python/sglang/kernels/ops/attention/dsa_metadata.py +++ b/python/sglang/kernels/ops/attention/dsa_metadata.py @@ -90,7 +90,9 @@ def _fused_dsa_decode_metadata_kernel( # fused decode CUDA graph drops it and consumes real_page_table alone. if HAS_PAGE_TABLE_1: tl.store( - page_table_1 + row * page_table_stride_0 + offs_n * page_table_stride_1, + page_table_1 + + row.to(tl.int64) * page_table_stride_0 + + offs_n * page_table_stride_1, vals, mask=mask, ) @@ -100,7 +102,7 @@ def _fused_dsa_decode_metadata_kernel( real_cols = offs_n // real_page_size tl.store( real_page_table - + row * real_page_table_stride_0 + + row.to(tl.int64) * real_page_table_stride_0 + real_cols * real_page_table_stride_1, vals // real_page_size, mask=real_mask, @@ -320,7 +322,9 @@ def _fused_dsa_target_verify_metadata_kernel( # fused_dsa_decode_metadata for the optional-page_table_1 contract). if HAS_PAGE_TABLE_1: tl.store( - page_table_1 + out_row * page_table_stride_0 + offs_n * page_table_stride_1, + page_table_1 + + out_row.to(tl.int64) * page_table_stride_0 + + offs_n * page_table_stride_1, vals, mask=mask, ) @@ -330,7 +334,7 @@ def _fused_dsa_target_verify_metadata_kernel( real_cols = offs_n // real_page_size tl.store( real_page_table - + out_row * real_page_table_stride_0 + + out_row.to(tl.int64) * real_page_table_stride_0 + real_cols * real_page_table_stride_1, vals // real_page_size, mask=real_mask, @@ -592,7 +596,7 @@ def _fused_dsa_draft_extend_metadata_kernel( if HAS_PAGE_TABLE_1: tl.store( page_table_1 - + out_rows[:, None] * page_table_stride_0 + + out_rows.to(tl.int64)[:, None] * page_table_stride_0 + offs_n[None, :] * page_table_stride_1, vals[None, :], mask=mask, @@ -603,7 +607,7 @@ def _fused_dsa_draft_extend_metadata_kernel( real_cols = offs_n // real_page_size tl.store( real_page_table - + out_rows[:, None] * real_page_table_stride_0 + + out_rows.to(tl.int64)[:, None] * real_page_table_stride_0 + real_cols[None, :] * real_page_table_stride_1, (vals // real_page_size)[None, :], mask=real_mask, diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index 605583a31..9975619bc 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -29,7 +29,6 @@ from sglang.kernels.ops.attention.fla.utils import ( check_shared_mem, is_intel, is_nvidia, - is_tf32_supported, ) if is_intel: @@ -742,7 +741,7 @@ def recompute_w_u_fwd( BT=BT, STORE_KG=kg is not None, IS_VARLEN=cu_seqlens is not None, - DOT_PRECISION="tf32" if is_tf32_supported else "ieee", + DOT_PRECISION="ieee", **(static_config or {}), ) return w, u, kg @@ -751,8 +750,8 @@ def recompute_w_u_fwd( @triton.autotune( configs=[ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [64] - for BV in [64] + for BK in [32, 64] + for BV in [64, 128] for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], @@ -863,7 +862,7 @@ def chunk_gla_fwd_kernel_o( # [BT, BT] b_A = tl.load(p_A, boundary_check=(0, 1)) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) - b_o += tl.dot(b_A, b_v) + b_o += tl.dot(b_A, b_v, allow_tf32=False) tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index f2b0f7073..0b6fe3b97 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -564,6 +564,7 @@ def handle_deterministic_inference(server_args: Any): "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", "Glm4MoeLiteForCausalLM", + "Glm5NextForConditionalGeneration", ] except Exception: pass diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index 4b3a5dfba..a15ffb300 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/python/sglang/srt/arg_groups/cuda_graph_hook.py @@ -275,10 +275,17 @@ def disable_breakable_cudagraph_if_incompatible(server_args: Any): """ cfg = resolving_view(server_args) - from sglang.srt.configs.model_config import is_deepseek_v4 + from sglang.srt.configs.model_config import ( + is_deepseek_v4, + uses_kda_attention, + ) from sglang.srt.layers.cp.bcg import supports_prefill_cp_bcg rules = [ + ( + "KDA hybrid linear attention", + lambda: uses_kda_attention(model_config_of(server_args).hf_config), + ), # DSV4 is BCG-compatible but introduces heavy memory pressure: the # c4 indexer scratch is pinned in the capture pool and OOMs. Disable. ( diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 7d0e109cc..f11033d6a 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -183,6 +183,7 @@ def handle_model_specific_adjustments(server_args: Any): "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "Glm5NextForConditionalGeneration", "HYV4ForCausalLM", "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py index ee192eb88..47e5793b5 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v2.py @@ -1,6 +1,6 @@ """Config-time override declarations for deepseek_v2. -Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, GlmMoeDsaForCausalLM, HYV4ForCausalLM, HYV4ForCausalLMNextN, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration. +Architectures: DeepseekV32ForCausalLM, DeepseekV3ForCausalLM, Dots3NoteForCausalLM, Glm5NextForConditionalGeneration, GlmMoeDsaForCausalLM, HYV4ForCausalLM, HYV4ForCausalLMNextN, KimiK25ForConditionalGeneration, LongcatFlashForCausalLM, LongcatFlashForCausalLMNextN, MistralLarge3ForCausalLM, PixtralForConditionalGeneration. """ import logging @@ -24,6 +24,7 @@ logger = logging.getLogger(__name__) "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "Glm5NextForConditionalGeneration", "HYV4ForCausalLM", "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 0415d8f65..6d4e463aa 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -512,6 +512,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset( "Lfm2ForCausalLM", "Lfm2MoeForCausalLM", "ZayaForCausalLM", + "Glm5NextForConditionalGeneration", } ) @@ -533,6 +534,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( "BailingMoeV3ForCausalLM", "FalconH1ForCausalLM", "GraniteMoeHybridForCausalLM", + "Glm5NextForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", # KDA-based: same MambaPool ping-pong machinery as GDN; requires the @@ -792,6 +794,7 @@ _DEEPSEEK_FAMILY_ARCHS = frozenset( "MistralLarge3ForCausalLM", "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", + "Glm5NextForConditionalGeneration", "HYV4ForCausalLM", "HYV4ForCausalLMNextN", "LongcatFlashForCausalLM", diff --git a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py index a615e59ea..285d686a6 100644 --- a/python/sglang/srt/arg_groups/pd_disaggregation_hook.py +++ b/python/sglang/srt/arg_groups/pd_disaggregation_hook.py @@ -256,9 +256,10 @@ def handle_encoder_disaggregation(server_args: Any): "KimiK25ForConditionalGeneration", "KimiK3ForConditionalGeneration", "MiMoV2ForCausalLM", + "Glm5NextForConditionalGeneration", ]: raise ValueError( f"Model type {model_arch} is not supported for encoder disaggregation. " f"Supported architectures: Qwen2VL, Qwen3VL, Qwen3.5, InternS2, " - f"Qwen2Audio, Qwen2.5Omni, Dots3-Note, Kimi, MiMoV2." + f"Qwen2Audio, Qwen2.5Omni, Dots3-Note, Kimi, MiMoV2, GLM5Next." ) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index f0d3b94d0..093f29125 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -16,6 +16,7 @@ from sglang.srt.configs.dots_ocr import DotsOCRConfig from sglang.srt.configs.dots_vlm import DotsVLMConfig from sglang.srt.configs.exaone import ExaoneConfig from sglang.srt.configs.falcon_h1 import FalconH1Config +from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig from sglang.srt.configs.hy_v4 import HYV4Config from sglang.srt.configs.inkling import ( @@ -102,6 +103,8 @@ __all__ = [ "Olmo3Config", "MuseGlimmerConfig", "MuseGlimmerAssistantConfig", + "Glm5NextConfig", + "Glm5NextTextConfig", "KimiLinearConfig", "KimiK3Config", "KimiK25Config", diff --git a/python/sglang/srt/configs/glm5_next.py b/python/sglang/srt/configs/glm5_next.py new file mode 100644 index 000000000..17c149ff0 --- /dev/null +++ b/python/sglang/srt/configs/glm5_next.py @@ -0,0 +1,341 @@ +from typing import List, Optional, Union + +from transformers.configuration_utils import PretrainedConfig +from transformers.models.glm_ocr.configuration_glm_ocr import GlmOcrVisionConfig + +from sglang.srt.configs.mamba_utils import KimiLinearCacheParams, KimiLinearStateShape +from sglang.srt.runtime_context import get_parallel + +_GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS = ( + "architectures", + "vocab_size", + "hidden_size", + "head_dim", + "intermediate_size", + "moe_intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "hidden_act", + "max_position_embeddings", + "rms_norm_eps", + "use_cache", + "pad_token_id", + "bos_token_id", + "eos_token_id", + "rope_theta", + "rope_scaling", + "rope_parameters", + "partial_rotary_factor", + "tie_word_embeddings", + "attention_bias", + "attention_dropout", + "n_routed_experts", + "num_experts_per_tok", + "n_shared_experts", + "n_group", + "topk_group", + "norm_topk_prob", + "routed_scaling_factor", + "scoring_func", + "topk_method", + "first_k_dense_replace", + "moe_layer_freq", + "q_lora_rank", + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "swiglu_limit", + "mhc", + "hc_mult", + "hc_sinkhorn_iters", + "hc_eps", + "num_nextn_predict_layers", + "linear_attn_config", + "linear_head_dim", + "linear_num_heads", + "linear_conv_kernel_dim", + "linear_lower_bound", + "gate_lower_bound", + "index_head_dim", + "index_topk", + "index_kpool", + "index_kpool_always_select_tail", + "index_kpool_compress", + "index_n_heads", + "index_topk_freq", + "index_topk_pattern", + "index_skip_topk_offset", + "index_share_for_mtp_iteration", + "indexer_rope_interleave", + "layer_types", + "mlp_layer_types", + "quantization_config", +) + + +class Glm5NextTextConfig(PretrainedConfig): + model_type = "glm5_next_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 154880, + hidden_size: int = 4096, + head_dim: Optional[int] = None, + intermediate_size: int = 12288, + moe_intermediate_size: int = 2048, + num_hidden_layers: int = 45, + num_attention_heads: int = 64, + num_key_value_heads: Optional[int] = None, + hidden_act: str = "silu", + max_position_embeddings: int = 1013760, + rms_norm_eps: float = 1e-5, + use_cache: bool = True, + pad_token_id: Optional[int] = None, + bos_token_id: Optional[int] = None, + eos_token_id: Optional[Union[int, List[int]]] = None, + rope_theta: float = 800000.0, + rope_scaling: Optional[dict] = None, + rope_parameters: Optional[dict] = None, + partial_rotary_factor: float = 1.0, + tie_word_embeddings: bool = False, + attention_bias: bool = False, + attention_dropout: float = 0.0, + n_routed_experts: Optional[int] = 288, + num_experts_per_tok: int = 7, + n_shared_experts: int = 1, + n_group: int = 1, + topk_group: int = 1, + norm_topk_prob: bool = True, + routed_scaling_factor: float = 2.5, + scoring_func: str = "sigmoid", + topk_method: str = "noaux_tc", + first_k_dense_replace: int = 3, + moe_layer_freq: int = 1, + q_lora_rank: Optional[int] = 1536, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 256, + qk_rope_head_dim: int = 0, + v_head_dim: int = 256, + swiglu_limit: Optional[float] = None, + mhc: bool = False, + hc_mult: int = 4, + hc_sinkhorn_iters: int = 20, + hc_eps: float = 1e-6, + num_nextn_predict_layers: int = 1, + linear_attn_config: Optional[dict] = None, + linear_head_dim: int = 128, + linear_num_heads: int = 64, + linear_conv_kernel_dim: int = 4, + linear_lower_bound: Optional[float] = None, + gate_lower_bound: Optional[float] = None, + index_head_dim: int | None = None, + index_topk: int | None = None, + index_n_heads: int | None = None, + index_topk_freq: int = 1, + index_topk_pattern: Optional[str] = None, + index_skip_topk_offset: Optional[int] = None, + **kwargs, + ): + if rope_scaling is None and rope_parameters is not None: + rope_scaling = rope_parameters + if rope_parameters is not None: + rope_theta = rope_parameters.get("rope_theta", rope_theta) + partial_rotary_factor = rope_parameters.get( + "partial_rotary_factor", partial_rotary_factor + ) + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = head_dim + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.partial_rotary_factor = partial_rotary_factor + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.n_routed_experts = n_routed_experts + self.num_experts_per_tok = num_experts_per_tok + self.n_shared_experts = n_shared_experts + self.n_group = n_group + self.topk_group = topk_group + self.norm_topk_prob = norm_topk_prob + self.routed_scaling_factor = routed_scaling_factor + self.scoring_func = scoring_func + self.topk_method = topk_method + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.swiglu_limit = swiglu_limit + self.mhc = mhc + self.hc_mult = hc_mult + self.hc_sinkhorn_iters = hc_sinkhorn_iters + self.hc_eps = hc_eps + + self.num_nextn_predict_layers = num_nextn_predict_layers + self.linear_head_dim = linear_head_dim + self.linear_num_heads = linear_num_heads + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_lower_bound = linear_lower_bound + self.gate_lower_bound = ( + gate_lower_bound if gate_lower_bound is not None else linear_lower_bound + ) + if linear_attn_config is None: + layer_types = kwargs.get("layer_types") + if layer_types is None: + kda_layers = [ + layer_idx + for layer_idx in range(num_hidden_layers) + if layer_idx % 4 != 3 + ] + else: + kda_layers = [ + layer_idx + for layer_idx, layer_type in enumerate(layer_types) + if layer_type == "linear_attention" + ] + kda_layer_set = set(kda_layers) + linear_attn_config = { + "full_attn_layers": [ + layer_idx + for layer_idx in range(num_hidden_layers) + if layer_idx not in kda_layer_set + ], + "head_dim": linear_head_dim, + "kda_layers": kda_layers, + "num_heads": linear_num_heads, + "short_conv_kernel_size": linear_conv_kernel_dim, + "gate_lower_bound": self.gate_lower_bound, + } + self.linear_attn_config = linear_attn_config + self.index_head_dim = index_head_dim + self.index_topk = index_topk + self.index_n_heads = index_n_heads + self.index_topk_freq = index_topk_freq + self.index_topk_pattern = index_topk_pattern + self.index_skip_topk_offset = index_skip_topk_offset + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + if rope_parameters is not None or rope_scaling is not None: + self.rope_parameters = rope_parameters or rope_scaling + + def is_kda_layer(self, layer_idx: int): + return ( + self.linear_attn_config is not None + and layer_idx in self.linear_attn_config["kda_layers"] + ) + + @property + def linear_layer_ids(self): + return [i for i in range(self.num_hidden_layers) if self.is_kda_layer(i)] + + @property + def nextn_layer_ids(self): + num_nextn_layers = self.num_nextn_predict_layers or 0 + return [self.num_hidden_layers + i for i in range(num_nextn_layers)] + + @property + def full_attention_layer_ids(self): + return [i for i in range(self.num_hidden_layers) if not self.is_kda_layer(i)] + + @property + def mamba2_cache_params(self) -> KimiLinearCacheParams: + shape = KimiLinearStateShape.create( + tp_world_size=get_parallel().attn_tp_size, + num_heads=self.linear_attn_config["num_heads"], + head_dim=self.linear_attn_config["head_dim"], + conv_kernel_size=self.linear_attn_config["short_conv_kernel_size"], + ) + + return KimiLinearCacheParams(shape=shape, layers=self.linear_layer_ids) + + +class Glm5NextVisionConfig(GlmOcrVisionConfig): + def __init__( + self, + swiglu_limit: float, + **kwargs, + ): + super().__init__(**kwargs) + self.swiglu_limit = swiglu_limit + + +class Glm5NextConfig(PretrainedConfig): + model_type = "glm5_next" + sub_configs = { + "vision_config": Glm5NextVisionConfig, + "text_config": Glm5NextTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id: int = 59280, + video_token_id: int = 59281, + image_start_token_id: int = 59256, + image_end_token_id: int = 59257, + video_start_token_id: int = 59258, + video_end_token_id: int = 59259, + **kwargs, + ): + top_level_text_config = { + key: kwargs[key] + for key in _GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS + if key in kwargs + } + + if isinstance(text_config, dict): + text_config = {**top_level_text_config, **text_config} + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"](**top_level_text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = None + else: + if isinstance(vision_config, dict): + vision_config = dict(vision_config) + else: + vision_config = vision_config.to_dict() + self.vision_config = self.sub_configs["vision_config"](**vision_config) + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.image_start_token_id = image_start_token_id + self.image_end_token_id = image_end_token_id + self.video_start_token_id = video_start_token_id + self.video_end_token_id = video_end_token_id + + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(**kwargs) + for key in _GLM5_NEXT_TOP_LEVEL_CONFIG_KEYS: + if hasattr(self.text_config, key): + setattr(self, key, getattr(self.text_config, key)) diff --git a/python/sglang/srt/configs/hybrid_arch.py b/python/sglang/srt/configs/hybrid_arch.py index 5f1b3cc3c..af2f0105f 100644 --- a/python/sglang/srt/configs/hybrid_arch.py +++ b/python/sglang/srt/configs/hybrid_arch.py @@ -113,6 +113,16 @@ def kimi_linear_config(model_config: ModelConfig): return None +def glm5_next_config(model_config: ModelConfig): + hf_config = model_config.hf_config + if ( + getattr(hf_config, "model_type", None) == "glm5_next" + and not model_config.is_draft_model + ): + return hf_config.get_text_config() + return None + + def linear_attn_model_spec(model_config: ModelConfig): result = _get_linear_attn_registry_result(model_config) return result[0] if result else None @@ -123,6 +133,7 @@ def mambaish_config(model_config: ModelConfig): mamba2_config(model_config) or hybrid_gdn_config(model_config) or kimi_linear_config(model_config) + or glm5_next_config(model_config) or hybrid_lightning_config(model_config) ) if existing: diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 18cf06d5a..43c6be020 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -138,6 +138,8 @@ def is_deepseek_dsa(config) -> bool: "PixtralForConditionalGeneration", "GlmMoeDsaForCausalLM", "GlmMoeDsaForCausalLMNextN", + "Glm5NextForConditionalGenerationNextN", + "Glm5NextForConditionalGeneration", "LongcatFlashForCausalLM", "LongcatFlashForCausalLMNextN", "Dots3NoteForCausalLM", @@ -156,6 +158,31 @@ def is_kimi_k3(config) -> bool: ) +def uses_kda_attention(config) -> bool: + configs = [config] + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + configs.append(get_text_config()) + else: + text_config = _hf_attr(config, "text_config") + if text_config is not None: + configs.append(text_config) + for config in configs: + linear_attn_config = _hf_attr(config, "linear_attn_config") + if isinstance(linear_attn_config, dict) and linear_attn_config.get( + "kda_layers" + ): + return True + layer_types = _hf_attr(config, "layer_types") or [] + if ( + "linear_attention" in layer_types + and _hf_attr(config, "linear_num_heads") is not None + and _hf_attr(config, "linear_head_dim") is not None + ): + return True + return False + + def is_dspark_draft(config) -> bool: return _hf_arch(config) == "DSparkDraftModel" @@ -285,6 +312,21 @@ def get_dsa_index_n_heads(config: PretrainedConfig) -> int: return config.index_n_heads +def get_dsa_index_kpool(config: PretrainedConfig) -> int: + return getattr(config, "index_kpool", 1) + + +def get_dsa_mtp_topk_width(config: PretrainedConfig) -> int: + """MTP seeds include index_topk pooled tokens plus up to index_kpool - 1 tail tokens.""" + index_kpool = get_dsa_index_kpool(config) + assert index_kpool >= 1, f"index_kpool must be positive, got {index_kpool}" + return config.index_topk + index_kpool - 1 + + +def get_dsa_index_kpool_compress(config: PretrainedConfig) -> bool: + return getattr(config, "index_kpool_compress", False) + + REQUANTIZATION_METHODS = ["quark_mxfp4"] @@ -731,6 +773,15 @@ class ModelConfig: ): self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN" + if ( + is_draft_model + and self.hf_config.architectures[0] == "Glm5NextForConditionalGeneration" + ): + self.hf_config.architectures[0] = "Glm5NextForConditionalGenerationNextN" + self.hf_text_config.architectures = list(self.hf_config.architectures) + self.hf_text_config.num_nextn_predict_layers = 1 + self.hf_text_config.linear_attn_config = None + if is_draft_model and self.hf_config.architectures[0] in [ "GlmOcrForConditionalGeneration", ]: @@ -968,6 +1019,8 @@ class ModelConfig: or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures or "GlmMoeDsaForCausalLM" in self.hf_config.architectures or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures + or "Glm5NextForConditionalGeneration" in self.hf_config.architectures + or "Glm5NextForConditionalGenerationNextN" in self.hf_config.architectures or "LongcatFlashForCausalLM" in self.hf_config.architectures or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures or "HYV4ForCausalLM" in self.hf_config.architectures @@ -1085,7 +1138,10 @@ class ModelConfig: self.v_head_dim = self.hf_config.v_head_dim self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) - elif "SarvamMLAForCausalLM" in self.hf_config.architectures: + elif ( + "SarvamMLAForCausalLM" in self.hf_config.architectures + or "Glm5NextForConditionalGeneration" in self.hf_config.architectures + ): self.head_dim = ( self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim ) @@ -1138,9 +1194,19 @@ class ModelConfig: self.num_key_value_heads = self.num_attention_heads self.hidden_size = self.hf_text_config.hidden_size hc_mult = getattr(self.hf_text_config, "hc_mult", 1) - self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size( - self.hf_config, self.hidden_size, hc_mult + is_glm5_next = getattr(self.hf_config, "model_type", None) == "glm5_next" or ( + getattr(self.hf_text_config, "model_type", None) == "glm5_next_text" ) + if is_glm5_next and not getattr(self.hf_text_config, "mhc", False): + hc_mult = 1 + if is_glm5_next: + # mHC-flattened hidden size; None when not running an mHC model. + self.hc_hidden_size = self.hidden_size * hc_mult if hc_mult > 1 else None + self.spec_hidden_size = self.hidden_size + else: + self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size( + self.hf_config, self.hidden_size, hc_mult + ) self.num_hidden_layers = self.hf_text_config.num_hidden_layers self.num_attention_layers = self.num_hidden_layers if "LongcatFlashForCausalLM" in self.hf_config.architectures: @@ -1950,6 +2016,7 @@ multimodal_model_archs = [ "Gemma4UnifiedForConditionalGeneration", "Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", + "Glm5NextForConditionalGeneration", "GlmOcrForConditionalGeneration", "GlmAsrForConditionalGeneration", "GlmImageForConditionalGeneration", @@ -2017,6 +2084,7 @@ piecewise_cuda_graph_disabled_model_archs = [ "DeepseekV4ForCausalLMNextN", "DeepseekV4ForCausalLMDSpark", "Qwen3NextForCausalLM", + "Glm5NextForConditionalGeneration", "BailingMoeV2_5ForCausalLM", "LLaDAModelLM", ] diff --git a/python/sglang/srt/configs/qwen3_asr.py b/python/sglang/srt/configs/qwen3_asr.py index cb3252b3d..bf2ce3a07 100644 --- a/python/sglang/srt/configs/qwen3_asr.py +++ b/python/sglang/srt/configs/qwen3_asr.py @@ -164,5 +164,5 @@ class Qwen3ASRConfig(PretrainedConfig): return self.thinker_config.text_config -AutoConfig.register("qwen3_asr", Qwen3ASRConfig) -AutoConfig.register("qwen3_asr_thinker", Qwen3ASRThinkerConfig) +AutoConfig.register("qwen3_asr", Qwen3ASRConfig, exist_ok=True) +AutoConfig.register("qwen3_asr_thinker", Qwen3ASRThinkerConfig, exist_ok=True) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 4996cc6eb..358ff95c4 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -18,6 +18,9 @@ class StateType(str, enum.Enum): MAMBA = "mamba" SWA = "swa" DSA = "dsa" + # DSA kpool-compress tail: one per-request ring row. The indices encode + # only the live subrange of that row for the current open pool. + DSA_TAIL = "dsa_tail" MINIMAX_INDEX_K = "minimax_index_k" # DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot # (req_pool_idx * ring_stride + pos % ring_stride), needs its own component. diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 163f45fee..39763024a 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1009,7 +1009,7 @@ class CommonKVManager(BaseKVManager): returned unchanged. """ start_layer = self.kv_args.prefill_start_layer - end_layer = getattr(self.kv_args, "prefill_end_layer", None) + end_layer = self.kv_args.prefill_end_layer assert end_layer is not None, ( "KVArgs.prefill_end_layer must be set when using compressed-MLA PD with PP" ) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 03cd52a33..a36674982 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -54,6 +54,7 @@ from sglang.srt.disaggregation.utils import ( _is_fake_transfer, build_kv_layer_ids, build_staging_slot_metadata, + get_dsa_tail_state_indices, get_dsv4_c128_state_indices, get_kv_class, is_dsv4_c128_online_enabled, @@ -1411,6 +1412,13 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): device_page_size = self.token_to_kv_pool.page_size return kv_to_page_indices(kv_indices_full, device_page_size) + def _dsa_tail_payload(): + return get_dsa_tail_state_indices( + self.token_to_kv_pool, + decode_req.req.kv.req_pool_idx, + seq_len, + ) + def _swa_ring_payload(): # Mirror of prefill _swa_ring_payload using this side's req_pool_idx. # Same window positions and order -> positional match with prefill. @@ -1443,6 +1451,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): StateType.MAMBA: _mamba_payload, StateType.SWA: _swa_payload, StateType.DSA: _full_kv_pages_payload, + StateType.DSA_TAIL: _dsa_tail_payload, StateType.MINIMAX_INDEX_K: _full_kv_pages_payload, StateType.SWA_RING: _swa_ring_payload, StateType.C128_STATE: _c128_state_payload, diff --git a/python/sglang/srt/disaggregation/encoder/preprocessor.py b/python/sglang/srt/disaggregation/encoder/preprocessor.py index 364747746..2fca0d798 100644 --- a/python/sglang/srt/disaggregation/encoder/preprocessor.py +++ b/python/sglang/srt/disaggregation/encoder/preprocessor.py @@ -11,6 +11,7 @@ import concurrent.futures import functools import logging from dataclasses import dataclass +from http import HTTPStatus from typing import Callable, List, Optional, Tuple, Union import numpy as np @@ -26,6 +27,17 @@ from sglang.srt.multimodal.encoder_preprocessing import ( EncoderPreprocessOutput, invoke_encoder_preprocessor, ) +from sglang.srt.multimodal.processors.glm4v import ( + _glm_effective_presize_budget, + glm_budget_kwargs, + glm_decode_frames_at, + glm_max_image_tokens_from_configs, + glm_processor_video_config, + glm_sample_and_decode_sync, + glm_sample_frame_indices, + preprocess_video_frames_sync, + split_glm_video_items, +) from sglang.srt.multimodal.processors.qwen_vl import preprocess_video from sglang.srt.runtime_context import ( get_device, @@ -220,9 +232,12 @@ class EncoderPreprocessor: self.vision_config[modality_str]["device"] = self.device if modality_str == "video": - video_defaults = {"fps": 2.0, "max_frames": 768, "min_frames": 4} - for k, v in video_defaults.items(): - self.vision_config["video"].setdefault(k, v) + # GLM reads its own defaults from the HF video processor + # (max_frames=2048); applying the Qwen values here would clobber them. + if "glm" not in self.model_type: + video_defaults = {"fps": 2.0, "max_frames": 768, "min_frames": 4} + for k, v in video_defaults.items(): + self.vision_config["video"].setdefault(k, v) if modality_str == "audio": if "return_attention_mask" not in self.vision_config["audio"]: @@ -334,7 +349,17 @@ class EncoderPreprocessor: } return img elif modality == Modality.VIDEO: - return load_video(data, frame_count_limit) + vid = load_video(data, frame_count_limit) + if ( + media_metadata + and self.encoder_media_processor_config.preserve_media_metadata + ): + return { + "type": "video", + "video": vid, + **media_metadata, + } + return vid elif modality == Modality.AUDIO: return load_audio(data, self.model_audio_sr) @@ -399,10 +424,109 @@ class EncoderPreprocessor: async def _flatten_and_load_images(self, mm_items): return await self._flatten_and_load_data_by_modality(mm_items, Modality.IMAGE) + @staticmethod + def _close_video_decoders(video_items) -> None: + for video in video_items or []: + close = getattr(video, "close", None) + if callable(close): + close() + + async def _dp_sharded_decode_single_video( + self, + vr, + video_config, + *, + tp_rank: int, + tp_size: int, + video_processor_kwargs: dict, + precomputed_indices: Optional[List[int]] = None, + ): + video_config = video_config or {} + video_fps = vr.avg_fps + duration = len(vr) / video_fps if video_fps else 0 + global_indices = precomputed_indices or glm_sample_frame_indices( + len(vr), + video_fps, + duration, + target_fps=video_config.get("fps"), + max_frame_count=video_config.get("max_frames"), + ) + n_units = len(global_indices) // 2 + base, remainder = divmod(n_units, tp_size) + gpu_sample_counts = [ + base + (1 if rank < remainder else 0) for rank in range(tp_size) + ] + start = sum(gpu_sample_counts[:tp_rank]) + count = gpu_sample_counts[tp_rank] + local_indices = global_indices[2 * start : 2 * (start + count)] + + local_error = None + frames = None + try: + frames = await asyncio.get_running_loop().run_in_executor( + self.io_executor, + glm_decode_frames_at, + vr, + local_indices, + video_config, + ) + except Exception as exc: + local_error = exc + + # All ranks must either enter the later ViT all-gather or fail before + # it. A rank-local decoder error must therefore be agreed globally. + ok = torch.tensor([0 if local_error else 1], dtype=torch.int32) + if tp_size > 1: + torch.distributed.all_reduce( + ok, + op=torch.distributed.ReduceOp.MIN, + group=get_parallel().attn_tp_group.cpu_group, + ) + if not int(ok.item()): + if local_error is not None: + raise local_error + from sglang.srt.disaggregation.encoder.server import MMError + + raise MMError( + "peer encoder rank failed during sharded video decode", + code=HTTPStatus.SERVICE_UNAVAILABLE, + ) + if frames is None: + height, width = vr.frame_shape + frames = np.zeros((0, height, width, 3), dtype=np.uint8) + + video_processor_kwargs["do_sample_frames"] = False + video_processor_kwargs["return_metadata"] = True + # Preserve the same per-frame spatial budget as the unsharded request. + if global_indices and local_indices: + budget = video_config.get("max_image_tokens") + if budget is None: + budget = getattr(self.video_processor, "max_image_tokens", None) + if budget is not None: + video_processor_kwargs["max_image_tokens"] = max( + 1, int(int(budget) * len(local_indices) / len(global_indices)) + ) + video_processor_kwargs["_dp_meta"] = { + "global_indices": list(global_indices), + "fps": video_fps, + "n_units": n_units, + "gpu_sample_counts": gpu_sample_counts, + } + return [frames], video_processor_kwargs + async def _flatten_and_load_videos(self, mm_items): if not isinstance(mm_items, (list, tuple)): mm_items = [mm_items] + video_configs = [{} for _ in mm_items] + if "glm" in self.model_type: + mm_items, video_configs = split_glm_video_items(mm_items) + defaults = glm_processor_video_config(self.video_processor) + defaults.update(self.vision_config.get("video", {})) + video_configs = [ + {**defaults, **dict(config or {})} for config in video_configs + ] + futures, _ = self._submit_data_loading_tasks( mm_items, [Modality.VIDEO] * len(mm_items) ) @@ -422,10 +546,87 @@ class EncoderPreprocessor: if video_metadata: video_processor_kwargs["video_metadata"] = video_metadata return videos, video_processor_kwargs - else: - raise NotImplementedError( - f"Video processing is not supported for {self.model_type} model." + + if "glm" in self.model_type: + budget_kwargs = glm_budget_kwargs( + self.video_processor, + user_max_image_tokens=glm_max_image_tokens_from_configs(video_configs), + count=len(video_items), + split=True, ) + if budget_kwargs is not None: + video_processor_kwargs.update(budget_kwargs) + video_configs = [ + _glm_effective_presize_budget( + config, budget_kwargs.get("max_image_tokens") + ) + for config in video_configs + ] + + framed = any(isinstance(video, list) for video in video_items) + if framed: + processed = await asyncio.gather( + *[ + asyncio.get_running_loop().run_in_executor( + self.io_executor, preprocess_video_frames_sync, video + ) + for video in video_items + ] + ) + else: + parallel = get_parallel() + tp_size = parallel.attn_tp_size + sampled = None + if len(video_items) == 1: + vr = video_items[0] + config = video_configs[0] + sampled = glm_sample_frame_indices( + len(vr), + vr.avg_fps, + len(vr) / vr.avg_fps if vr.avg_fps else 0, + target_fps=config.get("fps"), + max_frame_count=config.get("max_frames"), + ) + if ( + self.server_args.mm_enable_dp_encoder + and tp_size > 1 + and sampled is not None + and len(sampled) >= max(32, tp_size * 2) + ): + result = await self._dp_sharded_decode_single_video( + video_items[0], + video_configs[0], + tp_rank=parallel.attn_tp_rank, + tp_size=tp_size, + video_processor_kwargs=video_processor_kwargs, + precomputed_indices=sampled, + ) + self._close_video_decoders(video_items) + return result + + processed = await asyncio.gather( + *[ + asyncio.get_running_loop().run_in_executor( + self.io_executor, + glm_sample_and_decode_sync, + video, + video_configs[index], + ) + for index, video in enumerate(video_items) + ] + ) + videos, video_metadata = map(list, zip(*processed)) + video_processor_kwargs["do_sample_frames"] = False + video_processor_kwargs["return_metadata"] = True + if video_metadata: + video_processor_kwargs["video_metadata"] = video_metadata + self._close_video_decoders(video_items) + return videos, video_processor_kwargs + + self._close_video_decoders(video_items) + raise NotImplementedError( + f"Video processing is not supported for {self.model_type} model." + ) async def _flatten_and_load_audios(self, mm_items): return await self._flatten_and_load_data_by_modality(mm_items, Modality.AUDIO) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 2c03adead..d8914b736 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -47,9 +47,11 @@ from sglang.srt.disaggregation.mooncake.utils import ( ) from sglang.srt.disaggregation.utils import ( DisaggregationMode, + build_dsa_tail_transfer_blocks, build_transfer_entry_pairs, compute_mamba_state_slice_byte_blocks, resolve_dcp_dst_entry_indices, + slice_dsa_tail_dst_ptrs_for_pp, ) from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs @@ -1406,10 +1408,25 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): ) or rc ) + elif st == StateType.DSA_TAIL: + rc = ( + self._send_slot_state( + req, + src_data_ptrs, + src_item_lens, + dst_data_ptrs, + dst_item_lens, + list(indices), + list(dst_indices), + st.value, + ) + or rc + ) elif self._is_generic_kvcache_state_type(st): if ( target_rank_registration_info is not None and not self.is_mla_backend + and not self.is_hybrid_mla_backend and self.attn_tp_size != target_rank_registration_info.dst_attn_tp_size ): @@ -1491,6 +1508,43 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): ) return rc + def _send_slot_state( + self, + req: TransferInfo, + src_ptrs: list[int], + src_item_lens: list[int], + dst_ptrs: list[int], + dst_item_lens: list[int], + src_indices: list[int], + dst_indices: list[int], + label: str, + ) -> int: + try: + dst_ptrs = slice_dsa_tail_dst_ptrs_for_pp( + src_ptrs, + dst_ptrs, + self.kv_args.prefill_start_layer, + self.kv_args.prefill_end_layer, + ) + dst_item_lens = slice_dsa_tail_dst_ptrs_for_pp( + src_ptrs, + dst_item_lens, + self.kv_args.prefill_start_layer, + self.kv_args.prefill_end_layer, + ) + transfer_blocks = build_dsa_tail_transfer_blocks( + src_ptrs, + src_item_lens, + dst_ptrs, + src_indices, + dst_indices, + dst_item_lens, + ) + except ValueError as exc: + logger.error("%s: %s", label, exc) + return -1 + return self._transfer_data(req.mooncake_session_id, transfer_blocks) + def _send_mamba_state( self, req: TransferInfo, @@ -1549,7 +1603,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): attn_tp_size, we slice the state accordingly. GDN conv_state is the concatenation [query | key | value] with each sub-block head-sharded independently, so on the scatter path it is sliced per sub-block via - ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks). + ``src_state_conv_shard_groups`` (see + compute_mamba_state_slice_byte_blocks). """ logger.warning_once( "Using Mamba state slice transfer for different TP sizes between prefill and decode. " diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 8e1ceb7fd..a7fc753ad 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -41,9 +41,11 @@ from sglang.srt.disaggregation.common.utils import ( ) from sglang.srt.disaggregation.utils import ( DisaggregationMode, + build_dsa_tail_transfer_blocks, build_transfer_entry_pairs, compute_mamba_state_slice_byte_blocks, resolve_dcp_dst_entry_indices, + slice_dsa_tail_dst_ptrs_for_pp, ) from sglang.srt.environ import envs from sglang.srt.runtime_context import get_parallel, get_schedule @@ -2086,6 +2088,60 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): raise Exception("KVSender failed to post transfer") return xfer_handle + def _send_slot_state( + self, + peer_name: str, + src_data_ptrs: list[int], + src_item_lens: list[int], + dst_data_ptrs: list[int], + dst_item_lens: list[int], + src_indices: list[int], + dst_indices: list[int], + dst_gpu_id: int, + notif: str, + ): + dst_data_ptrs = slice_dsa_tail_dst_ptrs_for_pp( + src_data_ptrs, + dst_data_ptrs, + self.kv_args.prefill_start_layer, + self.kv_args.prefill_end_layer, + ) + dst_item_lens = slice_dsa_tail_dst_ptrs_for_pp( + src_data_ptrs, + dst_item_lens, + self.kv_args.prefill_start_layer, + self.kv_args.prefill_end_layer, + ) + transfer_blocks = build_dsa_tail_transfer_blocks( + src_data_ptrs, + src_item_lens, + dst_data_ptrs, + src_indices, + dst_indices, + dst_item_lens, + ) + if not transfer_blocks: + return None + + src_addrs = [ + (src_addr, length, self.kv_args.gpu_id) + for src_addr, _, length in transfer_blocks + ] + dst_addrs = [ + (dst_addr, length, dst_gpu_id) for _, dst_addr, length in transfer_blocks + ] + src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM") + dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") + xfer_handle = self.agent.initialize_xfer( + "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii") + ) + if not xfer_handle: + raise Exception("KVSender failed to create dsa_tail transfer") + state = self.agent.transfer(xfer_handle) + if state == "ERR": + raise Exception("KVSender failed to post dsa_tail transfer") + return xfer_handle + def _send_mamba_state( self, peer_name: str, @@ -2169,7 +2225,8 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): accordingly, mirroring Mooncake's _send_mamba_state_slice. GDN conv_state is [query | key | value] with each sub-block head-sharded independently, so on the scatter path it is sliced per sub-block via - ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks). + ``src_state_conv_shard_groups`` (see + compute_mamba_state_slice_byte_blocks). """ logger.warning_once( "Using Mamba state slice transfer for different TP sizes. " @@ -2308,7 +2365,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): src_indices = ( prefill_state_indices[i] if i < len(prefill_state_indices) else None ) - if src_indices is None or len(src_indices) == 0: + if src_indices is None or ( + len(src_indices) == 0 and st != StateType.DSA_TAIL + ): continue src_ptrs = src_state_data_ptrs[i] if i < len(src_state_data_ptrs) else [] src_lens = src_state_item_lens[i] if i < len(src_state_item_lens) else [] @@ -2369,12 +2428,37 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): src_layer_ids=src_lids, dst_layer_ids=dst_lids, ) - elif st in ( - StateType.SWA, - StateType.DSA, - StateType.SWA_RING, - StateType.C128_STATE, - ): + elif st == StateType.DSA_TAIL: + h = self._send_slot_state( + peer_name, + src_ptrs, + src_lens, + dst_ptrs, + dst_lens, + list(src_indices), + list(dst_indices), + dst_gpu_id, + comp_notif, + ) + elif st == StateType.DSA: + if len(src_indices) != len(dst_indices): + raise RuntimeError( + f"State index length mismatch at component {i}: " + f"prefill={len(src_indices)}, dst={len(dst_indices)}" + ) + h = self._send_kvcache_generic( + peer_name=peer_name, + src_data_ptrs=src_ptrs, + dst_data_ptrs=dst_ptrs, + item_lens=src_lens, + prefill_data_indices=np.array(src_indices, dtype=np.int32), + dst_data_indices=np.array(dst_indices, dtype=np.int32), + dst_gpu_id=dst_gpu_id, + notif=comp_notif, + state_type=st, + force_flat=True, + ) + elif st in (StateType.SWA, StateType.SWA_RING, StateType.C128_STATE): if not self.is_mla_backend and self.attn_tp_size != decode_tp_size: raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 67248b3e2..d4163ea88 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -45,6 +45,7 @@ from sglang.srt.disaggregation.utils import ( TransferBackend, build_kv_layer_ids, build_staging_slot_metadata, + get_dsa_tail_state_indices, get_dsv4_c128_state_indices, get_kv_class, is_aborted, @@ -906,6 +907,7 @@ class SchedulerDisaggregationPrefillMixin: can_run_cuda_graph=can_run_cuda_graph, dp_cooperation_info=batch.dp_cooperation_info, ) + self.maybe_send_health_check_signal() @scheduler_stage_method(SCHEDULER_STAGE_PROCESS_QUEUE) def process_disagg_prefill_inflight_queue( @@ -1306,6 +1308,13 @@ class SchedulerDisaggregationPrefillMixin: ] return kv_to_page_indices(kv_indices_full, page_size) + def _dsa_tail_payload(): + return get_dsa_tail_state_indices( + self.token_to_kv_pool_allocator.get_kvcache(), + req.kv.req_pool_idx, + seq_len, + ) + def _swa_ring_payload(): # Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride) # for the last `window` positions, in ascending position order so @@ -1342,6 +1351,7 @@ class SchedulerDisaggregationPrefillMixin: StateType.MAMBA: _mamba_payload, StateType.SWA: _swa_payload, StateType.DSA: _full_kv_pages_payload, + StateType.DSA_TAIL: _dsa_tail_payload, StateType.MINIMAX_INDEX_K: _full_kv_pages_payload, StateType.SWA_RING: _swa_ring_payload, StateType.C128_STATE: _c128_state_payload, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 9309b3b6a..56aaec3d3 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -19,7 +19,7 @@ import numpy as np import torch import torch.distributed as dist -from sglang.srt.configs.model_config import get_dsa_index_topk +from sglang.srt.configs.model_config import get_dsa_mtp_topk_width from sglang.srt.disaggregation.base import KVPoll from sglang.srt.environ import envs from sglang.srt.runtime_context import ( @@ -71,7 +71,7 @@ def get_dsa_seed_metadata_dim(hf_config) -> int: """Return the model-defined PD seed width, independent of local spec mode.""" if not getattr(hf_config, "index_share_for_mtp_iteration", False): return 0 - return get_dsa_index_topk(hf_config) + return get_dsa_mtp_topk_width(hf_config) def is_dsv4_c128_online_enabled() -> bool: @@ -1067,6 +1067,191 @@ def append_state_component( kv_args.state_layer_ids.append(layer_ids or []) +def get_dsa_tail_state_indices(pool, req_pool_idx: int, seq_len: int) -> List[int]: + if getattr(pool, "use_dsa", False): + pool = pool.full_kv_pool + if not pool.kpool_use_compress: + return [] + + pool_size = int(pool.index_kpool) + tail_size = pool_size + int(getattr(pool, "tail_extra_slots", 0)) + if pool_size <= 1 or tail_size < pool_size: + raise ValueError( + "DSA kpool-compress requires pool_size > 1 and " + f"tail_size >= pool_size, got pool_size={pool_size}, " + f"tail_size={tail_size}" + ) + + n_valid = int(seq_len) % pool_size + if n_valid == 0: + return [] + start_phys = (int(seq_len) - n_valid) % tail_size + first_n = min(n_valid, tail_size - start_phys) + second_n = n_valid - first_n + return [ + int(req_pool_idx), + start_phys, + first_n, + 0, + second_n, + tail_size, + ] + + +def slice_dsa_tail_dst_ptrs_for_pp( + src_ptrs: List[int], + dst_ptrs: List[int], + start_layer: int, + end_layer: Optional[int], +) -> List[int]: + if len(src_ptrs) == len(dst_ptrs): + return list(dst_ptrs) + if len(src_ptrs) % 2 != 0 or len(dst_ptrs) % 2 != 0: + raise ValueError( + "DSA tail pointer lists must contain equal key/score halves, got " + f"src={len(src_ptrs)}, dst={len(dst_ptrs)}" + ) + + src_layers = len(src_ptrs) // 2 + dst_layers = len(dst_ptrs) // 2 + expected_end = start_layer + src_layers + if end_layer is not None and end_layer - start_layer == src_layers: + expected_end = end_layer + if start_layer < 0 or expected_end > dst_layers: + raise ValueError( + "DSA tail pointer count mismatch: " + f"src={len(src_ptrs)}, dst={len(dst_ptrs)}, " + f"prefill_layers=[{start_layer}, {expected_end})" + ) + + return list(dst_ptrs[start_layer:expected_end]) + list( + dst_ptrs[dst_layers + start_layer : dst_layers + expected_end] + ) + + +def build_dsa_tail_transfer_blocks( + src_ptrs: List[int], + src_item_lens: List[int], + dst_ptrs: List[int], + src_indices: List[int], + dst_indices: List[int], + dst_item_lens: Optional[List[int]] = None, +) -> List[Tuple[int, int, int]]: + """Remap live DSA tail tokens between rings with different speculative-slot counts.""" + if not src_indices and not dst_indices: + return [] + if not src_indices or not dst_indices: + raise ValueError( + f"DSA tail slot index missing: src={src_indices}, dst={dst_indices}" + ) + if len(src_indices) != 6 or len(dst_indices) != 6: + raise ValueError( + "DSA tail slot indices must be 6-tuples, " + f"got src={src_indices}, dst={dst_indices}" + ) + if dst_item_lens is None: + dst_item_lens = src_item_lens + if not (len(src_ptrs) == len(dst_ptrs) == len(src_item_lens) == len(dst_item_lens)): + raise ValueError( + "DSA tail pointer metadata mismatch: " + f"src_ptrs={len(src_ptrs)}, dst_ptrs={len(dst_ptrs)}, " + f"src_item_lens={len(src_item_lens)}, " + f"dst_item_lens={len(dst_item_lens)}" + ) + + src_tail_size = int(src_indices[5]) + dst_tail_size = int(dst_indices[5]) + if src_tail_size <= 0 or dst_tail_size <= 0: + raise ValueError( + "DSA tail ring sizes must be positive: " + f"src={src_tail_size}, dst={dst_tail_size}" + ) + + def parse_segments(indices: List[int], tail_size: int, side: str): + segments = [] + for seg in (1, 2): + off = int(indices[seg * 2 - 1]) + n = int(indices[seg * 2]) + if min(off, n) < 0: + raise ValueError( + f"DSA tail {side} offsets and lengths must be non-negative" + ) + if off + n > tail_size: + raise ValueError( + f"DSA tail {side} segment {seg} exceeds ring size " + f"{tail_size}: ({off}, {n})" + ) + if n: + segments.append((off, n)) + return segments + + src_segments = parse_segments(src_indices, src_tail_size, "source") + dst_segments = parse_segments(dst_indices, dst_tail_size, "destination") + src_count = sum(n for _, n in src_segments) + dst_count = sum(n for _, n in dst_segments) + if src_count != dst_count: + raise ValueError( + f"DSA tail live-token count mismatch: src={src_count}, dst={dst_count}" + ) + + src_idx = int(src_indices[0]) + dst_idx = int(dst_indices[0]) + if src_idx < 0 or dst_idx < 0: + raise ValueError("DSA tail request row indices must be non-negative") + + transfer_blocks = [] + for src_ptr, src_row_bytes, dst_ptr, dst_row_bytes in zip( + src_ptrs, src_item_lens, dst_ptrs, dst_item_lens + ): + src_row_bytes = int(src_row_bytes) + dst_row_bytes = int(dst_row_bytes) + if src_row_bytes == 0 and dst_row_bytes == 0: + continue + if src_row_bytes <= 0 or src_row_bytes % src_tail_size != 0: + raise ValueError( + f"DSA source tail row size {src_row_bytes} is not divisible by " + f"{src_tail_size}" + ) + if dst_row_bytes <= 0 or dst_row_bytes % dst_tail_size != 0: + raise ValueError( + f"DSA destination tail row size {dst_row_bytes} is not " + f"divisible by {dst_tail_size}" + ) + src_slot_bytes = src_row_bytes // src_tail_size + dst_slot_bytes = dst_row_bytes // dst_tail_size + if src_slot_bytes != dst_slot_bytes: + raise ValueError( + "DSA tail slot-size mismatch: " + f"src={src_slot_bytes}, dst={dst_slot_bytes}" + ) + + slot_bytes = src_slot_bytes + src_row_base = int(src_ptr) + src_row_bytes * src_idx + dst_row_base = int(dst_ptr) + dst_row_bytes * dst_idx + src_seg_idx = dst_seg_idx = 0 + src_consumed = dst_consumed = 0 + while src_seg_idx < len(src_segments): + src_off, src_n = src_segments[src_seg_idx] + dst_off, dst_n = dst_segments[dst_seg_idx] + n = min(src_n - src_consumed, dst_n - dst_consumed) + transfer_blocks.append( + ( + src_row_base + (src_off + src_consumed) * slot_bytes, + dst_row_base + (dst_off + dst_consumed) * slot_bytes, + n * slot_bytes, + ) + ) + src_consumed += n + dst_consumed += n + if src_consumed == src_n: + src_seg_idx += 1 + src_consumed = 0 + if dst_consumed == dst_n: + dst_seg_idx += 1 + dst_consumed = 0 + return transfer_blocks + + def setup_state_kv_args( kv_args: KVArgs, token_to_kv_pool, @@ -1100,6 +1285,19 @@ def setup_state_kv_args( kv_args.is_hybrid_mla_backend = False kv_args.state_conv_shard_groups = [] + def append_dsa_tail(pool) -> None: + if not pool.kpool_use_compress: + return + tail_ptrs, tail_lens, tail_item_lens = pool.get_compress_tail_buf_infos() + if tail_ptrs: + append_state_component( + kv_args, + StateType.DSA_TAIL, + tail_ptrs, + tail_lens, + tail_item_lens, + ) + if isinstance(token_to_kv_pool, MHATokenToKVPoolMXFP8): append_state_component( kv_args, @@ -1203,7 +1401,25 @@ def setup_state_kv_args( slice_outer_counts, layer_ids, ) + # Hybrid DSA pools keep their index cache and kpool tail in the + # full-attention sub-pool rather than in the Mamba state above. + if getattr(token_to_kv_pool, "use_dsa", False): + dsa_pool = token_to_kv_pool.full_kv_pool + dsa_ptrs, dsa_lens, dsa_item_lens = dsa_pool.get_state_buf_infos() + append_state_component( + kv_args, + StateType.DSA, + dsa_ptrs, + dsa_lens, + dsa_item_lens, + ) + append_dsa_tail(dsa_pool) elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)): + tail_ptrs, tail_lens, tail_item_lens = [], [], [] + if isinstance(token_to_kv_pool, DSATokenToKVPool): + tail_ptrs, tail_lens, tail_item_lens = ( + token_to_kv_pool.get_compress_tail_buf_infos() + ) if draft_token_to_kv_pool is not None and isinstance( draft_token_to_kv_pool, DSATokenToKVPool ): @@ -1215,6 +1431,12 @@ def setup_state_kv_args( data_ptrs = data_ptrs + draft_data_ptrs data_lens = data_lens + draft_data_lens item_lens = item_lens + draft_item_lens + draft_tail_ptrs, draft_tail_lens, draft_tail_item_lens = ( + draft_token_to_kv_pool.get_compress_tail_buf_infos() + ) + tail_ptrs = tail_ptrs + draft_tail_ptrs + tail_lens = tail_lens + draft_tail_lens + tail_item_lens = tail_item_lens + draft_tail_item_lens if isinstance(token_to_kv_pool, NPUMLATokenToKVPool): kv_args.kv_buf_groups = ( len(kv_args.kv_data_ptrs) // token_to_kv_pool.layer_num @@ -1224,6 +1446,14 @@ def setup_state_kv_args( append_state_component( kv_args, StateType.DSA, data_ptrs, data_lens, item_lens ) + if tail_ptrs: + append_state_component( + kv_args, + StateType.DSA_TAIL, + tail_ptrs, + tail_lens, + tail_item_lens, + ) if is_npu() and isinstance(token_to_kv_pool, DSV4NPUTokenToKVPool): from sglang.srt.disaggregation.ascend.conn import AscendStateType diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index c91320b24..d893fba92 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -355,7 +355,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): # NPUMHATokenToKVPool stores buffers as # (num_pages, page_size, head_num, head_dim) # use_fia=False # (num_pages*page_size, 1, head_num, head_dim) # use_fia=True - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): torch.npu.synchronize() buf_of_layers = [] for local_layer_id in range(self.layer_num): @@ -370,7 +370,9 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): torch.npu.synchronize() return kv_cache_cpu - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): torch.npu.synchronize() chunk_size = self.cpu_offloading_chunk_size for local_layer_id in range(self.layer_num): @@ -743,7 +745,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): out.append(layer_chunks) return out - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): torch.npu.synchronize() buf_of_layers = [] has_ik = self.index_head_dim is not None @@ -761,7 +763,9 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): torch.npu.synchronize() return kv_cache_cpu - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): torch.npu.synchronize() chunk_size = self.cpu_offloading_chunk_size has_ik = self.index_head_dim is not None diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index bbf45f359..20a7e2003 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -7,6 +7,7 @@ from sglang.srt.arg_groups.overrides import ( resolved_view, ) from sglang.srt.configs.hybrid_arch import ( + glm5_next_config, hybrid_gdn_config, hybrid_lightning_config, kimi_linear_config, @@ -505,6 +506,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac hybrid_backend_cls = AscendKDAHybridLinearAttnBackend else: linear_attn_backend = KDAAttnBackend(runner) + elif glm5_next_config(runner.model_config) is not None: + linear_attn_backend = KDAAttnBackend(runner) elif hybrid_lightning_config(runner.model_config) is not None: linear_attn_backend = LightningAttentionBackend(runner) else: diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py new file mode 100644 index 000000000..acf04fcb7 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.attention.dsa.kpool_plan import ( + init_kpool_extend_metadata, + init_kpool_write_plan, + init_kpool_write_plan_capture, + init_pooled_paged_mqa_metadata, + update_kpool_write_plan, + update_pooled_paged_mqa_metadata, +) + +if TYPE_CHECKING: + from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( + PrecomputedMetadata, + ) + from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod + from sglang.srt.layers.attention.dsa_backend import _DSA_IMPL_T, DSAMetadata + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode + + +@dataclass +class _KPoolForwardInputs: + full_real_page_table: Optional[torch.Tensor] = None + full_seqlens_expanded: Optional[torch.Tensor] = None + + +class DeepseekSparseAttnBackendKPoolMixin: + """KPool-specific metadata and tail handling for the DSA backend.""" + + def _check_kpool_tail_backend( + self, + topk_indices: Optional[torch.Tensor], + dsa_impl: _DSA_IMPL_T, + phase: str, + ) -> None: + if ( + topk_indices is None + or self.dsa_index_kpool <= 1 + or dsa_impl in ("fa3", "tilelang", "trtllm") + ): + return + raise NotImplementedError( + "index_kpool > 1 appends tail tokens to topk_indices and is " + f"currently only supported by the FA3/TileLang/TRTLLM DSA {phase} " + "backend." + ) + + def _resolve_kpool_tail_backend( + self, + topk_indices: Optional[torch.Tensor], + dsa_impl: _DSA_IMPL_T, + ) -> _DSA_IMPL_T: + if ( + topk_indices is None + or self.dsa_index_kpool <= 1 + or dsa_impl != "flashmla_sparse" + ): + return dsa_impl + if self.device_sm_major >= 10: + return "trtllm" + if self.device_sm_major == 9: + return "fa3" + return dsa_impl + + def _kpool_slots_per_page(self) -> int: + return getattr(self.token_to_kv_pool, "slots_per_page", self.real_page_size) + + def _build_kpool_paged_mqa_schedule_metadata(self) -> bool: + if self.device_sm_major == 9: + return self.num_q_heads in (32, 64) + return True + + def _init_kpool_metadata( + self, + metadata: DSAMetadata, + forward_batch: ForwardBatch, + topk_transform_method: Optional[TopkTransformMethod] = None, + kpool_inputs: Optional[_KPoolForwardInputs] = None, + ) -> DSAMetadata: + if self.dsa_index_kpool <= 1: + return metadata + + forward_mode = forward_batch.forward_mode + slots_per_page = self._kpool_slots_per_page() + build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata() + if forward_mode.is_extend_without_speculative(): + assert topk_transform_method is not None + assert kpool_inputs is not None + return init_kpool_extend_metadata( + metadata, + forward_batch, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + slots_per_page=slots_per_page, + topk_transform_method=topk_transform_method, + full_real_page_table=kpool_inputs.full_real_page_table, + full_seqlens_expanded=kpool_inputs.full_seqlens_expanded, + ) + + if forward_mode.is_decode_or_idle(): + metadata = init_pooled_paged_mqa_metadata( + metadata, + metadata.cache_seqlens_int32, + forward_mode, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + return init_kpool_write_plan( + metadata, + forward_batch, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + real_page_table=metadata.real_page_table, + num_draft_tokens=1, + write_start=(forward_batch.seq_lens - 1).to(torch.int32), + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + + if forward_mode.is_target_verify(): + return init_kpool_write_plan( + metadata, + forward_batch, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + real_page_table=metadata.real_page_table, + num_draft_tokens=self.speculative_num_draft_tokens, + write_start=forward_batch.seq_lens.to(torch.int32), + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + + if forward_mode.is_draft_extend_v2(): + spec_info = forward_batch.spec_info + effective_n_per_batch = ( + spec_info.num_accept_tokens + if spec_info is not None + and getattr(spec_info, "num_accept_tokens", None) is not None + else None + ) + return init_kpool_write_plan( + metadata, + forward_batch, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + real_page_table=metadata.real_page_table, + num_draft_tokens=self.speculative_num_draft_tokens, + write_start=( + forward_batch.seq_lens - self.speculative_num_draft_tokens + ).to(torch.int32), + slots_per_page=slots_per_page, + effective_n_per_batch=effective_n_per_batch, + build_schedule_metadata=build_schedule_metadata, + ) + + return metadata + + def _init_kpool_metadata_capture( + self, metadata: DSAMetadata, bs: int, forward_mode: ForwardMode + ) -> DSAMetadata: + if self.dsa_index_kpool <= 1: + return metadata + + slots_per_page = self._kpool_slots_per_page() + build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata() + if forward_mode.is_decode_or_idle(): + metadata = init_pooled_paged_mqa_metadata( + metadata, + metadata.cache_seqlens_int32, + forward_mode, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + + if ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend_v2() + ): + is_decode = forward_mode.is_decode_or_idle() + is_v2 = forward_mode.is_draft_extend_v2() + metadata = init_kpool_write_plan_capture( + metadata, + max_bs=bs, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + num_draft_tokens=( + 1 if is_decode else self.speculative_num_draft_tokens + ), + device=self.device, + is_verify=not is_decode, + slots_per_page=slots_per_page, + is_v2=is_v2, + build_schedule_metadata=build_schedule_metadata, + ) + + return metadata + + def _update_kpool_metadata_replay( + self, + metadata: DSAMetadata, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + forward_mode: ForwardMode, + effective_n_per_batch: Optional[torch.Tensor] = None, + ) -> None: + if self.dsa_index_kpool <= 1: + return + + slots_per_page = self._kpool_slots_per_page() + build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata() + if forward_mode.is_decode_or_idle(): + update_pooled_paged_mqa_metadata( + metadata, + metadata.cache_seqlens_int32, + forward_mode, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + + if not ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend_v2() + ): + return + + is_decode = forward_mode.is_decode_or_idle() + is_v2 = forward_mode.is_draft_extend_v2() + if is_decode: + write_start = seq_lens.to(torch.int32) - 1 + elif is_v2: + write_start = seq_lens.to(torch.int32) - self.speculative_num_draft_tokens + else: + # Target verify: write_start == seq_lens exactly; the plan kernel + # casts on load, so skip the per-replay int32 alloc + conversion. + write_start = seq_lens + update_kpool_write_plan( + metadata, + write_start=write_start, + req_pool_indices=req_pool_indices, + real_page_table=metadata.real_page_table, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + num_draft_tokens=(1 if is_decode else self.speculative_num_draft_tokens), + forward_mode=forward_mode, + slots_per_page=slots_per_page, + effective_n_per_batch=effective_n_per_batch, + ) + + def _update_kpool_metadata_from_precomputed( + self, + metadata: DSAMetadata, + precomputed: PrecomputedMetadata, + forward_mode: ForwardMode, + ) -> None: + if self.dsa_index_kpool <= 1: + return + + slots_per_page = self._kpool_slots_per_page() + build_schedule_metadata = self._build_kpool_paged_mqa_schedule_metadata() + if forward_mode.is_decode_or_idle(): + update_pooled_paged_mqa_metadata( + metadata, + precomputed.cache_seqlens, + forward_mode, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + slots_per_page=slots_per_page, + build_schedule_metadata=build_schedule_metadata, + ) + + if not (forward_mode.is_decode_or_idle() or forward_mode.is_target_verify()): + return + + is_verify = forward_mode.is_target_verify() + write_start = precomputed.cache_seqlens.to(torch.int32) + write_start = ( + write_start - self.speculative_num_draft_tokens + if is_verify + else write_start - 1 + ) + update_kpool_write_plan( + metadata, + write_start=write_start, + req_pool_indices=precomputed.req_pool_indices, + real_page_table=metadata.real_page_table, + pool_size=self.dsa_index_kpool, + real_page_size=self.real_page_size, + num_draft_tokens=self.speculative_num_draft_tokens if is_verify else 1, + forward_mode=forward_mode, + slots_per_page=slots_per_page, + ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index e4b90b7ef..6950e4a34 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -35,6 +35,7 @@ class PrecomputedMetadata: # Basic seqlens cache_seqlens: torch.Tensor # int32, [bs] cu_seqlens_k: torch.Tensor # int32, [bs+1] + req_pool_indices: torch.Tensor # int64, [bs] # Page table page_indices: torch.Tensor # int32, [bs, max_len] or [expanded_bs, max_len] @@ -121,7 +122,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: """Precompute metadata for normal decode mode.""" max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1] - if _is_cuda and not _is_hip: + if _is_cuda and not _is_hip and self.dsa_index_kpool <= 1: from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, ) @@ -173,6 +174,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: return PrecomputedMetadata( cache_seqlens=cache_seqlens, cu_seqlens_k=cu_seqlens_k, + req_pool_indices=req_pool_indices, page_indices=page_indices, real_page_table=real_page_table, seqlens_expanded=seqlens_expanded, @@ -193,7 +195,9 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: # Compute DSA seqlens dsa_cache_seqlens = compute_dsa_seqlens( - cache_seqlens, dsa_index_topk=self.dsa_index_topk + cache_seqlens, + dsa_index_topk=self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) seqlens_expanded = cache_seqlens seqlens_expanded_size = seqlens_expanded.shape[0] @@ -218,6 +222,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: return PrecomputedMetadata( cache_seqlens=cache_seqlens, cu_seqlens_k=cu_seqlens_k, + req_pool_indices=req_pool_indices, page_indices=page_indices, real_page_table=real_page_table, seqlens_expanded=seqlens_expanded, @@ -240,7 +245,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1] seqlens_expanded_size = bs * self.speculative_num_draft_tokens - if _is_cuda and not _is_hip: + if _is_cuda and not _is_hip and self.dsa_index_kpool <= 1: from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_target_verify_metadata, ) @@ -305,6 +310,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: return PrecomputedMetadata( cache_seqlens=cache_seqlens, cu_seqlens_k=cu_seqlens_k, + req_pool_indices=req_pool_indices, page_indices=page_indices, real_page_table=real_page_table, seqlens_expanded=seqlens_expanded, @@ -342,7 +348,11 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: ) # Compute DSA seqlens - dsa_cache_seqlens = compute_dsa_seqlens(seqlens_expanded, self.dsa_index_topk) + dsa_cache_seqlens = compute_dsa_seqlens( + seqlens_expanded, + self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, + ) seqlens_expanded_size = seqlens_expanded.shape[0] # DSA cumsum @@ -365,6 +375,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: return PrecomputedMetadata( cache_seqlens=cache_seqlens, cu_seqlens_k=cu_seqlens_k, + req_pool_indices=req_pool_indices, page_indices=page_indices, real_page_table=real_page_table, seqlens_expanded=seqlens_expanded, diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 37b1bfa4e..a7e834743 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -486,7 +486,6 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): dim=-1, ) with torch.cuda.stream(self.alt_stream): - # TODO we should also put DeepGEMM half SM here? if self.use_dsa_indexer_fusion: key, weights_raw = self._fused_k_weights(x) else: @@ -876,10 +875,18 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): if use_dg_native: seqlens_32_2d = ctx_2d + elif ctx_2d is not None: + if ctx_2d.size(1) == 1: + seqlens_32_2d = ctx_2d + else: + seqlens_32_2d = ctx_2d.reshape(-1).contiguous().view(-1, 1) elif seqlens_32.dim() == 2: - seqlens_32_2d = seqlens_32 + if seqlens_32.size(1) == 1: + seqlens_32_2d = seqlens_32.contiguous() + else: + seqlens_32_2d = seqlens_32.reshape(-1).contiguous().view(-1, 1) else: - seqlens_32_2d = seqlens_32.unsqueeze(-1) + seqlens_32_2d = seqlens_32.contiguous().view(-1, 1) if _is_cuda: if schedule_metadata is None: schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( @@ -896,6 +903,52 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): assert len(weights.shape) == 3 weights = weights.squeeze(2) + # SM100 DeepGEMM paged MQA requires batch_size <= num_sms; chunk larger batches. + def _chunked_fp8_paged_mqa_logits( + q: torch.Tensor, + kv_cache: torch.Tensor, + w: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + mqa_schedule_metadata: torch.Tensor, + max_len: int, + clean_logits: bool = False, + ) -> torch.Tensor: + batch_size, chunk_next_n = q.shape[:2] + if batch_size == 0: + return torch.empty((0, max_len), dtype=torch.float32, device=q.device) + if batch_size <= self.sm_count: + return deep_gemm.fp8_paged_mqa_logits( + q, + kv_cache, + w, + context_lens, + block_table, + mqa_schedule_metadata, + max_len, + clean_logits=clean_logits, + ) + logits_chunks = [] + for start in range(0, batch_size, self.sm_count): + end = min(start + self.sm_count, batch_size) + chunk_context_lens = context_lens[start:end] + chunk_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + chunk_context_lens, blocksize, self.sm_count + ) + logits_chunks.append( + deep_gemm.fp8_paged_mqa_logits( + q[start:end], + kv_cache, + w[start * chunk_next_n : end * chunk_next_n], + chunk_context_lens, + block_table[start:end], + chunk_schedule_metadata, + max_len, + clean_logits=clean_logits, + ) + ) + return torch.cat(logits_chunks, dim=0) + if self.paged_mqa_logits_backend.is_aiter(): logits = aiter_paged_mqa_logits( q_fp8, @@ -928,7 +981,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): ) elif use_dg_native: logits = deepgemm_paged_mqa_logits_native( - deep_gemm.fp8_paged_mqa_logits, + _chunked_fp8_paged_mqa_logits, q_fp8, kv_cache_fp8, weights, @@ -942,7 +995,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): ) else: logits = deepgemm_paged_mqa_logits_split( - deep_gemm.fp8_paged_mqa_logits, + _chunked_fp8_paged_mqa_logits, q_fp8, kv_cache_fp8, weights, diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py new file mode 100644 index 000000000..4d24d61ff --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py @@ -0,0 +1,1520 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers import PretrainedConfig + +from sglang.srt.layers.attention.dsa.dsa_indexer import ( + DUAL_STREAM_TOKEN_THRESHOLD, + BaseIndexerMetadata, + rotate_activation, +) +from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod +from sglang.srt.layers.layernorm import LayerNorm +from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu + +if is_cuda(): + try: + import deep_gemm + except ImportError as e: + deep_gemm = e + +if is_npu(): + import custom_ops # noqa: F401 + +from sglang.srt.environ import envs +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding import get_rope_wrapper +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, + get_req_to_token_pool, + get_token_to_kv_pool, +) +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.runtime_context import get_device + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + + +class IndexerKPool(MultiPlatformOp): + def __init__( + self, + hidden_size: int, + index_n_heads: int, + index_head_dim: int, + rope_head_dim: int, + index_topk: int, + q_lora_rank: int, + max_position_embeddings: int, + rope_theta: float, + layer_id: int, + scale_fmt: Optional[str], + block_size: int = 128, + rope_scaling: Optional[Dict[str, Any]] = None, + is_neox_style: bool = True, + prefix: str = "", + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + skip_rope: bool = False, + config: Optional[PretrainedConfig] = None, + ): + super().__init__() + self.hidden_size = hidden_size + self.n_heads = index_n_heads + self.head_dim = index_head_dim + self.rope_head_dim = rope_head_dim + self.index_topk = index_topk + self.q_lora_rank = q_lora_rank + self.layer_id = layer_id + self.alt_stream = alt_stream + self.compress_gate_stream = None + self.skip_rope = skip_rope + + self.index_kpool = config.index_kpool + self.index_kpool_always_select_tail = config.index_kpool_always_select_tail + self.index_kpool_compress = config.index_kpool_compress + + assert ( + self.index_kpool > 1 + and self.index_kpool_compress + and self.index_kpool_always_select_tail + ) + + assert self.index_topk % self.index_kpool == 0, ( + f"index_topk ({self.index_topk}) must be divisible by " + f"index_kpool ({self.index_kpool})" + ) + assert 64 % self.index_kpool == 0, ( + f"index_kpool ({self.index_kpool}) must divide page_size (64)" + ) + + self.index_kpool_compress_ape = nn.Parameter( + torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32) + ) + self.index_kpool_compress_gate = nn.Parameter( + torch.empty(self.head_dim, self.hidden_size, dtype=torch.bfloat16) + ) + + if is_cuda() and self.alt_stream is not None: + self.compress_gate_stream = torch.cuda.Stream() + + if is_cuda(): + self.sm_count = deep_gemm.get_num_sms() + self.half_device_sm_count = ceil_align(self.sm_count // 2, 8) + + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + ) + + self.wk = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wk", prefix), + ) + # Keep weights_proj in FP32 because its accumulation/scale path requires FP32; + # checkpoint BF16 weights are cast on load. + self.weights_proj = ReplicatedLinear( + self.hidden_size, + self.n_heads, + bias=False, + params_dtype=torch.float32, + prefix=add_prefix("weights_proj", prefix), + ) + self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32) + if not self.skip_rope: + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, # type: ignore + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + device=get_device().device, + ) + self.block_size = block_size + self.scale_fmt = scale_fmt + self.softmax_scale = self.head_dim**-0.5 + + @torch.compile(dynamic=True) + def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): + weights, _ = self.weights_proj(x.float()) + weights = weights * self.n_heads**-0.5 + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale + return weights + + @staticmethod + def _get_index_k_read_buffer(pool, layer_id: int) -> torch.Tensor: + if hasattr(pool, "get_broadcastable_index_k_with_scale_buffer"): + return pool.get_broadcastable_index_k_with_scale_buffer(layer_id) + if hasattr(pool, "_get_broadcastable_index_buffer"): + return pool._get_broadcastable_index_buffer(layer_id) + return pool.get_index_k_with_scale_buffer(layer_id=layer_id) + + def _write_compressed_pooled_index_cache( + self, + slot_k, + slot_score, + write_locs, + forward_batch, + layer_id, + write_mask=None, + return_compressed: bool = False, + write_cache: bool = True, + ): + if slot_k.shape[0] == 0: + if return_compressed: + return ( + torch.empty( + (0, self.head_dim), + dtype=torch.float8_e4m3fn, + device=slot_k.device, + ), + torch.empty((0,), dtype=torch.float32, device=slot_k.device), + ) + return None + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_softmax_rotate_write_cache, + ) + + pool = get_token_to_kv_pool() + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + return kpool_softmax_rotate_write_cache( + pool=pool, + buf=buf, + slot_k=slot_k, + slot_score=slot_score, + ape=self.index_kpool_compress_ape, + loc=write_locs.contiguous(), + write_mask=write_mask.contiguous() if write_mask is not None else None, + round_scale=self.scale_fmt is not None, + return_compressed=return_compressed, + write_cache=write_cache, + ) + + @staticmethod + def _write_returned_compressed_pooled_index_cache( + forward_batch, layer_id, write_locs, k_fp8, k_scale + ): + if k_fp8.shape[0] == 0: + return + get_token_to_kv_pool().set_index_k_scale_buffer( + layer_id=layer_id, + loc=write_locs.contiguous(), + index_k=k_fp8.contiguous(), + index_k_scale=k_scale.contiguous(), + ) + + def _compress_write_decode( + self, + key, + gate_score, + positions, + forward_batch, + layer_id, + metadata, + ): + batch = key.shape[0] + if batch == 0: + return + + pool = get_token_to_kv_pool() + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + + pool.kpool_decode_update_index_cache( + layer_id=layer_id, + key=key, + slot_score=gate_score, + ape=self.index_kpool_compress_ape, + block_tables=metadata.get_page_table_64(), + req_pool_indices=forward_batch.req_pool_indices[:batch], + positions=positions[:batch], + seq_lens=metadata.get_seqlens_int32()[:batch], + out_cache_loc=forward_batch.out_cache_loc[:batch], + round_scale=self.scale_fmt is not None, + ) + + def _compress_write_extend( + self, + key, + gate_score, + positions, + forward_batch, + layer_id, + metadata, + return_compressed: bool = False, + write_cache: bool = True, + ): + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + attn_metadata = metadata.attn_metadata + plan = attn_metadata.kpool_extend_plan + if plan is not None: + import os + + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_assemble_softmax_rotate_write_cache, + scatter_kpool_tail_updates, + ) + + pool = get_token_to_kv_pool() + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + + writes, tails = plan.writes, plan.tails + if writes.is_empty and tails.is_empty: + return None + + tail_k_buf, tail_score_buf = pool.get_compress_tail_buffers(layer_id) + if os.environ.get("SGLANG_DSA_KPOOL_DEBUG_BOUNDS") == "1": + if not writes.is_empty: + active = torch.ones_like(writes.req, dtype=torch.bool) + if bool(active.any().item()): + active_req = writes.req[active] + active_write_loc = writes.write_loc[active] + active_chunk_max = writes.chunk_src[active].to( + torch.long + ) + torch.clamp( + self.index_kpool + - writes.n_from_tail[active].to(torch.long) + - 1, + min=0, + ) + max_chunk = int(active_chunk_max.max().item()) + min_chunk = int(writes.chunk_src[active].min().item()) + min_req = int(active_req.min().item()) + max_req = int(active_req.max().item()) + min_loc = int(active_write_loc.min().item()) + max_loc = int(active_write_loc.max().item()) + max_page = max_loc // int(pool.slots_per_page) + if ( + min_chunk < 0 + or max_chunk >= key.shape[0] + or max_chunk >= gate_score.shape[0] + or min_req < 0 + or max_req >= tail_k_buf.shape[0] + or min_loc < 0 + or max_page + >= pool.get_index_k_with_scale_buffer( + layer_id=layer_id + ).shape[0] + ): + raise RuntimeError( + "DSA kpool write plan out of bounds: " + f"{key.shape=}, {gate_score.shape=}, {tail_k_buf.shape=}, " + f"{pool.get_index_k_with_scale_buffer(layer_id=layer_id).shape=}, " + f"{min_chunk=}, {max_chunk=}, {min_req=}, {max_req=}, " + f"{min_loc=}, {max_loc=}, {pool.slots_per_page=}" + ) + if not tails.is_empty: + tail_chunk_max = tails.chunk_src.to(torch.long) + torch.clamp( + tails.n_write.to(torch.long) - 1, min=0 + ) + max_tail_chunk = int(tail_chunk_max.max().item()) + min_tail_chunk = int(tails.chunk_src.min().item()) + min_tail_req = int(tails.req.min().item()) + max_tail_req = int(tails.req.max().item()) + if ( + min_tail_chunk < 0 + or max_tail_chunk >= key.shape[0] + or min_tail_req < 0 + or max_tail_req >= tail_k_buf.shape[0] + ): + raise RuntimeError( + "DSA kpool tail plan out of bounds: " + f"{key.shape=}, {tail_k_buf.shape=}, " + f"{min_tail_chunk=}, {max_tail_chunk=}, " + f"{min_tail_req=}, {max_tail_req=}" + ) + if not writes.is_empty: + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + kpool_assemble_softmax_rotate_write_cache( + pool=pool, + buf=buf, + chunk_k=key, + chunk_score=gate_score, + tail_k=tail_k_buf, + tail_score=tail_score_buf, + req_pool_idx=writes.req, + n_from_tail=writes.n_from_tail, + chunk_src_start=writes.chunk_src, + tail_logical_base=writes.tail_logical_base, + ape=self.index_kpool_compress_ape, + loc=writes.write_loc, + round_scale=self.scale_fmt is not None, + ) + + if not tails.is_empty: + scatter_kpool_tail_updates( + pool=pool, + chunk_k=key, + chunk_score=gate_score, + tail_k=tail_k_buf, + tail_score=tail_score_buf, + req_pool_idx=tails.req, + dst_logical_start=tails.dst_logical_start, + chunk_src_start=tails.chunk_src, + n_write=tails.n_write, + ) + return None + + kpool = self.index_kpool + block_tables = metadata.get_page_table_64() + kpool_write_locs = getattr(attn_metadata, "kpool_extend_write_locs", None) + q_offset = 0 + compressed_by_batch = ( + [None for _ in range(forward_batch.batch_size)] + if return_compressed + else None + ) + + for i in range(forward_batch.batch_size): + q_len = int(forward_batch.extend_seq_lens_cpu[i]) + if q_len == 0: + continue + + req_pool_idx = forward_batch.req_pool_indices[i].to(torch.long) + key_chunk = key[q_offset : q_offset + q_len] + score_chunk = gate_score[q_offset : q_offset + q_len] + seq_len = int(forward_batch.seq_lens_cpu[i].item()) + first_pos = seq_len - q_len + first_slot = first_pos % kpool + + if first_slot != 0: + raise NotImplementedError( + "index_kpool_compress extend requires kpool-aligned chunk " + "starts. Set chunked_prefill_size % index_kpool == 0 and " + "avoid non-aligned prefix reuse." + ) + + pool_start_id = first_pos // kpool + page_size = get_token_to_kv_pool().page_size + use_returned_compressed = return_compressed and ( + not write_cache or pool_start_id == 0 + ) + n_pools = q_len // kpool + n_drain = n_pools * kpool + if n_pools > 0: + slot_k = key_chunk[:n_drain].view(n_pools, kpool, self.head_dim) + slot_score = score_chunk[:n_drain].view(n_pools, kpool, self.head_dim) + write_locs = ( + kpool_write_locs[i] if kpool_write_locs is not None else None + ) + if write_locs is None: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + compute_pooled_write_locs, + ) + + num_token_pages = (seq_len + page_size - 1) // page_size + token_page_table = block_tables[i, :num_token_pages].contiguous() + pool_ids = pool_start_id + torch.arange( + n_pools, dtype=torch.int64, device=key.device + ) + write_locs = compute_pooled_write_locs( + token_page_table, pool_ids, kpool + ) + compressed = self._write_compressed_pooled_index_cache( + slot_k, + slot_score, + write_locs, + forward_batch, + layer_id, + return_compressed=use_returned_compressed, + write_cache=write_cache, + ) + else: + compressed = self._write_compressed_pooled_index_cache( + slot_k, + slot_score, + write_locs, + forward_batch, + layer_id, + return_compressed=use_returned_compressed, + write_cache=write_cache, + ) + if compressed_by_batch is not None and compressed is not None: + compressed_by_batch[i] = ( + pool_start_id, + compressed[0], + compressed[1], + write_locs.contiguous(), + ) + + n_remain = q_len - n_drain + get_token_to_kv_pool().set_compress_tail_for_request( + layer_id=layer_id, + req_pool_idx=req_pool_idx, + key_tail=key_chunk[n_drain:] if n_remain > 0 else key_chunk[:0], + score_tail=score_chunk[n_drain:] if n_remain > 0 else score_chunk[:0], + n_remain=n_remain, + dst_logical_start=first_pos + n_drain, + ) + q_offset += q_len + + return compressed_by_batch + + def _compress_write( + self, + x, + key, + positions, + forward_batch, + layer_id, + metadata, + gate_score: Optional[torch.Tensor] = None, + return_compressed: bool = False, + write_cache: bool = True, + ): + if key.shape[0] == 0: + return None + + if gate_score is None: + gate_score = F.linear(x, self.index_kpool_compress_gate) + + if forward_batch.forward_mode.is_decode_or_idle(): + self._compress_write_decode( + key=key, + gate_score=gate_score, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + metadata=metadata, + ) + elif forward_batch.forward_mode.is_extend(): + return self._compress_write_extend( + key=key, + gate_score=gate_score, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + metadata=metadata, + return_compressed=return_compressed, + write_cache=write_cache, + ) + else: + raise NotImplementedError( + "index_kpool_compress currently supports decode and extend only." + ) + return None + + def _compute_gate_score_if_missing( + self, x: torch.Tensor, gate_score: Optional[torch.Tensor] + ) -> torch.Tensor: + if gate_score is not None: + return gate_score + return F.linear(x, self.index_kpool_compress_gate) + + def _get_q_k_bf16( + self, + q_lora: torch.Tensor, + x: torch.Tensor, + positions: torch.Tensor, + enable_dual_stream: bool, + forward_batch: ForwardBatch, + precompute_compress_gate: bool = False, + ): + gate_score = None + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + if precompute_compress_gate: + assert self.compress_gate_stream is not None + self.compress_gate_stream.wait_stream(current_stream) + + with deep_gemm_wrapper.configure_deep_gemm_num_sms( + self.half_device_sm_count + ): + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + with torch.cuda.stream(self.alt_stream): + key, _ = self.wk(x) + key = self.k_norm(key) + + k_rope, _ = torch.split( + key, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + + if precompute_compress_gate: + with torch.cuda.stream(self.compress_gate_stream): + gate_score = F.linear(x, self.index_kpool_compress_gate) + + current_stream.wait_stream(self.alt_stream) + else: + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + if not self.skip_rope: + q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope) + + query[..., : self.rope_head_dim] = q_rope + key[..., : self.rope_head_dim] = k_rope + + query = rotate_activation(query) + + return query, key, gate_score + + def _get_k_bf16( + self, + x: torch.Tensor, + positions: torch.Tensor, + ): + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + if not self.skip_rope: + _, k_rope = self.rotary_emb(positions, k_rope, k_rope) + key[..., : self.rope_head_dim] = k_rope + return key + + def _full_topk_for_short_sequence( + self, metadata: BaseIndexerMetadata, device: torch.device + ) -> torch.Tensor: + seq_lens_expanded = metadata.get_seqlens_expanded() + dummy_logits = torch.zeros( + seq_lens_expanded.shape[0], + self.index_topk, + dtype=torch.float32, + device=device, + ) + topk_full = metadata.topk_transform(dummy_logits, self.index_topk) + if self.index_kpool == 1: + return topk_full + padding = torch.full( + (topk_full.shape[0], self.index_kpool - 1), + -1, + dtype=topk_full.dtype, + device=topk_full.device, + ) + return torch.cat([topk_full, padding], dim=1) + + def _topk_from_kpool_logits( + self, + logits: torch.Tensor, + pool_lens: torch.Tensor, + seq_lens: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + topk_offsets: Optional[torch.Tensor] = None, + row_starts: Optional[torch.Tensor] = None, + out_rows: Optional[int] = None, + page_table_row_index: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + topk_from_pooled_history_logits, + ) + + n_rows = logits.shape[0] + if ( + page_table is not None + and page_table_row_index is None + and page_table.shape[0] != n_rows + ): + page_table = page_table[:n_rows] + if topk_offsets is not None and topk_offsets.shape[0] != n_rows: + topk_offsets = topk_offsets[:n_rows] + if page_table_row_index is not None and page_table_row_index.shape[0] != n_rows: + page_table_row_index = page_table_row_index[:n_rows] + + return topk_from_pooled_history_logits( + logits=logits, + group_lengths=pool_lens, + pool_size=self.index_kpool, + topk=self.index_topk, + page_table=page_table, + topk_offsets=topk_offsets, + seq_lens=seq_lens, + row_starts=row_starts, + out_rows=out_rows, + page_table_row_index=page_table_row_index, + ) + + def _get_kpool_decode_metadata( + self, + metadata: BaseIndexerMetadata, + block_tables: torch.Tensor, + seqlens_32: torch.Tensor, + blocksize: int, + build_schedule_metadata: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + build_pooled_page_table_64, + ) + + attn_metadata = metadata.attn_metadata + plan = attn_metadata.kpool_write_plan + if plan is not None and plan.pool_seqlens_per_q is not None: + pool_seqlens = plan.pool_seqlens_per_q[: seqlens_32.shape[0]] + pool_context_lens = pool_seqlens.contiguous().view(-1, 1) + pool_block_tables = build_pooled_page_table_64( + block_tables, self.index_kpool + ).contiguous() + pool_schedule_metadata = plan.pool_schedule_metadata + if pool_schedule_metadata is None and build_schedule_metadata: + pool_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + pool_context_lens.clamp(min=1), blocksize, self.sm_count + ) + return ( + pool_seqlens, + pool_context_lens, + pool_block_tables, + pool_schedule_metadata, + ) + + pool_seqlens = attn_metadata.pooled_cache_seqlens_int32 + pool_block_tables = attn_metadata.pooled_real_page_table + pool_schedule_metadata = attn_metadata.pooled_paged_mqa_schedule_metadata + + if ( + pool_seqlens is None + or pool_block_tables is None + or attn_metadata.pooled_index_kpool != self.index_kpool + ): + pool_seqlens = torch.div( + seqlens_32, self.index_kpool, rounding_mode="floor" + ).to(torch.int32) + pool_block_tables = build_pooled_page_table_64( + block_tables, self.index_kpool + ).contiguous() + pool_schedule_metadata = None + else: + pool_seqlens = pool_seqlens[: seqlens_32.shape[0]] + pool_block_tables = pool_block_tables[ + : block_tables.shape[0], + : (block_tables.shape[1] + self.index_kpool - 1) // self.index_kpool, + ] + + pool_context_lens = pool_seqlens.contiguous().view(-1, 1) + if pool_schedule_metadata is None and build_schedule_metadata: + pool_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + pool_context_lens.clamp(min=1), blocksize, self.sm_count + ) + + return ( + pool_seqlens, + pool_context_lens, + pool_block_tables, + pool_schedule_metadata, + ) + + @staticmethod + def _kpool_fused_topk_mapping( + metadata: BaseIndexerMetadata, + paged_page_table: Optional[torch.Tensor] = None, + paged_page_table_row_index: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + if not envs.SGLANG_DSA_FUSE_TOPK.get(): + return None, None, None + + topk_method = metadata.topk_transform_method + attn_metadata = metadata.attn_metadata + if topk_method == TopkTransformMethod.PAGED: + page_table_1 = ( + paged_page_table + if paged_page_table is not None + else attn_metadata.page_table_1 + ) + assert page_table_1 is not None + row_index = paged_page_table_row_index + return page_table_1, None, row_index + if topk_method == TopkTransformMethod.RAGGED: + return None, attn_metadata.topk_indices_offset, None + return None, None, None + + @staticmethod + def _should_use_tilelang_paged_mqa_logits(q_fp8: torch.Tensor) -> bool: + if not is_cuda(): + return False + arch_major, _ = torch.cuda.get_device_capability(q_fp8.device) + num_heads = q_fp8.shape[2] + return arch_major == 9 and num_heads not in (32, 64) + + def _get_topk_paged( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + pool = get_token_to_kv_pool() + page_size = pool.page_size + # DeepGEMM paged-MQA requires 64-token pages. + assert page_size == 64, "only support page size 64" + + block_tables = metadata.get_page_table_64() + + kv_cache_fp8 = self._get_index_k_read_buffer(pool, layer_id) + + blocksize = page_size + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + seqlens_32 = metadata.get_seqlens_expanded() + else: + seqlens_32 = metadata.get_seqlens_int32() + assert len(q_fp8.shape) == 3 + num_q_padded = q_fp8.shape[0] + n_real = seqlens_32.shape[0] + if n_real < num_q_padded: + q_fp8 = q_fp8[:n_real] + weights = weights[:n_real] + q_fp8 = q_fp8.unsqueeze(1) # the next_n dim is 1 now + assert len(kv_cache_fp8.shape) == 2 + block_kv = 64 + num_heads_kv = 1 + head_dim_with_sf = 132 + kv_cache_fp8 = kv_cache_fp8.view( + kv_cache_fp8.shape[0], block_kv, num_heads_kv, head_dim_with_sf + ) + assert len(weights.shape) == 3 + weights = weights.squeeze(2) + use_tilelang_paged_mqa = self._should_use_tilelang_paged_mqa_logits(q_fp8) + + pool_seqlens, pool_context_lens, pool_block_tables, pool_schedule_metadata = ( + self._get_kpool_decode_metadata( + metadata, + block_tables, + seqlens_32, + blocksize, + build_schedule_metadata=not use_tilelang_paged_mqa, + ) + ) + pool_max_seq_len = pool_block_tables.shape[1] * blocksize + if use_tilelang_paged_mqa: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import ( + tilelang_fp8_paged_mqa_logits, + ) + + logits = tilelang_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + pool_seqlens, + pool_block_tables, + pool_schedule_metadata, + pool_max_seq_len, + clean_logits=False, + ) + else: + logits = deep_gemm.fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + pool_context_lens, + pool_block_tables, + pool_schedule_metadata, + pool_max_seq_len, + clean_logits=False, + ) + + page_table_1, topk_offsets, _ = self._kpool_fused_topk_mapping(metadata) + topk_result = self._topk_from_kpool_logits( + logits, + pool_seqlens, + seq_lens=seqlens_32, + page_table=page_table_1, + topk_offsets=topk_offsets, + out_rows=num_q_padded if num_q_padded != n_real else None, + ) + return topk_result + + def _should_chunk_mqa_logits( + self, num_q: int, num_k: int, device: torch.device + ) -> Tuple[bool, int]: + if num_q * num_k < 8_000_000: + return False, 0 + + free_mem, total_mem = torch.cuda.mem_get_info(device) + bytes_per_elem = 4 + logits_bytes = num_q * num_k * bytes_per_elem + + need_chunk = (logits_bytes * 2 > free_mem) or (logits_bytes > total_mem * 0.3) + return need_chunk, free_mem + + def _get_topk_ragged_kpool_plan( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + ) -> torch.Tensor: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + gather_index_k_scale_prefix_into, + ) + + plan = metadata.attn_metadata.kpool_extend_plan + assert plan is not None, "kpool extend plan is required" + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + + device = q_fp8.device + total_q = q_fp8.shape[0] + seq_lens_expanded = plan.seq_lens_expanded + pool_lens = plan.pooled_seq_lens_expanded + ks_per_q = plan.ragged_q_ks + ke_per_q = plan.ragged_q_ke + total_k_rows = plan.ragged_total_k_rows + + n_real = seq_lens_expanded.shape[0] + assert n_real <= total_q, ( + f"plan has more real rows ({n_real}) than q_fp8 ({total_q})" + ) + + if total_k_rows > 0: + k_u8 = plan.ragged_k_u8 + k_scale = plan.ragged_k_scale + assert k_u8 is not None and k_scale is not None + pool = get_token_to_kv_pool() + gather_index_k_scale_prefix_into( + pool=pool, + buf=self._get_index_k_read_buffer(pool, layer_id), + page_indices=plan.ragged_concat_page_table, + seq_len=total_k_rows, + k_out=k_u8, + scale_out=k_scale, + ) + k_fp8 = k_u8.view(torch.float8_e4m3fn) + logits = deep_gemm.fp8_mqa_logits( + q_fp8[:n_real].contiguous(), + (k_fp8.contiguous(), k_scale.contiguous()), + weights[:n_real].contiguous(), + ks_per_q, + ke_per_q, + clean_logits=True, + ) + else: + logits = torch.empty((n_real, 0), dtype=torch.float32, device=device) + + topk_method = metadata.topk_transform_method + attn_metadata = metadata.attn_metadata + page_table_all = None + page_table_row_index_all = None + topk_offsets_all = None + if envs.SGLANG_DSA_FUSE_TOPK.get(): + if topk_method == TopkTransformMethod.PAGED: + page_table_all = plan.ragged_paged_page_table + page_table_row_index_all = plan.ragged_paged_page_table_row_index + elif topk_method == TopkTransformMethod.RAGGED: + topk_offsets_all = attn_metadata.topk_indices_offset + + return self._topk_from_kpool_logits( + logits, + pool_lens, + seq_lens=seq_lens_expanded, + page_table=page_table_all, + topk_offsets=topk_offsets_all, + row_starts=ks_per_q, + out_rows=total_q, + page_table_row_index=page_table_row_index_all, + ) + + def _get_topk_ragged_kpool( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + extend_pooled_cache: Optional[ + List[Optional[Tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]]] + ] = None, + ) -> torch.Tensor: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + build_pooled_page_table_64, + gather_index_k_scale_prefix_into, + ) + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + + pool_size = self.index_kpool + page_size = get_token_to_kv_pool().page_size + token_nums = q_fp8.shape[0] + tail_pool = pool_size - 1 + topk_result = torch.empty( + (token_nums, self.index_topk + tail_pool), + device=q_fp8.device, + dtype=torch.int32, + ) + block_tables = metadata.get_page_table_64() + seq_lens_expanded = metadata.get_seqlens_expanded() + topk_method = metadata.topk_transform_method + attn_metadata = metadata.attn_metadata + topk_offsets = attn_metadata.topk_indices_offset + pooled_page_tables = getattr( + attn_metadata, "kpool_extend_pooled_page_tables", None + ) + zero_starts_by_batch = getattr(attn_metadata, "kpool_extend_zero_starts", None) + pooled_seq_lens_expanded = getattr( + attn_metadata, "kpool_seqlens_expanded", None + ) + + q_offset = 0 + cache_write_stream = None + pending_cache_req_pool_indices = set() + pending_cache_tensors = [] + for i in range(forward_batch.batch_size): + q_len = int(forward_batch.extend_seq_lens_cpu[i]) + if q_len == 0: + continue + + seq_len = int(forward_batch.seq_lens_cpu[i].item()) + pool_seq_len = seq_len // pool_size + req_pool_idx = int(forward_batch.req_pool_indices[i].item()) + if ( + cache_write_stream is not None + and req_pool_idx in pending_cache_req_pool_indices + ): + torch.cuda.current_stream().wait_stream(cache_write_stream) + cache_write_stream = None + pending_cache_req_pool_indices.clear() + pending_cache_tensors.clear() + + q_slice = slice(q_offset, q_offset + q_len) + local_seqlens = seq_lens_expanded[q_slice] + if pooled_seq_lens_expanded is None: + local_pool_lens = torch.div( + local_seqlens, pool_size, rounding_mode="floor" + ).to(torch.int32) + else: + local_pool_lens = pooled_seq_lens_expanded[q_slice] + + if pool_seq_len > 0: + cached_current = ( + extend_pooled_cache[i] if extend_pooled_cache is not None else None + ) + deferred_cache_write = None + if cached_current is not None: + curr_pool_start, curr_k_fp8, curr_k_scale = cached_current[:3] + curr_pool_len = curr_k_fp8.shape[0] + if len(cached_current) == 4: + deferred_cache_write = ( + cached_current[3], + curr_k_fp8, + curr_k_scale, + ) + else: + curr_pool_start, curr_pool_len = pool_seq_len, 0 + if ( + pool_seq_len > 0 + and deferred_cache_write is not None + and self.alt_stream is not None + and not get_is_capture_mode() + ): + write_locs, write_k_fp8, write_k_scale = deferred_cache_write + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + pending_cache_tensors.append( + (write_locs, write_k_fp8, write_k_scale) + ) + with torch.cuda.stream(self.alt_stream): + self._write_returned_compressed_pooled_index_cache( + forward_batch, + layer_id, + write_locs, + write_k_fp8, + write_k_scale, + ) + cache_write_stream = self.alt_stream + pending_cache_req_pool_indices.add(req_pool_idx) + deferred_cache_write = None + + if ( + cached_current is not None + and curr_pool_start == 0 + and curr_pool_len == pool_seq_len + ): + k_fp8 = curr_k_fp8 + k_scale = curr_k_scale + elif ( + cached_current is not None + and curr_pool_start >= 0 + and curr_pool_start + curr_pool_len == pool_seq_len + ): + if curr_pool_start > 0: + pooled_page_table = ( + pooled_page_tables[i] + if pooled_page_tables is not None + else None + ) + if pooled_page_table is None: + num_token_pages = (seq_len + page_size - 1) // page_size + token_page_table = block_tables[ + i, :num_token_pages + ].contiguous() + pool_pages = (curr_pool_start + page_size - 1) // page_size + pooled_page_table = build_pooled_page_table_64( + token_page_table, pool_size + )[:pool_pages].contiguous() + k_u8 = torch.empty( + (pool_seq_len, self.head_dim), + dtype=torch.uint8, + device=q_fp8.device, + ) + k_scale = torch.empty( + (pool_seq_len,), dtype=torch.float32, device=q_fp8.device + ) + pool = get_token_to_kv_pool() + gather_index_k_scale_prefix_into( + pool=pool, + buf=self._get_index_k_read_buffer(pool, layer_id), + page_indices=pooled_page_table, + seq_len=curr_pool_start, + k_out=k_u8, + scale_out=k_scale, + ) + k_u8[curr_pool_start:pool_seq_len].copy_( + curr_k_fp8.view(torch.uint8) + ) + k_scale[curr_pool_start:pool_seq_len].copy_(curr_k_scale) + k_fp8 = k_u8.view(torch.float8_e4m3fn) + else: + k_fp8 = curr_k_fp8 + k_scale = curr_k_scale + else: + pooled_page_table = ( + pooled_page_tables[i] + if pooled_page_tables is not None + else None + ) + if pooled_page_table is None: + num_token_pages = (seq_len + page_size - 1) // page_size + token_page_table = block_tables[ + i, :num_token_pages + ].contiguous() + pool_pages = (pool_seq_len + page_size - 1) // page_size + pooled_page_table = build_pooled_page_table_64( + token_page_table, pool_size + )[:pool_pages].contiguous() + seq_len_t = torch.tensor( + [pool_seq_len], dtype=torch.int32, device=q_fp8.device + ) + k_fp8, k_scale = get_token_to_kv_pool().get_index_k_scale_buffer( + layer_id, + seq_len_t, + pooled_page_table.unsqueeze(0), + pool_seq_len, + pool_seq_len, + ) + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + k_scale = k_scale.view(torch.float32).squeeze(-1) + row_starts = ( + zero_starts_by_batch[i] + if zero_starts_by_batch is not None + and zero_starts_by_batch[i] is not None + else torch.zeros((q_len,), dtype=torch.int32, device=q_fp8.device) + ) + local_logits = deep_gemm.fp8_mqa_logits( + q_fp8[q_slice].contiguous(), + (k_fp8.contiguous(), k_scale.contiguous()), + weights[q_slice].contiguous(), + row_starts, + local_pool_lens, + clean_logits=True, + ) + else: + local_logits = torch.empty( + (q_len, 0), dtype=torch.float32, device=q_fp8.device + ) + + page_table_local = None + topk_offsets_local = None + if ( + envs.SGLANG_DSA_FUSE_TOPK.get() + and topk_method == TopkTransformMethod.PAGED + ): + page_table_local = ( + get_req_to_token_pool() + .req_to_token[req_pool_idx, :seq_len] + .to(torch.int32) + ) + page_table_local = page_table_local.unsqueeze(0).expand(q_len, -1) + elif ( + envs.SGLANG_DSA_FUSE_TOPK.get() + and topk_method == TopkTransformMethod.RAGGED + and topk_offsets is not None + ): + topk_offsets_local = topk_offsets[q_slice] + + local_topk = self._topk_from_kpool_logits( + local_logits, + local_pool_lens, + seq_lens=local_seqlens, + page_table=page_table_local, + topk_offsets=topk_offsets_local, + ) + + topk_result[q_slice] = local_topk + if pool_seq_len > 0 and deferred_cache_write is not None: + write_locs, curr_k_fp8, curr_k_scale = deferred_cache_write + self._write_returned_compressed_pooled_index_cache( + forward_batch, + layer_id, + write_locs, + curr_k_fp8, + curr_k_scale, + ) + q_offset += q_len + + if cache_write_stream is not None: + torch.cuda.current_stream().wait_stream(cache_write_stream) + pending_cache_tensors.clear() + + return topk_result + + def _get_topk_ragged( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + kpool_extend_cache: Optional[ + List[Optional[Tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]]] + ] = None, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + assert forward_batch.forward_mode.is_extend_without_speculative() + + page_size = get_token_to_kv_pool().page_size + assert page_size == 64, "only support page size 64" + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + if metadata.attn_metadata.kpool_extend_plan is not None: + return self._get_topk_ragged_kpool_plan( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + ) + return self._get_topk_ragged_kpool( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + extend_pooled_cache=kpool_extend_cache, + ) + + def _forward_cuda_skip_logits( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + metadata: BaseIndexerMetadata, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + assert forward_batch.forward_mode.is_extend_without_speculative() + + key = self._get_k_bf16(x, positions) + self._compress_write( + x=x, + key=key, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + metadata=metadata, + ) + + if not return_indices: + return None + + return self._full_topk_for_short_sequence(metadata, x.device) + + def _forward_cuda_target_verify( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + metadata: BaseIndexerMetadata, + enable_dual_stream: bool, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + assert is_cuda(), "DSA kpool target_verify is CUDA-only" + plan = metadata.attn_metadata.kpool_write_plan + assert plan is not None, "DSA kpool target_verify requires kpool_write_plan" + num_draft_tokens = plan.num_draft_tokens + + query, key, gate_score_maybe = self._get_q_k_bf16( + q_lora, + x, + positions, + enable_dual_stream=enable_dual_stream, + forward_batch=forward_batch, + precompute_compress_gate=( + enable_dual_stream and self.compress_gate_stream is not None + ), + ) + + pool = get_token_to_kv_pool() + tail_k_buf, tail_score_buf = pool.get_compress_tail_buffers(layer_id) + + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_write_tail_and_maybe_compress, + ) + + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + + def _compress_write() -> None: + kpool_write_tail_and_maybe_compress( + pool=pool, + buf=buf, + key=key, + score=self._compute_gate_score_if_missing(x, gate_score_maybe), + tail_k=tail_k_buf, + tail_score=tail_score_buf, + ape=self.index_kpool_compress_ape, + req_pool_indices=plan.req, + write_start=plan.write_start, + tail_logical_start=plan.tail_logical_start, + write_loc=plan.write_loc, + out_cache_loc=forward_batch.out_cache_loc, + num_draft_tokens=num_draft_tokens, + round_scale=self.scale_fmt is not None, + effective_n_per_batch=plan.effective_n_per_batch, + ) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + if gate_score_maybe is not None: + assert self.compress_gate_stream is not None + self.alt_stream.wait_stream(self.compress_gate_stream) + if return_indices: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + weights = self._get_logits_head_gate(x, q_scale) + with torch.cuda.stream(self.alt_stream): + _compress_write() + current_stream.wait_stream(self.alt_stream) + else: + _compress_write() + if return_indices: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + weights = self._get_logits_head_gate(x, q_scale) + + if not return_indices: + return None + return self._get_topk_paged(forward_batch, layer_id, q_fp8, weights, metadata) + + def forward_cuda( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + if is_hip(): + from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant + elif not is_npu(): + from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant + + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) + + enable_dual_stream = ( + self.alt_stream is not None + and get_is_capture_mode() + and q_lora.shape[0] > 0 + and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + ) + + # Skip DSA if the attention backend chooses to skip this batch. + if metadata is None: + return None + + assert forward_batch.seq_lens_cpu is not None + mode = forward_batch.forward_mode + if mode.is_idle() or len(forward_batch.seq_lens_cpu) == 0: + return torch.full( + (x.shape[0], self.index_topk + self.index_kpool - 1), + -1, + dtype=torch.int, + device=x.device, + ) + + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + return self._forward_cuda_target_verify( + x=x, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + act_quant=act_quant, + metadata=metadata, + enable_dual_stream=enable_dual_stream, + return_indices=return_indices, + ) + + skip_logits_computation = False + if forward_batch.forward_mode.is_extend_without_speculative(): + if forward_batch.seq_lens_cpu is not None: + max_kv_len = forward_batch.seq_lens_cpu.max().item() + skip_logits_computation = max_kv_len <= self.index_topk + + if skip_logits_computation: + return self._forward_cuda_skip_logits( + x, + positions, + forward_batch, + layer_id, + act_quant, + metadata, + return_indices, + ) + + precompute_compress_gate = ( + self.index_kpool > 1 + and self.index_kpool_compress + and enable_dual_stream + and forward_batch.forward_mode.is_decode_or_idle() + and self.compress_gate_stream is not None + ) + query, key, gate_score = self._get_q_k_bf16( + q_lora, + x, + positions, + enable_dual_stream, + forward_batch=forward_batch, + precompute_compress_gate=precompute_compress_gate, + ) + + weights = None + kpool_extend_cache = None + if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + if gate_score is not None: + self.alt_stream.wait_stream(self.compress_gate_stream) + with torch.cuda.stream(self.alt_stream): + self._compress_write( + x=x, + key=key, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + metadata=metadata, + gate_score=gate_score, + ) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + weights = self._get_logits_head_gate(x, q_scale) + current_stream.wait_stream(self.alt_stream) + else: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + has_kpool_extend_plan = metadata.attn_metadata.kpool_extend_plan is not None + defer_kpool_cache_write = ( + forward_batch.forward_mode.is_extend_without_speculative() + and return_indices + and not has_kpool_extend_plan + ) + kpool_extend_cache = self._compress_write( + x=x, + key=key, + positions=positions, + forward_batch=forward_batch, + layer_id=layer_id, + metadata=metadata, + gate_score=gate_score, + return_compressed=( + forward_batch.forward_mode.is_extend_without_speculative() + and return_indices + ), + write_cache=not defer_kpool_cache_write, + ) + if ( + forward_batch.forward_mode.is_extend_without_speculative() + and not return_indices + ): + return None + + if weights is None: + weights = self._get_logits_head_gate(x, q_scale) + + if is_cuda(): + if ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + topk_result = self._get_topk_paged( + forward_batch, layer_id, q_fp8, weights, metadata + ) + else: + if has_kpool_extend_plan: + topk_result = self._get_topk_ragged_kpool_plan( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + ) + else: + topk_result = self._get_topk_ragged( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + kpool_extend_cache=kpool_extend_cache, + ) + else: + raise NotImplementedError("kpool indexer is only supported on CUDA") + return topk_result diff --git a/python/sglang/srt/layers/attention/dsa/kpool_plan.py b/python/sglang/srt/layers/attention/dsa/kpool_plan.py index ea2d3a3fb..43d0f9c01 100644 --- a/python/sglang/srt/layers/attention/dsa/kpool_plan.py +++ b/python/sglang/srt/layers/attention/dsa/kpool_plan.py @@ -731,7 +731,6 @@ def update_kpool_write_plan( forward_mode: ForwardMode, slots_per_page: int, effective_n_per_batch: Optional[torch.Tensor] = None, - include_deep_gemm_schedule: bool = True, ) -> None: if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_cuda(): return @@ -767,9 +766,7 @@ def update_kpool_write_plan( effective_n_per_batch.to(torch.int32) ) - # In-graph replay updates plan lengths too late for host schedule construction; - # the caller rebuilds the schedule from raw seq_lens out of graph. - if include_deep_gemm_schedule and plan.pool_schedule_metadata is not None: + if plan.pool_schedule_metadata is not None: new_schedule = _compute_pool_schedule_metadata( plan.pool_seqlens_per_q, slots_per_page=slots_per_page, @@ -778,25 +775,6 @@ def update_kpool_write_plan( plan.pool_schedule_metadata.copy_(new_schedule) -def refresh_kpool_pool_schedule_from( - metadata: DSAMetadata, - pool_seqlens_per_q: torch.Tensor, - *, - slots_per_page: int, -) -> None: - """Use an explicit source because the captured plan buffer remains stale - until replay.""" - plan = metadata.kpool_write_plan - if plan is None or plan.pool_schedule_metadata is None: - return - new_schedule = _compute_pool_schedule_metadata( - pool_seqlens_per_q, - slots_per_page=slots_per_page, - ) - if new_schedule is not None: - plan.pool_schedule_metadata.copy_(new_schedule) - - def init_kpool_write_plan( metadata: DSAMetadata, forward_batch: ForwardBatch, diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index b1ce79611..92755d103 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -71,8 +71,18 @@ if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch -def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int): - return original_seq_lens.clamp(max=dsa_index_topk) +def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int, index_kpool: int = 1): + if index_kpool <= 1: + return original_seq_lens.clamp(max=dsa_index_topk) + + # Clamp only complete pools; the unfinished tail must remain selectable + # outside the pooled top-k budget. + full_pool_tokens = ( + torch.div(original_seq_lens, index_kpool, rounding_mode="floor") * index_kpool + ) + selected_history_tokens = full_pool_tokens.clamp(max=dsa_index_topk) + tail_tokens = original_seq_lens - full_pool_tokens + return selected_history_tokens + tail_tokens def should_remap_pd_dsa_seed_to_local_slots() -> bool: diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index a7b4e0e31..80ea57740 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -31,18 +31,34 @@ from sglang.kernels.ops.attention.dsa.dequant_k_cache import ( ) from sglang.kernels.ops.attention.dsa.quant_k_cache import quantize_k_cache from sglang.kernels.ops.attention.dsa.transform_index import ( + prepare_trtllm_nope_sparse_metadata, transform_index_page_table_decode, transform_index_page_table_prefill, ) +from sglang.kernels.ops.attention.dsa_metadata import ( + fused_dsa_decode_metadata, + fused_dsa_draft_extend_metadata, + fused_dsa_target_verify_metadata, +) from sglang.kernels.ops.attention.utils import ( concat_mla_absorb_q_general, mla_quantize_and_rope_for_fp8, + mla_quantize_for_fp8_no_rope, q8kv8_topk_length_from_indices, seqlens_expand_triton, ) from sglang.kernels.ops.kvcache.cache_ops import concat_and_cast_q_fp8_pad +from sglang.srt.configs.model_config import ( + get_dsa_index_kpool, + get_dsa_index_topk, + is_deepseek_dsa, +) from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsa.dsa_backend_kpool import ( + DeepseekSparseAttnBackendKPoolMixin, + _KPoolForwardInputs, +) from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( DeepseekSparseAttnBackendMTPPrecomputeMixin, PrecomputedMetadata, @@ -53,6 +69,10 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import ( DSATopKBackend, TopkTransformMethod, ) +from sglang.srt.layers.attention.dsa.kpool_plan import ( + KPoolExtendPlan, + KPoolWritePlan, +) from sglang.srt.layers.attention.dsa.utils import ( can_dsa_prefill_cp_round_robin_split, compute_dsa_seqlens, @@ -75,6 +95,7 @@ from sglang.srt.layers.utils.cp_utils import ( cp_split_and_rebuild_position, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_buffer, get_exec, get_parallel, get_spec from sglang.srt.utils import ( get_bool_env_var, is_cuda, @@ -83,6 +104,8 @@ from sglang.srt.utils import ( print_warning_once, ) +logger = logging.getLogger(__name__) + # Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton # per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @ # TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the @@ -253,6 +276,13 @@ class DSAMetadata: # batch index for each token. token_to_batch_idx: Optional[torch.Tensor] = None + pooled_index_kpool: int = 1 + pooled_cache_seqlens_int32: Optional[torch.Tensor] = None + pooled_real_page_table: Optional[torch.Tensor] = None + pooled_paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + kpool_extend_plan: Optional[KPoolExtendPlan] = None + kpool_write_plan: Optional[KPoolWritePlan] = None + @torch.compile def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: @@ -287,7 +317,9 @@ _DSA_IMPL_T: TypeAlias = Literal[ class DeepseekSparseAttnBackend( - DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend + DeepseekSparseAttnBackendKPoolMixin, + DeepseekSparseAttnBackendMTPPrecomputeMixin, + AttentionBackend, ): # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch # can never carry more seqs than the pool. @@ -314,12 +346,15 @@ class DeepseekSparseAttnBackend( self.num_splits = ( 1 if get_exec().deterministic.enable_deterministic_inference else 0 ) - self.use_dsa = is_deepseek_dsa(model_runner.model_config.hf_config) + hf_config = model_runner.model_config.hf_config + self.use_dsa = is_deepseek_dsa(hf_config) assert self.use_dsa, "DSA backend only supports DeepSeek DSA" self.dsa_kv_cache_store_fp8 = ( model_runner.token_to_kv_pool.dsa_kv_cache_store_fp8 ) - self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config) + self.dsa_index_topk = get_dsa_index_topk(hf_config) + self.dsa_index_kpool = get_dsa_index_kpool(hf_config) + self.needs_cpu_seq_lens = self.dsa_index_kpool > 1 self.max_context_len = model_runner.model_config.context_len self.num_q_heads = ( model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size @@ -350,6 +385,23 @@ class DeepseekSparseAttnBackend( self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto" self._sink_pad_cache: dict[tuple[int, int], torch.Tensor] = {} + # Hoisted per-call imports of set_dsa_prefill_impl. Module-scope + # imports would cycle through model_executor (which imports the + # attention backends); backend init runs after those modules are + # fully imported, so binding the function refs here is cycle-safe. + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, + ) + from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + is_in_tc_piecewise_cuda_graph, + ) + from sglang.srt.utils import get_device_sm, is_blackwell + + self._is_in_breakable_cuda_graph = is_in_breakable_cuda_graph + self._is_in_tc_piecewise_cuda_graph = is_in_tc_piecewise_cuda_graph + self._get_device_sm = get_device_sm + self._is_blackwell = is_blackwell + self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) if _is_hip: @@ -654,6 +706,21 @@ class DeepseekSparseAttnBackend( "num_kv_splits": self.aiter_dsa_max_split_per_batch, } + def _pad_trtllm_sparse_page_table( + self, page_table_1: torch.Tensor + ) -> Tuple[torch.Tensor, int]: + sparse_mla_top_k = page_table_1.shape[1] + padded_top_k = ((sparse_mla_top_k + 3) // 4) * 4 + if padded_top_k == sparse_mla_top_k: + return page_table_1, sparse_mla_top_k + padding = torch.full( + (page_table_1.shape[0], padded_top_k - sparse_mla_top_k), + -1, + dtype=page_table_1.dtype, + device=page_table_1.device, + ) + return torch.cat([page_table_1, padding], dim=1), padded_top_k + def _build_paged_mqa_schedule_2d_ctx_lens( self, forward_mode: ForwardMode, @@ -737,6 +804,9 @@ class DeepseekSparseAttnBackend( ) return self._arange_buf[:length] + def update_verify_buffers_to_fill_after_draft(self, *args, **kwargs): + return None + def _graph_page_table_width(self, metadata: DSAMetadata) -> int: """Column count to scan req_to_token during graph replay. Reads the wide page_table_1 width when present, else req_to_token's width (the wide table @@ -826,6 +896,16 @@ class DeepseekSparseAttnBackend( # seq_len_cpu of selected sequences indexer_seq_lens_cpu = forward_batch.seq_lens_cpu indexer_seq_lens = forward_batch.seq_lens + use_kpool = self.dsa_index_kpool > 1 + if use_kpool: + assert ( + self.real_page_size == 64 + and self.real_page_size % self.dsa_index_kpool == 0 + ), ( + f"kpool path requires page_size == 64 and page_size % pool_size == 0; " + f"got page_size={self.real_page_size}, pool_size={self.dsa_index_kpool}." + ) + kpool_inputs = _KPoolForwardInputs() if forward_batch.forward_mode.is_decode_or_idle(): extend_seq_lens_cpu = [1] * batch_size @@ -925,6 +1005,11 @@ class DeepseekSparseAttnBackend( ) ] ) + if use_kpool: + kpool_inputs.full_real_page_table = self._transform_table_1_to_real( + page_table + ) + kpool_inputs.full_seqlens_expanded = seqlens_expanded if can_dsa_prefill_cp_round_robin_split(forward_batch): if is_cp_v2_active(forward_batch): @@ -1011,6 +1096,7 @@ class DeepseekSparseAttnBackend( dsa_cache_seqlens_int32 = compute_dsa_seqlens( original_seq_lens=seqlens_expanded, dsa_index_topk=self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) dsa_cache_seqlens_int32 = pad_dsa_cache_seqlens( forward_batch, dsa_cache_seqlens_int32 @@ -1072,6 +1158,12 @@ class DeepseekSparseAttnBackend( token_to_batch_idx=token_to_batch_idx, topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), ) + metadata = self._init_kpool_metadata( + metadata, + forward_batch, + topk_transform_method=topk_transform_method, + kpool_inputs=kpool_inputs, + ) self.forward_metadata = metadata def _cal_indexer_k_start_end( @@ -1183,6 +1275,8 @@ class DeepseekSparseAttnBackend( and self.real_page_size > 1 and self.hisparse_coordinator is None and not self.speculative_num_draft_tokens + # kpool's PAGED fused-topk mapping still reads page_table_1. + and self.dsa_index_kpool <= 1 and self.use_fused_topk and self.dsa_topk_backend.should_use_topk_v2() and self.dsa_index_topk is not None @@ -1275,7 +1369,9 @@ class DeepseekSparseAttnBackend( # NOTE(dark): this is always arange, since we are decoding cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1] dsa_cache_seqlens_int32 = compute_dsa_seqlens( - cache_seqlens_int32, dsa_index_topk=self.dsa_index_topk + cache_seqlens_int32, + dsa_index_topk=self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) seqlens_expanded = cache_seqlens_int32 @@ -1338,7 +1434,9 @@ class DeepseekSparseAttnBackend( ] ) dsa_cache_seqlens_int32 = compute_dsa_seqlens( - seqlens_expanded, dsa_index_topk=self.dsa_index_topk + seqlens_expanded, + dsa_index_topk=self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) dsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens @@ -1400,6 +1498,7 @@ class DeepseekSparseAttnBackend( dsa_extend_seq_lens_list=dsa_extend_seq_lens_list, topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), ) + metadata = self._init_kpool_metadata_capture(metadata, bs, forward_mode) self.decode_cuda_graph_metadata[bs] = metadata self.forward_metadata = metadata @@ -1434,24 +1533,22 @@ class DeepseekSparseAttnBackend( ) return + metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs] + self.set_dsa_prefill_impl(forward_batch=None) seq_lens = seq_lens[:bs] req_pool_indices = req_pool_indices[:bs] # Normal Decode - metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs] used_fused_metadata_generation = False target_verify_ctx_lens_written = False + if forward_mode.is_decode_or_idle(): # Normal Decode max_len = self._graph_page_table_width(metadata) - if is_cuda() and not _is_hip: - from sglang.kernels.ops.attention.dsa_metadata import ( - fused_dsa_decode_metadata, - ) - + if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1: fused_dsa_decode_metadata( seq_lens=seq_lens, req_pool_indices=req_pool_indices, @@ -1482,18 +1579,16 @@ class DeepseekSparseAttnBackend( page_indices = self.req_to_token[req_pool_indices, :max_len] metadata.page_table_1[:, :max_len].copy_(page_indices) dsa_cache_seqlens = compute_dsa_seqlens( - cache_seqlens, dsa_index_topk=self.dsa_index_topk + cache_seqlens, + dsa_index_topk=self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) seqlens_expanded = cache_seqlens elif forward_mode.is_target_verify(): max_seqlen_k = self._graph_page_table_width(metadata) - if is_cuda() and not _is_hip: - from sglang.kernels.ops.attention.dsa_metadata import ( - fused_dsa_target_verify_metadata, - ) - + if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1: paged_mqa_ctx_lens_2d = None if ( self.speculative_num_draft_tokens >= 2 @@ -1565,7 +1660,9 @@ class DeepseekSparseAttnBackend( ) metadata.dsa_seqlens_expanded.copy_(seqlens_expanded) dsa_cache_seqlens = compute_dsa_seqlens( - seqlens_expanded, self.dsa_index_topk + seqlens_expanded, + self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) elif forward_mode.is_draft_extend_v2(): @@ -1587,11 +1684,7 @@ class DeepseekSparseAttnBackend( device=self.device, ) - if is_cuda() and not _is_hip: - from sglang.kernels.ops.attention.dsa_metadata import ( - fused_dsa_draft_extend_metadata, - ) - + if is_cuda() and not _is_hip and self.dsa_index_kpool <= 1: fused_dsa_draft_extend_metadata( seq_lens=seq_lens, extend_seq_lens=extend_seq_lens, @@ -1642,7 +1735,9 @@ class DeepseekSparseAttnBackend( seqlens_expanded ) dsa_cache_seqlens = compute_dsa_seqlens( - seqlens_expanded, self.dsa_index_topk + seqlens_expanded, + self.dsa_index_topk, + index_kpool=self.dsa_index_kpool, ) metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) @@ -1696,6 +1791,19 @@ class DeepseekSparseAttnBackend( else: assert metadata.real_page_table is metadata.page_table_1 + effective_n_per_batch = None + if forward_mode.is_draft_extend_v2() and spec_info is not None: + effective_n_per_batch = getattr(spec_info, "num_accept_tokens", None) + if effective_n_per_batch is not None: + effective_n_per_batch = effective_n_per_batch[:bs] + self._update_kpool_metadata_replay( + metadata, + seq_lens, + req_pool_indices, + forward_mode, + effective_n_per_batch=effective_n_per_batch, + ) + if self.dsa_decode_impl == "flashmla_kv": flashmla_metadata = metadata.flashmla_metadata.slice( slice(0, seqlens_expanded_size + 1) @@ -1848,10 +1956,8 @@ class DeepseekSparseAttnBackend( flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) flashmla_metadata.copy_(precomputed.flashmla_metadata) - # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of - # this replay (the captured graph holds stale data otherwise, which can - # deadlock the kernel when the runtime work decomposition diverges from - # the captured one). + # Refresh the schedule because stale shape decomposition can deadlock + # DeepGEMM paged MQA. if is_cuda(): if forward_mode.is_decode_or_idle(): seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) @@ -1869,6 +1975,10 @@ class DeepseekSparseAttnBackend( else: metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) + self._update_kpool_metadata_from_precomputed( + metadata, precomputed, forward_mode + ) + self.forward_metadata = metadata def forward_extend( @@ -1906,6 +2016,17 @@ class DeepseekSparseAttnBackend( f"Learnable attention sinks require flashmla_sparse, got {dsa_impl}" ) + phase = ( + "decode" + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ) + else "prefill" + ) + dsa_impl = self._resolve_kpool_tail_backend(topk_indices, dsa_impl) + self._check_kpool_tail_backend(topk_indices, dsa_impl, phase) + if dsa_impl == "trtllm" and not self.use_mha: return self._forward_trtllm( q, @@ -1962,7 +2083,9 @@ class DeepseekSparseAttnBackend( if q_rope is not None: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + q_nope.shape[0], + layer.tp_q_head_num, + layer.head_dim - layer.v_head_dim, ) else: q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) @@ -2209,13 +2332,15 @@ class DeepseekSparseAttnBackend( metadata = self.forward_metadata assert causal, "DSA is causal only" - if attn_sink is not None and self.dsa_decode_impl != "flashmla_sparse": + dsa_impl = self._resolve_kpool_tail_backend(topk_indices, self.dsa_decode_impl) + self._check_kpool_tail_backend(topk_indices, dsa_impl, "decode") + + if attn_sink is not None and dsa_impl != "flashmla_sparse": raise RuntimeError( - "Learnable attention sinks require flashmla_sparse, got " - f"{self.dsa_decode_impl}" + f"Learnable attention sinks require flashmla_sparse, got {dsa_impl}" ) - if self.dsa_decode_impl == "trtllm": + if dsa_impl == "trtllm": return self._forward_trtllm( q, k, @@ -2252,7 +2377,9 @@ class DeepseekSparseAttnBackend( if q_rope is not None: q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) q_rope = q_rope.view( - -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + q_nope.shape[0], + layer.tp_q_head_num, + layer.head_dim - layer.v_head_dim, ) # Caller passed split q_nope / q_rope; we'll need to concat below if # the chosen impl wants q_all. @@ -2285,7 +2412,7 @@ class DeepseekSparseAttnBackend( page_size=1, ) - if self.dsa_decode_impl == "flashmla_sparse": + if dsa_impl == "flashmla_sparse": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashmla_sparse( @@ -2297,7 +2424,7 @@ class DeepseekSparseAttnBackend( topk_length=metadata.dsa_cache_seqlens_int32, attn_sink=attn_sink, ) - elif self.dsa_decode_impl == "flashinfer_sparse_mla": + elif dsa_impl == "flashinfer_sparse_mla": if q_all is None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashinfer_sparse_mla( @@ -2308,7 +2435,7 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), ) - elif self.dsa_decode_impl == "flashmla_kv": + elif dsa_impl == "flashmla_kv": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_flashmla_kv( @@ -2321,7 +2448,7 @@ class DeepseekSparseAttnBackend( metadata=metadata, page_table_1=page_table_1, ) - elif self.dsa_decode_impl == "tilelang": + elif dsa_impl == "tilelang": # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all # has already been set to a zero-copy view of q in the else branch # above and we can reuse it directly. The `not _is_hip` clause keeps @@ -2335,7 +2462,7 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) - elif self.dsa_decode_impl == "fa3": + elif dsa_impl == "fa3": return self._forward_fa3( q_rope=q_rope, kv_cache=kv_cache, @@ -2350,7 +2477,7 @@ class DeepseekSparseAttnBackend( logit_cap=layer.logit_cap, page_size=1, ) - elif self.dsa_decode_impl == "aiter": + elif dsa_impl == "aiter": if q_all is None or not _is_hip: q_all = torch.cat([q_nope, q_rope], dim=-1) return self._forward_aiter( @@ -2363,7 +2490,7 @@ class DeepseekSparseAttnBackend( ) else: - assert False, f"Unsupported {self.dsa_decode_impl = }" + assert False, f"Unsupported {dsa_impl = }" def _forward_fa3( self, @@ -2383,13 +2510,21 @@ class DeepseekSparseAttnBackend( k_rope_cache = kv_cache[:, :, v_head_dim:] c_kv_cache = kv_cache[:, :, :v_head_dim] qk_rope_dim = k_rope_cache.shape[-1] - k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim) - c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim) + num_blocks = kv_cache.shape[0] // page_size + only_qv = qk_rope_dim == 0 + if only_qv: + k_rope_cache = None + else: + k_rope_cache = k_rope_cache.view(num_blocks, page_size, 1, qk_rope_dim) + c_kv_cache = c_kv_cache.view(num_blocks, page_size, 1, v_head_dim) + if self.dsa_index_kpool > 1: + page_table = page_table.clamp(min=0) o = flash_attn_with_kvcache( - q=q_rope, + q=None if only_qv else q_rope, k_cache=k_rope_cache, v_cache=c_kv_cache, qv=q_nope, + only_qv=only_qv, page_table=page_table, cache_seqlens=cache_seqlens, cu_seqlens_q=cu_seqlens_q, @@ -2955,6 +3090,19 @@ class DeepseekSparseAttnBackend( ) -> torch.Tensor: from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + # KPool appends up to index_kpool - 1 live tail tokens to the fixed + # index_topk columns. TileLang processes indices in 64-column blocks, + # so mask-pad the tail-extended table to the next complete block. + padding = (-page_table_1.shape[-1]) % 64 + if padding: + page_table_1 = torch.cat( + ( + page_table_1, + page_table_1.new_full((*page_table_1.shape[:-1], padding), -1), + ), + dim=-1, + ) + return tilelang_sparse_fwd( q=q_all, kv=kv_cache, @@ -3154,45 +3302,54 @@ class DeepseekSparseAttnBackend( metadata = self.forward_metadata - merge_query = q_rope is not None + # The BF16 no-RoPE path passes a zero-width q_rope tensor. + merge_query = q_rope is not None and self.qk_rope_head_dim > 0 if self.kv_cache_dtype == torch.float8_e4m3fn: # For FP8 path, we quantize the query and rope parts and merge them into a single tensor # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend assert q_rope is not None, "For FP8 path q_rope should not be None." assert k_rope is not None, "For FP8 path k_rope should not be None." - assert cos_sin_cache is not None, ( - "For FP8 path cos_sin_cache should not be None." - ) + if cos_sin_cache is None: + q, k, k_rope = mla_quantize_for_fp8_no_rope( + q, + q_rope, + k.squeeze(1), + k_rope.squeeze(1), + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + else: + rope_positions = forward_batch.positions + if dsa_use_prefill_cp(forward_batch): + if is_cp_v2_active(forward_batch): + rope_positions = get_cp_strategy().shard_position_ids( + rope_positions, forward_batch + ) + else: + rope_positions = cp_split_and_rebuild_position( + forward_batch, rope_positions + ) - rope_positions = forward_batch.positions - if dsa_use_prefill_cp(forward_batch): - if is_cp_v2_active(forward_batch): - rope_positions = get_cp_strategy().shard_position_ids( - rope_positions, forward_batch - ) - else: - rope_positions = cp_split_and_rebuild_position( - forward_batch, rope_positions - ) - - q, k, k_rope = mla_quantize_and_rope_for_fp8( - q, - q_rope, - k.squeeze(1), - k_rope.squeeze(1), - rope_positions, - cos_sin_cache, - is_neox, - self.kv_lora_rank, - self.qk_rope_head_dim, - ) - if save_kv_cache and dsa_use_prefill_cp(forward_batch): - if is_cp_v2_active(forward_batch): - k, k_rope = get_cp_strategy().all_gather_dsa_trtllm_fp8_kv( - forward_batch, k, k_rope - ) - else: - k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope) + q, k, k_rope = mla_quantize_and_rope_for_fp8( + q, + q_rope, + k.squeeze(1), + k_rope.squeeze(1), + rope_positions, + cos_sin_cache, + is_neox, + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + if save_kv_cache and dsa_use_prefill_cp(forward_batch): + if is_cp_v2_active(forward_batch): + k, k_rope = get_cp_strategy().all_gather_dsa_trtllm_fp8_kv( + forward_batch, k, k_rope + ) + else: + k, k_rope = _all_gather_dsa_trtllm_fp8_kv( + forward_batch, k, k_rope + ) merge_query = False # Save KV cache if requested @@ -3244,6 +3401,12 @@ class DeepseekSparseAttnBackend( topk_indices=topk_indices, page_size=1, ) + page_table_1, sparse_mla_top_k = self._pad_trtllm_sparse_page_table( + page_table_1 + ) + sparse_mla_top_k_lens = None + if self.qk_rope_head_dim == 0: + sparse_mla_top_k_lens = prepare_trtllm_nope_sparse_metadata(page_table_1) q_scale = 1.0 k_scale = ( @@ -3289,10 +3452,11 @@ class DeepseekSparseAttnBackend( block_tables=block_tables, seq_lens=seq_lens, max_seq_len=metadata.max_seq_len_k, - sparse_mla_top_k=self.dsa_index_topk, + sparse_mla_top_k=sparse_mla_top_k, bmm1_scale=bmm1_scale, backend="trtllm-gen", skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + sparse_mla_top_k_lens=sparse_mla_top_k_lens, multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) @@ -3327,13 +3491,11 @@ class DeepseekSparseAttnBackend( """ Decide all attention prefill dispatch strategies for this batch. """ - from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( - is_in_breakable_cuda_graph, - ) - from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - is_in_tc_piecewise_cuda_graph, - ) - from sglang.srt.utils import get_device_sm, is_blackwell + # Hoisted in __init__ (import cost is per-call otherwise). + is_in_breakable_cuda_graph = self._is_in_breakable_cuda_graph + is_in_tc_piecewise_cuda_graph = self._is_in_tc_piecewise_cuda_graph + get_device_sm = self._get_device_sm + is_blackwell = self._is_blackwell # Decide MHA vs MLA if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): @@ -3613,6 +3775,16 @@ class DeepseekSparseAttnMultiStepBackend: precomputed.seqlens_expanded_size, ) + for backend, metadata in zip( + self.attn_backends[:3], + (metadata0, metadata1, metadata2), + strict=True, + ): + backend._update_kpool_metadata_from_precomputed( + metadata, precomputed, ForwardMode.DECODE + ) + backend.forward_metadata = metadata + # Copy remaining backends one by one (if > 3 backends) for i in range(3, self.speculative_num_steps - 1): self.attn_backends[ @@ -3641,7 +3813,7 @@ class DeepseekSparseAttnMultiStepBackend: forward_mode=ForwardMode.DECODE, ) else: - # Less than 3 backends: copy to each backend individually + # Copy to each backend and refresh its derived metadata independently. for i in range(self.speculative_num_steps - 1): self.attn_backends[ i diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 772ccd142..fe25a46a4 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -96,6 +96,9 @@ class MambaAttnBackendBase(AttentionBackend): self.cached_cuda_graph_decode_query_start_loc: torch.Tensor = None self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None self.conv_states_shape: tuple[int, int] = None + # Constant (== 1) for mamba-like backends; hoisted so the replay path + # skips the per-cycle method dispatch. + self._graph_seq_len_fill_value = self.get_cuda_graph_seq_len_fill_value() @property def mamba_chunk_size(self) -> int: @@ -210,9 +213,13 @@ class MambaAttnBackendBase(AttentionBackend): new_vals[inv] = next_for_valid.to(write_pos_buf.dtype) write_pos_buf[uniq_slots] = new_vals elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True): - if forward_batch.forward_mode.is_draft_extend_v2(): - # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; - # skip mamba metadata. + has_extend_meta = ( + forward_batch.extend_start_loc is not None + and forward_batch.extend_seq_lens is not None + ) + if forward_batch.forward_mode.is_draft_extend_v2() and not has_extend_meta: + # Draft-extend-v2 may omit linear metadata when the draft runs only + # full-attention layers. query_start_loc = None elif forward_batch.forward_mode.is_target_verify(): ragged_layout = forward_batch.spec_info.ragged_verify_layout @@ -623,8 +630,9 @@ class MambaAttnBackendBase(AttentionBackend): num_padding = 0 else: num_padding = torch.count_nonzero( - seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() + seq_lens_cpu == self._graph_seq_len_fill_value ) + num_padding = int(num_padding) if self._fused_state_indices_ok and self.replayssm_write_pos_list is None: # Single-launch fast path: mapping gather + padding sentinel + store # into the static buffer, plus zeroing padded req_pool_indices rows — @@ -731,6 +739,7 @@ class MambaAttnBackendBase(AttentionBackend): ) new_vals[inv] = next_for_valid.to(write_pos_buf.dtype) write_pos_buf[uniq_slots] = new_vals + is_target_verify = forward_mode.is_target_verify() if forward_mode.is_decode_or_idle(): if num_padding == 0: self.query_start_loc_list[bs - 1].copy_( @@ -769,8 +778,9 @@ class MambaAttnBackendBase(AttentionBackend): ) else: raise ValueError(f"Invalid forward mode: {forward_mode=}") + qsl_buf = self.query_start_loc_list[bs - 1] - if forward_mode.is_target_verify() and self.topk > 1: + if is_target_verify and self.topk > 1: if ( spec_info is not None and getattr(spec_info, "retrieve_next_token", None) is not None @@ -783,7 +793,7 @@ class MambaAttnBackendBase(AttentionBackend): spec_info.retrieve_next_sibling ) return ForwardMetadata( - query_start_loc=self.query_start_loc_list[bs - 1], + query_start_loc=qsl_buf, mamba_cache_indices=self.state_indices_list[bs - 1], mamba_track_indices=track_buf, retrieve_next_token=self.retrieve_next_token_list[bs - 1], @@ -794,7 +804,7 @@ class MambaAttnBackendBase(AttentionBackend): ) else: return ForwardMetadata( - query_start_loc=self.query_start_loc_list[bs - 1], + query_start_loc=qsl_buf, mamba_cache_indices=self.state_indices_list[bs - 1], mamba_track_indices=track_buf, replayssm_write_pos=replayssm_write_pos, @@ -1067,9 +1077,15 @@ class HybridLinearAttnBackend(AttentionBackend): and self.linear_attn_backend.supports_ragged_verify_graph ) + @property + def use_mha(self) -> bool: + return getattr(self.full_attn_backend, "use_mha", False) + @property def kv_cache_dtype(self): - return self.full_attn_backend.kv_cache_dtype + # Expose the full-attention backend's cache dtype because fused DSA/NSA RoPE + # reads it from this wrapper. + return getattr(self.full_attn_backend, "kv_cache_dtype", None) def _is_full_attn( self, layer: Optional[RadixAttention], layer_id: Optional[int] = None @@ -1136,6 +1152,9 @@ class HybridLinearAttnBackend(AttentionBackend): if init is not None: init(forward_batch, disable_flashinfer_ragged) + def get_indexer_metadata(self, layer_id, forward_batch): + return self.full_attn_backend.get_indexer_metadata(layer_id, forward_batch) + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): for attn_backend in self.attn_backend_list: attn_backend.init_cuda_graph_state(max_bs, max_num_tokens) diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 5b877c9c1..44acc6600 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -151,7 +151,7 @@ class KDAKernelDispatcher: ) cutedsl_kernel = CuteDSLKDAKernel() - if getattr(cutedsl_kernel, "supports_prefill", False): + if cutedsl_kernel.supports_prefill: # SM100 chunk prefill pipeline. self.extend_kernel = cutedsl_kernel else: @@ -255,14 +255,17 @@ class KDAKernelDispatcher: ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: - kernel = self.decode_kernel - if kwargs.get("lower_bound") is not None and not getattr( - kernel, "supports_safe_gate", True + if lower_bound is not None and not isinstance( + self.decode_kernel, TritonKDAKernel ): - kernel = self.triton_kernel - return kernel.decode( + raise NotImplementedError( + f"lower_bound (safe gate) is only supported by TritonKDAKernel; " + f"got {self.decode_kernel.__class__.__name__}." + ) + return self.decode_kernel.decode( q, k, v, @@ -273,6 +276,7 @@ class KDAKernelDispatcher: ssm_states=ssm_states, cache_indices=cache_indices, query_start_loc=query_start_loc, + lower_bound=lower_bound, **kwargs, ) @@ -292,13 +296,20 @@ class KDAKernelDispatcher: intermediate_states_buffer: torch.Tensor, intermediate_state_indices: torch.Tensor, cache_steps: int, - retrieve_parent_token: torch.Tensor, + retrieve_parent_token: Optional[torch.Tensor], lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: """MTP / speculative-decode verify, routed to ``self.verify_kernel`` (FlashInfer decode -> recurrent_kda; Triton / CuTe DSL decode -> the Triton fused KDA verify).""" + if lower_bound is not None and not isinstance( + self.verify_kernel, TritonKDAKernel + ): + raise NotImplementedError( + "lower_bound (safe gate) target verify is only supported by " + f"TritonKDAKernel; got {self.verify_kernel.__class__.__name__}." + ) return self.verify_kernel.target_verify( A_log=A_log, dt_bias=dt_bias, @@ -380,10 +391,8 @@ class KDAAttnBackend(MambaAttnBackendBase): # to its dense layout, so ragged verify graphs are supported. supports_ragged_verify_graph: bool = True - # Read by decide_needs_cpu_seq_lens. Decode/verify metadata is GPU-only - # (graph replay already passes seq_lens_cpu=None), extend reads - # extend_seq_lens_cpu from schedule, mamba track indices rebuild from req - # objects, and the replayssm seq_lens_cpu force-flush is GDN-only. + # KDA gets graph padding explicitly and never uses ReplaySSM's host-seqlen + # force-flush path. needs_cpu_seq_lens: bool = False def __init__(self, model_runner: ModelRunner): @@ -721,9 +730,11 @@ class KDAAttnBackend(MambaAttnBackendBase): conv_state_indices=cache_indices, ) - # The packed kernel assumes one token per request. Assert the dispatch - # invariant before taking the fused path. - if self.kernel_dispatcher.supports_packed_decode: + # The packed kernel assumes one token per request. + if ( + self.kernel_dispatcher.supports_packed_decode + and getattr(layer, "lower_bound", None) is None + ): assert qkv.shape[0] == cache_indices.shape[0], ( "KDA packed decode requires one token per sequence (T=1): " f"got {qkv.shape[0]} tokens for {cache_indices.shape[0]} requests." @@ -805,6 +816,13 @@ class KDAAttnBackend(MambaAttnBackendBase): ) has_initial_state = forward_batch.extend_prefix_lens > 0 + physical_num_tokens = mixed_qkv.shape[0] + logical_num_tokens = int(query_start_loc[-1]) + if logical_num_tokens < physical_num_tokens: + mixed_qkv = mixed_qkv[:logical_num_tokens] + a = a[:, :logical_num_tokens] + b = b[:, :logical_num_tokens] + if self.forward_metadata.has_mamba_track_mask: # Snapshot the conv sliding window at the last track-aligned chunk # boundary into the ping-pong track slots (the prefix-cache restore @@ -815,55 +833,29 @@ class KDAAttnBackend(MambaAttnBackendBase): self.forward_metadata.conv_states_mask_indices ] = mixed_qkv[self.forward_metadata.track_conv_indices] - splits = [layer.q_dim, layer.k_dim, layer.v_dim] - q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0) - q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split( - splits, dim=0 - ) - q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2) - if layer.bias is not None: - q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0) - else: - q_bias, k_bias, v_bias = None, None, None - - q = causal_conv1d_fn( - q, - q_conv_weight, - q_bias, + # Depthwise conv is channel-independent, so one packed call over the + # full qkv width matches the decode path and saves two kernel launches. + qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer.conv_weights, + layer.bias, activation="silu", - conv_states=q_conv_state, - has_initial_state=has_initial_state, - cache_indices=cache_indices, - query_start_loc=query_start_loc, - seq_lens_cpu=forward_batch.extend_seq_lens_cpu, - ).transpose(0, 1) - k = causal_conv1d_fn( - k, - k_conv_weight, - k_bias, - activation="silu", - conv_states=k_conv_state, - has_initial_state=has_initial_state, - cache_indices=cache_indices, - query_start_loc=query_start_loc, - seq_lens_cpu=forward_batch.extend_seq_lens_cpu, - ).transpose(0, 1) - v = causal_conv1d_fn( - v, - v_conv_weight, - v_bias, - activation="silu", - conv_states=v_conv_state, + conv_states=conv_states, has_initial_state=has_initial_state, cache_indices=cache_indices, query_start_loc=query_start_loc, seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ).transpose(0, 1) + q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) # n (h d) -> 1 n h d k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) # n (h d) -> 1 n h d v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) # n (h d) -> 1 n h d + gate_was_flat = a.ndim == 3 + if gate_was_flat: + a = a.unflatten(-1, (-1, layer.head_k_dim)) + track_ssm = self.forward_metadata.has_mamba_track_mask core_attn_out = self.kernel_dispatcher.extend( q=q, @@ -877,6 +869,7 @@ class KDAAttnBackend(MambaAttnBackendBase): A_log=layer.A_log, dt_bias=layer.dt_bias, lower_bound=layer.lower_bound, + beta_is_raw=gate_was_flat, extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, # draft_extend_v2 must stay rollback-able, so kernels that commit state # in place (e.g. FlashKDA) must not run for it. @@ -898,6 +891,13 @@ class KDAAttnBackend(MambaAttnBackendBase): forward_batch, h, ssm_states, self.forward_metadata ) + if logical_num_tokens < physical_num_tokens: + pad = core_attn_out.new_zeros( + (1, physical_num_tokens - logical_num_tokens) + + tuple(core_attn_out.shape[2:]) + ) + core_attn_out = torch.cat((core_attn_out, pad), dim=1) + if ( self.accept_lens_pool is not None and not forward_batch.forward_mode.is_draft_extend_v2() diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 01258f90f..89a6d9b79 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -683,6 +683,10 @@ class TritonAttnBackend(AttentionBackend): window_num_kv_splits=None, window_kv_offsets=None, swa_attn_logits=self.cuda_graph_swa_attn_logits, + lean_Mp=self.cuda_graph_lean_Mp, + lean_Lp=self.cuda_graph_lean_Lp, + lean_Op=self.cuda_graph_lean_Op, + lean_locks=self.cuda_graph_lean_locks, ) return diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 459394b4f..21ed8c17b 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -250,12 +250,17 @@ class AttentionInputs: hidden_states: torch.Tensor, forward_batch: ForwardBatch, qkv_latent_func: Callable, + *, + is_pre_gathered: bool = False, ): self.hidden_states_local = hidden_states self.forward_batch = forward_batch self.qkv_latent_func = qkv_latent_func self.hidden_states_ = None self.qkv_latent_ = None + # When True, hidden_states_local is already attn_tp-gathered upstream + # (e.g. by MHC's prepare_attn for DSA). fetch_* must NOT gather again. + self.is_pre_gathered = is_pre_gathered def tp_all_gather_hidden_states(self, hidden_states, forward_batch): total_tokens = forward_batch.input_ids.shape[0] @@ -270,7 +275,7 @@ class AttentionInputs: self.qkv_latent_ = self.qkv_latent_func( self.hidden_states_local, self.forward_batch ) - if get_attn_tp_context().input_scattered: + if get_attn_tp_context().input_scattered and not self.is_pre_gathered: self.qkv_latent_ = self.tp_all_gather_hidden_states( self.qkv_latent_, self.forward_batch ) @@ -280,7 +285,7 @@ class AttentionInputs: if self.hidden_states_ is not None: return self.hidden_states_ self.hidden_states_ = self.hidden_states_local - if get_attn_tp_context().input_scattered: + if get_attn_tp_context().input_scattered and not self.is_pre_gathered: self.hidden_states_ = self.tp_all_gather_hidden_states( self.hidden_states_, self.forward_batch ) @@ -292,13 +297,15 @@ class AttnTpContext: self.allow_input_scattered = False self.is_dsa = False - def init_context(self, q_lora_rank, is_dsa): + def init_context(self, q_lora_rank, is_dsa, is_mhc=False): + # Only MHC pre-gathers hidden states before DSA attention, so non-MHC DSA + # cannot use scattered inputs. self.is_dsa = is_dsa self.allow_input_scattered = ( get_parallel().enable_attn_tp_input_scattered and (_is_cuda or _is_npu) and q_lora_rank is not None - and not is_dsa + and (is_mhc or not is_dsa) and get_parallel().tp_size > 1 and not is_dp_attention_enabled() and get_moe_a2a_backend().is_none() @@ -330,6 +337,11 @@ class AttnTpContext: def set_attn_inputs(self, attn_inputs: AttentionInputs): get_forward().set("attn_inputs", attn_inputs) + def set_hidden_states_local(self, hidden_states: torch.Tensor) -> None: + attn_inputs = get_forward().attn_inputs + if attn_inputs is not None: + attn_inputs.hidden_states_local = hidden_states + def fetch_qkv_latent(self): attn_inputs = get_forward().attn_inputs assert attn_inputs is not None @@ -486,6 +498,26 @@ def enable_dwdp(): return get_parallel().dwdp_size > 1 +def tp_reduce_scatter( + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + context: "CommunicateContext", +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Module-level so MHC communicators can reuse it without a + ``LayerCommunicator`` instance.""" + if hidden_states.shape[0] == 0: + return hidden_states, hidden_states + assert hidden_states.shape[0] % context.tp_size == 0, ( + f"Expected total tokens {hidden_states.shape[0]} % tp_size {context.tp_size} to be 0" + ) + local_tokens = hidden_states.shape[0] // context.tp_size + output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:]) + get_tp_group().reduce_scatter_tensor(output, hidden_states) + if residual is not None: + residual = residual.tensor_split(context.tp_size)[context.tp_rank] + return output, residual + + class LayerCommunicator: def __init__( self, @@ -826,19 +858,7 @@ class LayerCommunicator: hidden_states: torch.Tensor, residual: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - if hidden_states.shape[0] == 0: - return hidden_states, hidden_states - assert hidden_states.shape[0] % self._context.tp_size == 0, ( - f"Expected total tokens {hidden_states.shape[0]} % tp_size {self._context.tp_size} to be 0" - ) - local_tokens = hidden_states.shape[0] // self._context.tp_size - output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:]) - get_tp_group().reduce_scatter_tensor(output, hidden_states) - if residual is not None: - residual = residual.tensor_split(self._context.tp_size)[ - self._context.tp_rank - ] - return output, residual + return tp_reduce_scatter(hidden_states, residual, self._context) def prepare_mlp( self, @@ -862,6 +882,13 @@ class LayerCommunicator: context=self._context, ) + def maybe_prefetch_next_full_attention_kv( + self, + forward_batch: ForwardBatch, + next_full_attention_layer_id: Optional[int], + ) -> None: + return + def postprocess_layer( self, hidden_states: torch.Tensor, diff --git a/python/sglang/srt/layers/communicator_mhc.py b/python/sglang/srt/layers/communicator_mhc.py new file mode 100644 index 000000000..8a61902fe --- /dev/null +++ b/python/sglang/srt/layers/communicator_mhc.py @@ -0,0 +1,557 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +from dataclasses import dataclass +from functools import partial +from typing import Callable, Optional + +import torch + +from sglang.kernels.ops.layernorm.mhc import hc_contract, hc_expand +from sglang.srt.distributed import get_tp_group +from sglang.srt.distributed.communication_op import ( + attention_tensor_model_parallel_all_reduce, +) +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.communicator import ( + AttentionInputs, + CommunicateContext, + CommunicateSimpleFn, + CommunicateSummableTensorPairFn, + CommunicateWithAllReduceAndLayerNormFn, + LayerCommunicator, + LayerScatterModes, + ScatterMode, + get_attn_tp_context, + tp_reduce_scatter, +) +from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather_into_tensor, + attn_tp_reduce_scatter_tensor, + dp_gather_replicate, + dp_reduce_scatter_tensor, + dp_scatter, + get_dp_global_num_tokens, + get_global_dp_buffer, + get_local_dp_buffer_mhc, + is_allocation_symmetric, +) +from sglang.srt.layers.moe import should_use_dp_reduce_scatterv +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +def tp_all_gather_hidden_states(hidden_states, forward_batch): + assert get_attn_tp_context().input_scattered, ( + "Input scattered guarantees same num tokens in TP group." + ) + total_tokens = forward_batch.input_ids.shape[0] + output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1])) + get_tp_group().all_gather_into_tensor(output, hidden_states) + + return output + + +@dataclass +class MHCState: + """Parameters belong to the owning layer; this state only holds scratch + shared across communication stages.""" + + hc_mult: int + hc_attn_pre: Callable + hc_ffn_pre: Callable + hc_post: Callable + h_res: Optional[torch.Tensor] = None + h_post: Optional[torch.Tensor] = None + + @staticmethod + def _resolve_out_norm(out_norm): + if out_norm is None: + return None, None + return out_norm.weight.data, out_norm.variance_epsilon + + def attn_split(self, hidden_states, out_norm: Optional[torch.nn.Module] = None): + residual = hidden_states + out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) + hidden_states, self.h_res, self.h_post, norm_fused = self.hc_attn_pre( + hidden_states, out_norm_weight, out_norm_eps + ) + if out_norm is not None and not norm_fused and hidden_states.shape[0] != 0: + hidden_states = out_norm(hidden_states) + return hidden_states, residual + + def attn_to_mlp( + self, hidden_states, residual, out_norm: Optional[torch.nn.Module] = None + ): + hidden_states = self.hc_post(hidden_states, residual, self.h_res, self.h_post) + residual = hidden_states + out_norm_weight, out_norm_eps = self._resolve_out_norm(out_norm) + hidden_states, self.h_res, self.h_post, norm_fused = self.hc_ffn_pre( + hidden_states, out_norm_weight, out_norm_eps + ) + if out_norm is not None and not norm_fused and hidden_states.shape[0] != 0: + hidden_states = out_norm(hidden_states) + return hidden_states, residual + + def mlp_combine(self, hidden_states, residual): + return self.hc_post(hidden_states, residual, self.h_res, self.h_post) + + def reset_aux(self): + self.h_res = None + self.h_post = None + + +class MHCCommunicateWithAllReduceAndLayerNormFn(CommunicateWithAllReduceAndLayerNormFn): + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + hidden_states_output_mode: ScatterMode, + residual_output_mode: ScatterMode, + context: CommunicateContext, + ): + fn = CommunicateWithAllReduceAndLayerNormFn.get_fn( + hidden_states_input_mode, + residual_input_mode, + hidden_states_output_mode, + residual_output_mode, + context, + ) + replacements = { + CommunicateWithAllReduceAndLayerNormFn._simple: MHCCommunicateWithAllReduceAndLayerNormFn._simple, + CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual: MHCCommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual, + CommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual: MHCCommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual, + } + if isinstance(fn, partial): + return partial( + replacements.get(fn.func, fn.func), + *fn.args, + **(fn.keywords or {}), + ) + return replacements.get(fn, fn) + + @staticmethod + def _scatter_hidden_states_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + residual_input_mode, + mhc: MHCState, + ): + input_hidden_states = hidden_states + hidden_states = hidden_states.tensor_split(context.attn_tp_size)[ + context.attn_tp_rank + ] + attn_tp_reduce_scatter_tensor(hidden_states, input_hidden_states) + if residual_input_mode == ScatterMode.TP_ATTN_FULL: + residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank] + mhc.h_res = mhc.h_res.tensor_split(context.attn_tp_size)[ + context.attn_tp_rank + ] + mhc.h_post = mhc.h_post.tensor_split(context.attn_tp_size)[ + context.attn_tp_rank + ] + + hidden_states, residual = mhc.attn_to_mlp( + hidden_states, residual, out_norm=layernorm + ) + return hidden_states, residual + + @staticmethod + def _simple( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + mhc: MHCState, + ): + hidden_states, residual = mhc.attn_to_mlp( + hidden_states, residual, out_norm=layernorm + ) + return hidden_states, residual + + @staticmethod + def _tp_all_reduce_with_scattered_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + mhc: MHCState, + ): + if hidden_states.shape[0] == 0: + return hidden_states, hidden_states + + scatter_states = hidden_states.tensor_split(context.tp_size)[context.tp_rank] + get_tp_group().reduce_scatter_tensor(scatter_states, hidden_states) + + scatter_states, residual = mhc.attn_to_mlp( + scatter_states, residual, out_norm=layernorm + ) + + attn_tp_all_gather_into_tensor(hidden_states, scatter_states) + + return hidden_states, residual + + @staticmethod + def _gather_hidden_states_and_residual( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + layernorm: torch.nn.Module, + context: CommunicateContext, + *, + residual_input_mode, + mhc: MHCState, + ): + if get_attn_tp_context().input_scattered: + return MHCCommunicateWithAllReduceAndLayerNormFn._tp_all_reduce_with_scattered_residual( + hidden_states, + residual, + layernorm, + context, + mhc=mhc, + ) + + if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1: + raise NotImplementedError( + "Unsupported: h_res/h_post allgather not implemented." + ) + + hidden_states = attention_tensor_model_parallel_all_reduce(hidden_states) + if context.attn_dp_size != 1: + if hidden_states.shape[0] != 0: + with use_symmetric_memory( + get_tp_group(), + disabled=not is_allocation_symmetric(), + ): + hidden_states, residual = mhc.attn_to_mlp( + hidden_states, residual, out_norm=layernorm + ) + else: + hidden_states, residual = mhc.attn_to_mlp(hidden_states, residual) + + hidden_states, local_hidden_states = ( + get_global_dp_buffer(get_tp_group()), + hidden_states, + ) + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) + else: + hidden_states, residual = mhc.attn_to_mlp( + hidden_states, residual, out_norm=layernorm + ) + return hidden_states, residual + + +class MHCCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): + @staticmethod + def get_fn( + hidden_states_input_mode: ScatterMode, + residual_input_mode: ScatterMode, + output_mode: ScatterMode, + context: CommunicateContext, + ): + fn = CommunicateSummableTensorPairFn.get_fn( + hidden_states_input_mode, + residual_input_mode, + output_mode, + context, + ) + replacements = { + CommunicateSummableTensorPairFn._trivial: MHCCommunicateSummableTensorPairFn._trivial, + CommunicateSummableTensorPairFn._scatter_hidden_states: MHCCommunicateSummableTensorPairFn._scatter_hidden_states, + CommunicateSummableTensorPairFn._gather: MHCCommunicateSummableTensorPairFn._gather, + CommunicateSummableTensorPairFn._scatter: MHCCommunicateSummableTensorPairFn._scatter, + } + return replacements.get(fn, fn) + + @staticmethod + def _trivial( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + *, + mhc: MHCState, + is_last_layer: bool, + **kwargs, + ): + if get_attn_tp_context().input_scattered: + hidden_states, _ = tp_reduce_scatter(hidden_states, None, context) + + hidden_states = mhc.mlp_combine(hidden_states, residual) + if not is_last_layer: + return hidden_states, None + + hidden_states = hc_contract(hidden_states, mhc.hc_mult) + if get_attn_tp_context().input_scattered: + local_states = hidden_states + hidden_states = local_states.new_empty( + local_states.shape[0] * context.tp_size, *local_states.shape[1:] + ) + get_tp_group().all_gather_into_tensor(hidden_states, local_states) + + return hidden_states, None + + @staticmethod + def _scatter_hidden_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + allow_reduce_scatter: bool = False, + *, + mhc: MHCState, + is_last_layer: bool, + **kwargs, + ): + hidden_states, global_hidden_states = ( + get_local_dp_buffer_mhc(get_tp_group(), 1), + hidden_states, + ) + # MoE skips its post-expert all-reduce with reduce_scatterv, so this + # scatter must reduce while combining local-expert partial sums. + if should_use_dp_reduce_scatterv(): + get_tp_group().reduce_scatterv( + global_hidden_states, + output=hidden_states, + sizes=get_dp_global_num_tokens(), + ) + elif allow_reduce_scatter and forward_batch.dp_padding_mode.is_max_len(): + dp_reduce_scatter_tensor(hidden_states, global_hidden_states) + else: + dp_scatter(hidden_states, global_hidden_states, forward_batch) + + hidden_states = mhc.mlp_combine(hidden_states, residual) + if not is_last_layer: + return hidden_states, None + + hidden_states = hc_contract(hidden_states, mhc.hc_mult) + return hidden_states, None + + @staticmethod + def _gather( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + *, + mhc: MHCState, + is_last_layer: bool, + **kwargs, + ): + hidden_states = mhc.mlp_combine(hidden_states, residual) + if is_last_layer: + hidden_states = hc_contract(hidden_states, mhc.hc_mult) + + hidden_states, local_hidden_states = ( + get_local_dp_buffer_mhc( + get_tp_group(), 1 if is_last_layer else mhc.hc_mult + ), + hidden_states, + ) + + attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states) + return hidden_states, None + + @staticmethod + def _scatter( + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + context: CommunicateContext, + *, + mhc: MHCState, + is_last_layer: bool, + **kwargs, + ): + hidden_states = hidden_states.tensor_split(context.attn_tp_size)[ + context.attn_tp_rank + ] + residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank] + + hidden_states = mhc.mlp_combine(hidden_states, residual) + + return hidden_states, None + + +class MHCLayerCommunicator(LayerCommunicator): + def __init__( + self, + layer_scatter_modes: LayerScatterModes, + input_layernorm: torch.nn.Module, + post_attention_layernorm: torch.nn.Module, + allow_reduce_scatter: bool = False, + is_last_layer: bool = False, + qkv_latent_func: Optional[Callable] = None, + *, + is_first_layer: bool, + hc_mult: int, + hc_attn_pre: Callable, + hc_ffn_pre: Callable, + hc_post: Callable, + ): + self.is_first_layer = is_first_layer + self.mhc = MHCState( + hc_mult=hc_mult, + hc_attn_pre=hc_attn_pre, + hc_ffn_pre=hc_ffn_pre, + hc_post=hc_post, + ) + + super().__init__( + layer_scatter_modes, + input_layernorm, + post_attention_layernorm, + allow_reduce_scatter, + is_last_layer, + qkv_latent_func, + ) + + def _post_init_communicate(self): + # Base MOE_FULL callables do not accept ``mhc``, so reject this + # combination at construction. + if self.layer_scatter_modes.mlp_mode == ScatterMode.MOE_FULL: + raise NotImplementedError( + "MHCLayerCommunicator does not support MOE_FULL " + "(moe_dp_size < attention_context_parallel_size). Increase " + "moe_dp_size to match attention_context_parallel_size." + ) + self._communicate_simple_fn = CommunicateSimpleFn.get_fn( + input_mode=self.layer_scatter_modes.layer_input_mode, + output_mode=self.layer_scatter_modes.attn_mode, + context=self._context, + ) + self._communicate_with_all_reduce_and_layer_norm_fn = ( + MHCCommunicateWithAllReduceAndLayerNormFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.attn_mode, + residual_input_mode=self.layer_scatter_modes.layer_input_mode, + hidden_states_output_mode=self.layer_scatter_modes.mlp_mode, + residual_output_mode=self.layer_scatter_modes.middle_residual_mode, + context=self._context, + ) + ) + self._communicate_summable_tensor_pair_fn = ( + MHCCommunicateSummableTensorPairFn.get_fn( + hidden_states_input_mode=self.layer_scatter_modes.mlp_mode, + residual_input_mode=self.layer_scatter_modes.middle_residual_mode, + output_mode=self.layer_scatter_modes.layer_output_mode, + context=self._context, + ) + ) + + def prepare_attn( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + ): + if self.is_first_layer: + if get_attn_tp_context().input_scattered: + hidden_states, _ = tp_reduce_scatter( + hidden_states, + None, + self._context, + ) + hidden_states = hc_expand(hidden_states, self.mhc.hc_mult) + + hidden_states, residual = self.mhc.attn_split( + hidden_states, out_norm=self.input_layernorm + ) + + hidden_states = self._communicate_simple_fn( + hidden_states=hidden_states, + forward_batch=forward_batch, + context=self._context, + ) + + # DSA and attention without a QKV hook consume full hidden states, so + # gather them before attention. + ctx = get_attn_tp_context() + dsa_pre_gather = ctx.input_scattered and ctx.is_dsa + no_qkv_latent_pre_gather = ctx.input_scattered and self.qkv_latent_func is None + if dsa_pre_gather or no_qkv_latent_pre_gather: + hidden_states = tp_all_gather_hidden_states(hidden_states, forward_batch) + + if self.qkv_latent_func is not None: + attn_inputs = AttentionInputs( + hidden_states, + forward_batch, + self.qkv_latent_func, + is_pre_gathered=dsa_pre_gather, + ) + ctx.set_attn_inputs(attn_inputs) + + return hidden_states, residual + + def prepare_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + forward_batch: ForwardBatch, + cache=None, + ): + if cache is not None: + self._context.cache = cache + + hidden_states, residual = self._communicate_with_all_reduce_and_layer_norm_fn( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + layernorm=self.post_attention_layernorm, + context=self._context, + mhc=self.mhc, + ) + + return hidden_states, residual + + def postprocess_layer(self, hidden_states, residual, forward_batch): + hidden_states, residual = self._communicate_summable_tensor_pair_fn( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + context=self._context, + allow_reduce_scatter=self.allow_reduce_scatter, + mhc=self.mhc, + is_last_layer=self.is_last_layer, + ) + self.mhc.reset_aux() + + return hidden_states, residual + + def should_fuse_mlp_allreduce_with_next_layer(self, forward_batch): + return False + + def should_use_reduce_scatter(self, forward_batch: ForwardBatch): + if not self.allow_reduce_scatter: + return False + if ( + self._communicate_summable_tensor_pair_fn + is MHCCommunicateSummableTensorPairFn._scatter_hidden_states + ): + # reduce_scatterv already combines expert outputs; returning False + # would make RowParallelLinear perform an extra all-reduce. + if should_use_dp_reduce_scatterv(): + return True + if forward_batch.dp_padding_mode.is_max_len(): + return True + + if get_attn_tp_context().input_scattered: + return True + return False diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 86e2131f2..7340070c9 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -213,6 +213,21 @@ class _DpGatheredBufferWrapper: ) return buffer + @classmethod + def get_local_dp_buffer_mhc( + cls, group: GroupCoordinator, n: int = 1 + ) -> torch.Tensor: + from sglang.srt.runtime_context import get_flags + + dp = get_flags().dp + with use_symmetric_memory(group, disabled=not cls._dp_max_padding): + buffer = torch.empty( + (cls._local_dp_buffer_len, dp.buffer_hidden_size * n), + dtype=dp.buffer_dtype, + device=dp.buffer_device, + ) + return buffer + @classmethod def get_global_dp_buffer_len(cls) -> int: return cls._global_dp_buffer_len @@ -277,6 +292,10 @@ def get_local_dp_buffer(group: GroupCoordinator) -> torch.Tensor: return _DpGatheredBufferWrapper.get_local_dp_buffer(group=group) +def get_local_dp_buffer_mhc(group: GroupCoordinator, n: int = 1) -> torch.Tensor: + return _DpGatheredBufferWrapper.get_local_dp_buffer_mhc(group=group, n=n) + + def get_global_dp_buffer_len() -> int: return _DpGatheredBufferWrapper.get_global_dp_buffer_len() diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index e43857ae1..5f352f156 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -1650,6 +1650,7 @@ def _situ_mul_quant_contig_kernel( def _apply_swiglu_limit( gateup_output: torch.Tensor, swiglu_limit: float ) -> torch.Tensor: + """Clamp the contiguous runner's owned GEMM workspace in place.""" assert swiglu_limit == 10 num_tokens, hidden_size_x2 = gateup_output.shape @@ -1659,12 +1660,12 @@ def _apply_swiglu_limit( assert gate.shape == (num_tokens, hidden_size_x2 // 2) assert up.shape == (num_tokens, hidden_size_x2 // 2) - up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit) - gate = torch.clamp(gate, max=swiglu_limit) - - out = torch.cat([gate, up], dim=-1) - assert out.shape == (num_tokens, hidden_size_x2) - return out + # Both halves are views of a fresh GEMM output. Avoid separate clamped + # copies and their concatenation: large compact prefills need that + # headroom for the activation and down-projection workspaces. + up.clamp_(min=-swiglu_limit, max=swiglu_limit) + gate.clamp_(max=swiglu_limit) + return gateup_output @register_pre_permute("deepep_v2", "deep_gemm") diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index db3d02c29..204485c10 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1877,7 +1877,9 @@ class Req(ReqDllmMixin): ) self.kv.retraction_backup = RetractionBackup( cpu_tensors=token_to_kv_pool_allocator.get_cpu_copy( - token_indices, mamba_indices=self.kv.mamba_pool_idx + token_indices, + mamba_indices=self.kv.mamba_pool_idx, + req_pool_index=self.kv.req_pool_idx, ), mamba_cpu=( mamba_pool.get_cpu_copy(self.kv.mamba_pool_idx.unsqueeze(0)) @@ -1901,6 +1903,7 @@ class Req(ReqDllmMixin): self.kv.retraction_backup.cpu_tensors, token_indices, mamba_indices=self.kv.mamba_pool_idx, + req_pool_index=self.kv.req_pool_idx, ) self.kv.retraction_backup = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 33c08792d..a19e59287 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -16,6 +16,7 @@ import dataclasses import faulthandler import logging +import math import os import signal import sys @@ -681,6 +682,7 @@ class Scheduler( # Init prefill kv split size when deterministic inference is enabled with various attention backends self.init_deterministic_inference_config() + self.init_dsa_kpool_truncation_align() self.init_weight_updater() @@ -1118,6 +1120,16 @@ class Scheduler( if self.server_args.is_startup_weight_load_overlap: self.tp_worker.finalize_startup_weight_load() + # Adaptive/speculative graphs and post-capture KV sizing can consume + # the headroom seen by the initial DeepGEMM layout budget. Refresh it + # after these allocations, before elastic EP rejoins healthy ranks + # that do not participate in this startup collective. + from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( + refresh_deep_gemm_layout_memory_budget, + ) + + refresh_deep_gemm_layout_memory_budget(model_runner, only_if_initialized=True) + if ( get_exec().moe.elastic_ep_backend is not None and get_exec().moe.ep_join_mode == "recover" @@ -1672,6 +1684,28 @@ class Scheduler( get_int_env_var(env_var, default_size) if env_var else None ) + def init_dsa_kpool_truncation_align(self): + """Kpool compress-write asserts chunked extends start on pool boundaries. + Use the LCM to preserve any existing deterministic-inference alignment.""" + from sglang.srt.configs.model_config import ( + get_dsa_index_kpool, + is_deepseek_dsa, + ) + + if not is_deepseek_dsa(self.model_config.hf_config): + return + + dsa_index_kpool = get_dsa_index_kpool(self.model_config.hf_config) + if dsa_index_kpool <= 1: + return + + if self.truncation_align_size is None: + self.truncation_align_size = dsa_index_kpool + else: + self.truncation_align_size = math.lcm( + self.truncation_align_size, dsa_index_kpool + ) + def init_request_dispatcher(self): self._request_dispatcher = TypeBasedDispatcher( [ diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index d5efd3b37..ffbc4b96d 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -121,10 +121,12 @@ class BaseTokenToKVPoolAllocator(abc.ABC): virtual-id pools must override.""" return kv_indices - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError() - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): raise NotImplementedError() def alloc_extend(self, *args, **kwargs): diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 0d5ec1d9c..4d77446c9 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -340,10 +340,19 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.staged_pages: list[torch.Tensor] = [] self.num_staged_pages = 0 - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): + return self._kvcache.get_cpu_copy( + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, + ) - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices + kv_cache_cpu, + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, ) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 4b5c61812..ba4a3ae7c 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -471,12 +471,21 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.free_group = None self.swa_free_group = [] - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): + return self._kvcache.get_cpu_copy( + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, + ) - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices + kv_cache_cpu, + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, ) diff --git a/python/sglang/srt/mem_cache/allocator/token.py b/python/sglang/srt/mem_cache/allocator/token.py index fd0f32682..60aa4af5d 100644 --- a/python/sglang/srt/mem_cache/allocator/token.py +++ b/python/sglang/srt/mem_cache/allocator/token.py @@ -74,10 +74,19 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): else: self.free_group.append(self._copy_for_free_group(free_index)) - def get_cpu_copy(self, indices, mamba_indices=None): - return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): + return self._kvcache.get_cpu_copy( + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, + ) - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): return self._kvcache.load_cpu_copy( - kv_cache_cpu, indices, mamba_indices=mamba_indices + kv_cache_cpu, + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, ) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 8287abbf8..72f59e72a 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -256,10 +256,12 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): loc = self.translate_loc_to_hisparse_device(loc) return super().set_key_buffer_fused(layer_id, loc, cache_k) - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError("HiSparseC4DevicePool does not support get_cpu_copy") - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): raise NotImplementedError("HiSparseC4DevicePool does not support load_cpu_copy") diff --git a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py index a4861f55a..0f387edae 100644 --- a/python/sglang/srt/mem_cache/dsa_cache_layer_split.py +++ b/python/sglang/srt/mem_cache/dsa_cache_layer_split.py @@ -551,7 +551,7 @@ class LayerSplitDSATokenToKVPool(DSATokenToKVPool): # ---- HiCache CPU offload: skip empty (non-owned) layers --------------- - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): from sglang.srt.utils import current_platform current_platform.synchronize() @@ -569,9 +569,18 @@ class LayerSplitDSATokenToKVPool(DSATokenToKVPool): kv_cache_cpu[-1].append(kv_cpu) current_platform.synchronize() - return {"kv": kv_cache_cpu, "index_k": self.index_key_cache.cpu_copy(indices)} + return { + "kv": kv_cache_cpu, + "index_k": self.index_key_cache.cpu_copy(indices), + } - def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None): + def load_cpu_copy( + self, + kv_cache_cpu_dict, + indices, + mamba_indices=None, + req_pool_index=None, + ): from sglang.srt.utils import current_platform kv_cache_cpu = kv_cache_cpu_dict["kv"] diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 1bf865ca2..4808dadd3 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -40,6 +40,11 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool): kv_cache_dim: int, start_layer: Optional[int] = None, end_layer: Optional[int] = None, + index_kpool: int = 1, + index_kpool_compress: bool = False, + tail_extra_slots: int = 0, + max_running_requests: Optional[int] = None, + skip_topk_layers: Optional[list[bool]] = None, host_to_device_ratio: int = 2, ): super().__init__( @@ -56,6 +61,11 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool): start_layer=start_layer, end_layer=end_layer, index_buf_size=size * host_to_device_ratio, + index_kpool=index_kpool, + index_kpool_compress=index_kpool_compress, + tail_extra_slots=tail_extra_slots, + max_running_requests=max_running_requests, + skip_topk_layers=skip_topk_layers, ) self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize @@ -115,8 +125,10 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool): num_layers=self.layer_num, ) - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError("HiSparseDevicePool does not support get_cpu_copy") - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy") diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index b3ed6f84d..2d9e5ad77 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -788,8 +788,11 @@ def build_hybrid_mamba_stack( ) -> tuple[HostPoolGroup, HybridCacheController]: transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping) mamba_allocator = params.req_to_token_pool.mamba_allocator + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + mtp_draft_device_pools = tuple( - pool.full_kv_pool for pool in params.mtp_draft_device_pools + pool.full_kv_pool if isinstance(pool, HybridLinearKVPool) else pool + for pool in params.mtp_draft_device_pools ) kv_host_size, mamba_host_size = None, 0 if get_memory().hicache_size > 0: diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 65b591acd..92adb20dc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -25,6 +25,7 @@ from typing import TYPE_CHECKING from sglang.srt.arg_groups.overrides import resolving_view from sglang.srt.configs.hybrid_arch import ( + glm5_next_config, hybrid_gdn_config, hybrid_lightning_config, kimi_linear_config, @@ -126,6 +127,7 @@ def uses_ssm_state(model_config) -> bool: or mamba2_config(model_config) is not None or (spec.uses_mamba_radix_cache if spec is not None else False) or kimi_linear_config(model_config) is not None + or glm5_next_config(model_config) is not None or hybrid_lightning_config(model_config) is not None ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1c8df4911..0c293e90d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -19,6 +19,8 @@ from sglang.srt.configs.model_config import ( ModelConfig, dsa_layer_skips_topk, get_dsa_index_head_dim, + get_dsa_index_kpool, + get_dsa_index_kpool_compress, get_minimax_sparse_attention_config, get_minimax_sparse_disable_value_layer_ids, get_minimax_sparse_layer_ids, @@ -999,6 +1001,18 @@ class KVCacheConfigurator: ) return req_to_token_pool + def _get_mamba_layer_ids_for_req_pool(self) -> list: + mamba_layer_ids = [ + i + for i in self.mambaish_config.mamba2_cache_params.layers + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ] + if max_speculative_num_draft_tokens(): + for layer_id in getattr(self.mambaish_config, "nextn_layer_ids", []): + if layer_id not in mamba_layer_ids: + mamba_layer_ids.append(layer_id) + return mamba_layer_ids + def _build_hybrid_mamba_decode_req_pool( self, *, @@ -1016,13 +1030,7 @@ class KVCacheConfigurator: device=self.device, enable_memory_saver=get_exec().features.enable_memory_saver, cache_params=self.mambaish_config.mamba2_cache_params, - mamba_layer_ids=( - [ - i - for i in self.mambaish_config.mamba2_cache_params.layers - if self.layer_info.start_layer <= i < self.layer_info.end_layer - ] - ), + mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(), speculative_num_draft_tokens=max_speculative_num_draft_tokens(), speculative_eagle_topk=get_spec().speculative_eagle_topk, enable_mamba_extra_buffer=mamba_extra_buffer_enabled(), @@ -1095,13 +1103,7 @@ class KVCacheConfigurator: device=self.device, enable_memory_saver=get_exec().features.enable_memory_saver, cache_params=self.mambaish_config.mamba2_cache_params, - mamba_layer_ids=( - [ - i - for i in self.mambaish_config.mamba2_cache_params.layers - if self.layer_info.start_layer <= i < self.layer_info.end_layer - ] - ), + mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(), enable_mamba_extra_buffer=mamba_extra_buffer_enabled(), enable_mamba_extra_buffer_lazy=mamba_extra_buffer_lazy_enabled(), # A PD prefill server never runs TARGET_VERIFY, so skip the @@ -1230,9 +1232,10 @@ class KVCacheConfigurator: swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, is_dsa_model=is_dsa_model, ) - elif self.use_mla_backend and is_dsa_model: + elif self.use_mla_backend and is_dsa_model and not self.mambaish_config: token_to_kv_pool = self._build_dsa_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, + max_running_requests=sizes.max_running_requests, ) elif self.use_mla_backend and not self.mambaish_config: assert not is_dsa_model @@ -1521,7 +1524,9 @@ class KVCacheConfigurator: ) return token_to_kv_pool - def _build_dsa_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: + def _build_dsa_kv_pool( + self, *, max_total_num_tokens: int, max_running_requests: int + ) -> KVCache: from sglang.srt.layers.cp.utils import get_glm_dsa_cp_layer_shard_info ( @@ -1570,6 +1575,12 @@ class KVCacheConfigurator: start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), + index_kpool=get_dsa_index_kpool(self.model_config.hf_config), + index_kpool_compress=get_dsa_index_kpool_compress( + self.model_config.hf_config + ), + tail_extra_slots=(max_speculative_num_draft_tokens() or 0), + max_running_requests=max_running_requests, **pool_kwargs, ) return token_to_kv_pool @@ -1760,12 +1771,6 @@ class KVCacheConfigurator: req_to_token_pool: ReqToTokenPool, mha_pool_class: type, ) -> KVCache: - extra_args = {} - if self.use_mla_backend: - extra_args = { - "kv_lora_rank": self.model_config.kv_lora_rank, - "qk_rope_head_dim": self.model_config.qk_rope_head_dim, - } full_attention_layer_ids = ( [0] if self.is_draft_worker @@ -1775,6 +1780,39 @@ class KVCacheConfigurator: if self.layer_info.start_layer <= i < self.layer_info.end_layer ] ) + extra_args = {} + if self.use_mla_backend: + extra_args = { + "kv_lora_rank": self.model_config.kv_lora_rank, + "qk_rope_head_dim": self.model_config.qk_rope_head_dim, + } + if is_deepseek_dsa(self.model_config.hf_config): + dsa_index_kpool = get_dsa_index_kpool(self.model_config.hf_config) + extra_args.update( + use_dsa=True, + index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), + kv_cache_dim=calculate_mla_kv_cache_dim( + model_config=self.model_config, + kv_cache_dtype=self.kv_cache_dtype, + ), + index_kpool=dsa_index_kpool, + index_kpool_compress=get_dsa_index_kpool_compress( + self.model_config.hf_config + ), + skip_topk_layers=( + None + if self.is_draft_worker + else [ + dsa_layer_skips_topk(self.model_config.hf_config, layer_id) + for layer_id in full_attention_layer_ids + ] + ), + ) + if dsa_index_kpool > 1: + extra_args.update( + tail_extra_slots=(max_speculative_num_draft_tokens() or 0), + max_running_requests=(req_to_token_pool.req_to_token.shape[0]), + ) quant_method = self._build_mha_quant_method( num_layers=len(full_attention_layer_ids) ) @@ -2430,14 +2468,21 @@ def calculate_mla_kv_cache_dim( if not is_dsa_model: return kv_cache_dim - # TRTLLM backend does not override kv_cache_dim for MLA kv cache - # Assuming dsa prefill and decode backends are the same when using trtllm MLA backend, - # since it is not compatible for trtllm and other mla attn backend due to the different - # kv cache layout. - if ( - get_exec().kernel.dsa_prefill_backend == "trtllm" - or get_exec().kernel.dsa_decode_backend == "trtllm" - ): + # TRTLLM uses the raw MLA KV layout. In disaggregated serving only the + # backend for the local role determines the local pool layout; the + # inactive role may legitimately have a different default backend. + disaggregation_mode = get_disagg().disaggregation_mode + if disaggregation_mode == "decode": + uses_trtllm_kv_layout = get_exec().kernel.dsa_decode_backend == "trtllm" + elif disaggregation_mode == "prefill": + uses_trtllm_kv_layout = get_exec().kernel.dsa_prefill_backend == "trtllm" + else: + uses_trtllm_kv_layout = ( + get_exec().kernel.dsa_prefill_backend == "trtllm" + or get_exec().kernel.dsa_decode_backend == "trtllm" + ) + + if uses_trtllm_kv_layout: return kv_cache_dim # On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 59cec0fb5..ebc8cc3d2 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1406,6 +1406,19 @@ class HybridReqToTokenPool(ReqToTokenPool): def mamba2_layer_cache(self, layer_id: int): return self.mamba_pool.mamba2_layer_cache(self.mamba2_layer_index(layer_id)) + def copy_mamba_state( + self, src_index: torch.Tensor, dst_index: torch.Tensor + ) -> None: + if src_index.numel() == 0: + return + if ( + self.layer_transfer_counter is not None + and self.layer_transfer_counter.consumer_index >= 0 + ): + last_mamba_layer = max(self.mamba_map) + self.layer_transfer_counter.wait_until(last_mamba_layer - self.start_layer) + self.mamba_pool.copy_from(src_index, dst_index) + def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState: return self.mamba_pool.get_speculative_mamba2_params_all_layers() @@ -1554,6 +1567,9 @@ class HybridReqToTokenPool(ReqToTokenPool): req.kv.mamba_ping_pong_track_buffer = None req.kv.mamba_next_track_idx = None req.kv.mamba_last_track_idx = None + req.kv.mamba_last_track_seqlen = None + req.kv.mamba_cow_src_index = None + req.kv.mamba_needs_clear = False def clear(self): logger.info("Reset HybridReqToTokenPool") @@ -1761,10 +1777,12 @@ class KVCache(abc.ABC): def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): self.layer_transfer_counter = layer_transfer_counter - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError() - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): raise NotImplementedError() def get_kv_cache_quant_method(self) -> Any: @@ -2265,7 +2283,7 @@ class MHATokenToKVPool(KVCache): item_lens = [d.item_len_bytes(self.page_size) for d in self._kv_buffer_descs] return ptrs, lens, item_lens - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): assert not self.use_hnd, ( "CPU KV offload indexes by slot (NHD); HND KV cache " "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet." @@ -2287,7 +2305,9 @@ class MHATokenToKVPool(KVCache): current_platform.synchronize() return kv_cache_cpu - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): assert not self.use_hnd, ( "CPU KV offload indexes by slot (NHD); HND KV cache " "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet." @@ -3224,13 +3244,15 @@ class PageMajorMHATokenToKVPool(MHATokenToKVPool): "with a page-aware transfer scheme)." ) - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError( "CPU offloading is unsupported under the page-major layout " "(TODO: split token ids into page/slot for the 4-D index)." ) - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): raise NotImplementedError( "CPU offloading is unsupported under the page-major layout " "(TODO: split token ids into page/slot for the 4-D index)." @@ -3509,7 +3531,7 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): ) return self.k_scale_buffer[idx][loc], self.v_scale_buffer[idx][loc] - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): # The scales travel with their fp8 payload; a restored slot dequantizes # against mismatched exponents without them. assert not self.use_hnd, ( @@ -3539,7 +3561,9 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): current_platform.synchronize() return kv_cache_cpu - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): assert not self.use_hnd, ( "CPU KV offload indexes by slot (NHD); HND KV cache " "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet." @@ -3625,6 +3649,14 @@ class HybridLinearKVPool(KVCache): use_mla: bool = False, kv_lora_rank: int = None, qk_rope_head_dim: int = None, + use_dsa: bool = False, + index_head_dim: Optional[int] = None, + kv_cache_dim: Optional[int] = None, + index_kpool: int = 1, + index_kpool_compress: bool = False, + tail_extra_slots: int = 0, + max_running_requests: Optional[int] = None, + skip_topk_layers: Optional[List[bool]] = None, start_layer: Optional[int] = None, full_kv_pool_class: Optional[type] = None, quant_method=None, @@ -3648,6 +3680,7 @@ class HybridLinearKVPool(KVCache): # `load_cpu_copy`, the only readers, so those ids never arrive here. self._mamba_translate = lambda ids: ids self.use_mla = use_mla + self.use_dsa = use_dsa if full_kv_pool is not None: # Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool # aliasing the shared byte buffer. @@ -3693,6 +3726,30 @@ class HybridLinearKVPool(KVCache): **quant_method_kwarg, **post_capture_kwargs, ) + elif use_dsa: + # DSA sparse full-attention layers share the MLA latent layout and + # additionally keep a paged index_k cache. Only full-attn layer count + # is allocated here; the wrapper translates global layer_id to dense. + assert index_head_dim is not None and kv_cache_dim is not None, ( + "HybridLinearKVPool with use_dsa requires index_head_dim and kv_cache_dim" + ) + self.full_kv_pool = DSATokenToKVPool( + size=size, + page_size=self.page_size, + kv_lora_rank=kv_lora_rank, + dtype=dtype, + qk_rope_head_dim=qk_rope_head_dim, + layer_num=self.full_layer_nums, + device=device, + index_head_dim=index_head_dim, + enable_memory_saver=enable_memory_saver, + kv_cache_dim=kv_cache_dim, + index_kpool=index_kpool, + index_kpool_compress=index_kpool_compress, + tail_extra_slots=tail_extra_slots, + max_running_requests=max_running_requests, + skip_topk_layers=skip_topk_layers, + ) else: TokenToKVPoolClass = MLATokenToKVPool @@ -3737,6 +3794,38 @@ class HybridLinearKVPool(KVCache): self.full_kv_pool._finalize_backing_tokens(config.max_total_num_tokens) self.size = int(config.max_total_num_tokens) + @property + def dsa_kv_cache_store_fp8(self) -> bool: + return getattr(self.full_kv_pool, "dsa_kv_cache_store_fp8", False) + + @property + def kv_cache_dim(self): + return getattr(self.full_kv_pool, "kv_cache_dim", None) + + @property + def index_head_dim(self) -> Optional[int]: + return getattr(self.full_kv_pool, "index_head_dim", None) + + @property + def quant_block_size(self) -> Optional[int]: + return getattr(self.full_kv_pool, "quant_block_size", None) + + @property + def index_kpool(self) -> int: + return getattr(self.full_kv_pool, "index_kpool", 1) + + @property + def index_kpool_compress(self) -> bool: + return bool(getattr(self.full_kv_pool, "index_kpool_compress", False)) + + @property + def tail_extra_slots(self) -> int: + return getattr(self.full_kv_pool, "tail_extra_slots", 0) + + @property + def slots_per_page(self) -> int: + return getattr(self.full_kv_pool, "slots_per_page", self.page_size) + def get_kv_size_bytes(self): return self.full_kv_pool.get_kv_size_bytes() @@ -3901,8 +3990,8 @@ class HybridLinearKVPool(KVCache): def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): self.full_kv_pool.move_kv_cache(tgt_loc, src_loc) - def get_cpu_copy(self, indices, mamba_indices=None): - kv_cpu = self.full_kv_pool.get_cpu_copy(indices) + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): + kv_cpu = self.full_kv_pool.get_cpu_copy(indices, req_pool_index=req_pool_index) # mamba_pool stores PHYSICAL ids; translate the (unified-pool virtual) ids first. mamba_cpu = ( self.mamba_pool.get_cpu_copy(self._mamba_translate(mamba_indices)) @@ -3911,9 +4000,11 @@ class HybridLinearKVPool(KVCache): ) return kv_cpu, mamba_cpu - def load_cpu_copy(self, cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): kv_cpu, mamba_cpu = cache_cpu - self.full_kv_pool.load_cpu_copy(kv_cpu, indices) + self.full_kv_pool.load_cpu_copy(kv_cpu, indices, req_pool_index=req_pool_index) if mamba_cpu is not None and mamba_indices is not None: self.mamba_pool.load_cpu_copy( mamba_cpu, self._mamba_translate(mamba_indices) @@ -3951,6 +4042,145 @@ class HybridLinearKVPool(KVCache): with self._transfer_id_context(layer): return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype) + def set_index_k_scale_buffer( + self, + layer_id: int, + loc: torch.Tensor, + index_k: torch.Tensor, + index_k_scale: torch.Tensor, + ) -> None: + assert self.use_dsa, "set_index_k_scale_buffer called when use_dsa is False" + layer_id = self._transfer_full_attention_id(layer_id) + self.full_kv_pool.set_index_k_scale_buffer( + layer_id, loc, index_k, index_k_scale + ) + + def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: + assert self.use_dsa, ( + "get_index_k_with_scale_buffer called when use_dsa is False" + ) + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_index_k_with_scale_buffer(layer_id) + + def get_broadcastable_index_k_with_scale_buffer( + self, layer_id: int + ) -> torch.Tensor: + assert self.use_dsa, ( + "get_broadcastable_index_k_with_scale_buffer called when use_dsa is False" + ) + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + if hasattr(self.full_kv_pool, "_get_broadcastable_index_buffer"): + return self.full_kv_pool._get_broadcastable_index_buffer(layer_id) + return self.full_kv_pool.get_index_k_with_scale_buffer(layer_id) + + def invalidate_index_buffer_for_layer(self, layer_id: int) -> None: + if not self.use_dsa or not hasattr( + self.full_kv_pool, "invalidate_index_buffer_for_layer" + ): + return + layer_id = self._transfer_full_attention_id(layer_id) + self.full_kv_pool.invalidate_index_buffer_for_layer(layer_id) + + def get_index_k_continuous( + self, + layer_id: int, + seq_len: int, + page_indices: torch.Tensor, + ): + assert self.use_dsa, "get_index_k_continuous called when use_dsa is False" + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_index_k_continuous(layer_id, seq_len, page_indices) + + def get_index_k_scale_continuous( + self, + layer_id: int, + seq_len: int, + page_indices: torch.Tensor, + ): + assert self.use_dsa, "get_index_k_scale_continuous called when use_dsa is False" + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_index_k_scale_continuous( + layer_id, seq_len, page_indices + ) + + def get_index_k_scale_buffer( + self, + layer_id: int, + seq_len_tensor: torch.Tensor, + page_indices: torch.Tensor, + seq_len_sum: int, + max_seq_len: int, + ): + assert self.use_dsa, "get_index_k_scale_buffer called when use_dsa is False" + self._wait_for_layer(layer_id) + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_index_k_scale_buffer( + layer_id, seq_len_tensor, page_indices, seq_len_sum, max_seq_len + ) + + def get_compress_tail_buffers( + self, layer_id: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert self.use_dsa, "get_compress_tail_buffers called when use_dsa is False" + layer_id = self._transfer_full_attention_id(layer_id) + return self.full_kv_pool.get_compress_tail_buffers(layer_id) + + def kpool_decode_update_index_cache( + self, + layer_id: int, + key: torch.Tensor, + slot_score: torch.Tensor, + ape: torch.Tensor, + block_tables: torch.Tensor, + req_pool_indices: torch.Tensor, + positions: torch.Tensor, + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, + round_scale: bool = False, + ) -> None: + assert self.use_dsa, ( + "kpool_decode_update_index_cache called when use_dsa is False" + ) + layer_id = self._transfer_full_attention_id(layer_id) + self.full_kv_pool.kpool_decode_update_index_cache( + layer_id=layer_id, + key=key, + slot_score=slot_score, + ape=ape, + block_tables=block_tables, + req_pool_indices=req_pool_indices, + positions=positions, + seq_lens=seq_lens, + out_cache_loc=out_cache_loc, + round_scale=round_scale, + ) + + def set_compress_tail_for_request( + self, + layer_id: int, + req_pool_idx: torch.Tensor, + key_tail: torch.Tensor, + score_tail: torch.Tensor, + n_remain: int, + dst_logical_start: int, + ) -> None: + assert self.use_dsa, ( + "set_compress_tail_for_request called when use_dsa is False" + ) + layer_id = self._transfer_full_attention_id(layer_id) + self.full_kv_pool.set_compress_tail_for_request( + layer_id=layer_id, + req_pool_idx=req_pool_idx, + key_tail=key_tail, + score_tail=score_tail, + n_remain=n_remain, + dst_logical_start=dst_logical_start, + ) + class MLATokenToKVPool(KVCache): def __init__( @@ -4163,10 +4393,12 @@ class MLATokenToKVPool(KVCache): else: if cache_k_nope.dtype != self.dtype: cache_k_nope = cache_k_nope.to(self.dtype) - cache_k_rope = cache_k_rope.to(self.dtype) + if cache_k_rope is not None and cache_k_rope.numel() > 0: + cache_k_rope = cache_k_rope.to(self.dtype) if self.store_dtype != self.dtype: cache_k_nope = cache_k_nope.view(self.store_dtype) - cache_k_rope = cache_k_rope.view(self.store_dtype) + if cache_k_rope is not None and cache_k_rope.numel() > 0: + cache_k_rope = cache_k_rope.view(self.store_dtype) self._scatter_mla_rows(dst_buffer, loc, cache_k_nope, cache_k_rope) @@ -4213,11 +4445,14 @@ class MLATokenToKVPool(KVCache): dtype=dst_dtype, device=kv_buffer.device, ) - cache_k_rope = torch.empty( - (loc.shape[0], 1, self.qk_rope_head_dim), - dtype=dst_dtype, - device=kv_buffer.device, - ) + if self.qk_rope_head_dim == 0: + cache_k_rope = None + else: + cache_k_rope = torch.empty( + (loc.shape[0], 1, self.qk_rope_head_dim), + dtype=dst_dtype, + device=kv_buffer.device, + ) get_mla_kv_buffer_triton(kv_buffer, loc, cache_k_nope, cache_k_rope) return cache_k_nope, cache_k_rope @@ -4235,12 +4470,14 @@ class MLATokenToKVPool(KVCache): for kv_cache in self.kv_buffer: kv_cache[tgt_loc_flat] = kv_cache[src_loc_flat] - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): current_platform.synchronize() kv_cache_cpu = [] chunk_size = self.cpu_offloading_chunk_size for layer_id in range(self.layer_num): kv_cache_cpu.append([]) + if self.kv_buffer[layer_id].shape[0] == 0: + continue for i in range(0, len(indices), chunk_size): chunk_indices = indices[i : i + chunk_size] kv_cpu = self.kv_buffer[layer_id][chunk_indices].to( @@ -4250,10 +4487,14 @@ class MLATokenToKVPool(KVCache): current_platform.synchronize() return kv_cache_cpu - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): current_platform.synchronize() chunk_size = self.cpu_offloading_chunk_size for layer_id in range(self.layer_num): + if self.kv_buffer[layer_id].shape[0] == 0: + continue for i in range(0, len(indices), chunk_size): chunk_indices = indices[i : i + chunk_size] kv_cpu = kv_cache_cpu[layer_id][i // chunk_size] @@ -4381,13 +4622,18 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): cache_k_nope_fp4, cache_k_nope_fp4_sf = ( FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_nope) ) - cache_k_rope_fp4, cache_k_rope_fp4_sf = ( - FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_rope) - ) + if cache_k_rope is not None and cache_k_rope.numel() > 0: + cache_k_rope_fp4, cache_k_rope_fp4_sf = ( + FP4MXBlock16KVQuantizeUtil.batched_quantize(cache_k_rope) + ) + else: + cache_k_rope_fp4 = None + cache_k_rope_fp4_sf = None if self.store_dtype != self.dtype: cache_k_nope = cache_k_nope.view(self.store_dtype) - cache_k_rope = cache_k_rope.view(self.store_dtype) + if cache_k_rope is not None and cache_k_rope.numel() > 0: + cache_k_rope = cache_k_rope.view(self.store_dtype) self._scatter_mla_rows( self.kv_buffer[layer_id - self.start_layer], @@ -4423,6 +4669,10 @@ class DSATokenToKVPool(MLATokenToKVPool): start_layer: Optional[int] = None, end_layer: Optional[int] = None, index_buf_size: Optional[int] = None, + index_kpool: int = 1, + index_kpool_compress: bool = False, + tail_extra_slots: int = 0, + max_running_requests: Optional[int] = None, skip_topk_layers: Optional[List[bool]] = None, ): override_dim = ( @@ -4446,6 +4696,10 @@ class DSATokenToKVPool(MLATokenToKVPool): # self.index_k_dtype = torch.float8_e4m3fn # self.index_k_scale_dtype = torch.float32 self.index_head_dim = index_head_dim + self.index_kpool = index_kpool + self.index_kpool_compress = index_kpool_compress + self.tail_extra_slots = tail_extra_slots + self.slots_per_page = self.page_size if index_buf_size is None: index_buf_size = size self.index_buf_size = index_buf_size @@ -4471,19 +4725,172 @@ class DSATokenToKVPool(MLATokenToKVPool): else: assert self.page_size == 64 self.index_key_cache = self._create_index_key_cache() + self._init_kpool_compress_tail_buffers( + index_kpool=index_kpool, + index_kpool_compress=index_kpool_compress, + tail_extra_slots=tail_extra_slots, + index_head_dim=index_head_dim, + layer_num=layer_num, + device=device, + max_running_requests=max_running_requests, + ) self._finalize_allocation_log(size) def _create_index_key_cache(self) -> IndexKeyCache: return IndexKeyCache(self, self.index_buf_size) + def _should_allocate_index_layer(self, local_layer_idx: int) -> bool: + return not self.skip_topk_layers[local_layer_idx] + @property def index_k_with_scale_buffer(self): # Preserve direct HiCache access while storage lives behind the facade. return self.index_key_cache.buffer + def _init_kpool_compress_tail_buffers( + self, + index_kpool: int, + index_kpool_compress: bool, + tail_extra_slots: int, + index_head_dim: int, + layer_num: int, + device: str, + max_running_requests: Optional[int], + ) -> None: + """Keep request tails on the pool so they follow the index-cache lifecycle.""" + self.kpool_use_compress = index_kpool > 1 and index_kpool_compress + + if not self.kpool_use_compress: + self._compress_tail_k = None + self._compress_tail_score = None + return + + assert max_running_requests is not None, ( + "DSATokenToKVPool with kpool compress requires max_running_requests" + ) + # +1 mirrors req_to_token_pool.size + 1 used by the indexer to + # provide an extra slot for invalid / sentinel req indices. + req_pool_size = max_running_requests + 1 + tail_dtype = torch.bfloat16 + tail_width = index_kpool + tail_extra_slots + with ( + torch.cuda.use_mem_pool(self.custom_mem_pool) + if self.custom_mem_pool + else nullcontext() + ): + self._compress_tail_k: Optional[List[torch.Tensor]] = [ + torch.zeros( + req_pool_size if self._should_allocate_index_layer(i) else 0, + tail_width, + index_head_dim, + dtype=tail_dtype, + device=device, + ) + for i in range(layer_num) + ] + self._compress_tail_score: Optional[List[torch.Tensor]] = [ + torch.zeros( + req_pool_size if self._should_allocate_index_layer(i) else 0, + tail_width, + index_head_dim, + dtype=tail_dtype, + device=device, + ) + for i in range(layer_num) + ] + + def get_compress_tail_buffers( + self, layer_id: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert self.kpool_use_compress, ( + "get_compress_tail_buffers called when kpool compress is disabled" + ) + idx = layer_id - self.start_layer + return ( + self._compress_tail_k[idx], + self._compress_tail_score[idx], + ) + + def get_compress_tail_buf_infos(self): + if not self.kpool_use_compress: + return [], [], [] + transfer_layer_ids = list(range(self.layer_num)) + # Keep zero-row indexShare entries in the pointer list so layer offsets + # stay aligned across PD peers; item_len=0 makes transfer backends skip them. + tail_buffers = [self._compress_tail_k[i] for i in transfer_layer_ids] + [ + self._compress_tail_score[i] for i in transfer_layer_ids + ] + data_ptrs = [buf.data_ptr() for buf in tail_buffers] + data_lens = [buf.nbytes for buf in tail_buffers] + item_lens = [buf[0].nbytes if buf.shape[0] > 0 else 0 for buf in tail_buffers] + return data_ptrs, data_lens, item_lens + + def kpool_decode_update_index_cache( + self, + layer_id: int, + key: torch.Tensor, + slot_score: torch.Tensor, + ape: torch.Tensor, + block_tables: torch.Tensor, + req_pool_indices: torch.Tensor, + positions: torch.Tensor, + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, + round_scale: bool = False, + ) -> None: + from sglang.srt.layers.attention.dsa.kpool_fp8_index import ( + kpool_decode_update_and_maybe_write_cache, + ) + + assert self.kpool_use_compress, ( + "kpool_decode_update_index_cache called when kpool compress is disabled" + ) + idx = layer_id - self.start_layer + buf = self.get_index_k_with_scale_buffer(layer_id) + kpool_decode_update_and_maybe_write_cache( + pool=self, + buf=buf, + tail_k=self._compress_tail_k[idx], + tail_score=self._compress_tail_score[idx], + key=key, + slot_score=slot_score, + ape=ape, + block_tables=block_tables, + req_pool_indices=req_pool_indices, + positions=positions, + seq_lens=seq_lens, + out_cache_loc=out_cache_loc, + round_scale=round_scale, + ) + + def set_compress_tail_for_request( + self, + layer_id: int, + req_pool_idx: torch.Tensor, + key_tail: torch.Tensor, + score_tail: torch.Tensor, + n_remain: int, + dst_logical_start: int, + ) -> None: + """Leave the ring untouched at a pool boundary; no tail carries over.""" + assert self.kpool_use_compress, ( + "set_compress_tail_for_request called when kpool compress is disabled" + ) + idx = layer_id - self.start_layer + if n_remain > 0: + slots = ( + torch.arange(n_remain, device=key_tail.device, dtype=torch.long) + + int(dst_logical_start) + ) % self._compress_tail_k[idx].shape[1] + self._compress_tail_k[idx][req_pool_idx, slots] = key_tail + self._compress_tail_score[idx][req_pool_idx, slots] = score_tail + def _clear_buffers(self): super()._clear_buffers() self.index_key_cache.clear() + if hasattr(self, "_compress_tail_k") and self._compress_tail_k is not None: + del self._compress_tail_k + del self._compress_tail_score def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): """Move latent KV and the DSA indexer cache (key + scale) in lockstep.""" @@ -4532,15 +4939,78 @@ class DSATokenToKVPool(MLATokenToKVPool): ) -> None: self.index_key_cache.store_quantized(layer_id, loc, index_k, index_k_scale) - def get_cpu_copy(self, indices, mamba_indices=None): - kv_cache_cpu = super().get_cpu_copy(indices, mamba_indices=mamba_indices) - return {"kv": kv_cache_cpu, "index_k": self.index_key_cache.cpu_copy(indices)} + def _get_compress_tail_cpu_copy(self, req_pool_index): + if not self.kpool_use_compress or req_pool_index is None: + return None - def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None): + tail_k_cpu = [] + tail_score_cpu = [] + for tail_k, tail_score in zip(self._compress_tail_k, self._compress_tail_score): + if tail_k.shape[0] == 0: + tail_k_cpu.append(None) + tail_score_cpu.append(None) + continue + tail_k_cpu.append(tail_k[req_pool_index].to("cpu", non_blocking=True)) + tail_score_cpu.append( + tail_score[req_pool_index].to("cpu", non_blocking=True) + ) + return tail_k_cpu, tail_score_cpu + + def _load_compress_tail_cpu_copy(self, tail_k_cpu, tail_score_cpu, req_pool_index): + if ( + not self.kpool_use_compress + or req_pool_index is None + or tail_k_cpu is None + or tail_score_cpu is None + ): + return + + for tail_k, tail_score, saved_k, saved_score in zip( + self._compress_tail_k, + self._compress_tail_score, + tail_k_cpu, + tail_score_cpu, + ): + if tail_k.shape[0] == 0 or saved_k is None or saved_score is None: + continue + tail_k[req_pool_index] = saved_k.to(tail_k.device, non_blocking=True) + tail_score[req_pool_index] = saved_score.to( + tail_score.device, non_blocking=True + ) + + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): + # Retraction reuses index-cache pages; offload index/scale with KV so resume cannot read another request's entries. + kv_cache_cpu = super().get_cpu_copy(indices, mamba_indices=mamba_indices) + cpu_copy = { + "kv": kv_cache_cpu, + "index_k": self.index_key_cache.cpu_copy(indices), + } + compress_tail = self._get_compress_tail_cpu_copy(req_pool_index) + if compress_tail is not None: + cpu_copy["tail_k"], cpu_copy["tail_score"] = compress_tail + torch.cuda.synchronize() + return cpu_copy + + def load_cpu_copy( + self, + kv_cache_cpu_dict, + indices, + mamba_indices=None, + req_pool_index=None, + ): super().load_cpu_copy( - kv_cache_cpu_dict["kv"], indices, mamba_indices=mamba_indices + kv_cache_cpu_dict["kv"], + indices, + mamba_indices=mamba_indices, + req_pool_index=req_pool_index, ) self.index_key_cache.load_cpu_copy(kv_cache_cpu_dict["index_k"], indices) + self._load_compress_tail_cpu_copy( + kv_cache_cpu_dict.get("tail_k"), + kv_cache_cpu_dict.get("tail_score"), + req_pool_index, + ) + torch.cuda.synchronize() def get_state_buf_infos(self): return self.index_key_cache.state_buf_infos() diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index fb05c77cd..120616847 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -366,10 +366,12 @@ class SWAKVPool(BaseSWAKVPool): filtered.append(filtered_layer) return filtered - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): # For SWA, we need to copy KV cache from both full and SWA pools # The indices are for the full pool, and we use mapping to get SWA indices - full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices) + full_kv_cpu = self.full_kv_pool.get_cpu_copy( + indices, req_pool_index=req_pool_index + ) swa_mask = None if self.full_to_swa_index_mapping is not None: @@ -388,14 +390,18 @@ class SWAKVPool(BaseSWAKVPool): return {"full": full_kv_cpu, "swa": swa_kv_cpu, "swa_mask": swa_mask} - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): # Load KV cache back from CPU to both full and SWA pools # Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices full_kv_cpu = kv_cache_cpu["full"] swa_kv_cpu = kv_cache_cpu["swa"] # Load full KV cache to the new indices - self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices) + self.full_kv_pool.load_cpu_copy( + full_kv_cpu, indices, req_pool_index=req_pool_index + ) # Load SWA KV cache if it exists if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None: diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 9683330d8..0b7870046 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1621,23 +1621,29 @@ class UnifiedSWAKVPool(SWAKVPool): phys_pages = allocator.virtual_to_physical[virt_pages] return phys_pages * ps + offsets - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): assert self._full_allocator is not None assert self._swa_allocator is not None # `indices` are virtual TOKEN ids; translate per sub-pool. full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator) swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator) - full_cpu = self.full_kv_pool.get_cpu_copy(full_phys) + full_cpu = self.full_kv_pool.get_cpu_copy( + full_phys, req_pool_index=req_pool_index + ) valid = swa_phys >= 0 swa_cpu = None if bool(valid.any().item()): swa_cpu = self.swa_kv_pool.get_cpu_copy(swa_phys[valid]) return {"full": full_cpu, "swa": swa_cpu} - def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): + def load_cpu_copy( + self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None + ): assert self._full_allocator is not None full_phys = self._virt_tokens_to_phys_tokens(indices, self._full_allocator) - self.full_kv_pool.load_cpu_copy(kv_cache_cpu["full"], full_phys) + self.full_kv_pool.load_cpu_copy( + kv_cache_cpu["full"], full_phys, req_pool_index=req_pool_index + ) if kv_cache_cpu.get("swa") is not None: assert self._swa_allocator is not None swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index c1b8e1475..99a9c4096 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -656,7 +656,11 @@ def build_decode_registry( # init_new -- they leave the GLOBAL None and set the replicated LOCAL # count directly, so carry that through. if fb.global_num_token_non_padded is None: - buf.copy_(fb.num_token_non_padded) + # DFLASH's dense draft can omit both optional counts, even + # when EP on the target enables this slot. Preserve the + # registry's skip-missing-field behavior for that path. + if fb.num_token_non_padded is not None: + buf.copy_(fb.num_token_non_padded) return sharded = not enable_prefill_cp and attn_tp_sharded_fn( ctx.padded_num_tokens diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index f880e93a0..5872abf93 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -56,6 +56,7 @@ from sglang.srt.runtime_context import ( get_lora, get_parallel, ) +from sglang.srt.speculative.spec_info import SpecInputType from sglang.srt.utils import ( is_cpu, is_cuda, @@ -1739,6 +1740,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): logits_output.hidden_states = logits_output.hidden_states[ :num_tokens ] + elif ( + self.spec_info.spec_input_type == SpecInputType.EAGLE_DRAFT_EXTEND + and not self.forward_mode.is_draft_extend_v2() + ): + if self.spec_info.num_correct_drafts is not None: + self.spec_info.num_correct_drafts = ( + self.spec_info.num_correct_drafts[:bs] + ) + if self.spec_info.num_accept_tokens is not None: + self.spec_info.num_accept_tokens = self.spec_info.num_accept_tokens[ + :bs + ] + if self.extend_seq_lens is not None: + self.extend_seq_lens = self.extend_seq_lens[:bs] + if logits_output.next_token_logits is not None: + logits_output.next_token_logits = logits_output.next_token_logits[ + :bs + ] + if logits_output.hidden_states is not None: + logits_output.hidden_states = logits_output.hidden_states[:bs] elif self.forward_mode.is_draft_extend_v2(): # draft extend_v2 bs = bs * self.spec_info.num_tokens_per_req if logits_output.next_token_logits is not None: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 353f30878..ba6743da1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1741,7 +1741,7 @@ class ModelRunner: ) else: # mamba_pool is a pure PHYSICAL store; translate both COW slot ids. - pool.mamba_pool.copy_from( + pool.copy_mamba_state( pool.translate_mamba_indices(forward_batch.mamba_cow_src_indices), pool.translate_mamba_indices(forward_batch.mamba_cow_dst_indices), ) diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index aee5f1821..c8eba8c71 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -53,6 +53,7 @@ if TYPE_CHECKING: from sglang.srt.model_executor.runner.base_runner import BaseRunner logger = logging.getLogger(__name__) +_deep_gemm_layout_memory_budget_initialized = False def _align_pipeline_layers(layers: list, layer_model) -> list: @@ -154,6 +155,75 @@ class CudaGraphsCapture(msgspec.Struct, frozen=True, kw_only=True): ) +def refresh_deep_gemm_layout_memory_budget( + model_runner: ModelRunner, *, only_if_initialized: bool = False +) -> None: + """Set the all-rank budget before capture, then refresh after startup.""" + global _deep_gemm_layout_memory_budget_initialized + if ( + model_runner.device != "cuda" + or envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() != "auto" + ): + return + + if only_if_initialized: + # Target and draft share the budget. Its pre-capture initialization + # already used a world-wide collective, so this guard is rank-uniform + # and also covers a draft-only DeepGEMM backend outside draft context. + if not _deep_gemm_layout_memory_budget_initialized: + return + else: + if model_runner.is_draft_worker: + moe_runner_backend = ( + get_spec().speculative_moe_runner_backend + or get_exec().moe.moe_runner_backend + ) + moe_a2a_backend = ( + get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend + ) + else: + moe_runner_backend = get_exec().moe.moe_runner_backend + moe_a2a_backend = get_exec().moe.moe_a2a_backend + + uses_deep_gemm_moe_runner = moe_runner_backend == "deep_gemm" + if moe_runner_backend == "auto" and model_runner.model_config.quantization in ( + "fp8", + "mxfp8", + ): + from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend + from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod + + uses_deep_gemm_moe_runner = ( + Fp8MoEMethod.is_deepgemm_moe_runner_backend_enabled( + MoeRunnerBackend(moe_runner_backend), + MoeA2ABackend(moe_a2a_backend), + ) + ) + if not uses_deep_gemm_moe_runner: + return + + from sglang.srt.layers.moe.moe_runner.deep_gemm import ( + set_masked_standard_layout_memory_budget, + ) + + world_group = get_world_group() + available_memory_gb = get_available_gpu_memory( + model_runner.device, + model_runner.gpu_id, + distributed=world_group.world_size > 1, + cpu_group=world_group.cpu_group, + ) + budget_bytes = set_masked_standard_layout_memory_budget( + int(available_memory_gb * (1 << 30)) + ) + _deep_gemm_layout_memory_budget_initialized = True + logger.info( + "DeepGEMM masked layout budget: %.2f GiB from %.2f GiB free.", + budget_bytes / (1 << 30), + available_memory_gb, + ) + + def capture_cuda_graphs( *, model_runner: ModelRunner, capture_decode_cuda_graph: bool = True ) -> CudaGraphsCapture: @@ -176,56 +246,7 @@ def capture_cuda_graphs( # runners point at it) and the eager fallback when a cg runner can't run a # batch. eager_runner = EagerRunner(model_runner) - - if model_runner.is_draft_worker: - moe_runner_backend = ( - get_spec().speculative_moe_runner_backend - or get_exec().moe.moe_runner_backend - ) - moe_a2a_backend = ( - get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend - ) - else: - moe_runner_backend = get_exec().moe.moe_runner_backend - moe_a2a_backend = get_exec().moe.moe_a2a_backend - - uses_deep_gemm_moe_runner = moe_runner_backend == "deep_gemm" - if moe_runner_backend == "auto" and model_runner.model_config.quantization in ( - "fp8", - "mxfp8", - ): - from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend - from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod - - uses_deep_gemm_moe_runner = Fp8MoEMethod.is_deepgemm_moe_runner_backend_enabled( - MoeRunnerBackend(moe_runner_backend), - MoeA2ABackend(moe_a2a_backend), - ) - - if ( - model_runner.device == "cuda" - and envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower() == "auto" - and uses_deep_gemm_moe_runner - ): - from sglang.srt.layers.moe.moe_runner.deep_gemm import ( - set_masked_standard_layout_memory_budget, - ) - - world_group = get_world_group() - available_memory_gb = get_available_gpu_memory( - model_runner.device, - model_runner.gpu_id, - distributed=world_group.world_size > 1, - cpu_group=world_group.cpu_group, - ) - budget_bytes = set_masked_standard_layout_memory_budget( - int(available_memory_gb * (1 << 30)) - ) - logger.info( - "DeepGEMM masked layout budget: %.2f GiB from %.2f GiB free.", - budget_bytes / (1 << 30), - available_memory_gb, - ) + refresh_deep_gemm_layout_memory_budget(model_runner) # cuda-graph capture: prefill before decode, so both coalesce onto the # eager buffer allocated above. (capture_prefill_graph routes prefill diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py index 76a4a0386..950761105 100644 --- a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py +++ b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py @@ -57,6 +57,7 @@ def resolve_spec_aux_hidden_state_config( _resolve_eagle_aux_hidden_state( config=config, server_args=server_args, + model_config=model_config, spec_algorithm=spec_algorithm, is_draft_worker=is_draft_worker, ) @@ -74,55 +75,54 @@ def _resolve_eagle_aux_hidden_state( *, config: SpecAuxHiddenStateConfig, server_args: ServerArgs, + model_config: ModelConfig, spec_algorithm: SpeculativeAlgorithm, is_draft_worker: bool, ) -> None: - if ( + if not ( (spec_algorithm.is_eagle() or spec_algorithm.is_standalone()) and not is_draft_worker - and get_spec().speculative_draft_model_path ): - # Load draft config to get layer count for KV cache sizing + return + + draft_model_config = model_config + if get_spec().speculative_draft_model_path: draft_model_config = ModelConfig.from_server_args( server_args, model_path=get_spec().speculative_draft_model_path, model_revision=get_spec().speculative_draft_model_revision, is_draft_model=True, ) - num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers - if num_nextn_predict_layers is not None: - config.eagle_draft_num_layers = int(num_nextn_predict_layers) - else: - config.eagle_draft_num_layers = int( - max( - draft_model_config.num_hidden_layers, - draft_model_config.num_attention_layers, - ) + num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers + if num_nextn_predict_layers is not None: + config.eagle_draft_num_layers = int(num_nextn_predict_layers) + elif get_spec().speculative_draft_model_path: + config.eagle_draft_num_layers = int( + max( + draft_model_config.num_hidden_layers, + draft_model_config.num_attention_layers, ) + ) + else: + return - if ( - draft_model_config.is_hybrid_swa - and not draft_model_config.is_deepseek_v4_arch - ): - config.eagle_draft_swa_num_layers = len( - draft_model_config.swa_attention_layer_ids + if draft_model_config.is_hybrid_swa and not draft_model_config.is_deepseek_v4_arch: + config.eagle_draft_swa_num_layers = len( + draft_model_config.swa_attention_layer_ids + ) + + if spec_algorithm.is_eagle3(): + config.eagle_use_aux_hidden_state = True + try: + eagle_config = getattr(draft_model_config.hf_config, "eagle_config", None) + config.eagle_use_aux_hidden_state = eagle_config.get( + "use_aux_hidden_state", True ) - - if spec_algorithm.is_eagle3(): - config.eagle_use_aux_hidden_state = True - try: - eagle_config = getattr( - draft_model_config.hf_config, "eagle_config", None - ) - config.eagle_use_aux_hidden_state = eagle_config.get( - "use_aux_hidden_state", True - ) - config.eagle_aux_hidden_state_layer_ids = eagle_config[ - "eagle_aux_hidden_state_layer_ids" - ] - except: - # if there is no aux layer, set to None - config.eagle_aux_hidden_state_layer_ids = None + config.eagle_aux_hidden_state_layer_ids = eagle_config[ + "eagle_aux_hidden_state_layer_ids" + ] + except Exception: + config.eagle_aux_hidden_state_layer_ids = None def _resolve_dflash_aux_hidden_state( diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 3a5e7d3da..df8258da3 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -114,6 +114,23 @@ def _dflash_draft_cell_size(kvc: KVCacheConfigurator) -> int: return int(cell_size) * get_parallel().attn_dcp_size +def _get_dsa_cache_layer_ids(kvc: KVCacheConfigurator, num_layers: int) -> list[int]: + """Global layer ids represented by the local DSA pool's dense layer slots.""" + if kvc.mambaish_config and not kvc.is_draft_worker: + layer_ids = [ + layer_id + for layer_id in kvc.mambaish_config.full_attention_layer_ids + if kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer + ] + else: + layer_ids = list(range(kvc.layer_info.start_layer, kvc.layer_info.end_layer)) + # Draft pools and a few platform-specific pools may expose a synthetic layer + # count. They do not use indexShare, so only the length matters for sizing. + if len(layer_ids) != num_layers: + return list(range(num_layers)) + return layer_ids + + def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]: dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower() if dtype_name in ("float32", "fp32"): @@ -280,8 +297,10 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): get_glm_dsa_layer_split_effective_num_layers, ) - effective_num_layers = get_glm_dsa_layer_split_effective_num_layers( - kvc, num_layers + effective_num_layers = ( + num_layers + if kvc.server_args.enable_hisparse + else get_glm_dsa_layer_split_effective_num_layers(kvc, num_layers) ) kv_size = torch._utils._element_size(kv_cache_dtype) @@ -412,18 +431,13 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): _should_elide_dsa_index_k, ) - if allocate_all_layers or not _should_elide_dsa_index_k( - is_draft_worker=kvc.is_draft_worker + if ( + allocate_all_layers + or kvc.server_args.enable_hisparse + or not _should_elide_dsa_index_k(is_draft_worker=kvc.is_draft_worker) ): num_indexer_layers = num_layers else: - active_indexer_layers = [ - layer_id - for layer_id in range( - kvc.layer_info.start_layer, kvc.layer_info.end_layer - ) - if not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id) - ] from sglang.srt.layers.cp.utils import ( get_glm_dsa_cp_layer_shard_info, get_layer_shard_range, @@ -431,6 +445,16 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): _, shard_size = get_glm_dsa_cp_layer_shard_info(kvc) if shard_size > 1: + # Preserve the existing LayerSplit sizing semantics. GLM-5.3 + # hybrid-layer support is intentionally limited to the normal + # (non-LayerSplit) pool below. + active_indexer_layers = [ + layer_id + for layer_id in range( + kvc.layer_info.start_layer, kvc.layer_info.end_layer + ) + if not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id) + ] active_set = set(active_indexer_layers) max_owned = 0 for rank in range(shard_size): @@ -444,7 +468,10 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): ) num_indexer_layers = max_owned + 1 else: - num_indexer_layers = len(active_indexer_layers) + num_indexer_layers = sum( + not dsa_layer_skips_topk(kvc.model_config.hf_config, layer_id) + for layer_id in _get_dsa_cache_layer_ids(kvc, num_layers) + ) return int( indexer_size_per_token * num_indexer_layers * element_size * indexer_ratio diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index e1d746743..c934e0f58 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1068,14 +1068,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): and stash the returned per-bucket metadata object; otherwise fall back to the generic eager init that BCG/TC_PIECEWISE use today.""" attn_backend = self.model_runner.attn_backend - if not self.use_captured_attn_metadata: - attn_backend.init_forward_metadata(forward_batch) - return - metadata = attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture( - forward_batch - ) - assert self.attn_metadata_buffers is not None - self.attn_metadata_buffers[num_tokens] = metadata + with forward_context(ForwardContext(attn_backend=attn_backend)): + if not self.use_captured_attn_metadata: + attn_backend.init_forward_metadata(forward_batch) + return + metadata = ( + attn_backend.init_forward_metadata_for_breakable_cuda_graph_capture( + forward_batch + ) + ) + assert self.attn_metadata_buffers is not None + self.attn_metadata_buffers[num_tokens] = metadata def _prepare_forward_metadata_for_replay( self, diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index f5f2797fe..0f4847ca2 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py @@ -154,16 +154,18 @@ def _uninstall_wait_stream_hook(): def _weak_ref_if_tensor(x): - """Return a weak-ref tensor view (shared storage, no refcount) for tensors; - recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing - captured args lets the shared mempool reclaim per-layer intermediates - between segments — storage stays alive for each segment CUDAGraph's - lifetime via its pool use_count. + """Return a weak-ref view for nonempty accelerator tensors; recurse into + tuples/lists and keep CPU, empty, and non-tensor values unchanged. + Weak-ref'ing captured args lets the shared mempool reclaim per-layer + intermediates between segments — storage stays alive for each segment + CUDAGraph's lifetime via its pool use_count. weak_ref_tensors is imported lazily because it hard-raises on platforms without a CUDA/HIP/NPU backend; we only reach this code during an active Breakable capture, which runs only on those backends.""" if torch.is_tensor(x): + if x.numel() == 0 or x.device.type == "cpu": + return x from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors return weak_ref_tensors(x) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 5fd3f4a40..13284e46d 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -259,7 +259,11 @@ def _get_quantization_config( f"method {model_config.quantization}. Supported dtypes: " f"{supported_dtypes}" ) - hf_to_sglang_mapper = getattr(model_class, "hf_to_sglang_mapper", None) + get_hf_to_sglang_mapper = getattr(model_class, "get_hf_to_sglang_mapper", None) + if get_hf_to_sglang_mapper is not None: + hf_to_sglang_mapper = get_hf_to_sglang_mapper(model_config.hf_config) + else: + hf_to_sglang_mapper = getattr(model_class, "hf_to_sglang_mapper", None) # pass mappings by reference to quant_config if hf_to_sglang_mapper is not None and quant_config is not None: quant_config.apply_weight_name_mapper(hf_to_sglang_mapper) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 8b907de1a..98da6a4cb 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -1571,11 +1571,18 @@ def row_parallel_weight_loader( LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] -def sharded_weight_loader(shard_axis: int) -> LoaderFunction: +def sharded_weight_loader( + shard_axis: int, + tp_rank_getter=None, +) -> LoaderFunction: """Create a weight loader that shards the weights along the given axis""" def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: - tp_rank = get_parallel().attn_tp_rank + tp_rank = ( + tp_rank_getter() + if tp_rank_getter is not None + else get_parallel().attn_tp_rank + ) shard_size = param.data.shape[shard_axis] start_idx = tp_rank * shard_size diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index e0004378d..821455aee 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -118,9 +118,9 @@ def _handle_attention_backend(attn, forward_batch, backend_name): return _dispatch_mla_subtype(attn, forward_batch) sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch) - disable_ragged = ( - backend_name in ["flashinfer", "flashmla"] - ) and attn.flashinfer_mla_disable_ragged + disable_ragged = (backend_name in ["flashinfer", "flashmla"]) and ( + attn.flashinfer_mla_disable_ragged or attn.qk_rope_head_dim == 0 + ) if ( not disable_ragged diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 2b19c0185..13b18488b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -677,9 +677,22 @@ class DeepseekMHAForwardMixin: def _concat_and_cast_mha_k( self: DeepseekV2AttentionMLA, k_nope: torch.Tensor, - k_pe: torch.Tensor, + k_pe: torch.Tensor | None, forward_batch: ForwardBatch, ): + if self.qk_rope_head_dim == 0: + assert k_pe is None or k_pe.shape[-1] == 0 + k = k_nope.contiguous() + if ( + _is_cuda + and self.current_attention_backend == "fa3" + and self.kv_cache_dtype != "auto" + ): + # fa3 requires k in the pool dtype when KV cache is fp8; the + # concat branch below does the same cast for roped models. + k = k.to(get_token_to_kv_pool().dtype) + return k + # Temporary for DeepSeek V3/R1 only, but can generalize if needed k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim) if ( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index e535ee386..a5e89882b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -960,6 +960,8 @@ class DeepseekMLAForwardMixin: """ Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path. """ + if self.rotary_emb is None: + return False if self.current_attention_backend in ("dsa", "nsa"): return ( get_exec().kernel.dsa_decode_backend == "trtllm" diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index cd916cf60..b6c830fea 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -570,7 +570,7 @@ class DeepseekV2WeightLoaderMixin: for name in weight_names: if "kv_b_proj" in name: layer_id = int(name.split(".")[2]) - if layer_id < self.config.num_hidden_layers: + if self.model.start_layer <= layer_id < self.model.end_layer: layer_ids.add(layer_id) for layer_id in layer_ids: @@ -580,6 +580,9 @@ class DeepseekV2WeightLoaderMixin: else self.model.decoder.self_attn ) + if not hasattr(self_attn, "kv_b_proj"): + continue + if hasattr(self_attn.kv_b_proj, "qweight"): # awq compatible, dequantize the weight if supported awq_dequantize_f = awq_dequantize_func() diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index ad16178c1..b3ca28468 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -182,6 +182,7 @@ class DeepseekModelNextN(nn.Module): is_nextn=True, prefix=add_prefix(layer_name, prefix), alt_stream=self.alt_stream, + skip_rope=config.qk_rope_head_dim == 0, dsa_enable_prefill_cp=self.dsa_enable_prefill_cp, mla_enable_prefill_cp=self.mla_enable_prefill_cp, ) @@ -220,9 +221,28 @@ class DeepseekModelNextN(nn.Module): ) if input_embeds is None: - hidden_states = self.embed_tokens(input_ids) - else: - hidden_states = input_embeds + # MM positions in input_ids hold MM_PAD_SHIFT_VALUE+hash sentinels + # (far above vocab_size). Use target-produced mm_input_embeds for + # these positions and only call embed_tokens on the appended + # next-token to avoid embed OOB. + input_embeds = forward_batch.mm_input_embeds + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.contains_mm_inputs() + and not forward_batch.forward_mode.is_draft_extend_v2() + ): + assert input_embeds is not None + last_indices = ( + forward_batch.extend_start_loc + + forward_batch.extend_seq_lens + - 1 + ).long() + input_embeds[last_indices] = self.embed_tokens( + input_ids[last_indices] + ) + if input_embeds is None: + input_embeds = self.embed_tokens(input_ids) + hidden_states = input_embeds if hidden_states.shape[0] > 0: previous_hidden_states = forward_batch.spec_info.hidden_states @@ -320,6 +340,10 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): }, ) + @classmethod + def get_hf_to_sglang_mapper(cls, config) -> WeightsMapper: + return cls.hf_to_sglang_mapper + def _resolve_nextn_quant_config(self, config, quant_config): if quant_config is None or quant_config.get_name() != "quark": return quant_config @@ -327,7 +351,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): from sglang.srt.layers.quantization.quark.utils import should_ignore_layer ckpt_prefix = f"model.layers.{config.num_hidden_layers}" - mapped_prefix = self.hf_to_sglang_mapper._map_name(ckpt_prefix) + mapped_prefix = self.get_hf_to_sglang_mapper(config)._map_name(ckpt_prefix) if should_ignore_layer(mapped_prefix, quant_config.exclude_layers): return None return quant_config diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 45a84a6fb..7d0a46848 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -45,6 +45,7 @@ from sglang.srt.configs.model_config import ( compute_mla_mscale_scaling, dsa_layer_skips_topk, get_dsa_index_head_dim, + get_dsa_index_kpool, get_dsa_index_n_heads, get_dsa_index_topk, is_deepseek_dsa, @@ -62,6 +63,7 @@ from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.amx_utils import PackWeightMethod from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer +from sglang.srt.layers.attention.dsa.dsa_indexer_kpool import IndexerKPool from sglang.srt.layers.attention.dsa.utils import ( can_dsa_cp_split, dsa_use_prefill_cp, @@ -1860,7 +1862,10 @@ class DeepseekV2AttentionMLA( if not self.skip_topk or is_nextn: is_neox_style = not getattr(config, "indexer_rope_interleave", False) - self.indexer = Indexer( + indexer_cls = ( + IndexerKPool if get_dsa_index_kpool(config) > 1 else Indexer + ) + indexer_kwargs = dict( hidden_size=hidden_size, index_n_heads=get_dsa_index_n_heads(config), index_head_dim=get_dsa_index_head_dim(config), @@ -1879,6 +1884,9 @@ class DeepseekV2AttentionMLA( alt_stream=alt_stream, config=config, ) + if indexer_cls is IndexerKPool: + indexer_kwargs["skip_rope"] = skip_rope + self.indexer = indexer_cls(**indexer_kwargs) self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, @@ -1902,7 +1910,7 @@ class DeepseekV2AttentionMLA( ) self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) - if not skip_rope: + if not skip_rope and qk_rope_head_dim > 0: is_neox_style = not getattr(config, "rope_interleave", True) self.rotary_emb = get_rope_wrapper( qk_rope_head_dim, @@ -2318,6 +2326,7 @@ class DeepseekV2DecoderLayer(nn.Module): is_nextn: bool = False, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + skip_rope: bool = False, dsa_enable_prefill_cp: bool = False, mla_enable_prefill_cp: bool = False, ) -> None: @@ -2340,6 +2349,10 @@ class DeepseekV2DecoderLayer(nn.Module): self.mla_enable_prefill_cp = mla_enable_prefill_cp self.layer_id = layer_id self.is_nextn = is_nextn + if is_nextn and getattr(config, "mla_nope", False): + # The NextN draft must match the NoPE target layers, or its Q/K + # and the KV it verifies against live in different spaces. + skip_rope = True self.self_attn = DeepseekV2AttentionMLA( config=config, hidden_size=self.hidden_size, @@ -2359,6 +2372,7 @@ class DeepseekV2DecoderLayer(nn.Module): reduce_results=False, prefix=add_prefix("self_attn", prefix), alt_stream=alt_stream, + skip_rope=skip_rope, is_nextn=is_nextn, dsa_enable_prefill_cp=dsa_enable_prefill_cp, mla_enable_prefill_cp=mla_enable_prefill_cp, @@ -2689,6 +2703,7 @@ class DeepseekV2Model(nn.Module): quant_config=quant_config, prefix=prefix, alt_stream=self.alt_stream, + skip_rope=config.qk_rope_head_dim == 0, dsa_enable_prefill_cp=self.dsa_enable_prefill_cp, mla_enable_prefill_cp=self.mla_enable_prefill_cp, ), diff --git a/python/sglang/srt/models/glm5_next.py b/python/sglang/srt/models/glm5_next.py new file mode 100644 index 000000000..f5253cfdf --- /dev/null +++ b/python/sglang/srt/models/glm5_next.py @@ -0,0 +1,1630 @@ +import logging +from contextlib import nullcontext +from functools import partial +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import functional as F + +from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated +from sglang.kernels.ops.layernorm.mhc import hc_contract +from sglang.kernels.ops.layernorm.mhc import hc_post as _hc_post_fn +from sglang.kernels.ops.layernorm.mhc import hc_pre as _hc_pre_fn +from sglang.srt.batch_overlap.two_batch_overlap import ( + model_forward_maybe_tbo, +) +from sglang.srt.configs.glm5_next import Glm5NextConfig, Glm5NextTextConfig +from sglang.srt.configs.model_config import is_deepseek_dsa +from sglang.srt.distributed.parallel_state import get_pp_group +from sglang.srt.distributed.utils import divide +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import ( + get_global_expert_distribution_recorder, +) +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.layers.attention import vision_utils +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, + get_attn_tp_context, +) +from sglang.srt.layers.communicator_mhc import MHCLayerCommunicator +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelBatchedLinear, + ColumnParallelLinear, + MergedColumnParallelLinear, + MergedColumnParallelRepeatedLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, + is_shared_experts_fusion_disabled, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils.common import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, + get_embedding_tp_kwargs, +) +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, +) +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + PPProxyTensors, +) +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( + DeepseekV2WeightLoaderMixin, +) +from sglang.srt.models.deepseek_common.utils import ( + _device_sm, + _is_cuda, + _use_aiter_gfx95, +) +from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA +from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Glm5NextMLP +from sglang.srt.models.deepseek_v2 import DeepseekV2MoE as Glm5NextMoE +from sglang.srt.models.glm_ocr import ( + GlmOcrRMSNorm, + GlmOcrVisionBlock, + GlmOcrVisionMLP, + GlmOcrVisionModel, + GlmOcrVisionPatchEmbed, + GlmOcrVisionPatchMerger, +) +from sglang.srt.multimodal.mm_utils import ( + run_dp_presharded_mrope_vision_model, + run_dp_sharded_mrope_vision_model, +) +from sglang.srt.runtime_context import get_forward, get_mm, get_parallel, get_spec +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils.common import ( + BumpAllocator, + LazyValue, + add_prefix, + log_info_on_rank0, + make_layers, + set_weight_attrs, +) + +if _use_aiter_gfx95: + from sglang.srt.layers.rocm_linear_utils import ( + get_dsv3_gemm_output_zero_allocator_size, + ) + +logger = logging.getLogger(__name__) + + +@torch.compile +def swiglu_clamped(y: torch.Tensor, limit: float): + gate, up = torch.chunk(y, 2, dim=-1) + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + return F.silu(gate) * up + + +class Glm5NextVisionMLP(GlmOcrVisionMLP): + def __init__( + self, + in_features: int, + hidden_features: int, + swiglu_limit: float, + bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + ): + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + bias=bias, + quant_config=quant_config, + prefix=prefix, + use_data_parallel=use_data_parallel, + ) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor): + gate_up, _ = self.gate_up_proj(x) + x = swiglu_clamped(gate_up, self.swiglu_limit) + x, _ = self.down_proj(x) + return x + + +class Glm5NextVisionPatchMerger(GlmOcrVisionPatchMerger): + def __init__( + self, + d_model: int, + context_dim: int, + swiglu_limit: float, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__( + d_model=d_model, + context_dim=context_dim, + quant_config=quant_config, + bias=bias, + prefix=prefix, + use_data_parallel=use_data_parallel, + ) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor): + x, _ = self.proj(x) + x = self.extra_activation_func(self.post_projection_norm(x)) + gate_up, _ = self.gate_up_proj(x) + x = swiglu_clamped(gate_up, self.swiglu_limit) + x, _ = self.down_proj(x) + return x + + +class Glm5NextVisionBlock(GlmOcrVisionBlock): + def __init__( + self, + dim: int, + intermediate_dim: int, + num_heads: int, + swiglu_limit: float, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + attn_qkv_bias: bool = True, + num_dummy_heads: int = 0, + rms_norm_eps: float = 1e-5, + use_data_parallel: bool = False, + ) -> None: + nn.Module.__init__(self) + self.norm1 = RMSNorm(dim, eps=rms_norm_eps) + self.norm2 = RMSNorm(dim, eps=rms_norm_eps) + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + qkv_bias=attn_qkv_bias, + proj_bias=True, + qk_normalization_by_head_size=True, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + num_dummy_heads=num_dummy_heads, + use_data_parallel=use_data_parallel, + ) + self.mlp = Glm5NextVisionMLP( + dim, + intermediate_dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + use_data_parallel=use_data_parallel, + swiglu_limit=swiglu_limit, + ) + + +class Glm5NextVisionModel(GlmOcrVisionModel): + def __init__( + self, + vision_config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + nn.Module.__init__(self) + self.hidden_size = vision_config.hidden_size + self.num_heads = vision_config.num_heads + self.patch_size = vision_config.patch_size + self.spatial_merge_size = vision_config.spatial_merge_size + self.out_hidden_size = vision_config.out_hidden_size + self.intermediate_size = vision_config.intermediate_size + self.use_data_parallel = use_data_parallel + + self.patch_embed = GlmOcrVisionPatchEmbed( + patch_size=vision_config.patch_size, + temporal_patch_size=vision_config.temporal_patch_size, + in_channels=vision_config.in_channels, + hidden_size=self.hidden_size, + ) + + head_dim = self.hidden_size // self.num_heads + self.rotary_pos_emb = get_rope( + head_size=head_dim, + rotary_dim=head_dim // 2, + max_position=8192, + base=10000.0, + is_neox_style=True, + ) + + self.blocks = nn.ModuleList( + [ + Glm5NextVisionBlock( + dim=self.hidden_size, + intermediate_dim=self.intermediate_size, + num_heads=self.num_heads, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{layer_idx}", prefix), + rms_norm_eps=vision_config.rms_norm_eps, + attn_qkv_bias=vision_config.attention_bias, + use_data_parallel=use_data_parallel, + swiglu_limit=vision_config.swiglu_limit, + ) + for layer_idx in range(vision_config.depth) + ] + ) + projection_intermediate_size = getattr( + vision_config, "projection_intermediate_size", None + ) + self.merger = Glm5NextVisionPatchMerger( + d_model=vision_config.out_hidden_size, + context_dim=( + projection_intermediate_size + if projection_intermediate_size is not None + else vision_config.intermediate_size + ), + quant_config=quant_config, + bias=False, + prefix=add_prefix("merger", prefix), + use_data_parallel=use_data_parallel, + swiglu_limit=vision_config.swiglu_limit, + ) + + self.downsample = nn.Conv2d( + in_channels=vision_config.hidden_size, + out_channels=vision_config.out_hidden_size, + kernel_size=vision_config.spatial_merge_size, + stride=vision_config.spatial_merge_size, + ) + self.post_layernorm = GlmOcrRMSNorm( + vision_config.hidden_size, eps=vision_config.rms_norm_eps + ) + + +class Glm5NextLinearAttention(nn.Module): + def __init__( + self, + layer_idx: int, + hidden_size: int, + config: Glm5NextTextConfig, + quant_config: Optional[QuantizationConfig] = None, + rms_norm_eps: float = 1e-5, + prefix: str = "", + reduce_results: bool = False, + **kwargs, + ) -> None: + super().__init__() + self.tp_size = get_parallel().tp_size + head_shard_size = get_parallel().attn_tp_size + head_shard_rank = get_parallel().attn_tp_rank + _head_shard_rank_getter = partial(getattr, get_parallel(), "attn_tp_rank") + + self.hidden_size = hidden_size + self.config = config + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.num_k_heads = config.linear_attn_config["num_heads"] + self.num_v_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = config.linear_attn_config["head_dim"] + self.head_v_dim = config.linear_attn_config["head_dim"] + self.layer_idx = layer_idx + self.prefix = prefix + assert self.num_heads % head_shard_size == 0 + self.local_num_heads = divide(self.num_heads, head_shard_size) + + projection_size = self.head_dim * self.num_heads + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + + self.do_fuse_qkvbfg = quant_config is None and head_shard_size == self.tp_size + if self.do_fuse_qkvbfg: + self.qkvb_sizes = [ + projection_size, + projection_size, + projection_size, + self.num_heads, + ] + self.fg_sizes = [self.head_dim, self.head_dim] + + self.fused_qkvbfg_a_proj = MergedColumnParallelRepeatedLinear( + self.hidden_size, + self.qkvb_sizes, + self.fg_sizes, + quant_config=quant_config, + prefix=f"{prefix}.fused_qkvbfg_a_proj", + ) + self.split_sizes = [ + 3 * projection_size // head_shard_size, + self.num_heads // head_shard_size, + 2 * self.head_dim, + ] + fused_dtype = ( + getattr(config, "dtype", None) + or getattr(config, "torch_dtype", None) + or torch.get_default_dtype() + ) + self.fused_fg_b_proj = ColumnParallelBatchedLinear( + 2, self.head_dim, projection_size, dtype=fused_dtype + ) + else: + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.num_heads, + self.num_k_heads, + bias=False, + quant_config=quant_config, + tp_rank=head_shard_rank, + tp_size=head_shard_size, + prefix=f"{prefix}.qkv_proj", + ) + + self.f_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_a_proj", + ) + + self.f_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.f_b_proj", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) + + self.b_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.b_proj", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) + + self.g_a_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_a_proj", + ) + self.g_b_proj = ColumnParallelLinear( + self.head_dim, + projection_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.g_b_proj", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) + + self.dt_bias = nn.Parameter( + torch.empty(divide(projection_size, head_shard_size), dtype=torch.float32) + ) + + set_weight_attrs( + self.dt_bias, + {"weight_loader": sharded_weight_loader(0, _head_shard_rank_getter)}, + ) + + self.qkv_conv1d = MergedColumnParallelLinear( + input_size=self.conv_size, + output_sizes=[projection_size, projection_size, projection_size], + bias=False, + params_dtype=torch.float32, + prefix=f"{prefix}.qkv_conv1d", + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) + # ColumnParallelLinear's loader cannot reshape conv1d weights, so add the + # singleton dimension after construction. + self.qkv_conv1d.weight.data = self.qkv_conv1d.weight.data.unsqueeze(1) + + self.A_log = nn.Parameter( + torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) + ) + set_weight_attrs( + self.A_log, + {"weight_loader": sharded_weight_loader(2, _head_shard_rank_getter)}, + ) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=rms_norm_eps, activation="sigmoid" + ) + self.o_proj = RowParallelLinear( + projection_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + reduce_results=reduce_results, + tp_rank=head_shard_rank, + tp_size=head_shard_size, + ) + + conv_weights = self.qkv_conv1d.weight.squeeze(1) + bias = self.qkv_conv1d.bias + + self.attn = RadixLinearAttention( + layer_id=self.layer_idx, + num_q_heads=self.local_num_heads, + num_k_heads=self.local_num_heads, + num_v_heads=self.local_num_heads, + head_q_dim=self.head_k_dim, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_weights=conv_weights, + bias=bias, + A_log=self.A_log, + dt_bias=self.dt_bias, + ) + + self.attn.lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + + def forward_qkvbfg(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + qkv, _ = self.qkv_proj(hidden_states) + + beta = self.b_proj(hidden_states)[0] + forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0] + g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0] + + return ( + qkv, + beta, + forget_gate, + g_proj_states, + ) + + def forward_qkvbfg_fused( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ): + fused_states = self.fused_qkvbfg_a_proj(hidden_states) + + qkv, beta, fg_a_states = torch.split(fused_states, self.split_sizes, dim=-1) + + forget_gate, g_proj_states = self.fused_fg_b_proj( + fg_a_states.view(-1, 2, self.head_dim).transpose(0, 1) + ) + + return ( + qkv, + beta, + forget_gate, + g_proj_states, + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + if forward_batch.forward_mode.is_idle(): + return hidden_states + + if self.do_fuse_qkvbfg: + mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg_fused( + hidden_states, forward_batch + ) + else: + mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg( + hidden_states, forward_batch + ) + + if not forward_batch.forward_mode.is_decode(): + forget_gate = forget_gate.unsqueeze(0) + beta = beta.unsqueeze(0) + + core_attn_out = self.attn( + forward_batch, + mixed_qkv=mixed_qkv, + a=forget_gate, + b=beta, + ) + + norm_gate = g_proj_states.unflatten(-1, (-1, self.head_dim)) + core_attn_out = self.o_norm(core_attn_out, norm_gate) + core_attn_out = core_attn_out.squeeze(0).flatten(-2) + + return self.o_proj(core_attn_out)[0] + + +class Glm5NextDecoderLayer(nn.Module): + def __init__( + self, + config: Glm5NextTextConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + moe_quant_config_override: Optional[QuantizationConfig] = None, + is_nextn: bool = False, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + rope_theta = config.rope_theta + rope_scaling = config.rope_scaling + max_position_embeddings = config.max_position_embeddings + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + get_spec().speculative_algorithm + ) + self.layer_id = layer_id + self.is_nextn = is_nextn + self.is_linear_attn = config.is_kda_layer(layer_id) + + if self.is_linear_attn: + self.self_attn = Glm5NextLinearAttention( + layer_idx=layer_id, + hidden_size=config.hidden_size, + config=config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + rms_norm_eps=config.rms_norm_eps, + reduce_results=False, + ) + else: + self.self_attn = DeepseekV2AttentionMLA( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + layer_id=layer_id, + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + is_nextn=is_nextn, + skip_rope=True, + ) + + if config.q_lora_rank is None and envs.SGLANG_USE_AG_AFTER_QLORA.get(): + raise ValueError( + "SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank" + ) + + self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) + is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False) + is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=1 if is_nextn else config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + + if self.is_layer_sparse: + self.mlp = Glm5NextMoE( + config=config, + quant_config=moe_quant_config_override or quant_config, + prefix=add_prefix("mlp", prefix), + layer_id=self.layer_id, + alt_stream=alt_stream, + is_nextn=is_nextn, + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = Glm5NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + swiglu_limit=config.swiglu_limit, + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + if self.config.mhc: + hc_mult = config.hc_mult + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * config.hidden_size + + # Keep mHC params directly on the decoder layer so their names match + # the checkpoint verbatim. + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self.hc_attn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32) + ) + + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self.hc_ffn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32) + ) + + shared_kwargs: Dict[str, Any] = dict( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=( + is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) + ), + qkv_latent_func=( + self.self_attn.prepare_qkv_latent if not self.is_linear_attn else None + ), + ) + + if self.config.mhc: + mhc_kwargs: Dict[str, Any] = dict( + is_first_layer=(self.layer_id == 0), + hc_mult=config.hc_mult, + hc_attn_pre=self.hc_attn_pre, + hc_ffn_pre=self.hc_ffn_pre, + hc_post=self.hc_post, + ) + self.layer_communicator = MHCLayerCommunicator( + **shared_kwargs, + **mhc_kwargs, + ) + else: + self.layer_communicator = LayerCommunicator(**shared_kwargs) + + def _hc_pre( + self, hc_fn, hc_scale, hc_base, hidden_states, out_norm_weight, out_norm_eps + ): + return _hc_pre_fn( + x=hidden_states, + hc_fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + hc_mult=self.config.hc_mult, + rms_eps=self.config.rms_norm_eps, + hc_eps=self.config.hc_eps, + sinkhorn_iters=self.config.hc_sinkhorn_iters, + post_mult_value=2.0, + hc_norm_weight=None, + out_norm_weight=out_norm_weight, + out_norm_eps=out_norm_eps, + ) + + def hc_attn_pre(self, hidden_states, out_norm_weight, out_norm_eps): + return self._hc_pre( + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + hidden_states, + out_norm_weight, + out_norm_eps, + ) + + def hc_ffn_pre(self, hidden_states, out_norm_weight, out_norm_eps): + return self._hc_pre( + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + hidden_states, + out_norm_weight, + out_norm_eps, + ) + + def hc_post(self, hidden_states, residual, h_res, h_post): + assert self.config.mhc, "hc_post is only valid when config.mhc=True" + return _hc_post_fn( + x=hidden_states, + residual=residual, + h_post=h_post, + h_res=h_res, + hc_mult=self.config.hc_mult, + ) + + def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool: + return is_nextn or ( + self.config.n_routed_experts is not None + and layer_id >= self.config.first_k_dense_replace + and layer_id % self.config.moe_layer_freq == 0 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: Optional[BumpAllocator] = None, + gemm_output_zero_allocator: BumpAllocator = None, + prev_topk_indices: Optional[torch.Tensor] = None, + ): + hidden_states_orig = hidden_states + + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, + residual, + forward_batch, + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + layer_scatter_modes=self.layer_scatter_modes, + prev_topk_indices=prev_topk_indices, + ) + if isinstance(hidden_states, tuple): + hidden_states, topk_indices = hidden_states + else: + topk_indices = None + get_attn_tp_context().clear_attn_inputs() + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, + residual, + forward_batch, + ) + + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + if isinstance(self.mlp, Glm5NextMLP): + gemm_output_zero_allocator = None + + if ( + isinstance(self.mlp, Glm5NextMoE) + and not self.mlp.experts.moe_runner_config.inplace + and not torch.compiler.is_compiling() + ): + from sglang.srt.layers.moe.moe_runner.base import moe_output_buffer_ctx + + _mlp_ctx = moe_output_buffer_ctx(hidden_states_orig) + else: + _mlp_ctx = nullcontext() + + with get_forward().scoped( + fuse_mlp_allreduce=should_allreduce_fusion, + mlp_reduce_scatter=use_reduce_scatter, + ): + with _mlp_ctx: + hidden_states = self.mlp( + hidden_states, + forward_batch, + gemm_output_zero_allocator, + ) + + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + + if not should_allreduce_fusion: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, + residual, + forward_batch, + ) + + return hidden_states, residual, topk_indices + + +class Glm5NextModel(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: Glm5NextTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + self.first_k_dense_replace = config.first_k_dense_replace + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + **get_embedding_tp_kwargs(), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.alt_stream = ( + torch.cuda.Stream() + if ( + _is_cuda + or envs.SGLANG_NPU_USE_MULTI_STREAM.get() + or envs.SGLANG_ROCM_USE_MULTI_STREAM.get() + ) + else None + ) + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Glm5NextDecoderLayer( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=self.alt_stream, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + self.gemm_output_zero_allocator_size = 0 + if ( + _use_aiter_gfx95 + and config.n_routed_experts == 256 + and self.embed_tokens.embedding_dim == 7168 + ): + num_moe_layers = sum( + [ + 1 + for i in range(len(self.layers)) + if isinstance(self.layers[i].mlp, Glm5NextMoE) + ] + ) + + allocate_size = 0 + for i in range(len(self.layers)): + if isinstance(self.layers[i].mlp, Glm5NextMoE): + a2a_backend = get_moe_a2a_backend() + is_a2a_moe = ( + a2a_backend.is_deepep() + or a2a_backend.is_mori() + or a2a_backend.is_mooncake() + ) + tp_size = 1 if is_a2a_moe else get_parallel().tp_size + intermediate_size = ( + config.moe_intermediate_size * config.n_shared_experts + ) + share_expert_output_size_per_partition = divide( + intermediate_size * 2, tp_size + ) + allocate_size = share_expert_output_size_per_partition + break + + self.gemm_output_zero_allocator_size = ( + get_dsv3_gemm_output_zero_allocator_size( + config.n_routed_experts, + num_moe_layers, + allocate_size, + self.embed_tokens.embedding_dim, + ) + ) + self.layers_to_capture = [] + self.dflash_capture = False + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + self.enable_a2a_moe = True + else: + self.enable_a2a_moe = False + + def get_input_embeddings(self) -> torch.Tensor: + return self.embed_tokens + + def _prepare_aux_hidden_state( + self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] + ) -> torch.Tensor: + # mHC folds the residual into widened hidden state, so residual remains None + # until hc_contract merges it; only plain residual streams are added here. + aux_hidden_state = ( + hidden_states if residual is None else hidden_states + residual + ) + if self.dflash_capture and self.config.mhc: + aux_hidden_state = hc_contract(aux_hidden_state, self.config.hc_mult) + return aux_hidden_state + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + total_num_layers = self.end_layer - self.start_layer + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + device = hidden_states.device + zero_allocator = BumpAllocator( + buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), + dtype=torch.float32, + device=device, + ) + + has_gemm_output_zero_allocator = hasattr( + self, "gemm_output_zero_allocator_size" + ) + + gemm_output_zero_allocator = ( + BumpAllocator( + buffer_size=self.gemm_output_zero_allocator_size, + dtype=torch.float32, + device=device, + ) + if has_gemm_output_zero_allocator + and self.gemm_output_zero_allocator_size > 0 + else None + ) + + normal_start_layer = self.start_layer + normal_end_layer = self.end_layer + if forward_batch.can_run_tbo and not self.dflash_capture: + if ( + self.first_k_dense_replace > normal_start_layer + and self.first_k_dense_replace < normal_end_layer + ): + normal_end_layer = self.first_k_dense_replace + elif self.first_k_dense_replace < normal_start_layer: + normal_end_layer = normal_start_layer = 0 + aux_hidden_states = [] + topk_indices = None + for i in range(normal_start_layer, normal_end_layer): + # NOTE: torch dynamo does not support graph break in context manager + ctx = ( + nullcontext() + if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) + else get_global_expert_distribution_recorder().with_current_layer(i) + ) + with ctx: + if i in self.layers_to_capture: + aux_hidden_state = self._prepare_aux_hidden_state( + hidden_states, residual + ) + if self.enable_a2a_moe and i > self.first_k_dense_replace: + aux_hidden_state = get_parallel().attn_tp_group.all_gather( + aux_hidden_state, dim=0 + ) + aux_hidden_states.append(aux_hidden_state) + layer = self.layers[i] + hidden_states, residual, topk_indices = layer( + positions, + hidden_states, + forward_batch, + residual, + zero_allocator, + gemm_output_zero_allocator, + prev_topk_indices=topk_indices, + ) + + if normal_end_layer != self.end_layer: + hidden_states, residual = model_forward_maybe_tbo( + layers=self.layers[normal_end_layer : self.end_layer], + enable_tbo=True, + positions=positions, + forward_batch=forward_batch, + hidden_states=hidden_states, + residual=residual, + input_data_scatter_mode=self.layers[ + normal_end_layer - 1 + ].layer_scatter_modes.layer_output_mode, + zero_allocator=zero_allocator, + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + else: + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) == 0: + return hidden_states + return hidden_states, aux_hidden_states + + +class Glm5NextForConditionalGeneration(nn.Module): + packed_modules_mapping = { + "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"], + "fused_qkvbfg_a_proj": [ + "q_proj", + "k_proj", + "v_proj", + "b_proj", + "f_a_proj", + "g_a_proj", + ], + "fused_fg_b_proj": ["f_b_proj", "g_b_proj"], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "qkv_conv1d": ["q_conv1d", "k_conv1d", "v_conv1d"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + fall_back_to_pt_during_load = False + + def __init__( + self, + config: Glm5NextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + vision_utils.update_vit_attn_dummy_heads_config(config) + self.mm_config = config + text_config = config.text_config + self.encoder_only = bool(getattr(config, "encoder_only", False)) + self.language_only = bool(getattr(config, "language_only", False)) + + self.fuse_qkv_a_proj = ( + not self.encoder_only + and getattr(text_config, "q_lora_rank", None) is not None + ) + + self.pp_group = get_pp_group() + self.config = text_config + self.tp_size = get_parallel().tp_size + self.quant_config = quant_config + self.use_dsa = is_deepseek_dsa(text_config) + self.num_fused_shared_experts = 0 + self.model = None + self.lm_head = None + self.logits_processor = None + if not self.encoder_only: + self.determine_num_fused_shared_experts() + self.model = Glm5NextModel( + text_config, quant_config, prefix=add_prefix("model", prefix) + ) + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and text_config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + text_config.vocab_size, + text_config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(text_config) + + self._routed_experts_weights_of_layer = LazyValue( + lambda: ( + { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Glm5NextMoE) + } + if self.model is not None + else {} + ) + ) + self.capture_aux_hidden_states = False + + if not self.encoder_only: + get_attn_tp_context().init_context( + getattr(text_config, "q_lora_rank", None), + self.use_dsa, + text_config.mhc, + ) + + self.use_data_parallel = get_mm().mm_enable_dp_encoder + self.visual = None + if not self.language_only: + self.visual = Glm5NextVisionModel( + config.vision_config, + quant_config=quant_config, + prefix=add_prefix("visual", prefix), + use_data_parallel=self.use_data_parallel, + ) + self.is_mrope_enabled = not self.encoder_only and "mrope_section" in ( + self.config.rope_scaling or {} + ) + + def get_input_embeddings(self) -> nn.Embedding: + if self.model is None: + raise AttributeError( + "get_input_embeddings() is not available in encoder-only mode" + ) + return self.model.embed_tokens + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + @classmethod + def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): + # Kept in lockstep with the wrapper gate below: a divergence drops the + # shared-expert weights and runs the fused slot uninitialized. + text_config = getattr(hf_config, "text_config", hf_config) + if not getattr(text_config, "n_shared_experts", None): + return "No shared experts are defined in the config." + if not _is_cuda: + return "Shared experts fusion currently requires CUDA devices." + if _device_sm is not None and _device_sm < 80: + return "Shared experts fusion requires SM80 or newer GPUs." + if get_parallel().moe_ep_size > 1: + return ( + "Shared experts fusion is not supported together with expert " + "parallelism yet." + ) + if get_moe_a2a_backend().is_deepep(): + return ( + "Shared experts fusion is not supported when Deepep MoE backend " + "is enabled." + ) + return None + + def determine_num_fused_shared_experts(self): + self.num_fused_shared_experts = ( + 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts + ) + if self.num_fused_shared_experts == 0: + return + assert self.num_fused_shared_experts == 1, ( + f"Only 1 fused shared expert is supported for {type(self).__name__}" + ) + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): + if not self.pp_group.is_last_rank: + return + + if layer_ids is None: + self.capture_aux_hidden_states = True + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] + else: + self.capture_aux_hidden_states = True + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + def set_dflash_layers_to_capture(self, layer_ids: List[int]): + if not self.pp_group.is_last_rank: + return + + if layer_ids is None: + raise ValueError( + "DFLASH requires explicit layer_ids for aux hidden capture." + ) + + self.capture_aux_hidden_states = True + self.model.dflash_capture = True + # Capturing before layer k + 1 gives the completed output of layer k. + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + self.visual.dtype + ) + image_grid_thw = torch.concat([item.image_grid_thw for item in items], dim=0) + assert pixel_values.dim() == 2, pixel_values.dim() + assert image_grid_thw.dim() == 2, image_grid_thw.dim() + + if self.use_data_parallel: + return run_dp_sharded_mrope_vision_model( + self.visual, pixel_values, image_grid_thw.tolist(), rope_type="rope_3d" + ) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + return image_embeds + + def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + self.visual.dtype + ) + video_grid_thw = torch.concat([item.video_grid_thw for item in items], dim=0) + + temp_frames_hw = [] + for t, h, w in video_grid_thw: + repeated_row = ( + torch.tensor([1, h.item(), w.item()]).unsqueeze(0).repeat(t, 1) + ) + temp_frames_hw.append(repeated_row) + flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0) + + assert pixel_values.dim() == 2, pixel_values.dim() + assert video_grid_thw.dim() == 2, video_grid_thw.dim() + if items and getattr(items[0], "dp_decode_sharded", False): + if len(items) != 1: + raise ValueError("DP-sharded video decode requires one video item") + dp_meta = items[0].dp_meta + height = int(video_grid_thw[0][1]) + width = int(video_grid_thw[0][2]) + global_grid = [[1, height, width]] * int(dp_meta["n_units"]) + return run_dp_presharded_mrope_vision_model( + self.visual, + pixel_values, + flattened_video_grid_thw.tolist(), + global_grid, + dp_meta["gpu_sample_counts"], + ) + + if self.use_data_parallel: + return run_dp_sharded_mrope_vision_model( + self.visual, + pixel_values, + flattened_video_grid_thw.tolist(), + rope_type="rope_3d", + ) + video_embeds = self.visual(pixel_values, grid_thw=flattened_video_grid_thw) + return video_embeds + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + if self.is_mrope_enabled: + positions = forward_batch.mrope_positions + + with get_attn_tp_context().maybe_input_scattered(forward_batch): + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + positions=positions, + pp_proxy_tensors=pp_proxy_tensors, + ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + else: + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): + if is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + assert num_nextn_layers == 1, "Only 1 nextn layer is supported" + nextn_layer_id = ( + 0 + if self.config.num_hidden_layers == 1 + else self.config.num_hidden_layers + ) + else: + raise ValueError("num_nextn_predict_layers is not in the config") + + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + (".fused_qkvbfg_a_proj", ".q_proj", 0), + (".fused_qkvbfg_a_proj", ".k_proj", 1), + (".fused_qkvbfg_a_proj", ".v_proj", 2), + (".fused_qkvbfg_a_proj", ".b_proj", 3), + (".fused_qkvbfg_a_proj", ".f_a_proj", 4), + (".fused_qkvbfg_a_proj", ".g_a_proj", 5), + (".fused_fg_b_proj", ".f_b_proj", 0), + (".fused_fg_b_proj", ".g_b_proj", 1), + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".qkv_conv1d", ".q_conv1d", 0), + (".qkv_conv1d", ".k_conv1d", 1), + (".qkv_conv1d", ".v_conv1d", 2), + ] + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + + if is_nextn: + nextn_layer_prefix = f"model.layers.{nextn_layer_id}" + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + + fuse_qkv_a_proj = getattr(self, "fuse_qkv_a_proj", False) + cached_a_proj: dict[str, torch.Tensor] = {} if fuse_qkv_a_proj else None + qc = self.quant_config + if qc is not None and qc.get_name() in {"awq", "awq_marlin", "moe_wna16"}: + fused_cat_dim = 1 + else: + fused_cat_dim = 0 + + params_dict = dict(self.named_parameters()) + weight_names = [] + for name, loaded_weight in weights: + is_visual_weight = "visual" in name + if getattr(self, "encoder_only", False) and not is_visual_weight: + continue + if getattr(self, "language_only", False) and is_visual_weight: + continue + + if "language_model." in name: + name = name.replace("language_model.", "") + if "model.visual." in name: + name = name.replace("model.visual.", "visual.") + + if "visual" in name: + name = name.replace("attn.qkv.", "attn.qkv_proj.") + loaded_weight = vision_utils.pad_vit_attn_dummy_heads( + self.mm_config, name, loaded_weight + ) + + weight_names.append(name) + + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts}", + ) + + if not is_nextn: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + else: + if not name.startswith(nextn_layer_prefix): + continue + + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + if "rotary_emb.inv_freq" in name: + continue + if "hc_head" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + candidate = name.replace(weight_name, param_name) + if ( + param_name + in { + ".fused_qkvbfg_a_proj", + ".fused_fg_b_proj", + ".qkv_proj", + ".qkv_conv1d", + } + and candidate not in params_dict + ): + continue + name = candidate + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + is_expert_weight = True + name = name.replace(weight_name, param_name) + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if is_expert_weight: + continue + if name.endswith(".bias") and name not in params_dict: + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + fused_weight = torch.cat( + [ + cached_a_proj[q_a_proj_name], + cached_a_proj[kv_a_proj_name], + ], + dim=fused_cat_dim, + ) + target = ( + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") + if "q_a_proj" in name + else name.replace( + "kv_a_proj_with_mqa", + "fused_qkv_a_proj_with_mqa", + ) + ) + if target in params_dict: + param = params_dict[target] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name, None) + cached_a_proj.pop(kv_a_proj_name, None) + continue + + if name not in params_dict: + continue + + if name.endswith(".A_log") and loaded_weight.dim() == 1: + loaded_weight = loaded_weight.view(1, 1, -1, 1) + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + if getattr(self, "encoder_only", False): + run_post = False + elif is_nextn: + decoder_attn = getattr(self.model.decoder, "self_attn", None) + run_post = decoder_attn is not None and hasattr(decoder_attn, "kv_b_proj") + else: + run_post = True + if run_post: + DeepseekV2WeightLoaderMixin.post_load_weights( + self, is_nextn=is_nextn, weight_names=weight_names + ) + + def post_load_weights(self, is_nextn: bool = False, weight_names=None): + if self.encoder_only: + return + DeepseekV2WeightLoaderMixin.post_load_weights( + self, is_nextn=is_nextn, weight_names=weight_names + ) + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + if self.model is None: + raise AttributeError( + "load_kv_cache_scales() is not available in encoder-only mode" + ) + if callable(getattr(self.model, "load_kv_cache_scales", None)): + self.model.load_kv_cache_scales(quantization_param_path) + else: + logger.warning( + f"{self.model.__class__} does not support loading scaling factors." + ) + + def get_embed_and_head(self): + if self.model is None or self.lm_head is None: + raise AttributeError( + "get_embed_and_head() is not available in encoder-only mode" + ) + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + if self.model is None or self.lm_head is None: + raise AttributeError( + "set_embed_and_head() is not available in encoder-only mode" + ) + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @classmethod + def get_model_config_for_expert_location(cls, config): + config = getattr(config, "text_config", config) + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.n_routed_experts, + num_groups=None, + ) + + +EntryClass = [Glm5NextForConditionalGeneration] diff --git a/python/sglang/srt/models/glm5_next_nextn.py b/python/sglang/srt/models/glm5_next_nextn.py new file mode 100644 index 000000000..fb7cf5554 --- /dev/null +++ b/python/sglang/srt/models/glm5_next_nextn.py @@ -0,0 +1,80 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import logging + +from sglang.srt.models.deepseek_nextn import DeepseekV3ForCausalLMNextN +from sglang.srt.models.glm5_next import Glm5NextForConditionalGeneration +from sglang.srt.models.utils import WeightsMapper + +logger = logging.getLogger(__name__) + + +class Glm5NextForConditionalGenerationNextN(DeepseekV3ForCausalLMNextN): + @classmethod + def get_hf_to_sglang_mapper(cls, config) -> WeightsMapper: + text_config = getattr(config, "text_config", config) + return WeightsMapper( + orig_to_new_substr={ + f"model.layers.{text_config.num_hidden_layers}": "model.decoder", + }, + ) + + def _resolve_nextn_quant_config(self, config, quant_config): + """Mixed checkpoints list the BF16 NextN block in ``quantization_config.ignore``; + inheriting global FP8 quantization would corrupt its QKV weights.""" + raw_quant_config = getattr(config, "quantization_config", None) or {} + if hasattr(raw_quant_config, "to_dict"): + raw_quant_config = raw_quant_config.to_dict() + ignored = ( + raw_quant_config.get("ignore", []) + if isinstance(raw_quant_config, dict) + else [] + ) + nextn_layer_pattern = f"model.layers.{config.num_hidden_layers}.*" + if nextn_layer_pattern in ignored: + logger.warning( + "GLM5 NextN layer %s is checkpoint-declared unquantized; " + "using BF16 draft modules", + nextn_layer_pattern, + ) + return None + return super()._resolve_nextn_quant_config(config, quant_config) + + def __init__(self, config, quant_config=None, prefix: str = "") -> None: + super().__init__( + getattr(config, "text_config", config), + quant_config=quant_config, + prefix=prefix, + ) + + def load_weights(self, weights): + if not hasattr(self, "fuse_qkv_a_proj"): + self.fuse_qkv_a_proj = getattr(self.config, "q_lora_rank", None) is not None + layer_id = self.config.num_hidden_layers + layer_prefixes = ( + f"model.layers.{layer_id}.", + f"model.language_model.layers.{layer_id}.", + ) + nextn_weights = ( + (name, weight) + for name, weight in weights + if name.startswith(layer_prefixes) + ) + return Glm5NextForConditionalGeneration.load_weights( + self, nextn_weights, is_nextn=True + ) + + +EntryClass = [Glm5NextForConditionalGenerationNextN] diff --git a/python/sglang/srt/models/glm_ocr.py b/python/sglang/srt/models/glm_ocr.py index dfd00f31b..15fb7d6c9 100644 --- a/python/sglang/srt/models/glm_ocr.py +++ b/python/sglang/srt/models/glm_ocr.py @@ -26,7 +26,6 @@ import torch.nn as nn from einops import rearrange from transformers.models.glm_ocr.configuration_glm_ocr import ( GlmOcrConfig, - GlmOcrTextConfig, GlmOcrVisionConfig, ) @@ -158,7 +157,6 @@ class GlmOcrVisionModel(Glm4vVisionModel): def __init__( self, vision_config: GlmOcrVisionConfig, - text_config: GlmOcrTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", use_data_parallel: bool = False, @@ -209,9 +207,16 @@ class GlmOcrVisionModel(Glm4vVisionModel): for layer_idx in range(depth) ] ) + projection_intermediate_size = getattr( + vision_config, "projection_intermediate_size", None + ) self.merger = GlmOcrVisionPatchMerger( d_model=vision_config.out_hidden_size, - context_dim=text_config.intermediate_size, + context_dim=( + projection_intermediate_size + if projection_intermediate_size is not None + else vision_config.intermediate_size + ), quant_config=quant_config, bias=False, prefix=add_prefix("merger", prefix), @@ -285,7 +290,6 @@ class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration): self.use_data_parallel = get_mm().mm_enable_dp_encoder self.visual = GlmOcrVisionModel( vision_config=config.vision_config, - text_config=config.text_config, quant_config=quant_config, prefix=add_prefix("visual", prefix), use_data_parallel=self.use_data_parallel, diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 97cab3284..075d742f7 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -823,3 +823,67 @@ def run_dp_sharded_mrope_vision_model( current_idx += count out_embeddings = torch.cat(original_order_embeddings, dim=0) return out_embeddings + + +def run_dp_presharded_mrope_vision_model( + vision_model: torch.nn.Module, + pixel_values_local: torch.Tensor, + local_grid_thw_list: list, + global_grid_thw_list: list, + gpu_sample_counts: list, +) -> torch.Tensor: + """Rank-local shards are contiguous, so rank-order concatenation restores global video order.""" + parallel = get_parallel() + tp_size = parallel.attn_tp_size + patches_per_unit = [math.prod(grid) for grid in global_grid_thw_list] + grouped_patch_counts = [] + offset = 0 + for rank in range(tp_size): + count = gpu_sample_counts[rank] + grouped_patch_counts.append(sum(patches_per_unit[offset : offset + count])) + offset += count + + merge_factor = vision_model.spatial_merge_size**2 + grouped_output_lengths = [ + patch_count // merge_factor for patch_count in grouped_patch_counts + ] + max_output_length = max(grouped_output_lengths) + try: + model_device = vision_model.device + model_dtype = vision_model.dtype + except AttributeError: + parameter = next(vision_model.parameters()) + model_device, model_dtype = parameter.device, parameter.dtype + + if pixel_values_local.shape[0] > 0: + pixel_values_local = pixel_values_local.to( + device=model_device, dtype=model_dtype + ) + local_embeddings = vision_model( + pixel_values_local, + grid_thw=torch.tensor(local_grid_thw_list), + ) + else: + local_embeddings = torch.empty( + (0, vision_model.out_hidden_size), + device=model_device, + dtype=model_dtype, + ) + + if local_embeddings.shape[0] < max_output_length: + padding = torch.empty( + ( + max_output_length - local_embeddings.shape[0], + local_embeddings.shape[1], + ), + device=local_embeddings.device, + dtype=local_embeddings.dtype, + ) + local_embeddings = torch.cat([local_embeddings, padding], dim=0) + + gathered = parallel.attn_tp_group.all_gather(local_embeddings, dim=0) + pieces = [] + for rank, output_length in enumerate(grouped_output_lengths): + start = rank * max_output_length + pieces.append(gathered[start : start + output_length]) + return torch.cat(pieces, dim=0) diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 0deffbfe9..952ddae27 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -965,11 +965,14 @@ class BaseMultimodalProcessor(ABC): img, _ = load_image(data, cls.gpu_image_decode) if isinstance(img, torch.Tensor): return img # JPEG already decoded on GPU by nvJPEG + # PIL decodes lazily; do it here in the io worker so the decode + # doesn't run later on the event-loop thread. if discard_alpha_channel: if cls.smart_rgb_conversion: return smart_to_rgb(img) if img.mode != "RGB": return img.convert("RGB") + img.load() return img elif modality == Modality.VIDEO: return load_video(data, frame_count_limit) diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py index 67b81c245..105814e25 100644 --- a/python/sglang/srt/multimodal/processors/glm4v.py +++ b/python/sglang/srt/multimodal/processors/glm4v.py @@ -1,4 +1,5 @@ import asyncio +import json import math from typing import List, Tuple, Union @@ -77,14 +78,81 @@ def split_glm_video_items(mm_data): return urls, configs +def glm_budget_kwargs(processor, user_max_image_tokens=None, count=1, split=False): + if processor is None: + return None + default_max = getattr(processor, "max_image_tokens", None) + if not default_max: + return None + + if user_max_image_tokens is not None: + budget = int(user_max_image_tokens) + elif split: + budget = int(default_max) + else: + return None + + count = max(int(count or 1), 1) + effective = max(1, budget // count if split and count > 1 else budget) + if effective == default_max and user_max_image_tokens is None: + return None + return {"max_image_tokens": effective} + + +def glm_max_image_tokens_from_configs(configs): + values = [ + int(config["max_image_tokens"]) + for config in configs or [] + if isinstance(config, dict) and config.get("max_image_tokens") is not None + ] + return min(values) if values else None + + def glm_processor_video_config(processor): if processor is None: return {} - return { + config = { key: value for key in GLM_MEDIA_CONFIG_KEYS if (value := getattr(processor, key, None)) is not None } + budget = _glm_processor_resize_budget(processor) + if budget is not None: + config["_presize_budget"] = budget + return config + + +def _glm_processor_resize_budget(processor): + """Use token limits because Glm5NextVideoProcessor.size.longest_edge is only a sentinel.""" + max_image_tokens = getattr(processor, "max_image_tokens", None) + if not max_image_tokens: + return None + patch_size = getattr(processor, "patch_size", None) or GLM_VIDEO_PATCH_SIZE + merge_size = getattr(processor, "merge_size", None) or GLM_VIDEO_MERGE_SIZE + expand_factor = getattr(processor, "patch_expand_factor", None) or 1 + temporal_factor = getattr(processor, "temporal_patch_size", None) or 2 + pixels_per_token = int(temporal_factor * (patch_size * merge_size) ** 2) + return { + "factor": int(patch_size * merge_size * expand_factor), + "temporal_factor": int(temporal_factor), + "pixels_per_token": pixels_per_token, + "min_pixels": int(getattr(processor, "min_image_tokens", None) or 0) + * pixels_per_token, + "max_pixels": int(max_image_tokens) * pixels_per_token, + "resize_mode": getattr(processor, "resize_mode", None) or "resize", + } + + +def _glm_effective_presize_budget(video_config, effective_max_image_tokens): + budget = video_config.get("_presize_budget") if video_config else None + if not budget or effective_max_image_tokens is None: + return video_config + config = dict(video_config) + config["_presize_budget"] = { + **budget, + "max_pixels": int(effective_max_image_tokens) * budget["pixels_per_token"], + } + return config def _merge_glm_video_configs(default_config, item_configs): @@ -204,6 +272,122 @@ def _resize_frames_to_max_tokens(frames, max_tokens_per_frame): return nchw.permute(0, 2, 3, 1).contiguous() +def preprocess_video_frames_sync(frame_list: List[dict]): + total_num_frames = len(frame_list) + if total_num_frames == 0: + raise ValueError("GLM video frame list must not be empty") + duration = 0.0 + if frame_list[0].get("detail") is not None: + details = json.loads(frame_list[0]["detail"]) + duration = float(details.get("video_duration", 0)) + if duration == 0: + base_ts = float(frame_list[0].get("timestamp", 0) or 0) + duration = float(frame_list[-1].get("timestamp", base_ts) or base_ts) - base_ts + images = [frame["frame_image"] for frame in frame_list] + if isinstance(images[0], torch.Tensor): + images = torch.stack(images).permute(0, 2, 3, 1).contiguous() + else: + images = [np.asarray(image) for image in images] + fps = total_num_frames / duration if duration else 0 + return images, _glm_video_metadata( + total_num_frames, fps, duration, range(total_num_frames) + ) + + +GLM_VIDEO_PRE_RESIZE_CHUNK = 64 + + +def _vendor_smart_resize_canvas( + num_frames, height, width, *, temporal_factor, factor, min_pixels, max_pixels +): + """Replica of the vendor Glm5Next smart_resize (align-ceil + budget search).""" + + def align(value): + return math.ceil(value / factor) * factor + + def fit_within_budget(aligned_frames): + low, high = 1, height + best_height, best_width = factor, factor + while low <= high: + content_height = (low + high) // 2 + content_width = max(1, math.floor(width * content_height / height)) + candidate_height = align(content_height) + candidate_width = align(content_width) + if aligned_frames * candidate_height * candidate_width <= max_pixels: + best_height, best_width = candidate_height, candidate_width + low = content_height + 1 + else: + high = content_height - 1 + return best_height, best_width + + aligned_frames = max( + temporal_factor, round(num_frames / temporal_factor) * temporal_factor + ) + canvas_height, canvas_width = align(height), align(width) + if aligned_frames * canvas_height * canvas_width > max_pixels: + canvas_height, canvas_width = fit_within_budget(aligned_frames) + elif aligned_frames * canvas_height * canvas_width < min_pixels: + scale = math.sqrt(min_pixels / (num_frames * height * width)) + canvas_height = align(max(1, math.ceil(height * scale))) + canvas_width = align(max(1, math.ceil(width * scale))) + if aligned_frames * canvas_height * canvas_width > max_pixels: + canvas_height, canvas_width = fit_within_budget(aligned_frames) + return canvas_height, canvas_width + + +def _pre_resize_frames_for_processor( + frames, + *, + factor, + temporal_factor, + pixels_per_token, + min_pixels, + max_pixels, + resize_mode, +): + """Pre-resize in chunks to avoid HF's native-resolution float32 intermediate while preserving its output grid.""" + import torchvision.transforms.functional as TF + + if not isinstance(frames, torch.Tensor): + frames = torch.from_numpy(np.asarray(frames)) + nchw = frames.permute(0, 3, 1, 2) + num_frames, _, height, width = nchw.shape + canvas_height, canvas_width = _vendor_smart_resize_canvas( + num_frames, + height, + width, + temporal_factor=temporal_factor, + factor=factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + if resize_mode == "resize": + content_height, content_width = canvas_height, canvas_width + else: + scale = min(canvas_height / height, canvas_width / width) + if num_frames * height * width >= min_pixels: + scale = min(1.0, scale) + content_height = max(1, min(canvas_height, math.floor(height * scale))) + content_width = max(1, min(canvas_width, math.floor(width * scale))) + if (content_height, content_width) != (height, width): + nchw = torch.cat( + [ + TF.resize( + chunk, + [content_height, content_width], + interpolation=TF.InterpolationMode.BICUBIC, + antialias=True, + ) + for chunk in nchw.split(GLM_VIDEO_PRE_RESIZE_CHUNK) + ] + ) + if (content_height, content_width) != (canvas_height, canvas_width): + nchw = torch.nn.functional.pad( + nchw, (0, canvas_width - content_width, 0, canvas_height - content_height) + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + def glm_decode_frames_at(vr, indices, video_config=None): indices = list(indices) if not indices: @@ -216,6 +400,8 @@ def glm_decode_frames_at(vr, indices, video_config=None): max_tokens_per_frame = video_config.get("max_tokens_per_frame") if max_tokens_per_frame is not None: frames = _resize_frames_to_max_tokens(frames, max_tokens_per_frame) + elif budget := video_config.get("_presize_budget"): + frames = _pre_resize_frames_for_processor(frames, **budget) return frames @@ -372,12 +558,21 @@ class Glm4vImageProcessor(SGLangBaseProcessor): ) video_metadata = None + videos_kwargs = None if base_output.videos and not isinstance(base_output.videos[0], dict): + videos_kwargs = glm_budget_kwargs( + video_processor, + user_max_image_tokens=glm_max_image_tokens_from_configs(video_configs), + count=len(base_output.videos), + split=True, + ) + effective_max_image_tokens = (videos_kwargs or {}).get("max_image_tokens") loop = asyncio.get_running_loop() decode_tasks = [] for index, video in enumerate(base_output.videos): - video_config = ( - video_configs[index] if index < len(video_configs) else {} + video_config = _glm_effective_presize_budget( + video_configs[index] if index < len(video_configs) else {}, + effective_max_image_tokens, ) if isinstance(video, VideoDecoderWrapper): decode_tasks.append( @@ -389,6 +584,14 @@ class Glm4vImageProcessor(SGLangBaseProcessor): video_processor, ) ) + elif isinstance(video, list) and ( + not video or isinstance(video[0], dict) + ): + decode_tasks.append( + loop.run_in_executor( + self.io_executor, preprocess_video_frames_sync, video + ) + ) else: decode_tasks.append( asyncio.sleep( @@ -407,18 +610,23 @@ class Glm4vImageProcessor(SGLangBaseProcessor): close = getattr(video, "close", None) if callable(close): close() - base_output.videos, video_metadata = map(list, zip(*videos_processed)) + base_output.videos, metadata = map(list, zip(*videos_processed)) + if metadata and all(item is not None for item in metadata): + video_metadata = metadata combine_kwargs = {} if video_metadata is not None: # Skip HF resampling because these frames already carry their original indices. combine_kwargs["video_metadata"] = video_metadata combine_kwargs["do_sample_frames"] = False - combine_kwargs["processor_video_config"] = { + processor_video_config = { key: value for key, value in self.video_config.items() if key not in {"fps", "max_frames", "max_tokens_per_frame"} } + if videos_kwargs is not None: + processor_video_config.update(videos_kwargs) + combine_kwargs["processor_video_config"] = processor_video_config mm_items, input_ids, ret = await self.process_and_combine_mm_data_async( base_output, self.mm_tokens, **combine_kwargs diff --git a/python/sglang/srt/speculative/adaptive_spec_params.py b/python/sglang/srt/speculative/adaptive_spec_params.py index 96997c1ff..6d3353499 100644 --- a/python/sglang/srt/speculative/adaptive_spec_params.py +++ b/python/sglang/srt/speculative/adaptive_spec_params.py @@ -25,7 +25,7 @@ logger = logging.getLogger(__name__) DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = { "1": { - "candidate_steps": [1, 3, 7], + "candidate_steps": [1, 3, 5, 7], "up_hysteresis": 0.0, "down_hysteresis": -0.25, "ceiling_coeff": 0, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 1ca19cbf6..0afb522d9 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -237,7 +237,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): dsa_seed_topk = ( torch.zeros( - (self.max_bs, self.eagle_worker.dsa_index_topk), + (self.max_bs, self.eagle_worker.dsa_seed_topk_width), dtype=torch.int32, device=model_runner.device, ) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 0ef10e2b3..b24e0aa5d 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -55,6 +55,18 @@ if TYPE_CHECKING: from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker +def resolve_draft_extend_seq_len_fill_value( + attn_backend, captured_req_width: int +) -> int: + """Pad synthetic history past the fixed draft-width subtraction and KPool offset.""" + fill_value = attn_backend.get_cuda_graph_seq_len_fill_value() + full_attn_backend = getattr(attn_backend, "full_attn_backend", attn_backend) + dsa_index_kpool = getattr(full_attn_backend, "dsa_index_kpool", 1) + if dsa_index_kpool > 1: + fill_value = max(fill_value, captured_req_width + dsa_index_kpool) + return fill_value + + @dataclass class EagleDraftExtendInputBuffers(ForwardInputBuffers): input_ids: torch.Tensor @@ -142,8 +154,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): self.draft_extend_attn_backend.init_cuda_graph_state( self.max_bs, self.max_num_token ) - self.seq_len_fill_value = ( - self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value() + self.seq_len_fill_value = resolve_draft_extend_seq_len_fill_value( + self.draft_extend_attn_backend, self.captured_req_width ) self.extend_seq_lens_cpu = [self.captured_req_width] * self.max_bs @@ -176,8 +188,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): if _hidden_size is not None else None ) - self.seq_len_fill_value = ( - self.draft_extend_attn_backend.get_cuda_graph_seq_len_fill_value() + self.seq_len_fill_value = resolve_draft_extend_seq_len_fill_value( + self.draft_extend_attn_backend, self.captured_req_width ) seq_lens = torch.full( (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64 @@ -243,7 +255,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): dsa_seed_topk_capture = ( torch.full( - (self.max_num_token, self.eagle_worker.dsa_index_topk), + (self.max_num_token, self.eagle_worker.dsa_seed_topk_width), -1, dtype=torch.int32, device=model_runner.device, diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index f679b5050..99875779b 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -507,21 +507,43 @@ def get_draft_recurrent_hidden_state_spec( ) +_PREPARE_FOR_VERIFY_DEPS = None + + def eagle_prepare_for_verify( verify_input: EagleVerifyInput, req_to_token_pool: ReqToTokenPool, batch: ScheduleBatch, target_worker: TpModelWorker, ): - from sglang.kernels.ops.speculative.cache_locs import ( + # Imports must stay lazy (import-cycle safety) but only need to resolve + # once, not on every decode cycle of this hot path. + global _PREPARE_FOR_VERIFY_DEPS + if _PREPARE_FOR_VERIFY_DEPS is None: + from sglang.kernels.ops.speculative.cache_locs import ( + assign_extend_cache_locs_uniform_func, + ) + from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, + ) + from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify + + _PREPARE_FOR_VERIFY_DEPS = ( + assign_extend_cache_locs_uniform_func, + CaptureHiddenMode, + ForwardBatch, + ForwardMode, + prepare_mamba_track_for_verify, + ) + ( assign_extend_cache_locs_uniform_func, - ) - from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, - ) - from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify + prepare_mamba_track_for_verify, + ) = _PREPARE_FOR_VERIFY_DEPS if not batch.forward_mode.is_idle(): # Assign cache locations diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index afc843408..010811b41 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -495,6 +495,10 @@ def run_eagle_verify( # Batch 1: Target verify # Prepare for target verify in a separate stream with plan_stream_ctx: + if plan_stream is not None: + # Verify prep copies draft-produced tree metadata on the plan stream, + # so it must not start before the draft frontier. + plan_stream.wait_stream(fwd_stream) verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify( verify_input, req_to_token_pool, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index e89e61155..b0c79668a 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -7,6 +7,7 @@ from typing import List, Optional import torch from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess +from sglang.srt.configs.model_config import get_dsa_mtp_topk_width from sglang.srt.distributed import get_pp_group from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.environ import envs @@ -269,8 +270,13 @@ class EagleDraftWorker(EagleDraftWorkerBase): # GLM-5.2 MTP IndexShare: seed reused indexer top-k from draft-extend # (last verified token), not draft-decode step 0. self.dsa_index_topk = getattr(hf_config, "index_topk", None) + self.dsa_seed_topk_width = ( + get_dsa_mtp_topk_width(hf_config) + if self.index_share_for_mtp_iteration and self.dsa_index_topk is not None + else None + ) self.seed_dsa_topk_from_draft_extend = ( - self.index_share_for_mtp_iteration and self.dsa_index_topk is not None + self.index_share_for_mtp_iteration and self.dsa_seed_topk_width is not None ) def init_token_map(self): @@ -801,16 +807,29 @@ class EagleDraftWorker(EagleDraftWorkerBase): if not batch.forward_mode.is_idle(): # Chunked-prefill-aware tail tokens (see PR #26329). tail_tokens = _eagle_prefill_tail_tokens(batch, next_token_ids) + new_input_ids = torch.empty_like(batch.input_ids) + if mm_input_embeds is not None: + # Rotate mm embeddings the same way as input_ids: shift left by + # one per request so they stay aligned with the rotated ids. The + # last position per request is filled by the draft model's own + # embed_tokens lookup on next_token_ids (see DeepseekModelNextN). + rotated_mm = torch.empty_like(mm_input_embeds) pt = 0 for i, extend_len in enumerate(batch.extend_lens): input_ids = batch.input_ids[pt : pt + extend_len] new_input_ids[pt : pt + extend_len].copy_( torch.cat((input_ids[1:], tail_tokens[i].reshape(1))) ) + if mm_input_embeds is not None: + rotated_mm[pt : pt + extend_len - 1].copy_( + mm_input_embeds[pt + 1 : pt + extend_len] + ) pt += extend_len assert pt == batch.input_ids.numel() batch.input_ids = new_input_ids + if mm_input_embeds is not None: + mm_input_embeds = rotated_mm # Draft-extend spec_info for the extend forward; carries only # hidden_states + shape info. @@ -897,11 +916,10 @@ class EagleDraftWorker(EagleDraftWorkerBase): ) def _get_dsa_extend_topk_buf(self, num_tokens: int) -> torch.Tensor: - """Lazily-grown int32 [num_tokens, index_topk] eager draft-extend seed buffer.""" buf = self.dsa_extend_topk_buf if buf is None or buf.shape[0] < num_tokens: buf = torch.full( - (num_tokens, self.dsa_index_topk), + (num_tokens, self.dsa_seed_topk_width), -1, dtype=torch.int32, device=self.device, diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 30f211fe5..11da6592a 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -10,6 +10,7 @@ import torch from sglang.srt.arg_groups.overrides import resolving_view from sglang.srt.runtime_context import get_spec as get_spec_config from sglang.srt.speculative.spec_registry import ( + _RESERVED_ALIASES, CustomSpecAlgo, ServerArgsValidator, WorkerFactory, @@ -57,6 +58,8 @@ class SpeculativeAlgorithm(Enum): return cls[upper] except KeyError: pass + if upper in _RESERVED_ALIASES: + return cls.EAGLE spec = _get_registered_spec(upper) if spec is not None: return spec diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index a270cb278..5f1aaab95 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -108,6 +108,7 @@ class StandaloneDraftWorker(EagleDraftWorker): and self.topk == 1 ) self.dsa_index_topk = None + self.dsa_seed_topk_width = None self.seed_dsa_topk_from_draft_extend = False self.dsa_extend_topk_buf = None diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 3998c2f82..9eb70a498 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -90,7 +90,7 @@ import torch import torch.distributed as dist import triton from packaging import version as pkg_version -from PIL import Image, UnidentifiedImageError +from PIL import Image, ImageOps, UnidentifiedImageError from starlette.routing import Mount from torch import nn from torch.library import Library @@ -1813,6 +1813,7 @@ def smart_to_rgb( if not isinstance(image, Image.Image): return image + image = ImageOps.exif_transpose(image) if image.mode in ("RGBA", "LA") or "transparency" in image.info: image = image.convert("RGBA") width, height = image.size @@ -1954,6 +1955,8 @@ def load_image( image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode) else: raise ValueError(f"Invalid image: {image_file}") + if image_size is not None and isinstance(image, Image.Image): + image_size = (image.width, image.height) return image, image_size diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 13e4f897c..6b3eea06d 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -37,6 +37,8 @@ from sglang.srt.configs import ( DotsVLMConfig, ExaoneConfig, FalconH1Config, + Glm5NextConfig, + Glm5NextTextConfig, GraniteMoeHybridConfig, HYV4Config, InklingAudioConfig, @@ -120,6 +122,8 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { MuseGlimmerConfig, MuseGlimmerAssistantConfig, KimiK3Config, + Glm5NextConfig, + Glm5NextTextConfig, KimiLinearConfig, Qwen3NextConfig, FalconH1Config, diff --git a/python/sglang/srt/utils/video_decoder.py b/python/sglang/srt/utils/video_decoder.py index 0dfec642f..3b398b58e 100644 --- a/python/sglang/srt/utils/video_decoder.py +++ b/python/sglang/srt/utils/video_decoder.py @@ -87,9 +87,11 @@ class VideoDecoderWrapper: return len(self._decoder) def __getitem__(self, idx): - """Return single frame as numpy NHWC uint8.""" + """Return one NHWC uint8 frame (numpy on CPU, tensor on CUDA).""" if _BACKEND == "torchcodec": - return self._decoder[idx].numpy() + frame = self._decoder[idx] + data = frame.data if hasattr(frame, "data") else frame + return data if data.is_cuda else data.numpy() else: frame = self._decoder[idx] return frame.asnumpy() if hasattr(frame, "asnumpy") else np.array(frame) @@ -101,11 +103,22 @@ class VideoDecoderWrapper: else: return self._decoder.get_avg_fps() - def get_frames_at(self, indices: list) -> np.ndarray: - """Return frames at given indices as numpy array with shape (N, H, W, C).""" + @property + def frame_shape(self) -> tuple[int, int]: + if _BACKEND == "torchcodec": + metadata = self._decoder.metadata + height = getattr(metadata, "height", None) + width = getattr(metadata, "width", None) + if height and width: + return int(height), int(width) + shape = self[0].shape + return int(shape[-3]), int(shape[-2]) + + def get_frames_at(self, indices: list): + """Return NHWC uint8 frames (numpy on CPU, tensor on CUDA).""" if _BACKEND == "torchcodec": batch = self._decoder.get_frames_at(indices) - return batch.data.numpy() + return batch.data if batch.data.is_cuda else batch.data.numpy() else: return self._decoder.get_batch(indices).asnumpy() @@ -127,7 +140,7 @@ class VideoDecoderWrapper: if _BACKEND == "torchcodec": batch = self._decoder.get_frames_at(indices) - return batch.data.pin_memory() + return batch.data if batch.data.is_cuda else batch.data.pin_memory() else: arr = self._decoder.get_batch(indices).asnumpy() return torch.from_numpy(arr).pin_memory() @@ -141,8 +154,15 @@ class VideoDecoderWrapper: chunks = [list(c) for c in np.array_split(indices, num_threads) if len(c) > 0] source = self._source kwargs = self._tc_kwargs + cuda_device = None + if kwargs.get("device") == "cuda": + cuda_device = torch.cuda.current_device() def _decode_chunk(chunk): + # CUDA's current device is thread-local. Without this, decoder + # workers created by TP rank > 0 silently default to GPU 0. + if cuda_device is not None: + torch.cuda.set_device(cuda_device) d = VideoDecoder(source, **kwargs) return d.get_frames_at(chunk).data @@ -156,7 +176,8 @@ class VideoDecoderWrapper: idx = future_to_idx[future] results[idx] = future.result() - return torch.cat(results, dim=0).pin_memory() + output = torch.cat(results, dim=0) + return output if output.is_cuda else output.pin_memory() @property def source_bytes(self) -> bytes | None: @@ -171,8 +192,11 @@ class VideoDecoderWrapper: return None def close(self): - """Explicitly clean up temporary files.""" - if self._tmp_path is not None: + self._decoder = None + self._source = None + self._source_bytes = None + self._source_path = None + if getattr(self, "_tmp_path", None) is not None: if os.path.exists(self._tmp_path): os.unlink(self._tmp_path) self._tmp_path = None diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index df1c52082..6d966922d 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsa_indexer.py @@ -219,6 +219,7 @@ class MockModelRunner: self.sliding_window_size = None self.page_size = self.config["page_size"] + self.max_running_requests = max_batch_size # Create req_to_token_pool self.req_to_token_pool = type( @@ -1240,6 +1241,7 @@ class TestDSAIndexer(CustomTestCase): backend.use_fused_topk = True backend.dsa_topk_backend = topk_backend backend.dsa_index_topk = 2048 + backend.dsa_index_kpool = 1 backend.dsa_decode_impl = "fa3" backend.req_to_token = torch.empty( 2, 4096, dtype=torch.int32, device=self.device diff --git a/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py b/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py index 30e89f8d5..fd5791f85 100644 --- a/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py +++ b/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py @@ -26,7 +26,7 @@ try: except ImportError: KERNELS_AVAILABLE = False -register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True) @@ -234,5 +234,40 @@ def test_mtp_single_step_decode(N: int): assert state_fail_rate < 0.01, f"State mismatch: fail_rate={state_fail_rate:.2f}%" +@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available") +def test_verify_scratch_pitch_uses_allocated_steps(): + # Gear below the allocated step dim must not spill into the neighbor block. + N, T, ALLOCATED = 2, 4, 8 + H, HV, K, V = 16, 32, 128, 128 + + A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors( + N, T, H, HV, K, V + ) + buffer = torch.full( + (N + 1, ALLOCATED, HV, V, K), float("nan"), dtype=torch.float32, device="cuda" + ) + + run_fused_mtp( + A_log, + dt_bias, + q, + k, + v, + a, + b, + state, + indices, + cu_seqlens, + disable_state_update=True, + intermediate_states_buffer=buffer, + intermediate_state_indices=indices, + cache_steps=T, + ) + + assert not torch.isnan(buffer[:N, :T]).any() + assert torch.isnan(buffer[N:]).all() + assert torch.isnan(buffer[:N, T:]).all() + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernels/test_lean_attention.py b/test/registered/kernels/test_lean_attention.py index 0f444c0b9..550353a44 100644 --- a/test/registered/kernels/test_lean_attention.py +++ b/test/registered/kernels/test_lean_attention.py @@ -197,12 +197,12 @@ def _run_pair_fp8(H_Q, H_KV, D, B, S, fp8_dtype, dev="cuda", seed=0): def _run_pair_paged( H_Q, H_KV, D, B, S, page_size, dev="cuda", dt=torch.float16, seed=0 ): - """Standard vs Lean on a **paged** 4-D KV buffer ``[num_pages, page_size, head, dim]``. + """Standard vs Lean with page-aware addressing over a dense 3-D KV buffer. - The KV cache is stored in pages and addressed through scattered slot ids in ``kv_indices`` - (a permutation), so the kernel's page-aware address math (``kv_loc // page_size`` / - ``kv_loc % page_size``) is genuinely exercised — not the contiguous fast path. Both arms read - the identical buffer + indices, so their outputs must agree. Returns (o_std, o_lean). + The dense ``[max_slots, head, dim]`` cache is addressed through scattered slot ids in + ``kv_indices`` (a permutation). With ``page_size > 1``, the kernel still exercises its + page-aware address math (``kv_loc // page_size`` / ``kv_loc % page_size``). Both arms read the + identical buffer + indices, so their outputs must agree. Returns (o_std, o_lean). """ torch.manual_seed(seed) D_V = D @@ -212,11 +212,9 @@ def _run_pair_paged( assert tot % page_size == 0, ( "test setup: total tokens must be a multiple of page_size" ) - num_pages = tot // page_size - - # 4-D paged KV buffers [num_pages, page_size, head, dim] (the shared-pool layout). - k = torch.randn(num_pages, page_size, H_KV, D, dtype=dt, device=dev) - v = torch.randn(num_pages, page_size, H_KV, D_V, dtype=dt, device=dev) + # Unified memory exposes dense 3-D KV views even when the allocator uses pages. + k = torch.randn(tot, H_KV, D, dtype=dt, device=dev) + v = torch.randn(tot, H_KV, D_V, dtype=dt, device=dev) kv_indptr = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32) # Scatter slots across pages so page_id/tok_in_p vary within every BLOCK_N tile. @@ -312,9 +310,9 @@ class TestLeanAttentionParity(CustomTestCase): ) def test_paged_kv_parity(self): - # Lean must read a paged 4-D KV buffer the same way the standard kernel does. Guards - # the page-aware address math (kv_loc // page_size, kv_loc % page_size); a regression - # to the contiguous-only form would scramble reads and drop cos well below 1. + # Lean must apply page-aware address math to dense KV views the same way the standard + # kernel does. A regression in kv_loc // page_size or kv_loc % page_size would scramble + # the scattered reads and drop cos well below 1. for name, H_Q, H_KV, D in GQA_SHAPES: for page_size in (16, 64): with self.subTest(model=name, page_size=page_size): diff --git a/test/registered/models_e2e/test_glm53_flash_b200.py b/test/registered/models_e2e/test_glm53_flash_b200.py new file mode 100644 index 000000000..cd9993dd0 --- /dev/null +++ b/test/registered/models_e2e/test_glm53_flash_b200.py @@ -0,0 +1,138 @@ +"""B200 per-commit coverage for the GLM-5.3-Flash serving recipes. + +Runs the Low Latency, DFlash2, and High Throughput TP4/EP4 recipes on four +B200 GPUs. All recipes must retain GSM8K accuracy; the Low Latency recipe also +checks EAGLE speculative acceptance and single-request decode performance. +""" + +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + _wait_for_gpu_idle_in_ci, + popen_launch_server, + try_cached_model, +) + +register_cuda_ci(est_time=2400, stage="base-c", runner_config="4-gpu-b200") + +MODEL_PATH = "zai-org/GLM-5.3-Flash" +DFLASH2_DRAFT_MODEL_PATH = "incoai/GLM-5.3-Flash-DFlash2" +SERVER_LAUNCH_TIMEOUT = 3600 +GPU_IDLE_TIMEOUT = 120 + +COMMON_SERVER_ARGS = [ + "--tp-size", + "4", + "--ep-size", + "4", + "--dsa-prefill-backend", + "trtllm", + "--dsa-decode-backend", + "trtllm", + "--kv-cache-dtype", + "fp8_e4m3", + "--moe-runner-backend", + "deep_gemm", + "--reasoning-parser", + "glm45", + "--tool-call-parser", + "glm47", +] + + +def _stop_server(process): + if process: + kill_process_tree(process.pid) + _wait_for_gpu_idle_in_ci(timeout=GPU_IDLE_TIMEOUT) + + +class _GLM53FlashB200Base(CustomTestCase): + server_args: list[str] + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(MODEL_PATH) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = None + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=cls.server_args, + ) + + @classmethod + def tearDownClass(cls): + _stop_server(getattr(cls, "process", None)) + + +class TestGLM53FlashB200LowLatency( + SpecDecodingMixin, + GSM8KMixin, + _GLM53FlashB200Base, +): + gsm8k_score_threshold = 0.93 + # Match the established DSA+MTP accuracy workload. The generic 200-question, + # 5-shot defaults leave a single question worth 0.5 percentage points and + # make this tight quality floor unnecessarily sensitive to kernel numerics. + gsm8k_num_examples = 500 + gsm8k_num_shots = 20 + accept_length_thres = 4.0 + bs_1_speed_thres = 250 + server_args = [ + *COMMON_SERVER_ARGS, + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "5", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "6", + "--speculative-adaptive", + ] + + +class TestGLM53FlashB200HighThroughput( + GSM8KMixin, + _GLM53FlashB200Base, +): + gsm8k_score_threshold = 0.93 + gsm8k_num_examples = 500 + gsm8k_num_shots = 20 + server_args = [ + *COMMON_SERVER_ARGS, + "--enable-dp-attention", + "--dp-size", + "4", + "--moe-a2a-backend", + "deepep", + ] + + +class TestGLM53FlashB200DFlash2( + GSM8KMixin, + _GLM53FlashB200Base, +): + gsm8k_score_threshold = 0.93 + gsm8k_num_examples = 500 + gsm8k_num_shots = 20 + server_args = [ + *COMMON_SERVER_ARGS, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + DFLASH2_DRAFT_MODEL_PATH, + "--speculative-draft-attention-backend", + "fa4", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/models_e2e/test_glm53_flash_h200.py b/test/registered/models_e2e/test_glm53_flash_h200.py new file mode 100644 index 000000000..4105d21ac --- /dev/null +++ b/test/registered/models_e2e/test_glm53_flash_h200.py @@ -0,0 +1,119 @@ +"""H200 per-commit coverage for the GLM-5.3-Flash serving recipes. + +Runs the Low Latency and High Throughput TP8/EP8 recipes on eight H200 GPUs. +Both recipes must retain GSM8K accuracy; the Low Latency recipe also checks +EAGLE speculative acceptance and single-request decode performance. +""" + +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + _wait_for_gpu_idle_in_ci, + popen_launch_server, + try_cached_model, +) + +register_cuda_ci(est_time=2400, stage="extra-b", runner_config="8-gpu-h200") + +MODEL_PATH = "zai-org/GLM-5.3-Flash" +SERVER_LAUNCH_TIMEOUT = 3600 +GPU_IDLE_TIMEOUT = 120 + +COMMON_SERVER_ARGS = [ + "--tp-size", + "8", + "--ep-size", + "8", + "--dsa-prefill-backend", + "tilelang", + "--dsa-decode-backend", + "tilelang", + "--kv-cache-dtype", + "bf16", + "--moe-runner-backend", + "deep_gemm", + "--reasoning-parser", + "glm45", + "--tool-call-parser", + "glm47", +] + + +def _stop_server(process): + if process: + kill_process_tree(process.pid) + _wait_for_gpu_idle_in_ci(timeout=GPU_IDLE_TIMEOUT) + + +class _GLM53FlashH200Base(CustomTestCase): + server_args: list[str] + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(MODEL_PATH) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = None + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=cls.server_args, + ) + + @classmethod + def tearDownClass(cls): + _stop_server(getattr(cls, "process", None)) + + +class TestGLM53FlashH200LowLatency( + SpecDecodingMixin, + GSM8KMixin, + _GLM53FlashH200Base, +): + gsm8k_score_threshold = 0.93 + # Match the established DSA+MTP accuracy workload. The generic 200-question, + # 5-shot defaults leave a single question worth 0.5 percentage points and + # make this tight quality floor unnecessarily sensitive to kernel numerics. + gsm8k_num_examples = 500 + gsm8k_num_shots = 20 + accept_length_thres = 4.0 + bs_1_speed_thres = 200 + server_args = [ + *COMMON_SERVER_ARGS, + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "5", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "6", + "--speculative-adaptive", + ] + + +class TestGLM53FlashH200HighThroughput( + GSM8KMixin, + _GLM53FlashH200Base, +): + gsm8k_score_threshold = 0.93 + gsm8k_num_examples = 500 + gsm8k_num_shots = 20 + server_args = [ + *COMMON_SERVER_ARGS, + "--enable-dp-attention", + "--dp-size", + "8", + "--moe-a2a-backend", + "deepep", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py index 2d9bd5778..cc36276bc 100644 --- a/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py +++ b/test/registered/unit/disaggregation/test_prefill_abort_result_cleanup.py @@ -60,6 +60,7 @@ class _Scheduler(SchedulerDisaggregationPrefillMixin): self.send_kv_chunk = Mock() self.output_streamer = Mock() self.metrics_reporter = SimpleNamespace(report_prefill_stats=Mock()) + self.maybe_send_health_check_signal = Mock() self.req_to_metadata_buffer_idx_allocator = Mock() self.enable_hicache_storage = True self.chunked_req = None diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 9f561ad58..c5aefc8b8 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1520,6 +1520,7 @@ if _HAS_MLX: self.kv = ReqKvInfo() self.mamba_branching_seqlen = None self.inflight_middle_chunks = 0 + self.mamba_branching_seqlen = None class FakeTpWorker: def __init__(self, next_token_ids): diff --git a/test/registered/unit/layers/attention/test_mla_decode_forced_splits.py b/test/registered/unit/layers/attention/test_mla_decode_forced_splits.py index b7d1c7c01..49825f704 100644 --- a/test/registered/unit/layers/attention/test_mla_decode_forced_splits.py +++ b/test/registered/unit/layers/attention/test_mla_decode_forced_splits.py @@ -74,16 +74,10 @@ def _inputs(seq_lens, head_num, page_size, max_kv_splits, seed): total = sum(seq_lens) n_slots = total + 64 - if page_size == 1: - pool = torch.randn( - n_slots, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen - ) - else: - n_pages = (n_slots + page_size - 1) // page_size - pool = torch.randn( - n_pages, page_size, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen - ) - n_slots = n_pages * page_size + if page_size > 1: + n_slots = ((n_slots + page_size - 1) // page_size) * page_size + # Unified memory exposes dense 3-D KV views even when the allocator uses pages. + pool = torch.randn(n_slots, 1, LK, dtype=torch.bfloat16, device=dev, generator=gen) kv_indptr = torch.zeros(batch + 1, dtype=torch.int32, device=dev) kv_indptr[1:] = torch.cumsum( diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index cf3876daa..6626b2a93 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_generation_auxiliary_output.py @@ -476,6 +476,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): disagg_prefill_inflight_queue=[], send_kv_chunk=Mock(), metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()), + maybe_send_health_check_signal=Mock(), ) with patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req"): @@ -492,6 +493,7 @@ def test_disaggregated_prefill_consumes_auxiliary_output_after_commit(): host_output, [0], ) + scheduler.maybe_send_health_check_signal.assert_called_once_with() def test_logprob_only_reuses_preprocessing_without_observer_lifecycle(): diff --git a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py index 7d90e77ef..20829388a 100644 --- a/test/registered/unit/mem_cache/test_retraction_mamba_backup.py +++ b/test/registered/unit/mem_cache/test_retraction_mamba_backup.py @@ -30,10 +30,12 @@ class _Allocator: def get_kvcache(self): return self._kv - def get_cpu_copy(self, indices, mamba_indices=None): + def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): return "kv" - def load_cpu_copy(self, cpu_tensors, indices, mamba_indices=None): + def load_cpu_copy( + self, cpu_tensors, indices, mamba_indices=None, req_pool_index=None + ): self.loaded_kv = cpu_tensors diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index b76f5420d..783ff64f4 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -124,9 +124,11 @@ def _make_model_runner( mc.get_num_kv_heads = lambda tp_size, dcp_size=1: num_kv_heads mc.get_swa_num_kv_heads = lambda tp_size: swa_num_kv_heads or num_kv_heads mc.hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"]) + mc.hf_config.model_type = "llama" mc.hf_config.get_text_config = lambda: mc.hf_config mc.linear_attn_registry_result = None mc.context_len = 8192 + mc.is_draft_model = False mr.model_config = mc mr.kv_cache_dtype = "fake_bf16" diff --git a/test/registered/unit/models/test_deepseek_nextn_mm_embed.py b/test/registered/unit/models/test_deepseek_nextn_mm_embed.py new file mode 100644 index 000000000..4890c7f37 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_nextn_mm_embed.py @@ -0,0 +1,210 @@ +"""Regression test for sgl-project/sglang#37548. + +DeepseekModelNextN.forward must use forward_batch.mm_input_embeds for multimodal +positions (where input_ids hold MM_PAD_SHIFT_VALUE+hash sentinels far above +vocab_size) instead of calling embed_tokens on those sentinel values, which +causes a CUDA index-out-of-bounds gather. +""" + +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.managers.schedule_batch import MM_PAD_SHIFT_VALUE +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=30, suite="base-a-test-cpu") + +VOCAB_SIZE = 154880 +HIDDEN_SIZE = 64 # tiny for CPU test + + +def _make_forward_batch( + input_ids: torch.Tensor, + mm_input_embeds: torch.Tensor = None, + extend_seq_lens: torch.Tensor = None, + extend_start_loc: torch.Tensor = None, + has_mm: bool = True, +): + """Build a minimal mock ForwardBatch for DeepseekModelNextN.forward.""" + fb = MagicMock() + fb.mm_input_embeds = mm_input_embeds + fb.contains_mm_inputs.return_value = has_mm + fb.forward_mode.is_extend.return_value = True + fb.forward_mode.is_draft_extend_v2.return_value = False + fb.forward_mode.is_idle.return_value = False + fb.extend_seq_lens = extend_seq_lens + fb.extend_start_loc = extend_start_loc + fb.spec_info.hidden_states = torch.randn(input_ids.shape[0], HIDDEN_SIZE) + return fb + + +def _make_model_nextn(vocab_size: int, hidden_size: int): + """Build a mock DeepseekModelNextN with a real embed_tokens layer.""" + from sglang.srt.models.deepseek_nextn import DeepseekModelNextN + + model = DeepseekModelNextN.__new__(DeepseekModelNextN) + torch.nn.Module.__init__(model) + # Minimal attributes needed by forward + model.vocab_size = vocab_size + model.embed_tokens = torch.nn.Embedding(vocab_size, hidden_size) + model.enorm = torch.nn.RMSNorm(hidden_size) + model.hnorm = torch.nn.RMSNorm(hidden_size) + model.eh_proj = torch.nn.Linear(2 * hidden_size, hidden_size, bias=False) + model.rot_weight = None + model.alt_stream = None + model.quant_config = None + model.cp_rank = None + model.cp_size = None + model.dsa_enable_prefill_cp = False + model.mla_enable_prefill_cp = False + model.mtp_block = MagicMock(side_effect=lambda **kw: (kw["hidden_states"], None)) + return model + + +class TestDeepseekNextNMmEmbed(CustomTestCase): + """DeepseekModelNextN must not call embed_tokens on MM sentinel token ids.""" + + def test_mm_sentinel_ids_do_not_cause_oob(self): + """input_ids containing MM_PAD_SHIFT_VALUE+hash must not reach embed_tokens.""" + num_tokens = 10 + mm_start, mm_end = 3, 7 # MM sentinel positions + + input_ids = torch.arange(num_tokens, dtype=torch.long) + # Insert MM sentinel values + for i in range(mm_start, mm_end): + input_ids[i] = MM_PAD_SHIFT_VALUE + i + + # Build mm_input_embeds matching the target-produced embeddings + mm_embeds = torch.randn(num_tokens, HIDDEN_SIZE) + extend_seq_lens = torch.tensor([num_tokens]) + extend_start_loc = torch.tensor([0]) + + fb = _make_forward_batch( + input_ids, + mm_input_embeds=mm_embeds.clone(), + extend_seq_lens=extend_seq_lens, + extend_start_loc=extend_start_loc, + ) + + model = _make_model_nextn(VOCAB_SIZE, HIDDEN_SIZE) + + # Use MagicMock to track embed_tokens calls + mock_embed = MagicMock(side_effect=model.embed_tokens) + object.__setattr__(model, "embed_tokens", mock_embed) + + with ( + patch( + "sglang.srt.models.deepseek_nextn.is_cp_v2_active", return_value=False + ), + patch( + "sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp", + return_value=False, + ), + patch( + "sglang.srt.models.deepseek_nextn.mla_use_prefill_cp", + return_value=False, + ), + patch( + "sglang.srt.models.deepseek_nextn.fused_eh_norm", + side_effect=lambda h, p, ew, hw, eps: torch.cat( + [model.enorm(h), model.hnorm(p)], dim=-1 + ), + ), + patch( + "sglang.srt.models.deepseek_nextn.get_global_expert_distribution_recorder" + ), + patch("sglang.srt.models.deepseek_nextn.is_cuda", False), + patch("sglang.srt.models.deepseek_nextn.is_npu", False), + patch("sglang.srt.models.deepseek_nextn.envs") as mock_envs, + patch("sglang.srt.models.deepseek_nextn.get_model") as mock_get_model, + patch("sglang.srt.models.deepseek_nextn.get_parallel") as mock_get_parallel, + patch("sglang.srt.models.deepseek_nextn.get_spec") as mock_get_spec, + ): + mock_envs.SGLANG_NPU_USE_MULTI_STREAM.get.return_value = False + mock_get_model.return_value.quantization = None + + positions = torch.arange(num_tokens, dtype=torch.long) + try: + model.forward(input_ids, positions, fb) + except Exception: + pass # We only care about embed_tokens call args + + # embed_tokens should only be called for last_indices (the appended + # next-token), not with the full input_ids containing MM sentinels. + for call in mock_embed.call_args_list: + call_ids = call[0][0] + max_id = call_ids.max().item() + self.assertLess( + max_id, + VOCAB_SIZE, + f"embed_tokens was called with id {max_id} >= vocab_size " + f"{VOCAB_SIZE}. MM sentinel values (MM_PAD_SHIFT_VALUE+hash) " + f"must not reach embed_tokens.", + ) + + def test_no_mm_falls_back_to_embed_tokens(self): + """Without mm_input_embeds, embed_tokens is called normally.""" + num_tokens = 5 + input_ids = torch.arange(num_tokens, dtype=torch.long) + + fb = _make_forward_batch( + input_ids, + mm_input_embeds=None, + extend_seq_lens=torch.tensor([num_tokens]), + extend_start_loc=torch.tensor([0]), + has_mm=False, + ) + + model = _make_model_nextn(VOCAB_SIZE, HIDDEN_SIZE) + mock_embed = MagicMock(side_effect=model.embed_tokens) + object.__setattr__(model, "embed_tokens", mock_embed) + embed_calls = mock_embed.call_args_list + + with ( + patch( + "sglang.srt.models.deepseek_nextn.is_cp_v2_active", return_value=False + ), + patch( + "sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp", + return_value=False, + ), + patch( + "sglang.srt.models.deepseek_nextn.mla_use_prefill_cp", + return_value=False, + ), + patch( + "sglang.srt.models.deepseek_nextn.fused_eh_norm", + side_effect=lambda h, p, ew, hw, eps: torch.cat( + [model.enorm(h), model.hnorm(p)], dim=-1 + ), + ), + patch( + "sglang.srt.models.deepseek_nextn.get_global_expert_distribution_recorder" + ), + patch("sglang.srt.models.deepseek_nextn.is_cuda", False), + patch("sglang.srt.models.deepseek_nextn.is_npu", False), + patch("sglang.srt.models.deepseek_nextn.envs") as mock_envs, + patch("sglang.srt.models.deepseek_nextn.get_model") as mock_get_model, + patch("sglang.srt.models.deepseek_nextn.get_parallel") as mock_get_parallel, + patch("sglang.srt.models.deepseek_nextn.get_spec") as mock_get_spec, + ): + mock_envs.SGLANG_NPU_USE_MULTI_STREAM.get.return_value = False + mock_get_model.return_value.quantization = None + + positions = torch.arange(num_tokens, dtype=torch.long) + try: + model.forward(input_ids, positions, fb) + except Exception: + pass + + # embed_tokens should be called with the full input_ids + self.assertTrue(mock_embed.call_count > 0, "embed_tokens should be called") + full_ids_call = mock_embed.call_args_list[0][0][0] + self.assertEqual(full_ids_call.numel(), num_tokens) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_glm5_next_dflash_capture.py b/test/registered/unit/models/test_glm5_next_dflash_capture.py new file mode 100644 index 000000000..4d896b4e9 --- /dev/null +++ b/test/registered/unit/models/test_glm5_next_dflash_capture.py @@ -0,0 +1,40 @@ +"""Regression for DFLASH aux-hidden capture on mHC models. + +GLM-5.3-Flash runs with mhc=True. MHCLayerCommunicator folds the residual +into the widened hidden state and returns residual=None, so CUDA-graph +capture used to crash on `hidden_states + residual`. DFLASH also has to +contract that widened state back to the draft hidden size; skipping the +contract is a silent shape/quality bug the crash-guard alone would miss. +""" + +import unittest +from types import SimpleNamespace + +import torch +from torch import nn + +from sglang.srt.models.glm5_next import Glm5NextModel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestGlm5NextDflashCapture(CustomTestCase): + def test_dflash_contracts_mhc_hidden_state_without_residual(self): + model = Glm5NextModel.__new__(Glm5NextModel) + nn.Module.__init__(model) + model.config = SimpleNamespace(mhc=True, hc_mult=4) + model.dflash_capture = True + + hidden_states = torch.arange(24, dtype=torch.float32).reshape(2, 12) + + actual = model._prepare_aux_hidden_state(hidden_states, None) + expected = hidden_states.unflatten(-1, (4, -1)).mean(dim=-2) + + torch.testing.assert_close(actual, expected) + self.assertEqual(tuple(actual.shape), (2, 3)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/parser/test_jinja_template_utils.py b/test/registered/unit/parser/test_jinja_template_utils.py index 790bcd39a..716e356cc 100644 --- a/test/registered/unit/parser/test_jinja_template_utils.py +++ b/test/registered/unit/parser/test_jinja_template_utils.py @@ -29,7 +29,6 @@ class TestTemplateContentFormatDetection(CustomTestCase): select_template = ( "{{ messages | selectattr('tool_call_id', 'equalto', 'call-a') | list }}" ) - self.assertTrue(jinja_template_may_reorder_tool_results(attribute_template)) self.assertTrue(jinja_template_may_reorder_tool_results(item_template)) self.assertTrue(jinja_template_may_reorder_tool_results(get_template)) diff --git a/test/registered/unit/spec/test_adaptive_spec_params.py b/test/registered/unit/spec/test_adaptive_spec_params.py index 6cb71f1fb..f9bcf80e8 100644 --- a/test/registered/unit/spec/test_adaptive_spec_params.py +++ b/test/registered/unit/spec/test_adaptive_spec_params.py @@ -8,12 +8,13 @@ from sglang.srt.speculative.adaptive_spec_params import ( resolve_candidate_steps_from_config, ) from sglang.test.ci.ci_register import register_cpu_ci, register_xpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=6, suite="base-a-test-cpu") register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu") -class TestAdaptiveStepSlot(unittest.TestCase): +class TestAdaptiveStepSlot(CustomTestCase): def _make_params_from_config(self, initial_steps: int, config: dict): return AdaptiveStepSlot(initial_steps=initial_steps, cfg=config) @@ -244,11 +245,11 @@ class TestAdaptiveStepSlot(unittest.TestCase): self.assertEqual(params.ceiling_coeff, 0) -class TestAdaptiveSpeculativeParams(unittest.TestCase): +class TestAdaptiveSpeculativeParams(CustomTestCase): def test_default_config_loads(self): params = AdaptiveSpeculativeParams(initial_steps=3) self.assertEqual(params._bs_list, [1, 8, 32, 64]) - self.assertEqual(params._slots[1].candidate_steps, [1, 3, 7]) + self.assertEqual(params._slots[1].candidate_steps, [1, 3, 5, 7]) self.assertEqual(params._slots[8].candidate_steps, [0, 1, 3]) self.assertEqual(params._slots[32].candidate_steps, [0, 1]) self.assertEqual(params._slots[64].candidate_steps, [0]) @@ -339,18 +340,18 @@ class TestAdaptiveSpeculativeParams(unittest.TestCase): self.assertEqual(params._slots[1].up_hysteresis, 0.1) -class TestBatchSizeRouting(unittest.TestCase): +class TestBatchSizeRouting(CustomTestCase): """BS-aware routing: batch size selects the slot, CUDA-graph BS pads first.""" def _params(self): - # Slots: bs=1 -> [1,3,7], bs=8 -> [1,3], bs=32 -> [1]. + # Slots: bs=1 -> [1,3,5,7], bs=8 -> [0,1,3], bs=32 -> [0,1]. return AdaptiveSpeculativeParams(initial_steps=3) def test_routes_to_floor_slot_without_cuda_graph(self): params = self._params() # A batch maps to the largest slot BS <= batch (floor), capped at the top slot. - self.assertEqual(params._route(1).candidate_steps, [1, 3, 7]) - self.assertEqual(params._route(7).candidate_steps, [1, 3, 7]) + self.assertEqual(params._route(1).candidate_steps, [1, 3, 5, 7]) + self.assertEqual(params._route(7).candidate_steps, [1, 3, 5, 7]) self.assertEqual(params._route(8).candidate_steps, [0, 1, 3]) self.assertEqual(params._route(31).candidate_steps, [0, 1, 3]) self.assertEqual(params._route(32).candidate_steps, [0, 1]) @@ -373,6 +374,8 @@ class TestBatchSizeRouting(unittest.TestCase): self.assertEqual(params.cuda_graph_bs_for_step(1), [4, 8, 16, 32]) # step=3 lives in the bs=1 and bs=8 slots: graphs 4,8,16 floor into them. self.assertEqual(params.cuda_graph_bs_for_step(3), [4, 8, 16]) + # step=5 lives only in the bs=1 slot: only graph BS 4 floors into it. + self.assertEqual(params.cuda_graph_bs_for_step(5), [4]) # step=7 lives only in the bs=1 slot: only graph BS 4 floors into it. self.assertEqual(params.cuda_graph_bs_for_step(7), [4]) @@ -392,12 +395,10 @@ class TestBatchSizeRouting(unittest.TestCase): self.assertEqual(params.get_steps_for_batch(32), 1) -class TestResolveCandidateSteps(unittest.TestCase): +class TestResolveCandidateSteps(CustomTestCase): def test_default_config(self): steps = resolve_candidate_steps_from_config() - self.assertIn(1, steps) - self.assertIn(3, steps) - self.assertIn(7, steps) + self.assertEqual(steps, [0, 1, 3, 5, 7]) def test_config_file(self): with tempfile.NamedTemporaryFile("w", suffix=".json") as f: diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index 2ce9e2d89..ecdc0612d 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -204,6 +204,7 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase): worker.speculative_num_steps = 1 worker.speculative_num_draft_tokens = 2 worker.device = DEVICE + worker.plan_stream = None worker.tree_mask_mode = None worker.seed_dsa_topk_from_draft_extend = seed_enabled worker.index_share_for_mtp_iteration = True