[sgl atom] Enable pcp in sglang atom in glm5.2 by ZLkanyo009 · Pull Request #1680 · ROCm/ATOM

ZLkanyo009 · 2026-07-23T08:04:36Z

Motivation

Enable sgl atom pcp in glm5.2.Also can run in model based on deepseekv2 attention backend like deepseekv3.2.

Command

GLM5.2:

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 CUDA_VISIBLE_DEVICES=0,1,2,3
MODEL_PATH=/models/GLM-5.2-FP8


TP=4
PORT=9000
MODEL_LOADER_EXTRA_CONFIG='{"online_quant_config":{"global_quant_config":"ptpc_fp8","layer_quant_config":{"model.layers.*.mlp.experts":"per_block_fp8"},"exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate"]}}'

export ATOM_SGLANG_PCP_SIZE=4
export ATOM_PCP_MOE_MERGE=1

TORCHINDUCTOR_COMPILE_THREADS=128 \
python3 -m sglang.launch_server \
    --model-path "${MODEL_PATH}" \
    --host localhost \
    --port "${PORT}" \
    --trust-remote-code \
    --tp-size "${TP}" \
    --enable-nsa-prefill-context-parallel \
    --attn-cp-size $ATOM_SGLANG_PCP_SIZE \
    --nsa-prefill-cp-mode round-robin-split \
    --max-prefill-tokens 30720 \
    --chunked-prefill-size 30720 \
    --mem-fraction-static 0.85 \
    --disable-radix-cache \
    --kv-cache-dtype fp8_e4m3 \
    --model-loader-extra-config "${MODEL_LOADER_EXTRA_CONFIG}"

Acc

glm:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     3|exact_match|↑  |0.9424|±  |0.0064|
|     |       |strict-match    |     3|exact_match|↑  |0.9409|±  |0.0065|

Performance

glm:
TP4:

============ Serving Benchmark Result ============
Successful requests:                     192       
Benchmark duration (s):                  778.15    
Total input tokens:                      11796480  
Total generated tokens:                  115200    
Request throughput (req/s):              0.25      
Output token throughput (tok/s):         148.04    
Total Token throughput (tok/s):          15307.70  
---------------Time to First Token----------------
Mean TTFT (ms):                          197197.47 
Median TTFT (ms):                        223083.68 
P99 TTFT (ms):                           254167.43 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          40.08     
Median TPOT (ms):                        38.71     
P99 TPOT (ms):                           62.80     
---------------Inter-token Latency----------------
Mean ITL (ms):                           40.08     
Median ITL (ms):                         18.88     
P99 ITL (ms):                            20.23     
----------------End-to-end Latency----------------
Mean E2EL (ms):                          221207.79 
Median E2EL (ms):                        241903.43 
P99 E2EL (ms):                           283149.80 
==================================================

PCP4 TP1:

============ Serving Benchmark Result ============
Successful requests:                     192       
Benchmark duration (s):                  782.97    
Total input tokens:                      11796480  
Total generated tokens:                  115200    
Request throughput (req/s):              0.25      
Output token throughput (tok/s):         147.13    
Total Token throughput (tok/s):          15213.40  
---------------Time to First Token----------------
Mean TTFT (ms):                          196649.36 
Median TTFT (ms):                        221870.01 
P99 TTFT (ms):                           253955.05 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          42.26     
Median TPOT (ms):                        40.98     
P99 TPOT (ms):                           63.27     
---------------Inter-token Latency----------------
Mean ITL (ms):                           42.26     
Median ITL (ms):                         22.84     
P99 ITL (ms):                            23.89     
----------------End-to-end Latency----------------
Mean E2EL (ms):                          221964.00 
Median E2EL (ms):                        241971.16 
P99 E2EL (ms):                           283670.18 
==================================================

github-actions · 2026-07-23T08:05:13Z

🏷️ 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 1680 --add-label <label>

yitingw1 · 2026-07-23T10:01:42Z

            hidden_states.shape[1] == self.experts.hidden_size
        ), f"Hidden states dimension {hidden_states.shape[1]} does not match expected {self.experts.hidden_size}"

+        if self._pcp_moe_merge_active():


Maybe ATOM GLM5.2 need PCP all-gather?

github-actions · 2026-07-24T08:33:05Z

@@ -40,6 +40,8 @@ class PluginConfig:
    sglang_aiter_rank_id: int = 0
    sglang_dist_init_addr: Optional[str] = None


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

