diff --git a/.github/benchmark/sglang_benchmark_models.json b/.github/benchmark/sglang_benchmark_models.json index dfc2d9ce42..4108707a45 100644 --- a/.github/benchmark/sglang_benchmark_models.json +++ b/.github/benchmark/sglang_benchmark_models.json @@ -6,6 +6,8 @@ "qwen_reasoning": "--mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", "deepseek_v4_runtime": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.9 --swa-full-tokens-ratio 0.1 --max-running-requests 256 --page-size 256 --disable-radix-cache --disable-shared-experts-fusion --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4", "deepseek_v4_prefix_cache_runtime": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.85 --swa-full-tokens-ratio 0.1 --max-running-requests 256 --page-size 256 --enable-cache-report --disable-shared-experts-fusion --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4", + "glm52_fp8_runtime": "--trust-remote-code --tp-size 4 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.8 --disable-radix-cache --model-loader-extra-config '{\"online_quant_config\":{\"global_quant_config\":\"ptpc_fp8\",\"layer_quant_config\":{\"model.layers.*.mlp.experts\":\"mxfp8\"},\"exclude_layer\":[\"lm_head\",\"model.embed_tokens\",\"*.mlp.gate\"]}}'", + "glm52_fp4_runtime": "--trust-remote-code --tp-size 4 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.8 --disable-radix-cache --model-loader-extra-config '{\"online_quant_config\":{\"global_quant_config\":\"ptpc_fp8\",\"exclude_layer\":[\"lm_head\",\"model.embed_tokens\",\"*.mlp.gate\",\"*expert*\"]}}'", "mtp1_common": "--speculative-draft-model-path SGLang/DeepSeek-R1-NextN --speculative-algorithm NEXTN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 --max-running-requests 256 --cuda-graph-bs 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 160 192 224 256", "mtp3_common": "--speculative-draft-model-path SGLang/DeepSeek-R1-NextN --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --cuda-graph-bs 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 160 192 224 256" }, @@ -16,6 +18,7 @@ "deepseek_mtp_dp_common": "SGLANG_DEFAULT_SERVER_ARGS=\nAITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1\nSGLANG_ENABLE_SPEC_V2=1\nMORI_SHMEM_MODE=ISOLATION\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nSGLANG_ENABLE_TORCH_COMPILE=1\nTORCHINDUCTOR_COMPILE_THREADS=128", "deepseek_v4_common": "SGLANG_DEFAULT_SERVER_ARGS=\nAITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1\nSGLANG_DEFAULT_THINKING=1\nSGLANG_DSV4_REASONING_EFFORT=max\nSGLANG_USE_AITER=1\nSGLANG_DSV4_FP4_EXPERTS=true\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nTORCHINDUCTOR_COMPILE_THREADS=128", "deepseek_v4_prefix_cache_common": "SGLANG_DEFAULT_SERVER_ARGS=\nAITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1\nSGLANG_DEFAULT_THINKING=1\nSGLANG_DSV4_REASONING_EFFORT=max\nSGLANG_USE_AITER=1\nSGLANG_DSV4_FP4_EXPERTS=true\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nSGLANG_ENABLE_TORCH_COMPILE=1\nTORCHINDUCTOR_COMPILE_THREADS=128", + "glm52_common": "SGLANG_DEFAULT_SERVER_ARGS=\nAITER_QUICK_REDUCE_QUANTIZATION=INT4\nAITER_USE_FLYDSL_MOE_SORTING=1\nSGLANG_USE_AITER=1\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nTORCHINDUCTOR_COMPILE_THREADS=128", "qwen_common": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0" } }, @@ -946,6 +949,32 @@ "SGLANG_AITER_FP8_PREFILL_ATTN=0" ] }, + { + "display": "GLM-5.2 FP8 TP4", + "dashboard_model": "GLM-5.2-FP8-tp4", + "workload_label": "SGLang-OOB", + "source_path": "zai-org/GLM-5.2-FP8", + "path": "zai-org/GLM-5.2-FP8", + "prefix": "glm-5-2-fp8-tp4", + "extra_args": "glm52_fp8_runtime", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", + "env_vars": "glm52_common" + }, + { + "display": "GLM-5.2 FP4 TP4", + "dashboard_model": "GLM-5.2-FP4-tp4", + "workload_label": "SGLang-OOB", + "source_path": "amd/GLM-5.2-MXFP4", + "path": "amd/GLM-5.2-MXFP4", + "prefix": "glm-5-2-fp4-tp4", + "extra_args": "glm52_fp4_runtime", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", + "env_vars": "glm52_common" + }, { "display": "Qwen3.5-397B-A17B-FP8 TP4 MI308", "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4 MI308", diff --git a/.github/workflows/atom-sglang-benchmark.yaml b/.github/workflows/atom-sglang-benchmark.yaml index bb1fbf5cc3..5a3a1f8d1e 100644 --- a/.github/workflows/atom-sglang-benchmark.yaml +++ b/.github/workflows/atom-sglang-benchmark.yaml @@ -52,14 +52,16 @@ on: - "qwen3-32b-fp8-tp8-mi308 (1024x1024/8192x1024: [4,8,16,32,64])" - "deepseek-v3-2-fp8-tp8 (1024x1024/8192x1024: [4,8,16,32,64])" - "glm-5-1-fp8-tp8 (1024x1024/8192x1024: [4,8,16,32,64])" - - "all-deepseek (19 DeepSeek configs)" - - "all-deepseek-non-mtp (13 DeepSeek non-MTP configs)" + - "glm-5-2-fp8-tp4 (1024x1024/8192x1024: [4,8,16,32,64])" + - "glm-5-2-fp4-tp4 (1024x1024/8192x1024: [4,8,16,32,64])" + - "all-deepseek (20 DeepSeek configs)" + - "all-deepseek-non-mtp (14 DeepSeek non-MTP configs)" - "all-deepseek-mtp (6 DeepSeek MTP configs x 10 default params)" - "all-deepseek-v3-2 (4 DeepSeek-V3.2 FP8 OOB configs x 10 default params)" - "all-deepseek-v3-2_glm_qwen (DeepSeek-V3.2 + GLM + Qwen OOB configs x 10 default params)" - - "all-qwen (7 Qwen configs x 10 default params)" - - "all-glm (1 GLM config x 10 default params)" - - "all-oob (27 SGLang-OOB configs x 10 default params)" + - "all-qwen (8 Qwen configs x 10 default params)" + - "all-glm (3 GLM configs x 10 default params)" + - "all-oob (31 SGLang-OOB configs x 10 default params)" default: "none (do not run SGLang-OOB models)" mesh_config_preset: description: "SGLang-Mesh config subset (ignored for SGLang-OOB)" @@ -715,6 +717,8 @@ jobs: "deepseek-v4-pro-mtp1-tp8", "deepseek-v4-pro-mtp3-tp8", "glm-5-1-fp8-tp8", + "glm-5-2-fp8-tp4", + "glm-5-2-fp4-tp4", ] OOB_P1_PREFIX_ORDER = [ "deepseek-v3-2-fp8-tp4", diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ea846ebed3..4b49121009 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -2,6 +2,7 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import logging +import os from dataclasses import dataclass from functools import partial as functools_partial from typing import Optional @@ -290,6 +291,52 @@ def __init__( "aiter or set ATOM_MLA_PAGE_SIZE=1." ) + # Tiny graph-safe scratch for GLM target-verify debugging. During CUDA + # graph replay Python does not run, but captured device copies still + # update this buffer; the SGLang graph wrapper reads it after replay. + self.register_buffer( + "_atom_glm52_attn_debug", + torch.zeros((4, 8), dtype=torch.float32), + persistent=False, + ) + + def _atom_glm52_debug_layer_enabled(self) -> bool: + configured = os.environ.get("ATOM_GLM52_ATTENTION_DEBUG_LAYERS") + if not configured: + return False + if configured.strip().lower() in ("all", "*"): + return True + try: + layers = { + int(item) + for item in configured.replace(" ", ",").split(",") + if item.strip() + } + except ValueError: + return False + return int(self.layer_num) in layers + + def _atom_glm52_write_attn_debug(self, stage: int, tensor: torch.Tensor) -> None: + if not self._atom_glm52_debug_layer_enabled() or not torch.is_tensor(tensor): + return + if tensor.numel() == 0: + return + try: + flat = tensor.reshape(tensor.shape[0], -1).float() + row = flat[0] + dim = int(row.numel()) + buf = self._atom_glm52_attn_debug + out = buf[int(stage)] + out[0].fill_(float(self.layer_num)) + out[1].fill_(float(stage)) + out[2].fill_(float(tensor.shape[0])) + out[3 : 3 + min(5, dim)].copy_(row[: min(5, dim)]) + if dim < 5: + out[3 + dim :].zero_() + except Exception: + if not torch.cuda.is_current_stream_capturing(): + logger.exception("Failed to write GLM52 attention debug buffer") + def _seg_kv_cache_view(self, kv_cache: torch.Tensor) -> torch.Tensor: """Reshape the KV cache buffer into the page-level flat seg layout ``[num_blocks, page_size*(kv_lora_rank + qk_rope_head_dim)]`` that the @@ -511,6 +558,8 @@ def _forward_prefill_cached_single_pass( getattr(attn_metadata, "shuffle_kv_block_indptr", None), getattr(attn_metadata, "shuffle_kv_block_indices", None), ) + self._atom_glm52_write_attn_debug(2, k_full) + self._atom_glm52_write_attn_debug(3, v_full) output = flash_attn_varlen_func( q=prefill_q, k=k_full, @@ -605,7 +654,7 @@ def _forward_prefill_cached_chunked( n = MLAAttention._chunked_prefill_calls = ( getattr(MLAAttention, "_chunked_prefill_calls", 0) + 1 ) - if n == 1 or n % 500 == 0: + if (n == 1 or n % 500 == 0) and not torch.cuda.is_current_stream_capturing(): logger.info( "MLA chunked-prefill #%d: layer=%d num_chunks=%d " "total_kv=%s cu_seqlens_q[-1]=%d", @@ -1149,6 +1198,7 @@ def forward_impl( ) prefill_q_pe = prefill_q[..., self.qk_nope_head_dim :] self.rotary_emb(positions, prefill_q_pe, k_rope) + self._atom_glm52_write_attn_debug(0, prefill_q) if kv_cache.numel() > 0: if envs.ATOM_USE_TRITON_MLA and envs.ATOM_USE_TRITON_MLA_SHUFFLE_KV: @@ -1205,6 +1255,7 @@ def forward_impl( output = self._forward_prefill_mha( prefill_q, k_nope, k_rope, kv_cache, attn_metadata ) + self._atom_glm52_write_attn_debug(1, output) else: q_nope, q_rope = self._q_proj_and_k_up_proj(q, x_scale=q_scale) @@ -1292,11 +1343,13 @@ def forward_impl( is_nope_first=True, ) # q_out = self.fused_kv_bmm(q, q_scale, k_nope, k_rope, positions, kv_cache, attn_metadata) + self._atom_glm52_write_attn_debug(2, q_out) if context.is_prefill: output = self._forward_prefill_mla(q_out, kv_cache, attn_metadata) else: output = self._forward_decode(q_out, kv_cache, attn_metadata) + self._atom_glm52_write_attn_debug(3, output) return output diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index dff057eb8e..4a2828792a 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -170,6 +170,16 @@ def compute_logits( logits = mtp_layer.shared_head.head(mtp_layer.shared_head(hidden_states)) return logits + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + normed = mtp_layer.shared_head(hidden_states) + return mtp_layer.shared_head.head.compute_argmax_token(normed) + @support_torch_compile class DeepSeekMTP(nn.Module): @@ -259,6 +269,13 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model.compute_draft_token(hidden_states, spec_step_idx) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 3471b4155c..5c3d8db3f7 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -19,6 +19,15 @@ logger = logging.getLogger("atom") + +def _is_current_stream_capturing(torch_mod) -> bool: + try: + cuda_mod = getattr(torch_mod, "cuda", None) + is_capturing = getattr(cuda_mod, "is_current_stream_capturing", None) + return bool(is_capturing()) if is_capturing is not None else False + except Exception: + return False + _ATOM_SUPPORTED_MODELS = { "Qwen3ForCausalLM": Qwen3ForCausalLM, "Qwen3MoeForCausalLM": Qwen3MoeForCausalLM, @@ -76,6 +85,10 @@ def _register_custom_attention_to_sglang() -> None: from atom.plugin.sglang.attention_backend.deepseek_v4_backend import ( ATOMDeepseekV4BackendForSgl, ) + from atom.plugin.sglang.attention_backend.glm52_dsa_backend import ( + ATOMGLM52DSABackendForSgl, + ) + from atom.plugin.sglang.runtime import is_glm52_dsa_config # here register the custom attention backend with the name "aiter" # as sglang defines the fixed attention backend choices, which must be @@ -90,12 +103,18 @@ def _register_custom_attention_to_sglang() -> None: @register_attention_backend("aiter") def create_atom_backend(runner): - arches = getattr(runner.model_config.hf_config, "architectures", None) or [] + hf_config = runner.model_config.hf_config + arches = getattr(hf_config, "architectures", None) or [] if any("DeepseekV4" in str(arch) for arch in arches): logger.info( "Use ATOMDeepseekV4BackendForSgl for DeepSeek-V4 through SGLang aiter backend choice" ) return ATOMDeepseekV4BackendForSgl(runner) + if is_glm52_dsa_config(hf_config): + logger.info( + "Use ATOMGLM52DSABackendForSgl for GLM-5.2 through SGLang aiter backend choice" + ) + return ATOMGLM52DSABackendForSgl(runner) return ATOMAttnBackendForSgl(runner) @register_attention_backend("dsv4") @@ -105,6 +124,18 @@ def create_dsv4_backend(runner): ) return ATOMDeepseekV4BackendForSgl(runner) + @register_attention_backend("nsa") + def create_atom_nsa_backend(runner): + hf_config = runner.model_config.hf_config + if is_glm52_dsa_config(hf_config): + logger.info( + "Use ATOMGLM52DSABackendForSgl for GLM-5.2 through SGLang nsa backend choice" + ) + return ATOMGLM52DSABackendForSgl(runner) + from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend + + return NativeSparseAttnBackend(runner) + def _patch_sglang_dsv4_draft_backends() -> None: """Route SGLang's hard-coded DSV4 speculative factories to ATOM. @@ -156,6 +187,7 @@ def _patch_sglang_dsv4_spec_cuda_graph() -> None: """ try: + import torch from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, @@ -163,10 +195,16 @@ def _patch_sglang_dsv4_spec_cuda_graph() -> None: from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, ) - from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker + from sglang.srt.speculative.eagle_worker_v2 import ( + EAGLEWorkerV2, + EagleDraftWorker, + ) except Exception as exc: logger.debug("Skip patching SGLang DSV4 spec cuda graph: %s", exc) return + from atom.plugin.sglang.runtime.model_arch import ( + is_glm52_dsa_config as _is_glm52_dsa_config, + ) def _is_dsv4_nextn_runner(runner) -> bool: try: @@ -182,6 +220,24 @@ def _is_dsv4_nextn_runner(runner) -> bool: except Exception: return False + def _is_glm52_nextn_runner(runner) -> bool: + try: + arches = ( + getattr( + getattr(getattr(runner, "model_config", None), "hf_config", None), + "architectures", + None, + ) + or [] + ) + return any( + "GlmMoeDsaForCausalLMNextN" in str(arch) + or "DeepseekV3ForCausalLMNextN" in str(arch) + for arch in arches + ) + except Exception: + return False + def _is_dsv4_runner(runner) -> bool: try: arches = ( @@ -196,6 +252,24 @@ def _is_dsv4_runner(runner) -> bool: except Exception: return False + def _is_glm52_runner(runner) -> bool: + try: + hf_config = getattr(getattr(runner, "model_config", None), "hf_config", None) + arches = getattr(hf_config, "architectures", None) or [] + model_path = str( + getattr(getattr(runner, "server_args", None), "model_path", "") + or getattr(getattr(runner, "model_config", None), "path", "") + ) + return ( + _is_glm52_dsa_config(hf_config) + or any("GlmMoeDsa" in str(arch) for arch in arches) + or "GLM-5.2" in model_path + or "glm-5.2" in model_path.lower() + ) + except Exception: + return False + + def _flatten_spec_hidden_states(forward_batch): spec_info = getattr(forward_batch, "spec_info", None) hidden_states = getattr(spec_info, "hidden_states", None) @@ -245,6 +319,37 @@ def _draft_extend_graph_enabled(runner) -> bool: _is_dsv4_nextn_runner(runner) and _is_dsv4_flash_runner(runner) ) + def _tensor_head(tensor, limit=12, as_float=False): + if not torch.is_tensor(tensor): + return None + value = tensor.reshape(-1)[: min(limit, int(tensor.numel()))].detach() + if as_float: + value = value.float() + return value.cpu().tolist() + + def _hidden_probe(hidden_states, rows=4): + if not torch.is_tensor(hidden_states): + return None + hidden_rows = hidden_states[: min(rows, int(hidden_states.shape[0]))].detach().float() + dim = int(hidden_rows.shape[-1]) + checksum_slices = [] + for start in (0, 256, 1024, 2048, 4096, max(0, dim - 256)): + end = min(dim, start + 256) + if start < end: + checksum_slices.append(hidden_rows[:, start:end].sum(dim=-1)) + checksum = ( + torch.stack(checksum_slices, dim=-1) + if checksum_slices + else hidden_rows.new_empty((hidden_rows.shape[0], 0)) + ) + return { + "shape": tuple(hidden_states.shape), + "norm": hidden_rows.norm(dim=-1).cpu().tolist(), + "absmax": hidden_rows.abs().amax(dim=-1).cpu().tolist(), + "mean": hidden_rows.mean(dim=-1).cpu().tolist(), + "checksum": checksum.cpu().tolist(), + } + def _target_verify_graph_enabled() -> bool: return _env_flag("ATOM_SGLANG_V4_ENABLE_TARGET_VERIFY_CG") and not _env_flag( "ATOM_SGLANG_V4_DISABLE_TARGET_VERIFY_CG" @@ -257,6 +362,45 @@ def _safe_spec_graph_bs(original_bs, env_name: str): allowed = {int(x) for x in configured.replace(" ", ",").split(",") if x.strip()} return [bs for bs in original_bs if int(bs) in allowed] + def _runner_probe(obj): + if obj is None: + return None + return { + "type": f"{type(obj).__module__}.{type(obj).__qualname__}", + "has_load_batch": callable(getattr(obj, "load_batch", None)), + "has_execute": callable(getattr(obj, "execute", None)), + "has_replay": callable(getattr(obj, "replay", None)), + "has_can_run": callable(getattr(obj, "can_run", None)), + "has_can_run_graph": callable(getattr(obj, "can_run_graph", None)), + "has_capture_bs": hasattr(obj, "capture_bs"), + "has_output_buffers": hasattr(obj, "output_buffers"), + "has_buffers": hasattr(obj, "buffers"), + } + + def _log_glm52_graph_runner_probe(where: str, model_runner, self_runner=None) -> None: + if not _env_flag("ATOM_GLM52_GRAPH_RUNNER_PROBE"): + return + try: + attrs = {} + for name in ( + "graph_runner", + "cuda_graph_runner", + "decode_cuda_graph_runner", + "draft_cuda_graph_runner", + "cuda_graph_runner_for_draft_extend", + ): + if hasattr(model_runner, name): + attrs[name] = _runner_probe(getattr(model_runner, name, None)) + logger.warning( + "GLM52 graph runner probe: where=%s self=%s model_runner=%s attrs=%s", + where, + _runner_probe(self_runner), + _runner_probe(model_runner), + attrs, + ) + except Exception: + logger.exception("Failed to log GLM52 graph runner probe") + if not getattr(CudaGraphRunner, "_atom_dsv4_init_patched", False): original_target_init = CudaGraphRunner.__init__ @@ -268,6 +412,12 @@ def __init__(self, model_runner, *args, **kwargs): if server_args is not None else None ) + should_force_glm_hidden = False + original_enable_return_hidden_states = ( + getattr(server_args, "enable_return_hidden_states", None) + if server_args is not None + else None + ) try: should_cap = _is_dsv4_runner(model_runner) and bool( getattr( @@ -281,8 +431,20 @@ def __init__(self, model_runner, *args, **kwargs): and not getattr(model_runner, "is_draft_worker", False) and _target_verify_graph_enabled() ) + should_force_glm_hidden = ( + _is_glm52_runner(model_runner) + and bool( + getattr( + getattr(model_runner, "spec_algorithm", None), + "is_speculative", + lambda: False, + )() + ) + and not getattr(model_runner, "is_draft_worker", False) + ) except Exception: should_cap = False + should_force_glm_hidden = False try: if should_cap and server_args is not None and original_cuda_graph_bs: @@ -290,7 +452,16 @@ def __init__(self, model_runner, *args, **kwargs): original_cuda_graph_bs, "ATOM_SGLANG_V4_TARGET_VERIFY_CG_BS", ) + if should_force_glm_hidden and server_args is not None: + # GLM MTP needs verifier hidden states to seed the next + # draft step. Capture target graph in FULL hidden mode + # from startup so graph replay matches eager semantics. + server_args.enable_return_hidden_states = True original_target_init(self, model_runner, *args, **kwargs) + if should_force_glm_hidden: + _log_glm52_graph_runner_probe( + "CudaGraphRunner.__init__", model_runner, self + ) finally: if ( should_cap @@ -298,6 +469,14 @@ def __init__(self, model_runner, *args, **kwargs): and original_cuda_graph_bs is not None ): server_args.cuda_graph_bs = original_cuda_graph_bs + if ( + should_force_glm_hidden + and server_args is not None + and original_enable_return_hidden_states is not None + ): + server_args.enable_return_hidden_states = ( + original_enable_return_hidden_states + ) CudaGraphRunner.__init__ = __init__ CudaGraphRunner._atom_dsv4_init_patched = True @@ -333,11 +512,232 @@ def can_run(self, forward_batch): CudaGraphRunner.can_run = can_run CudaGraphRunner._atom_dsv4_spec_can_run_patched = True + if not getattr(CudaGraphRunner, "_atom_glm52_io_debug_patched", False): + original_load_batch = getattr(CudaGraphRunner, "load_batch", None) + original_execute = getattr(CudaGraphRunner, "execute", None) + original_replay = getattr(CudaGraphRunner, "replay", None) + + def _mode_is_target_verify(forward_batch) -> bool: + mode = getattr(forward_batch, "forward_mode", None) + return bool(mode is not None and getattr(mode, "is_target_verify", lambda: False)()) + + def _should_log_glm52_target_graph(self, forward_batch) -> bool: + return ( + ( + _env_flag("ATOM_GLM52_VERIFY_DEBUG") + or _env_flag("ATOM_GLM52_ATTENTION_DEBUG_LOG") + ) + and _is_glm52_runner(getattr(self, "model_runner", None)) + and _mode_is_target_verify(forward_batch) + ) + + def _log_glm52_target_graph_input(self, forward_batch, where: str) -> None: + buffers = getattr(self, "buffers", None) + logger.info( + "GLM52 target graph input debug: where=%s raw_bs=%s bs=%s " + "raw_tokens=%s input_ids=%s positions=%s out_cache=%s " + "seq_lens=%s req_pool=%s spec_positions=%s", + where, + getattr(self, "raw_bs", None), + getattr(self, "bs", None), + getattr(self, "raw_num_token", None), + _tensor_head(getattr(buffers, "input_ids", None)), + _tensor_head(getattr(buffers, "positions", None)), + _tensor_head(getattr(buffers, "out_cache_loc", None)), + _tensor_head(getattr(buffers, "seq_lens", None)), + _tensor_head(getattr(buffers, "req_pool_indices", None)), + _tensor_head( + getattr(getattr(forward_batch, "spec_info", None), "positions", None) + ), + ) + + def _log_glm52_target_graph_output(out, where: str) -> None: + logits = getattr(out, "next_token_logits", None) + logger.info( + "GLM52 target graph output debug: where=%s hidden=%s logits_shape=%s " + "logits_head=%s", + where, + _hidden_probe(getattr(out, "hidden_states", None)), + tuple(logits.shape) if torch.is_tensor(logits) else None, + _tensor_head(logits, limit=8, as_float=True), + ) + + def _log_glm52_attention_debug_buffers( + model_runner, where: str, forward_batch=None, runner=None + ) -> None: + if not _env_flag("ATOM_GLM52_ATTENTION_DEBUG_LOG"): + return + try: + buffers = getattr(runner, "buffers", None) + spec_info = getattr(forward_batch, "spec_info", None) + context = { + "input_ids": _tensor_head(getattr(forward_batch, "input_ids", None)), + "positions": _tensor_head( + getattr(spec_info, "positions", None) + if getattr(spec_info, "positions", None) is not None + else getattr(forward_batch, "positions", None) + ), + "out_cache_loc": _tensor_head( + getattr(forward_batch, "out_cache_loc", None) + ), + "seq_lens": _tensor_head(getattr(forward_batch, "seq_lens", None)), + "req_pool_indices": _tensor_head( + getattr(forward_batch, "req_pool_indices", None) + ), + "buffer_input_ids": _tensor_head(getattr(buffers, "input_ids", None)) + if buffers is not None + else None, + "buffer_positions": _tensor_head(getattr(buffers, "positions", None)) + if buffers is not None + else None, + "buffer_out_cache_loc": _tensor_head( + getattr(buffers, "out_cache_loc", None) + ) + if buffers is not None + else None, + "buffer_seq_lens": _tensor_head(getattr(buffers, "seq_lens", None)) + if buffers is not None + else None, + "buffer_req_pool_indices": _tensor_head( + getattr(buffers, "req_pool_indices", None) + ) + if buffers is not None + else None, + "raw_bs": getattr(runner, "raw_bs", None), + "bs": getattr(runner, "bs", None), + "raw_tokens": getattr(runner, "raw_num_token", None), + } + positions_head = context.get("positions") or [] + input_ids_head = context.get("input_ids") or [] + raw_bs = int(context.get("raw_bs") or 0) + if raw_bs <= 0 or ( + positions_head + and all(int(x) == 0 for x in positions_head) + and input_ids_head + and all(int(x) == 0 for x in input_ids_head) + ): + return + configured = os.environ.get("ATOM_GLM52_ATTENTION_DEBUG_LAYERS", "") + configured_layers = None + if configured.strip().lower() not in ("all", "*"): + configured_layers = { + int(item) + for item in configured.replace(" ", ",").split(",") + if item.strip() + } + model = getattr(model_runner, "model", None) + if model is None or not hasattr(model, "modules"): + return + collected = [] + for module in model.modules(): + buf = getattr(module, "_atom_glm52_attn_debug", None) + if not torch.is_tensor(buf): + continue + layer = getattr(module, "layer_num", None) + if ( + configured_layers is not None + and int(layer) not in configured_layers + ): + continue + sparse_buf = getattr(module, "sparse_kv_indices_buffer", None) + sparse_info = None + if torch.is_tensor(sparse_buf): + flat_sparse = sparse_buf.reshape(-1) + sparse_info = { + "shape": tuple(sparse_buf.shape), + "head": flat_sparse[: min(16, int(flat_sparse.numel()))] + .detach() + .cpu() + .tolist(), + "tail": flat_sparse[ + max(0, int(flat_sparse.numel()) - 16) : + ] + .detach() + .cpu() + .tolist(), + } + collected.append( + { + "layer": int(layer) if layer is not None else None, + "values": buf.detach().cpu().tolist(), + "sparse": sparse_info, + } + ) + if collected: + logger.info( + "GLM52 attention layer debug: where=%s context=%s values=%s", + where, + context, + collected, + ) + except Exception: + logger.exception("Failed to log GLM52 attention layer debug") + + def load_batch(self, forward_batch, *args, **kwargs): + ret = original_load_batch(self, forward_batch, *args, **kwargs) + try: + if _should_log_glm52_target_graph(self, forward_batch): + _log_glm52_target_graph_input(self, forward_batch, "load_batch") + except Exception: + logger.exception("Failed to log GLM52 target graph input debug") + return ret + + def execute(self, forward_batch, *args, **kwargs): + out = original_execute(self, forward_batch, *args, **kwargs) + try: + if _should_log_glm52_target_graph(self, forward_batch): + _log_glm52_target_graph_output(out, "execute") + except Exception: + logger.exception("Failed to log GLM52 target graph output debug") + return out + + def replay(self, forward_batch, *args, **kwargs): + should_log = False + try: + should_log = _should_log_glm52_target_graph(self, forward_batch) + if should_log: + _log_glm52_graph_runner_probe( + "CudaGraphRunner.replay", + getattr(self, "model_runner", None), + self, + ) + _log_glm52_target_graph_input(self, forward_batch, "replay") + except Exception: + logger.exception("Failed to log GLM52 target graph replay input debug") + out = original_replay(self, forward_batch, *args, **kwargs) + try: + if should_log: + _log_glm52_target_graph_output(out, "replay") + _log_glm52_attention_debug_buffers( + getattr(self, "model_runner", None), + "replay", + forward_batch=forward_batch, + runner=self, + ) + except Exception: + logger.exception("Failed to log GLM52 target graph replay output debug") + return out + + if original_load_batch is not None and original_execute is not None: + CudaGraphRunner.load_batch = load_batch + CudaGraphRunner.execute = execute + if original_replay is not None: + CudaGraphRunner.replay = replay + if ( + (original_load_batch is not None and original_execute is not None) + or original_replay is not None + ): + CudaGraphRunner._atom_glm52_io_debug_patched = True + if not getattr(EAGLEDraftCudaGraphRunner, "_atom_dsv4_replay_patched", False): original_draft_replay = EAGLEDraftCudaGraphRunner.replay + original_draft_execute = getattr(EAGLEDraftCudaGraphRunner, "execute", None) def replay(self, forward_batch): - if not _is_dsv4_nextn_runner(getattr(self, "model_runner", None)): + if not ( + _is_dsv4_nextn_runner(getattr(self, "model_runner", None)) + or _is_glm52_nextn_runner(getattr(self, "model_runner", None)) + ): return original_draft_replay(self, forward_batch) if _env_flag("ATOM_SGLANG_V4_DISABLE_DRAFT_CG"): raise RuntimeError( @@ -346,14 +746,341 @@ def replay(self, forward_batch): ) original_hidden_states = _flatten_spec_hidden_states(forward_batch) try: - return original_draft_replay(self, forward_batch) + try: + if os.path.exists("/tmp/atom_glm52_draft_debug_on"): + import torch + + spec_info = getattr(forward_batch, "spec_info", None) + logger.info( + "GLM52 draft graph seed debug: bs=%s topk_index_head=%s " + "topk_p_head=%s hidden_probe=%s out_cache_loc_head=%s " + "positions_head=%s", + getattr(forward_batch, "batch_size", None), + _tensor_head(getattr(spec_info, "topk_index", None)), + _tensor_head(getattr(spec_info, "topk_p", None), as_float=True), + _hidden_probe(getattr(spec_info, "hidden_states", None)), + _tensor_head(getattr(forward_batch, "out_cache_loc", None)), + _tensor_head(getattr(forward_batch, "positions", None)), + ) + except Exception: + logger.exception("Failed to log GLM52 draft graph seed debug") + out = original_draft_replay(self, forward_batch) + try: + if os.path.exists("/tmp/atom_glm52_draft_debug_on"): + if len(out) == 4: + parent_list, top_scores_index, draft_tokens, draft_probs = out + else: + parent_list, top_scores_index, draft_tokens = out + draft_probs = None + + def _head(tensor, rows=4): + if not torch.is_tensor(tensor): + return None + return ( + tensor[: min(rows, int(tensor.shape[0]))] + .detach() + .cpu() + .tolist() + ) + + logger.info( + "GLM52 draft graph replay debug: parent_head=%s " + "top_scores_index_head=%s draft_tokens_head=%s " + "draft_probs_head=%s", + _head(parent_list), + _head(top_scores_index), + _head(draft_tokens), + _head(draft_probs), + ) + except Exception: + logger.exception("Failed to log GLM52 draft graph replay debug") + return out finally: if original_hidden_states is not None: forward_batch.spec_info.hidden_states = original_hidden_states EAGLEDraftCudaGraphRunner.replay = replay + if original_draft_execute is not None: + + def execute(self, forward_batch): + out = original_draft_execute(self, forward_batch) + try: + if os.path.exists("/tmp/atom_glm52_draft_debug_on"): + parent_list, top_scores_index, draft_tokens, draft_probs = out + + def _head(tensor, rows=4): + if not torch.is_tensor(tensor): + return None + return ( + tensor[: min(rows, int(tensor.shape[0]))] + .detach() + .cpu() + .tolist() + ) + + logger.info( + "GLM52 draft graph execute debug: raw_bs=%s bs=%s " + "parent_head=%s top_scores_index_head=%s " + "draft_tokens_head=%s draft_probs_head=%s", + getattr(self, "raw_bs", None), + getattr(self, "bs", None), + _head(parent_list), + _head(top_scores_index), + _head(draft_tokens), + _head(draft_probs), + ) + except Exception: + logger.exception("Failed to log GLM52 draft graph execute debug") + return out + + EAGLEDraftCudaGraphRunner.execute = execute EAGLEDraftCudaGraphRunner._atom_dsv4_replay_patched = True + try: + from sglang.srt.speculative.eagle_info import EagleVerifyInput + except Exception: + EagleVerifyInput = None + + if EagleVerifyInput is not None and not getattr( + EagleVerifyInput, "_atom_glm52_sample_debug_patched", False + ): + original_sample = EagleVerifyInput.sample + + def sample(self, batch, logits_output, vocab_mask=None): + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ( + "1", + "true", + "True", + ): + try: + import torch + + bs = len(batch.seq_lens) + next_token_logits = logits_output.next_token_logits + target_predict = torch.argmax(next_token_logits, dim=-1).reshape( + bs, self.draft_token_num + ) + candidates = self.draft_token.reshape(bs, self.draft_token_num) + capturing = _is_current_stream_capturing(torch) + logits_probe = "" + cand_logits_probe = "" + hidden_probe = "" + metadata_probe = "" + if not capturing: + probe_rows = min( + int(next_token_logits.shape[0]), + max(1, min(2, bs)) * int(self.draft_token_num), + ) + top_vals, top_ids = torch.topk( + next_token_logits[:probe_rows], k=3, dim=-1 + ) + cand_flat = candidates.reshape(-1)[:probe_rows].to( + next_token_logits.device + ) + cand_logits = next_token_logits[:probe_rows].gather( + 1, cand_flat[:, None] + ) + logits_probe = { + "top_ids": top_ids.detach().cpu().tolist(), + "top_vals": top_vals.detach().float().cpu().tolist(), + } + cand_logits_probe = cand_logits.detach().float().cpu().tolist() + hidden_states = getattr(logits_output, "hidden_states", None) + if torch.is_tensor(hidden_states): + hidden_rows = hidden_states[:probe_rows].detach().float() + dim = int(hidden_rows.shape[-1]) + checksum_slices = [] + for start in (0, 256, 1024, 2048, 4096, max(0, dim - 256)): + end = min(dim, start + 256) + if start < end: + checksum_slices.append( + hidden_rows[:, start:end].sum(dim=-1) + ) + checksum = ( + torch.stack(checksum_slices, dim=-1) + if checksum_slices + else hidden_rows.new_empty((hidden_rows.shape[0], 0)) + ) + sample_cols = [ + c for c in (0, 1, 2, 3, 7, 31, 127, 511, 1023, 2047, 4095, dim - 1) + if 0 <= c < dim + ] + hidden_probe = { + "shape": tuple(hidden_states.shape), + "norm": hidden_rows.norm(dim=-1).cpu().tolist(), + "absmax": hidden_rows.abs().amax(dim=-1).cpu().tolist(), + "mean": hidden_rows.mean(dim=-1).cpu().tolist(), + "checksum": checksum.cpu().tolist(), + "sample_cols": sample_cols, + "sample_vals": hidden_rows[:, sample_cols] + .cpu() + .tolist() + if sample_cols + else [], + } + else: + hidden_probe = None + metadata_probe = { + "counter": getattr( + self, "_atom_glm52_verify_counter", None + ), + "row_probe": getattr(self, "_atom_glm52_row_probe", None), + } + logger.info( + "GLM52 verify sample debug: bs=%s draft_token_num=%s " + "logits_shape=%s candidates_head=%s target_predict_head=%s " + "seq_lens_head=%s top3_probe=%s cand_logits_probe=%s " + "hidden_probe=%s metadata_probe=%s", + bs, + self.draft_token_num, + tuple(next_token_logits.shape), + "" + if capturing + else candidates[: min(2, bs)].detach().cpu().tolist(), + "" + if capturing + else target_predict[: min(2, bs)].detach().cpu().tolist(), + ( + "" + if capturing + else batch.seq_lens[: min(8, int(batch.seq_lens.numel()))] + .detach() + .cpu() + .tolist() + ) + if torch.is_tensor(batch.seq_lens) + else None, + logits_probe, + cand_logits_probe, + hidden_probe, + metadata_probe, + ) + except Exception: + logger.exception("Failed to log GLM52 verify sample debug") + return original_sample(self, batch, logits_output, vocab_mask) + + EagleVerifyInput.sample = sample + EagleVerifyInput._atom_glm52_sample_debug_patched = True + + if not getattr(EAGLEWorkerV2, "_atom_glm52_verify_kv_debug_patched", False): + original_verify = EAGLEWorkerV2.verify + + def verify(self, batch): + forced_tokens = os.environ.get("ATOM_GLM52_FORCE_VERIFY_DRAFT_TOKENS", "") + if forced_tokens.strip(): + try: + target_runner = getattr( + getattr(self, "target_worker", None), "model_runner", None + ) + if _is_glm52_runner(target_runner): + tokens = [ + int(item) + for item in forced_tokens.replace(" ", ",").split(",") + if item.strip() + ] + verify_input = getattr(batch, "spec_info", None) + draft_token = getattr(verify_input, "draft_token", None) + draft_token_num = int( + getattr(verify_input, "draft_token_num", 0) + or len(tokens) + ) + bs = int(getattr(batch, "batch_size", 0) or len(batch.seq_lens)) + if ( + tokens + and torch.is_tensor(draft_token) + and draft_token_num > 0 + and int(draft_token.numel()) >= bs * draft_token_num + ): + row = torch.tensor( + tokens[:draft_token_num], + dtype=draft_token.dtype, + device=draft_token.device, + ) + if int(row.numel()) < draft_token_num: + row = torch.nn.functional.pad( + row, + (0, draft_token_num - int(row.numel())), + value=int(row[-1].item()), + ) + draft_token.view(bs, draft_token_num)[:, :].copy_( + row[None, :] + ) + logger.info( + "GLM52 forced verify draft tokens: bs=%s " + "draft_token_num=%s tokens=%s", + bs, + draft_token_num, + row.detach().cpu().tolist(), + ) + except Exception: + logger.exception("Failed to force GLM52 verify draft tokens") + out = original_verify(self, batch) + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ( + "1", + "true", + "True", + ): + try: + req_pool_indices = getattr(batch, "req_pool_indices", None) + req_to_token_pool = getattr(self, "req_to_token_pool", None) + req_to_token = getattr(req_to_token_pool, "req_to_token", None) + seq_lens = getattr(batch, "seq_lens", None) + new_seq_lens = getattr(out, "new_seq_lens", None) + tail_probe = [] + if ( + torch.is_tensor(req_pool_indices) + and torch.is_tensor(req_to_token) + and torch.is_tensor(seq_lens) + ): + probe_bs = min(4, int(req_pool_indices.numel())) + for row in range(probe_bs): + req_idx = int(req_pool_indices[row].detach().cpu()) + old_len = int(seq_lens[row].detach().cpu()) + new_len = ( + int(new_seq_lens[row].detach().cpu()) + if torch.is_tensor(new_seq_lens) + else old_len + ) + start = max(0, old_len - 8) + end = min(req_to_token.shape[1], new_len + 4) + tail_probe.append( + { + "row": row, + "req": req_idx, + "old_len": old_len, + "new_len": new_len, + "tokens": req_to_token[req_idx, start:end] + .detach() + .cpu() + .tolist(), + } + ) + logits_output = getattr(out, "logits_output", None) + next_draft_input = getattr(out, "next_draft_input", None) + logger.info( + "GLM52 verify kv debug: accept_lens=%s new_seq_lens=%s " + "next_token_ids_head=%s tail_probe=%s logits_hidden=%s " + "next_draft_hidden=%s next_draft_topk=%s next_draft_p=%s " + "can_run_cuda_graph=%s", + _tensor_head(getattr(out, "accept_lens", None)), + _tensor_head(new_seq_lens), + _tensor_head(getattr(out, "next_token_ids", None)), + tail_probe, + _hidden_probe(getattr(logits_output, "hidden_states", None)), + _hidden_probe(getattr(next_draft_input, "hidden_states", None)), + _tensor_head(getattr(next_draft_input, "topk_index", None)), + _tensor_head( + getattr(next_draft_input, "topk_p", None), as_float=True + ), + getattr(out, "can_run_cuda_graph", None), + ) + except Exception: + logger.exception("Failed to log GLM52 verify kv debug") + return out + + EAGLEWorkerV2.verify = verify + EAGLEWorkerV2._atom_glm52_verify_kv_debug_patched = True + if not getattr(EAGLEDraftExtendCudaGraphRunner, "_atom_dsv4_replay_patched", False): original_extend_replay = EAGLEDraftExtendCudaGraphRunner.replay original_extend_can_run = EAGLEDraftExtendCudaGraphRunner.can_run @@ -489,9 +1216,11 @@ def replay(self, forward_batch): def _draft_extend_for_decode(self, batch, batch_result): try: + is_fixed_nextn = _is_dsv4_nextn_runner( + getattr(self, "draft_runner", None) + ) or _is_glm52_nextn_runner(getattr(self, "draft_runner", None)) if ( - not _is_dsv4_nextn_runner(getattr(self, "draft_runner", None)) - or getattr(self, "cuda_graph_runner_for_draft_extend", None) is None + not is_fixed_nextn ): return original_draft_extend_for_decode(self, batch, batch_result) @@ -507,10 +1236,13 @@ def _draft_extend_for_decode(self, batch, batch_result): if num_draft_tokens <= 0: return original_draft_extend_for_decode(self, batch, batch_result) - if not _dsv4_draft_extend_graph_layout_ok( - self.cuda_graph_runner_for_draft_extend + draft_extend_graph_runner = getattr( + self, "cuda_graph_runner_for_draft_extend", None + ) + if draft_extend_graph_runner is not None and not _dsv4_draft_extend_graph_layout_ok( + draft_extend_graph_runner ): - runner = self.cuda_graph_runner_for_draft_extend + runner = draft_extend_graph_runner self.cuda_graph_runner_for_draft_extend = None try: return original_draft_extend_for_decode( @@ -523,10 +1255,11 @@ def _draft_extend_for_decode(self, batch, batch_result): if not torch.is_tensor(accept_lens): return original_draft_extend_for_decode(self, batch, batch_result) - # DRAFT_EXTEND_V2 materializes exactly `num_draft_tokens` slots - # per sequence. `accept_lens` includes the target bonus token, - # so the value can be `num_draft_tokens + 1`; using that directly - # in the fixed-layout index points one slot past the graph output. + # DRAFT_EXTEND_V2 materializes a fixed `num_draft_tokens` slots + # per sequence. SGLang's default compact `cumsum(accept_lens)-1` + # index aliases rows from neighboring requests in this layout. + # `accept_lens` includes the target bonus token, so clamp before + # converting it to a fixed-layout per-request row offset. graph_accept_lens = accept_lens.clamp(min=1, max=num_draft_tokens) draft_input = EagleDraftInput( @@ -548,7 +1281,7 @@ def _draft_extend_for_decode(self, batch, batch_result): batch_result.next_token_ids, num_draft_tokens, self.draft_runner, - self.cuda_graph_runner_for_draft_extend, + draft_extend_graph_runner, ) ) @@ -564,12 +1297,12 @@ def _draft_extend_for_decode(self, batch, batch_result): forward_batch.spec_info.num_accept_tokens = graph_accept_lens can_cuda_graph = ( - self.cuda_graph_runner_for_draft_extend - and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) + draft_extend_graph_runner + and draft_extend_graph_runner.can_run(forward_batch) ) if can_cuda_graph: draft_logits_output = ( - self.cuda_graph_runner_for_draft_extend.replay(forward_batch) + draft_extend_graph_runner.replay(forward_batch) ) else: draft_logits_output = self.draft_runner.forward( @@ -590,7 +1323,7 @@ def _draft_extend_for_decode(self, batch, batch_result): output_len = int(draft_logits_output.next_token_logits.shape[0]) if max_index >= output_len: raise RuntimeError( - "DSV4 DRAFT_EXTEND_V2 output/index layout mismatch: " + "ATOM DRAFT_EXTEND_V2 output/index layout mismatch: " f"max_index={max_index}, output_len={output_len}, " f"batch={len(batch.seq_lens)}, " f"num_draft_tokens={num_draft_tokens}, " @@ -609,6 +1342,35 @@ def _draft_extend_for_decode(self, batch, batch_result): probs = torch.softmax(selected_logits, dim=-1) ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ( + "1", + "true", + "True", + ): + try: + logger.info( + "GLM52 draft_extend fixed debug: is_glm=%s " + "accept_lens=%s graph_accept_lens=%s select_index=%s " + "target_hidden=%s next_token_ids_head=%s " + "output_shape=%s selected_hidden=%s ret_topk=%s ret_p=%s " + "can_cuda_graph=%s", + _is_glm52_nextn_runner( + getattr(self, "draft_runner", None) + ), + accept_lens.detach().cpu().tolist(), + graph_accept_lens.detach().cpu().tolist(), + select_index.detach().cpu().tolist(), + _hidden_probe(batch_result.logits_output.hidden_states), + _tensor_head(batch_result.next_token_ids), + tuple(draft_logits_output.next_token_logits.shape), + _hidden_probe(selected_hidden_states), + _tensor_head(ret_topk_index), + _tensor_head(ret_topk_p, as_float=True), + bool(can_cuda_graph), + ) + except Exception: + logger.exception("Failed to log GLM52 draft_extend fixed debug") + next_draft_input = batch_result.next_draft_input ( next_draft_input.topk_p, @@ -695,6 +1457,333 @@ def init_cuda_graphs(self): EagleDraftWorker._atom_dsv4_init_cuda_graphs_patched = True +def _patch_sglang_eagle_v2_draft_argmax() -> None: + """Use ATOM draft distributed argmax for SGLang EAGLE topk=1 drafting.""" + if os.getenv("ATOM_SGLANG_DRAFT_ARGMAX", "1").lower() in ("0", "false", "no"): + return + try: + import torch + from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker + from sglang.srt.speculative.spec_utils import ( + maybe_detect_nan, + maybe_detect_oob, + select_top_k_tokens, + ) + except Exception as exc: + logger.debug("Skip patching SGLang EAGLE draft argmax: %s", exc) + return + + if getattr(EagleDraftWorker, "_atom_sglang_draft_argmax_patched", False): + return + + def draft_forward(self, forward_batch): + spec_info = forward_batch.spec_info + out_cache_loc = forward_batch.out_cache_loc + topk_p, topk_index, hidden_states = ( + spec_info.topk_p, + spec_info.topk_index, + spec_info.hidden_states, + ) + + maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") + + if self.hot_token_id is not None: + topk_index = self.hot_token_id[topk_index] + + draft_debug = os.environ.get("ATOM_GLM52_DRAFT_DEBUG", "0") in ( + "1", + "true", + "True", + ) or os.path.exists("/tmp/atom_glm52_draft_debug_on") + if draft_debug and not _is_current_stream_capturing(torch): + try: + hidden_probe = None + if torch.is_tensor(hidden_states): + hidden_rows = hidden_states[ + : min(4, int(hidden_states.shape[0])) + ].detach().float() + dim = int(hidden_rows.shape[-1]) + checksum_slices = [] + for start in ( + 0, + 256, + 1024, + 2048, + 4096, + max(0, dim - 256), + ): + end = min(dim, start + 256) + if start < end: + checksum_slices.append( + hidden_rows[:, start:end].sum(dim=-1) + ) + checksum = ( + torch.stack(checksum_slices, dim=-1) + if checksum_slices + else hidden_rows.new_empty((hidden_rows.shape[0], 0)) + ) + hidden_probe = { + "shape": tuple(hidden_states.shape), + "norm": hidden_rows.norm(dim=-1).cpu().tolist(), + "absmax": hidden_rows.abs().amax(dim=-1).cpu().tolist(), + "mean": hidden_rows.mean(dim=-1).cpu().tolist(), + "checksum": checksum.cpu().tolist(), + } + logger.info( + "GLM52 draft_forward debug: bs=%s topk=%s steps=%s " + "topk_index_shape=%s topk_index_head=%s topk_p_head=%s " + "hidden_shape=%s hidden_probe=%s out_cache_loc_shape=%s", + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + tuple(topk_index.shape), + topk_index.reshape(-1)[: min(12, int(topk_index.numel()))] + .detach() + .cpu() + .tolist(), + topk_p.reshape(-1)[: min(12, int(topk_p.numel()))] + .detach() + .cpu() + .tolist(), + tuple(hidden_states.shape) if torch.is_tensor(hidden_states) else None, + hidden_probe, + tuple(out_cache_loc.shape) if torch.is_tensor(out_cache_loc) else None, + ) + except Exception: + logger.exception("Failed to log GLM52 draft_forward debug") + + out_cache_loc = out_cache_loc.reshape( + forward_batch.batch_size, self.topk, self.speculative_num_steps + ) + out_cache_loc = out_cache_loc.permute((2, 0, 1)).reshape( + self.speculative_num_steps, -1 + ) + + score_list = [] + token_list = [] + parents_list = [] + scores = None + + use_argmax = self.topk == 1 + for i in range(self.speculative_num_steps): + input_ids, hidden_states, scores, tree_info = select_top_k_tokens( + i, topk_p, topk_index, hidden_states, scores, self.topk + ) + score_list.append(tree_info[0]) + token_list.append(tree_info[1]) + parents_list.append(tree_info[2]) + + if i == self.speculative_num_steps - 1: + break + + forward_batch.input_ids = input_ids + forward_batch.out_cache_loc = out_cache_loc[i] + forward_batch.attn_backend = self.draft_attn_backend.attn_backends[i] + forward_batch._atom_use_draft_argmax = use_argmax + spec_info.hidden_states = hidden_states + + if draft_debug and not _is_current_stream_capturing(torch): + try: + logger.info( + "GLM52 draft step debug: step=%s input_ids_shape=%s " + "input_ids_head=%s hidden_shape=%s scores_shape=%s", + i, + tuple(input_ids.shape), + input_ids.reshape(-1)[: min(12, int(input_ids.numel()))] + .detach() + .cpu() + .tolist(), + tuple(hidden_states.shape) + if torch.is_tensor(hidden_states) + else None, + tuple(scores.shape) if torch.is_tensor(scores) else None, + ) + except Exception: + logger.exception("Failed to log GLM52 draft step debug") + + logits_output = self.draft_runner.forward( + forward_batch, skip_attn_backend_init=True + ).logits_output + + draft_token_ids = None + customized_info = getattr(logits_output, "customized_info", None) or {} + if use_argmax: + draft_token_ids = customized_info.get("draft_token_ids") + + if draft_token_ids is not None: + topk_index = draft_token_ids.reshape(-1, 1) + topk_p = torch.ones( + (topk_index.shape[0], 1), + dtype=torch.float32, + device=topk_index.device, + ) + else: + maybe_detect_nan( + logits_output.next_token_logits, f"draft_forward step {i}" + ) + probs = torch.softmax(logits_output.next_token_logits, dim=-1) + from sglang.srt.utils.common import fast_topk + + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + maybe_detect_oob( + topk_index, + 0, + logits_output.next_token_logits.shape[-1], + f"draft_forward step {i}: topk_index OOB vs vocab_size={logits_output.next_token_logits.shape[-1]}", + ) + + if self.hot_token_id is not None: + topk_index = self.hot_token_id[topk_index] + if draft_debug and not _is_current_stream_capturing(torch): + try: + logger.info( + "GLM52 draft step output debug: step=%s draft_token_ids_head=%s " + "topk_index_head=%s topk_p_head=%s logits_shape=%s hidden_shape=%s", + i, + draft_token_ids.reshape(-1)[ + : min(12, int(draft_token_ids.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(draft_token_ids) + else None, + topk_index.reshape(-1)[: min(12, int(topk_index.numel()))] + .detach() + .cpu() + .tolist(), + topk_p.reshape(-1)[: min(12, int(topk_p.numel()))] + .detach() + .cpu() + .tolist(), + tuple(logits_output.next_token_logits.shape) + if torch.is_tensor(logits_output.next_token_logits) + else None, + tuple(logits_output.hidden_states.shape) + if torch.is_tensor(logits_output.hidden_states) + else None, + ) + except Exception: + logger.exception("Failed to log GLM52 draft step output debug") + hidden_states = logits_output.hidden_states + forward_batch.positions.add_(1) + + score_list = torch.cat(score_list, dim=1).flatten(1) + ss_token_list = torch.cat(token_list, dim=1) + top_scores = torch.topk( + score_list, self.speculative_num_draft_tokens - 1, dim=-1 + ) + top_scores_index = torch.sort(top_scores.indices).values + maybe_detect_oob( + top_scores_index, + 0, + ss_token_list.shape[1], + "draft_forward: top_scores_index OOB for gather on ss_token_list", + ) + draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) + if draft_debug and not _is_current_stream_capturing(torch): + try: + logger.info( + "GLM52 draft final debug: parent_shape=%s top_scores_index_head=%s " + "draft_tokens_head=%s score_head=%s", + tuple(parent_list.shape) if "parent_list" in locals() else None, + top_scores_index[: min(4, int(top_scores_index.shape[0]))] + .detach() + .cpu() + .tolist(), + draft_tokens[: min(4, int(draft_tokens.shape[0]))] + .detach() + .cpu() + .tolist(), + score_list[: min(4, int(score_list.shape[0]))] + .detach() + .float() + .cpu() + .tolist(), + ) + except Exception: + logger.exception("Failed to log GLM52 draft final debug") + + if len(parents_list) > 1: + parent_list = torch.cat(parents_list[:-1], dim=1) + else: + batch_size = parents_list[0].shape[0] + parent_list = torch.empty(batch_size, 0, device=parents_list[0].device) + + return parent_list, top_scores_index, draft_tokens + + EagleDraftWorker.draft_forward = draft_forward + EagleDraftWorker._atom_sglang_draft_argmax_patched = True + logger.info("Patched SGLang EAGLE draft_forward for ATOM distributed argmax") + + +def _patch_sglang_glm52_logits_norm_debug() -> None: + if os.environ.get("ATOM_GLM52_LOGITS_NORM_DEBUG", "0") not in ( + "1", + "true", + "True", + ): + return + try: + import torch + from sglang.srt.layers.logits_processor import LogitsProcessor + except Exception as exc: + logger.debug("Skip patching GLM52 logits norm debug: %s", exc) + return + + if getattr(LogitsProcessor, "_atom_glm52_logits_norm_debug_patched", False): + return + + original_get_logits = LogitsProcessor._get_logits + + def _get_logits(self, hidden_states, lm_head, logits_metadata): + should_log = False + try: + forward_mode = getattr(logits_metadata, "forward_mode", None) + should_log = bool( + forward_mode is not None + and getattr(forward_mode, "is_target_verify", lambda: False)() + and not _is_current_stream_capturing(torch) + ) + if should_log: + probe = hidden_states[: min(8, int(hidden_states.shape[0]))].detach() + logger.info( + "GLM52 logits norm pre: mode=%s hidden_shape=%s " + "hidden_norm=%s hidden_absmax=%s hidden_mean=%s", + forward_mode, + tuple(hidden_states.shape), + probe.float().norm(dim=-1).cpu().tolist(), + probe.float().abs().amax(dim=-1).cpu().tolist(), + probe.float().mean(dim=-1).cpu().tolist(), + ) + except Exception: + logger.exception("Failed to log GLM52 logits norm pre debug") + + logits = original_get_logits(self, hidden_states, lm_head, logits_metadata) + + if should_log: + try: + probe = logits[: min(8, int(logits.shape[0]))].detach().float() + top_vals, top_ids = torch.topk(probe, k=3, dim=-1) + logger.info( + "GLM52 logits norm post: logits_shape=%s logits_norm=%s " + "logits_absmax=%s logits_mean=%s top_ids=%s top_vals=%s", + tuple(logits.shape), + probe.norm(dim=-1).cpu().tolist(), + probe.abs().amax(dim=-1).cpu().tolist(), + probe.mean(dim=-1).cpu().tolist(), + top_ids.cpu().tolist(), + top_vals.cpu().tolist(), + ) + except Exception: + logger.exception("Failed to log GLM52 logits norm post debug") + return logits + + LogitsProcessor._get_logits = _get_logits + LogitsProcessor._atom_glm52_logits_norm_debug_patched = True + logger.info("Patched SGLang LogitsProcessor for GLM52 norm debug") + + def register_ops_to_sglang(atom_config: Config) -> None: """ Register custom ops to sglang, including attention @@ -702,6 +1791,8 @@ def register_ops_to_sglang(atom_config: Config) -> None: _register_custom_attention_to_sglang() _patch_sglang_dsv4_draft_backends() _patch_sglang_dsv4_spec_cuda_graph() + _patch_sglang_eagle_v2_draft_argmax() + _patch_sglang_glm52_logits_norm_debug() def set_attn_cls() -> None: diff --git a/atom/plugin/sglang/attention_backend/glm52_dsa_backend.py b/atom/plugin/sglang/attention_backend/glm52_dsa_backend.py new file mode 100644 index 0000000000..4d40f9e65d --- /dev/null +++ b/atom/plugin/sglang/attention_backend/glm52_dsa_backend.py @@ -0,0 +1,719 @@ +"""SGLang backend shim for ATOM-owned GLM-5.2 native MLA attention.""" + +from __future__ import annotations + +import logging +import os +from types import SimpleNamespace + +import torch +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend + +logger = logging.getLogger("atom.plugin.sglang.attention_backend.glm52_dsa") + + +class ATOMGLM52DSABackendForSgl(AttentionBackend): + """Publish fixed-address ATOM GLM-5.2 metadata for SGLang CUDA graphs.""" + + needs_cpu_seq_lens = True + _last_atom_glm52_graph_metadata = None + + def __init__(self, model_runner, *args, **kwargs): + del args + logger.info("Initializing ATOMGLM52DSABackendForSgl") + self.model_runner = model_runner + self.device = torch.device(model_runner.device) + self.token_to_kv_pool = model_runner.token_to_kv_pool + self.req_to_token_pool = model_runner.req_to_token_pool + self.forward_metadata = None + self.atom_glm52_graph_metadata = None + self._cuda_graph_seq_len_fill_value = 1 + self._spec_graph_metadata_cache = {} + speculative_num_steps = int(kwargs.pop("speculative_num_steps", 0) or 0) + self.attn_backends = [self] * max(1, speculative_num_steps) + + @staticmethod + def get_name() -> str: + return "atom_glm52_dsa" + + def init_forward_metadata(self, forward_batch): + self.forward_metadata = forward_batch + self.atom_glm52_graph_metadata = None + + def _build_decode_graph_metadata(self, forward_batch, positions=None, max_bs=None): + if not forward_batch.forward_mode.is_decode_or_idle(): + self.atom_glm52_graph_metadata = None + return None + if positions is None: + positions = getattr(forward_batch, "positions", None) + if positions is None: + positions = ( + forward_batch.seq_lens[: int(forward_batch.batch_size)].to(torch.int64) + - 1 + ).clamp_min_(0) + + from atom.config import get_current_atom_config + from atom.plugin.sglang.glm52_dsa_bridge import ( + build_atom_glm52_decode_graph_metadata_from_sglang, + ) + + atom_config = get_current_atom_config() + max_context_len = int(self.req_to_token_pool.req_to_token.shape[1]) + self.atom_glm52_graph_metadata = ( + build_atom_glm52_decode_graph_metadata_from_sglang( + forward_batch, + positions, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token_pool=self.req_to_token_pool, + atom_config=atom_config, + max_bs=max_bs, + max_context_len=max_context_len, + ) + ) + forward_batch.atom_glm52_graph_metadata = self.atom_glm52_graph_metadata + ATOMGLM52DSABackendForSgl._last_atom_glm52_graph_metadata = ( + self.atom_glm52_graph_metadata + ) + self.forward_metadata = forward_batch + return self.atom_glm52_graph_metadata + + @staticmethod + def _is_spec_extend_mode(forward_mode) -> bool: + return bool( + forward_mode is not None + and ( + forward_mode.is_target_verify() + or getattr(forward_mode, "is_draft_extend", lambda **kwargs: False)( + include_v2=True + ) + ) + ) + + @staticmethod + def _spec_graph_key(forward_batch, positions) -> tuple[str, int, int]: + cache_bs = int( + getattr(forward_batch, "_graph_cache_bs", int(forward_batch.batch_size)) + ) + cache_rows = int( + getattr( + forward_batch, + "_graph_cache_rows", + int(positions.numel()) if positions is not None else 0, + ) + ) + return ( + str(forward_batch.forward_mode), + cache_bs, + cache_rows, + ) + + @staticmethod + def _copy_graph_metadata_in_place(dst, src): + def _copy_tensor(dst_tensor, src_tensor): + if dst_tensor.shape == src_tensor.shape and dst_tensor.dtype == src_tensor.dtype: + dst_tensor.copy_(src_tensor) + return True + if ( + dst_tensor.dim() == src_tensor.dim() + and dst_tensor.dtype == src_tensor.dtype + and all(o >= v for o, v in zip(dst_tensor.shape, src_tensor.shape)) + ): + dst_tensor.zero_() + slices = tuple(slice(0, int(v)) for v in src_tensor.shape) + dst_tensor[slices].copy_(src_tensor) + return True + if dst_tensor.numel() >= src_tensor.numel() and dst_tensor.dtype == src_tensor.dtype: + dst_tensor.reshape(-1)[: src_tensor.numel()].copy_(src_tensor.reshape(-1)) + if dst_tensor.numel() > src_tensor.numel(): + dst_tensor.reshape(-1)[src_tensor.numel() :].zero_() + return True + return False + + for name, value in vars(src).items(): + old = getattr(dst, name, None) + if torch.is_tensor(value) and torch.is_tensor(old): + if _copy_tensor(old, value): + continue + if name == "mla_chunk_meta" and old is not None and value is not None: + for field_name, field_value in vars(value).items(): + field_old = getattr(old, field_name, None) + if isinstance(field_value, list) and isinstance(field_old, list): + for old_item, new_item in zip(field_old, field_value): + if torch.is_tensor(old_item) and torch.is_tensor(new_item): + _copy_tensor(old_item, new_item) + if len(field_old) == len(field_value): + continue + if torch.is_tensor(field_value) and torch.is_tensor(field_old): + if _copy_tensor(field_old, field_value): + continue + setattr(old, field_name, field_value) + continue + setattr(dst, name, value) + return dst + + def _build_spec_graph_metadata(self, forward_batch, positions=None): + if not self._is_spec_extend_mode(forward_batch.forward_mode): + self.atom_glm52_graph_metadata = None + return None + tokens_per_req = int( + getattr(getattr(forward_batch, "spec_info", None), "num_tokens_per_req", 0) + or getattr(getattr(forward_batch, "spec_info", None), "draft_token_num", 0) + or 1 + ) + bs = int(forward_batch.batch_size) + total_rows = bs * max(1, tokens_per_req) + if positions is None: + positions = getattr(forward_batch, "positions", None) + if positions is None: + if forward_batch.forward_mode.is_target_verify(): + base = forward_batch.seq_lens[:bs].to(torch.int64) + else: + base = ( + forward_batch.seq_lens[:bs].to(torch.int64) - tokens_per_req + ).clamp_min_(0) + offsets = torch.arange(tokens_per_req, dtype=torch.int64, device=self.device) + positions = (base[:, None] + offsets[None, :]).reshape(-1) + elif int(positions.numel()) < total_rows: + if forward_batch.forward_mode.is_target_verify(): + base = forward_batch.seq_lens[:bs].to(torch.int64) + else: + base = ( + forward_batch.seq_lens[:bs].to(torch.int64) - tokens_per_req + ).clamp_min_(0) + offsets = torch.arange(tokens_per_req, dtype=torch.int64, device=self.device) + padded_positions = (base[:, None] + offsets[None, :]).reshape(-1) + padded_positions[: int(positions.numel())].copy_(positions) + positions = padded_positions + elif int(positions.numel()) > total_rows: + positions = positions[:total_rows] + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + if out_cache_loc is None or int(out_cache_loc.numel()) < total_rows: + scratch_slot = max(0, int(getattr(self.token_to_kv_pool, "size", 1)) - 1) + values = dict(getattr(forward_batch, "__dict__", {})) + padded_out_cache_loc = torch.full( + (total_rows,), + scratch_slot, + dtype=torch.int64, + device=self.device, + ) + if torch.is_tensor(out_cache_loc) and int(out_cache_loc.numel()) > 0: + padded_out_cache_loc[: int(out_cache_loc.numel())].copy_(out_cache_loc) + values["out_cache_loc"] = padded_out_cache_loc + forward_batch = SimpleNamespace(**values) + out_cache_loc = padded_out_cache_loc + elif int(out_cache_loc.numel()) > total_rows: + values = dict(getattr(forward_batch, "__dict__", {})) + values["out_cache_loc"] = out_cache_loc[:total_rows] + forward_batch = SimpleNamespace(**values) + out_cache_loc = values["out_cache_loc"] + + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ("1", "true", "True"): + try: + seq_lens = getattr(forward_batch, "seq_lens", None) + input_ids = getattr(forward_batch, "input_ids", None) + logger.info( + "GLM52 verify graph metadata: mode=%s bs=%s tokens_per_req=%s " + "positions_shape=%s positions_head=%s seq_lens_head=%s " + "out_cache_loc_shape=%s out_cache_loc_head=%s input_ids_shape=%s " + "input_ids_head=%s raw_replay_bs=%s raw_replay_tokens=%s", + forward_batch.forward_mode, + bs, + tokens_per_req, + tuple(positions.shape) if torch.is_tensor(positions) else None, + positions[: min(12, int(positions.numel()))].detach().cpu().tolist() + if torch.is_tensor(positions) + else None, + seq_lens[: min(8, int(seq_lens.numel()))].detach().cpu().tolist() + if torch.is_tensor(seq_lens) + else None, + tuple(out_cache_loc.shape) + if torch.is_tensor(out_cache_loc) + else None, + out_cache_loc[: min(12, int(out_cache_loc.numel()))] + .detach() + .cpu() + .tolist() + if torch.is_tensor(out_cache_loc) + else None, + tuple(input_ids.shape) if torch.is_tensor(input_ids) else None, + input_ids[: min(12, int(input_ids.numel()))].detach().cpu().tolist() + if torch.is_tensor(input_ids) + else None, + int(getattr(getattr(self, "_replay_forward_batch", None), "batch_size", 0) or 0), + int(getattr(getattr(getattr(self, "_replay_forward_batch", None), "input_ids", None), "numel", lambda: 0)()), + ) + except Exception: + logger.exception("Failed to log GLM52 verify graph metadata debug") + + from atom.config import get_current_atom_config + from atom.plugin.sglang.glm52_dsa_bridge import ( + build_atom_glm52_attention_metadata_from_sglang, + ) + + atom_config = get_current_atom_config() + new_metadata = build_atom_glm52_attention_metadata_from_sglang( + forward_batch, + positions, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token_pool=self.req_to_token_pool, + atom_config=atom_config, + ) + key = self._spec_graph_key(forward_batch, positions) + cached_metadata = self._spec_graph_metadata_cache.get(key) + if cached_metadata is None: + self._spec_graph_metadata_cache[key] = new_metadata + self.atom_glm52_graph_metadata = new_metadata + else: + self.atom_glm52_graph_metadata = self._copy_graph_metadata_in_place( + cached_metadata, new_metadata + ) + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ("1", "true", "True"): + try: + md = self.atom_glm52_graph_metadata + row_probe = None + try: + probe_bs = min(2, int(bs)) + positions_rows = ( + positions.reshape(bs, tokens_per_req)[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor(positions) + and int(positions.numel()) >= bs * tokens_per_req + else None + ) + out_rows = ( + out_cache_loc.reshape(bs, tokens_per_req)[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor(out_cache_loc) + and int(out_cache_loc.numel()) >= bs * tokens_per_req + else None + ) + req_rows = ( + forward_batch.req_pool_indices[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(forward_batch, "req_pool_indices", None)) + else None + ) + kv_indptr_cpu = ( + md.kv_indptr.detach().cpu().tolist() + if torch.is_tensor(getattr(md, "kv_indptr", None)) + else [] + ) + kv_indices = getattr(md, "kv_indices", None) + kv_ranges = [] + for row in range(min(probe_bs, max(0, len(kv_indptr_cpu) - 1))): + start = int(kv_indptr_cpu[row]) + end = int(kv_indptr_cpu[row + 1]) + if torch.is_tensor(kv_indices): + head = ( + kv_indices[start : min(end, start + 8)] + .detach() + .cpu() + .tolist() + ) + tail = ( + kv_indices[max(start, end - 8) : end] + .detach() + .cpu() + .tolist() + ) + else: + head = tail = None + kv_ranges.append( + { + "row": row, + "start": start, + "end": end, + "len": end - start, + "head": head, + "tail": tail, + } + ) + sparse_indptr = getattr(md, "sparse_kv_indptr", None) + sparse_probe = ( + sparse_indptr[ + : min( + int(sparse_indptr.numel()), + probe_bs * tokens_per_req + 1, + ) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(sparse_indptr) + else None + ) + token_to_seq = getattr(md, "token_to_seq_idxs", None) + token_to_seq_probe = ( + token_to_seq[: min(int(token_to_seq.numel()), probe_bs * tokens_per_req)] + .detach() + .cpu() + .tolist() + if torch.is_tensor(token_to_seq) + else None + ) + row_probe = { + "positions": positions_rows, + "out_cache_loc": out_rows, + "req_pool_indices": req_rows, + "kv_ranges": kv_ranges, + "sparse_kv_indptr": sparse_probe, + "token_to_seq": token_to_seq_probe, + } + spec_info = getattr(forward_batch, "spec_info", None) + if spec_info is not None: + counter = int( + getattr(spec_info, "_atom_glm52_verify_counter", 0) or 0 + ) + 1 + setattr(spec_info, "_atom_glm52_verify_counter", counter) + setattr(spec_info, "_atom_glm52_row_probe", row_probe) + except Exception: + logger.exception("Failed to build GLM52 verify graph row probe") + logger.info( + "GLM52 verify graph md: key=%s max_q=%s max_k=%s total_kv=%s " + "has_cached=%s cu_q=%s kv_indptr=%s sparse_cu=%s " + "sparse_kv_indptr=%s token_to_seq=%s row_probe=%s", + key, + getattr(md, "max_seqlen_q", None), + getattr(md, "max_seqlen_k", None), + getattr(md, "total_kv", None), + getattr(md, "has_cached", None), + md.cu_seqlens_q[: min(8, int(md.cu_seqlens_q.numel()))] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(md, "cu_seqlens_q", None)) + else None, + md.kv_indptr[: min(8, int(md.kv_indptr.numel()))] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(md, "kv_indptr", None)) + else None, + md.sparse_cu_seqlens_q[ + : min(8, int(md.sparse_cu_seqlens_q.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(md, "sparse_cu_seqlens_q", None)) + else None, + md.sparse_kv_indptr[ + : min(8, int(md.sparse_kv_indptr.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(md, "sparse_kv_indptr", None)) + else None, + md.token_to_seq_idxs[ + : min(12, int(md.token_to_seq_idxs.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(getattr(md, "token_to_seq_idxs", None)) + else None, + row_probe, + ) + except Exception: + logger.exception("Failed to log GLM52 verify graph md debug") + forward_batch.atom_glm52_graph_metadata = self.atom_glm52_graph_metadata + ATOMGLM52DSABackendForSgl._last_atom_glm52_graph_metadata = ( + self.atom_glm52_graph_metadata + ) + self.forward_metadata = forward_batch + return self.atom_glm52_graph_metadata + + def _build_graph_metadata(self, forward_batch, positions=None, max_bs=None): + if self._is_spec_extend_mode(forward_batch.forward_mode): + return self._build_spec_graph_metadata(forward_batch, positions=positions) + return self._build_decode_graph_metadata( + forward_batch, positions=positions, max_bs=max_bs + ) + + def init_forward_metadata_out_graph(self, forward_batch, in_capture: bool = False): + if not (in_capture or hasattr(forward_batch, "actual_forward_mode")): + self.forward_metadata = forward_batch + self.atom_glm52_graph_metadata = None + return + positions = getattr(forward_batch, "positions", None) + if positions is None: + graph_runner = getattr(self.model_runner, "graph_runner", None) + buffers = getattr(graph_runner, "buffers", None) + positions = getattr(buffers, "positions", None) + self._build_graph_metadata(forward_batch, positions=positions) + + def init_forward_metadata_capture_cuda_graph(self, *args, **kwargs): + if len(args) == 1 and not kwargs and hasattr(args[0], "forward_mode"): + return self.init_forward_metadata_out_graph(args[0], in_capture=True) + bs = kwargs.get("bs", args[0] if len(args) > 0 else None) + req_pool_indices = kwargs.get( + "req_pool_indices", args[2] if len(args) > 2 else None + ) + seq_lens = kwargs.get("seq_lens", args[3] if len(args) > 3 else None) + forward_mode = kwargs.get("forward_mode", args[5] if len(args) > 5 else None) + spec_info = kwargs.get("spec_info", args[6] if len(args) > 6 else None) + if bs is None or req_pool_indices is None or seq_lens is None: + self.atom_glm52_graph_metadata = None + return + forward_batch = SimpleNamespace( + forward_mode=forward_mode, + actual_forward_mode=forward_mode, + batch_size=int(bs), + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens.detach().cpu(), + out_cache_loc=None, + spec_info=spec_info, + ) + self._build_graph_metadata(forward_batch, max_bs=int(bs)) + + def init_forward_metadata_replay_cuda_graph(self, *args, **kwargs): + if len(args) == 2 and hasattr(args[0], "forward_mode"): + forward_batch, bs = args + values = dict(getattr(forward_batch, "__dict__", {})) + if self._is_spec_extend_mode(getattr(forward_batch, "forward_mode", None)): + active_bs = int(getattr(forward_batch, "batch_size", int(bs)) or int(bs)) + tokens_per_req = int( + getattr(getattr(forward_batch, "spec_info", None), "num_tokens_per_req", 0) + or getattr(getattr(forward_batch, "spec_info", None), "draft_token_num", 0) + or 1 + ) + values["batch_size"] = active_bs + values["_graph_cache_bs"] = int(bs) + values["_graph_cache_rows"] = int(bs) * max(1, tokens_per_req) + else: + values["batch_size"] = int(bs) + replay_batch = SimpleNamespace(**values) + return self._build_graph_metadata( + replay_batch, + positions=getattr(forward_batch, "positions", None), + max_bs=int(bs), + ) + + bs = kwargs.get("bs", args[0] if len(args) > 0 else None) + req_pool_indices = kwargs.get( + "req_pool_indices", args[1] if len(args) > 1 else None + ) + seq_lens = kwargs.get("seq_lens", args[2] if len(args) > 2 else None) + forward_mode = kwargs.get("forward_mode", args[5] if len(args) > 5 else None) + spec_info = kwargs.get("spec_info", args[6] if len(args) > 6 else None) + seq_lens_cpu = kwargs.get("seq_lens_cpu", args[7] if len(args) > 7 else None) + out_cache_loc = kwargs.get("out_cache_loc", args[8] if len(args) > 8 else None) + replay_batch = getattr(self, "_replay_forward_batch", None) + if out_cache_loc is None: + out_cache_loc = getattr(replay_batch, "out_cache_loc", None) + raw_seq_lens = getattr(replay_batch, "seq_lens", None) + raw_seq_lens_cpu = getattr(replay_batch, "seq_lens_cpu", None) + is_spec_replay = self._is_spec_extend_mode(forward_mode) + active_bs = int(bs) + if is_spec_replay: + active_bs = int(getattr(replay_batch, "batch_size", 0) or int(bs)) + raw_req_pool_indices = getattr(replay_batch, "req_pool_indices", None) + if raw_req_pool_indices is not None: + req_pool_indices = raw_req_pool_indices + if bs is None or req_pool_indices is None or seq_lens is None: + self.atom_glm52_graph_metadata = None + return + tokens_per_req = int( + getattr(spec_info, "num_tokens_per_req", 0) + or getattr(spec_info, "draft_token_num", 0) + or getattr(getattr(replay_batch, "spec_info", None), "num_tokens_per_req", 0) + or getattr(getattr(replay_batch, "spec_info", None), "draft_token_num", 0) + or 1 + ) + forward_batch = SimpleNamespace( + forward_mode=forward_mode, + actual_forward_mode=getattr(replay_batch, "forward_mode", forward_mode), + batch_size=active_bs, + _graph_cache_bs=int(bs), + _graph_cache_rows=int(bs) * max(1, tokens_per_req), + req_pool_indices=req_pool_indices, + seq_lens=( + raw_seq_lens + if is_spec_replay and torch.is_tensor(raw_seq_lens) + else seq_lens + ), + seq_lens_cpu=( + raw_seq_lens_cpu + if is_spec_replay and raw_seq_lens_cpu is not None + else seq_lens_cpu + ), + out_cache_loc=out_cache_loc, + positions=getattr(replay_batch, "positions", None), + spec_info=spec_info or getattr(replay_batch, "spec_info", None), + ) + self._build_graph_metadata( + forward_batch, + positions=getattr(forward_batch, "positions", None), + max_bs=int(bs), + ) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + bs = int(max_bs) + tokens_per_req = max(1, int(max_num_tokens) // max(1, bs)) + is_target_verify_graph = bool( + getattr( + getattr(self.model_runner, "spec_algorithm", None), + "is_speculative", + lambda: False, + )() + and not getattr(self.model_runner, "is_draft_worker", False) + ) + is_draft_extend_graph = bool( + getattr( + getattr(self.model_runner, "spec_algorithm", None), + "is_speculative", + lambda: False, + )() + and getattr(self.model_runner, "is_draft_worker", False) + and tokens_per_req > 1 + ) + is_graph_extend = is_target_verify_graph or is_draft_extend_graph + seq_lens = torch.ones(bs, dtype=torch.int32, device=self.device) + req_pool_indices = torch.arange(bs, dtype=torch.int64, device=self.device) + forward_mode = ( + ForwardMode.TARGET_VERIFY + if is_target_verify_graph + else ( + ForwardMode.DRAFT_EXTEND_V2 + if is_draft_extend_graph + else ForwardMode.DECODE + ) + ) + if is_graph_extend: + fill_override = int(os.environ.get("ATOM_GLM52_TV_CG_SEQ_LEN_FILL", "0") or 0) + self._cuda_graph_seq_len_fill_value = ( + max(tokens_per_req, fill_override) + if fill_override > 0 + else max(tokens_per_req, 1024) + ) + seq_lens.fill_(self._cuda_graph_seq_len_fill_value) + offsets = torch.arange(tokens_per_req, dtype=torch.int64, device=self.device) + if is_target_verify_graph: + base = seq_lens[:bs].to(torch.int64) + else: + base = (seq_lens[:bs].to(torch.int64) - tokens_per_req).clamp_min_(0) + positions = (base[:, None] + offsets[None, :]).reshape(-1) + if os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ("1", "true", "True"): + logger.info( + "GLM52 graph init seq_len fill: mode=%s bs=%s tokens_per_req=%s " + "fill=%s override=%s", + forward_mode, + bs, + tokens_per_req, + self._cuda_graph_seq_len_fill_value, + fill_override, + ) + else: + self._cuda_graph_seq_len_fill_value = 1 + positions = torch.zeros(bs, dtype=torch.int64, device=self.device) + forward_batch = SimpleNamespace( + forward_mode=forward_mode, + actual_forward_mode=forward_mode, + batch_size=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens.detach().cpu(), + out_cache_loc=None, + spec_info=SimpleNamespace( + num_tokens_per_req=tokens_per_req, + draft_token_num=tokens_per_req, + ), + ) + self._build_graph_metadata( + forward_batch, + positions=positions, + max_bs=bs, + ) + del max_num_tokens + return None + + def get_cuda_graph_seq_len_fill_value(self): + return int(self._cuda_graph_seq_len_fill_value) + + def get_verify_buffers_to_fill_after_draft(self): + graph_runner = getattr(self.model_runner, "graph_runner", None) + buffers = getattr(graph_runner, "buffers", None) + if buffers is None: + return [None, None] + # SGLang fills this captured mask buffer after draft graph replay. The + # verifier positions are copied in update_verify_buffers_to_fill_after_draft(). + return [getattr(buffers, "custom_mask", None), None] + + def update_verify_buffers_to_fill_after_draft(self, spec_info, cuda_graph_bs): + if cuda_graph_bs is None: + return + graph_runner = getattr(self.model_runner, "graph_runner", None) + buffers = getattr(graph_runner, "buffers", None) + if buffers is None: + return + + tokens_per_req = int( + getattr( + spec_info, + "num_tokens_per_req", + getattr(spec_info, "draft_token_num", 1), + ) + or 1 + ) + total = int(cuda_graph_bs) * tokens_per_req + + positions = getattr(spec_info, "positions", None) + active_total = int(positions.numel()) if torch.is_tensor(positions) else total + active_bs = max(1, active_total // max(1, tokens_per_req)) + if torch.is_tensor(positions): + copy_n = min(int(positions.numel()), total) + if copy_n: + buffers.positions[:copy_n].copy_(positions[:copy_n]) + if total > copy_n: + buffers.positions[copy_n:total].zero_() + positions = buffers.positions[:active_total] + else: + positions = buffers.positions[: active_bs * tokens_per_req] + + custom_mask = getattr(spec_info, "custom_mask", None) + graph_custom_mask = getattr(buffers, "custom_mask", None) + if ( + torch.is_tensor(custom_mask) + and torch.is_tensor(graph_custom_mask) + and custom_mask.data_ptr() != graph_custom_mask.data_ptr() + ): + graph_custom_mask[: custom_mask.numel()].copy_(custom_mask) + + forward_mode = getattr( + getattr(self, "forward_metadata", None), "forward_mode", None + ) + if forward_mode is None: + return + seq_lens_cpu = getattr(buffers, "seq_lens_cpu", None) + forward_batch = SimpleNamespace( + forward_mode=forward_mode, + actual_forward_mode=forward_mode, + batch_size=active_bs, + _graph_cache_bs=int(cuda_graph_bs), + _graph_cache_rows=total, + req_pool_indices=buffers.req_pool_indices[:active_bs], + seq_lens=buffers.seq_lens[:active_bs], + seq_lens_cpu=( + seq_lens_cpu[:active_bs] if seq_lens_cpu is not None else None + ), + out_cache_loc=buffers.out_cache_loc[:active_total], + positions=positions, + spec_info=spec_info, + ) + self._build_graph_metadata(forward_batch, positions=positions) + + def forward_decode(self, *args, **kwargs): + raise RuntimeError("ATOM GLM-5.2 SGLang bridge should use ATOM attention") + + def forward_extend(self, *args, **kwargs): + raise RuntimeError("ATOM GLM-5.2 SGLang bridge should use ATOM attention") diff --git a/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py b/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py index a91e54d537..0410a03de4 100644 --- a/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py +++ b/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py @@ -27,6 +27,26 @@ from atom.utils.custom_register import direct_register_custom_op +_static_i32_cache: dict[tuple[str, str, int], torch.Tensor] = {} + + +def _get_cached_i32_arange(n: int, device: torch.device) -> torch.Tensor: + key = (str(device), "arange", int(n)) + value = _static_i32_cache.get(key) + if value is None: + value = torch.arange(n, dtype=torch.int32, device=device) + _static_i32_cache[key] = value + return value + + +def _get_cached_i32_ones(n: int, device: torch.device) -> torch.Tensor: + key = (str(device), "ones", int(n)) + value = _static_i32_cache.get(key) + if value is None: + value = torch.ones(n, dtype=torch.int32, device=device) + _static_i32_cache[key] = value + return value + @triton.jit def _convert_req_index_to_global_index_kernel( @@ -114,6 +134,135 @@ def triton_convert_req_index_to_global_index( ) +@triton.jit +def _target_verify_query_ranges_kernel( + seq_lens_ptr, + starts_ptr, + ends_ptr, + bs, + draft_token_num: tl.constexpr, + BLOCK_BS: tl.constexpr, + BLOCK_DRAFT: tl.constexpr, +): + req = tl.arange(0, BLOCK_BS) + valid_req = req < bs + seq = tl.load(seq_lens_ptr + req, mask=valid_req, other=0).to(tl.int32) + kv_lens = seq + draft_token_num + base = tl.cumsum(kv_lens, axis=0) - kv_lens + + draft = tl.arange(0, BLOCK_DRAFT) + valid_draft = draft < draft_token_num + row = req[:, None] * draft_token_num + draft[None, :] + mask = valid_req[:, None] & valid_draft[None, :] + tl.store(starts_ptr + row, base[:, None], mask=mask) + tl.store(ends_ptr + row, base[:, None] + seq[:, None] + draft[None, :] + 1, mask=mask) + + +def _build_target_verify_query_ranges_triton( + seq_lens: torch.Tensor, + bs: int, + draft_token_num: int, +) -> tuple[torch.Tensor, torch.Tensor]: + device = seq_lens.device + num_tokens = bs * draft_token_num + starts = torch.empty(num_tokens, dtype=torch.int32, device=device) + ends = torch.empty(num_tokens, dtype=torch.int32, device=device) + block_bs = triton.next_power_of_2(bs) + block_draft = triton.next_power_of_2(draft_token_num) + _target_verify_query_ranges_kernel[(1,)]( + seq_lens, + starts, + ends, + bs, + draft_token_num, + BLOCK_BS=block_bs, + BLOCK_DRAFT=block_draft, + ) + return starts, ends + + +@triton.jit +def _target_verify_req_indptr_kernel( + seq_lens_ptr, + req_id_ptr, + indptr_ptr, + bs, + draft_token_num: tl.constexpr, + topk_tokens: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, +): + offs = tl.arange(0, BLOCK_TOKENS) + total = bs * draft_token_num + mask = offs < total + req = offs // draft_token_num + draft = offs - req * draft_token_num + seq = tl.load(seq_lens_ptr + req, mask=mask, other=0).to(tl.int32) + counts = tl.minimum(seq + draft + 1, topk_tokens) + counts = tl.where(mask, counts, 0) + csum = tl.cumsum(counts, axis=0) + tl.store(req_id_ptr + offs, req, mask=mask) + tl.store(indptr_ptr + offs + 1, csum, mask=mask) + tl.store(indptr_ptr, tl.full((), 0, tl.int32)) + + +def _build_target_verify_req_indptr_triton( + seq_lens: torch.Tensor, + bs: int, + draft_token_num: int, + topk_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + device = seq_lens.device + num_tokens = bs * draft_token_num + req_id = torch.empty(num_tokens, dtype=torch.int32, device=device) + indptr = torch.empty(num_tokens + 1, dtype=torch.int32, device=device) + block_tokens = triton.next_power_of_2(num_tokens) + _target_verify_req_indptr_kernel[(1,)]( + seq_lens, + req_id, + indptr, + bs, + draft_token_num, + topk_tokens, + BLOCK_TOKENS=block_tokens, + ) + return req_id, indptr + + +@triton.jit +def _localize_topk_indices_kernel( + topk_ptr, + starts_ptr, + num_rows, + topk_tokens: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + offs = tile * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (row < num_rows) & (offs < topk_tokens) + ptr = topk_ptr + row * topk_tokens + offs + vals = tl.load(ptr, mask=mask, other=-1) + start = tl.load(starts_ptr + row, mask=row < num_rows, other=0) + vals = tl.where(vals >= 0, vals - start, vals) + tl.store(ptr, vals, mask=mask) + + +def _localize_topk_indices_triton( + topk_indices: torch.Tensor, + cu_starts: torch.Tensor, +) -> None: + num_rows, topk_tokens = topk_indices.shape + block_n = 256 + grid = (num_rows, triton.cdiv(topk_tokens, block_n)) + _localize_topk_indices_kernel[grid]( + topk_indices, + cu_starts, + num_rows, + topk_tokens, + BLOCK_N=block_n, + ) + + def _parse_layer_id_from_indexer_prefix(prefix: str) -> int: match = re.search(r"\.layers\.(\d+)\.", prefix) if match is None: @@ -139,21 +288,10 @@ def _build_sglang_query_ranges(forward_batch) -> tuple[torch.Tensor, torch.Tenso ) if draft_token_num <= 0: raise RuntimeError("TARGET_VERIFY sparse MLA requires draft_token_num") - seq_lens = forward_batch.seq_lens[:bs].to(dtype=torch.int32) - kv_lens = seq_lens + draft_token_num - base_offsets = torch.cumsum( - torch.cat([torch.zeros(1, dtype=torch.int32, device=device), kv_lens[:-1]]), - dim=0, - ) - starts = torch.repeat_interleave(base_offsets, draft_token_num) - per_req_end_base = torch.repeat_interleave( - base_offsets + seq_lens, draft_token_num - ) - draft_offsets = torch.arange( - 1, draft_token_num + 1, dtype=torch.int32, device=device - ).repeat(bs) - return starts.to(dtype=torch.int32), (per_req_end_base + draft_offsets).to( - dtype=torch.int32 + return _build_target_verify_query_ranges_triton( + forward_batch.seq_lens[:bs].to(dtype=torch.int32), + bs, + draft_token_num, ) query_lens = getattr(forward_batch, "extend_seq_lens", None) @@ -305,9 +443,24 @@ def forward_sparse_mla_for_sglang( topk_tokens = topk_indices.shape[1] page_size = int(getattr(forward_batch.token_to_kv_pool, "page_size", 1)) - req_id_per_token = _build_sparse_req_id_per_token_for_sglang( - forward_batch, q.device - ) + target_verify = forward_batch.forward_mode.is_target_verify() + if target_verify: + bs = int(forward_batch.batch_size) + draft_token_num = int( + getattr(getattr(forward_batch, "spec_info", None), "draft_token_num", 0) + or 0 + ) + req_id_per_token, paged_kv_indptr = _build_target_verify_req_indptr_triton( + forward_batch.seq_lens[:bs].to(dtype=torch.int32), + bs, + draft_token_num, + topk_indices.shape[1], + ) + else: + req_id_per_token = _build_sparse_req_id_per_token_for_sglang( + forward_batch, q.device + ) + paged_kv_indptr = None block_table = _build_sglang_block_table(forward_batch, page_size).to( dtype=torch.int32 ) @@ -324,10 +477,13 @@ def forward_sparse_mla_for_sglang( ) fp8_sparse_mla = q.dtype == dtypes.fp8 or k_buffer.dtype == dtypes.fp8 - seq_len = (topk_indices != -1).sum(dim=-1).to(dtype=torch.int32) - paged_kv_indptr = torch.empty((num_tokens + 1,), dtype=torch.int32, device=q.device) - paged_kv_indptr[0].zero_() - torch.cumsum(seq_len, dim=0, out=paged_kv_indptr[1:]) + if paged_kv_indptr is None: + seq_len = (topk_indices != -1).sum(dim=-1).to(dtype=torch.int32) + paged_kv_indptr = torch.empty( + (num_tokens + 1,), dtype=torch.int32, device=q.device + ) + paged_kv_indptr[0].zero_() + torch.cumsum(seq_len, dim=0, out=paged_kv_indptr[1:]) paged_kv_indices = torch.empty( (num_tokens * topk_tokens,), dtype=torch.int32, device=q.device ) @@ -341,8 +497,8 @@ def forward_sparse_mla_for_sglang( NUM_TOPK_TOKENS=topk_tokens, ) - qo_indptr = torch.arange(num_tokens + 1, dtype=torch.int32, device=q.device) - last_page_len = torch.ones(num_tokens, dtype=torch.int32, device=q.device) + qo_indptr = _get_cached_i32_arange(num_tokens + 1, q.device) + last_page_len = _get_cached_i32_ones(num_tokens, q.device) work_metadata = None work_indptr = None @@ -598,9 +754,7 @@ def sparse_attn_indexer_sglang_plugin_mode( stride0=logits.stride(0), stride1=logits.stride(1), ) - topk_indices.copy_( - torch.where(topk_indices >= 0, topk_indices - cu_starts[:, None], topk_indices) - ) + _localize_topk_indices_triton(topk_indices, cu_starts) return weights diff --git a/atom/plugin/sglang/glm52_dsa_bridge.py b/atom/plugin/sglang/glm52_dsa_bridge.py new file mode 100644 index 0000000000..c4041589cb --- /dev/null +++ b/atom/plugin/sglang/glm52_dsa_bridge.py @@ -0,0 +1,888 @@ +"""Bridge SGLang ForwardBatch metadata to ATOM GLM-5.2 sparse MLA.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import numpy as np +import torch +from aiter import dtypes, get_mla_metadata_info_v1, get_mla_metadata_v1 +from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton + +from atom.plugin.sglang.runtime.model_arch import is_glm52_dsa_config + +_DECODE_GRAPH_BUFFERS_ATTR = "_atom_glm52_decode_graph_buffers" +_EMPTY_VALUE_CACHE_ATTR = "_atom_glm52_empty_value_cache" +_INDEXER_PAGE_SIZE_ATTR = "_atom_glm52_indexer_page_size" +_ATTENTION_PAGE_SIZE_ATTR = "_atom_glm52_attention_page_size" +_SHARED_SPARSE_INDICES_ATTR = "_atom_glm52_shared_sparse_kv_indices" +_CHUNK_K_WORKSPACE_ATTR = "_atom_glm52_chunk_k_workspace" +_CHUNK_V_WORKSPACE_ATTR = "_atom_glm52_chunk_v_workspace" + + +def is_glm52_dsa_arch(config) -> bool: + return is_glm52_dsa_config(config) + + +def maybe_get_glm52_dsa_pools_from_sglang_backend(forward_batch=None): + if forward_batch is not None: + token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) + req_to_token_pool = getattr(forward_batch, "req_to_token_pool", None) + if token_to_kv_pool is not None and req_to_token_pool is not None: + return token_to_kv_pool, req_to_token_pool + return None, None + + +def _get_seq_lens_cpu(forward_batch, bs: int) -> np.ndarray: + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + if seq_lens_cpu is None: + seq_lens_cpu = forward_batch.seq_lens.detach().cpu() + if torch.is_tensor(seq_lens_cpu): + seq_lens_cpu = seq_lens_cpu.detach().cpu().numpy() + return np.asarray(seq_lens_cpu[:bs], dtype=np.int32) + + +def _get_extend_lens_cpu(forward_batch, positions: torch.Tensor, bs: int) -> np.ndarray: + extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_lens is None: + extend_lens = getattr(forward_batch, "extend_seq_lens", None) + if extend_lens is not None: + if torch.is_tensor(extend_lens): + extend_lens = extend_lens.detach().cpu().numpy() + return np.asarray(extend_lens[:bs], dtype=np.int32) + + tokens_per_req = getattr( + getattr(forward_batch, "spec_info", None), "num_tokens_per_req", None + ) + if tokens_per_req is None: + tokens_per_req = max(1, int(positions.numel()) // max(1, bs)) + return np.full(bs, int(tokens_per_req), dtype=np.int32) + + +def _build_token_table( + forward_batch, + req_to_token_pool, + *, + seq_lens: np.ndarray, + extend_lens: np.ndarray | None, + page_size: int, +) -> torch.Tensor: + bs = int(forward_batch.batch_size) + prefix_lens = None + if extend_lens is not None and not forward_batch.forward_mode.is_decode_or_idle(): + # CUDA graph capture can present allocation-shaped dummy inputs where + # seq_lens is smaller than the fixed draft width. The token table still + # needs enough columns for the verifier/draft slots. + prefix_lens = np.maximum(seq_lens - extend_lens, 0).astype(np.int32) + table_lens = np.maximum(seq_lens, prefix_lens + extend_lens) + else: + table_lens = seq_lens + max_seq_len = int(table_lens.max(initial=1)) + req_pool_indices = forward_batch.req_pool_indices[:bs] + token_table = req_to_token_pool.req_to_token[req_pool_indices, :max_seq_len].clone() + + if extend_lens is not None and not forward_batch.forward_mode.is_decode_or_idle(): + offset = 0 + for req_idx, (prefix_len, query_len) in enumerate( + zip(prefix_lens, extend_lens) + ): + prefix_len = int(prefix_len) + query_len = int(query_len) + if query_len > 0: + token_table[req_idx, prefix_len : prefix_len + query_len] = ( + forward_batch.out_cache_loc[offset : offset + query_len] + ) + offset += query_len + + if page_size == 1: + return token_table.to(dtype=torch.int32).contiguous() + return (token_table[:, ::page_size] // page_size).to(dtype=torch.int32).contiguous() + + +def _flatten_kv_indices(token_table: torch.Tensor, lengths: np.ndarray) -> torch.Tensor: + pieces = [] + for row, length in enumerate(lengths): + if int(length) > 0: + pieces.append(token_table[row, : int(length)]) + if not pieces: + return torch.empty(0, dtype=torch.int32, device=token_table.device) + return torch.cat(pieces).to(dtype=torch.int32).contiguous() + + +def _counts_to_indptr(counts: np.ndarray, device: torch.device) -> torch.Tensor: + indptr = np.zeros(len(counts) + 1, dtype=np.int32) + if len(counts): + indptr[1:] = np.cumsum(counts, dtype=np.int32) + return torch.from_numpy(indptr).to(device=device) + + +def _get_index_topk(atom_config) -> int: + topk = getattr(atom_config.hf_config, "index_topk", None) + if topk is None: + raise RuntimeError("GLM-5.2 DSA bridge requires hf_config.index_topk") + return int(topk) + + +def _local_num_attention_heads(atom_config) -> int: + hf_config = atom_config.hf_config + num_heads = int(getattr(hf_config, "num_attention_heads")) + tp_size = int(getattr(atom_config, "tensor_parallel_size", 1)) + return max(1, num_heads // max(1, tp_size)) + + +def _metadata_dtype(atom_config): + kv_dtype = getattr(atom_config, "kv_cache_dtype", "bf16") + if str(kv_dtype).startswith("fp8"): + return dtypes.fp8 + return getattr(dtypes, "d_dtypes", {}).get(kv_dtype, torch.bfloat16) + + +def _make_mla_work_buffers( + *, + cu_seqlens_q: torch.Tensor, + kv_indptr: torch.Tensor, + kv_last_page_lens: torch.Tensor, + num_heads: int, + dtype_q, + dtype_kv, + page_size: int, +) -> dict[str, torch.Tensor]: + num_seqs = max(1, int(cu_seqlens_q.numel()) - 1) + max_q_len = 1 + if cu_seqlens_q.numel() > 1: + q_counts = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + max_q_len = max(1, int(q_counts.max().item())) + padded_heads = max(num_heads, 16) + ( + (work_meta_data_size, work_meta_data_type), + (work_indptr_size, work_indptr_type), + (work_info_set_size, work_info_set_type), + (reduce_indptr_size, reduce_indptr_type), + (reduce_final_map_size, reduce_final_map_type), + (reduce_partial_map_size, reduce_partial_map_type), + ) = get_mla_metadata_info_v1( + num_seqs, + max_q_len, + padded_heads, + dtype_q, + dtype_kv, + is_sparse=True, + fast_mode=True, + ) + device = cu_seqlens_q.device + work = { + "work_meta_data": torch.empty( + work_meta_data_size, dtype=work_meta_data_type, device=device + ), + "work_indptr": torch.empty( + work_indptr_size, dtype=work_indptr_type, device=device + ), + "work_info_set": torch.empty( + work_info_set_size, dtype=work_info_set_type, device=device + ), + "reduce_indptr": torch.empty( + reduce_indptr_size, dtype=reduce_indptr_type, device=device + ), + "reduce_final_map": torch.empty( + reduce_final_map_size, dtype=reduce_final_map_type, device=device + ), + "reduce_partial_map": torch.empty( + reduce_partial_map_size, dtype=reduce_partial_map_type, device=device + ), + } + get_mla_metadata_v1( + cu_seqlens_q, + kv_indptr, + kv_last_page_lens, + padded_heads, + 1, + True, + work["work_meta_data"], + work["work_info_set"], + work["work_indptr"], + work["reduce_indptr"], + work["reduce_final_map"], + work["reduce_partial_map"], + page_size=page_size, + dtype_q=dtype_q, + dtype_kv=dtype_kv, + kv_granularity=max(page_size, 16), + max_seqlen_qo=max_q_len, + uni_seqlen_qo=max_q_len, + fast_mode=True, + ) + return work + + +def _ensure_shared_sparse_buffer( + token_to_kv_pool, + *, + num_tokens: int, + topk: int, + device: torch.device, +) -> torch.Tensor: + required = max(1, int(num_tokens) * int(topk)) + buffer = getattr(token_to_kv_pool, _SHARED_SPARSE_INDICES_ATTR, None) + if ( + buffer is None + or buffer.device != device + or buffer.dtype != torch.int32 + or buffer.numel() < required + ): + buffer = torch.empty(required, dtype=torch.int32, device=device) + setattr(token_to_kv_pool, _SHARED_SPARSE_INDICES_ATTR, buffer) + return buffer[:required] + + +def _ensure_chunk_workspace( + token_to_kv_pool, + *, + num_tokens: int, + num_heads: int, + qk_head_dim: int, + v_head_dim: int, + dtype, + device: torch.device, +): + required = max(1, int(num_tokens)) + k_buf = getattr(token_to_kv_pool, _CHUNK_K_WORKSPACE_ATTR, None) + v_buf = getattr(token_to_kv_pool, _CHUNK_V_WORKSPACE_ATTR, None) + if ( + k_buf is None + or k_buf.device != device + or k_buf.dtype != dtype + or k_buf.shape[0] < required + or k_buf.shape[1] != int(num_heads) + or k_buf.shape[2] != int(qk_head_dim) + ): + k_buf = torch.empty( + (required, int(num_heads), int(qk_head_dim)), + dtype=dtype, + device=device, + ) + setattr(token_to_kv_pool, _CHUNK_K_WORKSPACE_ATTR, k_buf) + if ( + v_buf is None + or v_buf.device != device + or v_buf.dtype != dtype + or v_buf.shape[0] < required + or v_buf.shape[1] != int(num_heads) + or v_buf.shape[2] != int(v_head_dim) + ): + v_buf = torch.empty( + (required, int(num_heads), int(v_head_dim)), + dtype=dtype, + device=device, + ) + setattr(token_to_kv_pool, _CHUNK_V_WORKSPACE_ATTR, v_buf) + return k_buf[:required], v_buf[:required] + + +class _GLM52DecodeGraphBuffers: + def __init__( + self, + *, + max_bs: int, + max_context_len: int, + indexer_page_size: int, + attention_page_size: int, + index_topk: int, + num_heads: int, + dtype_q, + dtype_kv, + device: torch.device, + ) -> None: + self.max_bs = int(max_bs) + self.max_context_len = int(max_context_len) + self.indexer_page_size = int(indexer_page_size) + self.attention_page_size = int(attention_page_size) + self.index_topk = int(index_topk) + self.device = device + + max_blocks = max( + 1, + (self.max_context_len + self.indexer_page_size - 1) + // self.indexer_page_size, + ) + self.cu_q = torch.arange(self.max_bs + 1, dtype=torch.int32, device=device) + self.kv_indptr = torch.zeros(self.max_bs + 1, dtype=torch.int32, device=device) + self.sparse_kv_indptr = torch.zeros( + self.max_bs + 1, dtype=torch.int32, device=device + ) + self.kv_indices = torch.empty( + self.max_bs * self.max_context_len, dtype=torch.int32, device=device + ) + self.kv_last_page_lens = torch.ones( + self.max_bs, dtype=torch.int32, device=device + ) + self.block_tables = torch.empty( + self.max_bs, max_blocks, dtype=torch.int32, device=device + ) + self.context_lens = torch.zeros(self.max_bs, dtype=torch.int32, device=device) + self.slot_mapping = torch.zeros(self.max_bs, dtype=torch.int64, device=device) + self.shared_sparse = torch.empty( + self.max_bs * self.index_topk, dtype=torch.int32, device=device + ) + + work = _make_mla_work_buffers( + cu_seqlens_q=self.cu_q, + kv_indptr=self.sparse_kv_indptr, + kv_last_page_lens=self.kv_last_page_lens, + num_heads=num_heads, + dtype_q=dtype_q, + dtype_kv=dtype_kv, + page_size=self.attention_page_size, + ) + self.work_meta_data = work["work_meta_data"] + self.work_indptr = work["work_indptr"] + self.work_info_set = work["work_info_set"] + self.reduce_indptr = work["reduce_indptr"] + self.reduce_final_map = work["reduce_final_map"] + self.reduce_partial_map = work["reduce_partial_map"] + + def stage_block_tables(self, req_to_token_pool, req_pool_indices, bs: int) -> None: + req_to_token = req_to_token_pool.req_to_token + live = req_to_token[ + req_pool_indices[:bs], + : self.max_context_len : self.indexer_page_size, + ] + self.block_tables[:bs, : live.shape[1]].copy_( + (live // self.indexer_page_size).to(torch.int32) + ) + + +def _get_or_create_decode_graph_buffers( + token_to_kv_pool, + *, + max_bs: int, + max_context_len: int, + indexer_page_size: int, + attention_page_size: int, + atom_config, + device: torch.device, +) -> _GLM52DecodeGraphBuffers: + topk = _get_index_topk(atom_config) + dtype_q = _metadata_dtype(atom_config) + bufs = getattr(token_to_kv_pool, _DECODE_GRAPH_BUFFERS_ATTR, None) + if ( + bufs is None + or bufs.max_bs < int(max_bs) + or bufs.max_context_len < int(max_context_len) + or bufs.indexer_page_size != int(indexer_page_size) + or bufs.attention_page_size != int(attention_page_size) + or bufs.index_topk != int(topk) + or bufs.device != device + ): + bufs = _GLM52DecodeGraphBuffers( + max_bs=max_bs, + max_context_len=max_context_len, + indexer_page_size=indexer_page_size, + attention_page_size=attention_page_size, + index_topk=topk, + num_heads=_local_num_attention_heads(atom_config), + dtype_q=dtype_q, + dtype_kv=dtype_q, + device=device, + ) + setattr(token_to_kv_pool, _DECODE_GRAPH_BUFFERS_ATTR, bufs) + setattr(token_to_kv_pool, _SHARED_SPARSE_INDICES_ATTR, bufs.shared_sparse) + return bufs + + +def _validate_page_size(token_to_kv_pool, atom_config) -> int: + page_size = int(getattr(token_to_kv_pool, "page_size", 1)) + from atom.utils import envs + + atom_config.kv_cache_block_size = page_size + setattr(token_to_kv_pool, _INDEXER_PAGE_SIZE_ATTR, page_size) + setattr(token_to_kv_pool, _ATTENTION_PAGE_SIZE_ATTR, int(envs.ATOM_MLA_PAGE_SIZE)) + return page_size + + +def _attention_page_size(token_to_kv_pool) -> int: + return int(getattr(token_to_kv_pool, _ATTENTION_PAGE_SIZE_ATTR, 1)) + + +def _build_decode_metadata( + forward_batch, + positions: torch.Tensor, + *, + token_to_kv_pool, + req_to_token_pool, + atom_config, +): + from atom.utils.forward_context import AttentionMetaData, AttnState + + del positions + device = forward_batch.seq_lens.device + bs = int(forward_batch.batch_size) + seq_lens = _get_seq_lens_cpu(forward_batch, bs) + topk = _get_index_topk(atom_config) + page_size = _validate_page_size(token_to_kv_pool, atom_config) + + cu_q = torch.arange(bs + 1, dtype=torch.int32, device=device) + block_tables = _build_token_table( + forward_batch, + req_to_token_pool, + seq_lens=seq_lens, + extend_lens=None, + page_size=page_size, + ) + token_table = _build_token_table( + forward_batch, + req_to_token_pool, + seq_lens=seq_lens, + extend_lens=None, + page_size=1, + ) + kv_indptr = _counts_to_indptr(seq_lens, device) + kv_indices = _flatten_kv_indices(token_table, seq_lens) + # Sparse decode consumes a compact list of selected physical token ids, so + # each selected entry behaves as a single-token page even when the backing + # cache is stored in page64/segmented layout. + kv_last_page_lens = torch.ones(bs, dtype=torch.int32, device=device) + sparse_counts = np.minimum(seq_lens, topk).astype(np.int32) + sparse_kv_indptr = _counts_to_indptr(sparse_counts, device) + + _ensure_shared_sparse_buffer( + token_to_kv_pool, + num_tokens=bs, + topk=topk, + device=device, + ) + dtype_q = _metadata_dtype(atom_config) + work = _make_mla_work_buffers( + cu_seqlens_q=cu_q, + kv_indptr=sparse_kv_indptr, + kv_last_page_lens=kv_last_page_lens, + num_heads=_local_num_attention_heads(atom_config), + dtype_q=dtype_q, + dtype_kv=dtype_q, + page_size=_attention_page_size(token_to_kv_pool), + ) + + md = AttentionMetaData( + cu_seqlens_q=cu_q, + cu_seqlens_k=kv_indptr, + max_seqlen_q=1, + max_seqlen_k=int(seq_lens.max(initial=1)), + slot_mapping=forward_batch.out_cache_loc[:bs], + context_lens=forward_batch.seq_lens[:bs].to(dtype=torch.int32), + block_tables=block_tables, + state=AttnState.DECODE, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_last_page_lens=kv_last_page_lens, + sparse_kv_indptr=sparse_kv_indptr, + **work, + ) + md.dtype_q = dtype_q + return md + + +def build_atom_glm52_decode_graph_metadata_from_sglang( + forward_batch, + positions: torch.Tensor, + *, + token_to_kv_pool, + req_to_token_pool, + atom_config, + max_bs: int | None = None, + max_context_len: int | None = None, +): + from atom.utils.forward_context import AttentionMetaData, AttnState + + del positions + device = forward_batch.seq_lens.device + bs = int(forward_batch.batch_size) + seq_lens = forward_batch.seq_lens[:bs].to(dtype=torch.int32) + if max_context_len is None: + req_to_token = req_to_token_pool.req_to_token + max_context_len = int(req_to_token.shape[1]) + if max_bs is None: + max_bs = max(bs, int(getattr(req_to_token_pool, "size", bs))) + + indexer_page_size = _validate_page_size(token_to_kv_pool, atom_config) + attention_page_size = _attention_page_size(token_to_kv_pool) + topk = _get_index_topk(atom_config) + dtype_q = _metadata_dtype(atom_config) + + bufs = _get_or_create_decode_graph_buffers( + token_to_kv_pool, + max_bs=max_bs, + max_context_len=max_context_len, + indexer_page_size=indexer_page_size, + attention_page_size=attention_page_size, + atom_config=atom_config, + device=device, + ) + + bufs.kv_indptr.zero_() + bufs.kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + bufs.sparse_kv_indptr.zero_() + bufs.sparse_kv_indptr[1 : bs + 1] = torch.cumsum( + torch.clamp(seq_lens, max=topk), dim=0 + ) + bufs.context_lens[:bs].copy_(seq_lens) + bufs.kv_last_page_lens[:bs].fill_(1) + + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + if torch.is_tensor(out_cache_loc): + copy_n = min(bs, int(out_cache_loc.numel())) + if copy_n: + bufs.slot_mapping[:copy_n].copy_(out_cache_loc[:copy_n]) + if bs > copy_n: + scratch_slot = max(0, int(getattr(token_to_kv_pool, "size", 1)) - 1) + bufs.slot_mapping[copy_n:bs].fill_(scratch_slot) + else: + scratch_slot = max(0, int(getattr(token_to_kv_pool, "size", 1)) - 1) + bufs.slot_mapping[:bs].fill_(scratch_slot) + + create_flashinfer_kv_indices_triton[(bs,)]( + req_to_token_pool.req_to_token, + forward_batch.req_pool_indices[:bs], + seq_lens, + bufs.kv_indptr[: bs + 1], + None, + bufs.kv_indices, + req_to_token_pool.req_to_token.stride(0), + ) + bufs.stage_block_tables(req_to_token_pool, forward_batch.req_pool_indices, bs) + + get_mla_metadata_v1( + bufs.cu_q[: bs + 1], + bufs.sparse_kv_indptr[: bs + 1], + bufs.kv_last_page_lens[:bs], + max(_local_num_attention_heads(atom_config), 16), + 1, + True, + bufs.work_meta_data, + bufs.work_info_set, + bufs.work_indptr, + bufs.reduce_indptr, + bufs.reduce_final_map, + bufs.reduce_partial_map, + page_size=attention_page_size, + dtype_q=dtype_q, + dtype_kv=dtype_q, + kv_granularity=max(attention_page_size, 16), + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=True, + ) + + setattr(token_to_kv_pool, _SHARED_SPARSE_INDICES_ATTR, bufs.shared_sparse) + md = AttentionMetaData( + cu_seqlens_q=bufs.cu_q[: bs + 1], + cu_seqlens_k=bufs.kv_indptr[: bs + 1], + max_seqlen_q=1, + max_seqlen_k=int(seq_lens.max().item()) if bs else 1, + slot_mapping=bufs.slot_mapping[:bs], + context_lens=bufs.context_lens[:bs], + block_tables=bufs.block_tables[:bs], + state=AttnState.DECODE, + kv_indptr=bufs.kv_indptr[: bs + 1], + kv_indices=bufs.kv_indices, + kv_last_page_lens=bufs.kv_last_page_lens[:bs], + sparse_kv_indptr=bufs.sparse_kv_indptr[: bs + 1], + work_meta_data=bufs.work_meta_data, + work_indptr=bufs.work_indptr, + work_info_set=bufs.work_info_set, + reduce_indptr=bufs.reduce_indptr, + reduce_final_map=bufs.reduce_final_map, + reduce_partial_map=bufs.reduce_partial_map, + ) + md.dtype_q = dtype_q + return md + + +def _build_prefill_metadata( + forward_batch, + positions: torch.Tensor, + *, + token_to_kv_pool, + req_to_token_pool, + atom_config, +): + from atom.utils.forward_context import AttentionMetaData, AttnState + + device = positions.device + bs = int(forward_batch.batch_size) + seq_lens = _get_seq_lens_cpu(forward_batch, bs) + is_target_verify = forward_batch.forward_mode.is_target_verify() + if is_target_verify: + draft_token_num = int( + getattr(getattr(forward_batch, "spec_info", None), "draft_token_num", 0) + or 0 + ) + if draft_token_num <= 0: + raise RuntimeError("GLM-5.2 DSA target_verify requires draft_token_num") + extend_lens = np.full(bs, draft_token_num, dtype=np.int32) + position_rows = positions.detach().cpu().numpy().astype(np.int32) + if position_rows.size < bs * draft_token_num: + raise RuntimeError( + "GLM-5.2 DSA target_verify positions are shorter than " + f"bs*draft_token_num: positions={position_rows.size}, " + f"bs={bs}, draft_token_num={draft_token_num}" + ) + # In graph replay, forward_batch.seq_lens may be a padded/capture tensor. + # The verifier positions are the reliable SGLang token-layout truth: + # [prefix, prefix+1, ...] per request. + seq_lens = position_rows[: bs * draft_token_num : draft_token_num] + # TARGET_VERIFY appends verifier draft slots after the committed prefix. + # The attention metadata must expose those slots as part of the KV range. + seq_lens = seq_lens + extend_lens + force_total_kv = int(os.environ.get("ATOM_GLM52_TV_FORCE_TOTAL_KV", "0") or 0) + if force_total_kv > 0: + seq_lens = np.maximum(seq_lens, force_total_kv).astype(np.int32) + else: + extend_lens = _get_extend_lens_cpu(forward_batch, positions, bs) + if getattr( + forward_batch.forward_mode, "is_draft_extend", lambda **kwargs: False + )(include_v2=True): + num_position_rows = int(positions.numel()) + if int(extend_lens.sum()) != num_position_rows and bs > 0: + if num_position_rows % bs != 0: + raise RuntimeError( + "GLM-5.2 DSA draft_extend positions cannot be evenly " + f"distributed: positions={num_position_rows}, bs={bs}, " + f"extend_lens_sum={int(extend_lens.sum())}" + ) + extend_lens = np.full( + bs, num_position_rows // bs, dtype=np.int32 + ) + topk = _get_index_topk(atom_config) + page_size = _validate_page_size(token_to_kv_pool, atom_config) + cached_lens = np.maximum(seq_lens - extend_lens, 0).astype(np.int32) + seq_lens = np.maximum(seq_lens, cached_lens + extend_lens).astype(np.int32) + + q_np = np.zeros(bs + 1, dtype=np.int32) + q_np[1:] = np.cumsum(extend_lens, dtype=np.int32) + cu_q = torch.from_numpy(q_np).to(device=device) + kv_indptr = _counts_to_indptr(seq_lens, device) + block_tables = _build_token_table( + forward_batch, + req_to_token_pool, + seq_lens=seq_lens, + extend_lens=extend_lens, + page_size=page_size, + ) + token_table = _build_token_table( + forward_batch, + req_to_token_pool, + seq_lens=seq_lens, + extend_lens=extend_lens, + page_size=1, + ) + kv_indices = _flatten_kv_indices(token_table, seq_lens) + has_cached = bool(np.any(cached_lens > 0)) + state = AttnState.PREFILL_PREFIX if has_cached else AttnState.PREFILL_NATIVE + total_tokens = int(extend_lens.sum()) + kv_last_page_lens = torch.ones(bs, dtype=torch.int32, device=device) + + md = AttentionMetaData( + cu_seqlens_q=cu_q, + cu_seqlens_k=kv_indptr if has_cached else cu_q, + max_seqlen_q=int(extend_lens.max(initial=1)), + max_seqlen_k=int(seq_lens.max(initial=1)), + slot_mapping=forward_batch.out_cache_loc[:total_tokens], + context_lens=( + torch.from_numpy((seq_lens - extend_lens).astype(np.int32)).to( + device=device + ) + if is_target_verify + else forward_batch.seq_lens[:bs].to(dtype=torch.int32) + ), + block_tables=block_tables, + state=state, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_last_page_lens=kv_last_page_lens, + has_cached=has_cached, + total_kv=int(seq_lens.sum()), + num_cached_tokens=torch.from_numpy(cached_lens).to(device=device), + seq_starts=torch.from_numpy(cached_lens).to(device=device), + ) + dtype_q = _metadata_dtype(atom_config) + md.dtype_q = dtype_q + + if ( + is_target_verify + and has_cached + and os.environ.get("ATOM_GLM52_TV_FORCE_CHUNKED", "0").lower() + in ("1", "true", "yes", "on") + ): + prefix_lens = cached_lens.astype(np.int32) + prefix_indptr = _counts_to_indptr(prefix_lens, device) + prefix_indices = _flatten_kv_indices(token_table, prefix_lens) + num_heads = _local_num_attention_heads(atom_config) + qk_head_dim = int(getattr(atom_config.hf_config, "qk_nope_head_dim")) + int( + getattr(atom_config.hf_config, "qk_rope_head_dim") + ) + v_head_dim = int(getattr(atom_config.hf_config, "v_head_dim")) + k_workspace, v_workspace = _ensure_chunk_workspace( + token_to_kv_pool, + num_tokens=int(prefix_lens.sum()), + num_heads=num_heads, + qk_head_dim=qk_head_dim, + v_head_dim=v_head_dim, + dtype=getattr(atom_config, "torch_dtype", torch.bfloat16), + device=device, + ) + md.mla_chunk_meta = SimpleNamespace( + kv_indptr=[prefix_indptr], + kv_indices=[prefix_indices], + cu_seqlens_k=[prefix_indptr], + total_tokens=[int(prefix_lens.sum())], + max_seqlen_k=[int(prefix_lens.max(initial=0))], + num_chunks=1, + k_workspace=k_workspace, + v_workspace=v_workspace, + shuffle_kv_block_indptr=None, + shuffle_kv_block_indices=None, + ) + + if md.max_seqlen_k > topk: + counts = extend_lens.astype(np.int32) + local_offsets = np.concatenate( + [np.arange(int(count), dtype=np.int32) for count in counts] + ) + if has_cached: + seq_starts = kv_indptr[:-1].detach().cpu().numpy().astype(np.int32) + repeated_seq_starts = np.repeat(seq_starts, counts) + repeated_cached_lens = np.repeat(cached_lens, counts) + cu_ks = repeated_seq_starts + cu_ke = repeated_seq_starts + repeated_cached_lens + local_offsets + 1 + sparse_counts = repeated_cached_lens + local_offsets + 1 + else: + cu_ks = np.repeat(q_np[:bs], counts) + cu_ke = np.arange(total_tokens, dtype=np.int32) + 1 + sparse_counts = local_offsets + 1 + + sparse_cu = torch.arange(total_tokens + 1, dtype=torch.int32, device=device) + sparse_kv_indptr = _counts_to_indptr( + np.minimum(sparse_counts, topk).astype(np.int32), device + ) + sparse_last_page_lens = torch.ones( + total_tokens, dtype=torch.int32, device=device + ) + md.cu_seqlen_ks = torch.from_numpy(cu_ks.astype(np.int32)).to(device=device) + md.cu_seqlen_ke = torch.from_numpy(cu_ke.astype(np.int32)).to(device=device) + md.sparse_cu_seqlens_q = sparse_cu + md.sparse_kv_indptr = sparse_kv_indptr + md.kv_last_page_lens = sparse_last_page_lens + md.token_to_seq_idxs = torch.repeat_interleave( + torch.arange(bs, dtype=torch.int32, device=device), + torch.from_numpy(counts.astype(np.int64)).to(device=device), + ) + _ensure_shared_sparse_buffer( + token_to_kv_pool, + num_tokens=total_tokens, + topk=topk, + device=device, + ) + sparse_work = _make_mla_work_buffers( + cu_seqlens_q=sparse_cu, + kv_indptr=sparse_kv_indptr, + kv_last_page_lens=sparse_last_page_lens, + num_heads=_local_num_attention_heads(atom_config), + dtype_q=dtype_q, + dtype_kv=dtype_q, + page_size=_attention_page_size(token_to_kv_pool), + ) + for key, value in sparse_work.items(): + setattr(md, f"sparse_prefill_{key}", value) + else: + _ensure_shared_sparse_buffer( + token_to_kv_pool, + num_tokens=max(1, total_tokens), + topk=topk, + device=device, + ) + + return md + + +def build_atom_glm52_attention_metadata_from_sglang( + forward_batch, + positions: torch.Tensor, + *, + token_to_kv_pool, + req_to_token_pool, + atom_config, +): + if forward_batch.forward_mode.is_decode_or_idle(): + return _build_decode_metadata( + forward_batch, + positions, + token_to_kv_pool=token_to_kv_pool, + req_to_token_pool=req_to_token_pool, + atom_config=atom_config, + ) + return _build_prefill_metadata( + forward_batch, + positions, + token_to_kv_pool=token_to_kv_pool, + req_to_token_pool=req_to_token_pool, + atom_config=atom_config, + ) + + +def bind_glm52_dsa_cache_views(model, token_to_kv_pool) -> bool: + if token_to_kv_pool is None or not hasattr(token_to_kv_pool, "get_key_buffer"): + return False + if not hasattr(token_to_kv_pool, "get_index_k_with_scale_buffer"): + return False + + from atom.config import KVCacheTensor + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + from atom.utils.forward_context import get_forward_context, set_kv_cache_data + + shared_sparse = getattr(token_to_kv_pool, _SHARED_SPARSE_INDICES_ATTR, None) + if shared_sparse is None: + return False + + page_size = int( + getattr( + token_to_kv_pool, + _INDEXER_PAGE_SIZE_ATTR, + getattr(token_to_kv_pool, "page_size", 1), + ) + ) + empty_value_cache = getattr(token_to_kv_pool, _EMPTY_VALUE_CACHE_ATTR, None) + if empty_value_cache is None or empty_value_cache.device != shared_sparse.device: + empty_value_cache = torch.empty(0, device=shared_sparse.device) + setattr(token_to_kv_pool, _EMPTY_VALUE_CACHE_ATTR, empty_value_cache) + kv_cache_data = {} + for module in model.modules(): + if not isinstance(module, DeepseekV2MLAAttention): + continue + + layer_id = int(module.layer_num) + kv_cache_data[f"layer_{layer_id}"] = KVCacheTensor( + layer_num=layer_id, + k_cache=token_to_kv_pool.get_key_buffer(layer_id), + v_cache=empty_value_cache, + k_scale=getattr(module.mla_attn, "_k_scale", None), + v_scale=getattr(module.mla_attn, "_k_scale", None), + ) + + indexer = getattr(module, "indexer", None) + if indexer is not None: + index_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) + index_entry_dim = int(getattr(indexer, "head_dim")) + 4 + indexer.k_cache.kv_cache[0] = index_cache.view( + -1, page_size, index_entry_dim + ) + indexer.sparse_kv_indices_buffer = shared_sparse + + if hasattr(module.mla_attn, "sparse_kv_indices_buffer"): + module.mla_attn.sparse_kv_indices_buffer = shared_sparse + + if not kv_cache_data: + return False + + set_kv_cache_data(kv_cache_data) + get_forward_context().kv_cache_data = kv_cache_data + return True diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index ff9d9c7554..cdeb95ce06 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -8,6 +8,7 @@ import inspect import logging +import os from typing import Any, Iterable, Optional, Tuple, Union import torch @@ -210,24 +211,8 @@ def forward( input_embeds=input_embeds, set_forward_context=not self.model_arch_spec.wrapper_binds_gdn_context, ) as runtime: - if self.model_arch == "DeepseekV4ForCausalLM": - from atom.plugin.sglang.deepseek_v4_bridge import ( - bind_deepseek_v4_proxy_cache_views, - maybe_get_proxy_pool_from_sglang_backend, - reset_deepseek_v4_state_slots, - ) - - proxy_pool, _ = maybe_get_proxy_pool_from_sglang_backend() - if not bind_deepseek_v4_proxy_cache_views(self.model, proxy_pool): - raise RuntimeError( - "DeepSeek-V4 SGLang proxy KV pool is not initialized" - ) - from atom.utils.forward_context import get_forward_context - - reset_slots = getattr( - get_forward_context().attn_metadata, "reset_slots", None - ) - reset_deepseek_v4_state_slots(self.model, reset_slots) + if self.model_arch_spec.bind_cache_views is not None: + self.model_arch_spec.bind_cache_views(self.model, runtime) metadata = SGLangForwardBatchMetadata.build( runtime.forward_batch, @@ -269,6 +254,197 @@ def forward( **self._filter_model_forward_kwargs(model_call_kwargs) ) + try: + mode = getattr(forward_batch, "forward_mode", None) + if ( + os.environ.get("ATOM_GLM52_ATTENTION_DEBUG_LOG", "0") + in ("1", "true", "True") + and mode is not None + and bool(getattr(mode, "is_target_verify", lambda: False)()) + and "GLM" in str(self.model_arch).upper() + and not torch.cuda.is_current_stream_capturing() + ): + dummy_context = ( + torch.is_tensor(positions) + and int(positions.numel()) > 0 + and bool(torch.all(positions == 0).item()) + and torch.is_tensor(input_ids) + and int(input_ids.numel()) > 0 + and bool(torch.all(input_ids == 0).item()) + ) + collected = [] + if dummy_context: + # Skip capture/dummy target-verify forwards; they use + # all-zero rows and otherwise drown real request logs. + configured_layers = set() + else: + configured = os.environ.get( + "ATOM_GLM52_ATTENTION_DEBUG_LAYERS", "" + ) + configured_layers = None + if configured.strip().lower() not in ("all", "*"): + configured_layers = { + int(item) + for item in configured.replace(" ", ",").split(",") + if item.strip() + } + for module in self.model.modules(): + buf = getattr(module, "_atom_glm52_attn_debug", None) + if not torch.is_tensor(buf): + continue + layer = getattr(module, "layer_num", None) + if ( + configured_layers is not None + and int(layer) not in configured_layers + ): + continue + sparse_buf = getattr(module, "sparse_kv_indices_buffer", None) + sparse_info = None + if torch.is_tensor(sparse_buf): + flat_sparse = sparse_buf.reshape(-1) + sparse_info = { + "shape": tuple(sparse_buf.shape), + "head": flat_sparse[ + : min(16, int(flat_sparse.numel())) + ] + .detach() + .cpu() + .tolist(), + "tail": flat_sparse[ + max(0, int(flat_sparse.numel()) - 16) : + ] + .detach() + .cpu() + .tolist(), + } + collected.append( + { + "layer": int(layer) if layer is not None else None, + "values": buf.detach().cpu().tolist(), + "sparse": sparse_info, + } + ) + if collected: + context = { + "input_ids": input_ids[ + : min(12, int(input_ids.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(input_ids) + else None, + "positions": positions[ + : min(12, int(positions.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor(positions) + else None, + "out_cache_loc": ( + forward_batch.out_cache_loc[ + : min( + 12, + int(forward_batch.out_cache_loc.numel()), + ) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(forward_batch, "out_cache_loc", None) + ) + else None + ), + "seq_lens": ( + forward_batch.seq_lens[ + : min(12, int(forward_batch.seq_lens.numel())) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(forward_batch, "seq_lens", None) + ) + else None + ), + "req_pool_indices": ( + forward_batch.req_pool_indices[ + : min( + 12, + int(forward_batch.req_pool_indices.numel()), + ) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(forward_batch, "req_pool_indices", None) + ) + else None + ), + } + logger.info( + "GLM52 attention layer debug: where=eager context=%s values=%s", + context, + collected, + ) + if ( + os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") + in ("1", "true", "True") + and mode is not None + and bool(getattr(mode, "is_target_verify", lambda: False)()) + and "GLM" in str(self.model_arch).upper() + and torch.is_tensor(hidden_states) + and not torch.cuda.is_current_stream_capturing() + ): + rows = hidden_states[ + : min(4, int(hidden_states.shape[0])) + ].detach().float() + dim = int(rows.shape[-1]) + checksum_slices = [] + for start in ( + 0, + 256, + 1024, + 2048, + 4096, + max(0, dim - 256), + ): + end = min(dim, start + 256) + if start < end: + checksum_slices.append(rows[:, start:end].sum(dim=-1)) + checksum = ( + torch.stack(checksum_slices, dim=-1) + if checksum_slices + else rows.new_empty((rows.shape[0], 0)) + ) + logger.info( + "GLM52 wrapper target output debug: raw_shape=%s " + "norm=%s absmax=%s mean=%s checksum=%s input_ids=%s " + "positions=%s", + tuple(hidden_states.shape), + rows.norm(dim=-1).cpu().tolist(), + rows.abs().amax(dim=-1).cpu().tolist(), + rows.mean(dim=-1).cpu().tolist(), + checksum.cpu().tolist(), + input_ids[: min(12, int(input_ids.numel()))] + .detach() + .cpu() + .tolist() + if torch.is_tensor(input_ids) + else None, + positions[: min(12, int(positions.numel()))] + .detach() + .cpu() + .tolist() + if torch.is_tensor(positions) + else None, + ) + except Exception: + logger.exception("Failed to log GLM52 wrapper target output debug") + hidden_states = runtime.trim_output(hidden_states) logits_input_ids = input_ids try: diff --git a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py index 9fc734a0dd..d79156813d 100644 --- a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py +++ b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py @@ -14,7 +14,7 @@ from torch import nn from sglang.srt.distributed import get_pp_group -from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import get_global_server_args @@ -376,6 +376,16 @@ def forward( if self.pp_group.is_last_rank: hidden_states = runtime.trim_output(hidden_states) + if ( + getattr(forward_batch, "_atom_use_draft_argmax", False) + and hasattr(self.model, "compute_draft_token") + ): + draft_token_ids = self.model.compute_draft_token(hidden_states) + return LogitsProcessorOutput( + next_token_logits=None, + hidden_states=hidden_states, + customized_info={"draft_token_ids": draft_token_ids}, + ) return self.logits_processor( input_ids, hidden_states, diff --git a/atom/plugin/sglang/models/glm52_dsa.py b/atom/plugin/sglang/models/glm52_dsa.py new file mode 100644 index 0000000000..eb55a31784 --- /dev/null +++ b/atom/plugin/sglang/models/glm52_dsa.py @@ -0,0 +1,64 @@ +"""Model-level GLM-5.2 DSA adaptation for SGLang plugin mode.""" + +from __future__ import annotations + +from typing import Any + +from atom.plugin.sglang.models.deepseek_mla import _align_qknorm_fusion_for_sglang +from atom.plugin.sglang.models.deepseek_mla_forward import ( + _patch_attention_projs_for_sglang_mxfp4, +) +from atom.plugin.sglang.models.glm52_dsa_attention import ( + SGLangATOMGLM52MLAAttention, +) + + +def setup_glm52_dsa_for_sglang(model: Any) -> None: + """Patch GLM-5.2 for native ATOM sparse MLA under SGLang. + + This deliberately does not install ``SGLangDeepseekMLAAttention``. GLM-5.2 + should keep ATOM's native ``MLAAttention`` frontend so full-index layers run + the ATOM indexer into a shared physical-index buffer and shared-index layers + reuse that buffer. + """ + + if not hasattr(model, "atom_config"): + from atom.config import get_current_atom_config + + model.atom_config = get_current_atom_config() + + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + from sglang.srt.configs.model_config import is_deepseek_nsa + from sglang.srt.layers.communicator import get_attn_tp_context + + config = model.config + get_attn_tp_context().init_context(config.q_lora_rank, is_deepseek_nsa(config)) + + last_full_index_seen = False + for module in model.modules(): + if not isinstance(module, DeepseekV2MLAAttention): + continue + + _align_qknorm_fusion_for_sglang(module) + _patch_attention_projs_for_sglang_mxfp4(module) + + if not isinstance(module.mla_attn, SGLangATOMGLM52MLAAttention): + raise RuntimeError( + "GLM-5.2 SGLang native DSA setup expected " + "SGLangATOMGLM52MLAAttention. Ensure the GLM construction " + "context is installed before model initialization." + ) + + if getattr(module, "is_v32", False): + owns_active_indexer = getattr( + module, "indexer", None + ) is not None and not getattr(module, "skip_topk", False) + if owns_active_indexer: + last_full_index_seen = True + elif not last_full_index_seen: + raise RuntimeError( + "GLM-5.2 IndexShare cannot start with a shared-index layer; " + f"layer={getattr(module, 'prefix', '')!r}" + ) + + model._atom_sglang_uses_glm52_native_dsa = True diff --git a/atom/plugin/sglang/models/glm52_dsa_attention.py b/atom/plugin/sglang/models/glm52_dsa_attention.py new file mode 100644 index 0000000000..533144a78a --- /dev/null +++ b/atom/plugin/sglang/models/glm52_dsa_attention.py @@ -0,0 +1,135 @@ +"""GLM-5.2 native MLA attention frontend for SGLang plugin mode.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any + +import torch + +from atom.model_ops.attention_mla import MLAAttention + + +@contextmanager +def glm52_native_mla_attention_construction(): + """Temporarily make GLM sparse MLA layers construct native ATOM attention.""" + + import atom.models.deepseek_v2 as deepseek_v2 + + previous = deepseek_v2.Attention + + def _build_glm52_native_mla_attention(*args: Any, **kwargs: Any): + mla_modules = kwargs.get("mla_modules", None) + if ( + kwargs.get("use_mla", False) + and mla_modules is not None + and getattr(mla_modules, "is_sparse", False) + ): + return SGLangATOMGLM52MLAAttention(*args, **kwargs) + return previous(*args, **kwargs) + + deepseek_v2.Attention = _build_glm52_native_mla_attention + try: + yield + finally: + deepseek_v2.Attention = previous + + +class SGLangATOMGLM52MLAAttention(MLAAttention): + """Use ATOM native ``MLAAttention`` under SGLang plugin runtime. + + ``DeepseekV2MLAAttention.forward`` calls ``self.mla_attn`` with the ATOM + model-side argument order, while ``MLAAttention.forward`` has the lower-level + backend argument order. This frontend keeps the model-side call contract and + delegates directly to ``forward_impl``. + """ + + def __init__(self, *args: Any, prefix: str | None = None, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.layer_name = ( + prefix if prefix is not None else f"GLM52_MLA_{self.layer_num}" + ) + from atom.config import get_current_atom_config + + get_current_atom_config().compilation_config.static_forward_context[ + self.layer_name + ] = self + + def _get_forward_batch(self): + from atom.plugin.sglang.runtime import get_current_forward_batch + + forward_batch = get_current_forward_batch() + if forward_batch is None: + raise RuntimeError( + "forward_batch is required for SGLang GLM-5.2 native MLA attention" + ) + return forward_batch + + def _infer_total_tokens(self, forward_batch, tensor: torch.Tensor) -> int: + if hasattr(forward_batch, "input_ids") and forward_batch.input_ids is not None: + return int(forward_batch.input_ids.shape[0]) + if hasattr(forward_batch, "positions") and forward_batch.positions is not None: + return int(forward_batch.positions.shape[0]) + if hasattr(forward_batch, "seq_lens_sum"): + return int(forward_batch.seq_lens_sum) + return int(tensor.shape[0]) + + def _maybe_all_gather( + self, + tensor: torch.Tensor | None, + *, + total_tokens: int, + input_scattered: bool, + ): + if tensor is None or not input_scattered: + return tensor + from sglang.srt.distributed import get_tp_group + + output = tensor.new_empty((total_tokens, *tensor.shape[1:])) + get_tp_group().all_gather_into_tensor(output, tensor) + return output + + def forward( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + forward_batch = self._get_forward_batch() + + from sglang.srt.layers.communicator import get_attn_tp_context + + attn_tp_context = get_attn_tp_context() + with attn_tp_context.maybe_input_scattered(forward_batch): + total_tokens = self._infer_total_tokens(forward_batch, q_input) + q_input = self._maybe_all_gather( + q_input, + total_tokens=total_tokens, + input_scattered=attn_tp_context.input_scattered, + ) + kv_c_normed = self._maybe_all_gather( + kv_c_normed, + total_tokens=total_tokens, + input_scattered=attn_tp_context.input_scattered, + ) + k_pe = self._maybe_all_gather( + k_pe, + total_tokens=total_tokens, + input_scattered=attn_tp_context.input_scattered, + ) + positions = self._maybe_all_gather( + positions, + total_tokens=total_tokens, + input_scattered=attn_tp_context.input_scattered, + ) + q_scale = self._maybe_all_gather( + q_scale, + total_tokens=total_tokens, + input_scattered=attn_tp_context.input_scattered, + ) + + return self.forward_impl(q_input, kv_c_normed, k_pe, positions, q_scale) diff --git a/atom/plugin/sglang/prepare.py b/atom/plugin/sglang/prepare.py index fd813628c1..905f4dca6a 100644 --- a/atom/plugin/sglang/prepare.py +++ b/atom/plugin/sglang/prepare.py @@ -2,6 +2,7 @@ import inspect import logging +from contextlib import nullcontext from typing import Any from atom.plugin.prepare import _set_framework_backbone @@ -82,13 +83,19 @@ def prepare_model(config: Any): apply_graph_capture_patch() - init_params = inspect.signature(model_cls.__init__).parameters - if "atom_config" in init_params: - model = model_cls(atom_config=atom_config) - elif "config" in init_params: - model = model_cls(config=atom_config) - else: - model = model_cls(atom_config) + construct_context = ( + model_adapter.construction_context() + if model_adapter.construction_context is not None + else nullcontext() + ) + with construct_context: + init_params = inspect.signature(model_cls.__init__).parameters + if "atom_config" in init_params: + model = model_cls(atom_config=atom_config) + elif "config" in init_params: + model = model_cls(config=atom_config) + else: + model = model_cls(atom_config) if not hasattr(model, "atom_config"): model.atom_config = atom_config return model diff --git a/atom/plugin/sglang/runtime/__init__.py b/atom/plugin/sglang/runtime/__init__.py index 455a09fcf8..4a3f128524 100644 --- a/atom/plugin/sglang/runtime/__init__.py +++ b/atom/plugin/sglang/runtime/__init__.py @@ -9,16 +9,19 @@ ) from atom.plugin.sglang.runtime.forward_context import SGLangPluginRuntime from atom.plugin.sglang.runtime.model_arch import ( + GLM52_DSA_ARCH, MODEL_ADAPTER_SPECS, MODEL_ARCH_SPECS, SGLangModelAdapterSpec, get_model_arch_spec, + is_glm52_dsa_config, ) apply_load_config_patch() __all__ = [ "apply_load_config_patch", + "GLM52_DSA_ARCH", "MODEL_ADAPTER_SPECS", "MODEL_ARCH_SPECS", "SGLangForwardBatchMetadata", @@ -27,5 +30,6 @@ "bind_current_forward_batch", "get_current_forward_batch", "get_model_arch_spec", + "is_glm52_dsa_config", "plugin_runtime_scope", ] diff --git a/atom/plugin/sglang/runtime/forward_context.py b/atom/plugin/sglang/runtime/forward_context.py index 432a62b657..1f0fdfba3e 100644 --- a/atom/plugin/sglang/runtime/forward_context.py +++ b/atom/plugin/sglang/runtime/forward_context.py @@ -4,6 +4,7 @@ import copy import logging +import os from contextlib import ExitStack from dataclasses import dataclass, field from typing import Any, Optional @@ -214,6 +215,250 @@ def _slice_attr(name: str, n: int): return md +def _is_current_stream_capturing() -> bool: + try: + return bool(torch.cuda.is_current_stream_capturing()) + except Exception: + return False + + +def _get_sglang_attention_backend(): + try: + from sglang.srt.model_executor.forward_context import get_attn_backend + + return get_attn_backend() + except Exception: + return None + + +def _build_glm52_dsa_metadata( + atom_config: Any, + forward_batch: ForwardBatch, + positions: torch.Tensor, +): + hf_config = getattr(atom_config, "hf_config", None) + if _is_dummy_forward(forward_batch) or hf_config is None: + return None + + from atom.plugin.sglang.runtime.model_arch import is_glm52_dsa_config + + if not is_glm52_dsa_config(hf_config): + return None + + from atom.plugin.sglang.glm52_dsa_bridge import ( + build_atom_glm52_attention_metadata_from_sglang, + maybe_get_glm52_dsa_pools_from_sglang_backend, + ) + + attn_metadata = getattr(forward_batch, "atom_glm52_graph_metadata", None) + if attn_metadata is None: + backend = _get_sglang_attention_backend() + attn_metadata = getattr(backend, "atom_glm52_graph_metadata", None) + + is_capture_batch = _is_current_stream_capturing() + if attn_metadata is None and is_capture_batch: + from atom.plugin.sglang.attention_backend.glm52_dsa_backend import ( + ATOMGLM52DSABackendForSgl, + ) + + attn_metadata = ATOMGLM52DSABackendForSgl._last_atom_glm52_graph_metadata + + token_to_kv_pool, req_to_token_pool = maybe_get_glm52_dsa_pools_from_sglang_backend( + forward_batch + ) + if ( + attn_metadata is None + and token_to_kv_pool is not None + and req_to_token_pool is not None + ): + if is_capture_batch: + raise RuntimeError( + "ATOM GLM-5.2 CUDA graph metadata was not initialized before capture" + ) + attn_metadata = build_atom_glm52_attention_metadata_from_sglang( + forward_batch, + positions, + token_to_kv_pool=token_to_kv_pool, + req_to_token_pool=req_to_token_pool, + atom_config=atom_config, + ) + if ( + os.environ.get("ATOM_GLM52_VERIFY_DEBUG", "0") in ("1", "true", "True") + and getattr(forward_batch.forward_mode, "is_target_verify", lambda: False)() + ): + try: + spec_info = getattr(forward_batch, "spec_info", None) + tokens_per_req = int( + getattr(spec_info, "draft_token_num", 0) + or getattr(spec_info, "num_tokens_per_req", 0) + or 1 + ) + bs = int(getattr(forward_batch, "batch_size", 0) or 0) + probe_bs = min(2, bs) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + kv_indptr = getattr(attn_metadata, "kv_indptr", None) + kv_indices = getattr(attn_metadata, "kv_indices", None) + kv_indptr_cpu = ( + kv_indptr.detach().cpu().tolist() + if torch.is_tensor(kv_indptr) + else [] + ) + kv_ranges = [] + for row in range(min(probe_bs, max(0, len(kv_indptr_cpu) - 1))): + start = int(kv_indptr_cpu[row]) + end = int(kv_indptr_cpu[row + 1]) + if torch.is_tensor(kv_indices): + head = ( + kv_indices[start : min(end, start + 8)] + .detach() + .cpu() + .tolist() + ) + tail = ( + kv_indices[max(start, end - 8) : end] + .detach() + .cpu() + .tolist() + ) + else: + head = tail = None + kv_ranges.append( + { + "row": row, + "start": start, + "end": end, + "len": end - start, + "head": head, + "tail": tail, + } + ) + row_probe = { + "positions": ( + positions.reshape(bs, tokens_per_req)[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor(positions) + and bs > 0 + and int(positions.numel()) >= bs * tokens_per_req + else None + ), + "out_cache_loc": ( + out_cache_loc.reshape(bs, tokens_per_req)[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor(out_cache_loc) + and bs > 0 + and int(out_cache_loc.numel()) >= bs * tokens_per_req + else None + ), + "req_pool_indices": ( + forward_batch.req_pool_indices[:probe_bs] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(forward_batch, "req_pool_indices", None) + ) + else None + ), + "kv_ranges": kv_ranges, + "sparse_kv_indptr": ( + attn_metadata.sparse_kv_indptr[ + : min( + int(attn_metadata.sparse_kv_indptr.numel()), + probe_bs * tokens_per_req + 1, + ) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(attn_metadata, "sparse_kv_indptr", None) + ) + else None + ), + "token_to_seq": ( + attn_metadata.token_to_seq_idxs[ + : min( + int(attn_metadata.token_to_seq_idxs.numel()), + probe_bs * tokens_per_req, + ) + ] + .detach() + .cpu() + .tolist() + if torch.is_tensor( + getattr(attn_metadata, "token_to_seq_idxs", None) + ) + else None + ), + } + counter = int( + getattr(spec_info, "_atom_glm52_verify_counter", 0) or 0 + ) + 1 + setattr(spec_info, "_atom_glm52_verify_counter", counter) + setattr(spec_info, "_atom_glm52_row_probe", row_probe) + logger.info("GLM52 eager metadata row_probe=%s", row_probe) + except Exception: + logger.exception("Failed to log GLM52 eager metadata row probe") + return attn_metadata + + +def _build_deepseek_v4_metadata(forward_batch: ForwardBatch, positions: torch.Tensor): + backend = None + attn_metadata = getattr(forward_batch, "atom_v4_graph_metadata", None) + from atom.plugin.sglang.deepseek_v4_bridge import ( + build_atom_v4_attention_metadata_from_sglang, + maybe_get_proxy_pool_from_sglang_backend, + ) + + if attn_metadata is None: + backend = _get_sglang_attention_backend() + attn_metadata = getattr(backend, "atom_v4_graph_metadata", None) + + if attn_metadata is None: + backend = getattr(forward_batch, "attn_backend", None) + attn_metadata = getattr(backend, "atom_v4_graph_metadata", None) + + if attn_metadata is None and backend is not None: + backend_forward_batch = getattr(backend, "forward_metadata", None) + attn_metadata = getattr(backend_forward_batch, "atom_v4_graph_metadata", None) + + proxy_pool, req_to_token_pool = maybe_get_proxy_pool_from_sglang_backend() + + is_capture_batch = _is_current_stream_capturing() + if attn_metadata is None and is_capture_batch: + try: + from atom.plugin.sglang.attention_backend.deepseek_v4_backend import ( + ATOMDeepseekV4BackendForSgl, + ) + + attn_metadata = ATOMDeepseekV4BackendForSgl._last_atom_v4_graph_metadata + if attn_metadata is not None: + attn_metadata = _slice_v4_graph_metadata_for_capture( + attn_metadata, + num_tokens=int(positions.shape[0]), + bs=int(forward_batch.batch_size), + ) + except Exception: + attn_metadata = None + + if attn_metadata is None and getattr(proxy_pool, "is_atom_v4_proxy_pool", False): + if is_capture_batch: + raise RuntimeError( + "ATOM DeepSeek-V4 CUDA graph metadata was not initialized before capture" + ) + attn_metadata = build_atom_v4_attention_metadata_from_sglang( + forward_batch, + positions, + proxy_pool=proxy_pool, + req_to_token_pool=req_to_token_pool, + ) + return attn_metadata + + def _set_atom_forward_context( atom_config: Any, forward_batch: ForwardBatch, @@ -232,70 +477,24 @@ def _set_atom_forward_context( max_seqlen_q = 1 if forward_mode.is_decode_or_idle() else 0 attn_metadata = None try: - attn_metadata = getattr(forward_batch, "atom_v4_graph_metadata", None) - from atom.plugin.sglang.deepseek_v4_bridge import ( - build_atom_v4_attention_metadata_from_sglang, - maybe_get_proxy_pool_from_sglang_backend, + attn_metadata = _build_glm52_dsa_metadata( + atom_config, + forward_batch, + positions, ) - - if attn_metadata is None: - try: - from sglang.srt.model_executor.forward_context import get_attn_backend - - backend = get_attn_backend() - attn_metadata = getattr(backend, "atom_v4_graph_metadata", None) - except Exception: - attn_metadata = None - - if attn_metadata is None: - backend = getattr(forward_batch, "attn_backend", None) - attn_metadata = getattr(backend, "atom_v4_graph_metadata", None) - - if attn_metadata is None and backend is not None: - backend_forward_batch = getattr(backend, "forward_metadata", None) - attn_metadata = getattr( - backend_forward_batch, "atom_v4_graph_metadata", None - ) - - proxy_pool, req_to_token_pool = maybe_get_proxy_pool_from_sglang_backend() - try: - is_capture_batch = bool(torch.cuda.is_current_stream_capturing()) - except Exception: - is_capture_batch = False - if attn_metadata is None and is_capture_batch: - try: - from atom.plugin.sglang.attention_backend.deepseek_v4_backend import ( - ATOMDeepseekV4BackendForSgl, - ) - - attn_metadata = ATOMDeepseekV4BackendForSgl._last_atom_v4_graph_metadata - if attn_metadata is not None: - attn_metadata = _slice_v4_graph_metadata_for_capture( - attn_metadata, - num_tokens=int(positions.shape[0]), - bs=int(forward_batch.batch_size), - ) - except Exception: - attn_metadata = None - if attn_metadata is None and getattr( - proxy_pool, "is_atom_v4_proxy_pool", False - ): - if is_capture_batch: - raise RuntimeError( - "ATOM DeepSeek-V4 CUDA graph metadata was not initialized before capture" - ) - else: - attn_metadata = build_atom_v4_attention_metadata_from_sglang( - forward_batch, - positions, - proxy_pool=proxy_pool, - req_to_token_pool=req_to_token_pool, - ) except Exception as exc: raise RuntimeError( - "Failed to build ATOM DeepSeek-V4 metadata for SGLang" + "Failed to build ATOM GLM-5.2 DSA metadata for SGLang" ) from exc + if attn_metadata is None: + try: + attn_metadata = _build_deepseek_v4_metadata(forward_batch, positions) + except Exception as exc: + raise RuntimeError( + "Failed to build ATOM DeepSeek-V4 metadata for SGLang" + ) from exc + if attn_metadata is None: attn_metadata = AttentionMetaData(max_seqlen_q=max_seqlen_q) batch_size = int(forward_batch.batch_size) diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py index f3620a57d5..bd0e39a429 100644 --- a/atom/plugin/sglang/runtime/model_arch.py +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -3,8 +3,22 @@ from __future__ import annotations from dataclasses import dataclass +from contextlib import AbstractContextManager from typing import Any, Callable, Optional +GLM52_DSA_ARCH = "GlmMoeDsaForCausalLM" +GLM52_DSA_MODEL_TYPE = "glm_moe_dsa" + + +def is_glm52_dsa_config(config: Any) -> bool: + """Return whether an HF config describes GLM-5.2 DSA.""" + + archs = getattr(config, "architectures", None) or [] + return ( + any(GLM52_DSA_ARCH in str(arch) for arch in archs) + or getattr(config, "model_type", None) == GLM52_DSA_MODEL_TYPE + ) + @dataclass(frozen=True) class SGLangModelAdapterSpec: @@ -18,7 +32,9 @@ class SGLangModelAdapterSpec: wrapper_binds_gdn_context: bool = False uses_context_only_forward: bool = False prepare_config: Optional[Callable[[Any, str], None]] = None + construction_context: Optional[Callable[[], AbstractContextManager[Any]]] = None install_adapters: Optional[Callable[[Any], None]] = None + bind_cache_views: Optional[Callable[[Any, Any], None]] = None def _prepare_qwen35_config(atom_config: Any, model_arch: str) -> None: @@ -70,12 +86,53 @@ def _prepare_minimax_m3_config(atom_config: Any, model_arch: str) -> None: ) +def _prepare_glm52_dsa_config(atom_config: Any, model_arch: str) -> None: + from atom.models.deepseek_v2 import GlmMoeDsaForCausalLM + + quant_config = getattr(atom_config, "quant_config", None) + if quant_config is not None: + quant_config.remap_layer_name( + atom_config.hf_config, + packed_modules_mapping=getattr( + GlmMoeDsaForCausalLM, "packed_modules_mapping", {} + ), + weights_mapper=getattr(GlmMoeDsaForCausalLM, "hf_to_atom_mapper", {}), + quant_exclude_name_mapping=getattr( + GlmMoeDsaForCausalLM, "quant_exclude_name_mapping", {} + ), + ) + default_excludes = getattr( + GlmMoeDsaForCausalLM, "quant_default_exclude_layers", [] + ) + if default_excludes: + quant_config.apply_default_exclude_layers(default_excludes) + + # SGLang's DSA pool uses page64/preshuffle for GLM/DeepSeek-family DSA. + # Keep ATOM's config aligned for the native GLM indexer, while + # ATOM_MLA_PAGE_SIZE can remain 1 so sparse MLA reads selected physical ids. + atom_config.kv_cache_block_size = 64 + + def _install_deepseek_mla_adapters(model: Any) -> None: from atom.plugin.sglang.models.deepseek_mla import setup_deepseek_for_sglang setup_deepseek_for_sglang(model) +def _glm52_dsa_construction_context(): + from atom.plugin.sglang.models.glm52_dsa_attention import ( + glm52_native_mla_attention_construction, + ) + + return glm52_native_mla_attention_construction() + + +def _install_glm52_dsa_native_adapters(model: Any) -> None: + from atom.plugin.sglang.models.glm52_dsa import setup_glm52_dsa_for_sglang + + setup_glm52_dsa_for_sglang(model) + + def _install_deepseek_v4_adapters(model: Any) -> None: # DeepSeek-V4 in SGLang plugin mode follows the proxy-KV bridge path: # SGLang owns scheduling/allocation, while ATOM owns the model, cache views, @@ -91,6 +148,40 @@ def _install_deepseek_v4_adapters(model: Any) -> None: patch_deepseek_v4_attention_for_sglang(module) +def _bind_deepseek_v4_cache_views(model: Any, runtime: Any) -> None: + del runtime + from atom.plugin.sglang.deepseek_v4_bridge import ( + bind_deepseek_v4_proxy_cache_views, + maybe_get_proxy_pool_from_sglang_backend, + reset_deepseek_v4_state_slots, + ) + + proxy_pool, _ = maybe_get_proxy_pool_from_sglang_backend() + if not bind_deepseek_v4_proxy_cache_views(model, proxy_pool): + raise RuntimeError("DeepSeek-V4 SGLang proxy KV pool is not initialized") + + from atom.utils.forward_context import get_forward_context + + reset_slots = getattr(get_forward_context().attn_metadata, "reset_slots", None) + reset_deepseek_v4_state_slots(model, reset_slots) + + +def _bind_glm52_dsa_cache_views(model: Any, runtime: Any) -> None: + if getattr(runtime.forward_batch.forward_mode, "is_idle", lambda: False)(): + return + + from atom.plugin.sglang.glm52_dsa_bridge import ( + bind_glm52_dsa_cache_views, + maybe_get_glm52_dsa_pools_from_sglang_backend, + ) + + token_to_kv_pool, _ = maybe_get_glm52_dsa_pools_from_sglang_backend( + runtime.forward_batch + ) + if not bind_glm52_dsa_cache_views(model, token_to_kv_pool): + raise RuntimeError("GLM-5.2 SGLang DSA KV pool is not initialized") + + def _install_minimax_m3_adapters(model: Any) -> None: from atom.plugin.sglang.models.minimax_m3 import setup_minimax_m3_for_sglang @@ -106,8 +197,11 @@ def _install_minimax_m3_adapters(model: Any) -> None: install_adapters=_install_deepseek_mla_adapters, uses_context_only_forward=True, ), - "GlmMoeDsaForCausalLM": SGLangModelAdapterSpec( - install_adapters=_install_deepseek_mla_adapters, + GLM52_DSA_ARCH: SGLangModelAdapterSpec( + prepare_config=_prepare_glm52_dsa_config, + construction_context=_glm52_dsa_construction_context, + install_adapters=_install_glm52_dsa_native_adapters, + bind_cache_views=_bind_glm52_dsa_cache_views, uses_context_only_forward=True, ), "KimiK25ForConditionalGeneration": SGLangModelAdapterSpec( @@ -131,6 +225,7 @@ def _install_minimax_m3_adapters(model: Any) -> None: ), "DeepseekV4ForCausalLM": SGLangModelAdapterSpec( install_adapters=_install_deepseek_v4_adapters, + bind_cache_views=_bind_deepseek_v4_cache_views, ), "MiniMaxM3SparseForCausalLM": SGLangModelAdapterSpec( uses_context_only_forward=True, @@ -152,7 +247,7 @@ def _install_minimax_m3_adapters(model: Any) -> None: for key in ( "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", - "GlmMoeDsaForCausalLM", + GLM52_DSA_ARCH, "Qwen3ForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM",