[ATOM SGLang] GLM 5.2 MTP by ZhiweiYan-96 · Pull Request #1694 · ROCm/ATOM

ZhiweiYan-96 · 2026-07-24T11:19:09Z

Motivation

Enable GLM 5.2 MTP at sglang plugin

Technical Details

Test Plan

export AITER_QUICK_REDUCE_QUANTIZATION=INT4
export AITER_USE_FLYDSL_MOE_SORTING=1
export SGLANG_USE_AITER=1
export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
export PYTHONPATH="${ROOT}:/app/sglang"
export CUDA_VISIBLE_DEVICES=4,5,6,7
export TORCHINDUCTOR_COMPILE_THREADS=128

export ATOM_SPEC_DECODE_FAST_LOADER=0
export ATOM_LOADER_NUM_THREADS=1
export ATOM_GLM52_DRAFT_ATTN_MODE=full
export ATOM_SGLANG_DRAFT_ARGMAX=1
export ATOM_REPLICATE_VOCAB_EMBED=1
export ATOM_GLM52_MTP_FUSED_INPUT=1
export ATOM_GLM52_MTP_WRAPPER_MODE=current

python3 -m sglang.launch_server \
  --model-path /workspace/shared/data/amd_int/models/GLM-5.2-MXFP4 \
  --host 0.0.0.0 \
  --port 8016 \
  --trust-remote-code \
  --tp-size 4 \
  --mem-fraction-static 0.8 \
  --disable-radix-cache \
  --max-running-requests 128 \
  --enable-return-hidden-states \
  --kv-cache-dtype fp8_e4m3 \
  --model-loader-extra-config '{"online_quant_config":{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate","*expert*"]}}' \
  --speculative-algorithm EAGLE \
  --speculative-draft-attention-backend aiter \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  2>&1 | tee log.glm52_mtp_server.log

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Add the plugin-owned GLM-5.2 MTP metadata bridge and fixed-shape CUDA graph support for draft decode, target verify, and draft extend. Keep SGLang request/KV ownership and route the draft model through the general ATOM AITER attention backend.

Keep the PR branch limited to the general ATOM draft attention route by removing hybrid native-attention cache and layer-id compatibility.

Keep the GLM-5.2 PR route fixed on the validated plugin-owned draft-extend implementation and remove the upstream-handler environment switch.

Remove GLM-5.2 phase and capacity environment switches. Let SGLang's native disable_cuda_graph server argument control graph capture, with all speculative graph phases enabled by default and an internal fixed metadata capacity.

github-actions · 2026-07-24T11:20:09Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1694 --add-label <label>

github-actions

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 520 in 65173c1

work_metadata: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 521 in 65173c1

work_indptr: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 522 in 65173c1

work_info_set: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 523 in 65173c1

reduce_indptr: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 524 in 65173c1

reduce_final_map: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 525 in 65173c1

reduce_partial_map: Optional[torch.Tensor]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 777 in 65173c1

) -> Optional[SparseMLAGraphBuffers]:

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 1067 in 65173c1

input_dtype: Optional[torch.dtype] = None,

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Line 1068 in 65173c1

q_scale: Optional[torch.Tensor] = None,

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_dsa_bridge.py