Suggested change

sglang_dist_init_addr: Optional[str] = None

sglang_dist_init_addr: str | None = None

github-actions · 2026-07-24T08:33:05Z

+        )
+
+        pcp_size = get_pcp_world_size()
+    except Exception:


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

github-actions · 2026-07-24T08:33:05Z

+        )
+
+        pcp_size = get_pcp_world_size()
+    except Exception:


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

github-actions · 2026-07-24T08:33:05Z

 def _build_sparse_req_id_per_token_for_sglang(
    forward_batch,
    device: torch.device,
+    num_tokens: Optional[int] = None,


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

Suggested change

num_tokens: Optional[int] = None,

num_tokens: int | None = None,

github-actions · 2026-07-24T08:33:05Z

+            pcp_pad_len,
+            pcp_round_robin_query_indices,
+        )
+    except Exception:


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

github-actions · 2026-07-24T08:33:07Z

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

ATOM/atom/models/deepseek_v2.py

Lines 26 to 114 in b77c59b

import logging

from typing import Optional, Tuple, Union

import torch

from aiter import (

QuantType,

cp_gather_indexer_k_quant_cache,

dtypes,

fused_qk_rmsnorm,

gemm_a8w8_blockscale_bpreshuffle,

get_hip_quant,

indexer_k_quant_and_cache,

indexer_qk_rope_quant_and_cache,

top_k_per_row_decode,

top_k_per_row_prefill,

)

from aiter.dist.communication_op import tensor_model_parallel_all_reduce

from aiter.dist.parallel_state import get_pp_group, get_tensor_model_parallel_world_size

from aiter.jit.utils.torch_guard import torch_compile_guard

from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits

from aiter.ops.triton.fused_fp8_quant import fused_reduce_rms_fp8_group_quant

from aiter.ops.triton.fused_mxfp4_quant import (

fused_reduce_rms_mxfp4_quant,

fused_rms_mxfp4_quant,

)

from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits

from aiter.rotary_embedding import get_rope

from atom.config import Config, QuantizationConfig, get_current_atom_config

from atom.model_ops.activation import SiluAndMul

from atom.model_ops.attention_mla import (

MLAModules,

is_rocm_aiter_fp4bmm_enabled,

triton_convert_req_index_to_global_index,

triton_convert_req_index_to_global_index_dsa_prefill,

triton_gather_kv_indices_sparse,

)

from atom.model_ops.base_attention import Attention

from atom.model_ops.embed_head import (

ParallelLMHead,

ReplicatedEmbedding,

VocabParallelEmbedding,

)

from atom.model_ops.layernorm import LayerNorm, RMSNorm

from atom.model_ops.linear import (

ColumnParallelLinear,

MergedColumnParallelLinear,

MergedReplicatedLinear,

ReplicatedLinear,

RowParallelLinear,

use_fp4_non_shuffle_triton_gemm,

use_triton_gemm,

)

from atom.model_ops.moe import FusedMoE

from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled

from atom.model_ops.utils import MXFP4_QUANT_BLOCK_SIZE, atom_parameter

from atom.models.utils import (

IntermediateTensors,

PPMissingLayer,

make_empty_intermediate_tensors_factory,

make_layers,

maybe_prefix,

)

from atom.quant_spec import should_skip_online_quant

from atom.utils import envs

from atom.utils.custom_register import direct_register_custom_op

# Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward`,

# shared with deepseek_v4. DeepseekV2MoE.forward dispatches via this op when

# `_use_dual_stream` is True so torch.compile/Dynamo treats stream code as opaque.

from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401

from atom.distributed.pcp_utils import (

pcp_all_reduce,

pcp_allgather_rankmajor,

get_pcp_world_size,

pcp_allgather_rerange,

pcp_is_enabled,

pcp_pad_dense,

pcp_pad_len,

pcp_reduce_scatter,

pcp_round_robin_split,

)

from atom.utils.decorators import mark_trace, support_torch_compile

from atom.utils.forward_context import get_forward_context

from atom.plugin.vllm.attention.layer_sparse_mla import (

IndexerDecoratorForPluginMode,

DeepseekV32IndexerCacheDecoratorForPluginMode,

)

from torch import nn

from transformers import PretrainedConfig

github-actions · 2026-07-24T08:33:08Z

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

ATOM/atom/plugin/config.py

Lines 327 to 342 in b77c59b