Lines 17 to 49 in 65173c1

    
           from atom.plugin.sglang.glm52_mtp.cache_bind import bind_glm52_dsa_cache_views 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               ATTENTION_PAGE_SIZE_ATTR, 
        
               DECODE_GRAPH_BUFFERS_ATTR, 
        
               DRAFT_SUB_STEP_ATTR, 
        
               EMPTY_VALUE_CACHE_ATTR, 
        
               INDEXER_PAGE_SIZE_ATTR, 
        
               SHARED_SPARSE_INDICES_ATTR, 
        
               is_glm52_dsa_arch, 
        
               maybe_get_glm52_dsa_pools_from_sglang_backend, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.decode_graph import ( 
        
               build_atom_glm52_decode_graph_metadata_from_sglang, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.dispatcher import ( 
        
               build_atom_glm52_attention_metadata_from_sglang, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.draft_decode import ( 
        
               build_decode_metadata as _build_decode_metadata, 
        
               clear_draft_decode_sub_step, 
        
               get_draft_decode_sub_step, 
        
               set_draft_decode_sub_step, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.draft_extend import ( 
        
               build_mtp_draft_extend_decode_metadata as _build_mtp_draft_extend_decode_metadata, 
        
               draft_extend_token_num, 
        
               draft_extend_token_num as _draft_extend_token_num, 
        
               should_use_mtp_draft_extend_decode_path as _should_use_mtp_draft_extend_decode_path, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.target_verify import ( 
        
               build_mtp_verify_decode_metadata as _build_mtp_verify_decode_metadata, 
        
               should_use_mtp_verify_prefill_path as _should_use_mtp_verify_prefill_path, 
        
           )

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/__init__.py

Lines 23 to 38 in 65173c1

    
           from atom.plugin.sglang.glm52_mtp.cache_bind import bind_glm52_dsa_cache_views 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               maybe_get_glm52_dsa_pools_from_sglang_backend, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.dispatcher import ( 
        
               build_atom_glm52_attention_metadata_from_sglang, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.decode_graph import ( 
        
               build_atom_glm52_decode_graph_metadata_from_sglang, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.draft_decode import ( 
        
               clear_draft_decode_sub_step, 
        
               get_draft_decode_sub_step, 
        
               set_draft_decode_sub_step, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.draft_extend import draft_extend_token_num

⚠️ [ruff] <B009> _{reported by reviewdog 🐶}
Do not call getattr with a constant attribute value. It is not any safer than normal property access.

ATOM/atom/plugin/sglang/glm52_mtp/cache_bind.py

Line 57 in 65173c1

index_entry_dim = int(getattr(indexer, "head_dim")) + 4

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/common.py

Line 35 in 65173c1

from atom.plugin.sglang.glm52_mtp.draft_decode import is_draft_decode_metadata

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

ATOM/atom/plugin/sglang/glm52_mtp/common.py

Line 40 in 65173c1

except Exception:

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

ATOM/atom/plugin/sglang/glm52_mtp/common.py

Line 50 in 65173c1

except Exception:

⚠️ [ruff] <PLR1730> _{reported by reviewdog 🐶}
Replace if statement with min call

ATOM/atom/plugin/sglang/glm52_mtp/common.py

Lines 176 to 177 in 65173c1

    
           if extend_end > int(req_to_token_row.numel()): 
        
               extend_end = int(req_to_token_row.numel())

⚠️ [ruff] <B009> _{reported by reviewdog 🐶}
Do not call getattr with a constant attribute value. It is not any safer than normal property access.

ATOM/atom/plugin/sglang/glm52_mtp/common.py

Line 409 in 65173c1

num_heads = int(getattr(hf_config, "num_attention_heads"))

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/decode_graph.py

Lines 3 to 18 in 65173c1

    
           from __future__ import annotations 
        
           import torch 
        
           from aiter import get_mla_metadata_v1 
        
           from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               DECODE_GRAPH_BUFFERS_ATTR, 
        
               SHARED_SPARSE_INDICES_ATTR, 
        
               get_index_topk, 
        
               local_num_attention_heads, 
        
               make_mla_work_buffers, 
        
               metadata_dtype, 
        
               validate_page_size, 
        
               attention_page_size, 
        
           )

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
atom.plugin.sglang.glm52_mtp.common.compute_mtp_sparse_per_token_kv_lens imported but unused

ATOM/atom/plugin/sglang/glm52_mtp/draft_decode.py

Line 17 in 65173c1

compute_mtp_sparse_per_token_kv_lens,

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/draft_extend.py

Lines 3 to 16 in 65173c1

    
           from __future__ import annotations 
        
           import os 
        
           import numpy as np 
        
           import torch 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               get_extend_lens_cpu, 
        
               get_extend_prefix_lens_cpu, 
        
               get_seq_lens_cpu, 
        
               is_draft_extend_mode, 
        
           ) 
        
           from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata

⚠️ [ruff] <SIM103> _{reported by reviewdog 🐶}
Return the condition directly

ATOM/atom/plugin/sglang/glm52_mtp/draft_extend.py

Lines 83 to 85 in 65173c1

    
           if int(getattr(forward_batch, "_graph_cache_bs", 0) or 0) > 0: 
        
               return True 
        
           return False

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/multi_token.py

Lines 3 to 25 in 65173c1

    
           from __future__ import annotations 
        
           import logging 
        
           import os 
        
           import numpy as np 
        
           import torch 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               attention_page_size, 
        
               build_token_table, 
        
               compute_mtp_sparse_per_token_kv_lens, 
        
               counts_to_indptr, 
        
               ensure_shared_sparse_buffer, 
        
               flatten_kv_indices, 
        
               get_index_topk, 
        
               is_draft_extend_mode, 
        
               local_num_attention_heads, 
        
               make_mla_work_buffers, 
        
               make_sparse_mtp_work_buffers, 
        
               metadata_dtype, 
        
               validate_page_size, 
        
           )

⚠️ [ruff] <SIM103> _{reported by reviewdog 🐶}
Return the negated condition directly

ATOM/atom/plugin/sglang/glm52_mtp/prefill.py

Lines 46 to 48 in 65173c1

    
           if num_tokens_per_req is not None and int(num_tokens_per_req) != 1: 
        
               return False 
        
           return True

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/target_verify.py

Lines 3 to 10 in 65173c1

    
           from __future__ import annotations 
        
           import os 
        
           import numpy as np 
        
           import torch 
        
           from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable bs is assigned to but never used

ATOM/atom/plugin/sglang/glm52_mtp/target_verify.py

Line 60 in 65173c1

bs = int(forward_batch.batch_size)

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
contextlib imported but unused

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 9 in 65173c1

import contextlib

⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
Import from collections.abc instead: Iterable

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 13 in 65173c1

from typing import Iterable, Optional, Tuple

⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
typing.Tuple is deprecated, use tuple instead

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 13 in 65173c1

from typing import Iterable, Optional, Tuple

⚠️ [ruff] <RET501> _{reported by reviewdog 🐶}
Do not explicitly return None in function if it is the only possible return value

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 241 in 65173c1

return None

⚠️ [ruff] <RET501> _{reported by reviewdog 🐶}
Do not explicitly return None in function if it is the only possible return value

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 255 in 65173c1

return None

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

ATOM/atom/plugin/sglang/runtime/forward_context.py

Line 319 in 65173c1

except Exception:

⚠️ [ruff] <B010> _{reported by reviewdog 🐶}
Do not call setattr with a constant attribute value. It is not any safer than normal property access.

ATOM/atom/plugin/sglang/runtime/forward_context.py

Lines 324 to 328 in 65173c1

    
           setattr( 
        
               token_to_kv_pool, 
        
               "_atom_glm52_draft_decode_graph_metadata", 
        
               graph_cache, 
        
           )

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/runtime/model_arch.py

Lines 3 to 9 in 65173c1

    
           from __future__ import annotations 
        
           import logging 
        
           import os 
        
           from dataclasses import dataclass 
        
           from contextlib import AbstractContextManager, nullcontext 
        
           from typing import Any, Callable, Optional

⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
Import from collections.abc instead: Callable

ATOM/atom/plugin/sglang/runtime/model_arch.py

Line 9 in 65173c1

from typing import Any, Callable, Optional

github-actions · 2026-07-24T11:20:14Z

 import time
 from glob import glob
-from typing import Generator, Tuple
+from typing import Callable, Generator, Tuple


⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
Import from collections.abc instead: Callable, Generator

Suggested change

from typing import Callable, Generator, Tuple

from typing import Tuple

from collections.abc import Callable, Generator

github-actions · 2026-07-24T11:20:14Z

 import time
 from glob import glob
-from typing import Generator, Tuple
+from typing import Callable, Generator, Tuple


⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
typing.Tuple is deprecated, use tuple instead

github-actions · 2026-07-24T11:20:14Z

+    *,
+    shard_files: list[str] | None = None,
+    key_filter: Callable[[str], bool] | None = None,
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:


⚠️ [ruff] <UP006> _{reported by reviewdog 🐶}
Use tuple instead of Tuple for type annotation

Suggested change

) -> Generator[Tuple[str, torch.Tensor], None, None]:

) -> Generator[tuple[str, torch.Tensor], None, None]:

github-actions · 2026-07-24T11:20:14Z

                    yield name, param
        else:
            with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
                for name in f.keys():


⚠️ [ruff] <SIM118> _{reported by reviewdog 🐶}
Use key in dict instead of key in dict.keys()

Suggested change

for name in f.keys():

for name in f:

github-actions · 2026-07-24T11:20:14Z

+                if remapped is None:
+                    return False
+                name = remapped
+            for mapping_part in weights_mapping.keys():


⚠️ [ruff] <SIM118> _{reported by reviewdog 🐶}
Use key in dict instead of key in dict.keys()

Suggested change

for mapping_part in weights_mapping.keys():

for mapping_part in weights_mapping:

github-actions · 2026-07-24T11:20:15Z

+    last_page_len: torch.Tensor
+    use_fast_metadata: bool
+    work_metadata: Optional[torch.Tensor] = None
+    work_indptr: Optional[torch.Tensor] = None


⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

work_indptr: Optional[torch.Tensor] = None

work_indptr: torch.Tensor | None = None

github-actions · 2026-07-24T11:20:15Z

+    use_fast_metadata: bool
+    work_metadata: Optional[torch.Tensor] = None
+    work_indptr: Optional[torch.Tensor] = None
+    work_info_set: Optional[torch.Tensor] = None


⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

work_info_set: Optional[torch.Tensor] = None

work_info_set: torch.Tensor | None = None

github-actions · 2026-07-24T11:20:15Z

+    work_metadata: Optional[torch.Tensor] = None
+    work_indptr: Optional[torch.Tensor] = None
+    work_info_set: Optional[torch.Tensor] = None
+    reduce_indptr: Optional[torch.Tensor] = None


⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

reduce_indptr: Optional[torch.Tensor] = None

reduce_indptr: torch.Tensor | None = None

github-actions · 2026-07-24T11:20:15Z

+    work_indptr: Optional[torch.Tensor] = None
+    work_info_set: Optional[torch.Tensor] = None
+    reduce_indptr: Optional[torch.Tensor] = None
+    reduce_final_map: Optional[torch.Tensor] = None


⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

reduce_final_map: Optional[torch.Tensor] = None

reduce_final_map: torch.Tensor | None = None

github-actions · 2026-07-24T11:20:15Z

+    work_info_set: Optional[torch.Tensor] = None
+    reduce_indptr: Optional[torch.Tensor] = None
+    reduce_final_map: Optional[torch.Tensor] = None
+    reduce_partial_map: Optional[torch.Tensor] = None


⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

reduce_partial_map: Optional[torch.Tensor] = None

reduce_partial_map: torch.Tensor | None = None

github-actions · 2026-07-24T11:20:20Z

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/model_loader/loader.py

Lines 4 to 22 in 65173c1

import concurrent.futures

import json

import os

import logging

import re

import threading

import time

from glob import glob

from typing import Callable, Generator, Tuple

from collections.abc import Iterable, Mapping

from dataclasses import dataclass, field

from typing import Any, Optional

import safetensors

import safetensors.torch

import torch

from torch import nn

from tqdm import tqdm

from transformers import AutoConfig

github-actions · 2026-07-24T11:20:21Z

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/attention_backend/sparse_mla_indexer.py

Lines 6 to 33 in 65173c1

from __future__ import annotations

from dataclasses import dataclass

import logging

import os

import re

from typing import Optional

import numpy as np

import torch

from aiter import (

cp_gather_indexer_k_quant_cache,

dtypes,

get_mla_metadata_info_v1,

get_mla_metadata_v1,

indexer_k_quant_and_cache,

indexer_qk_rope_quant_and_cache,

top_k_per_row_decode,

top_k_per_row_prefill,

)

from aiter.mla import mla_decode_fwd

from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits

from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits

import triton

import triton.language as tl

from atom.utils.custom_register import direct_register_custom_op

from atom.plugin.sglang.glm52_mtp.common import GLM52_GRAPH_SEQ_LEN_CAPACITY

github-actions · 2026-07-24T11:20:22Z

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Lines 8 to 33 in 65173c1

import copy

import contextlib

import logging

import os

import re

from typing import Iterable, Optional, Tuple

import torch

from torch import nn

from sglang.srt.distributed import get_pp_group

from sglang.srt.layers.logits_processor import LogitsProcessor

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

from atom.config import QuantizationConfig as AtomQuantizationConfig

from atom.config import SpeculativeConfig

from atom.plugin.config import generate_atom_config_for_plugin_mode

from atom.plugin.sglang.models.deepseek_mla import (

setup_deepseek_for_sglang,

)

from atom.plugin.sglang.runtime import (

SGLangPluginRuntime,

plugin_runtime_scope,

)

github-actions · 2026-07-24T11:20:23Z

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Lines 205 to 210 in 65173c1

from atom.plugin.register import (

init_aiter_dist,

register_ops_to_sglang,

set_attn_cls,

)

from atom.models.deepseek_mtp import DeepSeekMTP

Keep target attention on the native ATOM GLM-5.2 backend and draft attention on the general AITER backend without a runtime mode switch. Remove debug metadata hooks, allocation pointer logs, and route-only logging wrappers.

github-actions

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/multi_token.py

Lines 3 to 22 in 7bf17aa

    
           from __future__ import annotations 
        
           import numpy as np 
        
           import torch 
        
           from atom.plugin.sglang.glm52_mtp.common import ( 
        
               attention_page_size, 
        
               build_token_table, 
        
               compute_mtp_sparse_per_token_kv_lens, 
        
               counts_to_indptr, 
        
               ensure_shared_sparse_buffer, 
        
               flatten_kv_indices, 
        
               get_index_topk, 
        
               is_draft_extend_mode, 
        
               local_num_attention_heads, 
        
               make_mla_work_buffers, 
        
               make_sparse_mtp_work_buffers, 
        
               metadata_dtype, 
        
               validate_page_size, 
        
           )

⚠️ [ruff] <SIM103> _{reported by reviewdog 🐶}
Return the negated condition directly

ATOM/atom/plugin/sglang/glm52_mtp/prefill.py

Lines 46 to 48 in 7bf17aa

    
           if num_tokens_per_req is not None and int(num_tokens_per_req) != 1: 
        
               return False 
        
           return True

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/sglang/glm52_mtp/target_verify.py

Lines 3 to 10 in 7bf17aa

    
           from __future__ import annotations 
        
           import os 
        
           import numpy as np 
        
           import torch 
        
           from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable bs is assigned to but never used

ATOM/atom/plugin/sglang/glm52_mtp/target_verify.py

Line 60 in 7bf17aa

bs = int(forward_batch.batch_size)

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
contextlib imported but unused

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 9 in 7bf17aa

import contextlib

⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
Import from collections.abc instead: Iterable

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 13 in 7bf17aa

from typing import Iterable, Optional, Tuple

⚠️ [ruff] <UP035> _{reported by reviewdog 🐶}
typing.Tuple is deprecated, use tuple instead

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 13 in 7bf17aa

from typing import Iterable, Optional, Tuple

⚠️ [ruff] <RET501> _{reported by reviewdog 🐶}
Do not explicitly return None in function if it is the only possible return value

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 235 in 7bf17aa

return None

⚠️ [ruff] <RET501> _{reported by reviewdog 🐶}
Do not explicitly return None in function if it is the only possible return value

ATOM/atom/plugin/sglang/models/deepseek_nextn_wrapper.py

Line 249 in 7bf17aa

return None

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

ATOM/atom/plugin/sglang/runtime/forward_context.py

Line 319 in 7bf17aa

except Exception:

⚠️ [ruff] <B010> _{reported by reviewdog 🐶}
Do not call setattr with a constant attribute value. It is not any safer than normal property access.

ATOM/atom/plugin/sglang/runtime/forward_context.py

Lines 324 to 328 in 7bf17aa

    
           setattr( 
        
               token_to_kv_pool, 
        
               "_atom_glm52_draft_decode_graph_metadata", 
        
               graph_cache, 
        
           )

github-actions · 2026-07-25T07:44:36Z

+        from atom.plugin.sglang.glm52_mtp.common import (
+            GLM52_GRAPH_SEQ_LEN_CAPACITY,
+        )
    except Exception as exc:


⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

github-actions · 2026-07-25T07:44:36Z

@@ -204,6 +213,43 @@ def _is_dsv4_nextn_runner(runner) -> bool:
        except Exception:


⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

github-actions · 2026-07-25T07:44:36Z

+                or "DeepseekV3ForCausalLMNextN" in str(arch)
+                for arch in arches
+            )
+        except Exception:


⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

github-actions · 2026-07-25T07:44:36Z

+            return model_type == "glm_moe_dsa" or any(
+                "GlmMoeDsaForCausalLM" in str(arch) for arch in arches
+            )
+        except Exception:


⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

github-actions · 2026-07-25T07:44:36Z

            except Exception:
                pass


⚠️ [ruff] <S110> _{reported by reviewdog 🐶}
try-except-pass detected, consider logging the exception

github-actions · 2026-07-25T07:44:37Z

+
+def local_num_attention_heads(atom_config) -> int:
+    hf_config = atom_config.hf_config
+    num_heads = int(getattr(hf_config, "num_attention_heads"))


⚠️ [ruff] <B009> _{reported by reviewdog 🐶}
Do not call getattr with a constant attribute value. It is not any safer than normal property access.

Suggested change

num_heads = int(getattr(hf_config, "num_attention_heads"))

num_heads = int(hf_config.num_attention_heads)

github-actions · 2026-07-25T07:44:37Z

+from __future__ import annotations
+
+import torch
+from aiter import get_mla_metadata_v1
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+
+from atom.plugin.sglang.glm52_mtp.common import (
+    DECODE_GRAPH_BUFFERS_ATTR,
+    SHARED_SPARSE_INDICES_ATTR,
+    get_index_topk,
+    local_num_attention_heads,
+    make_mla_work_buffers,
+    metadata_dtype,
+    validate_page_size,
+    attention_page_size,
+)
+
+
+class GLM52DecodeGraphBuffers:


⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

Suggested change

from __future__ import annotations

import torch

from aiter import get_mla_metadata_v1

from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton

from atom.plugin.sglang.glm52_mtp.common import (

DECODE_GRAPH_BUFFERS_ATTR,

SHARED_SPARSE_INDICES_ATTR,

get_index_topk,

local_num_attention_heads,

make_mla_work_buffers,

metadata_dtype,

validate_page_size,

attention_page_size,

)

class GLM52DecodeGraphBuffers:

from __future__ import annotations

import torch

from aiter import get_mla_metadata_v1

from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton

from atom.plugin.sglang.glm52_mtp.common import (

DECODE_GRAPH_BUFFERS_ATTR,

SHARED_SPARSE_INDICES_ATTR,

attention_page_size,

get_index_topk,

local_num_attention_heads,

make_mla_work_buffers,

metadata_dtype,

validate_page_size,

)

class GLM52DecodeGraphBuffers:

github-actions · 2026-07-25T07:44:37Z

+from atom.plugin.sglang.glm52_mtp.common import (
+    DRAFT_SUB_STEP_ATTR,
+    attention_page_size,
+    build_draft_decode_token_table,
+    build_token_table,
+    compute_mtp_sparse_per_token_kv_lens,
+    counts_to_indptr,
+    ensure_shared_sparse_buffer,
+    flatten_kv_indices,
+    get_index_topk,
+    get_seq_lens_cpu,
+    local_num_attention_heads,
+    make_mla_work_buffers,
+    make_sparse_mtp_work_buffers,
+    metadata_dtype,
+    resolve_draft_decode_context_lens,
+    resolve_speculative_num_steps,
+    validate_page_size,
+)


⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
atom.plugin.sglang.glm52_mtp.common.compute_mtp_sparse_per_token_kv_lens imported but unused

Suggested change

from atom.plugin.sglang.glm52_mtp.common import (

DRAFT_SUB_STEP_ATTR,

attention_page_size,

build_draft_decode_token_table,

build_token_table,

compute_mtp_sparse_per_token_kv_lens,

counts_to_indptr,

ensure_shared_sparse_buffer,

flatten_kv_indices,

get_index_topk,

get_seq_lens_cpu,

local_num_attention_heads,

make_mla_work_buffers,

make_sparse_mtp_work_buffers,

metadata_dtype,

resolve_draft_decode_context_lens,

resolve_speculative_num_steps,

validate_page_size,

)

from atom.plugin.sglang.glm52_mtp.common import (

DRAFT_SUB_STEP_ATTR,

attention_page_size,

build_draft_decode_token_table,

build_token_table,

counts_to_indptr,

ensure_shared_sparse_buffer,

flatten_kv_indices,

get_index_topk,

get_seq_lens_cpu,

local_num_attention_heads,

make_mla_work_buffers,

make_sparse_mtp_work_buffers,

metadata_dtype,

resolve_draft_decode_context_lens,

resolve_speculative_num_steps,

validate_page_size,

)

github-actions · 2026-07-25T07:44:38Z

+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+
+from atom.plugin.sglang.glm52_mtp.common import (
+    get_extend_lens_cpu,
+    get_extend_prefix_lens_cpu,
+    get_seq_lens_cpu,
+    is_draft_extend_mode,
+)
+from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata
+
+
+def draft_extend_k_only() -> bool:


⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

Suggested change

from __future__ import annotations

import os

import numpy as np

import torch

from atom.plugin.sglang.glm52_mtp.common import (

get_extend_lens_cpu,

get_extend_prefix_lens_cpu,

get_seq_lens_cpu,

is_draft_extend_mode,

)

from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata

def draft_extend_k_only() -> bool:

from __future__ import annotations

import os

import numpy as np

import torch

from atom.plugin.sglang.glm52_mtp.common import (

get_extend_lens_cpu,

get_extend_prefix_lens_cpu,

get_seq_lens_cpu,

is_draft_extend_mode,

)

from atom.plugin.sglang.glm52_mtp.multi_token import (

build_mtp_multi_token_decode_metadata,

)

def draft_extend_k_only() -> bool:

github-actions · 2026-07-25T07:44:38Z

+    if int(getattr(forward_batch, "_graph_cache_bs", 0) or 0) > 0:
+        return True
+    return False


⚠️ [ruff] <SIM103> _{reported by reviewdog 🐶}
Return the condition directly

Suggested change

if int(getattr(forward_batch, "_graph_cache_bs", 0) or 0) > 0:

return True

return False

return int(getattr(forward_batch, "_graph_cache_bs", 0) or 0) > 0

github-actions · 2026-07-25T07:44:42Z

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

ATOM/atom/plugin/register.py

Lines 187 to 197 in 7bf17aa

from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner

from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (

EAGLEDraftCudaGraphRunner,

)

from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (

EAGLEDraftExtendCudaGraphRunner,

)

from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker

from atom.plugin.sglang.glm52_mtp.common import (

GLM52_GRAPH_SEQ_LEN_CAPACITY,

)

ZhiweiYan-96 added 5 commits July 23, 2026 02:18

eager success

f1bc9f4

cleanup(sglang): remove native draft compatibility

98506c8

Keep the PR branch limited to the general ATOM draft attention route by removing hybrid native-attention cache and layer-id compatibility.

cleanup(sglang): remove draft extend A/B bypass

13c0760

Keep the GLM-5.2 PR route fixed on the validated plugin-owned draft-extend implementation and remove the upstream-handler environment switch.

cleanup(sglang): follow server CUDA graph defaults

65173c1

Remove GLM-5.2 phase and capacity environment switches. Let SGLang's native disable_cuda_graph server argument control graph capture, with all speculative graph phases enabled by default and an internal fixed metadata capacity.

github-actions Bot reviewed Jul 24, 2026

View reviewed changes

cleanup(sglang): fix GLM-5.2 attention routing

7bf17aa

Keep target attention on the native ATOM GLM-5.2 backend and draft attention on the general AITER backend without a runtime mode switch. Remove debug metadata hooks, allocation pointer logs, and route-only logging wrappers.

github-actions Bot reviewed Jul 25, 2026

View reviewed changes

	from atom.plugin.sglang.glm52_mtp.cache_bind import bind_glm52_dsa_cache_views
	from atom.plugin.sglang.glm52_mtp.common import (
	ATTENTION_PAGE_SIZE_ATTR,
	DECODE_GRAPH_BUFFERS_ATTR,
	DRAFT_SUB_STEP_ATTR,
	EMPTY_VALUE_CACHE_ATTR,
	INDEXER_PAGE_SIZE_ATTR,
	SHARED_SPARSE_INDICES_ATTR,
	is_glm52_dsa_arch,
	maybe_get_glm52_dsa_pools_from_sglang_backend,
	)
	from atom.plugin.sglang.glm52_mtp.decode_graph import (
	build_atom_glm52_decode_graph_metadata_from_sglang,
	)
	from atom.plugin.sglang.glm52_mtp.dispatcher import (
	build_atom_glm52_attention_metadata_from_sglang,
	)
	from atom.plugin.sglang.glm52_mtp.draft_decode import (
	build_decode_metadata as _build_decode_metadata,
	clear_draft_decode_sub_step,
	get_draft_decode_sub_step,
	set_draft_decode_sub_step,
	)
	from atom.plugin.sglang.glm52_mtp.draft_extend import (
	build_mtp_draft_extend_decode_metadata as _build_mtp_draft_extend_decode_metadata,
	draft_extend_token_num,
	draft_extend_token_num as _draft_extend_token_num,
	should_use_mtp_draft_extend_decode_path as _should_use_mtp_draft_extend_decode_path,
	)
	from atom.plugin.sglang.glm52_mtp.target_verify import (
	build_mtp_verify_decode_metadata as _build_mtp_verify_decode_metadata,
	should_use_mtp_verify_prefill_path as _should_use_mtp_verify_prefill_path,
	)

	if extend_end > int(req_to_token_row.numel()):
	extend_end = int(req_to_token_row.numel())

	from __future__ import annotations

	import torch
	from aiter import get_mla_metadata_v1
	from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton

	from atom.plugin.sglang.glm52_mtp.common import (
	DECODE_GRAPH_BUFFERS_ATTR,
	SHARED_SPARSE_INDICES_ATTR,
	get_index_topk,
	local_num_attention_heads,
	make_mla_work_buffers,
	metadata_dtype,
	validate_page_size,
	attention_page_size,
	)

	from __future__ import annotations

	import os

	import numpy as np
	import torch

	from atom.plugin.sglang.glm52_mtp.common import (
	get_extend_lens_cpu,
	get_extend_prefix_lens_cpu,
	get_seq_lens_cpu,
	is_draft_extend_mode,
	)
	from atom.plugin.sglang.glm52_mtp.multi_token import build_mtp_multi_token_decode_metadata

	if int(getattr(forward_batch, "_graph_cache_bs", 0) or 0) > 0:
	return True
	return False

	from __future__ import annotations

	import logging
	import os

	import numpy as np
	import torch

	from atom.plugin.sglang.glm52_mtp.common import (
	attention_page_size,
	build_token_table,
	compute_mtp_sparse_per_token_kv_lens,
	counts_to_indptr,
	ensure_shared_sparse_buffer,
	flatten_kv_indices,
	get_index_topk,
	is_draft_extend_mode,
	local_num_attention_heads,
	make_mla_work_buffers,
	make_sparse_mtp_work_buffers,
	metadata_dtype,
	validate_page_size,
	)

	if num_tokens_per_req is not None and int(num_tokens_per_req) != 1:
	return False
	return True

	setattr(
	token_to_kv_pool,
	"_atom_glm52_draft_decode_graph_metadata",
	graph_cache,
	)

	from typing import Callable, Generator, Tuple
	from typing import Tuple
	from collections.abc import Callable, Generator

	) -> Generator[Tuple[str, torch.Tensor], None, None]:
	) -> Generator[tuple[str, torch.Tensor], None, None]:

	for mapping_part in weights_mapping.keys():
	for mapping_part in weights_mapping:

	work_indptr: Optional[torch.Tensor] = None
	work_indptr: torch.Tensor \| None = None

	work_info_set: Optional[torch.Tensor] = None
	work_info_set: torch.Tensor \| None = None

	reduce_indptr: Optional[torch.Tensor] = None
	reduce_indptr: torch.Tensor \| None = None

	reduce_final_map: Optional[torch.Tensor] = None
	reduce_final_map: torch.Tensor \| None = None

	reduce_partial_map: Optional[torch.Tensor] = None
	reduce_partial_map: torch.Tensor \| None = None

	import concurrent.futures
	import json
	import os
	import logging
	import re
	import threading
	import time
	from glob import glob
	from typing import Callable, Generator, Tuple
	from collections.abc import Iterable, Mapping
	from dataclasses import dataclass, field
	from typing import Any, Optional

	import safetensors
	import safetensors.torch
	import torch
	from torch import nn
	from tqdm import tqdm
	from transformers import AutoConfig

	from __future__ import annotations

	from dataclasses import dataclass
	import logging
	import os
	import re
	from typing import Optional

	import numpy as np
	import torch
	from aiter import (
	cp_gather_indexer_k_quant_cache,
	dtypes,
	get_mla_metadata_info_v1,
	get_mla_metadata_v1,
	indexer_k_quant_and_cache,
	indexer_qk_rope_quant_and_cache,
	top_k_per_row_decode,
	top_k_per_row_prefill,
	)
	from aiter.mla import mla_decode_fwd
	from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
	from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
	import triton
	import triton.language as tl

	from atom.utils.custom_register import direct_register_custom_op
	from atom.plugin.sglang.glm52_mtp.common import GLM52_GRAPH_SEQ_LEN_CAPACITY

	import copy
	import contextlib
	import logging
	import os
	import re
	from typing import Iterable, Optional, Tuple

	import torch
	from torch import nn

	from sglang.srt.distributed import get_pp_group
	from sglang.srt.layers.logits_processor import LogitsProcessor
	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

	from atom.config import QuantizationConfig as AtomQuantizationConfig
	from atom.config import SpeculativeConfig
	from atom.plugin.config import generate_atom_config_for_plugin_mode
	from atom.plugin.sglang.models.deepseek_mla import (
	setup_deepseek_for_sglang,
	)
	from atom.plugin.sglang.runtime import (
	SGLangPluginRuntime,
	plugin_runtime_scope,
	)

	from atom.plugin.register import (
	init_aiter_dist,
	register_ops_to_sglang,
	set_attn_cls,
	)
	from atom.models.deepseek_mtp import DeepSeekMTP

Uh oh!

Conversation

ZhiweiYan-96 commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Jul 24, 2026

🏷️ CI Guide

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 25, 2026

ZhiweiYan-96 commented Jul 24, 2026 •

edited

Loading

		@@ -204,6 +213,43 @@ def _is_dsv4_nextn_runner(runner) -> bool:
		except Exception:

	num_heads = int(getattr(hf_config, "num_attention_heads"))
	num_heads = int(hf_config.num_attention_heads)

	from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
	from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
	EAGLEDraftCudaGraphRunner,
	)
	from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
	EAGLEDraftExtendCudaGraphRunner,
	)
	from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
	from atom.plugin.sglang.glm52_mtp.common import (
	GLM52_GRAPH_SEQ_LEN_CAPACITY,
	)