from sglang.srt.distributed import get_tensor_model_parallel_rank

from sglang.srt.layers.dp_attention import (

get_attention_cp_rank,

get_attention_cp_size,

get_attention_tp_rank,

get_attention_tp_size,

)

from sglang.srt.server_args import (

get_global_server_args,

PortArgs,

ZMQ_TCP_PORT_DELTA,

)

from sglang.srt.configs.model_config import ModelConfig as SglangModelConfig

from sglang.srt.configs.modelopt_config import ModelOptConfig

from sglang.srt.configs.load_config import LoadConfig

from atom.config import Config, ParallelConfig, CompilationConfig

github-actions · 2026-07-24T08:33:08Z

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

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

Lines 186 to 194 in b77c59b

from aiter import dtypes

from atom.model_ops.attention_mla import fused_qk_rope_concat_and_cache_mla

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

_get_sglang_radix_attn,

mla_absorbed_bmm,

mla_v_up_proj,

)

from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp

from atom.models.deepseek_v2 import _pcp_active

Route PCP MoE merge through an opaque custom op so dummy warmup cannot bake away runtime prefill and decode collectives.

github-actions · 2026-07-24T09:07:42Z

@@ -94,11 +94,14 @@
 # `_use_dual_stream` is True so torch.compile/Dynamo treats stream code as opaque.
 from atom.model_ops import module_dispatch_ops as _module_dispatch_ops  # noqa: F401


⚠️ [ruff] <RUF100> _{reported by reviewdog 🐶}
Unused noqa directive (unused: F401)

Suggested change

from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401

from atom.model_ops import module_dispatch_ops as _module_dispatch_ops

github-actions · 2026-07-24T09:07:42Z

+        if context is None or bool(getattr(context, "is_dummy_run", False)):
+            return False
+        return True


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

Suggested change

if context is None or bool(getattr(context, "is_dummy_run", False)):

return False

return True

return not (context is None or bool(getattr(context, "is_dummy_run", False)))

ZLkanyo009 marked this pull request as ready for review July 23, 2026 08:04

ZLkanyo009 added ci:atom ci:sglang labels Jul 23, 2026

junyyang-amd previously approved these changes Jul 23, 2026

View reviewed changes

yitingw1 reviewed Jul 23, 2026

View reviewed changes

zufayu requested a review from ZhangLirong-amd July 24, 2026 01:32

[sgl atom] Enable pcp in sglang atom in glm5.2

b77c59b

ZLkanyo009 dismissed junyyang-amd’s stale review via b77c59b July 24, 2026 08:32

ZLkanyo009 force-pushed the lingzha/sgl_glm_pcp_v2 branch from faa2420 to b77c59b Compare July 24, 2026 08:32

github-actions Bot reviewed Jul 24, 2026

View reviewed changes

[sgl atom] Keep V2 PCP MoE merge opaque to compile

0c437d2

Route PCP MoE merge through an opaque custom op so dummy warmup cannot bake away runtime prefill and decode collectives.

github-actions Bot reviewed Jul 24, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[sgl atom] Enable pcp in sglang atom in glm5.2#1680

[sgl atom] Enable pcp in sglang atom in glm5.2#1680
ZLkanyo009 wants to merge 2 commits into
mainfrom
lingzha/sgl_glm_pcp_v2

ZLkanyo009 commented Jul 23, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 23, 2026

Uh oh!

yitingw1 Jul 23, 2026 •

edited

Loading

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

github-actions Bot Jul 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -40,6 +40,8 @@ class PluginConfig:
		sglang_aiter_rank_id: int = 0
		sglang_dist_init_addr: Optional[str] = None

	sglang_dist_init_addr: Optional[str] = None
	sglang_dist_init_addr: str \| None = None

	num_tokens: Optional[int] = None,
	num_tokens: int \| None = None,

	import logging
	from typing import Optional, Tuple, Union

	import torch
	from aiter import (
	QuantType,
	cp_gather_indexer_k_quant_cache,
	dtypes,
	fused_qk_rmsnorm,
	gemm_a8w8_blockscale_bpreshuffle,
	get_hip_quant,
	indexer_k_quant_and_cache,
	indexer_qk_rope_quant_and_cache,
	top_k_per_row_decode,
	top_k_per_row_prefill,
	)
	from aiter.dist.communication_op import tensor_model_parallel_all_reduce
	from aiter.dist.parallel_state import get_pp_group, get_tensor_model_parallel_world_size
	from aiter.jit.utils.torch_guard import torch_compile_guard
	from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
	from aiter.ops.triton.fused_fp8_quant import fused_reduce_rms_fp8_group_quant
	from aiter.ops.triton.fused_mxfp4_quant import (
	fused_reduce_rms_mxfp4_quant,
	fused_rms_mxfp4_quant,
	)
	from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
	from aiter.rotary_embedding import get_rope
	from atom.config import Config, QuantizationConfig, get_current_atom_config
	from atom.model_ops.activation import SiluAndMul
	from atom.model_ops.attention_mla import (
	MLAModules,
	is_rocm_aiter_fp4bmm_enabled,
	triton_convert_req_index_to_global_index,
	triton_convert_req_index_to_global_index_dsa_prefill,
	triton_gather_kv_indices_sparse,
	)
	from atom.model_ops.base_attention import Attention
	from atom.model_ops.embed_head import (
	ParallelLMHead,
	ReplicatedEmbedding,
	VocabParallelEmbedding,
	)
	from atom.model_ops.layernorm import LayerNorm, RMSNorm
	from atom.model_ops.linear import (
	ColumnParallelLinear,
	MergedColumnParallelLinear,
	MergedReplicatedLinear,
	ReplicatedLinear,
	RowParallelLinear,
	use_fp4_non_shuffle_triton_gemm,
	use_triton_gemm,
	)
	from atom.model_ops.moe import FusedMoE
	from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled
	from atom.model_ops.utils import MXFP4_QUANT_BLOCK_SIZE, atom_parameter
	from atom.models.utils import (
	IntermediateTensors,
	PPMissingLayer,
	make_empty_intermediate_tensors_factory,
	make_layers,
	maybe_prefix,
	)
	from atom.quant_spec import should_skip_online_quant
	from atom.utils import envs
	from atom.utils.custom_register import direct_register_custom_op

	# Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward`,
	# shared with deepseek_v4. DeepseekV2MoE.forward dispatches via this op when
	# `_use_dual_stream` is True so torch.compile/Dynamo treats stream code as opaque.
	from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401
	from atom.distributed.pcp_utils import (
	pcp_all_reduce,
	pcp_allgather_rankmajor,
	get_pcp_world_size,
	pcp_allgather_rerange,
	pcp_is_enabled,
	pcp_pad_dense,
	pcp_pad_len,
	pcp_reduce_scatter,
	pcp_round_robin_split,
	)
	from atom.utils.decorators import mark_trace, support_torch_compile
	from atom.utils.forward_context import get_forward_context
	from atom.plugin.vllm.attention.layer_sparse_mla import (
	IndexerDecoratorForPluginMode,
	DeepseekV32IndexerCacheDecoratorForPluginMode,
	)
	from torch import nn
	from transformers import PretrainedConfig

	from sglang.srt.distributed import get_tensor_model_parallel_rank
	from sglang.srt.layers.dp_attention import (
	get_attention_cp_rank,
	get_attention_cp_size,
	get_attention_tp_rank,
	get_attention_tp_size,
	)
	from sglang.srt.server_args import (
	get_global_server_args,
	PortArgs,
	ZMQ_TCP_PORT_DELTA,
	)
	from sglang.srt.configs.model_config import ModelConfig as SglangModelConfig
	from sglang.srt.configs.modelopt_config import ModelOptConfig
	from sglang.srt.configs.load_config import LoadConfig
	from atom.config import Config, ParallelConfig, CompilationConfig

	from aiter import dtypes
	from atom.model_ops.attention_mla import fused_qk_rope_concat_and_cache_mla
	from atom.plugin.sglang.models.deepseek_mla_forward import (
	_get_sglang_radix_attn,
	mla_absorbed_bmm,
	mla_v_up_proj,
	)
	from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
	from atom.models.deepseek_v2 import _pcp_active

		@@ -94,11 +94,14 @@
		# `_use_dual_stream` is True so torch.compile/Dynamo treats stream code as opaque.
		from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401

Uh oh!

Conversation

ZLkanyo009 commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Command

Acc

Performance

Uh oh!

github-actions Bot commented Jul 23, 2026

🏷️ CI Guide

Uh oh!

yitingw1 Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ZLkanyo009 commented Jul 23, 2026 •

edited

Loading

yitingw1 Jul 23, 2026 •

edited

Loading