From d3475867c021a98099103fd140bdd3ccf25e1d5a Mon Sep 17 00:00:00 2001 From: Pleaplusone Date: Fri, 8 May 2026 15:21:12 +0800 Subject: [PATCH 01/76] enable configurable weight bpreshuffle for fp8 blockscale gemm (#694) * enable configurable weight bpreshuffle for fp8 blockscale gemm Signed-off-by: ganyi * add moe for configurable bpreshuffle Signed-off-by: ganyi --------- Signed-off-by: ganyi Co-authored-by: wuhuikx --- atom/model_ops/layernorm.py | 22 +++++++--- atom/model_ops/linear.py | 39 ++++++++++++----- atom/model_ops/moe.py | 29 +++++++++---- .../triton_fused_sigmoid_mul_quant.py | 42 +++++++++++++------ atom/utils/envs.py | 6 +++ 5 files changed, 101 insertions(+), 37 deletions(-) diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 030bcaae7a..7a687e5243 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -21,6 +21,7 @@ from atom.model_ops.utils import atom_parameter from atom.quant_spec import LayerQuantConfig from atom.utils.decorators import mark_trace +from atom.utils import envs from torch import Tensor, nn from torch.overrides import handle_torch_function, has_torch_function_unary @@ -347,8 +348,9 @@ def __init__( # Extract group size from quant type if quant_type == QuantType.per_1x128: self.group_size_quant = 128 - # per_1x128 blockscale GEMM requires transposed scale layout - self.transpose_scale = True + # preshuffle GEMM expects column-major x_scale; + # non-preshuffle GEMM expects row-major x_scale + self.transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE elif quant_type == QuantType.per_1x32: self.group_size_quant = 32 self.transpose_scale = False @@ -557,15 +559,23 @@ def _forward_fused_fp8(self, x, residual=None): from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant from aiter.utility.dtypes import fp8 + transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE group_size = 128 M = x.shape[0] N = x.shape[1] num_groups = N // group_size out_fp8 = torch.empty((M, N), dtype=fp8, device=x.device) - out_scale = torch.empty( - (num_groups, M), dtype=torch.float32, device=x.device - ).view(M, num_groups) + if transpose_scale: + # column-major: allocate (num_groups, M) then view as (M, num_groups) + out_scale = torch.empty( + (num_groups, M), dtype=torch.float32, device=x.device + ).view(M, num_groups) + else: + # row-major: allocate (M, num_groups) directly + out_scale = torch.empty( + (M, num_groups), dtype=torch.float32, device=x.device + ) out_bf16 = ( torch.empty((M, N), dtype=x.dtype, device=x.device) if self.write_bf16 @@ -583,7 +593,7 @@ def _forward_fused_fp8(self, x, residual=None): q_res_out=res_out, q_residual=residual, group_size=group_size, - transpose_scale=True, + transpose_scale=transpose_scale, gemma_norm=True, ) if residual is not None: diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index a0cb99179e..c3a09e829e 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -13,6 +13,7 @@ gemm_a8w8, gemm_a8w8_blockscale_bpreshuffle, gemm_a8w8_bpreshuffle, + gemm_a8w8_blockscale, get_hip_quant, ) @@ -379,10 +380,14 @@ def process_weights_after_loading(self): shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) else: - if ( + need_shuffle = ( self.quant_type == QuantType.per_Token and self.params_dtype == dtypes.fp8 - ) or (self.quant_type in [QuantType.per_1x32, QuantType.per_1x128]): + ) or self.quant_type == QuantType.per_1x32 + # per_1x128 only needs shuffle when using the preshuffle GEMM path + if not need_shuffle and self.quant_type == QuantType.per_1x128: + need_shuffle = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + if need_shuffle: if self.weight.dim() == 2: shuffle_weights(self.weight) # self.weight_scale.data = fp4_utils.e8m0_shuffle(self.weight_scale.data) @@ -405,8 +410,11 @@ def forward( if x_scale is None: quant_func = self.quant_func if self.quant_type.value == QuantType.per_1x128.value: + # preshuffle GEMM expects column-major x_scale; + # non-preshuffle GEMM expects row-major x_scale quant_func = functools_partial( - self.quant_func, transpose_scale=True + self.quant_func, + transpose_scale=envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE, ) if self.quant_type.value != QuantType.per_1x32.value: x, x_scale = quant_func( @@ -444,14 +452,23 @@ def forward( if self.bias is not None: y += self.bias elif self.quant_type.value == QuantType.per_1x128.value: - y = gemm_a8w8_blockscale_preshuffle_impl( - x, - self.weight, - x_scale, - self.weight_scale, - dtype=otype, - prefix=self.prefix, - ) + if envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE: + y = gemm_a8w8_blockscale_preshuffle_impl( + x, + self.weight, + x_scale, + self.weight_scale, + dtype=otype, + prefix=self.prefix, + ) + else: + y = gemm_a8w8_blockscale( + x, + self.weight, + x_scale, + self.weight_scale, + dtype=otype, + ) if self.bias is not None: y += self.bias elif self.quant_type.value == QuantType.per_1x32.value: diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 9f6e0a0f43..5a42b12b82 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -1351,13 +1351,22 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Update scale to single max scale per expert [E] layer.w13_weight_scale = atom_parameter(max_w13_scales) - # Shuffle weights for asm moe (moved from inference to load time for better performance) - if w13.dtype in [ - torch.int8, - torch.uint8, - torch.float8_e4m3fnuz, - torch.float8_e4m3fn, - ]: + # Shuffle weights for asm moe (moved from inference to load time for better performance). + # For per_1x128 blockscale (block_quant), only shuffle when the preshuffle GEMM + # path is enabled — the non-preshuffle kernel expects the un-shuffled layout. + skip_shuffle_for_block = ( + self.block_quant and not envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + ) + if ( + w13.dtype + in [ + torch.int8, + torch.uint8, + torch.float8_e4m3fnuz, + torch.float8_e4m3fn, + ] + and not skip_shuffle_for_block + ): from aiter.ops.shuffle import shuffle_weight w13.data = shuffle_weight(w13.data) @@ -1667,7 +1676,11 @@ def _process_block_quant(self, layer: nn.Module) -> None: layer.w2_weight = atom_parameter(layer.w2_weight.data) layer.w2_weight_scale = atom_parameter(layer.w2_weight_scale.data) - shuffle_weights(layer.w13_weight, layer.w2_weight) + # per_1x128 blockscale MoE only needs weight bpreshuffle when the + # preshuffle GEMM path is enabled. Skip it to match the non-preshuffle + # kernel's expected weight layout. + if envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE: + shuffle_weights(layer.w13_weight, layer.w2_weight) def _process_channel_quant(self, layer: nn.Module) -> None: """PTPTC""" diff --git a/atom/model_ops/triton_fused_sigmoid_mul_quant.py b/atom/model_ops/triton_fused_sigmoid_mul_quant.py index 7829540521..baf9c6e1cf 100644 --- a/atom/model_ops/triton_fused_sigmoid_mul_quant.py +++ b/atom/model_ops/triton_fused_sigmoid_mul_quant.py @@ -120,6 +120,7 @@ def fused_sigmoid_mul_fp8_quant( attn_output: Tensor, gate: Tensor, group_size: int = 128, + transpose_scale: bool | None = None, ) -> tuple[Tensor, Tensor]: """Fused sigmoid(gate) * attn_output + FP8 per-group quantization. @@ -127,13 +128,20 @@ def fused_sigmoid_mul_fp8_quant( attn_output: [M, N] bf16/fp16 — attention output tensor. gate: [M, N] bf16/fp16 — gating tensor (pre-sigmoid). group_size: Quantization group size (default 128, matching per_1x128). + transpose_scale: If True, produce column-major x_scale (for preshuffle GEMM). + If False, produce row-major x_scale (for non-preshuffle GEMM). + If None (default), follows ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE env var. Returns: (x_fp8, x_scale): x_fp8: [M, N] FP8 — quantized sigmoid(gate) * attn_output. - x_scale: [M, N // group_size] float32 — per-group scales in - transpose_scale layout (column-major storage). + x_scale: [M, N // group_size] float32 — per-group scales. """ + if transpose_scale is None: + from atom.utils import envs + + transpose_scale = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + M, N = attn_output.shape assert ( N % group_size == 0 @@ -144,11 +152,21 @@ def fused_sigmoid_mul_fp8_quant( out_fp8 = torch.empty((M, N), dtype=fp8_dtype, device=attn_output.device) num_scale_cols = N // group_size - # transpose_scale layout: allocate as (num_scale_cols, M) contiguous, - # then view as (M, num_scale_cols) after kernel - out_scale = torch.empty( - (num_scale_cols, M), dtype=torch.float32, device=attn_output.device - ) + if transpose_scale: + # column-major: allocate as (num_scale_cols, M) contiguous, + # then view as (M, num_scale_cols) after kernel + out_scale = torch.empty( + (num_scale_cols, M), dtype=torch.float32, device=attn_output.device + ) + stride_scale_m = out_scale.stride(1) + stride_scale_n = out_scale.stride(0) + else: + # row-major: allocate as (M, num_scale_cols) contiguous + out_scale = torch.empty( + (M, num_scale_cols), dtype=torch.float32, device=attn_output.device + ) + stride_scale_m = out_scale.stride(0) + stride_scale_n = out_scale.stride(1) grid = (M, N // BLOCK_SIZE_N) _fused_sigmoid_mul_fp8_group_quant_kernel[grid]( @@ -162,9 +180,8 @@ def fused_sigmoid_mul_fp8_quant( gate.stride(1), out_fp8.stride(0), out_fp8.stride(1), - # transpose_scale: swap strides so kernel writes in column-major - out_scale.stride(1), # stride_scale_m = stride along dim 1 (M dim) - out_scale.stride(0), # stride_scale_n = stride along dim 0 (scale_col dim) + stride_scale_m, + stride_scale_n, N=N, BLOCK_SIZE_N=BLOCK_SIZE_N, QUANT_BLOCK_SIZE=group_size, @@ -172,7 +189,8 @@ def fused_sigmoid_mul_fp8_quant( FP8_MIN=DTYPE_MIN, ) - # View transposed buffer back to (M, num_scale_cols) shape - out_scale = out_scale.view(M, num_scale_cols) + if transpose_scale: + # View transposed buffer back to (M, num_scale_cols) shape + out_scale = out_scale.view(M, num_scale_cols) return out_fp8, out_scale diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 269fa0c20a..51fee833a2 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -93,6 +93,12 @@ # Enable gradient tracking on model parameters. Default "0" (disabled) # is correct for inference; set to "1" only for training / fine-tuning. "ATOM_REQUIRES_GRAD": lambda: os.getenv("ATOM_REQUIRES_GRAD", "0") == "1", + # --- Bpreshuffle for weight --- + # Preshuffle weight. Default "1" (enabled) + "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE": lambda: os.getenv( + "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE", "1" + ) + == "1", # --- V4 Attention Backend Refactor (PR-A: kill .item(), unlock CUDAGraph) --- # `legacy` (default) keeps the per-seq Python dispatch loop with .item() # syncs in deepseek_v4.py. `new` routes through V4AttentionBackend with From 68d0fbee77b9e0812d8f5555f4e92610ea70de95 Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Fri, 8 May 2026 16:32:31 +0800 Subject: [PATCH 02/76] Support torch compile in deepseek v4 (#705) * Support deepseek v4 torch compile * remove hc head custom op --- atom/models/deepseek_v4.py | 77 +++++++++++++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 6 deletions(-) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 002511913c..1e41e72404 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -35,6 +35,7 @@ get_hip_quant, indexer_k_quant_and_cache, ) +from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.dist.communication_op import tensor_model_parallel_all_reduce from aiter.dist.parallel_state import get_tensor_model_parallel_world_size from aiter.ops.topk import top_k_per_row_decode, top_k_per_row_prefill @@ -66,7 +67,7 @@ hc_split_sinkhorn, ) from atom.model_ops.utils import atom_parameter -from atom.utils import envs +from atom.utils import envs, mark_spliting_op # Side-effect import: registers `torch.ops.aiter.maybe_dual_stream_forward` # (shared with deepseek_v2) and `torch.ops.aiter.indexer_score_topk` (V4-only). @@ -86,6 +87,7 @@ update_compressor_states, ) from atom.utils.forward_context import get_forward_context +from atom.utils.decorators import support_torch_compile # --------------------------------------------------------------------------- # Classical KV cache scatter / gather helpers (PR3-pre2c-B). @@ -120,6 +122,53 @@ def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: return rmsnorm2d_fwd_(x, ones, eps, dim) +def _hc_head_reduce_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + return torch.empty(x.shape[0], x.shape[-1], dtype=x.dtype, device=x.device) + + +@torch_compile_guard(gen_fake=_hc_head_reduce_fake, mutates_args=[]) +def _hc_head_reduce( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + x_flat = x.flatten(-2) + x_normed = _rmsnorm_nw(x_flat, norm_eps, x_flat.shape[-1]) + mixes = F.linear(x_normed.float(), hc_fn) + pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps + y = torch.sum(pre.unsqueeze(-1) * x, dim=-2) + return y.to(x.dtype) + + +def _v4_attention_fake( + x: torch.Tensor, + positions: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + return torch.empty_like(x) + + +@mark_spliting_op(is_custom=True, gen_fake=_v4_attention_fake, mutates_args=[]) +def v4_attention_with_output( + x: torch.Tensor, + positions: torch.Tensor, + layer_name: str, +) -> torch.Tensor: + atom_config = get_current_atom_config() + self = atom_config.compilation_config.static_forward_context[layer_name] + return self.forward_impl(x, positions) + + # --------------------------------------------------------------------------- # Config wrapper # --------------------------------------------------------------------------- @@ -1291,6 +1340,10 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): self.indexer.rotary_emb = self.rotary_emb self.indexer.compressor.rotary_emb = self.rotary_emb + self.layer_name = prefix + atom_config = get_current_atom_config() + atom_config.compilation_config.static_forward_context[self.layer_name] = self + def process_weights_after_loading(self) -> None: """Dequant wo_a (FP8 + e8m0 block scale) → BF16 in place. @@ -1342,6 +1395,13 @@ def process_weights_after_loading(self) -> None: self.wo_a.quant_type = QuantType.No def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.aiter.v4_attention_with_output(x, positions, self.layer_name) + + def forward_impl( self, x: torch.Tensor, # [num_tokens, dim] flat ragged-batch hidden state positions: torch.Tensor, # [num_tokens] int absolute token positions @@ -2128,10 +2188,14 @@ def hc_head( """Reduce mHC residual `[num_tokens, hc, dim]` → `[num_tokens, dim]` via Sigmoid-gated weighted sum (vs Block.hc_pre's Sinkhorn variant). """ - _, _, y = aiter.mhc_pre( - x, hc_fn, hc_scale, hc_base, self.norm_eps, self.hc_eps, sinkhorn_repeat=0 + return _hc_head_reduce( + x, + hc_fn, + hc_scale, + hc_base, + self.norm_eps, + self.hc_eps, ) - return y def forward( self, @@ -2219,6 +2283,7 @@ def forward( ) +@support_torch_compile class DeepseekV4Model(nn.Module): """Full model: embed -> expand to hc_mult copies -> N blocks -> hc_head -> logits. @@ -2234,6 +2299,7 @@ class DeepseekV4Model(nn.Module): def __init__( self, *, + atom_config: Config, args: DeepseekV4Args, ): super().__init__() @@ -2285,7 +2351,6 @@ def __init__( self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) - @torch.inference_mode() def forward( self, input_ids: torch.Tensor, # [num_tokens] int flat ragged-batch token ids @@ -2382,7 +2447,7 @@ def __init__(self, config: Config, prefix: str = "") -> None: # config lacks `quantization_config` (e.g. dummy / toy validation), # this still works — base spec is QuantType.No. self.args.quant_config = make_v4_quant_config(self.hf_config) - self.model = DeepseekV4Model(args=self.args) + self.model = DeepseekV4Model(atom_config=config, args=self.args) def forward( self, From 5d34e9e0127ff935e5ff23ef41760d019deb1da0 Mon Sep 17 00:00:00 2001 From: Jiayun Date: Fri, 8 May 2026 18:02:00 +0800 Subject: [PATCH 03/76] Glm mtp test (#709) * sparse attn mtp wip * maybe I need these * wip * remove logs * remove logs * fix mem * run pass * rm logs * fix format * add test * fix draft model * dsa: remove block_table_convert_triton in dsa * commit * commit * clear code * function pass * fix zeros * fix format * Fix merge isssue --------- Co-authored-by: chenjun Co-authored-by: la <46212055+junhaha666@users.noreply.github.com> --- .github/benchmark/models_accuracy.json | 14 +- atom/config.py | 3 +- atom/model_engine/model_runner.py | 2 + atom/model_ops/attention_mla.py | 164 ++++++++++++--- atom/model_ops/attentions/aiter_mla.py | 271 ++++++++++++++++++++----- atom/models/deepseek_mtp.py | 45 +++- atom/models/deepseek_v2.py | 13 +- atom/models/mimo_v2_flash_mtp.py | 9 +- atom/spec_decode/eagle.py | 18 +- 9 files changed, 448 insertions(+), 91 deletions(-) diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 99efa88f28..a86ca34914 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -179,6 +179,18 @@ "accuracy_baseline_model": "zai-org/GLM-5.1", "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" }, + { + "model_name": "GLM-5.1-MXFP4 MTP", + "model_path": "amd/GLM-5.1-MXFP4", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --default-chat-template-kwargs '{\"enable_thinking\":false}' --method mtp --num-speculative-tokens 3", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "pr", + "accuracy_threshold": 0.87, + "accuracy_baseline": 0.9545, + "accuracy_baseline_model": "zai-org/GLM-5.1", + "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" + }, { "model_name": "GLM-5.1-MXFP4", "model_path": "amd/GLM-5.1-MXFP4", @@ -263,4 +275,4 @@ "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows MXFP4=0.9256, baseline=0.9401" } -] \ No newline at end of file +] diff --git a/atom/config.py b/atom/config.py index 2f42471389..d5de8e112b 100644 --- a/atom/config.py +++ b/atom/config.py @@ -727,6 +727,7 @@ class SpeculativeConfig: # model_type → mtp_model_type mapping _MTP_TYPE_MAP: ClassVar[dict[str, str]] = { "deepseek_v3": "deepseek_mtp", + "deepseek_v32": "deepseek_mtp", "glm_moe_dsa": "deepseek_mtp", "qwen3_next": "qwen3_next_mtp", "qwen3_5": "qwen3_5_mtp", @@ -745,7 +746,7 @@ class SpeculativeConfig: def __post_init__(self): if self.draft_model_hf_config is None: - self.draft_model_hf_config = AutoConfig.from_pretrained( + self.draft_model_hf_config = get_hf_config( self.model, trust_remote_code=True ) # For multimodal models, extract text_config diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index ce607625df..5199912d0b 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1877,6 +1877,8 @@ def capture_cudagraph(self): positions = self.forward_vars["positions"].gpu outputs = self.forward_vars["outputs"] self.forward_vars["kv_indptr"].gpu.zero_() + if self.is_deepseek_v32 and "sparse_kv_indptr" in self.forward_vars: + self.forward_vars["sparse_kv_indptr"].gpu.zero_() self.graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] = dict() self.graph_logits: dict[tuple[int, int], torch.Tensor] = dict() diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index b2a8aa67d1..556ec8cf07 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -484,6 +484,7 @@ def _forward_prefill_mla( attn_metadata.block_tables, attn_metadata.cu_seqlens_k, NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + PAGE_SIZE=get_current_atom_config().kv_cache_block_size, ) paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q paged_kv_indptr = attn_metadata.sparse_kv_indptr @@ -547,35 +548,65 @@ def _forward_decode( ) kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + paged_cu_seqlens_q = attn_metadata.cu_seqlens_q paged_kv_indptr = attn_metadata.kv_indptr paged_kv_indices = attn_metadata.kv_indices + paged_kv_last_page_lens = attn_metadata.kv_last_page_lens + max_q_len = attn_metadata.max_seqlen_q if self.topk_indices_buffer is not None: - paged_kv_indptr = attn_metadata.sparse_kv_indptr - paged_kv_indices = triton_convert_req_index_to_global_index( - attn_metadata.cu_seqlens_q, - attn_metadata.kv_indptr, - paged_kv_indptr, - attn_metadata.kv_indices, - self.topk_indices_buffer[:B], - NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], - ) - - # q_scale = kv_scale = None - # if self.kv_cache_dtype.startswith("fp8"): - # q = q.to(dtypes.fp8) - # q_scale = kv_scale = self.one_scale + if attn_metadata.max_seqlen_q > 1: + # MTP verify: per-token layout with max_q_len=1. + # Persistent metadata is per-token (from _set_mla_persistent_worker_buffers_sparse_mtp). + paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens + # Gather physical page indices from kv_indices using topk positions. + # block_tables contains large-block IDs (block_ratio > 1) that + # need expansion; kv_indices already has per-token page indices. + paged_kv_indices = triton_gather_kv_indices_sparse( + paged_kv_indptr, + attn_metadata.token_to_seq_idxs, + self.topk_indices_buffer[:B], + attn_metadata.kv_indices, + attn_metadata.kv_indptr, + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) + max_q_len = 1 + else: + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_indices = triton_convert_req_index_to_global_index( + attn_metadata.cu_seqlens_q, + attn_metadata.kv_indptr, + paged_kv_indptr, + attn_metadata.kv_indices, + self.topk_indices_buffer[:B], + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) dp_size = get_dp_group().world_size use_persistent_mode = not (dp_size > 1) + # Sparse layers in MTP verify use separate persistent metadata + # (per-token, max_seqlen_qo=1) while dense layers use normal metadata + # (max_seqlen_qo=2). + is_sparse_mtp = ( + self.topk_indices_buffer is not None and attn_metadata.max_seqlen_q > 1 + ) + if not use_persistent_mode: - # DP : disable persistent mode to avoid overflow work_meta_data = None work_indptr = None work_info_set = None reduce_indptr = None reduce_final_map = None reduce_partial_map = None + elif is_sparse_mtp: + work_meta_data = attn_metadata.sparse_mtp_work_meta_data + work_indptr = attn_metadata.sparse_mtp_work_indptr + work_info_set = attn_metadata.sparse_mtp_work_info_set + reduce_indptr = attn_metadata.sparse_mtp_reduce_indptr + reduce_final_map = attn_metadata.sparse_mtp_reduce_final_map + reduce_partial_map = attn_metadata.sparse_mtp_reduce_partial_map else: work_meta_data = attn_metadata.work_meta_data work_indptr = attn_metadata.work_indptr @@ -588,11 +619,11 @@ def _forward_decode( q, kv_buffer.view(-1, 1, 1, q.shape[-1]), o, - attn_metadata.cu_seqlens_q, + paged_cu_seqlens_q, paged_kv_indptr, paged_kv_indices, - attn_metadata.kv_last_page_lens, - attn_metadata.max_seqlen_q, + paged_kv_last_page_lens, + max_q_len, num_kv_splits=16, sm_scale=self.scale, work_meta_data=work_meta_data, @@ -878,7 +909,7 @@ def triton_convert_req_index_to_global_index( kv_indices_c = kv_indices.contiguous() token_indices_c = token_indices.contiguous() page_kv_indptr_c = page_kv_indptr.contiguous() - # TODO: not support mtp + # NOTE: MTP (max_seqlen_q > 1) uses triton_convert_req_index_to_global_index_dsa_prefill instead new_kv_indices = torch.empty_like(kv_indices) # Strides in elements @@ -916,7 +947,7 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel( out_kv_indices, # int32 # shapes (compile-time where possible) NUM_TOPK_TOKENS: tl.constexpr, - BLOCK_SIZE: tl.constexpr, + PAGE_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) ti_stride0: tl.int64, # topk_indices stride 0 @@ -941,14 +972,19 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel( ) # int32 pre_seqlens_q = tl.load(cu_seqlens_q + req_id) + seq_token_idx = indice - pre_seqlens_q + block_id = seq_token_idx // PAGE_SIZE + inblock_offset = seq_token_idx % PAGE_SIZE + # Guard block_table access store_mask = (col_id < kv_len) & (col_id < NUM_TOPK_TOKENS) valid_mask = store_mask & (indice >= 0) - out_val = tl.load( - block_table + req_id * bt_stride0 + (indice - pre_seqlens_q) * bt_stride1, + physical_block = tl.load( + block_table + req_id * bt_stride0 + block_id * bt_stride1, mask=valid_mask, other=-1, ) + out_val = tl.where(valid_mask, physical_block * PAGE_SIZE + inblock_offset, -1) # Store results out_ptr_ij = out_kv_indices + kv_start + col_id @@ -967,7 +1003,7 @@ def triton_convert_req_index_to_global_index_dsa_prefill( block_table: torch.Tensor, # int32 [num_req, max_num_blocks_per_req] cu_seqlens_q: torch.Tensor, # int32 [num_tokens + 1] # dsa_kv_indices: torch.Tensor, # int32 [total_kv_seqlen] -->>> output for this kernel - PAGE_SIZE: int = 1, # page_block_size = 1 for now + PAGE_SIZE: int = 1, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 1024, # tile width along columns ): @@ -1010,3 +1046,85 @@ def triton_convert_req_index_to_global_index_dsa_prefill( bt_stride1, ) return new_kv_indices + + +@triton.jit +def _gather_kv_indices_sparse_kernel( + sparse_kv_indptr, + token_to_seq_idxs, + topk_indices, + kv_indices, + kv_indptr, + out_kv_indices, + NUM_TOPK_TOKENS: tl.constexpr, + BLOCK_N: tl.constexpr, + ti_stride0: tl.int64, + ti_stride1: tl.constexpr, +): + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + col_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + req_id = tl.load(token_to_seq_idxs + token_id) + + out_start = tl.load(sparse_kv_indptr + token_id) + out_end = tl.load(sparse_kv_indptr + token_id + 1) + kv_len = out_end - out_start + + pos = tl.load(topk_indices + token_id * ti_stride0 + col_id * ti_stride1) + + kv_base = tl.load(kv_indptr + req_id) + kv_end = tl.load(kv_indptr + req_id + 1) + req_kv_len = kv_end - kv_base + + store_mask = (col_id < kv_len) & (col_id < NUM_TOPK_TOKENS) + valid_mask = store_mask & (pos >= 0) & (pos < req_kv_len) + + out_val = tl.load( + kv_indices + kv_base + pos, + mask=valid_mask, + other=0, + ) + + tl.store( + out_kv_indices + out_start + col_id, + out_val, + mask=store_mask, + ) + + +def triton_gather_kv_indices_sparse( + sparse_kv_indptr: torch.Tensor, + token_to_seq_idxs: torch.Tensor, + topk_indices: torch.Tensor, + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 1024, +): + assert topk_indices.shape[1] == NUM_TOPK_TOKENS + assert NUM_TOPK_TOKENS % BLOCK_N == 0 + + num_tokens = token_to_seq_idxs.shape[0] + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + out = torch.empty( + num_tokens * NUM_TOPK_TOKENS, dtype=torch.int32, device=topk_indices.device + ) + + ti_stride0, ti_stride1 = topk_indices.stride() + grid = (num_tokens, tiles_per_row) + + _gather_kv_indices_sparse_kernel[grid]( + sparse_kv_indptr, + token_to_seq_idxs, + topk_indices, + kv_indices, + kv_indptr, + out, + NUM_TOPK_TOKENS, + BLOCK_N, + ti_stride0, + ti_stride1, + ) + return out diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 1f1aa0f86d..68dad8f636 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -143,7 +143,54 @@ def __init__(self, model_runner): mla_metadata["sparse_kv_last_page_lens"].np[:] = 1 mla_metadata["sparse_kv_last_page_lens"].copy_to_gpu() + if self.is_sparse and max_seqlen_qo > 1: + # Allocate a second set of persistent work buffers for sparse MTP + # per-token layout: max_bs*max_seqlen_qo virtual seqs, each q_len=1. + smt_max_bs = self.max_bs * max_seqlen_qo + ( + (smt_wmd_size, smt_wmd_type), + (smt_wi_size, smt_wi_type), + (smt_wis_size, smt_wis_type), + (smt_ri_size, smt_ri_type), + (smt_rfm_size, smt_rfm_type), + (smt_rpm_size, smt_rpm_type), + ) = get_mla_metadata_info_v1( + smt_max_bs, + 1, # max_seqlen_qo=1 for per-token + self.padded_num_attention_heads, + self.dtype_q, + self.dtype_kv, + is_sparse=True, + fast_mode=True, + ) + mla_metadata["sparse_mtp_work_meta_data"] = torch.empty( + smt_wmd_size, dtype=smt_wmd_type, device=self.device + ) + mla_metadata["sparse_mtp_work_indptr"] = torch.empty( + smt_wi_size, dtype=smt_wi_type, device=self.device + ) + mla_metadata["sparse_mtp_work_info_set"] = torch.empty( + smt_wis_size, dtype=smt_wis_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_indptr"] = torch.empty( + smt_ri_size, dtype=smt_ri_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_final_map"] = torch.empty( + smt_rfm_size, dtype=smt_rfm_type, device=self.device + ) + mla_metadata["sparse_mtp_reduce_partial_map"] = torch.empty( + smt_rpm_size, dtype=smt_rpm_type, device=self.device + ) + self.model_runner.forward_vars.update(mla_metadata) + + if self.is_sparse: + self._token_to_seq_idxs_gpu = torch.zeros( + self.max_num_batched_tokens, + dtype=torch.int32, + device=self.device, + ) + # Per-ubatch buffers for CUDAGraph TBO if config.enable_tbo: self._allocate_ubatch_buffers( @@ -213,7 +260,6 @@ def _allocate_ubatch_buffers( **i32_kwargs, ) var[f"{p}cu_seqlens_q"] = CpuGpuBuffer(ub_max_bs + 1, **i32_kwargs) - # cu_seqlens_q for decode is constant [0, q, 2q, ..., bs*q] var[f"{p}cu_seqlens_q"].cpu.copy_( torch.arange( 0, @@ -267,6 +313,60 @@ def prep_stream(self): # return self.model_runner.tokenID_processor.async_copy_stream return self.model_runner.async_execute_stream + def _set_mla_persistent_worker_buffers_sparse_mtp( + self, + num_tokens: int, + ): + """Compute persistent metadata for sparse MTP per-token layout. + + B = batch_size * max_seqlen_q tokens are treated as B independent + virtual sequences each with q_len=1. cu_seqlens_q = [0,1,...,B], + kv_indptr = per-token sparse_kv_indptr, kv_last_page_lens = all 1s. + + Uses separate sparse_mtp_* buffers so dense layers can keep + their own persistent metadata (max_seqlen_qo=2) intact. + """ + var = self.model_runner.forward_vars + split_params = { + "kv_granularity": max(self.block_size, 16), + "max_seqlen_qo": 1, + "uni_seqlen_qo": 1, + "fast_mode": 1, + "max_split_per_batch": 16, + } + work_meta_data = var["sparse_mtp_work_meta_data"] + work_info_set = var["sparse_mtp_work_info_set"] + work_indptr = var["sparse_mtp_work_indptr"] + reduce_indptr = var["sparse_mtp_reduce_indptr"] + reduce_final_map = var["sparse_mtp_reduce_final_map"] + reduce_partial_map = var["sparse_mtp_reduce_partial_map"] + get_mla_metadata_v1( + var["sparse_cu_seqlens_q"].gpu[: num_tokens + 1], + var["sparse_kv_indptr"].gpu[: num_tokens + 1], + var["sparse_kv_last_page_lens"].gpu[:num_tokens], + self.padded_num_attention_heads, + 1, # nhead_kv + True, + work_meta_data, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + page_size=self.block_size, + dtype_q=self.dtype_q, + dtype_kv=self.dtype_kv, + **split_params, + ) + return { + "sparse_mtp_work_meta_data": work_meta_data, + "sparse_mtp_work_info_set": work_info_set, + "sparse_mtp_work_indptr": work_indptr, + "sparse_mtp_reduce_indptr": reduce_indptr, + "sparse_mtp_reduce_final_map": reduce_final_map, + "sparse_mtp_reduce_partial_map": reduce_partial_map, + } + def set_mla_persistent_worker_buffers( self, bs: int, @@ -288,14 +388,19 @@ def set_mla_persistent_worker_buffers( reduce_indptr = var["reduce_indptr"] reduce_final_map = var["reduce_final_map"] reduce_partial_map = var["reduce_partial_map"] + # Dense layers use kv_indptr (full KV lengths per seq). + # sparse_kv_indptr is per-token in MTP mode and must NOT be + # indexed with [:bs+1] here — that misinterprets the per-token + # cumsum as per-seq, producing wrong KV lengths and OOB metadata. + kv_indptr_for_metadata = ( + var["sparse_kv_indptr"].gpu[: bs + 1] + if self.is_sparse and max_q_len == 1 + else var["kv_indptr"].gpu[: bs + 1] + ) if only_update: decode_update_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], - ( - var["sparse_kv_indptr"].gpu[: bs + 1] - if self.is_sparse - else var["kv_indptr"].gpu[: bs + 1] - ), + kv_indptr_for_metadata, var["kv_last_page_lens"].gpu[:bs], self.padded_num_attention_heads, 1, # nhead_kv, @@ -316,11 +421,7 @@ def set_mla_persistent_worker_buffers( else: get_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], - ( - var["sparse_kv_indptr"].gpu[: bs + 1] - if self.is_sparse - else var["kv_indptr"].gpu[: bs + 1] - ), + kv_indptr_for_metadata, var["kv_last_page_lens"].gpu[:bs], self.padded_num_attention_heads, 1, # nhead_kv, @@ -356,7 +457,14 @@ def prepare_mtp_decode( var = self.model_runner.forward_vars kv_indptr = var["kv_indptr"].gpu[: bs + 1] if self.is_sparse: - assert False, "TODO: MTP decode is not supported for sparse attention yet" + # Update dense kv_indptr (needed for kv_indices generation and slot_mapping) + kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] + # Recompute sparse_kv_indptr: per-seq sparse count = min(dense_kv_count, index_topk) + sparse_kv_indptr = var["sparse_kv_indptr"].gpu[: bs + 1] + kv_counts = kv_indptr[1 : bs + 1] - kv_indptr[:bs] + sparse_counts = torch.clamp(kv_counts, max=self.index_topk) + sparse_kv_indptr[0] = 0 + sparse_kv_indptr[1 : bs + 1] = torch.cumsum(sparse_counts, dim=0) else: assert self.block_size == 1 kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] @@ -368,9 +476,12 @@ def prepare_mtp_decode( self.block_ratio, max_seqlen_k, ) - return self.set_mla_persistent_worker_buffers( + result = self.set_mla_persistent_worker_buffers( bs, max_seqlen_q, only_update, num_reject_tokens ) + if self.is_sparse: + result["sparse_kv_indptr"] = sparse_kv_indptr + return result def compute_block_bytes(self) -> int: """MLA per-block bytes: single 576-dim packed tensor per layer @@ -378,7 +489,7 @@ def compute_block_bytes(self) -> int: V cache or kv_scale). DeepSeek-V3.2 sparse variant adds an indexer cache contribution - (`hf_config.num_hidden_layers` rows, aligned-to-16 fp8). + for every bound layer, including draft/MTP layers. """ runner = self.model_runner config = runner.config @@ -391,7 +502,7 @@ def compute_block_bytes(self) -> int: index_dim = hf_config.index_head_dim + 4 aligned_index_dim = ((index_dim + 15) // 16) * 16 block_bytes += ( - hf_config.num_hidden_layers + total_num_layers * runner.block_size * aligned_index_dim * dtypes.fp8.itemsize @@ -429,7 +540,7 @@ def allocate_kv_cache_tensors( aligned = ((index_dim + 15) // 16) * 16 out["aligned_index_dim"] = aligned out["index_cache"] = torch.zeros( - hf_config.num_hidden_layers, + total_num_layers, runner.num_physical_kvcache_blocks, runner.physical_block_size, aligned, @@ -489,14 +600,6 @@ def prepare_prefill(self, batch: ScheduledBatch): if attn_metadata.block_tables is None: self.prepare_block_tables(batch) attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs) - if self.block_ratio > 1: - block_table_convert_triton( - var["block_tables"].gpu[:bs], - var["block_tables_converted"].gpu[:bs], - var["context_lens"].gpu[:bs], - self.block_ratio, - ) - attn_metadata.block_tables = var["block_tables_converted"].gpu[:bs] counts = var["cu_seqlens_q"].np[1 : bs + 1] - var["cu_seqlens_q"].np[:bs] if attn_metadata.has_cached: # Full context (cached + new): use cu_seqlens_k for indexer @@ -624,6 +727,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): var["slot_mapping"].np[:sum_scheduled_tokens] = slot_mapping var["positions"].np[:sum_scheduled_tokens] = positions var["context_lens"].np[:scheduled_bs] = context_lens + var["context_lens"].np[scheduled_bs:bs] = 0 if any(batch.is_first_decode_without_local_prefill): num_blocks_per_seq = [ @@ -666,15 +770,39 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): if self.is_sparse: index_topk = self.index_topk - sparse_context_lens = np.clip(var["context_lens"].np[:bs], None, index_topk) - var["sparse_kv_indptr"].np[1 : bs + 1] = np.cumsum( - sparse_context_lens, dtype=np.int32 - ) - var["sparse_kv_indptr"].np[scheduled_bs : bs + 1] = var[ - "sparse_kv_indptr" - ].np[scheduled_bs] - vars_used.append(("sparse_kv_indptr", bs + 1)) - metadata_deps.add("sparse_kv_indptr") + if max_seqlen_q > 1: + # MTP verify: per-token sparse metadata + # Each token at offset j in seq s sees (context_lens[s] - max_seqlen_q + j + 1) KV entries + per_token_kv_lens = ( + np.repeat(context_lens[:scheduled_bs], max_seqlen_q) + - max_seqlen_q + + np.tile( + np.arange(1, max_seqlen_q + 1, dtype=np.int32), scheduled_bs + ) + ) + sparse_per_token_lens = np.clip(per_token_kv_lens, 0, index_topk) + var["sparse_kv_indptr"].np[1 : sum_scheduled_tokens + 1] = np.cumsum( + sparse_per_token_lens, dtype=np.int32 + ) + sum_tokens = bs * max_seqlen_q + var["sparse_kv_indptr"].np[ + sum_scheduled_tokens + 1 : sum_tokens + 1 + ] = var["sparse_kv_indptr"].np[sum_scheduled_tokens] + vars_used.append(("sparse_kv_indptr", sum_tokens + 1)) + vars_used.append(("sparse_cu_seqlens_q", sum_tokens + 1)) + metadata_deps.add("sparse_kv_indptr") + else: + sparse_context_lens = np.clip( + var["context_lens"].np[:bs], None, index_topk + ) + var["sparse_kv_indptr"].np[1 : bs + 1] = np.cumsum( + sparse_context_lens, dtype=np.int32 + ) + var["sparse_kv_indptr"].np[scheduled_bs : bs + 1] = var[ + "sparse_kv_indptr" + ].np[scheduled_bs] + vars_used.append(("sparse_kv_indptr", bs + 1)) + metadata_deps.add("sparse_kv_indptr") prep_stream = self.prep_stream vars_for_metadata = [(el, num) for el, num in vars_used if el in metadata_deps] @@ -698,21 +826,22 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): max_seqlen_k, ) + is_sparse_mtp = self.is_sparse and max_seqlen_q > 1 # metadata copies on main_stream positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) ctx.update({el: var[el].copy_to_gpu(num) for el, num in vars_for_metadata}) - ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + + if is_sparse_mtp: + sum_tokens = bs * max_seqlen_q + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + ctx_mla_ps_sparse = self._set_mla_persistent_worker_buffers_sparse_mtp( + sum_tokens + ) + else: + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q) + ctx_mla_ps_sparse = None ctx.update(ctx_mla_ps) current_stream.wait_stream(prep_stream) - # if self.block_ratio > 1: - # if "block_tables" in ctx: - # block_table_convert_triton( - # var["block_tables"].gpu[:bs], - # var["block_tables_converted"].gpu[:bs], - # var["context_lens"].gpu[:bs], - # self.block_ratio, - # ) - # ctx["block_tables_converted"] = var["block_tables_converted"].gpu[:bs] attn_metadata = AttentionMetaData( dropout_p=dropout_p, max_seqlen_q=max_seqlen_q, @@ -721,6 +850,24 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ) attn_metadata.dtype_q = self.dtype_q + if ctx_mla_ps_sparse is not None: + for k, v in ctx_mla_ps_sparse.items(): + setattr(attn_metadata, k, v) + + if is_sparse_mtp: + sum_tokens = bs * max_seqlen_q + attn_metadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[ + : sum_tokens + 1 + ] + attn_metadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:sum_tokens] + self._token_to_seq_idxs_gpu[:sum_scheduled_tokens] = torch.arange( + scheduled_bs, dtype=torch.int32, device=self.device + ).repeat_interleave(max_seqlen_q) + self._token_to_seq_idxs_gpu[sum_scheduled_tokens:sum_tokens] = 0 + attn_metadata.token_to_seq_idxs = self._token_to_seq_idxs_gpu[:sum_tokens] + # Use bs (graph_bs) >= 2 instead of scheduled_bs >= 2 to avoid accuracy issue: if self.model_runner.config.enable_tbo_decode and bs >= 2: self._prepare_ubatch_decode( @@ -884,9 +1031,19 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: var = self.model_runner.forward_vars sparse_kv_indptr = var["sparse_kv_indptr"].gpu if self.is_sparse else None max_q_len = var["mtp_k"] + 1 if "mtp_k" in var else 1 - ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + sum_tokens = bs * max_q_len + is_sparse_mtp = self.is_sparse and max_q_len > 1 + if is_sparse_mtp: + # Two sets: normal for dense layers, sparse_mtp for sparse layers + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + ctx_mla_ps_sparse = self._set_mla_persistent_worker_buffers_sparse_mtp( + sum_tokens + ) + else: + ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_q_len) + ctx_mla_ps_sparse = None attn_matadata = AttentionMetaData( - slot_mapping=var["slot_mapping"].gpu[: bs * max_q_len], + slot_mapping=var["slot_mapping"].gpu[:sum_tokens], context_lens=var["context_lens"].gpu[:bs], block_tables=var["block_tables"].gpu[:bs], max_seqlen_q=max_q_len, @@ -895,15 +1052,27 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: kv_indices=var["kv_indices"].gpu, kv_last_page_lens=var["kv_last_page_lens"].gpu[:bs], sparse_kv_indptr=sparse_kv_indptr, - block_tables_converted=( - var["block_tables_converted"].gpu[:bs] - if "block_tables_converted" in var - else None - ), **ctx_mla_ps, ) attn_matadata.dtype_q = self.dtype_q - positions = var["positions"].copy_to_gpu(bs * max_q_len) + if ctx_mla_ps_sparse is not None: + for k, v in ctx_mla_ps_sparse.items(): + setattr(attn_matadata, k, v) + if is_sparse_mtp: + attn_matadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[ + : sum_tokens + 1 + ] + attn_matadata.sparse_kv_indptr = var["sparse_kv_indptr"].gpu[ + : sum_tokens + 1 + ] + attn_matadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:sum_tokens] + self._token_to_seq_idxs_gpu[:sum_tokens] = torch.arange( + bs, dtype=torch.int32, device=self.device + ).repeat_interleave(max_q_len) + attn_matadata.token_to_seq_idxs = self._token_to_seq_idxs_gpu[:sum_tokens] + positions = var["positions"].copy_to_gpu(sum_tokens) context = Context( positions=positions, is_prefill=False, batch_size=bs, graph_bs=bs ) diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index 924608403d..ec7abb2139 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -8,6 +8,7 @@ from atom.config import Config, QuantizationConfig from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ReplicatedLinear from atom.model_ops.moe import FusedMoE from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled from atom.models.utils import IntermediateTensors @@ -40,7 +41,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class DeepSeekMultiTokenPredictorLayer(nn.Module): - def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: + def __init__( + self, + atom_config: Config, + prefix: str, + layer_idx: int, + topk_indices_buffer: Optional[torch.Tensor] = None, + ) -> None: super().__init__() config = atom_config.hf_config @@ -48,7 +55,13 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.eh_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + quant_config=atom_config.quant_config, + prefix=maybe_prefix(prefix, "eh_proj"), + ) self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=atom_config.quant_config @@ -63,6 +76,7 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: quant_config=quant_config, layer_num=layer_idx, is_mtp_block=True, + topk_indices_buffer=topk_indices_buffer, ) def forward( @@ -95,7 +109,13 @@ def forward( class DeepSeekMultiTokenPredictor(nn.Module): - def __init__(self, *, atom_config: Config, prefix: str = ""): + def __init__( + self, + *, + atom_config: Config, + prefix: str = "", + topk_indices_buffer: Optional[torch.Tensor] = None, + ): super().__init__() config = atom_config.hf_config self.mtp_start_layer_idx = config.num_hidden_layers @@ -104,7 +124,10 @@ def __init__(self, *, atom_config: Config, prefix: str = ""): self.layers = torch.nn.ModuleDict( { str(idx): DeepSeekMultiTokenPredictorLayer( - atom_config, f"{prefix}.layers.{idx}", layer_idx=idx + atom_config, + f"{prefix}.layers.{idx}", + layer_idx=idx, + topk_indices_buffer=topk_indices_buffer, ) for idx in range( self.mtp_start_layer_idx, @@ -167,8 +190,20 @@ def __init__(self, atom_config: Config, prefix: str = ""): "up_proj": ("gate_up_proj", 1), } + if hasattr(self.config, "index_topk"): + topk_indices_buffer = torch.empty( + atom_config.max_num_batched_tokens, + self.config.index_topk, + dtype=torch.int32, + device="cuda", + ) + else: + topk_indices_buffer = None + self.model = DeepSeekMultiTokenPredictor( - atom_config=atom_config, prefix=maybe_prefix(prefix, "model") + atom_config=atom_config, + prefix=maybe_prefix(prefix, "model"), + topk_indices_buffer=topk_indices_buffer, ) def remap_mtp_weight_name(self, name: str) -> str | None: diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index c9ad8e53f4..c0daa899a8 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -980,13 +980,19 @@ def sparse_attn_indexer( if forward_context.context.is_dummy_run: # dummy runner return weights - num_decode_tokens = context.batch_size if not context.is_prefill else 0 + # For MTP verify decode, max_seqlen_q > 1 so total decode tokens = batch_size * max_seqlen_q + num_decode_tokens = ( + context.batch_size * attn_metadata.max_seqlen_q if not context.is_prefill else 0 + ) + runner_block_size = get_current_atom_config().kv_cache_block_size + kv_cache = kv_cache.view(-1, runner_block_size, kv_cache.shape[-1]) indexer_k_quant_and_cache( k, kv_cache, slot_mapping, quant_block_size, scale_fmt, + preshuffle=True, ) if context.is_prefill: if attn_metadata.max_seqlen_k <= topk_indices_buffer.shape[1]: @@ -1018,6 +1024,7 @@ def sparse_attn_indexer( if prefill_metadata.has_cached else prefill_metadata.cu_seqlens_q ), + preshuffle=True, ) cu_seqlen_ks = prefill_metadata.cu_seqlen_ks cu_seqlen_ke = prefill_metadata.cu_seqlen_ke @@ -1069,6 +1076,8 @@ def sparse_attn_indexer( decode_metadata.context_lens, attn_metadata.block_tables, max_model_len, + KVBlockSize=runner_block_size, + Preshuffle=True, ) num_rows = logits.shape[0] assert topk_tokens == 2048, "top_k_per_row assumes size 2048" @@ -1161,7 +1170,7 @@ def __init__( self.weights_proj = ReplicatedLinear( hidden_size, self.n_head, - quant_config=quant_config, + quant_config=None, prefix=f"{prefix}.weights_proj", ) self.softmax_scale = self.head_dim**-0.5 diff --git a/atom/models/mimo_v2_flash_mtp.py b/atom/models/mimo_v2_flash_mtp.py index 1da461eca5..bb61741955 100644 --- a/atom/models/mimo_v2_flash_mtp.py +++ b/atom/models/mimo_v2_flash_mtp.py @@ -10,6 +10,7 @@ from atom.config import Config from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ReplicatedLinear from atom.models.utils import IntermediateTensors, maybe_prefix from atom.utils.decorators import support_torch_compile @@ -113,7 +114,13 @@ def __init__(self, atom_config: Config, prefix: str, layer_idx: int) -> None: self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) - self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.eh_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + quant_config=atom_config.quant_config, + prefix=maybe_prefix(prefix, "eh_proj"), + ) self.mtp_block = MiMoV2FlashMTPLayer( atom_config=atom_config, diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index aa3aa4c8b1..b06ab8fded 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -207,10 +207,14 @@ def propose( if use_mla: kv_last_page_lens = var["kv_last_page_lens"].gpu[:bs] attn_metadata.kv_last_page_lens = kv_last_page_lens - else: - # MHA needs block_tables and context_lens - attn_metadata.block_tables = var["block_tables"].gpu[:bs] - attn_metadata.context_lens = var["context_lens"].gpu[:bs] + # block_tables, context_lens, and sparse_kv_indptr are + # needed by both MHA and MLA+sparse attention + attn_metadata.block_tables = var["block_tables"].gpu[:bs] + attn_metadata.context_lens = var["context_lens"].gpu[:bs] + if "sparse_kv_indptr" in var: + attn_metadata.sparse_kv_indptr = var[ + "sparse_kv_indptr" + ].gpu[: bs + 1] cu_seqlens_q[: bs + 1] = self.arrange_bs[: bs + 1] if use_mla: # MLA: block_size=1, kv_indptr tracks tokens @@ -222,9 +226,9 @@ def propose( # update metadata attn_metadata.max_seqlen_k += 1 - if not use_mla: - # MHA: update context_lens for this draft step - attn_metadata.context_lens[:bs] += 1 + # Update context_lens for each draft step (needed by both + # MHA attention and MLA+sparse indexer) + attn_metadata.context_lens[:bs] += 1 workinfos = self.runner.attn_metadata_builder.prepare_mtp_decode( bs, ( From 247e9b1331a9062b1fe19fdd18068b2823135106 Mon Sep 17 00:00:00 2001 From: wuhuikx Date: Fri, 8 May 2026 23:09:33 +0800 Subject: [PATCH 04/76] Add the benchmark flow for ATOM vLLM plugin (#514) * Add the benchmark flow for OOT * Update the benchmark * Add attention-backend * Update the vLLM-ATOM GPT-OSS recipe * Change the recipe of kimi-k2 * revert the changes in recipe * Update atom-vllm-benchmark-guide.md --- .claude/commands/atom-vllm-benchmark-guide.md | 283 ++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 .claude/commands/atom-vllm-benchmark-guide.md diff --git a/.claude/commands/atom-vllm-benchmark-guide.md b/.claude/commands/atom-vllm-benchmark-guide.md new file mode 100644 index 0000000000..a28562f4b8 --- /dev/null +++ b/.claude/commands/atom-vllm-benchmark-guide.md @@ -0,0 +1,283 @@ +# ATOM vLLM Benchmark Guide + +Use this guide when the user asks to benchmark ATOM vLLM Plugin performance, compare with the upstream vLLM (optional), or validate a performance claim. + +## What this guide should do + +1. Confirm benchmark goal and variables (throughput, TTFT, TPOT, E2EL, etc.). +2. Use plugin-on as default setup; optionally benchmark upstream vLLM for performance comparison. +3. Run repeatable measurements and summarize median results. +4. Report reproducible commands, key metrics, and known risks. +5. Prefer model-specific launch settings from `recipes/atom_vllm/` when available; otherwise use the standard workflow in this guide. + +## Inputs To Collect First + +Before running any benchmark, confirm: + +- Model path (HF id or local path) +- Serving backend (`vllm`) +- Hardware shape (GPU type/count) +- `tensor_parallel_size`, KV cache dtype, and scheduler flags +- Request mix (ISL/OSL/concurrency/request rate) +- Comparison target: + - default candidate (plugin on) + - optional baseline/control (plugin off) + - optional runtime baseline (upstream vLLM nightly) + +If any of these are missing, ask before execution. + +## Environment Checklist + +For each benchmark round: + +- Kill old server processes and verify VRAM is released. +- Clear cache when startup behavior is inconsistent: + - `rm -rf /root/.cache/atom/*` + - `rm -rf /root/.cache/vllm/*` +- Set stable logging: + - `export AITER_LOG_LEVEL=WARNING` +- Confirm server is actually ready: + - `curl -sf http://localhost:8000/v1/models` + - `rocm-smi --showmemuse` and verify allocated memory is greater than 0. + +## Benchmark Workflow (Plugin On by Default) + +Run plugin-on by default. + +### Recipe-first execution policy + +When running ATOM vLLM benchmark: + +1. First check whether the target model has a matching recipe under `recipes/atom_vllm/`. +2. If a matching recipe exists, use its model-specific environment variables and `vllm serve` launch args as the default benchmark setup. +3. If no matching recipe exists, follow the standard launch template and workflow in this guide. + +Keep benchmark methodology unchanged (smoke check, concurrency isolation, and fixed comparison variables) regardless of whether the launch config comes from a recipe or from the default template. + +- **Default (candidate):** plugin enabled +- **Optional control (baseline):** plugin disabled (only when A/B comparison is requested) +- **Default ATOM vLLM Plugin image:** `rocm/atom-dev:vllm-latest` + +If both conditions are run, only one variable is allowed to change between them. + +### 1) Launch server + +Use the same launch command template. Keep plugin-on as default, and only change plugin toggle env vars when running the optional plugin-off control. + +```bash +vllm serve \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size \ + --attention-backend ROCM_AITER_FA \ + --trust-remote-code \ + --kv-cache-dtype fp8 \ + --gpu_memory_utilization 0.9 \ + --async-scheduling \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --no-enable-prefix-caching +``` + +Common plugin toggles: + +- Default plugin-on: do not set `ATOM_DISABLE_VLLM_PLUGIN` or `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION` (or set them to `0`) +- Optional plugin-off control: set `ATOM_DISABLE_VLLM_PLUGIN=1` +- Optional attention-off control: set `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1` +- Always set `VLLM_ROCM_USE_AITER=1` in both ATOM plugin and upstream vLLM runs. + +### 2) Run smoke validation + +```bash +curl -sf http://localhost:8000/v1/models +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model \ + --dataset-name random \ + --random-input-len 32 \ + --random-output-len 32 \ + --temperature 0.0 \ + --max-concurrency \ + --num-prompts \ + --num-warmups \ + --request-rate 1 \ + --disable-tqdm +``` + +Use `smoke_conc_x10 = smoke_conc * 10` and `smoke_conc_x2 = smoke_conc * 2`. + +If smoke test fails, do not continue with performance comparison. + +### 3) Run benchmark workload + +```bash +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model= \ + --dataset-name=random \ + --random-input-len= \ + --random-output-len= \ + --random-range-ratio=0.8 \ + --temperature=0.0 \ + --num-prompts= \ + --num-warmups= \ + --max-concurrency= \ + --trust-remote-code \ + --request-rate=inf \ + --ignore-eos \ + --disable-tqdm \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" +``` + +### 4) Concurrency sweep isolation (required) + +When sweeping multiple concurrency points, each point must run in a fresh container lifecycle: + +1. Start a fresh container. +2. Start vLLM server in that container. +3. Run smoke + benchmark for one `--max-concurrency=`. +4. Stop vLLM server. +5. Exit and stop/remove the container. +6. Start a new container for the next concurrency point. + +Do not reuse the same running server/container across different concurrency points. + +Example pattern: + +```bash +for conc in 4 8 16 32 64 128; do + docker run -dt --name "atom_bench_c${conc}" + docker exec "atom_bench_c${conc}" bash -lc 'vllm serve ...' + docker exec "atom_bench_c${conc}" bash -lc "vllm bench serve ... --max-concurrency=${conc} ..." + docker exec "atom_bench_c${conc}" bash -lc 'pkill -f "vllm serve" || true' + docker stop "atom_bench_c${conc}" && docker rm "atom_bench_c${conc}" +done +``` + +### 5) Minimum benchmark matrix + +| Scenario | ISL | OSL | Concurrency | +|---|---:|---:|---:| +| Short in / short out | 1024 | 1024 | 4, 8, 16, 32, 64, 128 | +| Short in / long out | 1024 | 8192 | 4, 8, 16, 32, 64, 128 | +| Long in / short out | 8192 | 1024 | 4, 8, 16, 32, 64, 128 | + +### 6) Repetition and statistics + +- Run each scenario at least 1 times. +- Use median values for comparison. +- If variance is large, increase repeat count. + +## Optional Upstream vLLM Benchmark + +Use this optional path when you need ATOM vs upstream runtime comparison. + +Reference images: + +- ATOM vLLM plugin default image: `rocm/atom-dev:vllm-lastest` +- Upstream vLLM comparison image: `vllm/vllm-openai-rocm:nightly` + +### 1) Pull latest upstream ROCm image + +```bash +docker pull vllm/vllm-openai-rocm:nightly +``` + +### 2) Run the same workflow and config + +For upstream runs, keep everything identical to ATOM runs: + +- Same model path +- Same GPU set and `tensor_parallel_size` +- Same KV cache dtype and scheduler flags +- Same request mix (ISL/OSL/concurrency/request rate) +- Same concurrency sweep isolation policy (restart container per concurrency point) + +Only runtime image is allowed to change. + +### 3) Upstream execution note + +- Use the same smoke and benchmark commands from this guide. +- Do not set ATOM plugin toggle env vars for upstream runs. +- Always set `VLLM_ROCM_USE_AITER=1` for upstream runs. +- Keep result naming explicit (for example, include `upstream-nightly` in filenames). + +## Regression Profiling Workflow + +When regression is observed: + +```bash +# 1) Start (or restart) server with profiler enabled +vllm serve \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size \ + --attention-backend ROCM_AITER_FA \ + --trust-remote-code \ + --kv-cache-dtype fp8 \ + --profiler-config '' + +# 2) Run profiled benchmark workload +vllm bench serve \ + --backend vllm \ + --base-url http://localhost:8000 \ + --endpoint /v1/completions \ + --model \ + --dataset-name random \ + --random-input-len \ + --random-output-len 20 \ + --temperature 0.0 \ + --max-concurrency \ + --num-prompts \ + --num-warmups \ + --request-rate inf \ + --profile \ + --save-result +``` + +- `--profile` requires server-side `--profiler-config`. +- Keep profiler config identical between plugin-off and plugin-on runs. +- For reproducibility, keep benchmark client semantics fixed to completion mode: + `--backend vllm --endpoint /v1/completions` (this path uses `max_tokens`). + +## Reporting Template + +Use this format in final response: + +```markdown +### ATOM vLLM Benchmark Result +- Goal: +- Setup: + - Model: + - Hardware: + - Fixed args: + - Isolation: +- Validation: +- Benchmark: + - Default (plugin on): + - Optional control (plugin off): + - Optional runtime baseline (upstream nightly): + - Delta: +- Confidence: + - Runs: + - Variance: +- Risk: + - +- Repro: + - + - + - +``` + +## Guardrails + +- Never compare runs with different model, TP, KV dtype, or request mix. +- For concurrency sweep, stop server and restart container between concurrency points. +- For ATOM vs upstream comparison, only runtime image may differ. +- Keep `VLLM_ROCM_USE_AITER=1` fixed for both ATOM and upstream runs. +- If GPU memory is not allocated, treat benchmark results as invalid. +- If plugin-on and plugin-off outputs differ semantically, run accuracy checks before claiming performance gains. From e87b2f30e3c2e1a08e21ccf19af0be61802069fe Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Sat, 9 May 2026 00:58:14 +0800 Subject: [PATCH 05/76] perf(deepseek_v4): fused_compress kernel optimization + DualRMSNorm fusion (#723) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(deepseek_v4): fused_compress kernel + DualRMSNorm fusion + decode buffer cleanup Fused-compress kernel (atom/model_ops/v4_kernels/fused_compress.py) - Split K-loop into two phases (state-only `[0, window_len)` then input-only `[window_len, K)`): on AMD CDNA masked tl.load still issues the LD instruction, predicate only suppresses the register write. Issuing only the live side's loads cuts HBM traffic ~40% in the bandwidth-bound regime. Microbench (median per launch, BATCH=50, 200 iters): HCA (ratio=128 K=128): 35.07us → 28.76us (-18%) CSA (ratio=4 K=8 ): unchanged (launch-overhead floor at K=8) - Padding invariant verified: window_len = K - min(j_in_seq+1, K) bounds K-1-position, so padding (`s < 0`) lies entirely in the state phase — input phase needs no padding mask. - Eviction hints: ragged kv_in/score_in marked evict_first (single-use per program), ape evict_last (small + reused). - FP8 quant fusion path: `tl.clamp + plain .to(fp8)` (aiter style; avoids the slow `fp_downcast_rounding="rtne"` path on AMD that bypasses `v_cvt_pk_fp8_f32`), with UE8M0 scale + MFMA 16x16 preshuffle + .cs streaming stores. - Bit-exact match vs `fused_compress_attn_reference` (verified at BF16 precision; ≤1 BF16 ULP due to `tl.exp` HW vs libm). DualRMSNorm fusion (atom/model_ops/layernorm.py, atom/models/deepseek_v4.py) - q_norm2 + kv_norm (per-head Q + KV, both head_dim=128) routed through existing `DualRMSNorm` + `_fused_qk_norm_single_kernel`. - q_norm2 carries no learnable weight — added `_make_weightless_rmsnorm` factory (`del weight; weight = None`) so the parameter is absent from `state_dict` (no loader warning), and a `Q_HAS_WEIGHT` constexpr in the fused kernel skips the load when the weight is None. - `DualRMSNorm._eps` resolved once with explicit None-check fallback (handles both `variance_epsilon` and `eps` attribute names). Compressor refactor (atom/models/deepseek_v4.py) - Side-effecting `forward()` returns None — the prior caller-visible BF16 return was vestigial (paged_decode/paged_prefill read scattered entries directly from `unified_kv` / FP8 indexer pool). - `cache_scale` strided fp32 view binding for the Indexer-inner Compressor (FP8 scale region of the same allocation). - Auto-detects quant via `kv_cache.dtype != bfloat16`. CompressPlan decode capacity (atom/model_ops/v4_kernels/compress_plan.py) - New `decode_capacity_per_ratio` arg: when supplied, the returned `compress_plan_gpu` slice has fixed length = decode-tight bound (`max_decode_tokens // ratio + max_bs`) instead of prefill worst case (~13× larger), with sentinel-fill of trailing rows so the captured kernel grid is decode-sized but address-stable. - Empty-fwd path now also produces CompressPlans pointing at the pre-allocated buffers (sentinel-filled), so capture-time and replay-time addresses match even on a zero-token fwd. V4 attention builder (atom/model_ops/attentions/deepseek_v4_attn.py) - `_decode_compress_cap[ratio]` plumbing for CG decode path. - Indexer-inner Compressor `cache_scale` view bound from `runner.v4_csa_idx_kv` per-layer. - Removed `_build_indexer_compress_slot_mapping` and `compress_slot_mapping_gpu` (Indexer-inner now uses block_tables directly). - Dropped `v4_indexer_decode_logits` and `v4_indexer_decode_topk_indices` from the metadata pool — these are write-once GPU scratch with no CPU mirror; allocated per-fwd via `torch.empty` in `Indexer._score_topk_decode`. Under CG capture the allocations land in the graph's private pool and replay reuses the same address (saves ~2 MiB pinned host + ~2 MiB GPU on the prior `CpuGpuBuffer` overhead). mark_trace typing (atom/utils/decorators.py) - `@overload` + ParamSpec: pyright/pylance no longer flags DualRMSNorm-style decorated callables as "not callable". Triton MoE block_m default (atom/model_ops/fused_moe_triton.py) - `ATOM_TRITON_MOE_BLOCK_M` default 64 → 32 (better MI355X tile occupancy at typical MoE shapes). GSM8K nshot=5 (DeepSeek-V4-Pro, --level 0, ATOM_USE_TRITON_MOE=1): flexible-extract 0.9522 ±0.0059 / strict-match 0.9530 ±0.0058 (baseline 0.953/0.954 — within 1σ, no regression) * test remove toch compile * remove level0 --- atom/model_ops/attentions/deepseek_v4_attn.py | 236 +++++------- atom/model_ops/fused_moe_triton.py | 2 +- atom/model_ops/layernorm.py | 65 +++- atom/model_ops/v4_kernels/compress_plan.py | 56 ++- atom/model_ops/v4_kernels/fused_compress.py | 359 ++++++++++++------ atom/models/deepseek_v4.py | 219 ++++++----- atom/utils/decorators.py | 23 +- 7 files changed, 567 insertions(+), 393 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 80a668748e..6b4f6bfc08 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -147,18 +147,17 @@ class AttentionMetaData_DSV4(AttentionMetaData): indexer_meta: Optional[Dict[str, Any]] = None """dict — `Indexer.forward_batched` per-fwd GPU tensors. Notable keys: cu_committed_gpu [bs+1] int32 per-seq committed cumsum - compress_slot_mapping_gpu [num_compress] int64 FP8 cache write slots seq_base_per_token_gpu [T] int32 prefill subtract base (also aliased as cu_starts_gpu for fp8_mqa_logits) cu_ends_gpu [T] int32 per-token end offset for fp8_mqa_logits (causal cap) - decode_logits_gpu [max_bs, ...] pre-allocated decode buffer - decode_topk_indices_gpu [max_bs, index_topk] int32 decode topk - output buffer (raw seq-local - with -1 sentinels in tail) total_committed int sum of n_committed_csa_per_seq + Note: decode logits / topk-indices scratch are allocated per-fwd inside + `Indexer._score_topk_decode` (write-once, no CPU mirror, CG-stable via + the captured graph's private memory pool). + The indexer's downstream contract: `_score_topk_*` returns RAW seq-local `[T, index_topk] int32` with kernel-native -1 in tail cols (cells past the per-token visibility cap). `csa_translate_pack` consumes this layout @@ -618,22 +617,38 @@ def build_kv_cache_tensor(self, layer_id: int, module): is_indexer_inner = "indexer" in parts ratio = module.compress_ratio - # Per-kind shared compress output buffer (CUDAGraph: stable - # data pointer + fixed shape across captures of the same kind). - # Read from forward_vars (allocated in _alloc_v4_metadata_buffers). - compress_out_buffers = self.model_runner.forward_vars if is_indexer_inner: assert ratio == 4, "Indexer-inner Compressor only on CSA layers" pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_idx_kv_state[pos] module.score_state = runner.v4_csa_idx_score_state[pos] # Inner compressor writes target the SAME storage as the - # outer Indexer.kv_cache (csa_idx_kv). Same 3D shape — write - # via slot_mapping is shape-agnostic (flat indexing), but we - # keep [NB, k1_csa, aligned_dim] for symmetry with the read - # binding above. - module.kv_cache = runner.v4_csa_idx_kv[pos] - module.compress_out = compress_out_buffers["v4_csa_idx_compress_out"] + # outer Indexer.kv_cache (csa_idx_kv). Same [NB, k1_csa, + # aligned_dim] FP8 shape — `Compressor.forward` resolves + # slot via block_table+ci internally (no flat slot_mapping + # needed; matches CSA Main's path). + idx_kv = runner.v4_csa_idx_kv[pos] + module.kv_cache = idx_kv + # FP8 quant path: bind a strided fp32 view of the per-block + # scale region. Layout per block: [k1*head_dim FP8 region] + # then [k1 fp32 scale region] then padding (cache_kernels.cu + # :1209-1239). Strides expressed in fp32 elements. + nb, k1, aligned_dim = idx_kv.shape + head_dim = self.index_head_dim + assert ( + k1 * aligned_dim + ) % 4 == 0, f"per-block bytes ({k1 * aligned_dim}) must be 4-aligned" + block_fp32_stride = (k1 * aligned_dim) // 4 + scale_fp32_offset = (k1 * head_dim) // 4 + module.cache_scale = ( + idx_kv.view(torch.float32) + .view(-1) + .as_strided( + size=(nb, k1), + stride=(block_fp32_stride, 1), + storage_offset=scale_fp32_offset, + ) + ) elif ratio == 4: pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_main_kv_state[pos] @@ -647,7 +662,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k1_csa, self.head_dim ) - module.compress_out = compress_out_buffers["v4_csa_main_compress_out"] elif ratio == 128: pos = self.layer_id_to_hca_pos[layer_id_from_prefix] module.kv_state = runner.v4_hca_main_kv_state[pos] @@ -657,7 +671,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.kv_cache = unified[swa_pages:].view( num_blocks, self.k2_hca, self.head_dim ) - module.compress_out = compress_out_buffers["v4_hca_main_compress_out"] else: raise ValueError( f"Unknown V4 compress_ratio={ratio} on Compressor at " @@ -685,18 +698,12 @@ def _attach_v4_indexer_meta( """Build and attach the CSA Indexer per-fwd GPU metadata. Hoists per-CSA-layer H2D calls (batch_id_per_token / cu_committed / - n_committed / seq_base_per_token / cu_ends / compress_slot_mapping) - into a single per-fwd build. compress_plans[4] (CSA, ratio=4) carries - the compress rows that the indexer-FP8 path needs to derive its - write-side slot_mapping. None for warmup or empty fwd; - `_build_v4_indexer_meta` handles both. + n_committed / seq_base_per_token / cu_ends) into a single per-fwd + build. None for warmup or empty fwd; `_build_v4_indexer_meta` + handles both. """ import numpy as np - csa_compress_plan_cpu = None - plans = getattr(attn_metadata, "compress_plans", None) or {} - if 4 in plans: - csa_compress_plan_cpu = plans[4].compress_plan_cpu attn_metadata.indexer_meta = self._build_v4_indexer_meta( attn_metadata=attn_metadata, cu_seqlens_q_np=cu_seqlens_q_np, @@ -705,7 +712,6 @@ def _attach_v4_indexer_meta( scheduled_bs=scheduled_bs, total_tokens=total_tokens, device=self.device, - csa_compress_plan_cpu=csa_compress_plan_cpu, ) def _build_v4_indexer_meta( @@ -718,7 +724,6 @@ def _build_v4_indexer_meta( scheduled_bs: int, total_tokens: int, device, - csa_compress_plan_cpu, ): """Build per-fwd GPU index tensors consumed by `Indexer.forward_batched`. @@ -734,21 +739,16 @@ def _build_v4_indexer_meta( (int64) and `v4_indexer_n_committed_per_seq` (int64) — one extra H2D each per fwd. Now we read int32 → cast to int64 on GPU. - The FP8-cache write side (`indexer_k_quant_and_cache`) needs a flat - `compress_slot_mapping` int64 tensor — one entry per row in - `csa_compress_plan_cpu` mapping the compress entry to a global slot - in the `[n_csa, NB, k1, aligned_dim=144]` cache pool (layer-major, - FP8 + 4-byte fp32 scale per row, 16B-aligned). Computed here host-side - because the plan rows + block_tables_cpu are both already on host. + The FP8 indexer K-cache write happens inside `fused_compress_attn` + (the unified Indexer-inner Compressor path) via the same block_tables + that CSA Main uses; no separate slot_mapping is built here. """ - from atom.models.deepseek_v4 import _V4_BLOCK_SIZE # Caller contract: scheduled_bs >= 1, total_tokens >= 1 (same # invariants as `_attach_v4_per_fwd_meta` — guaranteed by every # prepare_*/CG-capture path). bs = scheduled_bs ratio = 4 # CSA - k_per_block = _V4_BLOCK_SIZE // ratio # 32 cu_seqlens_q_arr = cu_seqlens_q_np[: bs + 1].astype(np.int64) token_num_per_seq = (cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:bs]).astype( np.int64 @@ -779,14 +779,6 @@ def _build_v4_indexer_meta( cu_committed_cpu[-1] = max(int(cu_committed_cpu[-1]), 1) total_committed = int(cu_committed_cpu[-1]) - # FP8 write-side slot_mapping (independent of `total_committed` — - # written even when there's nothing to read because num_compress can - # be > 0 on the same fwd that crosses a compress boundary for the - # FIRST committed entry of a fresh seq). - compress_slot_mapping_gpu = self._build_indexer_compress_slot_mapping( - csa_compress_plan_cpu, scheduled_bs, k_per_block, ratio - ) - # batch_id_per_token + n_committed_csa: reuse the shared GPU # tensors set in `_attach_v4_per_fwd_meta` (which MUST run before # this helper — see prepare_decode/prefill ordering). int64 @@ -822,23 +814,10 @@ def _build_v4_indexer_meta( seq_base_per_token_gpu + visible_end_gpu ) # [total_tokens] int32 — fp8_mqa_logits per-token end offset - # Pre-allocated decode-path buffers (full [max_bs, ...] views). Decode - # helper slices each to [:total_tokens]. Returned full because - # `_build_v4_indexer_meta` is called for both prefill and decode - # batches; prefill's `total_tokens` may exceed `max_bs` so builder - # can't pre-slice. Prefill path doesn't read these. - decode_logits_gpu = self.model_runner.forward_vars[ - "v4_indexer_decode_logits" - ].gpu # [max_bs, max_model_len_idx] fp32 - decode_topk_indices_gpu = self.model_runner.forward_vars[ - "v4_indexer_decode_topk_indices" - ].gpu # [max_bs, index_topk] int32 - return { "total_committed": total_committed, "cu_committed_gpu": cu_committed_gpu, "n_committed_per_seq_gpu": n_committed_per_seq_gpu, # int32, [bs] - "compress_slot_mapping_gpu": compress_slot_mapping_gpu, "batch_id_per_token_gpu": batch_id_per_token_gpu, # int64, [total_tokens] # Prefill-only fields below — decode never consults them. NOT # in pre-allocated buffers (per-fwd derived); CG capture path @@ -847,46 +826,8 @@ def _build_v4_indexer_meta( "seq_base_per_token_gpu": seq_base_per_token_gpu, "cu_starts_gpu": seq_base_per_token_gpu, # alias for fp8_mqa_logits "cu_ends_gpu": cu_ends_gpu, - "decode_logits_gpu": decode_logits_gpu, - "decode_topk_indices_gpu": decode_topk_indices_gpu, } - def _build_indexer_compress_slot_mapping( - self, - csa_compress_plan_cpu, - scheduled_bs: int, - k_per_block: int, - ratio: int, - ): - """Compute the per-compress-row flat slot in `v4_csa_idx_kv` pool. - - For each row in `csa_compress_plan_cpu` (= `(ragged_id, batch_id, - position, window_len)`): - ci = position // ratio # compress entry idx in seq - block_in_seq = ci // k_per_block - slot_in_block = ci % k_per_block - physical_block = block_tables_cpu[batch_id, block_in_seq] - slot = physical_block * k_per_block + slot_in_block - - Returns None when the plan is empty (no boundary crossed this fwd) — - the caller skips the `indexer_k_quant_and_cache` write entirely. - """ - if csa_compress_plan_cpu is None or csa_compress_plan_cpu.shape[0] == 0: - return None - var = self.model_runner.forward_vars - block_tables_np = var["block_tables"].np[:scheduled_bs] - bid = csa_compress_plan_cpu[:, 1] - pos = csa_compress_plan_cpu[:, 2] - ci = pos // ratio - block_in_seq = ci // k_per_block - slot_in_block = ci % k_per_block - physical_block = block_tables_np[bid, block_in_seq] - # int64 — `indexer_k_quant_and_cache` kernel signature is `int64_t*` - # (matches V3.2's `attn_metadata.slot_mapping` dtype). int32 caused - # GPU memory access faults from 2x stride mis-stepping. - slot_mapping_np = physical_block.astype(np.int64) * k_per_block + slot_in_block - return self._stage("v4_indexer_compress_slot_mapping", slot_mapping_np) - def prepare_decode(self, batch: ScheduledBatch, bs: int): """V4-style decode prep: populates positions, cu_seqlens_q, block_tables, and state_slot_mapping. @@ -970,7 +911,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): # plan.compress_plan_cpu to derive its FP8 write-side slot_mapping. extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device + extend_lens_np, context_lens_np, positions.device, for_decode_cg=True ) # CG capture path: model_runner pads cu_seqlens_q to graph_bs. The # captured kernels still iterate `bs * max_q_len` tokens at replay, @@ -1053,7 +994,7 @@ def prepare_prefill(self, batch: ScheduledBatch): np.int32 ) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device + extend_lens_np, context_lens_np, positions.device, for_decode_cg=False ) # Prefill goes through eager (no CG): defaults make padded_total_tokens # collapse to total_tokens — no padding logic kicks in. Must still run @@ -1721,14 +1662,22 @@ def _build_paged_prefill_meta( ).to(device) attn_metadata.swa_pages = swa_pages - def _build_compress_plans(self, extend_lens_np, seq_lens_np, device): + def _build_compress_plans( + self, extend_lens_np, seq_lens_np, device, *, for_decode_cg: bool + ): """Build per-ratio CompressPlan dict consumed by batched compressor. - Reuse this from both prepare_decode and prepare_prefill — caller - supplies extend_lens / seq_lens (np int32) and target device. Plan - tensors are written into the pre-allocated `v4_compress_plan_{ratio}` - / `v4_write_plan_{ratio}` CpuGpuBuffers (fixed pointers for - CUDAGraph capture); the kernels skip sentinel-marked tail rows. + Reuse this from prepare_decode / prepare_prefill / prepare_capture — + caller supplies extend_lens / seq_lens (np int32) and target device. + Plan tensors are written into the pre-allocated + `v4_compress_plan_{ratio}` / `v4_write_plan_{ratio}` CpuGpuBuffers + (fixed pointers for CUDAGraph capture); the kernels skip + sentinel-marked tail rows. + + `for_decode_cg`: True for decode runtime AND decode CG capture — + the returned plan_gpu is sliced to a fixed `_decode_compress_cap` + per ratio so capture/replay shapes match. False for eager prefill — + the plan_gpu is sliced to the actual `n_compress` (smallest grid). """ from atom.model_ops.v4_kernels import make_compress_plans @@ -1753,6 +1702,9 @@ def _build_compress_plans(self, extend_lens_np, seq_lens_np, device): self._unique_compress_ratios_overlap, device, plan_buffers=plan_buffers, + decode_capacity_per_ratio=( + self._decode_compress_cap if for_decode_cg else None + ), ) def _populate_block_tables( @@ -1817,13 +1769,17 @@ def build_for_cudagraph_capture( buffers in `forward_vars`. Replay-time prepare_decode writes into the SAME buffers — captured graph reads stable addresses. - NOTE on dynamic-shape kernels (fused_compress_attn / update_compressor_states / - swa_write): these currently use variable kernel grids (`grid=(num_compress,)`), + NOTE on dynamic-shape kernels (`update_compressor_states` / `swa_write`): + these currently use variable kernel grids (`grid=(num_compress,)`), which CUDAGraph capture rejects. A follow-up PR converts them to fixed grid + sentinel masking. Until then, capture itself can succeed (the helpers run on CPU + small H2D), but model.forward inside torch.cuda.graph will likely fail at the first such kernel launch — the user can detect - this via capture log output. + this via capture log output. (`fused_compress_attn` is already + CG-safe: launches at the decode-tight slice + (`_decode_compress_cap[ratio]`, baked at capture) and + sentinel-skips inactive rows internally for both BF16 Main and FP8 + Indexer paths.) """ device = self.model_runner.device var = self.model_runner.forward_vars @@ -1882,7 +1838,7 @@ def build_for_cudagraph_capture( # helpers used at runtime — guarantees addresses match. extend_lens_np = np.full(bs, max_q_len, dtype=np.int32) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, device + extend_lens_np, context_lens_np, device, for_decode_cg=True ) # Capture: padded_bs == scheduled_bs == bs (synthetic batch is full). # Must run BEFORE `_attach_v4_indexer_meta` so the indexer-side meta @@ -2008,33 +1964,18 @@ def _alloc_v4_metadata_buffers(self) -> None: # for cu_seq_lens. Also reused as cu_starts/cu_ends for fp8_mqa_logits # (which accepts both int32 and int64). bufs["v4_indexer_cu_committed"] = CpuGpuBuffer(bs + 1, **i32) - # Decode-path logits buffer for `deepgemm_fp8_paged_mqa_logits`. - # Sized [max_bs, max_model_len_idx] fp32 — assumes V4-Pro next_n=1. - # deepgemm writes valid cols [0, n_committed_per_seq[batch]) per row; - # padding cols carry stale data but `top_k_per_row_decode` honors - # `n_committed_per_seq` per-row so unwritten cols are never selected. - bufs["v4_indexer_decode_logits"] = CpuGpuBuffer( - bs, self.max_model_len_idx, dtype=torch.float32, device=self.device - ) - # Decode-path top-k indices buffer (consumed by `top_k_per_row_decode`). - # Sized [max_bs, index_topk] int32 — V4-Pro next_n=1; multi-token decode - # would scale rows by next_n. Replaces the per-fwd torch.topk allocation. - bufs["v4_indexer_decode_topk_indices"] = CpuGpuBuffer( - bs, self.index_topk, dtype=torch.int32, device=self.device - ) + # NOTE: decode-path `logits` ([T, max_model_len_idx] fp32) and + # `topk_indices` ([T, index_topk] int32) are NOT pre-allocated — + # they are write-once GPU scratch with no CPU mirror, allocated + # per-fwd inside `Indexer._score_topk_decode` via `torch.empty`. + # Under CUDAGraph capture they land in the graph's private pool + # and replay reuses the same address; eager keeps the standard + # caching-allocator fast path. # Per-token write offset consumed by `csa_translate_pack` (decode path # fills with `window_size`; prefill path will fill with per-token # prior_swa_count once the new dual-source kernel lands). Allocated # `mnbt` (worst-case prefill) so prefill writes don't overflow. bufs["v4_skip_prefix_len_csa"] = CpuGpuBuffer(mnbt, **i32) - # FP8 cache write-side slot mapping (one entry per compress row). - # Bound = ⌈mnbt/4⌉ + bs (worst-case num_compress for ratio=4 CSA; - # matches v4_compress_plan_4 row count). - # int64 — `indexer_k_quant_and_cache` requires int64 slot_mapping. - idx_compress_bound = mnbt // 4 + bs - bufs["v4_indexer_compress_slot_mapping"] = CpuGpuBuffer( - idx_compress_bound, **i64 - ) # Compress plan buffers (per-ratio) — pre-allocated for CUDAGraph # plan-tensor address stability. `make_compress_plans(..., plan_buffers=)` @@ -2042,12 +1983,21 @@ def _alloc_v4_metadata_buffers(self) -> None: # sizes: num_compress ≤ ⌈mnbt/ratio⌉ + bs (one boundary per seq plus # alignment slack); num_write ≤ bs * STATE_SIZE (per-seq ring window # carries STATE_SIZE rows per fwd at most). - max_compress_per_ratio = {} + # + # The decode CG path uses a much tighter capacity than the prefill + # worst case — the kernel grid is dictated by the slice of this + # buffer that we hand to the kernel, and decode only ever needs + # `max_decode_tokens // ratio + max_bs` rows (vs `mnbt // ratio + bs` + # for prefill, which is ~13× larger at typical config). We still + # allocate the full prefill capacity (eager prefill needs it), but + # both decode capture and replay slice down to `_decode_compress_cap` + # so the captured grid is the decode-tight bound. capture and + # replay MUST use the same value (CG kernel call args are baked). + self._decode_compress_cap: dict[int, int] = {} for ratio, is_overlap in self._unique_compress_ratios_overlap: state_size = (2 if is_overlap else 1) * ratio max_compress = mnbt // ratio + bs max_write = min(mnbt, bs * state_size) - max_compress_per_ratio[ratio] = max_compress bufs[f"v4_compress_plan_{ratio}"] = CpuGpuBuffer(max_compress, 4, **i32) bufs[f"v4_write_plan_{ratio}"] = CpuGpuBuffer(max_write, 4, **i32) # Pre-fill with sentinel so capture-time buffer state is valid @@ -2056,24 +2006,12 @@ def _alloc_v4_metadata_buffers(self) -> None: bufs[f"v4_compress_plan_{ratio}"].copy_to_gpu() bufs[f"v4_write_plan_{ratio}"].cpu.fill_(-1) bufs[f"v4_write_plan_{ratio}"].copy_to_gpu() - - # Compressor output buffers (one per kind, shared across same-kind - # layers within a single fwd — Compressor outputs are consumed - # immediately by the layer's sparse_attn before the next layer runs). - # Sized to (max_compress_per_ratio, head_dim) so fused_compress_attn - # can launch with full-capacity grid + sentinel-skip; output rows - # past the actual num_compress carry stale data but are never read - # (consumer slices via `cu_compress_cpu`). - bf16 = {"dtype": torch.bfloat16, "device": self.device} - if 4 in max_compress_per_ratio: - mc = max_compress_per_ratio[4] - bufs["v4_csa_main_compress_out"] = torch.empty((mc, self.head_dim), **bf16) - bufs["v4_csa_idx_compress_out"] = torch.empty( - (mc, self.index_head_dim), **bf16 - ) - if 128 in max_compress_per_ratio: - mc = max_compress_per_ratio[128] - bufs["v4_hca_main_compress_out"] = torch.empty((mc, self.head_dim), **bf16) + # Decode-tight bound. Worst case = total_tokens_decode boundaries + # all firing simultaneously, each in its own ratio-aligned slot. + # `total_tokens_decode = max_decode_tokens` (= max_bs * (1+spec)). + # The `+ max_bs` covers per-seq alignment slack (each seq can hit + # at most one extra boundary when extend_len isn't ratio-aligned). + self._decode_compress_cap[ratio] = self.max_decode_tokens // ratio + bs self.model_runner.forward_vars.update(bufs) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 96aafcfe9b..851cb3dce5 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -73,7 +73,7 @@ def _amd_smem_safe_tile(): # Defaults chosen so BLOCK_M*BLOCK_N stays ≤ 16384 entries (64 KiB FP32 # acc), comfortably fitting MI355X's register file. Override via env if # a future compiler/kernel update relaxes the budget. - block_m = int(os.getenv("ATOM_TRITON_MOE_BLOCK_M", "64")) + block_m = int(os.getenv("ATOM_TRITON_MOE_BLOCK_M", "32")) block_n = int(os.getenv("ATOM_TRITON_MOE_BLOCK_N", "256")) update_opt_flags_constraints({"block_m": block_m, "block_n": block_n}) try: diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 7a687e5243..86488dd5d5 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -635,6 +635,7 @@ def _fused_qk_norm_single_kernel( num_q_heads, num_k_heads, ADD_UNIT_OFFSET: tl.constexpr, + Q_HAS_WEIGHT: tl.constexpr, RBLOCK: tl.constexpr, XBLOCK: tl.constexpr, ): @@ -664,15 +665,22 @@ def _fused_qk_norm_single_kernel( mask = xmask & col_mask - # Weight: load both, select via is_q - qw = tl.load( - q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" - ).to(tl.float32) + # Weight: load both (or use ones for Q when Q_HAS_WEIGHT=False), select via is_q. + # Q_HAS_WEIGHT=False is for callers whose Q-side norm has implicit identity + # weight (e.g. V4's per-head Q normalization, equivalent to the prior + # `_rmsnorm_nw` helper) — saves a load + register row. + if Q_HAS_WEIGHT: + qw = tl.load( + q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" + ).to(tl.float32) + if ADD_UNIT_OFFSET: + qw = qw + 1.0 + else: + qw = tl.full((RBLOCK,), 1.0, tl.float32) kw = tl.load( k_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" ).to(tl.float32) if ADD_UNIT_OFFSET: - qw = qw + 1.0 kw = kw + 1.0 w = tl.where(is_q, qw, kw) @@ -714,7 +722,7 @@ def _fused_qk_norm_single_kernel( def fused_qk_norm( q: torch.Tensor, k: torch.Tensor, - q_weight: torch.Tensor, + q_weight: Optional[torch.Tensor], k_weight: torch.Tensor, eps: float, add_unit_offset: bool = False, @@ -724,11 +732,18 @@ def fused_qk_norm( Args: q: [num_tokens, num_heads, head_dim] k: [num_tokens, num_kv_heads, head_dim] - q_weight, k_weight: [head_dim] norm weights + q_weight: [head_dim] norm weight, or None for an implicit ones weight + (skips the Q-side weight load — for callers whose Q norm + is the identity, e.g. V4's per-head Q normalization). + k_weight: [head_dim] norm weight (always required) eps: epsilon for numerical stability add_unit_offset: True for GemmaRMSNorm (w+1), False for standard """ - head_dim = q_weight.shape[0] + head_dim = k_weight.shape[0] + if q_weight is not None: + assert ( + q_weight.shape[0] == head_dim + ), f"q_weight head_dim {q_weight.shape[0]} != k_weight {head_dim}" num_tokens = q.shape[0] num_q_heads = q.shape[1] num_k_heads = k.shape[1] @@ -745,12 +760,15 @@ def fused_qk_norm( # num_warps=1 is universally optimal for head_dim=256 workloads on MI355X. XBLOCK = 2 if total_rows > 8192 else 1 NUM_WARPS = 1 + # When q_weight is None pass k_weight as a placeholder pointer (the + # kernel won't load from it — Q_HAS_WEIGHT=False gates the load). + q_weight_arg = q_weight if q_weight is not None else k_weight _fused_qk_norm_single_kernel[((total_rows + XBLOCK - 1) // XBLOCK,)]( q, k, q_out, k_out, - q_weight, + q_weight_arg, k_weight, eps, num_tokens, @@ -762,6 +780,7 @@ def fused_qk_norm( num_q_heads, num_k_heads, ADD_UNIT_OFFSET=add_unit_offset, + Q_HAS_WEIGHT=q_weight is not None, RBLOCK=RBLOCK, XBLOCK=XBLOCK, num_warps=NUM_WARPS, @@ -773,6 +792,12 @@ class DualRMSNorm: """Fused Q/K RMSNorm — single Triton kernel launch. Not an nn.Module. References existing q_norm/k_norm for weights. + + Q-side weightless mode: when `q_norm.weight is None`, the Q-side norm + is treated as the identity (implicit ones weight). The kernel skips + the q_weight load entirely (Q_HAS_WEIGHT=False). Use this when the + checkpoint does not ship a Q weight (e.g. V4's per-head Q normalization, + equivalent to the prior `_rmsnorm_nw` helper). """ def __init__( @@ -789,7 +814,22 @@ def __init__( self.num_q_heads = num_q_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim - self.add_unit_offset = isinstance(q_norm, GemmaRMSNorm) + # add_unit_offset only applies when q has a real weight; the kernel's + # weightless path (q_norm.weight is None) emits qw=1.0 directly with + # no offset, so the GemmaRMSNorm test is gated on weight existence. + self.add_unit_offset = getattr( + q_norm, "weight", None + ) is not None and isinstance(q_norm, GemmaRMSNorm) + # Resolve eps once. Different RMSNorm implementations name the + # attribute differently — `variance_epsilon` (HF/Gemma style) or + # `eps` (ATOM RMSNorm). Cache to avoid the lookup per forward call. + self._eps = getattr(q_norm, "variance_epsilon", None) or getattr( + q_norm, "eps", None + ) + assert self._eps is not None, ( + f"q_norm {type(q_norm).__name__} must expose `eps` or " + f"`variance_epsilon`" + ) self.prefix = prefix @mark_trace @@ -803,12 +843,15 @@ def __call__( Returns: (q_normed, k_normed) same shapes as input """ + # `self.q_norm.weight` is None for weightless q_norm (e.g. V4 + # `q_norm2`); fused_qk_norm forwards that None to the kernel which + # takes the Q_HAS_WEIGHT=False fast path. q, k = fused_qk_norm( q.view(-1, self.num_q_heads, self.head_dim), k.view(-1, self.num_kv_heads, self.head_dim), self.q_norm.weight, self.k_norm.weight, - self.q_norm.variance_epsilon, + self._eps, add_unit_offset=self.add_unit_offset, ) return ( diff --git a/atom/model_ops/v4_kernels/compress_plan.py b/atom/model_ops/v4_kernels/compress_plan.py index b7cde3b423..082c84a5aa 100644 --- a/atom/model_ops/v4_kernels/compress_plan.py +++ b/atom/model_ops/v4_kernels/compress_plan.py @@ -62,6 +62,7 @@ def make_compress_plans( device: torch.device, *, plan_buffers: dict | None = None, + decode_capacity_per_ratio: dict[int, int] | None = None, ) -> dict[int, CompressPlan]: """Build a CompressPlan per (ratio, overlap) variant. @@ -80,6 +81,18 @@ def make_compress_plans( the returned `compress_plan_gpu` / `write_plan_gpu` views have stable data pointers across calls. When None, the legacy fresh-allocation path is used (eager only). + decode_capacity_per_ratio: optional dict[ratio] -> int. Slice length + for the returned `compress_plan_gpu`. When PROVIDED: + the slice is exactly that fixed value (independent of + `n_compress`) — required by the decode CUDAGraph path + so capture and replay produce kernel calls with + identical tensor shapes. Caller must size this to the + decode worst case (`max_decode_tokens // ratio + bs`). + When NONE: the slice is exactly `n_compress` (tight, + eager-only) — gives the smallest possible kernel grid. + The slice is contiguous-from-base so the data pointer + is stable; only `shape[0]` shrinks. The underlying + buffer is always sized to the prefill worst case. Returns: dict[ratio] -> CompressPlan. Empty dict if `extend_lens_cpu.sum() == 0` @@ -96,15 +109,27 @@ def make_compress_plans( if total == 0 or bs == 0: if plan_buffers is None: return out - # Capture path with empty fwd: fully sentinel-fill the buffers and - # return CompressPlans so addresses are stable. Skipped via num_*=0. + # Empty fwd: produce CompressPlans pointing at the pre-allocated + # buffers so capture-time addresses match replay-time addresses + # even on a zero-token fwd. Skipped via num_*=0. + # + # CG path (decode_capacity_per_ratio provided): slice = fixed cap. + # Eager path (None): slice = 0 — kernel grid is empty, no work + # dispatched, no sentinel fill required. for ratio, _ in unique_ratios_overlap: cbuf = plan_buffers[ratio]["compress"] wbuf = plan_buffers[ratio]["write"] - cbuf.np[:].fill(-1) + cap = ( + decode_capacity_per_ratio.get(ratio) + if decode_capacity_per_ratio is not None + else 0 + ) + if cap > 0: + cbuf.np[:cap].fill(-1) + compress_plan_gpu = cbuf.copy_to_gpu(cap if cap > 0 else 0) wbuf.np[:].fill(-1) out[ratio] = CompressPlan( - compress_plan_gpu=cbuf.copy_to_gpu(), + compress_plan_gpu=compress_plan_gpu, write_plan_gpu=wbuf.copy_to_gpu(), num_compress=0, num_write=0, @@ -163,9 +188,19 @@ def make_compress_plans( if plan_buffers is not None: cbuf = plan_buffers[ratio]["compress"] wbuf = plan_buffers[ratio]["write"] - assert n_compress <= cbuf.np.shape[0], ( - f"ratio={ratio} num_compress={n_compress} exceeds buffer " - f"capacity {cbuf.np.shape[0]}; bump in builder __init__." + full_cap = cbuf.np.shape[0] + # CG path: fixed slice (capture / replay must match). Eager + # path: slice = n_compress (smallest possible kernel grid). + cap = ( + decode_capacity_per_ratio.get(ratio) + if decode_capacity_per_ratio is not None + else None + ) + slice_cap = n_compress if cap is None else cap + assert n_compress <= slice_cap <= full_cap, ( + f"ratio={ratio} num_compress={n_compress}, slice={slice_cap}, " + f"buffer={full_cap}: invariant violated. CG path requires " + f"n_compress ≤ decode_cap; eager path uses n_compress as slice." ) assert n_write <= wbuf.np.shape[0], ( f"ratio={ratio} num_write={n_write} exceeds buffer " @@ -173,11 +208,14 @@ def make_compress_plans( ) if n_compress > 0: cbuf.np[:n_compress] = compress_plan - cbuf.np[n_compress:].fill(-1) # sentinel + # Sentinel only within the slice we hand to the kernel; rows + # beyond `slice_cap` are unreachable from this launch. + if slice_cap > n_compress: + cbuf.np[n_compress:slice_cap].fill(-1) if n_write > 0: wbuf.np[:n_write] = write_plan wbuf.np[n_write:].fill(-1) # sentinel - compress_plan_gpu = cbuf.copy_to_gpu() + compress_plan_gpu = cbuf.copy_to_gpu(slice_cap) write_plan_gpu = wbuf.copy_to_gpu() else: compress_plan_gpu = torch.from_numpy( diff --git a/atom/model_ops/v4_kernels/fused_compress.py b/atom/model_ops/v4_kernels/fused_compress.py index 0e26b7eb8a..9f5eeaf8d6 100644 --- a/atom/model_ops/v4_kernels/fused_compress.py +++ b/atom/model_ops/v4_kernels/fused_compress.py @@ -10,12 +10,14 @@ Each compression boundary across the entire fwd is one row in `compress_plan_gpu` — a packed `[num_compress, 4] int32` tensor where each row is `[ragged_id, batch_id, position, window_len]`. The kernel grid is - `num_compress` (zero waste, no early-exit), and each program does ONE 4×i32 - load to get all the metadata it needs: + the caller-supplied slice length (decode CG: `_decode_compress_cap[ratio]` + / eager prefill: `n_compress`); inactive plan rows are sentinel-marked + (`position == -1`) and bail at the top of the kernel. Each program does + ONE 4×i32 load to get all the metadata it needs: ragged_id → row index in the ragged kv_in / score_in stream batch_id → seq index → state_slot_mapping[batch_id], block_table[batch_id] - position → absolute token position (drives RoPE + paged scatter) + position → absolute token position (drives RoPE + paged scatter); -1 = skip window_len → number of leading K-loop iterations that read state cache (instead of the ragged input). K = STATE_SIZE. @@ -32,18 +34,21 @@ the caller MUST invoke this kernel BEFORE `update_compressor_states` runs (which would overwrite the historic positions this kernel needs). -Output: - Returns `[num_compress, head_dim]` BF16 tensor. Rows are in plan order - (= ragged_id ascending = per-seq grouped). Caller slices per seq via - `plan.cu_compress_cpu[i]:plan.cu_compress_cpu[i+1]`. - - Also writes valid rows into the paged kv_cache at compressed index - `position // RATIO` (when `block_table` is provided). - -TODO: FP8/FP4 quant fusion. Currently stores raw BF16 (no act_quant). The -`scale_fmt="ue8m0"` round-to-even f32→e8m0 path requires porting aiter's -`f32_to_e8m0` bit manipulation into Triton (~20 lines per quant block). -Skipping for this PR; follow-up PR will add it. +Quant modes (constexpr-selected by the `quant` arg of the Python wrapper): + - quant=False (CSA Main / HCA Main): writes BF16 rows into the paged BF16 + `kv_cache` at compressed slot `position // ratio`. + - quant=True (CSA Indexer-inner): per-row amax → ue8m0 (or raw) scale → + fp8 cast → preshuffled (MFMA 16x16 tile) write into the FP8 `kv_cache`, + plus fp32 scale into the per-block scale region (`cache_scale` — a + strided view of the same allocation built by the V4 builder). Bit-exact + match for `indexer_k_quant_and_cache` / + `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + +Output: side-effecting only — cache scatter IS the only output. The earlier +caller-visible `[num_compress, head_dim]` BF16 return tensor was vestigial +(paged_decode/paged_prefill read the scattered compress entries directly +from `unified_kv` (Main) or the FP8 indexer pool, not from the kernel +return). """ from typing import Optional @@ -83,15 +88,14 @@ def _fused_compress_attn_kernel( sin_cache_ptr, cos_sin_pos_stride, # = rope_head_dim // 2 # ── KV cache scatter (paged) ──────────────────────────────────────── - kv_cache_ptr, # [num_blocks, k_per_block, head_dim] bf16 + kv_cache_ptr, # bf16: [NB, k_per_block, head_dim] / fp8: [NB, k_per_block, head_dim] kv_cache_block_stride, kv_cache_token_stride, + cache_scale_ptr, # fp32 [NB, k_per_block] (QUANT path only; dummy ptr otherwise) + cache_scale_block_stride, # fp32 elements per block (= k_per_block * cache_stride / 4) block_table_ptr, # [bs, max_blocks_per_seq] int32 block_table_seq_stride, # row stride k_per_block, - # ── output to caller (post norm+rope BF16 for sparse_attn input) ──── - out_ptr, # [num_compress, head_dim] bf16 - out_token_stride, head_dim, rope_head_dim, # ── constexpr ─────────────────────────────────────────────────────── @@ -102,10 +106,19 @@ def _fused_compress_attn_kernel( STATE_SIZE: tl.constexpr, # = 2*RATIO if OVERLAP else RATIO K: tl.constexpr, # = STATE_SIZE (softmax-pool reduce dim) HAS_BLOCK_TABLE: tl.constexpr, + QUANT: tl.constexpr, # 0 = raw BF16 (CSA/HCA Main), 1 = FP8 e4m3 + ue8m0 scale (Indexer) + USE_UE8M0: tl.constexpr, # round scale to power-of-2 (only when QUANT == 1) + PRESHUFFLE: tl.constexpr, # MFMA 16x16 preshuffled FP8 layout (only when QUANT == 1) + # E4M3 max (=448 for E4M3FN, =240 for E4M3FNUZ). Constexpr so the clamp + # bounds and reciprocal fold to compile-time constants. Ignored when + # QUANT == 0 (caller passes 1.0 as a placeholder). + FP8_MAX: tl.constexpr = 1.0, ): - """One program per boundary in the plan. Grid = plan capacity (CUDAGraph- - safe fixed grid); inactive rows are sentinel-marked (position == -1) and - bail before any load/store/scatter.""" + """One program per boundary in the plan. Grid = caller-supplied slice + length (decode CG: `_decode_compress_cap[ratio]` for capture/replay + address stability; eager prefill: tight `n_compress`). Inactive rows + are sentinel-marked (position == -1) and bail before any load / + store / scatter.""" pid = tl.program_id(0) plan_base = plan_ptr + pid * 4 ragged_id = tl.load(plan_base + 0) @@ -122,75 +135,98 @@ def _fused_compress_attn_kernel( d_mask = d < head_dim # ── 1. Per-source-position load + online softmax-pool ────────────── + # Two-phase split (vs. the older single masked loop): for each program + # `is_input = k_static >= window_len` partitions the K iterations into + # a leading state-cache-only run and a trailing input-only run. Issuing + # only the live side's loads (rather than masked-off both) cuts HBM + # bandwidth ~40% on AMD CDNA where masked tl.load still issues the LD + # instruction (predicate only suppresses register write-back). + # + # Padding invariant: padding (`s < 0` ⟺ `k_static < K-1-position`) lies + # entirely within `k_static < window_len` because + # `window_len = K - min(j_in_seq+1, K) ≥ K - 1 - j_in_seq ≥ K - 1 - position` + # (`position = prefix_len + j_in_seq`, `prefix_len ≥ 0`). The input phase + # therefore needs no padding mask. NEG_INF: tl.constexpr = float("-inf") m_acc = tl.full([BLOCK_D], NEG_INF, tl.float32) kv_acc = tl.zeros([BLOCK_D], tl.float32) w_acc = tl.zeros([BLOCK_D], tl.float32) - # Dynamic loop (NOT unrolled) — K=128 (HCA) would otherwise blow up hsaco. - for k_static in tl.range(K): + # ── Phase 1: state cache (k_static ∈ [0, window_len)) ── + # Dynamic bound (window_len is per-program, not constexpr) — Triton + # cannot static-unroll this. Loop body issues only state-side loads. + for k_static in tl.range(0, window_len): s = position - K + 1 + k_static is_padding = s < 0 - is_input = k_static >= window_len - # is_state = (~is_input) & (~is_padding) - - # B-side (k < RATIO): cols [:head_dim] (= col_off=0) - # A-side (k >= RATIO): cols [head_dim:] (= col_off=head_dim) - # HCA (no overlap, K=RATIO): col_off=0 always (k_static < RATIO). + # B-side (k >= RATIO): cols [head_dim:]; A-side (k < RATIO): cols [:head_dim]. + # HCA (no overlap, K=RATIO): col_off=0 always. col_off = (k_static >= RATIO) * head_dim if OVERLAP else 0 - ape_row = k_static % RATIO # [0, RATIO) — same for B/A sides - # Input source: row index in ragged stream. - # k_static = K-1 corresponds to s = position (the boundary token itself, - # = ragged_id row). Earlier k_static values map to earlier ragged rows. - in_row = ragged_id - (K - 1 - k_static) - kv_a = tl.load( - kv_in_ptr + in_row * kv_in_row_stride + col_off + d, - mask=is_input & d_mask, - other=0.0, - ) - - # State cache source: ring slot indexed by absolute s. s_safe = tl.maximum(s, 0) ring = s_safe % STATE_SIZE + state_row_off = ( + slot * kv_state_slot_stride + ring * kv_state_pos_stride + col_off + ) kv_b = tl.load( - kv_state_ptr - + slot * kv_state_slot_stride - + ring * kv_state_pos_stride + kv_state_ptr + state_row_off + d, + mask=(~is_padding) & d_mask, + other=0.0, + ) + score_b = tl.load( + score_state_ptr + + slot * score_state_slot_stride + + ring * score_state_pos_stride + col_off + d, - mask=(~is_input) & (~is_padding) & d_mask, - other=0.0, + mask=(~is_padding) & d_mask, + other=NEG_INF, ) - kv_k = kv_a + kv_b # exactly one path active per source pos - # ── score_k load ── + m_new = tl.maximum(m_acc, score_b) + scale = tl.where(m_acc == NEG_INF, 0.0, tl.exp(m_acc - m_new)) + # Padding lanes have score_b = NEG_INF → w_k = 0, contributes nothing. + w_k = tl.where(score_b == NEG_INF, 0.0, tl.exp(score_b - m_new)) + kv_acc = kv_acc * scale + w_k * kv_b + w_acc = w_acc * scale + w_k + m_acc = m_new + + # ── Phase 2: ragged input (k_static ∈ [window_len, K)) ── + # No padding here (per invariant above). All loads unconditional in + # the position dimension; only the head_dim mask remains. + for k_static in tl.range(window_len, K): + col_off = (k_static >= RATIO) * head_dim if OVERLAP else 0 + ape_row = k_static % RATIO + # k_static = K-1 → s = position (the boundary token itself, + # = ragged_id row). Earlier k_static map to earlier ragged rows. + in_row = ragged_id - (K - 1 - k_static) + + # kv_in / score_in: single-use per program → evict_first to keep + # state-cache lines (small, possibly shared across programs) hot. + kv_a = tl.load( + kv_in_ptr + in_row * kv_in_row_stride + col_off + d, + mask=d_mask, + other=0.0, + eviction_policy="evict_first", + ) score_a = tl.load( score_in_ptr + in_row * score_in_row_stride + col_off + d, - mask=is_input & d_mask, - other=NEG_INF, + mask=d_mask, + other=0.0, + eviction_policy="evict_first", ) ape_v = tl.load( ape_ptr + ape_row * dim_full + col_off + d, - mask=is_input & d_mask, + mask=d_mask, other=0.0, + eviction_policy="evict_last", ) - score_b = tl.load( - score_state_ptr - + slot * score_state_slot_stride - + ring * score_state_pos_stride - + col_off - + d, - mask=(~is_input) & (~is_padding) & d_mask, - other=NEG_INF, - ) - score_k = tl.where(is_input, score_a + ape_v, score_b) + score_k = score_a + ape_v - # ── Online softmax-pool accumulate ── m_new = tl.maximum(m_acc, score_k) scale = tl.where(m_acc == NEG_INF, 0.0, tl.exp(m_acc - m_new)) - w_k = tl.where(score_k == NEG_INF, 0.0, tl.exp(score_k - m_new)) - kv_acc = kv_acc * scale + w_k * kv_k + # score_k always finite in input phase → no NEG_INF guard needed. + w_k = tl.exp(score_k - m_new) + kv_acc = kv_acc * scale + w_k * kv_a w_acc = w_acc * scale + w_k m_acc = m_new @@ -231,26 +267,80 @@ def _fused_compress_attn_kernel( new_odd = odd_v * cos_per_pair + even_v * sin_per_pair rotated = tl.interleave(new_even, new_odd) # [BLOCK_D] fp32 - # ── 4. Cast to BF16 + store ──────────────────────────────────────── - rotated_bf16 = rotated.to(tl.bfloat16) - - # Output: row pid (plan order = per-seq grouped, caller uses cu_compress). - tl.store(out_ptr + pid * out_token_stride + d, rotated_bf16, mask=d_mask) - - # KV cache scatter (paged): block_table[batch_id, ci // k_per_block][ci % k_per_block] + # ── 4. Cache scatter (paged) ─────────────────────────────────────── + # The Compressor's BF16 return value was historically consumed by sparse + # attention but is now vestigial — paged_decode/paged_prefill read the + # scattered compress entries directly from `unified_kv` (Main) or the FP8 + # indexer pool. So no caller-visible `out` write; the cache scatter IS + # the only output. if HAS_BLOCK_TABLE: + # block_table[batch_id, ci // k_per_block] → physical_block. + # slot_in_block = ci % k_per_block. Same address resolution as + # `indexer_k_quant_and_cache` (which used a pre-flattened slot_mapping + # = physical_block * k_per_block + slot_in_block). ci = position // RATIO block_in_seq = ci // k_per_block slot_in_block = ci % k_per_block physical_block = tl.load( block_table_ptr + batch_id * block_table_seq_stride + block_in_seq ).to(tl.int64) - cache_addr = ( - physical_block * kv_cache_block_stride - + slot_in_block * kv_cache_token_stride - + d - ) - tl.store(kv_cache_ptr + cache_addr, rotated_bf16, mask=d_mask) + + if QUANT: + # FP8 e4m3 quantization: per-row amax → ue8m0 (or raw) scale → + # clamp+cast to fp8 → write preshuffled (MFMA 16x16 tile) or + # linear layout into the FP8 region; write the fp32 scale into + # the per-block scale region (separate `cache_scale_ptr` view). + # Bit-exact match for `indexer_k_quant_and_cache` / + # `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + # + # Quant pattern follows aiter's `_fp8_quant_op` / + # `_fused_rms_gated_fp8_group_quant_kernel` — the explicit + # `fp_downcast_rounding="rtne"` is intentionally NOT used here: + # on AMD it forces a slow software-RTNE path and bypasses the + # `v_cvt_pk_fp8_f32` HW intrinsic (which already rounds RTNE). + # The pre-cast `tl.clamp` saturates intermediate overflow and + # mirrors aiter's hand-tuned HIP `aiter::scaled_cast`. + rotated_for_amax = tl.where(d_mask, tl.abs(rotated), 0.0) + amax = tl.max(rotated_for_amax, axis=0) # scalar (per row) + scale = tl.maximum(amax, 1e-4) * (1.0 / FP8_MAX) + if USE_UE8M0: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + inv_scale = 1.0 / scale + scaled = tl.clamp(rotated * inv_scale, -FP8_MAX, FP8_MAX) + fp8_val = scaled.to(kv_cache_ptr.dtype.element_ty) + if PRESHUFFLE: + TILE: tl.constexpr = 16 + token_tile_id = slot_in_block // TILE + token_in_tile = slot_in_block % TILE + col_tile_id = d // TILE + col_in_tile = d % TILE + fp8_offset = ( + physical_block * kv_cache_block_stride + + token_tile_id * (TILE * head_dim) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile + ) + else: + fp8_offset = ( + physical_block * kv_cache_block_stride + + slot_in_block * head_dim + + d + ) + # Streaming write — these slots aren't reread inside this kernel. + tl.store( + kv_cache_ptr + fp8_offset, fp8_val, mask=d_mask, cache_modifier=".cs" + ) + # Scale: one fp32 per row, packed at end of block. + scale_offset = physical_block * cache_scale_block_stride + slot_in_block + tl.store(cache_scale_ptr + scale_offset, scale, cache_modifier=".cs") + else: + cache_addr = ( + physical_block * kv_cache_block_stride + + slot_in_block * kv_cache_token_stride + + d + ) + tl.store(kv_cache_ptr + cache_addr, rotated.to(tl.bfloat16), mask=d_mask) def fused_compress_attn( @@ -270,7 +360,9 @@ def fused_compress_attn( cos_cache: torch.Tensor, # [max_seq, ..., rope_head_dim/2] bf16/fp16 sin_cache: torch.Tensor, # same shape # KV cache scatter - kv_cache: Optional[torch.Tensor], # [num_blocks, k_per_block, head_dim] bf16 + kv_cache: Optional[ + torch.Tensor + ], # bf16: [NB, k_per_block, head_dim] / fp8: same shape, fp8 block_tables: Optional[torch.Tensor], # [bs, max_blocks_per_seq] int32 k_per_block: int, # Geometry @@ -278,32 +370,42 @@ def fused_compress_attn( ratio: int, head_dim: int, rope_head_dim: int, - out_dtype: torch.dtype = torch.bfloat16, - out: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - """Batched fused per-source-position pool + RMSNorm + RoPE + bf16 kv_cache - scatter, dispatched via SGLang-style packed plan. - - Returns `[num_compress, head_dim]` BF16 tensor in plan order (= ragged_id - ascending = per-seq grouped). Caller uses `plan.cu_compress_cpu[i:i+2]` to - slice per-seq chunks. - - Returns None if `plan.num_compress == 0` AND no `out` buffer is provided. - When `out` is supplied (CUDAGraph path), kernel ALWAYS launches at full - plan capacity — inactive rows are sentinel-skipped inside the kernel — - and `out` is returned unchanged when `num_compress == 0`. + # FP8 quant fusion (Indexer-inner Compressor path) + quant: bool = False, + cache_scale: Optional[ + torch.Tensor + ] = None, # fp32 [NB, k_per_block]; required when quant=True + use_ue8m0: bool = True, # round scale to power-of-2 (UE8M0); only when quant=True + preshuffle: bool = True, # MFMA 16x16 preshuffled FP8 layout; only when quant=True + fp8_max: Optional[float] = None, # E4M3 max; required when quant=True +) -> None: + """Batched fused per-source-position pool + RMSNorm + RoPE + cache scatter, + dispatched via SGLang-style packed plan. + + Two scatter modes (constexpr-selected by `quant`): + + - quant=False (CSA Main / HCA Main): writes BF16 rows into the paged + BF16 kv_cache (compressed slot = `position // ratio`). + + - quant=True (CSA Indexer-inner): per-row amax → ue8m0 scale → fp8 cast + → preshuffled (MFMA 16x16 tile) write into the FP8 kv_cache, plus + fp32 scale into `cache_scale` (the per-block scale region of the + same allocation). Bit-exact match for `indexer_k_quant_and_cache` / + `cp_gather_indexer_k_quant_cache` (cache_kernels.cu:1145+). + + Side-effecting: cache scatter IS the only output (Main path's BF16 + return tensor was vestigial — paged_decode/paged_prefill read directly + from `unified_kv` and the indexer FP8 pool, not from the kernel return). + Grid is always `plan_capacity` (CUDAGraph-safe); inactive plan rows are + sentinel-skipped (`position == -1`) inside the kernel. Caller MUST invoke BEFORE `update_compressor_states` (state cache reads must see previous-fwd data). """ - device = kv_in.device - num_compress = plan.num_compress plan_capacity = plan.compress_plan_gpu.shape[0] + num_compress = plan.num_compress if plan_capacity == 0: - return out # nothing to do; out (or None) returned as-is. - if num_compress == 0 and out is None: - # Eager path: nothing to do, no caller buffer to fill. - return None + return # nothing to do — no plan rows ever populated. # Validate shapes dim_full = (2 if overlap else 1) * head_dim @@ -345,30 +447,45 @@ def fused_compress_attn( else: bt_seq_stride = 0 - if out is None: - # Eager path: allocate output sized to the actual num_compress so the - # returned tensor has the legacy [num_compress, head_dim] shape. - out = torch.empty(num_compress, head_dim, dtype=out_dtype, device=device) - else: - # CUDAGraph path: caller-provided buffer of capacity ≥ plan_capacity. - # Validate shape; kernel will write into rows [0:num_compress) and - # leave [num_compress:plan_capacity) untouched (those plan rows are - # sentinel-skipped). Inactive output rows therefore carry stale data - # but no consumer reads them (caller slices via cu_compress_cpu). + # Quant validation. quant=True is FP8 cache write (Indexer-inner path); + # requires a valid block_tables (slot resolution) AND a paired fp32 scale + # view of the same allocation (see Compressor.forward for how to slice it). + if quant: + assert has_bt, "quant=True requires block_tables for slot resolution" + assert ( + kv_cache.dtype != torch.bfloat16 + ), f"quant=True expects an FP8/uint8 kv_cache; got {kv_cache.dtype}" assert ( - out.shape[0] >= plan_capacity and out.shape[1] == head_dim - ), f"out {tuple(out.shape)}, expected ≥ ({plan_capacity}, {head_dim})" - assert out.dtype == out_dtype - assert out.is_contiguous() or out.stride(1) == 1 + cache_scale is not None and cache_scale.dtype == torch.float32 + ), "quant=True requires `cache_scale` (fp32 [NB, k_per_block])" + assert cache_scale.dim() == 2 and cache_scale.shape[0] == kv_cache.shape[0] + assert fp8_max is not None and fp8_max > 0 + if preshuffle: + assert ( + head_dim % 16 == 0 + ), f"preshuffle requires head_dim%16==0, got {head_dim}" + assert ( + k_per_block % 16 == 0 + ), f"preshuffle requires k_per_block%16==0, got {k_per_block}" BLOCK_D = triton.next_power_of_2(head_dim) HALF_ROPE = rope_head_dim // 2 K = state_size + # Cache-scale args (only consumed by the quant path; pass placeholders + # otherwise so the constexpr branch is never taken). + if quant: + cache_scale_arg = cache_scale + cache_scale_block_stride_arg = cache_scale.stride(0) + fp8_max_arg = float(fp8_max) + else: + cache_scale_arg = state_slot_mapping # placeholder int32 ptr (unused) + cache_scale_block_stride_arg = 0 + fp8_max_arg = 1.0 # placeholder; FP8_MAX is constexpr, must be > 0 + # Fixed grid for CUDAGraph compat: launch one program per plan row; - # sentinel rows (position=-1) skipped inside the kernel. When out is - # eager-allocated above, grid still shrinks to num_compress (no pad). - grid = (plan_capacity if out.shape[0] >= plan_capacity else num_compress,) + # sentinel rows (position=-1) skip inside the kernel. + grid = (plan_capacity,) _fused_compress_attn_kernel[grid]( kv_in, kv_in.stride(0), @@ -392,11 +509,11 @@ def fused_compress_attn( kv_cache if has_bt else cos_cache, # placeholder when no scatter kv_cache.stride(0) if has_bt else 0, kv_cache.stride(1) if has_bt else 0, + cache_scale_arg, + cache_scale_block_stride_arg, block_tables if has_bt else state_slot_mapping, # placeholder bt_seq_stride, k_per_block, - out, - out.stride(0), head_dim, rope_head_dim, BLOCK_D=BLOCK_D, @@ -406,10 +523,12 @@ def fused_compress_attn( STATE_SIZE=state_size, K=K, HAS_BLOCK_TABLE=int(has_bt), + QUANT=int(quant), + FP8_MAX=fp8_max_arg, + USE_UE8M0=int(use_ue8m0), + PRESHUFFLE=int(preshuffle), ) - return out # [num_compress, head_dim] - def fused_compress_attn_reference( *, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 1e41e72404..58ab089183 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -33,7 +33,6 @@ cp_gather_indexer_k_quant_cache, dtypes, get_hip_quant, - indexer_k_quant_and_cache, ) from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.dist.communication_op import tensor_model_parallel_all_reduce @@ -50,7 +49,7 @@ ) from atom.model_loader.loader import WeightsMapper from atom.model_ops.embed_head import VocabParallelEmbedding -from atom.model_ops.layernorm import RMSNorm, rmsnorm2d_fwd_ +from atom.model_ops.layernorm import DualRMSNorm, RMSNorm, rmsnorm2d_fwd_ from atom.model_ops.triton_rmsnorm_nw import rmsnorm_nw from atom.model_ops.linear import ( ColumnParallelLinear, @@ -122,6 +121,22 @@ def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: return rmsnorm2d_fwd_(x, ones, eps, dim) +def _make_weightless_rmsnorm(dim: int, eps: float) -> RMSNorm: + """Build an `RMSNorm(dim, eps)` whose `.weight` is `None`. + + Drops the learnable Parameter so `state_dict()` is empty for this + submodule and `load_model` doesn't expect a weight from disk (no + "unloaded parameter" warning). `DualRMSNorm` recognizes the + sentinel `weight is None` and instructs the fused kernel to skip + both the q_weight load and the multiply (Q_HAS_WEIGHT=False), + matching the prior `_rmsnorm_nw(x, eps, dim)` behavior. + """ + norm = RMSNorm(dim, eps) + del norm.weight # remove Parameter from `_parameters` + norm.weight = None # sentinel for downstream consumers + return norm + + def _hc_head_reduce_fake( x: torch.Tensor, hc_fn: torch.Tensor, @@ -664,12 +679,10 @@ def __init__( # External tensors — assigned by the owning Attention / Indexer at first forward. self.kv_cache: Optional[torch.Tensor] = None self.rotary_emb: Optional[_V4RoPE] = None - # Shared compress-output buffer (set by V4 builder.build_kv_cache_tensor - # to a per-kind torch.empty pre-allocation in `forward_vars`). When - # present, fused_compress_attn writes into it via `out=` so the GPU - # pointer stays stable across captures (CUDAGraph requirement); when - # None, falls back to the eager fresh-allocation path. - self.compress_out: Optional[torch.Tensor] = None + # FP8 quant path only: strided fp32 view of the per-block scale region + # of `self.kv_cache`. Bound by the V4 builder when `kv_cache.dtype` is + # FP8 (Indexer-inner Compressor); None for BF16 cache (Main path). + self.cache_scale: Optional[torch.Tensor] = None # State cache (per paper §3.6.1 "uncompressed tail + B-side overlap # window" portion). Indexed as a single ring buffer of size @@ -711,45 +724,41 @@ def forward( plan: "CompressPlan", state_slot_mapping: torch.Tensor, # [bs] int32 block_tables: Optional[torch.Tensor] = None, # [bs, max_blocks_per_seq] int32 - out_dtype: Optional[torch.dtype] = None, - idx_slot_mapping: Optional[torch.Tensor] = None, # [num_compress] int64 - ) -> Optional[torch.Tensor]: # [num_compress, head_dim] BF16 (None if plan empty) + ) -> None: """Batched plan-style compress: one fused kernel call for the whole fwd's batch (across all seqs). - Single fused Triton kernel does pool + RMSNorm + RoPE + bf16 kv_cache - scatter in one launch. Each compression boundary across the batch is - one row in `plan.compress_plan_gpu`. State cache update fires after - (write order critical — fused kernel reads state-cache-as-of-previous- - fwd; update_compressor_states overwrites for next fwd). - - TODO: quant (FP8 ue8m0 for CSA Main, FP4 + Hadamard for Indexer) is - NOT applied for the Main path; see class docstring. The Indexer-inner - path *does* apply FP8 quantization via `indexer_k_quant_and_cache` - when `idx_slot_mapping` is provided. + Single fused Triton kernel does pool + RMSNorm + RoPE + cache scatter + in one launch. Each compression boundary across the batch is one row + in `plan.compress_plan_gpu`. State cache update fires after (write + order critical — fused kernel reads state-cache-as-of-previous-fwd; + `update_compressor_states` overwrites for next fwd). + + Quant mode is auto-selected by `self.kv_cache.dtype`: + - BF16 cache (CSA Main / HCA Main): raw BF16 row write into + `self.kv_cache` (consumed by paged_decode/paged_prefill via + `unified_kv` per-fwd indices). + - FP8 cache (Indexer-inner): per-row amax → ue8m0 scale → fp8 cast + → preshuffled (MFMA 16x16 tile) write into `self.kv_cache`, plus + fp32 scale into `self.cache_scale` (a strided view of the same + allocation built by the V4 builder). Bit-exact with + `indexer_k_quant_and_cache` / `cp_gather_indexer_k_quant_cache` + (cache_kernels.cu:1145+). + + Side-effecting only — no return value (cache scatter IS the output). + + TODO: QAT for the BF16 Main path (FP8 round-trip per Compressor + docstring) is not yet fused. End-to-end accuracy unaffected today + because the input act_quant simulation is applied upstream. Args: x: [num_tokens, dim] flat ragged batch hidden state. plan: CompressPlan from attn_metadata.compress_plans[ratio] (or a synthetic bs=1 plan during warmup). - state_slot_mapping: [bs] int32 — per-seq state cache slot - (attn_metadata.state_slot_mapping). + state_slot_mapping: [bs] int32 — per-seq state cache slot. block_tables: [bs, max_blocks_per_seq] int32 — physical block IDs per seq; None during warmup (skips kv_cache scatter). - out_dtype: override kernel output dtype. Defaults to bf16. - idx_slot_mapping: [num_compress] int64 — global flat slot in the - FP8 indexer cache pool, one per compress row (kernel - signature requires int64). When provided, the fused - kernel skips its BF16 scatter and we instead call - `indexer_k_quant_and_cache` to FP8-quantize+write each - row into `self.kv_cache` (interleaved [head_dim] FP8 + - 4-byte fp32 scale per row, dtypes.fp8 storage). - Only set by the CSA-inner-Compressor caller. - - Returns: - Compressed KV `[num_compress, head_dim]` post-norm post-rope BF16, - in plan order (= per-seq grouped by `plan.cu_compress_cpu`). - Returns None if `plan.num_compress == 0`. + Required for the Indexer FP8 path (slot resolution). """ assert self.rotary_emb is not None, "compressor.rotary_emb must be set by owner" assert ( @@ -759,8 +768,6 @@ def forward( overlap = self.overlap d = self.head_dim rd = self.rope_head_dim - if out_dtype is None: - out_dtype = torch.bfloat16 if x.dtype == torch.float32 else x.dtype # Single fused BF16 GEMM via tgemm; otype=fp32 keeps softmax-pool # accumulator in fp32. torch.split returns zero-copy strided views; @@ -775,29 +782,21 @@ def forward( # PREVIOUS-fwd. `update_compressor_states` overwrites them with this # fwd's data for the NEXT fwd's overlap — must run AFTER the fused # kernel. - # - # The kernel does pool + RMSNorm + RoPE + (optional) BF16 store to - # kv_cache. For the Indexer-inner path (`idx_slot_mapping is not None`) - # we bypass the kernel scatter and instead call - # `indexer_k_quant_and_cache` separately, which FP8-quantizes the - # post-process K row and writes the (FP8 + scale) interleaved layout - # into the uint8 indexer pool. cos_cache, sin_cache = self.rotary_emb._caches(x.device) - is_idx_fp8 = idx_slot_mapping is not None - # Warmup runs through the same path: kv_cache is None until the V4 - # builder binds it post-warmup, so the kernel skips the scatter. - # Indexer-FP8 path always skips the kernel scatter (separate write). - scatter_kv_cache = ( - None - if is_idx_fp8 - else (self.kv_cache if block_tables is not None else None) - ) - scatter_block_tables = ( - None - if is_idx_fp8 - else (block_tables if scatter_kv_cache is not None else None) - ) - out = fused_compress_attn( + # Quant path triggers when the bound cache is FP8 (Indexer-inner). + # `self.cache_scale` is bound alongside `self.kv_cache` by the V4 + # builder when the cache is FP8 (strided fp32 view of the per-block + # scale region). + is_quant = self.kv_cache is not None and self.kv_cache.dtype != torch.bfloat16 + # Skip the kernel's cache scatter during warmup (kv_cache/block_tables + # not yet bound). + if block_tables is None or self.kv_cache is None: + scatter_kv_cache = None + scatter_block_tables = None + else: + scatter_kv_cache = self.kv_cache + scatter_block_tables = block_tables + fused_compress_attn( kv_in=kv, score_in=score, kv_state=self.kv_state, @@ -816,21 +815,12 @@ def forward( ratio=ratio, head_dim=d, rope_head_dim=rd, - out_dtype=out_dtype, - out=self.compress_out, + quant=is_quant, + cache_scale=self.cache_scale if is_quant else None, + use_ue8m0=(self.scale_fmt == "ue8m0"), + preshuffle=True, + fp8_max=(torch.finfo(self.kv_cache.dtype).max if is_quant else None), ) - # Indexer-FP8 path: quantize the post-RoPE BF16 K rows into the uint8 - # FP8+scale layout of `self.kv_cache`. `out` is [num_compress, head_dim] - # BF16 in plan order (matches `idx_slot_mapping` row order). - if is_idx_fp8 and out is not None and self.kv_cache is not None: - indexer_k_quant_and_cache( - out, - self.kv_cache, - idx_slot_mapping, - d, # quant_block_size = head_dim (one fp32 scale per row) - self.scale_fmt, - preshuffle=True, - ) update_compressor_states( kv, score, @@ -843,7 +833,6 @@ def forward( ratio=ratio, overlap=overlap, ) - return out class Indexer(nn.Module): @@ -1134,13 +1123,18 @@ def _score_topk_decode( kv_cache_4d = self.kv_cache.unsqueeze( -2 ) # [num_blocks, k1_csa, 1, head_dim+scale_dim] uint8 - # Pre-allocated decode logits buffer (sized [max_bs, max_model_len_idx] - # in builder, sliced here to [total_tokens, max_model_len_idx]). - # No `fill_(-inf)` needed — `top_k_per_row_decode` bounds each row - # by `n_committed_per_seq[batch]` so unwritten cols are never picked. - logits = indexer_meta["decode_logits_gpu"][ - :total_tokens - ] # [total_tokens, max_model_len_idx] fp32 + # Per-fwd write-once GPU scratch — no CPU mirror, no cross-fwd state. + # Under CUDAGraph capture, torch allocates from the graph's private + # memory pool and the address is stable across replays at this + # captured `total_tokens`. No `fill_(-inf)` needed — + # `top_k_per_row_decode` bounds each row by `n_committed_per_seq[batch]` + # so unwritten cols are never picked. + logits = torch.empty( + total_tokens, + self._max_model_len_idx, + dtype=torch.float32, + device=q_fp8.device, + ) deepgemm_fp8_paged_mqa_logits( q_4d, kv_cache_4d, @@ -1152,13 +1146,12 @@ def _score_topk_decode( KVBlockSize=self.kv_cache.size(1), # k1_csa = 32 Preshuffle=True, ) - # Pre-allocated indices buffer — fixed [max_bs, index_topk] int32 in - # builder, sliced to [total_tokens, index_topk] for this fwd. Kernel - # writes exactly `index_topk` ints per row (valid seq-local indices - # then -1 sentinels), no overflow risk. - topk_local = indexer_meta["decode_topk_indices_gpu"][ - :total_tokens - ] # [total_tokens, index_topk] int32 + # Per-fwd write-once int32 scratch. Kernel writes exactly `index_topk` + # ints per row (valid seq-local indices then -1 sentinels). CG-safe + # for the same reason as `logits` above. + topk_local = torch.empty( + total_tokens, self.index_topk, dtype=torch.int32, device=q_fp8.device + ) top_k_per_row_decode( logits, next_n, @@ -1243,7 +1236,12 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): prefix=f"{p}.wqkv_a", ) self.q_norm = RMSNorm(self.q_lora_rank, self.eps) - self.q_norm2 = RMSNorm(self.head_dim, self.eps) + # q_norm2: per-head Q normalization. The checkpoint has no + # `q_norm2.weight` entry, so the module is constructed with + # `weight=None` (no learnable Parameter) — `DualRMSNorm` reads the + # sentinel and tells the kernel to skip the q_weight load entirely. + # Behavior is identical to the prior `_rmsnorm_nw(q, eps, head_dim)`. + self.q_norm2 = _make_weightless_rmsnorm(self.head_dim, self.eps) self.wq_b = ColumnParallelLinear( self.q_lora_rank, self.n_heads * self.head_dim, @@ -1252,6 +1250,19 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): prefix=f"{p}.wq_b", ) self.kv_norm = RMSNorm(self.head_dim, self.eps) + # Fused per-head Q RMSNorm (identity weight) + KV RMSNorm — both have + # feature dim = head_dim, single Triton kernel via DualRMSNorm. + # `q_norm2.weight is None` tells the kernel's Q_HAS_WEIGHT=False + # path to skip the q_weight load (Q side has no learnable weight; + # equivalent to the prior `_rmsnorm_nw` helper). + self.qk_norm = DualRMSNorm( + self.q_norm2, + self.kv_norm, + num_q_heads=self.n_local_heads, + num_kv_heads=1, + head_dim=self.head_dim, + prefix=f"{p}.qk_norm", + ) # wo_a: grouped LoRA — V4QuantConfig forces this BF16 even though disk is FP8. # The grouped einsum (`bsgd,grd->bsgr`) needs BF16 weights; aiter has no FP8 einsum. self.wo_a = ColumnParallelLinear( @@ -1444,11 +1455,15 @@ def forward_impl( if _V4_FORCE_UE8M0_QUANT: qr = qr.clone() act_quant_inplace(qr, 128, "ue8m0") - q = self.wq_b(qr).view(num_tokens, self.n_local_heads, self.head_dim) - q = _rmsnorm_nw(q, self.eps, self.head_dim) + # Flat q_flat is [num_tokens, n_local_heads * head_dim]; DualRMSNorm + # views internally to per-head shape and returns flat. Single Triton + # launch fuses per-head Q RMSNorm (weightless) + KV RMSNorm (both + # head_dim=128), replacing the prior `_rmsnorm_nw + kv_norm` pair. + q_flat = self.wq_b(qr) + q_flat, kv = self.qk_norm(q_flat, kv_pre) + q = q_flat.view(num_tokens, self.n_local_heads, self.head_dim) # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes # to (1, num_tokens, ...) for aiter's per-position rope kernel. - kv = self.kv_norm(kv_pre) self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) if _V4_USE_REF_QUANT: act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) @@ -1460,7 +1475,6 @@ def forward_impl( # compress_plans, swa_write_indices, ...) is well-typed for pyright. attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) compress_plans = attn_md.compress_plans - v4_indexer_meta = attn_md.indexer_meta swa_write_indices = attn_md.swa_write_indices v4_batch_id_per_token = attn_md.batch_id_per_token block_tables_gpu = attn_md.block_tables @@ -1475,7 +1489,7 @@ def forward_impl( # the scattered compress entries from `unified_kv` via per-fwd indices. # Called purely for its side-effect (scatter into self.kv_cache). plan_for_layer = compress_plans[ratio] if ratio else None - if ratio: + if self.compressor is not None: # i.e. CSA or HCA layer, compress_ratio > 0 self.compressor( x, plan=plan_for_layer, @@ -1483,17 +1497,18 @@ def forward_impl( block_tables=block_tables_gpu, ) if self.indexer is not None: # i.e. CSA layer, compress_ratio == 4 - # Indexer's inner compressor populates the indexer kv_cache (FP8 - # via indexer_k_quant_and_cache when idx_slot_mapping is set); - # the outer `forward_batched` reads `v4_indexer_meta` (built once - # per fwd) so it has zero per-layer H2D / CPU index math. - idx_slot_mapping = v4_indexer_meta.get("compress_slot_mapping_gpu") + # Indexer's inner Compressor populates the FP8 indexer kv_cache + # via the unified `fused_compress_attn` quant path (auto-detected + # by cache dtype). `block_tables_gpu` is the same as the Main + # Compressor's; the kernel resolves `physical_block * k_per_block + # + slot_in_block` internally — no separate slot_mapping needed. + # The outer `forward_batched` consumes `v4_indexer_meta` (built + # once per fwd) so per-layer H2D / CPU index math stays zero. self.indexer.compressor( x, plan=plan_for_layer, state_slot_mapping=state_slot_mapping, block_tables=block_tables_gpu, - idx_slot_mapping=idx_slot_mapping, ) indexer_topk_batched = self.indexer.forward_batched( x_full=x, diff --git a/atom/utils/decorators.py b/atom/utils/decorators.py index eb16cd2463..e262030321 100644 --- a/atom/utils/decorators.py +++ b/atom/utils/decorators.py @@ -1,7 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Callable, Optional, TypeVar, Union +from typing import ( + Callable, + Optional, + ParamSpec, + TypeVar, + Union, + overload as _overload, +) import inspect import os import sys @@ -23,6 +30,8 @@ # from atom.utils import start_monitoring_torch_compile _T = TypeVar("_T", bound=type[nn.Module]) +_P = ParamSpec("_P") +_R = TypeVar("_R") context_manager = None torch_compile_start_time: float = 0.0 @@ -167,6 +176,18 @@ def wrapped(*args, **kwargs): return wrapped +# Overloads preserve the decorated callable's signature for pyright. +# Without them, `@mark_trace` instances were inferred as plain `Callable`, +# and class instances whose `__call__` was wrapped (DualRMSNorm, LinearBase, +# ...) were flagged as "not callable" at the call site. +@_overload +def mark_trace(func: Callable[_P, _R], /) -> Callable[_P, _R]: ... +@_overload +def mark_trace( + *, + torch_compile: bool = ..., + prefix: Optional[str] = ..., +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: ... def mark_trace( func: Optional[Callable] = None, *, From 9997e89d644d9b820d01e7a42ba867a211ce1224 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Sat, 9 May 2026 15:11:55 +0800 Subject: [PATCH 06/76] CI: gate ATOM tests on Pre Checkin workflow status (#724) * CI: add runner model-cache diagnostics in atom tests. * CI: gate ATOM tests on Pre Checkin workflow status Replace the signal artifact handoff with direct workflow run status checks so downstream CI avoids artifact lookup latency and pagination issues. --- .github/scripts/check_signal.sh | 164 +++++++++++++++++++----- .github/workflows/atom-sglang-test.yaml | 3 +- .github/workflows/atom-test.yaml | 2 +- .github/workflows/atom-vllm-test.yaml | 2 +- .github/workflows/pre-checks.yaml | 37 ------ 5 files changed, 132 insertions(+), 76 deletions(-) diff --git a/.github/scripts/check_signal.sh b/.github/scripts/check_signal.sh index 0e24f38e47..1037ff0cbe 100644 --- a/.github/scripts/check_signal.sh +++ b/.github/scripts/check_signal.sh @@ -1,52 +1,146 @@ #!/usr/bin/env bash +# Gate a downstream workflow on the Pre Checkin workflow run for the current commit. +# Exit 0 on success, 78 (neutral/skip) on any other conclusion, 1 if the run +# cannot be located within the retry budget. + set -euo pipefail -TARGET_SHA="${GITHUB_SHA:-${1:-}}" -if [ -z "${TARGET_SHA}" ]; then - echo "GITHUB_SHA is not set and no SHA argument was provided." - exit 1 -fi +CHECKS_WORKFLOW_NAME="${CHECKS_WORKFLOW_NAME:-Pre Checkin}" +MAX_RETRIES="${MAX_RETRIES:-5}" +RETRY_INTERVAL_SECONDS="${RETRY_INTERVAL_SECONDS:-30}" +REPO="${GITHUB_REPOSITORY:-}" + +get_target_branch() { + if [ -n "${GITHUB_HEAD_REF:-}" ]; then + printf '%s\n' "${GITHUB_HEAD_REF}" + return + fi + + if [ -n "${GITHUB_REF_NAME:-}" ]; then + printf '%s\n' "${GITHUB_REF_NAME}" + return + fi + + python3 - <<'PY' +import json +import os + +event_path = os.environ.get("GITHUB_EVENT_PATH") +if event_path and os.path.exists(event_path): + with open(event_path, encoding="utf-8") as fh: + data = json.load(fh) + print(data.get("pull_request", {}).get("head", {}).get("ref", "")) +else: + print("") +PY +} + +get_target_head_sha() { + case "${GITHUB_EVENT_NAME:-}" in + pull_request|pull_request_target) + python3 - <<'PY' +import json +import os + +event_path = os.environ.get("GITHUB_EVENT_PATH") +if event_path and os.path.exists(event_path): + with open(event_path, encoding="utf-8") as fh: + data = json.load(fh) + print(data.get("pull_request", {}).get("head", {}).get("sha", "")) +else: + print("") +PY + ;; + *) + printf '%s\n' "${GITHUB_SHA:-}" + ;; + esac +} + +find_checks_run_id() { + local target_branch target_head_sha + target_branch="$(get_target_branch)" + target_head_sha="$(get_target_head_sha)" + + if [ -z "${REPO}" ]; then + echo "GITHUB_REPOSITORY is required to locate the ${CHECKS_WORKFLOW_NAME} workflow run." >&2 + return 1 + fi + + if [ -z "${target_head_sha}" ]; then + echo "Could not determine the target head SHA for the ${CHECKS_WORKFLOW_NAME} workflow run." >&2 + return 1 + fi + + local -a gh_args=( + run list + --repo "${REPO}" + --workflow "${CHECKS_WORKFLOW_NAME}" + --limit 20 + --json databaseId,headSha,headBranch,event,createdAt,status + ) + + if [ -n "${target_branch}" ]; then + gh_args+=(--branch "${target_branch}") + fi + + # Nightly and reusable workflows reuse the Pre Checkin result from the + # original push or pull_request run on the same SHA. + if [ -n "${GITHUB_EVENT_NAME:-}" ] \ + && [ "${GITHUB_EVENT_NAME}" != "schedule" ] \ + && [ "${GITHUB_EVENT_NAME}" != "workflow_call" ]; then + gh_args+=(--event "${GITHUB_EVENT_NAME}") + fi + + gh "${gh_args[@]}" \ + --jq "(map(select(.headSha == \"${target_head_sha}\")) | first | .databaseId) // empty" +} -ARTIFACT_NAME="checks-signal-${TARGET_SHA}" -MAX_RETRIES="${MAX_RETRIES:-10}" -RETRY_DELAY_SECONDS="${RETRY_DELAY_SECONDS:-30}" +# Echoes "\t" for the given run. +fetch_run_state() { + local run_id="$1" + gh api "repos/${REPO}/actions/runs/${run_id}" \ + --jq '[.status, (.conclusion // "")] | @tsv' +} + +print_failed_jobs() { + local run_id="$1" + gh api -X GET "repos/${REPO}/actions/runs/${run_id}/jobs" --paginate \ + --jq '.jobs[] | select(.conclusion != null and .conclusion != "success" and .conclusion != "skipped") | "FAILED: \(.name) (\(.conclusion))"' \ + || true +} for i in $(seq 1 "${MAX_RETRIES}"); do - ATTEMPT_DIR="$(mktemp -d)" - echo "Attempt ${i}: downloading artifact ${ARTIFACT_NAME}..." - - if gh run download \ - --repo "${GITHUB_REPOSITORY}" \ - --name "${ARTIFACT_NAME}" \ - --dir "${ATTEMPT_DIR}"; then - SIGNAL_FILE="${ATTEMPT_DIR}/checks_signal.txt" - if [ -f "${SIGNAL_FILE}" ]; then - echo "Artifact ${ARTIFACT_NAME} downloaded successfully." - SIGNAL="$(head -n 1 "${SIGNAL_FILE}")" - if [ "${SIGNAL}" = "success" ]; then + echo "Attempt ${i}: Locating ${CHECKS_WORKFLOW_NAME} workflow run..." + + RUN_ID="$(find_checks_run_id || true)" + if [ -z "${RUN_ID}" ]; then + echo "Attempt ${i}: Matching ${CHECKS_WORKFLOW_NAME} run not found yet." + else + STATE="$(fetch_run_state "${RUN_ID}" || true)" + STATUS="${STATE%%$'\t'*}" + CONCLUSION="${STATE#*$'\t'}" + # If STATE is unexpected treat that as "no conclusion yet". + [ "${CONCLUSION}" = "${STATE}" ] && CONCLUSION="" + + echo "Attempt ${i}: run ${RUN_ID} status=${STATUS:-unknown} conclusion=${CONCLUSION:-}" + + if [ "${STATUS}" = "completed" ]; then + if [ "${CONCLUSION}" = "success" ]; then echo "Pre Checkin passed, continuing workflow." - rm -rf "${ATTEMPT_DIR}" exit 0 fi - if [ "${SIGNAL}" = "failure" ]; then - echo "Pre Checkin failed, skipping workflow. Details:" - tail -n +2 "${SIGNAL_FILE}" || true - rm -rf "${ATTEMPT_DIR}" - exit 78 - fi - - echo "Unknown signal '${SIGNAL}' in ${SIGNAL_FILE}." - rm -rf "${ATTEMPT_DIR}" - exit 1 + echo "Pre Checkin did not pass (conclusion: ${CONCLUSION:-unknown}), skipping workflow." + print_failed_jobs "${RUN_ID}" + exit 78 # 78 = neutral/skip fi fi - rm -rf "${ATTEMPT_DIR}" - echo "Artifact not found yet, retrying in ${RETRY_DELAY_SECONDS}s..." - sleep "${RETRY_DELAY_SECONDS}" + echo "Pre Checkin not ready yet, retrying in ${RETRY_INTERVAL_SECONDS}s..." + sleep "${RETRY_INTERVAL_SECONDS}" done -echo "Failed to download ${ARTIFACT_NAME} after ${MAX_RETRIES} attempts." +echo "Failed to read Pre Checkin status after ${MAX_RETRIES} attempts. Exiting workflow." exit 1 diff --git a/.github/workflows/atom-sglang-test.yaml b/.github/workflows/atom-sglang-test.yaml index c0ac04638a..4cd447d1da 100644 --- a/.github/workflows/atom-sglang-test.yaml +++ b/.github/workflows/atom-sglang-test.yaml @@ -35,7 +35,7 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow run: bash ./.github/scripts/check_signal.sh env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -434,4 +434,3 @@ jobs: rm -f Dockerfile.mod || true docker rmi "atom_sglang_base:ci" || true docker rmi "atom_sglang:ci" || true - diff --git a/.github/workflows/atom-test.yaml b/.github/workflows/atom-test.yaml index 769d34cc4b..6b6a830176 100644 --- a/.github/workflows/atom-test.yaml +++ b/.github/workflows/atom-test.yaml @@ -47,7 +47,7 @@ jobs: if: ${{ github.event_name != 'workflow_dispatch' }} uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow if: ${{ github.event_name != 'workflow_dispatch' }} run: bash ./.github/scripts/check_signal.sh env: diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index 798ec69622..9df3fa486a 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.github/workflows/atom-vllm-test.yaml @@ -35,7 +35,7 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 - - name: Download and check pre-checkin signal + - name: Wait for Pre Checkin workflow run: bash ./.github/scripts/check_signal.sh env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/pre-checks.yaml b/.github/workflows/pre-checks.yaml index 4d4e799e2d..8f6426afb1 100644 --- a/.github/workflows/pre-checks.yaml +++ b/.github/workflows/pre-checks.yaml @@ -62,40 +62,3 @@ jobs: -reporter=github-pr-review \ -filter-mode=diff_context \ -fail-on-error=true - - upload-success-artifact: - name: Upload Success Signal - runs-on: ubuntu-latest - needs: [black, ruff] - if: ${{ needs.black.result == 'success' && needs.ruff.result == 'success' }} - steps: - - name: Create success signal file - run: echo "success" > checks_signal.txt - - - name: Upload success artifact - uses: actions/upload-artifact@v4 - with: - name: checks-signal-${{ github.sha }} - path: checks_signal.txt - - upload-failure-artifact: - name: Upload Failure Signal - runs-on: ubuntu-latest - needs: [black, ruff] - if: ${{ always() && (needs.black.result != 'success' || needs.ruff.result != 'success') }} - steps: - - name: Create failure signal file with failed jobs - run: | - echo "failure" > checks_signal.txt - if [ "${{ needs.black.result }}" != "success" ]; then - echo "FAILED: black (${{ needs.black.result }})" >> checks_signal.txt - fi - if [ "${{ needs.ruff.result }}" != "success" ]; then - echo "FAILED: ruff (${{ needs.ruff.result }})" >> checks_signal.txt - fi - - - name: Upload failure artifact - uses: actions/upload-artifact@v4 - with: - name: checks-signal-${{ github.sha }} - path: checks_signal.txt From abdb093d30d2bfccb09717a5a720c7631a9567ce Mon Sep 17 00:00:00 2001 From: PerryZhang01 Date: Sat, 9 May 2026 20:21:16 +0800 Subject: [PATCH 07/76] [fix](gpt-oss): change the accuracy test for gpt-oss (#720) * [fix](gpt-oss): change the accuracy test for gpt-oss * [fix](ci): change gpt-oss accuracy test in ci * [feat](ci): add client command for gpt-oss --------- Co-authored-by: perzhang --- .claude/commands/ci-pr-guide.md | 4 +- .github/benchmark/models.json | 2 +- .github/benchmark/models_accuracy.json | 14 +-- .github/benchmark/oot_benchmark_models.json | 2 +- .github/benchmark/oot_models_accuracy.json | 10 +- .github/scripts/atom_oot_test.sh | 95 +++++++++++++++++-- .github/scripts/atom_test.sh | 63 ++++++++++-- .github/workflows/atom-test.yaml | 2 + .../atom-vllm-accuracy-validation.yaml | 8 +- .github/workflows/atom-vllm-test.yaml | 4 +- recipes/GPT-OSS.md | 8 +- recipes/atom_vllm/GPT-OSS.md | 2 +- 12 files changed, 177 insertions(+), 37 deletions(-) diff --git a/.claude/commands/ci-pr-guide.md b/.claude/commands/ci-pr-guide.md index cc8226e91c..2916b94a3c 100644 --- a/.claude/commands/ci-pr-guide.md +++ b/.claude/commands/ci-pr-guide.md @@ -37,8 +37,8 @@ PR/Push → pre-checks (black + ruff) | DeepSeek-R1-0528 MTP | `--kv_cache_dtype fp8 -tp 8 --method mtp` | 0.94 | Yes | 8gpu | | DeepSeek-R1-0528 MXFP4 | `--kv_cache_dtype fp8 -tp 8` | 0.93 | Yes | 8gpu | | DeepSeek-R1-0528 MXFP4 MTP | `--kv_cache_dtype fp8 -tp 8 --method mtp` | 0.93 | Yes | 8gpu | -| gpt-oss-120b | `--kv_cache_dtype fp8 --gpu-memory-utilization 0.3` | 0.38 | Yes | mi355-1 | -| gpt-oss-120b (2 GPU) | `-tp 2 --enable-dp-attention --enable-expert-parallel` | 0.38 | No | mi355-4 | +| gpt-oss-120b | `--kv_cache_dtype fp8 --gpu-memory-utilization 0.5` | 0.88 | Yes | mi355-1 | +| gpt-oss-120b (2 GPU) | `-tp 2 --enable-dp-attention --enable-expert-parallel` | 0.88 | No | mi355-4 | | Qwen3-235B FP8 | `-tp 8 --enable-expert-parallel` + `ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1` | 0.87 | Yes | 8gpu | | Qwen3-235B MXFP4 | `-tp 8 --enable-expert-parallel` + `ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1` | 0.87 | No | 8gpu | | Qwen3-Next-80B-A3B | `-tp 8` | 0.65 | Yes | 8gpu | diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index b542d3cf53..001c430900 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -73,7 +73,7 @@ "display": "gpt-oss-120b", "path": "openai/gpt-oss-120b", "prefix": "gpt-oss-120b", - "args": "--kv_cache_dtype fp8", + "args": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index a86ca34914..29b8e6bcd3 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -50,12 +50,13 @@ { "model_name": "gpt-oss-120b", "model_path": "openai/gpt-oss-120b", - "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.3", + "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", "env_vars": "", "runner": "linux-atom-mi35x-1", "test_level": "pr", - "accuracy_threshold": 0.39, - "accuracy_baseline": 0.41, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b", "_baseline_note": "No public GSM8K baseline available" }, @@ -122,12 +123,13 @@ { "model_name": "gpt-oss-120b (2 GPUs)", "model_path": "openai/gpt-oss-120b", - "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.3", + "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", "env_vars": "", "runner": "linux-atom-mi35x-4", "test_level": "main", - "accuracy_threshold": 0.39, - "accuracy_baseline": 0.41, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b", "_baseline_note": "No public GSM8K baseline available" }, diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index d6feac4c6e..5bc5585e84 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -49,7 +49,7 @@ "path": "openai/gpt-oss-120b", "prefix": "gpt-oss-120b", "nightly_group": "C", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index 4a259549a4..fdbeb1bd9a 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -152,22 +152,24 @@ "model_name": "gpt-oss-120b TP1", "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 1", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-1", "test_level": "nightly", - "accuracy_threshold": 0.38, - "accuracy_baseline": 0.38, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b" }, { "model_name": "gpt-oss-120b TP2", "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 2", + "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", - "accuracy_threshold": 0.38, - "accuracy_baseline": 0.38, + "accuracy_threshold": 0.88, + "accuracy_baseline": 0.90, "accuracy_baseline_model": "openai/gpt-oss-120b" }, { diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index fe2984b871..107f95f999 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -49,6 +49,7 @@ KEEP_SERVER_ALIVE_ON_EXIT=${KEEP_SERVER_ALIVE_ON_EXIT:-0} EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} +EXPLICIT_CLIENT_COMMAND=${OOT_CLIENT_COMMAND:-} OOT_DOCKER_IMAGE=${OOT_DOCKER_IMAGE:-} LM_EVAL_NUM_FEWSHOT=${LM_EVAL_NUM_FEWSHOT:-3} LAST_VLLM_LOG_LINE=0 @@ -64,7 +65,7 @@ if [[ -n "${EXPLICIT_MODEL_NAME}" || -n "${EXPLICIT_MODEL_PATH}" || -n "${EXPLIC echo "OOT_MODEL_NAME and OOT_MODEL_PATH must both be set when using explicit model overrides." exit 2 fi - ACTIVE_MODELS=("${EXPLICIT_MODEL_NAME}|${EXPLICIT_MODEL_PATH}|${EXPLICIT_EXTRA_ARGS}") + ACTIVE_MODELS=("${EXPLICIT_MODEL_NAME}|${EXPLICIT_MODEL_PATH}|${EXPLICIT_EXTRA_ARGS}|${EXPLICIT_CLIENT_COMMAND}") else echo "${MODE} mode requires OOT_MODEL_NAME and OOT_MODEL_PATH env vars from the workflow." exit 2 @@ -81,6 +82,14 @@ resolve_model_path() { fi } +# gpt-oss OpenAI weights require Chat Completions (messages); local-completions sends "prompt" and vLLM returns 400. +is_gpt_oss_model() { + local name_lc path_lc + name_lc="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" + path_lc="$(printf '%s' "$2" | tr '[:upper:]' '[:lower:]')" + [[ "${name_lc}" == *gpt-oss* || "${path_lc}" == *gpt-oss* ]] +} + emit_new_vllm_logs() { if [[ "${STREAM_VLLM_LOGS}" != "1" || ! -f "${VLLM_LOG_FILE}" ]]; then return 0 @@ -210,6 +219,7 @@ accuracy_one_model() { local model_name="$1" local model_path="$2" local extra_args="$3" + local client_command="${4:-}" local flat_result_file="" local resolved_model_path @@ -231,11 +241,80 @@ accuracy_one_model() { echo "Model name: ${model_name}" echo "Few-shot count: ${LM_EVAL_NUM_FEWSHOT}" - lm_eval --model local-completions \ - --model_args model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ - --tasks gsm8k \ - --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ - --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + if [[ -n "${client_command}" ]]; then + local -a client_command_args=() + while IFS= read -r -d '' token; do + client_command_args+=("${token}") + done < <( + CLIENT_COMMAND="${client_command}" \ + MODEL_PATH_VALUE="${resolved_model_path}" \ + OUTPUT_PATH_VALUE="${output_path}" \ + LM_EVAL_NUM_FEWSHOT_VALUE="${LM_EVAL_NUM_FEWSHOT}" \ + VLLM_PORT_VALUE="${VLLM_PORT}" \ + python3 - <<'PY' +import os +import shlex +import sys + +client_command = os.environ["CLIENT_COMMAND"] +replacements = { + "${MODEL_PATH}": os.environ["MODEL_PATH_VALUE"], + "$MODEL_PATH": os.environ["MODEL_PATH_VALUE"], + "${OUTPUT_PATH}": os.environ["OUTPUT_PATH_VALUE"], + "$OUTPUT_PATH": os.environ["OUTPUT_PATH_VALUE"], + "${LM_EVAL_NUM_FEWSHOT}": os.environ["LM_EVAL_NUM_FEWSHOT_VALUE"], + "$LM_EVAL_NUM_FEWSHOT": os.environ["LM_EVAL_NUM_FEWSHOT_VALUE"], + "${VLLM_PORT}": os.environ["VLLM_PORT_VALUE"], + "$VLLM_PORT": os.environ["VLLM_PORT_VALUE"], +} +for src, dst in replacements.items(): + client_command = client_command.replace(src, dst) + +for token in shlex.split(client_command): + sys.stdout.write(token) + sys.stdout.write("\0") +PY + ) + + if [[ ${#client_command_args[@]} -eq 0 ]]; then + echo "ERROR: client_command is set but empty after parsing." + return 2 + fi + + for arg in "${client_command_args[@]}"; do + if [[ "${arg}" =~ \$\{[A-Z0-9_]+\} ]] || [[ "${arg}" =~ \$[A-Z_][A-Z0-9_]* ]]; then + echo "ERROR: client_command contains unresolved placeholder after expansion: ${arg}" + return 2 + fi + done + + echo "Using custom lm-eval command from client_command: ${client_command}" + "${client_command_args[@]}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + else + local lm_args=( + --model_args + model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True + ) + if is_gpt_oss_model "${model_name}" "${model_path}"; then + echo "Using chat completions + apply_chat_template for gpt-oss (OpenAI-compatible messages API)." + lm_args=( + --model_args + model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/chat/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True + ) + lm_eval --model local-chat-completions \ + --apply_chat_template \ + "${lm_args[@]}" \ + --tasks gsm8k \ + --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ + --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + else + lm_eval --model local-completions \ + "${lm_args[@]}" \ + --tasks gsm8k \ + --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ + --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" + fi + fi # lm-eval output layout differs across versions: output_path may be a file # or a directory containing one/more JSON files. Follow native CI style: @@ -328,7 +407,7 @@ run_for_models() { local matched=0 for entry in "${ACTIVE_MODELS[@]}"; do - IFS='|' read -r model_name model_path extra_args <<< "${entry}" + IFS='|' read -r model_name model_path extra_args client_command <<< "${entry}" if [[ -n "${SELECTED_MODEL}" && "${SELECTED_MODEL}" != "${model_name}" ]]; then continue @@ -342,7 +421,7 @@ run_for_models() { # accuracy mode: launch + evaluate each selected model, then stop server. launch_one_model "${model_name}" "${model_path}" "${extra_args}" - accuracy_one_model "${model_name}" "${model_path}" "${extra_args}" + accuracy_one_model "${model_name}" "${model_path}" "${extra_args}" "${client_command}" stop_server done diff --git a/.github/scripts/atom_test.sh b/.github/scripts/atom_test.sh index 78f6c22d7c..c36eee2971 100755 --- a/.github/scripts/atom_test.sh +++ b/.github/scripts/atom_test.sh @@ -121,12 +121,63 @@ if [ "$TYPE" == "accuracy" ]; then RUN_TAG=$(date +%Y%m%d%H%M%S) OUTPUT_PATH=accuracy_test_results/${RUN_TAG} FLAT_RESULT_FILE=accuracy_test_results/${RUN_TAG}.json - lm_eval --model local-completions \ - --model_args model="$MODEL_PATH",base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=3,tokenized_requests=False,trust_remote_code=True \ - --tasks gsm8k \ - --num_fewshot 3 \ - --output_path "${OUTPUT_PATH}" \ - 2>&1 | tee "$ATOM_CLIENT_LOG" + CLIENT_COMMAND="${CLIENT_COMMAND:-}" + if [[ "${CLIENT_COMMAND}" == "null" ]]; then + CLIENT_COMMAND="" + fi + + if [[ -n "${CLIENT_COMMAND}" ]]; then + CLIENT_COMMAND_ARGS=() + while IFS= read -r -d '' token; do + CLIENT_COMMAND_ARGS+=("${token}") + done < <( + CLIENT_COMMAND="${CLIENT_COMMAND}" \ + MODEL_PATH_VALUE="${MODEL_PATH}" \ + OUTPUT_PATH_VALUE="${OUTPUT_PATH}" \ + python3 - <<'PY' +import os +import shlex +import sys + +client_command = os.environ["CLIENT_COMMAND"] +replacements = { + "${MODEL_PATH}": os.environ["MODEL_PATH_VALUE"], + "$MODEL_PATH": os.environ["MODEL_PATH_VALUE"], + "${OUTPUT_PATH}": os.environ["OUTPUT_PATH_VALUE"], + "$OUTPUT_PATH": os.environ["OUTPUT_PATH_VALUE"], +} +for src, dst in replacements.items(): + client_command = client_command.replace(src, dst) + +for token in shlex.split(client_command): + sys.stdout.write(token) + sys.stdout.write("\0") +PY + ) + + if [[ ${#CLIENT_COMMAND_ARGS[@]} -eq 0 ]]; then + echo "ERROR: CLIENT_COMMAND is set but empty after parsing." + exit 2 + fi + + for arg in "${CLIENT_COMMAND_ARGS[@]}"; do + if [[ "${arg}" =~ \$\{[A-Z0-9_]+\} ]] || [[ "${arg}" =~ \$[A-Z_][A-Z0-9_]* ]]; then + echo "ERROR: CLIENT_COMMAND contains unresolved placeholder after expansion: ${arg}" + exit 2 + fi + done + + echo "Using custom lm-eval command from client_command: ${CLIENT_COMMAND}" + "${CLIENT_COMMAND_ARGS[@]}" 2>&1 | tee "$ATOM_CLIENT_LOG" + else + echo "Using default lm-eval command." + lm_eval --model local-completions \ + --model_args "model=${MODEL_PATH},base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot 3 \ + --output_path "${OUTPUT_PATH}" \ + 2>&1 | tee "$ATOM_CLIENT_LOG" + fi RESULT_FILENAME=$( python3 - < Date: Sat, 9 May 2026 20:23:17 +0800 Subject: [PATCH 08/76] Add triton fallback for deepseek & gptoss (#721) * add triton fallback for ds & gptoss * fix format * add triton mha layout * remove useless * update limit * refactor block table --- atom/model_loader/loader.py | 73 +++++---- atom/model_ops/attention_mha.py | 185 ++++++++++++++-------- atom/model_ops/attention_mla.py | 198 ++++++++++++++---------- atom/model_ops/attentions/triton_mha.py | 131 ++++++++++++++++ atom/model_ops/attentions/triton_mla.py | 101 ++++++++++++ atom/model_ops/fused_moe_triton.py | 69 ++++++--- atom/model_ops/moe.py | 19 ++- atom/utils/envs.py | 10 ++ atom/utils/selector.py | 13 +- 9 files changed, 574 insertions(+), 225 deletions(-) create mode 100644 atom/model_ops/attentions/triton_mha.py create mode 100644 atom/model_ops/attentions/triton_mla.py diff --git a/atom/model_loader/loader.py b/atom/model_loader/loader.py index 987b05e18a..df2fbf02ad 100644 --- a/atom/model_loader/loader.py +++ b/atom/model_loader/loader.py @@ -318,6 +318,20 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: with concurrent.futures.ThreadPoolExecutor() as executor: futures = [] + use_threadpool = envs.ATOM_LOADER_USE_THREADPOOL + if use_threadpool: + executor = concurrent.futures.ThreadPoolExecutor() + else: + executor = None + futures = [] + + def _submit(fn, *args): + if executor is not None: + futures.append(executor.submit(fn, *args)) + else: + fn(*args) + + try: disable_mmap = envs.ATOM_DISABLE_MMAP for name, weight_tensor in safetensors_weights_iterator( model_name_or_path, disable_mmap=disable_mmap @@ -389,11 +403,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: ) continue weight_loader = getattr(param, "weight_loader") - futures.append( - executor.submit( - weight_loader, param, weight_tensor, shard_idx - ) - ) + _submit(weight_loader, param, weight_tensor, shard_idx) loaded_weights_record.add(prefix + param_name) else: # Checkpoint has separate weights, load into fused param @@ -407,12 +417,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: dropped_ckpt_keys.append((_orig_ckpt_name, param_name)) break weight_loader = getattr(param, "weight_loader") - # weight_loader(param, weight_tensor, shard_id) - futures.append( - executor.submit( - weight_loader, param, weight_tensor, shard_id - ) - ) + _submit(weight_loader, param, weight_tensor, shard_id) loaded_weights_record.add(prefix + param_name) break else: @@ -482,15 +487,13 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: matched = True break weight_loader = getattr(param, "weight_loader") - futures.append( - executor.submit( - weight_loader, - param, - weight_tensor, - name, - shard_id, - expert_id, - ) + _submit( + weight_loader, + param, + weight_tensor, + name, + shard_id, + expert_id, ) loaded_weights_record.add(prefix + name) matched = True @@ -508,15 +511,13 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit( - weight_loader, - param, - weight_tensor, - "", # use merged moe loader - "", - expert_id, - ) + _submit( + weight_loader, + param, + weight_tensor, + "", # use merged moe loader + "", + expert_id, ) loaded_weights_record.add(prefix + name) try: @@ -527,9 +528,7 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - futures.append( - executor.submit(weight_loader, param, weight_tensor) - ) + _submit(weight_loader, param, weight_tensor) loaded_weights_record.add(prefix + name) else: # Model doesn't have expert mapping, use generic loading @@ -541,12 +540,12 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: weight_loader = getattr( param, "weight_loader", default_weight_loader ) - # weight_loader(param, weight_tensor) - futures.append(executor.submit(weight_loader, param, weight_tensor)) + _submit(weight_loader, param, weight_tensor) loaded_weights_record.add(prefix + name) - # Wait for all tasks to complete and raise any exceptions. - for future in concurrent.futures.as_completed(futures): - future.result() + finally: + if executor is not None: + concurrent.futures.wait(futures) + executor.shutdown(wait=True) # Verify every model parameter actually got loaded from the checkpoint. # Without this check, weights_mapping bugs (e.g. a substring rule diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index c3ebd525a8..cea626ccf5 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -77,6 +77,10 @@ def __init__( self.rotary_emb = rotary_emb self.q_norm = q_norm self.k_norm = k_norm + # Set by the attention backend's build_kv_cache_tensor when KV cache is + # allocated in flash layout [num_blocks, block_size, num_kv_heads, head_dim] + # for aiter triton unified_attention. AiterBackend keeps this False. + self.use_flash_layout = False # for plugin mode(vllm), the query quant is disabled for now if is_vllm(): @@ -229,7 +233,7 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): k_scale, v_scale, self.rotary_emb.is_neox_style, - flash_layout=False, + flash_layout=self.use_flash_layout, apply_scale=self.kv_cache_dtype.startswith("fp8"), offs=None, q_out=q, @@ -372,71 +376,101 @@ def paged_attention_triton( o = torch.empty_like(q) num_seqs = attn_metadata.context_lens.shape[0] - _, num_q_heads_total, head_size = q.shape - num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape - # assume all query have same length - query_group_size = attn_metadata.max_seqlen_q * ( - num_q_heads_total // num_kv_heads - ) - assert num_q_heads_total % num_kv_heads == 0 - max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) + if self.use_flash_layout: + sliding_window = ( + (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) + ) - context_partition_size = 256 - if self.sliding_window > 0: - max_context_partition_num = 1 - context_partition_size = 128 + # KV cache is already in flash layout (4D), allocated by + # TritonMHAMetadataBuilder.build_kv_cache_tensor. + nkv = k_cache.shape[2] + descale_shape = (num_seqs, nkv) - # Output buffers (same as Triton) - intermediate_shape = ( - num_seqs, - num_kv_heads, - max_context_partition_num, - query_group_size, - ) - exp_sums = torch.empty(intermediate_shape, dtype=torch.float32, device=q.device) - max_logits = torch.empty( - intermediate_shape, dtype=torch.float32, device=q.device - ) - temporary_output = torch.empty( - *intermediate_shape, - head_size, - dtype=q.dtype, - device=q.device, - ) + unified_attention( + q, + k_cache, + v_cache, + o, + cu_seqlens_q=attn_metadata.cu_seqlens_q, + seqused_k=attn_metadata.context_lens, + max_seqlen_q=attn_metadata.max_seqlen_q, + max_seqlen_k=attn_metadata.max_seqlen_k, + softmax_scale=self.scale, + causal=True, + alibi_slopes=None, + window_size=sliding_window, + block_table=attn_metadata.block_tables, + softcap=0, + q_descale=None, + k_descale=self.kv_scale.expand(descale_shape), + v_descale=self.kv_scale.expand(descale_shape), + sinks=self.sinks, + ) + else: + _, num_q_heads_total, head_size = q.shape + num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape + query_group_size = attn_metadata.max_seqlen_q * ( + num_q_heads_total // num_kv_heads + ) + assert num_q_heads_total % num_kv_heads == 0 - if k_scale is not None and k_scale.numel() > 1: - k_scale = k_scale.unsqueeze(-1) - v_scale = v_scale.unsqueeze(-1) + max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) - compute_type = ( - torch.bfloat16 - if self.kv_cache_dtype == "bf16" # or per_tensor - else aiter.dtypes.fp8 - ) - torch.ops.aiter.pa_decode_gluon( - o, - q, - k_cache, - v_cache, - attn_metadata.context_lens, - attn_metadata.block_tables, - self.scale, - attn_metadata.max_seqlen_q, - max_context_partition_num, - context_partition_size, - compute_type, - None, # q_scale - None if self.kv_cache_dtype == "bf16" else k_scale, - None if self.kv_cache_dtype == "bf16" else v_scale, - exp_sums=exp_sums, - max_logits=max_logits, - temporary_output=temporary_output, - alibi_slopes=None, - sinks=self.sinks, - sliding_window=self.sliding_window, - ps=True, - ) + context_partition_size = 256 + if self.sliding_window > 0: + max_context_partition_num = 1 + context_partition_size = 128 + + intermediate_shape = ( + num_seqs, + num_kv_heads, + max_context_partition_num, + query_group_size, + ) + exp_sums = torch.empty( + intermediate_shape, dtype=torch.float32, device=q.device + ) + max_logits = torch.empty( + intermediate_shape, dtype=torch.float32, device=q.device + ) + temporary_output = torch.empty( + *intermediate_shape, + head_size, + dtype=q.dtype, + device=q.device, + ) + + if k_scale is not None and k_scale.numel() > 1: + k_scale = k_scale.unsqueeze(-1) + v_scale = v_scale.unsqueeze(-1) + + compute_type = ( + torch.bfloat16 if self.kv_cache_dtype == "bf16" else aiter.dtypes.fp8 + ) + torch.ops.aiter.pa_decode_gluon( + o, + q, + k_cache, + v_cache, + attn_metadata.context_lens, + attn_metadata.block_tables, + self.scale, + attn_metadata.max_seqlen_q, + max_context_partition_num, + context_partition_size, + compute_type, + None, # q_scale + None if self.kv_cache_dtype == "bf16" else k_scale, + None if self.kv_cache_dtype == "bf16" else v_scale, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + alibi_slopes=None, + sinks=self.sinks, + sliding_window=self.sliding_window, + ps=True, + ) return o @@ -540,19 +574,34 @@ def prefill_attention_triton( # value: [num_blocks, 1, num_kv_heads, head_size] attn_metadata = fwd_ctx.attn_metadata - block_tables = attn_metadata.block_tables o = torch.empty_like(q) - descale_shape = (attn_metadata.cu_seqlens_q.shape[0] - 1, k.shape[1]) + num_seqs = attn_metadata.cu_seqlens_q.shape[0] - 1 + descale_shape = (num_seqs, k.shape[1]) sliding_window = ( (self.sliding_window - 1, 0) if self.sliding_window is not None else (-1, -1) ) + + # `block_tables` is always populated by TritonMHAMetadataBuilder. + # For pure prefill (no cached tokens) it is the fake table built in + # prepare_prefill that maps seq i to token indices + # [cu_seqlens_k[i], ..., cu_seqlens_k[i+1]-1], paired with raw K/V + # treated as kv_cache with block_size=1. + if attn_metadata.has_cached: + k_for_attn = k_cache + v_for_attn = v_cache + else: + # k: [total_tokens, num_kv_heads, head_size] + # -> [total_tokens, 1, num_kv_heads, head_size] + k_for_attn = k.unsqueeze(1) + v_for_attn = v.unsqueeze(1) + unified_attention( q, - k_cache, - v_cache, + k_for_attn, + v_for_attn, o, cu_seqlens_q=attn_metadata.cu_seqlens_q, seqused_k=attn_metadata.context_lens, @@ -562,7 +611,7 @@ def prefill_attention_triton( causal=True, alibi_slopes=None, window_size=sliding_window, - block_table=block_tables, + block_table=attn_metadata.block_tables, softcap=0, q_descale=None, k_descale=self.kv_scale.expand(descale_shape), @@ -577,9 +626,11 @@ def dispatch_backend(self, fwd_ctx: ForwardContext): ctx = fwd_ctx.context if ctx.is_prefill: + if self.use_flash_layout: + return self.prefill_attention_triton return self.prefill_attention else: - if self.use_triton_attn: + if self.use_triton_attn or self.use_flash_layout: return self.paged_attention_triton else: # Only use pa persistent when block_size == 1024 diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 556ec8cf07..3939532217 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -547,94 +547,124 @@ def _forward_decode( device=q.device, ) - kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) - paged_cu_seqlens_q = attn_metadata.cu_seqlens_q - paged_kv_indptr = attn_metadata.kv_indptr - paged_kv_indices = attn_metadata.kv_indices - paged_kv_last_page_lens = attn_metadata.kv_last_page_lens - max_q_len = attn_metadata.max_seqlen_q - if self.topk_indices_buffer is not None: - if attn_metadata.max_seqlen_q > 1: - # MTP verify: per-token layout with max_q_len=1. - # Persistent metadata is per-token (from _set_mla_persistent_worker_buffers_sparse_mtp). - paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q - paged_kv_indptr = attn_metadata.sparse_kv_indptr - paged_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens - # Gather physical page indices from kv_indices using topk positions. - # block_tables contains large-block IDs (block_ratio > 1) that - # need expansion; kv_indices already has per-token page indices. - paged_kv_indices = triton_gather_kv_indices_sparse( - paged_kv_indptr, - attn_metadata.token_to_seq_idxs, - self.topk_indices_buffer[:B], - attn_metadata.kv_indices, - attn_metadata.kv_indptr, - NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], - ) - max_q_len = 1 - else: - paged_kv_indptr = attn_metadata.sparse_kv_indptr - paged_kv_indices = triton_convert_req_index_to_global_index( - attn_metadata.cu_seqlens_q, - attn_metadata.kv_indptr, - paged_kv_indptr, - attn_metadata.kv_indices, - self.topk_indices_buffer[:B], - NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], - ) + if hasattr(attn_metadata, "triton_block_table"): + from aiter.ops.triton.attention.mla_decode import decode_attention_fwd - dp_size = get_dp_group().world_size - use_persistent_mode = not (dp_size > 1) + k_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + v_buffer = k_buffer[..., : self.kv_lora_rank] + page_size = k_buffer.shape[1] - # Sparse layers in MTP verify use separate persistent metadata - # (per-token, max_seqlen_qo=1) while dense layers use normal metadata - # (max_seqlen_qo=2). - is_sparse_mtp = ( - self.topk_indices_buffer is not None and attn_metadata.max_seqlen_q > 1 - ) + q_for_triton = ( + q.to(torch.bfloat16) + if q.dtype.is_floating_point and q.element_size() == 1 + else q + ) - if not use_persistent_mode: - work_meta_data = None - work_indptr = None - work_info_set = None - reduce_indptr = None - reduce_final_map = None - reduce_partial_map = None - elif is_sparse_mtp: - work_meta_data = attn_metadata.sparse_mtp_work_meta_data - work_indptr = attn_metadata.sparse_mtp_work_indptr - work_info_set = attn_metadata.sparse_mtp_work_info_set - reduce_indptr = attn_metadata.sparse_mtp_reduce_indptr - reduce_final_map = attn_metadata.sparse_mtp_reduce_final_map - reduce_partial_map = attn_metadata.sparse_mtp_reduce_partial_map + # Use pre-built dense block_table from prepare_decode() + decode_attention_fwd( + q_for_triton, + k_buffer, + v_buffer, + o, + attn_metadata.triton_lse, + attn_metadata.triton_block_table, + attn_metadata.context_lens, + attn_metadata.triton_attn_logits, + 4, # num_kv_splits + self.scale, + page_size, + k_scale=self._k_scale, + v_scale=self._k_scale, + ) else: - work_meta_data = attn_metadata.work_meta_data - work_indptr = attn_metadata.work_indptr - work_info_set = attn_metadata.work_info_set - reduce_indptr = attn_metadata.reduce_indptr - reduce_final_map = attn_metadata.reduce_final_map - reduce_partial_map = attn_metadata.reduce_partial_map - - mla_decode_fwd( - q, - kv_buffer.view(-1, 1, 1, q.shape[-1]), - o, - paged_cu_seqlens_q, - paged_kv_indptr, - paged_kv_indices, - paged_kv_last_page_lens, - max_q_len, - num_kv_splits=16, - sm_scale=self.scale, - work_meta_data=work_meta_data, - work_indptr=work_indptr, - work_info_set=work_info_set, - reduce_indptr=reduce_indptr, - reduce_final_map=reduce_final_map, - reduce_partial_map=reduce_partial_map, - q_scale=self._q_scale, - kv_scale=self._k_scale, - ) + kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2) + paged_cu_seqlens_q = attn_metadata.cu_seqlens_q + paged_kv_indptr = attn_metadata.kv_indptr + paged_kv_indices = attn_metadata.kv_indices + paged_kv_last_page_lens = attn_metadata.kv_last_page_lens + max_q_len = attn_metadata.max_seqlen_q + if self.topk_indices_buffer is not None: + if attn_metadata.max_seqlen_q > 1: + # MTP verify: per-token layout with max_q_len=1. + # Persistent metadata is per-token (from _set_mla_persistent_worker_buffers_sparse_mtp). + paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens + # Gather physical page indices from kv_indices using topk positions. + # block_tables contains large-block IDs (block_ratio > 1) that + # need expansion; kv_indices already has per-token page indices. + paged_kv_indices = triton_gather_kv_indices_sparse( + paged_kv_indptr, + attn_metadata.token_to_seq_idxs, + self.topk_indices_buffer[:B], + attn_metadata.kv_indices, + attn_metadata.kv_indptr, + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) + max_q_len = 1 + else: + paged_kv_indptr = attn_metadata.sparse_kv_indptr + paged_kv_indices = triton_convert_req_index_to_global_index( + attn_metadata.cu_seqlens_q, + attn_metadata.kv_indptr, + paged_kv_indptr, + attn_metadata.kv_indices, + self.topk_indices_buffer[:B], + NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1], + ) + + dp_size = get_dp_group().world_size + use_persistent_mode = not (dp_size > 1) + + # Sparse layers in MTP verify use separate persistent metadata + # (per-token, max_seqlen_qo=1) while dense layers use normal metadata + # (max_seqlen_qo=2). + is_sparse_mtp = ( + self.topk_indices_buffer is not None and attn_metadata.max_seqlen_q > 1 + ) + + if not use_persistent_mode: + work_meta_data = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + elif is_sparse_mtp: + work_meta_data = attn_metadata.sparse_mtp_work_meta_data + work_indptr = attn_metadata.sparse_mtp_work_indptr + work_info_set = attn_metadata.sparse_mtp_work_info_set + reduce_indptr = attn_metadata.sparse_mtp_reduce_indptr + reduce_final_map = attn_metadata.sparse_mtp_reduce_final_map + reduce_partial_map = attn_metadata.sparse_mtp_reduce_partial_map + else: + work_meta_data = attn_metadata.work_meta_data + work_indptr = attn_metadata.work_indptr + work_info_set = attn_metadata.work_info_set + reduce_indptr = attn_metadata.reduce_indptr + reduce_final_map = attn_metadata.reduce_final_map + reduce_partial_map = attn_metadata.reduce_partial_map + + mla_decode_fwd( + q, + kv_buffer.view(-1, 1, 1, q.shape[-1]), + o, + paged_cu_seqlens_q, + paged_kv_indptr, + paged_kv_indices, + paged_kv_last_page_lens, + max_q_len, + num_kv_splits=16, + sm_scale=self.scale, + work_meta_data=work_meta_data, + work_indptr=work_indptr, + work_info_set=work_info_set, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + q_scale=self._q_scale, + kv_scale=self._k_scale, + ) if self.head_repeat_factor > 1: o = o[:, :: self.head_repeat_factor, :].contiguous() diff --git a/atom/model_ops/attentions/triton_mha.py b/atom/model_ops/attentions/triton_mha.py new file mode 100644 index 0000000000..90b2babce5 --- /dev/null +++ b/atom/model_ops/attentions/triton_mha.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import logging +from typing import Type + +import torch + +import atom.model_ops as ops +from atom.config import KVCacheTensor +from atom.model_engine.scheduler import ScheduledBatch +from atom.model_ops.attention_mha import PagedAttentionImpl +from atom.model_ops.paged_attention import PagedAttention + +from .aiter_attention import AiterAttentionMetadataBuilder +from .backends import AttentionBackend + +logger = logging.getLogger("atom") + + +class TritonMHABackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "ROCM_TRITON_MHA" + + @staticmethod + def get_builder_cls() -> Type["TritonMHAMetadataBuilder"]: + return TritonMHAMetadataBuilder + + @staticmethod + def get_impl_cls(): + attn_cls = ops.Attention + if attn_cls == PagedAttention: + return PagedAttentionImpl + raise NotImplementedError( + f"TritonMHABackend does not support attention class {attn_cls!r}" + ) + + +class TritonMHAMetadataBuilder(AiterAttentionMetadataBuilder): + """MHA metadata builder that allocates KV cache in flash layout. + + Flash layout: K/V both [num_blocks, block_size, num_kv_heads, head_dim]. + Consumed directly by aiter triton `unified_attention` for prefill+decode. + """ + + def prepare_prefill(self, batch: ScheduledBatch): + attn_metadata, positions = super().prepare_prefill(batch) + + # When there are no cached tokens, the base builder leaves + # `block_tables=None` because AiterBackend's prefill consumes raw q/k/v + # via flash_attn_varlen_func. The unified_attention path used by + # TritonMHABackend instead requires a block_table even for pure prefill, + # so build a fake one here that treats raw K/V as a kv_cache with + # block_size=1: row i = [cu_seqlens_k[i], ..., cu_seqlens_k[i]+max-1]. + if attn_metadata.block_tables is None: + cu_k = attn_metadata.cu_seqlens_k + num_seqs = cu_k.shape[0] - 1 + offsets = cu_k[:num_seqs] + attn_metadata.block_tables = offsets.unsqueeze(1) + torch.arange( + attn_metadata.max_seqlen_k, dtype=torch.int32, device=cu_k.device + ) + + return attn_metadata, positions + + def build_kv_cache_tensor(self, layer_id: int, module): + if not ( + hasattr(module, "base_attention") + and hasattr(module, "use_mla") + and not module.use_mla + ): + return None + + runner = self.model_runner + config = runner.config + hf_config = config.hf_config + + if runner.is_mimo_v2(): + raise NotImplementedError( + "TritonMHABackend does not support MiMo-V2 (per-layer alloc path)" + ) + + impl = getattr(module, "impl", None) + if impl is not None and ( + getattr(impl, "rotary_emb", None) is not None + and getattr(impl, "q_norm", None) is not None + and getattr(impl, "k_norm", None) is not None + ): + raise NotImplementedError( + "TritonMHABackend is incompatible with the fused qk_norm+rope+shuffle " + "cache path; use AiterBackend for this model." + ) + + if runner.is_qwen_next(): + mtp_start = runner.mtp_start_layer_idx + if layer_id < mtp_start: + attn_idx = layer_id // runner.full_attention_interval + else: + attn_idx = runner.num_full_attn + (layer_id - mtp_start) + else: + attn_idx = layer_id + + k_cache = runner.kv_cache[0, attn_idx].view( + runner.num_physical_kvcache_blocks, + runner.physical_block_size, + runner.num_kv_heads, + hf_config.head_dim, + ) + v_cache = runner.kv_cache[1, attn_idx].view( + runner.num_physical_kvcache_blocks, + runner.physical_block_size, + runner.num_kv_heads, + hf_config.head_dim, + ) + if config.kv_cache_dtype == "fp8": + module.k_scale = runner.kv_scale[0, attn_idx] + module.v_scale = runner.kv_scale[1, attn_idx] + + module.max_model_len = config.max_model_len + module.k_cache = k_cache + module.v_cache = v_cache + if impl is not None: + impl.use_flash_layout = True + + return KVCacheTensor( + layer_num=layer_id, + k_cache=k_cache, + v_cache=v_cache, + k_scale=module.k_scale, + v_scale=module.v_scale, + ) diff --git a/atom/model_ops/attentions/triton_mla.py b/atom/model_ops/attentions/triton_mla.py new file mode 100644 index 0000000000..3a52dfd4f0 --- /dev/null +++ b/atom/model_ops/attentions/triton_mla.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import logging +from typing import Type + +import torch +from aiter.ops.triton.attention.mla_decode import csr_to_dense_block_table +from atom.model_engine.scheduler import ScheduledBatch +from atom.model_ops.attention_mla import MLAAttention +from atom.utils.forward_context import AttentionMetaData + +from .aiter_mla import AiterMLAMetadataBuilder +from .backends import AttentionBackend + +logger = logging.getLogger("atom") + + +class TritonMLABackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "ROCM_TRITON_MLA" + + @staticmethod + def get_builder_cls() -> Type["TritonMLAMetadataBuilder"]: + return TritonMLAMetadataBuilder + + @staticmethod + def get_impl_cls() -> Type["MLAAttention"]: + return MLAAttention + + +class TritonMLAMetadataBuilder(AiterMLAMetadataBuilder): + + def __init__(self, model_runner): + super().__init__(model_runner) + + hf = model_runner.config.hf_config + kv_lora_rank = hf.kv_lora_rank + num_kv_splits = 4 + triton_mla_buffers = { + "triton_block_table": torch.zeros( + self.max_bs, + self.max_num_blocks_per_seq, + dtype=torch.int32, + device=self.device, + ), + "triton_attn_logits": torch.empty( + self.max_bs, + self.padded_num_attention_heads, + num_kv_splits, + kv_lora_rank + 1, + dtype=torch.float32, + device=self.device, + ), + "triton_lse": torch.empty( + self.max_bs, + self.padded_num_attention_heads, + dtype=torch.float32, + device=self.device, + ), + } + self.model_runner.forward_vars.update(triton_mla_buffers) + + def set_mla_persistent_worker_buffers( + self, bs, max_q_len, only_update=False, num_reject_tokens=None + ): + # Triton MLA does not use aiter persistent worker buffers + return {} + + def prepare_decode(self, batch: ScheduledBatch, bs: int): + attn_metadata, positions = super().prepare_decode(batch, bs) + + scheduled_bs = batch.total_seqs_num_decode + max_seqlen_k = attn_metadata.max_seqlen_k + var = self.model_runner.forward_vars + + triton_bt = var["triton_block_table"][:scheduled_bs, :max_seqlen_k] + triton_bt.zero_() + csr_to_dense_block_table( + attn_metadata.kv_indices, + attn_metadata.kv_indptr, + triton_bt, + max_seqlen_k, + scheduled_bs, + ) + attn_metadata.triton_block_table = triton_bt + attn_metadata.triton_attn_logits = var["triton_attn_logits"][:scheduled_bs] + attn_metadata.triton_lse = var["triton_lse"][:scheduled_bs] + + return attn_metadata, positions + + def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: + attn_metadata, context = super().build_for_cudagraph_capture(bs) + + var = self.model_runner.forward_vars + attn_metadata.triton_block_table = var["triton_block_table"][:bs] + attn_metadata.triton_attn_logits = var["triton_attn_logits"][:bs] + attn_metadata.triton_lse = var["triton_lse"][:bs] + + return attn_metadata, context diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 851cb3dce5..18c513e809 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -24,6 +24,7 @@ from typing import Any import logging from math import prod +from aiter import ActivationType from aiter.jit.utils.chip_info import get_gfx from atom.model_ops.utils import has_triton_kernels @@ -32,9 +33,14 @@ if has_triton_kernels(): try: - from triton_kernels.matmul_ogs import matmul_ogs + import triton_kernels.swiglu + from triton_kernels.matmul_ogs import ( + FnSpecs, + FusedActivation, + PrecisionConfig, + matmul_ogs, + ) from triton_kernels.routing import routing - from triton_kernels.matmul_ogs import PrecisionConfig except (AttributeError, ImportError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " @@ -283,28 +289,45 @@ def triton_kernel_fused_experts( ) with _amd_smem_safe_tile(): - matmul_ogs( - hidden_states, - w1, - w1_bias, - routing_data, - gather_indx=gather_indx, - precision_config=w13_precision_config, - gammas=gammas if apply_router_weight_on_input else None, - y=raw_intermediate, - ) - - # Standard SiLU/SwiGLU activation: silu(gate) * up - # With optional swiglu_limit clamping (V4: limit=10.0) - raw_2d = raw_intermediate.view(M * topk, N) - gate = raw_2d[:, :half_N] - up = raw_2d[:, half_N:] - if swiglu_limit > 0: - gate = gate.clamp(max=swiglu_limit) - up = up.clamp(-swiglu_limit, swiglu_limit) - intermediate_cache[0] = torch.nn.functional.silu(gate) * up + if activation == ActivationType.Swiglu: + # SwiGLU (GPT OSS): fused activation with interleaved [gate, up] layout + act = FusedActivation( + FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")), + (swiglu_alpha, swiglu_limit), + 2, + ) + matmul_ogs( + hidden_states, + w1, + w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w13_precision_config, + gammas=gammas if apply_router_weight_on_input else None, + fused_activation=act, + y=intermediate_cache, + ) + else: + # SiLU (DeepSeek): concatenated [gate | up] layout, manual activation + raw_intermediate = matmul_ogs( + hidden_states, + w1, + w1_bias, + routing_data, + gather_indx=gather_indx, + precision_config=w13_precision_config, + gammas=gammas if apply_router_weight_on_input else None, + ) + raw_2d = raw_intermediate.view(M * topk, N) + gate = raw_2d[:, :half_N] + up = raw_2d[:, half_N:] + if swiglu_limit > 0: + gate = gate.clamp(max=swiglu_limit) + up = up.clamp(-swiglu_limit, swiglu_limit) + intermediate_cache = intermediate_cache.view(M * topk, half_N) + intermediate_cache.copy_(torch.nn.functional.silu(gate) * up) + intermediate_cache = intermediate_cache.view(batch_dim, M * topk, half_N) - with _amd_smem_safe_tile(): matmul_ogs( intermediate_cache.view(M * topk, half_N), w2, diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 5a42b12b82..dd100bd67b 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -686,11 +686,14 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): or self.quant_type == QuantType.per_1x32 ) gfx = get_gfx() - self.use_triton = ( - gfx.startswith("gfx94") - or (gfx.startswith("gfx95") and envs.ATOM_USE_TRITON_GEMM) - or os.environ.get("ATOM_USE_TRITON_MOE") == "1" - ) + if envs.is_set("ATOM_USE_TRITON_MOE"): + self.use_triton = envs.ATOM_USE_TRITON_MOE + else: + self.use_triton = ( + gfx.startswith("gfx94") + or gfx.startswith("gfx12") + or (gfx.startswith("gfx95") and envs.ATOM_USE_TRITON_GEMM) + ) if self.use_triton: from atom.model_ops.utils import has_triton_kernels @@ -975,11 +978,13 @@ def apply( num_fused_shared_experts=layer.num_fused_shared_experts, routed_scaling_factor=layer.routed_scaling_factor, ) + n_expts_act = topk_weights.shape[1] # Convert to triton routing data structures n_expts_tot = router_logits.shape[-1] if global_num_experts > 0: n_expts_tot = global_num_experts + n_expts_tot = n_expts_tot + layer.num_fused_shared_experts routing_data, gather_idx, scatter_idx = routing_from_topk( topk_weights, topk_ids, n_expts_tot @@ -994,7 +999,7 @@ def apply( routing_data, gather_idx, scatter_idx, - topk=top_k, + topk=n_expts_act, activation=activation, w13_precision_config=self.w13_precision_config, w2_precision_config=self.w2_precision_config, @@ -1002,7 +1007,7 @@ def apply( w2_bias=layer.w2_bias, swiglu_limit=getattr(layer, "swiglu_limit", 0.0), apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=global_num_experts, + global_num_experts=n_expts_tot, expert_map=expert_map, ) return _moe_result diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 51fee833a2..7282248b80 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -33,6 +33,8 @@ "ATOM_USE_TRITON_MXFP4_BMM": lambda: ( os.getenv("ATOM_USE_TRITON_MXFP4_BMM", "0") == "1" ), + "ATOM_USE_TRITON_MLA": lambda: os.getenv("ATOM_USE_TRITON_MLA", "0") == "1", + "ATOM_USE_TRITON_MOE": lambda: os.getenv("ATOM_USE_TRITON_MOE", "0") == "1", # --- Kernel Fusion Toggles --- # QK-norm-rope-cache-quant fusion for Qwen3-MoE; disabled by default. # Enable for Qwen3-MoE to get better performance. @@ -69,6 +71,14 @@ "ATOM_DISABLE_MMAP": lambda: ( os.getenv("ATOM_DISABLE_MMAP", "false").lower() == "true" ), + # Use a thread pool for weight loading instead of main-process sequential I/O. + # Set to 0 to disable if the thread pool causes hangs (e.g. on gfx1250). + "ATOM_LOADER_USE_THREADPOOL": lambda: os.getenv("ATOM_LOADER_USE_THREADPOOL", "1") + == "1", + # --- Attention Backend --- + # Use unified_attention (flash-style) for MHA paged/prefill attention instead + # of pa_decode_gluon. Set to 1 to enable the unified_attention path. + "ATOM_USE_UNIFIED_ATTN": lambda: os.getenv("ATOM_USE_UNIFIED_ATTN", "0") == "1", # --- Plugin Mode --- "ATOM_DISABLE_VLLM_PLUGIN": lambda: ( os.getenv("ATOM_DISABLE_VLLM_PLUGIN", "0").lower() == "1" diff --git a/atom/utils/selector.py b/atom/utils/selector.py index e87b1f8194..bb9b1c7d65 100644 --- a/atom/utils/selector.py +++ b/atom/utils/selector.py @@ -7,6 +7,7 @@ from atom.model_ops.attentions.backends import AttentionBackend from atom.utils import resolve_obj_by_qualname from atom.plugin.prepare import is_sglang, is_vllm +from atom.utils import envs def get_attn_backend( @@ -51,13 +52,9 @@ def get_attn_backend_cls( if use_v4: return "atom.model_ops.attentions.deepseek_v4_attn.DeepseekV4Backend" if use_mla: - # if block_size == 1: - return "atom.model_ops.attentions.aiter_mla.AiterMLABackend" # noqa: E501 - # else: - # raise ValueError( - # f" The selected backend" - # f"does not support block size {block_size}." - # "(currently only supports block size 1)") + if envs.ATOM_USE_TRITON_MLA: + return "atom.model_ops.attentions.triton_mla.TritonMLABackend" + return "atom.model_ops.attentions.aiter_mla.AiterMLABackend" if use_gdn: if use_vllm: return "atom.plugin.vllm.attention_backend.gdn_attn.GDNAttentionBackend" @@ -66,4 +63,6 @@ def get_attn_backend_cls( "atom.plugin.sglang.attention_backend.attention_gdn.GDNAttentionBackend" ) return "atom.model_ops.attentions.gdn_attn.GDNAttentionBackend" + if envs.ATOM_USE_UNIFIED_ATTN: + return "atom.model_ops.attentions.triton_mha.TritonMHABackend" return "atom.model_ops.attentions.aiter_attention.AiterBackend" # noqa: E501 From e74abf5fe235455f88b4df4cd261a9eff73bc655 Mon Sep 17 00:00:00 2001 From: gbyu-amd Date: Sat, 9 May 2026 21:32:30 +0800 Subject: [PATCH 09/76] [fix][acc] fix accuracy of fp8 attn weights model using ptpc quant recipe (#670) * correct the quant type based on recipe for fuse_qknorm_quant * remove debug print * use quant type Signed-off-by: zejunchen-zejun * use unified aiter fused_qk_norm interface Signed-off-by: zejunchen-zejun * change the import path for fused_qknorm Signed-off-by: zejunchen-zejun * fix qknorm quant compile path * document qknorm quant compile guard * simplify qknorm quant dispatch * use single rmsnorm quant entrypoint * update api * clean the qk rmsnorm code * fix corner case --------- Signed-off-by: zejunchen-zejun Co-authored-by: Guanbao Yu Co-authored-by: zejunchen-zejun Co-authored-by: wuhuikx Co-authored-by: XiaobingSuper --- atom/models/deepseek_v2.py | 110 ++++++++++++++++++------------------- 1 file changed, 52 insertions(+), 58 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index c0daa899a8..3a92fe7410 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -160,7 +160,7 @@ def _fuse_rmsnorm_fp4_quant_fake( return out1_quantized, out1_bs, out1_unquantized, out2, out_res1 -def _fused_rms_fp8_group_quant_fake( +def _fused_rms_fp8_quant_fake( x1: torch.Tensor, x1_weight: torch.Tensor, x1_epsilon: float, @@ -170,6 +170,7 @@ def _fused_rms_fp8_group_quant_fake( res1: Optional[torch.Tensor] = None, dtype_quant: torch.dtype = dtypes.fp8, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ) -> Tuple[ @@ -180,9 +181,18 @@ def _fused_rms_fp8_group_quant_fake( torch.Tensor, ]: m, n1 = x1.shape - out1_quantized = torch.empty((m, n1), dtype=dtype_quant, device=x1.device) - num_bs_cols = (n1 + group_size - 1) // group_size - out1_bs = torch.empty((m, num_bs_cols), dtype=torch.float32, device=x1.device) + no_quant = quant_type is None or quant_type == QuantType.No.value + if not no_quant: + out1_quantized = torch.empty((m, n1), dtype=dtype_quant, device=x1.device) + else: + out1_quantized = torch.empty_like(x1) + if no_quant: + out1_bs = None + elif quant_type == QuantType.per_Token.value: + out1_bs = torch.empty((m, 1), dtype=torch.float32, device=x1.device) + else: + num_bs_cols = (n1 + group_size - 1) // group_size + out1_bs = torch.empty((m, num_bs_cols), dtype=torch.float32, device=x1.device) out1_unquantized = torch.empty_like(x1) if output_unquantized_inp1 else None out2 = None if x2 is not None: @@ -236,8 +246,8 @@ def _fuse_rmsnorm_fp4_quant( return out1_quantized, out1_bs, out1_unquantized, out2, out_res1 -@torch_compile_guard(gen_fake=_fused_rms_fp8_group_quant_fake) -def _fused_rms_fp8_group_quant( +@torch_compile_guard(gen_fake=_fused_rms_fp8_quant_fake) +def _fused_rms_fp8_quant( x1: torch.Tensor, x1_weight: torch.Tensor, x1_epsilon: float, @@ -247,6 +257,7 @@ def _fused_rms_fp8_group_quant( res1: Optional[torch.Tensor] = None, dtype_quant: torch.dtype = dtypes.fp8, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ) -> Tuple[ @@ -257,7 +268,7 @@ def _fused_rms_fp8_group_quant( torch.Tensor, ]: out1_quantized, out1_bs, out1_unquantized, out2, out_res1 = ( - _fused_rms_fp8_group_quant_fake( + _fused_rms_fp8_quant_fake( x1, x1_weight, x1_epsilon, @@ -267,14 +278,18 @@ def _fused_rms_fp8_group_quant( res1, dtype_quant, group_size, + quant_type, output_unquantized_inp1, transpose_scale, ) ) - from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant + if quant_type is None: + quant_type = QuantType.No + else: + quant_type = QuantType(quant_type) - fused_qk_rmsnorm_group_quant( + fused_qk_rmsnorm( q_out_quantized=out1_quantized, q_out_scale=out1_bs, q=x1, @@ -287,6 +302,7 @@ def _fused_rms_fp8_group_quant( k_weight=x2_weight, k_epsilon=x2_epsilon, q_residual=res1, + quant_type=quant_type, group_size=group_size, transpose_scale=transpose_scale, ) @@ -306,6 +322,7 @@ def _fuse_rmsnorm_quant( shuffle: bool = True, scale_shuffle_padding: bool = False, group_size: int = 128, + quant_type: Optional[int] = None, output_unquantized_inp1: bool = False, transpose_scale: bool = False, ): @@ -324,9 +341,9 @@ def _fuse_rmsnorm_quant( output_unquantized_inp1, ) ) - elif dtype_quant == dtypes.fp8: + elif dtype_quant == dtypes.fp8 or dtype_quant == torch.bfloat16: out1_quantized, out1_bs, out1_unquantized, out2, out_res1 = ( - _fused_rms_fp8_group_quant( + _fused_rms_fp8_quant( x1, x1_weight, x1_epsilon, @@ -334,10 +351,11 @@ def _fuse_rmsnorm_quant( x2_weight, x2_epsilon, res1, - dtype_quant, - group_size, - output_unquantized_inp1, - transpose_scale, + dtype_quant=dtype_quant, + group_size=group_size, + quant_type=quant_type, + output_unquantized_inp1=output_unquantized_inp1, + transpose_scale=transpose_scale, ) ) else: @@ -703,36 +721,6 @@ def _fuse_qkv_a_proj_reduce_rmsnorm_quant( return q_c, q_c_scale, kv_c_normed, k_pe -def _fused_qk_rmsnorm_fake( - q_c: torch.Tensor, - q_a_layernorm_weight: torch.Tensor, - q_a_layernorm_variance_epsilon: float, - kv_c: torch.Tensor, - kv_a_layernorm_weight: torch.Tensor, - kv_a_layernorm_variance_epsilon: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - return torch.empty_like(q_c), torch.empty_like(kv_c) - - -@torch_compile_guard(gen_fake=_fused_qk_rmsnorm_fake) -def _fused_qk_rmsnorm( - q_c: torch.Tensor, - q_a_layernorm_weight: torch.Tensor, - q_a_layernorm_variance_epsilon: float, - kv_c: torch.Tensor, - kv_a_layernorm_weight: torch.Tensor, - kv_a_layernorm_variance_epsilon: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - return fused_qk_rmsnorm( - q_c, - q_a_layernorm_weight, - q_a_layernorm_variance_epsilon, - kv_c, - kv_a_layernorm_weight, - kv_a_layernorm_variance_epsilon, - ) - - class DeepseekV2MLP(nn.Module): def __init__( @@ -1298,6 +1286,14 @@ def __init__( layer_quant_dtype = quant_config.get_layer_quant_config( f"{prefix}.{q_a_proj_name}" ).quant_dtype + layer_quant_type = quant_config.get_layer_quant_config( + f"{prefix}.{q_a_proj_name}" + ).quant_type + # Keep a plain int on the module so Dynamo does not guard on a + # QuantType enum attribute. + layer_quant_type_value = ( + None if layer_quant_type is None else layer_quant_type.value + ) if layer_quant_dtype == dtypes.fp4x2: if not use_triton_gemm(): source_quant_dtype = None @@ -1473,7 +1469,8 @@ def __init__( # When ATOM_ENABLE_DS_QKNORM_QUANT_FUSION is turned on, self.fuse_qknorm_quant is turned on only if FP8 or (use_triton_gemm() and FP4), self.prefix = prefix - self.quant_dtype = None + self.quant_dtype = layer_quant_dtype + self.qknorm_quant_type = layer_quant_type_value self.fuse_qknorm_quant = False # always fuse qknorm self.fuse_qknorm = ENABLE_DS_QKNORM_FUSION @@ -1481,7 +1478,6 @@ def __init__( if layer_quant_dtype == dtypes.fp8 or ( layer_quant_dtype == dtypes.fp4x2 and use_triton_gemm() ): - self.quant_dtype = layer_quant_dtype self.fuse_qknorm_quant = True def forward( @@ -1527,7 +1523,7 @@ def forward( dim=-1, ) # fuse q_c norm + kv_c norm + quant of hidden_states_or_q_c - if self.fuse_qknorm_quant: + if self.fuse_qknorm_quant or self.fuse_qknorm: ( (hidden_states_or_q_c, hidden_states_or_q_c_scale), _, @@ -1545,19 +1541,10 @@ def forward( shuffle=False, scale_shuffle_padding=False, group_size=128, + quant_type=self.qknorm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) - elif self.fuse_qknorm: - hidden_states_or_q_c, kv_c_normed = _fused_qk_rmsnorm( - q_c, - self.q_a_layernorm.weight, - self.q_a_layernorm.eps, - kv_c, - self.kv_a_layernorm.weight, - self.kv_a_layernorm.eps, - ) - hidden_states_or_q_c_scale = None else: hidden_states_or_q_c = self.q_a_layernorm(q_c) else: @@ -1636,6 +1623,11 @@ def __init__( if quant_config is None else quant_config.get_layer_quant_config(prefix).quant_dtype ) + self.input_norm_quant_type = ( + None + if quant_config is None + else quant_config.get_layer_quant_config(prefix).quant_type.value + ) self.fuse_input_norm_quant = False self.fuse_ar_input_norm = ENABLE_ALLREDUCE_RMSNORM_FUSION if quant_config is not None and ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION: @@ -1719,6 +1711,7 @@ def forward( shuffle=True, scale_shuffle_padding=True, group_size=128, + quant_type=self.input_norm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) @@ -1737,6 +1730,7 @@ def forward( shuffle=True, scale_shuffle_padding=True, group_size=128, + quant_type=self.input_norm_quant_type, output_unquantized_inp1=False, transpose_scale=True, ) From 3f4c42545db1feab82d9ddb0af3fd24449a587ce Mon Sep 17 00:00:00 2001 From: amd-ruitang3 <145657428+amd-ruitang3@users.noreply.github.com> Date: Sat, 9 May 2026 23:06:33 +0800 Subject: [PATCH 10/76] Dsv4 moe flydsl (#718) * [DSV4]moe flydsl * update * update * refactor moe shuffle * update * update * modify atom env * update env --------- Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> --- .github/benchmark/models.json | 2 +- .github/benchmark/models_accuracy.json | 2 +- atom/model_ops/moe.py | 100 ++++++++++++------------- atom/models/deepseek_v4.py | 45 ++++++----- atom/models/gpt_oss.py | 34 +++++++++ atom/utils/envs.py | 3 + 6 files changed, 114 insertions(+), 72 deletions(-) diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 001c430900..39aed57bd1 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -17,7 +17,7 @@ "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", - "env_vars": "ATOM_USE_TRITON_MOE=1" + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1" }, { "display": "DeepSeek-R1-0528 MTP3", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 29b8e6bcd3..94d1dcce48 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -27,7 +27,7 @@ "model_name": "DeepSeek-V4-Pro", "model_path": "deepseek-ai/DeepSeek-V4-Pro", "extraArgs": "--kv_cache_dtype fp8 -tp 8", - "env_vars": "ATOM_USE_TRITON_MOE=1", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.92, diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index dd100bd67b..4244e0a3cf 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -13,13 +13,14 @@ from aiter.fused_moe import fused_moe from aiter.jit.utils.chip_info import get_gfx from aiter.jit.utils.torch_guard import torch_compile_guard -from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4 +from aiter.ops.shuffle import shuffle_weight, shuffle_scale from aiter.utility import fp4_utils from atom.config import ( Config, QuantizationConfig, get_current_atom_config, ) +from aiter.ops.flydsl.moe_common import GateMode from atom.quant_spec import LayerQuantConfig from atom.model_loader.weight_utils import set_weight_attrs from atom.model_ops.base_config import QuantizeMethodBase @@ -681,6 +682,7 @@ def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): self.quant_dtype = quant_config.quant_dtype self.quant_method = quant_config.quant_method or "" self.static_input_scales = not quant_config.is_dynamic + self.is_guinterleave = envs.ATOM_MOE_GU_ITLV self.block_quant = ( self.quant_type == QuantType.per_1x128 or self.quant_type == QuantType.per_1x32 @@ -856,40 +858,7 @@ def process_weights_after_loading(self, layer): layer.w13_weight_scale = None layer.w2_weight_scale = None return - elif layer.activation == ActivationType.Swiglu: - e, n, k = layer.w13_weight.shape - layer.w13_weight.view(torch.uint8).copy_( - layer.w13_weight.data.view(torch.uint8) - .view(e, n // 2, 2, k) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, k) - ) - layer.w13_weight_scale.data = ( - layer.w13_weight_scale.data.view(e, n // 2, 2, -1) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, -1) - ) - layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True) - shuffled_w13_scale = shuffle_scale_a16w4( - layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]), - self.num_experts, - True, - ) - layer.w2_weight.data = shuffle_weight_a16w4(layer.w2_weight, 16, False) - shuffled_w2_scale = shuffle_scale_a16w4( - layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]), - self.num_experts, - False, - ) - if layer.w13_bias is not None: - layer.w13_bias.data = ( - layer.w13_bias.data.view(-1, n // 2, 2) - .permute(0, 2, 1) - .contiguous() - .view(-1, n) - ) + # quark method for moe, split it out? elif self.quant_method == "quark": shuffle_weights(layer.w13_weight, layer.w2_weight) @@ -904,12 +873,35 @@ def process_weights_after_loading(self, layer): layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) return else: - shuffle_weights(layer.w13_weight, layer.w2_weight) - shuffled_w13_scale = fp4_utils.e8m0_shuffle( - layer.w13_weight_scale.view(self.num_experts, -1) + # shuffle weight + layer.w13_weight.data = shuffle_weight( + layer.w13_weight, + is_guinterleave=self.is_guinterleave, + gate_up=True, ) - shuffled_w2_scale = fp4_utils.e8m0_shuffle( - layer.w2_weight_scale.view(self.num_experts, -1) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight, + is_guinterleave=self.is_guinterleave, + gate_up=False, + ) + + # shuffle scale + if self.is_guinterleave: + w13_scale_2d = layer.w13_weight_scale.view( + -1, layer.w13_weight_scale.shape[-1] + ) + w2_scale_2d = layer.w2_weight_scale.view( + -1, layer.w2_weight_scale.shape[-1] + ) + else: + w13_scale_2d = layer.w13_weight_scale.view(self.num_experts, -1) + w2_scale_2d = layer.w2_weight_scale.view(self.num_experts, -1) + + shuffled_w13_scale = shuffle_scale( + w13_scale_2d, self.num_experts, self.is_guinterleave, True + ) + shuffled_w2_scale = shuffle_scale( + w2_scale_2d, self.num_experts, self.is_guinterleave, False ) layer.w13_weight_scale = atom_parameter(shuffled_w13_scale) @@ -1070,6 +1062,12 @@ def apply( intermediate_pad=self.intermediate_pad, bias1=layer.w13_bias, bias2=layer.w2_bias, + swiglu_limit=getattr(layer, "swiglu_limit", 0.0), + gate_mode=( + GateMode.INTERLEAVE.value + if self.is_guinterleave + else GateMode.SEPARATED.value + ), ) return self.fused_experts( hidden_states=x, @@ -1096,7 +1094,6 @@ def apply( # Refer to CompressedTensorsW8A8Fp8MoEMethod in vllm class CompressedTensorsFp8MoEMethod(FusedMoEMethodBase): - def __init__(self, quant_config: LayerQuantConfig, moe: FusedMoEConfig): super().__init__(moe) self.quant_config = quant_config @@ -1294,12 +1291,14 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2_input_scale = getattr(layer, "w2_input_scale", None) if self.need_normalize_e4m3fn_to_e4m3fnuz: - w13.data, w13_scale.data, w13_input_scale_data = ( - normalize_e4m3fn_to_e4m3fnuz( - w13.data, - w13_scale.data, - w13_input_scale.data if w13_input_scale is not None else None, - ) + ( + w13.data, + w13_scale.data, + w13_input_scale_data, + ) = normalize_e4m3fn_to_e4m3fnuz( + w13.data, + w13_scale.data, + w13_input_scale.data if w13_input_scale is not None else None, ) if w13_input_scale is not None and w13_input_scale_data is not None: w13_input_scale.data = w13_input_scale_data @@ -1731,9 +1730,10 @@ def _process_tensor_quant(self, layer: nn.Module) -> None: layer.w13_weight_scale[expert_id][shard_id], ) quant_func = get_hip_quant(self.quant_type) - layer.w13_weight[expert_id][start : start + shard_size, :], _ = ( - quant_func(dq_weight, max_w13_scales[expert_id]) - ) + ( + layer.w13_weight[expert_id][start : start + shard_size, :], + _, + ) = quant_func(dq_weight, max_w13_scales[expert_id]) start += shard_size shuffle_weights(layer.w13_weight, layer.w2_weight) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 58ab089183..d96b0698a8 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1798,10 +1798,25 @@ def __init__( # (set by ModelRunner). torch.compile silently drops NNModule # attribute mutation across the compile boundary, so stashing on # `self.foo` from inside forward is a no-op at runtime. - + assert args.n_shared_experts == 1 + # aiter can fuse shared expert into the routed FusedMoE kernel ONLY + # when shared and routed have the same quant dtype. + # `is_rocm_aiter_fusion_shared_expert_enabled` compares shared vs + # GLOBAL (=base dtype), but V4-Pro has routed=FP4 (per-layer override + # for `.ffn.experts`) and shared=FP8 (=global). The function returns + # True for V4 because shared matches global, missing the routed-shared + # mismatch. Direct comparison: get routed and shared specs and compare. + routed_dtype = qc.get_layer_quant_config(f"{prefix}.experts").quant_dtype + shared_dtype = qc.get_layer_quant_config(f"{prefix}.shared_experts").quant_dtype + self._fuse_shared_into_routed = ( + routed_dtype == shared_dtype + and is_rocm_aiter_fusion_shared_expert_enabled() + ) moe_cfg = SimpleNamespace( routed_scaling_factor=self.routed_scaling_factor, - n_shared_experts=args.n_shared_experts, + n_shared_experts=( + args.n_shared_experts if self._fuse_shared_into_routed else 0 + ), ) self.experts = FusedMoE( num_experts=self.n_routed_experts, @@ -1818,21 +1833,9 @@ def __init__( config=moe_cfg, ) self.experts.swiglu_limit = args.swiglu_limit - assert args.n_shared_experts == 1 - # aiter can fuse shared expert into the routed FusedMoE kernel ONLY - # when shared and routed have the same quant dtype. - # `is_rocm_aiter_fusion_shared_expert_enabled` compares shared vs - # GLOBAL (=base dtype), but V4-Pro has routed=FP4 (per-layer override - # for `.ffn.experts`) and shared=FP8 (=global). The function returns - # True for V4 because shared matches global, missing the routed-shared - # mismatch. Direct comparison: get routed and shared specs and compare. - routed_dtype = qc.get_layer_quant_config(f"{prefix}.experts").quant_dtype - shared_dtype = qc.get_layer_quant_config(f"{prefix}.shared_experts").quant_dtype - self._fuse_shared_into_routed = ( - routed_dtype == shared_dtype - and is_rocm_aiter_fusion_shared_expert_enabled() - ) + if not self._fuse_shared_into_routed: + # self.experts.num_fused_shared_experts = 0 self.shared_experts = Expert( args.dim, args.moe_inter_dim, @@ -2119,10 +2122,12 @@ def forward( ) -> torch.Tensor: # [num_tokens, hc, dim] updated residual stream # ----- Attention sub-layer with mHC mixing ----- residual = x # [num_tokens, hc, dim] - x, post, comb = ( - self.hc_pre( # [num_tokens, dim], [num_tokens, hc], [num_tokens, hc, hc] - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base - ) + ( + x, + post, + comb, + ) = self.hc_pre( # [num_tokens, dim], [num_tokens, hc], [num_tokens, hc, hc] + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base ) x = self.attn_norm(x) # [num_tokens, dim] x = self.attn(x, positions) # [num_tokens, dim] diff --git a/atom/models/gpt_oss.py b/atom/models/gpt_oss.py index eca04cb62f..66b2b22b0f 100644 --- a/atom/models/gpt_oss.py +++ b/atom/models/gpt_oss.py @@ -153,6 +153,37 @@ def forward( return output +def _interleave_swiglu_weights(experts: FusedMoE): + """Interleave gate/up weights, scales, and biases for Swiglu activation. + + Must run before Mxfp4MoEMethod.process_weights_after_loading (shuffle). + The loader calls module.process_weights_after_loading() before + quant_method.process_weights_after_loading(module), so this ordering + is guaranteed. + """ + e, n, k = experts.w13_weight.shape + experts.w13_weight.view(torch.uint8).copy_( + experts.w13_weight.data.view(torch.uint8) + .view(e, n // 2, 2, k) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, k) + ) + experts.w13_weight_scale.data = ( + experts.w13_weight_scale.data.view(e, n // 2, 2, -1) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, -1) + ) + if experts.w13_bias is not None: + experts.w13_bias.data = ( + experts.w13_bias.data.view(-1, n // 2, 2) + .permute(0, 2, 1) + .contiguous() + .view(-1, n) + ) + + class MLPBlock(torch.nn.Module): def __init__( self, @@ -202,6 +233,9 @@ def __init__( else: self.moe_hidden_pad = 0 + def process_weights_after_loading(self): + _interleave_swiglu_weights(self.experts) + def forward(self, x: torch.Tensor) -> torch.Tensor: num_tokens = x.shape[0] diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 7282248b80..41294f80bd 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -95,6 +95,9 @@ "ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD": lambda: int( os.getenv("ATOM_DUAL_STREAM_MOE_TOKEN_THRESHOLD", "1024") ), + # Gate/Up interleave mode for MoE weight preshuffle and kernel gate_mode. + # "0" (default) = SEPARATED layout; "1" = INTERLEAVE layout. + "ATOM_MOE_GU_ITLV": lambda: os.getenv("ATOM_MOE_GU_ITLV", "0") == "1", # --- MTP (relaxed mtp for quantized mtp) --- "ATOM_ENABLE_RELAXED_MTP": lambda: ( os.getenv("ATOM_ENABLE_RELAXED_MTP", "0").lower() == "1" From fc64fb65a68d1e9f34d152f51ed98079a8ec192a Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Sat, 9 May 2026 23:13:39 +0800 Subject: [PATCH 11/76] Optimize Deepseek V4 prepare decode (#728) --- atom/model_ops/attentions/deepseek_v4_attn.py | 279 +++++++++--------- 1 file changed, 134 insertions(+), 145 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 6b4f6bfc08..395567671c 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -290,6 +290,39 @@ def _build_window_topk_batched( return out[:total] +def _build_window_topk_np( + positions: np.ndarray, + start_pos_per_token: np.ndarray, + window_size: int, +) -> np.ndarray: + """CPU numpy equivalent of ``_build_window_topk_batched``. + + For small decode batches the ~25 GPU kernel launches dominate; this + numpy version runs in microseconds on CPU and feeds a single H2D copy. + """ + total = positions.shape[0] + arange_w = np.arange(window_size, dtype=np.int64).reshape(1, window_size) + pos_col = positions.reshape(total, 1) + sp_col = start_pos_per_token.reshape(total, 1) + + case_a = np.maximum(pos_col - window_size + 1, 0) + arange_w + case_a = np.where(case_a > pos_col, -1, case_a) + + case_b = np.broadcast_to(arange_w, (total, window_size)).copy() + case_b = np.where(arange_w > sp_col, -1, case_b) + + sp_mod = sp_col % window_size + case_c = (sp_mod + 1 + arange_w) % window_size + + sp_eq_0 = sp_col == 0 + sp_in_prefix = (sp_col > 0) & (sp_col < window_size - 1) + + res = case_c + res = np.where(sp_in_prefix, case_b, res) + res = np.where(sp_eq_0, case_a, res) + return res.astype(np.int32) + + class DeepseekV4Backend(AttentionBackend): """Backend selector entry for V4 hybrid attention. @@ -394,6 +427,10 @@ def __init__(self, model_runner): # `torch.as_tensor(arr)` allocations. self._alloc_v4_metadata_buffers() + @property + def prep_stream(self): + return self.model_runner.async_execute_stream + # ------------------------------------------------------------------ # # AttentionMetadataBuilder hooks (per-request cache abstraction). # # ------------------------------------------------------------------ # @@ -832,51 +869,69 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): """V4-style decode prep: populates positions, cu_seqlens_q, block_tables, and state_slot_mapping. - For PR3-main multi-seq: - - cu_seqlens_q gives V4 forward per-seq token-range slicing. - - block_tables[i] is the per-seq block list for compressor scatter - and sparse_attn gather. - - state_slot_mapping[i] is the per-seq state-cache slot (swa_kv + - Compressor.kv_state row index). Distinct from `slot_mapping` - which is per-token paged-KV-pool index. - - Also publishes CPU mirrors (`v4_*_cpu` attrs) so the V4 forward path - can read per-seq metadata without GPU→CPU `.tolist()`/.item() syncs - (PR-A Phase 2: required to unlock CUDAGraph). + Uses stream overlap (like AiterMLAMetadataBuilder) to hide H2D + latency behind CPU numpy work: basic H2D copies fire on + ``prep_stream`` while ``_build_compress_plans`` runs on the CPU. """ import numpy as np var = self.model_runner.forward_vars scheduled_bs = batch.total_seqs_num_decode - # Naming: `_np` suffix marks numpy arrays, `_gpu` marks GPU tensors. - # Both representations of context_lens coexist (CPU mirror feeds the - # plan builder; GPU copy goes into attn_metadata for kernels). context_lens_np = np.asarray(batch.context_lens, dtype=np.int32) max_seqlen_q = batch.num_spec_step + 1 positions_np = np.tile( np.arange(max_seqlen_q, dtype=np.int32), scheduled_bs ) + np.repeat(context_lens_np - max_seqlen_q, max_seqlen_q) sum_scheduled_tokens = batch.total_tokens_num_decode + var["positions"].np[:sum_scheduled_tokens] = positions_np - positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) - # cu_seqlens_q for decode: each seq has max_seqlen_q tokens. cu_seqlens_q_np = np.arange( 0, (scheduled_bs + 1) * max_seqlen_q, max_seqlen_q, dtype=np.int32 ) var["cu_seqlens_q"].np[: scheduled_bs + 1] = cu_seqlens_q_np - cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(scheduled_bs + 1) - # PR-A: V4 Compressor state-cache update reads context_lens to - # discriminate fresh prefill vs decode/prefix-cache. Parent's - # prepare_prefill pushes this; decode path must do the same. var["context_lens"].np[:scheduled_bs] = context_lens_np - context_lens_gpu = var["context_lens"].copy_to_gpu(scheduled_bs) - block_tables_gpu = self._populate_block_tables(batch, scheduled_bs) - state_slot_gpu, state_slot_np = self._populate_state_slot_mapping( - batch, scheduled_bs, return_cpu=True + # Inline block_tables CPU fill (H2D deferred to prep_stream). + block_tables_np = var["block_tables"].np + for i, block_table in enumerate(batch.block_tables[:scheduled_bs]): + block_tables_np[i] = 0 + block_tables_np[i, : len(block_table)] = block_table + + state_slot_np = np.asarray( + batch.per_req_cache_groups[:scheduled_bs], dtype=np.int32 + ) + if len(state_slot_np) < scheduled_bs: + state_slot_np = np.zeros(scheduled_bs, dtype=np.int32) + ss_buf = var["v4_meta_state_slot_groups"] + ss_buf.np[:scheduled_bs] = state_slot_np + + # ---- fire H2D on prep_stream ---- + prep_stream = self.prep_stream + current_stream = torch.cuda.current_stream() + prep_stream.wait_stream(current_stream) + with torch.cuda.stream(prep_stream): + positions = var["positions"].copy_to_gpu(sum_scheduled_tokens) + cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(scheduled_bs + 1) + context_lens_gpu = var["context_lens"].copy_to_gpu(scheduled_bs) + block_tables_gpu = var["block_tables"].copy_to_gpu(scheduled_bs) + state_slot_gpu = ss_buf.copy_to_gpu(scheduled_bs) + + # ---- CPU numpy work, overlapped with prep_stream H2D ---- + start_pos_per_seq_cpu = positions_np[cu_seqlens_q_np[:scheduled_bs]] + + extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) + compress_plans = self._build_compress_plans( + extend_lens_np, + context_lens_np, + torch.device(self.device), + for_decode_cg=True, ) + + # ---- sync, build attn_metadata, per-fwd meta ---- + current_stream.wait_stream(prep_stream) + attn_metadata = AttentionMetaData_DSV4( cu_seqlens_q=cu_seqlens_q_gpu, cu_seqlens_k=None, @@ -890,44 +945,18 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): block_tables=block_tables_gpu, context_lens=context_lens_gpu, ) - # Attach V4-specific per-seq metadata via dynamic attribute. - # AttentionMetaData is a regular class (not slotted), so this works - # without modifying its schema. attn_metadata.state_slot_mapping = state_slot_gpu - # PR-A Phase 2: CPU mirrors so the V4 forward path can read per-seq - # metadata without `.tolist()` / `.item()` GPU→CPU syncs. Generic - # naming (no `v4_` prefix) — these are not V4-specific concepts and - # other backends with stateful per-request buffers can adopt them. attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[ - cu_seqlens_q_np[:scheduled_bs] - ] - # Compress plans (per ratio) for batched fused_compress + update_states. - # Decode batch: extend_lens = max_seqlen_q for all seqs (uniform). - # `context_lens_np` is post-extend (from batch.context_lens, set by - # scheduler after appending this fwd's tokens) — this is what plan - # generation needs as `seq_lens`. Must run BEFORE - # `_attach_v4_indexer_meta` since the indexer consumes - # plan.compress_plan_cpu to derive its FP8 write-side slot_mapping. - extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) - attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device, for_decode_cg=True - ) - # CG capture path: model_runner pads cu_seqlens_q to graph_bs. The - # captured kernels still iterate `bs * max_q_len` tokens at replay, - # so per-fwd metadata MUST be sized to that padded T (otherwise - # padded slots read stale buffer entries → GPU memory access fault). - # `bs` parameter is the padded graph_bs (or actual bs in eager). - # NOTE: must run BEFORE `_attach_v4_indexer_meta` because the latter's - # `_build_v4_indexer_meta` consumes the per-fwd shared tensors - # (batch_id_per_token, n_committed_csa_per_seq) staged here. + attn_metadata.start_pos_per_seq_cpu = start_pos_per_seq_cpu + attn_metadata.compress_plans = compress_plans + padded_bs = int(bs) self._attach_v4_per_fwd_meta( attn_metadata, - positions, + positions_np, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, - attn_metadata.state_slot_mapping_cpu, + start_pos_per_seq_cpu, + state_slot_np, scheduled_bs, sum_scheduled_tokens, padded_bs=padded_bs, @@ -936,7 +965,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): self._attach_v4_indexer_meta( attn_metadata, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + start_pos_per_seq_cpu, scheduled_bs, sum_scheduled_tokens, positions_gpu=positions, @@ -1002,7 +1031,7 @@ def prepare_prefill(self, batch: ScheduledBatch): # reuse the shared GPU tensors (batch_id_per_token, n_committed_csa). self._attach_v4_per_fwd_meta( attn_metadata, - positions, + positions_np, cu_seqlens_q_np, attn_metadata.start_pos_per_seq_cpu, attn_metadata.state_slot_mapping_cpu, @@ -1035,7 +1064,7 @@ def prepare_prefill(self, batch: ScheduledBatch): def _attach_v4_per_fwd_meta( self, attn_metadata: AttentionMetaData_DSV4, - positions: torch.Tensor, + positions_np: np.ndarray, cu_seqlens_q_np, start_pos_per_seq_cpu, state_slot_mapping_cpu, @@ -1080,88 +1109,34 @@ def _attach_v4_per_fwd_meta( start_pos_per_seq_np = np.asarray( start_pos_per_seq_cpu[:scheduled_bs], dtype=np.int64 ) - # CG-padding-aware token count: padded_T == captured kernels' grid size. - # Real per-fwd data fills [0:total_tokens], rest is sentinel-padded so - # padded slot kernels skip cleanly. Eager / prefill paths leave both - # kwargs as None → no padding (padded_total_tokens == total_tokens). if padded_bs is None or max_q_len is None: padded_total_tokens = total_tokens else: padded_total_tokens = max(int(padded_bs) * int(max_q_len), total_tokens) - # ----- window_topk (per-token) ----- - # Builder-internal intermediate: consumed below by - # `_attach_v4_paged_decode_meta` to derive SWA paged offsets. Not - # exposed on attn_metadata — no kernel in V4Attention.forward reads it. - # Stored in a stable forward_vars buffer (CG requires fixed addresses). - start_pos_per_seq_gpu = self._stage( - "v4_meta_start_pos_per_seq", start_pos_per_seq_np - ) - token_num_per_seq_gpu = self._stage( - "v4_meta_token_num_per_seq", token_num_per_seq - ) - # Decode-only fast path: token_num_per_seq is all-1 so the - # repeat_interleave below is identity. is_decode_only = scheduled_bs == total_tokens and bool( (token_num_per_seq == 1).all() ) - if is_decode_only: - start_pos_per_token = start_pos_per_seq_gpu - else: - start_pos_per_token = torch.repeat_interleave( - start_pos_per_seq_gpu, token_num_per_seq_gpu - ) - window_topk_gpu = _build_window_topk_batched( - positions[:total_tokens].to(torch.long), - start_pos_per_token, - win, - out=self.model_runner.forward_vars["v4_meta_window_topk"].gpu, - ) - # ----- batch_id_per_token (single per-token mapping) ----- - # Built unconditionally. Consumed by swa_write, Phase B/C/E - # paged-decode kernels, AND the indexer (cast to long inline). - # CG: real data [0:total_tokens]; padded slots get -1 sentinel so - # consumer kernels (csa_translate_pack, swa_write) skip on `bid<0`. - # int64 dtype: PyTorch fancy-index requires int64 INDEX (used by - # the indexer in `_build_v4_indexer_meta`); triton kernels (swa_write, - # csa_translate_pack) read whatever dtype the pointer carries — int64 - # works for them too. Single buffer for all consumers, no dtype mirror. + var = self.model_runner.forward_vars + + # ---- CPU numpy work (all on main thread) ---- batch_id_per_token_np = np.full(padded_total_tokens, -1, dtype=np.int64) batch_id_per_token_np[:total_tokens] = np.repeat( np.arange(scheduled_bs, dtype=np.int64), token_num_per_seq ) - attn_metadata.batch_id_per_token = self._stage( - "v4_batch_id_per_token", batch_id_per_token_np - ) - # ----- n_committed_csa_per_seq (per-seq CSA committed count) ----- - # = ctx_len // 4. Built unconditionally so both csa_translate_pack - # (via batch_id_per_token lookup) and the indexer (uses raw per-seq - # vector) read from this single H2D-staged buffer. Kernel-side - # consumers do their own `min(., index_topk)` clamp via mask. ctx_per_seq_np = np.asarray( - self.model_runner.forward_vars["context_lens"].np[:scheduled_bs], + var["context_lens"].np[:scheduled_bs], dtype=np.int64, ) n_committed_csa_per_seq_np = (ctx_per_seq_np // 4).astype(np.int32) - attn_metadata.n_committed_csa_per_seq = self._stage( - "v4_n_committed_csa_per_seq", n_committed_csa_per_seq_np - ) - # ----- SWA write indices (last `win` tokens per seq) ----- - # `state_slot_mapping_cpu` is always sized to scheduled_bs by - # `_populate_state_slot_mapping` (zeros-fill for dummy/warmup batches); - # `token_num_per_seq.sum() == total_tokens >= 1` from the caller - # contract. write_starts/write_ends sum is `min(token_num, win).sum()`, - # also >= 1. No empty-write fallback needed. write_starts = cu_seqlens_q_arr[:scheduled_bs] + np.maximum( 0, token_num_per_seq - win ) write_ends = cu_seqlens_q_arr[1:] - # Build write_indices_np (compact, num_write entries). if is_decode_only: - # 1 tok/seq: every row is the last-win token of its seq. write_indices_np = np.arange(total_tokens, dtype=np.int64) else: write_indices_np = np.concatenate( @@ -1170,11 +1145,8 @@ def _attach_v4_per_fwd_meta( for s, e in zip(write_starts, write_ends) ] ) - # Sentinel-pad to `padded_total_tokens` so the kernel grid matches the - # captured size at replay (CG bs may exceed real bs → padded slots - # need write_indices=-1 to bail). num_write = int(write_indices_np.shape[0]) - wi_buf = self.model_runner.forward_vars["v4_meta_swa_write_indices"] + wi_buf = var["v4_meta_swa_write_indices"] cap = wi_buf.np.shape[0] assert ( padded_total_tokens <= cap @@ -1183,14 +1155,28 @@ def _attach_v4_per_fwd_meta( wi_buf.np[:num_write] = write_indices_np if num_write < padded_total_tokens: wi_buf.np[num_write:padded_total_tokens].fill(-1) + + # ---- Build window_topk on CPU (avoids ~25 small GPU kernels) ---- + if is_decode_only: + sp_per_token = start_pos_per_seq_np + else: + sp_per_token = np.repeat(start_pos_per_seq_np, token_num_per_seq) + positions_long = positions_np[:total_tokens].astype(np.int64) + window_topk_np = _build_window_topk_np(positions_long, sp_per_token, win) + wt_buf = var["v4_meta_window_topk"] + wt_buf.np[:total_tokens, :win] = window_topk_np + + # ---- Stage all buffers to GPU ---- + window_topk_gpu = wt_buf.copy_to_gpu(total_tokens) + + attn_metadata.batch_id_per_token = self._stage( + "v4_batch_id_per_token", batch_id_per_token_np + ) + attn_metadata.n_committed_csa_per_seq = self._stage( + "v4_n_committed_csa_per_seq", n_committed_csa_per_seq_np + ) attn_metadata.swa_write_indices = wi_buf.copy_to_gpu(padded_total_tokens) - # ----- Phase B: paged-decode metadata (layer-invariant) ----- - # Builds per-fwd kv_indices/kv_indptr buffers for SWA / CSA / HCA - # plus valid_count + block_tables_per_token consumed by Phase C - # (V4Attention.forward CSA per-layer translation) and Phase E - # (sparse_attn_v4_paged_decode dispatch). Skipped (None defaults) for - # non-pure-decode batches (prefill, mixed, fresh-prefill). self._attach_v4_paged_decode_meta( attn_metadata=attn_metadata, token_num_per_seq=token_num_per_seq, @@ -1340,26 +1326,29 @@ def _attach_v4_paged_decode_meta( # batch_id_per_token + n_committed_csa_per_seq already staged in # `_attach_v4_per_fwd_meta`. - # ----- HCA compress paged offsets (CPU numpy) ----- - # `paged_offset = swa_pages + block_tables[bid, idx]` for - # idx in [0, n_committed_hca[bid]). HCA is layer-invariant so we fully - # build the packed flat array here and stage once. + # ----- HCA compress paged offsets (CPU numpy, vectorized) ----- block_tables_np_full = var["block_tables"].np[:scheduled_bs] hca_total_indices = int(hca_indptr_np[T]) - hca_indices_np = np.empty(hca_total_indices, dtype=np.int32) - # Layout per token: [w0..w_{win-1}, h0..h_{n_committed_hca[bid]-1}]. - # Window prefix is overwritten from GPU swa_paged below; pre-fill -1 - # so any uninitialized prefix slot is a sentinel the kernel skips. - hca_indices_np.fill(-1) - for t in range(T): - bid = int(batch_id_per_token_np[t]) - n_h = int(n_committed_hca_per_seq[bid]) - if n_h == 0: - continue - base_t = int(hca_indptr_np[t]) + win - hca_indices_np[base_t : base_t + n_h] = ( - swa_pages + block_tables_np_full[bid, :n_h] + hca_indices_np = np.full(hca_total_indices, -1, dtype=np.int32) + n_h_per_token = n_committed_hca_per_seq[batch_id_per_token_np[:T]].astype( + np.int64 + ) + total_hca_entries = int(n_h_per_token.sum()) + if total_hca_entries > 0: + token_indices = np.repeat(np.arange(T, dtype=np.int32), n_h_per_token) + cu_n_h = np.zeros(T + 1, dtype=np.int64) + np.cumsum(n_h_per_token, out=cu_n_h[1:]) + entry_offsets = np.arange(total_hca_entries, dtype=np.int64) - np.repeat( + cu_n_h[:T], n_h_per_token ) + write_pos = ( + hca_indptr_np[token_indices].astype(np.int64) + win + entry_offsets + ) + bid_expanded = batch_id_per_token_np[token_indices] + hca_indices_np[write_pos] = ( + swa_pages + + block_tables_np_full[bid_expanded, entry_offsets.astype(np.intp)] + ).astype(np.int32) # Stage to GPU (HCA compress tail; window prefix scattered below). hca_indices_gpu = self._stage("v4_kv_indices_hca", hca_indices_np) @@ -1845,7 +1834,7 @@ def build_for_cudagraph_capture( # builder can reuse the shared per-fwd GPU tensors. self._attach_v4_per_fwd_meta( attn_metadata, - positions, + positions_np, cu_seqlens_q_np, attn_metadata.start_pos_per_seq_cpu, attn_metadata.state_slot_mapping_cpu, From 20528ee4a9ef4e0e407d41c78da95a1cce509ebb Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Sat, 9 May 2026 23:25:22 +0800 Subject: [PATCH 12/76] perf(deepseek_v4): use aiter silu_and_mul with in-kernel clamp (#731) Replace the chunk + double torch.clamp + F.silu * up sequence in Expert.forward with a single aiter.silu_and_mul(out, combined, limit) call. The new limit parameter folds the swiglu_limit clamp (gate <= limit, up in [-limit, limit]) into the kernel via the v_med3_f32 intrinsic, removing several launch-bound ops on the per-token critical path. Requires aiter PR ROCm/aiter#3104 (already merged), which adds the limit parameter and HAS_LIMIT compile-time specialization to silu_and_mul. Verified on DeepSeek-V4-Pro tp=8 --level 0: GSM8K nshot=5 (AITER_BF16_FP8_MOE_BOUND=0 + ATOM_MOE_GU_ITLV=1): run 1: 0.9515 / 0.9522 (flexible / strict) run 2: 0.9522 / 0.9530 Matches V4-Pro baseline (0.9522 / 0.9530), within 1 sigma stderr. --- atom/models/deepseek_v4.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index d96b0698a8..adf951901e 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1720,18 +1720,25 @@ def forward( x: torch.Tensor, # [num_tokens, dim] weights: Optional[torch.Tensor] = None, # [num_tokens, 1] optional gate ) -> torch.Tensor: # [num_tokens, dim] + from aiter import silu_and_mul as aiter_silu_and_mul + dtype = x.dtype - # Single fused GEMM, then chunk(2) along last dim so TP-sharded per-rank - # output (inter_dim_per_tp * 2) is split correctly regardless of tp_size. - combined = self.gate_up_proj(x).float() # [num_tokens, 2*inter_dim_per_tp] - gate, up = combined.chunk(2, dim=-1) # each [num_tokens, inter_dim_per_tp] - if self.swiglu_limit > 0: - up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) - gate = torch.clamp(gate, max=self.swiglu_limit) - x = F.silu(gate) * up # [num_tokens, inter_dim_per_tp] + # Single fused GEMM. Layout is [gate | up] concat on last dim — matches + # aiter silu_and_mul's split([d, d], dim=-1) contract. The kernel does + # silu/clamp/mul in fp32 internally regardless of input dtype, so we + # feed the bf16 GEMM output directly. + combined = self.gate_up_proj(x) # [num_tokens, 2*inter_dim_per_tp] + out = torch.empty( + (combined.shape[0], combined.shape[-1] // 2), + dtype=dtype, + device=combined.device, + ) + # limit > 0 enables in-kernel clamp (gate≤limit, up∈[-limit,limit]) via + # ROCm v_med3_f32 — same semantics as the prior torch.clamp pair. + aiter_silu_and_mul(out, combined, self.swiglu_limit) if weights is not None: - x = weights * x - return self.w2(x.to(dtype)) # [num_tokens, dim] + out = weights.to(dtype) * out + return self.w2(out) # [num_tokens, dim] class MoE(nn.Module): From 28f59f2f991652107ff20efb2a59f19dba3c0f3e Mon Sep 17 00:00:00 2001 From: gbyu-amd Date: Sun, 10 May 2026 00:38:46 +0800 Subject: [PATCH 13/76] fix quant mapping for kimi k25 (#732) Co-authored-by: Guanbao Yu --- atom/models/kimi_k25.py | 4 ++++ atom/plugin/vllm/models/kimi_k25.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/atom/models/kimi_k25.py b/atom/models/kimi_k25.py index eb4aaafb61..6084f74ebf 100644 --- a/atom/models/kimi_k25.py +++ b/atom/models/kimi_k25.py @@ -40,6 +40,10 @@ class KimiK25ForCausalLM(nn.Module): "vision_tower.", "mm_projector.", ] + quant_exclude_name_mapping = { + "language_model.model.": "model.", + "language_model.lm_head": "lm_head", + } def __init__( self, diff --git a/atom/plugin/vllm/models/kimi_k25.py b/atom/plugin/vllm/models/kimi_k25.py index 47cfb9ce22..a9f73b7b6c 100644 --- a/atom/plugin/vllm/models/kimi_k25.py +++ b/atom/plugin/vllm/models/kimi_k25.py @@ -173,6 +173,10 @@ class KimiK25ForConditionalGeneration_(vLLMKimiK25): "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } + quant_exclude_name_mapping = { + "language_model.model.": "model.language_model.model.", + "language_model.lm_head": "model.language_model.lm_head", + } hf_to_atom_mapper = WeightsMapper( orig_to_new_prefix={ "model.visual.": "visual.", From 3e06687aa47559d2bb0bf0fd363c6770a7910e77 Mon Sep 17 00:00:00 2001 From: gbyu-amd Date: Sun, 10 May 2026 10:07:01 +0800 Subject: [PATCH 14/76] [dockerfile] temp reinstall triton==3.6.0 for accuracy (#726) * temp reinstall triton==3.6.0 * reinstall triton for vLLM-ATOM * update comments for better tracking --------- Co-authored-by: Guanbao Yu --- docker/Dockerfile | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docker/Dockerfile b/docker/Dockerfile index 38f93a2152..964168bd6e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -155,6 +155,14 @@ RUN echo "========== [OOT] Restore base image triton ==========" && \ rm -f /tmp/triton_pip_show.txt && \ "${VENV_PYTHON}" -c "import triton; print(f'triton.__version__ = {triton.__version__}')" +# Qwen3-next encounter accuracy issue with the triton from base image +# we install verified good version triton==3.6.0 here as a temporary fix +# this should be removed after the triton issue is resolved +RUN echo "========== [OOT] Reinstall known good triton ==========" && \ + pip uninstall -y triton amd-triton && \ + pip install triton==3.6.0 +RUN pip show triton || true + RUN echo "========== [vLLM-ATOM] Final transformers version ==========" && \ "${VENV_PYTHON}" -c "import transformers; print(f'transformers.__version__ = {transformers.__version__}')" && \ "${VENV_PYTHON}" -m pip show transformers || true @@ -390,6 +398,13 @@ RUN cd /app/aiter-test && pip install -e . --no-build-isolation && \ RUN pip install rocm-trace-lite && \ rtl --version || true +# Qwen3-next encounter accuracy issue with the triton from base image +# we install verified good version triton==3.6.0 here as a temporary fix +# this should be removed after the triton issue is resolved +RUN pip uninstall -y triton amd-triton && \ + pip install triton==3.6.0 +RUN pip show triton || true + # ATOM: lightweight install (no compilation needed) # CACHEBUST invalidates only this layer so parallel stages stay cached ARG CACHEBUST=1 From 303a2e43e3347d7e13726f00ef53da06d177de15 Mon Sep 17 00:00:00 2001 From: Shao-Chun Lee Date: Sat, 9 May 2026 22:04:22 -0500 Subject: [PATCH 15/76] [Triton] add fused_routing_from_topk switch (#725) * add _aiter_fused_routing_from_topk switch * ruff fix * change import * black format --- atom/model_ops/fused_moe_triton.py | 53 ++++++++++++++++++++++++++++++ atom/model_ops/moe.py | 4 +-- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index 18c513e809..f535af6312 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -26,6 +26,9 @@ from math import prod from aiter import ActivationType from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.fusions.fused_routing_from_topk import ( + fused_routing_from_topk as _aiter_fused_routing_from_topk, +) from atom.model_ops.utils import has_triton_kernels logger = logging.getLogger("atom") @@ -114,6 +117,56 @@ def _swizzle_mxfp4(quant_tensor, scale): return quant_tensor, InFlexData(), scale +def fused_routing_from_topk_triton(topk_weights, topk_ids, n_expts_tot): + """Build matmul_ogs routing data via the AITER fused-routing kernel. + + Thin bridge over ``aiter.ops.triton.fused_routing_from_topk``: invokes + the single-CTA counting-sort kernel for small NK and packages the + resulting indices into the ``RoutingData`` / ``GatherIndx`` / + ``ScatterIndx`` structures consumed by + ``triton_kernels.matmul_ogs``. For ``NK = n_tokens * n_expts_act`` + above the kernel's single-CTA budget (prefill-shaped inputs), falls + back to the multi-kernel ``routing_from_topk`` reference defined + below — that path does the per-row sort + global stable argsort in + plain torch and is correctness-stable at any NK. + + Equivalence vs reference: the fused kernel skips the per-row sort, + so ``topk_indx`` / ``gate_indx`` differ at intra-expert ordering. + ``hist`` and the per-(token, expert, weight) bucket assignments + match exactly; ``matmul_ogs`` is commutative over per-expert slices + so the MoE output is unchanged (up to FP non-associativity). + """ + if not has_triton_kernels(): + return routing_from_topk(topk_weights, topk_ids, n_expts_tot) + + n_tokens, n_expts_act = topk_weights.shape + n_gates_pad = n_tokens * n_expts_act + + if n_gates_pad > 4096: + # Single-CTA design exceeded; fall back rather than degrading + # silently. Typically only hit during prefill. + return routing_from_topk(topk_weights, topk_ids, n_expts_tot) + + hist, topk_indx, gate_indx, gate_scal = _aiter_fused_routing_from_topk( + topk_weights, topk_ids, n_expts_tot + ) + + # Package as the matmul_ogs routing data structures. + from triton_kernels.routing import ( + RoutingData, + GatherIndx, + ScatterIndx, + compute_expt_data, + ) + + gather_indx = GatherIndx(src_indx=topk_indx, dst_indx=gate_indx) + scatter_indx = ScatterIndx(src_indx=gate_indx, dst_indx=topk_indx) + expt_data = compute_expt_data(hist, n_expts_tot, n_gates_pad) + + routing_data = RoutingData(gate_scal, hist, n_expts_tot, n_expts_act, expt_data) + return routing_data, gather_indx, scatter_indx + + def routing_from_topk(topk_weights, topk_ids, n_expts_tot): """Convert FusedMoE.select_experts output to triton routing data structures. diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 4244e0a3cf..b2aeac8652 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -941,7 +941,7 @@ def apply( from atom.model_ops.fused_moe_triton import ( triton_kernel_moe_forward, triton_kernel_fused_experts, - routing_from_topk, + fused_routing_from_topk_triton, ) # Check if the model needs custom routing that triton routing() @@ -978,7 +978,7 @@ def apply( n_expts_tot = global_num_experts n_expts_tot = n_expts_tot + layer.num_fused_shared_experts - routing_data, gather_idx, scatter_idx = routing_from_topk( + routing_data, gather_idx, scatter_idx = fused_routing_from_topk_triton( topk_weights, topk_ids, n_expts_tot ) From 28f9702de2a195018303dabae298f410c72a1446 Mon Sep 17 00:00:00 2001 From: Pleaplusone Date: Sun, 10 May 2026 12:29:07 +0800 Subject: [PATCH 16/76] Remove chunk split in Qwen3.5 and Qwen3Next (#682) * remove chunk split Signed-off-by: ganyi * fix accuracy Signed-off-by: ganyi * add None input for not quantize case Signed-off-by: ganyi * fix ci Signed-off-by: ganyi * maybe fix the ci benchmark crash Signed-off-by: ganyi * add env var for bare atom acc test Signed-off-by: ganyi * add fp8 kv cache for qwen3next bare atom ci Signed-off-by: ganyi * make fusion as default Signed-off-by: ganyi * disable non-persistent pa Signed-off-by: ganyi * add num tokens Signed-off-by: ganyi --------- Signed-off-by: ganyi --- .github/benchmark/models_accuracy.json | 12 +- atom/model_ops/attention_gdn.py | 21 +- .../model_ops/fla_ops/fused_sigmoid_gating.py | 16 +- atom/model_ops/layernorm.py | 32 +- atom/model_ops/linear.py | 338 +++++++++++++++--- atom/model_ops/mamba_ops/causal_conv1d.py | 163 ++++----- atom/models/qwen3_5.py | 36 +- atom/models/qwen3_next.py | 161 +++------ atom/plugin/attention_mha.py | 7 +- .../vllm/attention_backend/attention_gdn.py | 21 +- atom/utils/envs.py | 2 +- 11 files changed, 501 insertions(+), 308 deletions(-) diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 94d1dcce48..d175f91da2 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -111,8 +111,8 @@ { "model_name": "Qwen3-Next-80B-A3B-Thinking", "model_path": "Qwen/Qwen3-Next-80B-A3B-Thinking", - "extraArgs": "-tp 8", - "env_vars": "", + "extraArgs": "-tp 8 --kv_cache_dtype fp8", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.65, @@ -209,7 +209,7 @@ "model_name": "Qwen3.5-397B-A17B-FP8", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--kv_cache_dtype fp8 -tp 4", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", "accuracy_threshold": 0.85, @@ -221,7 +221,7 @@ "model_name": "Qwen3.5-397B-A17B-FP8 MTP", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--kv_cache_dtype fp8 -tp 4 --method mtp --num-speculative-tokens 3", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "main", "accuracy_threshold": 0.85, @@ -233,7 +233,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--kv_cache_dtype fp8 -tp 4", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "pr", "accuracy_threshold": 0.835, @@ -245,7 +245,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4 MTP", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--kv_cache_dtype fp8 -tp 4 --method mtp --num-speculative-tokens 3", - "env_vars": "", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "atom-mi355-8gpu.predownload", "test_level": "main", "accuracy_threshold": 0.835, diff --git a/atom/model_ops/attention_gdn.py b/atom/model_ops/attention_gdn.py index 26581188c3..d6230e9eba 100644 --- a/atom/model_ops/attention_gdn.py +++ b/atom/model_ops/attention_gdn.py @@ -31,17 +31,19 @@ def fused_gdn_gating_kernel( dt_bias, seq_len, NUM_HEADS: tl.constexpr, + stride_a_batch, + stride_b_batch, beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + out_off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) - blk_b = tl.load(b + off, mask=mask) + blk_a = tl.load(a + i_b * stride_a_batch + head_off, mask=mask) + blk_b = tl.load(b + i_b * stride_b_batch + head_off, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) # If the model is loaded in fp16, without the .float() here, A might be -inf x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) @@ -49,11 +51,13 @@ def fused_gdn_gating_kernel( beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x ) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + out_off, blk_g.to(g.dtype.element_ty), mask=mask) # compute beta_output = sigmoid(b) blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) tl.store( - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + beta_output + out_off, + blk_beta_output.to(beta_output.dtype.element_ty), + mask=mask, ) @@ -85,6 +89,8 @@ def fused_gdn_gating( dt_bias, seq_len, num_heads, + a.stride(0), + b.stride(0), beta, threshold, 8, @@ -161,6 +167,7 @@ def forward( fwd_ctx.attn_metadata, "gdn_metadata", None ) if gdn_metadata is None: + core_attn_out.zero_() return core_attn_out gdn_cache = fwd_ctx.kv_cache_data @@ -373,4 +380,8 @@ def forward( else: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + # Zero padding tail for CUDA graph replay safety + if num_actual_tokens < core_attn_out.shape[0]: + core_attn_out[num_actual_tokens:].zero_() + return core_attn_out diff --git a/atom/model_ops/fla_ops/fused_sigmoid_gating.py b/atom/model_ops/fla_ops/fused_sigmoid_gating.py index 5f4df00484..236504ddc1 100644 --- a/atom/model_ops/fla_ops/fused_sigmoid_gating.py +++ b/atom/model_ops/fla_ops/fused_sigmoid_gating.py @@ -48,6 +48,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + stride_a_token, + stride_b_token, stride_init_state_token: tl.constexpr, stride_final_state_token: tl.constexpr, stride_indices_seq: tl.constexpr, @@ -87,13 +89,13 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_A_log = A_log + i_hv if not IS_KDA: - p_a = a + bos * HV + i_hv + p_a = a + bos * stride_a_token + i_hv p_dt_bias = dt_bias + i_hv else: p_a = a + (bos * HV + i_hv) * K + o_k p_dt_bias = dt_bias + i_hv * K + o_k - p_b = b + bos * HV + i_hv + p_b = b + bos * stride_b_token + i_hv p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v mask_k = o_k < K @@ -175,8 +177,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( p_k += H * K p_o += HV * V p_v += HV * V - p_b += HV - p_a += HV + p_b += stride_b_token + p_a += stride_a_token def fused_sigmoid_gating_delta_rule_update( @@ -246,8 +248,8 @@ def fused_sigmoid_gating_delta_rule_update( grid = (NK, NV, N * HV) fused_sigmoid_gating_delta_rule_update_kernel[grid]( A_log=A_log, - a=a.contiguous(), - b=b.contiguous(), + a=a, + b=b, dt_bias=dt_bias, beta=beta, threshold=threshold, @@ -270,6 +272,8 @@ def fused_sigmoid_gating_delta_rule_update( V=V, BK=BK, BV=BV, + stride_a_token=a.stride(-2), + stride_b_token=b.stride(-2), stride_init_state_token=stride_init_state_token, stride_final_state_token=stride_final_state_token, stride_indices_seq=stride_indices_seq, diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 86488dd5d5..91c69083fd 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -549,12 +549,38 @@ def forward_cuda( x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - from atom.model_ops.triton_gemma_rmsnorm import gemma_rmsnorm_triton + # Use the aiter HIP fused_qk_rmsnorm_group_quant kernel in no-quant mode + # (q_out_scale=None) to perform Gemma RMSNorm + optional residual add. + # Same math as the Triton kernel: out = rmsnorm(x [+ residual]) * (1 + w), + # but executed by the aiter kernel for higher achieved bandwidth. + from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant + + ori_shape = x.shape + x_2d = x.view(-1, ori_shape[-1]) + + out = torch.empty_like(x_2d) + if residual is not None: + residual_2d = residual.view(-1, ori_shape[-1]) + res_out = torch.empty_like(x_2d) + else: + residual_2d = None + res_out = None - return gemma_rmsnorm_triton( - x, self.weight.data, self.variance_epsilon, residual + fused_qk_rmsnorm_group_quant( + q=x_2d, + q_weight=self.weight.data, + q_epsilon=self.variance_epsilon, + q_out_unquantized=out, + q_res_out=res_out, + q_residual=residual_2d, + gemma_norm=True, ) + out = out.view(ori_shape) + if residual is not None: + return out, res_out.view(ori_shape) + return out + def _forward_fused_fp8(self, x, residual=None): from aiter.ops.fused_qk_rmsnorm_group_quant import fused_qk_rmsnorm_group_quant from aiter.utility.dtypes import fp8 diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index c3a09e829e..9a2859c0a8 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -652,6 +652,17 @@ def weight_loader( class QKVZBAParallelLinear(ColumnParallelLinear): + """Fused QKVZBA linear with deinterleaved output layout. + + Output layout is always ``[q | k | v | z | b | a]`` contiguous on dim-0, + so the caller can use ``torch.split`` for zero-copy views. + + The weight_loader deinterleaves the Qwen3-Next interleaved per-k-head-group + checkpoint layout during loading. Qwen3.5 checkpoints (which have separate + per-component weights) are loaded via individual shard_ids and work without + deinterleaving. + """ + def __init__( self, input_size: int, @@ -672,11 +683,15 @@ def __init__( tp_size = get_tp_group().world_size self.num_k_heads = divide(self.num_k_heads, tp_size) self.num_v_heads = divide(self.num_v_heads, tp_size) - output_sizes = [ - (2 * head_k_dim * self.num_k_heads + 2 * head_v_dim * self.num_v_heads) - * tp_size, - 2 * self.num_v_heads * tp_size, - ] + # Output layout: [q_all | k_all | v_all | z_all | b_all | a_all] + q_size = self.num_k_heads * head_k_dim + k_size = self.num_k_heads * head_k_dim + v_size = self.num_v_heads * head_v_dim + z_size = self.num_v_heads * head_v_dim + b_size = self.num_v_heads + a_size = self.num_v_heads + self._section_sizes = [q_size, k_size, v_size, z_size, b_size, a_size] + output_sizes = [s * tp_size for s in self._section_sizes] super().__init__( input_size, output_sizes, @@ -686,71 +701,304 @@ def __init__( prefix=prefix, ) + # -- helpers for deinterleaving during weight loading -- + + def _deinterleave_qkvz(self, param_data: torch.Tensor, loaded_weight: torch.Tensor): + """Scatter interleaved qkvz checkpoint rows into [q|k|v|z] regions. + + Checkpoint layout per k-head group (``QKVZ_DIM_SIZE`` rows): + [q(head_k_dim) | k(head_k_dim) | v0..vR(R*head_v_dim) | z0..zR(R*head_v_dim)] + where R = num_v_heads / num_k_heads (KV_HEAD_RATIO). + """ + nk = self.num_k_heads + hk = self.head_k_dim + hv = self.head_v_dim + R = self.num_v_heads // nk # KV_HEAD_RATIO + group_size = 2 * hk + 2 * hv * R + + q_total = nk * hk + k_total = nk * hk + v_total = self.num_v_heads * hv + + # TP shard the source + src = loaded_weight.narrow( + self.tp_dim, self.tp_rank * nk * group_size, nk * group_size + ) + + for g in range(nk): + base = g * group_size + # q rows + param_data[g * hk : (g + 1) * hk] = src[base : base + hk] + # k rows + param_data[q_total + g * hk : q_total + (g + 1) * hk] = src[ + base + hk : base + 2 * hk + ] + # v sub-heads + for s in range(R): + v_src_start = base + 2 * hk + s * hv + v_dst_start = q_total + k_total + (g * R + s) * hv + param_data[v_dst_start : v_dst_start + hv] = src[ + v_src_start : v_src_start + hv + ] + # z sub-heads + for s in range(R): + z_src_start = base + 2 * hk + R * hv + s * hv + z_dst_start = q_total + k_total + v_total + (g * R + s) * hv + param_data[z_dst_start : z_dst_start + hv] = src[ + z_src_start : z_src_start + hv + ] + + def _deinterleave_ba(self, param_data: torch.Tensor, loaded_weight: torch.Tensor): + """Scatter interleaved ba checkpoint rows into [b|a] regions. + + Checkpoint layout per k-head group (2*R elements): + [b_sub0, b_sub1, ..., b_subR-1, a_sub0, a_sub1, ..., a_subR-1] + where R = num_v_heads / num_k_heads. + """ + nk = self.num_k_heads + nv = self.num_v_heads + R = nv // nk + ba_total = 2 * nv + + # TP shard the source + src = loaded_weight.narrow(self.tp_dim, self.tp_rank * ba_total, ba_total) + + qkvz_total = sum(self._section_sizes[:4]) + b_offset = qkvz_total + a_offset = qkvz_total + nv + + for g in range(nk): + group_base = g * 2 * R + # b sub-heads + for s in range(R): + param_data[b_offset + g * R + s] = src[group_base + s] + # a sub-heads + for s in range(R): + param_data[a_offset + g * R + s] = src[group_base + R + s] + def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str ): param_data = param.data assert loaded_shard_id in ["qkvz", "ba", "qkv", "z", "b", "a"] + + is_scale = param is getattr(self, "weight_scale", None) or param is getattr( + self, "input_scale", None + ) + + # For interleaved checkpoint shards ("qkvz", "ba"), deinterleave + # weight rows so output is [q|k|v|z|b|a] contiguous. + if loaded_shard_id == "qkvz" and not is_scale: + self._deinterleave_qkvz(param_data, loaded_weight) + return + if loaded_shard_id == "ba" and not is_scale: + self._deinterleave_ba(param_data, loaded_weight) + return + + # For individual shard_ids or scale params, use offset-based loading. + q_size, k_size, v_size, z_size, b_size, a_size = self._section_sizes if loaded_shard_id == "qkvz": - shard_size = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) + shard_size = q_size + k_size + v_size + z_size shard_offset = 0 - shard_rank = self.tp_rank elif loaded_shard_id == "qkv": - shard_size = ( - 2 * self.num_k_heads * self.head_k_dim - + self.num_v_heads * self.head_v_dim - ) + shard_size = q_size + k_size + v_size shard_offset = 0 - shard_rank = self.tp_rank elif loaded_shard_id == "z": - shard_size = self.num_v_heads * self.head_v_dim - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = z_size + shard_offset = q_size + k_size + v_size elif loaded_shard_id == "ba": - shard_size = 2 * self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = b_size + a_size + shard_offset = q_size + k_size + v_size + z_size elif loaded_shard_id == "b": - shard_size = self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - ) - shard_rank = self.tp_rank + shard_size = b_size + shard_offset = q_size + k_size + v_size + z_size elif loaded_shard_id == "a": - shard_size = self.num_v_heads - shard_offset = ( - 2 * self.num_k_heads * self.head_k_dim - + 2 * self.num_v_heads * self.head_v_dim - + self.num_v_heads - ) - shard_rank = self.tp_rank + shard_size = a_size + shard_offset = q_size + k_size + v_size + z_size + b_size - if param is getattr(self, "weight_scale", None) or param is getattr( - self, "input_scale", None - ): + if is_scale: if self.quant_type == QuantType.per_1x128: shard_offset = (shard_offset + 127) // 128 shard_size = (shard_size + 127) // 128 elif self.quant_type == QuantType.per_Tensor: loaded_weight = loaded_weight.view(1, 1).repeat(self.tp_size, 1) - shard_offset = ["qkvz", "ba"].index(loaded_shard_id) + shard_offset = ["qkvz", "ba", "qkv", "z", "b", "a"].index( + loaded_shard_id + ) shard_size = 1 - start_idx = shard_rank * shard_size + start_idx = self.tp_rank * shard_size param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size) loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size) param.weight_loader_process(param_data, loaded_weight) +class QKVZParallelLinear(ColumnParallelLinear): + """Deinterleaving linear for the separate ``in_proj_qkvz`` checkpoint. + + The Qwen3-Next checkpoint stores qkvz in per-k-head-group interleaved + layout:: + + [q(hk) | k(hk) | v0..vR(R*hv) | z0..zR(R*hv)] per group + + This class deinterleaves both weight rows **and** per-1x128 block scales + during loading so the output layout is ``[q_all | k_all | v_all | z_all]`` + contiguous, enabling zero-copy ``torch.split`` at runtime. + """ + + def __init__( + self, + input_size: int, + head_k_dim: int, + head_v_dim: int, + num_k_heads: int, + num_v_heads: int, + bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + source_quant_dtype: torch.dtype = None, + prefix: str = "", + **kwargs, + ): + self.head_k_dim = head_k_dim + self.head_v_dim = head_v_dim + tp_size = get_tp_group().world_size + self.nk = divide(num_k_heads, tp_size) + self.nv = divide(num_v_heads, tp_size) + self.R = num_v_heads // num_k_heads # v-heads per k-head group + output_size = 2 * num_k_heads * head_k_dim + 2 * num_v_heads * head_v_dim + super().__init__( + input_size, + output_size, + bias=bias, + quant_config=quant_config, + source_quant_dtype=source_quant_dtype, + prefix=prefix, + ) + + @staticmethod + def _deinterleave(param_data, src, nk, R, hk, hv): + """Scatter interleaved rows into [q|k|v|z] regions.""" + group_size = 2 * hk + 2 * hv * R + q_total = nk * hk + k_total = nk * hk + v_total = nk * R * hv + + for g in range(nk): + base = g * group_size + # q + param_data[g * hk : (g + 1) * hk] = src[base : base + hk] + # k + param_data[q_total + g * hk : q_total + (g + 1) * hk] = src[ + base + hk : base + 2 * hk + ] + # v sub-heads + for s in range(R): + v_src = base + 2 * hk + s * hv + v_dst = q_total + k_total + (g * R + s) * hv + param_data[v_dst : v_dst + hv] = src[v_src : v_src + hv] + # z sub-heads + for s in range(R): + z_src = base + 2 * hk + R * hv + s * hv + z_dst = q_total + k_total + v_total + (g * R + s) * hv + param_data[z_dst : z_dst + hv] = src[z_src : z_src + hv] + + @staticmethod + def _match_dtype(param_data, loaded_weight): + """View param_data as loaded_weight's dtype if they differ but share element size. + + This mirrors ``weight_loader_process`` behaviour for FP8 on ROCm where + the param is ``float8_e4m3fnuz`` but the checkpoint stores + ``float8_e4m3fn``. The normalisation happens later in + ``process_weights_after_loading``. + """ + if ( + param_data.dtype != loaded_weight.dtype + and param_data.element_size() == loaded_weight.element_size() + ): + return param_data.view(loaded_weight.dtype) + return param_data + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + # Load like ColumnParallelLinear (no deinterleave at load time). + super().weight_loader(param, loaded_weight) + + def process_weights_after_loading(self): + nk, R = self.nk, self.R + hk, hv = self.head_k_dim, self.head_v_dim + + # Deinterleave weight rows: interleaved → [q|k|v|z] + w = self.weight.data + dw = torch.empty_like(w) + self._deinterleave(dw, w, nk, R, hk, hv) + self.weight.data = dw + + # Deinterleave weight_scale rows (per_1x128 block scale) + ws = getattr(self, "weight_scale", None) + if ws is not None: + hk_s, hv_s = hk // 128, hv // 128 + s = ws.data + ds = torch.empty_like(s) + self._deinterleave(ds, s, nk, R, hk_s, hv_s) + self.weight_scale.data = ds + + super().process_weights_after_loading() + + +class BAParallelLinear(ColumnParallelLinear): + """Deinterleaving linear for the separate ``in_proj_ba`` checkpoint. + + The Qwen3-Next checkpoint stores ba in per-k-head-group layout:: + + [b_s0..b_sR | a_s0..a_sR] per group (R = num_v_heads / num_k_heads) + + This class deinterleaves during loading so the output layout is + ``[b_all | a_all]`` contiguous, enabling zero-copy ``torch.split``. + + ``in_proj_ba`` is always BF16 (listed in ``modules_to_not_convert``), + so no weight-scale handling is needed. + """ + + def __init__( + self, + input_size: int, + num_k_heads: int, + num_v_heads: int, + bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + source_quant_dtype: torch.dtype = None, + prefix: str = "", + **kwargs, + ): + tp_size = get_tp_group().world_size + self.nk = divide(num_k_heads, tp_size) + self.R = num_v_heads // num_k_heads + self.nv = self.nk * self.R + output_size = 2 * num_v_heads + super().__init__( + input_size, + output_size, + bias=bias, + quant_config=quant_config, + source_quant_dtype=source_quant_dtype, + prefix=prefix, + ) + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + # Load like ColumnParallelLinear (no deinterleave at load time). + super().weight_loader(param, loaded_weight) + + def process_weights_after_loading(self): + nk, R, nv = self.nk, self.R, self.nv + w = self.weight.data + dw = torch.empty_like(w) + group_size = 2 * R + for g in range(nk): + base = g * group_size + dw[g * R : (g + 1) * R] = w[base : base + R] + dw[nv + g * R : nv + (g + 1) * R] = w[base + R : base + 2 * R] + self.weight.data = dw + super().process_weights_after_loading() + + class QKVGParallelLinear(ColumnParallelLinear): """QKV + output-Gate parallel linear. diff --git a/atom/model_ops/mamba_ops/causal_conv1d.py b/atom/model_ops/mamba_ops/causal_conv1d.py index 94c74a1447..bb60a40664 100644 --- a/atom/model_ops/mamba_ops/causal_conv1d.py +++ b/atom/model_ops/mamba_ops/causal_conv1d.py @@ -91,33 +91,16 @@ def _causal_conv1d_fwd_kernel( # continuous batching # BLOCK_N elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - idx_feats_start = tl.program_id(1) * BLOCK_N - o_ptr = x_ptr - stride_o_token = stride_x_token - stride_o_dim = stride_x_dim - idx_feats_o = idx_feats - dim_limits = dim - # Assume the idx_feats can be fully divided by BLOCK_N - if idx_feats_start < k_start_dim: - o_ptr = query_ptr - stride_o_token = query_token_stride - stride_o_dim = query_dim_stride - idx_feats_o = idx_feats - dim_limits = k_dim_size - elif idx_feats_start < v_start_dim: - o_ptr = key_ptr - stride_o_token = key_token_stride - stride_o_dim = key_dim_stride - idx_feats_o = idx_feats - k_start_dim - dim_limits = k_dim_size - elif idx_feats_start < dim: - o_ptr = value_ptr - stride_o_token = value_token_stride - stride_o_dim = value_dim_stride - idx_feats_o = idx_feats - v_start_dim - dim_limits = v_dim_size - else: - return + + # Pre-compute per-output feature indices and block-level masks. + # BLOCK_N divides k_dim_size evenly, so each program block falls entirely + # within one of q/k/v — only one mask is all-true per block. + q_feat_idx = idx_feats + k_feat_idx = idx_feats - k_start_dim + v_feat_idx = idx_feats - v_start_dim + is_q_block = idx_feats < k_start_dim + is_k_block = (idx_feats >= k_start_dim) & (idx_feats < v_start_dim) + is_v_block = idx_feats >= v_start_dim if idx_seq == pad_slot_id: return @@ -505,17 +488,25 @@ def _causal_conv1d_fwd_kernel( # continuous batching if SILU_ACTIVATION: acc = acc / (1 + tl.exp(-acc)) - mask_1d = (idx_token < segment_len) & ( - idx_feats_o < dim_limits - ) # token-index # feature-index + token_pos = sequence_start_index + token_offset + idx_token + mask_token = idx_token < segment_len - o_ptrs = ( - o_ptr - + (sequence_start_index + token_offset + idx_token) * stride_o_token - + (idx_feats_o * stride_o_dim) + # Store to q/k/v via separate tl.store calls (avoids pointer + # reassignment that crashes AMD TritonAMDGPUCanonicalizePointers). + # Each block falls entirely within one output, so only one store + # per block writes real data. + q_ptrs = ( + query_ptr + token_pos * query_token_stride + q_feat_idx * query_dim_stride ) + tl.store(q_ptrs, acc, mask=mask_token & is_q_block) - tl.store(o_ptrs, acc, mask=mask_1d) + k_ptrs = key_ptr + token_pos * key_token_stride + k_feat_idx * key_dim_stride + tl.store(k_ptrs, acc, mask=mask_token & is_k_block) + + v_ptrs = ( + value_ptr + token_pos * value_token_stride + v_feat_idx * value_dim_stride + ) + tl.store(v_ptrs, acc, mask=mask_token & is_v_block) def causal_conv1d_fn( @@ -596,7 +587,6 @@ def causal_conv1d_fn( # Store original dtype to cast back at the end original_x_dtype = x.dtype x = x.to(conv_states.dtype) - out = torch.empty_like(x) if metadata is not None: nums_dict = metadata.nums_dict args = nums_dict @@ -656,15 +646,7 @@ def causal_conv1d_fn( stride_istate_seq = conv_states.stride(0) stride_istate_dim = conv_states.stride(1) stride_istate_token = conv_states.stride(2) - # Keep this aligned with upstream: the Triton kernel consumes the - # runtime conv_state strides directly instead of requiring dim-stride 1. - # assert stride_istate_dim == 1 - if out.dim() == 2: - stride_o_dim = out.stride(0) - stride_o_token = out.stride(1) - else: - stride_o_dim = out.stride(1) - stride_o_token = out.stride(2) + stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0 if validate_data: @@ -884,42 +866,18 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - idx_feats_start = tl.program_id(1) * BLOCK_N k_start_point: tl.constexpr = k_dim_size v_start_point: tl.constexpr = k_dim_size * 2 - stride_o_seq = stride_q_seq - - o_ptr = x_ptr - stride_o_token = stride_x_token - stride_o_dim = stride_x_dim - idx_feats_o = idx_feats - dim_limits = dim - if idx_feats_start < k_start_point: - o_ptr = query_ptr - # In VARLEN mode, tokens are indexed along dim-0 of the output tensor, - # so the token stride is stride_q_seq (not stride_q_token which is the - # last-dim stride of the [num_tokens, k_dim, 1] output). - stride_o_token = stride_q_seq if IS_VARLEN else stride_q_token - stride_o_dim = stride_q_dim - idx_feats_o = idx_feats - dim_limits = k_dim_size - stride_o_seq = stride_q_seq - elif idx_feats_start < v_start_point: - o_ptr = key_ptr - stride_o_token = stride_k_seq if IS_VARLEN else stride_k_token - stride_o_dim = stride_k_dim - idx_feats_o = idx_feats - k_start_point - dim_limits = k_dim_size - stride_o_seq = stride_k_seq - elif idx_feats_start < dim: - o_ptr = value_ptr - stride_o_token = stride_v_seq if IS_VARLEN else stride_v_token - stride_o_dim = stride_v_dim - idx_feats_o = idx_feats - v_start_point - dim_limits = v_dim_size - stride_o_seq = stride_v_seq - else: - return + + # Pre-compute per-output feature indices and block-level masks. + # BLOCK_N divides k_dim_size evenly, so each program block falls entirely + # within one of q/k/v — only one mask is all-true per block. + q_feat_idx = idx_feats + k_feat_idx = idx_feats - k_start_point + v_feat_idx = idx_feats - v_start_point + is_q_block = idx_feats < k_start_point + is_k_block = (idx_feats >= k_start_point) & (idx_feats < v_start_point) + is_v_block = (idx_feats >= v_start_point) & (idx_feats < dim) if IS_APC_ENABLED: # Get the state from the initial_state_idx @@ -946,12 +904,25 @@ def _causal_conv1d_update_kernel( state_len = state_len - (seqlen - (query_end_index - query_start_index)) seqlen = query_end_index - query_start_index x_offset = query_start_index * stride_x_token - o_offset = query_start_index * stride_o_token + # Varlen: output is [num_tokens, split_dim, 1], token_pos indexes dim-0 + # so the "token" stride is stride_q_seq (the dim-0 stride). + q_base = query_start_index * stride_q_seq + k_base = query_start_index * stride_k_seq + v_base = query_start_index * stride_v_seq + q_tok_stride = stride_q_seq + k_tok_stride = stride_k_seq + v_tok_stride = stride_v_seq else: query_start_index = idx_seq * seqlen query_end_index = query_start_index + seqlen x_offset = idx_seq * stride_x_seq - o_offset = idx_seq * stride_o_seq + # Non-varlen: output is [batch, split_dim, seqlen] + q_base = idx_seq * stride_q_seq + k_base = idx_seq * stride_k_seq + v_base = idx_seq * stride_v_seq + q_tok_stride = stride_q_token + k_tok_stride = stride_k_token + v_tok_stride = stride_v_token if query_start_index == query_end_index: return @@ -1184,14 +1155,24 @@ def _causal_conv1d_update_kernel( if SILU_ACTIVATION: acc = acc / (1 + tl.exp(-acc)) - mask_1d = (idx_token < seqlen) & ( - idx_feats_o < dim_limits - ) # token-index # feature-index - o_ptrs = ( - o_ptr + o_offset + idx_token * stride_o_token + (idx_feats_o * stride_o_dim) + mask_token = idx_token < seqlen + + # Store to q/k/v via separate tl.store calls (avoids pointer + # reassignment that crashes AMD TritonAMDGPUCanonicalizePointers). + # Each block falls entirely within one output, so only one store + # per block writes real data. + q_ptrs = ( + query_ptr + q_base + idx_token * q_tok_stride + q_feat_idx * stride_q_dim ) + tl.store(q_ptrs, acc, mask=mask_token & is_q_block) - tl.store(o_ptrs, acc, mask=mask_1d) + k_ptrs = key_ptr + k_base + idx_token * k_tok_stride + k_feat_idx * stride_k_dim + tl.store(k_ptrs, acc, mask=mask_token & is_k_block) + + v_ptrs = ( + value_ptr + v_base + idx_token * v_tok_stride + v_feat_idx * stride_v_dim + ) + tl.store(v_ptrs, acc, mask=mask_token & is_v_block) def causal_conv1d_update( @@ -1308,19 +1289,15 @@ def causal_conv1d_update( stride_k_seq, stride_k_dim, stride_k_token = key.stride() stride_v_seq, stride_v_dim, stride_v_token = value.stride() - out = x stride_w_dim, stride_w_width = weight.stride() if query_start_loc is None: # X (batch, dim, seqlen) stride_x_seq, stride_x_dim, stride_x_token = x.stride() - stride_o_seq, stride_o_dim, stride_o_token = out.stride() else: - # X (dim, cu_seqlen) + # X (num_tokens, dim) stride_x_token, stride_x_dim = x.stride() stride_x_seq = 0 - stride_o_token, stride_o_dim = out.stride() - stride_o_seq = 0 stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() stride_state_indices = ( @@ -1392,8 +1369,6 @@ def grid(META): USE_PAD_SLOT=pad_slot_id is not None, BLOCK_N=256, ) - if unsqueeze: - out = out.squeeze(-1) query = query.squeeze(-1) key = key.squeeze(-1) value = value.squeeze(-1) diff --git a/atom/models/qwen3_5.py b/atom/models/qwen3_5.py index 008e8877f7..7c3cccd35a 100644 --- a/atom/models/qwen3_5.py +++ b/atom/models/qwen3_5.py @@ -40,10 +40,6 @@ maybe_prefix, extract_layer_index, ) -from atom.model_ops.split_chunk import ( - fused_split_chunk_zeros, - fused_split_chunk_zeros_qwen3_5_qkvzba, -) if is_vllm(): from vllm.model_executor.layers.mamba.mamba_utils import ( @@ -251,35 +247,35 @@ def forward( 1. Input projection 2. Core attention (custom op) """ + num_tokens = hidden_states.size(0) # ============================================================ # Part 1: Input Projection # ============================================================ + v_heads_tp = self.num_v_heads // self.tp_size + qkv_size = self.conv_dim // self.tp_size + z_size = v_heads_tp * self.head_v_dim + b_size = v_heads_tp + a_size = v_heads_tp + if hasattr(self, "in_proj_qkvzba"): qkvzba = self.in_proj_qkvzba(hidden_states) - k_heads_after_tp = self.num_k_heads // self.tp_size - v_heads_after_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_zeros_qwen3_5_qkvzba( - qkvzba, - k_heads_after_tp, - v_heads_after_tp, - self.head_k_dim, - self.head_v_dim, + # Qwen3.5 layout is already contiguous [q|k|v|z|b|a] + mixed_qkv, z_flat, b, a = torch.split( + qkvzba, [qkv_size, z_size, b_size, a_size], dim=-1 ) else: if x_fp8 is not None: mixed_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) else: mixed_qkvz = self.in_proj_qkvz(hidden_states) - ba = self.in_proj_ba(hidden_states) + projected_ba = self.in_proj_ba(hidden_states) + # Qwen3.5 layout is already contiguous [q|k|v|z] and [b|a] + mixed_qkv, z_flat = torch.split(mixed_qkvz, [qkv_size, z_size], dim=-1) + b, a = torch.split(projected_ba, [b_size, a_size], dim=-1) - qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size - z_size = self.value_dim // self.tp_size - num_v_heads_tp = self.num_v_heads // self.tp_size - - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_zeros( - mixed_qkvz, ba, qkv_size, z_size, self.head_v_dim, num_v_heads_tp - ) + z = z_flat.view(num_tokens, v_heads_tp, self.head_v_dim) + core_attn_out = torch.empty_like(z) # ============================================================ # Part 2: Core Attention (Custom Op) diff --git a/atom/models/qwen3_next.py b/atom/models/qwen3_next.py index 0361602119..f8abcb8672 100644 --- a/atom/models/qwen3_next.py +++ b/atom/models/qwen3_next.py @@ -20,19 +20,17 @@ from atom.model_ops.layernorm import GemmaRMSNorm as Qwen3NextRMSNorm from atom.model_ops.layernorm import RMSNormGated from atom.model_ops.linear import ( + BAParallelLinear, ColumnParallelLinear, MergedColumnParallelLinear, MergedReplicatedLinear, QKVGParallelLinear, QKVParallelLinear, QKVZBAParallelLinear, + QKVZParallelLinear, RowParallelLinear, -) +) # noqa: F401 from atom.model_ops.moe import FusedMoE -from atom.model_ops.split_chunk import ( - fused_split_chunk_qwen_next_qkvz_ba, - fused_split_chunk_qwen_next_qkvzba, -) from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled from atom.model_ops.utils import atom_parameter from atom.models.utils import ( @@ -594,103 +592,27 @@ def create_qkvzba_proj(self, quant_config, prefix): prefix=f"{self.prefix}.in_proj_qkvzba", ) else: - self.in_proj_qkvz = self.create_qkvz_proj( - hidden_size=self.hidden_size, - key_dim=self.key_dim, - value_dim=self.value_dim, + # Non-fused path (FP8): deinterleave weights + scales at load + # time so the forward can use zero-copy torch.split. + self.in_proj_qkvz = QKVZParallelLinear( + input_size=self.hidden_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + num_k_heads=self.num_k_heads, + num_v_heads=self.num_v_heads, + bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.in_proj_qkvz", ) - - self.in_proj_ba = self.create_ba_proj( - hidden_size=self.hidden_size, + self.in_proj_ba = BAParallelLinear( + input_size=self.hidden_size, + num_k_heads=self.num_k_heads, num_v_heads=self.num_v_heads, + bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.in_proj_ba", ) - def create_qkvz_proj( - self, - hidden_size: int, - key_dim: int, - value_dim: int, - quant_config: QuantizationConfig | None, - prefix: str, - ) -> MergedColumnParallelLinear: - return MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[sum((key_dim, key_dim, value_dim, value_dim))], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.in_proj_qkvz", - ) - - def create_ba_proj( - self, - hidden_size: int, - num_v_heads: int, - quant_config: QuantizationConfig | None, - prefix: str, - ) -> MergedColumnParallelLinear: - return MergedColumnParallelLinear( - input_size=hidden_size, - output_sizes=[num_v_heads, num_v_heads], - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.in_proj_ba", - ) - - def fix_query_key_value_ordering( - self, - mixed_qkvz: torch.Tensor, - mixed_ba: torch.Tensor, - ): - """ - Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. - """ - new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + ( - self.num_k_heads // self.tp_size, - ( - self.head_k_dim - + self.head_k_dim - + (self.head_v_dim + self.head_v_dim) - * self.num_v_heads - // self.num_k_heads - ), - ) - new_tensor_shape_ba = mixed_qkvz.size()[:-1] + ( - self.num_k_heads // self.tp_size, - 2 * self.num_v_heads // self.num_k_heads, - ) - - mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz) - mixed_ba = mixed_ba.view(*new_tensor_shape_ba) - - split_arg_list_qkvz = [ - self.head_k_dim, - self.head_k_dim, - (self.num_v_heads // self.num_k_heads * self.head_v_dim), - (self.num_v_heads // self.num_k_heads * self.head_v_dim), - ] - split_arg_list_ba = [ - self.num_v_heads // self.num_k_heads, - self.num_v_heads // self.num_k_heads, - ] - - # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)] - # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn], - # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng] - query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2) - b, a = torch.split(mixed_ba, split_arg_list_ba, dim=2) - - # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn] - value = value.reshape(value.size(0), -1, self.head_v_dim) - z = z.reshape(z.size(0), -1, self.head_v_dim) - b = b.reshape(b.size(0), self.num_v_heads // self.tp_size) - a = a.reshape(a.size(0), self.num_v_heads // self.tp_size) - - return query, key, value, z, b, a - def rearrange_mixed_qkv(self, mixed_qkv): if mixed_qkv is None: return None, None, None @@ -723,43 +645,42 @@ def forward( 3. Output projection """ + num_tokens = hidden_states.shape[0] # ============================================================ # Part 1: Input Projection # ============================================================ + v_heads_tp = self.num_v_heads // self.tp_size + qkv_size = self.conv_dim // self.tp_size + z_size = v_heads_tp * self.head_v_dim + b_size = v_heads_tp + a_size = v_heads_tp + if hasattr(self, "in_proj_qkvzba"): - projected_states_qkvzba = self.in_proj_qkvzba(hidden_states) - k_heads_after_tp = self.num_k_heads // self.tp_size - v_heads_after_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_qwen_next_qkvzba( - projected_states_qkvzba, - k_heads_after_tp, - v_heads_after_tp, - self.head_k_dim, - self.head_v_dim, + projected = self.in_proj_qkvzba(hidden_states) + # Output layout is [q|k|v|z|b|a] contiguous (deinterleaved at load) + mixed_qkv, z_flat, b, a = torch.split( + projected, [qkv_size, z_size, b_size, a_size], dim=-1 ) + z = z_flat.view(num_tokens, v_heads_tp, self.head_v_dim) else: if x_fp8 is not None: - projected_states_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) + projected_qkvz = self.in_proj_qkvz(x_fp8, x_scale=x_scale) else: - projected_states_qkvz = self.in_proj_qkvz(hidden_states) - projected_states_ba = self.in_proj_ba(hidden_states) # always BF16 - # Use Triton kernel to process qkvz and ba - num_k_heads_tp = self.num_k_heads // self.tp_size - num_v_heads_tp = self.num_v_heads // self.tp_size - mixed_qkv, z, b, a, core_attn_out = fused_split_chunk_qwen_next_qkvz_ba( - projected_states_qkvz, - projected_states_ba, - num_k_heads_tp, - num_v_heads_tp, - self.head_k_dim, - self.head_v_dim, + projected_qkvz = self.in_proj_qkvz(hidden_states) + projected_ba = self.in_proj_ba(hidden_states) # always BF16 + # Weights deinterleaved at load → output is [q|k|v|z] and [b|a] + mixed_qkv = projected_qkvz[:, :qkv_size] + z = projected_qkvz[:, qkv_size:].view( + num_tokens, v_heads_tp, self.head_v_dim ) + b = projected_ba[:, :b_size] + a = projected_ba[:, b_size:] + + core_attn_out = torch.empty(z.shape, dtype=z.dtype, device=z.device) # ============================================================ # Part 2: Core Attention (Custom Op) # ============================================================ - # Note: we should not use torch.empty here like other attention backends, - # see discussions in https://github.com/vllm-project/vllm/pull/28182 core_attn_out = self.attn(mixed_qkv, b, a, core_attn_out) diff --git a/atom/plugin/attention_mha.py b/atom/plugin/attention_mha.py index ba2073d794..16c88949da 100644 --- a/atom/plugin/attention_mha.py +++ b/atom/plugin/attention_mha.py @@ -245,9 +245,10 @@ def paged_attention_triton_plugin_mode( assert num_q_heads_total % num_kv_heads == 0 context_partition_size = 256 - use_ps = self.adopt_persistent_kernel( - head_size, num_kv_heads, num_q_heads_total - ) + # use_ps = self.adopt_persistent_kernel( + # head_size, num_kv_heads, num_q_heads_total + # ) + use_ps = True if use_ps: max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) else: diff --git a/atom/plugin/vllm/attention_backend/attention_gdn.py b/atom/plugin/vllm/attention_backend/attention_gdn.py index bff5370215..b6158a0867 100644 --- a/atom/plugin/vllm/attention_backend/attention_gdn.py +++ b/atom/plugin/vllm/attention_backend/attention_gdn.py @@ -75,17 +75,19 @@ def fused_gdn_gating_kernel( dt_bias, seq_len, NUM_HEADS: tl.constexpr, + stride_a_batch, + stride_b_batch, beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + out_off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) - blk_b = tl.load(b + off, mask=mask) + blk_a = tl.load(a + i_b * stride_a_batch + head_off, mask=mask) + blk_b = tl.load(b + i_b * stride_b_batch + head_off, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) # If the model is loaded in fp16, without the .float() here, A might be -inf x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) @@ -93,11 +95,13 @@ def fused_gdn_gating_kernel( beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x ) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + out_off, blk_g.to(g.dtype.element_ty), mask=mask) # compute beta_output = sigmoid(b) blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) tl.store( - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + beta_output + out_off, + blk_beta_output.to(beta_output.dtype.element_ty), + mask=mask, ) @@ -129,6 +133,8 @@ def fused_gdn_gating( dt_bias, seq_len, num_heads, + a.stride(0), + b.stride(0), beta, threshold, 8, @@ -206,6 +212,7 @@ def forward( if attn_metadata is None: # V1 profile run + core_attn_out.zero_() return core_attn_out assert isinstance(attn_metadata, dict) @@ -437,4 +444,8 @@ def forward( elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) + # Zero padding tail for CUDA graph replay safety + if num_actual_tokens < core_attn_out.shape[0]: + core_attn_out[num_actual_tokens:].zero_() + return core_attn_out diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 41294f80bd..fb487bc106 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -39,7 +39,7 @@ # QK-norm-rope-cache-quant fusion for Qwen3-MoE; disabled by default. # Enable for Qwen3-MoE to get better performance. "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION": lambda: ( - os.getenv("ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION", "0") == "1" + os.getenv("ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION", "1") == "1" ), "ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION", "1") == "1" From 6a4db8e9c97304211f91065748ea4f721a86bb47 Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Sun, 10 May 2026 12:30:10 +0800 Subject: [PATCH 17/76] perf(deepseek_v4): unified aiter rmsnorm_quant + q_norm/wkv_gate fuses (#733) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related changes that together cut input-quant overhead on V4-Pro attention: 1. atom/model_ops/layernorm.py — single dispatch covering all dynamic quant types - Replace the mxfp4-only `mxfp4_rms_quant_fuse` (triton wrapper) with a unified `_aiter_rms_quant` helper that dispatches on quant_type to aiter.{rmsnorm_quant, add_rmsnorm_quant} (HIP). - Now supports per_1x32 (MXFP4 fp4x2 + UE8M0 scale + preshuffle), per_1x128 (FP8 block + transpose-aware scale), and per_Token (FP8 dynamic per-token). - `_aiter_transpose_scale` resolved once at __init__ instead of per-forward env lookup. 2. atom/models/deepseek_v4.py — fuse q_norm with wq_b input quant - q_norm now constructed with fused_quant=True, quant_config=qc; emits (qr_fp8, qr_scale) in one launch. - Outer wq_b (ColumnParallel) and Indexer.wq_b (Replicated) both consume the pre-quantized pair via `x_scale=` — saves two redundant per_1x128 input quants per layer. - Drop `otype=torch.float32` on wkv_gate (Compressor): kernel's internal fp32 accumulator already handles the upcast; bf16 intermediate halves the (kv, score) buffer bandwidth. 3. .github/workflows/atom-benchmark.yaml — fix env_vars expression injection - Multi-line `env_vars` JSON values (e.g. "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1") were inlined directly into a `for ev in ${{ ... }}` loop, breaking host bash with `syntax error near unexpected token`. - Switch both Start CI container blocks to the env-block + docker --env-file pattern already used by atom-test.yaml. - Drop the redundant inline `${{ matrix.cell.env_vars }}` prefix on the regression launch step (vars are now in the container env via --env-file) — also fixes a silent bug that dropped the first var. Verified on DeepSeek-V4-Pro tp=8 --level 0: GSM8K nshot=5 (AITER_BF16_FP8_MOE_BOUND=0 + ATOM_MOE_GU_ITLV=1): baseline : 0.9522 / 0.9530 + rmsnorm_quant refactor: 0.9507 / 0.9515 + q_norm fuse : 0.9538 / 0.9538 + drop wkv_gate fp32 : 0.9583 / 0.9591 ← this PR All within 1 sigma stderr (~0.0058) of baseline; no regression. 1024/1024 c=64 microbench (q_norm fuse vs prior commit): Total token throughput : 2735.80 → 2813.00 tok/s (+2.8%) Median TPOT : 41.28 → 38.83 ms (-5.9%) --- .github/workflows/atom-benchmark.yaml | 34 +++---- atom/model_ops/layernorm.py | 133 ++++++++++++++++++-------- atom/models/deepseek_v4.py | 45 ++++++--- 3 files changed, 143 insertions(+), 69 deletions(-) diff --git a/.github/workflows/atom-benchmark.yaml b/.github/workflows/atom-benchmark.yaml index 0466903ac5..fbd6111ab1 100644 --- a/.github/workflows/atom-benchmark.yaml +++ b/.github/workflows/atom-benchmark.yaml @@ -190,13 +190,11 @@ jobs: MODEL_MOUNT="" [ -d "/models" ] && MODEL_MOUNT="-v /models:/models" - # Build env var flags from model config - ENV_FLAGS="" - if [ -n "${{ matrix.model.env_vars }}" ]; then - for ev in ${{ matrix.model.env_vars }}; do - ENV_FLAGS="$ENV_FLAGS -e $ev" - done - fi + # Write env_vars via env block (avoids expression injection — newlines + # in `matrix.model.env_vars` would otherwise break the host bash). + # Empty file when MODEL_ENV_VARS unset is fine — docker --env-file + # accepts it. + printenv MODEL_ENV_VARS 2>/dev/null | grep -v '^$' > /tmp/atom_benchmark_env.txt || true docker run -dt --device=/dev/kfd $DEVICE_FLAG \ -v "${GITHUB_WORKSPACE:-$PWD}":/workspace $MODEL_MOUNT \ @@ -210,11 +208,12 @@ jobs: -e CONC=${{ env.CONC }} -e RANDOM_RANGE_RATIO=${{ env.RANDOM_RANGE_RATIO }} \ -e ENABLE_TORCH_PROFILER=${{ inputs.enable_profiler && '1' || '0' }} \ -e ENABLE_RTL_PROFILER=${{ inputs.enable_rtl && '1' || '0' }} \ - $ENV_FLAGS \ + --env-file /tmp/atom_benchmark_env.txt \ --name atom-benchmark \ ${{ inputs.image || 'rocm/atom-dev:latest' }} env: GITHUB_WORKSPACE: ${{ github.workspace }} + MODEL_ENV_VARS: ${{ matrix.model.env_vars }} - name: Collect GPU info (inside container) id: gpu-info @@ -559,13 +558,9 @@ jobs: MODEL_MOUNT="" [ -d "/models" ] && MODEL_MOUNT="-v /models:/models" - # Build env var flags from model config - ENV_FLAGS="" - if [ -n "${{ matrix.cell.env_vars }}" ]; then - for ev in ${{ matrix.cell.env_vars }}; do - ENV_FLAGS="$ENV_FLAGS -e $ev" - done - fi + # Write env_vars via env block (avoids expression injection — newlines + # in `matrix.cell.env_vars` would otherwise break the host bash). + printenv CELL_ENV_VARS 2>/dev/null | grep -v '^$' > /tmp/atom_regression_env.txt || true docker run -dt --device=/dev/kfd $DEVICE_FLAG \ -v "${GITHUB_WORKSPACE:-$PWD}":/workspace $MODEL_MOUNT \ @@ -575,11 +570,12 @@ jobs: --security-opt seccomp=unconfined \ --ulimit memlock=-1 --ulimit stack=67108864 --pull always \ -e ATOM_DISABLE_MMAP=true \ - $ENV_FLAGS \ + --env-file /tmp/atom_regression_env.txt \ --name atom-regression \ ${{ inputs.image || 'rocm/atom-dev:latest' }} env: GITHUB_WORKSPACE: ${{ github.workspace }} + CELL_ENV_VARS: ${{ matrix.cell.env_vars }} - name: Download model run: | @@ -595,8 +591,12 @@ jobs: if [ -d "/models" ]; then MODEL_DIR="/models/${{ matrix.cell.model_path }}" else MODEL_DIR="${{ matrix.cell.model_path }}"; fi + # cell.env_vars are already injected into the container env via + # --env-file in "Start CI container" — `docker exec` inherits them, + # so no need to re-prefix them on the command line (the previous + # inline form silently dropped the first var when env_vars contained + # a literal newline). docker exec atom-regression bash -lc "set -euo pipefail - ${{ matrix.cell.env_vars }} \ ENABLE_TORCH_PROFILER=1 \ .github/scripts/atom_test.sh launch $MODEL_DIR ${{ matrix.cell.server_args }}" diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 91c69083fd..14898b2001 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -119,51 +119,93 @@ def fused_add_rmsnorm_pad_( return fused_add_rmsnorm_pad(x, weight, epsilon, res, x_pad_to_multiple) -def mxfp4_rms_quant_fuse_fake( +# --------------------------------------------------------------------------- +# Aiter dynamic RMSNorm + quant — single dispatch covering per_1x32 (MXFP4), +# per_1x128 (FP8 block), and per_Token (FP8). All three reach +# aiter.{rmsnorm_quant, add_rmsnorm_quant} (HIP) which both normalizes and +# emits a freshly-computed scale, so callers must have x_scale=None +# (static-scale FP8 stays on its own branch). Per-quant params (out_dtype, +# scale shape, group_size, shuffle_scale) are derived from quant_type_value +# inside the fake helper so torch.compile's schema infer sees a single, +# stable signature. +# +# `mutates_args=[]` keeps torch.compile from functionalizing the out-buffers +# — same pattern as the legacy mxfp4 fuse helper. +# --------------------------------------------------------------------------- + +_QV_PER_1X32 = QuantType.per_1x32.value +_QV_PER_1X128 = QuantType.per_1x128.value +_QV_PER_TOKEN = QuantType.per_Token.value +_AITER_RMS_QUANT_TYPE_VALUES = frozenset({_QV_PER_1X32, _QV_PER_1X128, _QV_PER_TOKEN}) + + +def _aiter_rms_quant_fake( x: torch.Tensor, weight: torch.Tensor, eps: float, - shuffle: bool = False, + quant_type_value: int, + transpose_scale: bool, res1: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from aiter.utility.dtypes import fp8 + M, N = x.shape - out = torch.empty((M, N // 2), dtype=torch.float4_e2m1fn_x2, device=x.device) - MXFP4_QUANT_BLOCK_SIZE = 32 - SCALE_N_valid = (N + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE - use_scale_shuffle_padding = shuffle - if use_scale_shuffle_padding: - SCALE_M = ((M + 255) // 256) * 256 - SCALE_N = ((SCALE_N_valid + 7) // 8) * 8 - else: - SCALE_M = M - SCALE_N = SCALE_N_valid - scale = torch.empty( - (SCALE_M, SCALE_N), - dtype=torch.float8_e8m0fnu, - device=x.device, - ) - out_res1 = None - if res1 is not None: - out_res1 = torch.empty_like(res1) + if quant_type_value == _QV_PER_1X32: + # MXFP4: out=(M, N/2) fp4x2; scale=(⌈M/256⌉*256, ⌈⌈N/32⌉/8⌉*8) UE8M0 + # bytes (kernel writes one uint8 per group; passing fp8_e8m0fnu + # directly yields the matching byte layout — no fp32 view). + out = torch.empty((M, N // 2), dtype=torch.float4_e2m1fn_x2, device=x.device) + scale_m = ((M + 255) // 256) * 256 + scale_n = ((((N + 31) // 32) + 7) // 8) * 8 + scale = torch.empty( + (scale_m, scale_n), dtype=torch.float8_e8m0fnu, device=x.device + ) + elif quant_type_value == _QV_PER_1X128: + # FP8 per-block: scale=(M, ⌈N/128⌉) fp32. Preshuffle GEMM expects + # column-major; allocate (num_groups, M) row-major then view as + # (M, num_groups). Matches GemmaRMSNorm._forward_fused_fp8. + out = torch.empty((M, N), dtype=fp8, device=x.device) + num_groups = N // 128 + if transpose_scale: + scale = torch.empty( + (num_groups, M), dtype=torch.float32, device=x.device + ).view(M, num_groups) + else: + scale = torch.empty((M, num_groups), dtype=torch.float32, device=x.device) + else: # _QV_PER_TOKEN + out = torch.empty((M, N), dtype=fp8, device=x.device) + scale = torch.empty((M, 1), dtype=torch.float32, device=x.device) + out_res1 = torch.empty_like(res1) if res1 is not None else None return (out, scale, out_res1) -# It's important to use mutates_args=[] to avoid functionized_v2 op generation -@torch_compile_guard(gen_fake=mxfp4_rms_quant_fuse_fake, mutates_args=[]) -def mxfp4_rms_quant_fuse( +@torch_compile_guard(gen_fake=_aiter_rms_quant_fake, mutates_args=[]) +def _aiter_rms_quant( x: torch.Tensor, weight: torch.Tensor, eps: float, - shuffle: bool = False, + quant_type_value: int, + transpose_scale: bool, res1: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - from aiter.ops.triton.fused_mxfp4_quant import fused_rms_mxfp4_quant + from aiter import add_rmsnorm_quant, rmsnorm_quant - (x_quant, x_scale), _, _, residual_out = fused_rms_mxfp4_quant( - x, weight, eps, shuffle=shuffle, res1=res1 + out, scale, out_res1 = _aiter_rms_quant_fake( + x, weight, eps, quant_type_value, transpose_scale, res1 ) - - return x_quant, x_scale, residual_out + if quant_type_value == _QV_PER_1X32: + group_size, shuffle = 32, True + elif quant_type_value == _QV_PER_1X128: + group_size, shuffle = 128, transpose_scale + else: # _QV_PER_TOKEN + group_size, shuffle = 0, False + if res1 is None: + rmsnorm_quant(out, x, scale, weight, eps, group_size, shuffle) + else: + add_rmsnorm_quant( + out, x, res1, out_res1, scale, weight, eps, group_size, shuffle + ) + return out, scale, out_res1 class RMSNorm(nn.Module): @@ -195,6 +237,14 @@ def __init__( params_dtype = layer_quant_config.quant_dtype self.quant_type = quant_type self.params_dtype = params_dtype + # transpose_scale (column-major scale) only applies to per_1x128 with + # the preshuffle GEMM consumer; resolve the env once at init time so + # forward sees a hot static bool instead of an env lookup per call. + self._aiter_transpose_scale = ( + fused_quant + and quant_type.value == _QV_PER_1X128 + and envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + ) @mark_trace(prefix="rmsnorm", torch_compile=True) def forward( @@ -263,19 +313,24 @@ def forward( res1=residual, ) return (x, x_scale), residual - elif self.use_fused_quant and ( - x_scale is None and self.quant_type.value == QuantType.per_1x32.value + elif ( + self.use_fused_quant + and x_scale is None + and self.quant_type.value in _AITER_RMS_QUANT_TYPE_VALUES ): + # Dynamic-scale fused RMSNorm + quant via aiter HIP kernels. + # Static FP8 (x_scale provided) stays on the branch above. + x, x_scale, residual_out = _aiter_rms_quant( + x, + self.weight, + self.eps, + self.quant_type.value, + self._aiter_transpose_scale, + residual, + ) if residual is None: - x, x_scale, _ = mxfp4_rms_quant_fuse( - x, self.weight, self.eps, shuffle=True - ) return x, x_scale - else: - x, x_scale, residual = mxfp4_rms_quant_fuse( - x, self.weight, self.eps, shuffle=True, res1=residual - ) - return (x, x_scale), residual + return (x, x_scale), residual_out else: if residual is None: # return rmsnorm2d_fwd(x, self.weight, self.eps).view(ori_shape) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index adf951901e..ee0b7a06ae 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -769,12 +769,14 @@ def forward( d = self.head_dim rd = self.rope_head_dim - # Single fused BF16 GEMM via tgemm; otype=fp32 keeps softmax-pool - # accumulator in fp32. torch.split returns zero-copy strided views; - # downstream kernels (fused_compress_attn, update_compressor_states) - # accept strided kv/score (only inner stride must be 1). + # Single fused BF16 GEMM via tgemm. (Probing whether dropping the + # otype=fp32 upcast — relying on fused_compress_attn's internal fp32 + # accumulator instead — is accuracy-neutral.) torch.split returns + # zero-copy strided views; downstream kernels (fused_compress_attn, + # update_compressor_states) accept strided kv/score (only inner + # stride must be 1). coff_d = (1 + overlap) * d - combined = self.wkv_gate(x, otype=torch.float32) + combined = self.wkv_gate(x) kv, score = torch.split(combined, [coff_d, coff_d], dim=-1) # ====== Unified fused kernel path (CSA + Indexer) ====== @@ -920,8 +922,11 @@ def __init__(self, args: DeepseekV4Args, compress_ratio: int = 4, prefix: str = def forward_batched( self, x_full: torch.Tensor, # [total_tokens, dim] - qr_full: torch.Tensor, # [total_tokens, q_lora_rank] + qr_full: torch.Tensor, # [total_tokens, q_lora_rank] — fp8 when qr_full_scale given positions: torch.Tensor, # [total_tokens] + qr_full_scale: Optional[ + torch.Tensor + ] = None, # per_1x128 scale paired with qr_full ) -> torch.Tensor: """Q proj + RoPE + FP8-quant + weights compute (have module state), then dispatch to `torch.ops.aiter.indexer_score_topk`, which calls @@ -946,7 +951,9 @@ def forward_batched( # Q proj + RoPE + rotate (batched). rotary_emb internally reshapes # to (1, num_tokens, -1, rotary_dim) so the input doesn't need an # explicit batch dim. rotate_activation is last-dim-only. - q = self.wq_b(qr_full).view(total_tokens, self.n_heads, self.head_dim) + q = self.wq_b(qr_full, x_scale=qr_full_scale).view( + total_tokens, self.n_heads, self.head_dim + ) # self.rotary_emb(positions, q[..., -rd:]) # q = rotate_activation(q) cos, sin = self.rotary_emb._caches(q.device) @@ -1235,7 +1242,16 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): quant_config=qc, prefix=f"{p}.wqkv_a", ) - self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + # Fuse q_norm + per_1x128 FP8 quant: kernel emits (qr_fp8, qr_scale) + # in one launch, both wq_b consumers (outer ColumnParallel + Indexer + # ReplicatedLinear) skip their own input quant. + self.q_norm = RMSNorm( + self.q_lora_rank, + self.eps, + fused_quant=True, + quant_config=qc, + prefix=f"{p}.q_norm", + ) # q_norm2: per-head Q normalization. The checkpoint has no # `q_norm2.weight` entry, so the module is constructed with # `weight=None` (no learnable Parameter) — `DualRMSNorm` reads the @@ -1451,15 +1467,17 @@ def forward_impl( # Single fused FP8 GEMM for [wq_a; wkv]; torch.split returns zero-copy views. qkv_a = self.wqkv_a(x) q_lora, kv_pre = torch.split(qkv_a, [self.q_lora_rank, self.head_dim], dim=-1) - qr = self.q_norm(q_lora) # [num_tokens, q_lora_rank] shared with Indexer - if _V4_FORCE_UE8M0_QUANT: - qr = qr.clone() - act_quant_inplace(qr, 128, "ue8m0") + # q_norm is fused-quant: returns (qr_fp8, qr_scale) — shared by + # outer wq_b and Indexer.wq_b, both per_1x128 FP8 GEMMs. + assert ( + not _V4_FORCE_UE8M0_QUANT + ), "_V4_FORCE_UE8M0_QUANT incompatible with fused q_norm quant (qr is already FP8)" + qr, qr_scale = self.q_norm(q_lora) # Flat q_flat is [num_tokens, n_local_heads * head_dim]; DualRMSNorm # views internally to per-head shape and returns flat. Single Triton # launch fuses per-head Q RMSNorm (weightless) + KV RMSNorm (both # head_dim=128), replacing the prior `_rmsnorm_nw + kv_norm` pair. - q_flat = self.wq_b(qr) + q_flat = self.wq_b(qr, x_scale=qr_scale) q_flat, kv = self.qk_norm(q_flat, kv_pre) q = q_flat.view(num_tokens, self.n_local_heads, self.head_dim) # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes @@ -1513,6 +1531,7 @@ def forward_impl( indexer_topk_batched = self.indexer.forward_batched( x_full=x, qr_full=qr, + qr_full_scale=qr_scale, positions=positions, ) # raw seq-local [T, index_topk] int32, -1 sentinels in tail cols # Translate seq-local topk → physical paged offsets and write into From 0b0e0097144dadc394f0789a1c08524694226a4a Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Sun, 10 May 2026 21:53:38 +0800 Subject: [PATCH 18/76] refactor(deepseek_v4): ParallelHead inherits ParallelLMHead; raise default cudagraph capture to 512 (#737) ParallelHead refactor: - Delegate vocab-axis sharding, weight_loader, last-token slicing, bf16 a16w16 GEMM, and TP all-gather to ParallelLMHead. - Remove ~30 lines of duplicated logic. - Drop fp32 weight (CDNA3/CDNA4 bf16 MFMA accumulates in fp32 natively; no precision change). Saves ~1.85 GB cluster VRAM. - Hoist Expert.forward's local silu_and_mul import to top-level. CUDAGraph capture default: - Add 512 to default --cudagraph-capture-sizes list. Required to capture the full graph at max_num_seqs=512; previously the cap was 256 and size-512 graphs were filtered out by model_runner. Validated: GSM8K 1319q 0.9591 +/- 0.0058 (statistically equivalent to baseline ~0.952). c=128 1k/1k throughput unchanged at ~4855 tok/s. c=512 1k/1k throughput +50.7% (7427 -> 11195 tok/s) when paired with --max-num-seqs 512. --- atom/model_engine/arg_utils.py | 2 +- atom/models/deepseek_v4.py | 70 ++++++++++++++-------------------- 2 files changed, 30 insertions(+), 42 deletions(-) diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 1358c04052..5477a058ce 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -114,7 +114,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--cudagraph-capture-sizes", type=str, - default="[1,2,4,8,16,32,48,64,128,256]", + default="[1,2,4,8,16,32,48,64,128,256,512]", help="Sizes to capture cudagraph. Example: [1,2,4,8,16]", ) parser.add_argument( diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index ee0b7a06ae..7461253fff 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -33,10 +33,15 @@ cp_gather_indexer_k_quant_cache, dtypes, get_hip_quant, + silu_and_mul as aiter_silu_and_mul, ) from aiter.jit.utils.torch_guard import torch_compile_guard -from aiter.dist.communication_op import tensor_model_parallel_all_reduce -from aiter.dist.parallel_state import get_tensor_model_parallel_world_size +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, +) +from aiter.dist.parallel_state import ( + get_tensor_model_parallel_world_size, +) from aiter.ops.topk import top_k_per_row_decode, top_k_per_row_prefill from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits @@ -48,7 +53,7 @@ get_current_atom_config, ) from atom.model_loader.loader import WeightsMapper -from atom.model_ops.embed_head import VocabParallelEmbedding +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import DualRMSNorm, RMSNorm, rmsnorm2d_fwd_ from atom.model_ops.triton_rmsnorm_nw import rmsnorm_nw from atom.model_ops.linear import ( @@ -1739,7 +1744,6 @@ def forward( x: torch.Tensor, # [num_tokens, dim] weights: Optional[torch.Tensor] = None, # [num_tokens, 1] optional gate ) -> torch.Tensor: # [num_tokens, dim] - from aiter import silu_and_mul as aiter_silu_and_mul dtype = x.dtype # Single fused GEMM. Layout is [gate | up] concat on last dim — matches @@ -2171,58 +2175,43 @@ def forward( return x -class ParallelHead(nn.Module): - """LM head with mHC reduction. +class ParallelHead(ParallelLMHead): + """V4 LM head with mHC reduction; vocab-parallel sharded across TP ranks. - Port of inference/model.py:704-736. Unlike `Block.hc_pre` (which uses - Sinkhorn projection on the combination matrix), `hc_head` uses simple - `Sigmoid(mix*scale + base) + eps` weights to reduce the [num_tokens, hc, dim] - residual to [num_tokens, dim] before applying the LM head linear projection. + Port of inference/model.py:704-736. Inherits from `ParallelLMHead` so the + vocab-axis sharding, `weight_loader`, last-token slicing, bf16 a16w16 GEMM + (`tgemm.mm`), and TP all-gather come for free. V4 only adds: - `get_logits` projects only the last token of each sequence (decode mode); - for prefill the caller should slice the desired positions before passing - through. + - `forward(...)` taking the mHC residual + hc_head params + final norm + - `get_logits(...)` so `compute_logits` can call it directly on the + hidden-state output of `model.forward` (CUDAGraph contract) + - `hc_head(...)` Sigmoid-gated mHC reduction (vs `Block.hc_pre`'s Sinkhorn) + + Note on weight dtype: the V4 reference (model.py:713-714) keeps the LM head + in fp32 because the disk weight is bf16; on AMD CDNA3/CDNA4 the bf16 MFMA + instruction accumulates in fp32 natively, so a bf16 GEMM with the + bf16-on-disk weight has the same effective precision as the reference's + fp32 path while halving VRAM and using the faster a16w16 kernel. """ def __init__( self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6 ): - super().__init__() - self.vocab_size = vocab_size + super().__init__(vocab_size, dim, bias=False) self.dim = dim self.norm_eps = norm_eps self.hc_eps = hc_eps - # PR1 single-rank: full vocab on this rank. - self.weight = atom_parameter( - torch.empty(self.vocab_size, self.dim, dtype=torch.float32) - ) def get_logits( self, x: torch.Tensor # [num_tokens, dim] ) -> torch.Tensor: # [bs, vocab] - """Project the last-token-per-seq slice of `x` to vocab logits. - - Picks the last token of each sequence using `cu_seqlens_q` from the - forward_context. Falls back to `x[-1:]` (single-seq) when no - forward_context is set (warmup / standalone). + """Project to vocab logits via the inherited `ParallelLMHead.forward`, + which handles last-token slicing (prefill) + tgemm.mm + all-gather. """ assert ( x.dim() == 2 and x.shape[-1] == self.dim ), f"get_logits expects [num_tokens, {self.dim}], got {tuple(x.shape)}" - ctx = get_forward_context() - cu_seqlens_q = ( - ctx.attn_metadata.cu_seqlens_q - if ctx is not None and ctx.attn_metadata is not None - else None - ) - if cu_seqlens_q is not None and cu_seqlens_q.numel() >= 2: - # Last-token positions per seq: cu_seqlens_q[1:bs+1] - 1. - # block_tables tells us the actual scheduled bs (cu_seqlens_q - # may have trailing zeros from the CpuGpuBuffer pool). - bs = ctx.attn_metadata.block_tables.size(0) - last_idx = cu_seqlens_q[1 : bs + 1] - 1 - return F.linear(x.index_select(0, last_idx.long()).float(), self.weight) - return F.linear(x[-1:].float(), self.weight) + return super().forward(x) def hc_head( self, @@ -2252,9 +2241,8 @@ def forward( norm: nn.Module, ) -> torch.Tensor: # [bs, vocab] x = self.hc_head(x, hc_fn, hc_scale, hc_base) # [num_tokens, dim] - logits = self.get_logits(norm(x)) # [bs, vocab] - # PR1 single-rank: skip all_gather - return logits + # get_logits handles the per-rank vocab shard + all-gather internally. + return self.get_logits(norm(x)) # [bs, vocab] class MTPBlock(Block): From 679422de1f0b53f800b4839023a3317d9efcf47e Mon Sep 17 00:00:00 2001 From: honglie Date: Sun, 10 May 2026 22:56:31 +0800 Subject: [PATCH 19/76] [Kimi] support Eagle3 speculative decoding for Kimi K2.5 (#631) * [Kimi] support Eagle3 speculative decoding for Kimi K2.5 Adds Eagle3 spec decode for Kimi K2.5 (MLA target + standard MHA draft): - Eagle3LlamaModel: 1-layer Llama draft (dual-norm input, wide QKV, independent embed/lm_head) matching the lightseekorg/kimi-k2.5-eagle3 checkpoint - Eagle3DraftBuilder: implements the post-#659 builder protocol (compute_block_bytes / allocate_kv_cache_tensors / build_kv_cache_tensor) for the draft's independent non-MLA KV cache, attached to the runner from EagleProposer.__init__ via runner.eagle3_draft_builder. ModelRunner delegates KV pool sizing, allocation, and per-module binding through this hook with no eagle3-specific code in the runner KV path - Aux hidden state pipeline: target forward returns (hidden, aux_hidden_states), captured through CUDAGraph via graph_aux_hidden and fed to the draft's combine_hidden_states (fc) as input - SpeculativeConfig: --method eagle3 + --draft-model CLI; eagle3 vs MTP branching at construction time; fail-fast if draft is MLA - Scheduler: spec_stats only updated when speculation actually ran (matches vLLM's gating) - propose: draft-perspective predicate `draft_uses_mha = hasattr(runner, "eagle3_draft_builder")` drives both the metadata-flow special-cases (slot_mapping re-slice, context_lens += 1, tuple-unpack of the draft return value); is_eagle3 string comparison is gone from the hot path Result on Kimi-K2.5-MXFP4 + kimi-k2.5-eagle3, 8x MI355X, gsm8k 5-shot: acceptance 67.85%, accuracy 93.78%. Co-Authored-By: Claude Opus 4.7 (1M context) * ci: add nightly Eagle3 spec-decode accuracy test for Kimi-K2.5 Reuses the base Kimi-K2.5-MXFP4 model + lightseekorg/kimi-k2.5-eagle3 draft, runs at TP=8 (Eagle3 draft KV needs full 8-rank sharding) under nightly schedule. Local case_verify_v9_gluon measured GSM8K 5-shot flexible-extract = 0.9257 (vLLM = 0.9280); threshold set to 0.91 with ~1.5pp noise headroom. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> --- .github/benchmark/models_accuracy.json | 12 + atom/config.py | 28 +++ atom/model_engine/arg_utils.py | 30 ++- atom/model_engine/model_runner.py | 135 +++++++++-- atom/model_engine/scheduler.py | 13 +- atom/model_ops/attention_mha.py | 15 +- atom/model_ops/linear.py | 18 +- atom/models/deepseek_v2.py | 28 ++- atom/models/eagle3_llama.py | 301 +++++++++++++++++++++++++ atom/models/kimi_k25.py | 6 + atom/spec_decode/eagle.py | 232 +++++++++++++++++-- 11 files changed, 760 insertions(+), 58 deletions(-) create mode 100644 atom/models/eagle3_llama.py diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index d175f91da2..4887f3d126 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -157,6 +157,18 @@ "accuracy_baseline_model": "moonshotai/Kimi-K2.5", "_baseline_note": "HF: amd/Kimi-K2.5-MXFP4 card shows Kimi-K2.5 baseline=0.9409" }, + { + "model_name": "Kimi-K2.5-MXFP4 Eagle3", + "model_path": "amd/Kimi-K2.5-MXFP4", + "extraArgs": "--kv_cache_dtype fp8 -tp 8 --trust-remote-code --method eagle3 --draft-model lightseekorg/kimi-k2.5-eagle3 --num-speculative-tokens 3", + "env_vars": "HSA_NO_SCRATCH_RECLAIM=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.91, + "accuracy_baseline": 0.9257, + "accuracy_baseline_model": "amd/Kimi-K2.5-MXFP4 + lightseekorg/kimi-k2.5-eagle3", + "_baseline_note": "Eagle3 spec decode on Kimi-K2.5-MXFP4. Local case_verify_v9_gluon GSM8K 5-shot flexible-extract=0.9257 (vLLM=0.9280, within ±0.71% se). Threshold 0.91 leaves ~1.5pp headroom for noise. -tp 8 (vs base entry's tp=4) because Eagle3 draft KV needs the full 8-rank sharding." + }, { "model_name": "GLM-5-FP8", "model_path": "zai-org/GLM-5-FP8", diff --git a/atom/config.py b/atom/config.py index d5de8e112b..8c1d5c819a 100644 --- a/atom/config.py +++ b/atom/config.py @@ -723,6 +723,8 @@ class SpeculativeConfig: model: Optional[str] = None num_speculative_tokens: Optional[int] = None draft_model_hf_config: Optional[PretrainedConfig] = None + use_aux_hidden_state: bool = False + eagle3_aux_layer_ids: list[int] = field(default_factory=list) # model_type → mtp_model_type mapping _MTP_TYPE_MAP: ClassVar[dict[str, str]] = { @@ -754,8 +756,34 @@ def __post_init__(self): self.draft_model_hf_config = self.draft_model_hf_config.text_config self.hf_config_override(self.draft_model_hf_config) + if self.method == "eagle3": + if getattr(self.draft_model_hf_config, "kv_lora_rank", None): + raise NotImplementedError( + "Eagle3 draft model with MLA attention is not supported" + ) + # Aux hidden state layers: prefer the draft checkpoint's + # eagle_config; if absent or the list is empty, ModelRunner + # falls back to model.get_eagle3_aux_hidden_state_layers(), + # which defaults to 3 layers — early / middle / late + # (see DeepseekV2ForCausalLM.get_eagle3_aux_hidden_state_layers, + # returns `(2, num_layers // 2, num_layers - 3)`, aligned with vLLM). + eagle_cfg = getattr(self.draft_model_hf_config, "eagle_config", None) + if eagle_cfg: + self.use_aux_hidden_state = eagle_cfg.get("use_aux_hidden_state", False) + if self.use_aux_hidden_state and not self.eagle3_aux_layer_ids: + self.eagle3_aux_layer_ids = eagle_cfg.get( + "eagle_aux_hidden_state_layer_ids", [] + ) + else: + self.use_aux_hidden_state = True + @staticmethod def hf_config_override(hf_config: PretrainedConfig) -> None: + # Eagle3 architecture mapping (architecture-level, not model_type) + arch = (getattr(hf_config, "architectures", None) or [""])[0] + if arch == "LlamaForCausalLMEagle3": + hf_config.architectures = ["Eagle3LlamaModel"] + # Step 1: resolve model_type → mtp model_type mtp_type = SpeculativeConfig._MTP_TYPE_MAP.get(hf_config.model_type) if mtp_type is not None: diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 5477a058ce..f93c225ca1 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -51,6 +51,7 @@ class EngineArgs: method: Optional[str] = None num_speculative_tokens: int = 1 kv_transfer_config: str = "{}" + draft_model: Optional[str] = None mark_trace: bool = False @staticmethod @@ -163,7 +164,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--method", type=str, default=None, - choices=["mtp"], + choices=["mtp", "eagle3"], help="Speculative method", ) parser.add_argument( @@ -172,6 +173,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: default=1, help="Number of speculative tokens to generate per iteration (draft model runs this many times autoregressively)", ) + parser.add_argument( + "--draft-model", + type=str, + default=None, + help="Path to external Eagle3 draft model. Required when --method eagle3.", + ) parser.add_argument( "--max-num-batched-tokens", type=int, @@ -243,14 +250,25 @@ def _get_engine_kwargs(self) -> dict: ), ) if self.method and self.num_speculative_tokens > 0: - kwargs["speculative_config"] = SpeculativeConfig( - method=kwargs.pop("method"), - model=self.model, - num_speculative_tokens=kwargs.pop("num_speculative_tokens"), - ) + method = kwargs.pop("method") + num_spec_tokens = kwargs.pop("num_speculative_tokens") + draft_model = kwargs.pop("draft_model") + if method == "eagle3": + kwargs["speculative_config"] = SpeculativeConfig( + method=method, + model=draft_model, + num_speculative_tokens=num_spec_tokens, + ) + else: + kwargs["speculative_config"] = SpeculativeConfig( + method=method, + model=self.model, + num_speculative_tokens=num_spec_tokens, + ) else: kwargs.pop("method") kwargs.pop("num_speculative_tokens") + kwargs.pop("draft_model") kwargs["speculative_config"] = None # --enable-tbo [prefill|all] → enable_tbo + enable_tbo_decode diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 5199912d0b..d6a3f10017 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -559,6 +559,13 @@ def __init__(self, rank: int, config: Config): self.num_spec_tokens = ( self.config.speculative_config.num_speculative_tokens if use_spec else 0 ) + self.eagle3_mode = ( + self.config.speculative_config is not None + and self.config.speculative_config.method == "eagle3" + ) + + self.use_aux_hidden_state_outputs = False + self._aux_hidden_states = None self.tokenID_processor = tokenIDProcessor( self, self.config.max_num_batched_tokens, @@ -621,6 +628,18 @@ def __init__(self, rank: int, config: Config): torch.set_default_device(None) logger.info("Loading drafter model...") self.drafter.load_model(self.model) + + if self.eagle3_mode and self.config.speculative_config.use_aux_hidden_state: + aux_ids = self.config.speculative_config.eagle3_aux_layer_ids + if not aux_ids and hasattr( + self.model, "get_eagle3_aux_hidden_state_layers" + ): + aux_ids = list(self.model.get_eagle3_aux_hidden_state_layers()) + if aux_ids: + self.model.set_aux_hidden_state_layers(tuple(aux_ids)) + self.use_aux_hidden_state_outputs = True + logger.info(f"Eagle3 aux hidden state layers: {aux_ids}") + torch.set_default_device(self.device) self.async_execute_stream = torch.cuda.Stream(self.device) self.allocate_forward_vars() @@ -1075,24 +1094,35 @@ def _get_num_kv_heads(self): return 1 def _get_total_num_layers(self): - """Return total layer count including draft (MTP) layers.""" + """Return total layer count including draft (MTP) layers. + + Drafts that own an independent KV cache via their own builder + (e.g. Eagle3 MHA draft on an MLA target) account for their layers + through that builder, so they are NOT added here. Only MTP-style + drafts that share the target's KV pool contribute. + """ total = self.config.hf_config.num_hidden_layers if self.config.speculative_config and hasattr(self, "drafter"): - draft_hf = self.config.speculative_config.draft_model_hf_config - total += getattr(draft_hf, "num_nextn_predict_layers", 1) + if not hasattr(self, "eagle3_draft_builder"): + draft_hf = self.config.speculative_config.draft_model_hf_config + total += getattr(draft_hf, "num_nextn_predict_layers", 1) return total def _compute_block_bytes(self): """Per-block bytes for the unified KV pool budget. - Delegates to the attention builder, which knows its own tensor - layout (MLA 576-dim packed, GDN-hybrid full-attn-only, MiMo-V2 - per-layer-type, standard MHA split-K/V). Mirror of - `attn_metadata_builder.allocate_kv_cache_tensors()` so the budget - math matches what's actually allocated. Per-request cache bytes - are accounted for separately via `compute_per_req_cache_bytes()`. + Sum across all attention builders attached to this runner: the + target builder always, plus an optional `eagle3_draft_builder` + when a heterogeneous spec-decode draft owns its own KV pool. Each + builder knows its own tensor layout (MLA 576-dim packed, GDN-hybrid + full-attn-only, MiMo-V2 per-layer-type, standard MHA split-K/V, + Eagle3 independent MHA). Per-request cache bytes are accounted + for separately via `compute_per_req_cache_bytes()`. """ - return self.attn_metadata_builder.compute_block_bytes() + block_bytes = self.attn_metadata_builder.compute_block_bytes() + if hasattr(self, "eagle3_draft_builder"): + block_bytes += self.eagle3_draft_builder.compute_block_bytes() + return block_bytes def _estimate_cudagraph_overhead(self): """Estimate GPU memory consumed by CUDA graph capture. @@ -1255,13 +1285,24 @@ def allocate_kv_cache(self, num_kvcache_blocks): num_draft_layers = 0 if self.config.speculative_config and hasattr(self, "drafter"): draft_hf_config = self.config.speculative_config.draft_model_hf_config - # For MTP, use num_nextn_predict_layers instead of num_hidden_layers - num_draft_layers = getattr(draft_hf_config, "num_nextn_predict_layers", 1) - total_num_layers += num_draft_layers - logger.info( - f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " - f"{num_draft_layers} draft (MTP) layers = {total_num_layers} total layers" - ) + if hasattr(self, "eagle3_draft_builder"): + # Heterogeneous draft (e.g. Eagle3 MHA on MLA target) owns + # its own KV pool via its builder; don't add to target's count. + num_draft_layers = draft_hf_config.num_hidden_layers + logger.info( + f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " + f"{num_draft_layers} Eagle3 draft layers (separate non-MLA cache)" + ) + else: + # For MTP, use num_nextn_predict_layers instead of num_hidden_layers + num_draft_layers = getattr( + draft_hf_config, "num_nextn_predict_layers", 1 + ) + total_num_layers += num_draft_layers + logger.info( + f"Allocating KV cache for {hf_config.num_hidden_layers} target layers + " + f"{num_draft_layers} draft (MTP) layers = {total_num_layers} total layers" + ) # Primary KV cache allocation (model-agnostic, delegated to the # attention builder). Each builder owns its tensor layout: MLA → @@ -1277,6 +1318,16 @@ def allocate_kv_cache(self, num_kvcache_blocks): for name, value in main_kv.items(): setattr(self, name, value) + # Heterogeneous draft (e.g. Eagle3 MHA alongside an MLA target) owns + # its own KV pool through a sibling builder; same protocol as above, + # tensors land under namespaced keys (eagle3_kv_cache, eagle3_kv_scale). + if hasattr(self, "eagle3_draft_builder"): + draft_kv = self.eagle3_draft_builder.allocate_kv_cache_tensors( + num_kv_heads, num_draft_layers + ) + for name, value in draft_kv.items(): + setattr(self, name, value) + # Per-request cache allocation (model-agnostic, delegated to the # attention metadata builder). For GDN this returns # `{"mamba_k_cache": ..., "mamba_v_cache": ...}`; for stateless @@ -1302,10 +1353,12 @@ def allocate_kv_cache(self, num_kvcache_blocks): kv_cache_tensors = [] layer_id = 0 # Promote to self so the attention builder's build_kv_cache_tensor() - # can access it without recomputing from drafter state. + # can access it without recomputing from drafter state. Heterogeneous + # drafts (Eagle3) own their own layer space via their builder, so + # leave mtp_start_layer_idx at hf_config.num_hidden_layers in that mode. self.mtp_start_layer_idx = ( self.drafter.model.model.mtp_start_layer_idx - if hasattr(self, "drafter") + if hasattr(self, "drafter") and not hasattr(self, "eagle3_draft_builder") else hf_config.num_hidden_layers ) for model_name, model in models_to_bind: @@ -1314,6 +1367,18 @@ def allocate_kv_cache(self, num_kvcache_blocks): ) for module in model.modules(): + # Drafts that own an independent KV pool (Eagle3) bind through + # their sibling builder first; for unrecognized modules it + # returns None and we fall through to the target builder. + if model_name == "draft" and hasattr(self, "eagle3_draft_builder"): + kv_cache_tensor = self.eagle3_draft_builder.build_kv_cache_tensor( + layer_id, module + ) + if kv_cache_tensor is not None: + kv_cache_tensors.append(kv_cache_tensor) + layer_id += 1 + continue + # Per-attention-type binding is owned by the attention # metadata builder; ModelRunner only walks modules and # collects the resulting KVCacheTensor entries. The builder @@ -1625,7 +1690,12 @@ def run_model( label += f" tok={batch.total_tokens_num} ctx={ctx_str}" label += "]" with record_function(label): - hidden_states = self.model(input_ids, positions) + model_output = self.model(input_ids, positions) + if self.use_aux_hidden_state_outputs: + hidden_states, self._aux_hidden_states = model_output + else: + hidden_states = model_output + self._aux_hidden_states = None logits = self.model.compute_logits(hidden_states) else: # decode[bs=128 tok=128 d=128] or decode[bs=128 tok=128 p=2 d=126 spec=3] @@ -1645,6 +1715,12 @@ def run_model( self.graphs[graph_key].replay() num_tokens = context.batch_size * max_q_len hidden_states = self.forward_vars["outputs"][:num_tokens] + if graph_key in self.graph_aux_hidden: + self._aux_hidden_states = [ + aux[:num_tokens] for aux in self.graph_aux_hidden[graph_key] + ] + else: + self._aux_hidden_states = None if self.logits_in_graph: logits = self.graph_logits[graph_key][:num_tokens] else: @@ -1833,6 +1909,7 @@ def propose_draft_token_ids( num_reject_tokens=num_reject_tokens, next_token_ids=next_token_ids, last_token_indices=last_token_indices, + aux_hidden_states=self._aux_hidden_states, ) return self.tokenID_processor.prepare_draft_ids(batch, draft_token) @@ -1882,6 +1959,7 @@ def capture_cudagraph(self): self.graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] = dict() self.graph_logits: dict[tuple[int, int], torch.Tensor] = dict() + self.graph_aux_hidden: dict[tuple[int, int], list[torch.Tensor]] = dict() self.graph_pool = None is_tbo = self.config.enable_tbo and isinstance(self.model, UBatchWrapper) # TBO graphs don't capture compute_logits, so disable logits_in_graph. @@ -1932,9 +2010,13 @@ def capture_cudagraph(self): ) # Warmup - outputs[:num_tokens] = self.model( + model_output = self.model( input_ids[:num_tokens], positions[:num_tokens] ) + if self.use_aux_hidden_state_outputs: + outputs[:num_tokens] = model_output[0] + else: + outputs[:num_tokens] = model_output if self.logits_in_graph: self.model.compute_logits(outputs[:num_tokens]) @@ -1953,13 +2035,20 @@ def capture_cudagraph(self): gc.stream, output_buffer=outputs[:num_tokens], ) + graph_aux = None else: # Standard single-stream capture graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, self.graph_pool, stream=gc.stream): - outputs[:num_tokens] = self.model( + model_output = self.model( input_ids[:num_tokens], positions[:num_tokens] ) + if self.use_aux_hidden_state_outputs: + outputs[:num_tokens] = model_output[0] + graph_aux = model_output[1] + else: + outputs[:num_tokens] = model_output + graph_aux = None if self.logits_in_graph: graph_logits = self.model.compute_logits( outputs[:num_tokens] @@ -1969,6 +2058,8 @@ def capture_cudagraph(self): self.graphs[(bs, max_q_len)] = graph if self.logits_in_graph and ubatch_slices is None: self.graph_logits[(bs, max_q_len)] = graph_logits + if graph_aux is not None: + self.graph_aux_hidden[(bs, max_q_len)] = graph_aux torch.cuda.synchronize() self.graph_bs.sort(reverse=False) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 9e9acd6ced..3d65369d4a 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -676,12 +676,21 @@ def postprocess( continue token_ids = prev_token_ids[idx] num_new_token = len(token_ids) - if self.spec_stats: - self.spec_stats.update(num_new_token) if is_deferred_out or self.use_spec: num_rejected = fwd_output.num_rejected[idx] num_bonus = fwd_output.num_bonus[idx] offset = 0 if (num_new_token + num_rejected) == 1 else self.mtp_k + # Align stats with vLLM: only count steps that actually ran + # speculation (drafts proposed and validated). Skip the + # prefill-only step where no draft tokens were scored against + # the target — vLLM gates this via + # `if scheduled_spec_token_ids and generated_token_ids`. + if ( + self.spec_stats + and num_new_token > 0 + and (num_new_token + num_rejected) > 1 + ): + self.spec_stats.update(num_new_token) seq.num_rejected = num_rejected seq.num_bonus_tokens = num_bonus for i, el in enumerate(token_ids): diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index cea626ccf5..5fbcaca0c5 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -131,7 +131,12 @@ def rope_cache(self, q, k, v, qkv, position, fwd_ctx: ForwardContext): k_scale = kv_cache_data[f"layer_{self.layer_num}"].k_scale v_scale = kv_cache_data[f"layer_{self.layer_num}"].v_scale - use_triton_attn = self.sliding_window != -1 or self.head_dim != 128 + # MTP MHA must go through triton/gluon; aiter ASM non-persistent path may have some unexpected behavior. + use_triton_attn = ( + self.sliding_window != -1 + or self.head_dim != 128 + or self.num_heads == self.num_kv_heads + ) self.use_triton_attn = use_triton_attn if ( @@ -535,9 +540,7 @@ def prefill_attention( # variable lenth attention use key value as input attn_metadata = fwd_ctx.attn_metadata sliding_window = ( - (self.sliding_window, 0, 0) - if self.sliding_window is not None - else (-1, -1, 0) + (self.sliding_window, 0, 0) if self.sliding_window > 0 else (-1, -1, 0) ) o = aiter.flash_attn_varlen_func( q, @@ -579,9 +582,7 @@ def prefill_attention_triton( num_seqs = attn_metadata.cu_seqlens_q.shape[0] - 1 descale_shape = (num_seqs, k.shape[1]) sliding_window = ( - (self.sliding_window - 1, 0) - if self.sliding_window is not None - else (-1, -1) + (self.sliding_window - 1, 0) if self.sliding_window > 0 else (-1, -1) ) # `block_tables` is always populated by TritonMHAMetadataBuilder. diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 9a2859c0a8..e016de7032 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -318,11 +318,19 @@ def weight_loader_process( loaded_weight: torch.Tensor, post_process_func: Callable = lambda a: a, ): - if ( - param.data.dtype != loaded_weight.dtype - and param.data.element_size() == loaded_weight.element_size() - ): - param.data = param.data.view(loaded_weight.dtype) + if param.data.dtype != loaded_weight.dtype: + if param.data.element_size() == loaded_weight.element_size(): + # Same byte-width: use view for raw-bit-compatible pairs + # (e.g. fp8 variants) but convert for semantically different + # formats (float16 ↔ bfloat16) where bit reinterpretation + # would corrupt values. + incompatible = {torch.float16, torch.bfloat16} + if {param.data.dtype, loaded_weight.dtype} == incompatible: + loaded_weight = loaded_weight.to(param.data.dtype) + else: + param.data = param.data.view(loaded_weight.dtype) + else: + loaded_weight = loaded_weight.to(param.data.dtype) loaded_weight = post_process_func(loaded_weight) if ( loaded_weight.shape != param.data.shape diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 3a92fe7410..252319a2da 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1840,6 +1840,8 @@ def __init__( ) else: self.norm = PPMissingLayer() + self.aux_hidden_state_layers: tuple[int, ...] = tuple() + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -1853,7 +1855,9 @@ def forward( positions: torch.Tensor, intermediate_tensors: Optional[IntermediateTensors], inputs_embeds: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, IntermediateTensors]: + ) -> Union[ + torch.Tensor, IntermediateTensors, Tuple[torch.Tensor, list[torch.Tensor]] + ]: if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -1865,7 +1869,13 @@ def forward( hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - for layer in self.layers[self.start_layer : self.end_layer]: + aux_hidden_states = [] + for idx in range(self.start_layer, self.end_layer): + layer = self.layers[idx] + if idx in self.aux_hidden_state_layers: + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) hidden_states, residual = layer(positions, hidden_states, residual) if not get_pp_group().is_last_rank: @@ -1874,6 +1884,9 @@ def forward( ) hidden_states, _ = self.norm(hidden_states, residual) + + if aux_hidden_states: + return hidden_states, aux_hidden_states return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: @@ -1974,6 +1987,17 @@ def make_empty_intermediate_tensors( } ) + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.model.aux_hidden_state_layers = layers + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + """Default Eagle3 aux hidden-state layer ids: early / middle / late + of the target model. Aligned with vLLM's default (see + vllm/model_executor/models/deepseek_v2.py). + """ + num_layers = len(self.model.layers) + return (2, num_layers // 2, num_layers - 3) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/atom/models/eagle3_llama.py b/atom/models/eagle3_llama.py new file mode 100644 index 0000000000..67837907b4 --- /dev/null +++ b/atom/models/eagle3_llama.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Eagle3 draft model (Llama full-attention) for speculative decoding. + +Implements the Eagle3 draft model matching the lightseekorg/kimi-k2.5-eagle3 +checkpoint layout: + + embed_tokens.weight — independent embedding + fc.weight — aux fusion projection (hidden*3 -> hidden) + midlayer.* — single decoder layer (dual-norm, wide QKV) + norm.weight — final RMSNorm + lm_head.weight — independent lm_head + +Weight keys map directly to model attribute paths; no key rewriting needed. +""" + +import torch +from aiter.dist.parallel_state import get_tensor_model_parallel_world_size +from aiter.rotary_embedding import get_rope +from atom.config import Config +from atom.model_ops.activation import SiluAndMul +from atom.model_ops.base_attention import Attention +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from atom.utils.decorators import support_torch_compile +from torch import nn + + +class Eagle3LlamaAttention(nn.Module): + """Llama full-attention with input_size = hidden_size * 2. + + The QKV projection accepts the concatenation of normalized embeddings + and fc output, hence input_size is doubled compared to standard Llama. + """ + + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + cache_config: str = "bf16", + prefix: str = "", + layer_num: int = 0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = hidden_size // self.total_num_heads + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + # QKV input_size = hidden_size * 2 (concat of embed + fc_output) + attn_input_size = hidden_size * 2 + self.qkv_proj = QKVParallelLinear( + hidden_size=attn_input_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=False, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + ) + + rope_theta = getattr(config, "rope_theta", 10000) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + is_neox_style=True, + ) + + sliding_window = -1 + if getattr(config, "use_sliding_window", False) and getattr( + config, "sliding_window", None + ): + sliding_window = config.sliding_window + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + kv_cache_dtype=cache_config, + layer_num=layer_num, + prefix=f"{prefix}.attn", + rotary_emb=self.rotary_emb, + per_layer_sliding_window=sliding_window, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv = self.qkv_proj(hidden_states) + q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) + attn_output = self.attn(q, k, v, positions) + output = self.o_proj(attn_output) + return output + + +class Eagle3LlamaDecoderLayer(nn.Module): + """Single decoder layer for Eagle3 with dual-norm input. + + Unlike standard LlamaDecoderLayer, this layer has: + - input_layernorm: normalizes the embedding input + - hidden_norm: normalizes the fc output (projected aux hidden states) + - Attention input is concat(normed_embed, normed_hidden) -> [N, hidden*2] + """ + + def __init__( + self, + config, + cache_config: str = "bf16", + prefix: str = "", + layer_num: int = 0, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = Eagle3LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr( + config, "num_key_value_heads", config.num_attention_heads + ), + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_num=layer_num, + ) + + self.mlp = Eagle3LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + prefix=f"{prefix}.mlp", + ) + + # Dual norms matching checkpoint keys: midlayer.input_layernorm, midlayer.hidden_norm + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + normed_embeds = self.input_layernorm(embeds) + normed_hidden = self.hidden_norm(hidden_states) + # Concat for attention input: [N, hidden*2] + attn_input = torch.cat([normed_embeds, normed_hidden], dim=-1) + attn_output = self.self_attn(positions, attn_input) + # Residual connection on hidden_states + hidden_states = hidden_states + attn_output + # MLP with pre-norm + residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class Eagle3LlamaMLP(nn.Module): + """Simple Llama MLP (gate+up fused, silu activation, down projection).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x = self.down_proj(x) + return x + + +@support_torch_compile +class Eagle3LlamaModel(nn.Module): + """Eagle3 draft model (Llama full-attention, single decoder layer). + + Matches the lightseekorg/kimi-k2.5-eagle3 checkpoint layout: + embed_tokens.weight [163840, 7168] independent embedding + fc.weight [7168, 21504] aux fusion (hidden*3 -> hidden) + midlayer.* single decoder layer + norm.weight final RMSNorm + lm_head.weight [163840, 7168] independent lm_head + """ + + packed_modules_mapping = { + "q_proj": ("qkv_proj", "q"), + "k_proj": ("qkv_proj", "k"), + "v_proj": ("qkv_proj", "v"), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, atom_config: Config, prefix: str = "", layer_offset: int = 0): + super().__init__() + config = atom_config.hf_config + cache_config = atom_config.kv_cache_dtype + self.config = config + + # Independent embedding (vocab matches target model) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + + # Aux fusion: concatenated aux hidden states [N, hidden*3] -> [N, hidden] + self.fc = ReplicatedLinear( + config.hidden_size * 3, config.hidden_size, bias=False + ) + + # Draft attention layer_num must start from the target model's layer + # count so kv_cache_data["layer_N"] maps to the correct cache entry. + self.midlayer = Eagle3LlamaDecoderLayer( + config=config, + cache_config=cache_config, + prefix="midlayer", + layer_num=layer_offset, + ) + + # Final norm + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Independent lm_head (not shared with target model) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project concatenated aux hidden states through fc. + + Args: + hidden_states: [N, hidden_size * 3] (3 aux layers concatenated) + + Returns: + [N, hidden_size] projected hidden states + """ + return self.fc(hidden_states) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + embeds = self.embed_tokens(input_ids) + hidden_states = self.midlayer(positions, embeds, hidden_states) + hidden_states_prenorm = hidden_states + hidden_states = self.norm(hidden_states) + return hidden_states, hidden_states_prenorm + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) diff --git a/atom/models/kimi_k25.py b/atom/models/kimi_k25.py index 6084f74ebf..06c3b1412f 100644 --- a/atom/models/kimi_k25.py +++ b/atom/models/kimi_k25.py @@ -92,5 +92,11 @@ def compute_logits( ) -> Optional[torch.Tensor]: return self.language_model.compute_logits(hidden_states) + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.language_model.set_aux_hidden_state_layers(layers) + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + return self.language_model.get_eagle3_aux_hidden_state_layers() + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.language_model.get_expert_mapping() diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index b06ab8fded..d729701b94 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -1,10 +1,13 @@ +import copy import logging +from typing import Optional import numpy as np import torch import torch.nn as nn +from aiter import dtypes from aiter.dist.parallel_state import get_pp_group -from atom.config import CompilationLevel, Config +from atom.config import CompilationLevel, Config, KVCacheTensor from atom.model_loader.loader import load_model from atom.utils import CpuGpuBuffer, resolve_obj_by_qualname from atom.utils.forward_context import SpecDecodeMetadata, get_forward_context @@ -18,9 +21,137 @@ "Qwen3NextMTPModel": "atom.models.qwen3_next_mtp.Qwen3NextMTP", "MiMoV2FlashMTPModel": "atom.models.mimo_v2_flash_mtp.MiMoV2FlashMTP", "Qwen3_5MTPModel": "atom.models.qwen3_5_mtp.Qwen3_5MTP", + "Eagle3LlamaModel": "atom.models.eagle3_llama.Eagle3LlamaModel", } +class Eagle3DraftBuilder: + """KV cache subsystem for an Eagle3 MHA draft alongside a non-MHA target. + + Implements the same subset of `AttentionMetadataBuilder` hooks that + ModelRunner consults during KV pool sizing and per-module binding — + `compute_block_bytes`, `allocate_kv_cache_tensors`, and + `build_kv_cache_tensor` — so the draft's independent non-MLA cache + fits the post-#659 builder protocol without leaking into the target's + builder. The draft does NOT drive prepare_decode/prepare_prefill; + it piggybacks on the target builder's metadata flow during propose. + """ + + def __init__(self, model_runner, draft_hf): + self.model_runner = model_runner + self.draft_hf = draft_hf + self.block_size = model_runner.block_size + self.num_kv_heads = draft_hf.num_key_value_heads // model_runner.world_size + self.num_layers = draft_hf.num_hidden_layers + self.head_dim = draft_hf.head_dim + self._next_layer_id = 0 # consumed by build_kv_cache_tensor + self.num_blocks = 0 # set in allocate_kv_cache_tensors + + def compute_block_bytes(self) -> int: + """Per-block bytes for the draft's independent non-MLA KV cache.""" + kv_dtype_size = dtypes.d_dtypes[ + self.model_runner.config.kv_cache_dtype + ].itemsize + bb = ( + 2 + * self.num_layers + * self.block_size + * self.num_kv_heads + * self.head_dim + * kv_dtype_size + ) + if self.model_runner.config.kv_cache_dtype == "fp8": + # fp8 KV cache needs an extra per-(layer, block, kv_head) scale + # tensor (one fp32 per element) to dequantize fp8 → bf16 at + # attention time. Reserve that space alongside the cache. + bb += ( + 2 + * self.num_layers + * self.block_size + * self.num_kv_heads + * dtypes.fp32.itemsize + ) + return bb + + def allocate_kv_cache_tensors(self, num_kv_heads, num_draft_layers) -> dict: + """Allocate the draft's [2, L, blocks, block_size, kv_heads, head_dim] + cache and matching fp32 scale; ModelRunner setattr's both onto itself + under namespaced keys so they don't collide with the target builder's + `kv_cache` / `kv_scale`. + """ + runner = self.model_runner + config = runner.config + # Draft's block budget scales with the target pool: same total token + # capacity, just paged at the draft's own block size. + self.num_blocks = ( + config.num_kvcache_blocks * runner.block_size // self.block_size + ) + cache = torch.zeros( + 2, + self.num_layers, + self.num_blocks, + self.block_size, + self.num_kv_heads, + self.head_dim, + dtype=dtypes.d_dtypes[config.kv_cache_dtype], + device="cuda", + ) + scale = torch.zeros( + 2, + self.num_layers, + self.num_blocks, + self.num_kv_heads, + self.block_size, + dtype=dtypes.fp32, + device="cuda", + ) + logger.info(f"Allocated Eagle3 draft KV cache: {cache.shape}") + return {"eagle3_kv_cache": cache, "eagle3_kv_scale": scale} + + def build_kv_cache_tensor(self, layer_id: int, module): + """Bind one Eagle3 draft attention module to its slice of the + independent draft KV cache. Returns None for non-MHA modules so + ModelRunner falls through to the target builder. + """ + if not ( + hasattr(module, "base_attention") + and hasattr(module, "use_mla") + and not module.use_mla + ): + return None + runner = self.model_runner + idx = self._next_layer_id + self._next_layer_id += 1 + cache = runner.eagle3_kv_cache + x = 16 // cache.element_size() + k_cache = cache[0, idx].view( + self.num_blocks, + self.num_kv_heads, + self.head_dim // x, + self.block_size, + x, + ) + v_cache = cache[1, idx].view( + self.num_blocks, + self.num_kv_heads, + self.head_dim, + self.block_size, + ) + module.max_model_len = runner.config.max_model_len + if runner.config.kv_cache_dtype == "fp8": + module.k_scale = runner.eagle3_kv_scale[0, idx] + module.v_scale = runner.eagle3_kv_scale[1, idx] + module.k_cache = k_cache + module.v_cache = v_cache + return KVCacheTensor( + layer_num=layer_id, + k_cache=k_cache, + v_cache=v_cache, + k_scale=getattr(module, "k_scale", None), + v_scale=getattr(module, "v_scale", None), + ) + + class EagleProposer: def __init__( @@ -49,7 +180,30 @@ def __init__( self.device = device draft_model_hf_config = self.speculative_config.draft_model_hf_config model_class = resolve_obj_by_qualname(support_eagle_model_arch_dict[draft_model_hf_config.architectures[0]]) # type: ignore - self.model = model_class(self.config) + + if self.speculative_config.method == "eagle3": + # Eagle3 draft model has its own architecture (Llama, not MLA), + # so it must be constructed with the draft model's hf_config. + # Also disable torch.compile for the draft model to avoid + # Dynamo tracing issues with the separate KV cache binding. + draft_atom_config = copy.deepcopy(atom_config) + draft_atom_config.hf_config = draft_model_hf_config + draft_atom_config.compilation_config.level = CompilationLevel.NO_COMPILATION + # Draft attention layer_num must continue from the target model's + # layer count so it maps to the correct kv_cache_data entry. + self.model = model_class( + draft_atom_config, + layer_offset=atom_config.hf_config.num_hidden_layers, + ) + # Attach the draft's KV-cache builder to the runner. ModelRunner + # consults `runner.eagle3_draft_builder` from `_compute_block_bytes` + # / `allocate_kv_cache` to size + allocate + bind the draft's + # independent non-MLA cache through the standard builder protocol. + runner.eagle3_draft_builder = Eagle3DraftBuilder( + runner, draft_model_hf_config + ) + else: + self.model = model_class(self.config) i32_kwargs = {"dtype": torch.int32, "device": self.device} i64_kwargs = {"dtype": torch.int64, "device": self.device} @@ -78,6 +232,23 @@ def _share_if_not_loaded( setattr(owner, attr, source) def load_model(self, target_model: nn.Module) -> None: + if self.speculative_config.method == "eagle3": + # Eagle3: load from a separate draft model checkpoint with + # independent embed_tokens and lm_head (no sharing). + load_model( + self.model, + self.speculative_config.model, + self.speculative_config.draft_model_hf_config, + self.config.load_dummy, + False, + ) + logger.info( + "Eagle3 draft model loaded from %s (independent embed/lm_head)", + self.speculative_config.model, + ) + return + + # MTP: load from the target model checkpoint and share embeddings/lm_head. loaded = load_model( self.model, self.config.model, @@ -101,11 +272,6 @@ def load_model(self, target_model: nn.Module) -> None: ) del self.model.model.embed_tokens self.model.model.embed_tokens = target_base.model.embed_tokens - else: - logger.info( - "The EAGLE head's vocab embedding will be loaded separately" - " from the target model." - ) # Share lm_head from target if not loaded from checkpoint. # Case 1: per-layer shared_head.head (DeepSeek MTP) @@ -148,6 +314,7 @@ def propose( num_reject_tokens: torch.Tensor, next_token_ids: torch.Tensor, last_token_indices: torch.Tensor, + aux_hidden_states: Optional[list[torch.Tensor]] = None, ) -> torch.Tensor: forward_context = get_forward_context() @@ -161,21 +328,48 @@ def propose( # input_ids[last_token_indices] = next_token_ids input_ids.scatter_(0, last_token_indices, next_token_ids) positions = target_positions + 1 - hidden_states = target_hidden_states + + # Eagle3: project concatenated aux hidden states through fc + if aux_hidden_states is not None: + concat_aux = torch.cat(aux_hidden_states, dim=-1) + hidden_states = self.model.combine_hidden_states(concat_aux) + else: + hidden_states = target_hidden_states draft_token_ids = torch.empty( bs, self.mtp_k, dtype=next_token_ids.dtype, device=next_token_ids.device ) - # return draft_token_ids.fill_(1) # for debug var = self.runner.forward_vars - use_mla = self.runner.use_mla + target_uses_mla = self.runner.use_mla + # Eaale3 only support mha currently + draft_uses_mha = hasattr(self.runner, "eagle3_draft_builder") + + # Eagle3 MLA: re-slice slot_mapping to len(input_ids). + # Target's MLA prepare_decode sized it + # to bs*max_q_len; after rejection len(input_ids) may be smaller, + # and the MHA cache-write kernel asserts slot_mapping <= q. + # Other fields (block_tables, context_lens, slot_mapping values + # themselves) are already in a draft-compatible format because + # MLA's prepare_decode uses the same runner.block_size as the draft. + if draft_uses_mha: + attn_metadata.slot_mapping = var["slot_mapping"].gpu[: len(input_ids)] + for i in range(self.mtp_k): with record_function(f"draft[{i}/{self.mtp_k} bs={bs}]"): - ret_hidden_states = self.model( + model_output = self.model( input_ids=input_ids, positions=positions, hidden_states=hidden_states, ) + # Eagle3 draft (the only draft_uses_mha case under narrow + # semantics) returns (post_norm, pre_norm); MTP drafts return + # a single hidden tensor. + if draft_uses_mha: + ret_hidden_states, ret_hidden_prenorm = model_output + else: + ret_hidden_states = model_output + ret_hidden_prenorm = None + sample_hidden_states = ( torch.index_select(ret_hidden_states, 0, last_token_indices) if i == 0 @@ -204,7 +398,7 @@ def propose( attn_metadata.kv_indices = kv_indices attn_metadata.cu_seqlens_q = cu_seqlens_q attn_metadata.slot_mapping = slot_mapping - if use_mla: + if target_uses_mla: kv_last_page_lens = var["kv_last_page_lens"].gpu[:bs] attn_metadata.kv_last_page_lens = kv_last_page_lens # block_tables, context_lens, and sparse_kv_indptr are @@ -216,7 +410,7 @@ def propose( "sparse_kv_indptr" ].gpu[: bs + 1] cu_seqlens_q[: bs + 1] = self.arrange_bs[: bs + 1] - if use_mla: + if target_uses_mla: # MLA: block_size=1, kv_indptr tracks tokens kv_indptr[1 : bs + 1] -= torch.cumsum( num_reject_tokens, dim=0 @@ -243,9 +437,19 @@ def propose( for k, v in workinfos.items(): attn_metadata.__dict__[k] = v slot_mapping[:] = kv_indices[kv_indptr[1 : bs + 1] - 1] + input_ids = new_draft_ids positions += 1 - hidden_states = sample_hidden_states + if ret_hidden_prenorm is not None: + hidden_states = ( + torch.index_select( + ret_hidden_prenorm, 0, last_token_indices + ) + if i == 0 + else ret_hidden_prenorm + ) + else: + hidden_states = sample_hidden_states # self.runner.debug(f"final {draft_token_ids=}") # [batch_size, mtp_k] From a82e5cd839c79b2070a34c00eccde2abfeedc003 Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Mon, 11 May 2026 10:47:32 +0800 Subject: [PATCH 20/76] support three stream in Deepseek V4 (#736) * support three stream in Deepseek V4 * modify hca dual stream if no indexer * clean code logic --- atom/models/deepseek_v4.py | 145 ++++++++++++++++++++++++------------- 1 file changed, 94 insertions(+), 51 deletions(-) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 7461253fff..4d9a59e162 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1202,7 +1202,14 @@ class DeepseekV4Attention(nn.Module): - attn_sink: per-head learnable logit added only to softmax denominator """ - def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): + def __init__( + self, + layer_id: int, + args: DeepseekV4Args, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + compress_stream: Optional[torch.cuda.Stream] = None, + ): super().__init__() self.layer_id = layer_id self.dim = args.dim @@ -1372,6 +1379,12 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): self.indexer.rotary_emb = self.rotary_emb self.indexer.compressor.rotary_emb = self.rotary_emb + self.alt_stream = alt_stream + self.compress_stream = compress_stream + self._use_async_compress = ( + self.alt_stream is not None and self.compressor is not None + ) + self.layer_name = prefix atom_config = get_current_atom_config() atom_config.compilation_config.static_forward_context[self.layer_name] = self @@ -1426,6 +1439,31 @@ def process_weights_after_loading(self) -> None: # batched-FP8 kernel. See attention_mla.py:211 for reference. self.wo_a.quant_type = QuantType.No + def _launch_compressors_async(self, x, plan, state_slot_mapping, block_tables): + """Fire Compressor(s) on side streams, return immediately. + + Main Compressor → alt_stream (CSA + HCA). + Indexer Compressor → compress_stream (CSA only). + Waits resolve instantly: side streams ~25us, main Q/KV chain ~87us.""" + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + self.compressor( + x, + plan=plan, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables, + ) + if self.indexer is not None and self.compress_stream is not None: + self.compress_stream.wait_stream(current_stream) + with torch.cuda.stream(self.compress_stream): + self.indexer.compressor( + x, + plan=plan, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables, + ) + def forward( self, x: torch.Tensor, @@ -1469,11 +1507,24 @@ def forward_impl( if _V4_FORCE_UE8M0_QUANT: x = x.clone() act_quant_inplace(x, 128, "ue8m0") - # Single fused FP8 GEMM for [wq_a; wkv]; torch.split returns zero-copy views. + + attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) + compress_plans = attn_md.compress_plans + swa_write_indices = attn_md.swa_write_indices + v4_batch_id_per_token = attn_md.batch_id_per_token + block_tables_gpu = attn_md.block_tables + state_slot_mapping = attn_md.state_slot_mapping + plan_for_layer = compress_plans[ratio] if ratio else None + + # ===== Triple-stream: Q/KV path + both Compressors in parallel ===== + if self._use_async_compress: + self._launch_compressors_async( + x, plan_for_layer, state_slot_mapping, block_tables_gpu + ) + + # ----- Q/KV projections (main stream) ----- qkv_a = self.wqkv_a(x) q_lora, kv_pre = torch.split(qkv_a, [self.q_lora_rank, self.head_dim], dim=-1) - # q_norm is fused-quant: returns (qr_fp8, qr_scale) — shared by - # outer wq_b and Indexer.wq_b, both per_1x128 FP8 GEMMs. assert ( not _V4_FORCE_UE8M0_QUANT ), "_V4_FORCE_UE8M0_QUANT incompatible with fused q_norm quant (qr is already FP8)" @@ -1490,61 +1541,41 @@ def forward_impl( self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) if _V4_USE_REF_QUANT: act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) - # ===== Per-fwd metadata (built once in prepare_prefill/decode). ===== - # All per-fwd state read once. Production prepare_decode/prefill - # always populates these; warmup goes through the same path - # (`_populate_state_slot_mapping` falls back to slot 0). - # Cast to V4 typed metadata so V4-specific attribute access (v4_*, - # compress_plans, swa_write_indices, ...) is well-typed for pyright. - attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) - compress_plans = attn_md.compress_plans - swa_write_indices = attn_md.swa_write_indices - v4_batch_id_per_token = attn_md.batch_id_per_token - block_tables_gpu = attn_md.block_tables - state_slot_mapping = attn_md.state_slot_mapping - # ===== Batched compressor + Indexer (ONCE per layer) ===== - # State cache reset on fresh prefill is redundant: SWA's start_pos==0 - # path uses raw seq_kv, and fused_compress's state-cache reads are - # already masked by `s < 0` for fresh prefill (compress_plan.py:88 + - # fused_compress.py:124-127). - # Compressor's return value is unused — paged_decode/paged_prefill read - # the scattered compress entries from `unified_kv` via per-fwd indices. - # Called purely for its side-effect (scatter into self.kv_cache). - plan_for_layer = compress_plans[ratio] if ratio else None - if self.compressor is not None: # i.e. CSA or HCA layer, compress_ratio > 0 - self.compressor( - x, - plan=plan_for_layer, - state_slot_mapping=state_slot_mapping, - block_tables=block_tables_gpu, - ) - if self.indexer is not None: # i.e. CSA layer, compress_ratio == 4 - # Indexer's inner Compressor populates the FP8 indexer kv_cache - # via the unified `fused_compress_attn` quant path (auto-detected - # by cache dtype). `block_tables_gpu` is the same as the Main - # Compressor's; the kernel resolves `physical_block * k_per_block - # + slot_in_block` internally — no separate slot_mapping needed. - # The outer `forward_batched` consumes `v4_indexer_meta` (built - # once per fwd) so per-layer H2D / CPU index math stays zero. - self.indexer.compressor( - x, - plan=plan_for_layer, - state_slot_mapping=state_slot_mapping, - block_tables=block_tables_gpu, - ) + # ===== Compressor + Indexer ===== + if not self._use_async_compress: + if self.compressor is not None: + self.compressor( + x, + plan=plan_for_layer, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables_gpu, + ) + if self.indexer is not None: + self.indexer.compressor( + x, + plan=plan_for_layer, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables_gpu, + ) + if self.indexer is not None: + if self._use_async_compress: + torch.cuda.current_stream().wait_stream(self.alt_stream) + torch.cuda.current_stream().wait_stream(self.compress_stream) indexer_topk_batched = self.indexer.forward_batched( x_full=x, qr_full=qr, qr_full_scale=qr_scale, positions=positions, - ) # raw seq-local [T, index_topk] int32, -1 sentinels in tail cols + ) # Translate seq-local topk → physical paged offsets and write into # the CSA section of either: # - decode buffer `kv_indices_csa` (is_pure_decode) # - prefill buffer `kv_indices_prefix_csa` (otherwise) # `_fill_csa_paged_compress` dispatches internally on is_pure_decode. self._fill_csa_paged_compress(attn_md, indexer_topk_batched, num_tokens) + elif self._use_async_compress: + torch.cuda.current_stream().wait_stream(self.alt_stream) # ===== Sparse attention dispatch ===== # Two paths over the unified KV pool. The order of `swa_write` vs @@ -2024,11 +2055,18 @@ def __init__( args: DeepseekV4Args, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + compress_stream: Optional[torch.cuda.Stream] = None, ): super().__init__() self.layer_id = layer_id self.norm_eps = args.norm_eps - self.attn = DeepseekV4Attention(layer_id, args, prefix=f"{prefix}.attn") + self.attn = DeepseekV4Attention( + layer_id, + args, + prefix=f"{prefix}.attn", + alt_stream=alt_stream, + compress_stream=compress_stream, + ) self.ffn = MoE(layer_id, args, prefix=f"{prefix}.ffn", alt_stream=alt_stream) self.attn_norm = RMSNorm(args.dim, self.norm_eps) self.ffn_norm = RMSNorm(args.dim, self.norm_eps) @@ -2347,12 +2385,16 @@ def __init__( # equals nn.Embedding's [vocab_size, dim] so dummy state_dicts load # directly. At TP>1 each rank holds vocab_size/tp rows. self.embed = VocabParallelEmbedding(args.vocab_size, args.dim) - # alt_stream for dual-stream MoE (shared_experts // routed_experts). - # Allocated once and shared across all blocks; MoE.__init__ decides - # per-layer whether to actually use it (env-gated, mismatched-dtype-aware). + # alt_stream: dual-stream MoE (shared_experts // routed_experts) AND + # Main Compressor overlap. compress_stream: Indexer Compressor overlap. + # Both allocated once, shared across all blocks. Attention runs before + # MoE in each block, so attn and MoE never contend for alt_stream. self.alt_stream: Optional[torch.cuda.Stream] = ( torch.cuda.Stream() if torch.cuda.is_available() else None ) + self.compress_stream: Optional[torch.cuda.Stream] = ( + torch.cuda.Stream() if torch.cuda.is_available() else None + ) self.layers = nn.ModuleList( [ Block( @@ -2360,6 +2402,7 @@ def __init__( args, prefix=f"layers.{layer_id}", alt_stream=self.alt_stream, + compress_stream=self.compress_stream, ) for layer_id in range(args.n_layers) ] From 50461c86455a942f4f0609d749b8f7fcee757e18 Mon Sep 17 00:00:00 2001 From: wuhuikx Date: Mon, 11 May 2026 16:17:16 +0800 Subject: [PATCH 21/76] Update the vLLM-ATOM benchmark scope (#739) * Update the vLLM-ATOM benchmark scope * Update benchmark model tags and weekday scheduling. Normalize model tagging to MET/OOB/AW prefixes and labels, and replace nightly A/B/C date rotation with weekday-based grouping (Mon/Wed MET, Tue/Thu AW, Fri ALL, weekends skipped) for clearer benchmark cadence control. * Align AW Minimax serve args and warmup behavior. Add --kv-cache-dtype fp8 for MiniMax-M2.5 AW TP2/4/8 so startup flags match expected cache settings, and always pass --num-warmups=$((2 * CONC)) for vllm bench serve runs to keep warmup load consistent. * Normalize AW metadata and gpt-oss memory env settings. Add nightly_group=B for all AW model entries, align all gpt-oss variants to OOT_GPU_MEMORY_UTILIZATION=0.5, and restrict scheduled benchmark cron to weekdays to avoid weekend empty runs. --- .github/benchmark/oot_benchmark_models.json | 318 +++++++++----- .github/scripts/atom_oot_test.sh | 9 +- .github/workflows/atom-vllm-benchmark.yaml | 432 +++++++++++++------- 3 files changed, 504 insertions(+), 255 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 5bc5585e84..da04f0a446 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -1,9 +1,9 @@ [ { - "display": "DeepSeek-R1 FP8 TP8", + "display": "DeepSeek-R1 FP8 TP8 (MET)", "source_path": "deepseek-ai/DeepSeek-R1-0528", "path": "deepseek-ai/DeepSeek-R1-0528", - "prefix": "deepseek-r1-fp8", + "prefix": "deepseek-r1-fp8-met", "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", @@ -11,10 +11,10 @@ "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "DeepSeek-R1 MXFP4 TP8", + "display": "DeepSeek-R1 MXFP4 TP8 (MET)", "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "prefix": "deepseek-r1-mxfp4", + "prefix": "deepseek-r1-mxfp4-met", "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", @@ -22,56 +22,37 @@ "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "DeepSeek-V3.2 FP8 TP8", - "dashboard_model": "DeepSeek-V3.2-FP8", - "source_path": "deepseek-ai/DeepSeek-V3.2", - "path": "deepseek-ai/DeepSeek-V3.2", - "prefix": "deepseek-v3-2-fp8", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "MiniMax-M2.5 TP2", - "dashboard_model": "MiniMax-M2.5", - "source_path": "MiniMaxAI/MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "MiniMax-M2.5-tp2", - "extra_args": "--trust-remote-code --kv-cache-dtype fp8 --tensor-parallel-size 2", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - }, - { - "display": "gpt-oss-120b TP1", + "display": "gpt-oss-120b TP1 (MET)", "source_path": "openai/gpt-oss-120b", "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b", + "prefix": "gpt-oss-120b-met", "nightly_group": "C", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" }, { - "display": "GLM-5.1-FP8 TP8", + "display": "GLM-5.1-FP8 TP8 (OOB)", "dashboard_model": "GLM-5.1-FP8", "source_path": "zai-org/GLM-5.1-FP8", "path": "zai-org/GLM-5.1-FP8", - "prefix": "glm-5-1-fp8-tp8", + "prefix": "glm-5-1-fp8-tp8-oob", "nightly_group": "C", + "excluded_input_output_pairs": [ + "1024x8192" + ], "extra_args": "--trust-remote-code --tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}' --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "Kimi-K2-Thinking-MXFP4 TP4", + "display": "Kimi-K2-Thinking-MXFP4 TP4 (MET)", "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp4", "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp4", + "prefix": "kimi-k2-thinking-mxfp4-tp4-met", "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", @@ -79,11 +60,11 @@ "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "Kimi-K2-Thinking-MXFP4 TP8", + "display": "Kimi-K2-Thinking-MXFP4 TP8 (MET)", "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp8", "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp8", + "prefix": "kimi-k2-thinking-mxfp4-tp8-met", "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", @@ -91,23 +72,11 @@ "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "Kimi-K2.5-MXFP4 TP2", - "dashboard_model": "Kimi-K2.5-MXFP4-tp2", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp2", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 2", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP4", + "display": "Kimi-K2.5-MXFP4 TP4 (MET)", "dashboard_model": "Kimi-K2.5-MXFP4-tp4", "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp4", + "prefix": "kimi-k25-mxfp4-tp4-met", "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", @@ -115,111 +84,252 @@ "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { - "display": "Kimi-K2.5-MXFP4 TP8", - "dashboard_model": "Kimi-K2.5-MXFP4", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP8", + "display": "Qwen3.5-397B-A17B-FP8 TP8 (OOB)", "dashboard_model": "Qwen3.5-397B-A17B-FP8", "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8", + "prefix": "qwen3-5-397b-a17b-fp8-oob", "nightly_group": "A", + "excluded_input_output_pairs": [ + "1024x8192" + ], "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" }, { - "display": "Qwen3.5-397B-A17B-MXFP4 TP4", - "dashboard_model": "Qwen3.5-397B-A17B-MXFP4-tp4", - "source_path": "amd/Qwen3.5-397B-A17B-MXFP4", - "path": "amd/Qwen3.5-397B-A17B-MXFP4", - "prefix": "qwen3-5-397b-a17b-mxfp4", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP4", + "display": "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8-tp4", + "prefix": "qwen3-5-397b-a17b-fp8-tp4-oob", "nightly_group": "A", + "excluded_input_output_pairs": [ + "1024x8192" + ], "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" }, { - "display": "Qwen3.5-397B-A17B TP8", + "display": "Qwen3.5-397B-A17B TP8 (OOB)", "dashboard_model": "Qwen3.5-397B-A17B", "source_path": "Qwen/Qwen3.5-397B-A17B", "path": "Qwen/Qwen3.5-397B-A17B", - "prefix": "qwen3-5-397b-a17b", + "prefix": "qwen3-5-397b-a17b-oob", "nightly_group": "A", + "excluded_input_output_pairs": [ + "1024x8192" + ], "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" }, { - "display": "Qwen3.5-397B-A17B TP4", - "dashboard_model": "Qwen3.5-397B-A17B-tp4", - "source_path": "Qwen/Qwen3.5-397B-A17B", - "path": "Qwen/Qwen3.5-397B-A17B", - "prefix": "qwen3-5-397b-a17b-tp4", + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", + "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1-met", "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" }, { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4-met", "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", "bench_args": "", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" }, { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp2", + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp1", "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp2", - "nightly_group": "A", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp2", + "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", + "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", + "bench_args": "--random-range-ratio 1", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" }, { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp4", "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4", - "nightly_group": "A", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", + "nightly_group": "B", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "display": "DeepSeek-V3.2 FP8 TP4 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", + "source_path": "deepseek-ai/DeepSeek-V3.2", + "path": "deepseek-ai/DeepSeek-V3.2", + "prefix": "deepseek-v3-2-fp8-aw-tp4", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "DeepSeek-V3.2 FP8 TP8 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", + "source_path": "deepseek-ai/DeepSeek-V3.2", + "path": "deepseek-ai/DeepSeek-V3.2", + "prefix": "deepseek-v3-2-fp8-aw-tp8", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "GLM-4.7-FP8 TP4 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp4", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "prefix": "glm-4-7-fp8-aw-tp4", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "GLM-4.7-FP8 TP8 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp8", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "prefix": "glm-4-7-fp8-aw-tp8", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "gpt-oss-120b TP1 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp1", + "source_path": "openai/gpt-oss-120b", + "path": "openai/gpt-oss-120b", + "prefix": "gpt-oss-120b-aw-tp1", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "display": "gpt-oss-120b TP2 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp2", + "source_path": "openai/gpt-oss-120b", + "path": "openai/gpt-oss-120b", + "prefix": "gpt-oss-120b-aw-tp2", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "display": "gpt-oss-120b TP8 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp8", + "source_path": "openai/gpt-oss-120b", + "path": "openai/gpt-oss-120b", + "prefix": "gpt-oss-120b-aw-tp8", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "display": "Kimi-K2.5-MXFP4 TP4 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp4", + "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "prefix": "kimi-k2-5-mxfp4-aw-tp4", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "Kimi-K2.5-MXFP4 TP8 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp8", + "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "prefix": "kimi-k2-5-mxfp4-aw-tp8", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "display": "MiniMax-M2.5 TP2 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp2", + "source_path": "MiniMaxAI/MiniMax-M2.5", + "path": "MiniMaxAI/MiniMax-M2.5", + "prefix": "minimax-m2-5-aw-tp2", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "display": "MiniMax-M2.5 TP4 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp4", + "source_path": "MiniMaxAI/MiniMax-M2.5", + "path": "MiniMaxAI/MiniMax-M2.5", + "prefix": "minimax-m2-5-aw-tp4", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", + "runner": "atom-mi355-8gpu-oot-benchmark", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "display": "MiniMax-M2.5 TP8 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp8", + "source_path": "MiniMaxAI/MiniMax-M2.5", + "path": "MiniMaxAI/MiniMax-M2.5", + "prefix": "minimax-m2-5-aw-tp8", + "nightly_group": "B", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "bench_args": "--random-range-ratio 1", "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" } ] diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index 107f95f999..081500247f 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -50,6 +50,7 @@ EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} EXPLICIT_CLIENT_COMMAND=${OOT_CLIENT_COMMAND:-} +OOT_GPU_MEMORY_UTILIZATION=${OOT_GPU_MEMORY_UTILIZATION:-0.9} OOT_DOCKER_IMAGE=${OOT_DOCKER_IMAGE:-} LM_EVAL_NUM_FEWSHOT=${LM_EVAL_NUM_FEWSHOT:-3} LAST_VLLM_LOG_LINE=0 @@ -59,6 +60,11 @@ if ! [[ "${LM_EVAL_NUM_FEWSHOT}" =~ ^[0-9]+$ ]]; then exit 2 fi +if ! [[ "${OOT_GPU_MEMORY_UTILIZATION}" =~ ^[0-9]*\.?[0-9]+$ ]]; then + echo "Invalid OOT_GPU_MEMORY_UTILIZATION: ${OOT_GPU_MEMORY_UTILIZATION}. Expected a numeric value." + exit 2 +fi + declare -a ACTIVE_MODELS=() if [[ -n "${EXPLICIT_MODEL_NAME}" || -n "${EXPLICIT_MODEL_PATH}" || -n "${EXPLICIT_EXTRA_ARGS}" ]]; then if [[ -z "${EXPLICIT_MODEL_NAME}" || -z "${EXPLICIT_MODEL_PATH}" ]]; then @@ -180,6 +186,7 @@ PY echo "Model name: ${model_name}" echo "Model path: ${resolved_model_path}" echo "Extra args: ${extra_args}" + echo "GPU memory utilization: ${OOT_GPU_MEMORY_UTILIZATION}" export SAFETENSORS_FAST_GPU=1 export VLLM_RPC_TIMEOUT=1800000 @@ -206,7 +213,7 @@ PY --trust-remote-code \ --kv-cache-dtype fp8 \ "${extra_arg_array[@]}" \ - --gpu-memory-utilization 0.9 \ + --gpu-memory-utilization "${OOT_GPU_MEMORY_UTILIZATION}" \ --no-enable-prefix-caching \ > "${VLLM_LOG_FILE}" 2>&1 & echo $! > "${VLLM_PID_FILE}" diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index 455cdb0276..f539ee8255 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -8,79 +8,115 @@ concurrency: on: schedule: # Nightly at 23:00 Beijing time (15:00 UTC) - - cron: '0 15 * * *' + - cron: '0 15 * * 0-4' workflow_dispatch: inputs: - deepseek-r1-fp8: - description: "DeepSeek-R1 FP8 TP8" + deepseek-r1-fp8-met: + description: "DeepSeek-R1 FP8 TP8 (MET)" type: boolean default: false - deepseek-r1-mxfp4: - description: "DeepSeek-R1 MXFP4 TP8" + deepseek-r1-mxfp4-met: + description: "DeepSeek-R1 MXFP4 TP8 (MET)" type: boolean default: false - deepseek-v3-2-fp8: - description: "DeepSeek-V3.2 FP8 TP8" + gpt-oss-120b-met: + description: "gpt-oss-120b TP1 (MET)" type: boolean default: false - gpt-oss-120b: - description: "gpt-oss-120b TP1" + glm-5-1-fp8-tp8-oob: + description: "GLM-5.1-FP8 TP8 (OOB)" type: boolean default: false - glm-5-1-fp8-tp8: - description: "GLM-5.1-FP8 TP8" + kimi-k2-thinking-mxfp4-tp4-met: + description: "Kimi-K2-Thinking-MXFP4 TP4 (MET)" type: boolean default: false - kimi-k2-thinking-mxfp4-tp4: - description: "Kimi-K2-Thinking-MXFP4 TP4" + kimi-k2-thinking-mxfp4-tp8-met: + description: "Kimi-K2-Thinking-MXFP4 TP8 (MET)" type: boolean default: false - kimi-k2-thinking-mxfp4-tp8: - description: "Kimi-K2-Thinking-MXFP4 TP8" + kimi-k25-mxfp4-tp4-met: + description: "Kimi-K2.5-MXFP4 TP4 (MET)" type: boolean default: false - kimi-k2-5-mxfp4-tp2: - description: "Kimi-K2.5-MXFP4 TP2" + qwen3-5-397b-a17b-oob: + description: "Qwen3.5-397B-A17B TP8 (OOB)" type: boolean default: false - kimi-k2-5-mxfp4-tp4: - description: "Kimi-K2.5-MXFP4 TP4" + qwen3-5-397b-a17b-fp8-tp4-oob: + description: "Qwen3.5-397B-A17B-FP8 TP4 (OOB)" type: boolean default: false - kimi-k2-5-mxfp4-tp8: - description: "Kimi-K2.5-MXFP4 TP8" + qwen3-5-397b-a17b-fp8-oob: + description: "Qwen3.5-397B-A17B-FP8 TP8 (OOB)" type: boolean default: false - qwen3-5-397b-a17b-tp4: - description: "Qwen3.5-397B-A17B TP4" + qwen3-next-80b-a3b-instruct-fp8-tp1-met: + description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)" type: boolean default: false - qwen3-5-397b-a17b: - description: "Qwen3.5-397B-A17B TP8" + qwen3-next-80b-a3b-instruct-fp8-tp4-met: + description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)" type: boolean default: false - qwen3-5-397b-a17b-fp8-tp4: - description: "Qwen3.5-397B-A17B-FP8 TP4" + qwen3-next-80b-a3b-instruct-fp8-aw-tp1: + description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)" type: boolean default: false - qwen3-5-397b-a17b-fp8: - description: "Qwen3.5-397B-A17B-FP8 TP8" + qwen3-next-80b-a3b-instruct-fp8-aw-tp2: + description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)" type: boolean default: false - qwen3-5-397b-a17b-mxfp4: - description: "Qwen3.5-397B-A17B-MXFP4 TP4" + qwen3-next-80b-a3b-instruct-fp8-aw-tp4: + description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)" type: boolean default: false - qwen3-next-80b-a3b-instruct-fp8-tp1: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1" + deepseek-v3-2-fp8-aw-tp4: + description: "DeepSeek-V3.2 FP8 TP4 (AW)" type: boolean default: false - qwen3-next-80b-a3b-instruct-fp8-tp2: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP2" + deepseek-v3-2-fp8-aw-tp8: + description: "DeepSeek-V3.2 FP8 TP8 (AW)" type: boolean default: false - qwen3-next-80b-a3b-instruct-fp8-tp4: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4" + glm-4-7-fp8-aw-tp4: + description: "GLM-4.7-FP8 TP4 (AW)" + type: boolean + default: false + glm-4-7-fp8-aw-tp8: + description: "GLM-4.7-FP8 TP8 (AW)" + type: boolean + default: false + gpt-oss-120b-aw-tp1: + description: "gpt-oss-120b TP1 (AW)" + type: boolean + default: false + gpt-oss-120b-aw-tp2: + description: "gpt-oss-120b TP2 (AW)" + type: boolean + default: false + gpt-oss-120b-aw-tp8: + description: "gpt-oss-120b TP8 (AW)" + type: boolean + default: false + kimi-k2-5-mxfp4-aw-tp4: + description: "Kimi-K2.5-MXFP4 TP4 (AW)" + type: boolean + default: false + kimi-k2-5-mxfp4-aw-tp8: + description: "Kimi-K2.5-MXFP4 TP8 (AW)" + type: boolean + default: false + minimax-m2-5-aw-tp2: + description: "MiniMax-M2.5 TP2 (AW)" + type: boolean + default: false + minimax-m2-5-aw-tp4: + description: "MiniMax-M2.5 TP4 (AW)" + type: boolean + default: false + minimax-m2-5-aw-tp8: + description: "MiniMax-M2.5 TP8 (AW)" type: boolean default: false oot_image: @@ -422,9 +458,9 @@ jobs: CONCURRENCY="${PARAMS[2]}" RANDOM_RANGE_RATIO="${PARAMS[3]}" case "${CONCURRENCY}" in - 4|8|16|32|64|128|256) ;; + 4|8|16|32|64|128|256|512) ;; *) - echo "Unsupported concurrency: ${CONCURRENCY}. Allowed values: 4,8,16,32,64,128,256" + echo "Unsupported concurrency: ${CONCURRENCY}. Allowed values: 4,8,16,32,64,128,256,512" exit 1 ;; esac @@ -453,31 +489,40 @@ jobs: - uses: actions/checkout@v6 - id: load env: - ENABLE_DEEPSEEK_R1_FP8: ${{ inputs.deepseek-r1-fp8 }} - ENABLE_DEEPSEEK_R1_MXFP4: ${{ inputs.deepseek-r1-mxfp4 }} - ENABLE_DEEPSEEK_V3_2_FP8: ${{ inputs.deepseek-v3-2-fp8 }} - ENABLE_GPT_OSS_120B: ${{ inputs.gpt-oss-120b }} - ENABLE_GLM_5_1_FP8_TP8: ${{ inputs.glm-5-1-fp8-tp8 }} - ENABLE_KIMI_K2_TP4: ${{ inputs.kimi-k2-thinking-mxfp4-tp4 }} - ENABLE_KIMI_K2_TP8: ${{ inputs.kimi-k2-thinking-mxfp4-tp8 }} - ENABLE_KIMI_K25_TP2: ${{ inputs.kimi-k2-5-mxfp4-tp2 }} - ENABLE_KIMI_K25_TP4: ${{ inputs.kimi-k2-5-mxfp4-tp4 }} - ENABLE_KIMI_K25_TP8: ${{ inputs.kimi-k2-5-mxfp4-tp8 }} - ENABLE_QWEN3_5_397B_A17B_TP4: ${{ inputs.qwen3-5-397b-a17b-tp4 }} - ENABLE_QWEN3_5_397B_A17B: ${{ inputs.qwen3-5-397b-a17b }} - ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }} - ENABLE_QWEN3_5_397B_A17B_FP8: ${{ inputs.qwen3-5-397b-a17b-fp8 }} - ENABLE_QWEN3_5_397B_A17B_MXFP4: ${{ inputs.qwen3-5-397b-a17b-mxfp4 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP2: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp2 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4 }} + ENABLE_DEEPSEEK_R1_FP8: ${{ inputs.deepseek-r1-fp8-met }} + ENABLE_DEEPSEEK_R1_MXFP4: ${{ inputs.deepseek-r1-mxfp4-met }} + ENABLE_GPT_OSS_120B: ${{ inputs.gpt-oss-120b-met }} + ENABLE_GLM_5_1_FP8_TP8: ${{ inputs.glm-5-1-fp8-tp8-oob }} + ENABLE_KIMI_K2_TP4: ${{ inputs.kimi-k2-thinking-mxfp4-tp4-met }} + ENABLE_KIMI_K2_TP8: ${{ inputs.kimi-k2-thinking-mxfp4-tp8-met }} + ENABLE_KIMI_K25_TP4: ${{ inputs.kimi-k25-mxfp4-tp4-met }} + ENABLE_QWEN3_5_397B_A17B: ${{ inputs.qwen3-5-397b-a17b-oob }} + ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4-oob }} + ENABLE_QWEN3_5_397B_A17B_FP8: ${{ inputs.qwen3-5-397b-a17b-fp8-oob }} + ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1-met }} + ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4-met }} + ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp1 }} + ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP2: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp2 }} + ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp4 }} + ENABLE_DEEPSEEK_V3_2_FP8_AW_TP4: ${{ inputs.deepseek-v3-2-fp8-aw-tp4 }} + ENABLE_DEEPSEEK_V3_2_FP8_AW_TP8: ${{ inputs.deepseek-v3-2-fp8-aw-tp8 }} + ENABLE_GLM_4_7_FP8_AW_TP4: ${{ inputs.glm-4-7-fp8-aw-tp4 }} + ENABLE_GLM_4_7_FP8_AW_TP8: ${{ inputs.glm-4-7-fp8-aw-tp8 }} + ENABLE_GPT_OSS_120B_AW_TP1: ${{ inputs.gpt-oss-120b-aw-tp1 }} + ENABLE_GPT_OSS_120B_AW_TP2: ${{ inputs.gpt-oss-120b-aw-tp2 }} + ENABLE_GPT_OSS_120B_AW_TP8: ${{ inputs.gpt-oss-120b-aw-tp8 }} + ENABLE_KIMI_K2_5_MXFP4_AW_TP4: ${{ inputs.kimi-k2-5-mxfp4-aw-tp4 }} + ENABLE_KIMI_K2_5_MXFP4_AW_TP8: ${{ inputs.kimi-k2-5-mxfp4-aw-tp8 }} + ENABLE_MINIMAX_M2_5_AW_TP2: ${{ inputs.minimax-m2-5-aw-tp2 }} + ENABLE_MINIMAX_M2_5_AW_TP4: ${{ inputs.minimax-m2-5-aw-tp4 }} + ENABLE_MINIMAX_M2_5_AW_TP8: ${{ inputs.minimax-m2-5-aw-tp8 }} GITHUB_TOKEN: ${{ github.token }} run: | python3 - <<'PY' >> "$GITHUB_OUTPUT" import json import os import sys - from datetime import date, datetime, timedelta, timezone + from datetime import datetime, timedelta, timezone from pathlib import Path from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen @@ -488,24 +533,33 @@ jobs: event = os.environ["GITHUB_EVENT_NAME"] manual_toggles = { - "deepseek-r1-fp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP8", "").lower() == "true", - "deepseek-r1-mxfp4": os.environ.get("ENABLE_DEEPSEEK_R1_MXFP4", "").lower() == "true", - "deepseek-v3-2-fp8": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8", "").lower() == "true", - "gpt-oss-120b": os.environ.get("ENABLE_GPT_OSS_120B", "").lower() == "true", - "glm-5-1-fp8-tp8": os.environ.get("ENABLE_GLM_5_1_FP8_TP8", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp4": os.environ.get("ENABLE_KIMI_K2_TP4", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp8": os.environ.get("ENABLE_KIMI_K2_TP8", "").lower() == "true", - "kimi-k25-mxfp4-tp2": os.environ.get("ENABLE_KIMI_K25_TP2", "").lower() == "true", - "kimi-k25-mxfp4-tp4": os.environ.get("ENABLE_KIMI_K25_TP4", "").lower() == "true", - "kimi-k25-mxfp4-tp8": os.environ.get("ENABLE_KIMI_K25_TP8", "").lower() == "true", - "qwen3-5-397b-a17b-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_TP4", "").lower() == "true", - "qwen3-5-397b-a17b": os.environ.get("ENABLE_QWEN3_5_397B_A17B", "").lower() == "true", - "qwen3-5-397b-a17b-fp8-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", "").lower() == "true", - "qwen3-5-397b-a17b-fp8": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8", "").lower() == "true", - "qwen3-5-397b-a17b-mxfp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_MXFP4", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp1": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp2": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP2", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp4": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4", "").lower() == "true", + "deepseek-r1-fp8-met": os.environ.get("ENABLE_DEEPSEEK_R1_FP8", "").lower() == "true", + "deepseek-r1-mxfp4-met": os.environ.get("ENABLE_DEEPSEEK_R1_MXFP4", "").lower() == "true", + "gpt-oss-120b-met": os.environ.get("ENABLE_GPT_OSS_120B", "").lower() == "true", + "glm-5-1-fp8-tp8-oob": os.environ.get("ENABLE_GLM_5_1_FP8_TP8", "").lower() == "true", + "kimi-k2-thinking-mxfp4-tp4-met": os.environ.get("ENABLE_KIMI_K2_TP4", "").lower() == "true", + "kimi-k2-thinking-mxfp4-tp8-met": os.environ.get("ENABLE_KIMI_K2_TP8", "").lower() == "true", + "kimi-k25-mxfp4-tp4-met": os.environ.get("ENABLE_KIMI_K25_TP4", "").lower() == "true", + "qwen3-5-397b-a17b-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B", "").lower() == "true", + "qwen3-5-397b-a17b-fp8-tp4-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", "").lower() == "true", + "qwen3-5-397b-a17b-fp8-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8", "").lower() == "true", + "qwen3-next-80b-a3b-instruct-fp8-tp1-met": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1", "").lower() == "true", + "qwen3-next-80b-a3b-instruct-fp8-tp4-met": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4", "").lower() == "true", + "qwen3-next-80b-a3b-instruct-fp8-aw-tp1": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP1", "").lower() == "true", + "qwen3-next-80b-a3b-instruct-fp8-aw-tp2": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP2", "").lower() == "true", + "qwen3-next-80b-a3b-instruct-fp8-aw-tp4": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP4", "").lower() == "true", + "deepseek-v3-2-fp8-aw-tp4": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8_AW_TP4", "").lower() == "true", + "deepseek-v3-2-fp8-aw-tp8": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8_AW_TP8", "").lower() == "true", + "glm-4-7-fp8-aw-tp4": os.environ.get("ENABLE_GLM_4_7_FP8_AW_TP4", "").lower() == "true", + "glm-4-7-fp8-aw-tp8": os.environ.get("ENABLE_GLM_4_7_FP8_AW_TP8", "").lower() == "true", + "gpt-oss-120b-aw-tp1": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP1", "").lower() == "true", + "gpt-oss-120b-aw-tp2": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP2", "").lower() == "true", + "gpt-oss-120b-aw-tp8": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP8", "").lower() == "true", + "kimi-k2-5-mxfp4-aw-tp4": os.environ.get("ENABLE_KIMI_K2_5_MXFP4_AW_TP4", "").lower() == "true", + "kimi-k2-5-mxfp4-aw-tp8": os.environ.get("ENABLE_KIMI_K2_5_MXFP4_AW_TP8", "").lower() == "true", + "minimax-m2-5-aw-tp2": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP2", "").lower() == "true", + "minimax-m2-5-aw-tp4": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP4", "").lower() == "true", + "minimax-m2-5-aw-tp8": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP8", "").lower() == "true", } def fetch_run_created_at() -> datetime: @@ -542,20 +596,30 @@ jobs: selected_group = "" if event == "schedule": beijing_tz = timezone(timedelta(hours=8)) - beijing_date = fetch_run_created_at().astimezone(beijing_tz).date() - # Rotation anchor: 2026-05-06 runs group A, then B, then C. - anchor_date = date(2026, 5, 6) - rotation = ("A", "B", "C") - selected_group = rotation[(beijing_date - anchor_date).days % len(rotation)] - selected = [ - model for model in models if model.get("nightly_group") == selected_group - ] - if not selected: + beijing_dt = fetch_run_created_at().astimezone(beijing_tz) + beijing_weekday = beijing_dt.weekday() # Monday=0 ... Sunday=6 + + if beijing_weekday in (0, 2): # Monday / Wednesday + selected_group = "A-MET" + selected = [ + model + for model in models + if str(model.get("prefix", "")).endswith("-met") + ] + elif beijing_weekday in (1, 3): # Tuesday / Thursday + selected_group = "B-AW" + selected = [ + model for model in models if "-aw-" in str(model.get("prefix", "")) + ] + elif beijing_weekday == 4: # Friday + selected_group = "C-ALL" + selected = list(models) + else: # Saturday / Sunday + selected_group = "SKIP-WEEKEND" + selected = [] print( - f"No models configured for nightly benchmark group {selected_group}.", - file=sys.stderr, + f"Weekend schedule in Beijing ({beijing_dt.date()}); no benchmark group configured." ) - sys.exit(1) else: selected = [ model for model in models if manual_toggles.get(model["prefix"], False) @@ -600,6 +664,13 @@ jobs: include = [] for model in models: + prefix = str(model.get("prefix", "")) + is_aw_model = "-aw-" in prefix + aw_pairs = ( + (1000, 100), + (5000, 500), + (10000, 1000), + ) supported_pairs = set(model.get("supported_input_output_pairs", [])) excluded_pairs = set(model.get("excluded_input_output_pairs", [])) model_params = [] @@ -608,11 +679,22 @@ jobs: extra_concurrency = (128, 256) for param in params: - variants = [param] - for concurrency in extra_concurrency: - variant = dict(param) - variant["concurrency"] = concurrency - variants.append(variant) + base_params = [param] + if is_aw_model: + base_params = [] + for input_length, output_length in aw_pairs: + aws_param = dict(param) + aws_param["input_length"] = input_length + aws_param["output_length"] = output_length + base_params.append(aws_param) + + variants = [] + for base_param in base_params: + variants.append(base_param) + for concurrency in extra_concurrency: + variant = dict(base_param) + variant["concurrency"] = concurrency + variants.append(variant) for variant in variants: key = ( @@ -775,6 +857,7 @@ jobs: MODEL_SOURCE_PATH: ${{ matrix.model.source_path || matrix.model.path }} MODEL_PATH: ${{ matrix.model.path || matrix.model.source_path }} OOT_EXTRA_ARGS: ${{ matrix.model.extra_args }} + BENCH_CLIENT: ${{ contains(matrix.model.prefix, '-aw-') && 'vllm-bench' || 'benchmark_serving.py' }} BENCH_EXTRA_ARGS: ${{ matrix.model.bench_args }} RESULT_PREFIX: ${{ matrix.model.prefix }} ISL: ${{ matrix.params.input_length }} @@ -816,24 +899,33 @@ jobs: fi case "${{ matrix.model.prefix }}" in - deepseek-r1-fp8) echo "enabled=${{ inputs.deepseek-r1-fp8 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-r1-mxfp4) echo "enabled=${{ inputs.deepseek-r1-mxfp4 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-v3-2-fp8) echo "enabled=${{ inputs.deepseek-v3-2-fp8 }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b) echo "enabled=${{ inputs.gpt-oss-120b }}" >> "$GITHUB_OUTPUT" ;; - glm-5-1-fp8-tp8) echo "enabled=${{ inputs.glm-5-1-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp4) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp8) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp2) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp2 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp4) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp8) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-tp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b) echo "enabled=${{ inputs.qwen3-5-397b-a17b }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-mxfp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-mxfp4 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp1) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp2) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp2 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp4) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; + deepseek-r1-fp8-met) echo "enabled=${{ inputs.deepseek-r1-fp8-met }}" >> "$GITHUB_OUTPUT" ;; + deepseek-r1-mxfp4-met) echo "enabled=${{ inputs.deepseek-r1-mxfp4-met }}" >> "$GITHUB_OUTPUT" ;; + gpt-oss-120b-met) echo "enabled=${{ inputs.gpt-oss-120b-met }}" >> "$GITHUB_OUTPUT" ;; + glm-5-1-fp8-tp8-oob) echo "enabled=${{ inputs.glm-5-1-fp8-tp8-oob }}" >> "$GITHUB_OUTPUT" ;; + kimi-k2-thinking-mxfp4-tp4-met) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp4-met }}" >> "$GITHUB_OUTPUT" ;; + kimi-k2-thinking-mxfp4-tp8-met) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp8-met }}" >> "$GITHUB_OUTPUT" ;; + kimi-k25-mxfp4-tp4-met) echo "enabled=${{ inputs.kimi-k25-mxfp4-tp4-met }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-oob }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-tp4-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4-oob }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-oob }}" >> "$GITHUB_OUTPUT" ;; + qwen3-next-80b-a3b-instruct-fp8-tp1-met) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1-met }}" >> "$GITHUB_OUTPUT" ;; + qwen3-next-80b-a3b-instruct-fp8-tp4-met) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4-met }}" >> "$GITHUB_OUTPUT" ;; + qwen3-next-80b-a3b-instruct-fp8-aw-tp1) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp1 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-next-80b-a3b-instruct-fp8-aw-tp2) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-next-80b-a3b-instruct-fp8-aw-tp4) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; + deepseek-v3-2-fp8-aw-tp4) echo "enabled=${{ inputs.deepseek-v3-2-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; + deepseek-v3-2-fp8-aw-tp8) echo "enabled=${{ inputs.deepseek-v3-2-fp8-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; + glm-4-7-fp8-aw-tp4) echo "enabled=${{ inputs.glm-4-7-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; + glm-4-7-fp8-aw-tp8) echo "enabled=${{ inputs.glm-4-7-fp8-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; + gpt-oss-120b-aw-tp1) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp1 }}" >> "$GITHUB_OUTPUT" ;; + gpt-oss-120b-aw-tp2) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; + gpt-oss-120b-aw-tp8) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; + kimi-k2-5-mxfp4-aw-tp4) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; + kimi-k2-5-mxfp4-aw-tp8) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; + minimax-m2-5-aw-tp2) echo "enabled=${{ inputs.minimax-m2-5-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; + minimax-m2-5-aw-tp4) echo "enabled=${{ inputs.minimax-m2-5-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; + minimax-m2-5-aw-tp8) echo "enabled=${{ inputs.minimax-m2-5-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; *) echo "enabled=true" >> "$GITHUB_OUTPUT" ;; esac @@ -1009,13 +1101,17 @@ jobs: - name: Prepare OOT benchmark runner in container if: steps.check.outputs.enabled == 'true' run: | - $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " - set -euo pipefail - rm -rf \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" - mkdir -p \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" - git clone --depth 1 \"${BENCH_SERVING_REPO_URL}\" \"${CONTAINER_BENCH_SERVING_DIR}\" - test -f \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" - " + if [[ "${BENCH_CLIENT}" == "vllm-bench" ]]; then + echo "Skipping bench_serving repo setup for vllm bench client." + else + $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" + mkdir -p \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" + git clone --depth 1 \"${BENCH_SERVING_REPO_URL}\" \"${CONTAINER_BENCH_SERVING_DIR}\" + test -f \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" + " + fi - name: Launch OOT benchmark server if: steps.check.outputs.enabled == 'true' @@ -1049,38 +1145,74 @@ jobs: fi { - echo "=== Benchmark config: ${MODEL_NAME} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} ===" - $CONTAINER_ENGINE exec \ - -e ISL="${ISL}" \ - -e OSL="${OSL}" \ - -e CONC="${CONC}" \ - -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ - -e RESULT_FILENAME="${RESULT_FILENAME}" \ - -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ - "$CONTAINER_NAME" bash -lc " - set -euo pipefail - rm -rf \"${CONTAINER_RESULT_DIR}\" - mkdir -p \"${CONTAINER_RESULT_DIR}\" - PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ - --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ - --backend=vllm \ - --base-url=http://127.0.0.1:8000 \ - --dataset-name=random \ - --random-input-len=\"${ISL}\" \ - --random-output-len=\"${OSL}\" \ - --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ - --num-prompts=\"$(( CONC * 10 ))\" \ - --max-concurrency=\"${CONC}\" \ - ${TRUST_REMOTE_CODE_ARG} \ - --num-warmups=\"$(( 2 * CONC ))\" \ - --request-rate=inf \ - --ignore-eos \ - --save-result \ - --percentile-metrics=\"ttft,tpot,itl,e2el\" \ - --result-dir=\"${CONTAINER_RESULT_DIR}\" \ - --result-filename=\"${RESULT_FILENAME}.json\" \ - ${BENCH_EXTRA_ARGS:-} - " + echo "=== Benchmark config: ${MODEL_NAME} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} CLIENT=${BENCH_CLIENT} ===" + if [[ "${BENCH_CLIENT}" == "vllm-bench" ]]; then + $CONTAINER_ENGINE exec \ + -e ISL="${ISL}" \ + -e OSL="${OSL}" \ + -e CONC="${CONC}" \ + -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ + -e RESULT_FILENAME="${RESULT_FILENAME}" \ + -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ + "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_RESULT_DIR}\" + mkdir -p \"${CONTAINER_RESULT_DIR}\" + vllm bench serve \ + --backend vllm \ + --host 127.0.0.1 \ + --port 8000 \ + --model \"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --dataset-name random \ + --random-input-len \"${ISL}\" \ + --random-output-len \"${OSL}\" \ + --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ + --max-concurrency \"${CONC}\" \ + --num-prompts \"$(( CONC * 10 ))\" \ + --num-warmups \"$(( 2 * CONC ))\" \ + --seed 42 \ + --ignore-eos \ + ${TRUST_REMOTE_CODE_ARG} \ + --percentile-metrics ttft,tpot,itl,e2el \ + --ready-check-timeout-sec 7200 \ + --save-result \ + --result-dir \"${CONTAINER_RESULT_DIR}\" \ + --result-filename \"${RESULT_FILENAME}.json\" \ + ${BENCH_EXTRA_ARGS:-} + " + else + $CONTAINER_ENGINE exec \ + -e ISL="${ISL}" \ + -e OSL="${OSL}" \ + -e CONC="${CONC}" \ + -e RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO}" \ + -e RESULT_FILENAME="${RESULT_FILENAME}" \ + -e BENCH_EXTRA_ARGS="${BENCH_EXTRA_ARGS}" \ + "$CONTAINER_NAME" bash -lc " + set -euo pipefail + rm -rf \"${CONTAINER_RESULT_DIR}\" + mkdir -p \"${CONTAINER_RESULT_DIR}\" + PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ + --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --backend=vllm \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len=\"${ISL}\" \ + --random-output-len=\"${OSL}\" \ + --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ + --num-prompts=\"$(( CONC * 10 ))\" \ + --max-concurrency=\"${CONC}\" \ + ${TRUST_REMOTE_CODE_ARG} \ + --num-warmups=\"$(( 2 * CONC ))\" \ + --request-rate=inf \ + --ignore-eos \ + --save-result \ + --percentile-metrics=\"ttft,tpot,itl,e2el\" \ + --result-dir=\"${CONTAINER_RESULT_DIR}\" \ + --result-filename=\"${RESULT_FILENAME}.json\" \ + ${BENCH_EXTRA_ARGS:-} + " + fi } > "./${RESULT_FILENAME}.client.log" 2>&1 - name: Copy OOT benchmark outputs From 7934d5e8ffd346b2e3f976e95a65f2867ee69f10 Mon Sep 17 00:00:00 2001 From: Zhu Yuhua Date: Mon, 11 May 2026 16:52:33 +0800 Subject: [PATCH 22/76] [fix][acc][sgl-atom] fix accuracy of fp8 attn weights model using ptpc quant recipe (#747) * [fix][acc] fix accuracy of fp8 attn weights model using ptpc quant recipe Signed-off-by: zhuyuhua-v --- .../attention_backend/sgl_attention_mla.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py b/atom/plugin/sglang/attention_backend/sgl_attention_mla.py index d0c2b3b07a..c60b12f313 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py +++ b/atom/plugin/sglang/attention_backend/sgl_attention_mla.py @@ -21,13 +21,13 @@ import torch from aiter import dtypes from aiter.dist.parallel_state import get_tensor_model_parallel_world_size, get_tp_group -from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant from atom.model_ops.base_attention import Attention from atom.model_ops.attention_mla import ( dynamic_per_batched_tensor_quant, fused_qk_rope_concat_and_cache_mla, ) from atom.models.utils import maybe_prefix +from atom.models.deepseek_v2 import _fuse_rmsnorm_quant # sglang imports from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context @@ -117,6 +117,11 @@ def _unwrap_linear_output(output: Any) -> torch.Tensor: return output +def _linear_quant_type_value(linear: Any) -> Optional[int]: + quant_type = getattr(linear, "quant_type", None) + return None if quant_type is None else getattr(quant_type, "value", quant_type) + + def _fuse_qk_rmsnorm_and_q_quant( attn: DeepseekV2MLAAttention, q: torch.Tensor, @@ -125,7 +130,6 @@ def _fuse_qk_rmsnorm_and_q_quant( output_unquantized_q: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: """Fuse q/k RMSNorm and q quant using ATOM's DeepSeek-V2 path.""" - from atom.models.deepseek_v2 import _fuse_rmsnorm_quant (q_quantized, q_scale), q_normed, k_nope_normed, _ = _fuse_rmsnorm_quant( q, @@ -139,6 +143,7 @@ def _fuse_qk_rmsnorm_and_q_quant( shuffle=False, scale_shuffle_padding=False, group_size=128, + quant_type=_linear_quant_type_value(attn.q_b_proj), output_unquantized_inp1=output_unquantized_q, transpose_scale=True, ) @@ -640,6 +645,7 @@ def forward_sgl_mha_prepare( hidden_states: torch.Tensor, **model_kwargs, ) -> SglMhaPrepareResult: + forward_batch = model_kwargs.get("forward_batch", None) if forward_batch is None: raise RuntimeError("forward_batch is required in forward_sgl_mha_prepare") @@ -681,16 +687,17 @@ def forward_sgl_mha_prepare( ) if _use_aiter_gfx95 and attn.q_b_proj.weight.dtype == torch.float8_e4m3fn: - (q, q_scale), _, _, _ = fused_rms_fp8_group_quant( + (q, q_scale), _, _, _ = _fuse_rmsnorm_quant( q, attn.q_a_layernorm.weight, attn.q_a_layernorm.eps, None, None, None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, res1=None, + dtype_quant=torch.float8_e4m3fn, + group_size=128, + quant_type=_linear_quant_type_value(attn.q_b_proj), output_unquantized_inp1=False, transpose_scale=True, ) @@ -715,16 +722,17 @@ def forward_sgl_mha_prepare( latent_cache = latent_cache.unsqueeze(1) if _use_aiter_gfx95 and attn.kv_b_proj.weight.dtype == torch.float8_e4m3fn: - (kv_a_quanted, kv_a_quanted_scale), kv_a, _, _ = fused_rms_fp8_group_quant( + (kv_a_quanted, kv_a_quanted_scale), kv_a, _, _ = _fuse_rmsnorm_quant( kv_a, attn.kv_a_layernorm.weight, attn.kv_a_layernorm.eps, None, None, None, - group_size=128, - dtype_quant=torch.float8_e4m3fn, res1=None, + dtype_quant=torch.float8_e4m3fn, + group_size=128, + quant_type=_linear_quant_type_value(attn.kv_b_proj), output_unquantized_inp1=True, transpose_scale=True, ) From c225b062bfd379b89fedde50c5264810b7b7bcca Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 12 May 2026 09:09:21 +0800 Subject: [PATCH 23/76] [atom-vllm benchmark] Refine atom-vllm benchmark (#729) * [atom-vllm benchmark] refine atom-vllm benchmark UI Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * add Signed-off-by: zejunchen-zejun * Restore main's (MET)/(OOB)/(AW) naming, weekday nightly rotation, and AW benchmark logic - Restore (MET)/(OOB)/(AW) suffix in all model display names and prefixes - Remove feature-branch-only TP variants, keep exactly main's 27 models - Add all 15 AW model variants from main with bench_args "--random-range-ratio 1" - Restore weekday-based nightly rotation: Mon/Wed A-MET, Tue/Thu B-AW, Fri C-ALL - Restore is_aw_model logic in build-benchmark-matrix (AW ISL/OSL pairs + extra concurrency) - Restore per-model BENCH_CLIENT override: AW models auto-use vLLM bench client Co-Authored-By: Claude Opus 4 * Fix dashboard upload defaults and AW model data accuracy - Enable both dashboard uploads (gh-pages + LLM-Booster) for nightly schedule runs - Set publish_to_dashboard default to true for manual runs (matching main) - Fix AW model benchmark_client in LLM-Booster data: read per-result client from enriched JSON instead of using global env (AW models override to vLLM bench) - Fix AW model random_range_ratio: set to "1" in matrix params so filename, benchmark command, and dashboard data are all consistent - Add dashboard_model for DeepSeek-R1 FP8/MXFP4 and gpt-oss-120b MET to avoid (MET) suffix leaking into LLM-Booster display names Co-Authored-By: Claude Opus 4 --------- Signed-off-by: zejunchen-zejun Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 700 +++++----- .github/scripts/atom_oot_test.sh | 1 + .github/workflows/atom-vllm-benchmark.yaml | 1263 +++++++++++++------ 3 files changed, 1263 insertions(+), 701 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index da04f0a446..7326bf05ce 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -1,335 +1,365 @@ -[ - { - "display": "DeepSeek-R1 FP8 TP8 (MET)", - "source_path": "deepseek-ai/DeepSeek-R1-0528", - "path": "deepseek-ai/DeepSeek-R1-0528", - "prefix": "deepseek-r1-fp8-met", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "DeepSeek-R1 MXFP4 TP8 (MET)", - "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "prefix": "deepseek-r1-mxfp4-met", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "gpt-oss-120b TP1 (MET)", - "source_path": "openai/gpt-oss-120b", - "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b-met", - "nightly_group": "C", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" - }, - { - "display": "GLM-5.1-FP8 TP8 (OOB)", - "dashboard_model": "GLM-5.1-FP8", - "source_path": "zai-org/GLM-5.1-FP8", - "path": "zai-org/GLM-5.1-FP8", - "prefix": "glm-5-1-fp8-tp8-oob", - "nightly_group": "C", - "excluded_input_output_pairs": [ - "1024x8192" - ], - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}' --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2-Thinking-MXFP4 TP4 (MET)", - "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp4", - "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp4-met", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2-Thinking-MXFP4 TP8 (MET)", - "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp8", - "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", - "prefix": "kimi-k2-thinking-mxfp4-tp8-met", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP4 (MET)", - "dashboard_model": "Kimi-K2.5-MXFP4-tp4", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k25-mxfp4-tp4-met", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP8 (OOB)", - "dashboard_model": "Qwen3.5-397B-A17B-FP8", - "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", - "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8-oob", - "nightly_group": "A", - "excluded_input_output_pairs": [ - "1024x8192" - ], - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", - "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", - "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", - "path": "Qwen/Qwen3.5-397B-A17B-FP8", - "prefix": "qwen3-5-397b-a17b-fp8-tp4-oob", - "nightly_group": "A", - "excluded_input_output_pairs": [ - "1024x8192" - ], - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3.5-397B-A17B TP8 (OOB)", - "dashboard_model": "Qwen3.5-397B-A17B", - "source_path": "Qwen/Qwen3.5-397B-A17B", - "path": "Qwen/Qwen3.5-397B-A17B", - "prefix": "qwen3-5-397b-a17b-oob", - "nightly_group": "A", - "excluded_input_output_pairs": [ - "1024x8192" - ], - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1-met", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4-met", - "nightly_group": "A", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp1", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp2", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" - }, - { - "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", - "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp4", - "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", - "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" - }, - { - "display": "DeepSeek-V3.2 FP8 TP4 (AW)", - "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", - "source_path": "deepseek-ai/DeepSeek-V3.2", - "path": "deepseek-ai/DeepSeek-V3.2", - "prefix": "deepseek-v3-2-fp8-aw-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "DeepSeek-V3.2 FP8 TP8 (AW)", - "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", - "source_path": "deepseek-ai/DeepSeek-V3.2", - "path": "deepseek-ai/DeepSeek-V3.2", - "prefix": "deepseek-v3-2-fp8-aw-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "GLM-4.7-FP8 TP4 (AW)", - "dashboard_model": "GLM-4.7-FP8-aw-tp4", - "source_path": "zai-org/GLM-4.7-FP8", - "path": "zai-org/GLM-4.7-FP8", - "prefix": "glm-4-7-fp8-aw-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "GLM-4.7-FP8 TP8 (AW)", - "dashboard_model": "GLM-4.7-FP8-aw-tp8", - "source_path": "zai-org/GLM-4.7-FP8", - "path": "zai-org/GLM-4.7-FP8", - "prefix": "glm-4-7-fp8-aw-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "gpt-oss-120b TP1 (AW)", - "dashboard_model": "gpt-oss-120b-aw-tp1", - "source_path": "openai/gpt-oss-120b", - "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b-aw-tp1", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" - }, - { - "display": "gpt-oss-120b TP2 (AW)", - "dashboard_model": "gpt-oss-120b-aw-tp2", - "source_path": "openai/gpt-oss-120b", - "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b-aw-tp2", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" - }, - { - "display": "gpt-oss-120b TP8 (AW)", - "dashboard_model": "gpt-oss-120b-aw-tp8", - "source_path": "openai/gpt-oss-120b", - "path": "openai/gpt-oss-120b", - "prefix": "gpt-oss-120b-aw-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" - }, - { - "display": "Kimi-K2.5-MXFP4 TP4 (AW)", - "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp4", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k2-5-mxfp4-aw-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "Kimi-K2.5-MXFP4 TP8 (AW)", - "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp8", - "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", - "prefix": "kimi-k2-5-mxfp4-aw-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" - }, - { - "display": "MiniMax-M2.5 TP2 (AW)", - "dashboard_model": "MiniMax-M2.5-aw-tp2", - "source_path": "MiniMaxAI/MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "minimax-m2-5-aw-tp2", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 2 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - }, - { - "display": "MiniMax-M2.5 TP4 (AW)", - "dashboard_model": "MiniMax-M2.5-aw-tp4", - "source_path": "MiniMaxAI/MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "minimax-m2-5-aw-tp4", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - }, - { - "display": "MiniMax-M2.5 TP8 (AW)", - "dashboard_model": "MiniMax-M2.5-aw-tp8", - "source_path": "MiniMaxAI/MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "minimax-m2-5-aw-tp8", - "nightly_group": "B", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", - "bench_args": "--random-range-ratio 1", - "runner": "atom-mi355-8gpu-oot-benchmark", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" - } -] +{ + "families": [ + { + "choice_label": "DeepSeek-R1 FP8", + "source_path": "deepseek-ai/DeepSeek-R1-0528", + "path": "deepseek-ai/DeepSeek-R1-0528", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 8, + "display": "DeepSeek-R1 FP8 TP8 (MET)", + "prefix": "deepseek-r1-fp8-met", + "dashboard_model": "DeepSeek-R1-FP8-tp8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "DeepSeek-R1 MXFP4", + "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 8, + "display": "DeepSeek-R1 MXFP4 TP8 (MET)", + "prefix": "deepseek-r1-mxfp4-met", + "dashboard_model": "DeepSeek-R1-MXFP4-tp8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "gpt-oss-120b", + "source_path": "openai/gpt-oss-120b", + "path": "openai/gpt-oss-120b", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "C", + "variants": [ + { + "tp_size": 1, + "display": "gpt-oss-120b TP1 (MET)", + "prefix": "gpt-oss-120b-met", + "dashboard_model": "gpt-oss-120b-tp1", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "tp_size": 1, + "display": "gpt-oss-120b TP1 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp1", + "prefix": "gpt-oss-120b-aw-tp1", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "tp_size": 2, + "display": "gpt-oss-120b TP2 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp2", + "prefix": "gpt-oss-120b-aw-tp2", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + }, + { + "tp_size": 8, + "display": "gpt-oss-120b TP8 (AW)", + "dashboard_model": "gpt-oss-120b-aw-tp8", + "prefix": "gpt-oss-120b-aw-tp8", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + } + ] + }, + { + "choice_label": "GLM-5.1-FP8", + "source_path": "zai-org/GLM-5.1-FP8", + "path": "zai-org/GLM-5.1-FP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "C", + "variants": [ + { + "tp_size": 8, + "display": "GLM-5.1-FP8 TP8 (OOB)", + "dashboard_model": "GLM-5.1-FP8", + "prefix": "glm-5-1-fp8-tp8-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}' --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Kimi-K2-Thinking-MXFP4", + "source_path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", + "path": "amd/Kimi-K2-Thinking-MXFP4-AttnFP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "Kimi-K2-Thinking-MXFP4 TP4 (MET)", + "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp4", + "prefix": "kimi-k2-thinking-mxfp4-tp4-met", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "Kimi-K2-Thinking-MXFP4 TP8 (MET)", + "dashboard_model": "Kimi-K2-Thinking-MXFP4-tp8", + "prefix": "kimi-k2-thinking-mxfp4-tp8-met", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Kimi-K2.5-MXFP4", + "source_path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "path": "amd/Kimi-K2.5-MXFP4-AttnFP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "Kimi-K2.5-MXFP4 TP4 (MET)", + "dashboard_model": "Kimi-K2.5-MXFP4-tp4", + "prefix": "kimi-k25-mxfp4-tp4-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 4, + "display": "Kimi-K2.5-MXFP4 TP4 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp4", + "prefix": "kimi-k2-5-mxfp4-aw-tp4", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "Kimi-K2.5-MXFP4 TP8 (AW)", + "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp8", + "prefix": "kimi-k2-5-mxfp4-aw-tp8", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "Qwen3.5-397B-A17B-FP8", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 4, + "display": "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", + "prefix": "qwen3-5-397b-a17b-fp8-tp4-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + }, + { + "tp_size": 8, + "display": "Qwen3.5-397B-A17B-FP8 TP8 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + } + ] + }, + { + "choice_label": "Qwen3.5-397B-A17B", + "source_path": "Qwen/Qwen3.5-397B-A17B", + "path": "Qwen/Qwen3.5-397B-A17B", + "bench_args": "", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 8, + "display": "Qwen3.5-397B-A17B TP8 (OOB)", + "dashboard_model": "Qwen3.5-397B-A17B", + "prefix": "qwen3-5-397b-a17b-oob", + "excluded_input_output_pairs": [ + "1024x8192" + ], + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + } + ] + }, + { + "choice_label": "Qwen3-Next-80B-A3B-Instruct-FP8", + "source_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "A", + "variants": [ + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4-met", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "tp_size": 2, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp2", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + } + ] + }, + { + "choice_label": "DeepSeek-V3.2 FP8", + "source_path": "deepseek-ai/DeepSeek-V3.2", + "path": "deepseek-ai/DeepSeek-V3.2", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "DeepSeek-V3.2 FP8 TP4 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", + "prefix": "deepseek-v3-2-fp8-aw-tp4", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "DeepSeek-V3.2 FP8 TP8 (AW)", + "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", + "prefix": "deepseek-v3-2-fp8-aw-tp8", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "GLM-4.7-FP8", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "GLM-4.7-FP8 TP4 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp4", + "prefix": "glm-4-7-fp8-aw-tp4", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "GLM-4.7-FP8 TP8 (AW)", + "dashboard_model": "GLM-4.7-FP8-aw-tp8", + "prefix": "glm-4-7-fp8-aw-tp8", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, + { + "choice_label": "MiniMax-M2.5", + "source_path": "MiniMaxAI/MiniMax-M2.5", + "path": "MiniMaxAI/MiniMax-M2.5", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 2, + "display": "MiniMax-M2.5 TP2 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp2", + "prefix": "minimax-m2-5-aw-tp2", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "tp_size": 4, + "display": "MiniMax-M2.5 TP4 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp4", + "prefix": "minimax-m2-5-aw-tp4", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + }, + { + "tp_size": 8, + "display": "MiniMax-M2.5 TP8 (AW)", + "dashboard_model": "MiniMax-M2.5-aw-tp8", + "prefix": "minimax-m2-5-aw-tp8", + "bench_args": "--random-range-ratio 1", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + } + ] + } + ] +} diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index 081500247f..dc67918590 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -46,6 +46,7 @@ RESULT_DIR=${RESULT_DIR:-/tmp/oot_accuracy_results} ACCURACY_LOG_FILE=${ACCURACY_LOG_FILE:-/tmp/oot_accuracy_output.txt} STREAM_VLLM_LOGS=${STREAM_VLLM_LOGS:-1} KEEP_SERVER_ALIVE_ON_EXIT=${KEEP_SERVER_ALIVE_ON_EXIT:-0} +OOT_GPU_MEMORY_UTILIZATION=${OOT_GPU_MEMORY_UTILIZATION:-0.9} EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index f539ee8255..c9efed1238 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -11,114 +11,293 @@ on: - cron: '0 15 * * 0-4' workflow_dispatch: inputs: - deepseek-r1-fp8-met: - description: "DeepSeek-R1 FP8 TP8 (MET)" - type: boolean - default: false - deepseek-r1-mxfp4-met: - description: "DeepSeek-R1 MXFP4 TP8 (MET)" - type: boolean - default: false - gpt-oss-120b-met: - description: "gpt-oss-120b TP1 (MET)" - type: boolean - default: false - glm-5-1-fp8-tp8-oob: - description: "GLM-5.1-FP8 TP8 (OOB)" - type: boolean - default: false - kimi-k2-thinking-mxfp4-tp4-met: - description: "Kimi-K2-Thinking-MXFP4 TP4 (MET)" - type: boolean - default: false - kimi-k2-thinking-mxfp4-tp8-met: - description: "Kimi-K2-Thinking-MXFP4 TP8 (MET)" - type: boolean - default: false - kimi-k25-mxfp4-tp4-met: - description: "Kimi-K2.5-MXFP4 TP4 (MET)" - type: boolean - default: false - qwen3-5-397b-a17b-oob: - description: "Qwen3.5-397B-A17B TP8 (OOB)" - type: boolean - default: false - qwen3-5-397b-a17b-fp8-tp4-oob: - description: "Qwen3.5-397B-A17B-FP8 TP4 (OOB)" - type: boolean - default: false - qwen3-5-397b-a17b-fp8-oob: - description: "Qwen3.5-397B-A17B-FP8 TP8 (OOB)" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-tp1-met: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-tp4-met: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-aw-tp1: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-aw-tp2: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)" - type: boolean - default: false - qwen3-next-80b-a3b-instruct-fp8-aw-tp4: - description: "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)" - type: boolean - default: false - deepseek-v3-2-fp8-aw-tp4: - description: "DeepSeek-V3.2 FP8 TP4 (AW)" - type: boolean - default: false - deepseek-v3-2-fp8-aw-tp8: - description: "DeepSeek-V3.2 FP8 TP8 (AW)" - type: boolean - default: false - glm-4-7-fp8-aw-tp4: - description: "GLM-4.7-FP8 TP4 (AW)" - type: boolean - default: false - glm-4-7-fp8-aw-tp8: - description: "GLM-4.7-FP8 TP8 (AW)" - type: boolean - default: false - gpt-oss-120b-aw-tp1: - description: "gpt-oss-120b TP1 (AW)" - type: boolean - default: false - gpt-oss-120b-aw-tp2: - description: "gpt-oss-120b TP2 (AW)" - type: boolean - default: false - gpt-oss-120b-aw-tp8: - description: "gpt-oss-120b TP8 (AW)" - type: boolean - default: false - kimi-k2-5-mxfp4-aw-tp4: - description: "Kimi-K2.5-MXFP4 TP4 (AW)" - type: boolean - default: false - kimi-k2-5-mxfp4-aw-tp8: - description: "Kimi-K2.5-MXFP4 TP8 (AW)" - type: boolean - default: false - minimax-m2-5-aw-tp2: - description: "MiniMax-M2.5 TP2 (AW)" - type: boolean - default: false - minimax-m2-5-aw-tp4: - description: "MiniMax-M2.5 TP4 (AW)" - type: boolean - default: false - minimax-m2-5-aw-tp8: - description: "MiniMax-M2.5 TP8 (AW)" - type: boolean - default: false + benchmark_client: + description: "Benchmark client command. `InferenceMax bench` uses the repo benchmark script. `vLLM bench` uses the vLLM CLI benchmark command." + type: choice + default: InferenceMax bench + options: + - InferenceMax bench + - vLLM bench + model_slot_1: + description: "Manual selection slot 1: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_2: + description: "Manual selection slot 2: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_3: + description: "Manual selection slot 3: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_4: + description: "Manual selection slot 4: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_5: + description: "Manual selection slot 5: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_6: + description: "Manual selection slot 6: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_7: + description: "Manual selection slot 7: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + model_slot_8: + description: "Manual selection slot 8: choose a model + TP size candidate." + type: choice + default: none + options: + - none + - DeepSeek-R1 FP8 TP8 (MET) + - DeepSeek-R1 MXFP4 TP8 (MET) + - gpt-oss-120b TP1 (MET) + - GLM-5.1-FP8 TP8 (OOB) + - Kimi-K2-Thinking-MXFP4 TP4 (MET) + - Kimi-K2-Thinking-MXFP4 TP8 (MET) + - Kimi-K2.5-MXFP4 TP4 (MET) + - Qwen3.5-397B-A17B-FP8 TP4 (OOB) + - Qwen3.5-397B-A17B-FP8 TP8 (OOB) + - Qwen3.5-397B-A17B TP8 (OOB) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP4 (AW) + - DeepSeek-V3.2 FP8 TP8 (AW) + - GLM-4.7-FP8 TP4 (AW) + - GLM-4.7-FP8 TP8 (AW) + - gpt-oss-120b TP1 (AW) + - gpt-oss-120b TP2 (AW) + - gpt-oss-120b TP8 (AW) + - Kimi-K2.5-MXFP4 TP4 (AW) + - Kimi-K2.5-MXFP4 TP8 (AW) + - MiniMax-M2.5 TP2 (AW) + - MiniMax-M2.5 TP4 (AW) + - MiniMax-M2.5 TP8 (AW) + isl_osl_pairs: + description: | + Benchmark ISL/OSL pairs. + Input as semicolon-separated `input_length,output_length` pairs. + Example: 1024,1024;1024,8192;8192,1024 + Additional copy-ready example: 1000,100;5000,500;10000,1000 + type: string + default: "1024,1024;1024,8192;8192,1024" + concurrency_values: + description: "Benchmark concurrency values. Use comma-separated values from 4,8,16,32,64,128,256,512. The workflow will benchmark exactly the values you provide." + type: string + default: "4,8,16,32,64,128,256,512" + random_range_ratios: + description: "Benchmark random range ratios. Use comma-separated values like 0.8,1.0." + type: string + default: "0.8" oot_image: description: "Optional OOT benchmark image override. Leave empty to resolve main to the same-digest published nightly tag behind vllm-latest, or rebuild from the selected non-main branch." type: string @@ -127,15 +306,10 @@ on: description: "Upload benchmark results to the dashboard" type: boolean default: true - param_lists: - description: | - "Benchmark parameter lists. - Input as a single or multiple sets (comma-separated, semicolon between sets), - format: input_length,output_length,concurrency,random_range_ratio. - Example (single set): 1024,1024,64,0.8 - Example (multiple sets): 1024,1024,64,0.8;2048,1024,32,0.7" - type: string - default: "1024,1024,4,0.8;1024,1024,8,0.8;1024,1024,16,0.8;1024,1024,32,0.8;1024,1024,64,0.8;1024,8192,4,0.8;1024,8192,8,0.8;1024,8192,16,0.8;1024,8192,32,0.8;1024,8192,64,0.8;8192,1024,4,0.8;8192,1024,8,0.8;8192,1024,16,0.8;8192,1024,32,0.8;8192,1024,64,0.8" + upload_to_custom_dashboard: + description: "Upload benchmark results to the LLM-Booster dashboard" + type: boolean + default: true jobs: wait-for-older-manual-runs: @@ -267,6 +441,8 @@ jobs: oot_base_image: ${{ steps.resolve.outputs.oot_base_image }} oot_image_source: ${{ steps.resolve.outputs.oot_image_source }} publish_to_dashboard: ${{ steps.resolve.outputs.publish_to_dashboard }} + upload_to_custom_dashboard: ${{ steps.resolve.outputs.upload_to_custom_dashboard }} + benchmark_client: ${{ steps.resolve.outputs.benchmark_client }} steps: - name: Checkout selected branch uses: actions/checkout@v6 @@ -276,6 +452,8 @@ jobs: env: INPUT_OOT_IMAGE: ${{ inputs.oot_image || '' }} INPUT_PUBLISH_TO_DASHBOARD: ${{ inputs.publish_to_dashboard }} + INPUT_UPLOAD_TO_CUSTOM_DASHBOARD: ${{ inputs.upload_to_custom_dashboard }} + INPUT_BENCHMARK_CLIENT: ${{ inputs.benchmark_client || 'InferenceMax bench' }} run: | set -euo pipefail @@ -297,10 +475,14 @@ jobs: if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then PUBLISH_TO_DASHBOARD=true + UPLOAD_TO_CUSTOM_DASHBOARD=true else PUBLISH_TO_DASHBOARD="${INPUT_PUBLISH_TO_DASHBOARD:-true}" + UPLOAD_TO_CUSTOM_DASHBOARD="${INPUT_UPLOAD_TO_CUSTOM_DASHBOARD:-true}" fi + BENCHMARK_CLIENT="${INPUT_BENCHMARK_CLIENT:-InferenceMax bench}" + mapfile -t VLLM_META < <(python3 - <<'PY' import re from pathlib import Path @@ -372,6 +554,8 @@ jobs: echo "oot_base_image=${OOT_BASE_IMAGE}" echo "oot_image_source=${OOT_IMAGE_SOURCE}" echo "publish_to_dashboard=${PUBLISH_TO_DASHBOARD}" + echo "upload_to_custom_dashboard=${UPLOAD_TO_CUSTOM_DASHBOARD}" + echo "benchmark_client=${BENCHMARK_CLIENT}" } >> "$GITHUB_OUTPUT" printf '### OOT benchmark source\n- Repository: `%s`\n- Ref: `%s`\n- Image mode: `%s`\n' \ @@ -426,53 +610,140 @@ jobs: esac fi - printf -- '- Upload to dashboard: `%s`\n' \ - "${PUBLISH_TO_DASHBOARD}" >> "$GITHUB_STEP_SUMMARY" + printf -- '- Benchmark client: `%s`\n- Upload to dashboard: `%s`\n- Upload to custom dashboard: `%s`\n' \ + "${BENCHMARK_CLIENT}" \ + "${PUBLISH_TO_DASHBOARD}" \ + "${UPLOAD_TO_CUSTOM_DASHBOARD}" >> "$GITHUB_STEP_SUMMARY" if [[ "${PUBLISH_TO_DASHBOARD}" != "true" ]]; then printf '\nDashboard upload is disabled by workflow input, so results are kept in artifacts and run summaries only.\n' \ >> "$GITHUB_STEP_SUMMARY" fi + if [[ "${UPLOAD_TO_CUSTOM_DASHBOARD}" != "true" ]]; then + printf 'Custom dashboard payload upload is disabled by workflow input.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + parse-param-lists: - name: Parse parameter lists + name: Parse benchmark parameter inputs needs: [wait-for-older-manual-runs] runs-on: ubuntu-latest outputs: - matrix_json: ${{ steps.parse-param-lists.outputs.matrix_json }} - param_lists: ${{ steps.parse-param-lists.outputs.param_lists }} + matrix_json: ${{ steps.parse-benchmark-params.outputs.matrix_json }} + isl_osl_pairs: ${{ steps.parse-benchmark-params.outputs.isl_osl_pairs }} + concurrency_values: ${{ steps.parse-benchmark-params.outputs.concurrency_values }} + random_range_ratios: ${{ steps.parse-benchmark-params.outputs.random_range_ratios }} steps: - - name: Parse parameter lists - id: parse-param-lists + - name: Parse benchmark parameter inputs + id: parse-benchmark-params + env: + ISL_OSL_PAIRS: ${{ inputs.isl_osl_pairs || '1024,1024;1024,8192;8192,1024' }} + CONCURRENCY_VALUES: ${{ inputs.concurrency_values || '4,8,16,32,64,128,256,512' }} + RANDOM_RANGE_RATIOS: ${{ inputs.random_range_ratios || '0.8' }} run: | - PARAM_LISTS="${{ inputs.param_lists || '1024,1024,4,0.8;1024,1024,8,0.8;1024,1024,16,0.8;1024,1024,32,0.8;1024,1024,64,0.8;1024,8192,4,0.8;1024,8192,8,0.8;1024,8192,16,0.8;1024,8192,32,0.8;1024,8192,64,0.8;8192,1024,4,0.8;8192,1024,8,0.8;8192,1024,16,0.8;8192,1024,32,0.8;8192,1024,64,0.8' }}" - echo "Using param_lists: ${PARAM_LISTS}" - printf 'param_lists=%s\n' "${PARAM_LISTS}" >> "$GITHUB_OUTPUT" - IFS=';' read -ra SETS <<< "${PARAM_LISTS}" - MATRIX_JSON="[" - SEP="" - for SET in "${SETS[@]}"; do - IFS=',' read -ra PARAMS <<< "$SET" - INPUT_LEN="${PARAMS[0]}" - OUTPUT_LEN="${PARAMS[1]}" - CONCURRENCY="${PARAMS[2]}" - RANDOM_RANGE_RATIO="${PARAMS[3]}" - case "${CONCURRENCY}" in - 4|8|16|32|64|128|256|512) ;; - *) - echo "Unsupported concurrency: ${CONCURRENCY}. Allowed values: 4,8,16,32,64,128,256,512" - exit 1 - ;; - esac - MATRIX_JSON="${MATRIX_JSON}${SEP}{\"input_length\":${INPUT_LEN},\"output_length\":${OUTPUT_LEN},\"concurrency\":${CONCURRENCY},\"random_range_ratio\":${RANDOM_RANGE_RATIO}}" - SEP="," - done - MATRIX_JSON="${MATRIX_JSON}]" - echo "matrix_json=${MATRIX_JSON}" >> "$GITHUB_OUTPUT" + set -euo pipefail + echo "Using isl_osl_pairs: ${ISL_OSL_PAIRS}" + echo "Using concurrency_values: ${CONCURRENCY_VALUES}" + echo "Using random_range_ratios: ${RANDOM_RANGE_RATIOS:-}" + printf 'isl_osl_pairs=%s\n' "${ISL_OSL_PAIRS}" >> "$GITHUB_OUTPUT" + printf 'concurrency_values=%s\n' "${CONCURRENCY_VALUES}" >> "$GITHUB_OUTPUT" + printf 'random_range_ratios=%s\n' "${RANDOM_RANGE_RATIOS}" >> "$GITHUB_OUTPUT" + python3 - <<'PY' >> "$GITHUB_OUTPUT" + import json + import os + + allowed_concurrencies = {4, 8, 16, 32, 64, 128, 256, 512} + + def dedupe_preserve_order(values): + deduped = [] + seen = set() + for value in values: + marker = ( + json.dumps(value, sort_keys=True) if isinstance(value, dict) else value + ) + if marker in seen: + continue + seen.add(marker) + deduped.append(value) + return deduped + + pairs = [] + for raw_pair in os.environ["ISL_OSL_PAIRS"].split(";"): + pair = raw_pair.strip() + if not pair: + continue + parts = [part.strip() for part in pair.split(",")] + if len(parts) != 2 or not all(part.isdigit() for part in parts): + raise SystemExit( + f"Invalid ISL/OSL pair {pair!r}. Use semicolon-separated pairs like 1024,1024;2048,1024." + ) + pairs.append( + { + "input_length": int(parts[0]), + "output_length": int(parts[1]), + } + ) + if not pairs: + raise SystemExit("At least one ISL/OSL pair is required.") + pairs = dedupe_preserve_order(pairs) + + concurrencies = [] + for raw_value in os.environ["CONCURRENCY_VALUES"].split(","): + value = raw_value.strip() + if not value: + continue + if not value.isdigit(): + raise SystemExit( + f"Invalid concurrency {value!r}. Allowed values: 4,8,16,32,64,128,256,512" + ) + parsed = int(value) + if parsed not in allowed_concurrencies: + raise SystemExit( + f"Unsupported concurrency: {parsed}. Allowed values: 4,8,16,32,64,128,256,512" + ) + concurrencies.append(parsed) + if not concurrencies: + raise SystemExit("At least one concurrency value is required.") + concurrencies = dedupe_preserve_order(concurrencies) + + ratio_text = os.environ["RANDOM_RANGE_RATIOS"].strip() + if not ratio_text: + ratios = ["0.8"] + else: + ratios = [] + for raw_ratio in ratio_text.split(","): + ratio = raw_ratio.strip() + if not ratio: + continue + try: + float(ratio) + except ValueError as exc: + raise SystemExit( + f"Invalid random range ratio {ratio!r}. Use values like 0.8 or 1.0." + ) from exc + ratios.append(ratio) + ratios = dedupe_preserve_order(ratios) or ["0.8"] + + matrix = [] + for pair in pairs: + for concurrency in concurrencies: + for ratio in ratios: + matrix.append( + { + "input_length": pair["input_length"], + "output_length": pair["output_length"], + "concurrency": concurrency, + "random_range_ratio": ratio, + } + ) - - name: Print matrix JSON + print(f"matrix_json={json.dumps(matrix, separators=(',', ':'))}") + PY + + - name: Print benchmark parameter matrix JSON run: | - echo "Matrix JSON: ${{ steps.parse-param-lists.outputs.matrix_json }}" + echo "Matrix JSON: ${{ steps.parse-benchmark-params.outputs.matrix_json }}" load-models: name: Load OOT model configs @@ -489,141 +760,170 @@ jobs: - uses: actions/checkout@v6 - id: load env: - ENABLE_DEEPSEEK_R1_FP8: ${{ inputs.deepseek-r1-fp8-met }} - ENABLE_DEEPSEEK_R1_MXFP4: ${{ inputs.deepseek-r1-mxfp4-met }} - ENABLE_GPT_OSS_120B: ${{ inputs.gpt-oss-120b-met }} - ENABLE_GLM_5_1_FP8_TP8: ${{ inputs.glm-5-1-fp8-tp8-oob }} - ENABLE_KIMI_K2_TP4: ${{ inputs.kimi-k2-thinking-mxfp4-tp4-met }} - ENABLE_KIMI_K2_TP8: ${{ inputs.kimi-k2-thinking-mxfp4-tp8-met }} - ENABLE_KIMI_K25_TP4: ${{ inputs.kimi-k25-mxfp4-tp4-met }} - ENABLE_QWEN3_5_397B_A17B: ${{ inputs.qwen3-5-397b-a17b-oob }} - ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4-oob }} - ENABLE_QWEN3_5_397B_A17B_FP8: ${{ inputs.qwen3-5-397b-a17b-fp8-oob }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1-met }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4-met }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP1: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp1 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP2: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp2 }} - ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP4: ${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp4 }} - ENABLE_DEEPSEEK_V3_2_FP8_AW_TP4: ${{ inputs.deepseek-v3-2-fp8-aw-tp4 }} - ENABLE_DEEPSEEK_V3_2_FP8_AW_TP8: ${{ inputs.deepseek-v3-2-fp8-aw-tp8 }} - ENABLE_GLM_4_7_FP8_AW_TP4: ${{ inputs.glm-4-7-fp8-aw-tp4 }} - ENABLE_GLM_4_7_FP8_AW_TP8: ${{ inputs.glm-4-7-fp8-aw-tp8 }} - ENABLE_GPT_OSS_120B_AW_TP1: ${{ inputs.gpt-oss-120b-aw-tp1 }} - ENABLE_GPT_OSS_120B_AW_TP2: ${{ inputs.gpt-oss-120b-aw-tp2 }} - ENABLE_GPT_OSS_120B_AW_TP8: ${{ inputs.gpt-oss-120b-aw-tp8 }} - ENABLE_KIMI_K2_5_MXFP4_AW_TP4: ${{ inputs.kimi-k2-5-mxfp4-aw-tp4 }} - ENABLE_KIMI_K2_5_MXFP4_AW_TP8: ${{ inputs.kimi-k2-5-mxfp4-aw-tp8 }} - ENABLE_MINIMAX_M2_5_AW_TP2: ${{ inputs.minimax-m2-5-aw-tp2 }} - ENABLE_MINIMAX_M2_5_AW_TP4: ${{ inputs.minimax-m2-5-aw-tp4 }} - ENABLE_MINIMAX_M2_5_AW_TP8: ${{ inputs.minimax-m2-5-aw-tp8 }} + MODEL_SLOT_1: ${{ inputs.model_slot_1 || 'none' }} + MODEL_SLOT_2: ${{ inputs.model_slot_2 || 'none' }} + MODEL_SLOT_3: ${{ inputs.model_slot_3 || 'none' }} + MODEL_SLOT_4: ${{ inputs.model_slot_4 || 'none' }} + MODEL_SLOT_5: ${{ inputs.model_slot_5 || 'none' }} + MODEL_SLOT_6: ${{ inputs.model_slot_6 || 'none' }} + MODEL_SLOT_7: ${{ inputs.model_slot_7 || 'none' }} + MODEL_SLOT_8: ${{ inputs.model_slot_8 || 'none' }} GITHUB_TOKEN: ${{ github.token }} run: | + set -euo pipefail + + if [[ "${GITHUB_EVENT_NAME}" == "schedule" ]]; then + SCHEDULE_CREATED_AT="$( + python3 - <<'PY' + import json + import os + from urllib.error import HTTPError, URLError + from urllib.request import Request, urlopen + + api_url = os.environ.get("GITHUB_API_URL", "https://api.github.com") + repository = os.environ["GITHUB_REPOSITORY"] + run_id = os.environ["GITHUB_RUN_ID"] + token = os.environ.get("GITHUB_TOKEN", "") + request = Request( + f"{api_url}/repos/{repository}/actions/runs/{run_id}", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "User-Agent": "atom-vllm-benchmark-nightly-rotation", + }, + ) + try: + with urlopen(request, timeout=30) as response: + payload = json.load(response) + except HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace").strip() + raise RuntimeError( + f"Failed to fetch workflow run metadata: HTTP {exc.code} {detail}" + ) from exc + except URLError as exc: + raise RuntimeError( + f"Failed to reach GitHub API for workflow run metadata: {exc}" + ) from exc + + created_at = payload.get("created_at") + if not created_at: + raise RuntimeError("Workflow run metadata did not include created_at") + print(created_at) + PY + )" + export SCHEDULE_CREATED_AT + else + export SCHEDULE_CREATED_AT="" + fi + python3 - <<'PY' >> "$GITHUB_OUTPUT" import json import os import sys from datetime import datetime, timedelta, timezone from pathlib import Path - from urllib.error import HTTPError, URLError - from urllib.request import Request, urlopen - models = json.loads( - Path(".github/benchmark/oot_benchmark_models.json").read_text(encoding="utf-8") - ) - event = os.environ["GITHUB_EVENT_NAME"] + none_choice = "none" - manual_toggles = { - "deepseek-r1-fp8-met": os.environ.get("ENABLE_DEEPSEEK_R1_FP8", "").lower() == "true", - "deepseek-r1-mxfp4-met": os.environ.get("ENABLE_DEEPSEEK_R1_MXFP4", "").lower() == "true", - "gpt-oss-120b-met": os.environ.get("ENABLE_GPT_OSS_120B", "").lower() == "true", - "glm-5-1-fp8-tp8-oob": os.environ.get("ENABLE_GLM_5_1_FP8_TP8", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp4-met": os.environ.get("ENABLE_KIMI_K2_TP4", "").lower() == "true", - "kimi-k2-thinking-mxfp4-tp8-met": os.environ.get("ENABLE_KIMI_K2_TP8", "").lower() == "true", - "kimi-k25-mxfp4-tp4-met": os.environ.get("ENABLE_KIMI_K25_TP4", "").lower() == "true", - "qwen3-5-397b-a17b-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B", "").lower() == "true", - "qwen3-5-397b-a17b-fp8-tp4-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", "").lower() == "true", - "qwen3-5-397b-a17b-fp8-oob": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp1-met": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP1", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-tp4-met": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_TP4", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-aw-tp1": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP1", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-aw-tp2": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP2", "").lower() == "true", - "qwen3-next-80b-a3b-instruct-fp8-aw-tp4": os.environ.get("ENABLE_QWEN3_NEXT_80B_A3B_INSTRUCT_FP8_AW_TP4", "").lower() == "true", - "deepseek-v3-2-fp8-aw-tp4": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8_AW_TP4", "").lower() == "true", - "deepseek-v3-2-fp8-aw-tp8": os.environ.get("ENABLE_DEEPSEEK_V3_2_FP8_AW_TP8", "").lower() == "true", - "glm-4-7-fp8-aw-tp4": os.environ.get("ENABLE_GLM_4_7_FP8_AW_TP4", "").lower() == "true", - "glm-4-7-fp8-aw-tp8": os.environ.get("ENABLE_GLM_4_7_FP8_AW_TP8", "").lower() == "true", - "gpt-oss-120b-aw-tp1": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP1", "").lower() == "true", - "gpt-oss-120b-aw-tp2": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP2", "").lower() == "true", - "gpt-oss-120b-aw-tp8": os.environ.get("ENABLE_GPT_OSS_120B_AW_TP8", "").lower() == "true", - "kimi-k2-5-mxfp4-aw-tp4": os.environ.get("ENABLE_KIMI_K2_5_MXFP4_AW_TP4", "").lower() == "true", - "kimi-k2-5-mxfp4-aw-tp8": os.environ.get("ENABLE_KIMI_K2_5_MXFP4_AW_TP8", "").lower() == "true", - "minimax-m2-5-aw-tp2": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP2", "").lower() == "true", - "minimax-m2-5-aw-tp4": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP4", "").lower() == "true", - "minimax-m2-5-aw-tp8": os.environ.get("ENABLE_MINIMAX_M2_5_AW_TP8", "").lower() == "true", - } - - def fetch_run_created_at() -> datetime: - api_url = os.environ.get("GITHUB_API_URL", "https://api.github.com") - repository = os.environ["GITHUB_REPOSITORY"] - run_id = os.environ["GITHUB_RUN_ID"] - token = os.environ.get("GITHUB_TOKEN", "") - request = Request( - f"{api_url}/repos/{repository}/actions/runs/{run_id}", - headers={ - "Accept": "application/vnd.github+json", - "Authorization": f"Bearer {token}", - "User-Agent": "atom-vllm-benchmark-nightly-rotation", - }, + catalog = json.loads( + Path(".github/benchmark/oot_benchmark_models.json").read_text( + encoding="utf-8" ) - try: - with urlopen(request, timeout=30) as response: - payload = json.load(response) - except HTTPError as exc: - detail = exc.read().decode("utf-8", errors="replace").strip() - raise RuntimeError( - f"Failed to fetch workflow run metadata: HTTP {exc.code} {detail}" - ) from exc - except URLError as exc: - raise RuntimeError( - f"Failed to reach GitHub API for workflow run metadata: {exc}" - ) from exc - - created_at = payload.get("created_at") - if not created_at: - raise RuntimeError("Workflow run metadata did not include created_at") - return datetime.fromisoformat(created_at.replace("Z", "+00:00")) + ) + families = catalog.get("families") + if not isinstance(families, list): + raise SystemExit("Benchmark catalog must define a 'families' list.") + + def flatten_variant(family, variant): + flattened = { + key: value + for key, value in family.items() + if key not in {"choice_label", "variants"} + } + flattened.update(variant) + return flattened selected_group = "" + event = os.environ["GITHUB_EVENT_NAME"] + selected = [] + + all_variants = [] + for family in families: + for variant in family["variants"]: + all_variants.append(flatten_variant(family, variant)) + if event == "schedule": + created_at = os.environ.get("SCHEDULE_CREATED_AT", "") + if not created_at: + raise SystemExit( + "Scheduled selection requires a run creation timestamp." + ) beijing_tz = timezone(timedelta(hours=8)) - beijing_dt = fetch_run_created_at().astimezone(beijing_tz) - beijing_weekday = beijing_dt.weekday() # Monday=0 ... Sunday=6 + run_date = datetime.fromisoformat(created_at.replace("Z", "+00:00")) + beijing_dt = run_date.astimezone(beijing_tz) + beijing_weekday = beijing_dt.weekday() - if beijing_weekday in (0, 2): # Monday / Wednesday + if beijing_weekday in (0, 2): selected_group = "A-MET" selected = [ - model - for model in models - if str(model.get("prefix", "")).endswith("-met") + m for m in all_variants + if str(m.get("prefix", "")).endswith("-met") ] - elif beijing_weekday in (1, 3): # Tuesday / Thursday + elif beijing_weekday in (1, 3): selected_group = "B-AW" selected = [ - model for model in models if "-aw-" in str(model.get("prefix", "")) + m for m in all_variants + if "-aw-" in str(m.get("prefix", "")) ] - elif beijing_weekday == 4: # Friday + elif beijing_weekday == 4: selected_group = "C-ALL" - selected = list(models) - else: # Saturday / Sunday + selected = list(all_variants) + else: selected_group = "SKIP-WEEKEND" selected = [] print( - f"Weekend schedule in Beijing ({beijing_dt.date()}); no benchmark group configured." + f"Weekend schedule in Beijing ({beijing_dt.date()}); no benchmark group configured.", + file=sys.stderr, ) else: - selected = [ - model for model in models if manual_toggles.get(model["prefix"], False) - ] + variant_map = {} + for family in families: + for variant in family["variants"]: + variant_map[str(variant["display"])] = flatten_variant( + family, variant + ) + + seen_prefixes = set() + for slot_idx in range(1, 9): + label = os.environ.get(f"MODEL_SLOT_{slot_idx}", "").strip() + if not label or label == none_choice: + continue + selected_variant = variant_map.get(label) + if selected_variant is None: + available = ", ".join(sorted(variant_map)) + raise SystemExit( + f"Unknown benchmark model variant choice {label!r}. Available choices: {available}" + ) + prefix = str(selected_variant["prefix"]) + if prefix in seen_prefixes: + continue + selected.append(selected_variant) + seen_prefixes.add(prefix) + + if selected_group: + print( + f"Scheduled nightly benchmark group: {selected_group}", + file=sys.stderr, + ) + if not selected: + print("No benchmark model variants were selected.", file=sys.stderr) + else: + print("Selected benchmark model variants:", file=sys.stderr) + for model in selected: + print( + f" - {model['display']} ({model['prefix']})", + file=sys.stderr, + ) print(f"models_json={json.dumps(selected, separators=(',', ':'))}") print(f"selected_group={selected_group}") @@ -675,7 +975,6 @@ jobs: excluded_pairs = set(model.get("excluded_input_output_pairs", [])) model_params = [] seen = set() - # Run all benchmark cases at the two highest supported concurrency levels too. extra_concurrency = (128, 256) for param in params: @@ -686,6 +985,7 @@ jobs: aws_param = dict(param) aws_param["input_length"] = input_length aws_param["output_length"] = output_length + aws_param["random_range_ratio"] = "1" base_params.append(aws_param) variants = [] @@ -857,7 +1157,6 @@ jobs: MODEL_SOURCE_PATH: ${{ matrix.model.source_path || matrix.model.path }} MODEL_PATH: ${{ matrix.model.path || matrix.model.source_path }} OOT_EXTRA_ARGS: ${{ matrix.model.extra_args }} - BENCH_CLIENT: ${{ contains(matrix.model.prefix, '-aw-') && 'vllm-bench' || 'benchmark_serving.py' }} BENCH_EXTRA_ARGS: ${{ matrix.model.bench_args }} RESULT_PREFIX: ${{ matrix.model.prefix }} ISL: ${{ matrix.params.input_length }} @@ -876,6 +1175,8 @@ jobs: VLLM_COMMIT_USED: ${{ needs.resolve-atom-source.outputs.selected_vllm_commit }} VLLM_VERSION_USED: ${{ needs.resolve-atom-source.outputs.selected_vllm_version }} PUBLISH_TO_DASHBOARD: ${{ needs.resolve-atom-source.outputs.publish_to_dashboard }} + UPLOAD_TO_CUSTOM_DASHBOARD: ${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }} + BENCHMARK_CLIENT: ${{ contains(matrix.model.prefix, '-aw-') && 'vLLM bench' || needs.resolve-atom-source.outputs.benchmark_client }} steps: - name: Detect container engine run: | @@ -890,47 +1191,7 @@ jobs: exit 1 fi - - name: Check if model is enabled - id: check - run: | - if [ "${GITHUB_EVENT_NAME}" = "schedule" ]; then - echo "enabled=true" >> "$GITHUB_OUTPUT" - exit 0 - fi - - case "${{ matrix.model.prefix }}" in - deepseek-r1-fp8-met) echo "enabled=${{ inputs.deepseek-r1-fp8-met }}" >> "$GITHUB_OUTPUT" ;; - deepseek-r1-mxfp4-met) echo "enabled=${{ inputs.deepseek-r1-mxfp4-met }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b-met) echo "enabled=${{ inputs.gpt-oss-120b-met }}" >> "$GITHUB_OUTPUT" ;; - glm-5-1-fp8-tp8-oob) echo "enabled=${{ inputs.glm-5-1-fp8-tp8-oob }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp4-met) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp4-met }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-thinking-mxfp4-tp8-met) echo "enabled=${{ inputs.kimi-k2-thinking-mxfp4-tp8-met }}" >> "$GITHUB_OUTPUT" ;; - kimi-k25-mxfp4-tp4-met) echo "enabled=${{ inputs.kimi-k25-mxfp4-tp4-met }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-oob }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8-tp4-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4-oob }}" >> "$GITHUB_OUTPUT" ;; - qwen3-5-397b-a17b-fp8-oob) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-oob }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp1-met) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp1-met }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-tp4-met) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-tp4-met }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-aw-tp1) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp1 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-aw-tp2) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; - qwen3-next-80b-a3b-instruct-fp8-aw-tp4) echo "enabled=${{ inputs.qwen3-next-80b-a3b-instruct-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-v3-2-fp8-aw-tp4) echo "enabled=${{ inputs.deepseek-v3-2-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; - deepseek-v3-2-fp8-aw-tp8) echo "enabled=${{ inputs.deepseek-v3-2-fp8-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; - glm-4-7-fp8-aw-tp4) echo "enabled=${{ inputs.glm-4-7-fp8-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; - glm-4-7-fp8-aw-tp8) echo "enabled=${{ inputs.glm-4-7-fp8-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b-aw-tp1) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp1 }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b-aw-tp2) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; - gpt-oss-120b-aw-tp8) echo "enabled=${{ inputs.gpt-oss-120b-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-5-mxfp4-aw-tp4) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; - kimi-k2-5-mxfp4-aw-tp8) echo "enabled=${{ inputs.kimi-k2-5-mxfp4-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; - minimax-m2-5-aw-tp2) echo "enabled=${{ inputs.minimax-m2-5-aw-tp2 }}" >> "$GITHUB_OUTPUT" ;; - minimax-m2-5-aw-tp4) echo "enabled=${{ inputs.minimax-m2-5-aw-tp4 }}" >> "$GITHUB_OUTPUT" ;; - minimax-m2-5-aw-tp8) echo "enabled=${{ inputs.minimax-m2-5-aw-tp8 }}" >> "$GITHUB_OUTPUT" ;; - *) echo "enabled=true" >> "$GITHUB_OUTPUT" ;; - esac - - name: Clean up containers and workspace - if: steps.check.outputs.enabled == 'true' run: | echo "=== Cleaning up containers on $(hostname) ===" containers=$($CONTAINER_ENGINE ps -q) @@ -941,7 +1202,6 @@ jobs: $CONTAINER_ENGINE run --rm -v "${GITHUB_WORKSPACE:-$PWD}":/workspace -w /workspace --privileged docker.io/rocm/pytorch:latest bash -lc "shopt -s dotglob && ls -la /workspace/ && rm -rf /workspace/*" || true - name: Checkout benchmark ATOM source - if: steps.check.outputs.enabled == 'true' uses: actions/checkout@v6 with: repository: ${{ env.ATOM_SOURCE_REPOSITORY }} @@ -949,14 +1209,12 @@ jobs: fetch-depth: 1 - name: Record benchmark source revision - if: steps.check.outputs.enabled == 'true' run: | SOURCE_SHA="$(git rev-parse HEAD)" echo "ATOM_SOURCE_SHA=${SOURCE_SHA}" >> "$GITHUB_ENV" echo "Benchmarking ${ATOM_SOURCE_REPOSITORY}@${ATOM_SOURCE_REF} (${SOURCE_SHA}) with ${OOT_IMAGE_SOURCE} image ${OOT_IMAGE_TAG}" - name: Container Engine Login - if: steps.check.outputs.enabled == 'true' run: | # Podman requires an explicit registry when registries.conf has no search list; # Docker Hub short names (e.g. rocm/foo) must use docker.io here. @@ -968,11 +1226,9 @@ jobs: echo "${{ secrets.DOCKER_PASSWORD }}" | $CONTAINER_ENGINE login "$REG" -u "${{ secrets.DOCKER_USERNAME }}" --password-stdin - name: Set HF_TOKEN - if: steps.check.outputs.enabled == 'true' run: echo "HF_TOKEN=${HF_TOKEN:-${{ secrets.AMD_HF_TOKEN }}}" >> "$GITHUB_ENV" - name: Pull OOT benchmark image - if: steps.check.outputs.enabled == 'true' run: | set -euo pipefail IMG="${OOT_IMAGE_TAG}" @@ -991,7 +1247,6 @@ jobs: $CONTAINER_ENGINE pull "${IMG}" - name: Prepare model cache mount - if: steps.check.outputs.enabled == 'true' run: | MODEL_CACHE_MOUNT="" MODEL_CACHE_DESC="container-local /models (no host cache mount)" @@ -1013,7 +1268,6 @@ jobs: echo "MODEL_CACHE_DESC=${MODEL_CACHE_DESC}" >> "$GITHUB_ENV" - name: Start OOT benchmark container - if: steps.check.outputs.enabled == 'true' run: | if [ -f "/etc/podinfo/gha-render-devices" ]; then DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) @@ -1050,12 +1304,10 @@ jobs: GITHUB_WORKSPACE: ${{ github.workspace }} - name: Collect GPU info (inside container) - if: steps.check.outputs.enabled == 'true' id: gpu-info run: bash .github/scripts/collect_gpu_info.sh "$CONTAINER_NAME" "$CONTAINER_ENGINE" "${{ matrix.model.runner }}" - name: Download model if needed - if: steps.check.outputs.enabled == 'true' run: | model_dir="${MODEL_PATH}" if [[ "${model_dir}" = /models/* ]]; then @@ -1078,7 +1330,6 @@ jobs: fi - name: Resolve OOT model path - if: steps.check.outputs.enabled == 'true' run: | model_dir="${MODEL_PATH}" if [[ "${model_dir}" = /models/* ]]; then @@ -1098,12 +1349,9 @@ jobs: echo "Using model id: ${MODEL_SOURCE_PATH}" fi - - name: Prepare OOT benchmark runner in container - if: steps.check.outputs.enabled == 'true' + - name: Prepare benchmark client in container run: | - if [[ "${BENCH_CLIENT}" == "vllm-bench" ]]; then - echo "Skipping bench_serving repo setup for vllm bench client." - else + if [[ "${BENCHMARK_CLIENT}" == "InferenceMax bench" ]]; then $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " set -euo pipefail rm -rf \"${CONTAINER_BENCH_SERVING_DIR%/bench_serving}\" @@ -1111,10 +1359,14 @@ jobs: git clone --depth 1 \"${BENCH_SERVING_REPO_URL}\" \"${CONTAINER_BENCH_SERVING_DIR}\" test -f \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" " + else + $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc " + set -euo pipefail + command -v vllm + " fi - name: Launch OOT benchmark server - if: steps.check.outputs.enabled == 'true' env: MAX_WAIT_RETRIES: "120" STREAM_VLLM_LOGS: "1" @@ -1134,19 +1386,22 @@ jobs: - name: Run OOT benchmark client id: run_benchmark_client - if: steps.check.outputs.enabled == 'true' timeout-minutes: 240 continue-on-error: true run: | set -euo pipefail TRUST_REMOTE_CODE_ARG="" if [[ "${OOT_EXTRA_ARGS}" == *"--trust-remote-code"* ]]; then - TRUST_REMOTE_CODE_ARG="--trust-remote-code" + if [[ "${BENCHMARK_CLIENT}" == "vLLM bench" ]]; then + TRUST_REMOTE_CODE_ARG="--trust_remote_code" + else + TRUST_REMOTE_CODE_ARG="--trust-remote-code" + fi fi { - echo "=== Benchmark config: ${MODEL_NAME} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} CLIENT=${BENCH_CLIENT} ===" - if [[ "${BENCH_CLIENT}" == "vllm-bench" ]]; then + echo "=== Benchmark config: ${MODEL_NAME} CLIENT=${BENCHMARK_CLIENT} ISL=${ISL} OSL=${OSL} CONC=${CONC} RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO} ===" + if [[ "${BENCHMARK_CLIENT}" == "InferenceMax bench" ]]; then $CONTAINER_ENGINE exec \ -e ISL="${ISL}" \ -e OSL="${OSL}" \ @@ -1158,26 +1413,24 @@ jobs: set -euo pipefail rm -rf \"${CONTAINER_RESULT_DIR}\" mkdir -p \"${CONTAINER_RESULT_DIR}\" - vllm bench serve \ - --backend vllm \ - --host 127.0.0.1 \ - --port 8000 \ - --model \"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ - --dataset-name random \ - --random-input-len \"${ISL}\" \ - --random-output-len \"${OSL}\" \ + PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ + --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --backend=vllm \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len=\"${ISL}\" \ + --random-output-len=\"${OSL}\" \ --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ - --max-concurrency \"${CONC}\" \ - --num-prompts \"$(( CONC * 10 ))\" \ - --num-warmups \"$(( 2 * CONC ))\" \ - --seed 42 \ - --ignore-eos \ + --num-prompts=\"$(( CONC * 10 ))\" \ + --max-concurrency=\"${CONC}\" \ ${TRUST_REMOTE_CODE_ARG} \ - --percentile-metrics ttft,tpot,itl,e2el \ - --ready-check-timeout-sec 7200 \ + --num-warmups=\"$(( 2 * CONC ))\" \ + --request-rate=inf \ + --ignore-eos \ --save-result \ - --result-dir \"${CONTAINER_RESULT_DIR}\" \ - --result-filename \"${RESULT_FILENAME}.json\" \ + --percentile-metrics=\"ttft,tpot,itl,e2el\" \ + --result-dir=\"${CONTAINER_RESULT_DIR}\" \ + --result-filename=\"${RESULT_FILENAME}.json\" \ ${BENCH_EXTRA_ARGS:-} " else @@ -1192,31 +1445,34 @@ jobs: set -euo pipefail rm -rf \"${CONTAINER_RESULT_DIR}\" mkdir -p \"${CONTAINER_RESULT_DIR}\" - PYTHONDONTWRITEBYTECODE=1 python \"${CONTAINER_BENCH_SERVING_DIR}/benchmark_serving.py\" \ - --model=\"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ - --backend=vllm \ - --base-url=http://127.0.0.1:8000 \ - --dataset-name=random \ - --random-input-len=\"${ISL}\" \ - --random-output-len=\"${OSL}\" \ + vllm bench serve \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --endpoint /v1/completions \ + --model \"${OOT_RESOLVED_MODEL_PATH:-$MODEL_PATH}\" \ + --dataset-name random \ + --random-input-len \"${ISL}\" \ + --random-output-len \"${OSL}\" \ --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ - --num-prompts=\"$(( CONC * 10 ))\" \ - --max-concurrency=\"${CONC}\" \ + --temperature 0.0 \ + --num-prompts \"$(( CONC * 10 ))\" \ + --max-concurrency \"${CONC}\" \ ${TRUST_REMOTE_CODE_ARG} \ - --num-warmups=\"$(( 2 * CONC ))\" \ - --request-rate=inf \ + --num-warmups \"$(( 2 * CONC ))\" \ + --request-rate inf \ --ignore-eos \ + --disable-tqdm \ --save-result \ - --percentile-metrics=\"ttft,tpot,itl,e2el\" \ - --result-dir=\"${CONTAINER_RESULT_DIR}\" \ - --result-filename=\"${RESULT_FILENAME}.json\" \ + --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir \"${CONTAINER_RESULT_DIR}\" \ + --result-filename \"${RESULT_FILENAME}.json\" \ ${BENCH_EXTRA_ARGS:-} " fi } > "./${RESULT_FILENAME}.client.log" 2>&1 - name: Copy OOT benchmark outputs - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail if $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc "test -f \"${CONTAINER_RESULT_DIR}/${RESULT_FILENAME}.json\"" >/dev/null 2>&1; then @@ -1227,7 +1483,7 @@ jobs: fi - name: Enrich copied benchmark result - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail if [ ! -f "${RESULT_FILENAME}.json" ]; then @@ -1247,6 +1503,8 @@ jobs: OOT_IMAGE_SOURCE="${OOT_IMAGE_SOURCE}" \ OOT_IMAGE_TAG_USED="${OOT_IMAGE_TAG}" \ PUBLISH_TO_DASHBOARD="${PUBLISH_TO_DASHBOARD}" \ + UPLOAD_TO_CUSTOM_DASHBOARD="${UPLOAD_TO_CUSTOM_DASHBOARD}" \ + BENCHMARK_CLIENT="${BENCHMARK_CLIENT}" \ python3 - <<'PY' import json import os @@ -1278,16 +1536,20 @@ jobs: data["vllm_version"] = os.environ.get("VLLM_VERSION_USED", "") data["oot_image_source"] = os.environ.get("OOT_IMAGE_SOURCE", "") data["oot_image_tag"] = os.environ.get("OOT_IMAGE_TAG_USED", "") + data["benchmark_client"] = os.environ.get("BENCHMARK_CLIENT", "") data["dashboard_publish_allowed"] = ( os.environ.get("PUBLISH_TO_DASHBOARD", "false").lower() == "true" ) + data["custom_dashboard_publish_allowed"] = ( + os.environ.get("UPLOAD_TO_CUSTOM_DASHBOARD", "false").lower() == "true" + ) with open(result_path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) PY - name: Dump OOT benchmark diagnostics - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | set -euo pipefail RESULT_PATH="${RESULT_FILENAME}.json" @@ -1342,7 +1604,6 @@ jobs: echo "::endgroup::" - name: Inject GPU metadata into benchmark results - if: steps.check.outputs.enabled == 'true' run: | shopt -s nullglob for f in "${RESULT_PREFIX}"-*.json; do @@ -1365,7 +1626,6 @@ jobs: - name: Validate benchmark result completeness id: validate_benchmark_result - if: steps.check.outputs.enabled == 'true' run: | set -euo pipefail RESULT_PATH="${RESULT_FILENAME}.json" @@ -1424,7 +1684,7 @@ jobs: PY - name: Upload benchmark artifacts - if: always() && steps.check.outputs.enabled == 'true' && steps.validate_benchmark_result.outcome == 'success' + if: always() && steps.validate_benchmark_result.outcome == 'success' uses: actions/upload-artifact@v7 with: name: oot-benchmark-${{ env.RESULT_FILENAME }} @@ -1435,7 +1695,7 @@ jobs: ${{ env.RESULT_FILENAME }}.vllm.log - name: Clean up OOT benchmark container - if: always() && steps.check.outputs.enabled == 'true' + if: always() run: | $CONTAINER_ENGINE exec "$CONTAINER_NAME" bash -lc "if [ -f /tmp/vllm_oot.pid ]; then kill \$(cat /tmp/vllm_oot.pid) || true; fi" || true $CONTAINER_ENGINE stop "$CONTAINER_NAME" || true @@ -1477,16 +1737,28 @@ jobs: - name: Note summary behavior run: | + printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table.\n- Benchmark client: `%s`\n- Upload to dashboard: `%s`\n- Upload to custom dashboard: `%s`\n' \ + "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ + "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ + "${{ needs.resolve-atom-source.outputs.oot_image_source }}" \ + "${{ needs.resolve-atom-source.outputs.benchmark_client }}" \ + "${{ needs.resolve-atom-source.outputs.publish_to_dashboard }}" \ + "${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }}" >> "$GITHUB_STEP_SUMMARY" + if [[ "${{ needs.resolve-atom-source.outputs.publish_to_dashboard }}" == "true" ]]; then - printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table. Because dashboard upload is enabled for this run, this final Ubuntu summary job collects benchmark artifacts from all completed GPU cases and publishes the available successful results in one pass, then reconciles the completed run for regression comparison.\n' \ - "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ - "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ - "${{ needs.resolve-atom-source.outputs.oot_image_source }}" >> "$GITHUB_STEP_SUMMARY" + printf '\nDashboard upload is enabled for this run, so this final Ubuntu summary job also reconciles the completed run for regression comparison and publishes dashboard data.\n' \ + >> "$GITHUB_STEP_SUMMARY" else - printf '### Summary mode\nResults from `%s@%s` using `%s` OOT image mode are shown below as a workflow summary table only. Dashboard upload is disabled for this run.\n' \ - "${{ needs.resolve-atom-source.outputs.atom_repository }}" \ - "${{ needs.resolve-atom-source.outputs.atom_ref }}" \ - "${{ needs.resolve-atom-source.outputs.oot_image_source }}" >> "$GITHUB_STEP_SUMMARY" + printf '\nDashboard upload is disabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + + if [[ "${{ needs.resolve-atom-source.outputs.upload_to_custom_dashboard }}" == "true" ]]; then + printf 'LLM-Booster dashboard publish is enabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" + else + printf 'LLM-Booster dashboard publish is disabled for this run.\n' \ + >> "$GITHUB_STEP_SUMMARY" fi - name: Generate OOT benchmark summary table @@ -1517,6 +1789,265 @@ jobs: name: oot-benchmark-summary path: oot_benchmark_summary.json + - name: Note skipped LLM-Booster dashboard publish + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases == '0' + run: | + printf '\nNo successful benchmark result JSONs were produced, so LLM-Booster dashboard publish was skipped.\n' \ + >> "$GITHUB_STEP_SUMMARY" + + - name: Check LLM-Booster deploy key + id: llmbooster-key + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' + env: + LLM_BOOSTER_DEPLOY_KEY: ${{ secrets.LLM_BOOSTER_DEPLOY_KEY }} + run: | + if [[ -n "${LLM_BOOSTER_DEPLOY_KEY}" ]]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + echo "::warning::LLM_BOOSTER_DEPLOY_KEY is not configured; skipping LLM-Booster dashboard publish." + printf '\nLLM-Booster dashboard publish skipped because secret `LLM_BOOSTER_DEPLOY_KEY` is not configured.\n' \ + >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Checkout LLM-Booster dashboard repo + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + uses: actions/checkout@v6 + with: + repository: ROCm/LLM-Booster + ref: main + path: llm-booster-repo + ssh-key: ${{ secrets.LLM_BOOSTER_DEPLOY_KEY }} + fetch-depth: 1 + + - name: Build LLM-Booster dashboard files + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + env: + RUN_URL: https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }} + ATOM_SOURCE_REPOSITORY: ${{ needs.resolve-atom-source.outputs.atom_repository }} + ATOM_SOURCE_REF: ${{ needs.resolve-atom-source.outputs.atom_ref }} + OOT_IMAGE_SOURCE: ${{ needs.resolve-atom-source.outputs.oot_image_source }} + BENCHMARK_CLIENT: ${{ needs.resolve-atom-source.outputs.benchmark_client }} + run: | + python3 - <<'PY' + import json + import os + import re + import sys + import time + from collections import defaultdict + from datetime import datetime, timezone + from pathlib import Path + + root = Path(".") + llm_booster_dashboard = root / "llm-booster-repo" / "dashboard" + llm_booster_data_dir = llm_booster_dashboard / "data" / "vllm-atom" + llm_booster_data_dir.mkdir(parents=True, exist_ok=True) + + sys.path.insert(0, str(llm_booster_dashboard)) + from update_data import index_root_for, rebuild_index # noqa: E402 + + catalog = json.loads( + (root / ".github" / "benchmark" / "oot_benchmark_models.json").read_text( + encoding="utf-8" + ) + ) + + def slugify(text: str) -> str: + return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-") or "unknown" + + def precision_from_model(model: str): + upper = model.upper() + if "MXFP4" in upper or "FP4" in upper: + return "fp4" + if "FP8" in upper: + return "fp8" + return None + + def model_name_for_variant(variant: dict) -> str: + if variant.get("dashboard_model"): + return str(variant["dashboard_model"]) + display = str(variant["display"]) + return display.replace(" TP", "-tp").replace(" ", "-") + + prefix_map = {} + for family in catalog["families"]: + base = {k: v for k, v in family.items() if k not in {"choice_label", "variants"}} + for variant in family["variants"]: + item = dict(base) + item.update(variant) + prefix_map[str(variant["prefix"])] = item + + def parse_result_filename(path: Path): + stem = path.stem + prefix, isl, osl, concurrency, ratio = stem.rsplit("-", 4) + return prefix, int(isl), int(osl), int(concurrency), float(ratio) + + default_benchmark_client = os.environ["BENCHMARK_CLIENT"] + + def resolve_client(payload): + bc = str(payload.get("benchmark_client", default_benchmark_client)) + slug = "vllmbench" if bc == "vLLM bench" else "inferencemax" + label = "vllm bench" if bc == "vLLM bench" else "inferencemax" + return slug, label + + grouped = defaultdict(list) + for path in sorted(root.glob("*.json")): + if path.name in {"oot_benchmark_summary.json", "regression_report.json"}: + continue + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + continue + if "output_throughput" not in payload: + continue + custom_dashboard_flag = str( + payload.get("custom_dashboard_publish_allowed", True) + ).strip().lower() + if custom_dashboard_flag in {"0", "false", "no"}: + continue + + prefix, isl, osl, concurrency, ratio = parse_result_filename(path) + variant = prefix_map.get(prefix) + if variant is None: + raise SystemExit(f"Unknown benchmark prefix in result file {path.name}: {prefix}") + + model = model_name_for_variant(variant) + precision = precision_from_model(model) + atom_source_sha = str(payload.get("atom_source_sha", "")) or "unknown" + sha7 = atom_source_sha[:7] if atom_source_sha else "unknown" + day = time.strftime("%Y-%m-%d", time.gmtime()) + client_slug, _ = resolve_client(payload) + grouped[(day, sha7, model, isl, osl, client_slug)].append((path, payload, precision, concurrency, ratio)) + + written_files = set() + for (day, sha7, model, isl, osl, client_slug), entries in sorted(grouped.items()): + entries.sort(key=lambda item: (item[3], item[4])) + client_label = "vllm bench" if client_slug == "vllmbench" else "inferencemax" + timestamp_ms = int(time.time() * 1000) + precision = entries[0][2] + first_payload = entries[0][1] + tp = first_payload.get("tensor_parallel_size") + tp_label = f"_tp{int(tp)}" if tp not in (None, "", 0) else "" + model_slug = slugify(model) + run_id = ( + f"atom_vllm_{day.replace('-', '')}_{model_slug}_isl{isl}_osl{osl}_" + f"{sha7}_{os.environ['GITHUB_RUN_ID']}_{client_slug}" + ) + config_label = f"atom-vllm_{model_slug}" + if precision: + config_label += f"_{precision}" + if tp not in (None, "", 0): + config_label += tp_label + if client_slug == "inferencemax": + config_label += "_inferencemax" + + points = [] + for _, payload, precision, concurrency, ratio in entries: + output_tput = payload.get("output_throughput") + total_tput = payload.get("total_token_throughput") + total_gpu = payload.get("tensor_parallel_size") + try: + total_gpu = int(total_gpu) if total_gpu is not None else None + except (TypeError, ValueError): + total_gpu = None + points.append( + { + "isl": isl, + "osl": osl, + "concurrency": concurrency, + "ratio": ratio, + "client_bench": client_label, + "ttft_ms": payload.get("mean_ttft_ms"), + "ttft_p99": payload.get("p99_ttft_ms"), + "tpot_ms": payload.get("mean_tpot_ms"), + "tpot_p99": payload.get("p99_tpot_ms"), + "itl_ms": payload.get("mean_itl_ms"), + "e2el_ms": payload.get("mean_e2el_ms"), + "interactivity": (output_tput / concurrency) if output_tput and concurrency else None, + "tput_per_gpu": (total_tput / total_gpu) if total_tput and total_gpu else None, + "output_tput_per_gpu": (output_tput / total_gpu) if output_tput and total_gpu else None, + "input_tput_per_gpu": None, + "output_tput": output_tput, + "total_tput": total_tput, + "req_tput": payload.get("request_throughput"), + "completed": payload.get("completed"), + "duration": payload.get("duration"), + "num_prompts": payload.get("num_prompts"), + "prefill_tp": None, + "decode_tp": None, + "tensor_parallel_size": total_gpu, + "num_prefill_gpu": None, + "num_decode_gpu": None, + "total_gpu": total_gpu, + "gpu_name": payload.get("gpu_name"), + "image": payload.get("oot_image_tag") or payload.get("docker_image"), + "rocm": payload.get("rocm_version"), + "run_url": os.environ["RUN_URL"], + "model_path": payload.get("model_id"), + "point_date": day, + "precision": precision, + "hardware": "mi355x", + } + ) + + run = { + "run_id": run_id, + "date": day, + "timestamp": timestamp_ms, + "model": model, + "backend": "vLLM-ATOM", + "config_label": config_label, + "precision": precision, + "source": "ATOM Dashboard", + "client_bench": client_label, + "hardware": "mi355x", + "commit": { + "id": first_payload.get("atom_source_sha"), + "message": f"ATOM benchmark run {os.environ['GITHUB_RUN_ID']}", + "timestamp": datetime.now(timezone.utc).isoformat(), + "url": f"https://github.com/{os.environ['ATOM_SOURCE_REPOSITORY']}/commit/{first_payload.get('atom_source_sha', '')}", + "author": {}, + "committer": {}, + }, + "points": points, + "gsm8k": None, + } + + run_file = llm_booster_data_dir / f"{run_id}.json" + run_file.write_text(json.dumps(run, indent=2) + "\n", encoding="utf-8") + written_files.add(run_file) + + rebuild_index(index_root_for(llm_booster_data_dir)) + print(f"Wrote {len(written_files)} LLM-Booster ATOM vLLM run files.") + PY + + - name: Upload LLM-Booster dashboard preview artifact + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + uses: actions/upload-artifact@v7 + with: + name: oot-benchmark-llm-booster-preview + path: | + llm-booster-repo/dashboard/data/index.json + llm-booster-repo/dashboard/data/vllm-atom/*.json + + - name: Push LLM-Booster dashboard data + if: needs.resolve-atom-source.outputs.upload_to_custom_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' && steps.llmbooster-key.outputs.available == 'true' + working-directory: llm-booster-repo + env: + GIT_AUTHOR_NAME: github-actions[bot] + GIT_AUTHOR_EMAIL: github-actions[bot]@users.noreply.github.com + GIT_COMMITTER_NAME: github-actions[bot] + GIT_COMMITTER_EMAIL: github-actions[bot]@users.noreply.github.com + run: | + git add dashboard/data/index.json dashboard/data/vllm-atom + if git diff --cached --quiet; then + echo "No LLM-Booster dashboard data changes to push." + exit 0 + fi + git commit -m "Update ATOM vLLM benchmark dashboard data" + git push origin HEAD:main + - name: Download baseline from previous successful main run if: needs.resolve-atom-source.outputs.publish_to_dashboard == 'true' && steps.summary-stats.outputs.passed_cases != '0' id: baseline @@ -1526,7 +2057,7 @@ jobs: mapfile -t CANDIDATE_RUN_IDS < <( { gh run list \ - --workflow="ATOM vLLM OOT Benchmark" \ + --workflow="ATOM vLLM Benchmark" \ --branch=main \ --event=workflow_dispatch \ --status=success \ From 247bb23865f8478249b0a2f20c6183116f6ec16a Mon Sep 17 00:00:00 2001 From: Faa Diallo <107946068+amdfaa@users.noreply.github.com> Date: Mon, 11 May 2026 21:57:04 -0400 Subject: [PATCH 24/76] Add upload step for runner fleet report (#754) * Add upload step for runner fleet report Added a step to upload the runner fleet report as an artifact. * Add error handling for missing files in artifact upload --- .github/workflows/amd-ci-job-monitor.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/amd-ci-job-monitor.yml b/.github/workflows/amd-ci-job-monitor.yml index b223c8716b..92e0d7aaca 100644 --- a/.github/workflows/amd-ci-job-monitor.yml +++ b/.github/workflows/amd-ci-job-monitor.yml @@ -219,3 +219,11 @@ jobs: --summary | tee runner-fleet-report.md cat runner-fleet-report.md >> "$GITHUB_STEP_SUMMARY" + + - name: Upload runner fleet report + uses: actions/upload-artifact@v4 + with: + name: runner-fleet-report + path: runner-fleet-report.md + if-no-files-found: error + retention-days: 7 From 23f22726befd42c6a63f6fad7411e1fcfc27233d Mon Sep 17 00:00:00 2001 From: la <46212055+junhaha666@users.noreply.github.com> Date: Tue, 12 May 2026 10:33:22 +0800 Subject: [PATCH 25/76] Remove lru_cache in get_per_token_exponential func (#742) --- atom/model_ops/sampler.py | 1 - 1 file changed, 1 deletion(-) diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index 14276c3c1c..b1b5800149 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -32,7 +32,6 @@ SAMPLER_EPS = 1e-10 -@lru_cache(maxsize=1) def get_per_token_exponential(vocab_size: int, device) -> torch.Tensor: """Returns a tensor of shape (1, vocab_size) filled with exponential random values. This is key to deterministic inference, as it ensures that the same random values are used for each token across different runs. From fb7cef6407b930bea0649cefad1d6de9145a9ad5 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Tue, 12 May 2026 10:41:47 +0800 Subject: [PATCH 26/76] [MoE] Align generic MXFP4 shuffle layout (#744) * [MoE] align generic MXFP4 shuffle layout Use the shared separated-layout weight shuffle path for non-interleaved MXFP4 MoE weights and keep the interleaved scale shuffle explicit, matching AITER GateMode layout expectations. Co-authored-by: Cursor * [MoE] mark shuffled MXFP4 weights * update code --------- Co-authored-by: Cursor --- atom/model_ops/moe.py | 68 +++++++++++++++---------------------------- 1 file changed, 24 insertions(+), 44 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index b2aeac8652..4fb928a1dd 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -14,7 +14,6 @@ from aiter.jit.utils.chip_info import get_gfx from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.ops.shuffle import shuffle_weight, shuffle_scale -from aiter.utility import fp4_utils from atom.config import ( Config, QuantizationConfig, @@ -859,51 +858,32 @@ def process_weights_after_loading(self, layer): layer.w2_weight_scale = None return - # quark method for moe, split it out? - elif self.quant_method == "quark": - shuffle_weights(layer.w13_weight, layer.w2_weight) - s0, s1, _ = layer.w13_weight_scale.shape - w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) - w13_weight_scale = fp4_utils.e8m0_shuffle(w13_weight_scale) - layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) - - s0, s1, _ = layer.w2_weight_scale.shape - w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) - w2_weight_scale = fp4_utils.e8m0_shuffle(w2_weight_scale) - layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) - return - else: - # shuffle weight - layer.w13_weight.data = shuffle_weight( - layer.w13_weight, - is_guinterleave=self.is_guinterleave, - gate_up=True, - ) - layer.w2_weight.data = shuffle_weight( - layer.w2_weight, - is_guinterleave=self.is_guinterleave, - gate_up=False, - ) - - # shuffle scale - if self.is_guinterleave: - w13_scale_2d = layer.w13_weight_scale.view( - -1, layer.w13_weight_scale.shape[-1] - ) - w2_scale_2d = layer.w2_weight_scale.view( - -1, layer.w2_weight_scale.shape[-1] - ) - else: - w13_scale_2d = layer.w13_weight_scale.view(self.num_experts, -1) - w2_scale_2d = layer.w2_weight_scale.view(self.num_experts, -1) + # shuffle weight + layer.w13_weight.data = shuffle_weight( + layer.w13_weight, + is_guinterleave=self.is_guinterleave, + gate_up=True, + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight, + is_guinterleave=self.is_guinterleave, + gate_up=False, + ) + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True - shuffled_w13_scale = shuffle_scale( - w13_scale_2d, self.num_experts, self.is_guinterleave, True - ) - shuffled_w2_scale = shuffle_scale( - w2_scale_2d, self.num_experts, self.is_guinterleave, False - ) + # shuffle scale + w13_scale_2d = layer.w13_weight_scale.reshape( + -1, layer.w13_weight_scale.shape[-1] + ) + w2_scale_2d = layer.w2_weight_scale.reshape(-1, layer.w2_weight_scale.shape[-1]) + shuffled_w13_scale = shuffle_scale( + w13_scale_2d, self.num_experts, self.is_guinterleave, True + ) + shuffled_w2_scale = shuffle_scale( + w2_scale_2d, self.num_experts, self.is_guinterleave, False + ) layer.w13_weight_scale = atom_parameter(shuffled_w13_scale) layer.w2_weight_scale = atom_parameter(shuffled_w2_scale) From c88261eef852c05de7171ba5eedc7b061e8d5066 Mon Sep 17 00:00:00 2001 From: gbyu-amd Date: Tue, 12 May 2026 11:35:08 +0800 Subject: [PATCH 27/76] fix kimi k25 config error with trust-remote-code (#751) Co-authored-by: Guanbao Yu --- atom/model_config/kimi_k25.py | 61 ++++++++++++++++++++++++----- atom/plugin/vllm/models/kimi_k25.py | 7 +++- 2 files changed, 58 insertions(+), 10 deletions(-) diff --git a/atom/model_config/kimi_k25.py b/atom/model_config/kimi_k25.py index f9b32dc573..fb85d257b1 100644 --- a/atom/model_config/kimi_k25.py +++ b/atom/model_config/kimi_k25.py @@ -11,6 +11,13 @@ from transformers.configuration_utils import PretrainedConfig +def _first_non_none(*values): + for value in values: + if value is not None: + return value + return None + + class KimiK25VisionConfig(PretrainedConfig): model_type = "kimi_k25_vision" @@ -22,10 +29,10 @@ def __init__( init_pos_emb_width: int = 64, init_pos_emb_time: int = 4, pos_emb_type: str = "divided_fixed", - num_attention_heads: int = 16, - num_hidden_layers: int = 27, - hidden_size: int = 1152, - intermediate_size: int = 4304, + num_attention_heads: int | None = None, + num_hidden_layers: int | None = None, + hidden_size: int | None = None, + intermediate_size: int | None = None, merge_kernel_size: tuple[int, int] = (2, 2), video_attn_type: str = "spatial_temporal", merge_type: str = "sd2_tpool", @@ -34,9 +41,23 @@ def __init__( mm_hidden_size: int | None = None, projector_hidden_act: str = "gelu", projector_ln_eps: float = 1e-5, + # moonshotai/Kimi-K2.5 remote-code field names + vt_num_attention_heads: int | None = None, + vt_num_hidden_layers: int | None = None, + vt_hidden_size: int | None = None, + vt_intermediate_size: int | None = None, + text_hidden_size: int | None = None, **kwargs, ): super().__init__(**kwargs) + num_attention_heads = _first_non_none( + num_attention_heads, vt_num_attention_heads, 16 + ) + num_hidden_layers = _first_non_none(num_hidden_layers, vt_num_hidden_layers, 27) + hidden_size = _first_non_none(hidden_size, vt_hidden_size, 1152) + intermediate_size = _first_non_none( + intermediate_size, vt_intermediate_size, 4304 + ) # Vision Tower self.patch_size = patch_size self.init_pos_emb_height = init_pos_emb_height @@ -47,6 +68,11 @@ def __init__( self.num_hidden_layers = num_hidden_layers self.hidden_size = hidden_size self.intermediate_size = intermediate_size + # Preserve the repo-specific aliases so either schema can be consumed. + self.vt_num_attention_heads = num_attention_heads + self.vt_num_hidden_layers = num_hidden_layers + self.vt_hidden_size = hidden_size + self.vt_intermediate_size = intermediate_size self.merge_kernel_size = merge_kernel_size self.video_attn_type = video_attn_type self.merge_type = merge_type @@ -58,6 +84,25 @@ def __init__( self.mm_hidden_size = hidden_size self.projector_hidden_act = projector_hidden_act self.projector_ln_eps = projector_ln_eps + if text_hidden_size is not None: + self.text_hidden_size = text_hidden_size + + @classmethod + def from_remote_config( + cls, + vision_config: dict | PretrainedConfig | None, + ) -> "KimiK25VisionConfig": + if vision_config is None: + return cls() + if isinstance(vision_config, cls): + return vision_config + if isinstance(vision_config, dict): + config_dict = dict(vision_config) + elif hasattr(vision_config, "to_dict"): + config_dict = dict(vision_config.to_dict()) + else: + config_dict = dict(vars(vision_config)) + return cls(**config_dict) class KimiK25Config(PretrainedConfig): @@ -89,11 +134,7 @@ def __init__( **kwargs, ): # Vision config - if vision_config is None: - vision_config = KimiK25VisionConfig() - elif isinstance(vision_config, dict): - vision_config = KimiK25VisionConfig(**vision_config) - self.vision_config: KimiK25VisionConfig = vision_config + self.vision_config = KimiK25VisionConfig.from_remote_config(vision_config) # Text config if text_config is None: @@ -105,6 +146,8 @@ def __init__( # Set mm_hidden_size to text hidden size if not explicitly set if self.vision_config.mm_hidden_size == self.vision_config.hidden_size: self.vision_config.mm_hidden_size = self.text_config.hidden_size + if getattr(self.vision_config, "text_hidden_size", None) is None: + self.vision_config.text_hidden_size = self.text_config.hidden_size # Other config self.ignore_index = ignore_index diff --git a/atom/plugin/vllm/models/kimi_k25.py b/atom/plugin/vllm/models/kimi_k25.py index a9f73b7b6c..f3b16dd641 100644 --- a/atom/plugin/vllm/models/kimi_k25.py +++ b/atom/plugin/vllm/models/kimi_k25.py @@ -80,6 +80,7 @@ def __init__( ) else: self.norm = PPMissingLayer() + self.aux_hidden_state_layers: tuple[int, ...] = tuple() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) @@ -191,7 +192,11 @@ class KimiK25ForConditionalGeneration_(vLLMKimiK25): def __init__(self, atom_config: Config, prefix: str = "model"): # protocols have not __init__ method, so we need to use nn.Module.__init__ nn.Module.__init__(self) - config: KimiK25Config = atom_config.hf_config + hf_config = getattr(atom_config, "hf_config", None) + assert hf_config is not None, "hf_config is not found in atom_config" + vision_config = getattr(hf_config, "vision_config", None) + text_config = getattr(hf_config, "text_config", None) + config = KimiK25Config(vision_config, text_config) vllm_config = atom_config.plugin_config.vllm_config # quant_config from vLLM ignores exclude_layers in model's quantization config From 2b43d39eb96b06c2ccf7ffb8032dc132b737e8a5 Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Tue, 12 May 2026 14:49:58 +0800 Subject: [PATCH 28/76] support dp attention in deepseek v4 (#745) --- atom/models/deepseek_v4.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 4d9a59e162..a1dd9e01d1 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1952,6 +1952,23 @@ def _hash_topk( fwd_input_ids is not None ), "forward_context.context.input_ids is None — caller must invoke DeepseekV4ForCausalLM.forward, not DeepseekV4Model.forward directly." ids = fwd_input_ids.flatten() + # Under DP-attention gather/scatter, forward_impl_graph gathers + # hidden_states and gating_output across DP ranks before calling + # select_experts → _hash_topk. input_ids is still local (not + # gathered), so topk_ids would be [local_bs, topk] while the + # gathered gating_output is [local_bs*dp, n_experts]. Gather + # input_ids to match. + num_tokens = gating_output.shape[0] + if ids.shape[0] < num_tokens: + from atom.model_ops.moe import pad_for_all_gather + + ids_2d = ids.unsqueeze(-1) + ids_2d, _ = pad_for_all_gather(ids_2d) + from aiter.dist.parallel_state import get_dp_group + + ids_2d = get_dp_group().all_gather(ids_2d, dim=0) + ids = ids_2d[:num_tokens].flatten() + ids = ids.clamp(0, self.gate.tid2eid.shape[0] - 1) topk_ids = self.gate.tid2eid[ids].to(torch.int32) # [N, topk] scores = torch.nn.functional.softplus(gating_output.float()).sqrt() topk_weights = scores.gather(dim=-1, index=topk_ids.long()) From 7452ea955159e87ec01bfe16899792619b7127cc Mon Sep 17 00:00:00 2001 From: gbyu-amd Date: Tue, 12 May 2026 15:43:06 +0800 Subject: [PATCH 29/76] [recipe][vLLM-ATOM] update kimi k2/k2.5 recipe (#730) * update kimi k2/k2.5 recipe * update command --------- Co-authored-by: Guanbao Yu --- .github/benchmark/oot_benchmark_models.json | 2 +- recipes/atom_vllm/Kimi-K2-Thinking.md | 17 ++++++++++++++--- recipes/atom_vllm/Kimi-K2.5.md | 19 +++++++++++-------- 3 files changed, 26 insertions(+), 12 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 7326bf05ce..116d710144 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -141,7 +141,7 @@ "dashboard_model": "Kimi-K2.5-MXFP4-tp4", "prefix": "kimi-k25-mxfp4-tp4-met", "bench_args": "", - "extra_args": "--trust-remote-code --tensor-parallel-size 4", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { diff --git a/recipes/atom_vllm/Kimi-K2-Thinking.md b/recipes/atom_vllm/Kimi-K2-Thinking.md index 949955b30b..5118129c2c 100644 --- a/recipes/atom_vllm/Kimi-K2-Thinking.md +++ b/recipes/atom_vllm/Kimi-K2-Thinking.md @@ -14,12 +14,13 @@ docker pull rocm/atom-dev:vllm-latest ## Step 2: Launch vLLM Server The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and general usage flow compatible with upstream vLLM. For general server options and API usage, refer to the [official vLLM documentation](https://docs.vllm.ai/en/latest/). +We adopt [amd/Kimi-K2-Thinking-MXFP4-AttnFP8](https://huggingface.co/amd/Kimi-K2-Thinking-MXFP4-AttnFP8) for better performance by leveraging FP8 weights for attention layers. ```bash # use quick allreduce to reduce TTFT export AITER_QUICK_REDUCE_QUANTIZATION=INT4 -vllm serve amd/Kimi-K2-Thinking-MXFP4 \ +vllm serve amd/Kimi-K2-Thinking-MXFP4-AttnFP8 \ --host localhost \ --port 8000 \ --trust-remote-code \ @@ -28,6 +29,8 @@ vllm serve amd/Kimi-K2-Thinking-MXFP4 \ --gpu_memory_utilization 0.9 \ --async-scheduling \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --max-num-batched-tokens 16384 \ + --max-model-len 16384 \ --no-enable-prefix-caching ``` @@ -37,7 +40,7 @@ Users can use the default vllm bench command for performance benchmarking. vllm bench serve \ --host localhost \ --port 8000 \ - --model amd/Kimi-K2-Thinking-MXFP4 \ + --model amd/Kimi-K2-Thinking-MXFP4-AttnFP8 \ --dataset-name random \ --random-input-len 8000 \ --random-output-len 1000 \ @@ -52,7 +55,15 @@ vllm bench serve \ ```bash lm_eval --model local-completions \ - --model_args model=amd/Kimi-K2-Thinking-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=16,max_retries=3,tokenized_requests=False \ + --model_args model=amd/Kimi-K2-Thinking-MXFP4-AttnFP8,base_url=http://localhost:8000/v1/completions,num_concurrent=16,max_retries=3,tokenized_requests=False \ --tasks gsm8k \ --num_fewshot 3 ``` +The reference values of corresponding metrics: +```bash +local-completions ({'model': 'amd/Kimi-K2-Thinking-MXFP4-AttnFP8', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 16, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 +|Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| +|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| +|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9340|± |0.0068| +| | |strict-match | 3|exact_match|↑ |0.9325|± |0.0069| +``` \ No newline at end of file diff --git a/recipes/atom_vllm/Kimi-K2.5.md b/recipes/atom_vllm/Kimi-K2.5.md index aeaaa6ad1d..f94a139d92 100644 --- a/recipes/atom_vllm/Kimi-K2.5.md +++ b/recipes/atom_vllm/Kimi-K2.5.md @@ -14,12 +14,13 @@ docker pull rocm/atom-dev:vllm-latest ## Step 2: Launch vLLM Server The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and general usage flow compatible with upstream vLLM. For general server options and API usage, refer to the [official vLLM documentation](https://docs.vllm.ai/en/latest/). +We adopt [amd/Kimi-K2.5-MXFP4-AttnFP8](https://huggingface.co/amd/Kimi-K2.5-MXFP4-AttnFP8) for better performance by leveraging FP8 weights for attention layers. ```bash # use quick allreduce to reduce TTFT export AITER_QUICK_REDUCE_QUANTIZATION=INT4 -vllm serve amd/Kimi-K2.5-MXFP4 \ +vllm serve amd/Kimi-K2.5-MXFP4-AttnFP8 \ --host localhost \ --port 8000 \ --async-scheduling \ @@ -28,6 +29,8 @@ vllm serve amd/Kimi-K2.5-MXFP4 \ --gpu_memory_utilization 0.9 \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ --kv-cache-dtype fp8 \ + --max-num-batched-tokens 16384 \ + --max-model-len 16384 \ --no-enable-prefix-caching ``` @@ -51,7 +54,7 @@ IMAGE_BASE64=$(base64 -w 0 ATOM/recipes/atom_vllm/dog.png) curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ -d '{ - "model": "amd/Kimi-K2.5-MXFP4", + "model": "amd/Kimi-K2.5-MXFP4-AttnFP8", "messages": [ { "role": "user", @@ -82,7 +85,7 @@ The expected response: "id": "chatcmpl-941a3736cc5cce95", "object": "chat.completion", "created": 1774365813, - "model": "amd/Kimi-K2.5-MXFP4", + "model": "amd/Kimi-K2.5-MXFP4-AttnFP8", "choices": [ { "index": 0, @@ -122,7 +125,7 @@ Users can use the default vllm bench command for performance benchmarking. vllm bench serve \ --host localhost \ --port 8000 \ - --model amd/Kimi-K2.5-MXFP4 \ + --model amd/Kimi-K2.5-MXFP4-AttnFP8 \ --dataset-name random \ --random-input-len 8000 \ --random-output-len 1000 \ @@ -137,15 +140,15 @@ vllm bench serve \ The accuracy can be verified on gsm8k dataset with command: ```bash lm_eval --model local-completions \ - --model_args model=amd/Kimi-K2.5-MXFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \ + --model_args model=amd/Kimi-K2.5-MXFP4-AttnFP8,base_url=http://localhost:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False \ --tasks gsm8k \ --num_fewshot 3 ``` The reference values of corresponding metrics: ```bash -local-completions ({'model': '/workspace/shared/data/models/Kimi-K2.5-MXFP4', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 64, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 +local-completions ({'model': 'amd/Kimi-K2.5-MXFP4-AttnFP8', 'base_url': 'http://localhost:8000/v1/completions', 'num_concurrent': 64, 'max_retries': 3, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: 3, batch_size: 1 |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| -|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9401|± |0.0065| -| | |strict-match | 3|exact_match|↑ |0.9386|± |0.0066| +|gsm8k| 3|flexible-extract| 3|exact_match|↑ |0.9325|± |0.0061| +| | |strict-match | 3|exact_match|↑ |0.9240|± |0.0062| ``` From 48068ebaa3981d84f4c627cd2898b4cfef6a8cb1 Mon Sep 17 00:00:00 2001 From: PerryZhang01 Date: Tue, 12 May 2026 16:44:37 +0800 Subject: [PATCH 30/76] [fix](gpt-oss): fix gpt-oss model accuracy bug (#748) * [fix](gpt-oss): fix gpt-oss model accuracy bug * [fix](server): fix gpu-memory-utilization args --------- Co-authored-by: perzhang --- .github/benchmark/models.json | 2 +- .github/benchmark/models_accuracy.json | 4 +- .github/benchmark/oot_models_accuracy.json | 4 +- .github/scripts/atom_oot_test.sh | 45 +++++-------------- .../atom-vllm-accuracy-validation.yaml | 4 +- .github/workflows/atom-vllm-test.yaml | 5 ++- recipes/GPT-OSS.md | 2 + recipes/atom_vllm/GPT-OSS.md | 4 +- 8 files changed, 25 insertions(+), 45 deletions(-) diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 39aed57bd1..bef8125bf2 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -77,7 +77,7 @@ "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", - "env_vars": "" + "env_vars": "ATOM_MOE_GU_ITLV=1" }, { "display": "Kimi-K2.5-MXFP4", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 4887f3d126..46db2c0496 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -52,7 +52,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", - "env_vars": "", + "env_vars": "ATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-1", "test_level": "pr", "accuracy_threshold": 0.88, @@ -125,7 +125,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", - "env_vars": "", + "env_vars": "ATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-4", "test_level": "main", "accuracy_threshold": 0.88, diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index fdbeb1bd9a..f02df88450 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -153,7 +153,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 1", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-1", "test_level": "nightly", "accuracy_threshold": 0.88, @@ -165,7 +165,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.88, diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index dc67918590..69f51436e0 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -51,7 +51,6 @@ EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} EXPLICIT_CLIENT_COMMAND=${OOT_CLIENT_COMMAND:-} -OOT_GPU_MEMORY_UTILIZATION=${OOT_GPU_MEMORY_UTILIZATION:-0.9} OOT_DOCKER_IMAGE=${OOT_DOCKER_IMAGE:-} LM_EVAL_NUM_FEWSHOT=${LM_EVAL_NUM_FEWSHOT:-3} LAST_VLLM_LOG_LINE=0 @@ -61,11 +60,6 @@ if ! [[ "${LM_EVAL_NUM_FEWSHOT}" =~ ^[0-9]+$ ]]; then exit 2 fi -if ! [[ "${OOT_GPU_MEMORY_UTILIZATION}" =~ ^[0-9]*\.?[0-9]+$ ]]; then - echo "Invalid OOT_GPU_MEMORY_UTILIZATION: ${OOT_GPU_MEMORY_UTILIZATION}. Expected a numeric value." - exit 2 -fi - declare -a ACTIVE_MODELS=() if [[ -n "${EXPLICIT_MODEL_NAME}" || -n "${EXPLICIT_MODEL_PATH}" || -n "${EXPLICIT_EXTRA_ARGS}" ]]; then if [[ -z "${EXPLICIT_MODEL_NAME}" || -z "${EXPLICIT_MODEL_PATH}" ]]; then @@ -89,14 +83,6 @@ resolve_model_path() { fi } -# gpt-oss OpenAI weights require Chat Completions (messages); local-completions sends "prompt" and vLLM returns 400. -is_gpt_oss_model() { - local name_lc path_lc - name_lc="$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" - path_lc="$(printf '%s' "$2" | tr '[:upper:]' '[:lower:]')" - [[ "${name_lc}" == *gpt-oss* || "${path_lc}" == *gpt-oss* ]] -} - emit_new_vllm_logs() { if [[ "${STREAM_VLLM_LOGS}" != "1" || ! -f "${VLLM_LOG_FILE}" ]]; then return 0 @@ -187,7 +173,6 @@ PY echo "Model name: ${model_name}" echo "Model path: ${resolved_model_path}" echo "Extra args: ${extra_args}" - echo "GPU memory utilization: ${OOT_GPU_MEMORY_UTILIZATION}" export SAFETENSORS_FAST_GPU=1 export VLLM_RPC_TIMEOUT=1800000 @@ -214,7 +199,6 @@ PY --trust-remote-code \ --kv-cache-dtype fp8 \ "${extra_arg_array[@]}" \ - --gpu-memory-utilization "${OOT_GPU_MEMORY_UTILIZATION}" \ --no-enable-prefix-caching \ > "${VLLM_LOG_FILE}" 2>&1 & echo $! > "${VLLM_PID_FILE}" @@ -249,6 +233,10 @@ accuracy_one_model() { echo "Model name: ${model_name}" echo "Few-shot count: ${LM_EVAL_NUM_FEWSHOT}" + if [[ "${client_command}" == "null" ]]; then + client_command="" + fi + if [[ -n "${client_command}" ]]; then local -a client_command_args=() while IFS= read -r -d '' token; do @@ -299,29 +287,16 @@ PY echo "Using custom lm-eval command from client_command: ${client_command}" "${client_command_args[@]}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" else + echo "Using default lm-eval command." local lm_args=( --model_args model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True ) - if is_gpt_oss_model "${model_name}" "${model_path}"; then - echo "Using chat completions + apply_chat_template for gpt-oss (OpenAI-compatible messages API)." - lm_args=( - --model_args - model="${resolved_model_path}",base_url="http://127.0.0.1:${VLLM_PORT}/v1/chat/completions",num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True - ) - lm_eval --model local-chat-completions \ - --apply_chat_template \ - "${lm_args[@]}" \ - --tasks gsm8k \ - --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ - --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" - else - lm_eval --model local-completions \ - "${lm_args[@]}" \ - --tasks gsm8k \ - --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ - --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" - fi + lm_eval --model local-completions \ + "${lm_args[@]}" \ + --tasks gsm8k \ + --num_fewshot "${LM_EVAL_NUM_FEWSHOT}" \ + --output_path "${output_path}" 2>&1 | tee -a "${ACCURACY_LOG_FILE}" fi # lm-eval output layout differs across versions: output_path may be a file diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index ba4ad21248..dcc814da36 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -288,7 +288,7 @@ jobs: "extra_args": "--tensor-parallel-size 1", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "accuracy_test_threshold": 0.88, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-1", }, { @@ -298,7 +298,7 @@ jobs: "extra_args": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "accuracy_test_threshold": 0.88, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", "runner": "linux-atom-mi35x-4", }, { diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index eadb39b46d..6000713dcf 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.github/workflows/atom-vllm-test.yaml @@ -213,7 +213,8 @@ jobs: model_name: "gpt-oss-120b" model_path: "openai/gpt-oss-120b" extra_args: "--tensor-parallel-size 1 --gpu-memory-utilization 0.5" - env_vars: "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" + client_command: "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}" + env_vars: "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1" accuracy_test_threshold: 0.88 runner: linux-atom-mi35x-1 - display_name: "Kimi-K2-Thinking-MXFP4 TP4" @@ -468,6 +469,7 @@ jobs: OOT_MODEL_NAME: ${{ matrix.model_name }} OOT_MODEL_PATH: ${{ matrix.model_path }} OOT_EXTRA_ARGS: ${{ matrix.extra_args }} + OOT_CLIENT_COMMAND: ${{ matrix.client_command || '' }} OOT_ENV_VARS: ${{ matrix.env_vars }} LM_EVAL_NUM_FEWSHOT: ${{ matrix.lm_eval_num_fewshot }} MAX_WAIT_RETRIES: "120" @@ -477,6 +479,7 @@ jobs: -e OOT_MODEL_NAME="${OOT_MODEL_NAME}" \ -e OOT_MODEL_PATH="${OOT_RESOLVED_MODEL_PATH:-$OOT_MODEL_PATH}" \ -e OOT_EXTRA_ARGS="${OOT_EXTRA_ARGS}" \ + -e OOT_CLIENT_COMMAND="${OOT_CLIENT_COMMAND}" \ -e OOT_ENV_VARS="${OOT_ENV_VARS}" \ -e LM_EVAL_NUM_FEWSHOT="${LM_EVAL_NUM_FEWSHOT}" \ -e MAX_WAIT_RETRIES="${MAX_WAIT_RETRIES}" \ diff --git a/recipes/GPT-OSS.md b/recipes/GPT-OSS.md index a65449dac2..748dc4cdf8 100644 --- a/recipes/GPT-OSS.md +++ b/recipes/GPT-OSS.md @@ -17,6 +17,7 @@ All the operations below will be executed inside the container. GPT-OSS-120B fits on a single MI300X/MI355X GPU: ```bash +ATOM_MOE_GU_ITLV=1 \ python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 \ @@ -28,6 +29,7 @@ python -m atom.entrypoints.openai_server \ Scale across 2 GPUs with data-parallel attention and expert parallelism: ```bash +ATOM_MOE_GU_ITLV=1 \ python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 -tp 2 \ diff --git a/recipes/atom_vllm/GPT-OSS.md b/recipes/atom_vllm/GPT-OSS.md index 6f4e751c10..f06ba19ac2 100644 --- a/recipes/atom_vllm/GPT-OSS.md +++ b/recipes/atom_vllm/GPT-OSS.md @@ -15,8 +15,8 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener GPT-OSS-120B is a single-GPU model, so `--tensor-parallel-size` defaults to 1 and can be omitted. ```bash -export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 - +ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 \ +ATOM_MOE_GU_ITLV=1 \ vllm serve openai/gpt-oss-120b \ --host localhost \ --port 8000 \ From 2b6d7eb632fdd37c2611e74eae1e9b52ed7ae110 Mon Sep 17 00:00:00 2001 From: yzhou103 Date: Tue, 12 May 2026 18:02:00 +0800 Subject: [PATCH 31/76] fix(deepseek_v4): switch MoE routing to topk_gating (#746) Update DeepSeek V4 MoE routing to call aiter topk_gating with explicit sqrtsoftplus scoring, and remove the forced fp32 cast on gate output so the router logits path follows the model default dtype behavior. Co-authored-by: Cursor --- atom/model_ops/moe.py | 5 +++-- atom/models/deepseek_v4.py | 4 +--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index 4fb928a1dd..e38388de46 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -8,7 +8,7 @@ from typing import Callable, List, Optional, Tuple import torch -from aiter import ActivationType, QuantType, dtypes, get_hip_quant, topk_softplus +from aiter import ActivationType, QuantType, dtypes, get_hip_quant, topk_gating from aiter.dist.parallel_state import get_dp_group, get_tp_group from aiter.fused_moe import fused_moe from aiter.jit.utils.chip_info import get_gfx @@ -2644,13 +2644,14 @@ def select_experts( topk_weights = torch.empty( tokens_num, top_k, dtype=torch.float32, device=router_logits.device ) - topk_softplus( + topk_gating( topk_weights, topk_ids, router_logits, e_score_correction_bias, renormalize, routed_scaling_factor, + score_func="sqrtsoftplus", ) else: raise ValueError( diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index a1dd9e01d1..a53eb3a50e 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1989,9 +1989,7 @@ def routed_expert_forward( `forward_context.context.input_ids` before each forward, and `_hash_topk` (FusedMoE's custom_routing_function) reads it there. """ - router_logits = self.gate( - x, otype=torch.float32 - ) # [num_tokens, n_routed_experts] + router_logits = self.gate(x) # [num_tokens, n_routed_experts] return self.experts(hidden_states=x, router_logits=router_logits) def combine_outputs( From 8307a37c9d177b38a9867577f22bbd0252dd47fd Mon Sep 17 00:00:00 2001 From: PerryZhang01 Date: Tue, 12 May 2026 20:48:46 +0800 Subject: [PATCH 32/76] [feat](gpt-oss): enable gpt-oss moe a4w4 (#764) Co-authored-by: perzhang --- .github/benchmark/models.json | 2 +- .github/benchmark/models_accuracy.json | 4 ++-- .github/benchmark/oot_models_accuracy.json | 4 ++-- .github/workflows/atom-vllm-accuracy-validation.yaml | 4 ++-- .github/workflows/atom-vllm-test.yaml | 2 +- recipes/GPT-OSS.md | 2 -- recipes/atom_vllm/GPT-OSS.md | 1 - 7 files changed, 8 insertions(+), 11 deletions(-) diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index bef8125bf2..39aed57bd1 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -77,7 +77,7 @@ "bench_args": "", "suffix": "", "runner": "atom-mi355-8gpu.predownload", - "env_vars": "ATOM_MOE_GU_ITLV=1" + "env_vars": "" }, { "display": "Kimi-K2.5-MXFP4", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 46db2c0496..4887f3d126 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -52,7 +52,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--kv_cache_dtype fp8 --gpu-memory-utilization 0.5", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_MOE_GU_ITLV=1", + "env_vars": "", "runner": "linux-atom-mi35x-1", "test_level": "pr", "accuracy_threshold": 0.88, @@ -125,7 +125,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_MOE_GU_ITLV=1", + "env_vars": "", "runner": "linux-atom-mi35x-4", "test_level": "main", "accuracy_threshold": 0.88, diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index f02df88450..fdbeb1bd9a 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -153,7 +153,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 1", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-1", "test_level": "nightly", "accuracy_threshold": 0.88, @@ -165,7 +165,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.88, diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index dcc814da36..ba4ad21248 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -288,7 +288,7 @@ jobs: "extra_args": "--tensor-parallel-size 1", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "accuracy_test_threshold": 0.88, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-1", }, { @@ -298,7 +298,7 @@ jobs: "extra_args": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "accuracy_test_threshold": 0.88, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, { diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index 6000713dcf..ca9d3e3dd8 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.github/workflows/atom-vllm-test.yaml @@ -214,7 +214,7 @@ jobs: model_path: "openai/gpt-oss-120b" extra_args: "--tensor-parallel-size 1 --gpu-memory-utilization 0.5" client_command: "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}" - env_vars: "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_MOE_GU_ITLV=1" + env_vars: "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" accuracy_test_threshold: 0.88 runner: linux-atom-mi35x-1 - display_name: "Kimi-K2-Thinking-MXFP4 TP4" diff --git a/recipes/GPT-OSS.md b/recipes/GPT-OSS.md index 748dc4cdf8..a65449dac2 100644 --- a/recipes/GPT-OSS.md +++ b/recipes/GPT-OSS.md @@ -17,7 +17,6 @@ All the operations below will be executed inside the container. GPT-OSS-120B fits on a single MI300X/MI355X GPU: ```bash -ATOM_MOE_GU_ITLV=1 \ python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 \ @@ -29,7 +28,6 @@ python -m atom.entrypoints.openai_server \ Scale across 2 GPUs with data-parallel attention and expert parallelism: ```bash -ATOM_MOE_GU_ITLV=1 \ python -m atom.entrypoints.openai_server \ --model openai/gpt-oss-120b \ --kv_cache_dtype fp8 -tp 2 \ diff --git a/recipes/atom_vllm/GPT-OSS.md b/recipes/atom_vllm/GPT-OSS.md index f06ba19ac2..414b5abc39 100644 --- a/recipes/atom_vllm/GPT-OSS.md +++ b/recipes/atom_vllm/GPT-OSS.md @@ -16,7 +16,6 @@ GPT-OSS-120B is a single-GPU model, so `--tensor-parallel-size` defaults to 1 an ```bash ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 \ -ATOM_MOE_GU_ITLV=1 \ vllm serve openai/gpt-oss-120b \ --host localhost \ --port 8000 \ From 30b91373aa98bce5f862220dc619ad2c42cecabc Mon Sep 17 00:00:00 2001 From: HaonanWang98 Date: Tue, 12 May 2026 21:41:37 +0800 Subject: [PATCH 33/76] Add chat encode for deepseek v4 without chat template jinja (#760) * add chat encoder for deepseek v4 * align simple inference --- atom/entrypoints/openai/api_server.py | 14 +-- atom/entrypoints/openai/chat_encoders.py | 107 +++++++++++++++++++++++ atom/examples/simple_inference.py | 45 ++-------- 3 files changed, 124 insertions(+), 42 deletions(-) create mode 100644 atom/entrypoints/openai/chat_encoders.py diff --git a/atom/entrypoints/openai/api_server.py b/atom/entrypoints/openai/api_server.py index 01ef080815..bcad5e7083 100644 --- a/atom/entrypoints/openai/api_server.py +++ b/atom/entrypoints/openai/api_server.py @@ -31,6 +31,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from transformers import AutoTokenizer +from .chat_encoders import apply_chat_template, load_custom_message_encoder from .protocol import ( ChatCompletionRequest, CompletionRequest, @@ -66,6 +67,7 @@ tokenizer: Optional[AutoTokenizer] = None model_name: str = "" default_chat_template_kwargs: Dict[str, Any] = {} +custom_message_encoder: Optional[Any] = None _stream_queues: Dict[str, asyncio.Queue] = {} _seq_id_to_request_id: Dict[int, str] = {} _stream_loops: Dict[str, AbstractEventLoop] = {} @@ -590,14 +592,12 @@ async def chat_completions(request: ChatCompletionRequest): merged_kwargs = dict(default_chat_template_kwargs) if request.chat_template_kwargs: merged_kwargs.update(request.chat_template_kwargs) - merged_kwargs["tokenize"] = False - merged_kwargs["add_generation_prompt"] = True - # Pass tools so the chat template can inject tool declarations - if request.tools: - merged_kwargs["tools"] = request.tools - prompt = tokenizer.apply_chat_template( + prompt = apply_chat_template( + tokenizer, + custom_message_encoder, [msg.to_template_dict() for msg in messages], + tools=request.tools, **merged_kwargs, ) @@ -825,6 +825,7 @@ async def stop_profile(): def main(): """Main entry point for the server.""" global engine, tokenizer, model_name, default_chat_template_kwargs, _request_logger + global custom_message_encoder parser = argparse.ArgumentParser(description="ATOM OpenAI API Server") EngineArgs.add_cli_args(parser) @@ -869,6 +870,7 @@ def main(): logger.info(f"Loading tokenizer from {args.model}...") tokenizer = _load_tokenizer(args.model, args.trust_remote_code) model_name = args.model + custom_message_encoder = load_custom_message_encoder(args.model) logger.info(f"Initializing engine with model {args.model}...") engine_args = EngineArgs.from_cli_args(args) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py new file mode 100644 index 0000000000..aefde69ec2 --- /dev/null +++ b/atom/entrypoints/openai/chat_encoders.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Chat-template dispatch for the OpenAI chat endpoint. + +Some models (e.g. DeepSeek V4) ship a Python encoder under ``/encoding/`` +instead of a Jinja ``chat_template``. This module discovers such encoders at +server startup and provides :func:`apply_chat_template`, a single entry point +that the request handler calls — it transparently routes to the custom encoder +when one was found, or to ``tokenizer.apply_chat_template`` otherwise. +""" + +import glob +import importlib.util +import logging +import os +from typing import Any, Callable, List, Optional + +logger = logging.getLogger("atom") + +MessageEncoder = Callable[..., str] + + +def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoder]: + """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. + + Returns ``None`` when the directory or matching file is absent (model uses + the standard Jinja path). Returns ``None`` and warns on ambiguity (multiple + matches) or load failures. + + """ + enc_dir = os.path.join(model_path, "encoding") + if not os.path.isdir(enc_dir): + return None + + candidates = sorted(glob.glob(os.path.join(enc_dir, "encoding_*.py"))) + if not candidates: + return None + if len(candidates) > 1: + logger.warning( + f"Multiple encoding_*.py found in {enc_dir}, refusing to guess: " + f"{[os.path.basename(p) for p in candidates]}" + ) + return None + + enc_path = candidates[0] + module_name = os.path.splitext(os.path.basename(enc_path))[0] + try: + spec = importlib.util.spec_from_file_location(module_name, enc_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + raw = getattr(mod, "encode_messages") + except Exception as e: + logger.warning(f"Failed to load encoder from {enc_path}: {e}") + return None + + # also valid is "chat" (non-thinking short-form). May need to add as an option. + # Revisit when a second model ships an encode_*.py — the default may need to be per-model. + def encode(messages, **kwargs): + kwargs.setdefault("thinking_mode", "thinking") + return raw(messages, **kwargs) + + logger.info(f"Loaded message encoder from {enc_path}") + return encode + + +def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: + """Probe ``model_path`` once at startup for a custom message encoder. + + Returns the encoder, or ``None`` when the model uses the standard Jinja + ``chat_template`` path. Result should be cached by the caller — this does + filesystem IO and a Python import. + """ + return _load_encoder_from_dir(model_path) + + +def apply_chat_template( + tokenizer: Any, + custom_encoder: Optional[MessageEncoder], + messages: List[dict], + *, + tools: Optional[List[dict]] = None, + **kwargs: Any, +) -> str: + """Render ``messages`` to a prompt string. + + Dispatches to ``custom_encoder`` if one was discovered for this model, + otherwise to ``tokenizer.apply_chat_template``. Jinja-only kwargs + (``tokenize``, ``add_generation_prompt``) are stripped on the custom + path; ``tools`` are forwarded only on the Jinja path (custom encoders + don't currently have a tools API — caller is warned and tools are + dropped). + """ + if custom_encoder is not None: + for k in ("tokenize", "add_generation_prompt"): + kwargs.pop(k, None) + if tools: + logger.warning( + "tools= is not supported with the custom message encoder; ignoring." + ) + return custom_encoder(messages, **kwargs) + + kwargs["tokenize"] = False + kwargs["add_generation_prompt"] = True + if tools: + kwargs["tools"] = tools + return tokenizer.apply_chat_template(messages, **kwargs) diff --git a/atom/examples/simple_inference.py b/atom/examples/simple_inference.py index fda0de7c63..73d51c2e69 100644 --- a/atom/examples/simple_inference.py +++ b/atom/examples/simple_inference.py @@ -2,9 +2,12 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import argparse -import os from atom import SamplingParams +from atom.entrypoints.openai.chat_encoders import ( + apply_chat_template, + load_custom_message_encoder, +) from atom.model_engine.arg_utils import EngineArgs from transformers import AutoTokenizer @@ -61,41 +64,11 @@ def main(): temperature=args.temperature, max_tokens=args.max_tokens ) - # Apply chat template. DeepSeek-V4 ships without a HuggingFace - # chat_template; use its custom encoding_dsv4.py if available. - if getattr(tokenizer, "chat_template", None): - prompts = [ - tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], - tokenize=False, - add_generation_prompt=True, - enable_thinking=True, - ) - for prompt in prompts - ] - else: - try: - import importlib.util - - enc_path = os.path.join(args.model, "encoding", "encoding_dsv4.py") - if os.path.exists(enc_path): - spec = importlib.util.spec_from_file_location("encoding_dsv4", enc_path) - enc_mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(enc_mod) - prompts = [ - enc_mod.encode_messages( - [{"role": "user", "content": p}], thinking_mode="chat" - ) - for p in prompts - ] - print( - f" (applied V4 encoding_dsv4, prompt tokens: " - f"{[len(tokenizer.encode(p)) for p in prompts]})" - ) - else: - print(" (tokenizer has no chat_template — feeding raw prompts as-is)") - except Exception as e: - print(f" (V4 encoding failed: {e} — feeding raw prompts as-is)") + custom_encoder = load_custom_message_encoder(args.model) + prompts = [ + apply_chat_template(tokenizer, custom_encoder, [{"role": "user", "content": p}]) + for p in prompts + ] print("This is prompts:", prompts) # print("Warming up...") # _ = llm.generate(["warmup"], sampling_params) From c615b35ba2c08dee1e7dfdde1f1b95be5564bb35 Mon Sep 17 00:00:00 2001 From: Zhu Yuhua Date: Tue, 12 May 2026 22:18:39 +0800 Subject: [PATCH 34/76] (ci)(recipe): Add DeepSeek-R1 FP4 TP4 validation and DS recipe for SGLang-ATOM (#614) * (ci)(recipe): Add DeepSeek-R1 FP4 TP4 validation and DS recipe for SGLang-ATOM Signed-off-by: zhuyuhua-v * update ci threshold Signed-off-by: zhuyuhua-v * update aiter whl download flow Signed-off-by: zhuyuhua-v * remove int flag from ci cases Signed-off-by: zhuyuhua-v * update recipe Signed-off-by: zhuyuhua-v * adjust threshold for fp4 tp4 to 0.91 Signed-off-by: zhuyuhua-v --------- Signed-off-by: zhuyuhua-v --- .../benchmark/sglang_benchmark_models.json | 10 +- .github/benchmark/sglang_models_accuracy.json | 40 ++++- .github/runner-config.yml | 4 + .../atom-sglang-accuracy-validation.yaml | 49 +++++- .github/workflows/atom-sglang-benchmark.yaml | 74 ++++++++- .github/workflows/atom-sglang-test.yaml | 157 ++++++++++++++---- recipes/atom_sglang/DeepSeek-R1.md | 146 ++++++++++++++++ 7 files changed, 438 insertions(+), 42 deletions(-) create mode 100644 recipes/atom_sglang/DeepSeek-R1.md diff --git a/.github/benchmark/sglang_benchmark_models.json b/.github/benchmark/sglang_benchmark_models.json index 8ddee78d0d..b06227844d 100644 --- a/.github/benchmark/sglang_benchmark_models.json +++ b/.github/benchmark/sglang_benchmark_models.json @@ -6,7 +6,7 @@ "prefix": "deepseek-r1-fp8-tp8", "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -17,7 +17,7 @@ "prefix": "deepseek-r1-fp8-tp4", "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -27,7 +27,7 @@ "prefix": "deepseek-r1-fp4-tp8", "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -38,7 +38,7 @@ "prefix": "deepseek-r1-fp4-tp4", "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -49,7 +49,7 @@ "prefix": "deepseek-r1-fp4-tp8-ep8", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --expert-parallel-size 8", "bench_args": "", - "runner": "atom-mi355-8gpu-oot-benchmark", + "runner": "atom-mi355-8gpu-aac-runner", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" } ] diff --git a/.github/benchmark/sglang_models_accuracy.json b/.github/benchmark/sglang_models_accuracy.json index 2580c133d3..1fa97d41fc 100644 --- a/.github/benchmark/sglang_models_accuracy.json +++ b/.github/benchmark/sglang_models_accuracy.json @@ -3,12 +3,48 @@ "model_name": "DeepSeek-R1-FP8 TP4", "model_path": "deepseek-ai/DeepSeek-R1-0528", "extraArgs": "--tensor-parallel-size 4", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", - "accuracy_threshold": 0.92, + "accuracy_threshold": 0.91, "accuracy_baseline": null, "accuracy_baseline_model": "deepseek-ai/DeepSeek-R1-0528", "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP8 TP8", + "model_path": "deepseek-ai/DeepSeek-R1-0528", + "extraArgs": "--tensor-parallel-size 8", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.93, + "accuracy_baseline": null, + "accuracy_baseline_model": "deepseek-ai/DeepSeek-R1-0528", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP4 TP4", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extraArgs": "--tensor-parallel-size 4", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.91, + "accuracy_baseline": null, + "accuracy_baseline_model": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "DeepSeek-R1-FP4 TP8", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extraArgs": "--tensor-parallel-size 8", + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.93, + "accuracy_baseline": null, + "accuracy_baseline_model": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." } ] diff --git a/.github/runner-config.yml b/.github/runner-config.yml index e0287eccf3..ed11880566 100644 --- a/.github/runner-config.yml +++ b/.github/runner-config.yml @@ -6,6 +6,10 @@ runners: gpu_arch: MI355 gpu_count: 8 + atom-mi355-8gpu-aac-runner: + gpu_arch: MI355 + gpu_count: 8 + linux-atom-mi355-1: gpu_arch: MI355 gpu_count: 1 diff --git a/.github/workflows/atom-sglang-accuracy-validation.yaml b/.github/workflows/atom-sglang-accuracy-validation.yaml index 7e6b9baa12..2c9eaaa92b 100644 --- a/.github/workflows/atom-sglang-accuracy-validation.yaml +++ b/.github/workflows/atom-sglang-accuracy-validation.yaml @@ -14,6 +14,21 @@ on: required: false type: boolean default: false + run_dsr1_fp8_tp8: + description: "DeepSeek-R1-FP8 TP8" + required: false + type: boolean + default: false + run_dsr1_fp4_tp4: + description: "DeepSeek-R1-FP4 TP4" + required: false + type: boolean + default: false + run_dsr1_fp4_tp8: + description: "DeepSeek-R1-FP4 TP8" + required: false + type: boolean + default: false upload_accuracy_to_dashboard: description: "Optional: upload SGLANG accuracy results to dashboard after this manual run" required: false @@ -55,6 +70,9 @@ jobs: id: meta env: RUN_DSR1_FP8_TP4: ${{ inputs.run_dsr1_fp8_tp4 }} + RUN_DSR1_FP8_TP8: ${{ inputs.run_dsr1_fp8_tp8 }} + RUN_DSR1_FP4_TP4: ${{ inputs.run_dsr1_fp4_tp4 }} + RUN_DSR1_FP4_TP8: ${{ inputs.run_dsr1_fp4_tp8 }} run: | set -euo pipefail @@ -72,10 +90,37 @@ jobs: "model_name": "DeepSeek-R1-FP8 TP4", "model_path": "deepseek-ai/DeepSeek-R1-0528", "extra_args": "--tensor-parallel-size 4", - "accuracy_test_threshold": 0.92, - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "accuracy_test_threshold": 0.91, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, + { + "toggle_env": "RUN_DSR1_FP8_TP8", + "model_name": "DeepSeek-R1-FP8 TP8", + "model_path": "deepseek-ai/DeepSeek-R1-0528", + "extra_args": "--tensor-parallel-size 8", + "accuracy_test_threshold": 0.93, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, + { + "toggle_env": "RUN_DSR1_FP4_TP4", + "model_name": "DeepSeek-R1-FP4 TP4", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extra_args": "--tensor-parallel-size 4", + "accuracy_test_threshold": 0.91, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_DSR1_FP4_TP8", + "model_name": "DeepSeek-R1-FP4 TP8", + "model_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "extra_args": "--tensor-parallel-size 8", + "accuracy_test_threshold": 0.93, + "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, ] selected = [] diff --git a/.github/workflows/atom-sglang-benchmark.yaml b/.github/workflows/atom-sglang-benchmark.yaml index 66720b9241..a925da28a6 100644 --- a/.github/workflows/atom-sglang-benchmark.yaml +++ b/.github/workflows/atom-sglang-benchmark.yaml @@ -5,6 +5,9 @@ concurrency: cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} on: + schedule: + # Nightly at 01:00 Beijing time (17:00 UTC on the previous day) + - cron: '0 17 * * *' workflow_dispatch: inputs: deepseek-r1-fp8-tp8: @@ -83,6 +86,10 @@ jobs: REBUILD_SGLANG_IMAGE=false SGLANG_IMAGE_SOURCE="prebuilt" USER_PROVIDED_SGLANG_IMAGE=false + RESOLVED_PREBUILT_ALIAS="" + RESOLVED_PREBUILT_DIGEST="" + PREBUILT_RESOLUTION_TYPE="" + RESOLVED_NIGHTLY_TAG="" if [[ -n "${INPUT_SGLANG_IMAGE}" ]]; then PREBUILT_SGLANG_IMAGE="${INPUT_SGLANG_IMAGE}" SGLANG_IMAGE_SOURCE="user-provided" @@ -116,6 +123,39 @@ jobs: SELECTED_SGLANG_REF="${VLLM_META[0]}" SELECTED_SGLANG_VERSION="${VLLM_META[1]}" + if [[ "${SGLANG_IMAGE_SOURCE}" == "prebuilt" ]]; then + if RESOLUTION_JSON="$( + python3 .github/scripts/resolve_atom_image.py \ + --repository rocm/atom-dev \ + --reference-tag sglang-latest \ + --preferred-version "${SELECTED_SGLANG_VERSION}" + )"; then + mapfile -t RESOLVED_IMAGE_META < <( + RESOLUTION_JSON="${RESOLUTION_JSON}" python3 - <<'PY' + import json + import os + + resolution = json.loads(os.environ["RESOLUTION_JSON"]) + print(resolution["resolved_image"]) + print(resolution["reference_image"]) + print(resolution["reference_digest"]) + print(resolution["match_type"]) + print(resolution.get("matched_tag", "")) + PY + ) + PREBUILT_SGLANG_IMAGE="${RESOLVED_IMAGE_META[0]}" + RESOLVED_PREBUILT_ALIAS="${RESOLVED_IMAGE_META[1]}" + RESOLVED_PREBUILT_DIGEST="${RESOLVED_IMAGE_META[2]}" + PREBUILT_RESOLUTION_TYPE="${RESOLVED_IMAGE_META[3]}" + RESOLVED_NIGHTLY_TAG="${RESOLVED_IMAGE_META[4]}" + else + echo "::warning::Failed to resolve rocm/atom-dev:sglang-latest to an immutable digest; falling back to the floating tag." + PREBUILT_SGLANG_IMAGE="rocm/atom-dev:sglang-latest" + RESOLVED_PREBUILT_ALIAS="rocm/atom-dev:sglang-latest" + PREBUILT_RESOLUTION_TYPE="resolution-failed" + fi + fi + { echo "atom_repository=${ATOM_REPOSITORY}" echo "atom_ref=${ATOM_REF}" @@ -153,6 +193,32 @@ jobs: "${PREBUILT_SGLANG_IMAGE}" \ "${SELECTED_SGLANG_VERSION}" \ "${SELECTED_SGLANG_REF}" >> "$GITHUB_STEP_SUMMARY" + if [[ -n "${RESOLVED_PREBUILT_ALIAS}" ]]; then + printf -- '- Floating alias: `%s`\n' \ + "${RESOLVED_PREBUILT_ALIAS}" >> "$GITHUB_STEP_SUMMARY" + fi + if [[ -n "${RESOLVED_PREBUILT_DIGEST}" ]]; then + printf -- '- Resolved digest: `%s`\n' \ + "${RESOLVED_PREBUILT_DIGEST}" >> "$GITHUB_STEP_SUMMARY" + fi + case "${PREBUILT_RESOLUTION_TYPE}" in + matched-nightly-tag) + printf -- '- Resolution: reverse-matched same digest to nightly tag `%s`\n' \ + "${RESOLVED_NIGHTLY_TAG}" >> "$GITHUB_STEP_SUMMARY" + ;; + digest-pinned-latest) + printf -- '- Resolution: no nightly tag with the same digest was found, so this run pins `sglang-latest` to its current digest instead.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + reverse-lookup-failed-digest-pinned-latest) + printf -- '- Resolution: reverse lookup could not complete, so this run pins `sglang-latest` to its current digest instead of falling back to the floating tag.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + resolution-failed) + printf -- '- Resolution: digest pinning itself failed, so this run falls back to floating `sglang-latest`.\n' \ + >> "$GITHUB_STEP_SUMMARY" + ;; + esac fi printf -- '- Upload to dashboard: `%s`\n' \ @@ -220,7 +286,8 @@ jobs: run: | MODELS_JSON="$(jq -c ' map(select( - (.prefix == "deepseek-r1-fp8-tp8" and env.ENABLE_DEEPSEEK_R1_FP8_TP8 == "true") + env.GITHUB_EVENT_NAME == "schedule" + or (.prefix == "deepseek-r1-fp8-tp8" and env.ENABLE_DEEPSEEK_R1_FP8_TP8 == "true") or (.prefix == "deepseek-r1-fp8-tp4" and env.ENABLE_DEEPSEEK_R1_FP8_TP4 == "true") or (.prefix == "deepseek-r1-fp4-tp8" and env.ENABLE_DEEPSEEK_R1_FP4_TP8 == "true") or (.prefix == "deepseek-r1-fp4-tp4" and env.ENABLE_DEEPSEEK_R1_FP4_TP4 == "true") @@ -469,6 +536,11 @@ jobs: - name: Check if model is enabled id: check run: | + if [ "${GITHUB_EVENT_NAME}" = "schedule" ]; then + echo "enabled=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + case "${{ matrix.model.prefix }}" in deepseek-r1-fp8-tp8) echo "enabled=${{ inputs.deepseek-r1-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp8-tp4) echo "enabled=${{ inputs.deepseek-r1-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; diff --git a/.github/workflows/atom-sglang-test.yaml b/.github/workflows/atom-sglang-test.yaml index 4cd447d1da..71e1b95688 100644 --- a/.github/workflows/atom-sglang-test.yaml +++ b/.github/workflows/atom-sglang-test.yaml @@ -50,55 +50,140 @@ jobs: aiter_artifact_id: ${{ steps.download.outputs.aiter_artifact_id }} aiter_wheel_name: ${{ steps.download.outputs.aiter_wheel_name }} steps: - - name: Find and download latest aiter wheel + - name: Prefer latest main aiter wheel manifest and fallback to artifact id: download run: | set -euo pipefail - echo "=== Finding latest aiter-whl-main artifact from ROCm/aiter ===" + echo "=== Trying latest main aiter wheel manifest from S3 first ===" + S3_MAIN_MANIFEST_URL="https://rocm.frameworks-nightlies.amd.com/whl-staging/gfx942-gfx950/main/latest.json" API_URL="https://api.github.com" AUTH_HEADER="Authorization: token ${{ secrets.GITHUB_TOKEN }}" AITER_TEST_WORKFLOW_ID=179476100 - RUNS=$(curl -s -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/workflows/$AITER_TEST_WORKFLOW_ID/runs?per_page=100&branch=main&event=push") - ARTIFACT_ID="" ARTIFACT_NAME="" - for RUN_ID in $(echo "$RUNS" | jq -r '.workflow_runs[].id'); do - ARTIFACT_JSON=$(curl -s -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/runs/$RUN_ID/artifacts" \ - | jq '[.artifacts[] | select(.name | startswith("aiter-whl-main")) | select(.expired == false)] | first') - - if [ "$ARTIFACT_JSON" != "null" ] && [ -n "$ARTIFACT_JSON" ]; then - ARTIFACT_ID=$(echo "$ARTIFACT_JSON" | jq -r '.id') - ARTIFACT_NAME=$(echo "$ARTIFACT_JSON" | jq -r '.name') - echo "Found artifact in run $RUN_ID: $ARTIFACT_NAME (ID: $ARTIFACT_ID)" - break + ARTIFACT_RUN_ID="" + ARTIFACT_RUN_SHA="" + ARTIFACT_RUN_CREATED_AT="" + + resolve_download_url() { + python3 -c 'import sys + from urllib.parse import quote, unquote, urlsplit, urlunsplit + parts = urlsplit(sys.argv[1]) + encoded_path = "/".join(quote(unquote(segment), safe="") for segment in parts.path.split("/")) + print(urlunsplit((parts.scheme, parts.netloc, encoded_path, parts.query, parts.fragment)))' "$1" + } + + find_latest_artifact() { + local runs_json artifact_json run_id + + if [ -n "$ARTIFACT_ID" ] && [ "$ARTIFACT_ID" != "null" ]; then + return 0 fi - done - if [ -z "$ARTIFACT_ID" ] || [ "$ARTIFACT_ID" = "null" ]; then - echo "ERROR: No aiter-whl-main artifact found in recent Aiter Test runs" - exit 1 - fi + echo "=== Finding latest aiter-whl-* artifact from ROCm/aiter ===" + runs_json=$(curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/workflows/$AITER_TEST_WORKFLOW_ID/runs?per_page=100&branch=main&event=push") + + for run_id in $(echo "$runs_json" | jq -r '.workflow_runs[].id'); do + artifact_json=$(curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/runs/$run_id/artifacts" \ + | jq '[.artifacts[] | select(.name | startswith("aiter-whl-")) | select(.expired == false)] | sort_by(.created_at) | last') + + if [ "$artifact_json" != "null" ] && [ -n "$artifact_json" ]; then + ARTIFACT_ID=$(echo "$artifact_json" | jq -r '.id') + ARTIFACT_NAME=$(echo "$artifact_json" | jq -r '.name') + ARTIFACT_RUN_ID="$run_id" + ARTIFACT_RUN_SHA=$(echo "$runs_json" | jq -r --arg run_id "$run_id" '.workflow_runs[] | select((.id | tostring) == $run_id) | .head_sha') + ARTIFACT_RUN_CREATED_AT=$(echo "$runs_json" | jq -r --arg run_id "$run_id" '.workflow_runs[] | select((.id | tostring) == $run_id) | .created_at') + echo "Found artifact in run $ARTIFACT_RUN_ID: $ARTIFACT_NAME (ID: $ARTIFACT_ID, SHA: $ARTIFACT_RUN_SHA)" + return 0 + fi + done + + return 1 + } + + download_from_s3_manifest() { + local manifest_file manifest_fetch_url manifest_branch manifest_timestamp manifest_commit wheel_name wheel_url resolved_wheel_url + + mkdir -p aiter-whl + rm -f aiter-whl/amd_aiter*.whl + + manifest_file=$(mktemp) + trap 'rm -f "$manifest_file"' RETURN + manifest_fetch_url="${S3_MAIN_MANIFEST_URL}?ts=$(date +%s)" + curl -fsSL -H "Cache-Control: no-cache" "$manifest_fetch_url" -o "$manifest_file" || return 1 + + manifest_branch=$(jq -r '.branch // empty' "$manifest_file") + manifest_timestamp=$(jq -r '.timestamp // empty' "$manifest_file") + manifest_commit=$(jq -r '.commit // empty' "$manifest_file") + wheel_name=$(jq -r '.wheel_name // empty' "$manifest_file") + wheel_url=$(jq -r '.wheel_url // empty' "$manifest_file") + + if [ "$manifest_branch" != "main" ] || [ -z "$manifest_timestamp" ] || [ -z "$manifest_commit" ] || [ -z "$wheel_name" ] || [ -z "$wheel_url" ]; then + echo "Invalid latest main wheel manifest" + return 1 + fi - echo "=== Downloading artifact ===" - mkdir -p aiter-whl - curl -s -L -H "$AUTH_HEADER" \ - "$API_URL/repos/ROCm/aiter/actions/artifacts/$ARTIFACT_ID/zip" \ - -o aiter-whl.zip - unzip -o aiter-whl.zip -d aiter-whl - rm -f aiter-whl.zip + if find_latest_artifact; then + if [ -n "$ARTIFACT_RUN_SHA" ] && [ "$manifest_commit" != "$ARTIFACT_RUN_SHA" ]; then + if [ -n "$ARTIFACT_RUN_CREATED_AT" ] && [[ "$manifest_timestamp" < "$ARTIFACT_RUN_CREATED_AT" ]]; then + echo "Manifest commit $manifest_commit is older than latest artifact run $ARTIFACT_RUN_ID ($ARTIFACT_RUN_SHA); treating manifest as stale" + return 1 + fi + echo "Manifest commit $manifest_commit differs from latest artifact run $ARTIFACT_RUN_ID ($ARTIFACT_RUN_SHA), but manifest timestamp is not older" + fi + else + echo "No GitHub fallback artifact found while checking manifest freshness" + fi + + resolved_wheel_url=$(resolve_download_url "$wheel_url") + + echo "Selected latest main wheel manifest: $S3_MAIN_MANIFEST_URL" + echo "Manifest timestamp: $manifest_timestamp" + echo "Manifest commit: $manifest_commit" + echo "Manifest wheel: $wheel_name" + echo "Downloading manifest-selected wheel: $resolved_wheel_url" + curl -fsSL "$resolved_wheel_url" -o "aiter-whl/$wheel_name" || return 1 + echo "Downloaded wheel from manifest: aiter-whl/$wheel_name" + + rm -f "$manifest_file" + trap - RETURN + } + + download_from_artifact() { + echo "=== Falling back to latest aiter-whl-* artifact from ROCm/aiter ===" + find_latest_artifact || { + echo "ERROR: No aiter-whl-* artifact found in recent Aiter Test runs" + return 1 + } + + mkdir -p aiter-whl + rm -f aiter-whl/amd_aiter*.whl + curl -fsSL -H "$AUTH_HEADER" \ + "$API_URL/repos/ROCm/aiter/actions/artifacts/$ARTIFACT_ID/zip" \ + -o aiter-whl.zip + unzip -o aiter-whl.zip -d aiter-whl + rm -f aiter-whl.zip + } + + if download_from_s3_manifest; then + echo "Using wheel from S3 main manifest" + else + echo "Main wheel manifest download failed, falling back to GitHub artifact" + download_from_artifact + fi AITER_WHL=$(ls -t aiter-whl/amd_aiter*.whl 2>/dev/null | head -1) if [ -z "$AITER_WHL" ]; then - echo "ERROR: No amd_aiter wheel found in artifact" - ls -la aiter-whl/ + echo "ERROR: No amd_aiter wheel available after S3/artifact attempts" + ls -la aiter-whl/ || true exit 1 fi - echo "Downloaded wheel: $AITER_WHL" + echo "Selected wheel: $AITER_WHL" echo "aiter_artifact_id=${ARTIFACT_ID}" >> "$GITHUB_OUTPUT" echo "aiter_wheel_name=$(basename "$AITER_WHL")" >> "$GITHUB_OUTPUT" @@ -121,11 +206,19 @@ jobs: model_path: "deepseek-ai/DeepSeek-R1-0528" extra_args: "--tensor-parallel-size 4" env_vars: | - AITER_QUICK_REDUCE_QUANTIZATION=INT4 SGLANG_AITER_FP8_PREFILL_ATTN=0 SGLANG_USE_AITER=1 ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 - accuracy_test_threshold: 0.92 + accuracy_test_threshold: 0.91 + runner: linux-atom-mi35x-4 + - model_name: "DeepSeek-R1-FP4 TP4" + model_path: "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4" + extra_args: "--tensor-parallel-size 4" + env_vars: | + SGLANG_AITER_FP8_PREFILL_ATTN=0 + SGLANG_USE_AITER=1 + ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 + accuracy_test_threshold: 0.91 runner: linux-atom-mi35x-4 runs-on: ${{ matrix.runner }} timeout-minutes: 180 diff --git a/recipes/atom_sglang/DeepSeek-R1.md b/recipes/atom_sglang/DeepSeek-R1.md new file mode 100644 index 0000000000..917959a7c8 --- /dev/null +++ b/recipes/atom_sglang/DeepSeek-R1.md @@ -0,0 +1,146 @@ +# DeepSeek-R1 with ATOM SGLang Backend + +This recipe shows how to run `deepseek-ai/DeepSeek-R1-0528` or `amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4` with the SGLang-ATOM backend. For background on the SGLang-ATOM integration, see [Introduce ATOM as external model package of SGLang](https://github.com/ROCm/ATOM/issues/359). + +## Step 1: Pull the SGLang-ATOM Docker + +```bash +docker pull rocm/atom-dev:sglang-latest +``` + +Launch a container from this image and run the remaining commands inside the container. + +## Step 2: Launch SGLang-ATOM Server + +The SGLang-ATOM backend keeps the standard SGLang CLI, server APIs, and general usage flow compatible with upstream SGLang. For general server options and API usage, users can refer to the [official SGLang documentation](https://docs.sglang.ai/). + +Before launching the server, export the same SGLang-ATOM settings used by the benchmark workflow: + +```bash +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export SGLANG_AITER_FP8_PREFILL_ATTN=0 +export SGLANG_USE_AITER=1 +export ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 +# Introduce ATOM as external model package of SGLang +export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models +``` + +### DeepSeek with FP8 (TP=8) + +Users can use this command to launch the FP8 server with the same settings as the SGLang benchmark workflow. + +```bash +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-0528 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with FP8 (TP=4) + +```bash +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-0528 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with MXFP4 (TP=8) + +AMD Instinct MI355X GPU supports MXFP4 computation instructions, and users can use the following command to launch the MXFP4 server on MI355X. For MXFP4 model weight, we suggest using the checkpoint quantized from AMD Quark. + +```bash +python3 -m sglang.launch_server \ + --model-path amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +### DeepSeek with MXFP4 (TP=4) + +```bash +python3 -m sglang.launch_server \ + --model-path amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --host localhost \ + --port 8000 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv-cache-dtype fp8_e4m3 \ + --mem-fraction-static 0.8 \ + --page-size 1 \ + --disable-radix-cache +``` + +## Step 3: Performance Benchmark + +The SGLang benchmark workflow uses the `bench_serving` client for performance benchmarking. The following example matches the workflow command pattern for an MXFP4 TP4 case. + +```bash +git clone --depth 1 https://github.com/kimbochen/bench_serving.git /tmp/bench_serving + +ISL=8192 +OSL=1024 +CONC=64 +RANDOM_RANGE_RATIO=0.8 +RESULT_DIR=./benchmark-results +RESULT_FILENAME=deepseek-r1-fp4-tp4-${ISL}-${OSL}-${CONC}-${RANDOM_RANGE_RATIO}.json + +python3 /tmp/bench_serving/benchmark_serving.py \ + --model=amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4 \ + --backend=sglang \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len="${ISL}" \ + --random-output-len="${OSL}" \ + --random-range-ratio "${RANDOM_RANGE_RATIO}" \ + --num-prompts="$(( CONC * 10 ))" \ + --max-concurrency="${CONC}" \ + --trust-remote-code \ + --num-warmups="$(( 2 * CONC ))" \ + --request-rate=inf \ + --ignore-eos \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" \ + --result-dir="${RESULT_DIR}" \ + --result-filename="${RESULT_FILENAME}" +``` + +For FP8 or TP8 cases, keep the same benchmark command and replace `--model` with the checkpoint used in Step 2. + +### Optional: Enable Profiling +If you want to collect profiling trace, set the SGLang profiling environment variables before launching the server, and add `--profile` to the benchmark client command. + +```bash +export SGLANG_PROFILE_RECORD_SHAPES=1 +export SGLANG_PROFILE_WITH_STACK=1 +export SGLANG_TORCH_PROFILER_DIR=./profile_sglang/ +``` + +Then append `--profile` to the `benchmark_serving.py` command in Step 3. + +## Step 4: Accuracy Validation + +```bash +lm_eval --model local-completions \ + --model_args model=amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4,base_url=http://localhost:8000/v1/completions,num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ + --tasks gsm8k \ + --num_fewshot 3 +``` From 86a074f4b3832f87e7944e77a0594112dd58670a Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Tue, 12 May 2026 23:35:19 +0800 Subject: [PATCH 35/76] fix(benchmark): use fixed ISL/OSL for AW model benchmarks with vLLM bench (#767) AW (Actual Workload) models use `vllm bench serve` which requires `--random-range-ratio` in [0, 1) where 0 means fixed length. The value was incorrectly set to 1 which is outside the valid range and causes a ValueError in vLLM's random dataset sampler. - Remove `--random-range-ratio 1` from `bench_args` in all AW model configs (vLLM bench defaults to 0 which is already fixed length) - Set AW model `random_range_ratio` to empty in the matrix builder so the flag is omitted entirely, letting vLLM use its default - Use bash parameter expansion to conditionally pass the flag only when a non-empty value is configured Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 30 ++++++++++----------- .github/workflows/atom-vllm-benchmark.yaml | 4 +-- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 116d710144..a028e76f88 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -57,7 +57,7 @@ "display": "gpt-oss-120b TP1 (AW)", "dashboard_model": "gpt-oss-120b-aw-tp1", "prefix": "gpt-oss-120b-aw-tp1", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" }, @@ -66,7 +66,7 @@ "display": "gpt-oss-120b TP2 (AW)", "dashboard_model": "gpt-oss-120b-aw-tp2", "prefix": "gpt-oss-120b-aw-tp2", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" }, @@ -75,7 +75,7 @@ "display": "gpt-oss-120b TP8 (AW)", "dashboard_model": "gpt-oss-120b-aw-tp8", "prefix": "gpt-oss-120b-aw-tp8", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" } @@ -149,7 +149,7 @@ "display": "Kimi-K2.5-MXFP4 TP4 (AW)", "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp4", "prefix": "kimi-k2-5-mxfp4-aw-tp4", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, @@ -158,7 +158,7 @@ "display": "Kimi-K2.5-MXFP4 TP8 (AW)", "dashboard_model": "Kimi-K2.5-MXFP4-aw-tp8", "prefix": "kimi-k2-5-mxfp4-aw-tp8", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" } @@ -247,7 +247,7 @@ "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp1", "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" }, @@ -256,7 +256,7 @@ "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp2", "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" }, @@ -265,7 +265,7 @@ "display": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-aw-tp4", "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" } @@ -283,7 +283,7 @@ "display": "DeepSeek-V3.2 FP8 TP4 (AW)", "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", "prefix": "deepseek-v3-2-fp8-aw-tp4", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, @@ -292,7 +292,7 @@ "display": "DeepSeek-V3.2 FP8 TP8 (AW)", "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", "prefix": "deepseek-v3-2-fp8-aw-tp8", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" } @@ -310,7 +310,7 @@ "display": "GLM-4.7-FP8 TP4 (AW)", "dashboard_model": "GLM-4.7-FP8-aw-tp4", "prefix": "glm-4-7-fp8-aw-tp4", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, @@ -319,7 +319,7 @@ "display": "GLM-4.7-FP8 TP8 (AW)", "dashboard_model": "GLM-4.7-FP8-aw-tp8", "prefix": "glm-4-7-fp8-aw-tp8", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" } @@ -337,7 +337,7 @@ "display": "MiniMax-M2.5 TP2 (AW)", "dashboard_model": "MiniMax-M2.5-aw-tp2", "prefix": "minimax-m2-5-aw-tp2", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" }, @@ -346,7 +346,7 @@ "display": "MiniMax-M2.5 TP4 (AW)", "dashboard_model": "MiniMax-M2.5-aw-tp4", "prefix": "minimax-m2-5-aw-tp4", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" }, @@ -355,7 +355,7 @@ "display": "MiniMax-M2.5 TP8 (AW)", "dashboard_model": "MiniMax-M2.5-aw-tp8", "prefix": "minimax-m2-5-aw-tp8", - "bench_args": "--random-range-ratio 1", + "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1" } diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index c9efed1238..81ff787470 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -985,7 +985,7 @@ jobs: aws_param = dict(param) aws_param["input_length"] = input_length aws_param["output_length"] = output_length - aws_param["random_range_ratio"] = "1" + aws_param["random_range_ratio"] = "" base_params.append(aws_param) variants = [] @@ -1453,7 +1453,7 @@ jobs: --dataset-name random \ --random-input-len \"${ISL}\" \ --random-output-len \"${OSL}\" \ - --random-range-ratio \"${RANDOM_RANGE_RATIO}\" \ + ${RANDOM_RANGE_RATIO:+--random-range-ratio \"${RANDOM_RANGE_RATIO}\"} \ --temperature 0.0 \ --num-prompts \"$(( CONC * 10 ))\" \ --max-concurrency \"${CONC}\" \ From 88dd3d052b67a8b799004c7e2cf46524da4edcc5 Mon Sep 17 00:00:00 2001 From: wuhuikx Date: Wed, 13 May 2026 13:32:21 +0800 Subject: [PATCH 36/76] [atom-vllm] fix atom-vllm benchmark issue (#768) * fix(benchmark): safely print selected model JSON in workflow Move selected_group/models_json into environment variables before printing so JSON quotes and parentheses do not break shell parsing in scheduled runs. * fix(benchmark): add --gpu-memory-utilization 0.5 for gpt-oss AW TP2/TP8 The env var OOT_GPU_MEMORY_UTILIZATION=0.5 was set but never consumed by the server launch script, so AW variants ran with vLLM's default 0.9. Pass the flag explicitly via extra_args to match the MET variant behavior. Co-Authored-By: Claude Opus 4 * fix(benchmark): handle empty random_range_ratio in LLM-Booster data builder AW model benchmarks use an empty random_range_ratio, producing result filenames ending with a trailing dash (e.g. `prefix-1000-100-4-`). The parse_result_filename function called float("") on the empty trailing segment, causing a ValueError. Default to 0.0 when ratio is empty. Co-Authored-By: Claude Opus 4 * fix(benchmark): move nightly schedule from 23:00 to 21:00 Beijing time Co-Authored-By: Claude Opus 4 --------- Co-authored-by: XiaobingSuper Co-authored-by: zejunchen-zejun Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 4 ++-- .github/workflows/atom-vllm-benchmark.yaml | 15 +++++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index a028e76f88..6e5bff8dcd 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -67,7 +67,7 @@ "dashboard_model": "gpt-oss-120b-aw-tp2", "prefix": "gpt-oss-120b-aw-tp2", "bench_args": "", - "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 2 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" }, { @@ -76,7 +76,7 @@ "dashboard_model": "gpt-oss-120b-aw-tp8", "prefix": "gpt-oss-120b-aw-tp8", "bench_args": "", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" } ] diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index 81ff787470..a12b914832 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -7,8 +7,8 @@ concurrency: on: schedule: - # Nightly at 23:00 Beijing time (15:00 UTC) - - cron: '0 15 * * 0-4' + # Nightly at 21:00 Beijing time (13:00 UTC) + - cron: '0 13 * * 0-4' workflow_dispatch: inputs: benchmark_client: @@ -931,11 +931,14 @@ jobs: PY - name: Print selected models + env: + SELECTED_GROUP: ${{ steps.load.outputs.selected_group }} + MODELS_JSON: ${{ steps.load.outputs.models_json }} run: | - if [ -n "${{ steps.load.outputs.selected_group }}" ]; then - echo "Scheduled nightly benchmark group: ${{ steps.load.outputs.selected_group }}" + if [[ -n "${SELECTED_GROUP}" ]]; then + echo "Scheduled nightly benchmark group: ${SELECTED_GROUP}" fi - echo "Selected models: ${{ steps.load.outputs.models_json }}" + printf 'Selected models: %s\n' "${MODELS_JSON}" build-benchmark-matrix: name: Build OOT benchmark matrix @@ -1881,7 +1884,7 @@ jobs: def parse_result_filename(path: Path): stem = path.stem prefix, isl, osl, concurrency, ratio = stem.rsplit("-", 4) - return prefix, int(isl), int(osl), int(concurrency), float(ratio) + return prefix, int(isl), int(osl), int(concurrency), float(ratio) if ratio else 0.0 default_benchmark_client = os.environ["BENCHMARK_CLIENT"] From 7f0fbf43b6f03527e27de5240c1a633e7d05fc34 Mon Sep 17 00:00:00 2001 From: HaonanWang98 Date: Wed, 13 May 2026 15:47:08 +0800 Subject: [PATCH 37/76] fix hf model path (#770) --- atom/entrypoints/openai/chat_encoders.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/atom/entrypoints/openai/chat_encoders.py b/atom/entrypoints/openai/chat_encoders.py index aefde69ec2..d3fb0462ce 100644 --- a/atom/entrypoints/openai/chat_encoders.py +++ b/atom/entrypoints/openai/chat_encoders.py @@ -16,11 +16,22 @@ import os from typing import Any, Callable, List, Optional +from huggingface_hub import snapshot_download + logger = logging.getLogger("atom") MessageEncoder = Callable[..., str] +def _resolve_model_path(model: str) -> str: + if os.path.isdir(model): + return model + try: + return snapshot_download(model, local_files_only=True, allow_patterns=[]) + except Exception: + return model + + def _load_encoder_from_dir(model_path: str) -> Optional[MessageEncoder]: """Look for ``/encoding/encoding_*.py`` and load ``encode_messages``. @@ -71,7 +82,7 @@ def load_custom_message_encoder(model_path: str) -> Optional[MessageEncoder]: ``chat_template`` path. Result should be cached by the caller — this does filesystem IO and a Python import. """ - return _load_encoder_from_dir(model_path) + return _load_encoder_from_dir(_resolve_model_path(model_path)) def apply_chat_template( From 93158846612ff0916100e35a0d58c1dcb2f9d992 Mon Sep 17 00:00:00 2001 From: Shao-Chun Lee Date: Wed, 13 May 2026 02:48:23 -0500 Subject: [PATCH 38/76] [Triton] DSV4 fusions phase 1 (#704) * dsv4 fusions * update * get rid of bf16 to fp4 upcast * remove host side sync per decode step * ruff clean up * black * clean * rms_quant fix * remove single stream hack * fused_clamp_act_mul post MOE * move kv_norm into q_norm fusion * clean * fix * dummy --- atom/model_ops/attentions/deepseek_v4_attn.py | 2 +- atom/model_ops/fused_moe_triton.py | 14 +- atom/model_ops/v4_kernels/fused_compress.py | 4 +- atom/models/deepseek_v4.py | 226 ++++++++++++++++-- 4 files changed, 218 insertions(+), 28 deletions(-) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 395567671c..d7fd6dea5d 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -264,7 +264,7 @@ def _build_window_topk_batched( ) pos_col = positions.view(total, 1) sp_col = start_pos_per_token.view(total, 1) - neg1 = torch.tensor(-1, device=device, dtype=positions.dtype) + neg1 = positions.new_full((), -1) # Case A: sp == 0 (fresh prefill) — abs positions [pos-win+1, pos] clamped. case_a = (pos_col - window_size + 1).clamp(min=0) + arange_w diff --git a/atom/model_ops/fused_moe_triton.py b/atom/model_ops/fused_moe_triton.py index f535af6312..8bb1be4c30 100644 --- a/atom/model_ops/fused_moe_triton.py +++ b/atom/model_ops/fused_moe_triton.py @@ -29,6 +29,7 @@ from aiter.ops.triton.fusions.fused_routing_from_topk import ( fused_routing_from_topk as _aiter_fused_routing_from_topk, ) +from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul from atom.model_ops.utils import has_triton_kernels logger = logging.getLogger("atom") @@ -372,13 +373,14 @@ def triton_kernel_fused_experts( gammas=gammas if apply_router_weight_on_input else None, ) raw_2d = raw_intermediate.view(M * topk, N) - gate = raw_2d[:, :half_N] - up = raw_2d[:, half_N:] - if swiglu_limit > 0: - gate = gate.clamp(max=swiglu_limit) - up = up.clamp(-swiglu_limit, swiglu_limit) intermediate_cache = intermediate_cache.view(M * topk, half_N) - intermediate_cache.copy_(torch.nn.functional.silu(gate) * up) + fused_clamp_act_mul( + raw_2d, + out=intermediate_cache, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=None, + ) intermediate_cache = intermediate_cache.view(batch_dim, M * topk, half_N) matmul_ogs( diff --git a/atom/model_ops/v4_kernels/fused_compress.py b/atom/model_ops/v4_kernels/fused_compress.py index 9f5eeaf8d6..a8b52e5f75 100644 --- a/atom/model_ops/v4_kernels/fused_compress.py +++ b/atom/model_ops/v4_kernels/fused_compress.py @@ -207,13 +207,13 @@ def _fused_compress_attn_kernel( mask=d_mask, other=0.0, eviction_policy="evict_first", - ) + ).to(tl.float32) score_a = tl.load( score_in_ptr + in_row * score_in_row_stride + col_off + d, mask=d_mask, other=0.0, eviction_policy="evict_first", - ) + ).to(tl.float32) ape_v = tl.load( ape_ptr + ape_row * dim_full + col_off + d, mask=d_mask, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index a53eb3a50e..c82a1f1d8d 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -44,6 +44,12 @@ ) from aiter.ops.topk import top_k_per_row_decode, top_k_per_row_prefill from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits +from aiter.ops.triton.fusions.fused_clamp_act_mul import ( + fused_clamp_act_mul, +) +from aiter.ops.triton.fusions.fused_reduce_qk_norm_rope_swa_write import ( + fused_reduce_qk_norm_rope_swa_write, +) from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits from atom.config import ( Config, @@ -117,6 +123,9 @@ # forward burns syscalls (V4-Pro: 64 layers × multiple sites per call). _V4_FORCE_UE8M0_QUANT = os.environ.get("V4_FORCE_UE8M0_QUANT", "0") == "1" _V4_USE_REF_QUANT = os.environ.get("V4_USE_REF_QUANT", "0") == "1" +# Fused-kernel switches. Default off; flip via env to A/B against the eager path. +_V4_USE_TRITON_FUSION = os.environ.get("ATOM_V4_USE_TRITON_FUSION", "0") == "1" +ENABLE_DS_QKNORM_QUANT_FUSION = envs.ATOM_ENABLE_DS_QKNORM_QUANT_FUSION def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: @@ -126,6 +135,111 @@ def _rmsnorm_nw(x: torch.Tensor, eps: float, dim: int) -> torch.Tensor: return rmsnorm2d_fwd_(x, ones, eps, dim) +def _fused_qk_norm_rope_swa_write_fake( + q: torch.Tensor, + kv: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + n_local_heads: int, + head_dim: int, + rope_head_dim: int, + kv_weight: torch.Tensor, + eps: float, + win: int, + swa_write_indices: Optional[torch.Tensor] = None, + batch_id_per_token: Optional[torch.Tensor] = None, + state_slot_mapping: Optional[torch.Tensor] = None, + swa_kv: Optional[torch.Tensor] = None, +) -> torch.Tensor: + M = q.shape[0] + return torch.empty( + (M, n_local_heads, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + + +@torch_compile_guard(gen_fake=_fused_qk_norm_rope_swa_write_fake) +def fused_qk_norm_rope_swa_write( + q: torch.Tensor, + kv: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + n_local_heads: int, + head_dim: int, + rope_head_dim: int, + kv_weight: torch.Tensor, + eps: float, + win: int, + swa_write_indices: Optional[torch.Tensor] = None, + batch_id_per_token: Optional[torch.Tensor] = None, + state_slot_mapping: Optional[torch.Tensor] = None, + swa_kv: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Fused wq_b GEMM (a8w8 1x128 blockscale) + per-head RMSNorm-nw + RoPE on + q/kv tail (+ SWA write) in a single triton kernel. + + `kv` must already be kv_norm-applied; the kernel does not weight-norm kv. + `kv` is RoPE-mutated in place. When all SWA tensors are provided, the + kernel also writes the windowed kv into `swa_kv`; the caller must then + skip the standalone `swa_write` for those rows. + """ + num_tokens = q.shape[0] + if num_tokens <= 64: + q_out = torch.empty( + (num_tokens, n_local_heads, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + fused_reduce_qk_norm_rope_swa_write( + q, + kv, + None, + kv_weight, + eps, + eps, + rope_head_dim, + cos_cache, + sin_cache, + positions, + q_out=q_out, + is_neox=False, + dtype=torch.bfloat16, + write_indices=swa_write_indices, + batch_id_per_token=batch_id_per_token, + state_slot_mapping=state_slot_mapping, + swa_kv=swa_kv, + win=win, + ) + else: + q = q.view(num_tokens, n_local_heads, head_dim) + q_out = _rmsnorm_nw(q, eps, head_dim) + kv = rmsnorm2d_fwd_(kv, kv_weight, eps, kv.shape[-1]) + aiter.rope_cached_positions_2c_fwd_inplace( + q_out[..., -rope_head_dim:].view(1, num_tokens, -1, rope_head_dim), + kv[..., -rope_head_dim:].view(1, num_tokens, -1, rope_head_dim), + cos_cache, + sin_cache, + positions.view(1, num_tokens), + 1, + reuse_freqs_front_part=True, + nope_first=False, + ) + if swa_write_indices is not None: + swa_write( + kv, + swa_write_indices, + positions, + batch_id_per_token, + state_slot_mapping, + swa_kv, + win, + ) + return q_out + + def _make_weightless_rmsnorm(dim: int, eps: float) -> RMSNorm: """Build an `RMSNorm(dim, eps)` whose `.weight` is `None`. @@ -1385,6 +1499,8 @@ def __init__( self.alt_stream is not None and self.compressor is not None ) + self.use_fuse_qk_norm_rope_swa_write = _V4_USE_TRITON_FUSION + self.layer_name = prefix atom_config = get_current_atom_config() atom_config.compilation_config.static_forward_context[self.layer_name] = self @@ -1500,6 +1616,28 @@ def forward_impl( ratio = self.compress_ratio rd = self.rope_head_dim + # Idempotent one-time plumb of rotary_emb into compressor / indexer + # (and the indexer's inner compressor). `rotary_emb` is set by the + # owning layer after __init__, so this can't move into __init__. + if self.compress_ratio and self.compressor.rotary_emb is None: + self.compressor.rotary_emb = self.rotary_emb + if self.indexer is not None: + self.indexer.rotary_emb = self.rotary_emb + self.indexer.compressor.rotary_emb = self.rotary_emb + + # ===== Per-fwd metadata (built once in prepare_prefill/decode). ===== + # All per-fwd state read once. Production prepare_decode/prefill + # always populates these; warmup goes through the same path + # (`_populate_state_slot_mapping` falls back to slot 0). + # Cast to V4 typed metadata so V4-specific attribute access (v4_*, + # compress_plans, swa_write_indices, ...) is well-typed for pyright. + attn_md = cast("AttentionMetaData_DSV4", get_forward_context().attn_metadata) + compress_plans = attn_md.compress_plans + swa_write_indices = attn_md.swa_write_indices + v4_batch_id_per_token = attn_md.batch_id_per_token + block_tables_gpu = attn_md.block_tables + state_slot_mapping = attn_md.state_slot_mapping + # ----- Batched ops on full flat tensors ----- # `_V4_FORCE_UE8M0_QUANT` (module-level): round-trip x/qr to ue8m0-FP8 # to mirror the reference's `act_quant(scale_fmt="ue8m0")` Linear-input @@ -1529,16 +1667,51 @@ def forward_impl( not _V4_FORCE_UE8M0_QUANT ), "_V4_FORCE_UE8M0_QUANT incompatible with fused q_norm quant (qr is already FP8)" qr, qr_scale = self.q_norm(q_lora) - # Flat q_flat is [num_tokens, n_local_heads * head_dim]; DualRMSNorm - # views internally to per-head shape and returns flat. Single Triton - # launch fuses per-head Q RMSNorm (weightless) + KV RMSNorm (both - # head_dim=128), replacing the prior `_rmsnorm_nw + kv_norm` pair. - q_flat = self.wq_b(qr, x_scale=qr_scale) - q_flat, kv = self.qk_norm(q_flat, kv_pre) - q = q_flat.view(num_tokens, self.n_local_heads, self.head_dim) - # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes - # to (1, num_tokens, ...) for aiter's per-position rope kernel. - self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) + q = self.wq_b(qr, x_scale=qr_scale) + if attn_md.is_pure_decode and self.use_fuse_qk_norm_rope_swa_write: + # Fused: wq_b GEMM (a8w8 1x128 blockscale) + per-head RMSNorm-nw + # + RoPE on q tail + RoPE on kv tail (+ SWA write) in one triton + # kernel. KV RMSNorm stays out (kernel doesn't apply weighted norm + # to kv); the standalone swa_write below is gated off when this + # path runs since the kernel already wrote the window slots. + cos_cache, sin_cache = self.rotary_emb._caches(x.device) + if swa_write_indices is not None: + write_indices = swa_write_indices + batch_id = v4_batch_id_per_token + slot_map = state_slot_mapping + swa_kv_buf = self.swa_kv + else: + write_indices = None + batch_id = None + slot_map = None + swa_kv_buf = None + q = fused_qk_norm_rope_swa_write( + q, + kv_pre, + cos_cache, + sin_cache, + positions, + self.n_local_heads, + self.head_dim, + rd, + self.kv_norm.weight, + self.eps, + win, + swa_write_indices=write_indices, + batch_id_per_token=batch_id, + state_slot_mapping=slot_map, + swa_kv=swa_kv_buf, + ) + else: + # Flat q_flat is [num_tokens, n_local_heads * head_dim]; DualRMSNorm + # views internally to per-head shape and returns flat. Single Triton + # launch fuses per-head Q RMSNorm (weightless) + KV RMSNorm (both + # head_dim=128), replacing the prior `_rmsnorm_nw + kv_norm` pair. + q_flat, kv = self.qk_norm(q, kv_pre) + q = q_flat.view(num_tokens, self.n_local_heads, self.head_dim) + # q [S, H, D] / kv [S, head_dim] — rotary_emb internally unsqueezes + # to (1, num_tokens, ...) for aiter's per-position rope kernel. + self.rotary_emb(positions, q[..., -rd:], kv[..., -rd:]) if _V4_USE_REF_QUANT: act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) @@ -1598,15 +1771,16 @@ def forward_impl( # ring read; swa_write order is irrelevant in that subcase. q_sa = q.contiguous() if attn_md.is_pure_decode: - swa_write( - kv, - swa_write_indices, - positions, - v4_batch_id_per_token, - state_slot_mapping, - self.swa_kv, - win, - ) + if not self.use_fuse_qk_norm_rope_swa_write: + swa_write( + kv, + swa_write_indices, + positions, + v4_batch_id_per_token, + state_slot_mapping, + self.swa_kv, + win, + ) if ratio == 0: kv_indices = attn_md.kv_indices_swa kv_indptr = attn_md.kv_indptr_swa @@ -1769,6 +1943,10 @@ def __init__( prefix=f"{prefix}.w2", ) self.swiglu_limit = swiglu_limit + # Switch: route clamp + silu(gate)*up [+ weights] + per-token FP8 1x128 + # quant through a single aiter triton kernel. The fused kernel emits + # FP8 + scale; w2 accepts `x_scale` and skips its own quant step. + self.use_fused_clamp_act_mul = _V4_USE_TRITON_FUSION def forward( self, @@ -1782,6 +1960,16 @@ def forward( # silu/clamp/mul in fp32 internally regardless of input dtype, so we # feed the bf16 GEMM output directly. combined = self.gate_up_proj(x) # [num_tokens, 2*inter_dim_per_tp] + if self.use_fused_clamp_act_mul: + x_fp8, x_scale = fused_clamp_act_mul( + combined, + swiglu_limit=self.swiglu_limit, + activation="silu", + weights=weights, + dtype_quant=dtypes.fp8, + transpose_scale=True, + ) + return self.w2(x_fp8, x_scale=x_scale) out = torch.empty( (combined.shape[0], combined.shape[-1] // 2), dtype=dtype, From 954a66f8ce8e6d0c70129f11bfbf5817c71f888d Mon Sep 17 00:00:00 2001 From: honglie Date: Wed, 13 May 2026 17:28:40 +0800 Subject: [PATCH 39/76] [bugfix] Fix/mtp entry proj quant bf16 mismatch (#755) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [Qwen3-Next MTP] fix three independent root causes blocking FP8 spec decode Qwen3-Next-FP8 + speculative decoding (--method mtp) silently produced ~0% MTP accept rate end-to-end. Three independent root causes had to be fixed for MTP to work; an ablation table proves none can substitute for the others. (A) Startup AttributeError: hf_config has no `n_routed_experts` (atom/config.py) SpeculativeConfig.hf_config_override only synthesized `n_routed_experts` for `qwen3_5_mtp` model_type, leaving Qwen3-Next MTP without the field. Downstream code (loader.py, qwen3_next.py expert wiring, etc.) then AttributeError'd during MTP drafter load. Generalize the synthesis: prefer the field's own value if it exists (DeepSeek/GLM ship it natively as 256), else fall back to `num_experts` (Qwen3.5/Next carry that as 256/512), else leave unset (non-MoE / unknown). Apply the same priority chain to `n_shared_experts` so the explicit Qwen-only allowlist isn't needed. (B) MTP `fc` constructed without prefix= (atom/models/qwen3_next_mtp.py) Qwen3-Next-FP8 ckpts ship `mtp.fc.weight` as bf16 with no weight_scale, and HF quantization_config correctly lists `mtp.fc` in modules_to_not_convert. But the ColumnParallelLinear constructor was missing `prefix=`, so LinearBase queried the empty string against exclude_layers, missed, fell back to global FP8 spec, and built the layer with an uninitialized weight_scale + a dtype-cast path that silently turned the bf16 weight into garbage FP8. Mirror qwen3_5_mtp.py:50 which already passes prefix correctly. (C) packed_modules_mapping missing two entries inherited from base (atom/models/qwen3_next_mtp.py) Qwen3NextForCausalLM packs `.gate.` and `shared_expert_gate` into a shared fused gate parameter (qwen3_next.py:1022-1030). Qwen3NextMTP.packed_modules_mapping omitted both entries, so the loader silently dropped `mtp.layers.0.mlp.shared_expert_gate.weight` — the MTP shared-expert gate stayed at its torch.empty random value, separately enough to flatten accept rate to 0%. Copy the two entries from the base class. End-to-end on Qwen3-Next-80B-A3B-Instruct-FP8 / TP=4 / MI355X×4 with all three fixes: MTP accept rate goes from 0.00% (0/11000) to 60.14% (4210/7000), avg toks/fwd from 1.00 to 1.60, smoke wallclock from 69.8s to 20.3s (3.5× throughput). Repro logs in repro_main_logs/. Co-Authored-By: Claude Opus 4.7 (1M context) * [DeepSeek MTP] auto-exclude eh_proj from quantization when ckpt has no weight_scale DeepSeek MTP checkpoints (R1, R1-0528, V3, V3.2 in FP8) store eh_proj as bf16 with no companion weight_scale tensor, even though their HF quantization_config does NOT list eh_proj in modules_to_not_convert. After commit 5d34e9e wrapped eh_proj in ReplicatedLinear with the global quant_config, these layers fell back to global FP8 spec, got built with an uninitialized weight_scale, and weight_loader_process silently cast the bf16 weight into the FP8 slot — the resulting GEMM consumed garbage scales and MTP accept rate collapsed to ~0%. Looking at HF/quark/compressed-tensors metadata is not reliable: DeepSeek FP8 omits the exclude entry while GLM ships it; GLM-5.1-MXFP4 actually quantizes eh_proj and includes weight_scale on disk. The authoritative signal is the checkpoint itself. Add `ckpt_has_tensor_suffix(model_path, suffix)` in atom/models/utils.py that scans the safetensors index (or the file's metadata for single-file ckpts) for tensor keys ending with the requested suffix. DeepSeekMTP calls it for `eh_proj.weight_scale` at construction time; when missing it calls `apply_default_exclude_layers(["*.eh_proj"])` so ReplicatedLinear sees a no_quant LayerQuantConfig and stays in bf16. Read failures log a warning and return False — the safer default, since if the ckpt actually IS quantized the bf16 ReplicatedLinear will fail loudly at weight load on dtype/shape mismatch instead of silently producing 0% accept. Behavior matrix: DeepSeek R1/V3/V3.2 (FP8 + Quark): no scale -> excluded -> bf16 ✓ fix GLM-5/4.7/5.1 FP8 : no scale -> excluded (HF dup, no-op) ✓ GLM-5.1-MXFP4 : has scale -> not excluded -> MXFP4 ✓ Any bf16-base model : no quant_config side-effect ✓ The same helper is reused by MiMoV2FlashMTP in a follow-up commit. Co-Authored-By: Claude Opus 4.7 (1M context) * [MiMo MTP] fix o_proj BF16 fallback + restrict to single-step draft Two issues surfaced when running end-to-end openai_server --method mtp on MiMo-V2-Flash, plus a guardrail to keep the working configuration sticky. eh_proj BF16 fallback (sibling of the DeepSeek MTP fix in the previous commit). MiMo's MTP eh_proj is also BF16 with no weight_scale on disk while the global quant_config is FP8. Reuse `ckpt_has_tensor_suffix` from atom/models/utils.py: if no eh_proj.weight_scale tensor exists, call apply_default_exclude_layers(["*.eh_proj"]) so the ReplicatedLinear stays in BF16. Same corruption chain as DeepSeek — silent FP8 cast of BF16 weight, garbage GEMM, ~0% accept. self_attn.o_proj BF16 fallback (MiMo-specific). MiMo's checkpoint stores every self_attn.o_proj as BF16 with no weight_scale across all 51 layers (48 base + 3 MTP). Its HF quantization_config.ignored_layers correctly lists all 48 base-model o_proj entries (model.layers.0..47.self_attn.o_proj) but forgets the 3 MTP-layer entries. Without an exclude, MTP o_proj is constructed by ATOM as an FP8 ReplicatedLinear with an uninitialized weight_scale; the weight_loader silently casts the BF16 weight into the FP8 slot. Mirror the eh_proj fix shape: scan the ckpt for any *.self_attn.o_proj.weight_scale[_inv]; if absent, call apply_default_exclude_layers(["*.self_attn.o_proj"]). HF's existing base-layer entries deduplicate harmlessly. Restrict to --num-speculative-tokens=1. MiMo-V2-Flash ships 3 trained MTP layers (multi-head independent, paper accept rate needs sglang's per-layer extend-mode decode + per-layer KV pools), but ATOM's single-token decode loop only writes one slot per layer per propose, so layer i's KV develops gaps at slots that step j (j != i) wrote and accept rate craters to ~40%. The current ATOM path can only drive a single-step draft (mtp_k=1, equivalent to vLLM's _MIMO_V2_FLASH_NUM_MTP_LAYERS = 1). Assert num_speculative_tokens == 1 in MiMoV2FlashMTP.__init__ so users get a clear startup error instead of garbage outputs or a downstream AttributeError mid-decode. Note: num_mtp_layers is left as the original `getattr(config, "num_nextn_predict_layers", 1)` — MiMo's BASE hf_config does not define this field, so the default of 1 fires naturally and matches the assert above. The architecture cap is enforced by the assert; the model file is otherwise untouched. Verification (TP=8 on /data/amd_int/models/MiMo-V2-Flash): Pre-fix: load_model emitted "1 NOT loaded" (mtp_block o_proj weight_scale at layer 48) on startup. Post-fix: 0 load_model warnings on the layers we keep; MTP layer 0 constructed with eh_proj + o_proj on the BF16 path. NOTE on accept rate: end-to-end MTP accept rate on MiMo-V2-Flash remains ~0% even with all fixes applied. This is unrelated to MTP — the MiMo BASE model alone (server started without --method mtp) also produces incoherent output. MiMo base support in ATOM has at least one further unidentified issue outside the entry-projection quant scope. The fixes in this commit are still correct for what they cover (verified by the disappearance of all load_model warnings), they just can't be observed end-to-end until the base-model issue is also resolved. Co-Authored-By: Claude Opus 4.7 (1M context) * [CI] add MiMo-V2-Flash base + MTP1 nightly accuracy entries GSM8K 5-shot flexible-extract = 0.8294 (±0.0104) measured locally on FP8 KV + MTP1 + tp=4 (see elog.atom). Threshold set 2pp below the measurement to absorb full-eval noise. Base-only entry reuses the same baseline since MTP greedy decoding produces identical text (drafts are verified by the target). MTP entry pinned to --num-speculative-tokens 1 and -tp 4: ATOM hardcodes MiMo MTP to a single layer (matching vLLM's _MIMO_V2_FLASH_NUM_MTP_LAYERS = 1); asking for more spec tokens routes through MTP layers 1/2 whose KV is never populated, collapsing accept rate. Co-Authored-By: Claude Opus 4.7 (1M context) * [MTP config] derive n_shared_experts from ckpt instead of defaulting to 1 When the HF config doesn't ship n_shared_experts (Qwen3-Next, Qwen3.5-MoE) the previous fallback synthesized n_shared_experts=1 whenever n_routed_experts was set. That happened to match every shipped MTP model by coincidence, but it would silently fabricate a phantom shared block on any future pure-routed MoE that gains MTP support — the same failure mode that produced the R1 MTP eh_proj BF16/FP8 mismatch. Replace the coincidence with a checkpoint scan: read the safetensors index for shared_expert keys, parse a count from indexed layouts, default to 1 for the flat-block layout used by every released MTP checkpoint, and leave the field unset when nothing is found. Helper lives next to ckpt_has_tensor_suffix in atom/models/utils.py. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- .github/benchmark/models_accuracy.json | 24 +++++ atom/config.py | 37 +++++-- atom/models/deepseek_mtp.py | 16 +++- atom/models/mimo_v2_flash_mtp.py | 36 ++++++- atom/models/qwen3_next_mtp.py | 3 + atom/models/utils.py | 128 +++++++++++++++++++++++++ 6 files changed, 236 insertions(+), 8 deletions(-) diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 4887f3d126..2254dfd3fc 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -288,5 +288,29 @@ "accuracy_baseline": 0.9401, "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows MXFP4=0.9256, baseline=0.9401" + }, + { + "model_name": "MiMo-V2-Flash", + "model_path": "XiaomiMiMo/MiMo-V2-Flash", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --trust-remote-code --gpu-memory-utilization 0.8", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.81, + "accuracy_baseline": 0.8294, + "accuracy_baseline_model": "XiaomiMiMo/MiMo-V2-Flash", + "_baseline_note": "CI measured GSM8K 5-shot flexible-extract = 0.8294 (±0.0104) under MTP1 greedy decoding (see MiMo-V2-Flash MTP entry); base-only number is identical because MTP only proposes speculative drafts that the target verifies. tp pinned to 4 to match the MTP entry's setup (no separate base measurement)." + }, + { + "model_name": "MiMo-V2-Flash MTP", + "model_path": "XiaomiMiMo/MiMo-V2-Flash", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --trust-remote-code --gpu-memory-utilization 0.8 --method mtp --num-speculative-tokens 1", + "env_vars": "", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.81, + "accuracy_baseline": 0.8294, + "accuracy_baseline_model": "XiaomiMiMo/MiMo-V2-Flash", + "_baseline_note": "CI measured GSM8K 5-shot flexible-extract = 0.8294 (±0.0104) on FP8 KV + MTP1. tp MUST be 4 and num-speculative-tokens MUST be 1: ATOM only constructs MTP layer 0 (matches vLLM _MIMO_V2_FLASH_NUM_MTP_LAYERS=1); driving more spec tokens would route through layers 1/2 whose KV is never populated and accept rate craters." } ] diff --git a/atom/config.py b/atom/config.py index 8c1d5c819a..ac201601c7 100644 --- a/atom/config.py +++ b/atom/config.py @@ -754,7 +754,7 @@ def __post_init__(self): # For multimodal models, extract text_config if hasattr(self.draft_model_hf_config, "text_config"): self.draft_model_hf_config = self.draft_model_hf_config.text_config - self.hf_config_override(self.draft_model_hf_config) + self.hf_config_override(self.draft_model_hf_config, self.model) if self.method == "eagle3": if getattr(self.draft_model_hf_config, "kv_lora_rank", None): @@ -778,7 +778,9 @@ def __post_init__(self): self.use_aux_hidden_state = True @staticmethod - def hf_config_override(hf_config: PretrainedConfig) -> None: + def hf_config_override( + hf_config: PretrainedConfig, model_path: Optional[str] = None + ) -> None: # Eagle3 architecture mapping (architecture-level, not model_type) arch = (getattr(hf_config, "architectures", None) or [""])[0] if arch == "LlamaForCausalLMEagle3": @@ -806,10 +808,33 @@ def hf_config_override(hf_config: PretrainedConfig) -> None: "num_nextn_predict_layers": n_predict, "architectures": [arch], } - # Qwen3.5 MTP needs expert counts for MoE layer construction - if hf_config.model_type == "qwen3_5_mtp": - updates["n_shared_experts"] = 1 - updates["n_routed_experts"] = getattr(hf_config, "num_experts", 0) + # Naming differs across families: + # DeepSeek / GLM → already have `n_routed_experts` + # Qwen3.5 / Qwen3-Next → only carry `num_experts` + # non-MoE / unknown → leave unset (no MoE = no field) + n_routed = getattr( + hf_config, + "n_routed_experts", + getattr(hf_config, "num_experts", None), + ) + if n_routed is not None: + updates["n_routed_experts"] = n_routed + # n_shared_experts: prefer the field's own value if it exists + # (DeepSeek / GLM ship it natively); else scan the checkpoint + # for `shared_expert` weights and use the parsed count (1 for + # the flat-block layout every released model uses). Leaving + # the field unset for ckpts without shared experts avoids + # fabricating a phantom shared block (e.g. R1 MTP eh_proj + # would otherwise pull a stale BF16 weight_scale). + existing_n_shared = getattr(hf_config, "n_shared_experts", None) + if existing_n_shared is not None: + updates["n_shared_experts"] = existing_n_shared + else: + from atom.models.utils import ckpt_shared_expert_count + + n_shared = ckpt_shared_expert_count(model_path) + if n_shared > 0: + updates["n_shared_experts"] = n_shared hf_config.update(updates) diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index ec7abb2139..7062f419df 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -17,7 +17,7 @@ from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig from .deepseek_v2 import DeepseekV2DecoderLayer -from .utils import maybe_prefix +from .utils import ckpt_has_tensor_suffix, maybe_prefix class SharedHead(nn.Module): @@ -177,6 +177,20 @@ def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config + # Several MTP checkpoints (DeepSeek R1/V3/V3.2 FP8 + the Quark mixed + # MXFP4/FP8 variants) store eh_proj as BF16 with no weight_scale even + # though their HF quantization_config does not list eh_proj in the + # exclude set. Without this guard ReplicatedLinear is built with the + # global FP8/MXFP4 spec, the BF16 weight is cast into the FP8 slot + # against an uninitialized weight_scale, and MTP accept rate collapses. + # GLM-FP8 ckpts already list eh_proj explicitly (this becomes a no-op); + # GLM-5.1-MXFP4 truly quantizes eh_proj and ships weight_scale on disk + # so the check below leaves the global spec in effect. + if atom_config.quant_config is not None and not ckpt_has_tensor_suffix( + atom_config.model, "eh_proj.weight_scale" + ): + atom_config.quant_config.apply_default_exclude_layers(["*.eh_proj"]) + if hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None: self.packed_modules_mapping = { "q_a_proj": ("fused_qkv_a_proj", 0), diff --git a/atom/models/mimo_v2_flash_mtp.py b/atom/models/mimo_v2_flash_mtp.py index bb61741955..fef4b7a8f5 100644 --- a/atom/models/mimo_v2_flash_mtp.py +++ b/atom/models/mimo_v2_flash_mtp.py @@ -11,7 +11,7 @@ from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding from atom.model_ops.layernorm import RMSNorm from atom.model_ops.linear import ReplicatedLinear -from atom.models.utils import IntermediateTensors, maybe_prefix +from atom.models.utils import IntermediateTensors, ckpt_has_tensor_suffix, maybe_prefix from atom.utils.decorators import support_torch_compile from .mimo_v2_flash import MiMoV2Attention, MiMoV2MLP @@ -220,6 +220,40 @@ def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config self._draft_hf_config = atom_config.speculative_config.draft_model_hf_config + num_spec = atom_config.speculative_config.num_speculative_tokens + assert num_spec == 1, ( + f"MiMo-V2-Flash MTP only supports --num-speculative-tokens=1 now " + f"(got {num_spec})." + ) + + # See deepseek_mtp.DeepSeekMTP for the full rationale: MTP eh_proj is + # commonly stored as BF16 with no weight_scale even when the model's + # global quant_config is FP8/MXFP4. Skip quantization for eh_proj only + # when the checkpoint actually has no scale tensor for it. + if atom_config.quant_config is not None and not ckpt_has_tensor_suffix( + atom_config.model, "eh_proj.weight_scale" + ): + atom_config.quant_config.apply_default_exclude_layers(["*.eh_proj"]) + # MiMo additionally keeps every self_attn.o_proj as BF16. Its HF + # `ignored_layers` covers all 48 base-model o_proj entries but + # forgets the MTP-layer ones (model.mtp.layers.0..2.self_attn.o_proj). + # Without this exclude, MTP o_proj falls back to the global FP8 spec, + # weight_scale stays at torch.empty junk memory, and accept rate + # collapses — same corruption chain as the eh_proj case above. + # Detect from disk: if the ckpt has no o_proj.weight_scale[_inv] for + # any layer, exclude *.self_attn.o_proj globally (HF entries dedup). + if ( + atom_config.quant_config is not None + and not ckpt_has_tensor_suffix( + atom_config.model, "self_attn.o_proj.weight_scale_inv" + ) + and not ckpt_has_tensor_suffix( + atom_config.model, "self_attn.o_proj.weight_scale" + ) + ): + atom_config.quant_config.apply_default_exclude_layers( + ["*.self_attn.o_proj"] + ) self.model = MiMoV2FlashMultiTokenPredictor( atom_config=atom_config, prefix=maybe_prefix(prefix, "model") diff --git a/atom/models/qwen3_next_mtp.py b/atom/models/qwen3_next_mtp.py index 8fe57e7408..2a5f0737ee 100644 --- a/atom/models/qwen3_next_mtp.py +++ b/atom/models/qwen3_next_mtp.py @@ -43,6 +43,7 @@ def __init__(self, atom_config: Config, prefix: str = ""): self.config.hidden_size, bias=False, quant_config=quant_config, + prefix=f"{prefix}.fc", ) self.layers = torch.nn.ModuleList( @@ -107,6 +108,8 @@ class Qwen3NextMTP(nn.Module): "v_proj": ("qkv_proj", "v"), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), + ".gate.": (".gate.", 0), + "shared_expert_gate": ("gate", 1), } weights_mapping = {"mtp.": "model."} diff --git a/atom/models/utils.py b/atom/models/utils.py index e5d834af3a..afb5f0ea82 100644 --- a/atom/models/utils.py +++ b/atom/models/utils.py @@ -241,6 +241,134 @@ def maybe_prefix(prefix: str, name: str) -> str: return name if not prefix else f"{prefix}.{name}" +def ckpt_has_tensor_suffix(model_path: str, suffix: str) -> bool: + """Return True if the checkpoint at ``model_path`` contains a tensor whose + key ends with ``suffix``. + + Used by MTP modules to decide whether the entry projection (eh_proj/fc) is + actually quantized on disk: when the HF quantization_config's exclude list + is incomplete, falling back to the global quant spec creates an + uninitialized weight_scale that silently corrupts inference. Checking the + safetensors index for a sibling ``*.weight_scale`` is the only reliable + signal. + """ + import glob + import json + import os + + if not model_path or not os.path.isdir(model_path): + logger.warning( + "ckpt_has_tensor_suffix: model_path %r is not a directory; " + "assuming suffix %r is absent", + model_path, + suffix, + ) + return False + + index_files = glob.glob(os.path.join(model_path, "*.safetensors.index.json")) + if index_files: + try: + with open(index_files[0]) as f: + weight_map = json.load(f)["weight_map"] + except (OSError, json.JSONDecodeError, KeyError) as e: + logger.warning( + "ckpt_has_tensor_suffix: failed to read %s (%s); " + "assuming suffix %r is absent", + index_files[0], + e, + suffix, + ) + return False + return any(k.endswith(suffix) for k in weight_map) + + safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors")) + if safetensors_files: + try: + from safetensors import safe_open + + with safe_open(safetensors_files[0], framework="pt") as sf: + return any(k.endswith(suffix) for k in sf.keys()) + except Exception as e: + logger.warning( + "ckpt_has_tensor_suffix: failed to read %s (%s); " + "assuming suffix %r is absent", + safetensors_files[0], + e, + suffix, + ) + return False + + return False + + +def ckpt_shared_expert_count(model_path: str | None) -> int: + """Return the number of shared experts present in the local checkpoint. + + 0 means no `shared_expert` weights were found. 1 covers the flat-block + layout used by every MTP checkpoint shipped today (a single shared block + written without a numeric index, e.g. `...mlp.shared_experts.gate_proj`). + A value >1 is parsed from indexed keys (e.g. `shared_experts.0.`, + `shared_experts.1.`) and emits a warning, since no released model uses + that layout — the path exists to surface a mismatch instead of silently + loading the wrong count. + + Used by SpeculativeConfig to synthesize n_shared_experts when the HF + config doesn't carry it. Non-local paths or unreadable indexes return 0 + (leaving the field unset, safer than fabricating one). + """ + import glob + import json + import os + import re + + if not model_path or not os.path.isdir(model_path): + return 0 + + keys: list[str] | None = None + index_files = glob.glob(os.path.join(model_path, "*.safetensors.index.json")) + if index_files: + try: + with open(index_files[0]) as f: + keys = list(json.load(f)["weight_map"].keys()) + except (OSError, json.JSONDecodeError, KeyError): + return 0 + else: + safetensors_files = glob.glob(os.path.join(model_path, "*.safetensors")) + if safetensors_files: + try: + from safetensors import safe_open + + with safe_open(safetensors_files[0], framework="pt") as sf: + keys = list(sf.keys()) + except Exception: + return 0 + + if not keys: + return 0 + + shared_keys = [k for k in keys if "shared_expert" in k] + if not shared_keys: + return 0 + + indexed = re.compile(r"shared_experts?\.(\d+)\.") + max_idx = -1 + for k in shared_keys: + m = indexed.search(k) + if m: + max_idx = max(max_idx, int(m.group(1))) + if max_idx >= 0: + count = max_idx + 1 + logger.warning( + "Checkpoint at %r carries indexed shared_experts (parsed " + "n_shared_experts=%d). This layout is unprecedented in shipped " + "models; verify against the model definition before trusting it.", + model_path, + count, + ) + return count + return 1 + + def cast_overflow_tensors( tensors: torch.Tensor, offset: float = 1000, From 73d039d661379c3bab9650c02cbd4f443dcc233b Mon Sep 17 00:00:00 2001 From: Jiayun Date: Wed, 13 May 2026 17:51:19 +0800 Subject: [PATCH 40/76] Fpz/fix glm prefix (#762) * wip * a few fix * skip first * fix format --- atom/config.py | 12 +++++++++++- atom/model_engine/arg_utils.py | 8 +++++--- atom/model_engine/scheduler.py | 4 +++- atom/model_ops/attention_mla.py | 10 +++++----- atom/model_ops/attentions/aiter_mla.py | 4 ++-- atom/model_ops/base_attention.py | 6 ++---- 6 files changed, 28 insertions(+), 16 deletions(-) diff --git a/atom/config.py b/atom/config.py index ac201601c7..678d1e7806 100644 --- a/atom/config.py +++ b/atom/config.py @@ -878,7 +878,7 @@ class Config: kv_cache_block_size: int = 16 num_kvcache_blocks: int = -1 kv_cache_dtype: str = "bf16" - enable_prefix_caching: bool = False + enable_prefix_caching: bool = True port: int = 8006 torch_profiler_dir: str | None = field( default_factory=lambda: envs.ATOM_TORCH_PROFILER_DIR @@ -1021,6 +1021,16 @@ def __post_init__(self): v4_block_size = 128 if self.kv_cache_block_size != v4_block_size: self.kv_cache_block_size = v4_block_size + # TODO: V4's per-request SWA buffer cannot be restored from the classical + # KV pool on prefix cache hit, so disable prefix caching silently. + if self.enable_prefix_caching: + import logging + + logging.getLogger(__name__).warning( + "DeepSeek-V4 does not support prefix caching " + "(SWA buffer is not cacheable); disabling automatically." + ) + self.enable_prefix_caching = False def compute_hash(self) -> str: """ diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index f93c225ca1..7dae8df47c 100644 --- a/atom/model_engine/arg_utils.py +++ b/atom/model_engine/arg_utils.py @@ -31,7 +31,7 @@ class EngineArgs: tensor_parallel_size: int = 1 data_parallel_size: int = 1 enforce_eager: bool = False - enable_prefix_caching: bool = False + enable_prefix_caching: bool = True port: int = 8006 kv_cache_dtype: str = "bf16" block_size: int = 16 @@ -87,8 +87,10 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: ) parser.add_argument( "--enable_prefix_caching", - action="store_true", - help="Enable prefix caching.", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable prefix caching (default: enabled). " + "Use --no-enable_prefix_caching to disable.", ) parser.add_argument( "--port", diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 3d65369d4a..cfe874b275 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -558,9 +558,11 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: total_tokens_num_prefill = sum(num_scheduled_tokens) if num_seqs_prefill > 0: + cached_per_req = [s.num_cached_tokens for s in scheduled_seqs.values()] logger.info( f"Scheduled prefill batch: {num_seqs_prefill} reqs, " - f"{total_tokens_num_prefill} tokens, " + f"{total_tokens_num_prefill} new tokens " + f"(cached: {cached_per_req}, new: {num_scheduled_tokens}), " f"req_ids: {tuple(scheduled_seqs.keys())}" ) self.prev_prompt = True diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 3939532217..9e2faa889c 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -699,9 +699,7 @@ def forward_impl_server_mode( if context.is_prefill and not use_prefill_mla: use_prefix_cache = ( - attn_metadata.has_cached - and not is_rocm_aiter_fp4bmm_enabled() - and self.qk_nope_head_dim == self.v_head_dim + attn_metadata.has_cached and self.kv_b_proj.weight.dtype != dtypes.fp4x2 ) prefill_q = self.q_proj(q, x_scale=q_scale).view( @@ -736,7 +734,7 @@ def forward_impl_server_mode( ( attn_metadata.total_kv, self.num_heads, - self.qk_nope_head_dim, + self.v_head_dim, ), device=q.device, dtype=self.dtype, @@ -752,7 +750,9 @@ def forward_impl_server_mode( self.kv_b_proj.weight_scale, k_full, v_full, - weight_preshuffle=True, + weight_preshuffle=getattr( + self.kv_b_proj.weight, "is_shuffled", False + ), ) output = flash_attn_varlen_func( q=prefill_q, diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 68dad8f636..3f7ab3c83b 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -604,10 +604,10 @@ def prepare_prefill(self, batch: ScheduledBatch): if attn_metadata.has_cached: # Full context (cached + new): use cu_seqlens_k for indexer var["cu_seqlen_ks"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[:-1], counts + var["cu_seqlens_k"].np[:bs], counts ) var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[1:], counts + var["cu_seqlens_k"].np[1 : bs + 1], counts ) else: var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = ( diff --git a/atom/model_ops/base_attention.py b/atom/model_ops/base_attention.py index 3bf7760a03..baf6f4e898 100644 --- a/atom/model_ops/base_attention.py +++ b/atom/model_ops/base_attention.py @@ -100,10 +100,8 @@ def cp_mha_gather_cache_kernel( # per-tensor: one scale per ptr, no offset k_scale = tl.load(k_scale_ptr) v_scale = tl.load(v_scale_ptr) - k_dtype = k_reg.dtype - v_dtype = v_reg.dtype - k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype) - v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype) + k_reg = k_reg.to(tl.float32) * k_scale + v_reg = v_reg.to(tl.float32) * v_scale tl.store(key_ptr_offset + col_offsets, k_reg) tl.store(value_ptr_offset + col_offsets, v_reg) From 3a84df7f9596ac73f92a62ef16cb7417384a062f Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Wed, 13 May 2026 17:53:10 +0800 Subject: [PATCH 41/76] Support DeepSeek v4 TBO (#766) * Support DeepSeek v4 TBO * refactor some code --- atom/model_ops/attentions/deepseek_v4_attn.py | 116 ++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index d7fd6dea5d..86f5cfb12a 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -1061,6 +1061,122 @@ def prepare_prefill(self, batch: ScheduledBatch): ) return attn_metadata, positions + def build_ubatch_prefill_metadata( + self, + attn_metadata: AttentionMetaData, + ub_slice, + padded_bs: int, + ) -> AttentionMetaData_DSV4: + """Split prefill AttentionMetaData for V4 TBO micro-batches.""" + from atom.utils.tbo.ubatch_splitting import split_attn_metadata + + ub_attn = split_attn_metadata(attn_metadata, ub_slice, padded_bs) + ub_attn.__class__ = AttentionMetaData_DSV4 + + src = cast(AttentionMetaData_DSV4, attn_metadata) + rs = ub_slice.request_slice + ts = ub_slice.token_slice + ub_num_reqs = rs.stop - rs.start + ub_num_tokens = ts.stop - ts.start + + if src.state_slot_mapping is not None: + ub_attn.state_slot_mapping = src.state_slot_mapping[rs] + if src.state_slot_mapping_cpu is not None: + ub_attn.state_slot_mapping_cpu = src.state_slot_mapping_cpu[rs] + if src.start_pos_per_seq_cpu is not None: + ub_attn.start_pos_per_seq_cpu = src.start_pos_per_seq_cpu[rs] + + var = self.model_runner.forward_vars + positions_np = np.asarray(var["positions"].np[ts.start : ts.stop]) + full_cu = var["cu_seqlens_q"].np + ub_cu = np.asarray(full_cu[rs.start : rs.stop + 1], dtype=np.int32).copy() + ub_cu -= ub_cu[0] + + extend_lens_np = (ub_cu[1:] - ub_cu[:ub_num_reqs]).astype(np.int32) + context_lens_np = (ub_attn.start_pos_per_seq_cpu + extend_lens_np).astype( + np.int32 + ) + device = src.state_slot_mapping.device + + from atom.model_ops.v4_kernels import make_compress_plans + + if self._unique_compress_ratios_overlap: + ub_attn.compress_plans = make_compress_plans( + np.ascontiguousarray(extend_lens_np, dtype=np.int32), + np.ascontiguousarray(context_lens_np, dtype=np.int32), + self._unique_compress_ratios_overlap, + device, + plan_buffers=None, + decode_capacity_per_ratio=None, + ) + else: + ub_attn.compress_plans = {} + + # Multiple helpers read context_lens and block_tables from + # forward_vars by position [0:scheduled_bs]. For ubatch 1 the + # relevant rows live at [rs.start:rs.stop], not [0:ub_num_reqs]. + # Temporarily place the ubatch's slices at the front so helpers + # see the right values. + bt = var["block_tables"].np + saved_ctx = var["context_lens"].np[:ub_num_reqs].copy() + saved_bt = bt[:ub_num_reqs].copy() + try: + var["context_lens"].np[:ub_num_reqs] = context_lens_np + bt[:ub_num_reqs] = bt[rs.start : rs.stop].copy() + + self._attach_v4_per_fwd_meta( + ub_attn, + positions_np, + ub_cu, + ub_attn.start_pos_per_seq_cpu, + ub_attn.state_slot_mapping_cpu, + ub_num_reqs, + ub_num_tokens, + ) + + positions_gpu = var["positions"].gpu[ts.start : ts.stop] + self._attach_v4_indexer_meta( + ub_attn, + ub_cu, + ub_attn.start_pos_per_seq_cpu, + ub_num_reqs, + ub_num_tokens, + positions_gpu=positions_gpu, + ) + + self._build_paged_prefill_meta( + ub_attn, + positions_np, + ub_cu, + ub_attn.start_pos_per_seq_cpu, + ub_attn.state_slot_mapping_cpu, + ub_num_reqs, + ub_num_tokens, + ) + finally: + bt[:ub_num_reqs] = saved_bt + var["context_lens"].np[:ub_num_reqs] = saved_ctx + + # Clone all GPU tensors that are views into shared CpuGpuBuffers. + # Without this, building the next ubatch overwrites this ubatch's + # data via the same underlying buffer. + if ub_attn.batch_id_per_token is not None: + ub_attn.batch_id_per_token = ub_attn.batch_id_per_token.clone() + if ub_attn.n_committed_csa_per_seq is not None: + ub_attn.n_committed_csa_per_seq = ub_attn.n_committed_csa_per_seq.clone() + if ub_attn.swa_write_indices is not None: + ub_attn.swa_write_indices = ub_attn.swa_write_indices.clone() + if ub_attn.indexer_meta is not None: + im = ub_attn.indexer_meta + if im.get("cu_committed_gpu") is not None: + im["cu_committed_gpu"] = im["cu_committed_gpu"].clone() + if im.get("batch_id_per_token_gpu") is not None: + im["batch_id_per_token_gpu"] = im["batch_id_per_token_gpu"].clone() + if im.get("n_committed_per_seq_gpu") is not None: + im["n_committed_per_seq_gpu"] = im["n_committed_per_seq_gpu"].clone() + + return ub_attn + def _attach_v4_per_fwd_meta( self, attn_metadata: AttentionMetaData_DSV4, From 5fa47b3663687bda6c9ead1f6fb54967d52430d4 Mon Sep 17 00:00:00 2001 From: PerryZhang01 Date: Wed, 13 May 2026 21:32:07 +0800 Subject: [PATCH 42/76] [fix](gpt-oss): add int allreduce for gpt-oss (#776) Co-authored-by: perzhang --- .github/benchmark/models_accuracy.json | 2 +- .github/benchmark/oot_benchmark_models.json | 8 ++++---- .github/benchmark/oot_models_accuracy.json | 2 +- .github/scripts/atom_oot_test.sh | 1 - .github/workflows/atom-vllm-accuracy-validation.yaml | 2 +- 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 2254dfd3fc..ac440c08dd 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -125,7 +125,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --enable-dp-attention --enable-expert-parallel --enable-tbo all --gpu-memory-utilization 0.5", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://localhost:8000/v1/chat/completions,num_concurrent=65,max_retries=3,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot 3 --output_path ${OUTPUT_PATH}", - "env_vars": "", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4", "runner": "linux-atom-mi35x-4", "test_level": "main", "accuracy_threshold": 0.88, diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 6e5bff8dcd..1b54c6265e 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -50,7 +50,7 @@ "dashboard_model": "gpt-oss-120b-tp1", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" }, { "tp_size": 1, @@ -59,7 +59,7 @@ "prefix": "gpt-oss-120b-aw-tp1", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" }, { "tp_size": 2, @@ -68,7 +68,7 @@ "prefix": "gpt-oss-120b-aw-tp2", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" }, { "tp_size": 8, @@ -77,7 +77,7 @@ "prefix": "gpt-oss-120b-aw-tp8", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 8 --gpu-memory-utilization 0.5 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1\nOOT_GPU_MEMORY_UTILIZATION=0.5" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nGPTOSS_USE_GENERIC_SWIGLU_MXFP4_LAYOUT=1" } ] }, diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index fdbeb1bd9a..0b094337b6 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -165,7 +165,7 @@ "model_path": "openai/gpt-oss-120b", "extraArgs": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.88, diff --git a/.github/scripts/atom_oot_test.sh b/.github/scripts/atom_oot_test.sh index 69f51436e0..9546582342 100644 --- a/.github/scripts/atom_oot_test.sh +++ b/.github/scripts/atom_oot_test.sh @@ -46,7 +46,6 @@ RESULT_DIR=${RESULT_DIR:-/tmp/oot_accuracy_results} ACCURACY_LOG_FILE=${ACCURACY_LOG_FILE:-/tmp/oot_accuracy_output.txt} STREAM_VLLM_LOGS=${STREAM_VLLM_LOGS:-1} KEEP_SERVER_ALIVE_ON_EXIT=${KEEP_SERVER_ALIVE_ON_EXIT:-0} -OOT_GPU_MEMORY_UTILIZATION=${OOT_GPU_MEMORY_UTILIZATION:-0.9} EXPLICIT_MODEL_NAME=${OOT_MODEL_NAME:-} EXPLICIT_MODEL_PATH=${OOT_MODEL_PATH:-} EXPLICIT_EXTRA_ARGS=${OOT_EXTRA_ARGS:-} diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index ba4ad21248..0b2d7e7a0b 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -298,7 +298,7 @@ jobs: "extra_args": "--tensor-parallel-size 2", "client_command": "lm_eval --model local-chat-completions --apply_chat_template --model_args model=${MODEL_PATH},base_url=http://127.0.0.1:${VLLM_PORT}/v1/chat/completions,num_concurrent=65,max_retries=1,max_gen_toks=2048,tokenized_requests=False,trust_remote_code=True --tasks gsm8k --num_fewshot ${LM_EVAL_NUM_FEWSHOT} --output_path ${OUTPUT_PATH}", "accuracy_test_threshold": 0.88, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, { From 4d366135581378996ec43c216240f5456aa47ead Mon Sep 17 00:00:00 2001 From: carlushuang Date: Thu, 14 May 2026 00:28:58 +0800 Subject: [PATCH 43/76] feat(minimax): upgrade M2.5 to M2.7 + fix reasoning parser (#775) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(reasoning): handle without (MiniMax M2.7 pattern) MiniMax M2.7's chat template injects as part of the prompt, so the model output contains only (no start tag). The reasoning parser now splits at even without a preceding in both non-streaming (separate_reasoning) and streaming (ReasoningFilter) paths. * ci(benchmark): replace MiniMax M2.5 with M2.7 in CI Replace MiniMax-M2.5 → M2.7 and M2.5-MXFP4 → M2.7-MXFP4 across all benchmark and accuracy configs. Same architecture (MiniMaxM2ForCausalLM), M2.7 has better-trained weights. Updated accuracy baselines from M2.7 HF card: gsm8k=0.9181 (BF16), MXFP4=0.9189. MXFP4 model: amd/MiniMax-M2.7-MXFP4 (Quark quantized). Local perf verified on MI355X: M2.7 BF16 TP=2 matches M2.5 dashboard numbers within noise (817 vs 808 tok/s at c=4, 4745 vs 4685 at c=64). --- .github/benchmark/models.json | 12 +++++----- .github/benchmark/models_accuracy.json | 24 +++++++++---------- .github/benchmark/oot_models_accuracy.json | 12 +++++----- atom/entrypoints/openai/reasoning.py | 27 +++++++++++++++++++++- tests/entrypoints/test_reasoning.py | 14 +++++++++++ 5 files changed, 64 insertions(+), 25 deletions(-) diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 39aed57bd1..5f8c96e4e0 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -90,9 +90,9 @@ "env_vars": "" }, { - "display": "MiniMax-M2.5", - "path": "MiniMaxAI/MiniMax-M2.5", - "prefix": "MiniMax-M2.5", + "display": "MiniMax-M2.7", + "path": "MiniMaxAI/MiniMax-M2.7", + "prefix": "MiniMax-M2.7", "args": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "bench_args": "", "suffix": "", @@ -100,9 +100,9 @@ "env_vars": "" }, { - "display": "MiniMax-M2.5-MXFP4", - "path": "amd/MiniMax-M2.5-MXFP4", - "prefix": "MiniMax-M2.5-MXFP4", + "display": "MiniMax-M2.7-MXFP4", + "path": "amd/MiniMax-M2.7-MXFP4", + "prefix": "MiniMax-M2.7-MXFP4", "args": "--kv_cache_dtype fp8 -tp 1 --trust-remote-code", "bench_args": "", "suffix": "", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index ac440c08dd..92430ed674 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -266,28 +266,28 @@ "_baseline_note": "CI baseline=0.8605 (FP8 tp=4, 3-shot completions API, thinking mode active). HF card reports 0.9538 but uses chat API with reasoning_parser" }, { - "model_name": "MiniMax-M2.5", - "model_path": "MiniMaxAI/MiniMax-M2.5", + "model_name": "MiniMax-M2.7", + "model_path": "MiniMaxAI/MiniMax-M2.7", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "env_vars": "", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", - "accuracy_threshold": 0.92, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured: 0.9022 (gsm8k 3-shot flexible-extract). Threshold = baseline - 0.015." }, { - "model_name": "MiniMax-M2.5-MXFP4", - "model_path": "amd/MiniMax-M2.5-MXFP4", + "model_name": "MiniMax-M2.7-MXFP4", + "model_path": "amd/MiniMax-M2.7-MXFP4", "extraArgs": "--kv_cache_dtype fp8 --trust-remote-code", "env_vars": "", "runner": "atom-mi355-8gpu.predownload", "test_level": "nightly", - "accuracy_threshold": 0.91, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows MXFP4=0.9256, baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured BF16=0.9022 (gsm8k 3-shot flexible-extract). HF amd/MiniMax-M2.7-MXFP4: MXFP4=91.89, baseline=91.81 (percentage)." }, { "model_name": "MiMo-V2-Flash", diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index 0b094337b6..1050e3f856 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -103,16 +103,16 @@ "_baseline_note": "Reference value from recipes/atom_vllm/Kimi-K2.5.md" }, { - "model_name": "MiniMax-M2.5 TP2", - "model_path": "MiniMaxAI/MiniMax-M2.5", + "model_name": "MiniMax-M2.7 TP2", + "model_path": "MiniMaxAI/MiniMax-M2.7", "extraArgs": "--kv_cache_dtype fp8 -tp 2 --trust-remote-code", "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-8", "test_level": "nightly", - "accuracy_threshold": 0.92, - "accuracy_baseline": 0.9401, - "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.5", - "_baseline_note": "HF: amd/MiniMax-M2.5-MXFP4 card shows baseline=0.9401" + "accuracy_threshold": 0.8872, + "accuracy_baseline": 0.9022, + "accuracy_baseline_model": "MiniMaxAI/MiniMax-M2.7", + "_baseline_note": "ATOM CI measured: 0.9022 (gsm8k 3-shot flexible-extract). Threshold = baseline - 0.015." }, { "model_name": "DeepSeek-R1-FP8 TP8", diff --git a/atom/entrypoints/openai/reasoning.py b/atom/entrypoints/openai/reasoning.py index 7759a04816..8e5cd6f3d1 100644 --- a/atom/entrypoints/openai/reasoning.py +++ b/atom/entrypoints/openai/reasoning.py @@ -23,13 +23,22 @@ def separate_reasoning(text: str) -> Tuple[Optional[str], str]: Tuple of (reasoning_content, content). reasoning_content is None if no thinking block was found. """ - # Check for closed thinking block + # Check for closed thinking block: ... match = re.match(r"(.*?)\s*(.*)", text, flags=re.DOTALL) if match: reasoning = match.group(1).strip() content = match.group(2).strip() return (reasoning if reasoning else None, content) + # Check for without — models like MiniMax M2.7 don't + # generate (the chat template injects it as part of the prompt), + # so the model output starts with reasoning text directly. + if "" in text: + reasoning, _, content = text.partition("") + reasoning = reasoning.strip() + content = content.strip() + return (reasoning if reasoning else None, content) + # Check for unclosed thinking block (truncated response) match = re.match(r"(.*)", text, flags=re.DOTALL) if match: @@ -89,7 +98,23 @@ def process(self, text: str) -> list: elif self.buf: results.append(("reasoning_content", self.buf)) self.buf = "" + elif "" in self.buf: + # No but found — model started reasoning + # without tag (e.g., MiniMax M2.7 where the chat + # template injects as part of the prompt). + reasoning = self.buf.split("", 1)[0] + after = self.buf.split("", 1)[1].lstrip("\n") + if reasoning: + results.append(("reasoning_content", reasoning)) + self.state = 2 + self.buf = "" + if after: + results.extend(self._process_content(after)) elif len(self.buf) > 7 and "<" not in self.buf: + # No tag found — emit as content. For models that + # don't emit (MiniMax), streaming reasoning separation + # requires buffering the entire response, which is impractical. + # Non-streaming path handles this correctly via separate_reasoning(). results.append(("content", self.buf)) self.buf = "" diff --git a/tests/entrypoints/test_reasoning.py b/tests/entrypoints/test_reasoning.py index 0539528198..0961fe488e 100644 --- a/tests/entrypoints/test_reasoning.py +++ b/tests/entrypoints/test_reasoning.py @@ -82,6 +82,20 @@ def test_whitespace_after_thinking(self): reasoning, content = separate_reasoning(text) assert content == "The answer." + def test_no_think_start_tag(self): + """MiniMax M2.7 pattern: model doesn't generate , only . + The chat template injects as part of the prompt.""" + text = "The user wants hello world...\n\n\nprint('Hello')" + reasoning, content = separate_reasoning(text) + assert reasoning == "The user wants hello world..." + assert content == "print('Hello')" + + def test_no_think_start_tag_empty_content(self): + text = "Reasoning only\n" + reasoning, content = separate_reasoning(text) + assert reasoning == "Reasoning only" + assert content == "" + # ============================================================================ # ReasoningFilter (Streaming) Tests From 4d0aa8fcf2256432a366db11c979e66f79b4a2aa Mon Sep 17 00:00:00 2001 From: "Chuan (Richard) Li" Date: Wed, 13 May 2026 09:55:42 -0700 Subject: [PATCH 44/76] ci(dashboard): map `mi35x` runner hint to MI355X and add Radeon backfill tool (#756) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to #692. Two issues remained after that PR: 1. The accuracy-validation workflows use runner labels `linux-atom-mi35x-1` / `-4` / `-8` (family-style `mi35x`, not the chip-specific `mi355`). `collect_gpu_info.sh`'s case statement only matched `*mi355*`, so the runner-hint fallback never fired and these runners kept publishing `GPU: AMD Radeon Graphics` to the dashboard. The accuracy job for `gpt-oss-120b` alone has produced 64 mis-labeled data points since #692 merged. Add `*mi35x*` to the same case as `*mi355*`. 2. github-action-benchmark@v1 is append-only — every pre-#692 entry (and the post-#692 entries written before this fix) is permanently stuck on the dashboard with the wrong label. Add a small one-shot tool, `tools/backfill_dashboard_gpu_name.py`, that does a literal text replacement of `GPU: AMD Radeon Graphics` → `GPU: AMD Instinct MI355X` in `data.js`. Only the two labels have ever appeared in the file (verified: 6357 Radeon + 4321 MI355X, no MI300X / MI250X), so the substitution is safe and the action's exact JSON formatting is preserved. The script defaults to dry-run; pass `--write` to apply. The gh-pages backfill was already applied in a separate direct push so the dashboard reflects the fix immediately; the script is committed here so the rewrite is auditable and re-runnable if needed. Co-authored-by: Cursor Co-authored-by: XiaobingZhang --- .github/scripts/collect_gpu_info.sh | 8 +- tools/backfill_dashboard_gpu_name.py | 119 +++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 2 deletions(-) create mode 100644 tools/backfill_dashboard_gpu_name.py diff --git a/.github/scripts/collect_gpu_info.sh b/.github/scripts/collect_gpu_info.sh index d4a5723fdd..a98cd1c595 100755 --- a/.github/scripts/collect_gpu_info.sh +++ b/.github/scripts/collect_gpu_info.sh @@ -6,13 +6,17 @@ # 1. `amd-smi static --asic` MARKET_NAME # 2. `rocm-smi --showproductname` Card Series # 3. `rocminfo` Marketing Name -# 4. pattern match (mi355 / mi325 / mi300 / mi250) +# 4. pattern match (mi355 / mi35x / mi325 / mi300 / mi250) # # Step 4 is needed because on freshly-released ASICs (currently MI355X) every # in-container SMI tool can still report "Radeon Graphics" until the # marketing-name table is patched. The CI runner name is operator-asserted # and reliable, so we use it as the final tie-breaker for the dashboard label. # +# Note: `mi35x` is the family-style label used by accuracy-validation runners +# (`linux-atom-mi35x-1` / `-4` / `-8`). The MI35 series in ATOM CI currently +# only contains MI355X, so we map both `mi355` and `mi35x` to MI355X. +# # Usage: # collect_gpu_info.sh # local host # collect_gpu_info.sh # docker exec @@ -58,7 +62,7 @@ if { [ -z "${GPU_NAME:-}" ] || echo "$GPU_NAME" | grep -qi "Radeon Graphics"; } && [ -n "${RUNNER_HINT:-}" ]; then hint_lc=$(echo "$RUNNER_HINT" | tr '[:upper:]' '[:lower:]') case "$hint_lc" in - *mi355*) GPU_NAME="AMD Instinct MI355X" ;; + *mi355*|*mi35x*) GPU_NAME="AMD Instinct MI355X" ;; *mi325*) GPU_NAME="AMD Instinct MI325X" ;; *mi300x*|*mi300*) GPU_NAME="AMD Instinct MI300X" ;; *mi250x*|*mi250*) GPU_NAME="AMD Instinct MI250X" ;; diff --git a/tools/backfill_dashboard_gpu_name.py b/tools/backfill_dashboard_gpu_name.py new file mode 100644 index 0000000000..d9268fd3e4 --- /dev/null +++ b/tools/backfill_dashboard_gpu_name.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Backfill historical `AMD Radeon Graphics` labels in the dashboard data.js. + +`benchmark-action/github-action-benchmark@v1` (the action that powers +https://rocm.github.io/ATOM/benchmark-dashboard/) is append-only: every nightly +run pushes a new `{commit, date, benches}` entry to `benchmark-dashboard/data.js` +on the `gh-pages` branch and never rewrites history. + +That meant the entries written before PR #692 — and the entries written by +accuracy-validation workflows that use `linux-atom-mi35x-*` runners which a +follow-up patch fixed — are permanently stuck with the mis-detected +`GPU: AMD Radeon Graphics` label, even though every ATOM CI machine is +physically an MI355X. + +This script does a one-time, in-place rewrite of `data.js` to relabel those +entries. It is intentionally a literal text substitution (not a JSON +round-trip) so the action's exact formatting is preserved and the diff stays +inspectable. + +Assumptions (verified against the live data.js on 2026-05-11): + +* The dashboard only ever stores two GPU labels: `AMD Radeon Graphics` (6357 + occurrences, all mis-detections) and `AMD Instinct MI355X` (4321 correct). +* No MI300X / MI250X / MI210 entries have ever been written — the ATOM project + has only run on MI355X-class hardware since the dashboard was created. + +If either assumption stops holding (e.g. someone adds an MI300X runner), this +script needs to be tightened to gate the substitution on the runner hint. + +Usage: + # Inspect what would change (default — no writes): + python tools/backfill_dashboard_gpu_name.py path/to/data.js + + # Apply the rewrite in place: + python tools/backfill_dashboard_gpu_name.py path/to/data.js --write + + # Full end-to-end against the live gh-pages branch: + git fetch origin gh-pages + git worktree add /tmp/atom-gh-pages gh-pages + python tools/backfill_dashboard_gpu_name.py /tmp/atom-gh-pages/benchmark-dashboard/data.js --write + cd /tmp/atom-gh-pages + git add benchmark-dashboard/data.js + git commit -m "ci(dashboard): backfill MI355X label on historical Radeon entries" + git push origin gh-pages +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +WRONG = "GPU: AMD Radeon Graphics" +RIGHT = "GPU: AMD Instinct MI355X" + +# Defensive: only touch occurrences that appear inside a benchmark `extra` +# string. Every `extra` value the action writes starts with `Run: https://...`, +# so anchoring on `"extra": "Run:` makes accidental matches in commit messages +# or other free-form fields essentially impossible. +EXTRA_LINE_RE = re.compile(r'("extra":\s*"Run: [^"]*?)' + re.escape(WRONG)) + + +def rewrite(text: str) -> tuple[str, int]: + """Return (new_text, replacement_count).""" + new_text, count = EXTRA_LINE_RE.subn(r"\1" + RIGHT, text) + return new_text, count + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("data_js", type=Path, help="Path to benchmark-dashboard/data.js") + p.add_argument( + "--write", + action="store_true", + help="Rewrite the file in place (default is dry-run).", + ) + args = p.parse_args() + + if not args.data_js.is_file(): + print(f"error: {args.data_js} not found", file=sys.stderr) + return 2 + + original = args.data_js.read_text(encoding="utf-8") + pre_wrong = original.count(WRONG) + pre_right = original.count(RIGHT) + + new_text, n = rewrite(original) + post_wrong = new_text.count(WRONG) + post_right = new_text.count(RIGHT) + + print(f"file: {args.data_js}") + print(f"size: {len(original):,} bytes") + print(f"before: {pre_wrong:>6} Radeon, {pre_right:>6} MI355X") + print(f"after: {post_wrong:>6} Radeon, {post_right:>6} MI355X") + print(f"replacements made: {n:,}") + + if post_wrong != 0: + print( + f"warning: {post_wrong} Radeon mention(s) remain outside of " + f'`"extra": "Run: ..."` strings; left untouched.', + file=sys.stderr, + ) + + if not args.write: + print("\nDry-run only. Re-run with --write to apply.") + return 0 + + if n == 0: + print("\nNothing to write.") + return 0 + + args.data_js.write_text(new_text, encoding="utf-8") + print(f"\nWrote {args.data_js}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 72ebb0e9a4fb445b39eac5f24cd54eb7bc789e73 Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Thu, 14 May 2026 09:13:18 +0800 Subject: [PATCH 45/76] [Feat][Plugin] Enable MTP for vLLM Plugin (#557) * adapt mtp Signed-off-by: whx-sjtu * fix _is_draft_layer check Signed-off-by: whx-sjtu * fix review Signed-off-by: whx-sjtu * add mtp forward wrapper Signed-off-by: whx-sjtu * fix max_qo_len bug Signed-off-by: whx-sjtu --------- Signed-off-by: whx-sjtu --- atom/config.py | 24 +++ atom/model_loader/loader.py | 21 +- atom/models/deepseek_mtp.py | 17 +- atom/plugin/attention.py | 57 ++++-- .../vllm/attention_backend/mla_sparse.py | 8 +- atom/plugin/vllm/model_wrapper.py | 188 +++++++++++++++++- atom/plugin/vllm/register.py | 3 + atom/plugin/vllm/spec_decode_patch.py | 35 ++++ 8 files changed, 319 insertions(+), 34 deletions(-) create mode 100644 atom/plugin/vllm/spec_decode_patch.py diff --git a/atom/config.py b/atom/config.py index 678d1e7806..e39df36ac2 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1075,6 +1075,30 @@ def set_current_atom_config(atom_config: Config): _current_atom_config = atom_config +def _get_current_atom_config_from_vllm_forward_context() -> Optional[Config]: + # In vLLM plugin mode (especially speculative decode), main/draft models + # can coexist in one process. Resolve per-forward config first to avoid + # reading a stale global singleton. + try: + from vllm.forward_context import ( + get_forward_context as get_vllm_forward_context, + is_forward_context_available, + ) + except Exception: + return None + if not is_forward_context_available(): + return None + try: + return get_vllm_forward_context().additional_kwargs.get("atom_config") + except Exception: + return None + + def get_current_atom_config() -> Config: + # Try to get the atom config from forward context first in vLLM plugin mode. + if is_vllm(): + forward_atom_config = _get_current_atom_config_from_vllm_forward_context() + if forward_atom_config is not None: + return forward_atom_config assert _current_atom_config is not None, "Current atom config is not set" return _current_atom_config diff --git a/atom/model_loader/loader.py b/atom/model_loader/loader.py index df2fbf02ad..0763179f51 100644 --- a/atom/model_loader/loader.py +++ b/atom/model_loader/loader.py @@ -196,6 +196,8 @@ def load_model_in_plugin_mode( prefix: str = "", weights_mapper: WeightsMapper | None = None, load_fused_expert_weights_fn=None, + spec_decode: bool = False, + hf_config_override: AutoConfig | None = None, ) -> set[str]: # during loading model, the outplace operation may consume more @@ -216,17 +218,24 @@ def _empty_cache(): model_name_or_path = config.plugin_config.model_config.model_path _empty_cache() - config_for_loading = ( - config.hf_config.text_config - if hasattr(config.hf_config, "text_config") - else config.hf_config - ) + if hf_config_override is not None: + config_for_loading = getattr( + hf_config_override, "hf_config", hf_config_override + ) + if hasattr(config_for_loading, "text_config"): + config_for_loading = config_for_loading.text_config + else: + config_for_loading = ( + config.hf_config.text_config + if hasattr(config.hf_config, "text_config") + else config.hf_config + ) loaded_weights_record = load_model( model=model, model_name_or_path=model_name_or_path, hf_config=config_for_loading, load_dummy=config.load_dummy, - spec_decode=False, + spec_decode=spec_decode, prefix=prefix, is_plugin_mode=True, weights_mapper=weights_mapper, diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index 7062f419df..952c4d1039 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -63,6 +63,20 @@ def __init__( prefix=maybe_prefix(prefix, "eh_proj"), ) + if hasattr(config, "index_topk"): + max_num_batched_tokens = getattr( + atom_config, "max_num_batched_tokens", atom_config.max_num_seqs + ) + buffer_device = getattr(atom_config, "device", "cuda") + topk_indices_buffer = torch.empty( + max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=buffer_device, + ) + else: + topk_indices_buffer = None + self.shared_head = SharedHead( config=config, prefix=prefix, quant_config=atom_config.quant_config ) @@ -88,9 +102,6 @@ def forward( spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - # masked_inputs_embeds = torch.where( - # positions.unsqueeze(-1) == 0, 0, inputs_embeds - # ) masked_inputs_embeds = inputs_embeds inputs_embeds = self.enorm(masked_inputs_embeds) previous_hidden_states = self.hnorm(previous_hidden_states) diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index d6aaf4f11e..9c674c1cd4 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -840,12 +840,8 @@ def _build_decode( ) paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs] - max_qo_len = ( - (query_start_loc_cpu[-1] - query_start_loc_cpu[-2]).item() - if query_start_loc_cpu.numel() > 1 - else 1 - ) - + qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + max_qo_len = qo_len.max().item() kv_indices_generate_triton( block_table_tensor, self.paged_kv_indices, @@ -882,7 +878,9 @@ def _build_decode( self.qo_indptr[1 + num_reqs :] = num_decode_tokens qo_indptr = self.qo_indptr[: 1 + num_reqs] - ctx_mla_ps = self._set_mla_persistent_worker_buffers(num_reqs, qo_indptr, 1) + ctx_mla_ps = self._set_mla_persistent_worker_buffers( + num_reqs, qo_indptr, max_qo_len + ) self.mla_persistent_metadata.update(ctx_mla_ps) attn_metadata = AiterMLADecodeMetadataForPluginMode( @@ -936,7 +934,6 @@ def build( query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu seq_lens = common_attn_metadata.seq_lens dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, @@ -2136,6 +2133,7 @@ def init_method_under_plugin_mode( from vllm.utils.platform_utils import num_compute_units except ImportError: from vllm.utils.platform_utils import get_cu_count as num_compute_units + from atom.models.utils import extract_layer_index from vllm.v1.worker.cp_utils import get_total_cp_world_size from vllm.utils.math_utils import cdiv @@ -2146,16 +2144,30 @@ def init_method_under_plugin_mode( self.kv_cache_spec = kv_cache_spec self.device = device max_num_batched_tokens = config.scheduler_config.max_num_batched_tokens + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) self.max_prefill_buffer_size = get_max_prefill_buffer_size( self.model_config.max_model_len ) - self.num_speculative_tokens = ( - self.vllm_config.speculative_config.num_speculative_tokens - if self.vllm_config.speculative_config - else 0 - ) - self.reorder_batch_threshold += self.num_speculative_tokens + # Determine if this builder is for draft model layers (MTP). + # Draft model layers have layer indices >= num_hidden_layers. + # The draft model itself does not do speculative decoding, so + # num_speculative_tokens should be 0 for its builders. + _is_draft_layer = False + if layer_names: + num_hidden_layers = config.model_config.hf_config.num_hidden_layers + layer_indices = [ + extract_layer_index(layer_name) for layer_name in layer_names + ] + _is_draft_layer = all(idx >= num_hidden_layers for idx in layer_indices) + if _is_draft_layer: + self.num_speculative_tokens = 0 + else: + self.num_speculative_tokens = ( + self.vllm_config.speculative_config.num_speculative_tokens + if self.vllm_config.speculative_config + else 0 + ) sm_count = num_compute_units(self.device.index) self.num_sms = sm_count @@ -2222,6 +2234,7 @@ def init_method_under_plugin_mode( self.kv_cache_spec = kv_cache_spec self.device = device max_num_batched_tokens = config.scheduler_config.max_num_batched_tokens + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) parallel_config = config.parallel_config self.num_heads = self.model_config.get_num_attention_heads(parallel_config) @@ -2275,8 +2288,9 @@ def setup_mla_sparse_attn_metadata_builder_base_class_and_attributes(class_dict: generic_base = AttentionMetadataBuilder needs_generic = True - # align with vllm ROCMAiterMLASparseMetadataBuilder - class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + # The plugin sparse MLA path builds per-token ragged metadata, so uniform + # spec-decode batches from the main model can use full decode cudagraphs. + class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_BATCH # For indexer metadata, not present in vllm sparse MLA class_dict["reorder_batch_threshold"] = 1 @@ -2297,12 +2311,16 @@ def unified_attention_with_output_base_for_plugin_mode( use_mla: bool, qkv: torch.Tensor, ) -> torch.Tensor: - atom_config = get_current_atom_config() + current_atom_config = get_current_atom_config() + static_forward_context = ( + current_atom_config.compilation_config.static_forward_context + ) + if use_mla: # raise NotImplementedError("MLA is not supported for plugin mode for now") kv_c_normed = k k_pe = v - self = atom_config.compilation_config.static_forward_context[layer_name] + self = static_forward_context[layer_name] q = self.q_proj(q, q_scale) q = q.view(-1, self.num_heads, self.qk_head_dim) # Add head dim of 1 to k_pe @@ -2321,8 +2339,7 @@ def unified_attention_with_output_base_for_plugin_mode( ) return self.o_proj(output) else: - self = atom_config.compilation_config.static_forward_context[layer_name] - + self = static_forward_context[layer_name] # here is the standard vllm attention impl interface # when using fusion, we need to pass the qkv and positions through the q,k,v # [watch out] accept_output_buffer must be False for plugin mode diff --git a/atom/plugin/vllm/attention_backend/mla_sparse.py b/atom/plugin/vllm/attention_backend/mla_sparse.py index 595dbad965..d8d8e27b9f 100644 --- a/atom/plugin/vllm/attention_backend/mla_sparse.py +++ b/atom/plugin/vllm/attention_backend/mla_sparse.py @@ -23,7 +23,7 @@ class AiterMLASparseBackend(AttentionBackend): @staticmethod def get_name() -> str: - return "ROCM_AITER_MLA_SPARSE" if not is_plugin_mode() else "CUSTOM" + return "ROCM_AITER_MLA_SPARSE" if not is_plugin_mode() else "CUSTOM_MLA_SPARSE" @staticmethod def get_builder_cls() -> Type["AiterMLASparseMetadataBuilder"]: @@ -59,7 +59,11 @@ class AiterMLASparseIndexerBackend(AttentionBackend): @staticmethod def get_name() -> str: - return "ROCM_AITER_MLA_SPARSE_INDEXER" if not is_plugin_mode() else "CUSTOM" + return ( + "ROCM_AITER_MLA_SPARSE_INDEXER" + if not is_plugin_mode() + else "CUSTOM_MLA_SPARSE_INDEXER" + ) @staticmethod def get_builder_cls() -> Type["AiterMLASparseIndexerMetadataBuilder"]: diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index 1eeee3ddb9..c2b990c182 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -1,6 +1,8 @@ from collections.abc import Iterable +import functools import importlib +import types import torch import torch.nn as nn from aiter.dist.parallel_state import ( @@ -43,6 +45,7 @@ "DeepseekV32ForCausalLM": "atom.models.deepseek_v2:DeepseekV3ForCausalLM", "Glm4MoeForCausalLM": "atom.models.glm4_moe:Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM": "atom.models.deepseek_v2:GlmMoeDsaForCausalLM", + "DeepSeekMTPModel": "atom.models.deepseek_mtp:DeepSeekMTP", "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLM", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration_", "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration_", @@ -73,7 +76,46 @@ def _prepare_env(atom_config) -> None: init_aiter_dist(config=atom_config) +def _safe_get_first_arch(config_like) -> str | None: + if config_like is None: + return None + architectures = getattr(config_like, "architectures", None) + if isinstance(architectures, list) and len(architectures) > 0: + return architectures[0] + return None + + +def _select_model_arch(vllm_config: VllmConfig) -> str: + model_arch = _safe_get_first_arch(getattr(vllm_config, "model_config", None)) + if model_arch is None: + raise ValueError("Cannot determine model architecture from vLLM model_config") + speculative_config = getattr(vllm_config, "speculative_config", None) + draft_model_config = getattr(speculative_config, "draft_model_config", None) + draft_arch = _safe_get_first_arch(draft_model_config) + if draft_arch is None: + return model_arch + model_tag = None + try: + from vllm.compilation import backends as vllm_backends + + model_tag = getattr(vllm_backends, "model_tag", None) + except Exception: + pass + if model_tag is None: + model_tag = getattr( + getattr(vllm_config, "compilation_config", None), "model_tag", None + ) + if model_tag in {"eagle_head", "draft_model", "drafter"}: + logger.info( + f"Use draft model architecture {draft_arch} for speculative tag {model_tag}" + ) + return draft_arch + return model_arch + + class ATOMModelBase(nn.Module, VllmModel, SupportsQuant, SupportsPP): + # forced_model_arch: str | None = None + def __init_subclass__(cls, *args, **kwargs): super().__init_subclass__(*args, **kwargs) @@ -99,10 +141,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vllm_config = vllm_config self.atom_config = generate_atom_config_for_plugin_mode(vllm_config) + self.is_mtp = False + speculative_config = getattr(vllm_config, "speculative_config", None) + if speculative_config is not None: + spec_method = speculative_config.method + self.is_mtp = spec_method == "mtp" _prepare_env(atom_config=self.atom_config) - model_arch = vllm_config.model_config.architectures[0] + main_model_arch = vllm_config.model_config.architectures[0] + model_arch = _select_model_arch(vllm_config) + self.is_mtp_draft_model = self.is_mtp and model_arch != main_model_arch + self.model_arch = model_arch model_cls = _get_atom_model_cls(model_arch) module_remapping = getattr(model_cls, "packed_modules_mapping", {}) weights_mapper = getattr(model_cls, "hf_to_atom_mapper", {}) @@ -116,6 +166,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # to ATOM's format. It is invoked in ATOM's model_runner initialization, but # lacks correspondences in vLLM. So we invoke the translation here for vLLM OOT. exclude_mapping = getattr(model_cls, "quant_exclude_name_mapping", {}) + # add exclude mapping for mtp layer of GLM5. + if model_arch != main_model_arch and main_model_arch == "GlmMoeDsaForCausalLM": + exclude_mapping.update( + { + "indexers_proj": "indexer.weights_proj", + } + ) if exclude_mapping and self.atom_config.quant_config is not None: self.atom_config.quant_config.apply_exclude_name_mapping(exclude_mapping) @@ -125,6 +182,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): logger.info(f"Construct ATOM model {model_arch} for vLLM plugin mode") self.model = model_cls(self.atom_config) + self._adapt_mtp_layers_for_vllm() + # Mirror nested attributes required by vLLM speculative decoding. + self._expose_spec_decode_attrs() # For sparse MLA, register the Indexer's DeepseekV32IndexerCache as # a virtual subclass of vLLM's AttentionLayerBase so vLLM can discover @@ -141,6 +201,58 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.pp_group = get_pp_group() self.tp_group = get_tp_group() + # Attributes whose writes on the outer model must propagate to the + # inner model so vLLM's weight-sharing reaches the forward path. + _WEIGHT_SHARED_ATTRS = frozenset({"embed_tokens", "embedding", "lm_head"}) + + def _expose_spec_decode_attrs(self) -> None: + """Bridge the extra nesting level between vLLM and ATOM for spec decode. + + ATOM wraps the HF model with one extra level: + vLLM sees: wrapper.model (DeepSeekMTP) + forward uses: .model (DeepSeekMultiTokenPredictor) + + vLLM's EagleSpeculator reads/writes embed_tokens, lm_head, layers on + the outer model. The forward path reads them from the inner model. + + We need two things: + 1. Mirror inner → outer so vLLM can discover the attrs. + 2. When vLLM later *replaces* embed_tokens / lm_head with shared + target-model weights, propagate the write to the inner model + so the forward path picks up the shared tensor. + """ + model = self.model + inner = getattr(model, "model", None) + if inner is None: + if hasattr(model, "lm_head") and not hasattr(self, "lm_head"): + self.lm_head = model.lm_head + return + + # (1) Mirror: make attrs visible on the outer model for vLLM discovery. + for attr in (*self._WEIGHT_SHARED_ATTRS, "layers"): + if not hasattr(model, attr) and hasattr(inner, attr): + setattr(model, attr, getattr(inner, attr)) + + if not hasattr(self, "lm_head") and hasattr(model, "lm_head"): + self.lm_head = model.lm_head + + # (2) Propagate: future writes on the outer model sync to the inner + # model. We create a one-off subclass so the hook only affects + # this particular draft-model instance, not the base class. + shared = self._WEIGHT_SHARED_ATTRS + base_setattr = model.__class__.__setattr__ + + def _syncing_setattr(self_model, name, value): + base_setattr(self_model, name, value) + if name in shared and hasattr(inner, name): + base_setattr(inner, name, value) + + model.__class__ = type( + model.__class__.__name__, + (model.__class__,), + {"__setattr__": _syncing_setattr}, + ) + def _register_indexer_caches_with_vllm(self): """Register DeepseekV32IndexerCache instances with vLLM so that: 1. vLLM discovers them via isinstance(AttentionLayerBase) for KV cache @@ -206,6 +318,51 @@ def _register_indexer_caches_with_vllm(self): f"static_forward_context, skipping" ) + def _adapt_mtp_layers_for_vllm(self) -> None: + """Install vLLM-only MTP input masking without changing model code.""" + if not self.is_mtp_draft_model: + return + + inner_model = getattr(self.model, "model", None) + layers = ( + getattr(inner_model, "layers", None) if inner_model is not None else None + ) + if layers is None: + return + + layer_iter = layers.values() if isinstance(layers, nn.ModuleDict) else layers + for layer in layer_iter: + if getattr(layer, "_atom_vllm_mtp_masked", False): + continue + + layer.forward = types.MethodType( + self._make_vllm_mtp_layer_forward(layer.forward), + layer, + ) + layer._atom_vllm_mtp_masked = True + + @staticmethod + def _make_vllm_mtp_layer_forward(original_forward): + @functools.wraps(original_forward) + def masked_forward( + self_layer, + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + spec_step_index=0, + ): + inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds) + return original_forward( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + spec_step_index, + ) + + return masked_forward + def forward( self, input_ids: torch.Tensor | None, @@ -221,7 +378,12 @@ def forward( # pass positions from vLLM to OOT execution path via vLLM's per-forward context if is_forward_context_available(): - get_vllm_forward_context().additional_kwargs["atom_positions"] = positions + forward_context = get_vllm_forward_context() + forward_context.additional_kwargs["atom_positions"] = positions + # set atom_config into vLLM forward_context in order to + # make sure main model and draft model can get their specific + # static_forward_context from their own atom_config + forward_context.additional_kwargs["atom_config"] = self.atom_config elif "positions" in self.atom_config.compilation_config.static_forward_context: buf = self.atom_config.compilation_config.static_forward_context[ "positions" @@ -248,8 +410,28 @@ def load_weights( # prevent circular import from atom.model_loader.loader import load_model_in_plugin_mode + is_mtp_draft_model = self.model_arch in { + "DeepSeekMTPModel", + "Qwen3NextMTPModel", + } + draft_hf_config = None + if is_mtp_draft_model: + draft_model_config = getattr( + getattr(self.vllm_config, "speculative_config", None), + "draft_model_config", + None, + ) + if draft_model_config is not None: + draft_hf_config = getattr( + draft_model_config, "hf_config", draft_model_config + ) + loaded_weights_record = load_model_in_plugin_mode( - model=self.model, config=self.atom_config, prefix="model." + model=self.model, + config=self.atom_config, + prefix="model.", + spec_decode=is_mtp_draft_model, + hf_config_override=draft_hf_config, ) return loaded_weights_record diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 32308c1e4a..9ef76e601d 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -5,6 +5,7 @@ from atom.plugin.prepare import _set_framework_backbone from atom.utils import envs from atom.plugin.vllm.mla_patch import patch_vllm_mla_attention +from atom.plugin.vllm.spec_decode_patch import apply_vllm_spec_decode_patch logger = logging.getLogger("atom") @@ -27,6 +28,7 @@ "DeepseekV32ForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Glm4MoeForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "GlmMoeDsaForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, + "DeepSeekMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLMVllm", "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration", @@ -126,6 +128,7 @@ def register_model() -> None: _patch_vllm_attention_process_weights_after_loading(Attention) _patch_vllm_attention_process_weights_after_loading(MLAAttention) + apply_vllm_spec_decode_patch() # Patch vLLM graph_capture to also enter aiter's ca_comm.capture(), # avoiding hipMemcpyAsync in fused_allreduce_rmsnorm when model uses aiter collectives diff --git a/atom/plugin/vllm/spec_decode_patch.py b/atom/plugin/vllm/spec_decode_patch.py new file mode 100644 index 0000000000..857af6874e --- /dev/null +++ b/atom/plugin/vllm/spec_decode_patch.py @@ -0,0 +1,35 @@ +import functools +import logging + +logger = logging.getLogger("atom") + + +def apply_vllm_spec_decode_patch() -> None: + """Patch vLLM speculative decoding for ATOM metadata compatibility.""" + try: + from atom.utils.forward_context import ( + AttentionMetaData as AtomAttentionMetaData, + ) + from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer + except Exception as exc: + logger.debug("ATOM spec-decode patch skipped: %s", exc) + return + + original_init = SpecDecodeBaseProposer.__init__ + if getattr(original_init, "_atom_allowed_attn_types_patched", False): + return + + @functools.wraps(original_init) + def wrapped_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + allowed = getattr(self, "allowed_attn_types", None) + if allowed is not None and AtomAttentionMetaData not in allowed: + self.allowed_attn_types = (*allowed, AtomAttentionMetaData) + + setattr(wrapped_init, "_atom_allowed_attn_types_patched", True) + SpecDecodeBaseProposer.__init__ = wrapped_init + + logger.info( + "ATOM plugin: patched vLLM speculative decoder for " + "ATOM attention-metadata compatibility." + ) From d372640aeb648dd7604ea8b7fec77621bfaf2e46 Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Thu, 14 May 2026 13:41:40 +0800 Subject: [PATCH 46/76] Support (P/D) disaggregation on mooncake (#690) * Support (P/D) disaggregation on mooncake * fix ruff * revert fence code for the same partiton acc perfect * add recipes * refactor modelrunner * slove conflict * format * support defer_output in pd * refresh schedule to optimize part of performance * add mooncake build method * clean code * remove useless code * fix mori engine * fix: support mixed TP sizes in PD disaggregation (e.g. prefill TP4 + decode TP8) * remove test in tests --- README.md | 1 + atom/kv_transfer/disaggregation/aggregator.py | 62 +- atom/kv_transfer/disaggregation/base.py | 9 + atom/kv_transfer/disaggregation/factory.py | 24 +- .../disaggregation/mooncake/__init__.py | 0 .../mooncake/mooncake_connector.py | 1051 +++++++++++++++++ .../disaggregation/moriio/__init__.py | 12 + .../disaggregation/moriio/moriio_common.py | 178 +++ .../moriio_connector.py} | 567 +-------- .../disaggregation/moriio/moriio_engine.py | 360 ++++++ atom/kv_transfer/disaggregation/proxy.py | 99 +- atom/kv_transfer/disaggregation/types.py | 6 +- atom/kv_transfer/disaggregation/utils.py | 116 ++ atom/model_engine/model_runner.py | 14 +- atom/model_engine/scheduler.py | 83 +- recipes/pd_disaggregation_guide.md | 177 +++ .../test_kv_aggregator.py | 0 .../test_kv_connector_scheduler.py | 0 .../disaggregation => tests}/test_proxy.py | 0 .../test_transfer_engine.py | 0 .../disaggregation => tests}/test_types.py | 0 21 files changed, 2142 insertions(+), 617 deletions(-) create mode 100644 atom/kv_transfer/disaggregation/mooncake/__init__.py create mode 100644 atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py create mode 100644 atom/kv_transfer/disaggregation/moriio/__init__.py create mode 100644 atom/kv_transfer/disaggregation/moriio/moriio_common.py rename atom/kv_transfer/disaggregation/{kv_transfer_engine.py => moriio/moriio_connector.py} (67%) create mode 100644 atom/kv_transfer/disaggregation/moriio/moriio_engine.py create mode 100644 atom/kv_transfer/disaggregation/utils.py create mode 100644 recipes/pd_disaggregation_guide.md rename {test/kv_transfer/disaggregation => tests}/test_kv_aggregator.py (100%) rename {test/kv_transfer/disaggregation => tests}/test_kv_connector_scheduler.py (100%) rename {test/kv_transfer/disaggregation => tests}/test_proxy.py (100%) rename {test/kv_transfer/disaggregation => tests}/test_transfer_engine.py (100%) rename {test/kv_transfer/disaggregation => tests}/test_types.py (100%) diff --git a/README.md b/README.md index d6ae9bb8bd..2805acb7ec 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ ## 📢 News - **[2026/05]** [Dissecting DeepSeek V4 Compressor](https://rocm.github.io/ATOM/dissecting_dsv4_compressor) — interactive animation visualizing how the CSA/HCA compressor state cache works (overlap mechanism, prefill vs decode, bulk compression vs sequential accumulation). +- **[2026/05]** ATOM now supports **Prefill/Decode (P/D) disaggregation** with [Mooncake](https://github.com/kvcache-ai/Mooncake) RDMA push-mode KV cache transfer. See [PD disaggregation guide](docs/pd_disaggregation_guide.md). - **[2026/03]** ATOM now supports **Prefill/Decode (P/D) disaggregation** — run prefill and decode on separate GPU nodes with RDMA-based KV cache transfer via [MORI-IO](https://github.com/ROCm/mori). See [disaggregation docs](atom/kv_transfer/disaggregation/README.md). ## 🚀 Features diff --git a/atom/kv_transfer/disaggregation/aggregator.py b/atom/kv_transfer/disaggregation/aggregator.py index 4c4afe59d4..bbe74f0f0e 100644 --- a/atom/kv_transfer/disaggregation/aggregator.py +++ b/atom/kv_transfer/disaggregation/aggregator.py @@ -28,10 +28,10 @@ class KVOutputAggregator: """Aggregates :class:`KVConnectorOutput` from all TP workers. - Uses a countdown approach: when a request ID first appears, a counter - is initialized to ``world_size``. Each worker that reports it as - finished decrements the counter. When the counter reaches zero, - the request is considered globally complete and is emitted. + Tracks which unique worker indices have reported each request as + finished. A request is globally complete only when all + ``world_size`` workers have reported it — duplicate reports from + the same worker (e.g. from retried notifications) are ignored. Args: world_size: Number of TP workers to aggregate over. @@ -48,8 +48,8 @@ def __init__(self, world_size: int = 8) -> None: if world_size <= 0: raise ValueError(f"world_size must be positive, got {world_size}") self._world_size = world_size - self._remaining_sending: dict[str, int] = {} - self._remaining_recving: dict[str, int] = {} + self._seen_sending: dict[str, set[int]] = {} + self._seen_recving: dict[str, set[int]] = {} @property def world_size(self) -> int: @@ -60,6 +60,7 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu Args: worker_outputs: One :class:`KVConnectorOutput` per worker. + The list index is the worker index. Returns: A new :class:`KVConnectorOutput` containing only request IDs @@ -68,40 +69,29 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu if not worker_outputs: return KVConnectorOutput() - all_sending_ids: set[str] = set() - all_recving_ids: set[str] = set() - for wo in worker_outputs: - if wo.finished_sending: - all_sending_ids.update(wo.finished_sending) - if wo.finished_recving: - all_recving_ids.update(wo.finished_recving) - - for rid in all_sending_ids: - if rid not in self._remaining_sending: - self._remaining_sending[rid] = self._world_size - - for rid in all_recving_ids: - if rid not in self._remaining_recving: - self._remaining_recving[rid] = self._world_size - - for wo in worker_outputs: + for worker_idx, wo in enumerate(worker_outputs): if wo.finished_sending: for rid in wo.finished_sending: - if rid in self._remaining_sending: - self._remaining_sending[rid] -= 1 - + self._seen_sending.setdefault(rid, set()).add(worker_idx) if wo.finished_recving: for rid in wo.finished_recving: - if rid in self._remaining_recving: - self._remaining_recving[rid] -= 1 - - done_sending = {rid for rid, cnt in self._remaining_sending.items() if cnt <= 0} - done_recving = {rid for rid, cnt in self._remaining_recving.items() if cnt <= 0} + self._seen_recving.setdefault(rid, set()).add(worker_idx) + + done_sending = { + rid + for rid, workers in self._seen_sending.items() + if len(workers) >= self._world_size + } + done_recving = { + rid + for rid, workers in self._seen_recving.items() + if len(workers) >= self._world_size + } for rid in done_sending: - del self._remaining_sending[rid] + del self._seen_sending[rid] for rid in done_recving: - del self._remaining_recving[rid] + del self._seen_recving[rid] return KVConnectorOutput( finished_sending=done_sending, @@ -110,10 +100,10 @@ def aggregate(self, worker_outputs: list[KVConnectorOutput]) -> KVConnectorOutpu def reset(self) -> None: """Clear all internal tracking state.""" - self._remaining_sending.clear() - self._remaining_recving.clear() + self._seen_sending.clear() + self._seen_recving.clear() @property def pending_count(self) -> tuple[int, int]: """Return ``(num_pending_sending, num_pending_recving)``.""" - return len(self._remaining_sending), len(self._remaining_recving) + return len(self._seen_sending), len(self._seen_recving) diff --git a/atom/kv_transfer/disaggregation/base.py b/atom/kv_transfer/disaggregation/base.py index a798c04b3c..c32902237c 100644 --- a/atom/kv_transfer/disaggregation/base.py +++ b/atom/kv_transfer/disaggregation/base.py @@ -53,6 +53,15 @@ def get_finished(self) -> tuple[set, set]: """ ... + def get_finished_recv_blocks(self) -> list[int]: + """Return block IDs from recently completed receives for GPU memory fence. + + RDMA writes to HBM may not be immediately visible to GPU compute + kernels. Connectors using RDMA should override this to return + blocks that need a GPU-side read-write cycle to ensure coherence. + """ + return [] + class KVConnectorSchedulerBase(ABC): """Scheduler-side KV connector interface.""" diff --git a/atom/kv_transfer/disaggregation/factory.py b/atom/kv_transfer/disaggregation/factory.py index e6867c254d..d18d621c40 100644 --- a/atom/kv_transfer/disaggregation/factory.py +++ b/atom/kv_transfer/disaggregation/factory.py @@ -31,10 +31,10 @@ class KVConnectorFactory: # Registration (happens once, typically at import time) KVConnectorFactory.register( "moriio", - worker_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - worker_class="KVConnector", - scheduler_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - scheduler_class="KVConnectorScheduler", + worker_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + worker_class="MoRIIOConnector", + scheduler_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + scheduler_class="MoRIIOConnectorScheduler", ) # Instantiation (called from forward_context.py) @@ -121,8 +121,16 @@ def create_connector( KVConnectorFactory.register( "moriio", - worker_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - worker_class="KVConnector", - scheduler_module="atom.kv_transfer.disaggregation.kv_transfer_engine", - scheduler_class="KVConnectorScheduler", + worker_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + worker_class="MoRIIOConnector", + scheduler_module="atom.kv_transfer.disaggregation.moriio.moriio_connector", + scheduler_class="MoRIIOConnectorScheduler", +) + +KVConnectorFactory.register( + "mooncake", + worker_module="atom.kv_transfer.disaggregation.mooncake.mooncake_connector", + worker_class="MooncakeConnector", + scheduler_module="atom.kv_transfer.disaggregation.mooncake.mooncake_connector", + scheduler_class="MooncakeConnectorScheduler", ) diff --git a/atom/kv_transfer/disaggregation/mooncake/__init__.py b/atom/kv_transfer/disaggregation/mooncake/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py new file mode 100644 index 0000000000..1599ddf187 --- /dev/null +++ b/atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py @@ -0,0 +1,1051 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Worker-side and scheduler-side KV cache connectors for disaggregated P/D. + +Uses Mooncake TransferEngine for RDMA-based push (WRITE) transfers of +KV cache data from producer (prefill) to consumer (decode) nodes. +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import msgpack +import msgspec +import torch +import zmq + +from atom.config import Config +from atom.kv_transfer.disaggregation.base import ( + KVConnectorBase, + KVConnectorSchedulerBase, +) +from atom.kv_transfer.disaggregation.types import ( + ConnectorMetadata, + ReqId, + TransferId, +) +from atom.model_engine.sequence import Sequence +from atom.utils import get_open_port, make_zmq_path, zmq_socket_ctx +from atom.utils.network import get_ip +from aiter.dist.parallel_state import get_dp_group, get_tp_group + +logger = logging.getLogger("atom") + +# --------------------------------------------------------------------------- +# Mooncake availability check +# --------------------------------------------------------------------------- + +_MOONCAKE_AVAILABLE = False +try: + from mooncake.engine import TransferEngine + + _MOONCAKE_AVAILABLE = True + logger.info("Mooncake TransferEngine loaded successfully") +except ImportError: + logger.warning( + "Mooncake is not available — KV cache disaggregation via mooncake " + "will not work. Install the mooncake package to enable push-mode " + "RDMA transfers." + ) + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MOONCAKE_PING_INTERVAL_SECONDS = 5 +MOONCAKE_MAX_PING_RETRIES = 100 +MOONCAKE_DEFAULT_PROTOCOL = "rdma" +PREFILL_LOOKUP_TIMEOUT = 60 +PREFILL_LOOKUP_POLL_INTERVAL = 0.01 + +# ZMQ side-channel message types +MSG_WRITE_REQUEST = b"write_request" +MSG_WRITE_DONE = b"write_done" +MSG_GET_META = b"get_meta" + + +# --------------------------------------------------------------------------- +# Metadata struct for bootstrap handshake +# --------------------------------------------------------------------------- + + +class MooncakeAgentMetadata( + msgspec.Struct, + omit_defaults=True, + dict=True, + kw_only=True, +): + """Serializable metadata exchanged during the mooncake bootstrap.""" + + engine_id: str + rpc_port: int + kv_caches_base_addr: list[int] | None = None + num_blocks: int = 0 + block_len: int = 0 + + +def _port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: + return dp_rank * tp_size + tp_rank + + +# =================================================================== +# MooncakeConnectorScheduler — scheduler-side connector +# =================================================================== + + +class MooncakeConnectorScheduler(KVConnectorSchedulerBase): + def __init__(self, config: Config) -> None: + kv_transfer_config = config.kv_transfer_config + self.is_producer = ( + kv_transfer_config.get("kv_role", "kv_producer") == "kv_producer" + ) + self.handshake_port = get_open_port() + self.engine_id = "None" + self.tp_size = config.tensor_parallel_size + self.dp_size = config.parallel_config.data_parallel_size + self.dp_rank = config.parallel_config.data_parallel_rank + self.host_ip = get_ip() + + # Pending requests: req_id -> (Sequence, block_table) + self._reqs_need_recv: dict[ReqId, tuple[Any, list[int]]] = {} + self._reqs_need_save: dict[ReqId, tuple[Any, list[int]]] = {} + + # Bidirectional transfer_id <-> request_id mapping + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + + def get_num_new_matched_tokens(self, seq: Sequence) -> tuple[int, bool]: + params = seq.kv_transfer_params or {} + + if params.get("do_remote_prefill") and not getattr( + seq, "kv_async_tagged", False + ): + return len(seq.prompt_token_ids), True + + return 0, False + + def build_connector_meta(self) -> ConnectorMetadata: + meta = ConnectorMetadata() + meta.request_id_to_transfer_id = self.request_id_to_transfer_id + + for req_id, (req, block_ids) in self._reqs_need_recv.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_recv( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + # Producer side: pass completed prefill block_ids to worker + for req_id, (req, block_ids) in self._reqs_need_save.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_save( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + if self._reqs_need_recv or self._reqs_need_save: + logger.info( + "[SCHEDULER] build_connector_meta: %d recv, %d save, " "id_map=%s", + len(self._reqs_need_recv), + len(self._reqs_need_save), + meta.request_id_to_transfer_id, + ) + self._reqs_need_recv.clear() + self._reqs_need_save.clear() + return meta + + def update_state_after_alloc(self, seq: Sequence) -> None: + params = seq.kv_transfer_params or {} + + if not self.is_producer: + transfer_id = params.get("transfer_id") + if transfer_id is not None: + self.transfer_id_to_request_id[transfer_id] = seq.id + self.request_id_to_transfer_id[seq.id] = transfer_id + + # Consumer side: queue for remote KV loading + if params.get("do_remote_prefill"): + assert ( + not self.is_producer + ), "Only the decode (consumer) side handles do_remote_prefill" + self._reqs_need_recv[seq.id] = (seq, seq.block_table) + params["do_remote_prefill"] = False + logger.info( + "[SCHEDULER-CONSUMER] Queued req %s for remote KV recv " + "(%d blocks), transfer_id=%s, remote_host=%s, " + "remote_handshake_port=%s", + seq.id, + len(seq.block_table), + params.get("transfer_id"), + params.get("remote_host"), + params.get("remote_handshake_port"), + ) + + # Producer side: queue block_ids for the write listener to look up + if params.get("do_remote_decode"): + assert self.is_producer, "Only the producer side handles do_remote_decode" + self._reqs_need_save[seq.id] = (seq, seq.block_table) + logger.debug( + "Queued req %s for KV save (%d blocks)", + seq.id, + len(seq.block_table), + ) + + def request_finished(self, seq: Sequence) -> None: + first_token_id = seq.output_tokens[0] if seq.output_tokens else None + seq.kv_transfer_params_output = { + "do_remote_prefill": True, + "do_remote_decode": False, + "remote_block_ids": seq.block_table.copy(), + "remote_engine_id": self.engine_id, + "remote_host": self.host_ip, + "remote_port": self.handshake_port, + "tp_size": self.tp_size, + "dp_rank": self.dp_rank, + "transfer_id": seq.id, + "first_token_id": first_token_id, + } + + if not self.is_producer: + transfer_id = self.request_id_to_transfer_id.pop(seq.id, None) + if transfer_id is not None: + self.transfer_id_to_request_id.pop(transfer_id, None) + + +# =================================================================== +# MooncakeConnector — worker-side connector (runs inside each TP rank) +# =================================================================== + + +class MooncakeConnector(KVConnectorBase): + """Worker-side KV cache connector using Mooncake push-mode RDMA. + + Mooncake uses a push/WRITE model: the prefill (producer) node writes + KV cache data directly into the decode (consumer) node's registered + GPU memory via ``batch_transfer_sync_write``. + """ + + def __init__(self, config: Config) -> None: + self.tp_rank = get_tp_group().rank_in_group + self.dp_rank = get_dp_group().rank_in_group + self.tp_size = get_tp_group().world_size + self.dp_size = get_dp_group().world_size + + kv_transfer_config = config.kv_transfer_config + self.local_ip = get_ip() + self._local_ping_port = get_open_port() + + self.is_producer = ( + kv_transfer_config.get("kv_role", "kv_producer") == "kv_producer" + ) + self.is_consumer = not self.is_producer + + # Networking / service discovery config + self.http_port = kv_transfer_config.get("http_port", 8000) + self.proxy_ping_port = kv_transfer_config.get("proxy_ping_port", 36367) + self.proxy_ip = kv_transfer_config.get("proxy_ip") + self.request_address = f"{self.local_ip}:{self.http_port}" + self.protocol = kv_transfer_config.get("protocol", MOONCAKE_DEFAULT_PROTOCOL) + + # Side channel port (ZMQ) — deterministic from config for proxy relay + self.base_handshake_port = kv_transfer_config.get("handshake_port", 6301) + self._side_channel_port = self.base_handshake_port + _port_offset( + self.dp_rank, self.tp_rank, self.tp_size + ) + + # --- Mooncake TransferEngine initialization --- + if not _MOONCAKE_AVAILABLE: + raise RuntimeError( + "Mooncake is not installed but kv_connector='mooncake' was requested. " + "Install the mooncake package to use push-mode RDMA transfers." + ) + + # Determine which RDMA device this TP rank should use. + # On AMD MI300X, each GPU has a paired RDMA NIC: GPU N → rdmaN. + # Registering GPU memory with a non-local RDMA NIC fails with + # EINVAL. Pass the device name as a filter so Mooncake only + # creates a context for the local NIC. + ib_device = kv_transfer_config.get("ib_device", "") + if not ib_device: + ib_device = os.environ.get("ATOM_MOONCAKE_IB_DEVICE", "") + if not ib_device: + gpu_idx = torch.cuda.current_device() + ib_device = f"rdma{gpu_idx}" + logger.info( + "Auto-selecting RDMA device %s for GPU %d (tp_rank=%d)", + ib_device, + gpu_idx, + self.tp_rank, + ) + + self.transfer_engine = TransferEngine() + ret = self.transfer_engine.initialize( + self.local_ip, + "P2PHANDSHAKE", + self.protocol, + ib_device, + ) + if ret != 0: + raise RuntimeError( + f"Mooncake TransferEngine.initialize() failed (ret={ret}) " + f"on ip={self.local_ip}, protocol={self.protocol}, " + f"ib_device={ib_device}" + ) + self.rpc_port = self.transfer_engine.get_rpc_port() + self.engine_id = f"{self.local_ip}:{self.rpc_port}" + logger.info( + "Mooncake TransferEngine initialized: ip=%s, protocol=%s, " + "ib_device=%s, rpc_port=%d", + self.local_ip, + self.protocol, + ib_device, + self.rpc_port, + ) + + # --- KV cache state (populated in register_kv_caches) --- + self.kv_caches: dict[str, Any] | None = None + self.kv_caches_base_addr: list[int] = [] + self._per_block_bytes_list: list[int] = [] + self.kv_cache_shape: tuple[int, ...] | None = None + self.block_len: int = config.kv_cache_block_size + self.num_blocks: int = 0 + self._per_block_bytes: int = 0 + + # --- Producer: completed prefill block_ids cache --- + # Populated from ConnectorMetadata.reqs_to_save each step. + # The write listener looks up block_ids here when consumer requests a write. + self._completed_prefills: dict[ReqId, list[int]] = {} + self._completed_prefills_lock = threading.Lock() + self._completed_prefills_cv = threading.Condition(self._completed_prefills_lock) + self._transfer_refcount: dict[ReqId, int] = {} + self._transfer_refcount_lock = threading.Lock() + + # --- Consumer: pending receive tracking --- + self._pending_recv: set[ReqId] = set() + self._pending_recv_blocks: dict[ReqId, list[int]] = {} + self._notification_port = get_open_port() + + # --- Completion tracking --- + self.done_sending: set[str] = set() + self.done_recving: set[str] = set() + self._completion_lock = threading.Lock() + + # --- GPU memory fence: blocks pending coherence enforcement --- + self._blocks_pending_fence: list[int] = [] + self._fence_lock = threading.Lock() + + # --- Transfer ID mapping (worker side) --- + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + + # --- Producer: thread pool for RDMA writes --- + if self.is_producer: + self._send_executor = ThreadPoolExecutor( + max_workers=kv_transfer_config.get("num_worker_threads", 16), + thread_name_prefix="mooncake-send-worker", + ) + + # --- ZMQ for metadata exchange --- + self.zmq_context = zmq.Context() + + # --- Producer: persistent socket cache for write-done notifications --- + self._notify_sockets: dict[str, zmq.Socket] = {} + self._notify_sockets_lock = threading.Lock() + + # --- Msgspec encoder/decoder for bootstrap metadata --- + self._encoder = msgspec.msgpack.Encoder() + self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata) + + # --- Service discovery ping (rank 0 only) --- + if self.tp_rank == 0 and self.dp_rank == 0: + self._ping_thread = threading.Thread( + target=self._service_discovery_ping, + args=(self.zmq_context,), + daemon=True, + name="mooncake-ping", + ) + self._ping_thread.start() + + # ----------------------------------------------------------------- + # Service discovery + # ----------------------------------------------------------------- + + def _service_discovery_ping(self, zmq_context: zmq.Context) -> None: + """Periodically register with the proxy (rank 0 only).""" + grpc_endpoint = f"http://{self.request_address}/v1/completions" + role_code = "P" if self.is_producer else "D" + retry_count = 0 + msg_index = 1 + proxy_path = f"tcp://{self.proxy_ip}:{self.proxy_ping_port}" + + with zmq_context.socket(zmq.DEALER) as sock: + sock.connect(proxy_path) + + while True: + try: + registration_data = { + "type": "register", + "role": role_code, + "index": str(msg_index), + "request_address": grpc_endpoint, + "rpc_port": self.rpc_port, + "handshake_port": self.base_handshake_port, + "dp_size": self.dp_size, + "tp_size": self.tp_size, + "transfer_mode": "write", + } + sock.send(msgpack.dumps(registration_data)) + logger.debug( + "Ping #%d sent to %s (role=%s)", + msg_index, + proxy_path, + role_code, + ) + retry_count = 0 + + except ConnectionRefusedError: + logger.info( + "Proxy connection refused: %s -> %s", + self.local_ip, + proxy_path, + ) + retry_count += 1 + + except OSError as e: + logger.info("OS error during ping: %s", e) + retry_count += 1 + + except Exception as e: + logger.info("Unexpected ping error: %s", e) + retry_count += 1 + if retry_count >= MOONCAKE_MAX_PING_RETRIES: + logger.error( + "Ping failed after %d retries, aborting", + MOONCAKE_MAX_PING_RETRIES, + ) + raise RuntimeError( + f"Service discovery ping failed after " + f"{retry_count} retries" + ) from e + + finally: + time.sleep(MOONCAKE_PING_INTERVAL_SECONDS) + msg_index += 1 + + # ----------------------------------------------------------------- + # KVConnectorBase: register_kv_caches + # ----------------------------------------------------------------- + + # RDMA MR size limit — must stay under device max_mr_size (2 GiB on + # AMD ROCm RDMA). Mooncake truncates silently, leaving addresses + # beyond the limit unregistered. + _MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 + + def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: + """Register all KV cache tensors with the Mooncake TransferEngine. + + Must be called once after model loading and KV cache allocation. + Starts the side-channel listener threads for write coordination. + + Large tensors (> 2 GiB) are split into chunks for RDMA memory + registration to stay within the device max_mr_size. The logical + per-tensor base addresses (used for transfer offset computation) + remain unchanged. + + For fp8 KV caches, k_scale and v_scale tensors are also registered + so that quantization scale factors are transferred alongside data. + """ + self.kv_caches = kv_caches + + # Logical per-tensor info (for transfer address computation) + tensor_ptrs: list[int] = [] + tensor_sizes: list[int] = [] + + # Chunked regions (for RDMA registration) + reg_ptrs: list[int] = [] + reg_sizes: list[int] = [] + + def _register_tensor(tensor): + self.kv_caches_base_addr.append(tensor.data_ptr()) + sz = tensor.element_size() + if is_mla: + bpb = self.block_len * tensor.stride(0) * sz + else: + bpb = tensor.stride(0) * sz + self._per_block_bytes_list.append(bpb) + + base = tensor.data_ptr() + total = tensor.numel() * sz + tensor_ptrs.append(base) + tensor_sizes.append(total) + + # Chunk for RDMA registration + offset = 0 + while offset < total: + chunk = min(self._MAX_RDMA_CHUNK_BYTES, total - offset) + reg_ptrs.append(base + offset) + reg_sizes.append(chunk) + offset += chunk + + for layer_name, kv_cache in kv_caches.items(): + k_cache = kv_cache.k_cache + v_cache = kv_cache.v_cache + is_mla = v_cache is None + + if self.kv_cache_shape is None: + self.kv_cache_shape = k_cache.shape + + _register_tensor(k_cache) + + if not is_mla: + _register_tensor(v_cache) + + if kv_cache.k_scale is not None: + _register_tensor(kv_cache.k_scale) + if kv_cache.v_scale is not None: + _register_tensor(kv_cache.v_scale) + + logger.info( + "Registering %d tensors as %d RDMA chunks " "(max_chunk=%.2f GiB)", + len(tensor_ptrs), + len(reg_ptrs), + self._MAX_RDMA_CHUNK_BYTES / (1024**3), + ) + for i, (ptr, sz_bytes) in enumerate(zip(reg_ptrs, reg_sizes)): + logger.info( + " reg_chunk[%d] ptr=0x%x size=%d (%.2f GiB)", + i, + ptr, + sz_bytes, + sz_bytes / (1024**3), + ) + + ret = self.transfer_engine.batch_register_memory(reg_ptrs, reg_sizes) + if ret != 0: + logger.error( + "batch_register_memory FAILED (ret=%d). " + "Trying individual registration as fallback...", + ret, + ) + for i, (ptr, sz_bytes) in enumerate(zip(reg_ptrs, reg_sizes)): + r = self.transfer_engine.register_memory(ptr, sz_bytes) + logger.info( + " register_memory[%d] ptr=0x%x size=%d (%.2f GiB) " "=> ret=%d", + i, + ptr, + sz_bytes, + sz_bytes / (1024**3), + r, + ) + else: + logger.info( + "batch_register_memory OK (%d chunks from %d tensors)", + len(reg_ptrs), + len(tensor_ptrs), + ) + + # Block geometry from last layer + last_kv = list(kv_caches.values())[-1] + k = last_kv.k_cache + is_mla = last_kv.v_cache is None + sz = k.element_size() + + if is_mla: + self.num_blocks = k.shape[0] // self.block_len + self._per_block_bytes = self.block_len * k.stride(0) * sz + else: + self.num_blocks = k.shape[0] + self._per_block_bytes = k.stride(0) * sz + + logger.info( + "Mooncake registered %d tensors (%d RDMA chunks, %d blocks, " + "per_block_bytes=%s, block_len=%d, is_mla=%s, " + "k_shape=%s, k_stride=%s, elem_sz=%d, " + "tensor_sizes_MB=%s, chunk_sizes_MB=%s)", + len(self.kv_caches_base_addr), + len(reg_ptrs), + self.num_blocks, + sorted(set(self._per_block_bytes_list)), + self.block_len, + is_mla, + k.shape, + k.stride(), + sz, + [s // (1024 * 1024) for s in sorted(set(tensor_sizes))], + [s // (1024 * 1024) for s in sorted(set(reg_sizes))], + ) + + logger.info( + "Mooncake KV registration complete: role=%s, engine_id=%s, " + "rpc_port=%d, num_regions=%d, num_blocks=%d, " + "base_addrs=[%s]", + "PRODUCER" if self.is_producer else "CONSUMER", + self.engine_id, + self.rpc_port, + len(self.kv_caches_base_addr), + self.num_blocks, + ", ".join(f"0x{a:x}" for a in self.kv_caches_base_addr[:4]), + ) + + # Build metadata for bootstrap exchange + self._local_metadata = MooncakeAgentMetadata( + engine_id=self.engine_id, + rpc_port=self.rpc_port, + kv_caches_base_addr=self.kv_caches_base_addr, + num_blocks=self.num_blocks, + block_len=self.block_len, + ) + + # Start side channel threads + if self.is_producer: + self._write_listener_thread = threading.Thread( + target=self._write_listener, + daemon=True, + name="mooncake-write-listener", + ) + self._write_listener_thread.start() + else: + self._notification_listener_thread = threading.Thread( + target=self._notification_listener, + daemon=True, + name="mooncake-notify-listener", + ) + self._notification_listener_thread.start() + + # ----------------------------------------------------------------- + # KVConnectorBase: start_load_kv + # ----------------------------------------------------------------- + + def start_load_kv(self, metadata: ConnectorMetadata) -> None: + """Initiate KV transfers for pending requests. + + **Producer side**: Cache completed prefill block_ids from + ``metadata.reqs_to_save`` so the write listener can look them up. + + **Consumer side**: For each pending recv request, connect to the + producer's ZMQ side channel and send a write request with our + memory addresses and block allocation. + """ + if metadata is None: + return + + self.request_id_to_transfer_id = metadata.request_id_to_transfer_id + + # Producer: cache block_ids from completed prefills + if self.is_producer: + for req_id, meta in metadata.reqs_to_save.items(): + with self._completed_prefills_cv: + self._completed_prefills[req_id] = meta.local_block_ids + self._completed_prefills_cv.notify_all() + logger.info( + "[PRODUCER] Cached %d prefill blocks for req %s", + len(meta.local_block_ids), + req_id, + ) + return + + # Consumer: send write requests to producer + if not metadata.reqs_to_recv: + return + + logger.info( + "[CONSUMER] start_load_kv: %d reqs_to_recv, id_map=%s", + len(metadata.reqs_to_recv), + metadata.request_id_to_transfer_id, + ) + + for req_id, meta in metadata.reqs_to_recv.items(): + remote_tp_size = meta.tp_size + if remote_tp_size != self.tp_size: + remote_tp_rank = self.tp_rank % remote_tp_size + else: + remote_tp_rank = self.tp_rank + remote_port = meta.remote_handshake_port + _port_offset( + meta.remote_dp_rank, remote_tp_rank, remote_tp_size + ) + remote_addr = make_zmq_path("tcp", meta.remote_host, remote_port) + + unique_bpb = sorted(set(self._per_block_bytes_list)) + logger.info( + "[CONSUMER] Sending write_request for req %s (transfer_id=%s) " + "to %s (handshake_port=%d, dp_rank=%d, " + "local_tp=%d, remote_tp=%d/%d), " + "dst_block_ids=%s, num_regions=%d, " + "bytes/block=%s, num_blocks=%d", + req_id, + meta.transfer_id, + remote_addr, + meta.remote_handshake_port, + meta.remote_dp_rank, + self.tp_rank, + remote_tp_rank, + remote_tp_size, + meta.local_block_ids[:10], + len(self.kv_caches_base_addr), + unique_bpb, + self.num_blocks, + ) + + write_request = msgpack.dumps( + { + "request_id": req_id, + "transfer_id": meta.transfer_id, + "consumer_host": self.local_ip, + "consumer_rpc_port": self.rpc_port, + "consumer_base_addrs": self.kv_caches_base_addr, + "dst_block_ids": meta.local_block_ids, + "notify_host": self.local_ip, + "notify_port": self._notification_port, + "consumer_tp_size": self.tp_size, + } + ) + + with self._notify_sockets_lock: + sock = self._notify_sockets.get(remote_addr) + if sock is None: + sock = self.zmq_context.socket(zmq.DEALER) + sock.setsockopt(zmq.LINGER, 5000) + sock.setsockopt(zmq.SNDHWM, 0) + sock.connect(remote_addr) + self._notify_sockets[remote_addr] = sock + sock.send_multipart([MSG_WRITE_REQUEST, write_request]) + + self._pending_recv.add(req_id) + self._pending_recv_blocks[req_id] = list(meta.local_block_ids) + logger.info( + "[CONSUMER] write_request sent for req %s to %s", + req_id, + remote_addr, + ) + + # ----------------------------------------------------------------- + # KVConnectorBase: get_finished + # ----------------------------------------------------------------- + + def get_finished(self) -> tuple[set, set]: + """Return ``(done_sending, done_recving)`` and clear internal sets.""" + with self._completion_lock: + ds = self.done_sending.copy() + dr = self.done_recving.copy() + self.done_sending.clear() + self.done_recving.clear() + if ds or dr: + logger.info( + "[%s] get_finished: sending=%s, recving=%s", + "PRODUCER" if self.is_producer else "CONSUMER", + ds, + dr, + ) + return ds, dr + + def get_finished_recv_blocks(self) -> list[int]: + """Return block IDs from recently completed RDMA receives.""" + with self._fence_lock: + blocks = self._blocks_pending_fence + self._blocks_pending_fence = [] + return blocks + + # ----------------------------------------------------------------- + # Producer: write listener (ZMQ ROUTER) + # ----------------------------------------------------------------- + + def _write_listener(self) -> None: + """Accept write requests from consumers and dispatch RDMA writes.""" + path = make_zmq_path("tcp", "*", self._side_channel_port) + logger.info("Mooncake write listener bound to %s", path) + + with zmq_socket_ctx(path, zmq.ROUTER, bind=True) as sock: + while True: + parts = sock.recv_multipart() + identity, msg_type = parts[0], parts[1] + + if msg_type == MSG_GET_META: + encoded = self._encoder.encode(self._local_metadata) + sock.send_multipart([identity, b"", encoded]) + logger.debug("Sent metadata to peer") + + elif msg_type == MSG_WRITE_REQUEST: + request_data = msgpack.loads(parts[2]) + logger.info( + "[PRODUCER] Received write_request for req %s " + "(transfer_id=%s, consumer=%s:%s)", + request_data["request_id"], + request_data.get("transfer_id"), + request_data.get("consumer_host"), + request_data.get("consumer_rpc_port"), + ) + self._send_executor.submit(self._execute_transfer, request_data) + + else: + logger.error("Unknown message type: %s", msg_type) + + # ----------------------------------------------------------------- + # Producer: execute RDMA write + # ----------------------------------------------------------------- + + def _execute_transfer(self, request_data: dict) -> None: + """Compute offsets and perform RDMA write for a single request.""" + req_id = request_data["request_id"] + transfer_id = request_data.get("transfer_id", req_id) + consumer_host = request_data["consumer_host"] + consumer_rpc_port = request_data["consumer_rpc_port"] + consumer_base_addrs = request_data["consumer_base_addrs"] + dst_block_ids = request_data["dst_block_ids"] + notify_host = request_data["notify_host"] + notify_port = request_data["notify_port"] + consumer_tp_size = request_data.get("consumer_tp_size", self.tp_size) + consumers_per_rank = max(1, consumer_tp_size // self.tp_size) + + logger.info( + "[PRODUCER] _execute_transfer: req_id=%s, transfer_id=%s, " + "consumer=%s:%s, dst_blocks=%d", + req_id, + transfer_id, + consumer_host, + consumer_rpc_port, + len(dst_block_ids), + ) + + # Look up producer's block_ids by transfer_id (= prefill's seq.id), + # not request_id (= decode's seq.id). + src_block_ids = self._wait_for_prefill_blocks(transfer_id) + if src_block_ids is None: + logger.error( + "[PRODUCER] Timed out waiting for prefill blocks for " + "transfer_id=%s (req_id=%s). Available keys: %s", + transfer_id, + req_id, + list(self._completed_prefills.keys()), + ) + return + + target = f"{consumer_host}:{consumer_rpc_port}" + + # Verify the consumer's segment is discoverable via P2P handshake + if hasattr(self.transfer_engine, "get_first_buffer_address"): + remote_buf = self.transfer_engine.get_first_buffer_address(target) + logger.info( + "[PRODUCER] P2P segment probe for %s: " "get_first_buffer_address=0x%x", + target, + remote_buf, + ) + if remote_buf == 0: + logger.error( + "[PRODUCER] Consumer %s has NO registered buffers visible " + "via P2P handshake! Consumer batch_register_memory may " + "have failed, or segment descriptors haven't propagated.", + target, + ) + + src_addrs: list[int] = [] + dst_addrs: list[int] = [] + sizes: list[int] = [] + + num_regions = len(self.kv_caches_base_addr) + if num_regions != len(consumer_base_addrs): + logger.error( + "[PRODUCER] REGION COUNT MISMATCH: producer has %d regions, " + "consumer has %d regions", + num_regions, + len(consumer_base_addrs), + ) + if len(src_block_ids) != len(dst_block_ids): + logger.error( + "[PRODUCER] BLOCK COUNT MISMATCH: src has %d blocks, " + "dst has %d blocks", + len(src_block_ids), + len(dst_block_ids), + ) + for region_idx in range(num_regions): + src_base = self.kv_caches_base_addr[region_idx] + dst_base = consumer_base_addrs[region_idx] + bpb = self._per_block_bytes_list[region_idx] + for sb, db in zip(src_block_ids, dst_block_ids): + src_addrs.append(src_base + sb * bpb) + dst_addrs.append(dst_base + db * bpb) + sizes.append(bpb) + + unique_bpb = sorted(set(self._per_block_bytes_list)) + logger.info( + "[PRODUCER] Starting RDMA write: req=%s, target=%s, " + "%d regions × %d blocks, bytes/block=%s, " + "src_block_ids=%s, dst_block_ids=%s, " + "total_transfer_entries=%d, total_bytes=%d", + req_id, + target, + num_regions, + len(src_block_ids), + unique_bpb, + src_block_ids[:10], + dst_block_ids[:10], + len(src_addrs), + sum(sizes), + ) + + if dst_addrs: + logger.info( + "[PRODUCER] dst_addr range: 0x%x -- 0x%x (first region), " + "consumer_base_addrs[0]=0x%x", + min(dst_addrs[: len(dst_block_ids)]), + max(dst_addrs[: len(dst_block_ids)]) + sizes[0], + consumer_base_addrs[0], + ) + + max_entries_per_batch = 4096 + total_entries = len(src_addrs) + max_retries = 3 + + for chunk_start in range(0, total_entries, max_entries_per_batch): + chunk_end = min(chunk_start + max_entries_per_batch, total_entries) + chunk_src = src_addrs[chunk_start:chunk_end] + chunk_dst = dst_addrs[chunk_start:chunk_end] + chunk_sizes = sizes[chunk_start:chunk_end] + + retry_delay = 2.0 + for attempt in range(max_retries): + try: + ret = self.transfer_engine.batch_transfer_sync_write( + target, chunk_src, chunk_dst, chunk_sizes + ) + if ret == 0: + break + logger.error( + "[PRODUCER] RDMA write chunk error %d for req %s → %s " + "(entries %d-%d/%d, attempt %d/%d)", + ret, + req_id, + target, + chunk_start, + chunk_end, + total_entries, + attempt + 1, + max_retries, + ) + except Exception: + logger.exception( + "[PRODUCER] RDMA write chunk FAILED for req %s " + "(entries %d-%d/%d, attempt %d/%d)", + req_id, + chunk_start, + chunk_end, + total_entries, + attempt + 1, + max_retries, + ) + ret = -1 + if ret == 0: + break + if attempt < max_retries - 1: + time.sleep(retry_delay) + retry_delay *= 2 + else: + return + + logger.info( + "[PRODUCER] RDMA write completed for req %s → %s (%d entries)", + req_id, + target, + total_entries, + ) + + # Notify this consumer immediately — its data is written. + self._send_write_done(notify_host, notify_port, req_id) + + # Track how many consumers still need this transfer_id's blocks. + # Only mark done_sending (which triggers block deallocation on the + # scheduler) after ALL consumers sharing this producer rank have + # completed their transfers. + all_done = False + with self._transfer_refcount_lock: + if transfer_id not in self._transfer_refcount: + self._transfer_refcount[transfer_id] = consumers_per_rank + self._transfer_refcount[transfer_id] -= 1 + if self._transfer_refcount[transfer_id] <= 0: + self._transfer_refcount.pop(transfer_id) + all_done = True + + if all_done: + with self._completion_lock: + self.done_sending.add(transfer_id) + with self._completed_prefills_lock: + self._completed_prefills.pop(transfer_id, None) + logger.info( + "[PRODUCER] All %d consumers served for transfer_id=%s, " + "blocks released", + consumers_per_rank, + transfer_id, + ) + + def _wait_for_prefill_blocks(self, req_id: str) -> list[int] | None: + """Wait until prefill block_ids are available for this request.""" + with self._completed_prefills_cv: + ready = self._completed_prefills_cv.wait_for( + lambda: req_id in self._completed_prefills, + timeout=PREFILL_LOOKUP_TIMEOUT, + ) + if ready: + return self._completed_prefills[req_id] + return None + + def _send_write_done(self, host: str, port: int, req_id: str) -> None: + """Send write-done notification to consumer via persistent socket. + + Sends the notification multiple times for reliability — the consumer + uses a set so duplicates are harmless. + """ + path = make_zmq_path("tcp", host, port) + notification = msgpack.dumps({"request_id": req_id}) + with self._notify_sockets_lock: + sock = self._notify_sockets.get(path) + if sock is None: + sock = self.zmq_context.socket(zmq.DEALER) + sock.setsockopt(zmq.LINGER, 5000) + sock.setsockopt(zmq.SNDHWM, 0) + sock.connect(path) + self._notify_sockets[path] = sock + for _ in range(3): + sock.send_multipart([MSG_WRITE_DONE, notification]) + logger.info("[PRODUCER] write-done sent for req %s", req_id) + + # ----------------------------------------------------------------- + # Consumer: notification listener (ZMQ ROUTER) + # ----------------------------------------------------------------- + + def _notification_listener(self) -> None: + """Receive write-done notifications from producers.""" + path = make_zmq_path("tcp", "*", self._notification_port) + logger.info("Mooncake notification listener bound to %s", path) + + with zmq_socket_ctx(path, zmq.ROUTER, bind=True) as sock: + while True: + parts = sock.recv_multipart() + msg_type = parts[1] + + if msg_type == MSG_WRITE_DONE: + data = msgpack.loads(parts[2]) + req_id = data["request_id"] + dst_blocks = self._pending_recv_blocks.pop(req_id, None) + if dst_blocks: + with self._fence_lock: + self._blocks_pending_fence.extend(dst_blocks) + with self._completion_lock: + self.done_recving.add(req_id) + self._pending_recv.discard(req_id) + logger.info( + "[CONSUMER] Write-done received for req %s, " + "done_recving now: %s", + req_id, + self.done_recving, + ) + else: + logger.error("Unknown notification type: %s", msg_type) diff --git a/atom/kv_transfer/disaggregation/moriio/__init__.py b/atom/kv_transfer/disaggregation/moriio/__init__.py new file mode 100644 index 0000000000..8fbac11053 --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +from atom.kv_transfer.disaggregation.moriio.moriio_connector import ( + MoRIIOConnector, + MoRIIOConnectorScheduler, +) + +__all__ = [ + "MoRIIOConnector", + "MoRIIOConnectorScheduler", +] diff --git a/atom/kv_transfer/disaggregation/moriio/moriio_common.py b/atom/kv_transfer/disaggregation/moriio/moriio_common.py new file mode 100644 index 0000000000..40d7db2c8f --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/moriio_common.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Shared types, constants, enums, and helpers for the MoRIIO KV connector. + +This module has no heavy dependencies (no torch at import time, no RDMA +engine instances) so it can be imported freely by the other moriio +submodules. +""" + +from __future__ import annotations + +import logging +import threading +from enum import Enum +from typing import Optional + +import msgspec + +from atom.kv_transfer.disaggregation.utils import ( # noqa: F401 + MAX_RDMA_CHUNK_BYTES, + chunk_tensor_for_rdma, +) + +logger = logging.getLogger("atom") + +_chunk_tensor_for_rdma = chunk_tensor_for_rdma + +# --------------------------------------------------------------------------- +# MoRIIO availability check +# --------------------------------------------------------------------------- + +_MORIIO_AVAILABLE = False +try: + import mori.io # noqa: F401 + + _MORIIO_AVAILABLE = True + logger.info("MoRIIO RDMA library loaded successfully") +except ImportError: + logger.warning( + "MoRIIO is not available — KV cache disaggregation will not work. " + "Install the mori package to enable RDMA transfers." + ) + + +# --------------------------------------------------------------------------- +# Msgspec metadata structs +# --------------------------------------------------------------------------- + + +class MoRIIOAgentMetadata( + msgspec.Struct, + omit_defaults=True, + dict=True, + kw_only=True, +): + """Serializable metadata exchanged during the RDMA handshake.""" + + engine_id: str + agent_metadata: bytes + kv_caches_base_addr: Optional[list[int]] = None + num_blocks: int = 0 + block_len: int = 0 + attn_backend_name: str = "aiter" + + +# --------------------------------------------------------------------------- +# Enums & role management +# --------------------------------------------------------------------------- + + +class Role(Enum): + """Role of the current engine instance in the P/D architecture.""" + + PRODUCER = "producer" + CONSUMER = "consumer" + NOT_INITIALIZED = "not_initialized" + + +def convert_virtual_to_physical_pages( + virtual_pages: list[int], + virtual_block_size: int = 16, + physical_block_size: int = 1, +) -> list[int]: + """Expand virtual (coarse) block IDs into physical (fine-grained) page IDs. + + In paged-attention the scheduler works with *virtual* blocks of + ``virtual_block_size`` tokens, but the RDMA transfer operates at + ``physical_block_size`` granularity. + + Args: + virtual_pages: List of virtual block IDs. + virtual_block_size: Tokens per virtual block. + physical_block_size: Tokens per physical block. + + Returns: + Expanded list of physical page IDs. + """ + block_ratio = virtual_block_size // physical_block_size + physical_pages: list[int] = [] + for vp in virtual_pages: + start = vp * block_ratio + physical_pages.extend(range(start, start + block_ratio)) + return physical_pages + + +class _RoleManager: + """Thread-safe singleton that tracks the P/D role of this process. + + Use the module-level :func:`get_role` / :func:`set_role` helpers + instead of accessing this class directly. + """ + + _instance: Optional[_RoleManager] = None + _lock = threading.Lock() + + def __init__(self) -> None: + self._role: Role = Role.NOT_INITIALIZED + + @classmethod + def get_instance(cls) -> _RoleManager: + """Return the singleton, creating it on first call.""" + if cls._instance is None: + with cls._lock: + if cls._instance is None: + instance = object.__new__(cls) + instance.__init__() + cls._instance = instance + return cls._instance + + def set_role(self, role: Role) -> None: + with self._lock: + self._role = role + + @property + def role(self) -> Role: + return self._role + + +def set_role(role: Role) -> None: + """Set the global P/D role for this process.""" + _RoleManager.get_instance().set_role(role) + + +def get_role() -> Role: + """Get the global P/D role for this process.""" + return _RoleManager.get_instance().role + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + + +class MoRIIOConstants: + """Protocol constants for the MoRIIO-based KV connector.""" + + # ZMQ handshake message types + GET_META_MSG = b"get_meta_msg" + POP_DONE_RECV = b"pop_done_recv" + OVER = b"OVER" + COMPLETION_PREFIX = "cmpl" + + # Service discovery + PING_INTERVAL_SECONDS = 5 + MAX_PING_RETRIES = 100 + + # Networking + DEFAULT_HANDSHAKE_PORT = 6301 + DEFAULT_NOTIFY_PORT = "61005" + + # Timeouts + ABORT_REQUEST_TIMEOUT = 3600 + + +def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: + return (dp_rank) * tp_size + tp_rank diff --git a/atom/kv_transfer/disaggregation/kv_transfer_engine.py b/atom/kv_transfer/disaggregation/moriio/moriio_connector.py similarity index 67% rename from atom/kv_transfer/disaggregation/kv_transfer_engine.py rename to atom/kv_transfer/disaggregation/moriio/moriio_connector.py index 7b83833590..d287a135d8 100644 --- a/atom/kv_transfer/disaggregation/kv_transfer_engine.py +++ b/atom/kv_transfer/disaggregation/moriio/moriio_connector.py @@ -2,26 +2,10 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. """ -KV Cache Connector for Disaggregated Prefill-Decode (P/D) Architecture. +Worker-side and scheduler-side KV cache connectors for disaggregated P/D. -This module implements the KV cache transfer mechanism for disaggregated -inference, where prefill and decode stages run on separate GPU instances. -It uses RDMA-based zero-copy transfers via the MoRIIO library for efficient +Uses RDMA-based zero-copy transfers via the MoRIIO library for efficient KV cache migration between producer (prefill) and consumer (decode) nodes. - -Key Components: - - KVConnector: Worker-side connector managing RDMA transfers and handshakes. - - KVConnectorScheduler: Scheduler-side connector coordinating transfer state. - - MoRIIOWrapper: Abstraction over the MoRIIO RDMA engine. - - ConnectorMetadata: Transfer metadata exchanged between scheduler and workers. - -Transfer Modes: - - Read mode: The decode instance reads KV cache directly from prefill memory. - Prefill completes first, then decode reads the blocks via RDMA. - -Architecture:: - - Proxy (service discovery) <-> Prefill Instance <--RDMA--> Decode Instance """ from __future__ import annotations @@ -32,26 +16,29 @@ import time from collections import defaultdict from concurrent.futures import Future, ThreadPoolExecutor -from enum import Enum -from typing import TYPE_CHECKING, Any, Optional +from typing import Any import msgpack import msgspec import numpy as np import zmq -if TYPE_CHECKING: - import torch - from atom.config import Config from atom.kv_transfer.disaggregation.base import ( KVConnectorBase, KVConnectorSchedulerBase, ) +from atom.kv_transfer.disaggregation.moriio.moriio_common import ( + MoRIIOAgentMetadata, + MoRIIOConstants, + _MORIIO_AVAILABLE, + get_port_offset, +) +from atom.kv_transfer.disaggregation.utils import chunk_tensor_for_rdma +from atom.kv_transfer.disaggregation.moriio.moriio_engine import MoRIIOWrapper from atom.kv_transfer.disaggregation.types import ( ConnectorMetadata, EngineId, - RemoteAllocInfo, ReqId, ReqMeta, TransferId, @@ -60,542 +47,32 @@ from atom.utils import ( get_open_port, make_zmq_path, - make_zmq_socket, zmq_socket_ctx, ) from atom.utils.network import get_ip from aiter.dist.parallel_state import get_dp_group, get_tp_group -logger = logging.getLogger("atom") - -# --------------------------------------------------------------------------- -# MoRIIO availability check -# --------------------------------------------------------------------------- - -_MORIIO_AVAILABLE = False -try: +if _MORIIO_AVAILABLE: from mori.io import ( BackendType, - EngineDesc, IOEngine, IOEngineConfig, - MemoryDesc, - MemoryLocationType, PollCqMode, RdmaBackendConfig, ) - _MORIIO_AVAILABLE = True - logger.info("MoRIIO RDMA library loaded successfully") -except ImportError: - logger.warning( - "MoRIIO is not available — KV cache disaggregation will not work. " - "Install the mori package to enable RDMA transfers." - ) - - -# --------------------------------------------------------------------------- -# Msgspec metadata structs -# --------------------------------------------------------------------------- - - -class MoRIIOAgentMetadata( - msgspec.Struct, - omit_defaults=True, - dict=True, - kw_only=True, -): - """Serializable metadata exchanged during the RDMA handshake.""" - - engine_id: str - agent_metadata: bytes - kv_caches_base_addr: Optional[list[int]] = None - num_blocks: int = 0 - block_len: int = 0 - attn_backend_name: str = "aiter" - - -# --------------------------------------------------------------------------- -# Enums & role management -# --------------------------------------------------------------------------- - - -class Role(Enum): - """Role of the current engine instance in the P/D architecture.""" - - PRODUCER = "producer" - CONSUMER = "consumer" - NOT_INITIALIZED = "not_initialized" - - -def convert_virtual_to_physical_pages( - virtual_pages: list[int], - virtual_block_size: int = 16, - physical_block_size: int = 1, -) -> list[int]: - """Expand virtual (coarse) block IDs into physical (fine-grained) page IDs. - - In paged-attention the scheduler works with *virtual* blocks of - ``virtual_block_size`` tokens, but the RDMA transfer operates at - ``physical_block_size`` granularity. - - Args: - virtual_pages: List of virtual block IDs. - virtual_block_size: Tokens per virtual block. - physical_block_size: Tokens per physical block. - - Returns: - Expanded list of physical page IDs. - """ - block_ratio = virtual_block_size // physical_block_size - physical_pages: list[int] = [] - for vp in virtual_pages: - start = vp * block_ratio - physical_pages.extend(range(start, start + block_ratio)) - return physical_pages - - -class _RoleManager: - """Thread-safe singleton that tracks the P/D role of this process. - - Use the module-level :func:`get_role` / :func:`set_role` helpers - instead of accessing this class directly. - """ - - _instance: Optional[_RoleManager] = None - _lock = threading.Lock() - - def __init__(self) -> None: - self._role: Role = Role.NOT_INITIALIZED - - @classmethod - def get_instance(cls) -> _RoleManager: - """Return the singleton, creating it on first call.""" - if cls._instance is None: - with cls._lock: - if cls._instance is None: - instance = object.__new__(cls) - instance.__init__() - cls._instance = instance - return cls._instance - - def set_role(self, role: Role) -> None: - with self._lock: - self._role = role - - @property - def role(self) -> Role: - return self._role - - -def set_role(role: Role) -> None: - """Set the global P/D role for this process.""" - _RoleManager.get_instance().set_role(role) - - -def get_role() -> Role: - """Get the global P/D role for this process.""" - return _RoleManager.get_instance().role - - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- - - -class MoRIIOConstants: - """Protocol constants for the MoRIIO-based KV connector.""" - - # ZMQ handshake message types - GET_META_MSG = b"get_meta_msg" - POP_DONE_RECV = b"pop_done_recv" - OVER = b"OVER" - COMPLETION_PREFIX = "cmpl" - - # Service discovery - PING_INTERVAL_SECONDS = 5 - MAX_PING_RETRIES = 100 - - # Networking - DEFAULT_HANDSHAKE_PORT = 6301 - DEFAULT_NOTIFY_PORT = "61005" - - # Timeouts - ABORT_REQUEST_TIMEOUT = 3600 - - -def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: - return (dp_rank) * tp_size + tp_rank - - -MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 # just under 2 GiB - - -def _chunk_tensor_for_rdma( - tensor: torch.Tensor, block_size_in_dim0: int = 1 -) -> tuple[list[tuple[int, int]], int]: - """Split a tensor into <2 GiB RDMA-registrable chunks along dim 0. - - Args: - tensor: contiguous torch.Tensor whose dim-0 is the block (or - token) axis. - block_size_in_dim0: elements per logical block in dim 0. - Non-MLA: 1 (dim 0 = num_blocks). - MLA: block_size (dim 0 = num_blocks * block_size). - - Returns: - ``(chunks, blocks_per_chunk)`` where *chunks* is a list of - ``(data_ptr, size_bytes)`` pairs and *blocks_per_chunk* is - the number of logical blocks in each full chunk. - """ - elem_sz = tensor.element_size() - per_block_bytes = block_size_in_dim0 * tensor.stride(0) * elem_sz - total_blocks = tensor.shape[0] // block_size_in_dim0 - bpc = max(1, MAX_RDMA_CHUNK_BYTES // per_block_bytes) - chunks: list[tuple[int, int]] = [] - base = tensor.data_ptr() - for start in range(0, total_blocks, bpc): - end = min(start + bpc, total_blocks) - chunks.append((base + start * per_block_bytes, (end - start) * per_block_bytes)) - return chunks, bpc - - -# =================================================================== -# MoRIIOWrapper — thin abstraction over the RDMA engine -# =================================================================== - - -class MoRIIOWrapper: - """Low-level wrapper around a MoRIIO ``IOEngine``. - - Provides helper methods for memory registration, session management, - and asynchronous RDMA read/write operations. Both producer and - consumer code paths share this wrapper. - - Thread-safety: - ``transfer_status``, ``done_req_ids``, and ``done_write_cache_req_ids`` - are guarded by ``self.lock``. The ZMQ socket cache ``self._sockets`` - is *not* thread-safe — callers must ensure ``send_notify`` is invoked - from a single thread. - - Args: - moriio_engine: MoRIIO IOEngine instance. - tp_rank: Tensor parallel rank. - dp_rank: Data parallel rank. - """ - - def __init__( - self, - moriio_engine: Any = None, - tp_rank: int = 0, - dp_rank: int = 0, - ): - self.tp_rank = tp_rank - self.dp_rank = dp_rank - self.moriio_engine = moriio_engine - - self.remote_memory_metadata: Any = None - self.local_memory_registered: bool = False - # Raw-pointer registration produces multiple MemoryDesc per layer - # (e.g. K + V on MHA) to keep each ibv_reg_mr below the AINIC ~2 GiB - # limit. We retain handles here purely so they aren't GC'd; nothing - # in the active session/transfer path indexes into this list. - self.local_memory_descs: list[Any] = [] - self.transfer_status: list[Any] = [] - self.remote_engine_ip: str | None = None - self.notify_port: int | None = None - - self.lock = threading.Lock() - self.done_req_ids: list[str] = [] - self.done_write_cache_req_ids: list[str] = [] - self.notify_thread: threading.Thread | None = None - - # ZMQ socket cache keyed by endpoint path - self._sockets: dict[str, zmq.Socket] = {} - - def set_moriio_engine(self, moriio_engine: Any) -> None: - """Assign the MoRIIO engine (must not be None).""" - if moriio_engine is None: - raise ValueError("Cannot assign a None MoRIIO engine") - self.moriio_engine = moriio_engine - - def set_backend_type(self, backend_type, backend_config=None): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - if backend_config is not None: - self.moriio_engine.create_backend(backend_type, backend_config) - else: - self.moriio_engine.create_backend(backend_type) - - def get_agent_metadata(self): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - engine_metadata = self.moriio_engine.get_engine_desc() - engine_metadata_packed = engine_metadata.pack() - return engine_metadata_packed - - def register_remote_engine(self, remote_packed_engine_metadata): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata) - self.moriio_engine.register_remote_engine(consumer_engine_metadata) - logger.info( - "Registered remote engine with key: %s", consumer_engine_metadata.key - ) - return consumer_engine_metadata.key - - def register_local_buffer(self, ptr: int, size: int, device_id: int) -> bytes: - """Register one raw GPU memory region with MoRIIO. - - Using ``register_memory(ptr, size, ...)`` directly (instead of - ``register_torch_tensor``) lets callers split a single tensor into - multiple smaller regions, which is required to stay under the - AINIC ~2 GiB ``ibv_reg_mr`` limit. - """ - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - try: - desc = self.moriio_engine.register_memory( - ptr, size, device_id, MemoryLocationType.GPU - ) - assert desc is not None, "register_memory returned None" - packed = desc.pack() - except Exception as e: - raise ValueError(f"Failed to register local memory: {e}") from e - self.local_memory_descs.append(desc) - self.local_memory_registered = True - return packed - - def register_local_tensor(self, tensor): - """Back-compat helper: register an entire contiguous torch tensor. - - Prefer :meth:`register_local_buffer` from new code so that callers - explicitly choose how to chunk large tensors before registration. - """ - if not tensor.is_contiguous(): - raise RuntimeError("input tensor must be contiguous") - ptr = tensor.data_ptr() - size = tensor.numel() * tensor.element_size() - device_id = tensor.device.index if tensor.device.index is not None else -1 - return self.register_local_buffer(ptr, size, device_id) - - def get_unpack_memory_metadata(self, packed_memory_metadata): - return MemoryDesc.unpack(packed_memory_metadata) - - def build_session(self, local_memory_metadata, remote_memory_metadata): - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - tmp = self.moriio_engine.create_session( - local_memory_metadata, remote_memory_metadata - ) - - return tmp - - def read_remote_data( - self, transfer_size_byte, local_offset=0, remote_offset=0, session=None - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - transfer_status = session.batch_read( - local_offset, - remote_offset, - transfer_size_byte, - self.moriio_engine.allocate_transfer_uid(), - ) - return transfer_status - - def write_remote_data( - self, transfer_size_byte, local_offset=0, remote_offset=0, session=None - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - write_uid = self.moriio_engine.allocate_transfer_uid() - - transfer_status = session.batch_write( - local_offset, remote_offset, transfer_size_byte, write_uid - ) - with self.lock: - self.transfer_status.append(transfer_status) - - def write_remote_data_single( - self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0 - ): - assert self.local_memory_registered, "You have not register local memory data!" - assert self.moriio_engine is not None, "MoRIIO engine must be set first" - transfer_status = self.sessions[sess_idx].write( - local_offset, - remote_offset, - transfer_size_byte, - self.moriio_engine.allocate_transfer_uid(), - ) - with self.lock: - self.transfer_status.append(transfer_status) - - def waiting_for_transfer_complete(self): - if not self.transfer_status: - return - - transfers_to_wait = [] - with self.lock: - transfers_to_wait = self.transfer_status[:] - self.transfer_status.clear() - - for status in transfers_to_wait: - try: - status.Wait() - if not status.Succeeded(): - logger.error( - "Transfer failed: %s, Code: %s", status.Message(), status.Code() - ) - raise ValueError("MoRIIO transfer failed!") - except Exception as e: - logger.error("Transfer %s failed: %s", status, e) - raise - - def async_wait_reqid(self): - assert self.notify_port is not None, "Notify port cannot be None" - - if self.notify_thread is not None: - return - - def _async_wait(): - host = "*" - path = make_zmq_path("tcp", host, self.notify_port) - logger.info("Node starting to listen notify from path = %s", path) - - with _zmq_ctx(zmq.ROUTER, path) as sock: - while True: - try: - identity, msg = sock.recv_multipart() - self._dispatch_message(msg) - except Exception as e: - logger.error("Error processing message: %s", e) - raise ValueError(f"Error processing message: {e}") from e - - self.notify_thread = threading.Thread( - target=_async_wait, daemon=True, name="moriio-notify-listener" - ) - self.notify_thread.start() - - def _dispatch_message(self, msg: bytes) -> None: - """Route an incoming ZMQ message to the appropriate handler. - - Message formats: - - msgpack dict with ``req_id``: remote block allocation (producer only) - - UTF-8 string prefixed with ``cmpl``: transfer completion signal - """ - # Try msgpack structured message first (block allocation from decode) - try: - data = msgpack.loads(msg) - if isinstance(data, dict) and "req_id" in data: - self._handle_block_alloc_message(data) - return - except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException): - pass - - # Fall back to string-encoded completion message - try: - msg_str = msg.decode("utf-8") - except UnicodeDecodeError: - logger.warning("Received non-decodable message of %d bytes", len(msg)) - return - - if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX): - self._handle_completion_message(msg_str) - else: - raise ValueError(f"Unrecognized message format: {msg_str!r}") - - def _handle_block_alloc_message(self, data: dict) -> None: - """Process a remote block allocation notification (producer side).""" - assert ( - get_role() == Role.PRODUCER - ), "Only producer should receive block alloc messages" - req_id = data["req_id"] - block_notify_list = data.get("block_notify_list", []) - decode_dp_rank = data.get("decode_rank", 0) - assert ( - len(block_notify_list) > 0 - ), "block_notify_list cannot be empty in remote allocate message" - - with self.lock: - self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo( - block_ids=block_notify_list, decode_dp_rank=decode_dp_rank - ) - - def _handle_completion_message(self, msg: str) -> None: - """Record a transfer completion notification.""" - with self.lock: - if get_role() == Role.PRODUCER: - self.done_req_ids.append(msg) - else: - self.done_write_cache_req_ids.append(msg) - - def send_notify( - self, - req_ids: str | int | list[str | int], - remote_ip: str, - remote_port: str | int, - ) -> None: - """Notify a remote engine that transfer(s) have completed.""" - if not remote_ip or not remote_port: - logger.warning("Cannot send notification: missing remote_ip or remote_port") - return - - path = make_zmq_path("tcp", remote_ip, int(remote_port)) - - if path not in self._sockets: - ctx = zmq.Context.instance() - self._sockets[path] = make_zmq_socket( - ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False - ) - - id_list = req_ids if isinstance(req_ids, list) else [req_ids] - sock = self._sockets[path] - try: - for rid in id_list: - rid_str = str(rid) if isinstance(rid, int) else rid - if not isinstance(rid_str, str): - logger.warning("Skipping non-string req_id of type %s", type(rid)) - continue - sock.send_multipart( - [MoRIIOConstants.POP_DONE_RECV, rid_str.encode("utf-8")] - ) - except Exception as e: - logger.error("Failed to send notification to %s: %s", path, e) - self._sockets.pop(path, None) - raise - - def pop_finished_req_ids(self) -> set[str]: - """Return and clear the set of completed send-side request IDs.""" - with self.lock: - result = set(self.done_req_ids) - self.done_req_ids.clear() - return result - - def pop_finished_write_req_ids(self) -> set[str]: - """Return and clear the set of completed write-side request IDs.""" - with self.lock: - result = set(self.done_write_cache_req_ids) - self.done_write_cache_req_ids.clear() - return result - - def shutdown(self) -> None: - """Close all cached ZMQ sockets and release resources.""" - logger.debug( - "Shutting down MoRIIOWrapper, closing %d sockets", len(self._sockets) - ) - for path, sock in self._sockets.items(): - try: - sock.close(linger=0) - except Exception as e: - logger.warning("Error closing socket for %s: %s", path, e) - self._sockets.clear() +logger = logging.getLogger("atom") # =================================================================== -# KVConnector — worker-side connector (runs inside each TP rank) +# MoRIIOConnector — worker-side connector (runs inside each TP rank) # =================================================================== -class KVConnector(KVConnectorBase): +class MoRIIOConnector(KVConnectorBase): """Worker-side KV cache connector for disaggregated P/D inference. - Each tensor-parallel worker instantiates one ``KVConnector``. It is + Each tensor-parallel worker instantiates one ``MoRIIOConnector``. It is responsible for: 1. Registering local KV cache tensors for RDMA access. @@ -727,7 +204,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: transfers can occur. Each K (and V, when present) tensor is split into block-aligned - chunks of < 2 GiB via :func:`_chunk_tensor_for_rdma` and each + chunks of < 2 GiB via :func:`chunk_tensor_for_rdma` and each chunk is registered independently with ``ibv_reg_mr``. Per-layer metadata list layout: ``[k_chunk0, k_chunk1, ..., v_chunk0, v_chunk1, ...]``. @@ -754,7 +231,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: # equals the KV cache block_size (typically 16). # Non-MLA: dim 0 = num_blocks directly, so 1. bsd0 = self.kv_cache_block_size if is_mla else 1 - k_chunks, bpc = _chunk_tensor_for_rdma(cache_tensor, bsd0) + k_chunks, bpc = chunk_tensor_for_rdma(cache_tensor, bsd0) if self.blocks_per_chunk is None: self.blocks_per_chunk = bpc @@ -770,7 +247,7 @@ def register_kv_caches(self, kv_caches: dict[str, Any]) -> None: v_device_id = ( v_cache.device.index if v_cache.device.index is not None else -1 ) - v_chunks, _ = _chunk_tensor_for_rdma(v_cache, 1) + v_chunks, _ = chunk_tensor_for_rdma(v_cache, 1) for ptr, size in v_chunks: meta_list.append( self.moriio_wrapper.register_local_buffer( @@ -1454,11 +931,11 @@ def get_finished(self) -> tuple[set[int], set[str]]: # =================================================================== -# KVConnectorScheduler — scheduler-side connector +# MoRIIOConnectorScheduler — scheduler-side connector # =================================================================== -class KVConnectorScheduler(KVConnectorSchedulerBase): +class MoRIIOConnectorScheduler(KVConnectorSchedulerBase): """Scheduler-side KV connector that tracks transfer lifecycle. Runs in the scheduler process (not in TP workers). Responsible for: @@ -1564,6 +1041,7 @@ def request_finished(self, seq: Sequence) -> None: the transfer_id mapping. """ # Attach output metadata for the proxy to relay + first_token_id = seq.output_tokens[0] if seq.output_tokens else None seq.kv_transfer_params_output = { "do_remote_prefill": True, "do_remote_decode": False, @@ -1574,6 +1052,7 @@ def request_finished(self, seq: Sequence) -> None: "tp_size": self.tp_size, "dp_rank": self.dp_rank, "transfer_id": seq.id, + "first_token_id": first_token_id, } # Clean up transfer ID mapping on the consumer side diff --git a/atom/kv_transfer/disaggregation/moriio/moriio_engine.py b/atom/kv_transfer/disaggregation/moriio/moriio_engine.py new file mode 100644 index 0000000000..beda8dc46a --- /dev/null +++ b/atom/kv_transfer/disaggregation/moriio/moriio_engine.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Low-level RDMA wrapper around a MoRIIO ``IOEngine``. + +Provides helper methods for memory registration, session management, +and asynchronous RDMA read/write operations. Both producer and +consumer code paths share this wrapper. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +import msgpack +import zmq + +from atom.kv_transfer.disaggregation.moriio.moriio_common import ( + MoRIIOConstants, + Role, + _MORIIO_AVAILABLE, + get_role, +) +from atom.kv_transfer.disaggregation.types import RemoteAllocInfo +from atom.utils import make_zmq_path, make_zmq_socket + +if _MORIIO_AVAILABLE: + from mori.io import EngineDesc, MemoryDesc, MemoryLocationType + +logger = logging.getLogger("atom") + + +class MoRIIOWrapper: + """Low-level wrapper around a MoRIIO ``IOEngine``. + + Provides helper methods for memory registration, session management, + and asynchronous RDMA read/write operations. Both producer and + consumer code paths share this wrapper. + + Thread-safety: + ``transfer_status``, ``done_req_ids``, ``done_write_cache_req_ids``, + and ``done_remote_allocate_req_dict`` are guarded by ``self.lock``. The ZMQ socket cache ``self._sockets`` + is *not* thread-safe — callers must ensure ``send_notify`` is invoked + from a single thread. + + Args: + moriio_engine: MoRIIO IOEngine instance. + tp_rank: Tensor parallel rank. + dp_rank: Data parallel rank. + """ + + def __init__( + self, + moriio_engine: Any = None, + tp_rank: int = 0, + dp_rank: int = 0, + ): + self.tp_rank = tp_rank + self.dp_rank = dp_rank + self.moriio_engine = moriio_engine + + self.remote_memory_metadata: Any = None + self.local_memory_registered: bool = False + # Raw-pointer registration produces multiple MemoryDesc per layer + # (e.g. K + V on MHA) to keep each ibv_reg_mr below the AINIC ~2 GiB + # limit. We retain handles here purely so they aren't GC'd; nothing + # in the active session/transfer path indexes into this list. + self.local_memory_descs: list[Any] = [] + self.transfer_status: list[Any] = [] + self.remote_engine_ip: str | None = None + self.notify_port: int | None = None + + self.lock = threading.Lock() + self.done_req_ids: list[str] = [] + self.done_write_cache_req_ids: list[str] = [] + self.done_remote_allocate_req_dict: dict[str, Any] = {} + self.notify_thread: threading.Thread | None = None + + # ZMQ socket cache keyed by endpoint path + self._sockets: dict[str, zmq.Socket] = {} + + def set_moriio_engine(self, moriio_engine: Any) -> None: + """Assign the MoRIIO engine (must not be None).""" + if moriio_engine is None: + raise ValueError("Cannot assign a None MoRIIO engine") + self.moriio_engine = moriio_engine + + def set_backend_type(self, backend_type, backend_config=None): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + if backend_config is not None: + self.moriio_engine.create_backend(backend_type, backend_config) + else: + self.moriio_engine.create_backend(backend_type) + + def get_agent_metadata(self): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + engine_metadata = self.moriio_engine.get_engine_desc() + engine_metadata_packed = engine_metadata.pack() + return engine_metadata_packed + + def register_remote_engine(self, remote_packed_engine_metadata): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata) + self.moriio_engine.register_remote_engine(consumer_engine_metadata) + logger.info( + "Registered remote engine with key: %s", consumer_engine_metadata.key + ) + return consumer_engine_metadata.key + + def register_local_buffer(self, ptr: int, size: int, device_id: int) -> bytes: + """Register one raw GPU memory region with MoRIIO. + + Using ``register_memory(ptr, size, ...)`` directly (instead of + ``register_torch_tensor``) lets callers split a single tensor into + multiple smaller regions, which is required to stay under the + AINIC ~2 GiB ``ibv_reg_mr`` limit. + """ + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + try: + desc = self.moriio_engine.register_memory( + ptr, size, device_id, MemoryLocationType.GPU + ) + assert desc is not None, "register_memory returned None" + packed = desc.pack() + except Exception as e: + raise ValueError(f"Failed to register local memory: {e}") from e + self.local_memory_descs.append(desc) + self.local_memory_registered = True + return packed + + def register_local_tensor(self, tensor): + """Back-compat helper: register an entire contiguous torch tensor. + + Prefer :meth:`register_local_buffer` from new code so that callers + explicitly choose how to chunk large tensors before registration. + """ + if not tensor.is_contiguous(): + raise RuntimeError("input tensor must be contiguous") + ptr = tensor.data_ptr() + size = tensor.numel() * tensor.element_size() + device_id = tensor.device.index if tensor.device.index is not None else -1 + return self.register_local_buffer(ptr, size, device_id) + + def get_unpack_memory_metadata(self, packed_memory_metadata): + return MemoryDesc.unpack(packed_memory_metadata) + + def build_session(self, local_memory_metadata, remote_memory_metadata): + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + tmp = self.moriio_engine.create_session( + local_memory_metadata, remote_memory_metadata + ) + + return tmp + + def read_remote_data( + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + transfer_status = session.batch_read( + local_offset, + remote_offset, + transfer_size_byte, + self.moriio_engine.allocate_transfer_uid(), + ) + return transfer_status + + def write_remote_data( + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + write_uid = self.moriio_engine.allocate_transfer_uid() + + transfer_status = session.batch_write( + local_offset, remote_offset, transfer_size_byte, write_uid + ) + with self.lock: + self.transfer_status.append(transfer_status) + + def write_remote_data_single( + self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0 + ): + assert self.local_memory_registered, "You have not register local memory data!" + assert self.moriio_engine is not None, "MoRIIO engine must be set first" + transfer_status = self.sessions[sess_idx].write( + local_offset, + remote_offset, + transfer_size_byte, + self.moriio_engine.allocate_transfer_uid(), + ) + with self.lock: + self.transfer_status.append(transfer_status) + + def waiting_for_transfer_complete(self): + if not self.transfer_status: + return + + transfers_to_wait = [] + with self.lock: + transfers_to_wait = self.transfer_status[:] + self.transfer_status.clear() + + for status in transfers_to_wait: + try: + status.Wait() + if not status.Succeeded(): + logger.error( + "Transfer failed: %s, Code: %s", status.Message(), status.Code() + ) + raise ValueError("MoRIIO transfer failed!") + except Exception as e: + logger.error("Transfer %s failed: %s", status, e) + raise + + def async_wait_reqid(self): + assert self.notify_port is not None, "Notify port cannot be None" + + if self.notify_thread is not None: + return + + def _async_wait(): + from atom.kv_transfer.disaggregation.moriio.moriio_connector import ( + _zmq_ctx, + ) + + host = "*" + path = make_zmq_path("tcp", host, self.notify_port) + logger.info("Node starting to listen notify from path = %s", path) + + with _zmq_ctx(zmq.ROUTER, path) as sock: + while True: + try: + identity, msg = sock.recv_multipart() + self._dispatch_message(msg) + except Exception as e: + logger.error("Error processing message: %s", e) + raise ValueError(f"Error processing message: {e}") from e + + self.notify_thread = threading.Thread( + target=_async_wait, daemon=True, name="moriio-notify-listener" + ) + self.notify_thread.start() + + def _dispatch_message(self, msg: bytes) -> None: + """Route an incoming ZMQ message to the appropriate handler. + + Message formats: + - msgpack dict with ``req_id``: remote block allocation (producer only) + - UTF-8 string prefixed with ``cmpl``: transfer completion signal + """ + # Try msgpack structured message first (block allocation from decode) + try: + data = msgpack.loads(msg) + if isinstance(data, dict) and "req_id" in data: + self._handle_block_alloc_message(data) + return + except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException): + pass + + # Fall back to string-encoded completion message + try: + msg_str = msg.decode("utf-8") + except UnicodeDecodeError: + logger.warning("Received non-decodable message of %d bytes", len(msg)) + return + + if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX): + self._handle_completion_message(msg_str) + else: + raise ValueError(f"Unrecognized message format: {msg_str!r}") + + def _handle_block_alloc_message(self, data: dict) -> None: + """Process a remote block allocation notification (producer side).""" + assert ( + get_role() == Role.PRODUCER + ), "Only producer should receive block alloc messages" + req_id = data["req_id"] + block_notify_list = data.get("block_notify_list", []) + decode_dp_rank = data.get("decode_rank", 0) + assert ( + len(block_notify_list) > 0 + ), "block_notify_list cannot be empty in remote allocate message" + + with self.lock: + self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo( + block_ids=block_notify_list, decode_dp_rank=decode_dp_rank + ) + + def _handle_completion_message(self, msg: str) -> None: + """Record a transfer completion notification.""" + with self.lock: + if get_role() == Role.PRODUCER: + self.done_req_ids.append(msg) + else: + self.done_write_cache_req_ids.append(msg) + + def send_notify( + self, + req_ids: str | int | list[str | int], + remote_ip: str, + remote_port: str | int, + ) -> None: + """Notify a remote engine that transfer(s) have completed.""" + if not remote_ip or not remote_port: + logger.warning("Cannot send notification: missing remote_ip or remote_port") + return + + path = make_zmq_path("tcp", remote_ip, int(remote_port)) + + if path not in self._sockets: + ctx = zmq.Context.instance() + self._sockets[path] = make_zmq_socket( + ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False + ) + + id_list = req_ids if isinstance(req_ids, list) else [req_ids] + sock = self._sockets[path] + try: + for rid in id_list: + rid_str = str(rid) if isinstance(rid, int) else rid + if not isinstance(rid_str, str): + logger.warning("Skipping non-string req_id of type %s", type(rid)) + continue + sock.send_multipart( + [MoRIIOConstants.POP_DONE_RECV, rid_str.encode("utf-8")] + ) + except Exception as e: + logger.error("Failed to send notification to %s: %s", path, e) + self._sockets.pop(path, None) + raise + + def pop_finished_req_ids(self) -> set[str]: + """Return and clear the set of completed send-side request IDs.""" + with self.lock: + result = set(self.done_req_ids) + self.done_req_ids.clear() + return result + + def pop_finished_write_req_ids(self) -> set[str]: + """Return and clear the set of completed write-side request IDs.""" + with self.lock: + result = set(self.done_write_cache_req_ids) + self.done_write_cache_req_ids.clear() + return result + + def shutdown(self) -> None: + """Close all cached ZMQ sockets and release resources.""" + logger.debug( + "Shutting down MoRIIOWrapper, closing %d sockets", len(self._sockets) + ) + for path, sock in self._sockets.items(): + try: + sock.close(linger=0) + except Exception as e: + logger.warning("Error closing socket for %s: %s", path, e) + self._sockets.clear() diff --git a/atom/kv_transfer/disaggregation/proxy.py b/atom/kv_transfer/disaggregation/proxy.py index 81c062074f..f3c6d192d1 100644 --- a/atom/kv_transfer/disaggregation/proxy.py +++ b/atom/kv_transfer/disaggregation/proxy.py @@ -213,13 +213,42 @@ async def stream_decode_response( request_id: str, ): """Yield response chunks from a decode instance, then close the session.""" + chunk_count = 0 try: if response.status != 200: raise RuntimeError( f"Decode request {request_id} failed with status {response.status}" ) - async for chunk_bytes in response.content.iter_chunked(1024): + while True: + try: + chunk_bytes = await asyncio.wait_for( + response.content.readany(), timeout=120.0 + ) + except asyncio.TimeoutError: + logger.warning( + "[PROXY] stream_decode TIMEOUT for %s after %d chunks " + "(120s no data), closing", + request_id, + chunk_count, + ) + break + if not chunk_bytes: + break + chunk_count += 1 yield chunk_bytes + logger.info( + "[PROXY] stream_decode finished for %s, %d chunks forwarded", + request_id, + chunk_count, + ) + except Exception as e: + logger.exception( + "[PROXY] stream_decode ERROR for %s after %d chunks: %s", + request_id, + chunk_count, + e, + ) + raise finally: await session.close() @@ -254,6 +283,7 @@ async def handle_request(): req_data = await request.get_json() request_id = str(uuid.uuid4()) + logger.info("[PROXY] handle_request #%d, id=%s", request_nums, request_id) if not prefill_instances or not decode_instances: return await make_response( @@ -300,7 +330,10 @@ async def handle_request(): ) # --- Decode request --- - req_data["max_tokens"] -= 1 + # NOTE: Do NOT decrement max_tokens. The decode engine's first token + # (T0 re-prediction) is overridden with prefill's T0 via first_token_id, + # but that decode step still counts toward max_tokens. Keeping the + # original max_tokens ensures the total output length matches non-PD. req_data["kv_transfer_params"] = { "do_remote_decode": False, "do_remote_prefill": True, @@ -313,26 +346,30 @@ async def handle_request(): logger.info("Transfer type: %s", TRANSFER_TYPE) + prefill_response = await send_prefill_task + logger.info("Prefill response received for request %s", request_id) + prefill_kv = prefill_response["kv_transfer_params"] + req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] + if "first_token_id" in prefill_kv: + req_data["kv_transfer_params"]["first_token_id"] = prefill_kv[ + "first_token_id" + ] + + actual_dp_rank = prefill_kv.get("dp_rank", selected_prefill_dp_rank) + if actual_dp_rank is not None: + selected_prefill_dp_rank = actual_dp_rank + if TRANSFER_TYPE == "read": - prefill_response = await send_prefill_task - logger.info("Prefill response received for request %s", request_id) - prefill_kv = prefill_response["kv_transfer_params"] req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[ "remote_engine_id" ] req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[ "remote_block_ids" ] - req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] - - # Use the actual dp_rank from the prefill response (authoritative), - # falling back to the proxy's round-robin selection. - actual_dp_rank = prefill_kv.get("dp_rank", selected_prefill_dp_rank) - if actual_dp_rank is not None: - selected_prefill_dp_rank = actual_dp_rank req_data["kv_transfer_params"]["remote_dp_size"] = prefill_ep["dp_size"] req_data["kv_transfer_params"]["remote_tp_size"] = prefill_ep["tp_size"] + req_data["kv_transfer_params"]["tp_size"] = prefill_ep["tp_size"] if selected_prefill_dp_rank is not None: req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank @@ -341,12 +378,19 @@ async def handle_request(): decode_ep["request_address"], req_data, request_id ) stream_gen = stream_decode_response(session, decode_response, request_id) + logger.info( + "[PROXY] decode response status=%d for %s", + decode_response.status, + request_id, + ) response = await make_response(stream_gen) response.headers["Content-Type"] = "application/json; charset=utf-8" return response except Exception as e: - logger.exception("Error handling request: %s", e) + logger.exception( + "[PROXY] Error handling request #%d id=%s: %s", request_nums, request_id, e + ) resp = await make_response((f"Internal Server Error: {e!s}", 500)) resp.headers["Content-Type"] = "application/json; charset=utf-8" return resp @@ -355,11 +399,38 @@ async def handle_request(): _DEFAULT_DISCOVERY_PORT = 36367 +@app.route("/start_profile", methods=["POST"]) +@app.route("/stop_profile", methods=["POST"]) +async def proxy_profile(): + """Forward profiler start/stop to all registered prefill and decode instances.""" + path = request.path # "/start_profile" or "/stop_profile" + results = {} + async with aiohttp.ClientSession() as session: + for tag, instances in [ + ("prefill", prefill_instances), + ("decode", decode_instances), + ]: + for i, ep in enumerate(instances): + ip, port = _extract_ip_port(ep["request_address"]) + url = f"http://{ip}:{port}{path}" + try: + async with session.post( + url, timeout=aiohttp.ClientTimeout(total=600) + ) as resp: + results[f"{tag}_{i}"] = { + "status": resp.status, + "body": await resp.json(), + } + except Exception as e: + results[f"{tag}_{i}"] = {"status": "error", "detail": str(e)} + return results + + def main(port: int = 10001): """Launch the P/D proxy service.""" discovery_thread = start_service_discovery("0.0.0.0", _DEFAULT_DISCOVERY_PORT) - app.debug = True + app.debug = False app.config["BODY_TIMEOUT"] = 360000 app.config["RESPONSE_TIMEOUT"] = 360000 diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index fdcc6ab1ab..ca24b0f111 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -129,7 +129,11 @@ def _build_req_meta( remote_handshake_port=kv_transfer_params["remote_handshake_port"], remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), - tp_size=kv_transfer_params.get("tp_size", 1), + tp_size=( + kv_transfer_params.get("tp_size") + if "tp_size" in kv_transfer_params + else kv_transfer_params.get("remote_tp_size", 1) + ), transfer_id=kv_transfer_params.get("transfer_id", 0), ) diff --git a/atom/kv_transfer/disaggregation/utils.py b/atom/kv_transfer/disaggregation/utils.py new file mode 100644 index 0000000000..ea705d1ccc --- /dev/null +++ b/atom/kv_transfer/disaggregation/utils.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +""" +Shared utilities for KV cache disaggregation backends (MoRIIO, Mooncake, etc.). +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import triton +import triton.language as tl + +if TYPE_CHECKING: + import torch + +logger = logging.getLogger("atom") + +MAX_RDMA_CHUNK_BYTES = 2 * 1024 * 1024 * 1024 - 64 * 1024 # just under 2 GiB + + +# --------------------------------------------------------------------------- +# GPU memory fence for RDMA-written KV cache blocks +# +# Not needed when producer and consumer are in the same network partition +# (same-partition RDMA completes before decode starts). +# Enable for cross-partition deployments where higher RDMA latency +# can cause stale GPU reads (improves accuracy from ~0.85 to ~0.90). +# --------------------------------------------------------------------------- + + +@triton.jit +def _gpu_fence_kernel( + kv_ptr, + block_ids_ptr, + total_blocks: tl.int64, + elems_per_block: tl.int64, + NUM_CL: tl.constexpr, +): + blk_idx = tl.program_id(0) + layer_idx = tl.program_id(1).to(tl.int64) + blk = tl.load(block_ids_ptr + blk_idx).to(tl.int64) + base = (layer_idx * total_blocks + blk) * elems_per_block + cl_offsets = tl.arange(0, NUM_CL).to(tl.int64) * 32 + mask = cl_offsets < elems_per_block + tl.atomic_or( + kv_ptr + base + cl_offsets, tl.zeros([NUM_CL], dtype=tl.int32), mask=mask + ) + + +def gpu_memory_fence( + kv_cache: "torch.Tensor", + block_ids: list[int], + use_mla: bool, +) -> None: + """Force GPU memory coherence for RDMA-written KV cache blocks. + + Single Triton kernel launch -- one atomic_or(0) per cache line (128 bytes) + across all (layer, block) pairs. + + Args: + kv_cache: The full KV cache tensor. + block_ids: Physical block IDs that received RDMA writes. + use_mla: Whether the model uses Multi-head Latent Attention. + """ + import torch + + block_ids_t = torch.tensor(block_ids, dtype=torch.int32, device=kv_cache.device) + kv_flat = kv_cache.view(torch.int32) + if use_mla: + num_layers = kv_cache.shape[0] + total_blocks = kv_cache.shape[1] + else: + num_layers = 2 * kv_cache.shape[1] + total_blocks = kv_cache.shape[2] + elems_per_block = kv_flat.numel() // (num_layers * total_blocks) + num_cl = triton.cdiv(elems_per_block, 32) + NUM_CL = triton.next_power_of_2(num_cl) + _gpu_fence_kernel[(len(block_ids), num_layers)]( + kv_flat, + block_ids_t, + total_blocks, + elems_per_block, + NUM_CL=NUM_CL, + ) + + +def chunk_tensor_for_rdma( + tensor: torch.Tensor, block_size_in_dim0: int = 1 +) -> tuple[list[tuple[int, int]], int]: + """Split a tensor into <2 GiB RDMA-registrable chunks along dim 0. + + Args: + tensor: contiguous torch.Tensor whose dim-0 is the block (or + token) axis. + block_size_in_dim0: elements per logical block in dim 0. + Non-MLA: 1 (dim 0 = num_blocks). + MLA: block_size (dim 0 = num_blocks * block_size). + + Returns: + ``(chunks, blocks_per_chunk)`` where *chunks* is a list of + ``(data_ptr, size_bytes)`` pairs and *blocks_per_chunk* is + the number of logical blocks in each full chunk. + """ + elem_sz = tensor.element_size() + per_block_bytes = block_size_in_dim0 * tensor.stride(0) * elem_sz + total_blocks = tensor.shape[0] // block_size_in_dim0 + bpc = max(1, MAX_RDMA_CHUNK_BYTES // per_block_bytes) + chunks: list[tuple[int, int]] = [] + base = tensor.data_ptr() + for start in range(0, total_blocks, bpc): + end = min(start + bpc, total_blocks) + chunks.append((base + start * per_block_bytes, (end - start) * per_block_bytes)) + return chunks, bpc diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index d6a3f10017..0dcc56fa55 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -85,8 +85,8 @@ def __init__( """Asynchronously copy the sampled_token_ids tensor to the host.""" # Deferred output is disabled when running in P/D disaggregation mode # (kv_transfer_config is set), enabled otherwise. - # TODO: enable deferred output in P/D disaggregation mode - self.is_deferred_out = not bool(runner.config.kv_transfer_config) + # TODO: In P/D disaggregation mode, if have issue, we can disable it + self.is_deferred_out = True self.runner = runner device = runner.device @@ -313,6 +313,15 @@ def prepare_input_ids( self.input_ids.np[:total_tokens_decode] = token_ids return self.input_ids.copy_to_gpu(total_tokens_decode) + # PD consumer first decode: no prior prefill step initialized + # prev_batch, so use scheduled_tokens directly for this step. + if self.prev_batch is None: + token_ids = scheduled_tokens[ + total_tokens_prefill : total_tokens_prefill + total_tokens_decode + ] + self.input_ids.np[:total_tokens_decode] = token_ids + return self.input_ids.copy_to_gpu(total_tokens_decode) + """for decode: input ids are from prev_sampled_token_ids""" deferred_curr_indices, deferred_prev_indices, new_curr_indices, is_all_same = ( self.get_token_locations(batch) @@ -1879,6 +1888,7 @@ def async_proc_aggregation(self) -> KVConnectorOutput: if connector is None: return KVConnectorOutput(finished_sending=[], finished_recving=[]) done_sending, done_recving = connector.get_finished() + return KVConnectorOutput( finished_sending=done_sending, finished_recving=done_recving ) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index cfe874b275..4cd1b5d2b0 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -224,9 +224,13 @@ def __init__( is_dummy_run: bool = False, num_spec_step: int = 0, scheduled_spec_decode_tokens: dict[int, np.ndarray] | None = None, + remote_kv_block_ids: list[int] | None = None, + remote_kv_seq_blocks: dict[int, list[int]] | None = None, ): if scheduled_spec_decode_tokens is None: scheduled_spec_decode_tokens = {} + self.remote_kv_block_ids = remote_kv_block_ids or [] + self.remote_kv_seq_blocks = remote_kv_seq_blocks or {} self.req_ids = list(seqs.keys()) # self.scheduled_tokens = [ @@ -512,6 +516,27 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: self.kv_connector.get_num_new_matched_tokens(seq) ) + if waiting_remote_to_waiting_ready: + seq.status = SequenceStatus.RUNNING + seq.is_first_decode = True + first_token_id = (seq.kv_transfer_params or {}).get("first_token_id") + if first_token_id is not None: + seq.append_token(first_token_id) + seq._injected_t0 = first_token_id + logger.info( + "[PD-TRANSITION] seq %s: num_tokens=%d, " + "num_prompt=%d, blocks=%d, first_token=%s, " + "last_5_tids=%s", + seq.id, + seq.num_tokens, + seq.num_prompt_tokens, + len(seq.block_table), + first_token_id, + seq.token_ids[-5:], + ) + self.running.append(seq) + continue + num_new_tokens = seq.num_tokens - seq.num_cached_tokens if ( num_batched_tokens + num_new_tokens > self.max_num_batched_tokens @@ -520,8 +545,7 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: self.waiting.appendleft(seq) break - if not waiting_remote_to_waiting_ready: - self.block_manager.allocate(seq) + self.block_manager.allocate(seq) if self.kv_connector is not None: self.kv_connector.update_state_after_alloc(seq) @@ -531,12 +555,6 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: seq.status = SequenceStatus.WAITING_FOR_REMOTE_KVS continue - if waiting_remote_to_waiting_ready: - seq.status = SequenceStatus.RUNNING - seq.is_first_decode = True - self.running.append(seq) - continue - num_seqs_prefill += 1 num_new_tokens = seq.num_tokens - seq.num_cached_tokens if self.cache_stats: @@ -587,6 +605,8 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: # --- Decode scheduling --- num_seqs_decode = 0 num_new_tokens = self.mtp_k + 1 + remote_kv_blocks: set[int] = set() + remote_kv_seq_blocks: dict[int, list[int]] = {} while self.running and num_seqs_decode < self.max_num_seqs: seq = self.running.popleft() while not self.block_manager.can_append(seq, num_new_tokens): @@ -599,10 +619,31 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: if seq.spec_token_ids.size > 0: scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids num_seqs_decode += 1 - # Skip block append for the first decode step after remote - # prefill — blocks were already allocated during prefill. - if not getattr(seq, "is_first_decode", False): + # For PD first-decode: if T0 was injected, may_append is + # needed for the new position N. Without T0 injection, + # blocks were already allocated during prefill. + is_first = getattr(seq, "is_first_decode", False) + if is_first and seq.block_table: + remote_kv_blocks.update(seq.block_table) + remote_kv_seq_blocks[seq.id] = list(seq.block_table) + has_injected_t0 = ( + is_first + and (seq.kv_transfer_params or {}).get("first_token_id") is not None + ) + if not is_first or has_injected_t0: self.block_manager.may_append(seq, num_new_tokens) + if is_first: + logger.info( + "[PD-FIRST-DECODE] seq %s: num_tokens=%d, " + "blocks=%d, injected_t0=%s, " + "last_block_num=%d, context_will_be=%d", + seq.id, + seq.num_tokens, + len(seq.block_table), + has_injected_t0, + seq.last_block_num_tokens, + seq.num_tokens, + ) scheduled_seqs[seq.id] = seq seq.type = SequenceType.DECODE num_scheduled_tokens.append(num_new_tokens) @@ -628,6 +669,8 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: connector_meta_output=connector_meta_output, num_spec_step=self.mtp_k, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + remote_kv_block_ids=sorted(remote_kv_blocks) if remote_kv_blocks else [], + remote_kv_seq_blocks=remote_kv_seq_blocks, ) return (decode_batch, scheduled_seqs) @@ -645,6 +688,7 @@ def preempt(self, seq: Sequence): seq.num_rejected = 0 seq.num_bonus_tokens = 0 seq.spec_token_ids = np.array([], dtype=np.int32) + seq.is_first_decode = False self.block_manager.deallocate(seq) self.waiting.appendleft(seq) @@ -708,11 +752,26 @@ def postprocess( seq.append_token(token_id) new_tokens = token_ids + injected_t0 = getattr(seq, "_injected_t0", None) + if injected_t0 is not None: + new_tokens = [injected_t0] + list(new_tokens) + seq._injected_t0 = None + if self.mtp_k > 0: # idx already resolved above via get_idx seq.spec_token_ids = draft_token_ids[idx] - if seq.num_completion_tokens == 1 and seq.first_token_time == 0.0: + if seq.num_completion_tokens <= 3 and seq.kv_transfer_params: + logger.info( + "[PD-DECODE] seq %s: comp_tokens=%d, " + "new_token=%s, num_tokens=%d, blocks=%d", + seq.id, + seq.num_completion_tokens, + token_ids, + seq.num_tokens, + len(seq.block_table), + ) + if seq.num_completion_tokens >= 1 and seq.first_token_time == 0.0: seq.first_token_time = time.time() num_tokens = seq.num_tokens - self.mtp_k - num_rejected diff --git a/recipes/pd_disaggregation_guide.md b/recipes/pd_disaggregation_guide.md new file mode 100644 index 0000000000..8a21c7c661 --- /dev/null +++ b/recipes/pd_disaggregation_guide.md @@ -0,0 +1,177 @@ +# PD Disaggregation with Mooncake (RDMA) + +Prefill-Decode disaggregation splits inference into two stages on separate nodes: +- **Producer** (prefill): runs prompt prefill, pushes KV cache via RDMA +- **Consumer** (decode): receives KV cache, runs autoregressive decode + +## Prerequisites + +- Two nodes with AMD MI300X GPUs (8 GPUs each for TP=8) +- RDMA network connectivity between nodes (RoCE or InfiniBand) +- Mooncake package installed (`pip install mooncake`) +- Producer and consumer should be in the **same network partition** for best accuracy + +## Building Mooncake for ROCm + +If Mooncake is not pre-installed in your Docker image, build from source: + +### 1. System dependencies + +```bash +apt update && apt install -y \ + zip unzip wget gcc make libtool autoconf cmake \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core \ + openssh-server openmpi-bin openmpi-common libopenmpi-dev +``` + +### 2. Install Go 1.22.2 + +Mooncake's etcd wrapper requires Go. **Must use Go 1.22.2** — Go 1.23+ causes `mallocHeaderSize redeclared` errors. + +```bash +apt remove -y golang golang-go 2>/dev/null || true +rm -rf /usr/local/go +wget https://go.dev/dl/go1.22.2.linux-amd64.tar.gz +tar -C /usr/local -xzf go1.22.2.linux-amd64.tar.gz +rm go1.22.2.linux-amd64.tar.gz +export PATH=$PATH:/usr/local/go/bin +go version # expect: go1.22.2 +``` + +### 3. Clone and build + +Use the SGLang-validated commit: + +```bash +git clone https://github.com/kvcache-ai/Mooncake.git +cd Mooncake +git checkout b6a841dc78c707ec655a563453277d969fb8f38d +git submodule update --init --recursive + +# Install C++ dependencies (etcd, protobuf, gRPC, abseil) +bash dependencies.sh -y + +# CMake build +mkdir -p build && cd build +cmake .. -DUSE_HIP=ON -DUSE_ETCD=ON +make -j$(nproc) +make install +``` + +Key flags: `-DUSE_HIP=ON` for ROCm (default is CUDA), `-DUSE_ETCD=ON` for metadata store. + +### 4. Install Python package + +```bash +cd ../mooncake-wheel +pip install . + +# Copy compiled .so to Python package +MOONCAKE_DIR=$(python -c "import mooncake; print(mooncake.__path__[0])") +cp ../build/mooncake-integration/engine.cpython-*-linux-gnu.so "$MOONCAKE_DIR/" + +ldconfig +``` + +### 5. Verify + +```bash +python -c "from mooncake.engine import TransferEngine; print('Mooncake ROCm OK')" +``` + +If you see `libglog.so` or other missing library errors, install them (`apt install -y libgoogle-glog-dev`) and re-run `ldconfig`. + +> **Important**: Producer and consumer nodes must use the **same Mooncake build version**. Mismatched versions cause `Corrupted segment descriptor` errors due to incompatible metadata serialization formats. + +## Quick Start + +### Step 0: Find local IP + +On each node, find the network IP (not loopback): + +```bash +export LOCAL_IP=$(ip addr show | grep "inet " | grep -v 127.0.0.1 | awk '{print $2}' | cut -d/ -f1 | head -1) +echo "Local IP: ${LOCAL_IP}" +``` + +### Step 1: Start Proxy (on producer node) + +The proxy handles routing between producer and consumer: + +```bash +python -m atom.kv_transfer.disaggregation.proxy --port 10001 +``` + +### Step 2: Start Producer (prefill node) + +```bash +ATOM_DISABLE_MMAP=true \ +NCCL_SOCKET_IFNAME=lo \ +AITER_LOG_LEVEL=WARNING \ +python -m atom.entrypoints.openai_server \ + --model /data/models/DeepSeek-R1/ \ + --kv_cache_dtype fp8 \ + -tp 8 \ + --server-port 8003 \ + --kv-transfer-config '{ + "kv_role": "kv_producer", + "kv_connector": "mooncake", + "proxy_ip": "'"${LOCAL_IP}"'", + "proxy_ping_port": 36367, + "http_port": 8003 + }' \ + 2>&1 | tee producer.log +``` + +### Step 3: Start Consumer (decode node) + +Replace `PRODUCER_IP` with the producer node's IP: + +```bash +export PRODUCER_IP= + +ATOM_DISABLE_MMAP=true \ +NCCL_SOCKET_IFNAME=lo \ +AITER_LOG_LEVEL=WARNING \ +python -m atom.entrypoints.openai_server \ + --model /data/models/DeepSeek-R1/ \ + --kv_cache_dtype fp8 \ + -tp 8 \ + --server-port 8004 \ + --kv-transfer-config '{ + "kv_role": "kv_consumer", + "kv_connector": "mooncake", + "proxy_ip": "'"${PRODUCER_IP}"'", + "proxy_ping_port": 36367, + "http_port": 8004 + }' \ + 2>&1 | tee consumer.log +``` + +### Step 4: Validate Accuracy + +Run GSM8K evaluation against the consumer endpoint: + +```bash +lm_eval --model local-chat-completions \ + --model_args "model=DeepSeek-R1,base_url=http://${CONSUMER_IP}:8004/v1,tokenizer_backend=huggingface,pretrained=/data/models/DeepSeek-R1/" \ + --tasks gsm8k_cot \ + --batch_size 1 \ + --limit 100 \ + --apply_chat_template \ + --fewshot_as_multiturn \ + --predict_only \ + --log_samples \ + --output_path results/ \ + --gen_kwargs "max_tokens=8192,temperature=0.6" +``` + +Expected accuracy: ~0.95-0.96 (matching non-PD baseline). +``` +ewshot: None, batch_size: 1 +|Tasks|Version| Filter |n-shot| Metric | |Value| |Stderr| +|-----|------:|----------------|-----:|-----------|---|----:|---|-----:| +|gsm8k| 3|flexible-extract| 5|exact_match|↑ | 0.96|± | 0.028| +| | |strict-match | 5|exact_match|↑ | 0.96|± | 0.028| +``` \ No newline at end of file diff --git a/test/kv_transfer/disaggregation/test_kv_aggregator.py b/tests/test_kv_aggregator.py similarity index 100% rename from test/kv_transfer/disaggregation/test_kv_aggregator.py rename to tests/test_kv_aggregator.py diff --git a/test/kv_transfer/disaggregation/test_kv_connector_scheduler.py b/tests/test_kv_connector_scheduler.py similarity index 100% rename from test/kv_transfer/disaggregation/test_kv_connector_scheduler.py rename to tests/test_kv_connector_scheduler.py diff --git a/test/kv_transfer/disaggregation/test_proxy.py b/tests/test_proxy.py similarity index 100% rename from test/kv_transfer/disaggregation/test_proxy.py rename to tests/test_proxy.py diff --git a/test/kv_transfer/disaggregation/test_transfer_engine.py b/tests/test_transfer_engine.py similarity index 100% rename from test/kv_transfer/disaggregation/test_transfer_engine.py rename to tests/test_transfer_engine.py diff --git a/test/kv_transfer/disaggregation/test_types.py b/tests/test_types.py similarity index 100% rename from test/kv_transfer/disaggregation/test_types.py rename to tests/test_types.py From 51e9a6599a67e36ba634bfb2d26f985e3e2ee72f Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Thu, 14 May 2026 16:48:53 +0800 Subject: [PATCH 47/76] [Feat] enable dp attention in sgl+ATOM (#743) --- atom/plugin/config.py | 29 ++- atom/plugin/register.py | 2 + .../attention_backend/sgl_attention_mla.py | 10 + .../sglang/models/base_model_wrapper.py | 243 +++++++++++++++++- 4 files changed, 277 insertions(+), 7 deletions(-) diff --git a/atom/plugin/config.py b/atom/plugin/config.py index 07aafa4a52..6d37bd70d7 100644 --- a/atom/plugin/config.py +++ b/atom/plugin/config.py @@ -31,6 +31,7 @@ class PluginConfig: sglang_enable_torch_compile: bool = False sglang_disable_cuda_graph: bool = False sglang_enable_dp_attention: bool = False + sglang_aiter_rank_id: int = 0 sglang_dist_init_addr: Optional[str] = None sglang_port_args: Any = None @@ -148,6 +149,7 @@ def _generate_atom_config_from_sglang_config(config: Any): 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 @@ -212,7 +214,9 @@ def _generate_atom_config_from_sglang_config(config: Any): # sglang uses the atom parallel config sgl_parallel_config = ParallelConfig( data_parallel_size=atom_data_parallel_size, + data_parallel_size_local=atom_data_parallel_size, data_parallel_rank=atom_data_parallel_rank, + data_parallel_rank_local=atom_data_parallel_rank, ) # use sglang torch compile policy and cuda graph policy @@ -224,6 +228,26 @@ def _generate_atom_config_from_sglang_config(config: Any): cudagraph_mode=None, ) + sglang_dist_init_addr = server_args.dist_init_addr + # In single-node DP attention, synthesize the same TCP base address that + # SGLang uses for its DP-attention TCP port family. The primary purpose is + # to avoid calling PortArgs.init_new() again in ATOM plugin mode, because a + # second call would probe that fixed TCP range again and conflict with + # SGLang's existing allocation. In the current plugin path, this value + # should be treated as a compatibility/fallback hint rather than a + # guaranteed representation of the runtime default torch.distributed world + # rendezvous endpoint. + if ( + sglang_dist_init_addr is None + and server_args.enable_dp_attention + and server_args.nnodes == 1 + ): + sglang_dist_init_addr = f"127.0.0.1:{server_args.port + ZMQ_TCP_PORT_DELTA}" + + sglang_port_args = None + if sglang_dist_init_addr is None: + sglang_port_args = PortArgs.init_new(server_args) + plugin_config = PluginConfig( # common config model_config=sgl_model_config, @@ -237,8 +261,9 @@ def _generate_atom_config_from_sglang_config(config: Any): sglang_enable_torch_compile=server_args.enable_torch_compile, sglang_disable_cuda_graph=server_args.disable_cuda_graph, sglang_enable_dp_attention=server_args.enable_dp_attention, - sglang_dist_init_addr=server_args.dist_init_addr, - sglang_port_args=PortArgs.init_new(server_args), + sglang_aiter_rank_id=sglang_aiter_rank_id, + sglang_dist_init_addr=sglang_dist_init_addr, + sglang_port_args=sglang_port_args, ) # force max num batched tokens to 16K because sgl doesn't have diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 1af8811e50..0af553e476 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -94,6 +94,8 @@ def init_aiter_dist(config: Config) -> None: ) rank = config.plugin_config.rank + if getattr(config.plugin_config, "is_sglang", False): + rank = getattr(config.plugin_config, "sglang_aiter_rank_id", rank) tensor_parallel_size = config.tensor_parallel_size assert ( diff --git a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py b/atom/plugin/sglang/attention_backend/sgl_attention_mla.py index c60b12f313..1dae9349e3 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py +++ b/atom/plugin/sglang/attention_backend/sgl_attention_mla.py @@ -601,6 +601,16 @@ def forward_sgl_plugin_mode_mla( **model_kwargs, ) -> torch.Tensor: prepared = forward_sgl_prepare(attn, positions, hidden_states, **model_kwargs) + from atom.utils.forward_context import get_forward_context + + if get_forward_context().context.is_dummy_run: + base_hidden_states = ( + hidden_states[0] if isinstance(hidden_states, tuple) else hidden_states + ) + dummy_output = base_hidden_states.new_empty( + (base_hidden_states.shape[0], base_hidden_states.shape[-1]) + ) + return dummy_output return forward_sgl_core(attn, prepared) diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index afeaf6b5da..db75a67433 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -2,8 +2,12 @@ Registers model architecture classes via SGLANG_EXTERNAL_MODEL_PACKAGE, replacing sglang's built-in implementations with ATOM-optimized versions. + +To add a new model, append its architecture class name to _MODEL_NAMES. """ +import copy + import logging from contextlib import contextmanager from contextvars import ContextVar @@ -30,6 +34,187 @@ def get_current_forward_batch(): return _current_forward_batch.get() +def _is_dummy_forward(forward_batch: ForwardBatch) -> bool: + # SGLang's IDLE batch is the plugin-side equivalent of ATOM dummy run. + forward_mode = getattr(forward_batch, "forward_mode", None) + return bool( + forward_mode is not None + and hasattr(forward_mode, "is_idle") + and forward_mode.is_idle() + ) + + +def _pad_dummy_like( + tensor: Optional[torch.Tensor], + *, + length: int, + fill_value: int | float = 0, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + shape = (length, *tensor.shape[1:]) + return torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device) + + +def _materialize_atom_dummy_forward( + input_ids: Optional[torch.Tensor], + positions: Optional[torch.Tensor], + input_embeds: Optional[torch.Tensor], + forward_batch: ForwardBatch, +) -> tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + ForwardBatch, +]: + """Convert an empty SGLang IDLE batch into ATOM-style dummy forward inputs.""" + dummy_positions = positions.new_zeros((1,)) + dummy_input_ids = input_ids.new_zeros((1,)) + dummy_input_embeds = _pad_dummy_like(input_embeds, length=1, fill_value=0) + + model_forward_batch = copy.copy(forward_batch) + model_forward_batch.positions = dummy_positions + model_forward_batch.batch_size = 1 + model_forward_batch.seq_lens_sum = 1 + model_forward_batch.seq_lens = forward_batch.seq_lens.new_ones((1,)) + model_forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu.new_ones((1,)) + + return dummy_input_ids, dummy_positions, dummy_input_embeds, model_forward_batch + + +def _trim_hidden_states_for_output(hidden_states, num_tokens: int): + if torch.is_tensor(hidden_states): + return hidden_states[:num_tokens] + if isinstance(hidden_states, tuple): + return tuple( + tensor[:num_tokens] if torch.is_tensor(tensor) else tensor + for tensor in hidden_states + ) + return hidden_states + + +def _resolve_num_tokens_across_dp( + atom_config: Any, + forward_batch: ForwardBatch, + num_tokens: int, + is_dummy_run: bool, +) -> torch.Tensor: + """Resolve per-DP token counts for ATOM's CPU-side DPMetadata. + + Real SGLang dp-attention batches carry ``global_num_tokens_cpu`` from the + scheduler. That list is the source of truth for mixed prefill/decode/idle + batches, where token counts may look like [8, 1, 8, 8]. + + Some SGLang synthetic/static batches, especially CUDA graph capture batches, + only keep the global token buffer on GPU. ATOM's DPMetadata is CPU-side and + needs a CPU tensor before model forward, so avoid reading the GPU buffer back + to CPU. We only fallback when the batch advertises the same-shape DP buffer + layout (global_dp_buffer_len == local_num_tokens * dp_size), where the CPU + equivalent is exactly [local_num_tokens] * dp_size. + + IDLE batches are reported by SGLang as 0 tokens on the current rank, but + this wrapper materializes them as one local dummy token before entering + ATOM. Patch the current DP rank after resolving the distribution so + ``DPMetadata`` sees a local count that matches the actual ATOM input. + """ + global_num_tokens_cpu = getattr(forward_batch, "global_num_tokens_cpu", None) + if global_num_tokens_cpu is not None: + num_tokens_across_dp = torch.tensor( + global_num_tokens_cpu, dtype=torch.int32, device="cpu" + ) + else: + dp_size = atom_config.parallel_config.data_parallel_size + global_num_tokens_gpu = getattr(forward_batch, "global_num_tokens_gpu", None) + global_dp_buffer_len = getattr(forward_batch, "global_dp_buffer_len", None) + is_static_same_shape_batch = ( + global_num_tokens_gpu is not None + and global_dp_buffer_len == num_tokens * dp_size + ) + if not is_static_same_shape_batch: + raise RuntimeError( + "[SGL+ATOM] SGLang dp-attention requires " + "forward_batch.global_num_tokens_cpu unless the batch uses static " + "same-shape DP metadata." + ) + + # Static batches, such as CUDA graph capture batches, may only keep + # global token counts on GPU. Avoid GPU-to-CPU reads here and mirror + # their same-shape layout directly for ATOM's CPU DPMetadata. + num_tokens_across_dp = torch.full( + (dp_size,), num_tokens, dtype=torch.int32, device="cpu" + ) + + if is_dummy_run: + # SGLang reports idle ranks as 0 tokens, but ATOM materializes them + # as one local dummy token so collectives and DPMetadata stay aligned. + dp_rank = atom_config.parallel_config.data_parallel_rank + num_tokens_across_dp[dp_rank] = num_tokens + return num_tokens_across_dp + + +def _set_sglang_forward_context( + atom_config: Any, + forward_batch: ForwardBatch, + positions: torch.Tensor, +) -> None: + """Bridge SGLang batch metadata into ATOM's global forward context.""" + from atom.utils.forward_context import ( + AttentionMetaData, + Context, + set_forward_context, + ) + + forward_mode = forward_batch.forward_mode + # TODO: This max_seqlen_q is not the source of truth for prefill attention; + # SGLang plugin attention consumes forward_batch.attn_backend.forward_metadata + # directly. In this wrapper it is only needed by ATOM MoE padding: under + # dp-attention + TP (non-EP all_gather/reduce_scatter), decode/idle batches + # must use 1 so pad_for_all_gather keeps fixed-shape collectives aligned. + # Leaving it as 0 there can make active and dummy ranks send different + # shapes to DP all_gather and hang. + max_seqlen_q = 1 if forward_mode.is_decode_or_idle() else 0 + attn_metadata = AttentionMetaData(max_seqlen_q=max_seqlen_q) + batch_size = int(forward_batch.batch_size) + is_dummy_run = _is_dummy_forward(forward_batch) + is_prefill = forward_mode.is_prefill() + num_tokens = int(positions.shape[0]) + + enable_dp_attention = bool(atom_config.enable_dp_attention) + if enable_dp_attention: + # SGLang owns the cross-DP token distribution under dp-attention; ATOM + # uses it to derive graph_bs and fixed-size MoE gather/scatter buffers. + num_tokens_across_dp = _resolve_num_tokens_across_dp( + atom_config, forward_batch, num_tokens, is_dummy_run + ) + graph_bs = int(torch.max(num_tokens_across_dp).item()) + else: + # Without dp-attention, ATOM runs with local-rank shapes only. There is + # no cross-DP token distribution to pass into DPMetadata, so graph_bs + # follows the local prefill token count or decode batch size. + num_tokens_across_dp = None + graph_bs = num_tokens if is_prefill else batch_size + context = Context( + positions=positions, + is_prefill=is_prefill, + is_dummy_run=is_dummy_run, + batch_size=batch_size, + graph_bs=graph_bs, + ) + set_forward_context( + attn_metadata=attn_metadata, + atom_config=atom_config, + context=context, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ) + + +def _reset_sglang_forward_context() -> None: + from atom.utils.forward_context import reset_forward_context + + reset_forward_context() + + @dataclass(frozen=True) class SGLangForwardBatchMetadata: """Small context object for one SGLang model forward.""" @@ -137,13 +322,24 @@ def __init__( import atom + # TODO: prepare_model() currently handles model construction, config + # generation, attention backend registration, and distributed init. + # Refactor so this wrapper only dispatches the attention backend + # (register_ops_to_sglang + set_attn_cls), and let sglang handle + # model construction directly self.model = atom.prepare_model(config=config, engine="sglang") if self.model is None: raise ValueError( f"ATOM failed to create model for architecture {self.model_arch}" ) - self.logits_processor = LogitsProcessor(config) + # Under SGLang dp-attention, ATOM runtime interprets non-MoE modules + # like lm_head with tp=1 semantics, so plugin logits must not perform + # an extra TP all-gather after local lm_head matmul. + plugin_skip_all_gather = bool(self.model.atom_config.enable_dp_attention) + self.logits_processor = LogitsProcessor( + config, skip_all_gather=plugin_skip_all_gather + ) # Apply ds model-specific sglang patches (attn dispatch, weight hooks, etc.) # TODO: will remove this after sglang supports atom attention backend @@ -170,13 +366,38 @@ def forward( pp_proxy_tensors=pp_proxy_tensors, save_kv_cache=model_kwargs.get("save_kv_cache"), ) + + if _is_dummy_forward(forward_batch): + ( + model_input_ids, + model_positions, + model_input_embeds, + model_forward_batch, + ) = _materialize_atom_dummy_forward( + input_ids, + positions, + input_embeds, + forward_batch, + ) + else: + ( + model_input_ids, + model_positions, + model_input_embeds, + model_forward_batch, + ) = ( + input_ids, + positions, + input_embeds, + forward_batch, + ) model_inputs = dict( - input_ids=input_ids, - positions=positions, + input_ids=model_input_ids, + positions=model_positions, intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( pp_proxy_tensors, metadata ), - inputs_embeds=input_embeds, + inputs_embeds=model_input_embeds, ) with SGLangForwardBatchMetadata.bind(metadata): if self.model_arch_spec.wrapper_binds_gdn_context: @@ -187,9 +408,21 @@ def forward( with SGLangGDNForwardContext.bind(metadata): hidden_states = self.model(**model_inputs) else: - hidden_states = self.model(**model_inputs) + try: + _set_sglang_forward_context( + self.model.atom_config, model_forward_batch, model_positions + ) + hidden_states = self.model(**model_inputs) + finally: + _reset_sglang_forward_context() if self.pp_group.is_last_rank: + if _is_dummy_forward(forward_batch): + # TODO: Revisit if SGLang ever sends non-empty dummy batches. + # Today this path only runs when an empty IDLE batch is expanded + # to one ATOM dummy token, so the output boundary must trim back to + # the original SGLang-visible length: 0 tokens. + hidden_states = _trim_hidden_states_for_output(hidden_states, 0) return self.logits_processor( input_ids, hidden_states, From 90b292ae458a0c759024f16b071bcfbfad0f2cdb Mon Sep 17 00:00:00 2001 From: Yutao Xu Date: Thu, 14 May 2026 17:24:02 +0800 Subject: [PATCH 48/76] Support fused qknorm+allreduce for minimax m2 (#774) * Update minimax_m2.py * reorder * more general * fix lint --- atom/models/minimax_m2.py | 52 +++++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 19 deletions(-) diff --git a/atom/models/minimax_m2.py b/atom/models/minimax_m2.py index 7d960edd0d..d254405648 100644 --- a/atom/models/minimax_m2.py +++ b/atom/models/minimax_m2.py @@ -1,7 +1,10 @@ from typing import Optional, Union import torch -from aiter.dist.communication_op import tensor_model_parallel_all_reduce +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, + tensor_model_parallel_fused_qknorm_allreduce, +) from aiter.dist.parallel_state import ( get_pp_group, get_tensor_model_parallel_rank, @@ -224,31 +227,42 @@ def forward( hidden_states: torch.Tensor, ) -> torch.Tensor: qkv = self.qkv_proj(hidden_states) - q, k, v = torch.split(qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: # TP-aware RMSNorm: all-reduce variance across TP ranks so # normalization uses the global variance (over 6144/1024 dims) # rather than per-rank variance (768/128 dims). - orig_dtype = q.dtype - q = q.to(torch.float32) - k = k.to(torch.float32) - q_var = q.pow(2).mean(dim=-1, keepdim=True) - k_var = k.pow(2).mean(dim=-1, keepdim=True) - if self.tp_size > 1: - qk_var = torch.cat([q_var, k_var], dim=-1) - qk_var = tensor_model_parallel_all_reduce(qk_var) / self.tp_size - q_var, k_var = qk_var.chunk(2, dim=-1) - q = (q * torch.rsqrt(q_var + self.rms_norm_eps) * self.q_norm.weight).to( - orig_dtype - ) - k = (k * torch.rsqrt(k_var + self.rms_norm_eps) * self.k_norm.weight).to( - orig_dtype + if qkv.shape[0] <= 64 and self.tp_size > 1: + q, k, v = tensor_model_parallel_fused_qknorm_allreduce( + qkv, self.q_norm.weight, self.k_norm.weight, self.rms_norm_eps + ) + else: + q, k, v = torch.split( + qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1 + ) + orig_dtype = q.dtype + q = q.to(torch.float32) + k = k.to(torch.float32) + q_var = q.pow(2).mean(dim=-1, keepdim=True) + k_var = k.pow(2).mean(dim=-1, keepdim=True) + if self.tp_size > 1: + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / self.tp_size + q_var, k_var = qk_var.chunk(2, dim=-1) + q = ( + q * torch.rsqrt(q_var + self.rms_norm_eps) * self.q_norm.weight + ).to(orig_dtype) + k = ( + k * torch.rsqrt(k_var + self.rms_norm_eps) * self.k_norm.weight + ).to(orig_dtype) + else: + q, k, v = torch.split( + qkv, [self.q_size, self.kv_size, self.kv_size], dim=-1 ) - attn_output = self.attn( - query=q, key=k, value=v, positions=positions, q_scale=None, qkv=qkv - ) + attn_output = self.attn( + query=q, key=k, value=v, positions=positions, q_scale=None, qkv=qkv + ) output = self.o_proj(attn_output) return output From e77a5ce1afdf71bd4d3acb6324709b9ba21c3e80 Mon Sep 17 00:00:00 2001 From: Jiayun Date: Thu, 14 May 2026 22:16:42 +0800 Subject: [PATCH 49/76] fix prefix default on (#784) * fix all tokens hit issue * refine --- atom/model_engine/block_manager.py | 19 +++++++++++++++++++ atom/model_ops/attention_mla.py | 6 +----- tests/test_block_manager.py | 12 +++++++----- 3 files changed, 27 insertions(+), 10 deletions(-) diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index d62a2c413c..7e081a7253 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -118,6 +118,15 @@ def can_allocate(self, seq: Sequence) -> bool: block_id = self.hash_to_block_id.get(h, -1) if block_id == -1 or self.blocks[block_id].token_ids != token_ids: cache_miss = True + # If the entire prompt would be cached, force the last full block + # to recompute so prefill has at least one token to forward and + # produce logits for the next-token sampler. + if ( + not cache_miss + and i == seq.num_blocks - 1 + and len(token_ids) == self.block_size + ): + cache_miss = True if cache_miss: needed_free += 1 return ( @@ -142,6 +151,16 @@ def allocate(self, seq: Sequence): ) if block_id == -1 or self.blocks[block_id].token_ids != token_ids: cache_miss = True + # If the entire prompt would be cached, force the last full block + # to recompute so prefill has at least one token to forward and + # produce logits for the next-token sampler. Must mirror the same + # condition in can_allocate() so the block budget agrees. + if ( + not cache_miss + and i == seq.num_blocks - 1 + and len(token_ids) == self.block_size + ): + cache_miss = True if cache_miss: block_id = self._pop_free_block() block = self._allocate_block(block_id) diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 9e2faa889c..c92c231c0d 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -698,10 +698,6 @@ def forward_impl_server_mode( kv_cache = kv_cache_data[f"layer_{self.layer_num}"].k_cache if context.is_prefill and not use_prefill_mla: - use_prefix_cache = ( - attn_metadata.has_cached and self.kv_b_proj.weight.dtype != dtypes.fp4x2 - ) - prefill_q = self.q_proj(q, x_scale=q_scale).view( -1, self.num_heads, self.qk_head_dim ) @@ -718,7 +714,7 @@ def forward_impl_server_mode( scale=self._k_scale, ) - if use_prefix_cache: + if attn_metadata.has_cached: # k_full/v_full are used for attention compute; gather_kv_b_proj reads # fp8 from cache and dequantizes internally, so output must be model dtype k_full = torch.empty( diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index 60a6c10612..9ad059ff2d 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -301,10 +301,11 @@ def test_preempt_and_reschedule_reuses_cache(self, seq_factory): assert s1.num_cached_tokens == 0 assert s1.block_table == [] - # Re-allocate — should get cache hits on both blocks + # Re-allocate — first block is a cache hit; the last full block is + # force-recomputed so prefill has at least one token to forward. s1_retry = seq_factory([1, 2, 3, 4, 5, 6, 7, 8]) bm.allocate(s1_retry) - assert s1_retry.num_cached_tokens == 8 # both blocks cached + assert s1_retry.num_cached_tokens == 4 # ── Edge cases ─────────────────────────────────────────────────────────── @@ -325,8 +326,9 @@ def test_single_token_no_cache(self, seq_factory): # Partial block → hash is -1 → no caching assert s2.num_cached_tokens == 0 - def test_exact_block_size_fully_cached(self, seq_factory): - """Sequence with exactly block_size tokens — fully cached on reuse.""" + def test_exact_block_size_last_block_recomputed(self, seq_factory): + """Single-block prompt: last full block is force-recomputed on reuse so + prefill has at least one token to forward and produce logits.""" cfg = MockConfig( num_kvcache_blocks=4, kv_cache_block_size=4, enable_prefix_caching=True ) @@ -336,7 +338,7 @@ def test_exact_block_size_fully_cached(self, seq_factory): bm.deallocate(s1) s2 = seq_factory([1, 2, 3, 4]) bm.allocate(s2) - assert s2.num_cached_tokens == 4 + assert s2.num_cached_tokens == 0 def test_free_block_ids_set_consistent(self, block_manager, seq_factory): """free_block_ids_set stays consistent through allocate/deallocate.""" From c0c6bee20cdf2f14d4fa3ad98980b2b6e0cfd0f7 Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Fri, 15 May 2026 09:53:43 +0800 Subject: [PATCH 50/76] [Perf][vLLM-ATOM] Optimize Sparse MLA in vLLM-ATOM (#765) * preshuffle indexer cache Signed-off-by: kliuae * use persistent mode mla Signed-off-by: kliuae * Format sparse MLA plugin files * Remove --block-size 1 from DeepSeek-V3.2 benchmark and accuracy configs The sparse MLA plugin now defaults to block_size=64 for preshuffled indexer cache. The hardcoded --block-size 1 in CI configs would override this default and prevent the performance gains from taking effect. Co-Authored-By: Claude Opus 4 --------- Signed-off-by: kliuae Co-authored-by: XiaobingZhang Co-authored-by: zejunchen-zejun Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 4 +- .github/benchmark/oot_models_accuracy.json | 2 +- atom/plugin/attention.py | 126 +++++++++- atom/plugin/attention_mla_sparse.py | 250 +++++++++++++++++--- recipes/atom_vllm/GLM-5.md | 3 +- 5 files changed, 332 insertions(+), 53 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 1b54c6265e..f006161feb 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -284,7 +284,7 @@ "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp4", "prefix": "deepseek-v3-2-fp8-aw-tp4", "bench_args": "", - "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --kv-cache-dtype auto --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" }, { @@ -293,7 +293,7 @@ "dashboard_model": "DeepSeek-V3.2-FP8-aw-tp8", "prefix": "deepseek-v3-2-fp8-aw-tp8", "bench_args": "", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --block-size 1 --max-num-batched-tokens 16384 --max-model-len 16384", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --max-num-batched-tokens 16384 --max-model-len 16384", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" } ] diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index 1050e3f856..edcc58f78c 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -139,7 +139,7 @@ { "model_name": "DeepSeek-V3.2-FP8 TP8", "model_path": "deepseek-ai/DeepSeek-V3.2", - "extraArgs": "--tensor-parallel-size 8 --block-size 1", + "extraArgs": "--tensor-parallel-size 8", "env_vars": "", "runner": "linux-atom-mi35x-8", "test_level": "nightly", diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index 9c674c1cd4..d38da4a142 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -1744,11 +1744,20 @@ class AiterMLASparseMetadataForPluginMode: paged_kv_last_page_len: torch.Tensor paged_kv_indices: torch.Tensor paged_kv_indptr: torch.Tensor - paged_kv_indptr_rest: torch.Tensor + attn_out_dtype: torch.dtype block_size: int = 1 topk_tokens: int = 2048 + # Persistent MLA metadata (only populated when persistent mode is enabled, + # i.e. when the aiter sparse decode kernel supports work-stealing splits). + work_meta_data: torch.Tensor | None = None + work_indptr: torch.Tensor | None = None + work_info_set: torch.Tensor | None = None + reduce_indptr: torch.Tensor | None = None + reduce_final_map: torch.Tensor | None = None + reduce_partial_map: torch.Tensor | None = None + class vllmMLASparseAttentionMetadataBuilderMethods: def __init__(self): @@ -1766,18 +1775,61 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): ) # Zero-fill for cudagraphs self.req_id_per_token_buffer.fill_(0) + self.paged_kv_indices.fill_(0) + self.paged_kv_indptr.fill_(0) self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( req_id_per_token, non_blocking=True ) - self.paged_kv_indices.fill_(0) - self.paged_kv_indptr.fill_(0) + query_lens = ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + seq_lens = common_attn_metadata.seq_lens + + from atom.plugin.attention_mla_sparse import generate_sparse_seqlen_triton + + sparse_seqlen = generate_sparse_seqlen_triton( + query_lens, + seq_lens, + common_attn_metadata.query_start_loc, + self.topk_tokens, + num_tokens, + common_attn_metadata.max_query_len, + ) + + torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1]) + self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens]) req_id_per_token = self.req_id_per_token_buffer[:num_tokens] qo_indptr = self.qo_indptr[: num_tokens + 1] paged_kv_last_page_len = self.paged_kv_last_page_len[:num_tokens] paged_kv_indices = self.paged_kv_indices[: num_tokens * self.topk_tokens] paged_kv_indptr = self.paged_kv_indptr[: num_tokens + 1] - paged_kv_indptr_rest = self.paged_kv_indptr[num_tokens + 1 :] + + # ----- Compute persistent MLA metadata ----- + # The aiter sparse decode kernel uses qseqlen=1 (each query token is + # treated as its own batch entry), so persistent metadata can always + # be precomputed here. The kernel switches to the persistent + # work-stealing path automatically when work_meta_data is non-None. + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self.padded_num_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=True, + ) attn_metadata_for_plugin_mode = AiterMLASparseMetadataForPluginMode( num_reqs=common_attn_metadata.num_reqs, @@ -1790,12 +1842,18 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): block_table=common_attn_metadata.block_table_tensor, req_id_per_token=req_id_per_token, block_size=self.kv_cache_spec.block_size, + attn_out_dtype=self.model_dtype, topk_tokens=self.topk_tokens, qo_indptr=qo_indptr, paged_kv_last_page_len=paged_kv_last_page_len, paged_kv_indices=paged_kv_indices, paged_kv_indptr=paged_kv_indptr, - paged_kv_indptr_rest=paged_kv_indptr_rest, + work_meta_data=self._mla_work_meta_data, + work_indptr=self._mla_work_indptr, + work_info_set=self._mla_work_info_set, + reduce_indptr=self._mla_reduce_indptr, + reduce_final_map=self._mla_reduce_final_map, + reduce_partial_map=self._mla_reduce_partial_map, ) attn_metadata = AttentionMetaData( @@ -2063,14 +2121,12 @@ def __init__(self): @staticmethod def get_supported_kernel_block_sizes(): - return [1] + return [1, 64] @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: - # ATOM's sparse MLA plugin path assumes kernel/page block size == 1 - # throughout metadata construction and KV-cache indexing, so force that - # value when vLLM 0.19 negotiates the backend-specific cache block size. - return 1 + # Prefer block_size == 64 so the indexer's preshuffled path is taken + return 64 @staticmethod def get_kv_cache_shape( @@ -2231,6 +2287,7 @@ def init_method_under_plugin_mode( self.vllm_config = config self.model_config = config.model_config + self.model_dtype = self.model_config.dtype self.kv_cache_spec = kv_cache_spec self.device = device max_num_batched_tokens = config.scheduler_config.max_num_batched_tokens @@ -2241,9 +2298,6 @@ def init_method_under_plugin_mode( self.padded_num_heads = max(self.num_heads, _MLA_MIN_HEADS) self.mla_dims = get_mla_dims(self.model_config) self.topk_tokens = config.model_config.hf_config.index_topk - self.topk_tokens_tensor = torch.tensor( - [self.topk_tokens], device=device, dtype=torch.int32 - ) self.max_model_len_tensor = torch.tensor( [self.model_config.max_model_len], device=device, dtype=torch.int32 ) @@ -2275,6 +2329,52 @@ def init_method_under_plugin_mode( [max_num_batched_tokens + 1], dtype=torch.int32, device=device ) + # ----- Persistent MLA metadata buffers ----- + # The aiter sparse decode kernel supports a "persistent" path that + # uses precomputed work-splitting metadata for better load balancing + # across CUs. Mirrors the approach used in rocm_aiter_mla.py. + # + # In the sparse case each query token is its own "batch" entry in the + # qo_indptr (qo_indptr = [0, 1, 2, ..., num_tokens]) and max_qo_len=1. + # We pad get_mla_metadata_info_v1's batch_size to max_num_batched_tokens + # so the buffers are large enough for any decode shape we might see. + ( + (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( + max_num_batched_tokens, + 1, + self.padded_num_heads, + torch.bfloat16, + dtypes.d_dtypes[config.cache_config.cache_dtype], + is_sparse=True, + fast_mode=True, + ) + self._mla_work_meta_data = torch.empty( + work_meta_data_size, dtype=work_meta_data_type, device=device + ) + self._mla_work_indptr = torch.empty( + work_indptr_size, dtype=work_indptr_type, device=device + ) + self._mla_work_info_set = torch.empty( + work_info_set_size, dtype=work_info_set_type, device=device + ) + self._mla_reduce_indptr = torch.empty( + reduce_indptr_size, dtype=reduce_indptr_type, device=device + ) + self._mla_reduce_final_map = torch.empty( + reduce_final_map_size, dtype=reduce_final_map_type, device=device + ) + self._mla_reduce_partial_map = torch.empty( + reduce_partial_map_size, + dtype=reduce_partial_map_type, + device=device, + ) + return init_method_under_plugin_mode diff --git a/atom/plugin/attention_mla_sparse.py b/atom/plugin/attention_mla_sparse.py index 6a3e01635c..2e48acf3d8 100644 --- a/atom/plugin/attention_mla_sparse.py +++ b/atom/plugin/attention_mla_sparse.py @@ -46,6 +46,188 @@ logger = logging.getLogger("atom") +@triton.jit +def _convert_req_index_to_global_index_kernel( + req_id_ptr, # int32 [num_tokens] + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + cu_seqlens_ptr, # int32 [num_tokens + 1] + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + # shapes (compile-time where possible) + max_num_blocks_per_req: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # tile width along columns + # strides (in elements) + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, +): + # program_id(0) -> token_id (row) + # program_id(1) -> tile index along columns + token_id = tl.program_id(0) + tile_id = tl.program_id(1) + + # Each program covers BLOCK_N consecutive columns + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + # Load request id for this token (no mask: grid is exact) + req = tl.load(req_id_ptr + token_id) + + # Load cumulative sequence lengths to get starting index of this request + seq_start = tl.load(cu_seqlens_ptr + token_id) + seq_end = tl.load(cu_seqlens_ptr + token_id + 1) + + if tile_id * BLOCK_N + seq_start >= seq_end: + return + + # Load token indices for this tile + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 + tok = tl.load(ti_ptr) # int32 + + # Only token == -1 should propagate as -1 + is_invalid_tok = tok < 0 + + # Compute block id and in-block offset + block_id = tok // BLOCK_SIZE + inblock_off = tok % BLOCK_SIZE + + # Guard block_table access + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 + base = tl.load(bt_ptr, mask=valid_block, other=0) + + # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset + out_val = tl.where( + is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off + ) + out_ptr_ij = out_ptr + seq_start + indice_id + out_ptr_ij_mask = (seq_start + indice_id) < seq_end + + # store the results with mask + tl.store(out_ptr_ij, out_val, mask=out_ptr_ij_mask) + + +def triton_convert_req_index_to_global_index( + req_id: torch.Tensor, # int32 [num_tokens] + block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] + token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] + cu_seqlens: torch.Tensor, # int32 [num_tokens + 1] + paged_kv_indices: torch.Tensor, # int32 [num_tokens * topk] out_buffer + BLOCK_SIZE: int = 64, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 128, # tile width along columns +): + """ + out[token_id, indice_id] = + block_table[req_id[token_id], + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE + + token_indices[token_id, indice_id] % BLOCK_SIZE + + Only when token_indices[token_id, indice_id] == -1 do we output -1. + For safety, we also output -1 if the derived block_id would be + out-of-bounds. + """ + assert req_id.dtype == torch.int32 + assert block_table.dtype == torch.int32 + assert token_indices.dtype == torch.int32 + assert token_indices.shape[1] == NUM_TOPK_TOKENS + assert ( + NUM_TOPK_TOKENS % BLOCK_N == 0 + ), f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible byBLOCK_N ({BLOCK_N})" + # print("req_id: ", req_id, flush=True) + num_tokens = req_id.shape[0] + _, max_num_blocks_per_req = block_table.shape + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + # Ensure contiguous tensors on the same device + req_id_c = req_id.contiguous() + block_table_c = block_table.contiguous() + token_indices_c = token_indices.contiguous() + + # Strides in elements + bt_stride0, bt_stride1 = block_table_c.stride() + ti_stride0, ti_stride1 = token_indices_c.stride() + + # Exact 2D grid: tokens × column tiles + grid = (num_tokens, tiles_per_row) + + _convert_req_index_to_global_index_kernel[grid]( + req_id_c, + block_table_c, + token_indices_c, + cu_seqlens, + paged_kv_indices, + # shapes / constexprs + max_num_blocks_per_req, + BLOCK_SIZE, + BLOCK_N, + # strides + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + ) + return + + +@triton.jit +def generate_sparse_seqlen_kernel( + seq_len_ptr, # [num_seq] + cu_query_lens_ptr, # [num_seq] + out_ptr, # [num_query_tokens] + topk_token: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + seq_id = tl.program_id(0) + query_offset = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + query_start = tl.load(cu_query_lens_ptr + seq_id) + query_end = tl.load(cu_query_lens_ptr + seq_id + 1) + if query_start + tl.program_id(1) * BLOCK_SIZE > query_end: + return + query_len = query_end - query_start + query_mask = query_offset + query_start < query_end + seq_len = tl.load(seq_len_ptr + seq_id) + # Just return since the out_ptr is zero initialized. + if seq_len == 0: + return + context_start_point = seq_len - query_len + sparse_seqlen = context_start_point + query_offset + sparse_seqlen_masked = tl.where( + sparse_seqlen + 1 < topk_token, sparse_seqlen + 1, topk_token + ) + tl.store( + out_ptr + query_start + query_offset, sparse_seqlen_masked, mask=query_mask + ) + + +def generate_sparse_seqlen_triton( + query_lens: torch.Tensor, + seq_lens: torch.Tensor, + cu_query_lens: torch.Tensor, + topk_token: int, + num_tokens: int, + max_query_len: int, +): + num_seqs = query_lens.size(0) + # zero initialize the tensor to make sure invalid positions will be zero + out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device) + block_size = 64 + num_block_per_row = triton.cdiv(max_query_len, block_size) + grid = ( + num_seqs, + num_block_per_row, + ) + generate_sparse_seqlen_kernel[grid]( + seq_lens, + cu_query_lens, + out, + topk_token, + block_size, + ) + return out + + @triton.jit def fetch_id_to_ragged_kernel( in_tensor_ptr, # [num_seq, topk] @@ -101,11 +283,10 @@ def __init__(self): "It is only used as a method container for the decorator." ) - def _forward_sparse_bf16_kv( + def _forward_sparse_mla( self, q: torch.Tensor, # [sq, heads, d_qk] kv_cache: torch.Tensor, # [blocks, heads, d_qk] - topk_indices_global: torch.Tensor, # [sq, topk] attn_metadata, layer, ) -> torch.Tensor: @@ -115,25 +296,12 @@ def _forward_sparse_bf16_kv( output = torch.empty( [num_tokens, self.padded_num_heads, self.kv_lora_rank], # dtype=q.dtype, - dtype=torch.bfloat16, + dtype=sparse_meta.attn_out_dtype, device=q.device, ) - seq_len = (topk_indices_global != -1).sum(dim=-1) - torch.cumsum(seq_len, dim=0, out=sparse_meta.paged_kv_indptr[1:]) - sparse_meta.paged_kv_indptr_rest.fill_(sparse_meta.paged_kv_indptr[-1]) - fetch_id_to_ragged_triton( - topk_indices_global, - sparse_meta.paged_kv_indptr, - sparse_meta.paged_kv_indices, - sparse_meta.topk_tokens, - ) - kv_buffer = kv_cache.unsqueeze(2) - # TODO: Currently, persistent mode has memory access issues when input context is long - # in fp8 kv cache settings, so it is not used for now. - # Re-enable persistent mode when the issue is fixed. mla_decode_fwd( q, kv_buffer.view(-1, 1, 1, q.shape[-1]), @@ -147,6 +315,12 @@ def _forward_sparse_bf16_kv( q_scale=layer._q_scale, kv_scale=layer._k_scale, page_size=1, + work_meta_data=sparse_meta.work_meta_data, + work_indptr=sparse_meta.work_indptr, + work_info_set=sparse_meta.work_info_set, + reduce_indptr=sparse_meta.reduce_indptr, + reduce_final_map=sparse_meta.reduce_final_map, + reduce_partial_map=sparse_meta.reduce_partial_map, ) if self.head_repeat_factor > 1: @@ -288,22 +462,15 @@ def forward_impl_sparse_plugin_mode( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] - try: - from vllm.v1.attention.backends.mla.sparse_utils import ( - triton_convert_req_index_to_global_index, - ) - except ImportError: - from vllm.v1.attention.backends.mla.flashmla_sparse import ( - triton_convert_req_index_to_global_index, - ) - req_id_i32 = sparse_meta.req_id_per_token.to(dtype=torch.int32) block_table_i32 = sparse_meta.block_table.to(dtype=torch.int32) topk_indices_i32 = topk_indices.to(dtype=torch.int32) - topk_indices_global = triton_convert_req_index_to_global_index( - req_id_i32, # sparse_meta.req_id_per_token, - block_table_i32, # sparse_meta.block_table, - topk_indices_i32, # topk_indices, + triton_convert_req_index_to_global_index( + req_id_i32, + block_table_i32, + topk_indices_i32, + sparse_meta.paged_kv_indptr, + sparse_meta.paged_kv_indices, BLOCK_SIZE=sparse_meta.block_size, NUM_TOPK_TOKENS=sparse_meta.topk_tokens, ) @@ -316,9 +483,7 @@ def forward_impl_sparse_plugin_mode( layer._q_scale, ) q_out = q_flat.reshape(q_out.shape) - attn_out = self._forward_sparse_bf16_kv( - q_out, kv_cache, topk_indices_global, attn_metadata, layer - ) + attn_out = self._forward_sparse_mla(q_out, kv_cache, attn_metadata, layer) # V up-projection self._v_up_proj(attn_out, out=output[:num_actual_toks]) @@ -355,13 +520,13 @@ def MLASparseAttentionImplDecoratorForPluginMode(cls): Decorator that injects sparse MLA methods into the MLAAttentionImpl class. Applied alongside the regular MLAAttentionImplDecoratorForPluginMode. - Injects forward_impl_sparse_plugin_mode and _forward_sparse_bf16_kv. + Injects forward_impl_sparse_plugin_mode and _forward_sparse_mla. The patched forward_impl in register.py calls forward_impl_plugin_mode, which dispatches to forward_impl_sparse_plugin_mode when topk_indices_buffer is set. """ sparse_method_names = [ - "_forward_sparse_bf16_kv", + "_forward_sparse_mla", "forward_impl_sparse_plugin_mode", ] @@ -434,6 +599,8 @@ def sparse_attn_indexer_plugin_mode( has_decode = indexer_meta.num_decodes > 0 has_prefill = indexer_meta.num_prefills > 0 num_decode_tokens = indexer_meta.num_decode_tokens + kv_block_size = kv_cache.shape[1] + preshuffle_cache = kv_block_size != 1 indexer_k_quant_and_cache( k, @@ -441,6 +608,7 @@ def sparse_attn_indexer_plugin_mode( slot_mapping, quant_block_size, scale_fmt, + preshuffle=preshuffle_cache, ) topk_indices_buffer[: hidden_states.shape[0]] = -1 @@ -465,6 +633,7 @@ def sparse_attn_indexer_plugin_mode( k_scale.view(dtypes.fp8), chunk.block_table, chunk.cu_seq_lens, + preshuffle=preshuffle_cache, ) logits = fp8_mqa_logits( @@ -530,6 +699,10 @@ def sparse_attn_indexer_plugin_mode( decode_metadata.seq_lens, decode_metadata.block_table, max_model_len, + ChunkK=256, + Preshuffle=preshuffle_cache, + KVBlockSize=kv_block_size, + WavePerEU=2, ) num_rows = logits.shape[0] @@ -616,8 +789,15 @@ def new_init(self, *args, **kwargs): def _deepseek_v32_indexer_get_kv_cache_spec(self, vllm_config): from vllm.v1.kv_cache_interface import MLAAttentionSpec + # Use the negotiated cache_config.block_size so the indexer cache and the + # main MLA cache share a single block_size. With a uniform block_size, + # vLLM groups the two layer types via `UniformTypeKVCacheSpecs.from_specs` + # and allocates a separate KVCacheTensor per layer sized to its own + # page_size_bytes (576B/token MLA vs 132B/token indexer); otherwise the + # smaller indexer page does not divide the MLA page and + # `unify_kv_cache_spec_page_size` raises NotImplementedError. return MLAAttentionSpec( - block_size=1, # block_size = 1 for indexer on ROCm + block_size=vllm_config.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, diff --git a/recipes/atom_vllm/GLM-5.md b/recipes/atom_vllm/GLM-5.md index 6ee61c6305..1a8720dd54 100644 --- a/recipes/atom_vllm/GLM-5.md +++ b/recipes/atom_vllm/GLM-5.md @@ -23,8 +23,7 @@ vllm serve zai-org/GLM-5-FP8 \ --gpu_memory_utilization 0.9 \ --async-scheduling \ --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ - --no-enable-prefix-caching \ - --block-size 1 + --no-enable-prefix-caching ``` ## Step 3: Performance Benchmark From f99324530e37560fc9f74bbc8318e0a8544ce330 Mon Sep 17 00:00:00 2001 From: Zhu Yuhua Date: Fri, 15 May 2026 11:29:47 +0800 Subject: [PATCH 51/76] (ci)[SGLang-ATOM]: Add Qwen3.5 cases for ci, nightly and benchmark (#777) * add qwen3.5 cases for sglang-atom Signed-off-by: zhuyuhua-v * remove redundant env Signed-off-by: zhuyuhua-v --------- Signed-off-by: zhuyuhua-v --- .../benchmark/sglang_benchmark_models.json | 47 ++++++-- .github/benchmark/sglang_models_accuracy.json | 48 ++++++++ .github/scripts/atom_sglang_test.sh | 25 +++-- .../atom-sglang-accuracy-validation.yaml | 60 ++++++++++ .github/workflows/atom-sglang-benchmark.yaml | 104 ++++++++++++++---- .github/workflows/atom-sglang-test.yaml | 9 ++ recipes/atom_sglang/Qwen3_5.md | 8 +- 7 files changed, 259 insertions(+), 42 deletions(-) diff --git a/.github/benchmark/sglang_benchmark_models.json b/.github/benchmark/sglang_benchmark_models.json index b06227844d..d293b64110 100644 --- a/.github/benchmark/sglang_benchmark_models.json +++ b/.github/benchmark/sglang_benchmark_models.json @@ -7,6 +7,7 @@ "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -18,6 +19,7 @@ "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -28,6 +30,7 @@ "extra_args": "--trust-remote-code --tensor-parallel-size 8", "bench_args": "", "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { @@ -39,17 +42,43 @@ "extra_args": "--trust-remote-code --tensor-parallel-size 4", "bench_args": "", "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" }, { - "display": "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8", - "dashboard_model": "DeepSeek-R1-0528-MXFP4-tp8-ep8", - "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", - "prefix": "deepseek-r1-fp4-tp8-ep8", - "extra_args": "--trust-remote-code --tensor-parallel-size 8 --expert-parallel-size 8", - "bench_args": "", - "runner": "atom-mi355-8gpu-aac-runner", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" + "display": "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8", + "dashboard_model": "DeepSeek-R1-0528-MXFP4-tp8-ep8", + "source_path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "path": "amd/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4", + "prefix": "deepseek-r1-fp4-tp8-ep8", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --expert-parallel-size 8", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "A", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nSGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1" + }, + { + "display": "Qwen3.5-397B-A17B-FP8 TP4", + "dashboard_model": "Qwen3.5-397B-A17B-FP8-tp4", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-tp4", + "extra_args": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "B", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0" + }, + { + "display": "Qwen3.5-397B-A17B-FP8 TP8", + "dashboard_model": "Qwen3.5-397B-A17B-FP8", + "source_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "path": "Qwen/Qwen3.5-397B-A17B-FP8", + "prefix": "qwen3-5-397b-a17b-fp8-tp8", + "extra_args": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "bench_args": "", + "runner": "atom-mi355-8gpu-aac-runner", + "nightly_group": "B", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0" } ] diff --git a/.github/benchmark/sglang_models_accuracy.json b/.github/benchmark/sglang_models_accuracy.json index 1fa97d41fc..e2e8548b7c 100644 --- a/.github/benchmark/sglang_models_accuracy.json +++ b/.github/benchmark/sglang_models_accuracy.json @@ -11,6 +11,54 @@ "accuracy_baseline_model": "deepseek-ai/DeepSeek-R1-0528", "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." }, + { + "model_name": "Qwen3.5-35B-A3B-FP8 TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B-FP8", + "extraArgs": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.76, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-35B-A3B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-35B-A3B TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B", + "extraArgs": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-35B-A3B", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-397B-A17B-FP8 TP4", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extraArgs": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-397B-A17B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, + { + "model_name": "Qwen3.5-397B-A17B-FP8 TP8", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extraArgs": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.83, + "accuracy_baseline": null, + "accuracy_baseline_model": "Qwen/Qwen3.5-397B-A17B-FP8", + "_baseline_note": "Threshold aligned with the SGLANG accuracy validation workflow target for gsm8k." + }, { "model_name": "DeepSeek-R1-FP8 TP8", "model_path": "deepseek-ai/DeepSeek-R1-0528", diff --git a/.github/scripts/atom_sglang_test.sh b/.github/scripts/atom_sglang_test.sh index 1940e4418f..1f1cbeb338 100644 --- a/.github/scripts/atom_sglang_test.sh +++ b/.github/scripts/atom_sglang_test.sh @@ -12,6 +12,7 @@ set -euo pipefail # Optional environment variables: # SGLANG_EXTRA_ARGS # SGLANG_ENV_VARS +# SGLANG_DEFAULT_SERVER_ARGS # SGLANG_PORT # SGLANG_HOST # MAX_WAIT_RETRIES @@ -146,11 +147,6 @@ launch_server() { local resolved_model_path resolved_model_path=$(resolve_model_path "${MODEL_PATH}") - local -a extra_arg_array=() - if [[ -n "${MODEL_EXTRA_ARGS}" ]]; then - read -r -a extra_arg_array <<< "${MODEL_EXTRA_ARGS}" - fi - prepare_runtime_paths export AITER_QUICK_REDUCE_QUANTIZATION="${AITER_QUICK_REDUCE_QUANTIZATION:-INT4}" @@ -168,6 +164,19 @@ launch_server() { done <<< "$(printf '%b' "${MODEL_ENV_VARS}")" fi + local default_server_args + default_server_args=${SGLANG_DEFAULT_SERVER_ARGS---trust-remote-code --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.8 --page-size 1 --disable-radix-cache} + + local -a default_arg_array=() + if [[ -n "${default_server_args}" ]]; then + read -r -a default_arg_array <<< "${default_server_args}" + fi + + local -a extra_arg_array=() + if [[ -n "${MODEL_EXTRA_ARGS}" ]]; then + read -r -a extra_arg_array <<< "${MODEL_EXTRA_ARGS}" + fi + rm -rf /root/.cache rm -f "${SGLANG_PID_FILE}" "${SGLANG_LOG_FILE}" || true @@ -182,11 +191,7 @@ launch_server() { --model-path "${resolved_model_path}" \ --host "${SGLANG_HOST}" \ --port "${SGLANG_PORT}" \ - --trust-remote-code \ - --kv-cache-dtype fp8_e4m3 \ - --mem-fraction-static 0.8 \ - --page-size 1 \ - --disable-radix-cache \ + "${default_arg_array[@]}" \ "${extra_arg_array[@]}" \ > "${SGLANG_LOG_FILE}" 2>&1 & diff --git a/.github/workflows/atom-sglang-accuracy-validation.yaml b/.github/workflows/atom-sglang-accuracy-validation.yaml index 2c9eaaa92b..d51a3cff1d 100644 --- a/.github/workflows/atom-sglang-accuracy-validation.yaml +++ b/.github/workflows/atom-sglang-accuracy-validation.yaml @@ -14,6 +14,26 @@ on: required: false type: boolean default: false + run_qwen35_35b_a3b_fp8_tp2: + description: "Qwen3.5-35B-A3B-FP8 TP2" + required: false + type: boolean + default: false + run_qwen35_35b_a3b_tp2: + description: "Qwen3.5-35B-A3B TP2" + required: false + type: boolean + default: false + run_qwen35_397b_a17b_fp8_tp4: + description: "Qwen3.5-397B-A17B-FP8 TP4" + required: false + type: boolean + default: false + run_qwen35_397b_a17b_fp8_tp8: + description: "Qwen3.5-397B-A17B-FP8 TP8" + required: false + type: boolean + default: false run_dsr1_fp8_tp8: description: "DeepSeek-R1-FP8 TP8" required: false @@ -70,6 +90,10 @@ jobs: id: meta env: RUN_DSR1_FP8_TP4: ${{ inputs.run_dsr1_fp8_tp4 }} + RUN_QWEN35_35B_A3B_FP8_TP2: ${{ inputs.run_qwen35_35b_a3b_fp8_tp2 }} + RUN_QWEN35_35B_A3B_TP2: ${{ inputs.run_qwen35_35b_a3b_tp2 }} + RUN_QWEN35_397B_A17B_FP8_TP4: ${{ inputs.run_qwen35_397b_a17b_fp8_tp4 }} + RUN_QWEN35_397B_A17B_FP8_TP8: ${{ inputs.run_qwen35_397b_a17b_fp8_tp8 }} RUN_DSR1_FP8_TP8: ${{ inputs.run_dsr1_fp8_tp8 }} RUN_DSR1_FP4_TP4: ${{ inputs.run_dsr1_fp4_tp4 }} RUN_DSR1_FP4_TP8: ${{ inputs.run_dsr1_fp4_tp8 }} @@ -94,6 +118,42 @@ jobs: "env_vars": "SGLANG_AITER_FP8_PREFILL_ATTN=0\nSGLANG_USE_AITER=1\nATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, + { + "toggle_env": "RUN_QWEN35_35B_A3B_FP8_TP2", + "model_name": "Qwen3.5-35B-A3B-FP8 TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B-FP8", + "extra_args": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.76, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_35B_A3B_TP2", + "model_name": "Qwen3.5-35B-A3B TP2", + "model_path": "Qwen/Qwen3.5-35B-A3B", + "extra_args": "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_397B_A17B_FP8_TP4", + "model_name": "Qwen3.5-397B-A17B-FP8 TP4", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extra_args": "--tensor-parallel-size 4 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN35_397B_A17B_FP8_TP8", + "model_name": "Qwen3.5-397B-A17B-FP8 TP8", + "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", + "extra_args": "--tensor-parallel-size 8 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache", + "accuracy_test_threshold": 0.83, + "env_vars": "SGLANG_DEFAULT_SERVER_ARGS=\nSGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0", + "runner": "linux-atom-mi35x-8", + }, { "toggle_env": "RUN_DSR1_FP8_TP8", "model_name": "DeepSeek-R1-FP8 TP8", diff --git a/.github/workflows/atom-sglang-benchmark.yaml b/.github/workflows/atom-sglang-benchmark.yaml index a925da28a6..84680999b9 100644 --- a/.github/workflows/atom-sglang-benchmark.yaml +++ b/.github/workflows/atom-sglang-benchmark.yaml @@ -6,8 +6,8 @@ concurrency: on: schedule: - # Nightly at 01:00 Beijing time (17:00 UTC on the previous day) - - cron: '0 17 * * *' + # Weekday nightly at 23:00 Beijing time (15:00 UTC) + - cron: '0 15 * * 1-5' workflow_dispatch: inputs: deepseek-r1-fp8-tp8: @@ -30,6 +30,14 @@ on: description: "DeepSeek-R1-0528-MXFP4 FP4 TP8 EP8" type: boolean default: false + qwen3-5-397b-a17b-fp8-tp4: + description: "Qwen3.5-397B-A17B-FP8 TP4" + type: boolean + default: false + qwen3-5-397b-a17b-fp8-tp8: + description: "Qwen3.5-397B-A17B-FP8 TP8" + type: boolean + default: false sglang_image: description: "Optional SGLang benchmark image override. Leave empty to use sglang-latest on main or rebuild from the selected non-main branch." type: string @@ -274,6 +282,7 @@ jobs: outputs: models_json: ${{ steps.load.outputs.models_json }} has_enabled_models: ${{ steps.load.outputs.has_enabled_models }} + selected_group: ${{ steps.load.outputs.selected_group }} steps: - uses: actions/checkout@v6 - id: load @@ -283,25 +292,80 @@ jobs: ENABLE_DEEPSEEK_R1_FP4_TP8: ${{ inputs.deepseek-r1-fp4-tp8 }} ENABLE_DEEPSEEK_R1_FP4_TP4: ${{ inputs.deepseek-r1-fp4-tp4 }} ENABLE_DEEPSEEK_R1_FP4_TP8_EP8: ${{ inputs.deepseek-r1-fp4-tp8-ep8 }} + ENABLE_QWEN3_5_397B_A17B_FP8_TP4: ${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }} + ENABLE_QWEN3_5_397B_A17B_FP8_TP8: ${{ inputs.qwen3-5-397b-a17b-fp8-tp8 }} run: | - MODELS_JSON="$(jq -c ' - map(select( - env.GITHUB_EVENT_NAME == "schedule" - or (.prefix == "deepseek-r1-fp8-tp8" and env.ENABLE_DEEPSEEK_R1_FP8_TP8 == "true") - or (.prefix == "deepseek-r1-fp8-tp4" and env.ENABLE_DEEPSEEK_R1_FP8_TP4 == "true") - or (.prefix == "deepseek-r1-fp4-tp8" and env.ENABLE_DEEPSEEK_R1_FP4_TP8 == "true") - or (.prefix == "deepseek-r1-fp4-tp4" and env.ENABLE_DEEPSEEK_R1_FP4_TP4 == "true") - or (.prefix == "deepseek-r1-fp4-tp8-ep8" and env.ENABLE_DEEPSEEK_R1_FP4_TP8_EP8 == "true") - )) - ' .github/benchmark/sglang_benchmark_models.json)" - echo "models_json=${MODELS_JSON}" >> "$GITHUB_OUTPUT" - if [ "${MODELS_JSON}" = "[]" ]; then - echo "has_enabled_models=false" >> "$GITHUB_OUTPUT" - echo "No models selected for SGLang benchmark." - else - echo "has_enabled_models=true" >> "$GITHUB_OUTPUT" - echo "Selected models: ${MODELS_JSON}" + set -euo pipefail + + export BEIJING_WEEKDAY="$(TZ=Asia/Shanghai date +%u)" + python3 - <<'PY' >> "$GITHUB_OUTPUT" + import json + import os + import sys + from pathlib import Path + + models = json.loads( + Path(".github/benchmark/sglang_benchmark_models.json").read_text( + encoding="utf-8" + ) + ) + + event = os.environ["GITHUB_EVENT_NAME"] + selected_group = "" + + if event == "schedule": + weekday = int(os.environ["BEIJING_WEEKDAY"]) + if weekday in (1, 3): + selected_group = "A-DEEPSEEK" + selected = [m for m in models if m.get("nightly_group", "A") == "A"] + elif weekday in (2, 4): + selected_group = "B-QWEN35" + selected = [m for m in models if m.get("nightly_group") == "B"] + elif weekday == 5: + selected_group = "C-ALL" + selected = list(models) + else: + selected_group = "SKIP-WEEKEND" + selected = [] + else: + enabled_by_prefix = { + "deepseek-r1-fp8-tp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP8_TP8", ""), + "deepseek-r1-fp8-tp4": os.environ.get("ENABLE_DEEPSEEK_R1_FP8_TP4", ""), + "deepseek-r1-fp4-tp8": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP8", ""), + "deepseek-r1-fp4-tp4": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP4", ""), + "deepseek-r1-fp4-tp8-ep8": os.environ.get("ENABLE_DEEPSEEK_R1_FP4_TP8_EP8", ""), + "qwen3-5-397b-a17b-fp8-tp4": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP4", ""), + "qwen3-5-397b-a17b-fp8-tp8": os.environ.get("ENABLE_QWEN3_5_397B_A17B_FP8_TP8", ""), + } + selected = [ + model + for model in models + if enabled_by_prefix.get(str(model.get("prefix")), "").lower() == "true" + ] + + if selected_group: + print(f"Scheduled SGLang benchmark group: {selected_group}", file=sys.stderr) + if not selected: + print("No models selected for SGLang benchmark.", file=sys.stderr) + else: + print("Selected SGLang benchmark models:", file=sys.stderr) + for model in selected: + print(f" - {model['display']} ({model['prefix']})", file=sys.stderr) + + print(f"models_json={json.dumps(selected, separators=(',', ':'))}") + print(f"selected_group={selected_group}") + print(f"has_enabled_models={'true' if selected else 'false'}") + PY + + - name: Print selected models + env: + SELECTED_GROUP: ${{ steps.load.outputs.selected_group }} + MODELS_JSON: ${{ steps.load.outputs.models_json }} + run: | + if [[ -n "${SELECTED_GROUP}" ]]; then + echo "Scheduled SGLang benchmark group: ${SELECTED_GROUP}" fi + printf 'Selected models: %s\n' "${MODELS_JSON}" build-benchmark-matrix: name: Build SGLang benchmark matrix @@ -547,6 +611,8 @@ jobs: deepseek-r1-fp4-tp8) echo "enabled=${{ inputs.deepseek-r1-fp4-tp8 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp4-tp4) echo "enabled=${{ inputs.deepseek-r1-fp4-tp4 }}" >> "$GITHUB_OUTPUT" ;; deepseek-r1-fp4-tp8-ep8) echo "enabled=${{ inputs.deepseek-r1-fp4-tp8-ep8 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-tp4) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp4 }}" >> "$GITHUB_OUTPUT" ;; + qwen3-5-397b-a17b-fp8-tp8) echo "enabled=${{ inputs.qwen3-5-397b-a17b-fp8-tp8 }}" >> "$GITHUB_OUTPUT" ;; *) echo "enabled=true" >> "$GITHUB_OUTPUT" ;; esac diff --git a/.github/workflows/atom-sglang-test.yaml b/.github/workflows/atom-sglang-test.yaml index 71e1b95688..2a5d8572fb 100644 --- a/.github/workflows/atom-sglang-test.yaml +++ b/.github/workflows/atom-sglang-test.yaml @@ -220,6 +220,15 @@ jobs: ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 accuracy_test_threshold: 0.91 runner: linux-atom-mi35x-4 + - model_name: "Qwen3.5-35B-A3B-FP8 TP2" + model_path: "Qwen/Qwen3.5-35B-A3B-FP8" + extra_args: "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache" + env_vars: | + SGLANG_DEFAULT_SERVER_ARGS= + SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models + ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0 + accuracy_test_threshold: 0.76 + runner: linux-atom-mi35x-4 runs-on: ${{ matrix.runner }} timeout-minutes: 180 env: diff --git a/recipes/atom_sglang/Qwen3_5.md b/recipes/atom_sglang/Qwen3_5.md index 6baf617816..f9c1d17c04 100644 --- a/recipes/atom_sglang/Qwen3_5.md +++ b/recipes/atom_sglang/Qwen3_5.md @@ -27,6 +27,8 @@ Users can use this command to launch the FP8 server with the same settings as th ```bash +export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0 + model_path=[Qwen/Qwen3.5-397B-A17B-FP8](https://huggingface.co/Qwen/Qwen3.5-397B-A17B-FP8) # [Qwen/Qwen3-Next-80B-A3B-Instruct-FP8](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct-FP8) tp=4 @@ -35,7 +37,6 @@ python3 -m sglang.launch_server \ --port 8000 \ --tensor-parallel-size ${tp} \ --mem-fraction-static 0.9 \ - --cuda-graph-max-bs 256 \ --reasoning-parser qwen3 \ --disable-radix-cache ``` @@ -86,10 +87,9 @@ Then append `--profile` to the `sglang.bench_serving` command in Step 3. ```bash lm_eval --model local-completions \ - --model_args model=${model_path},base_url=http://localhost:30000/v1/completions,num_concurrent=256,max_retries=2,tokenized_requests=False,trust_remote_code=True \ + --model_args model=${model_path},base_url=http://localhost:30000/v1/completions,num_concurrent=65,max_retries=1,tokenized_requests=False,trust_remote_code=True \ --tasks gsm8k \ - --batch_size auto \ - --num_fewshot 5 \ + --num_fewshot 3 \ --trust_remote_code ``` From df1a638b8c7cfd0a88d0ad0f8ded5085bb88b2b8 Mon Sep 17 00:00:00 2001 From: JiaoliangYu Date: Fri, 15 May 2026 11:55:44 +0800 Subject: [PATCH 52/76] fix: qwen_3.5 bf16 accurancy (#782) * fix(qwen3_5): packed_modules_mapping for in_proj_qkvzba BF16 weight has a fused in_proj_qkvzba; FP8/MXFP4 ship pre-split in_proj_qkvz and in_proj_ba. Add a helper that rewrites the packed_modules_mapping when quant_dtype is bfloat16 so the loader slices the fused tensor along the qkv/z/b/a axes correctly. Apply at the two atom-native conditional-generation entrypoints; the vllm plugin path keeps its existing inline fixup. * ci(accuracy): add Qwen3.5-397B-A17B BF16 nightly entry Adds the BF16 weight variant alongside the existing FP8 / MXFP4 entries. Use --kv_cache_dtype fp8 to avoid the known MI355 BF16 KV cache accuracy issue. Threshold/baseline are placeholders copied from the FP8 entry; refresh after the first CI measurement. --------- Co-authored-by: JiaoliangYu --- .github/benchmark/models_accuracy.json | 12 ++++++++++++ atom/models/qwen3_5.py | 25 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 92430ed674..5c1f3ac34d 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -217,6 +217,18 @@ "accuracy_baseline_model": "zai-org/GLM-5.1", "_baseline_note": "CI uses 3-shot, not comparable to HF 5-shot baseline" }, + { + "model_name": "Qwen3.5-397B-A17B", + "model_path": "Qwen/Qwen3.5-397B-A17B", + "extraArgs": "--kv_cache_dtype fp8 -tp 8", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "nightly", + "accuracy_threshold": 0.85, + "accuracy_baseline": 0.9538, + "accuracy_baseline_model": "Qwen3.5-397B-A17B-FP8", + "_baseline_note": "BF16 weight, TP8. Placeholder threshold/baseline copied from FP8 entry; refresh after first CI measurement." + }, { "model_name": "Qwen3.5-397B-A17B-FP8", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", diff --git a/atom/models/qwen3_5.py b/atom/models/qwen3_5.py index 7c3cccd35a..88282ac739 100644 --- a/atom/models/qwen3_5.py +++ b/atom/models/qwen3_5.py @@ -490,6 +490,25 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ) +_BF16_IN_PROJ_MAPPING = { + "in_proj_qkv": ("in_proj_qkvzba", (0, 1, 2)), + "in_proj_z": ("in_proj_qkvzba", 3), + "in_proj_b": ("in_proj_qkvzba", 4), + "in_proj_a": ("in_proj_qkvzba", 5), +} + + +def _apply_bf16_in_proj_mapping(mapping: dict, atom_config: Config) -> dict: + if atom_config.quant_config.global_quant_config.quant_dtype != torch.bfloat16: + return mapping + + mapping.pop("in_proj_qkvz", None) + mapping.pop("in_proj_ba", None) + mapping["in_proj_qkvzba"] = ("in_proj_qkvzba", None) + mapping.update(_BF16_IN_PROJ_MAPPING) + return mapping + + class Qwen3_5ForConditionalGenerationTextOnly(nn.Module): packed_modules_mapping = { "q_proj": ("qkv_proj", "q"), @@ -517,6 +536,9 @@ class Qwen3_5ForConditionalGenerationTextOnly(nn.Module): def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config + self.packed_modules_mapping = _apply_bf16_in_proj_mapping( + dict(self.packed_modules_mapping), atom_config + ) self.visual = PPMissingLayer() self.language_model = Qwen3_5ForCausalLM(atom_config=atom_config, prefix="") self.make_empty_intermediate_tensors = ( @@ -554,6 +576,9 @@ class Qwen3_5MoeForConditionalGenerationTextOnly( def __init__(self, atom_config: Config, prefix: str = ""): nn.Module.__init__(self) self.config = atom_config.hf_config + self.packed_modules_mapping = _apply_bf16_in_proj_mapping( + dict(self.packed_modules_mapping), atom_config + ) self.visual = PPMissingLayer() self.language_model = Qwen3_5MoeForCausalLM(atom_config=atom_config, prefix="") self.make_empty_intermediate_tensors = ( From cc3539e50847f1db420ab1a8903a1b8606c1d7ff Mon Sep 17 00:00:00 2001 From: Pleaplusone Date: Fri, 15 May 2026 15:32:24 +0800 Subject: [PATCH 53/76] Qwen3Next MTP for vLLM plugin mode (#772) * mtp 1 acc right Signed-off-by: ganyi * add recipe for qwen3-next-mtp Signed-off-by: ganyi * modify some qwen3.5 recipe Signed-off-by: ganyi * black Signed-off-by: ganyi * remove redundant code Signed-off-by: ganyi * remove redundant code Signed-off-by: ganyi * add spec decode convert for vllm plugin Signed-off-by: ganyi * remove vllm related branch Signed-off-by: ganyi * use atom spec decode config for plugin loading Signed-off-by: ganyi * remove unnecessary changes in modeling Signed-off-by: ganyi * format Signed-off-by: ganyi * add qwen3next mtp into benchmark Signed-off-by: ganyi * [ci] disable FP8 blockscale weight preshuffle for Qwen3.5/Qwen3-Next Add ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 to all Qwen3.5 and Qwen3-Next model configs across benchmark, nightly accuracy, and recipe files. Co-Authored-By: Claude Opus 4 * [ci] fix Qwen3-Next MTP benchmark label from MET to AW Co-Authored-By: Claude Opus 4 * [docs] fix Qwen3.5 recipe: update env var count and add preshuffle doc Remove stale "three" count (now variable list), add ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 to both the Important section and Key Environment Variables section. Co-Authored-By: Claude Opus 4 --------- Signed-off-by: ganyi Co-authored-by: zejunchen-zejun Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 34 ++++++++--- .github/benchmark/oot_models_accuracy.json | 26 ++++++--- .../atom-vllm-accuracy-validation.yaml | 27 +++++++-- .github/workflows/atom-vllm-benchmark.yaml | 16 +++++ atom/models/qwen3_next.py | 10 +++- atom/models/qwen3_next_mtp.py | 12 +++- atom/plugin/attention.py | 36 +++++++++++- atom/plugin/attention_mha.py | 58 +++++++++++++------ atom/plugin/config.py | 45 ++++++++++++++ .../vllm/attention_backend/attention_gdn.py | 9 ++- atom/plugin/vllm/model_wrapper.py | 32 +++++----- atom/plugin/vllm/register.py | 1 + recipes/atom_vllm/Qwen3.5.md | 9 ++- recipes/atom_vllm/Qwen3Next.md | 25 ++++++-- 14 files changed, 272 insertions(+), 68 deletions(-) diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index f006161feb..25d203a26e 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -181,7 +181,7 @@ "1024x8192" ], "extra_args": "--trust-remote-code --tensor-parallel-size 4 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" }, { "tp_size": 8, @@ -192,7 +192,7 @@ "1024x8192" ], "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" } ] }, @@ -213,7 +213,7 @@ "1024x8192" ], "extra_args": "--trust-remote-code --tensor-parallel-size 8 --attention-backend ROCM_AITER_FA --gpu-memory-utilization 0.8 --max-num-batched-tokens 16384 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" } ] }, @@ -231,7 +231,7 @@ "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp1-met", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" }, { "tp_size": 4, @@ -240,7 +240,7 @@ "prefix": "qwen3-next-80b-a3b-instruct-fp8-tp4-met", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" }, { "tp_size": 1, @@ -249,7 +249,7 @@ "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp1", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" }, { "tp_size": 2, @@ -258,7 +258,7 @@ "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp2", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 2 --max-num-batched-tokens 32768 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" }, { "tp_size": 4, @@ -267,7 +267,25 @@ "prefix": "qwen3-next-80b-a3b-instruct-fp8-aw-tp4", "bench_args": "", "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0" + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 1, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-mtp-tp1", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-mtp-tp1-aw", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 1 --max-num-batched-tokens 32768 --max-model-len 16384 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" + }, + { + "tp_size": 4, + "display": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW)", + "dashboard_model": "Qwen3-Next-80B-A3B-Instruct-FP8-mtp-tp4", + "prefix": "qwen3-next-80b-a3b-instruct-fp8-mtp-tp4-aw", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-model-len 16384 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0" } ] }, diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index edcc58f78c..fb68d48360 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -3,7 +3,7 @@ "model_name": "Qwen3-235B-A22B-Instruct-2507-FP8 TP8+EP8", "model_path": "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8", "extraArgs": "--tensor-parallel-size 8 --enable-expert-parallel", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.87, @@ -14,7 +14,7 @@ "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8 TP4", "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extraArgs": "--tensor-parallel-size 4 --attention-backend ROCM_AITER_FA", - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.76, @@ -25,7 +25,7 @@ "model_name": "Qwen3.5-397B-A17B-FP8 TP8", "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extraArgs": "--tensor-parallel-size 8 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.83, @@ -36,7 +36,7 @@ "model_name": "Qwen3.5-397B-A17B TP8", "model_path": "Qwen/Qwen3.5-397B-A17B", "extraArgs": "--tensor-parallel-size 8 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", "test_level": "nightly", "accuracy_threshold": 0.83, @@ -47,7 +47,7 @@ "model_name": "Qwen3.5-397B-A17B-MXFP4 TP4", "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extraArgs": "--tensor-parallel-size 4 --attention-backend ROCM_AITER_FA", - "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.82, @@ -55,6 +55,18 @@ "accuracy_baseline_model": "Qwen/Qwen3-235B-A22B-Instruct-2507", "_baseline_note": "Using Qwen3-235B baseline as proxy; needs CI measurement for Qwen3.5 specific baseline" }, + { + "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4", + "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "extraArgs": "--tensor-parallel-size 4 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", + "runner": "linux-atom-mi35x-4", + "test_level": "nightly", + "accuracy_threshold": 0.8, + "accuracy_baseline": 0.81, + "accuracy_baseline_model": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "_baseline_note": "Qwen3-Next-80B-A3B-Instruct-FP8 baseline with TP4 (no MTP) as proxy; needs CI measurement for MTP-specific baseline" + }, { "model_name": "Llama-3.1-8B-Instruct TP1", "model_path": "meta-llama/Llama-3.1-8B-Instruct", @@ -157,7 +169,7 @@ "runner": "linux-atom-mi35x-1", "test_level": "nightly", "accuracy_threshold": 0.88, - "accuracy_baseline": 0.90, + "accuracy_baseline": 0.9, "accuracy_baseline_model": "openai/gpt-oss-120b" }, { @@ -169,7 +181,7 @@ "runner": "linux-atom-mi35x-4", "test_level": "nightly", "accuracy_threshold": 0.88, - "accuracy_baseline": 0.90, + "accuracy_baseline": 0.9, "accuracy_baseline_model": "openai/gpt-oss-120b" }, { diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index 0b2d7e7a0b..dc448212e1 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -24,6 +24,11 @@ on: required: false type: boolean default: false + run_qwen3_next_80b_mtp_tp4: + description: "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4" + required: false + type: boolean + default: false run_qwen35_397b_fp8_tp8: description: "Qwen3.5-397B-A17B-FP8 TP8" required: false @@ -137,6 +142,7 @@ jobs: RUN_QWEN3_MOE_TP8: ${{ inputs.run_qwen3_moe_tp8 }} RUN_QWEN3_NEXT_80B_TP1: ${{ inputs.run_qwen3_next_80b_tp1 }} RUN_QWEN3_NEXT_80B_TP4: ${{ inputs.run_qwen3_next_80b_tp4 }} + RUN_QWEN3_NEXT_80B_MTP_TP4: ${{ inputs.run_qwen3_next_80b_mtp_tp4 }} RUN_QWEN35_397B_FP8_TP8: ${{ inputs.run_qwen35_397b_fp8_tp8 }} RUN_QWEN35_397B_TP8: ${{ inputs.run_qwen35_397b_tp8 }} RUN_QWEN35_397B_FP4_TP4: ${{ inputs.run_qwen35_397b_fp4_tp4 }} @@ -169,7 +175,7 @@ jobs: "model_path": "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8", "extra_args": "--tensor-parallel-size 8 --enable-expert-parallel", "accuracy_test_threshold": 0.87, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -178,7 +184,7 @@ jobs: "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extra_args": "--tensor-parallel-size 1", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-1", }, { @@ -187,7 +193,16 @@ jobs: "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", "extra_args": "--tensor-parallel-size 4", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_USE_FLYDSL_GDR=1\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", + "runner": "linux-atom-mi35x-4", + }, + { + "toggle_env": "RUN_QWEN3_NEXT_80B_MTP_TP4", + "model_name": "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4", + "model_path": "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8", + "extra_args": "--tensor-parallel-size 4 --speculative-config '{\"num_speculative_tokens\":1, \"method\": \"mtp\"}'", + "accuracy_test_threshold": 0.80, + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", }, { @@ -196,7 +211,7 @@ jobs: "model_path": "Qwen/Qwen3.5-397B-A17B-FP8", "extra_args": "--tensor-parallel-size 8", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -205,7 +220,7 @@ jobs: "model_path": "Qwen/Qwen3.5-397B-A17B", "extra_args": "--tensor-parallel-size 8", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-8", }, { @@ -214,7 +229,7 @@ jobs: "model_path": "amd/Qwen3.5-397B-A17B-MXFP4", "extra_args": "--tensor-parallel-size 4", "accuracy_test_threshold": 0.83, - "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0", + "env_vars": "ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1\nATOM_USE_CUSTOM_ALL_GATHER=0\nATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0", "runner": "linux-atom-mi35x-4", }, { diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index a12b914832..d7a178d517 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -36,6 +36,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -69,6 +71,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -102,6 +106,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -135,6 +141,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -168,6 +176,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -201,6 +211,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -234,6 +246,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) @@ -267,6 +281,8 @@ on: - Qwen3.5-397B-A17B TP8 (OOB) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW) + - Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW) - Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW) diff --git a/atom/models/qwen3_next.py b/atom/models/qwen3_next.py index f8abcb8672..40ad2380a3 100644 --- a/atom/models/qwen3_next.py +++ b/atom/models/qwen3_next.py @@ -485,7 +485,8 @@ def __init__( self.config = config self.quant_config = quant_config - self.speculative_config = speculative_config + + self.speculative_config = speculative_config or atom_config.speculative_config self.num_spec = ( self.speculative_config.num_speculative_tokens if self.speculative_config @@ -863,7 +864,6 @@ def forward( residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention - if self.input_layernorm.use_fused_quant: if residual is None: residual = hidden_states @@ -1059,6 +1059,11 @@ def __init__( if self.config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight + # Expose embed_tokens at this level for vLLM MTP embedding sharing. + # vLLM's proposer accesses target_wrapper.model.embed_tokens, where + # target_wrapper.model = this class (Qwen3NextForCausalLM). + self.embed_tokens = self.model.embed_tokens + self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors ) @@ -1132,6 +1137,7 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( tp_size, hf_config.linear_num_key_heads, diff --git a/atom/models/qwen3_next_mtp.py b/atom/models/qwen3_next_mtp.py index 2a5f0737ee..5c547d9559 100644 --- a/atom/models/qwen3_next_mtp.py +++ b/atom/models/qwen3_next_mtp.py @@ -171,9 +171,19 @@ def compute_logits( 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) + # Mirror target's get_expert_mapping: when shared-expert fusion is on, + # the loader rewrites `mlp.shared_expert.*` to `mlp.experts.{N}.*` + # (where N == n_routed_experts), so the expert_mapping must include + # an extra slot for that fused shared-expert. Without this, MTP's + # shared_expert weights get silently dropped during loading. + from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled + + n_routed = getattr(self.config, "n_routed_experts", self.config.num_experts) + n_shared = getattr(self.config, "n_shared_experts", 1) return FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, + num_experts=n_routed + + (n_shared if is_rocm_aiter_fusion_shared_expert_enabled() else 0), ) diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index d38da4a142..70e444379a 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -283,6 +283,28 @@ def init_method_under_plugin_mode( i64_kwargs = {"dtype": torch.int64, "device": device} self.positions = CpuGpuBuffer(max_num_batched_tokens, **i64_kwargs) + # Bump reorder_batch_threshold so multi-token spec-decode requests + # (MTP / EAGLE) are routed through the decode path. Mirrors vLLM's + # AttentionMetadataBuilder._init_reorder_batch_threshold(supports_spec_as_decode=True). + speculative_config = getattr(config, "speculative_config", None) + if ( + getattr(self, "reorder_batch_threshold", None) is not None + and speculative_config is not None + and getattr(speculative_config, "num_speculative_tokens", None) is not None + ): + parallel_drafting = getattr(speculative_config, "parallel_drafting", False) + max_num_queries_for_spec = 1 + (2 if parallel_drafting else 1) * ( + speculative_config.num_speculative_tokens + ) + self.reorder_batch_threshold = max( + self.reorder_batch_threshold, max_num_queries_for_spec + ) + logger.info( + "Spec decode: bumped reorder_batch_threshold to %d (num_spec_tokens=%d)", + self.reorder_batch_threshold, + speculative_config.num_speculative_tokens, + ) + return init_method_under_plugin_mode @@ -300,7 +322,7 @@ def setup_attn_metadata_builder_base_class_and_attributes(class_dict: dict): needs_generic = True # align with vllm rocm aiter fa - class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + class_dict["_cudagraph_support"] = AttentionCGSupport.UNIFORM_BATCH class_dict["reorder_batch_threshold"] = 1 return base_class, generic_base, needs_generic, class_dict @@ -324,9 +346,12 @@ def build( from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends - # here assume the decode num token is 1 per request + # decode_threshold tracks reorder_batch_threshold so MTP/EAGLE + # multi-token verification (query_len > 1) routes through decode. + decode_threshold = getattr(self, "reorder_batch_threshold", 1) or 1 split_ret = split_decodes_prefills_and_extends( - common_attn_metadata=common_attn_metadata, decode_threshold=1 + common_attn_metadata=common_attn_metadata, + decode_threshold=decode_threshold, ) ( @@ -351,6 +376,11 @@ def build( query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu + # In async spec-decode mode (auto-enabled for MTP/EAGLE), vLLM sets + # _num_computed_tokens_cpu to None because the GPU seq_lens is the + # authoritative source. Reconstruct from CPU tensors we already have. + if num_computed_tokens_cpu is None: + num_computed_tokens_cpu = seq_lens - query_lens_cpu prefill_max_query_len = decode_max_query_len = ( common_attn_metadata.max_query_len diff --git a/atom/plugin/attention_mha.py b/atom/plugin/attention_mha.py index 16c88949da..100c492aba 100644 --- a/atom/plugin/attention_mha.py +++ b/atom/plugin/attention_mha.py @@ -234,15 +234,27 @@ def paged_attention_triton_plugin_mode( v_cache: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor, + num_decodes: int, out: torch.Tensor, attn_metadata: "AttentionMetaData", ps: bool = True, ): - o = out - num_seqs, num_q_heads_total, head_size = q.shape + # q.shape[0] == num_decodes * max_query_len for MTP (one row per decode + # token, query_len > 1). For non-MTP it equals num_decodes (query_len = 1). + # pa_decode_gluon handles multi-token causal masking internally when + # `query_length > 1` is passed; intermediate buffers must be sized + # `num_decodes` (not q.shape[0]) and `query_group_size` must include + # the max_qlen multiplier — mirroring server-mode `paged_attention_triton`. + _, num_q_heads_total, head_size = q.shape num_blocks, num_kv_heads, _, block_size, _ = k_cache.shape - query_group_size = num_q_heads_total // num_kv_heads + decode_metadata = attn_metadata.plugin_metadata.decode_metadata + max_qlen = decode_metadata.max_query_len if decode_metadata is not None else 1 assert num_q_heads_total % num_kv_heads == 0 + + seq_lens = attn_metadata.plugin_metadata.seq_lens[:num_decodes] + block_tables = attn_metadata.plugin_metadata.block_table[:num_decodes] + + query_group_size = max_qlen * (num_q_heads_total // num_kv_heads) context_partition_size = 256 # use_ps = self.adopt_persistent_kernel( @@ -250,7 +262,9 @@ def paged_attention_triton_plugin_mode( # ) use_ps = True if use_ps: - max_context_partition_num = get_recommended_splits(num_seqs, num_kv_heads) + max_context_partition_num = get_recommended_splits( + num_decodes, num_kv_heads + ) else: max_context_partition_num = _NO_PS_FIXED_SPLITS @@ -258,9 +272,8 @@ def paged_attention_triton_plugin_mode( max_context_partition_num = 1 context_partition_size = 128 - # Output buffers (same as Triton) intermediate_shape = ( - num_seqs, + num_decodes, num_kv_heads, max_context_partition_num, query_group_size, @@ -283,21 +296,19 @@ def paged_attention_triton_plugin_mode( k_scale = k_scale.unsqueeze(-1) v_scale = v_scale.unsqueeze(-1) - num_decode_seqs = q.shape[0] - seq_lens_decode = attn_metadata.plugin_metadata.seq_lens[:num_decode_seqs] - block_tables_decode = attn_metadata.plugin_metadata.block_table[ - :num_decode_seqs - ] - + # Kernel takes natural q layout [batch * query_length, num_q_heads, head_size]. + # Internally it derives batch_size = q.shape[0] // query_length and reshapes + # to [batch, query_length, num_kv_heads, group, head_size]. See + # aiter/aiter/ops/triton/gluon/pa_decode_gluon.py:5371-5377 and 5542-5544. torch.ops.aiter.pa_decode_gluon( - o, + out, q, k_cache, v_cache, - seq_lens_decode, - block_tables_decode, + seq_lens, + block_tables, self.scale, - 1, # query_lenth + max_qlen, # query_length — handles multi-token causal mask internally max_context_partition_num, context_partition_size, compute_type, @@ -312,8 +323,7 @@ def paged_attention_triton_plugin_mode( sliding_window=self.sliding_window, ps=use_ps, ) - - return o + return out def paged_attention_asm_plugin_mode( self, @@ -327,6 +337,11 @@ def paged_attention_asm_plugin_mode( attn_metadata: "AttentionMetaData", out: torch.Tensor, ): + decode_metadata = attn_metadata.plugin_metadata.decode_metadata + max_qlen = decode_metadata.max_query_len if decode_metadata is not None else 1 + qo_indptr = ( + decode_metadata.query_start_loc if decode_metadata is not None else None + ) aiter.pa_fwd_asm( Q=q, K=k_cache, @@ -336,9 +351,11 @@ def paged_attention_asm_plugin_mode( block_tables_stride0=attn_metadata.plugin_metadata.block_table[ :num_decodes ].stride(0), + max_qlen=max_qlen, K_QScale=k_scale, V_QScale=v_scale, out_=out[:num_decode_tokens], + qo_indptr=qo_indptr, high_precision=0, ) @@ -706,12 +723,13 @@ def forward_impl_plugin_mode( extend_tokens_slice = slice( num_decode_tokens, num_decode_tokens + num_extend_tokens ) + extend_reqs_slice = slice(num_decodes, num_decodes + num_extends) extend_querys = query[extend_tokens_slice] extend_keys = key[extend_tokens_slice] extend_values = value[extend_tokens_slice] extend_outputs = output[extend_tokens_slice] extend_block_table = attn_metadata.plugin_metadata.block_table[ - extend_tokens_slice + extend_reqs_slice ] extend_slot_mapping = attn_metadata.plugin_metadata.slot_mapping[ extend_tokens_slice @@ -745,6 +763,7 @@ def forward_impl_plugin_mode( v_cache=new_value_cache, k_scale=k_scale, v_scale=v_scale, + num_decodes=num_decodes, out=output_actual_tokens[:num_decode_tokens], attn_metadata=attn_metadata, ) @@ -757,6 +776,7 @@ def forward_impl_plugin_mode( v_cache=new_value_cache, k_scale=k_scale, v_scale=v_scale, + num_decodes=num_decodes, out=output_actual_tokens[:num_decode_tokens], attn_metadata=attn_metadata, ) diff --git a/atom/plugin/config.py b/atom/plugin/config.py index 6d37bd70d7..575948aab3 100644 --- a/atom/plugin/config.py +++ b/atom/plugin/config.py @@ -1,3 +1,4 @@ +import copy from typing import Any, Optional from dataclasses import dataclass @@ -72,6 +73,45 @@ def _normalize_sglang_parallel_config( return tp_size, 1, 0, tp_rank +def _build_atom_speculative_config_from_vllm(vllm_spec_config: Any): + """Translate vLLM's SpeculativeConfig into ATOM's SpeculativeConfig. + + Reuses vLLM's already-loaded draft hf_config (skips a second disk fetch + in ATOM SpeculativeConfig.__post_init__) but still runs ATOM's + hf_config_override on it — so MTP model_type remap, n_routed_experts + backfill (Qwen families), and architecture rewrite all land on the + draft config in one place. Mirrors how standalone ATOM MTP exposes + the draft hf_config via atom_config.speculative_config. + + The draft hf_config is deepcopied first because hf_config_override + mutates `architectures` to ATOM's standalone naming (e.g. + "Qwen3NextMTPModel"), which differs from vLLM's registry name + ("Qwen3NextMTP"). Mutating in place would make vLLM's later draft + architecture lookup fail. + """ + if vllm_spec_config is None: + return None + + from atom.config import SpeculativeConfig + + draft_model_config = getattr(vllm_spec_config, "draft_model_config", None) + draft_hf_config = getattr(draft_model_config, "hf_config", None) + if draft_hf_config is not None: + draft_hf_config = copy.deepcopy(draft_hf_config) + model_path = getattr(draft_model_config, "model", None) or getattr( + vllm_spec_config, "model", None + ) + + return SpeculativeConfig( + method=getattr(vllm_spec_config, "method", "") or "", + model=model_path, + num_speculative_tokens=getattr( + vllm_spec_config, "num_speculative_tokens", None + ), + draft_model_hf_config=draft_hf_config, + ) + + def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: from atom.config import Config, CompilationConfig @@ -118,6 +158,10 @@ def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: max_num_batched_tokens = vllm_scheduler_config.max_num_batched_tokens + atom_speculative_config = _build_atom_speculative_config_from_vllm( + getattr(config, "speculative_config", None) + ) + return Config( model=vllm_model_config.model, trust_remote_code=getattr(vllm_model_config, "trust_remote_code", False), @@ -141,6 +185,7 @@ def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: master_addr=None, enable_dp_attention=False, plugin_config=plugin_config, + speculative_config=atom_speculative_config, ) diff --git a/atom/plugin/vllm/attention_backend/attention_gdn.py b/atom/plugin/vllm/attention_backend/attention_gdn.py index b6158a0867..87a2f2f9f9 100644 --- a/atom/plugin/vllm/attention_backend/attention_gdn.py +++ b/atom/plugin/vllm/attention_backend/attention_gdn.py @@ -22,6 +22,7 @@ from atom.model_ops.fla_ops.fused_sigmoid_gating import ( fused_sigmoid_gating_delta_rule_update, ) + from atom.utils import envs from torch import nn @@ -385,7 +386,13 @@ def forward( ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to( ssm_state.dtype ) - core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + # Only write directly when there are no spec tokens. With spec + # decode active, mixed_qkv was index_select'd by non_spec_token_indx + # so core_attn_out_non_spec has fewer rows than num_actual_tokens. + # The merge below (index_copy_) handles the scatter back to the + # correct slot positions. + if spec_sequence_masks is None: + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) elif attn_metadata.num_decodes > 0: o = core_attn_out[: attn_metadata.num_decode_tokens] if USE_FLYDSL_GDR: diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index c2b990c182..4eada7d4ca 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -35,7 +35,9 @@ logger = logging.getLogger("atom") - +_MTP_MASK_INPUT_ARCH: set[str] = { + "DeepSeekMTPModel", +} _ATOM_MODEL_CLASSES: dict[str, str] = { "LlamaForCausalLM": "atom.models.llama:LlamaForCausalLM", "Qwen3ForCausalLM": "atom.models.qwen3:Qwen3ForCausalLM", @@ -47,6 +49,7 @@ "GlmMoeDsaForCausalLM": "atom.models.deepseek_v2:GlmMoeDsaForCausalLM", "DeepSeekMTPModel": "atom.models.deepseek_mtp:DeepSeekMTP", "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLM", + "Qwen3NextMTP": "atom.models.qwen3_next_mtp:Qwen3NextMTP", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration_", "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration_", "KimiK25ForConditionalGeneration": "atom.plugin.vllm.models.kimi_k25:KimiK25ForConditionalGeneration_", @@ -121,6 +124,7 @@ def __init_subclass__(cls, *args, **kwargs): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + from atom.config import get_current_atom_config _set_framework_backbone("vllm") @@ -140,19 +144,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.ignore_unexpected_suffixes: list[str] = [] self.vllm_config = vllm_config - self.atom_config = generate_atom_config_for_plugin_mode(vllm_config) self.is_mtp = False speculative_config = getattr(vllm_config, "speculative_config", None) if speculative_config is not None: spec_method = speculative_config.method self.is_mtp = spec_method == "mtp" - _prepare_env(atom_config=self.atom_config) - main_model_arch = vllm_config.model_config.architectures[0] model_arch = _select_model_arch(vllm_config) self.is_mtp_draft_model = self.is_mtp and model_arch != main_model_arch + if self.is_mtp_draft_model: + self.atom_config = get_current_atom_config() + else: + self.atom_config = generate_atom_config_for_plugin_mode(vllm_config) self.model_arch = model_arch + _prepare_env(atom_config=self.atom_config) model_cls = _get_atom_model_cls(model_arch) module_remapping = getattr(model_cls, "packed_modules_mapping", {}) weights_mapper = getattr(model_cls, "hf_to_atom_mapper", {}) @@ -182,9 +188,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): logger.info(f"Construct ATOM model {model_arch} for vLLM plugin mode") self.model = model_cls(self.atom_config) - self._adapt_mtp_layers_for_vllm() - # Mirror nested attributes required by vLLM speculative decoding. - self._expose_spec_decode_attrs() + + if model_arch in _MTP_MASK_INPUT_ARCH: + self._adapt_mtp_layers_for_vllm() + # Mirror nested attributes required by vLLM speculative decoding. + self._expose_spec_decode_attrs() # For sparse MLA, register the Indexer's DeepseekV32IndexerCache as # a virtual subclass of vLLM's AttentionLayerBase so vLLM can discover @@ -192,7 +200,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._register_indexer_caches_with_vllm() if self.model is None: - model_arch = vllm_config.model_config.architectures[0] raise ValueError( f"The model {model_arch} is not supported by model impl backend atom" ) @@ -309,8 +316,7 @@ def _register_indexer_caches_with_vllm(self): if prefix not in vllm_sfc: vllm_sfc[prefix] = module logger.info( - f"Registered indexer cache in vLLM static_forward_context: " - f"{prefix}" + f"Registered indexer cache in vLLM static_forward_context: {prefix}" ) else: logger.warning( @@ -397,7 +403,6 @@ def forward( inputs_embeds=inputs_embeds, **model_kwargs, ) - if not self.pp_group.is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) @@ -412,12 +417,12 @@ def load_weights( is_mtp_draft_model = self.model_arch in { "DeepSeekMTPModel", - "Qwen3NextMTPModel", + "Qwen3NextMTP", } draft_hf_config = None if is_mtp_draft_model: draft_model_config = getattr( - getattr(self.vllm_config, "speculative_config", None), + getattr(self.atom_config, "speculative_config", None), "draft_model_config", None, ) @@ -452,7 +457,6 @@ class ATOMMoEForCausalLM(ATOMModelBase, VllmModelForTextGeneration): ... class ATOMForConditionalGeneration( ATOMModelBase, VllmModelForTextGeneration, SupportsMultiModal, SupportsMRoPE ): - @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: """ diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 9ef76e601d..91e241e9ae 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -30,6 +30,7 @@ "GlmMoeDsaForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "DeepSeekMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLMVllm", + "Qwen3NextMTP": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration", "KimiK25ForConditionalGeneration": "atom.plugin.vllm.models.kimi_k25:KimiK25ForConditionalGeneration", diff --git a/recipes/atom_vllm/Qwen3.5.md b/recipes/atom_vllm/Qwen3.5.md index 94a900e074..4e3b8c077f 100644 --- a/recipes/atom_vllm/Qwen3.5.md +++ b/recipes/atom_vllm/Qwen3.5.md @@ -18,6 +18,7 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3.5-35B-A3B-FP8 \ --host localhost \ @@ -37,6 +38,7 @@ vllm serve Qwen/Qwen3.5-35B-A3B-FP8 \ export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3.5-397B-A17B-FP8 \ --host localhost \ @@ -56,6 +58,7 @@ vllm serve Qwen/Qwen3.5-397B-A17B-FP8 \ export AITER_QUICK_REDUCE_QUANTIZATION=INT4 export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve amd/Qwen3.5-397B-A17B-MXFP4 \ --host localhost \ @@ -69,10 +72,10 @@ vllm serve amd/Qwen3.5-397B-A17B-MXFP4 \ --no-enable-prefix-caching ``` -**Important**: The following three environment variables are required for Qwen3.5: +**Important**: The following environment variables are required for Qwen3.5: -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: Disables ATOM attention plugin to use vLLM's implementation for full attention layers (required because Qwen3.5 uses a hybrid architecture with both linear attention (GatedDeltaNet) and full attention layers) - `ATOM_USE_CUSTOM_ALL_GATHER=0`: Disables custom all-gather for compatibility with Qwen3.5 model architecture +- `ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0`: Disables FP8 blockscale weight preshuffle - `AITER_QUICK_REDUCE_QUANTIZATION=INT4`: **Performance optimization** - enables INT4 quantization for quick reduce operations, which can significantly improve TTFT (Time To First Token) performance. **Note**: This optimization may introduce a risk of accuracy degradation. For accuracy-critical workloads, consider validating with your specific use case. ## Step 3: Performance Benchmark @@ -133,8 +136,8 @@ Reference result (TP=4): ## Key Environment Variables -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: **Required** - disables ATOM attention plugin to use vLLM's implementation for full attention layers - `ATOM_USE_CUSTOM_ALL_GATHER=0`: **Required** - disables custom all-gather for compatibility with Qwen3.5 model architecture +- `ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0`: **Required** - disables FP8 blockscale weight preshuffle - `AITER_QUICK_REDUCE_QUANTIZATION=INT4`: **Performance optimization** - enables INT4 quantization for quick reduce operations - **Benefit**: Significantly improves TTFT (Time To First Token) performance by reducing communication overhead during tensor parallelism all-reduce operations - **Risk**: May cause slight accuracy degradation due to lower quantization precision diff --git a/recipes/atom_vllm/Qwen3Next.md b/recipes/atom_vllm/Qwen3Next.md index e22f80d1c1..019e297f65 100644 --- a/recipes/atom_vllm/Qwen3Next.md +++ b/recipes/atom_vllm/Qwen3Next.md @@ -17,6 +17,8 @@ The ATOM vLLM plugin backend keeps the standard vLLM CLI, server APIs, and gener ```bash export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 export ATOM_USE_CUSTOM_ALL_GATHER=0 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ --host localhost \ @@ -31,8 +33,26 @@ vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ --no-enable-prefix-caching ``` -**Important**: `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1` is required for Qwen3-Next because it uses a hybrid architecture with both linear attention (GatedDeltaNet) and full attention layers. This env var ensures full attention layers use vLLM's default implementation. +### Qwen3-Next-80B-A3B-Instruct-FP8 MTP (TP=1, MI355X) +```bash +export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 +export ATOM_USE_CUSTOM_ALL_GATHER=0 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0 +vllm serve Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 \ + --host localhost \ + --port 8000 \ + --tensor-parallel-size 1 \ + --kv-cache-dtype fp8 \ + --gpu_memory_utilization 0.9 \ + --async-scheduling \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --max-model-len 16384 \ + --max-num-batched-tokens 32768 \ + --speculative-config '{"num_speculative_tokens":1, "method": "mtp"}' \ + --no-enable-prefix-caching +``` ## Step 3: Performance Benchmark Users can use the default vllm bench commands for performance benchmarking. @@ -70,9 +90,6 @@ lm_eval --model local-completions \ --num_fewshot 3 ``` -## Key Environment Variables - -- `ATOM_DISABLE_VLLM_PLUGIN_ATTENTION=1`: **Required** - disables ATOM attention plugin to use vLLM's implementation for full attention layers ## Architecture Notes From 8aa730e5d21080b16c5a974cdccf97e3fb613db4 Mon Sep 17 00:00:00 2001 From: ZhangLirong Date: Fri, 15 May 2026 16:03:23 +0800 Subject: [PATCH 54/76] ci: branch-aware docker release + fix benchmark model selection (#785) * ci: branch-aware docker release + fix benchmark model selection docker-release: use github.ref_name for branch-aware builds so non-default branches produce tagged images (nightly_YYYYMMDDHHMM-{branch}) without overwriting the latest tag. atom-benchmark: fix model selection bug where MiniMax and GLM-5.1 models were always included regardless of checkbox state, caused by input name mismatches (MiniMax-M2.5 vs M2.7 prefix, missing glm-5.1-fp4 input). Co-Authored-By: Claude Opus 4 * support tag by self --------- Co-authored-by: root Co-authored-by: Claude Opus 4 --- .github/workflows/atom-benchmark.yaml | 8 +++-- .github/workflows/docker-release.yaml | 46 +++++++++++++++++++++++---- 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/.github/workflows/atom-benchmark.yaml b/.github/workflows/atom-benchmark.yaml index fbd6111ab1..289d116c05 100644 --- a/.github/workflows/atom-benchmark.yaml +++ b/.github/workflows/atom-benchmark.yaml @@ -22,6 +22,10 @@ on: description: "Benchmark GLM-5-FP8" type: boolean default: true + glm-5.1-fp4: + description: "Benchmark GLM-5.1-MXFP4" + type: boolean + default: true deepseek-r1-0528-mxfp4: description: "Benchmark DeepSeek-R1-0528 MXFP4 (+ MXFP4-MTP3)" type: boolean @@ -34,8 +38,8 @@ on: description: "Benchmark Kimi-K2.5-MXFP4" type: boolean default: true - MiniMax-M2.5: - description: "Benchmark MiniMax-M2.5" + MiniMax-M2.7: + description: "Benchmark MiniMax-M2.7 (+ MXFP4)" type: boolean default: true qwen35-397b-fp8: diff --git a/.github/workflows/docker-release.yaml b/.github/workflows/docker-release.yaml index ebe6041afb..4cca15e75e 100644 --- a/.github/workflows/docker-release.yaml +++ b/.github/workflows/docker-release.yaml @@ -67,6 +67,9 @@ on: description: "Release the OOT vLLM image (skips native image test/push). If both this and SGLang are checked, both will be released." type: boolean default: false + custom_tag: + description: "Custom tag for the native Docker image (e.g. 'my-feature'). Produces rocm/atom-dev:. When set, branch-based tagging and 'latest' update are skipped." + default: "" only_release_sglang: description: "Release the SGLang+ATOM image (skips native image test/push). If both this and OOT are checked, both will be released." type: boolean @@ -110,12 +113,14 @@ jobs: - name: Echo environment variables run: | echo "ATOM_REPO: ${{ inputs.atom_repo || env.ATOM_REPO }}" + echo "ATOM_BRANCH: ${{ github.ref_name }}" echo "ATOM_COMMIT: ${{ inputs.atom_commit || env.ATOM_COMMIT }}" echo "AITER_REPO: ${{ inputs.aiter_repo || env.AITER_REPO }}" echo "AITER_COMMIT: ${{ inputs.aiter_commit || env.AITER_COMMIT }}" echo "RCCL_REPO: ${{ inputs.rccl_repo || env.RCCL_REPO }}" echo "RCCL_BRANCH: ${{ inputs.rccl_branch || env.RCCL_BRANCH }}" echo "OOT_TAG_SUFFIX: ${{ inputs.oot_tag_suffix || '' }}" + echo "CUSTOM_TAG: ${{ inputs.custom_tag || '' }}" echo "SGLANG_REPO: ${{ inputs.sglang_repo || env.SGLANG_REPO }}" echo "SGLANG_REF: ${{ inputs.sglang_ref || env.SGLANG_REF }}" echo "SGLANG_VERSION: ${{ inputs.sglang_version || env.SGLANG_VERSION }}" @@ -123,12 +128,19 @@ jobs: - name: Build Docker image timeout-minutes: 180 run: | + ATOM_COMMIT="${{ inputs.atom_commit || env.ATOM_COMMIT }}" + if [ "${ATOM_COMMIT}" = "HEAD" ]; then + ATOM_CHECKOUT="${{ github.ref_name }}" + else + ATOM_CHECKOUT="${ATOM_COMMIT}" + fi + echo "ATOM checkout ref: ${ATOM_CHECKOUT} (branch: ${{ github.ref_name }})" DOCKER_BUILDKIT=1 docker build --pull --network=host -t atom_release:ci \ --target atom_image \ --build-arg BASE_IMAGE="${{ matrix.base_image }}" \ --build-arg GPU_ARCH="${{ env.GPU_ARCH }}" \ --build-arg ATOM_REPO="${{ inputs.atom_repo || env.ATOM_REPO }}" \ - --build-arg ATOM_COMMIT="${{ inputs.atom_commit || env.ATOM_COMMIT }}" \ + --build-arg ATOM_COMMIT="${ATOM_CHECKOUT}" \ --build-arg AITER_REPO="${{ inputs.aiter_repo || env.AITER_REPO }}" \ --build-arg AITER_COMMIT="${{ inputs.aiter_commit || env.AITER_COMMIT }}" \ --build-arg RCCL_REPO="${{ inputs.rccl_repo || env.RCCL_REPO }}" \ @@ -189,12 +201,32 @@ jobs: - name: Push Docker image if: ${{ success() && inputs.only_release_oot != true && inputs.only_release_sglang != true }} run: | - # Push both the versioned tag and update the 'latest' tag for the Docker image - TAG=nightly_$(date +%Y%m%d%H%M) - docker tag atom_release:ci rocm/atom-dev:latest - docker push rocm/atom-dev:latest - docker tag atom_release:ci rocm/atom-dev:${TAG} - docker push rocm/atom-dev:${TAG} + CUSTOM_TAG="${{ inputs.custom_tag || '' }}" + if [ -n "${CUSTOM_TAG}" ]; then + if [[ ! "${CUSTOM_TAG}" =~ ^[A-Za-z0-9_.-]+$ ]]; then + echo "Invalid custom_tag '${CUSTOM_TAG}'. Allowed characters: letters, numbers, '.', '_' and '-'." + exit 1 + fi + docker tag atom_release:ci rocm/atom-dev:${CUSTOM_TAG} + docker push rocm/atom-dev:${CUSTOM_TAG} + echo "Custom tag: rocm/atom-dev:${CUSTOM_TAG}. Skipping latest update." + else + BRANCH="${{ github.ref_name }}" + DEFAULT_BRANCH="${{ github.event.repository.default_branch }}" + TAG=nightly_$(date +%Y%m%d%H%M) + if [ "${BRANCH}" != "${DEFAULT_BRANCH}" ]; then + SAFE_BRANCH=$(echo "${BRANCH}" | sed 's/[^A-Za-z0-9._-]/-/g') + TAG="${TAG}-${SAFE_BRANCH}" + docker tag atom_release:ci rocm/atom-dev:${TAG} + docker push rocm/atom-dev:${TAG} + echo "Branch '${BRANCH}'; tag: rocm/atom-dev:${TAG}. Skipping latest update." + else + docker tag atom_release:ci rocm/atom-dev:latest + docker push rocm/atom-dev:latest + docker tag atom_release:ci rocm/atom-dev:${TAG} + docker push rocm/atom-dev:${TAG} + fi + fi - name: Build OOT Docker image if: ${{ success() && (inputs.only_release_oot == true || (inputs.only_release_sglang != true && (github.event_name == 'schedule' || inputs.build_oot_image == true))) }} From 8927f69426428f18cbc034c0439ac57b2549d847 Mon Sep 17 00:00:00 2001 From: kliuae <17350011+kliuae@users.noreply.github.com> Date: Fri, 15 May 2026 17:09:54 +0800 Subject: [PATCH 55/76] [Feat] Support GLM-4.7 MTP in vLLM-ATOM plugin (#722) * adapt mtp for glm5 (vllm plugin) Signed-off-by: whx-sjtu * add patch to support mtp>1 Signed-off-by: whx-sjtu * fix model load failure of draft model Signed-off-by: whx-sjtu * adapt full graph with mtp enabled Signed-off-by: whx-sjtu * fix MLA MTP acceptance issue Signed-off-by: whx-sjtu * fall back to vllm-style mtp position Signed-off-by: whx-sjtu * fix embedding sharing failure for mtp Signed-off-by: whx-sjtu * fix lint Signed-off-by: whx-sjtu * fix comment Signed-off-by: whx-sjtu * remove warnig log Signed-off-by: whx-sjtu * add mtp support for glm4 Signed-off-by: kliuae * fix rope double apply for mha Signed-off-by: kliuae * guard vllm forward context retrieval Signed-off-by: kliuae * add glm4.7 mtp to workflow Signed-off-by: kliuae * clean up Signed-off-by: kliuae * format: fix black formatting in glm4_moe_mtp.py Co-Authored-By: Claude Opus 4 --------- Signed-off-by: whx-sjtu Signed-off-by: kliuae Signed-off-by: kliuae-amd Co-authored-by: whx-sjtu Co-authored-by: kuanfliu@amd.com Co-authored-by: zejunchen-zejun Co-authored-by: Claude Opus 4 --- .github/benchmark/oot_benchmark_models.json | 27 +++ .github/benchmark/oot_models_accuracy.json | 11 + .../atom-vllm-accuracy-validation.yaml | 15 ++ .github/workflows/atom-vllm-benchmark.yaml | 16 ++ atom/models/glm4_moe_mtp.py | 205 ++++++++++++++++++ atom/plugin/attention.py | 12 +- atom/plugin/vllm/model_wrapper.py | 2 + atom/plugin/vllm/register.py | 1 + 8 files changed, 288 insertions(+), 1 deletion(-) create mode 100644 atom/models/glm4_moe_mtp.py diff --git a/.github/benchmark/oot_benchmark_models.json b/.github/benchmark/oot_benchmark_models.json index 25d203a26e..63140f7289 100644 --- a/.github/benchmark/oot_benchmark_models.json +++ b/.github/benchmark/oot_benchmark_models.json @@ -343,6 +343,33 @@ } ] }, + { + "choice_label": "GLM-4.7-FP8 MTP", + "source_path": "zai-org/GLM-4.7-FP8", + "path": "zai-org/GLM-4.7-FP8", + "runner": "atom-mi355-8gpu-oot-benchmark", + "nightly_group": "B", + "variants": [ + { + "tp_size": 4, + "display": "GLM-4.7-FP8 MTP TP4 (AW)", + "dashboard_model": "GLM-4.7-FP8-mtp-aw-tp4", + "prefix": "glm-4-7-fp8-mtp-aw-tp4", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 4 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + }, + { + "tp_size": 8, + "display": "GLM-4.7-FP8 MTP TP8 (AW)", + "dashboard_model": "GLM-4.7-FP8-mtp-aw-tp8", + "prefix": "glm-4-7-fp8-mtp-aw-tp8", + "bench_args": "", + "extra_args": "--trust-remote-code --tensor-parallel-size 8 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4" + } + ] + }, { "choice_label": "MiniMax-M2.5", "source_path": "MiniMaxAI/MiniMax-M2.5", diff --git a/.github/benchmark/oot_models_accuracy.json b/.github/benchmark/oot_models_accuracy.json index fb68d48360..2e7042077b 100644 --- a/.github/benchmark/oot_models_accuracy.json +++ b/.github/benchmark/oot_models_accuracy.json @@ -184,6 +184,17 @@ "accuracy_baseline": 0.9, "accuracy_baseline_model": "openai/gpt-oss-120b" }, + { + "model_name": "GLM-4.7-FP8 TP8", + "model_path": "zai-org/GLM-4.7-FP8", + "extraArgs": "--tensor-parallel-size 8 --default-chat-template-kwargs '{\"enable_thinking\":false}'", + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + "test_level": "nightly", + "accuracy_threshold": 0.92, + "accuracy_baseline": 0.9386, + "accuracy_baseline_model": "zai-org/GLM-4.7-FP8" + }, { "model_name": "GLM-5.1-FP8 TP8", "model_path": "zai-org/GLM-5.1-FP8", diff --git a/.github/workflows/atom-vllm-accuracy-validation.yaml b/.github/workflows/atom-vllm-accuracy-validation.yaml index dc448212e1..6d82397d9c 100644 --- a/.github/workflows/atom-vllm-accuracy-validation.yaml +++ b/.github/workflows/atom-vllm-accuracy-validation.yaml @@ -89,6 +89,11 @@ on: required: false type: boolean default: false + run_glm4_7_fp8_mtp_tp8: + description: "GLM-4.7-FP8 MTP TP8" + required: false + type: boolean + default: false run_glm5_1_fp8_tp8: description: "GLM-5.1-FP8 TP8" required: false @@ -155,6 +160,7 @@ jobs: RUN_DSV32_FP8_TP8: ${{ inputs.run_dsv32_fp8_tp8 }} RUN_GPT_OSS_120B_TP1: ${{ inputs.run_gpt_oss_120b_tp1 }} RUN_GPT_OSS_120B_TP2: ${{ inputs.run_gpt_oss_120b_tp2 }} + RUN_GLM4_7_FP8_MTP_TP8: ${{ inputs.run_glm4_7_fp8_mtp_tp8 }} RUN_GLM5_1_FP8_TP8: ${{ inputs.run_glm5_1_fp8_tp8 }} INPUT_OOT_IMAGE: ${{ inputs.oot_image || '' }} run: | @@ -316,6 +322,15 @@ jobs: "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", "runner": "linux-atom-mi35x-4", }, + { + "toggle_env": "RUN_GLM4_7_FP8_MTP_TP8", + "model_name": "GLM-4.7-FP8 MTP TP8", + "model_path": "zai-org/GLM-4.7-FP8", + "extra_args": "--tensor-parallel-size 8 --speculative-config.method mtp --speculative-config.num_speculative_tokens 1", + "accuracy_test_threshold": 0.92, + "env_vars": "AITER_QUICK_REDUCE_QUANTIZATION=INT4\nATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1", + "runner": "linux-atom-mi35x-8", + }, { "toggle_env": "RUN_GLM5_1_FP8_TP8", "model_name": "GLM-5.1-FP8 TP8", diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index d7a178d517..19f66ee996 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -45,6 +45,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -80,6 +82,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -115,6 +119,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -150,6 +156,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -185,6 +193,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -220,6 +230,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -255,6 +267,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) @@ -290,6 +304,8 @@ on: - DeepSeek-V3.2 FP8 TP8 (AW) - GLM-4.7-FP8 TP4 (AW) - GLM-4.7-FP8 TP8 (AW) + - GLM-4.7-FP8 MTP TP4 (AW) + - GLM-4.7-FP8 MTP TP8 (AW) - gpt-oss-120b TP1 (AW) - gpt-oss-120b TP2 (AW) - gpt-oss-120b TP8 (AW) diff --git a/atom/models/glm4_moe_mtp.py b/atom/models/glm4_moe_mtp.py new file mode 100644 index 0000000000..b546b28323 --- /dev/null +++ b/atom/models/glm4_moe_mtp.py @@ -0,0 +1,205 @@ +from typing import Optional + +import torch +import torch.nn as nn +from atom.config import Config, QuantizationConfig +from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.moe import FusedMoE +from atom.model_ops.topK import is_rocm_aiter_fusion_shared_expert_enabled +from atom.models.utils import IntermediateTensors +from atom.utils.decorators import support_torch_compile +from transformers import PretrainedConfig + +from .deepseek_mtp import rewrite_spec_layer_name + +from .glm4_moe import Glm4MoeDecoderLayer, get_spec_layer_idx_from_weight_name +from .utils import maybe_prefix + + +class SharedHead(nn.Module): + def __init__( + self, + config: PretrainedConfig, + prefix: str, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "head"), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + +class Glm4MoeMultiTokenPredictorLayer(nn.Module): + def __init__(self, atom_config: Config, prefix: str) -> None: + super().__init__() + + config = atom_config.hf_config + self.config = config + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + + self.shared_head = SharedHead( + config=config, prefix=prefix, quant_config=atom_config.quant_config + ) + + self.mtp_block = Glm4MoeDecoderLayer( + config=config, + atom_config=atom_config, + prefix=prefix, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None + masked_inputs_embeds = inputs_embeds + inputs_embeds = self.enorm(masked_inputs_embeds) + previous_hidden_states = self.hnorm(previous_hidden_states) + + hidden_states = self.eh_proj( + torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + ) + + hidden_states, residual = self.mtp_block( + positions=positions, hidden_states=hidden_states, residual=None + ) + hidden_states = residual + hidden_states + return hidden_states + + +class Glm4MoeMultiTokenPredictor(nn.Module): + def __init__(self, *, atom_config: Config, prefix: str = ""): + super().__init__() + config = atom_config.hf_config + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + # to map the exact layer index from weights + self.layers = torch.nn.ModuleDict( + { + str(idx): Glm4MoeMultiTokenPredictorLayer( + atom_config, f"{prefix}.layers.{idx}" + ) + for idx in range( + self.mtp_start_layer_idx, + self.mtp_start_layer_idx + self.num_mtp_layers, + ) + } + ) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)]( + input_ids, + positions, + previous_hidden_states, + inputs_embeds, + current_step_idx, + ) + + def compute_logits( + 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)] + logits = mtp_layer.shared_head.head(mtp_layer.shared_head(hidden_states)) + return logits + + +@support_torch_compile +class Glm4MoeMTP(nn.Module): + # ATOM format: checkpoint_weight_name -> (model_param_name, shard_id) + packed_modules_mapping = { + "q_proj": ("qkv_proj", "q"), + "k_proj": ("qkv_proj", "k"), + "v_proj": ("qkv_proj", "v"), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__() + self.config = atom_config.hf_config + + self.model = Glm4MoeMultiTokenPredictor( + atom_config=atom_config, prefix=maybe_prefix(prefix, "model") + ) + + def remap_mtp_weight_name(self, name: str) -> str | None: + # GLM-4 MoE MTP shares the rewrite rules with DeepSeek MTP: + # - shared scalars (embed_tokens) → top-level model.* + # - per-layer scalars (enorm/hnorm/eh_proj/shared_head) → kept verbatim + # - decoder block weights (self_attn/mlp/...) → model.layers.{i}.mtp_block.* + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) + if spec_layer is None: + return None + return rewrite_spec_layer_name(spec_layer, name) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(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) + num_experts = self.config.n_routed_experts + if ( + is_rocm_aiter_fusion_shared_expert_enabled() + and self.config.n_shared_experts + ): + num_experts += self.config.n_shared_experts + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index 70e444379a..9441eb495e 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -375,6 +375,9 @@ def build( query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + # The spec-decode draft path invalidates it between proposal steps when + # num_speculative_tokens > 1. + # Fall back to seq_lens - query_lens computed on already-CPU tensors. num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu # In async spec-decode mode (auto-enabled for MTP/EAGLE), vLLM sets # _num_computed_tokens_cpu to None because the GPU seq_lens is the @@ -2477,7 +2480,14 @@ def unified_attention_with_output_base_for_plugin_mode( # ATOM needs to handle all of the buffer here # Positions for compiled unified_attention are provided via vLLM ForwardContext # (atom_positions) in ATOMModelBase.forward, not via this Python arg. - if envs.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION: + # + # RoPE single-source-of-truth in plugin mode: when ATOM's + # PagedAttentionImpl is active, it applies RoPE inside + # rope_cache_plugin_mode. Applying RoPE here too would double apply + # it for any model hitting the "else" branch in rope_cache_plugin_mode when + # ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION is not set. + # if envs.ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION: + if current_atom_config.plugin_config.vllm_use_atom_attention: output = self.attn(q, k, v) else: # calculate the q and k with rotary embedding diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index 4eada7d4ca..fbdd5e94dd 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -48,6 +48,7 @@ "Glm4MoeForCausalLM": "atom.models.glm4_moe:Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM": "atom.models.deepseek_v2:GlmMoeDsaForCausalLM", "DeepSeekMTPModel": "atom.models.deepseek_mtp:DeepSeekMTP", + "Glm4MoeMTPModel": "atom.models.glm4_moe_mtp:Glm4MoeMTP", "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLM", "Qwen3NextMTP": "atom.models.qwen3_next_mtp:Qwen3NextMTP", "Qwen3_5MoeForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5MoeForConditionalGeneration_", @@ -418,6 +419,7 @@ def load_weights( is_mtp_draft_model = self.model_arch in { "DeepSeekMTPModel", "Qwen3NextMTP", + "Glm4MoeMTPModel", } draft_hf_config = None if is_mtp_draft_model: diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 91e241e9ae..25c75bba55 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -29,6 +29,7 @@ "Glm4MoeForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "GlmMoeDsaForCausalLM": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "DeepSeekMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, + "Glm4MoeMTPModel": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3NextForCausalLM": "atom.models.qwen3_next:Qwen3NextForCausalLMVllm", "Qwen3NextMTP": ATOM_MOE_CAUSAL_LM_MODEL_WRAPPER, "Qwen3_5ForConditionalGeneration": "atom.models.qwen3_5:Qwen3_5ForConditionalGeneration", From 7516968e9b8478e5251240aa36c3f88bf954f6ef Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Fri, 15 May 2026 19:31:00 +0900 Subject: [PATCH 56/76] Add inferencex-pr skill for ATOM benchmark comparison and PR creation (#795) * Add inferencex-pr skill for ATOM benchmark comparison and PR creation Co-Authored-By: Claude Sonnet 4.6 * Replace inferencex-pr skill with inferencex-sync Renames the skill to better reflect its purpose: comparing ATOM upstream benchmark results against InferenceX and reporting regression/improvement. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- .claude/commands/inferencex-sync.md | 68 +++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 .claude/commands/inferencex-sync.md diff --git a/.claude/commands/inferencex-sync.md b/.claude/commands/inferencex-sync.md new file mode 100644 index 0000000000..4f9f538884 --- /dev/null +++ b/.claude/commands/inferencex-sync.md @@ -0,0 +1,68 @@ +--- +name: inferencex-sync +description: Compare the performance of the latest successful run from ATOM upstream benchmark(https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml) and that of InferenceX(https://inferencex.semianalysis.com/api/v1/benchmarks?model=$MODEL) and report tput_per_gpu (Total token throughput / number of GPUs), report performance regression of each model on atom framework, create a PR if the performance of ATOM upstream benchamrk is better than that of InferenceX by updating docker image and atom serve arguments (python3 -m atom.entrypoints.openai_server) +memory: project +model: opus +--- + +# Compare ATOM upstream benchmark and InferneceX, and create a PR to InferenceX Instructions + +Use this skill to compare the througput of ATOM upstream benchamrk and InferenceX and create a PR to InferenceX with the udpated docker image and atom serve arguments. + +## Workflow + +```text +- [ ] 1) Create a throughput table from InferenceX benchmark +- [ ] 2) Create a throughput table from ATOM upstream benchmark +- [ ] 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark +``` + +## 1) Create a throughput table from InferenceX benchmark + +Check the InferenceX benchmark that is from these URLs and create a throughput performance table from ATOM upstream benchmark. +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=Kimi-K2.5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=DeepSeek-V4-Pro +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=DeepSeek-R1-0528 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=Qwen-3.5-397B-A17B +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=GLM-5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=MiniMax-M2.5 +- https://inferencex.semianalysis.com/api/v1/benchmarks?model=gpt-oss-120b + +Rules: +- 1. Use curl to fetch data from URLs +- 2. Use only "hardware":"mi355x","framework":"atom","disagg":false,"is_multinode":false,"spec_method":"none" +- 3. If muliple data of "hardware":"mi355x","framework":"atom","precision":"fp4","disagg":false,"is_multinode":false exist for a same "model", use only the latest "date" +- 4. Create a throughput performance table of each "tput_per_gpu" of each "model","isl", "osl","conc","num_decode_gpus"(number of GPUs),"image","precision" + +## 2) Create a throughput table from ATOM upstream benchmark + +Check the ATOM upstream benchmark that is from the latest successful gh action runs at https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml and append the throughput to the InferenceX performance table in Step 1) + +Rules: +- 1. Use main branch. Find the latest job of https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml where all jobs passed +- 2. "tput_per_gpu" (tput/GPU (tok/s)) is "Total token throughput" divided by TP(number of GPUs). TP(number of GPUs) is the value followed by -tp in https://github.com/ROCm/ATOM/blob/main/.github/benchmark/models.json +- 3. "conc" is "Max Concurrency". Compare same "conc" or "Max Concurrency" between InferenceX benchmark and Atom benchmark +- 4. Model name follows this mapping + + ```python + # inferenceX_model -> (atom_model, precision) + MODEL_MAP = { + "dsv4": [("DeepSeek-V4-Pro", "fp4")], + "dsr1": [("DeepSeek-R1-0528", "fp8"), + ("DeepSeek-R1-0528-MXFP4", "fp4")], + "kimik2.5": [("Kimi-K2.5-MXFP4", "fp4")], + "qwen3.5": [("Qwen3.5-397B-A17B-FP8", "fp8"), + ("Qwen3.5-397B-A17B-MXFP4", "fp4")], + "glm5": [("GLM-5-FP8", "fp8"), + ("GLM-5.1-MXFP4", "fp4")], + "minimaxm2.5": [("MiniMax-M2.7", "fp8"), + ("MiniMax-M2.7-MXFP4", "fp4")], + "gptoss120b": [("gpt-oss-120b", "fp4")], + } + ``` + +If the data is not available, stop and ask for any workaround + +## 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark + +Update the performance table to show regression in % of ATOM upstream benchmark from InferenceX benchmark \ No newline at end of file From 26204cc09e91961e44ef901246c82f9e1395a688 Mon Sep 17 00:00:00 2001 From: wuhuikx Date: Fri, 15 May 2026 19:42:35 +0800 Subject: [PATCH 57/76] fix: handle auto kv cache dtype in vLLM plugin (#804) Map vLLM cache dtype aliases before passing them into AITER metadata setup so the default auto value does not fail dtype parsing. Co-authored-by: XiaobingSuper Co-authored-by: Cursor --- atom/plugin/attention.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index 9441eb495e..ec7195c337 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -22,6 +22,17 @@ disable_vllm_plugin_attention = envs.ATOM_DISABLE_VLLM_PLUGIN_ATTENTION +def _get_aiter_kv_cache_dtype(config) -> torch.dtype: + kv_cache_dtype = config.cache_config.cache_dtype + if kv_cache_dtype == "auto": + kv_cache_dtype = "bf16" + elif kv_cache_dtype == "bfloat16": + kv_cache_dtype = "bf16" + elif kv_cache_dtype == "float16": + kv_cache_dtype = "fp16" + return dtypes.d_dtypes[kv_cache_dtype] + + @dataclass class AiterFlashAttentionPhaseMetadata: max_query_len: int @@ -1299,7 +1310,7 @@ def init_method_under_plugin_mode( 1, self.padded_num_attention_heads, torch.bfloat16, - dtypes.d_dtypes[config.cache_config.cache_dtype], + _get_aiter_kv_cache_dtype(config), is_sparse=False, # TODO: support sparse fast_mode=True, ) From 9d279f00deea4a7d33c4086997f90c60a79339b2 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Fri, 15 May 2026 21:23:26 +0800 Subject: [PATCH 58/76] [fix](mha): handle scalar kv scales in prefix gather (#793) Co-authored-by: Cursor --- atom/model_ops/attention_mha.py | 8 ++++++++ atom/model_ops/base_attention.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 5fbcaca0c5..9e0d5dbb0b 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -354,6 +354,13 @@ def _gather_prefix_and_concat_kv( block_size = k_cache.shape[1] block_tables = attn_metadata.block_tables + per_token_quant = ( + self.kv_cache_dtype.startswith("fp8") + and k_scale is not None + and v_scale is not None + and k_scale.numel() > 1 + and v_scale.numel() > 1 + ) cp_mha_gather_cache( key_cache=k_cache_gather, value_cache=v_cache_gather, @@ -368,6 +375,7 @@ def _gather_prefix_and_concat_kv( dequant=self.kv_cache_dtype.startswith("fp8"), kv_cache_layout="SHUFFLE" if use_shuffle else "NHD", total_tokens=total_tokens, + per_token_quant=per_token_quant, ) return q, k_full, v_full, k_cache, v_cache, k_scale, v_scale diff --git a/atom/model_ops/base_attention.py b/atom/model_ops/base_attention.py index baf6f4e898..8a7cd45c50 100644 --- a/atom/model_ops/base_attention.py +++ b/atom/model_ops/base_attention.py @@ -165,6 +165,12 @@ def cp_mha_gather_cache( ], "kv_cache_layout only support NHD, SHUFFLE" if dequant: assert k_scales is not None and v_scales is not None + if k_scales.numel() == 1 and v_scales.numel() == 1: + per_token_quant = False + else: + assert ( + k_scales.numel() > 1 and v_scales.numel() > 1 + ), "k_scales and v_scales must both be scalar or per-token" head_dim = key.shape[2] x = 16 // key_cache.element_size() From ebc515455999188ded964bbe48d4f10a5287b2cc Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Fri, 15 May 2026 23:11:32 +0800 Subject: [PATCH 59/76] fix: handle auto kv cache dtype in vLLM plugin for sparse mla (#806) --- atom/plugin/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atom/plugin/attention.py b/atom/plugin/attention.py index ec7195c337..e46441af37 100644 --- a/atom/plugin/attention.py +++ b/atom/plugin/attention.py @@ -2394,7 +2394,7 @@ def init_method_under_plugin_mode( 1, self.padded_num_heads, torch.bfloat16, - dtypes.d_dtypes[config.cache_config.cache_dtype], + _get_aiter_kv_cache_dtype(config), is_sparse=True, fast_mode=True, ) From d2c77638fac5387d2e38a7be47e24d4531b45890 Mon Sep 17 00:00:00 2001 From: XiaobingZhang Date: Fri, 15 May 2026 23:13:26 +0800 Subject: [PATCH 60/76] Add DeepSeek-V3.2 fused indexer path (#788) * [perf](deepseek): add fused indexer path Co-authored-by: Cursor * [fix](deepseek): clear fused indexer rope args Co-authored-by: Cursor * [fix](deepseek): simplify fused indexer path Co-authored-by: Cursor * [fix](deepseek): drop unrelated indexer weights change Co-authored-by: Cursor * [fix](deepseek): gate fused indexer path Co-authored-by: Cursor * [fix](deepseek): guard indexer wk fusion by dtype Co-authored-by: Cursor * [fix](deepseek): remove unused sparse import Co-authored-by: Cursor * [fix](deepseek): address fused indexer review feedback Keep indexer weight fusion decisions consistent with quant fallback paths and make dummy/profile behavior deterministic. Co-authored-by: Cursor * [fix](deepseek): align indexer review fixes with black Keep the PR compatible with the repository's Black formatting check after the review fixes. Co-authored-by: Cursor * [fix](deepseek): simplify indexer fusion guard Collapse the review helper stack while keeping a single model-level fusion decision for weight loading. Co-authored-by: Cursor * [fix](deepseek): move pending fp8 wk before dequant Ensure threaded checkpoint loading dequantizes pending FP8 indexer wk weights on the same device as their scales. Co-authored-by: Cursor * [fix](deepseek): tighten indexer fusion fallback Keep the env-disabled path aligned with the unfused baseline and fail clearly when fused wk weights are incomplete. Co-authored-by: Cursor * [fix](deepseek): decouple indexer wk fusion guard Keep wk/weights projection fusion capability independent from the QK RoPE cache fusion runtime switch. Co-authored-by: Cursor * [fix](deepseek): restrict fused indexer kernel configs Avoid enabling the fused indexer kernels for GLM configs with incompatible RoPE cache shapes while keeping compatible sparse indexer configs eligible. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- atom/models/deepseek_v2.py | 399 +++++++++++++++++++++++----- atom/plugin/attention_mla_sparse.py | 77 ++++-- atom/utils/envs.py | 15 +- 3 files changed, 407 insertions(+), 84 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 252319a2da..a2b8b5af83 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -35,6 +35,7 @@ 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, ) @@ -116,6 +117,48 @@ ENABLE_DS_QKNORM_FUSION = envs.ATOM_ENABLE_DS_QKNORM_FUSION ENABLE_ALLREDUCE_RMSNORM_FUSION = envs.ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION = envs.ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION +ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION = ( + envs.ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION +) +_FP8_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e4m3fnuz", None), + ) + if dtype is not None +) + + +def _supports_fused_indexer_kernel_config(config: PretrainedConfig) -> bool: + if not hasattr(config, "index_topk"): + return False + if getattr(config, "model_type", None) == "glm_moe_dsa": + return False + return ( + getattr(config, "index_head_dim", None) == 128 + and getattr(config, "qk_rope_head_dim", None) == 64 + ) + + +def _can_fuse_indexer_wk_weights_proj( + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], + indexer_prefixes: list[str], +) -> bool: + if not _supports_fused_indexer_kernel_config(config): + return False + if quant_config is None: + return True + + for indexer_prefix in indexer_prefixes: + wk_quant_config = quant_config.get_layer_quant_config(f"{indexer_prefix}.wk") + if ( + wk_quant_config.quant_type != QuantType.No + and wk_quant_config.quant_dtype != dtypes.fp8 + ): + return False + return True def _fuse_rmsnorm_fp4_quant_fake( @@ -722,7 +765,6 @@ def _fuse_qkv_a_proj_reduce_rmsnorm_quant( class DeepseekV2MLP(nn.Module): - def __init__( self, hidden_size: int, @@ -750,8 +792,7 @@ def __init__( ) if hidden_act != "silu": raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only silu is supported for now." + f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() @@ -763,7 +804,6 @@ def forward(self, x): class DeepseekV2MoE(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -932,7 +972,6 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: @DeepseekV32IndexerCacheDecoratorForPluginMode class DeepseekV32IndexerCache(nn.Module): - def __init__( self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: str ): @@ -948,7 +987,7 @@ def sparse_attn_indexer( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -958,6 +997,15 @@ def sparse_attn_indexer( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -967,21 +1015,50 @@ def sparse_attn_indexer( # Skip for dummy runs to avoid corrupting KV cache if forward_context.context.is_dummy_run: # dummy runner - return weights + return torch.zeros_like(weights, dtype=torch.float32) # For MTP verify decode, max_seqlen_q > 1 so total decode tokens = batch_size * max_seqlen_q num_decode_tokens = ( context.batch_size * attn_metadata.max_seqlen_q if not context.is_prefill else 0 ) runner_block_size = get_current_atom_config().kv_cache_block_size kv_cache = kv_cache.view(-1, runner_block_size, kv_cache.shape[-1]) - indexer_k_quant_and_cache( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, - preshuffle=True, - ) + if use_qk_rope_cache_fusion: + q_bf16 = q_input + q_fp8 = torch.empty_like(q_bf16, dtype=dtypes.fp8) + weights_out = torch.empty( + weights.shape, device=weights.device, dtype=torch.float32 + ) + indexer_qk_rope_quant_and_cache( + q_bf16, + q_fp8, + weights, + weights_out, + k, + kv_cache, + slot_mapping, + k_norm_weight, + k_norm_bias, + positions, + cos_cache, + sin_cache, + k_norm_eps, + quant_block_size, + scale_fmt, + weights_scale, + preshuffle=True, + is_neox=is_neox_style, + ) + weights = weights_out + else: + q_fp8 = q_input + indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + preshuffle=True, + ) if context.is_prefill: if attn_metadata.max_seqlen_k <= topk_indices_buffer.shape[1]: return weights @@ -1086,7 +1163,7 @@ def sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -1096,6 +1173,15 @@ def sparse_attn_indexer_fake( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # profile run # NOTE(Chen): create the max possible flattened_kv. So that @@ -1105,7 +1191,7 @@ def sparse_attn_indexer_fake( ) _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous() _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous() - return weights + return torch.empty(weights.shape, device=weights.device, dtype=torch.float32) direct_register_custom_op( @@ -1116,9 +1202,120 @@ def sparse_attn_indexer_fake( ) +def _dequant_fp8_block_to_bf16( + weight_fp8: torch.Tensor, + scale: torch.Tensor, + block_size: int = 128, +) -> torch.Tensor: + """Dequantize block-scaled FP8 weights to BF16 for BF16-only fused GEMMs.""" + out_dim, in_dim = weight_fp8.shape + if out_dim % block_size != 0 or in_dim % block_size != 0: + raise ValueError( + "FP8 block dequant expects dimensions divisible by " + f"{block_size}, got {tuple(weight_fp8.shape)}" + ) + weight = ( + weight_fp8.unflatten(0, (-1, block_size)) + .unflatten(-1, (-1, block_size)) + .float() + ) + scale = scale.float() + return (weight * scale[:, None, :, None]).flatten(2, 3).flatten(0, 1).bfloat16() + + +class IndexerWkWeightsProjLinear(MergedReplicatedLinear): + """Fused Indexer wk + weights projection with FP8 wk load support.""" + + def __init__( + self, + hidden_size: int, + head_dim: int, + n_head: int, + prefix: str = "", + ): + self._wk_pending_weight: Optional[torch.Tensor] = None + self._wk_pending_scale: Optional[torch.Tensor] = None + self._wk_loaded = False + super().__init__( + hidden_size, + [head_dim, n_head], + bias=False, + quant_config=None, + prefix=prefix, + ) + # Checkpoints may store indexer.wk as FP8 plus weight_scale_inv. The + # fused GEMM runs in BF16, so this parameter only receives the scale + # during loading and is not consumed in forward. + self.weight_scale = atom_parameter( + torch.empty( + ((head_dim + 127) // 128, (hidden_size + 127) // 128), + dtype=torch.float32, + ) + ) + self.weight_scale.weight_loader_process = self.weight_loader_process + self.weight_scale.weight_loader = self.weight_loader + + def _maybe_load_pending_wk(self) -> None: + if self._wk_pending_weight is None or self._wk_pending_scale is None: + return + wk_weight_fp8 = self._wk_pending_weight + if wk_weight_fp8.device != self._wk_pending_scale.device: + wk_weight_fp8 = wk_weight_fp8.to(self._wk_pending_scale.device) + wk_weight = _dequant_fp8_block_to_bf16( + wk_weight_fp8, + self._wk_pending_scale, + ) + super().weight_loader(self.weight, wk_weight, 0) + self._wk_pending_weight = None + self._wk_pending_scale = None + self._wk_loaded = True + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[int] = None, + ): + if param is self.weight_scale: + if loaded_shard_id == 0: + param.weight_loader_process(param.data, loaded_weight) + self._wk_pending_scale = param.data.detach().clone() + self._maybe_load_pending_wk() + return + + if ( + param is self.weight + and loaded_shard_id == 0 + and loaded_weight.dtype in _FP8_DTYPES + ): + self._wk_pending_weight = loaded_weight.detach().clone() + self._maybe_load_pending_wk() + return + + if param is self.weight and loaded_shard_id == 0: + self._wk_pending_weight = None + self._wk_pending_scale = None + self._wk_loaded = True + + super().weight_loader(param, loaded_weight, loaded_shard_id) + + def process_weights_after_loading(self): + if self._wk_pending_weight is not None or ( + self._wk_pending_scale is not None and not self._wk_loaded + ): + raise RuntimeError( + "Incomplete FP8 indexer.wk load: both weight and weight_scale " + "are required before building wk_weights_proj." + ) + if not self._wk_loaded: + raise RuntimeError( + "Missing indexer.wk load before building wk_weights_proj." + ) + super().process_weights_after_loading() + + @IndexerDecoratorForPluginMode class Indexer(nn.Module): - def __init__( self, atom_config: Config, @@ -1128,6 +1325,7 @@ def __init__( quant_config: Optional[QuantizationConfig], cache_config: str, topk_indices_buffer: Optional[torch.Tensor], + use_wk_weights_proj_fusion: bool = True, prefix: str = "", ): super().__init__() @@ -1147,25 +1345,41 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - self.wk = ReplicatedLinear( - hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wk", - ) + self.use_wk_weights_proj_fusion = use_wk_weights_proj_fusion + if self.use_wk_weights_proj_fusion: + self.wk_weights_proj = IndexerWkWeightsProjLinear( + hidden_size, + self.head_dim, + self.n_head, + prefix=f"{prefix}.wk_weights_proj", + ) + else: + self.wk = ReplicatedLinear( + hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wk", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) self.k_norm = LayerNorm(self.head_dim, eps=1e-6) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, - quant_config=None, - prefix=f"{prefix}.weights_proj", - ) self.softmax_scale = self.head_dim**-0.5 + self._weights_scale = self.softmax_scale * self.n_head**-0.5 self.scale_fmt = "ue8m0" self.quant_func = get_hip_quant(QuantType.per_1x128) self.quant_block_size = 128 # TODO: get from config + self.use_qk_rope_cache_fusion = ( + ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION + and _supports_fused_indexer_kernel_config(config) + and self.head_dim == self.quant_block_size + and self.rope_dim == self.head_dim // 2 + ) self.topk_indices_buffer = topk_indices_buffer # TODO (zyongye) change dim to fp8 later to (self.head_dim + 4) @@ -1192,36 +1406,45 @@ def forward( ) -> torch.Tensor: q = self.wq_b(qr, qr_scale) q = q.view(-1, self.n_head, self.head_dim) - q_pe, q_nope = torch.split( - q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - - k = self.wk(hidden_states) - k = self.k_norm(k) - k_pe, k_nope = torch.split( - k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 - ) - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) - - # we only quant q here since k quant is fused with cache insertion - q = q.view(-1, self.head_dim) - - q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) - q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) - q_scale = q_scale.view(-1, self.n_head, 1) - - weights = self.weights_proj(hidden_states) - weights = ( - weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 - ) - weights = weights.squeeze(-1) + if self.use_wk_weights_proj_fusion: + k, weights = torch.split( + self.wk_weights_proj(hidden_states), + [self.head_dim, self.n_head], + dim=-1, + ) + else: + k = self.wk(hidden_states) + weights = self.weights_proj(hidden_states) + k = k.contiguous() + weights = weights.contiguous() + + if not self.use_qk_rope_cache_fusion: + q_pe, _ = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k = self.k_norm(k) + k_pe, _ = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) + + q = q.view(-1, self.head_dim) + q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + weights = (weights.unsqueeze(-1) * q_scale * self._weights_scale).squeeze( + -1 + ) + q_input = q_fp8 + else: + q_input = q return self.sparse_attn_indexer_impl( hidden_states, self.k_cache.prefix, self.k_cache.kv_cache[0], - q_fp8, + q_input, k, weights, self.quant_block_size, @@ -1231,6 +1454,15 @@ def forward( self.max_model_len, self.max_total_seq_len, self.topk_indices_buffer, + self.k_norm.weight, + self.k_norm.bias, + self.k_norm.eps, + positions, + rotary_emb.cos_cache, + rotary_emb.sin_cache, + self._weights_scale, + rotary_emb.is_neox_style, + self.use_qk_rope_cache_fusion, ) @@ -1258,6 +1490,7 @@ def __init__( prefix: str = "", layer_num: int = 0, topk_indices_buffer: Optional[torch.Tensor] = None, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -1268,6 +1501,7 @@ def __init__( self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank + model_quant_config = quant_config self.num_heads = num_heads tp_size = get_tensor_model_parallel_world_size() @@ -1429,6 +1663,15 @@ def __init__( base_quant_config, cache_config, topk_indices_buffer, + ( + _can_fuse_indexer_wk_weights_proj( + config, + model_quant_config, + [f"{prefix}.indexer"], + ) + if use_indexer_wk_weights_proj_fusion is None + else use_indexer_wk_weights_proj_fusion + ), f"{prefix}.indexer", ) else: @@ -1576,7 +1819,6 @@ def forward( class DeepseekV2DecoderLayer(nn.Module): - def __init__( self, config: PretrainedConfig, @@ -1587,6 +1829,7 @@ def __init__( layer_num: int = 0, is_mtp_block: bool = False, alt_stream: Optional[torch.cuda.Stream] = None, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -1611,6 +1854,7 @@ def __init__( prefix=f"{prefix}.self_attn", layer_num=layer_num, topk_indices_buffer=topk_indices_buffer, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ) # When ATOM_ENABLE_DS_INPUT_RMSNORM_QUANT_FUSION is turned on self.fuse_input_norm_quant is turned on only if use_triton_gemm and (FP8 or FP4), @@ -1782,6 +2026,7 @@ def __init__( atom_config: Config, prefix: str = "", layer_type: type[nn.Module] = DeepseekV2DecoderLayer, + use_indexer_wk_weights_proj_fusion: Optional[bool] = None, ): super().__init__() @@ -1826,6 +2071,7 @@ def __init__( quant_config=quant_config, layer_num=layer_num, alt_stream=_alt_stream, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ), prefix=f"{prefix}.layers", layer_num_offset=0, @@ -1906,7 +2152,6 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: class DeepseekV2ForCausalLM(nn.Module): - def __init__( self, atom_config: Config, @@ -1919,6 +2164,23 @@ def __init__( self.config = config self.quant_config = quant_config + model_prefix = maybe_prefix(prefix, "model") + attn_module_list_cfg = getattr(config, "attn_module_list_cfg", None) + indexer_prefixes = [] + if isinstance(attn_module_list_cfg, (list, tuple)): + indexer_prefixes = [ + f"{model_prefix}.layers.{layer_idx}.self_attn.indexer" + for layer_idx, layer_cfg in enumerate(attn_module_list_cfg) + if isinstance(layer_cfg, dict) + and layer_cfg.get("attn_index") is not None + ] + if not indexer_prefixes: + indexer_prefixes = [f"{model_prefix}.layers.0.self_attn.indexer"] + use_indexer_wk_weights_proj_fusion = _can_fuse_indexer_wk_weights_proj( + config, + quant_config, + indexer_prefixes, + ) if hasattr(config, "q_lora_rank") and config.q_lora_rank is not None: self.packed_modules_mapping = { "q_a_proj": ("fused_qkv_a_proj", 0), @@ -1931,11 +2193,19 @@ def __init__( "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } + if use_indexer_wk_weights_proj_fusion: + self.packed_modules_mapping.update( + { + "indexer.wk": ("indexer.wk_weights_proj", 0), + "indexer.weights_proj": ("indexer.wk_weights_proj", 1), + } + ) self.model = DeepseekV2Model( atom_config=atom_config, - prefix=maybe_prefix(prefix, "model"), + prefix=model_prefix, layer_type=layer_type, + use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( @@ -2003,17 +2273,20 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM): - # DeepSeek-V3.2's indexer weights_proj are not quantized, but they are not listed in - # model's quant exclude mapping. So we add it to quant_default_exclude_layers. - quant_default_exclude_layers: list[str] = ["*.indexer.weights_proj"] + # DeepSeek-V3.2's indexer weights projection is BF16. Keep the original + # checkpoint path and the fused ATOM path excluded from default quantization. + quant_default_exclude_layers: list[str] = [ + "*.indexer.weights_proj", + "*.indexer.wk_weights_proj", + ] class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM): """GLM 5.0 MoE (structurally similar to DeepSeek v3.2). Reuses DeepseekV2 implementation.""" # GLM-5's HF quant config uses `indexers_proj` in modules_to_not_convert, but - # the ATOM module path is `indexer.weights_proj`. Declaring the mapping here - # keeps the translation co-located with the model and out of config.py. + # the unfused ATOM module path is `indexer.weights_proj`. Keep that path + # excluded so FP4/MXFP4 fallback does not quantize the BF16 projection. quant_exclude_name_mapping: dict[str, str] = { # HF quant config uses "indexers_proj" but the ATOM module path is # "indexer.weights_proj". str.replace translates each exclude entry. diff --git a/atom/plugin/attention_mla_sparse.py b/atom/plugin/attention_mla_sparse.py index 2e48acf3d8..c63dd09fa0 100644 --- a/atom/plugin/attention_mla_sparse.py +++ b/atom/plugin/attention_mla_sparse.py @@ -26,8 +26,8 @@ cp_gather_indexer_k_quant_cache, dtypes, indexer_k_quant_and_cache, + indexer_qk_rope_quant_and_cache, top_k_per_row_decode, - top_k_per_row_prefill, ) from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits @@ -558,7 +558,7 @@ def sparse_attn_indexer_plugin_mode( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -568,6 +568,15 @@ def sparse_attn_indexer_plugin_mode( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: try: from vllm.forward_context import ( @@ -584,12 +593,12 @@ def sparse_attn_indexer_plugin_mode( # During profile/dummy run the metadata dict may not contain # our layer or may be None. if attn_metadata_dict is None: - return weights + return torch.zeros_like(weights, dtype=torch.float32) if k_cache_prefix not in attn_metadata_dict: - return weights + return torch.zeros_like(weights, dtype=torch.float32) layer_meta = attn_metadata_dict[k_cache_prefix] if layer_meta is None: - return weights + return torch.zeros_like(weights, dtype=torch.float32) # In plugin mode, plugin_metadata is vllmDeepseekV32IndexerMetadata from # AiterMLASparseIndexerMetadataBuilder. @@ -602,14 +611,43 @@ def sparse_attn_indexer_plugin_mode( kv_block_size = kv_cache.shape[1] preshuffle_cache = kv_block_size != 1 - indexer_k_quant_and_cache( - k, - kv_cache, - slot_mapping, - quant_block_size, - scale_fmt, - preshuffle=preshuffle_cache, - ) + if use_qk_rope_cache_fusion: + q_bf16 = q_input + q_fp8 = torch.empty_like(q_bf16, dtype=dtypes.fp8) + weights_out = torch.empty( + weights.shape, device=weights.device, dtype=torch.float32 + ) + indexer_qk_rope_quant_and_cache( + q_bf16, + q_fp8, + weights, + weights_out, + k, + kv_cache, + slot_mapping, + k_norm_weight, + k_norm_bias, + positions, + cos_cache, + sin_cache, + k_norm_eps, + quant_block_size, + scale_fmt, + weights_scale, + preshuffle=preshuffle_cache, + is_neox=is_neox_style, + ) + weights = weights_out + else: + q_fp8 = q_input + indexer_k_quant_and_cache( + k, + kv_cache, + slot_mapping, + quant_block_size, + scale_fmt, + preshuffle=preshuffle_cache, + ) topk_indices_buffer[: hidden_states.shape[0]] = -1 # topk_indices_buffer[: num_actual_tokens] = -1 @@ -738,7 +776,7 @@ def sparse_attn_indexer_fake( hidden_states: torch.Tensor, k_cache_prefix: str, kv_cache: torch.Tensor, - q_fp8: torch.Tensor, + q_input: torch.Tensor, k: torch.Tensor, weights: torch.Tensor, quant_block_size: int, @@ -748,6 +786,15 @@ def sparse_attn_indexer_fake( max_model_len: int, total_seq_lens: int, topk_indices_buffer: torch.Tensor, + k_norm_weight: torch.Tensor, + k_norm_bias: torch.Tensor, + k_norm_eps: float, + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + weights_scale: float, + is_neox_style: bool, + use_qk_rope_cache_fusion: bool, ) -> torch.Tensor: # profile run # NOTE(Chen): create the max possible flattened_kv. So that @@ -757,7 +804,7 @@ def sparse_attn_indexer_fake( ) _k_fp8 = _flattened_kv[..., :head_dim].view(torch.float8_e4m3fn).contiguous() _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous() - return weights + return torch.empty(weights.shape, device=weights.device, dtype=torch.float32) direct_register_custom_op( diff --git a/atom/utils/envs.py b/atom/utils/envs.py index fb487bc106..105f109749 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -50,6 +50,9 @@ "ATOM_ENABLE_DS_QKNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_QKNORM_FUSION", "1") == "1" ), + "ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION": lambda: ( + os.getenv("ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION", "1") == "1" + ), "ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION", "1") == "1" ), @@ -73,8 +76,9 @@ ), # Use a thread pool for weight loading instead of main-process sequential I/O. # Set to 0 to disable if the thread pool causes hangs (e.g. on gfx1250). - "ATOM_LOADER_USE_THREADPOOL": lambda: os.getenv("ATOM_LOADER_USE_THREADPOOL", "1") - == "1", + "ATOM_LOADER_USE_THREADPOOL": lambda: ( + os.getenv("ATOM_LOADER_USE_THREADPOOL", "1") == "1" + ), # --- Attention Backend --- # Use unified_attention (flash-style) for MHA paged/prefill attention instead # of pa_decode_gluon. Set to 1 to enable the unified_attention path. @@ -108,10 +112,9 @@ "ATOM_REQUIRES_GRAD": lambda: os.getenv("ATOM_REQUIRES_GRAD", "0") == "1", # --- Bpreshuffle for weight --- # Preshuffle weight. Default "1" (enabled) - "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE": lambda: os.getenv( - "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE", "1" - ) - == "1", + "ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE": lambda: ( + os.getenv("ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE", "1") == "1" + ), # --- V4 Attention Backend Refactor (PR-A: kill .item(), unlock CUDAGraph) --- # `legacy` (default) keeps the per-seq Python dispatch loop with .item() # syncs in deepseek_v4.py. `new` routes through V4AttentionBackend with From 044d089b4b6de74949c0ecc444c0e818ce8f1bd4 Mon Sep 17 00:00:00 2001 From: junyyang-amd Date: Fri, 15 May 2026 23:29:37 +0800 Subject: [PATCH 61/76] [ci] Add MTP cases and run the case according to P0/P1/P2 (#803) * modify logic with running p1 daily and p1/p2 on weekends * Monday-Thuesday: P0; Fridy: P1&P2; Weekend: C-ALL * modify code * modify the running time logic * Add MTP cases and run the case according to P0/P1/P2 * Merge 'cron' * Change the scheduling time --------- Co-authored-by: root --- .github/workflows/atom-vllm-benchmark.yaml | 125 +++++++++++++++++---- 1 file changed, 103 insertions(+), 22 deletions(-) diff --git a/.github/workflows/atom-vllm-benchmark.yaml b/.github/workflows/atom-vllm-benchmark.yaml index 19f66ee996..1261f2f362 100644 --- a/.github/workflows/atom-vllm-benchmark.yaml +++ b/.github/workflows/atom-vllm-benchmark.yaml @@ -7,8 +7,10 @@ concurrency: on: schedule: - # Nightly at 21:00 Beijing time (13:00 UTC) - - cron: '0 13 * * 0-4' + # Mon–Sat 21:00 Beijing (13:00 UTC): Mon/Wed/Fri AWS P0; Tue/Thu AWS P1+Meta P0; Sat AWS P2+OOB (see load-models). + - cron: '0 13 * * 1-6' + # Sat/Sun 10:00 Beijing (02:00 UTC): Sat AWS P1+Meta P0; Sun Meta P1 + - cron: '0 2 * * 0,6' workflow_dispatch: inputs: benchmark_client: @@ -884,6 +886,54 @@ jobs: for variant in family["variants"]: all_variants.append(flatten_variant(family, variant)) + def select_by_display(variants, displays): + return [ + m for m in variants if str(m.get("display", "")) in displays + ] + + AWS_P0 = { + "MiniMax-M2.5 TP2 (AW)", + "MiniMax-M2.5 TP4 (AW)", + "DeepSeek-V3.2 FP8 TP8 (AW)", + "gpt-oss-120b TP2 (AW)", + "gpt-oss-120b TP8 (AW)", + } + AWS_P1 = { + "GLM-4.7-FP8 TP4 (AW)", + "Kimi-K2.5-MXFP4 TP4 (AW)", + "Kimi-K2.5-MXFP4 TP8 (AW)", + "GLM-4.7-FP8 MTP TP4 (AW)", + "GLM-4.7-FP8 MTP TP8 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP1 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8-MTP TP4 (AW)", + } + AWS_P2 = { + "DeepSeek-V3.2 FP8 TP4 (AW)", + "GLM-4.7-FP8 TP8 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP2 (AW)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (AW)", + "gpt-oss-120b TP1 (AW)", + } + META_P0 = { + "Kimi-K2.5-MXFP4 TP4 (MET)", + "gpt-oss-120b TP1 (MET)", + } + META_P1 = { + "Kimi-K2-Thinking-MXFP4 TP4 (MET)", + "Kimi-K2-Thinking-MXFP4 TP8 (MET)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP1 (MET)", + "Qwen3-Next-80B-A3B-Instruct-FP8 TP4 (MET)", + "Qwen3.5-397B-A17B-FP8 TP4 (OOB)", + "DeepSeek-R1 FP8 TP8 (MET)", + "DeepSeek-R1 MXFP4 TP8 (MET)", + } + OOB_DISPLAYS = { + str(m.get("display", "")) + for m in all_variants + if str(m.get("display", "")).endswith("(OOB)") + } + if event == "schedule": created_at = os.environ.get("SCHEDULE_CREATED_AT", "") if not created_at: @@ -894,29 +944,60 @@ jobs: run_date = datetime.fromisoformat(created_at.replace("Z", "+00:00")) beijing_dt = run_date.astimezone(beijing_tz) beijing_weekday = beijing_dt.weekday() + beijing_hour = beijing_dt.hour + + # Beijing weekday: Mon=0 .. Sun=6. Hour buckets use 15:00 to split + # ~10:00 daytime vs ~21:00 evening scheduled runs. + # Mon/Wed/Fri evening: AWS P0. Tue/Thu evening: AWS P1 + Meta P0. + # Fri evening: AWS P0. Sat daytime (~10:00): AWS P1 + Meta P0. + # Sat evening (~21:00): AWS P2 + OOB. Sun daytime (~10:00): Meta P1. + is_daytime_slot = beijing_hour < 15 + + if beijing_weekday == 5: # Saturday + if is_daytime_slot: + selected_group = "AWS-P1+META-P0" + target_displays = AWS_P1 | META_P0 + else: + selected_group = "AWS-P2+OOB" + target_displays = AWS_P2 | OOB_DISPLAYS + elif beijing_weekday == 6: # Sunday + if is_daytime_slot: + selected_group = "META-P1" + target_displays = META_P1 + else: + selected_group = "SKIP" + target_displays = set() + elif beijing_weekday in (0, 2, 4): # Mon / Wed / Fri + if is_daytime_slot: + selected_group = "SKIP" + target_displays = set() + else: + selected_group = "AWS-P0" + target_displays = AWS_P0 + elif beijing_weekday in (1, 3): # Tue / Thu + if is_daytime_slot: + selected_group = "SKIP" + target_displays = set() + else: + selected_group = "AWS-P1+META-P0" + target_displays = AWS_P1 | META_P0 + else: + selected_group = "SKIP" + target_displays = set() - if beijing_weekday in (0, 2): - selected_group = "A-MET" - selected = [ - m for m in all_variants - if str(m.get("prefix", "")).endswith("-met") - ] - elif beijing_weekday in (1, 3): - selected_group = "B-AW" - selected = [ - m for m in all_variants - if "-aw-" in str(m.get("prefix", "")) - ] - elif beijing_weekday == 4: - selected_group = "C-ALL" + if target_displays is None: selected = list(all_variants) else: - selected_group = "SKIP-WEEKEND" - selected = [] - print( - f"Weekend schedule in Beijing ({beijing_dt.date()}); no benchmark group configured.", - file=sys.stderr, - ) + selected = select_by_display(all_variants, target_displays) + missing = target_displays - { + str(m.get("display", "")) for m in selected + } + if missing: + print( + "Warning: scheduled targets missing from catalog: " + + ", ".join(sorted(missing)), + file=sys.stderr, + ) else: variant_map = {} for family in families: From 88a3ff139b7fe9aa334117c74b19a1091ca5e66e Mon Sep 17 00:00:00 2001 From: honglie Date: Sat, 16 May 2026 00:34:20 +0800 Subject: [PATCH 62/76] Add benchmark commit override input (#801) * Add benchmark checkout ref inputs Allow manual benchmark runs to target a specific ATOM branch or commit while preserving the default scheduled behavior. Co-authored-by: Cursor * Simplify benchmark ref input Keep only the commit override input for manual benchmark runs and place it after the regular benchmark parameters. Co-authored-by: Cursor --------- Co-authored-by: root Co-authored-by: Cursor --- .github/workflows/atom-benchmark.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/.github/workflows/atom-benchmark.yaml b/.github/workflows/atom-benchmark.yaml index 289d116c05..ebb67b5415 100644 --- a/.github/workflows/atom-benchmark.yaml +++ b/.github/workflows/atom-benchmark.yaml @@ -82,6 +82,10 @@ on: Example (multiple sets): 1024,1024,128,0.8;2048,1024,256,0.7" type: string default: "1024,1024,128,0.8" + atom_commit: + description: "ATOM commit SHA to checkout; leave empty to use the workflow ref" + type: string + default: "" jobs: parse-param-lists: @@ -91,6 +95,8 @@ jobs: matrix_json: ${{ steps.parse-param-lists.outputs.matrix_json }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Parse parameter lists id: parse-param-lists @@ -122,6 +128,8 @@ jobs: models_json: ${{ steps.load.outputs.models_json }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - id: load env: EVENT_NAME: ${{ github.event_name }} @@ -186,6 +194,8 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Start CI container run: | @@ -349,6 +359,8 @@ jobs: steps: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Download all benchmark results uses: actions/download-artifact@v8 @@ -497,6 +509,8 @@ jobs: has_matrix: ${{ steps.gen.outputs.has_matrix }} steps: - uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - uses: actions/download-artifact@v8 with: @@ -552,6 +566,8 @@ jobs: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Docker Login run: echo "${{ secrets.DOCKER_PASSWORD }}" | docker login -u ${{ secrets.DOCKER_USERNAME }} --password-stdin @@ -744,6 +760,8 @@ jobs: steps: - name: Checkout ATOM repo uses: actions/checkout@v6 + with: + ref: ${{ inputs.atom_commit || github.ref }} - name: Download regression traces if: needs.collect-regression-traces.result == 'success' From 45c31f220ff51ec9b6526c4f7f326b8eecf01f6b Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Sat, 16 May 2026 17:18:56 +0900 Subject: [PATCH 63/76] Add inferencex-sync skill for ATOM benchmark comparison and InferenceX PR creation (#802) * Add inferencex-pr skill for ATOM benchmark comparison and PR creation Co-Authored-By: Claude Sonnet 4.6 * Replace inferencex-pr skill with inferencex-sync Renames the skill to better reflect its purpose: comparing ATOM upstream benchmark results against InferenceX and reporting regression/improvement. Co-Authored-By: Claude Opus 4.6 * pr flow update Signed-off-by: seungrokj * (skill)[inferencex-sync]: Clarify all-jobs-passed condition for latest run Co-Authored-By: Claude Opus 4.6 --------- Signed-off-by: seungrokj Co-authored-by: Claude Sonnet 4.6 --- .claude/commands/inferencex-sync.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/.claude/commands/inferencex-sync.md b/.claude/commands/inferencex-sync.md index 4f9f538884..5a82eb1956 100644 --- a/.claude/commands/inferencex-sync.md +++ b/.claude/commands/inferencex-sync.md @@ -15,6 +15,7 @@ Use this skill to compare the througput of ATOM upstream benchamrk and Inference - [ ] 1) Create a throughput table from InferenceX benchmark - [ ] 2) Create a throughput table from ATOM upstream benchmark - [ ] 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark +- [ ] 4) Ask users what model needs a new PR ``` ## 1) Create a throughput table from InferenceX benchmark @@ -36,7 +37,7 @@ Rules: ## 2) Create a throughput table from ATOM upstream benchmark -Check the ATOM upstream benchmark that is from the latest successful gh action runs at https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml and append the throughput to the InferenceX performance table in Step 1) +Check the ATOM upstream benchmark that is from the latest successful gh action runs at https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml, which as all jobs passed. Append the throughput to the InferenceX performance table in Step 1) Rules: - 1. Use main branch. Find the latest job of https://github.com/ROCm/ATOM/actions/workflows/atom-benchmark.yaml where all jobs passed @@ -65,4 +66,19 @@ If the data is not available, stop and ask for any workaround ## 3) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark -Update the performance table to show regression in % of ATOM upstream benchmark from InferenceX benchmark \ No newline at end of file +Update the performance table to show regression in % of ATOM upstream benchmark from InferenceX benchmark + +## 4) Chek performance regression of ATOM upstream benchmark from InferenceX benchmark + +Based on Step 3) results, ask user what models need a PR to https://github.com/SemiAnalysisAI/InferenceX/pulls and create a PR to each model. + +Rules: +- 1. Find the docker image from the step 2) ATOM upstream benchmark gh action runs. for eample, "Docker: + rocm/atom-dev:nightly_202605141645" from https://github.com/ROCm/ATOM/actions/runs/25875375446/job/76041129934#step:6:26 +- 2. Create a PR and check https://github.com/ROCm/ATOM/blob/main/.github/benchmark/models.json options + +PR example: +- 1. https://github.com/SemiAnalysisAI/InferenceX/pull/1311/changes +- 2. https://github.com/SemiAnalysisAI/InferenceX/blob/main/benchmarks/single_node/dsv4_fp4_mi355x_atom.sh +- 3. https://github.com/SemiAnalysisAI/InferenceX/blob/main/.github/configs/amd-master.yaml#L1646 +- 4. https://github.com/SemiAnalysisAI/InferenceX/blob/main/perf-changelog.yaml#L1824-L1833 From 0643080d80dfae675fc3848088e5a10bb802d2b3 Mon Sep 17 00:00:00 2001 From: wuhuikx Date: Sun, 17 May 2026 10:02:16 +0800 Subject: [PATCH 64/76] Fix DeepSeek-V3.2 indexer RoPE cache layout (#809) * [fix](deepseek): handle padded sparse indexer inputs Keep fused indexer cache writes limited to actual vLLM plugin tokens and pass RoPE caches in the 2D shape expected by the AITER indexer kernel. Co-authored-by: Cursor * [fix](deepseek): pass unsliced sparse indexer inputs Rely on AITER's mapped-token bounds for padded graph inputs so the sparse indexer plugin can pass the original tensors through both fused and fallback cache update paths. * [fix](deepseek): keep sparse indexer plugin unchanged Leave the plugin sparse indexer path unchanged so this branch only carries the DeepSeek-V3.2 RoPE cache layout fix. --------- Co-authored-by: XiaobingSuper Co-authored-by: Cursor --- atom/models/deepseek_v2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index a2b8b5af83..92145c1cf8 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1458,8 +1458,8 @@ def forward( self.k_norm.bias, self.k_norm.eps, positions, - rotary_emb.cos_cache, - rotary_emb.sin_cache, + rotary_emb.cos_cache.squeeze(-2).squeeze(-2), + rotary_emb.sin_cache.squeeze(-2).squeeze(-2), self._weights_scale, rotary_emb.is_neox_style, self.use_qk_rope_cache_fusion, From beae6cd4227201adb27d50f84753fedbb0ec6e2e Mon Sep 17 00:00:00 2001 From: Yutao Xu Date: Mon, 18 May 2026 14:40:52 +0800 Subject: [PATCH 65/76] Update minimax_m2.py (#820) --- atom/models/minimax_m2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atom/models/minimax_m2.py b/atom/models/minimax_m2.py index d254405648..d928e54840 100644 --- a/atom/models/minimax_m2.py +++ b/atom/models/minimax_m2.py @@ -232,7 +232,7 @@ def forward( # TP-aware RMSNorm: all-reduce variance across TP ranks so # normalization uses the global variance (over 6144/1024 dims) # rather than per-rank variance (768/128 dims). - if qkv.shape[0] <= 64 and self.tp_size > 1: + if qkv.shape[0] <= 256 and self.tp_size > 1: q, k, v = tensor_model_parallel_fused_qknorm_allreduce( qkv, self.q_norm.weight, self.k_norm.weight, self.rms_norm_eps ) From aa6d2c8b83d77235b76afec114ad65a98ceec3d4 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Mon, 18 May 2026 15:44:11 +0800 Subject: [PATCH 66/76] CI: extend aiter wheel artifact retention (#823) --- .github/workflows/atom-sglang-test.yaml | 2 +- .github/workflows/atom-test.yaml | 2 +- .github/workflows/atom-vllm-test.yaml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/atom-sglang-test.yaml b/.github/workflows/atom-sglang-test.yaml index 2a5d8572fb..da23aa4423 100644 --- a/.github/workflows/atom-sglang-test.yaml +++ b/.github/workflows/atom-sglang-test.yaml @@ -192,7 +192,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 atom-sglang-accuracy: needs: [check-signal, download_aiter_wheel] diff --git a/.github/workflows/atom-test.yaml b/.github/workflows/atom-test.yaml index 2e0e2fdbfb..83cdb49fc5 100644 --- a/.github/workflows/atom-test.yaml +++ b/.github/workflows/atom-test.yaml @@ -199,7 +199,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 load-test-models: name: Load test model configs diff --git a/.github/workflows/atom-vllm-test.yaml b/.github/workflows/atom-vllm-test.yaml index ca9d3e3dd8..732741c255 100644 --- a/.github/workflows/atom-vllm-test.yaml +++ b/.github/workflows/atom-vllm-test.yaml @@ -192,7 +192,7 @@ jobs: with: name: aiter-whl path: aiter-whl/amd_aiter*.whl - retention-days: 1 + retention-days: 7 atom-vllm-oot: needs: [check-signal, download_aiter_wheel] From aa7c25af978003c2293537adc11984d4374cbb18 Mon Sep 17 00:00:00 2001 From: PerryZhang01 Date: Mon, 18 May 2026 17:05:45 +0800 Subject: [PATCH 67/76] [fix](gpt-oss): fix quark quantized model in moe bias (#787) Co-authored-by: perzhang --- atom/model_ops/moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index e38388de46..b53aa73371 100644 --- a/atom/model_ops/moe.py +++ b/atom/model_ops/moe.py @@ -756,7 +756,7 @@ def create_weights( if layer.has_bias: w13_bias = atom_parameter( - torch.empty( + torch.zeros( num_experts, 2 * intermediate_size_per_partition_after_pad, dtype=torch.bfloat16, @@ -793,7 +793,7 @@ def create_weights( if layer.has_bias: w2_bias = atom_parameter( - torch.empty( + torch.zeros( num_experts, hidden_size, dtype=torch.bfloat16, From a076ab3557848c9c44a340ea1d478dba03dcad26 Mon Sep 17 00:00:00 2001 From: Jiayun Date: Mon, 18 May 2026 20:48:52 +0800 Subject: [PATCH 68/76] rm mtp no prefix cache assert (#824) --- atom/models/qwen3_5_mtp.py | 2 -- atom/models/qwen3_next_mtp.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/atom/models/qwen3_5_mtp.py b/atom/models/qwen3_5_mtp.py index 1ec4df8a46..73e8b39116 100644 --- a/atom/models/qwen3_5_mtp.py +++ b/atom/models/qwen3_5_mtp.py @@ -136,8 +136,6 @@ def remap_mtp_weight_name(self, name: str) -> str | None: def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() config = get_qwen3_5_text_config(atom_config) - if atom_config.enable_prefix_caching: - raise ValueError("Qwen3_5MTP currently does not support prefix caching") self.config = config # Reindex MTP exclude entries on a copy: checkpoint uses 0-based diff --git a/atom/models/qwen3_next_mtp.py b/atom/models/qwen3_next_mtp.py index 5c547d9559..f865aa2dbc 100644 --- a/atom/models/qwen3_next_mtp.py +++ b/atom/models/qwen3_next_mtp.py @@ -131,8 +131,6 @@ def remap_mtp_weight_name(self, name: str) -> str | None: def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() config = atom_config.hf_config - if atom_config.enable_prefix_caching: - raise ValueError("Qwen3NextMTP currently does not support prefix caching") self.config = config self.model = Qwen3NextMultiTokenPredictor( atom_config=atom_config, prefix=maybe_prefix(prefix, "mtp") From 6619cc7d33e3bb1f5a8f2f9966b2c4cb3a20e494 Mon Sep 17 00:00:00 2001 From: Zhiwei Date: Mon, 18 May 2026 22:39:01 +0800 Subject: [PATCH 69/76] [ATOM SGL] Qwen3.5 mha support (#819) * [ATOM SGL] Qwen3.5 mha support * precheckin --------- Co-authored-by: wuhuikx --- atom/plugin/prepare.py | 8 +- .../attention_backend/sgl_attn_backend.py | 268 +++++++++++++++--- 2 files changed, 227 insertions(+), 49 deletions(-) diff --git a/atom/plugin/prepare.py b/atom/plugin/prepare.py index 297d9faa95..3c7b722d29 100644 --- a/atom/plugin/prepare.py +++ b/atom/plugin/prepare.py @@ -84,13 +84,7 @@ def prepare_model(config: Any, engine: str): apply_prepare_model_adaptations(atom_config, model_arch) - # Qwen3-Next and Qwen3.5 series models keep the upstream attention backend path. - if model_arch not in { - "Qwen3NextForCausalLM", - "Qwen3_5ForConditionalGeneration", - "Qwen3_5MoeForConditionalGeneration", - }: - register_ops_to_sglang(atom_config=atom_config) + register_ops_to_sglang(atom_config=atom_config) set_attn_cls() # init aiter dist for using aiter custom collective ops diff --git a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py b/atom/plugin/sglang/attention_backend/sgl_attn_backend.py index ad6723570d..57883e32a1 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py +++ b/atom/plugin/sglang/attention_backend/sgl_attn_backend.py @@ -20,7 +20,10 @@ import sglang.srt.layers.attention.aiter_backend as _sglang_aiter from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend -from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.attention.utils import ( + create_flashinfer_kv_indices_triton, + launch_reshape_and_cache_flash, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import get_bool_env_var @@ -35,8 +38,10 @@ dtypes, get_pa_metadata_info_v1, get_pa_metadata_v1, + mha_batch_prefill_func, pa_fwd_asm, pa_persistent_fwd, + paged_attention_ragged, ) except ImportError as e: raise ImportError( @@ -380,20 +385,11 @@ def _init_decode_mha(self, bs, kv_indptr, kv_indices, forward_batch): if seq_lens_cpu is None: seq_lens_cpu = forward_batch.seq_lens.cpu() - page_table_persistent = self.page_table - seq_lens_persistent = self.seq_lens - seq_lens_persistent.fill_(0) - page_table_persistent.fill_(0) - seq_lens_persistent[:bs].copy_(forward_batch.seq_lens, non_blocking=True) - max_seq_pages = ( - seq_lens_cpu.max().item() + self.page_size - 1 - ) // self.page_size + 1 - page_table = self.req_to_token[ - forward_batch.req_pool_indices[:, None], - self.strided_indices[:max_seq_pages][None, :], - ] - page_table_persistent[:bs, :max_seq_pages].copy_( - page_table // self.page_size, non_blocking=True + page_table, seq_lens = self._update_decode_page_table( + bs, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + seq_lens_cpu=seq_lens_cpu, ) self.forward_metadata = ForwardMetadata( kv_indptr, @@ -402,8 +398,8 @@ def _init_decode_mha(self, bs, kv_indptr, kv_indices, forward_batch): None, 1, None, - page_table_persistent[:bs, :max_seq_pages], - seq_lens_persistent[:bs], + page_table, + seq_lens, ) self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) else: @@ -519,8 +515,11 @@ def _init_extend_mha(self, bs, forward_batch): self.indices_updater_prefill.kv_indices, self.qo_indptr[: bs + 1], None, - self.indices_updater_prefill.max_q_len, - self.indices_updater_prefill.max_kv_len, + self._max_len( + forward_batch.extend_seq_lens_cpu, + forward_batch.extend_seq_lens, + ), + self._max_len(forward_batch.seq_lens_cpu, forward_batch.seq_lens), None, forward_batch.seq_lens, ) @@ -844,12 +843,28 @@ def init_forward_metadata_capture_cuda_graph( if self.use_mla: self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) else: - page_table = self.page_table[:bs, :] - self.seq_lens[:bs].copy_(seq_lens, non_blocking=True) - seq_lens_persistent = self.seq_lens[:bs] - self.forward_metadata = ForwardMetadata( - None, + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + static_columns=True, + ) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, None, None, 1, @@ -878,35 +893,128 @@ def init_forward_metadata_replay_cuda_graph( if self.use_mla: self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) else: - page_table_persistent = self.page_table - seq_lens_persistent = self.seq_lens - seq_lens_persistent.fill_(0) - page_table_persistent.fill_(0) - seq_lens_persistent[:bs].copy_(seq_lens, non_blocking=True) - max_seq_pages = ( - seq_lens_cpu.max().item() + self.page_size - 1 - ) // self.page_size + 1 - page_table = self.req_to_token[ - req_pool_indices[:, None], - self.strided_indices[:max_seq_pages][None, :], - ] - page_table_persistent[:bs, :max_seq_pages].copy_( - page_table // self.page_size, non_blocking=True + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + seq_lens_cpu=seq_lens_cpu, + static_columns=True, ) self.forward_metadata = ForwardMetadata( - None, - None, + kv_indptr, + kv_indices, None, None, 1, None, - page_table_persistent[:bs, :max_seq_pages], + page_table, seq_lens_persistent[:bs], ) if self.decode_using_pa_ps: self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + @staticmethod + def _max_len(cpu_values, device_values) -> int: + values = cpu_values if cpu_values is not None else device_values + if isinstance(values, torch.Tensor): + return int(values.max().item()) + return int(max(values)) + + def _update_decode_page_table( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor] = None, + static_columns: bool = False, + ): + page_table_persistent = self.page_table + seq_lens_persistent = self.seq_lens + seq_lens_persistent.fill_(0) + page_table_persistent.fill_(0) + seq_lens_persistent[:bs].copy_(seq_lens, non_blocking=True) + + if seq_lens_cpu is None: + max_seq_len = int(seq_lens.max().item()) + elif isinstance(seq_lens_cpu, torch.Tensor): + max_seq_len = int(seq_lens_cpu.max().item()) + else: + max_seq_len = int(max(seq_lens_cpu)) + + max_seq_pages = (max_seq_len + self.page_size - 1) // self.page_size + 1 + max_seq_pages = min(max_seq_pages, page_table_persistent.shape[1]) + page_table = self.req_to_token[ + req_pool_indices[:, None], + self.strided_indices[:max_seq_pages][None, :], + ] + page_table_persistent[:bs, :max_seq_pages].copy_( + page_table // self.page_size, non_blocking=True + ) + + if static_columns: + return page_table_persistent[:bs, :], seq_lens_persistent[:bs] + return page_table_persistent[:bs, :max_seq_pages], seq_lens_persistent[:bs] + + def _should_use_native_dense_mha(self, layer) -> bool: + sliding_window_size = getattr(layer, "sliding_window_size", None) + return ( + not self.use_mla + and not layer.is_cross_attention + and layer.head_dim == 256 + and layer.qk_head_dim == 256 + and layer.v_head_dim == 256 + and (sliding_window_size is None or sliding_window_size <= -1) + ) + + def _kv_descales(self, layer): + if self.kv_cache_dtype != dtypes.fp8: + return None, None + k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale + return k_descale, v_descale + + def _get_aiter_paged_ragged_kv_cache_dtype(self) -> str: + if self.kv_cache_dtype != dtypes.fp8: + return "auto" + return "fp8_e4m3" + + def _set_kv_buffer_native_dense(self, layer, cache_loc, k, v, forward_batch): + k_descale, v_descale = self._kv_descales(layer) + if self.kv_cache_dtype == dtypes.fp8: + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + launch_reshape_and_cache_flash( + k.view(-1, layer.tp_k_head_num, layer.qk_head_dim), + v.view(-1, layer.tp_v_head_num, layer.v_head_dim), + k_cache.view( + -1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim + ), + v_cache.view(-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim), + cache_loc, + k_scale=k_descale, + v_scale=v_descale, + ) + return + + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v, k_descale, v_descale + ) + def set_kv_buffer_with_layout_shuffle( self, cache_loc, @@ -947,11 +1055,16 @@ def forward_extend(self, q, k, v, layer, forward_batch, save_kv_cache=True): if not layer.is_cross_attention else forward_batch.encoder_out_cache_loc ) + use_native_dense_mha = self._should_use_native_dense_mha(layer) if k is not None: assert v is not None if save_kv_cache: - if self.use_mla: + if use_native_dense_mha: + self._set_kv_buffer_native_dense( + layer, cache_loc, k, v, forward_batch + ) + elif self.use_mla: forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v) else: k_buffer, v_buffer = forward_batch.token_to_kv_pool.get_kv_buffer( @@ -970,9 +1083,44 @@ def forward_extend(self, q, k, v, layer, forward_batch, save_kv_cache=True): if self.use_mla: return self._forward_extend_mla(q, k, v, layer, forward_batch) + if use_native_dense_mha: + return self._forward_extend_native_dense_mha(q, layer, forward_batch) else: return self._forward_extend_mha(q, k, v, layer, forward_batch) + def _forward_extend_native_dense_mha(self, q, layer, forward_batch): + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + q_descale, k_descale, v_descale = None, None, None + + if self.kv_cache_dtype == dtypes.fp8: + q = q.to(dtypes.fp8) + q_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale + v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale + + bs0 = forward_batch.batch_size + 1 + o = mha_batch_prefill_func( + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache, + v_cache, + self.forward_metadata.qo_indptr[:bs0], + self.forward_metadata.kv_indptr[:bs0], + self.forward_metadata.kv_indices, + self.forward_metadata.max_q_len, + self.forward_metadata.max_kv_len, + causal=True, + logits_soft_cap=layer.logit_cap, + alibi_slopes=None, + return_lse=False, + return_attn_probs=False, + window_size=(-1, -1), + sink_ptr=None, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + def _forward_extend_mha(self, q, k, v, layer, forward_batch): """Non-MLA extend path: standard MHA with flash_attn_varlen_func.""" seqlens_in_batch = forward_batch.seq_lens @@ -1401,6 +1549,14 @@ def forward_decode( return o # Non-MLA decode paths + use_native_dense_mha = self._should_use_native_dense_mha(layer) + if use_native_dense_mha: + if save_kv_cache: + self._set_kv_buffer_native_dense( + layer, forward_batch.out_cache_loc, k, v, forward_batch + ) + return self._forward_decode_native_dense_mha(q, layer, forward_batch) + o = q.new_empty((batch_size, layer.tp_q_head_num, head_dim_out)) if save_kv_cache: @@ -1476,3 +1632,31 @@ def forward_decode( ) return o.view(-1, layer.tp_q_head_num * head_dim_out) + + def _forward_decode_native_dense_mha(self, q, layer, forward_batch): + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) + aiter_kv_str = self._get_aiter_paged_ragged_kv_cache_dtype() + + o = torch.empty_like(q, dtype=self.input_dtype) + paged_attention_ragged( + o.view(-1, layer.tp_q_head_num, layer.head_dim), + self.workspace_buffer, + q.view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache.view(-1, 1, layer.tp_k_head_num, layer.head_dim), + v_cache.view(-1, 1, layer.tp_v_head_num, layer.v_head_dim), + layer.scaling, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.kv_last_page_len, + 1, + self.max_num_partitions, + None, + aiter_kv_str, + "NHD", + layer.logit_cap, + self.k_scale, + self.v_scale, + None, + getattr(_sglang_aiter, "_AITER_PARTITION_SIZE_ROCM", 256), + ) + return o From 0bb75144e859529514ce2b0d7243944fca44ec95 Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Tue, 19 May 2026 16:52:08 +0800 Subject: [PATCH 70/76] feat(deepseek_v4): MTP-K skeleton + SWA race/aliasing fixes + per-PR MTP-1 CI (#817) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(deepseek_v4_mtp): wrapper-owned MTP with share_with_target hook Redesign V4 MTP to match V2/V3/Qwen convention: - Target (DeepseekV4ForCausalLM) carries no MTP modules / weights - Wrapper (DeepseekV4MTP) self-contains MTP blocks + loads mtp.* via the standard load_model path with spec_decode=True - EagleProposer.share_with_target() rebinds embed/head from the target's already-loaded instances; no double KV / double weight load Loader: replace per-call `not spec_decode` mtp short-circuit with `need_load_mtp = spec_decode and any("mtp" in n for n in params_dict)` so target loads and Eagle3 drafts (no `mtp.*` params) skip cleanly. EagleProposer: add `share_with_target` escape valve for models whose embed/head naming diverges from the standard `embed_tokens`/`lm_head` (V4 uses `model.embed`/`model.head`). ModelRunner: thread `extra_output_dims` so V4's `[N, hc_mult, dim]` residual-stack output (hc_head + LM head deferred to compute_logits) flows correctly into the outputs buffer. SpeculativeConfig: register `deepseek_v4` -> `deepseek_v4_mtp` mapping and stamp `model_type=deepseek_v4_mtp` onto the draft hf_config; the runner's is_deepseek_v4() now accepts both so V4-specific paths fire on the draft too. * fix(deepseek_v4): SWA+MTP cache aliasing, num_rejected rollback, EOS truncation - SWA ring buffer: cache_size = window_size + max_spec_steps to prevent draft-vs-verified slot aliasing when k>=1 spec tokens share a row's window. - Per-row window_topk indexing collapses case_a/b/c into a single formula: (pos - W + 1 + w) % cache_size with abs<0 / abs>pos masked to -1. - prepare_decode rolls back context_lens by num_rejected (mirrors aiter_mla), fixing stale ctx after MTP rejections. - scheduler.postprocess truncates bonus tokens emitted after EOS / stop_token / stop_sequence (rejection_sampler greedy kernel does not check EOS), removing trailing in MTP outputs. - Cleanup: drop dead _build_window_topk_batched, drop start_pos_per_token param, rename kernel param win -> cache_size for clarity. Validated: GSM8K 0.9545 (flex/strict), MTP-1 acceptance 86.27%. * wip * feat(deepseek_v4): MTP-K>1 support via prepare_mtp_decode + backend-agnostic eagle mid-step - eagle.propose mid-step: gate flat kv_indices/kv_indptr reads on `has_flat_kv` (hoisted out of the loop) so V4 path skips them; MLA/MHA unchanged - DeepseekV4AttentionMetadataBuilder.prepare_mtp_decode: per-draft-step V4 region metadata rebuild for 1-token-per-seq shape, zero D2H (mirrors eagle's GPU context_lens += 1 on CPU side instead of reading back). Entry short-circuits when max_per_req_cache_slots is unset (warmup case; attr is set by ModelRunner.get_num_blocks which runs after warmup_model). - CommonAttentionBuilder.__init__: centralize num_attention_heads (was duplicated in V4, MLA, MHA builders). - start_simple_inference.sh: drop env-var translation for engine flags, all engine config now passes through EXTRA_ARGS as native CLI args. - stop_atom_server.sh: also SIGTERM lm_eval so the next launch starts clean. Validated: V4 simple_inference (MTP1 + MTP3), V4 GSM8K full-1319 flexible-extract MTP3=0.9568 / MTP1=0.9545 (within stderr, no regression), acceptance distribution {0: 6.5%, 1: 24.1%, 2: 39.7%, 3: 29.7%} for MTP3, avg toks/fwd 2.92 vs 1.87 for MTP1 (+56%). * fix(deepseek_v4): wire write_v4_paged_decode_indices kernel + tooling - _attach_v4_paged_decode_meta: replace 6 transient tensors + 2 x index_copy_ chain with the existing write_v4_paged_decode_indices Triton kernel that reads only persistent forward_vars buffers. Eliminates allocator-churn race surfaced as ASSERT_TRAP in index_copy_kernel_impl> under MTP-3 long prefill. - _stage: harden silent astype auto-cast into an assertion so dtype drift fails fast at the call site instead of corrupting downstream triton reads. - scripts: move run_benchmark{,_sweep}.sh to atom/scripts/ and add wait_infer_drain.sh (hang vs drain monitor used during the V4 MTP investigation). - skills: convert flat .claude/skills/*.md into the SKILL.md/ subdirectory layout that the Skill tool actually loads; add new capture-trace skill (English, correct script paths, structured to match peer skills). * refactor(deepseek_v4): rename seq_lens→context_lens, require plan_buffers, drop dead device arg - `_build_compress_plans` + `make_compress_plans`: rename `seq_lens_np` / `seq_lens_cpu` → `context_lens_np` / `context_lens_cpu`. Every caller already passes `context_lens`; the old name suggested "raw seq length" while the actual semantics is "absolute seq_len AFTER the new extend tokens" (= prefix + extend). Internal `prefix = context_lens - extend_lens` reconstruction unchanged. - `make_compress_plans`: `plan_buffers` is now required (was `Optional[dict]`). The legacy fresh-`torch.from_numpy(...).to(device)` fallback is removed, along with the now-unused `device` positional arg. That fallback reproduces the exact allocator-churn race we fixed in `write_v4_paged_decode_indices` (transient tensors handed to in-flight kernels). Forcing every caller through the pre-allocated CpuGpuBuffer pool keeps data pointers stable and prevents the race from being re-introduced. - `_build_compress_plans`: keep the `assert isinstance(..., np.ndarray)` guards added in the same series so callers can't silently pass torch tensors and trigger a hidden D2H sync. Verified on V4-Pro no-MTP GSM8K (TP=8, level=0, CG on): HEAD c454322f: 0.95 flexible-extract this commit: 0.9560 flexible-extract (within 1σ; 1319/1319) * docs(skill): debug-agent-locate-kernel — `--enforce-eager`/`--level 0` are optional fallbacks Reword pre-flight item 2 + the matching anti-pattern + the example launch command. Previous text claimed both flags were required ("Always pass both"). That's wrong: the debug agent runs fine under hipgraph in most cases. The flags only become useful when the symptom points at graph mode or Inductor: - `--enforce-eager` for faults that don't reproduce in eager / capture- replay crashes under the agent's no-caching-allocator behavior; - `--level 0` for AMD Inductor `cluster_dims` autotune crashes at warmup. Adding them blindly hides graph-mode-only bugs (current branch: V4-Pro no-MTP server-mode prefill hangs ONLY under `--enforce-eager`; CG path works fine — we'd never have spotted that with the old "always pass both" rule). CLAUDE.md `run_debug_agent.sh` row updated in lockstep. * fix(deepseek_v4): gate alt_stream async-compress on CUDAGraph capture eager mode triggered a hipStream deadlock on V4-Pro: 60 layers of main Compressor launches accumulate on alt_stream, and the first splitk GEMM workspace allocation (the third bf16gemm_*_splitk_clean kernel load) hits a caching-allocator race with the MoE shared_experts GEMM also targeting alt_stream. Verified hang on both small (~800 token) and large (>2k token) prefill batches in eager. Inside a CUDAGraph capture block this same pattern is safe: the graph records the fork-join edges and replay re-uses the same stream layout without per-launch allocator contention. So the fix is to gate `use_async_compress` on a new `ForwardContext.in_hipgraph` flag set only by `model_runner.capture_model`. Replay does not re-execute Python forward, so the flag is irrelevant there. While here: - Cache `torch.cuda.current_stream()` once per fwd in `ForwardContext.main_stream` and have V4 attention / MoE / launcher read it instead of querying repeatedly (eliminated per-call handle allocation; metadata builder still queries directly because it runs before `set_forward_context`). - Module-level cache `_CUDA_AVAILABLE` so the cuda-availability check doesn't fire per `set_forward_context`. - Fix `wait_infer_drain.sh` drain detection: once the eval client is gone and a single poll shows no new output, declare drain immediately instead of waiting STUCK_POLLS polls (the old logic added ~120 s of false-positive wait to every run). Verified end-to-end: - V4-Pro eager + full GSM8K (1319 reqs): 420s, 0.9553/0.9560 (was deadlocking; now matches CG baseline) - V4-Pro CG mode + full GSM8K: 240s, 0.9492/0.9500 - V4-Pro CG + MTP1 + full GSM8K: 180s, 0.9575/0.9583 - V4-Pro CG + MTP3 still crashes on large prefill batch (pre-existing MEMORY_VIOLATION unrelated to these changes) * fix(deepseek_v4): race-free swa_write_indices + unify is_pure_decode source Under concurrent serving the swa_write kernel hit MEMORY_VIOLATION on `tl.load(kv_ptr + src_id * head_dim + d_offsets)` because the shared pinned `v4_meta_swa_write_indices` buffer was rewritten by the next fwd's CPU side BEFORE the previous fwd's async H2D DMA actually fired, producing torn `src_id` values that exceed `kv.shape[0]`. Captured under rocm-debug-agent: all faulting waves at PC +1144 (`s_waitcnt vmcnt(0)`) in `_swa_write_kernel`, confirming the prior async kv load as the OOB source. The fix writes into the pre-allocated shared GPU buffer with a static GPU iota source for the pure-decode arange path and a fresh local numpy for prefill — never touches the pinned `.np` alias whose next-fwd rewrite would tear the in-flight DMA. Co-changes: - Hoist `is_pure_decode` to AttentionMetaData_DSV4 construction time (single source of truth: prepare_decode / prepare_prefill / build_for_cudagraph_capture each declare their own semantics). `_attach_v4_per_fwd_meta` and `_attach_v4_paged_decode_meta` now read instead of recomputing. - Replace the old `(token_num_per_seq == 1).all()` arange-shortcut proxy with `is_pure_decode` (MTP verifier batches no longer take the slow concat path). - Tighten swa-write grid: prefill (eager) uses `num_write` exactly, decode/MTP CG uses the existing `padded_bs * (1+max_spec_steps)` bucket. Long-prefill chunks no longer launch up to ~64x sentinel- bail programs. - Tighten `v4_meta_swa_write_indices` from `mnbt` to `max_bs * win` (~16x smaller; universal worst case across paths). - Pre-allocate static GPU iota `self._swa_iota` as the H2D-free source for the is_pure_decode arange. - scripts/wait_infer_drain.sh: add offline simple_inference mode (process-exit + no-fault = drain) and detect "Memory access fault by GPU" alongside MEMORY_VIOLATION / ASSERT_TRAP. Verified: eager mode gsm8k 5-shot conc=65, 0 MEMORY_VIOLATION, exact_match flexible 0.9553 / strict 0.9545 (baseline ~0.95+). * ci(deepseek_v4): add MTP-1 per-PR accuracy + benchmark entries * fix(deepseek_v4): ruff lint + restore prep_stream sync + cache_size in fused swa_write - Drop unused `device` / `lru_cache` / first `prep_stream` assignment - Restore `prep_stream.wait_stream` / `with torch.cuda.stream(prep_stream)` around prepare_prefill H2D staging (suspected root for CG+MTP-3 tail deadlock at recv_mtp_status_async) - fused_qk_norm_rope_swa_write: pass `cache_size` instead of undefined `win` so the ring stride matches `swa_kv.shape[1]` under MTP * feat(benchmark): custom message encoder fallback + wait_infer_drain client-log support - benchmark_serving.py: route `--use-chat-template` through atom.entrypoints.openai.chat_encoders so models without a Jinja chat_template in tokenizer_config.json (e.g. DeepSeek-V4-Pro) render via the model-shipped custom encoder - wait_infer_drain.sh: treat LOG_FILE mtime growth as a progress signal alongside the engine "output send" marker, so passing a client log (e.g. benchmark.log with tqdm output) no longer triggers a false- positive HANG * refactor(deepseek_v4): inline window_topk + ragged-packed decode kv_indices Two coupled simplifications to the V4 paged-decode index build (decode/MTP path); GSM8K 3-shot 0.949 and 1k/1k c=64 MTP1 4247 tok/s match pre-refactor. 1) Drop the [mnbt, win] `v4_meta_window_topk` CG buffer (~4 MB) and the CPU `_build_window_topk_np` function. `write_v4_paged_decode_indices` now derives `n = min(positions[t]+1, win)` and `ring = (pos - n + 1 + i) % cs` inline from `var["positions"].gpu` — no intermediate. 2) Decode `kv_indices_swa/csa/hca` now use ragged-packed per-token slot counts (`actual_swa_count[t] + n_compress[t]`), matching the prefill layout instead of the prior uniform `win + n_compress` with -1 sentinel padding. `skip_prefix_len_csa[t]` is now `actual_swa_count[t]` in both paths (was hard-coded `win` in decode). * refactor(deepseek_v4): metadata int32 sweep + drain log auto-discovery + run-atom-workload skill deepseek_v4_attn.py: - Drop all unnecessary int64 in CPU metadata paths (cu_seqlens, indptr cumsums, segment indices, n_committed_{csa,hca}_per_seq). Keeps int64 only where GPU fancy-index ABI requires it (batch_id_per_token, swa_write_indices). Saves widen/narrow churn; no perf change but cleaner mental model. - Hoist n_committed_{csa,hca}_per_seq_cpu as dataclass single source of truth, replacing 3 independent ctx//k recomputes in paged_decode_meta / paged_prefill_meta / v4_indexer_meta. - Drop dead `start_pos_per_seq_cpu` dataclass field — only one real consumer (_build_paged_prefill_meta); inline as local var, drop 3 dead setters (prepare_decode, CG capture, prepare_mtp_decode). - Drop dead `positions_np` + `start_pos_per_seq_cpu` parameters from _attach_v4_per_fwd_meta; update 5 call sites. - Drop `cu_seqlens_q_np` + `start_pos_per_seq` params from _attach_v4_indexer_meta (read from attn_metadata.n_committed_csa_per_seq_cpu). - Drop dead helper `_clear_v4_paged_decode_meta` (defaults already None). - Drop dead CpuGpuBuffers: v4_meta_start_pos_per_seq + v4_meta_token_num_per_seq (allocated but never staged). - Drop redundant `if True:` wrapper + 3 redundant `import numpy as np` inside methods (module-level import suffices). - net -178 lines, ruff clean. scripts/wait_infer_drain.sh (v1.2): - Auto-discover server log via readlink /proc//fd/1 of the atom.entrypoints process. Eliminates the recurring bug where callers passed the wrong LOG_FILE (e.g. lm_eval's silent gsm8k_eval.log) causing false HANG verdicts. - Caller-supplied LOG_FILE is now optional (was: defaulted to hardcoded /app/logs_claude/atom_server.log). Used only as supplementary signal for fault grep (dual-log scan) and tqdm-style mtime detection. - Path-portable across repo layouts; no hardcoded log paths. .claude/skills/run-atom-workload/SKILL.md (new): - Codifies the canonical 4-step ATOM workload flow (stop → start → workload-in-shell-bg → wait_infer_drain → stop) for accuracy eval (GSM8K), benchmark, sweep, offline simple_inference, and fault repro under rocm-debug-agent. - Pins model-family env vars (V4-Pro: AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1; Kimi: HSA_NO_SCRATCH_RECLAIM=1; etc.). - Hard rules block past failure modes: no wrapper scripts in /app/logs_claude/, no && chaining, no double-backgrounding the server, always step 5 teardown. - Uses project-relative scripts/ paths (repo-portable). Verification: - MTP-1 GSM8K 3-shot = 0.9492 +/- 0.006 (baseline 0.9545 +/- 0.0057, within 1 sigma, no regression). - MTP-1 benchmark 1024/1024 c=64 = 4748 Total tok/s, 24.3 ms TPOT (vs prior baseline 4666 tok/s, +1.8% / no regression). - Drain auto-discovery smoke-tested on live MTP-1 server: discovered /app/logs_claude/atom_server.log via /proc, detected DRAINED cleanly. - MTP-3 + CG still hangs on high-concurrency GSM8K (known issue per feedback_v4_cg_mtp_status_deadlock.md); this commit does not address it. --- .../SKILL.md} | 0 .claude/skills/capture-trace/SKILL.md | 208 ++++ .../skills/debug-agent-locate-kernel/SKILL.md | 248 ++++ .../SKILL.md} | 0 .claude/skills/run-atom-workload/SKILL.md | 239 ++++ .github/benchmark/models.json | 10 + .github/benchmark/models_accuracy.json | 12 + atom/benchmarks/benchmark_serving.py | 22 +- atom/config.py | 10 +- atom/model_engine/engine_core.py | 9 + atom/model_engine/model_runner.py | 25 +- atom/model_engine/scheduler.py | 169 ++- atom/model_loader/loader.py | 12 +- atom/model_ops/attentions/aiter_attention.py | 4 +- atom/model_ops/attentions/aiter_mla.py | 6 +- atom/model_ops/attentions/backends.py | 8 + atom/model_ops/attentions/deepseek_v4_attn.py | 1040 +++++++++-------- atom/model_ops/sampler.py | 1 - atom/model_ops/v4_kernels/__init__.py | 6 + atom/model_ops/v4_kernels/compress_plan.py | 118 +- atom/model_ops/v4_kernels/fused_compress.py | 22 +- .../v4_kernels/paged_decode_indices.py | 235 ++++ atom/model_ops/v4_kernels/state_writes.py | 72 +- atom/models/deepseek_v4.py | 229 ++-- atom/models/deepseek_v4_mtp.py | 267 +++++ atom/spec_decode/eagle.py | 35 +- atom/utils/forward_context.py | 29 + scripts/run_benchmark.sh | 92 ++ scripts/run_benchmark_sweep.sh | 35 + scripts/run_debug_agent.sh | 82 ++ scripts/run_gsm8k_eval.sh | 56 + scripts/start_atom_server.sh | 115 ++ scripts/start_simple_inference.sh | 132 +++ scripts/stop_atom_server.sh | 92 ++ scripts/wait_infer_drain.sh | 191 +++ scripts/wait_server_ready.sh | 41 + tests/conftest.py | 1 + tests/test_mxfp4_moe_has_bias.py | 35 +- tests/test_per_req_cache_decoupling.py | 1 + 39 files changed, 3129 insertions(+), 780 deletions(-) rename .claude/skills/{atom-patterns.md => atom-patterns/SKILL.md} (100%) create mode 100644 .claude/skills/capture-trace/SKILL.md create mode 100644 .claude/skills/debug-agent-locate-kernel/SKILL.md rename .claude/skills/{dump-bisect-debug.md => dump-bisect-debug/SKILL.md} (100%) create mode 100644 .claude/skills/run-atom-workload/SKILL.md create mode 100644 atom/model_ops/v4_kernels/paged_decode_indices.py create mode 100644 atom/models/deepseek_v4_mtp.py create mode 100755 scripts/run_benchmark.sh create mode 100755 scripts/run_benchmark_sweep.sh create mode 100755 scripts/run_debug_agent.sh create mode 100755 scripts/run_gsm8k_eval.sh create mode 100755 scripts/start_atom_server.sh create mode 100755 scripts/start_simple_inference.sh create mode 100755 scripts/stop_atom_server.sh create mode 100755 scripts/wait_infer_drain.sh create mode 100755 scripts/wait_server_ready.sh diff --git a/.claude/skills/atom-patterns.md b/.claude/skills/atom-patterns/SKILL.md similarity index 100% rename from .claude/skills/atom-patterns.md rename to .claude/skills/atom-patterns/SKILL.md diff --git a/.claude/skills/capture-trace/SKILL.md b/.claude/skills/capture-trace/SKILL.md new file mode 100644 index 0000000000..adff103055 --- /dev/null +++ b/.claude/skills/capture-trace/SKILL.md @@ -0,0 +1,208 @@ +--- +name: capture-trace +description: Capture a PyTorch profiler / kineto trace from a running ATOM server for a short benchmark window. Use when the user asks for "a trace", "profiler trace", "GPU trace", or "抓 trace" for performance investigation — what kernels ran, what's on the critical path, what's slow. Do NOT use for crashes (use debug-agent-locate-kernel) or numerical bugs (use dump-bisect-debug). +version: 1.1.0 +scope: ATOM on AMD ROCm (PyTorch kineto profiler, per-rank `*.pt.trace.json.gz`) +last_updated: 2026-05-17 +--- + +## When to use + +- User asks for a trace, profiler dump, or kineto dump +- Performance analysis: "what's eating the time", "is decode fused", "did kernel X get called", "is this kernel on the critical path" +- Verifying that a code path was actually exercised at runtime (search by kernel name) + +Do NOT use this skill for: + +- Crashes / `Memory access fault` → [[debug-agent-locate-kernel]] +- Wrong outputs / accuracy regression → [[dump-bisect-debug]] +- Suspicion of a hang (no progress) → `scripts/wait_infer_drain.sh` first + +## Critical pre-flight + +1. **Stop the existing server cleanly** — the profiler argument has to be on the launch command line. `start_atom_server.sh` auto-kills the prior atom workers, so just relaunching with the new args is enough. +2. **Pick a SHORT workload** — a trace from a long run is unreadable and OOMs the profiler exporter. Default to `CONC * 1` requests (one prompt per concurrent slot). Never use the production `PROMPT_MULTIPLIER=10` default for a profiling run. +3. **`ATOM_PROFILER_MORE` belongs on the server, not the benchmark client.** The profiler runs inside the model-runner worker processes; an env on the bench client does nothing. +4. **Trace dir must be empty** for a clean per-rank layout. `start_atom_server.sh` does NOT clear it — `rm -rf $TRACE_DIR` before relaunch if you're iterating. + +## Required tools + +```bash +ls /app/ATOM/scripts/start_atom_server.sh # launcher +ls /app/ATOM/scripts/run_benchmark.sh # bench driver (passes --profile when PROFILE=1) +ls /app/ATOM/scripts/wait_server_ready.sh # ready-poll +python3 -c "import torch.profiler" # kineto present +``` + +## Parameters + +Pull these out of the user's request; everything except `MODEL` has a sensible default. + +| Param | Meaning | Typical | +|---|---|---| +| `MODEL` | Model path under `/data/` | `/data/DeepSeek-V4-Pro` | +| `TP` | Tensor-parallel size | `8` (4 for Kimi, 1 for gpt-oss-120b) | +| `ISL` / `OSL` | Random input / output length | `1024 / 1024` | +| `CONC` | Concurrency the bench keeps in flight | `64` or `128` | +| `PROMPT_MULTIPLIER` | Total prompts = `CONC * this` | **`1` for trace runs** (override the script default of 10) | +| `ATOM_PROFILER_MORE` | `1` = shapes + stack + memory (large traces, OOM risk); `0` = kernel-name only | **`0`** unless asked | +| `TRACE_DIR` | Where the kineto `.pt.trace.json.gz` lands | `/app/logs_claude/traces/` | +| `EXTRA_ARGS` | Forwarded to the openai server (MTP, kv-cache, etc.) | See [[atom-patterns]] | + +## Workflow + +### Step 1: Launch the server with the profiler bound + +```bash +TRACE_DIR=/app/logs_claude/traces/ +mkdir -p "$TRACE_DIR" + +# ATOM_PROFILER_MORE on the server env — not the client. +ATOM_PROFILER_MORE=0 \ + bash /app/ATOM/scripts/start_atom_server.sh \ + "$MODEL" "$TP" 8000 \ + --torch-profiler-dir "$TRACE_DIR" \ + $EXTRA_ARGS +``` + +`start_atom_server.sh` blocks until either `Server is ready!` or `Server process died`. Check the tail line; if it died, no point profiling. + +### Step 2: Drive a SHORT bench with `PROFILE=1` + +```bash +bash /app/ATOM/scripts/run_benchmark.sh \ + "$MODEL" 8000 "$ISL" "$OSL" "$CONC" \ + 1 \ # PROMPT_MULTIPLIER — keep this at 1 for traces + 1 \ # PROFILE=1 flag + $BENCH_EXTRA_ARGS +``` + +Position 6 is `PROMPT_MULTIPLIER`; position 7 is `PROFILE`. The bench sends a `start` HTTP call before the run and a `stop` call after, which is what trips the kineto exporter on the server. + +### Step 3: Wait for the exporter to finish (asynchronous on the server) + +Kineto exports lazily on the worker side — `stop_profiler` returns immediately to the bench, but the per-rank `.json` write + gzip can take 10-60 seconds. Poll the output dir: + +```bash +for i in $(seq 1 60); do + GZ=$(find "$TRACE_DIR" -name "*.pt.trace.json.gz" | wc -l) + JSON=$(find "$TRACE_DIR" -name "*.pt.trace.json" -not -name "*.gz" | wc -l) + echo "[t=${i}0s] gz=$GZ json=$JSON" + # Done = expected gz count AND no orphan .json (the .json is deleted after gzip) + [ "$GZ" -ge "$TP" ] && [ "$JSON" -eq 0 ] && break + sleep 10 +done +``` + +The completion signal is **`.gz` exists AND the same-name `.json` is gone**. File size of the `.gz` alone is unreliable (per `feedback_trace_gz_truncated.md`) — the exporter writes the raw `.json`, then gzip + unlink, so an orphan `.json` means it crashed mid-export. + +### Step 4: Verify the layout + +```bash +find "$TRACE_DIR" -type f | xargs ls -la +``` + +Expected: + +- `/rank_0/`, `rank_1/`, …, `rank_/` — one dir per rank +- Each dir has exactly one `*.pt.trace.json.gz` +- `ATOM_PROFILER_MORE=0`: ~50-80 MB per rank +- `ATOM_PROFILER_MORE=1`: ~200-300 MB per rank + +### Step 5: Inspect + +For a quick "did kernel X run" check: + +```bash +zcat "$TRACE_DIR"/rank_0/*.gz | python3 -c " +import json, sys +events = json.load(sys.stdin)['traceEvents'] +names = {e.get('name','') for e in events} +for kw in ['', ...]: + hits = sorted(n for n in names if kw in n) + print(f'{kw}: {len(hits)} matches') + for h in hits[:5]: print(f' {h}') +" +``` + +For counts and aggregate time per kernel: + +```bash +zcat "$TRACE_DIR"/rank_0/*.gz | python3 -c " +import json, sys +events = json.load(sys.stdin)['traceEvents'] +def stat(kw): + m = [e for e in events if kw in e.get('name','') and 'dur' in e] + if m: print(f'{kw}: count={len(m)} total_us={sum(e[\"dur\"] for e in m)}') +stat('aiter::topk_softplus') +stat('aiter::moe_forward') +# ... +" +``` + +For a UI view, drop the `.gz` (decompressed `.json`) into or `chrome://tracing`. + +## `record_function` tag format + +ATOM annotates the critical path with `torch.profiler.record_function`. The tags follow a stable format used by `parse_trace.py`: + +- Prefill: `prefill[bs= tok= ctx=]` +- Decode: `decode[bs= tok= d= spec=]` +- Draft (eagle mid-step): `draft[i/k bs=]` + +Searching by these tags is far more reliable than searching by kernel name (which varies across PyTorch/Triton/AITER versions). + +## `ATOM_PROFILER_MORE` cost + +`ATOM_PROFILER_MORE=1` enables `record_shapes + with_stack + profile_memory`. This multiplies trace size ~3-4x and, more importantly, the **C++ kineto aggregation at `stop_profiler` time scales with the recorded event count × TP**. On 8-rank V4-Pro, a 60-second profile with `PROFILER_MORE=1` has been observed to OOM the worker processes during export. Rules of thumb: + +- Default `=0` (kernel names + durations only — enough for 95% of investigations) +- `=1` only when you specifically need shapes or Python stacks AND you've kept the window short (≤ 5 seconds of bench traffic, ≤ `CONC * 1` prompts) + +## Common model configs + +| Model | Path | TP | Server `EXTRA_ARGS` | Bench notes | +|---|---|---|---|---| +| DeepSeek-V4-Pro | `/data/DeepSeek-V4-Pro` | 8 | `--level 0` (CG not yet enabled) | tokenizer has no chat template — do NOT pass `--use-chat-template` | +| DeepSeek-V4-Pro MTP1 | `/data/DeepSeek-V4-Pro` | 8 | `--level 0 --method mtp --num-speculative-tokens 1` | same — no chat template | +| DeepSeek-R1-0528 | `/data/DeepSeek-R1-0528` | 8 | | | +| DeepSeek-R1-0528 MTP3 | `/data/DeepSeek-R1-0528` | 8 | `--method mtp --num-speculative-tokens 3` | `--use-chat-template` REQUIRED on bench | +| DeepSeek FP4 MTP3 | `/data/DeepSeek-R1-0528-MXFP4-MTP-MoEFP4` | 8 | `--method mtp --num-speculative-tokens 3` | `--use-chat-template` REQUIRED | +| GLM-5 FP8 | `/data/GLM-5-FP8` | 8 | | | +| gpt-oss-120b | `/data/openai/gpt-oss-120b` | 1 | drop `-tp` | | +| Kimi-K2.5-MXFP4 | `/data/Kimi-K2.5-MXFP4` | 4 | `--trust-remote-code` + `HSA_NO_SCRATCH_RECLAIM=1` env | | + +## Anti-patterns + +- **Setting `ATOM_PROFILER_MORE=1` on the bench client.** The profiler runs in the model-runner workers; client env is ignored. Set it on the server launch. +- **Using `PROMPT_MULTIPLIER=10` (the bench default) for a trace run.** That sends 10× `CONC` requests and produces a multi-GB trace that takes forever to export and can OOM the workers. +- **Trusting the `.gz` size to decide if the export finished.** Use `.gz exists AND no orphan .json`. The exporter writes `.json` first, then gzips + unlinks; a leftover `.json` means it crashed mid-export. +- **Forgetting `--use-chat-template` on MTP DeepSeek-R1 benchmarks.** Tokenizer mismatch silently degrades accuracy — the trace will look fine but the workload is wrong. +- **Adding `--mark-trace` or `ENABLE_TORCH_PROFILER=1`.** Neither is needed — `--torch-profiler-dir` on the server + `PROFILE=1` on the bench is the complete handshake. The extras either no-op or interfere. +- **Profiling under `--level 3` for V4-Pro right now.** V4-Pro Inductor + autotune hits a `cluster_dims` bug on AMD — start with `--level 0` until CG support lands. + +## Real example — May 17 2026 + +Captured an MTP-1 vs no-MTP V4-Pro trace at `1024/1024/c=64`, 64 prompts each, to compare whether `_hash_topk` uses the fused `aiter::topk_softplus` kernel or the native `softplus + sqrt` chain: + +```bash +mkdir -p /app/logs_claude/traces/v4_mtp1 +bash /app/ATOM/scripts/start_atom_server.sh /data/DeepSeek-V4-Pro 8 8000 \ + --level 0 --method mtp --num-speculative-tokens 1 \ + --torch-profiler-dir /app/logs_claude/traces/v4_mtp1 +bash /app/ATOM/scripts/run_benchmark.sh /data/DeepSeek-V4-Pro 8000 1024 1024 64 1 1 +# (wait loop on .gz / .json — see Step 3) +zcat /app/logs_claude/traces/v4_mtp1/rank_0/*.gz | python3 -c "..." +``` + +Finding: 12,938 `aiter::topk_softplus` calls (the `elif scoring_func == 'sqrtsoftplus'` +branch in `moe.py:2637` IS taken) plus 12,658 native `softplus_kernel` calls +(from `_hash_topk` in `deepseek_v4.py:1982` — the first 3 hash layers don't go +through the fused path). Trace files: `/app/logs_claude/traces/v4_mtp1/` and +`/app/logs_claude/traces/v4_nomtp/`. + +## Cross-references + +- [[debug-agent-locate-kernel]] — when the server crashes or hangs, this skill is the wrong tool +- [[dump-bisect-debug]] — when the trace shows correct kernels but outputs are wrong +- [[atom-patterns]] — V4 attention buffer/stream conventions referenced from trace kernel names diff --git a/.claude/skills/debug-agent-locate-kernel/SKILL.md b/.claude/skills/debug-agent-locate-kernel/SKILL.md new file mode 100644 index 0000000000..968d0dc490 --- /dev/null +++ b/.claude/skills/debug-agent-locate-kernel/SKILL.md @@ -0,0 +1,248 @@ +--- +name: debug-agent-locate-kernel +description: Use rocm-debug-agent to identify which GPU kernel is faulting/hanging when ATOM server hits a HIP memory access fault, MEMORY_VIOLATION, or silent infinite loop. The agent dumps wave registers, faulting PC, and (with --save-code-objects) the disassembled code object so you can name the exact kernel and trace it back to the Python call site. Use when: server crashes with "Memory access fault by GPU node-N", server hangs with GPU at 100% but no token output, or you need to identify a kernel asserting `s_trap`. Do NOT use for: numerical bugs (use dump-bisect-debug), compile errors, OOM. +version: 1.0.0 +scope: ATOM on AMD ROCm (debug-agent at /opt/rocm/lib/librocm-debug-agent.so.2) +last_updated: 2026-05-14 +--- + +## When to use + +Symptoms that point to this skill: + +- `Memory access fault by GPU node-N (Agent handle: 0x...) on address 0x...` in `atom_server.log` +- Server alive (`curl /v1/models` returns) but `rocm-smi --showuse` shows 100% GPU and no `Engine Core: output send` for >30s — silent kernel hang / livelock +- Workers stuck at `torch.cuda.synchronize()` per `py-spy dump --pid ` — prior kernel never completes +- `HIP_LAUNCH_BLOCKING=1` makes the bug disappear → you have an async race; agent will tell you which kernel +- Reproduces only at certain batch shapes (e.g. MTP-3 + long prefill) + +Do NOT use this skill for: precision bugs (use [[dump-bisect-debug]]), build/compile errors (use `/build-fix`), OOM. + +## Required tools (verify before starting) + +```bash +ls /opt/rocm/lib/librocm-debug-agent.so.2 # the agent itself +ls /opt/rocm/llvm/bin/llvm-objdump # for disassembling code objects +which py-spy && py-spy --version # for stack traces of stuck workers +``` + +If any missing: install `rocm-debug-agent`, `llvm`, `pip install py-spy`. Stop here if not available. + +## Critical pre-flight + +1. **`ulimit -c 0`** — disables gpucore dumps. ROCm fault dumps gpucore files of 30-50 GB each per rank; on 8-GPU TP this fills disk in seconds. The launcher script sets this for you. +2. **`--enforce-eager` / `--level 0` — optional fallbacks, not required.** Try the default launch first; the debug agent runs fine under hipgraph in most cases. Only reach for these flags when symptoms point at graph mode: + - **`--enforce-eager`** disables CUDAGraph capture. Try this when the agent reports faults that don't reproduce in eager mode, or when capture/replay itself crashes under the agent's no-caching-allocator behavior. + - **`--level 0`** disables Inductor. Try this on AMD when you hit the `cluster_dims` autotune bug or other Inductor-side crashes during warmup. + - They are independent — apply only the one(s) the symptom points at. The launcher script does NOT inject either; pass via `EXTRA_ARGS` when you want them. +3. **Clean GPU state** — kill any prior `spawn_main`/`openai_server` processes. Stale `KFD process` entries (`rocm-smi --showpids` showing UNKNOWN PIDs holding VRAM) cause the next launch to OOM at NCCL barrier. `scripts/start_atom_server.sh` does the standard cleanup. +4. **Model-specific env** — pass on the command line or export before calling. Examples: + - V4-Pro requires `ATOM_USE_TRITON_MOE=1` + - Kimi-K2.5-MXFP4 requires `--trust-remote-code` + `HSA_NO_SCRATCH_RECLAIM=1` + +## Launcher scripts (in repo) + +| Script | Purpose | +|--------|---------| +| `scripts/start_atom_server.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...]` | Standard launcher: clears GPU, clears compile cache, backgrounds server, redirects to `atom_server.log`. | +| `scripts/stop_atom_server.sh` | SIGTERM atom.entrypoints, force-kill spawn workers, wait for VRAM release. | +| `scripts/run_debug_agent.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...]` | Wraps `start_atom_server.sh` with `HSA_TOOLS_LIB=librocm-debug-agent.so.2 + --save-code-objects`. Server output goes to `atom_server.log`; code objects land in `/app/logs_claude/debug_run/`. | +| `scripts/run_debug_agent.sh --simple [MODEL] [TP] [EXTRA_ARGS...]` | Same wrapper but invokes `start_simple_inference.sh` (offline, no port). Default log: `/app/logs_claude/simple_inference_debug_agent.log` (override via `LOG_FILE=`). Use for offline batch repros (e.g. V4 MTP-3 prefill hang). | +| `scripts/wait_server_ready.sh [PORT] [MAX_MIN] [POLL] [LOG_FILE]` | Poll `/v1/models` until ready or startup error detected. Allow MAX_MIN ≥ 5 under the agent (3-5× slower than normal). | + +## Workflow + +### Step 1: Reproduce under the agent + +```bash +bash scripts/stop_atom_server.sh # ensure clean +ATOM_USE_TRITON_MOE=1 \ + bash scripts/run_debug_agent.sh \ + /data/DeepSeek-V4-Pro 8 8000 \ + --method mtp --num-speculative-tokens 3 & +# If launch fails / faults look graph-mode-specific, retry with +# `--enforce-eager` (and `--level 0` on AMD for the Inductor cluster_dims bug) +# appended to EXTRA_ARGS. + +bash scripts/wait_server_ready.sh 8000 5 30 # 2-4 min under agent +cd /app/logs_claude && python .py # smallest hang trigger +``` + +Server load is **3-5× slower** under the debug agent. Expect ready at 2-4 min, repro at 30-90s after first big batch. + +### Step 2: Find the fault wave dump + +```bash +grep -E "stopped, reason|Memory access fault|MEMORY_VIOLATION|Disassembly" \ + /app/logs_claude/atom_server.log | head -20 +``` + +Each fault produces a block like: + +``` +wave_27876: pc=0x7f20f5e534c4 (kernel_code_entry=0x7f20f5e52900 ) (stopped, reason: ) + +scalar registers: ... +vector registers: ... ← v0..v? show per-lane values; v6 often holds index values being processed +trap registers: ... +general registers: pc=... + +Disassembly for function : + code object: memory://#offset=&size= + loaded at: [-] + => : +``` + +The `` is the demangled kernel name. **That's the suspect kernel.** Common cases: + +| Kernel name fragment | What it actually is | +|----------------------|---------------------| +| `at::native::index_copy_kernel_impl>` | `Tensor.index_copy_(dim, idx, src)` for dtype with N-byte size (4=int32/float32, 8=int64/float64) | +| `at::native::scatter_kernel` | `Tensor.scatter_(dim, idx, src)` | +| `at::native::index_kernel_impl` | Advanced indexing READ `tensor[idx]` | +| `_swa_write_kernel` / `_update_compressor_states_kernel` | ATOM Triton kernel — name in `state_writes.py` | + +### Step 3: Read the trap reason + +| reason | what it means | +|--------|---------------| +| `ASSERT_TRAP` | Kernel hit `s_trap 2` — almost always a `CUDA_KERNEL_ASSERT(...)` failed device-side. For PyTorch `index_copy_`/`scatter_` this is the bound check `0 <= idx < self.size(dim)`. Recompile PyTorch with `TORCH_USE_HIP_DSA=1` for the assert text — usually unavailable, infer from kernel name. | +| `MEMORY_VIOLATION` | Real OOB load/store. The `pc` instruction is the access; back-trace the address from `s_*`/`v_*` registers. | +| `INVALID_OPCODE` | Corrupted code object — usually an allocator stomp on the kernel binary (very rare). | + +### Step 4: Disassemble the code object + +The trap dump points to `code object: memory://#offset=&size=`. The agent saved it under `/app/logs_claude/debug_run/`. Find it: + +```bash +ls /app/logs_claude/debug_run/ | grep "" | grep "size_" +# Returns e.g.: 7_memory___2188702_offset_0x546c3060_size_4026672 +``` + +Disassemble: + +```bash +/opt/rocm/llvm/bin/llvm-objdump --disassemble-all \ + /app/logs_claude/debug_run/ > /app/logs_claude/fault.s +grep -nE "|s_trap|s_endpgm" /app/logs_claude/fault.s | head -20 +``` + +The PC's surrounding instructions tell you what the kernel was doing. For `s_trap 2` followed by `s_endpgm` you've confirmed an assert (PyTorch `CUDA_KERNEL_ASSERT` lowering). For random other instructions it's a true memory violation — read the address from registers (e.g. `v[0:1]` typically holds the destination address being stored). + +### Step 5: Verify it's actually that kernel (PC can lie) + +Wave debugger PC reports can be **off** when the wave is mid-flight or when the trap fires from a sibling wave. Especially common with Triton — a swa_write trap might be a downstream kernel's fault attributed back. Cross-check: + +- Does the trap reproduce **only when this code path runs**? Disable the call (comment out in Python), retest. +- Does **`HIP_LAUNCH_BLOCKING=1`** make it disappear? Then it's an async race, not a static OOB; the PC kernel is the **victim**, not necessarily the root cause. Bisect for the racer (next step). +- Does inserting `torch.cuda.synchronize()` **right before** this kernel call eliminate the trap? Then root cause is upstream of this point on the same stream. + +### Step 6: Bisect the racer (when PC is racer-victim) + +1. **Comment out one suspect call at a time.** The one whose absence fixes it is the racer (or one of the racing parties). +2. **If neither alone but both together fail**: the race is between them sharing storage / launch slot. Add `torch.cuda.synchronize()` between them as a workaround, but THIS IS NOT A SHIPPABLE FIX — see Step 7. +3. **`py-spy dump --pid `** on stuck ranks: shows the Python frame waiting on the GPU. If it's at your inserted `synchronize()`, the racer is upstream of that line. + +### Step 7: Real fix vs workaround + +Per [[atom-patterns]] / DeepSeek V4 guidance, do not ship `cuda.synchronize()` workarounds without root-causing the race — they mask one workload and surface a worse hang on a larger one. Common real fixes: + +| Symptom | Real fix | +|---------|----------| +| Race involves freshly-allocated transient tensors (e.g. from `torch.where`, `arange`, `.reshape`, `.to(int64)`) | Pre-allocate them in `_alloc_v4_metadata_buffers` (ATOM) or as module-level scratch. Eliminates allocator churn entirely. | +| Multiple `index_copy_` / `scatter_` in sequence | Replace with a single Triton kernel that writes all destinations once. | +| Per-fwd kernel reads stale forward_vars from prior fwd | Switch H2D path off `prep_stream` to default stream (matches ATOM `prepare_mtp_decode` v2 pattern). | +| Cross-rank inconsistency causes one rank to OOB | Ensure all ranks see identical batch shapes before launching kernel; check `cu_seqlens_q` / `state_slot_mapping` parity. | + +## Recovery checklist (after agent run) + +1. `bash scripts/stop_atom_server.sh` — agent leaves zombie KFD entries; if you skip, next launch OOMs at NCCL barrier. +2. `pkill -9 -f spawn_main` — sometimes `stop_atom_server.sh` misses workers stuck in fault state. +3. Wait 10s, then `rocm-smi --showmemuse` — all GPUs must show 0% before relaunching. If not, `rocm-smi --showpids` to find lingering UNKNOWN PIDs (killed but KFD hasn't cleaned yet — wait or escalate). +4. `rm /app/logs_claude/debug_run/memory_*` — code objects are 4 MB each, accumulate fast across runs. +5. Drop `--save-code-objects` from production launches — disk pressure (~500 MB per run). + +## Anti-patterns + +- **Don't assume `--enforce-eager --level 0` is mandatory.** Default launch is fine for most agent runs; reach for these flags only when symptoms point at hipgraph or Inductor (see pre-flight item 2). Adding them blindly hides graph-mode-only bugs. +- **Don't grep `atom_server.log` for "error" or "Traceback"** — agent's wave dump has neither; grep `"stopped, reason"` instead. +- **Don't trust PC literally** — see Step 5. Especially Triton kernels are notorious for cross-wave PC misattribution. Bisect-confirm. +- **Don't leave `--save-code-objects` on for routine runs** — each run dumps ~500 MB. Only enable for the bisect run. +- **Don't add `torch.cuda.synchronize()` "fixes" and ship** — they mask the race for one workload and surface a worse hang (livelock) on a larger one. Find the allocator/stream root cause. + +## Real example (May 14 2026 — V4-Pro MTP-3 prefill hang) + +**Symptom**: `Memory access fault by GPU node-N` on all 8 ranks within 1s of "Scheduled prefill batch: 19 reqs, 9573 tokens". `HIP_LAUNCH_BLOCKING=1` and `ATOM_DUMP_SWA_WRITE=1` (which inserts `.cpu()` sync) both eliminated it. + +**Workflow**: +1. Step 1-2: ran `run_debug_agent.sh`, grep `"stopped, reason"` → got `ASSERT_TRAP` in `at::native::index_copy_kernel_impl>`. +2. Step 4: disassembled code object → confirmed `s_trap 2` followed by `s_endpgm` → PyTorch `CUDA_KERNEL_ASSERT` failed (bound check on index_copy). +3. Step 5: grep'd ATOM source for `index_copy_` of int32 dtype → only csa/hca calls in `deepseek_v4_attn.py:1505-1506`. +4. Step 6: bisected — neither alone trapped, both together did. Added `torch.cuda.synchronize()` between → small workload passed, GSM8K (1319 reqs) eventually livelocked at 1281, py-spy showed all ranks stuck in the synchronize call → confirmed root cause is allocator churn from transient tensors (`csa_win_pos`/`hca_win_pos`/`swa_paged_flat` built via `torch.where`/`arange`/`.reshape`/`.to(int64)`). +5. Real fix (in progress): replace 3 sequential `index_copy_` + transient tensor construction with a single Triton kernel reading only persistent forward_vars buffers. + +Wave dumps and code objects from this run: `/app/logs_claude/atom_server.log` (line 419+), `/app/logs_claude/debug_run/`, `/app/logs_claude/fault.s`. + +## Sample wave dump (what to expect in atom_server.log) + +Trimmed real example from the May 14 2026 V4-Pro MTP-3 run. Key fields are the +mangled function name in `kernel_code_entry=...`, the `stopped, reason: ...` +tag, the `code object: memory://#offset=&size=` line that +points at the saved file, and the `=> : ` arrow showing the +faulting PC. Vector registers (only v0/v1/v6 shown — full dump prints v0..v15 +and beyond) often reveal address / index values that pin down the operand. + +``` +[atom 15:31:09] Scheduled prefill batch: 19 reqs, 9573 tokens, req_ids: (1, 2, ..., 19) +... (some [aiter] type-hints chatter, then the agent's wave dump arrives) ... +-------------------------------------------------------- +wave_27876: pc=0x7f20f5e534c4 (kernel_code_entry=0x7f20f5e52900 >(at::TensorIterator&, long, long, long)::{lambda(int)#1}>(long, at::native::index_copy_kernel_impl >(at::TensorIterator&, long, long, long)::{lambda(int)#1})>) (stopped, reason: ASSERT_TRAP) + +scalar registers: + s0: ffffffff s1: ffffffff s2: 00000000 s3: f8000000 + ... + s32: 0ec00000 s33: 00000002 ... + +system registers: + mode: 000003f0 trapsts: 80000000 status: 80010041 + +trap registers: + ttmp4: 00006ce4 ttmp5: 00000000 ... + +vector registers: + v0: [0] 95f02814 [1] 95f02818 [2] 95f0281c [3] 95f02820 ... [58] 95f028fc [59] 00000a40 ... + v1: [0] 00007f20 [1] 00007f20 ... ← v0:v1 = per-lane dst address + v6: [0] 000080a6 [1] 000080a7 ... [58] 000080e0 [59] 000027ec ← per-lane src VALUES being stored + +general registers: + m0: 000103c0 + pc: 00007f20f5e534c4 exec: f800000000000000 + vcc: ffffffffffffffff + +Disassembly for function void at::native::index_elementwise_kernel<128, 4, at::native::index_copy_kernel_impl >(...)>: + code object: memory://2188702#offset=0x546c3060&size=4026672 + loaded at: [0x7f20f5e00000-0x7f20f615ff09] + => 0x7f20f5e534c4 <+3012>: s_endpgm + 0x7f20f5e534c8 <+3016>: v_cndmask_b32_e32 v0, s0, v0, vcc +``` + +How to read this: + +- **Kernel** = `at::native::index_copy_kernel_impl>` → PyTorch + `Tensor.index_copy_(dim, idx, src)` for 4-byte dtype (int32 / float32). +- **Reason** = `ASSERT_TRAP` → some lane's `index_value` failed + `0 <= idx < self.size(dim)`. Look at `v6` per-lane values to see what was + being processed (if v6 holds the stored value here; the relevant register + varies by kernel). +- **PC** lands on `s_endpgm` because the assert lowering is + `s_trap 2; s_endpgm` — the actual condition test was earlier (look ~10 + instructions back in the disassembly for `s_cbranch_*` + `s_trap`). +- **Code object** at `memory://2188702#offset=0x546c3060&size=4026672` → + saved as `7_memory___2188702_offset_0x546c3060_size_4026672` in + `/app/logs_claude/debug_run/`. Use `llvm-objdump --disassemble-all` on it. + +## Cross-references + +- [[dump-bisect-debug]] — for numerical bugs (wrong output, not crashes) +- [[capture-trace]] — for performance investigation (which kernels eat time) +- [[atom-patterns]] — V4 attention buffer / stream conventions diff --git a/.claude/skills/dump-bisect-debug.md b/.claude/skills/dump-bisect-debug/SKILL.md similarity index 100% rename from .claude/skills/dump-bisect-debug.md rename to .claude/skills/dump-bisect-debug/SKILL.md diff --git a/.claude/skills/run-atom-workload/SKILL.md b/.claude/skills/run-atom-workload/SKILL.md new file mode 100644 index 0000000000..68d329fdb4 --- /dev/null +++ b/.claude/skills/run-atom-workload/SKILL.md @@ -0,0 +1,239 @@ +--- +name: run-atom-workload +description: Run any ATOM workload — accuracy eval (GSM8K via lm_eval), performance benchmark, concurrency sweep, offline simple_inference, or fault repro under rocm-debug-agent. Use when the user asks to "test accuracy", "测精度", "跑 GSM8K", "跑 benchmark", "test performance", "run sweep", "repro the fault", "测一下 MTP1 精度", "跑 simple_inference" — anything that drives an ATOM workload. Encodes the canonical flow (stop → start → workload-in-shell-bg → wait_infer_drain → stop) and the model-family env vars. Same pattern works for both server-based workloads (lm_eval / benchmark client) and offline simple_inference. Do NOT use for profiling traces (use capture-trace). +version: 1.2.0 +scope: ATOM on AMD ROCm; `scripts/` orchestration scripts under repo root +last_updated: 2026-05-18 +--- + +## Path convention + +All script paths in this skill are **project-relative** (`scripts/foo.sh`), +not absolute. The skill lives inside the ATOM repo at `.claude/skills/`, so +it travels with the repo wherever it's cloned. CWD when invoking commands +should be the repo root (Claude Code's default). + +If the user is somewhere else, prefix with `bash $(git rev-parse --show-toplevel)/scripts/foo.sh` +or `cd` to the repo first. + +## Why this skill exists + +Every blocking ATOM workload follows the same 4-step shape: **stop → start → workload-in-bg → wait_infer_drain → stop**. The scripts in `scripts/` are orchestration-grade — chain them via **separate Bash tool calls**, not wrappers, not `&&`. + +Past failure modes this skill prevents (collected from many sessions): + +1. **Writing wrapper scripts in `/app/logs_claude/`** — `start_atom_server.sh` etc. ARE the orchestration layer. Wrapping them is pure noise. The dozens of `run_*.sh` / `start_*_safe.sh` files in `/app/logs_claude/` are session debris — **do not mimic them**. +2. **Chaining all steps with `&&` into one long command** — the user has explicitly forbidden this. Each step gets its own Bash tool call so logs are separate, errors abort cleanly, and the user can interrupt at any boundary. +3. **Double-backgrounding the server** — `start_atom_server.sh` already forks python in background internally and polls ready. Adding `&` after `start_atom_server.sh` is wrong; it blocks until ready or fail. +4. **Skipping `wait_infer_drain.sh`** — without it, GPU faults take the whole timeout to surface, and hangs go undetected. `wait_infer_drain.sh` exits in ~10s on fault and ~1min on hang, with tail-log attached. +5. **Using `curl /health` for liveness** — under heavy load it can false-negative. The flow uses `/v1/models` (start script) and `pgrep` + Engine Core marker (drain script). +6. **Forgetting model-family env vars** — V4-Pro silently regresses on accuracy without `AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1`. Pinned in the table below. +7. **Skipping drain for offline simple_inference** — `wait_infer_drain.sh` supports offline mode (process-exit detection + fault scan). Without it you lose early fault visibility. +8. **Passing the wrong LOG_FILE to drain** — historically callers had to know which log carries the "Engine Core: output send" marker. v1.2 of drain auto-discovers the server log via `/proc//fd/1`, so the user LOG_FILE is now only a **secondary** signal (fault scan + mtime progress for clients with tqdm output). Pass any log you want extra coverage on, or pass nothing — drain still works. + +## Backgrounding mechanism — shell `&`, NOT claude task + +In step 3 the workload must run concurrent with step 4's drain monitor. Use **shell-level `&`** (append to the bash command), NOT the Bash tool's `run_in_background: true`: + +- Shell `&`: bash starts workload, returns immediately, workload runs as orphan; drain finds it via `pgrep` — no claude task tracking dependency +- `run_in_background: true`: workload becomes a claude task accessed via TaskOutput — adds complexity, doesn't help drain since drain uses pgrep anyway + +Pattern (literal): +``` +# Step 3 — single Bash tool call, command ends with `&` +bash scripts/run_gsm8k_eval.sh /data/MODEL 30000 3 & + +# Step 4 — single Bash tool call, blocks +bash scripts/wait_infer_drain.sh 30000 30 10 +``` + +The Bash invocation for step 3 returns the instant `&` is processed (`bash -c 'cmd &'` exits as soon as cmd is backgrounded). Step 4 then runs as the next Bash call and blocks on drain. + +## Canonical 4-step flow + +Run each step as a **separate Bash tool call**. Never chain with `&&`. + +### Step 1 — clean GPU (always) + +```bash +bash scripts/stop_atom_server.sh +``` + +Idempotent. SIGTERM → SIGKILL → force-kill GPU PIDs, waits ≤60s for VRAM=0. Always run first even if you believe no server is up — clears orphaned multiprocessing children that hold GPU memory. + +### Step 2 — start workload host (blocks until ready / completion) + +**Server-based workloads** (GSM8K / benchmark / sweep / fault repro): + +```bash + bash scripts/start_atom_server.sh +``` + +- Self-contained: forks python in background internally, polls `/v1/models` + GPU VRAM, returns when ready or 1 on fail +- **Do NOT** add `&` — the script self-backgrounds python and waits for ready +- **Do NOT** chain `wait_server_ready.sh` after — already inlined inside start script +- Log: hard-coded `/app/logs_claude/atom_server.log` (`LOG_FILE` env is NOT respected by this script). Drain auto-discovers it via `/proc//fd/1` regardless of path + +**Offline workload** (simple_inference): step 2 is fused with step 3, see below. + +Model-family env vars (set as `VAR=val VAR=val bash ...` prefix): + +| Model | Required env vars | Required CLI args | +|---|---|---| +| DeepSeek-V4-Pro | `AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8 --level 0` | +| DeepSeek-R1-0528 (default) | `AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8` | +| Kimi-K2.5-MXFP4 | `HSA_NO_SCRATCH_RECLAIM=1 AITER_LOG_LEVEL=WARNING` | `--kv_cache_dtype fp8 --trust-remote-code` (tp=4) | + +MTP add-on (any supporting model): append `--method mtp --num-speculative-tokens N` to EXTRA_ARGS. V4-Pro: keep `--level 0`. + +### Step 3 — launch workload in shell background (`&`) + +The workload script must end with shell `&` so the Bash tool returns immediately and step 4 can start monitoring in parallel. + +**Server-based workloads** (PORT is needed): + +| Workload | Command (note trailing `&`) | Optional client log for drain | +|---|---|---| +| GSM8K accuracy | `bash scripts/run_gsm8k_eval.sh MODEL PORT NUM_FEWSHOT &` | `/app/logs_claude/gsm8k_eval.log` (lm_eval is silent during requests; drain's auto-discovered server log carries the engine markers — passing this log only helps fault grep coverage) | +| Single benchmark | `bash scripts/run_benchmark.sh MODEL PORT ISL OSL CONC [PROMPT_MULT] [PROFILE] &` | `/app/logs_claude/benchmark.log` (has tqdm progress, useful mtime signal) | +| Concurrency sweep | `bash scripts/run_benchmark_sweep.sh MODEL PORT ISL OSL "CONC1 CONC2 ..." &` | `/app/logs_claude/benchmark.log` (overwritten per step) | + +**Offline simple_inference** (no PORT; step 2 is skipped since this script IS the workload host): + +```bash + bash scripts/start_simple_inference.sh MODEL TP & +``` + +Optional client log for drain: `/app/logs_claude/simple_inference.log` (drain auto-discovers via /proc anyway; this only helps fault grep redundancy). + +Common workload knobs: +- GSM8K shots: 3 for fast/CI parity, 5 for thorough. Set `LIMIT=50` env for first-50-sample sanity. +- Benchmark `PROMPT_MULTIPLIER` default 10. Profiling: use 2 (CONC × 2 requests). +- MTP benchmark MUST add `--use-chat-template` via EXTRA_ARGS (tokenizer mismatch otherwise). +- Benchmark throughput metric: report **Total Token throughput (tok/s)**, NOT Output throughput. Total = input+output, which users care about. +- Never add `--mark-trace` or `ENABLE_TORCH_PROFILER=1` (handled by capture-trace skill). + +### Step 4 — wait_infer_drain (blocks, with early fault/hang detection) + +```bash +bash scripts/wait_infer_drain.sh PORT MAX_MIN POLL_SEC [LOG_FILE] [STUCK_POLLS] +``` + +Defaults: PORT=8000, MAX_MIN=30, POLL_SEC=10, LOG_FILE=empty (server log auto-discovered via `/proc//fd/1`), STUCK_POLLS=6. + +`LOG_FILE` is **optional** in v1.2+. The drain script discovers the server log itself from the running `atom.entrypoints` process. Pass an additional client/workload log only if you want: +- Extra fault grep coverage (drain scans both) +- Mtime-based progress detection for client tools that write tqdm to a file (benchmark, simple_inference) + +PORT is unused in offline mode but kept positional for API symmetry. + +How drain decides (auto-detects server vs offline by `SERVER_PATTERN` pgrep): +- **Server mode**: client gone (lm_eval / curl / benchmark process exited) + no new "Engine Core: output send" since last poll → exit 0 +- **Offline mode**: simple_inference process exited cleanly (no fault grep) → exit 0 +- **Either mode**: fault grep on auto-discovered server log + optional caller LOG_FILE → exit 2 in ≤10s +- **Server only**: no progress (engine output count flat + caller LOG_FILE mtime flat) AND client still running for STUCK_POLLS × POLL_SEC ≈ 1min → exit 1 (hang) +- **Either mode**: MAX_MIN elapsed without resolution → exit 4 + +If exit ≠ 0: read the printed tail, then run step 5 regardless. + +Typical wait windows: +- GSM8K (1319 samples): MAX_MIN=30 plenty for V4-Pro +- Single benchmark: MAX_MIN=30 +- Sweep (8 conc points): MAX_MIN=60 +- Simple_inference (default ~10 prompts): MAX_MIN=15 plenty +- Fault repro: MAX_MIN=10 (fault should land within first request) + +### Step 5 — teardown (always) + +```bash +bash scripts/stop_atom_server.sh +``` + +Same script as step 1. ALWAYS run, even on fault or for offline workloads — releases GPU for next attempt and kills any lingering multiprocessing children. + +## Reading results + +After step 4 returns 0: + +```bash +# GSM8K +grep -E "flexible-extract|strict-match" /app/logs_claude/gsm8k_eval.log | head -2 +# Benchmark +grep -E "Total Token throughput|Mean TPOT|Mean TTFT" /app/logs_claude/benchmark.log +# Simple_inference +grep -E "^Generated|^Output|tokens/s" /app/logs_claude/simple_inference.log +``` + +GSM8K format: `|gsm8k|3|flexible-extract|3|exact_match|↑|0.XXXX|±|0.00XX|` — flexible-extract is the headline number, ±value is the noise band. Anything within 1σ of baseline = no regression. + +## V4-Pro accuracy baselines (current, 3-shot, no `--limit`) + +- Non-MTP: ~0.9545 ± 0.0057 +- MTP-1: ~0.9492 ± 0.006 +- MTP-3: see `feedback_v4_cg_mtp_status_deadlock.md` — currently hangs on high-concurrency GSM8K + +## Hard rules (do not violate) + +1. **One Bash tool call per script.** No `&&` chains. User has explicitly forbidden chaining. +2. **No wrapper scripts in `/app/logs_claude/`.** Call `scripts/*` directly. +3. **Shell `&`, not `run_in_background: true`** for step 3 (drain finds workload via pgrep, no task tracking needed). +4. **Drain auto-discovers server log via `/proc//fd/1`** — you no longer need to know or pass the canonical server log path. Pass a client log only as supplementary signal. +5. **Always step 5 (`stop_atom_server.sh`)**, even after a fault, even for offline workloads. +6. **Never use `curl /health`** to verify ready — only `/v1/models` (already inlined in start script). +7. **No extra `&`, `wait_server_ready.sh`, `LOG_FILE=` env on `start_atom_server.sh`** — start script self-backgrounds, self-polls, hard-codes log path. + +## Reference: each script in one line + +| Script | What it does | Step | Blocks? | +|---|---|---|---| +| `stop_atom_server.sh` | Kill all atom + multiproc children, wait for VRAM=0 | 1, 5 | Yes ≤60s | +| `start_atom_server.sh MODEL TP PORT [ARGS...]` | Clean GPU, fork python in bg, poll ready | 2 (server) | Yes, until ready or fail | +| `start_simple_inference.sh MODEL TP [ARGS...]` | Offline inference (no server, runs prompts) — wrap with `&` for drain | 3 (offline) | Blocks unless `&` | +| `run_gsm8k_eval.sh MODEL PORT FEWSHOT` | lm_eval local-completions GSM8K — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `run_benchmark.sh MODEL PORT ISL OSL CONC [PMULT] [PROF]` | Single perf point — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `run_benchmark_sweep.sh MODEL PORT ISL OSL "CONCs"` | Loop run_benchmark — wrap with `&` for drain | 3 (server) | Blocks unless `&` | +| `wait_infer_drain.sh PORT MAX_MIN POLL [LOG] [STUCK]` | Monitor workload for drain / hang / fault (auto-discovers server log) | 4 | Yes, until exit code | +| `wait_server_ready.sh PORT MAX_MIN POLL LOG` | Standalone ready poller (rarely needed; start script self-polls) | (debug) | Yes | +| `run_debug_agent.sh [--simple] MODEL TP [PORT] [ARGS...]` | Server (or simple_inference) under rocm-debug-agent — fault repro | 2 (replaces start) | Yes, until ready or fault | + +## Worked example: V4-Pro MTP1 GSM8K accuracy + +``` +# Step 1 +bash scripts/stop_atom_server.sh + +# Step 2 +AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING \ + bash scripts/start_atom_server.sh /data/DeepSeek-V4-Pro 8 30000 \ + --kv_cache_dtype fp8 --method mtp --num-speculative-tokens 1 --level 0 + +# Step 3 — note trailing `&` +bash scripts/run_gsm8k_eval.sh /data/DeepSeek-V4-Pro 30000 3 & + +# Step 4 — drain auto-discovers server log; no LOG_FILE needed +bash scripts/wait_infer_drain.sh 30000 30 10 + +# Step 5 +bash scripts/stop_atom_server.sh + +# Read result +grep -E "flexible-extract|strict-match" /app/logs_claude/gsm8k_eval.log | head -2 +``` + +## Worked example: V4-Pro offline simple_inference + +``` +# Step 1 +bash scripts/stop_atom_server.sh + +# Step 2+3 fused (simple_inference IS the workload host) — note trailing `&` +AITER_BF16_FP8_MOE_BOUND=0 ATOM_MOE_GU_ITLV=1 AITER_LOG_LEVEL=WARNING \ + bash scripts/start_simple_inference.sh /data/DeepSeek-V4-Pro 8 \ + --kv_cache_dtype fp8 --level 0 & + +# Step 4 — drain auto-discovers via /proc; PORT unused +bash scripts/wait_infer_drain.sh 0 15 10 + +# Step 5 +bash scripts/stop_atom_server.sh +``` diff --git a/.github/benchmark/models.json b/.github/benchmark/models.json index 5f8c96e4e0..f9f2e28c84 100644 --- a/.github/benchmark/models.json +++ b/.github/benchmark/models.json @@ -19,6 +19,16 @@ "runner": "atom-mi355-8gpu.predownload", "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1" }, + { + "display": "DeepSeek-V4-Pro MTP1", + "path": "deepseek-ai/DeepSeek-V4-Pro", + "prefix": "deepseek-v4-pro", + "args": "--kv_cache_dtype fp8 -tp 8 --method mtp --num-speculative-tokens 1", + "bench_args": "--use-chat-template", + "suffix": "-mtp1", + "runner": "atom-mi355-8gpu.predownload", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1" + }, { "display": "DeepSeek-R1-0528 MTP3", "path": "deepseek-ai/DeepSeek-R1-0528", diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 5c1f3ac34d..fc22c02418 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -35,6 +35,18 @@ "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4-Pro", "_baseline_note": "Local 4-run average GSM8K-100 3-shot flexible-extract = 0.96 (runs: 0.96/0.98/0.96/0.94, stderr ~0.024). ATOM_USE_TRITON_MOE=1 is required — without it accuracy drops to ~0.6. Threshold set 4pp below local baseline to absorb full-eval (1319 samples) noise; refresh after first CI measurement." }, + { + "model_name": "DeepSeek-V4-Pro MTP", + "model_path": "deepseek-ai/DeepSeek-V4-Pro", + "extraArgs": "--kv_cache_dtype fp8 -tp 8 --method mtp --num-speculative-tokens 1", + "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1", + "runner": "atom-mi355-8gpu.predownload", + "test_level": "pr", + "accuracy_threshold": 0.92, + "accuracy_baseline": 0.96, + "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4-Pro", + "_baseline_note": "Same base model as DeepSeek-V4-Pro FP8 (MTP-1: 1 speculative token)" + }, { "model_name": "DeepSeek-R1-0528 MTP", "model_path": "deepseek-ai/DeepSeek-R1-0528", diff --git a/atom/benchmarks/benchmark_serving.py b/atom/benchmarks/benchmark_serving.py index 7f584444cd..35b4302c49 100644 --- a/atom/benchmarks/benchmark_serving.py +++ b/atom/benchmarks/benchmark_serving.py @@ -28,6 +28,7 @@ import argparse import asyncio import contextlib +import functools import gc import json import os @@ -37,12 +38,17 @@ from argparse import ArgumentParser as FlexibleArgumentParser from dataclasses import dataclass from datetime import datetime -from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple +from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple import numpy as np from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase +from atom.entrypoints.openai.chat_encoders import ( + apply_chat_template, + load_custom_message_encoder, +) + from .backend_request_func import ( ASYNC_REQUEST_FUNCS, RequestFuncInput, @@ -92,16 +98,15 @@ def sample_random_requests( range_ratio: float, tokenizer: PreTrainedTokenizerBase, use_chat_template: bool = False, + apply_chat_template_fn: Callable = lambda x: x, ) -> List[Tuple[str, int, int]]: prefix_token_ids = np.random.randint( 0, tokenizer.vocab_size, size=prefix_len ).tolist() if use_chat_template: - chat_template_dummy = tokenizer.apply_chat_template( + chat_template_dummy = apply_chat_template_fn( [{"role": "user", "content": "a"}], - add_generation_prompt=True, - tokenize=False, ) tokenized_chat_template_dummy = tokenizer.encode( chat_template_dummy, add_special_tokens=False @@ -143,10 +148,8 @@ def sample_uniform(seq_len): prompt = tokenizer.decode(prompt_token_ids) if use_chat_template: - prompt = tokenizer.apply_chat_template( + prompt = apply_chat_template_fn( [{"role": "user", "content": prompt}], - add_generation_prompt=True, - tokenize=False, ) prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) @@ -705,6 +708,10 @@ def main(args: argparse.Namespace): tokenizer_mode=tokenizer_mode, trust_remote_code=args.trust_remote_code, ) + custom_encoder = load_custom_message_encoder(model_id) + apply_chat_template_fn = functools.partial( + apply_chat_template, tokenizer, custom_encoder + ) if args.dataset_name == "random": input_requests = sample_random_requests( @@ -714,6 +721,7 @@ def main(args: argparse.Namespace): num_prompts=args.num_prompts, range_ratio=args.random_range_ratio, tokenizer=tokenizer, + apply_chat_template_fn=apply_chat_template_fn, use_chat_template=args.use_chat_template, ) diff --git a/atom/config.py b/atom/config.py index e39df36ac2..223f6f4557 100644 --- a/atom/config.py +++ b/atom/config.py @@ -730,6 +730,7 @@ class SpeculativeConfig: _MTP_TYPE_MAP: ClassVar[dict[str, str]] = { "deepseek_v3": "deepseek_mtp", "deepseek_v32": "deepseek_mtp", + "deepseek_v4": "deepseek_v4_mtp", "glm_moe_dsa": "deepseek_mtp", "qwen3_next": "qwen3_next_mtp", "qwen3_5": "qwen3_5_mtp", @@ -742,6 +743,7 @@ class SpeculativeConfig: # mtp_model_type → (n_predict_attr, architecture) _MTP_CONFIG: ClassVar[dict[str, tuple[str, str]]] = { "deepseek_mtp": ("num_nextn_predict_layers", "DeepSeekMTPModel"), + "deepseek_v4_mtp": ("num_nextn_predict_layers", "DeepseekV4MTPModel"), "qwen3_next_mtp": ("num_nextn_predict_layers", "Qwen3NextMTPModel"), "qwen3_5_mtp": ("mtp_num_hidden_layers", "Qwen3_5MTPModel"), } @@ -1017,7 +1019,13 @@ def __post_init__(self): # tokens. ATOM's BlockManager + slot_mapping math assume one global # block_size, so we override `kv_cache_block_size` here when V4 is # detected; the V4 attention builder enforces the same value. - if getattr(self.hf_config, "model_type", None) == "deepseek_v4": + # + # NOTE: cannot use `hf_config.model_type` for detection — `_CONFIG_REGISTRY` + # maps "deepseek_v4" → "deepseek_v3" so model_type reads as "deepseek_v3". + # Use the preserved `architectures` field (re-injected by get_hf_config, + # line 567) which keeps the original "DeepseekV4ForCausalLM[NextN]" name. + arches = getattr(self.hf_config, "architectures", None) or [] + if any("DeepseekV4" in str(a) for a in arches): v4_block_size = 128 if self.kv_cache_block_size != v4_block_size: self.kv_cache_block_size = v4_block_size diff --git a/atom/model_engine/engine_core.py b/atom/model_engine/engine_core.py index f220231343..ad57c56e66 100644 --- a/atom/model_engine/engine_core.py +++ b/atom/model_engine/engine_core.py @@ -192,6 +192,15 @@ def busy_loop(self): def _process_engine_step(self): result = self.scheduler.schedule() + + # Surface admit-rejected seqs (those `_unschedulable_reason` flags in + # the scheduler) through the same finished-seq path as normal seqs. + # Without this, `llm.generate()` blocks forever waiting for an output + # the rejected seq will never produce. + rejected = self.scheduler.take_rejected() + if rejected: + self.output_queue.put_nowait(rejected) + if result is None: return False scheduled_batch, seqs = result diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 0dcc56fa55..1adbf9df25 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -735,9 +735,18 @@ def is_qwen_next(self) -> bool: return False def is_deepseek_v4(self) -> bool: - if not hasattr(self.hf_text_config, "model_type"): - return False - return self.hf_text_config.model_type == "deepseek_v4" + # NOTE: `hf_text_config.model_type` reads "deepseek_v3" for V4 because + # `_CONFIG_REGISTRY` maps deepseek_v4 → deepseek_v3 (V4 reuses V3 schema). + # Use `architectures` (preserved by get_hf_config:567) instead. Covers + # both target (DeepseekV4ForCausalLM[NextN]) and draft (whose model_type + # SpeculativeConfig stamps as deepseek_v4_mtp). + arches = getattr(self.hf_text_config, "architectures", None) or [] + if any("DeepseekV4" in str(a) for a in arches): + return True + return getattr(self.hf_text_config, "model_type", None) in ( + "deepseek_v4", + "deepseek_v4_mtp", + ) def is_mimo_v2(self) -> bool: if not hasattr(self.hf_text_config, "model_type"): @@ -1082,8 +1091,15 @@ def allocate_forward_vars(self): "top_ks": CpuGpuBuffer(self.max_bs, **i32_kwargs), "top_ps": CpuGpuBuffer(self.max_bs, **f32_kwargs), # Keep enough space for MTP decode (max_q_len > 1). + # `extra_output_dims` lets a model insert dims between N and dim + # (e.g. DeepSeek-V4 returns the un-reduced mHC residual + # [N, hc_mult, dim] from forward, with hc_head + LM head deferred + # to compute_logits). Default `()` keeps the standard 2D layout. "outputs": torch.empty( - self.max_num_batched_tokens, hidden_size, dtype=hidden_type + self.max_num_batched_tokens, + *getattr(self.model, "extra_output_dims", ()), + hidden_size, + dtype=hidden_type, ), } if hasattr(self, "drafter"): @@ -2017,6 +2033,7 @@ def capture_cudagraph(self): num_tokens=num_tokens, num_tokens_across_dp=num_tokens_across_dp, ubatch_slices=ubatch_slices, + in_hipgraph=True, ) # Warmup diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 4cd1b5d2b0..a977c95eb4 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -372,6 +372,7 @@ class Scheduler: def __init__(self, config: Config): self.max_num_seqs = config.max_num_seqs self.max_num_batched_tokens = config.max_num_batched_tokens + self.max_model_len = config.max_model_len self.bos_token_id = config.bos_token_id self.eos_token_id = config.eos_token_id self.stop_token_ids = config.stop_token_ids @@ -380,6 +381,11 @@ def __init__(self, config: Config): self.running: deque[Sequence] = deque() self.config = config + # Admit-rejected seqs (those `_unschedulable_reason` flags). Drained + # by `take_rejected` each EngineCore step; routed through the same + # output_queue path as forward-finished seqs. + self._rejected: list[Sequence] = [] + # KV transfer bookkeeping self.finished_recving_kv_req_ids: list[int] = [] self.deferred_free_blocks: dict[int, Sequence] = {} @@ -409,7 +415,13 @@ def __init__(self, config: Config): self.kv_connector = get_kvconnector("scheduler", config) def is_finished(self): - return not self.waiting and not self.running + # `_rejected` must be considered too: if a batch of seqs is all + # oversized, schedule() moves them straight from `waiting` to + # `_rejected`, leaving both `waiting` and `running` empty. Without + # this check, busy_loop's `is_finished()` short-circuits to True + # before EngineCore drains `_rejected` via take_rejected(), and + # llm.generate() blocks forever. + return not self.waiting and not self.running and not self._rejected def add(self, seq: Sequence): self._warn_if_unschedulable(seq) @@ -420,52 +432,73 @@ def extend(self, seqs: list[Sequence]): self._warn_if_unschedulable(seq) self.waiting.extend(seqs) - def _warn_if_unschedulable(self, seq: Sequence) -> None: - """Detect requests that exceed static scheduler/KV-pool capacity. - - These requests would otherwise sit in `waiting` forever (head-of-line - blocking the prefill loop, which `break`s on the first oversized seq) - with no log output. Surface a single warning at submit time so the - caller knows the request will never be picked up. - - Three permanent failure modes: + def _unschedulable_reason(self, seq: Sequence) -> Optional[str]: + """Return a human-readable reason if `seq` is permanently unschedulable. + + Only checks static (configuration-time) capacity. Dynamic conditions + that can clear up as other seqs finish (e.g. transiently full + per-req-cache pool) are NOT checked here — they're warned at submit + time (`_warn_if_unschedulable`) but not eligible for permanent drop + at schedule time, since the prefill loop's existing `can_allocate` + check will retry them later. + + Permanent failure modes (each leaves the seq stuck in `waiting` + forever and would head-of-line block the prefill loop, which + `break`s on the first oversized seq): + - prompt longer than `max_model_len` → exceeds per-seq KV cache + geometry; attention backends size `block_tables` as + `max_model_len // block_size` cols and would crash with a + broadcast error at prepare-time. (Checked first since it's the + usual actionable cause.) - prompt longer than `max_num_batched_tokens` → no single prefill forward can ever fit it - prompt's KV blocks (+ per-req cache reservation) exceed the total pool size → never fits even on a fully empty pool - - request needs a per-req cache slot but the model was started with - zero slots (e.g. GDN/V4 with `max_num_seqs=0`) + + Called at submit time (`_warn_if_unschedulable`, which logs the + reason and adds extra dynamic warnings) and at schedule time + (drops the seq before it reaches the attention backend). """ num_tokens = seq.num_tokens + if num_tokens > self.max_model_len: + return ( + f"input tokens={num_tokens} > max_model_len={self.max_model_len}. " + f"Increase --max-model-len or shorten the prompt." + ) if num_tokens > self.max_num_batched_tokens: - logger.warning( - "Request %s will never be scheduled: input tokens=%d > " - "max_num_batched_tokens=%d. Increase --max-num-batched-tokens " - "or shorten the prompt.", - seq.id, - num_tokens, - self.max_num_batched_tokens, + return ( + f"input tokens={num_tokens} > max_num_batched_tokens=" + f"{self.max_num_batched_tokens}. Increase --max-num-batched-tokens " + f"or shorten the prompt." ) - return - bm = self.block_manager per_req_cost = bm.per_req_cache_equiv_blocks if seq.has_per_req_cache else 0 total_blocks = len(bm.blocks) if seq.num_blocks + per_req_cost > total_blocks: - logger.warning( - "Request %s will never be scheduled: needs %d KV blocks " - "(%d for %d input tokens + %d for per-req cache) > " - "total pool blocks=%d. Reduce prompt length, lower " - "--max-num-seqs, or raise --gpu-memory-utilization.", - seq.id, - seq.num_blocks + per_req_cost, - seq.num_blocks, - num_tokens, - per_req_cost, - total_blocks, + return ( + f"needs {seq.num_blocks + per_req_cost} KV blocks " + f"({seq.num_blocks} for {num_tokens} input tokens + " + f"{per_req_cost} for per-req cache) > total pool blocks=" + f"{total_blocks}. Reduce prompt length, lower --max-num-seqs, " + f"or raise --gpu-memory-utilization." ) - return + return None + def _warn_if_unschedulable(self, seq: Sequence) -> None: + """Log a single warning at submit time for permanently-unschedulable + sequences. The seq still enters `waiting`; the prefill scheduler drops + it later (see `schedule`). + + Also surfaces a dynamic configuration-time-only warning when the + model was started with zero per-req-cache slots (max_num_seqs=0) — + this is permanent if it holds at submit time, but is NOT eligible + for schedule-time drop (a future config change could create slots). + """ + reason = self._unschedulable_reason(seq) + if reason is not None: + logger.warning("Request %s will never be scheduled: %s", seq.id, reason) + return + bm = self.block_manager if ( seq.has_per_req_cache and not bm.free_per_req_cache_groups @@ -473,11 +506,23 @@ def _warn_if_unschedulable(self, seq: Sequence) -> None: ): logger.warning( "Request %s will never be scheduled: needs per-req cache slot " - "but no slots were allocated (max_num_seqs=0 for this model " - "type).", + "but no slots were allocated (max_num_seqs=0 for this model type).", seq.id, ) + def take_rejected(self) -> list[Sequence]: + """Pop and return any seqs the prefill scheduler dropped because + `_unschedulable_reason` flagged them (oversized prompt, exhausted + pool, etc.). Caller (EngineCore) pushes them onto the same + output_queue as forward-finished seqs so `llm.generate()` returns + an output for them instead of blocking forever. + """ + if not self._rejected: + return [] + out = self._rejected + self._rejected = [] + return out + def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: """Select the next batch of sequences for a forward pass. @@ -498,6 +543,22 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]: while self.waiting and num_seqs_prefill < self.max_num_seqs: seq = self.waiting.popleft() + # Drop seqs the static-capacity check at submit-time flagged as + # permanently unschedulable (oversized prompt, exhausted pool, + # etc.). They've already been warned; mark FINISHED + record the + # rejection reason and route them to `_rejected` so EngineCore + # surfaces them through the same output_queue as forward-finished + # seqs. Without this they'd reach the attention backend (where an + # oversized prompt crashes with a broadcast error) AND + # `llm.generate()` would block forever waiting for an output. + # Re-check here (not just at submit) since pool state may change. + unschedulable = self._unschedulable_reason(seq) + if unschedulable is not None: + seq.status = SequenceStatus.FINISHED + seq.leave_reason = f"unschedulable: {unschedulable}" + self._rejected.append(seq) + continue + # KV Transfer: skip request if still waiting for remote KVs waiting_remote_to_waiting_ready = False if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS: @@ -776,6 +837,17 @@ def postprocess( num_tokens = seq.num_tokens - self.mtp_k - num_rejected leave_reason = None + # MTP edge case: `rejection_sampler` does NOT inspect EOS — it + # only compares draft vs target_argmax for acceptance. So when + # the verified token is EOS the kernel still emits 1+ accepted + # bonus tokens after EOS (often BOS, since the model naturally + # starts a new sentence). Without truncating, those post-EOS + # tokens leak into the detokenized output (e.g. "...6."). + # Empirically confirmed via DIAG: `token_ids=[EOS=1, BOS=0]`, + # `eos_idx=0`, `num_new=2`, `num_rejected=0` for V4-Pro MTP-1. + # Track the earliest stop position so `num_tokens` can drop the + # spurious tail below. + stop_at_idx: Optional[int] = None # Check if sequence ends with any stop sequence for stop_seq in seq.stop_token_sequences: stop_len = len(stop_seq) @@ -785,6 +857,10 @@ def postprocess( offset = num_tokens - i if seq.token_ids[offset - stop_len : offset] == stop_seq: is_stop = True + # `i` counts back from the last sampled token + # (i=0 = last). Truncate to include this stop + # sequence (drop everything after it). + stop_at_idx = num_new_token - 1 - i break if is_stop: leave_reason = "stop_sequence" @@ -793,16 +869,31 @@ def postprocess( # Check the last token in the list for EOS if token_ids and not seq.ignore_eos and self.eos_token_id in token_ids: leave_reason = "eos" + stop_at_idx = token_ids.index(self.eos_token_id) elif not seq.ignore_eos and any( t in self.stop_token_ids for t in token_ids ): - first_stop_token = next( - t for t in token_ids if t in self.stop_token_ids + stop_at_idx = next( + i for i, t in enumerate(token_ids) if t in self.stop_token_ids ) - leave_reason = f"stop_{first_stop_token}" - elif seq.num_completion_tokens >= seq.max_tokens: + leave_reason = f"stop_{token_ids[stop_at_idx]}" + elif (num_tokens - seq.num_prompt_tokens) >= seq.max_tokens: + # Use the local `num_tokens` (= seq.num_tokens - mtp_k - + # num_rejected, set at line 716) instead of the property + # `seq.num_completion_tokens` which still reflects the + # raw mtp_k+1 placeholder bump from `prepare_decode`. The + # property over-counts by `mtp_k + num_rejected`, causing + # max_tokens to trip that many tokens early (visible as + # `output tokens=95` for max_tokens=100, mtp_k=3). Non-MTP + # path: mtp_k = num_rejected = 0 → behavior unchanged. leave_reason = "max_tokens" + # Drop accepted-draft tokens past the stop position (MTP only — + # for non-spec the sampler emits exactly 1 token so this is a + # no-op). + if stop_at_idx is not None and stop_at_idx < num_new_token - 1: + num_tokens -= (num_new_token - 1) - stop_at_idx + # Prepare stream output if stream_output_queue is not None and new_tokens: if self.kv_connector is not None and leave_reason is not None: diff --git a/atom/model_loader/loader.py b/atom/model_loader/loader.py index 0763179f51..057f007d97 100644 --- a/atom/model_loader/loader.py +++ b/atom/model_loader/loader.py @@ -296,6 +296,12 @@ def extract_expert_target_and_id(name: str) -> Tuple[str, int] | None: if weights_mapper is None: weights_mapper = getattr(model, "weights_mapper", None) params_dict = dict(model.named_parameters()) + # Load `mtp.*` ckpt entries only when this is a spec-decode draft model + # AND it actually carries `mtp.*` params. The two-condition gate handles + # both directions: target loads (spec_decode=False) skip mtp regardless, + # and spec drafts that don't use the standard `mtp.*` param naming + # (e.g. Eagle3 with its own arch) also skip cleanly. + need_load_mtp: bool = spec_decode and any("mtp" in n for n in params_dict) # Pre-index expert_mapping by weight_name_part for O(1) lookup. # Original code does O(N) scan of expert_mapping (768 entries) per tensor, @@ -353,7 +359,7 @@ def _submit(fn, *args): name = mapped_name if load_dummy: continue - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: continue if name.endswith("kv_scale") or "inv_freq" in name: continue @@ -485,7 +491,7 @@ def _submit(fn, *args): ) and name not in params_dict: matched = True break - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: matched = True break try: @@ -508,7 +514,7 @@ def _submit(fn, *args): matched = True break if not matched: - if "mtp" in name and not spec_decode: + if "mtp" in name and not need_load_mtp: continue if merged_target := extract_expert_target_and_id(name): fused_name, expert_id = merged_target diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index c6defd2fd6..e4e7cc47ef 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -71,9 +71,7 @@ def __init__( CommonAttentionBuilder.__init__(self, model_runner) config = model_runner.config hf_config = config.hf_config - self.num_attention_heads = ( - hf_config.num_attention_heads // get_tp_group().world_size - ) + # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. # For speculative decode (MTP), max_qlen = num_speculative_tokens + 1 if ( config.speculative_config is not None diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 3f7ab3c83b..03b44b2fb3 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -12,7 +12,6 @@ get_mla_metadata_info_v1, get_mla_metadata_v1, ) -from aiter.dist.parallel_state import get_tp_group from atom.model_engine.scheduler import ScheduledBatch from atom.model_ops.attention_mla import _MLA_MIN_HEADS, MLAAttention from atom.plugin.attention import ( @@ -22,7 +21,6 @@ from atom.plugin.prepare import is_plugin_mode from atom.utils import CpuGpuBuffer from atom.utils.block_convert import ( - block_table_convert_triton, kv_indices_generate_triton, ) from atom.utils.forward_context import AttentionMetaData, Context @@ -61,9 +59,7 @@ def __init__(self, model_runner): CommonAttentionBuilder.__init__(self, model_runner) config = model_runner.config hf_config = config.hf_config - self.num_attention_heads = ( - hf_config.num_attention_heads // get_tp_group().world_size - ) + # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. self.padded_num_attention_heads = max(self.num_attention_heads, _MLA_MIN_HEADS) self.is_sparse = model_runner.is_deepseek_v32 self.index_topk = hf_config.index_topk if self.is_sparse else -1 diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index ca1b77ddc7..9c7963c551 100644 --- a/atom/model_ops/attentions/backends.py +++ b/atom/model_ops/attentions/backends.py @@ -6,6 +6,7 @@ from typing import Any, Dict, Generic, Optional, Type, TypeVar import torch +from aiter.dist.parallel_state import get_tp_group from atom.model_engine.scheduler import ScheduledBatch from atom.model_ops.attention_mla import MLAModules from atom.utils import CpuGpuBuffer @@ -208,6 +209,13 @@ def __init__(self, model_runner): self.max_num_blocks_per_seq = ( config.max_model_len + self.block_size - 1 ) // self.block_size + # Per-rank attention head count. eagle.propose's mid-step path reads + # this to gate the `do_attn_metadata_update` branch. Subclasses that + # need a kernel-minimum-padded count set `self.padded_num_attention_heads` + # separately (it does NOT replace this attribute). + self.num_attention_heads = ( + hf_config.num_attention_heads // get_tp_group().world_size + ) i64_kwargs = {"dtype": torch.int64, "device": self.device} i32_kwargs = {"dtype": torch.int32, "device": self.device} diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index 86f5cfb12a..6fd548369c 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -47,8 +47,13 @@ AttentionMetadataBuilder, CommonAttentionBuilder, ) +from atom.model_ops.v4_kernels import write_v4_paged_decode_indices from atom.utils import CpuGpuBuffer -from atom.utils.forward_context import AttentionMetaData, Context +from atom.utils.forward_context import ( + AttentionMetaData, + Context, + get_forward_context, +) # --------------------------------------------------------------------------- # Typed metadata surface for V4. The base AttentionMetaData class is shared @@ -89,8 +94,15 @@ class AttentionMetaData_DSV4(AttentionMetaData): # ----- CPU mirrors (avoid GPU→CPU `.item()` / `.tolist()` syncs) ----- state_slot_mapping_cpu: Optional[Any] = None """[bs] np.int32 — per-seq state cache slot id (host copy).""" - start_pos_per_seq_cpu: Optional[Any] = None - """[bs] np.int64 — absolute start position per seq.""" + n_committed_csa_per_seq_cpu: Optional[Any] = None + """[bs] np.int32 — `ctx_len // 4` (CSA committed K per seq). Built once + in `_attach_v4_per_fwd_meta` from `var["context_lens"].np`; consumed by + `_attach_v4_paged_decode_meta`, `_build_paged_prefill_meta`, and + `_build_v4_indexer_meta` (indptr cumsums). Single source of truth so + those callers don't each re-read context_lens + divide.""" + n_committed_hca_per_seq_cpu: Optional[Any] = None + """[bs] np.int32 — `ctx_len // 128` (HCA committed compress entries per + seq). Same lifecycle as `n_committed_csa_per_seq_cpu`.""" # ----- Per-seq GPU scalars (single-source-of-truth, shared by kernels) ----- state_slot_mapping: Optional[torch.Tensor] = None @@ -110,9 +122,11 @@ class AttentionMetaData_DSV4(AttentionMetaData): kernels skip on `bid < 0`. All other per-token quantities resolved as `per_seq_data[batch_id_per_token[t]]` — no [T] aliases of seq data.""" swa_write_indices: Optional[torch.Tensor] = None - """[padded_T] int64 GPU — src row id into per-fwd KV for swa_write. - `[0:num_write]` = real (last `win` tokens per seq); - `[num_write:padded_T]` = -1 sentinel. `None` for warmup / empty fwd.""" + """[W] int64 GPU — src row id into per-fwd KV for swa_write. + `[0:num_write]` = real (last `win` tokens per seq); trailing entries + `[num_write:W]` = -1 sentinel (only present on decode/MTP CG paths + where `W = padded_bs * (1 + max_spec_steps)`; prefill is eager and + uses `W = num_write` exactly, no padding). `None` for warmup / empty fwd.""" compress_plans: Optional[Dict[int, Any]] = None """dict[ratio:int -> CompressPlan] — packed plan tensors per compress_ratio (4=CSA, 128=HCA).""" @@ -141,7 +155,9 @@ class AttentionMetaData_DSV4(AttentionMetaData): (= `win + n_committed_hca[bid]`). Padded tail = last value.""" swa_pages: int = 0 """Boundary in `unified_kv`: index < swa_pages → SWA region; index >= - swa_pages → compress region. Equal to `max_per_req_cache_slots * win`.""" + swa_pages → compress region. Equal to + `max_per_req_cache_slots * win_with_spec` (per-slot SWA region holds + `win + mtp_k` ring entries; reduces to `win` when MTP is off).""" # ----- Indexer / sparse-layout side metadata ----- indexer_meta: Optional[Dict[str, Any]] = None @@ -225,104 +241,15 @@ def _segment_indices( total = int(lens.sum()) if total == 0: return ( - np.empty(0, dtype=np.int64), - np.empty(0, dtype=np.int64), + np.empty(0, dtype=np.int32), + np.empty(0, dtype=np.int32), ) - token_seq_ids = np.repeat(seq_ids.astype(np.int64), lens) - cum = np.concatenate(([0], np.cumsum(lens.astype(np.int64))[:-1])) - local_pos = np.arange(total, dtype=np.int64) - np.repeat(cum, lens) + token_seq_ids = np.repeat(seq_ids.astype(np.int32), lens) + cum = np.concatenate(([0], np.cumsum(lens, dtype=np.int32)[:-1])) + local_pos = np.arange(total, dtype=np.int32) - np.repeat(cum, lens) return token_seq_ids, local_pos -def _build_window_topk_batched( - positions: torch.Tensor, # [total_tokens] int (abs token positions) - start_pos_per_token: torch.Tensor, # [total_tokens] int (each token's seq start_pos) - window_size: int, - out: Optional[ - torch.Tensor - ] = None, # [total_tokens, window_size] int32 — CG-stable buffer -) -> torch.Tensor: # [total_tokens, window_size] int32 - """Per-token sliding-window topk indices for the whole batch. - - Three-branch semantics: - - sp == 0 (fresh prefill): matrix entries = abs positions in the window - [pos-win+1, pos] clamped to [0, pos]; mask future via -1. - - 0 < sp < win-1 (prefix mode): all tokens in the seq share a single - matrix [0..sp, -1, ..., -1] (matches original semantics, including - MTP-N where the same start_pos is reused). - - sp >= win-1 (cyclic mode): cyclic ring offsets starting at sp+1 mod win. - - `out`: when supplied, the final int32 result is written into `out` - (sliced to `[total_tokens, window_size]`) so its storage address is - stable across captures — required for CUDAGraph replay correctness. - Intermediates remain transient (not consumed by the captured graph). - """ - device = positions.device - total = positions.size(0) - arange_w = torch.arange(window_size, device=device, dtype=positions.dtype).view( - 1, window_size - ) - pos_col = positions.view(total, 1) - sp_col = start_pos_per_token.view(total, 1) - neg1 = positions.new_full((), -1) - - # Case A: sp == 0 (fresh prefill) — abs positions [pos-win+1, pos] clamped. - case_a = (pos_col - window_size + 1).clamp(min=0) + arange_w - case_a = torch.where(case_a > pos_col, neg1, case_a) - - # Case B: 0 < sp < win-1 (prefix mode) — shared per-seq matrix. - case_b = arange_w.expand(total, window_size).clone() - case_b = torch.where(arange_w > sp_col, neg1, case_b) - - # Case C: sp >= win-1 (cyclic mode) — ring offsets. - sp_mod = sp_col % window_size - case_c = (sp_mod + 1 + arange_w) % window_size - - sp_eq_0 = sp_col == 0 - sp_in_prefix = (sp_col > 0) & (sp_col < window_size - 1) - - res = case_c - res = torch.where(sp_in_prefix.expand_as(res), case_b, res) - res = torch.where(sp_eq_0.expand_as(res), case_a, res) - if out is None: - return res.to(torch.int32) - out[:total].copy_(res, non_blocking=True) - return out[:total] - - -def _build_window_topk_np( - positions: np.ndarray, - start_pos_per_token: np.ndarray, - window_size: int, -) -> np.ndarray: - """CPU numpy equivalent of ``_build_window_topk_batched``. - - For small decode batches the ~25 GPU kernel launches dominate; this - numpy version runs in microseconds on CPU and feeds a single H2D copy. - """ - total = positions.shape[0] - arange_w = np.arange(window_size, dtype=np.int64).reshape(1, window_size) - pos_col = positions.reshape(total, 1) - sp_col = start_pos_per_token.reshape(total, 1) - - case_a = np.maximum(pos_col - window_size + 1, 0) + arange_w - case_a = np.where(case_a > pos_col, -1, case_a) - - case_b = np.broadcast_to(arange_w, (total, window_size)).copy() - case_b = np.where(arange_w > sp_col, -1, case_b) - - sp_mod = sp_col % window_size - case_c = (sp_mod + 1 + arange_w) % window_size - - sp_eq_0 = sp_col == 0 - sp_in_prefix = (sp_col > 0) & (sp_col < window_size - 1) - - res = case_c - res = np.where(sp_in_prefix, case_b, res) - res = np.where(sp_eq_0, case_a, res) - return res.astype(np.int32) - - class DeepseekV4Backend(AttentionBackend): """Backend selector entry for V4 hybrid attention. @@ -386,13 +313,6 @@ def __init__(self, model_runner): # max-density ratio (1 indexer slot per 4 source tokens). self.max_model_len_idx = model_runner.config.max_model_len // 4 - # Compressor state shape: [coff * ratio, coff * head_dim], fp32. - # CSA: ratio=4, overlap=True -> coff=2 -> [8, 2*head_dim] - # HCA: ratio=128, overlap=False -> coff=1 -> [128, head_dim] - self.csa_main_state_shape = (2 * 4, 2 * self.head_dim) - self.csa_idx_state_shape = (2 * 4, 2 * self.index_head_dim) - self.hca_main_state_shape = (128, self.head_dim) - # Classical KV pool geometry. block_size=128 original tokens means # each V4 block holds k1=128/4=32 CSA entries and k2=128/128=1 HCA # entry per layer (paper §3.6.1). @@ -415,7 +335,40 @@ def __init__(self, model_runner): self.max_spec_steps = ( int(model_runner.drafter.mtp_k) if hasattr(model_runner, "drafter") else 0 ) + + # Compressor state shape: [ring_size, coff * head_dim], fp32. + # ring_size = K_pool + (max_spec_steps + 1), where K_pool = coff * ratio. + # The extra (max_spec_steps + 1) slots prevent reject K/V written to the + # ring in round R from being borrow-read by the round-R+1 re-commit of + # the same boundary (slot index = pos % ring_size). With ring_size = + # K_pool the slot of pos P+K_pool aliases the slot of pos P, so a + # rejected K_{P+K_pool} written in R overwrites the K_P that R+1's + # commit pool window [P..P+K_pool-1] still needs. Bumping by one + # verify window's worth of positions decouples the two. + # CSA: ratio=4, overlap=True → K_pool=8; spec ring_size=8 + (mtp_k+1) + # HCA: ratio=128, overlap=False → K_pool=128; spec ring_size=128+(mtp_k+1) + # Non-spec (max_spec_steps=0) → ring_size = K_pool + 1 (effectively + # equivalent to old K_pool layout for correctness; one extra slot is + # the algebraic minimum and trivial in memory). + ring_extra = self.max_spec_steps + 1 + self.csa_main_state_shape = (2 * 4 + ring_extra, 2 * self.head_dim) + self.csa_idx_state_shape = (2 * 4 + ring_extra, 2 * self.index_head_dim) + self.hca_main_state_shape = (128 + ring_extra, self.head_dim) self.max_decode_tokens = self.max_bs * (1 + self.max_spec_steps) + # SWA ring-buffer slots per req. Distinct from `window_size`: + # * `window_size` = SWA attention window = topk count per token + # (each query attends to W consecutive K/V positions). + # * `win_with_spec` = `window_size + max_spec_steps` = ring-buffer + # slot count per req. With MTP-k the per-fwd writes the verified + # token + k draft tokens at positions [p_0..p_k]; if the cache + # were only sized W, draft slots `p_(i+1)..p_k` would alias into + # [p_0-W+1..p_0] and the verified query at `p_0` would read + # future tokens (silent correctness bug). MTP off → max_spec_steps + # == 0 → win_with_spec == window_size, identical bytes layout. + # Used as: SWA `unified_kv` per-slot stride, `swa_kv` ring-buffer dim, + # `swa_write` modulo, and the ring-index modulo `cs` in the V4 + # paged-decode index-write kernel. + self.win_with_spec = self.window_size + self.max_spec_steps # Worst-case HCA per-token committed compress count # (= max_model_len // 128 for V4-Pro = 8192 at 1M context). self.max_committed_hca = model_runner.config.max_model_len // 128 @@ -449,8 +402,9 @@ def compute_per_req_cache_bytes(self) -> int: csa_main = self._numel(self.csa_main_state_shape) * 2 * elem_state csa_idx = self._numel(self.csa_idx_state_shape) * 2 * elem_state hca_main = self._numel(self.hca_main_state_shape) * 2 * elem_state - # SWA window per layer. - swa_per_layer = self.window_size * self.head_dim * elem_swa + # SWA window per layer. Cache holds `win_with_spec = win + mtp_k` + # slots so MTP draft tokens don't alias verified-token slots. + swa_per_layer = self.win_with_spec * self.head_dim * elem_swa return ( len(self.csa_layers) * (csa_main + csa_idx) + len(self.hca_layers) * hca_main @@ -458,7 +412,9 @@ def compute_per_req_cache_bytes(self) -> int: ) def slots_per_req(self) -> int: - # No spec decoding lookahead in V4 (pre-PR3-main). + # State cache is one slot per req regardless of MTP. The MTP draft + # lookahead bytes are absorbed into per-slot SWA size via + # `win_with_spec` (above), not into a slots_per_req multiplier. return 1 def compute_block_bytes(self) -> int: @@ -518,10 +474,12 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: Per-layer `unified_kv` layout (decode-time paged_decode kernel reads a single base ptr; offsets `[0, swa_pages)` are SWA, `[swa_pages, ..)` - are compress): - Dense layer: [num_slots*win, head_dim] BF16 - CSA layer: [num_slots*win + NB*k1, head_dim] BF16 - HCA layer: [num_slots*win + NB*k2, head_dim] BF16 + are compress). Per-slot SWA region is `win_with_spec = win + mtp_k` + (extra slack so MTP draft tokens don't overwrite the verified token's + ring slot mid-fwd): + Dense layer: [num_slots*win_with_spec, head_dim] BF16 + CSA layer: [num_slots*win_with_spec + NB*k1, head_dim] BF16 + HCA layer: [num_slots*win_with_spec + NB*k2, head_dim] BF16 `build_kv_cache_tensor` slices per-layer views to bind into `attn.swa_kv` (SWA portion, reshape to [num_slots, win, head_dim]) @@ -544,7 +502,7 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: num_blocks = self.model_runner.num_physical_kvcache_blocks n_csa = len(self.csa_layers) n_hca = len(self.hca_layers) - swa_pages = num_slots * self.window_size + swa_pages = num_slots * self.win_with_spec head_dim = self.head_dim dtype = self._swa_dtype @@ -614,18 +572,20 @@ def build_kv_cache_tensor(self, layer_id: int, module): runner = self.model_runner num_slots = self.model_runner.max_per_req_cache_slots - swa_pages = num_slots * self.window_size + swa_pages = num_slots * self.win_with_spec if isinstance(module, _V4Attention): # Bind both: # - `attn.unified_kv`: the full per-layer pool (paged_decode reads). - # - `attn.swa_kv`: a [num_slots, win, head_dim] view onto the - # SWA prefix (compatibility with `swa_write` and the existing - # prefill/non-CG path that index_select's by state slot). + # - `attn.swa_kv`: a [num_slots, win_with_spec, head_dim] view + # onto the SWA prefix. Per-slot dim is `win + mtp_k` so MTP + # draft tokens have their own ring slots; `swa_write` modulo + # and `paged_decode` per-row case_c modulo both use this + # dim (= `swa_kv.shape[1]`). unified = runner.v4_unified_kv[module.layer_id] module.unified_kv = unified module.swa_kv = unified[:swa_pages].view( - num_slots, self.window_size, self.head_dim + num_slots, self.win_with_spec, self.head_dim ) return None @@ -726,8 +686,6 @@ def build_kv_cache_tensor(self, layer_id: int, module): def _attach_v4_indexer_meta( self, attn_metadata: AttentionMetaData_DSV4, - cu_seqlens_q_np, - start_pos_per_seq, scheduled_bs: int, total_tokens: int, positions_gpu=None, @@ -739,12 +697,8 @@ def _attach_v4_indexer_meta( build. None for warmup or empty fwd; `_build_v4_indexer_meta` handles both. """ - import numpy as np - attn_metadata.indexer_meta = self._build_v4_indexer_meta( attn_metadata=attn_metadata, - cu_seqlens_q_np=cu_seqlens_q_np, - start_pos_per_seq_np=np.asarray(start_pos_per_seq, dtype=np.int64), positions_gpu=positions_gpu, scheduled_bs=scheduled_bs, total_tokens=total_tokens, @@ -755,8 +709,6 @@ def _build_v4_indexer_meta( self, *, attn_metadata: AttentionMetaData_DSV4, - cu_seqlens_q_np, - start_pos_per_seq_np, positions_gpu, scheduled_bs: int, total_tokens: int, @@ -768,13 +720,13 @@ def _build_v4_indexer_meta( inline H2D path) or when CSA / Indexer is not on the model. CSA ratio is fixed at 4; we always build under that assumption. - Reuses two shared GPU tensors set by `_attach_v4_per_fwd_meta` - (which MUST be called before this helper): - - `attn_metadata.batch_id_per_token` [padded_T] int32 + Reads pre-computed `attn_metadata.n_committed_csa_per_seq_cpu` + (set by `_attach_v4_per_fwd_meta`, which MUST run first) for the + per-seq committed count and cumsums it on CPU. + + Reuses two shared GPU tensors also set by `_attach_v4_per_fwd_meta`: + - `attn_metadata.batch_id_per_token` [padded_T] int64 - `attn_metadata.n_committed_csa_per_seq` [bs] int32 - These were previously re-staged here as `v4_indexer_batch_id_per_token` - (int64) and `v4_indexer_n_committed_per_seq` (int64) — one extra H2D - each per fwd. Now we read int32 → cast to int64 on GPU. The FP8 indexer K-cache write happens inside `fused_compress_attn` (the unified Indexer-inner Compressor path) via the same block_tables @@ -785,18 +737,12 @@ def _build_v4_indexer_meta( # invariants as `_attach_v4_per_fwd_meta` — guaranteed by every # prepare_*/CG-capture path). bs = scheduled_bs - ratio = 4 # CSA - cu_seqlens_q_arr = cu_seqlens_q_np[: bs + 1].astype(np.int64) - token_num_per_seq = (cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:bs]).astype( - np.int64 - ) - start_pos_per_seq = start_pos_per_seq_np[:bs].astype(np.int64) - end_pos_per_seq = start_pos_per_seq + token_num_per_seq - n_committed_per_seq = end_pos_per_seq // ratio # host copy for cumsum + ratio = 4 # CSA — also referenced by `visible_end_gpu` below + n_committed_per_seq = attn_metadata.n_committed_csa_per_seq_cpu[:bs] cu_committed_cpu = np.concatenate( [ - np.zeros(1, dtype=np.int64), - np.cumsum(n_committed_per_seq, dtype=np.int64), + np.zeros(1, dtype=np.int32), + np.cumsum(n_committed_per_seq, dtype=np.int32), ] ) # Empty-batch guard: when no seq has committed K yet @@ -828,9 +774,7 @@ def _build_v4_indexer_meta( # cu_committed_gpu is consumed both as `cu_starts/cu_ends` for the # fp8_mqa_logits per-token range AND as `cu_seq_lens` for the # cp_gather_indexer_k_quant_cache call (per-seq cumulative committed K). - cu_committed_gpu = self._stage( - "v4_indexer_cu_committed", cu_committed_cpu.astype(np.int32) - ) + cu_committed_gpu = self._stage("v4_indexer_cu_committed", cu_committed_cpu) # Layer-invariant GPU derivations (each was previously rebuilt ~30x # per fwd inside the per-CSA-layer body). @@ -865,6 +809,90 @@ def _build_v4_indexer_meta( "cu_ends_gpu": cu_ends_gpu, } + def prepare_mtp_decode( + self, + bs: int, + max_seqlen_q: int, + max_seqlen_k: int, + only_update: bool = False, + num_reject_tokens: torch.Tensor = None, + ): + """Per-draft-step V4 region metadata rebuild for 1-token-per-seq shape. + + Called by EagleProposer.propose at mid-step iters (i < mtp_k - 1). + Eagle has already bumped attn_metadata.context_lens GPU by +1 and + max_seqlen_k by +1 before calling us. We mirror the +1 on CPU (zero + D2H: structural invariant of eagle's loop) and rebuild + v4_kv_indices_{swa,csa,hca}, batch_id_per_token, n_committed_csa_per_seq, + indexer meta, and compress_plans. + + ``only_update`` / ``num_reject_tokens`` are MLA-specific (no V4 + analog — V4 has no incremental-update kernel and the ctx rollback + is already applied at the top of prepare_decode for the verify + shape). Ignored. + """ + # `max_per_req_cache_slots` is set inside `model_runner.get_num_blocks`, + # which runs AFTER `warmup_model`. The full per-fwd meta rebuild below + # eventually reads it via `_attach_v4_paged_decode_meta`, so during + # warmup (attr unset) we no-op — warmup discards draft output anyway, + # and the verify-shape attn_metadata from the main forward stays valid + # for the rest of eagle.propose. + if not getattr(self.model_runner, "max_per_req_cache_slots", 0): + return {} + + var = self.model_runner.forward_vars + attn_metadata = get_forward_context().attn_metadata + + # 1. CPU mirror of eagle's GPU `context_lens[:bs] += 1` bump. Zero + # D2H: we know the offset by construction (+1 per call). + var["context_lens"].np[:bs] += 1 + context_lens_np = var["context_lens"].np[:bs] + + # 2. 1-token-per-seq shape: positions = ctx-1, cu_seqlens_q = arange. + positions_np = (context_lens_np - 1).astype(np.int32) + cu_seqlens_q_np = np.arange(bs + 1, dtype=np.int32) + var["positions"].np[:bs] = positions_np + var["cu_seqlens_q"].np[: bs + 1] = cu_seqlens_q_np + + # 3. H2D staging. context_lens already on GPU (eagle bumped it); + # CPU is now mirrored. + positions_gpu = var["positions"].copy_to_gpu(bs) + var["cu_seqlens_q"].copy_to_gpu(bs + 1) + + # 4. CPU numpy: extend_lens (=1 per seq). + extend_lens_np = np.ones(bs, dtype=np.int32) + + # 5. Rebuild V4 region metadata via existing helpers (numpy + H2D, + # no D2H — `_build_compress_plans` only triggers `.cpu()` when + # given torch tensors; we pass numpy below). + self._attach_v4_per_fwd_meta( + attn_metadata, + cu_seqlens_q_np, + extend_lens_np, # = np.ones(bs) — MTP draft step is 1 token per seq + state_slot_mapping_cpu=attn_metadata.state_slot_mapping_cpu, + scheduled_bs=bs, + total_tokens=bs, + padded_bs=bs, + max_q_len=1, + ) + self._attach_v4_indexer_meta( + attn_metadata, + scheduled_bs=bs, + total_tokens=bs, + positions_gpu=positions_gpu, + ) + + # 6. Compress plans for state-ring write of each new draft token. + attn_metadata.compress_plans = self._build_compress_plans( + extend_lens_np, + context_lens_np, + for_decode_cg=True, + ) + + # All updates done in-place on attn_metadata; eagle's + # `for k, v in workinfos.items(): __dict__[k] = v` loop is a no-op. + return {} + def prepare_decode(self, batch: ScheduledBatch, bs: int): """V4-style decode prep: populates positions, cu_seqlens_q, block_tables, and state_slot_mapping. @@ -873,12 +901,25 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): latency behind CPU numpy work: basic H2D copies fire on ``prep_stream`` while ``_build_compress_plans`` runs on the CPU. """ - import numpy as np - var = self.model_runner.forward_vars scheduled_bs = batch.total_seqs_num_decode context_lens_np = np.asarray(batch.context_lens, dtype=np.int32) max_seqlen_q = batch.num_spec_step + 1 + # MTP: roll back ctx by `num_rejected` so this fwd's positions overwrite + # last fwd's rejected-draft slots (matches aiter_mla.py:701 / + # aiter_attention.py:542). `batch.context_lens` = `seq.num_tokens` + # which the scheduler advances by `mtp_k - num_rejected` placeholders + # per fwd (scheduler.py:789); without this rollback, MTP-k positions + # would skip ahead by `num_rejected` and the rejected slots would + # never be overwritten with the corrected K/V. `num_rejected` is None + # on dummy runs and on the first fwd before any sampler output. + # Bound n_committed_csa/hca via the rolled-back ctx (n_committed_* = + # ctx // 4 / 128 in `_attach_v4_paged_decode_meta`), so block_tables + # truncation isn't needed here — the per-token kv_len already shrinks. + if not batch.is_dummy_run and max_seqlen_q > 1: + num_rejected = self.model_runner.tokenID_processor.num_rejected + if num_rejected is not None: + context_lens_np = context_lens_np - num_rejected.astype(np.int32) positions_np = np.tile( np.arange(max_seqlen_q, dtype=np.int32), scheduled_bs ) + np.repeat(context_lens_np - max_seqlen_q, max_seqlen_q) @@ -894,10 +935,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): var["context_lens"].np[:scheduled_bs] = context_lens_np # Inline block_tables CPU fill (H2D deferred to prep_stream). - block_tables_np = var["block_tables"].np - for i, block_table in enumerate(batch.block_tables[:scheduled_bs]): - block_tables_np[i] = 0 - block_tables_np[i, : len(block_table)] = block_table + self.prepare_block_tables(batch) state_slot_np = np.asarray( batch.per_req_cache_groups[:scheduled_bs], dtype=np.int32 @@ -908,6 +946,8 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ss_buf.np[:scheduled_bs] = state_slot_np # ---- fire H2D on prep_stream ---- + # NB: this runs inside attn_metadata_builder.build(), BEFORE + # set_forward_context() — can't read main_stream from the context yet. prep_stream = self.prep_stream current_stream = torch.cuda.current_stream() prep_stream.wait_stream(current_stream) @@ -919,13 +959,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): state_slot_gpu = ss_buf.copy_to_gpu(scheduled_bs) # ---- CPU numpy work, overlapped with prep_stream H2D ---- - start_pos_per_seq_cpu = positions_np[cu_seqlens_q_np[:scheduled_bs]] - extend_lens_np = np.full(scheduled_bs, max_seqlen_q, dtype=np.int32) compress_plans = self._build_compress_plans( extend_lens_np, context_lens_np, - torch.device(self.device), for_decode_cg=True, ) @@ -944,18 +981,17 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): num_cached_tokens=None, block_tables=block_tables_gpu, context_lens=context_lens_gpu, + is_pure_decode=True, ) attn_metadata.state_slot_mapping = state_slot_gpu attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = start_pos_per_seq_cpu attn_metadata.compress_plans = compress_plans padded_bs = int(bs) self._attach_v4_per_fwd_meta( attn_metadata, - positions_np, cu_seqlens_q_np, - start_pos_per_seq_cpu, + extend_lens_np, # = np.full(scheduled_bs, max_seqlen_q) for decode state_slot_np, scheduled_bs, sum_scheduled_tokens, @@ -964,8 +1000,6 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - start_pos_per_seq_cpu, scheduled_bs, sum_scheduled_tokens, positions_gpu=positions, @@ -983,14 +1017,17 @@ def prepare_prefill(self, batch: ScheduledBatch): Also publishes CPU mirrors (`v4_*_cpu`) consumed by the V4 forward path to avoid `.item()` / `.tolist()` syncs (PR-A Phase 2). """ - import numpy as np - base_md, positions = super().prepare_prefill(batch) # Promote to V4 typed metadata so V4-specific attribute assignments # below are well-typed. Safe because AttentionMetaData_DSV4 only adds # fields with defaults; the parent dataclass is non-slotted. base_md.__class__ = AttentionMetaData_DSV4 attn_metadata = cast(AttentionMetaData_DSV4, base_md) + # Prefill is by definition not pure decode (fresh-prefill seqs have + # start_pos == 0). Class-promotion via __class__ doesn't run the + # dataclass __init__, so V4-specific defaults aren't applied — set + # explicitly. + attn_metadata.is_pure_decode = False scheduled_bs = batch.total_seqs_num_prefill if attn_metadata.block_tables is None: attn_metadata.block_tables = self._populate_block_tables( @@ -1007,9 +1044,10 @@ def prepare_prefill(self, batch: ScheduledBatch): positions_np = np.asarray(var["positions"].np[:sum_scheduled_tokens]) cu_seqlens_q_np = np.asarray(var["cu_seqlens_q"].np[: scheduled_bs + 1]) attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[ - cu_seqlens_q_np[:scheduled_bs] - ] + # `start_pos_per_seq` = position of FIRST token of each seq in this fwd. + # Only consumed by `_build_paged_prefill_meta` below; not stashed on + # attn_metadata (no other reader, no inter-fwd reuse). + start_pos_per_seq_np = positions_np[cu_seqlens_q_np[:scheduled_bs]] # Compress plans (per ratio) for batched fused_compress + update_states. # Prefill batch: extend_lens read from cu_seqlens_q_np. # Must run BEFORE `_attach_v4_indexer_meta` (the indexer consumes @@ -1017,13 +1055,16 @@ def prepare_prefill(self, batch: ScheduledBatch): extend_lens_np = ( cu_seqlens_q_np[1 : scheduled_bs + 1] - cu_seqlens_q_np[:scheduled_bs] ).astype(np.int32) - # context_lens for prefill = positions[seq_start_in] + extend_lens - # (= absolute seq_len incl. this fwd's tokens). - context_lens_np = (attn_metadata.start_pos_per_seq_cpu + extend_lens_np).astype( - np.int32 + # context_lens already populated on host by `super().prepare_prefill` + # (backends.py: `var["context_lens"].np[:bs] = batch.context_lens`). + # Mathematically equals `start_pos + extend_lens` but reading the + # canonical buffer avoids drift if scheduler/batch semantics ever + # change. + context_lens_np = np.asarray( + var["context_lens"].np[:scheduled_bs], dtype=np.int32 ) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, positions.device, for_decode_cg=False + extend_lens_np, context_lens_np, for_decode_cg=False ) # Prefill goes through eager (no CG): defaults make padded_total_tokens # collapse to total_tokens — no padding logic kicks in. Must still run @@ -1031,17 +1072,14 @@ def prepare_prefill(self, batch: ScheduledBatch): # reuse the shared GPU tensors (batch_id_per_token, n_committed_csa). self._attach_v4_per_fwd_meta( attn_metadata, - positions_np, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, # = cu_seqlens_q[1:] - cu_seqlens_q[:bs] attn_metadata.state_slot_mapping_cpu, scheduled_bs, sum_scheduled_tokens, ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, scheduled_bs, sum_scheduled_tokens, positions_gpu=positions, @@ -1054,7 +1092,8 @@ def prepare_prefill(self, batch: ScheduledBatch): attn_metadata, positions_np, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, + start_pos_per_seq_np, attn_metadata.state_slot_mapping_cpu, scheduled_bs, sum_scheduled_tokens, @@ -1083,8 +1122,6 @@ def build_ubatch_prefill_metadata( ub_attn.state_slot_mapping = src.state_slot_mapping[rs] if src.state_slot_mapping_cpu is not None: ub_attn.state_slot_mapping_cpu = src.state_slot_mapping_cpu[rs] - if src.start_pos_per_seq_cpu is not None: - ub_attn.start_pos_per_seq_cpu = src.start_pos_per_seq_cpu[rs] var = self.model_runner.forward_vars positions_np = np.asarray(var["positions"].np[ts.start : ts.stop]) @@ -1093,9 +1130,13 @@ def build_ubatch_prefill_metadata( ub_cu -= ub_cu[0] extend_lens_np = (ub_cu[1:] - ub_cu[:ub_num_reqs]).astype(np.int32) - context_lens_np = (ub_attn.start_pos_per_seq_cpu + extend_lens_np).astype( - np.int32 - ) + # Slice the full batch's `var["context_lens"]` (populated by + # `super().prepare_prefill`) for this ubatch — mathematically equals + # `start_pos + extend_lens` but reads from the canonical source. + # `.copy()` so the swap-into-front below doesn't alias. + context_lens_np = np.asarray( + var["context_lens"].np[rs.start : rs.stop], dtype=np.int32 + ).copy() device = src.state_slot_mapping.device from atom.model_ops.v4_kernels import make_compress_plans @@ -1126,9 +1167,8 @@ def build_ubatch_prefill_metadata( self._attach_v4_per_fwd_meta( ub_attn, - positions_np, ub_cu, - ub_attn.start_pos_per_seq_cpu, + extend_lens_np, # ubatch's per-seq token counts ub_attn.state_slot_mapping_cpu, ub_num_reqs, ub_num_tokens, @@ -1137,18 +1177,19 @@ def build_ubatch_prefill_metadata( positions_gpu = var["positions"].gpu[ts.start : ts.stop] self._attach_v4_indexer_meta( ub_attn, - ub_cu, - ub_attn.start_pos_per_seq_cpu, ub_num_reqs, ub_num_tokens, positions_gpu=positions_gpu, ) + # start_pos = position of first token of each seq in this ubatch. + ub_start_pos_per_seq_np = positions_np[ub_cu[:ub_num_reqs]] self._build_paged_prefill_meta( ub_attn, positions_np, ub_cu, - ub_attn.start_pos_per_seq_cpu, + extend_lens_np, + ub_start_pos_per_seq_np, ub_attn.state_slot_mapping_cpu, ub_num_reqs, ub_num_tokens, @@ -1180,9 +1221,8 @@ def build_ubatch_prefill_metadata( def _attach_v4_per_fwd_meta( self, attn_metadata: AttentionMetaData_DSV4, - positions_np: np.ndarray, cu_seqlens_q_np, - start_pos_per_seq_cpu, + token_num_per_seq, state_slot_mapping_cpu, scheduled_bs: int, total_tokens: int, @@ -1197,10 +1237,14 @@ def _attach_v4_per_fwd_meta( them once per fwd saves ~64 redundant constructions for V4-Pro. Sets: - - `attn_metadata.swa_write_indices`: [padded_T] int64 row ids - into per-token KV. Real entries are the last `win` tokens per seq; - tail (entries beyond the per-fwd `num_write`) is sentinel-filled - with -1 so a fixed grid `padded_T` is always safe (CUDAGraph). + - `attn_metadata.swa_write_indices`: [W] int64 row ids into + per-token KV. Real entries are the last `win` tokens per seq. + For decode/MTP (CG-captured): `W = padded_bs * (1 + max_spec_steps)`, + trailing entries `[num_write:W]` get -1 sentinel so the fixed + grid is CUDAGraph-safe. For prefill (eager, no CG): + `W = num_write` exactly — no padding, no wasted grid programs + (long-prefill chunks would otherwise launch up to ~64× more + programs that all bail on `src_id < 0`). - `attn_metadata.batch_id_per_token`: [padded_T] int32 batch id per token (single per-token mapping; consumed by `swa_write`, Phase B/C/E paged-decode kernels, and the indexer). @@ -1218,21 +1262,36 @@ def _attach_v4_per_fwd_meta( """ win = self.window_size + # cu_seqlens_q_arr still needed for write_starts/write_ends below; + # token_num_per_seq is now passed in by caller (== batch.num_scheduled_tokens + # for prepare_decode/prefill, np.ones for MTP draft, etc.) — no longer + # re-derived from cu_seqlens_q here. cu_seqlens_q_arr = np.asarray( - cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int64 - ) - token_num_per_seq = cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:scheduled_bs] - start_pos_per_seq_np = np.asarray( - start_pos_per_seq_cpu[:scheduled_bs], dtype=np.int64 + cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int32 ) - if padded_bs is None or max_q_len is None: - padded_total_tokens = total_tokens + # is_pure_decode is set by the caller at AttentionMetaData_DSV4 + # construction time (single source of truth — prepare_decode / + # prepare_prefill / build_for_cudagraph_capture each know their + # own semantics). Consumed for: padded_total_tokens (CG bucket vs + # eager-tight), arange shortcut (write_indices), SWA grid bucketing, + # and by `_attach_v4_paged_decode_meta` for the index-buffer skip. + is_pure_decode = attn_metadata.is_pure_decode + + # padded_total_tokens: CG-captured decode/MTP pads to the fixed + # bucket `padded_bs * (1+max_spec_steps)` so per-token buffers + # (batch_id_per_token, swa_write_indices) have a stable shape across + # captures. Non-pure-decode (prefill / mixed / fresh-prefill) is + # eager and uses `total_tokens` exactly — no wasted padding (a long + # prefill chunk doesn't need to be padded up to a bucket that + # doesn't exist for it). + if is_pure_decode: + assert padded_bs is not None and max_q_len is not None, ( + "is_pure_decode requires padded_bs + max_q_len from caller " + "(CG bucket size — fixed at capture)" + ) + padded_total_tokens = int(padded_bs) * int(max_q_len) else: - padded_total_tokens = max(int(padded_bs) * int(max_q_len), total_tokens) - - is_decode_only = scheduled_bs == total_tokens and bool( - (token_num_per_seq == 1).all() - ) + padded_total_tokens = total_tokens var = self.model_runner.forward_vars @@ -1242,148 +1301,172 @@ def _attach_v4_per_fwd_meta( np.arange(scheduled_bs, dtype=np.int64), token_num_per_seq ) - ctx_per_seq_np = np.asarray( - var["context_lens"].np[:scheduled_bs], - dtype=np.int64, - ) - n_committed_csa_per_seq_np = (ctx_per_seq_np // 4).astype(np.int32) - - write_starts = cu_seqlens_q_arr[:scheduled_bs] + np.maximum( - 0, token_num_per_seq - win - ) - write_ends = cu_seqlens_q_arr[1:] - if is_decode_only: - write_indices_np = np.arange(total_tokens, dtype=np.int64) + # context_lens is int32 on the buffer; keep dtype through divide so + # n_committed_{csa,hca} stay int32 (max value ~max_model_len // 4 ≪ 2^31). + ctx_per_seq_np = var["context_lens"].np[:scheduled_bs] + # Single source of truth for n_committed_{csa,hca}_per_seq on CPU. + # Stashed on attn_metadata so paged_decode_meta / paged_prefill_meta / + # v4_indexer_meta can read instead of each re-running `ctx // k`. + n_committed_csa_per_seq_np = ctx_per_seq_np // 4 + n_committed_hca_per_seq_np = ctx_per_seq_np // 128 + attn_metadata.n_committed_csa_per_seq_cpu = n_committed_csa_per_seq_np + attn_metadata.n_committed_hca_per_seq_cpu = n_committed_hca_per_seq_np + + # Build write_indices for the SWA-write kernel. + # is_pure_decode ⇒ source is `self._swa_iota` (static GPU tensor) + + # GPU fill_(-1) tail. No CPU numpy needed; `num_write = total_tokens`, + # `swa_write_grid = padded_total_tokens` (CG bucket). + # else ⇒ per-seq concat tail (filter `last win tokens per seq`). + # Built as a fresh local numpy (NOT into the shared pinned + # `wi_buf.np`, which would race with the next fwd's rewrite during + # in-flight DMA). `swa_write_grid = num_write` (tight, no padding — + # prefill is eager). + if is_pure_decode: + write_indices_np = None + num_write = total_tokens + swa_write_grid = padded_total_tokens else: + write_starts = cu_seqlens_q_arr[:scheduled_bs] + np.maximum( + 0, token_num_per_seq - win + ) + write_ends = cu_seqlens_q_arr[1:] write_indices_np = np.concatenate( [ np.arange(s, e, dtype=np.int64) for s, e in zip(write_starts, write_ends) ] ) - num_write = int(write_indices_np.shape[0]) - wi_buf = var["v4_meta_swa_write_indices"] - cap = wi_buf.np.shape[0] - assert ( - padded_total_tokens <= cap - ), f"v4_meta_swa_write_indices too small: need {padded_total_tokens}, have {cap}" - if num_write > 0: - wi_buf.np[:num_write] = write_indices_np - if num_write < padded_total_tokens: - wi_buf.np[num_write:padded_total_tokens].fill(-1) - - # ---- Build window_topk on CPU (avoids ~25 small GPU kernels) ---- - if is_decode_only: - sp_per_token = start_pos_per_seq_np - else: - sp_per_token = np.repeat(start_pos_per_seq_np, token_num_per_seq) - positions_long = positions_np[:total_tokens].astype(np.int64) - window_topk_np = _build_window_topk_np(positions_long, sp_per_token, win) - wt_buf = var["v4_meta_window_topk"] - wt_buf.np[:total_tokens, :win] = window_topk_np + num_write = int(write_indices_np.shape[0]) + swa_write_grid = num_write # ---- Stage all buffers to GPU ---- - window_topk_gpu = wt_buf.copy_to_gpu(total_tokens) - + # window_topk used to be CPU-built here ([T, win] of ring indices with + # -1 sentinels) and staged via v4_meta_window_topk. Now the ring index + # is computed inline inside `write_v4_paged_decode_indices` kernel + # from `var["positions"].gpu` — saves O(T·win) numpy work + 4 MB + # staging buffer. The `positions` H2D is already done by the caller. attn_metadata.batch_id_per_token = self._stage( "v4_batch_id_per_token", batch_id_per_token_np ) - attn_metadata.n_committed_csa_per_seq = self._stage( - "v4_n_committed_csa_per_seq", n_committed_csa_per_seq_np + # Stage n_committed to GPU. For CG-replay safety: aiter + # `top_k_per_row_decode` iterates the CAPTURED grid (= padded_bs * + # next_n rows) and reads `rowEnds[batch_id]` for every row. Its + # per-row length formula is + # `row_len = rowEnds[bid] - next_n + (r % next_n) + 1` + # — for pad rows `bid ∈ [scheduled_bs, padded_bs)` the buffer slot + # carries a stale value from a prior fwd; if that stale value is + # `< next_n - 1` (easy with MTP3 next_n=4 if a prior fwd had a seq + # in early prefill with ctx ≤ 11), row_len becomes negative and the + # kernel's radix loop runs unbounded → GPU hang. The downstream + # `batch_id_per_token = -1` sentinel masks pad rows out of + # `csa_translate_pack`, so the value just needs to be "big enough" + # to keep row_len non-negative. Use `index_topk` (≥ 1024 ≫ next_n). + n_csa_buf = var["v4_n_committed_csa_per_seq"] + n_csa_buf.np[:scheduled_bs] = n_committed_csa_per_seq_np + if is_pure_decode and padded_bs is not None and padded_bs > scheduled_bs: + n_csa_buf.np[scheduled_bs:padded_bs] = self.index_topk + attn_metadata.n_committed_csa_per_seq = n_csa_buf.copy_to_gpu(padded_bs) + else: + attn_metadata.n_committed_csa_per_seq = n_csa_buf.copy_to_gpu(scheduled_bs) + # Race-free write into the shared `v4_meta_swa_write_indices` GPU + # buffer (stable pointer → CG-capture-safe). Bypasses `_stage` + # because `_stage`'s `buf.np[:n] = arr; copy_to_gpu(non_blocking=True)` + # pattern uses the shared pinned alias as the H2D source — the next + # fwd's CPU rewrite of `buf.np` can land mid-DMA and tear the GPU + # contents, producing torn `src_id` values that exceed `kv.shape[0]` + # and MEMORY_VIOLATION the `tl.load(kv_ptr + src_id * ...)` in + # `_swa_write_kernel`. + wi_gpu = var["v4_meta_swa_write_indices"].gpu + assert swa_write_grid <= wi_gpu.shape[0], ( + f"v4_meta_swa_write_indices too small: need {swa_write_grid}, " + f"have {wi_gpu.shape[0]}" ) - attn_metadata.swa_write_indices = wi_buf.copy_to_gpu(padded_total_tokens) + if is_pure_decode: + if num_write > 0: + wi_gpu[:num_write].copy_(self._swa_iota[:num_write]) + if num_write < swa_write_grid: + wi_gpu[num_write:swa_write_grid].fill_(-1) + elif num_write > 0: + wi_gpu[:num_write].copy_(torch.from_numpy(write_indices_np)) + attn_metadata.swa_write_indices = wi_gpu[:swa_write_grid] self._attach_v4_paged_decode_meta( attn_metadata=attn_metadata, token_num_per_seq=token_num_per_seq, - start_pos_per_seq_np=start_pos_per_seq_np, state_slot_mapping_cpu=state_slot_mapping_cpu, - window_topk_gpu=window_topk_gpu, scheduled_bs=scheduled_bs, total_tokens=total_tokens, padded_total_tokens=padded_total_tokens, ) - def _clear_v4_paged_decode_meta(self, attn_metadata) -> None: - """Set all Phase-B paged-decode fields to default (non-decode batches). - - Note: `batch_id_per_token` and `n_committed_csa_per_seq` are - intentionally NOT cleared here — both are also consumed by - `swa_write` / the indexer on prefill / mixed paths and are built - unconditionally in `_attach_v4_per_fwd_meta`. - """ - attn_metadata.is_pure_decode = False - attn_metadata.swa_pages = 0 - for k in ( - "kv_indices_swa", - "kv_indices_csa", - "kv_indices_hca", - "kv_indptr_swa", - "kv_indptr_csa", - "kv_indptr_hca", - ): - setattr(attn_metadata, k, None) - def _attach_v4_paged_decode_meta( self, attn_metadata, token_num_per_seq, - start_pos_per_seq_np, state_slot_mapping_cpu, - window_topk_gpu, scheduled_bs: int, total_tokens: int, padded_total_tokens: Optional[int] = None, ) -> None: """Phase B: build per-fwd paged-decode index buffers (layer-invariant). + All three per-token regions are RAGGED-PACKED — same layout family as + the prefill path (`_build_paged_prefill_meta`). Per-token slot count: + SWA: actual_swa_count = min(positions[t]+1, win) + CSA: actual_swa_count + min(n_committed_csa, index_topk) + HCA: actual_swa_count + n_committed_hca + Writes into stable forward_vars buffers (attn_metadata fields are the V4-namespaced counterparts on `AttentionMetaData_DSV4`): - - kv_indices_swa : SWA paged offsets, uniform stride win + - kv_indices_swa : per-token SWA paged offsets, ragged-packed - kv_indices_csa : SWA prefix written here; CSA compress section left UNINITIALIZED — V4Attention.forward fills it per-layer via csa_translate_pack (Phase C) - kv_indices_hca : SWA prefix + HCA compress section, both fully written (HCA is layer-invariant) - - kv_indptr_{swa,csa,hca} : 3 indptr cumsums (SWA uniform, CSA / - HCA packed; per doc §6.3). Padded tail repeats + - kv_indptr_{swa,csa,hca} : 3 ragged cumsums. Padded tail repeats last value → kv_len=0 sentinel for CG-padded slots. + - skip_prefix_len_csa : per-token actual_swa_count (offset where + csa_translate_pack starts writing CSA topk + within each token's region). Matches prefill + semantics (where it equals prefix_swa_count[t]). Reuses (built earlier in `_attach_v4_per_fwd_meta`): - batch_id_per_token : single per-token mapping (with -1 sentinel) - n_committed_csa_per_seq : per-seq `ctx_len // 4` - - Skipped (defaults via `_clear_v4_paged_decode_meta`) when not - is_pure_decode (prefill, mixed, fresh-prefill). + - var["positions"] : global token positions (already H2D-copied by + the caller; consumed by the index-write kernel + + CPU-side actual_swa_count cumsum here) + + Skipped when not is_pure_decode (prefill, mixed, fresh-prefill). The + Phase-B fields (kv_indices_*, kv_indptr_*, swa_pages) stay at their + dataclass defaults (None / 0) for non-decode batches; downstream + consumers gate on `is_pure_decode`. """ if scheduled_bs == 0 or total_tokens == 0: - self._clear_v4_paged_decode_meta(attn_metadata) - return + return # fields stay at dataclass defaults - # is_pure_decode = uniform tokens-per-seq AND no fresh prefill - # (per doc §7.4). Catches both regular decode (1 tok/seq) and MTP-1 - # (2 tok/seq); rejects mixed batches and fresh-prefill seqs. - tok_first = int(token_num_per_seq[0]) - is_uniform_tok = bool((token_num_per_seq == tok_first).all()) - no_fresh_prefill = bool((start_pos_per_seq_np > 0).all()) - is_pure_decode = is_uniform_tok and no_fresh_prefill + if not attn_metadata.is_pure_decode: + return # caller already marked non-decode; nothing to build - if not is_pure_decode or len(state_slot_mapping_cpu) < scheduled_bs: - self._clear_v4_paged_decode_meta(attn_metadata) + if len(state_slot_mapping_cpu) < scheduled_bs: + # Warmup / dummy_run carve-out: caller asserted pure_decode but + # state_slot_mapping is incomplete. Flip the flag so downstream + # consumers (V4Attention.forward) take the non-decode codepath + # instead of dereferencing the unbuilt kv_indices_* buffers. + attn_metadata.is_pure_decode = False return var = self.model_runner.forward_vars - win = self.window_size - # swa_pages = num_slots * win, layer-invariant; matches the boundary + win = self.window_size # per-token max SWA prefix slots + cs = self.win_with_spec # SWA region per-slot stride (W + mtp_k) + # swa_pages = num_slots * cs, layer-invariant; matches the boundary # between SWA and compress regions in unified_kv (Phase A). - swa_pages = self.model_runner.max_per_req_cache_slots * win + swa_pages = self.model_runner.max_per_req_cache_slots * cs T = total_tokens # ----- Per-seq scalars (CPU numpy) ----- - # context_lens / block_tables CPU mirrors are populated by the caller - # (prepare_decode / build_for_cudagraph_capture) BEFORE this runs. - ctx_per_seq = np.asarray(var["context_lens"].np[:scheduled_bs], dtype=np.int64) # The single per-token mapping. Built once in `_attach_v4_per_fwd_meta` # (so swa_write / indexer can also consume it). Pull the GPU view from # attn_metadata; recompute the np copy here only for cumsum math below. @@ -1392,21 +1475,31 @@ def _attach_v4_paged_decode_meta( ) # [T] int32 — host copy for indptr cumsums batch_id_per_token_gpu = attn_metadata.batch_id_per_token - n_committed_csa_per_seq = ctx_per_seq // 4 # host copy for cumsum below - n_committed_hca_per_seq = (ctx_per_seq // 128).astype(np.int64) - - # ----- 3 indptr cumsums (CPU numpy, per doc §6.3) ----- - # Per-token kv_len = f(per-seq quantity)[batch_id_per_token]. CSA - # length is clamped to index_topk because csa_translate_pack only - # writes that many rows per seq (kernel mask `(k < n_committed) & - # (k < index_topk)`); the host-side clamp here just keeps the indices - # buffer correctly sized — the staged GPU n_committed_csa stays raw. + # Read pre-computed `ctx // {4,128}` from attn_metadata — populated by + # `_attach_v4_per_fwd_meta` (always runs first). int32. + n_committed_csa_per_seq = attn_metadata.n_committed_csa_per_seq_cpu + n_committed_hca_per_seq = attn_metadata.n_committed_hca_per_seq_cpu + + # ----- 3 indptr cumsums (CPU numpy, ragged) ----- + # Per-token kv_len = actual_swa_count + n_compress. CSA length is + # clamped to index_topk because csa_translate_pack only writes that + # many rows per seq (kernel mask `(k < n_committed) & (k < index_topk)`); + # the host-side clamp here just keeps the indices buffer correctly + # sized — the staged GPU n_committed_csa stays raw. index_topk = self.index_topk n_committed_csa_clamped_per_token = np.minimum( n_committed_csa_per_seq[batch_id_per_token_np], index_topk ) n_committed_hca_per_token = n_committed_hca_per_seq[batch_id_per_token_np] + # actual_swa_count[t] = min(positions[t]+1, win). Matches the kernel's + # inline `n = tl.minimum(pos+1, win)` so SWA-prefix segment sizes line + # up perfectly. `var["positions"]` is the int64 CpuGpuBuffer populated + # + H2D-copied by the caller (prepare_decode / build_for_cudagraph_capture). + actual_swa_count_np = np.minimum(var["positions"].np[:T] + 1, win).astype( + np.int32 + ) + # CG-padding-aware T_for_indptr: indptr buffer must size to the # captured kernel grid (= padded_total_tokens) so padded slots see # `kv_len = indptr[t+1] - indptr[t] = 0` and the inner loop bails. @@ -1416,23 +1509,24 @@ def _attach_v4_paged_decode_meta( if T_pad < T: T_pad = T - # SWA: uniform stride for real [0:T+1]. Padded entries [T+1:T_pad+1] - # repeat the last value (T*win) → kv_len = 0 for padded tokens, kernel - # bails out of the inner loop without reading kv_indices. - swa_indptr_np = np.empty(T_pad + 1, dtype=np.int32) - swa_indptr_np[: T + 1] = np.arange(T + 1, dtype=np.int64) * win + # All three indptr cumsums output int32 directly. Values are bounded + # (T ≤ mnbt=8192, per-tok ≤ win + index_topk ≈ 2200 → max cumsum ~18M, + # well within int32). + # SWA: ragged, per-token len = actual_swa_count[t]. + swa_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) + swa_indptr_np[1 : T + 1] = np.cumsum(actual_swa_count_np, dtype=np.int32) if T_pad > T: - swa_indptr_np[T + 1 :].fill(int(T * win)) - # CSA: packed, per-token len = win + min(n_committed_csa, index_topk) - csa_per_tok = win + n_committed_csa_clamped_per_token.astype(np.int64) + swa_indptr_np[T + 1 :].fill(int(swa_indptr_np[T])) + # CSA: ragged, per-token len = actual_swa_count + min(n_csa, index_topk) + csa_per_tok = actual_swa_count_np + n_committed_csa_clamped_per_token csa_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) - csa_indptr_np[1 : T + 1] = np.cumsum(csa_per_tok).astype(np.int32) + csa_indptr_np[1 : T + 1] = np.cumsum(csa_per_tok, dtype=np.int32) if T_pad > T: csa_indptr_np[T + 1 :].fill(int(csa_indptr_np[T])) - # HCA: packed, per-token len = win + n_committed_hca - hca_per_tok = win + n_committed_hca_per_token + # HCA: ragged, per-token len = actual_swa_count + n_committed_hca + hca_per_tok = actual_swa_count_np + n_committed_hca_per_token hca_indptr_np = np.zeros(T_pad + 1, dtype=np.int32) - hca_indptr_np[1 : T + 1] = np.cumsum(hca_per_tok).astype(np.int32) + hca_indptr_np[1 : T + 1] = np.cumsum(hca_per_tok, dtype=np.int32) if T_pad > T: hca_indptr_np[T + 1 :].fill(int(hca_indptr_np[T])) @@ -1446,86 +1540,74 @@ def _attach_v4_paged_decode_meta( block_tables_np_full = var["block_tables"].np[:scheduled_bs] hca_total_indices = int(hca_indptr_np[T]) hca_indices_np = np.full(hca_total_indices, -1, dtype=np.int32) - n_h_per_token = n_committed_hca_per_seq[batch_id_per_token_np[:T]].astype( - np.int64 - ) + # n_committed_hca_per_seq is int32; gather stays int32. + n_h_per_token = n_committed_hca_per_seq[batch_id_per_token_np[:T]] total_hca_entries = int(n_h_per_token.sum()) if total_hca_entries > 0: token_indices = np.repeat(np.arange(T, dtype=np.int32), n_h_per_token) - cu_n_h = np.zeros(T + 1, dtype=np.int64) - np.cumsum(n_h_per_token, out=cu_n_h[1:]) - entry_offsets = np.arange(total_hca_entries, dtype=np.int64) - np.repeat( + cu_n_h = np.zeros(T + 1, dtype=np.int32) + np.cumsum(n_h_per_token, out=cu_n_h[1:], dtype=np.int32) + entry_offsets = np.arange(total_hca_entries, dtype=np.int32) - np.repeat( cu_n_h[:T], n_h_per_token ) + # HCA compress section starts at `actual_swa_count[t]` (was `win` + # under the old uniform layout) — ragged-packed offset matches + # what the kernel writes for the SWA prefix segment. write_pos = ( - hca_indptr_np[token_indices].astype(np.int64) + win + entry_offsets + hca_indptr_np[token_indices] + + actual_swa_count_np[token_indices] + + entry_offsets ) bid_expanded = batch_id_per_token_np[token_indices] hca_indices_np[write_pos] = ( - swa_pages - + block_tables_np_full[bid_expanded, entry_offsets.astype(np.intp)] + swa_pages + block_tables_np_full[bid_expanded, entry_offsets] ).astype(np.int32) # Stage to GPU (HCA compress tail; window prefix scattered below). hca_indices_gpu = self._stage("v4_kv_indices_hca", hca_indices_np) - # ----- SWA paged offsets (GPU-side from window_topk_gpu) ----- - # `state_slot_per_token = state_slot_per_seq[batch_id_per_token]` — - # done as a transient GPU gather (no persistent per-token buffer). - # Reuse `attn_metadata.state_slot_mapping` (= forward_vars[ - # "v4_meta_state_slot_groups"] gpu view, set by `_populate_state_slot_mapping`) - # — same data as the legacy `v4_meta_state_slot_i32` re-staging, no - # second H2D needed. - device = self.device - state_slot_per_seq_gpu = attn_metadata.state_slot_mapping - # Use only the real-data prefix of batch_id_per_token here — the - # padded tail carries -1 sentinels, and `swa_paged_2d` is computed - # against `window_topk_gpu` of shape [T, win] only. - state_slot_per_token_gpu = state_slot_per_seq_gpu[ - batch_id_per_token_gpu[:T].long() - ] # [T] int32 transient - # `window_topk_gpu` [T, win] int32 — passed in from caller - # (`_attach_v4_per_fwd_meta` builds it once per fwd). Builder-internal - # intermediate; not exposed on attn_metadata since no kernel reads it. - swa_paged_2d = torch.where( - window_topk_gpu >= 0, - state_slot_per_token_gpu[:, None] * win + window_topk_gpu, - torch.full_like(window_topk_gpu, -1), - ) # [T, win] int32 - swa_paged_flat = swa_paged_2d.reshape(-1) # [T*win] int32 - - # ----- Write SWA paged offsets to all 3 buffer heads (GPU) ----- - # SWA buffer: uniform stride win → simple flat copy. + # ----- Write SWA / CSA / HCA window-prefix paged offsets (1 kernel) ----- + # Kernel computes `n = min(positions[t]+1, win)` and ring-index + # `(positions[t] - n + 1 + i) % cs` inline — no window_topk staging. + # See `write_v4_paged_decode_indices` docstring and plan + # `sequential-noodling-turing.md` for the motivation. Reads only + # persistent forward_vars buffers — no allocator churn (the prior + # `index_copy_` chain raced under MTP-3 long-prefill; this kernel + # also fixes that, see skill `debug-agent-locate-kernel`). swa_indices_gpu = var["v4_kv_indices_swa"].gpu - swa_indices_gpu[: T * win].copy_(swa_paged_flat) - - # CSA / HCA buffers: window prefix at packed positions - # [indptr[t], indptr[t]+win) — scatter via index_copy_. - win_arange = torch.arange(win, device=device, dtype=torch.int64) - csa_win_pos = ( - csa_indptr_gpu[:T].to(torch.int64).unsqueeze(1) + win_arange.unsqueeze(0) - ).reshape(-1) - hca_win_pos = ( - hca_indptr_gpu[:T].to(torch.int64).unsqueeze(1) + win_arange.unsqueeze(0) - ).reshape(-1) csa_indices_gpu = var["v4_kv_indices_csa"].gpu - csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat) - hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat) - - # ----- skip_prefix_len_csa: decode = full SWA window per token ----- - # csa_translate_pack consumes this to know where the CSA section - # starts within each per-token region. Decode path fills [0:T] with - # `win`; padded tail [T:padded_T] = 0 (kernel bails on bid<0). - # Pre-allocated buffer gives a stable address for CG capture. - skip_csa_gpu = var["v4_skip_prefix_len_csa"].gpu - skip_csa_gpu[:T].fill_(win) - skip_csa_gpu[T:].zero_() + write_v4_paged_decode_indices( + state_slot_per_seq=attn_metadata.state_slot_mapping, + batch_id_per_token=batch_id_per_token_gpu, + positions=var["positions"].gpu, + swa_indptr=swa_indptr_gpu, + csa_indptr=csa_indptr_gpu, + hca_indptr=hca_indptr_gpu, + swa_indices=swa_indices_gpu, + csa_indices=csa_indices_gpu, + hca_indices=hca_indices_gpu, + T=T, + win=win, + cs=cs, + ) + + # ----- skip_prefix_len_csa: per-token actual SWA-prefix length ----- + # csa_translate_pack consumes this as the offset within each token's + # `kv_indices_csa` region where the CSA topk section starts (after + # the SWA prefix segment). Decode + prefill now share this semantics + # (was `win` in decode, `prefix_swa_count[t]` in prefill). + skip_csa_buf = var["v4_skip_prefix_len_csa"] + skip_csa_buf.np[:T] = actual_swa_count_np + skip_csa_buf.np[T:T_pad].fill(0) + skip_csa_gpu = skip_csa_buf.copy_to_gpu(T_pad) # ----- Stash on attn_metadata for V4Attention.forward consumption ----- # batch_id_per_token + n_committed_csa_per_seq already set in # `_attach_v4_per_fwd_meta` (single source of truth, also consumed by # swa_write / indexer outside the is_pure_decode branch). - attn_metadata.is_pure_decode = True - attn_metadata.kv_indices_swa = swa_indices_gpu[: T * win] + # is_pure_decode was set by the caller at AttentionMetaData_DSV4 + # construction time; we only flip it (True→False) above when the + # warmup carve-out fires (incomplete state_slot_mapping_cpu). + attn_metadata.kv_indices_swa = swa_indices_gpu[: int(swa_indptr_np[T])] attn_metadata.kv_indices_csa = csa_indices_gpu[: int(csa_indptr_np[T])] attn_metadata.kv_indices_hca = hca_indices_gpu # already exact len attn_metadata.kv_indptr_swa = swa_indptr_gpu @@ -1539,6 +1621,7 @@ def _build_paged_prefill_meta( attn_metadata: AttentionMetaData_DSV4, positions_np: np.ndarray, cu_seqlens_q_np: np.ndarray, + token_num_per_seq: np.ndarray, start_pos_per_seq_np: np.ndarray, state_slot_mapping_cpu: np.ndarray, scheduled_bs: int, @@ -1579,11 +1662,14 @@ def _build_paged_prefill_meta( - skip_prefix_len_csa = prefix_swa_count_per_token (per-token) - swa_pages """ - if scheduled_bs == 0 or total_tokens == 0: - return + assert scheduled_bs >= 1 and total_tokens >= 1, ( + "scheduled_bs and total_tokens must be positive for prefill meta " + "build (got scheduled_bs={scheduled_bs}, total_tokens={total_tokens})" + ) device = self.device - win = self.window_size + win = self.window_size # per-token topk count + cs = self.win_with_spec # SWA region per-slot stride (W + mtp_k) index_topk = self.index_topk T = total_tokens # warmup_model runs BEFORE allocate_kv_cache binds the paged pool @@ -1594,60 +1680,65 @@ def _build_paged_prefill_meta( num_slots = getattr(self.model_runner, "max_per_req_cache_slots", 0) if num_slots == 0: return - swa_pages = num_slots * win + swa_pages = num_slots * cs + var = self.model_runner.forward_vars # used for block_tables + plan buffers # ----- Per-token quantities (CPU numpy) ----- + # cu_seqlens_q_arr still needed below (ext_cu_q gather); + # token_num_per_seq is now passed in by caller — no longer re-derived. + # All quantities here fit in int32: positions ≤ max_model_len ≪ 2^31, + # counts (win, index_topk) are small, cumsums bounded by T·max_per_tok + # ≈ 18M. Keeping int32 throughout avoids needless widen/narrow churn. cu_seqlens_q_arr = np.asarray( - cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int64 + cu_seqlens_q_np[: scheduled_bs + 1], dtype=np.int32 ) - token_num_per_seq = ( - cu_seqlens_q_arr[1:] - cu_seqlens_q_arr[:scheduled_bs] - ).astype(np.int64) + token_num_per_seq = np.asarray(token_num_per_seq, dtype=np.int32) chunk_start_per_seq = np.asarray( - start_pos_per_seq_np[:scheduled_bs], dtype=np.int64 + start_pos_per_seq_np[:scheduled_bs], dtype=np.int32 ) # batch_id_per_token_np mirrors what _attach_v4_per_fwd_meta computed - # but we need a CPU copy for the cumsum / segment math below. + # but we need a CPU copy for the cumsum / segment math below. CPU-only + # — the GPU copy (int64 for PyTorch fancy index) is staged separately. batch_id_per_token_np = np.repeat( - np.arange(scheduled_bs, dtype=np.int64), token_num_per_seq - ) # [T] int64 - positions_arr = np.asarray(positions_np[:T], dtype=np.int64) + np.arange(scheduled_bs, dtype=np.int32), token_num_per_seq + ) # [T] int32 + positions_arr = np.asarray(positions_np[:T], dtype=np.int32) token_pos_in_chunk = ( positions_arr - chunk_start_per_seq[batch_id_per_token_np] - ) # [T] int64 + ) # [T] int32 # SWA window low bound (clamped at 0); used to derive prefix_swa_count # AND the absolute global positions inside the prefix SWA section. - swa_window_low_global = np.maximum(0, positions_arr - win + 1) # [T] int64 + swa_window_low_global = np.maximum(0, positions_arr - win + 1) # [T] int32 - extend_count_np = np.minimum(token_pos_in_chunk + 1, win).astype(np.int64) + extend_count_np = np.minimum(token_pos_in_chunk + 1, win).astype(np.int32) prefix_swa_count_np = np.maximum( 0, chunk_start_per_seq[batch_id_per_token_np] - swa_window_low_global - ).astype(np.int64) + ).astype(np.int32) - # n_committed_{csa,hca}[bid] from per-seq context_lens (already populated - # by parent). HCA: ctx_len // 128, CSA already exposed as int32 GPU tensor; - # rebuild host copy for cumsum below. - var = self.model_runner.forward_vars - ctx_per_seq_np = np.asarray( - var["context_lens"].np[:scheduled_bs], dtype=np.int64 - ) - n_committed_csa_per_seq_np = ctx_per_seq_np // 4 - n_committed_hca_per_seq_np = ctx_per_seq_np // 128 + # Read pre-computed `ctx // {4,128}` from attn_metadata — populated by + # `_attach_v4_per_fwd_meta` (always runs first). int32. + n_committed_csa_per_seq_np = attn_metadata.n_committed_csa_per_seq_cpu + n_committed_hca_per_seq_np = attn_metadata.n_committed_hca_per_seq_cpu n_csa_per_token = np.minimum( n_committed_csa_per_seq_np[batch_id_per_token_np], index_topk - ) # [T] int64 — clamped because csa_translate_pack writes at most + ) # [T] int32 — clamped because csa_translate_pack writes at most # index_topk per token (kernel mask `(k < n) & (k < index_topk)`) - n_hca_per_token = n_committed_hca_per_seq_np[batch_id_per_token_np] # [T] int64 + n_hca_per_token = n_committed_hca_per_seq_np[batch_id_per_token_np] # [T] int32 # ----- indptr cumsums (CPU) ----- - prefix_swa_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_swa_indptr_np[1:] = np.cumsum(prefix_swa_count_np) - prefix_csa_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_csa_indptr_np[1:] = np.cumsum(prefix_swa_count_np + n_csa_per_token) - prefix_hca_indptr_np = np.zeros(T + 1, dtype=np.int64) - prefix_hca_indptr_np[1:] = np.cumsum(prefix_swa_count_np + n_hca_per_token) - extend_indptr_np = np.zeros(T + 1, dtype=np.int64) - extend_indptr_np[1:] = np.cumsum(extend_count_np) + # All output int32 directly (downstream H2D stages int32 GPU buffers). + prefix_swa_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_swa_indptr_np[1:] = np.cumsum(prefix_swa_count_np, dtype=np.int32) + prefix_csa_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_csa_indptr_np[1:] = np.cumsum( + prefix_swa_count_np + n_csa_per_token, dtype=np.int32 + ) + prefix_hca_indptr_np = np.zeros(T + 1, dtype=np.int32) + prefix_hca_indptr_np[1:] = np.cumsum( + prefix_swa_count_np + n_hca_per_token, dtype=np.int32 + ) + extend_indptr_np = np.zeros(T + 1, dtype=np.int32) + extend_indptr_np[1:] = np.cumsum(extend_count_np, dtype=np.int32) # ----- Extend kv_indices (in `kv` tensor): per-token rows ----- # extend window for token t: kv rows @@ -1655,10 +1746,10 @@ def _build_paged_prefill_meta( # ... cu_seqlens_q[bid] + token_pos_in_chunk[t]]`. Build via segment # expansion. ext_seg_tok, ext_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), extend_count_np + np.arange(T, dtype=np.int32), extend_count_np ) # Pre-gather per-token vars for the segment positions (vectorised). - ext_cu_q = cu_seqlens_q_arr[:scheduled_bs][batch_id_per_token_np] # [T] int64 + ext_cu_q = cu_seqlens_q_arr[:scheduled_bs][batch_id_per_token_np] # [T] int32 ext_indices_np = ( ext_cu_q[ext_seg_tok] + ( @@ -1673,19 +1764,23 @@ def _build_paged_prefill_meta( # For each token's prefix SWA segment: positions # `[swa_window_low_global[t] + k for k in range(prefix_swa_count[t])]`, # paged into `unified_kv` SWA region: - # paged = state_slot[bid] * win + (global_pos % win) + # paged = state_slot[bid] * cs + (global_pos % cs) + # where cs = win_with_spec (per-slot ring size). MUST match the same + # stride/modulo used by `swa_write` and the decode-path + # `_attach_v4_paged_decode_meta`, otherwise prefill prefix reads would + # land in the wrong slot once MTP makes cs > win. swa_seg_tok, swa_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), prefix_swa_count_np + np.arange(T, dtype=np.int32), prefix_swa_count_np ) state_slot_arr = np.asarray( - state_slot_mapping_cpu[:scheduled_bs], dtype=np.int64 + state_slot_mapping_cpu[:scheduled_bs], dtype=np.int32 ) prefix_swa_global_pos = ( swa_window_low_global[swa_seg_tok] + swa_seg_k - ) # [sum prefix_swa_count] int64 + ) # [sum prefix_swa_count] int32 prefix_swa_paged_np = ( - state_slot_arr[batch_id_per_token_np[swa_seg_tok]] * win - + (prefix_swa_global_pos % win) + state_slot_arr[batch_id_per_token_np[swa_seg_tok]] * cs + + (prefix_swa_global_pos % cs) ).astype(np.int32) # ----- HCA compress paged offsets (layer-invariant, fully built here) ----- @@ -1693,7 +1788,7 @@ def _build_paged_prefill_meta( # mapped to `swa_pages + phys` (HCA block_capacity = 1). block_tables_np = var["block_tables"].np[:scheduled_bs] hca_seg_tok, hca_seg_k = _segment_indices( - np.arange(T, dtype=np.int64), n_hca_per_token + np.arange(T, dtype=np.int32), n_hca_per_token ) hca_phys = block_tables_np[ batch_id_per_token_np[hca_seg_tok], hca_seg_k @@ -1741,39 +1836,40 @@ def _build_paged_prefill_meta( # ----- One-shot H2D for everything ----- attn_metadata.kv_indices_extend = torch.from_numpy(ext_indices_np).to(device) - attn_metadata.kv_indptr_extend = torch.from_numpy( - extend_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_extend = torch.from_numpy(extend_indptr_np).to(device) attn_metadata.kv_indices_prefix_swa = torch.from_numpy( kv_indices_prefix_swa_np ).to(device) - attn_metadata.kv_indptr_prefix_swa = torch.from_numpy( - prefix_swa_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_swa = torch.from_numpy(prefix_swa_indptr_np).to( + device + ) attn_metadata.kv_indices_prefix_csa = torch.from_numpy( kv_indices_prefix_csa_np ).to(device) - attn_metadata.kv_indptr_prefix_csa = torch.from_numpy( - prefix_csa_indptr_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_csa = torch.from_numpy(prefix_csa_indptr_np).to( + device + ) attn_metadata.kv_indices_prefix_hca = torch.from_numpy( kv_indices_prefix_hca_np ).to(device) - attn_metadata.kv_indptr_prefix_hca = torch.from_numpy( - prefix_hca_indptr_np.astype(np.int32) - ).to(device) - attn_metadata.skip_prefix_len_csa = torch.from_numpy( - prefix_swa_count_np.astype(np.int32) - ).to(device) + attn_metadata.kv_indptr_prefix_hca = torch.from_numpy(prefix_hca_indptr_np).to( + device + ) + attn_metadata.skip_prefix_len_csa = torch.from_numpy(prefix_swa_count_np).to( + device + ) attn_metadata.swa_pages = swa_pages def _build_compress_plans( - self, extend_lens_np, seq_lens_np, device, *, for_decode_cg: bool + self, extend_lens_np, context_lens_np, *, for_decode_cg: bool ): """Build per-ratio CompressPlan dict consumed by batched compressor. Reuse this from prepare_decode / prepare_prefill / prepare_capture — - caller supplies extend_lens / seq_lens (np int32) and target device. + caller supplies extend_lens / context_lens (np int32). context_lens + is the absolute per-seq length AFTER the new extend tokens (i.e. + prefix + extend); `make_compress_plans` reads it as `context_lens_cpu` + and reconstructs prefix internally. Plan tensors are written into the pre-allocated `v4_compress_plan_{ratio}` / `v4_write_plan_{ratio}` CpuGpuBuffers (fixed pointers for CUDAGraph capture); the kernels skip @@ -1788,11 +1884,16 @@ def _build_compress_plans( if not self._unique_compress_ratios_overlap: return {} - # Ensure inputs are np int32 (callers may pass torch tensors / lists). - if isinstance(extend_lens_np, torch.Tensor): - extend_lens_np = extend_lens_np.cpu().numpy().astype(np.int32) - if isinstance(seq_lens_np, torch.Tensor): - seq_lens_np = seq_lens_np.cpu().numpy().astype(np.int32) + # Inputs MUST be numpy int32 — torch tensors would force a D2H sync. + # Callers are responsible for staging from forward_vars np mirrors. + assert isinstance(extend_lens_np, np.ndarray), ( + f"extend_lens_np must be np.ndarray, got {type(extend_lens_np).__name__} " + "— passing torch.Tensor here would trigger a hidden D2H sync" + ) + assert isinstance(context_lens_np, np.ndarray), ( + f"context_lens_np must be np.ndarray, got {type(context_lens_np).__name__} " + "— passing torch.Tensor here would trigger a hidden D2H sync" + ) var = self.model_runner.forward_vars plan_buffers = { ratio: { @@ -1802,10 +1903,9 @@ def _build_compress_plans( for ratio, _ in self._unique_compress_ratios_overlap } return make_compress_plans( - np.ascontiguousarray(extend_lens_np, dtype=np.int32), - np.ascontiguousarray(seq_lens_np, dtype=np.int32), + extend_lens_np, + context_lens_np, self._unique_compress_ratios_overlap, - device, plan_buffers=plan_buffers, decode_capacity_per_ratio=( self._decode_compress_cap if for_decode_cg else None @@ -1842,8 +1942,6 @@ def _populate_state_slot_mapping( copy is consumed by the V4 forward path to avoid `.tolist()` syncs (PR-A Phase 2). """ - import numpy as np - groups_np = np.asarray( batch.per_req_cache_groups[:scheduled_bs], dtype=np.int32 ) @@ -1886,7 +1984,6 @@ def build_for_cudagraph_capture( sentinel-skips inactive rows internally for both BF16 Main and FP8 Indexer paths.) """ - device = self.model_runner.device var = self.model_runner.forward_vars # Honor MTP at capture time: V4-Pro `mtp_k=1` → 2 tokens/req. The # outer `model_runner.capture_cudagraph` populates cu_seqlens_q with @@ -1922,6 +2019,9 @@ def build_for_cudagraph_capture( block_tables_gpu = var["block_tables"].copy_to_gpu(bs) state_slot_gpu = self._stage("v4_meta_state_slot_groups", state_slot_np) + # Synthetic decode batch: start_pos = win > 0 and uniform + # max_q_len tokens per seq, so is_pure_decode is True by + # construction (capture replays the decode codepath). attn_metadata = AttentionMetaData_DSV4( cu_seqlens_q=cu_seqlens_q_gpu, cu_seqlens_k=None, @@ -1934,25 +2034,24 @@ def build_for_cudagraph_capture( num_cached_tokens=None, block_tables=block_tables_gpu, context_lens=context_lens_gpu, + is_pure_decode=True, ) attn_metadata.state_slot_mapping = state_slot_gpu attn_metadata.state_slot_mapping_cpu = state_slot_np - attn_metadata.start_pos_per_seq_cpu = positions_np[cu_seqlens_q_np[:bs]] # Build compress_plans + per-fwd meta + indexer meta via the same # helpers used at runtime — guarantees addresses match. extend_lens_np = np.full(bs, max_q_len, dtype=np.int32) attn_metadata.compress_plans = self._build_compress_plans( - extend_lens_np, context_lens_np, device, for_decode_cg=True + extend_lens_np, context_lens_np, for_decode_cg=True ) # Capture: padded_bs == scheduled_bs == bs (synthetic batch is full). # Must run BEFORE `_attach_v4_indexer_meta` so the indexer-side meta # builder can reuse the shared per-fwd GPU tensors. self._attach_v4_per_fwd_meta( attn_metadata, - positions_np, cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, + extend_lens_np, # = np.full(bs, max_q_len) — synthetic uniform decode batch attn_metadata.state_slot_mapping_cpu, bs, total_tokens, @@ -1961,8 +2060,6 @@ def build_for_cudagraph_capture( ) self._attach_v4_indexer_meta( attn_metadata, - cu_seqlens_q_np, - attn_metadata.start_pos_per_seq_cpu, bs, total_tokens, positions_gpu=positions, @@ -1986,7 +2083,6 @@ def _alloc_v4_metadata_buffers(self) -> None: Bounds: - per-seq: max_bs - per-token: max_num_batched_tokens - - window_topk: max_num_batched_tokens * window_size - csa compress: max_num_batched_tokens * index_topk - hca compress: max_num_batched_tokens * max_num_blocks_per_seq - csa gather: max_bs * max_num_blocks_per_seq * (block_size // 4) @@ -2016,18 +2112,19 @@ def _alloc_v4_metadata_buffers(self) -> None: # `_populate_state_slot_mapping`); attn_metadata.state_slot_mapping # exposes that GPU view to all downstream consumers (no second # H2D-staged copy). - bufs["v4_meta_start_pos_per_seq"] = CpuGpuBuffer(bs, **i64) - bufs["v4_meta_token_num_per_seq"] = CpuGpuBuffer(bs, **i64) bufs["v4_meta_state_slot_groups"] = CpuGpuBuffer(bs, **i32) - bufs["v4_meta_swa_write_indices"] = CpuGpuBuffer(mnbt, **i64) - # Window-topk batched output buffer (per-token sliding-window indices). - # Sized [mnbt, win] int32 — `_build_window_topk_batched(out=)` writes - # into this; the [total_tokens, win] view feeds `_attach_v4_paged_decode_meta`'s - # SWA paged-offsets derivation. Builder-internal intermediate, but the - # buffer must be CG-capture-stable since it backs a tensor read by the - # captured graph (the SWA paged offsets derived from it land in - # `kv_indices_swa/csa/hca` which the captured kernels do read). - bufs["v4_meta_window_topk"] = CpuGpuBuffer(mnbt, win, **i32) + # swa_write_indices: tight bound = `max_bs * window_size`. Universal + # worst-case across paths — prefill compact write is `sum(min(num_i, + # win)) ≤ bs * win`; decode/MTP CG is `bs * (1+max_spec_steps) ≤ + # bs * win` (since `1+max_spec_steps ≪ win`). Legacy `mnbt` sizing + # was over-padded to the prefill total-token worst case. + bufs["v4_meta_swa_write_indices"] = CpuGpuBuffer(bs * win, **i64) + # Static GPU iota used by `_attach_v4_per_fwd_meta` for the + # is_pure_decode arange shortcut. Pre-allocated once so the per-fwd + # write into `v4_meta_swa_write_indices` is a pure GPU-to-GPU copy + # (no CPU intermediate → no pinned-buffer H2D race) and the source + # pointer stays stable for CUDAGraph capture. + self._swa_iota = torch.arange(bs * win, dtype=torch.int64, device=self.device) # Phase B: paged-decode index buffers (consumed by Phase C/E). # Sized to worst-case decode shape `T = max_bs * (1 + max_spec_steps)` @@ -2100,9 +2197,14 @@ def _alloc_v4_metadata_buffers(self) -> None: # replay MUST use the same value (CG kernel call args are baked). self._decode_compress_cap: dict[int, int] = {} for ratio, is_overlap in self._unique_compress_ratios_overlap: - state_size = (2 if is_overlap else 1) * ratio + # NOTE: this is the pool-window size (algorithm constant), NOT the + # state ring buffer size. The ring buffer is now K_pool + max_spec_steps + 1 + # to avoid R+1 re-commit borrow-reads (see csa_main_state_shape comment), + # but write_plan still emits ≤ K_pool rows per seq per fwd because + # `write_starts = max(0, context_lens - K_pool)` in make_compress_plans. + K_pool = (2 if is_overlap else 1) * ratio max_compress = mnbt // ratio + bs - max_write = min(mnbt, bs * state_size) + max_write = min(mnbt, bs * K_pool) bufs[f"v4_compress_plan_{ratio}"] = CpuGpuBuffer(max_compress, 4, **i32) bufs[f"v4_write_plan_{ratio}"] = CpuGpuBuffer(max_write, 4, **i32) # Pre-fill with sentinel so capture-time buffer state is valid @@ -2127,15 +2229,19 @@ def _stage(self, name: str, arr) -> torch.Tensor: """ buf = self.model_runner.forward_vars[name] n = arr.shape[0] if arr.ndim > 0 else 1 - if n == 0: - return buf.gpu[:0] + assert ( + n > 0 + ), f"Cannot stage empty array for {name!r} — ensure the input array has at least one element." cap = buf.np.shape[0] assert n <= cap, ( f"V4 buffer {name!r} too small: need {n}, have {cap}. " f"Increase the corresponding bound in _alloc_v4_metadata_buffers." ) - if arr.dtype != buf.np.dtype: - arr = arr.astype(buf.np.dtype, copy=False) + assert arr.dtype == buf.np.dtype, ( + f"V4 buffer {name!r} dtype mismatch: buffer is {buf.np.dtype}, " + f"but got arr with dtype {arr.dtype}. Cast arr to the correct " + f"dtype before calling _stage." + ) buf.np[:n] = arr return buf.copy_to_gpu(n) diff --git a/atom/model_ops/sampler.py b/atom/model_ops/sampler.py index b1b5800149..a1f52d47c4 100644 --- a/atom/model_ops/sampler.py +++ b/atom/model_ops/sampler.py @@ -2,7 +2,6 @@ # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. import warnings -from functools import lru_cache import torch from aiter import mixed_sample_outer_exponential diff --git a/atom/model_ops/v4_kernels/__init__.py b/atom/model_ops/v4_kernels/__init__.py index 46d356da1d..da0c9b08c3 100644 --- a/atom/model_ops/v4_kernels/__init__.py +++ b/atom/model_ops/v4_kernels/__init__.py @@ -31,6 +31,10 @@ sparse_attn_v4_paged_prefill_reference, ) from atom.model_ops.v4_kernels.inverse_rope import inverse_rope_inplace +from atom.model_ops.v4_kernels.paged_decode_indices import ( + write_v4_paged_decode_indices, + write_v4_paged_decode_indices_reference, +) from atom.model_ops.v4_kernels.state_writes import update_compressor_states, swa_write __all__ = [ @@ -48,4 +52,6 @@ "make_compress_plans", "inverse_rope_inplace", "scale_indexer_weights", + "write_v4_paged_decode_indices", + "write_v4_paged_decode_indices_reference", ] diff --git a/atom/model_ops/v4_kernels/compress_plan.py b/atom/model_ops/v4_kernels/compress_plan.py index 082c84a5aa..8e7fe2b4b2 100644 --- a/atom/model_ops/v4_kernels/compress_plan.py +++ b/atom/model_ops/v4_kernels/compress_plan.py @@ -57,30 +57,31 @@ class CompressPlan: def make_compress_plans( extend_lens_cpu: np.ndarray, - seq_lens_cpu: np.ndarray, + context_lens_cpu: np.ndarray, unique_ratios_overlap: Iterable[Tuple[int, bool]], - device: torch.device, *, - plan_buffers: dict | None = None, + plan_buffers: dict, decode_capacity_per_ratio: dict[int, int] | None = None, ) -> dict[int, CompressPlan]: """Build a CompressPlan per (ratio, overlap) variant. Args: - extend_lens_cpu: np[bs] int — number of tokens this fwd processes per seq. - seq_lens_cpu: np[bs] int — absolute seq_len per seq (= prefix + extend). + extend_lens_cpu: np[bs] int — number of tokens this fwd processes per seq. + context_lens_cpu: np[bs] int — absolute seq_len per seq AFTER the new + extend tokens (= prefix + extend). Internally + reconstructs prefix via `context_lens - extend_lens`. unique_ratios_overlap: iterable of (ratio, is_overlap) pairs. Typically {(4, True), (128, False)} for V4-Pro; a subset for models with only CSA or only HCA layers. - device: target GPU device for the GPU-side plan tensors. - plan_buffers: optional dict[ratio] -> {"compress": CpuGpuBuffer, + plan_buffers: required dict[ratio] -> {"compress": CpuGpuBuffer, "write": CpuGpuBuffer} of pre-allocated fixed-capacity - plan buffers (CUDAGraph path). When provided, the - function writes into the existing buffers and sentinel- - fills (-1) the trailing rows beyond the actual count, so - the returned `compress_plan_gpu` / `write_plan_gpu` views - have stable data pointers across calls. When None, the - legacy fresh-allocation path is used (eager only). + plan buffers. Function writes into the existing buffers + and sentinel-fills (-1) trailing rows beyond the actual + count, so the returned `compress_plan_gpu` / + `write_plan_gpu` views have stable data pointers across + calls (CUDAGraph requirement). Fresh per-call alloc is + not supported — that pattern caused allocator-churn + races (see `write_v4_paged_decode_indices` docstring). decode_capacity_per_ratio: optional dict[ratio] -> int. Slice length for the returned `compress_plan_gpu`. When PROVIDED: the slice is exactly that fixed value (independent of @@ -95,20 +96,17 @@ def make_compress_plans( buffer is always sized to the prefill worst case. Returns: - dict[ratio] -> CompressPlan. Empty dict if `extend_lens_cpu.sum() == 0` - AND no plan_buffers are provided. With plan_buffers, an empty fwd still - returns CompressPlans pointing at the pre-allocated buffers (fully - sentinel-filled), so capture-time addresses match replay-time addresses - even on a zero-token fwd. + dict[ratio] -> CompressPlan. On empty fwd (`extend_lens_cpu.sum() == 0`) + still returns CompressPlans pointing at the pre-allocated buffers + (fully sentinel-filled), so capture-time addresses match replay-time + addresses even on a zero-token fwd. """ bs = len(extend_lens_cpu) extend_lens_cpu = np.ascontiguousarray(extend_lens_cpu, dtype=np.int32) - seq_lens_cpu = np.ascontiguousarray(seq_lens_cpu, dtype=np.int32) + context_lens_cpu = np.ascontiguousarray(context_lens_cpu, dtype=np.int32) total = int(extend_lens_cpu.sum()) out: dict[int, CompressPlan] = {} if total == 0 or bs == 0: - if plan_buffers is None: - return out # Empty fwd: produce CompressPlans pointing at the pre-allocated # buffers so capture-time addresses match replay-time addresses # even on a zero-token fwd. Skipped via num_*=0. @@ -145,7 +143,7 @@ def make_compress_plans( cu_extend[0] = 0 np.cumsum(extend_lens_cpu, out=cu_extend[1:]) j_in_seq = ragged_ids - cu_extend[batch_ids] - prefix_lens = seq_lens_cpu - extend_lens_cpu + prefix_lens = context_lens_cpu - extend_lens_cpu positions = prefix_lens[batch_ids] + j_in_seq for ratio, is_overlap in unique_ratios_overlap: @@ -172,58 +170,50 @@ def make_compress_plans( # write: tokens whose absolute position falls in the per-seq # "last STATE_SIZE positions" window. STATE_SIZE = K. - # write_start[i] = max(0, seq_lens[i] - K) — uniform across overlap/non-overlap; + # write_start[i] = max(0, context_lens[i] - K) — uniform across overlap/non-overlap; # the SGLang formula `(seq_len // ratio) * ratio - (ratio if overlap else 0)` # is a stricter bound that includes only ratio-aligned writes; the looser - # `seq_len - K` is what ATOM's update_compressor_states already uses + # `context_len - K` is what ATOM's update_compressor_states already uses # (state_writes.py:152-154 docstring) and what the fused kernel's # state-cache reader expects. - write_starts = np.maximum(0, seq_lens_cpu - K).astype(np.int32) + write_starts = np.maximum(0, context_lens_cpu - K).astype(np.int32) write_mask = positions >= write_starts[batch_ids] write_plan = plan_rows[write_mask] n_compress = int(compress_plan.shape[0]) n_write = int(write_plan.shape[0]) - if plan_buffers is not None: - cbuf = plan_buffers[ratio]["compress"] - wbuf = plan_buffers[ratio]["write"] - full_cap = cbuf.np.shape[0] - # CG path: fixed slice (capture / replay must match). Eager - # path: slice = n_compress (smallest possible kernel grid). - cap = ( - decode_capacity_per_ratio.get(ratio) - if decode_capacity_per_ratio is not None - else None - ) - slice_cap = n_compress if cap is None else cap - assert n_compress <= slice_cap <= full_cap, ( - f"ratio={ratio} num_compress={n_compress}, slice={slice_cap}, " - f"buffer={full_cap}: invariant violated. CG path requires " - f"n_compress ≤ decode_cap; eager path uses n_compress as slice." - ) - assert n_write <= wbuf.np.shape[0], ( - f"ratio={ratio} num_write={n_write} exceeds buffer " - f"capacity {wbuf.np.shape[0]}; bump in builder __init__." - ) - if n_compress > 0: - cbuf.np[:n_compress] = compress_plan - # Sentinel only within the slice we hand to the kernel; rows - # beyond `slice_cap` are unreachable from this launch. - if slice_cap > n_compress: - cbuf.np[n_compress:slice_cap].fill(-1) - if n_write > 0: - wbuf.np[:n_write] = write_plan - wbuf.np[n_write:].fill(-1) # sentinel - compress_plan_gpu = cbuf.copy_to_gpu(slice_cap) - write_plan_gpu = wbuf.copy_to_gpu() - else: - compress_plan_gpu = torch.from_numpy( - np.ascontiguousarray(compress_plan) - ).to(device, non_blocking=True) - write_plan_gpu = torch.from_numpy(np.ascontiguousarray(write_plan)).to( - device, non_blocking=True - ) + cbuf = plan_buffers[ratio]["compress"] + wbuf = plan_buffers[ratio]["write"] + full_cap = cbuf.np.shape[0] + # CG path: fixed slice (capture / replay must match). Eager + # path: slice = n_compress (smallest possible kernel grid). + cap = ( + decode_capacity_per_ratio.get(ratio) + if decode_capacity_per_ratio is not None + else None + ) + slice_cap = n_compress if cap is None else cap + assert n_compress <= slice_cap <= full_cap, ( + f"ratio={ratio} num_compress={n_compress}, slice={slice_cap}, " + f"buffer={full_cap}: invariant violated. CG path requires " + f"n_compress ≤ decode_cap; eager path uses n_compress as slice." + ) + assert n_write <= wbuf.np.shape[0], ( + f"ratio={ratio} num_write={n_write} exceeds buffer " + f"capacity {wbuf.np.shape[0]}; bump in builder __init__." + ) + if n_compress > 0: + cbuf.np[:n_compress] = compress_plan + # Sentinel only within the slice we hand to the kernel; rows + # beyond `slice_cap` are unreachable from this launch. + if slice_cap > n_compress: + cbuf.np[n_compress:slice_cap].fill(-1) + if n_write > 0: + wbuf.np[:n_write] = write_plan + wbuf.np[n_write:].fill(-1) # sentinel + compress_plan_gpu = cbuf.copy_to_gpu(slice_cap) + write_plan_gpu = wbuf.copy_to_gpu() out[ratio] = CompressPlan( compress_plan_gpu=compress_plan_gpu, diff --git a/atom/model_ops/v4_kernels/fused_compress.py b/atom/model_ops/v4_kernels/fused_compress.py index a8b52e5f75..3005ab6b0f 100644 --- a/atom/model_ops/v4_kernels/fused_compress.py +++ b/atom/model_ops/v4_kernels/fused_compress.py @@ -103,8 +103,11 @@ def _fused_compress_attn_kernel( HALF_ROPE: tl.constexpr, # = rope_head_dim // 2 OVERLAP: tl.constexpr, RATIO: tl.constexpr, - STATE_SIZE: tl.constexpr, # = 2*RATIO if OVERLAP else RATIO - K: tl.constexpr, # = STATE_SIZE (softmax-pool reduce dim) + STATE_SIZE: tl.constexpr, # ring buffer modulo = kv_state.shape[1] (≥ K_pool; + # for spec decode is K_pool + max_spec_steps + 1 to avoid R+1 re-commit + # borrow-reads of prior round's reject K/V; for non-spec decode is K_pool) + K: tl.constexpr, # pool-window reduce dim (= 2*RATIO if OVERLAP else RATIO); + # ≤ STATE_SIZE; used for `s = position - K + 1 + k_static` loop bound HAS_BLOCK_TABLE: tl.constexpr, QUANT: tl.constexpr, # 0 = raw BF16 (CSA/HCA Main), 1 = FP8 e4m3 + ue8m0 scale (Indexer) USE_UE8M0: tl.constexpr, # round scale to power-of-2 (only when QUANT == 1) @@ -409,12 +412,17 @@ def fused_compress_attn( # Validate shapes dim_full = (2 if overlap else 1) * head_dim - state_size = (2 if overlap else 1) * ratio + K_pool = (2 if overlap else 1) * ratio # pool window size (algorithm-defined) + state_size = kv_state.shape[ + 1 + ] # ring buffer modulo (≥ K_pool; spec path enlarges to K_pool + max_spec_steps + 1) assert ( kv_in.dim() == 2 and kv_in.shape[1] == dim_full ), f"kv_in {kv_in.shape}, expected [*, {dim_full}]" assert score_in.shape == kv_in.shape - assert kv_state.shape[1] == state_size and kv_state.shape[2] == dim_full + assert ( + state_size >= K_pool and kv_state.shape[2] == dim_full + ), f"kv_state {kv_state.shape}, expected [*, ≥{K_pool}, {dim_full}]" assert score_state.shape == kv_state.shape assert ape.shape == (ratio, dim_full) assert rms_weight.shape == (head_dim,) @@ -470,7 +478,7 @@ def fused_compress_attn( BLOCK_D = triton.next_power_of_2(head_dim) HALF_ROPE = rope_head_dim // 2 - K = state_size + K = K_pool # pool window reduce-dim (constexpr; not equal to ring modulo) # Cache-scale args (only consumed by the quant path; pass placeholders # otherwise so the constexpr branch is never taken). @@ -559,8 +567,8 @@ def fused_compress_attn_reference( if plan.num_compress == 0: return None device = kv_in.device - K = (2 if overlap else 1) * ratio - state_size = K + K = (2 if overlap else 1) * ratio # pool window + state_size = kv_state.shape[1] # ring buffer modulo (≥ K) plan_cpu = plan.compress_plan_gpu.detach().cpu() slot_map_cpu = state_slot_mapping.detach().cpu() if block_tables is not None: diff --git a/atom/model_ops/v4_kernels/paged_decode_indices.py b/atom/model_ops/v4_kernels/paged_decode_indices.py new file mode 100644 index 0000000000..1403acbe86 --- /dev/null +++ b/atom/model_ops/v4_kernels/paged_decode_indices.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""V4 paged-decode index scatter — single Triton kernel writes SWA window- +prefix paged offsets into the three ragged-packed destination buffers +(`kv_indices_swa` / `kv_indices_csa` / `kv_indices_hca`). + +Replaces the prior chain (numpy `_build_window_topk_np` + `index_copy_`): + + window_topk_np = _build_window_topk_np(positions, win, cs) # [T, win] + swa_paged_2d = torch.where(window_topk >= 0, slot * cs + topk, -1) + swa_paged_flat = swa_paged_2d.reshape(-1) + swa_indices_gpu[:T*win].copy_(swa_paged_flat) + csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat) + hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat) + +Two simplifications vs the prior implementation (see plan +`sequential-noodling-turing.md` for details): + +1. The ring-index formula `ring = (pos - win + 1 + w) % cs` is computed + inline inside the kernel from `positions[t]`. The `[T, win]` + `window_topk` intermediate buffer (mnbt·win·4 = 4 MB at typical config) + is gone; no separate CPU build + H2D copy. +2. The destination layout is now ragged-packed (same as prefill): each + token's SWA prefix segment has length `n = min(positions[t]+1, win)` + (NOT a fixed `win` padded with `-1` sentinels). The caller's + `swa_indptr` / `csa_indptr` / `hca_indptr` reflect this ragged sizing. + +Bytewise correctness: for tokens with `position >= win-1` (all `n == win`), +the output is identical to the prior implementation. For shorter +positions, the prior layout wrote `(win - n)` leading `-1` entries that +the sparse-attention kernel masked out; the new layout omits those slots +entirely, saving sparse-attn loop iterations. + +Caller contract: +- Grid = T (one program per token). +- `batch_id_per_token[:T]` may carry `-1` sentinels in the CG-padded tail — + kernel checks and bails (matches `_attach_v4_per_fwd_meta` convention). +- `swa_indptr` / `csa_indptr` / `hca_indptr` must reflect the ragged-packed + sizing: per-token slot count = `min(positions[t]+1, win) + n_compress[t]` + where `n_compress[t]` is 0 for SWA, `min(n_committed_csa, index_topk)` + for CSA, `n_committed_hca` for HCA. +- `swa_indices` / `csa_indices` / `hca_indices` capacity ≥ corresponding + indptr[T]; this kernel only writes the SWA-prefix segment + `[indptr[t], indptr[t] + n)` per token. The compress-tail is filled + elsewhere (HCA: numpy fill in caller, CSA: `csa_translate_pack` per layer). +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _v4_paged_decode_indices_kernel( + state_slot_per_seq_ptr, # [bs] int32 + batch_id_per_token_ptr, # [T+pad] int — sentinel -1 in pad tail + positions_ptr, # [T+pad] int — global token position + swa_indptr_ptr, # [T+1] int32 — ragged SWA-prefix cumsum + csa_indptr_ptr, # [T+1] int32 — ragged (SWA + CSA topk) + hca_indptr_ptr, # [T+1] int32 — ragged (SWA + HCA committed) + swa_indices_ptr, # [swa_total] int32, output + csa_indices_ptr, # [csa_total] int32, output (writes SWA-prefix segment only) + hca_indices_ptr, # [hca_total] int32, output (writes SWA-prefix segment only) + cs, # win_with_spec — stride into unified_kv SWA region (paper §3.6.1) + win: tl.constexpr, # window_size — max SWA prefix slots + BLOCK_N: tl.constexpr, # next_pow2(win) +): + """One program per token. Writes `n = min(positions[t]+1, win)` paged + offsets to the SWA prefix segment of each of SWA/CSA/HCA index buffers. + + For token `t`: + bid = batch_id_per_token[t] # bail if -1 (CG pad) + slot = state_slot_per_seq[bid] + pos = positions[t] + n = min(pos + 1, win) + # Old -1 sentinels were at the leading `win - n` cols; reparameterize + # to skip them: i in [0, n) → abs_pos = pos - n + 1 + i ∈ [0, pos]. + for i in range(n): + abs_pos = pos - n + 1 + i + ring = abs_pos % cs + paged = slot * cs + ring + swa_indices[swa_indptr[t] + i] = paged + csa_indices[csa_indptr[t] + i] = paged + hca_indices[hca_indptr[t] + i] = paged + """ + t = tl.program_id(0) + bid = tl.load(batch_id_per_token_ptr + t) + if bid < 0: + return # CG-padded sentinel — leave outputs untouched + + slot = tl.load(state_slot_per_seq_ptr + bid) + pos = tl.load(positions_ptr + t) + # `n` = actual valid SWA prefix count. Cast to match `win` (compile-time + # int) — pos is i32/i64 from positions buffer. + n = tl.minimum(pos + 1, win) + swa_base = tl.load(swa_indptr_ptr + t) + csa_base = tl.load(csa_indptr_ptr + t) + hca_base = tl.load(hca_indptr_ptr + t) + + i = tl.arange(0, BLOCK_N) + mask = i < n + abs_pos = pos - n + 1 + i # ∈ [0, pos] for valid i + ring_idx = abs_pos % cs + paged = slot * cs + ring_idx + + tl.store(swa_indices_ptr + swa_base + i, paged, mask=mask) + tl.store(csa_indices_ptr + csa_base + i, paged, mask=mask) + tl.store(hca_indices_ptr + hca_base + i, paged, mask=mask) + + +def write_v4_paged_decode_indices( + *, + state_slot_per_seq: torch.Tensor, + batch_id_per_token: torch.Tensor, + positions: torch.Tensor, + swa_indptr: torch.Tensor, + csa_indptr: torch.Tensor, + hca_indptr: torch.Tensor, + swa_indices: torch.Tensor, + csa_indices: torch.Tensor, + hca_indices: torch.Tensor, + T: int, + win: int, + cs: int, +) -> None: + """In-place fill SWA / CSA / HCA window-prefix offsets via a single + Triton kernel. Replaces the prior `_build_window_topk_np` (CPU O(T·win)) + + `index_copy_` chain. All inputs are persistent forward_vars buffers — + no allocator churn. + + Args (all GPU tensors except T/win/cs): + state_slot_per_seq: [bs] int32 — per-seq state cache slot. + batch_id_per_token: [>=T] int — token→seq map; -1 sentinel skipped. + positions: [>=T] int — global token position + (forward_vars["positions"]); used to derive + `n = min(pos+1, win)` per token + the ring + index `(pos - n + 1 + i) % cs`. + swa_indptr: [>=T+1] int32 — ragged SWA-prefix cumsum, where + `swa_indptr[t+1] - swa_indptr[t] = + min(positions[t]+1, win)`. + csa_indptr: [>=T+1] int32 — ragged CSA buffer indptr (SWA + prefix + CSA topk per token). + hca_indptr: [>=T+1] int32 — ragged HCA buffer indptr (SWA + prefix + HCA committed per token). + swa_indices: [>=swa_indptr[T]] int32 OUT — fully written by + this kernel (no other source). + csa_indices: [>=csa_indptr[T]] int32 OUT — window-prefix + `[csa_indptr[t], +n)` written here; CSA + topk tail filled per-layer by + `csa_translate_pack`. + hca_indices: [>=hca_indptr[T]] int32 OUT — same semantics; HCA + compress tail filled in the caller via + numpy fill. + T: int — number of real tokens (grid size). + win: int — SWA window size (typically 128 for V4-Pro). + cs: int — `win_with_spec = window_size + max_spec_steps`, + stride into unified_kv SWA region per slot + AND modulo for ring-index wrap. + """ + if T == 0: + return + assert state_slot_per_seq.dim() == 1 + assert batch_id_per_token.dim() == 1 and batch_id_per_token.shape[0] >= T + assert positions.dim() == 1 and positions.shape[0] >= T + assert swa_indptr.dim() == 1 and swa_indptr.shape[0] >= T + 1 + assert csa_indptr.dim() == 1 and csa_indptr.shape[0] >= T + 1 + assert hca_indptr.dim() == 1 and hca_indptr.shape[0] >= T + 1 + assert swa_indices.dim() == 1 + assert csa_indices.dim() == 1 + assert hca_indices.dim() == 1 + + BLOCK_N = triton.next_power_of_2(win) + _v4_paged_decode_indices_kernel[(T,)]( + state_slot_per_seq, + batch_id_per_token, + positions, + swa_indptr, + csa_indptr, + hca_indptr, + swa_indices, + csa_indices, + hca_indices, + cs, + win=win, + BLOCK_N=BLOCK_N, + ) + + +def write_v4_paged_decode_indices_reference( + *, + state_slot_per_seq: torch.Tensor, + batch_id_per_token: torch.Tensor, + positions: torch.Tensor, + swa_indptr: torch.Tensor, + csa_indptr: torch.Tensor, + hca_indptr: torch.Tensor, + swa_indices: torch.Tensor, + csa_indices: torch.Tensor, + hca_indices: torch.Tensor, + T: int, + win: int, + cs: int, +) -> None: + """Pure-PyTorch reference equivalent of `write_v4_paged_decode_indices`. + For unit tests and bisect verification. Mirrors the kernel exactly: + per-token ragged-packed write, no -1 sentinels in output. + """ + if T == 0: + return + bid = batch_id_per_token[:T].long() + pos_t = positions[:T].long() + valid = bid >= 0 + # n = min(pos+1, win) per token; clamp invalid rows to 0 to skip writes. + n_per_tok = torch.minimum(pos_t + 1, torch.full_like(pos_t, win)) + n_per_tok = torch.where(valid, n_per_tok, torch.zeros_like(n_per_tok)) + slot = torch.where( + valid, state_slot_per_seq[bid.clamp(min=0)].long(), torch.zeros_like(bid) + ) + for t in range(T): + n = int(n_per_tok[t].item()) + if n == 0: + continue + p = int(pos_t[t].item()) + s = int(slot[t].item()) + i_arr = torch.arange(n, device=positions.device, dtype=torch.long) + abs_pos = p - n + 1 + i_arr # [n] + ring = abs_pos % cs + paged = (s * cs + ring).to(torch.int32) + swa_base = int(swa_indptr[t].item()) + csa_base = int(csa_indptr[t].item()) + hca_base = int(hca_indptr[t].item()) + swa_indices[swa_base : swa_base + n] = paged + csa_indices[csa_base : csa_base + n] = paged + hca_indices[hca_base : hca_base + n] = paged diff --git a/atom/model_ops/v4_kernels/state_writes.py b/atom/model_ops/v4_kernels/state_writes.py index 03d9f25cca..2ebfbc4577 100644 --- a/atom/model_ops/v4_kernels/state_writes.py +++ b/atom/model_ops/v4_kernels/state_writes.py @@ -9,13 +9,16 @@ Currently implemented: - `swa_write`: writes `swa_kv[state_slot_per_seq[batch_id_per_token[t]], - positions[t] % win, :] = kv[t, :]` for each src row id `t` selected by - `write_indices`. The kernel does ALL gathers (kv row, position, batch id, - state slot) itself — caller passes only stable forward_vars buffers (full - `kv`, full `positions`, full `batch_id_per_token`, per-seq `state_slot`). - Race-free for long prefill via the `write_indices` filter (only the last - `win` tokens per seq selected); CUDAGraph-safe via sentinel skip - (`write_indices[pid] < 0` → bail). + positions[t] % cache_size, :] = kv[t, :]` for each src row id `t` selected + by `write_indices`. The kernel does ALL gathers (kv row, position, batch + id, state slot) itself — caller passes only stable forward_vars buffers + (full `kv`, full `positions`, full `batch_id_per_token`, per-seq + `state_slot`). Race-free for long prefill via the `write_indices` filter + (only the last `cache_size` tokens per seq selected); CUDAGraph-safe via + sentinel skip (`write_indices[pid] < 0` → bail). `cache_size` is + `window_size + max_spec_steps` — for non-MTP this reduces to `window_size` + (no behavioral change); for MTP-k draft tokens get their own ring slots + separate from the verified token's slot. - `update_compressor_states`: unified in-place update of Compressor's per-request `kv_state` + `score_state` ring buffers, covering both prefill (B-side overlap context + tail) and decode (every token at `pos % STATE_SIZE` @@ -31,17 +34,19 @@ - `kv` [T, head_dim] flat — full per-fwd KV (forward_vars). - `write_indices` [W] int — src row ids into `kv` / `positions` / `batch_id_per_token`. Sentinel = -1 → kernel skips. - For long prefill (`seqlen > win`) builder pre-filters - to the last `win` rows per seq to avoid `pos % win` - collisions. For decode/MTP every row is written. + For long prefill (`seqlen > cache_size`) builder + pre-filters to the last `cache_size` rows per seq to + avoid `pos % cache_size` collisions. For decode/MTP + every row is written. - `positions` [T] int — full positions buffer (forward_vars). - `batch_id_per_token` [T] int — Phase-B `v4_batch_id_per_token` mapping; kernel does `state_slot_per_seq[batch_id]` for the per-seq state-cache slot. Single per-token mapping principle (no per-token slot alias). - `state_slot_per_seq` [bs] int — `state_slot_mapping_gpu_i32`. -- `swa_kv` [num_slots, win, head_dim] in-place buffer. -- `win` int sliding-window size (e.g. 128). +- `swa_kv` [num_slots, cache_size, head_dim] in-place buffer. +- `cache_size` int ring-slot count = `window_size + max_spec_steps` + (e.g. 128 + 0 = 128 non-MTP; 128 + 1 = 129 MTP-1). Grid = `write_indices.shape[0]`; each program processes one src row id. """ @@ -58,11 +63,11 @@ def _swa_write_kernel( positions_ptr, # [T] int — full positions batch_id_per_token_ptr, # [T] int — v4_batch_id_per_token state_slot_per_seq_ptr, # [bs] int — state_slot_mapping_gpu_i32 - swa_kv_ptr, # [num_slots, win, head_dim] - swa_kv_slot_stride, # = win * head_dim + swa_kv_ptr, # [num_slots, cache_size, head_dim] + swa_kv_slot_stride, # = cache_size * head_dim swa_kv_pos_stride, # = head_dim head_dim, - win, + cache_size, BLOCK_D: tl.constexpr, ): """One program per write_indices entry. Sentinel skip via write_indices < 0. @@ -78,7 +83,7 @@ def _swa_write_kernel( pos = tl.load(positions_ptr + src_id) bid = tl.load(batch_id_per_token_ptr + src_id) slot = tl.load(state_slot_per_seq_ptr + bid) - ring_idx = pos % win + ring_idx = pos % cache_size d_offsets = tl.arange(0, BLOCK_D) d_mask = d_offsets < head_dim @@ -103,10 +108,11 @@ def swa_write( batch_id_per_token: torch.Tensor, state_slot_per_seq: torch.Tensor, swa_kv: torch.Tensor, - win: int, + cache_size: int, ) -> None: - """In-place write `swa_kv[state_slot_per_seq[bid], pos % win, :] = kv[r, :]` - for each `r = write_indices[pid]` (skip pid where `write_indices[pid] < 0`). + """In-place write + `swa_kv[state_slot_per_seq[bid], pos % cache_size, :] = kv[r, :]` for + each `r = write_indices[pid]` (skip pid where `write_indices[pid] < 0`). Per-token quantities (`pos`, `bid`) are gathered inside the kernel via `positions[r]` / `batch_id_per_token[r]`; per-seq `state_slot_per_seq` @@ -122,21 +128,25 @@ def swa_write( positions: [T] int — full forward_vars["positions"]. batch_id_per_token: [T] int — v4_batch_id_per_token mapping. state_slot_per_seq: [bs] int — per-seq state cache slot. - swa_kv: [num_slots, win, head_dim] in-place ring buffer. - win: sliding-window size. + swa_kv: [num_slots, cache_size, head_dim] in-place ring buffer. + cache_size: ring-slot count = `window_size + max_spec_steps`. + For non-MTP this equals `window_size` and the kernel is bytewise + identical to the pre-MTP behavior. """ assert kv.dim() == 2, f"kv must be [T, D], got {kv.shape}" assert write_indices.dim() == 1 assert positions.dim() == 1 assert batch_id_per_token.dim() == 1 assert state_slot_per_seq.dim() == 1 - assert swa_kv.dim() == 3, f"swa_kv must be [S, W, D], got {swa_kv.shape}" + assert swa_kv.dim() == 3, f"swa_kv must be [S, C, D], got {swa_kv.shape}" T, head_dim = kv.shape assert positions.shape[0] >= T, f"positions {positions.shape[0]} < kv T={T}" assert ( batch_id_per_token.shape[0] >= T ), f"batch_id_per_token {batch_id_per_token.shape[0]} < kv T={T}" - assert swa_kv.shape[1] == win + assert ( + swa_kv.shape[1] == cache_size + ), f"swa_kv ring dim {swa_kv.shape[1]} != cache_size {cache_size}" assert swa_kv.shape[2] == head_dim assert kv.is_contiguous() and swa_kv.is_contiguous() @@ -148,6 +158,7 @@ def swa_write( # block per token covers it. Round up to the next power of two for tl. BLOCK_D = triton.next_power_of_2(head_dim) grid = (W,) + _swa_write_kernel[grid]( kv, write_indices, @@ -158,7 +169,7 @@ def swa_write( swa_kv.stride(0), swa_kv.stride(1), head_dim, - win, + cache_size, BLOCK_D=BLOCK_D, ) @@ -170,7 +181,7 @@ def swa_write_reference( batch_id_per_token: torch.Tensor, state_slot_per_seq: torch.Tensor, swa_kv: torch.Tensor, - win: int, + cache_size: int, ) -> None: """Pure-PyTorch reference equivalent of `swa_write`. For tests / dump-bisect. @@ -185,7 +196,7 @@ def swa_write_reference( src_pos = positions[src_ids] bids = batch_id_per_token[src_ids].long() slots = state_slot_per_seq[bids].long() - ring_idx = src_pos % win + ring_idx = src_pos % cache_size swa_kv[slots, ring_idx] = src_kv @@ -319,10 +330,11 @@ def update_compressor_states( assert kv.dim() == 2 and score.dim() == 2 assert kv.shape == score.shape, f"{kv.shape} vs {score.shape}" assert ape.dim() == 2 and ape.shape[0] == ratio - state_size = (2 if overlap else 1) * ratio + K_pool = (2 if overlap else 1) * ratio # pool window (lower bound) + state_size = kv_state.shape[1] # ring buffer modulo (≥ K_pool) assert ( - kv_state.shape[1] == state_size - ), f"kv_state.shape[1]={kv_state.shape[1]}, expected {state_size}" + state_size >= K_pool + ), f"kv_state.shape[1]={state_size}, must be ≥ K_pool={K_pool}" dim = kv.shape[1] assert write_plan.dim() == 2 and write_plan.shape[1] == 4 assert write_plan.dtype == torch.int32 @@ -379,7 +391,7 @@ def update_compressor_states_reference( `write_plan[i] = (ragged_id, batch_id, position, _)` — each row is one token to write. No mask (host filtered). """ - state_size = (2 if overlap else 1) * ratio + state_size = kv_state.shape[1] # ring buffer modulo (≥ (1+overlap)*ratio) plan_cpu = write_plan.detach().cpu() slot_map_cpu = state_slot_mapping.detach().cpu() for i in range(plan_cpu.shape[0]): diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index c82a1f1d8d..7bb2389416 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -805,8 +805,8 @@ def __init__( # State cache (per paper §3.6.1 "uncompressed tail + B-side overlap # window" portion). Indexed as a single ring buffer of size - # `coff * compress_ratio` by `pos % STATE_SIZE` per token — no - # segment switching, no roll. The `forward` softmax-pool consumer + # `ring_size` (≥ coff * compress_ratio) by `pos % ring_size` per token + # — no segment switching, no roll. The `forward` softmax-pool consumer # resolves A-side (current block) vs B-side (previous block) by # block-id parity (`comp_id % 2`). # @@ -814,9 +814,12 @@ def __init__( # runs before allocate_kv_cache → build_kv_cache_tensor) sees a # valid tensor; afterwards `DeepseekV4AttentionMetadataBuilder. # build_kv_cache_tensor` setattr-replaces these attributes with - # views of the per-request cache pool (shape - # `[max_num_seqs, coff*ratio, coff*head_dim]`). The 1-slot init - # buffers (≈9 MB total across all layers) are GC'd once replaced. + # views of the per-request cache pool whose second dim is the real + # ring_size = coff*ratio + max_spec_steps + 1 (spec) or coff*ratio + # (non-spec). The 1-slot init buffers (≈9 MB total across all layers) + # are GC'd once replaced before any real kernel call, so the + # placeholder's smaller second dim never actually flows through the + # kernel's `state_size >= K_pool` assertion. self.register_buffer( "kv_state", torch.zeros( @@ -1561,17 +1564,21 @@ def _launch_compressors_async(self, x, plan, state_slot_mapping, block_tables): Main Compressor → alt_stream (CSA + HCA). Indexer Compressor → compress_stream (CSA only). Waits resolve instantly: side streams ~25us, main Q/KV chain ~87us.""" - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - with torch.cuda.stream(self.alt_stream): - self.compressor( - x, - plan=plan, - state_slot_mapping=state_slot_mapping, - block_tables=block_tables, - ) + current_stream = get_forward_context().main_stream + if self.compressor is not None and self.alt_stream is not None: + self.alt_stream.wait_stream(current_stream) if self.indexer is not None and self.compress_stream is not None: self.compress_stream.wait_stream(current_stream) + + if self.compressor is not None and self.alt_stream is not None: + with torch.cuda.stream(self.alt_stream): + self.compressor( + x, + plan=plan, + state_slot_mapping=state_slot_mapping, + block_tables=block_tables, + ) + if self.indexer is not None and self.compress_stream is not None: with torch.cuda.stream(self.compress_stream): self.indexer.compressor( x, @@ -1612,7 +1619,24 @@ def forward_impl( if get_forward_context().context.is_dummy_run: return torch.zeros_like(x) num_tokens = x.size(0) - win = self.window_size + # Async-compress (alt_stream main Compressor + compress_stream + # indexer.compressor) is only safe inside CUDAGraph capture: graph + # records the fork-join edges and replay re-uses the same stream + # layout. In eager mode the side-stream launches accumulate across + # 60 layers and deadlock the hipStream queue when the first splitk + # GEMM kernel allocates its workspace — verified hang on both + # small (~800-token) and large (>2k-token) prefill batches in + # eager. Replay does not re-execute this Python code so the flag + # doesn't matter then. + use_async_compress = ( + self._use_async_compress and get_forward_context().in_hipgraph + ) + # SWA ring-slot count per req (= window_size + max_spec_steps for + # MTP-aware cache). Sourced from the bound cache to avoid threading + # `max_spec_steps` through V4Args; for non-MTP this equals + # `self.window_size`. Used as the modulo in `swa_write` (and matches + # the per-row case_c modulo in window_topk). + cache_size = self.swa_kv.shape[1] ratio = self.compress_ratio rd = self.rope_head_dim @@ -1655,7 +1679,8 @@ def forward_impl( plan_for_layer = compress_plans[ratio] if ratio else None # ===== Triple-stream: Q/KV path + both Compressors in parallel ===== - if self._use_async_compress: + current_stream = get_forward_context().main_stream + if use_async_compress: self._launch_compressors_async( x, plan_for_layer, state_slot_mapping, block_tables_gpu ) @@ -1696,7 +1721,7 @@ def forward_impl( rd, self.kv_norm.weight, self.eps, - win, + cache_size, swa_write_indices=write_indices, batch_id_per_token=batch_id, state_slot_mapping=slot_map, @@ -1716,7 +1741,7 @@ def forward_impl( act_quant_inplace(kv[..., :-rd], 64, self.scale_fmt) # ===== Compressor + Indexer ===== - if not self._use_async_compress: + if not use_async_compress: if self.compressor is not None: self.compressor( x, @@ -1732,9 +1757,9 @@ def forward_impl( block_tables=block_tables_gpu, ) if self.indexer is not None: - if self._use_async_compress: - torch.cuda.current_stream().wait_stream(self.alt_stream) - torch.cuda.current_stream().wait_stream(self.compress_stream) + if use_async_compress: + current_stream.wait_stream(self.compress_stream) + current_stream.wait_stream(self.alt_stream) indexer_topk_batched = self.indexer.forward_batched( x_full=x, qr_full=qr, @@ -1747,25 +1772,26 @@ def forward_impl( # - prefill buffer `kv_indices_prefix_csa` (otherwise) # `_fill_csa_paged_compress` dispatches internally on is_pure_decode. self._fill_csa_paged_compress(attn_md, indexer_topk_batched, num_tokens) - elif self._use_async_compress: - torch.cuda.current_stream().wait_stream(self.alt_stream) + elif use_async_compress: + current_stream.wait_stream(self.alt_stream) # ===== Sparse attention dispatch ===== # Two paths over the unified KV pool. The order of `swa_write` vs # `sparse_attn` differs because the two kernels read SWA differently: # # decode (is_pure_decode==True): - # `paged_decode` reads SWA from `unified_kv` ring slot `pos % win`. - # The current decode token's K must be present in the ring before - # attn fires, otherwise the token can't see its own K. So: + # `paged_decode` reads SWA from `unified_kv` ring slot + # `pos % cache_size`. The current decode token's K must be present + # in the ring before attn fires, otherwise the token can't see its + # own K. So: # swa_write → paged_decode # # prefill / mixed (is_pure_decode==False): # `paged_prefill` reads in-chunk K from per-fwd `kv` tensor (extend # region) and prior-chunk K from `unified_kv` ring (prefix region). - # `swa_write` writes the LAST `win` tokens of THIS fwd into ring - # slots `pos % win`, which OVERLAP with prior-chunk slots that the - # prefix SWA region wants to read (chunked prefill). So: + # `swa_write` writes the LAST `cache_size` tokens of THIS fwd into + # ring slots `pos % cache_size`, which OVERLAP with prior-chunk + # slots that the prefix SWA region wants to read (chunked prefill): # paged_prefill → swa_write # Pure prefill (chunk_start==0) has prefix_swa_count==0 so no prior # ring read; swa_write order is irrelevant in that subcase. @@ -1779,7 +1805,7 @@ def forward_impl( v4_batch_id_per_token, state_slot_mapping, self.swa_kv, - win, + cache_size, ) if ratio == 0: kv_indices = attn_md.kv_indices_swa @@ -1824,7 +1850,7 @@ def forward_impl( ) # [S, H, head_dim] # swa_write AFTER attn so chunked-prefill prefix SWA reads see # prior-chunk's ring contents (current swa_write would overwrite - # ring slots `pos % win` for positions in this chunk's tail). + # ring slots `pos % cache_size` for positions in this chunk's tail). swa_write( kv, swa_write_indices, @@ -1832,7 +1858,7 @@ def forward_impl( v4_batch_id_per_token, state_slot_mapping, self.swa_kv, - win, + cache_size, ) # Inverse RoPE on output's rope dims to remove absolute-position @@ -2210,11 +2236,11 @@ def dual_stream_moe_forward( independent; main stream waits on alt_stream's completion before combining. """ - current_stream = torch.cuda.current_stream() + current_stream = get_forward_context().main_stream self.alt_stream.wait_stream(current_stream) + routed = self.routed_expert_forward(x) with torch.cuda.stream(self.alt_stream): shared = self.shared_experts(x) - routed = self.routed_expert_forward(x) current_stream.wait_stream(self.alt_stream) return self.combine_outputs(routed, shared) @@ -2486,89 +2512,15 @@ def forward( return self.get_logits(norm(x)) # [bs, vocab] -class MTPBlock(Block): - """MTP block: V4 dense block + e_proj/h_proj/enorm/hnorm + own hc_head params + LM head. - - Port of inference/model.py:739-767. Subclass of Block reusing all HC + Attention + FFN - machinery; adds a token-embed projection (`e_proj`), a hidden-state projection - (`h_proj`), per-input RMSNorms, and its own `hc_head_fn/base/scale` parameters - for the final LM head reduction. - - `embed` and `head` are assigned externally by `DeepseekV4Model` (shared with - the main model's embedding and LM head). - """ - - def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): - super().__init__(layer_id, args, prefix=prefix) - # e_proj / h_proj are FP8 on disk per index; ATOM Linear with V4QuantConfig - # picks per_1x128 automatically. nn.Linear at construction works for the - # toy/dummy path; for real-checkpoint loading, switch to ReplicatedLinear. - qc = args.quant_config - if qc is None: - self.e_proj = nn.Linear(args.dim, args.dim, bias=False) - self.h_proj = nn.Linear(args.dim, args.dim, bias=False) - else: - self.e_proj = ReplicatedLinear( - args.dim, - args.dim, - bias=False, - quant_config=qc, - prefix=f"{prefix}.e_proj", - ) - self.h_proj = ReplicatedLinear( - args.dim, - args.dim, - bias=False, - quant_config=qc, - prefix=f"{prefix}.h_proj", - ) - self.enorm = RMSNorm(args.dim, args.norm_eps) - self.hnorm = RMSNorm(args.dim, args.norm_eps) - self.norm = RMSNorm(args.dim, args.norm_eps) - # Per-MTP hc_head params (distinct from Block's hc_attn/hc_ffn params). - hc_mult = args.hc_mult - hc_dim = hc_mult * args.dim - self.hc_head_fn = atom_parameter( - torch.empty(hc_mult, hc_dim, dtype=torch.float32) - ) - self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) - self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) - # Externally-assigned by DeepseekV4Model (shared with main model). - self.embed: Optional[nn.Module] = None - self.head: Optional[ParallelHead] = None - - def forward( - self, - x: torch.Tensor, # [num_tokens, hc, dim] residual stream from main model - positions: torch.Tensor, # [num_tokens] int absolute positions - input_ids: torch.Tensor, # [num_tokens] int - ) -> torch.Tensor: # [bs, vocab] last-token logits via self.head - assert ( - self.embed is not None and self.head is not None - ), "MTPBlock requires .embed and .head to be assigned by the parent model" - e = self.enorm(self.embed(input_ids)) # [num_tokens, dim] - x = self.hnorm(x) # [num_tokens, hc, dim] - # Mix token-embed + hidden into a fresh residual. e_proj output is - # [num_tokens, dim]; unsqueeze adds the hc axis so it broadcasts - # against h_proj(x) [num_tokens, hc, dim]. - x = self.e_proj(e).unsqueeze(-2) + self.h_proj(x) # [num_tokens, hc, dim] - x = super().forward(x, positions) # [num_tokens, hc, dim] - return self.head( - x, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm - ) - - @support_torch_compile class DeepseekV4Model(nn.Module): """Full model: embed -> expand to hc_mult copies -> N blocks -> hc_head -> logits. - Port of inference/model.py:Transformer (770-810). MTP blocks are constructed - and have their `.embed` and `.head` linked to the main model's, but they are - NOT called from the main forward path — PR5 will integrate them into ATOM's - EagleProposer via `self.mtp[k].forward(...)` from outside. - - PR1 single-rank: uses plain `nn.Embedding` for `self.embed` (state_dict-compatible - with reference's `ParallelEmbedding` since both store a single `weight` parameter). + Port of inference/model.py:Transformer (770-810). MTP blocks live in the + EagleProposer wrapper (`atom.models.deepseek_v4_mtp.DeepseekV4MTP`), not + on this target. The ckpt's `mtp.*` weights are filtered out at + `loader.load_model(spec_decode=False)` via the auto-detected + `need_load_mtp` flag (target has no `mtp.*` params). """ def __init__( @@ -2613,14 +2565,6 @@ def __init__( self.norm = RMSNorm(args.dim, self.norm_eps) self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) - # MTP blocks: constructed and linked, but only invoked externally (PR5). - self.mtp = nn.ModuleList() - for layer_id in range(args.n_mtp_layers): - blk = MTPBlock(args.n_layers + layer_id, args, prefix=f"mtp.{layer_id}") - blk.embed = self.embed - blk.head = self.head - self.mtp.append(blk) - # Top-level hc_head params used to reduce the final hc_mult residual stack # before the LM head linear projection. hc_mult = args.hc_mult @@ -2635,14 +2579,13 @@ def forward( self, input_ids: torch.Tensor, # [num_tokens] int flat ragged-batch token ids positions: torch.Tensor, # [num_tokens] int abs positions (required) - ) -> ( - torch.Tensor - ): # [num_tokens, dim] hidden_states (vocab head deferred to compute_logits) + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual stream """Forward over `num_tokens` flat ragged-batch tokens. - Returns hidden_states `[num_tokens, dim]` — the vocab projection is - deferred to `compute_logits` to satisfy ATOM's CUDAGraph contract - (capture buffer is sized to hidden_size, not vocab). + Returns the mHC residual stack `[num_tokens, hc, dim]` BEFORE hc_head + reduction — `hc_head + RMSNorm + LM head` are all deferred to + `compute_logits`. Returning the hc-shaped residual lets the (future) + MTP draft consume it without re-expanding from a dim-reduced state. """ assert input_ids.dim() == 1, f"input_ids must be 1D, got {input_ids.shape}" h = self.embed(input_ids) # [num_tokens, dim] @@ -2652,12 +2595,7 @@ def forward( for layer in self.layers: h = layer(h, positions) # [num_tokens, hc, dim] - # Reduce the mHC residual stack (hc_head + final RMSNorm); leave the - # vocab projection to compute_logits. - x_hc = self.head.hc_head( # [num_tokens, dim] - h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base - ) - return self.norm(x_hc) + return h class DeepseekV4ForCausalLM(nn.Module): @@ -2700,7 +2638,6 @@ class DeepseekV4ForCausalLM(nn.Module): "norm.weight": "model.norm.weight", "head.weight": "model.head.weight", "hc_head_": "model.hc_head_", - "mtp.": "model.mtp.", } ) weights_mapping = { @@ -2728,6 +2665,11 @@ def __init__(self, config: Config, prefix: str = "") -> None: # this still works — base spec is QuantType.No. self.args.quant_config = make_v4_quant_config(self.hf_config) self.model = DeepseekV4Model(atom_config=config, args=self.args) + # Tell ModelRunner to size the CG outputs buffer as + # [max_num_batched_tokens, hc_mult, hidden_size] instead of the + # default [max_num_batched_tokens, hidden_size]. forward returns + # the un-reduced mHC residual stack [N, hc, dim]. + self.extra_output_dims: tuple[int, ...] = (self.args.hc_mult,) def forward( self, @@ -2746,12 +2688,21 @@ def forward( def compute_logits( self, - hidden_states: torch.Tensor, # [num_tokens, dim] post hc_head + final norm + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head residual ) -> torch.Tensor: # [bs, vocab] - # Vocab projection is split off from `model.forward` so the latter - # returns hidden_size-shaped tensors — required by ATOM's CUDAGraph - # capture contract (outputs buffer is sized to hidden_size, not vocab). - return self.model.head.get_logits(hidden_states) + # mHC reduce + final RMSNorm + LM head are all here so `model.forward` + # can return the un-reduced [N, hc, dim] residual stream — the future + # MTP draft consumes it directly without re-expanding from a dim-reduced + # state. CG output buffer is sized [N, hc, dim] in ModelRunner via the + # `extra_output_dims = (hc_mult,)` hook on this class. + x = self.model.head.hc_head( + hidden_states, + self.model.hc_head_fn, + self.model.hc_head_scale, + self.model.hc_head_base, + ) + x = self.model.norm(x) + return self.model.head.get_logits(x) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: """Return (param_name, weight_name, expert_id, shard_id) tuples for FusedMoE. diff --git a/atom/models/deepseek_v4_mtp.py b/atom/models/deepseek_v4_mtp.py new file mode 100644 index 0000000000..79eaadbb0c --- /dev/null +++ b/atom/models/deepseek_v4_mtp.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 MTP wrapper for ATOM's EagleProposer. + +Mirrors the V2/V3 + Qwen MTP convention: + +1. The target (`DeepseekV4ForCausalLM`) constructs and loads ONLY the main + stack — no MTP modules, no MTP weights. + +2. This wrapper owns the MTP block(s) and loads `mtp.{i}.*` ckpt entries + into them via the standard `load_model` path (with `spec_decode=True`). + `weights_mapper`, `packed_modules_mapping`, `get_expert_mapping`, and + `remap_mtp_weight_name` are declared so the loader recognizes V4 MTP + ckpt naming + filters out target-only entries silently. + +3. `EagleProposer.load_model` then calls `share_with_target(...)` which + rebinds `self.model.{embed,head}` to point at the target's already-loaded + instances, and propagates them onto each MTPBlock (which requires + `embed`/`head` set externally). No second weight load, no double KV. + +4. V4's MTP block consumes the **un-reduced mHC residual stack** `[N, hc, dim]` + from the target — enabled by deferring `hc_head + RMSNorm + LM head` from + `DeepseekV4Model.forward` to `DeepseekV4ForCausalLM.compute_logits`. The + wrapper's `forward` contract matches ATOM's EagleProposer: + - input `hidden_states` is the target's `[N, hc, dim]` residual + - output is `[N, dim]` post-(MTP block + its own hc_head + norm) +""" + +from typing import Optional + +import torch +from torch import nn + +from atom.config import Config +from atom.model_loader.loader import WeightsMapper +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.linear import ReplicatedLinear +from atom.model_ops.moe import FusedMoE +from atom.model_ops.utils import atom_parameter +from atom.utils.forward_context import get_forward_context + +from .deepseek_v4 import ( + Block, + DeepseekV4Args, + ParallelHead, + make_v4_quant_config, +) + + +class MTPBlock(Block): + """MTP block: V4 dense block + e_proj/h_proj/enorm/hnorm + own hc_head params + LM head. + + Port of inference/model.py:739-767. Subclass of Block reusing all HC + Attention + FFN + machinery; adds a token-embed projection (`e_proj`), a hidden-state projection + (`h_proj`), per-input RMSNorms, and its own `hc_head_fn/base/scale` parameters + for the final LM head reduction. + + `embed` and `head` are assigned externally by the wrapper (shared with + the target's embedding and LM head via `share_with_target`). + """ + + def __init__(self, layer_id: int, args: DeepseekV4Args, prefix: str = ""): + super().__init__(layer_id, args, prefix=prefix) + # e_proj / h_proj are FP8 on disk per index; ATOM Linear with V4QuantConfig + # picks per_1x128 automatically. nn.Linear at construction works for the + # toy/dummy path; for real-checkpoint loading, switch to ReplicatedLinear. + qc = args.quant_config + if qc is None: + self.e_proj = nn.Linear(args.dim, args.dim, bias=False) + self.h_proj = nn.Linear(args.dim, args.dim, bias=False) + else: + self.e_proj = ReplicatedLinear( + args.dim, + args.dim, + bias=False, + quant_config=qc, + prefix=f"{prefix}.e_proj", + ) + self.h_proj = ReplicatedLinear( + args.dim, + args.dim, + bias=False, + quant_config=qc, + prefix=f"{prefix}.h_proj", + ) + self.enorm = RMSNorm(args.dim, args.norm_eps) + self.hnorm = RMSNorm(args.dim, args.norm_eps) + self.norm = RMSNorm(args.dim, args.norm_eps) + # Per-MTP hc_head params (distinct from Block's hc_attn/hc_ffn params). + hc_mult = args.hc_mult + hc_dim = hc_mult * args.dim + self.hc_head_fn = atom_parameter( + torch.empty(hc_mult, hc_dim, dtype=torch.float32) + ) + self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32)) + self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32)) + # Externally-assigned by the wrapper (shared with the target). + self.embed: Optional[nn.Module] = None + self.head: Optional[ParallelHead] = None + + def forward( + self, + x: torch.Tensor, # [num_tokens, hc, dim] residual stream from main model + positions: torch.Tensor, # [num_tokens] int absolute positions + input_ids: torch.Tensor, # [num_tokens] int + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual + """Run one MTP step. Returns the un-reduced mHC residual stack + `[num_tokens, hc, dim]` — same shape contract as + `DeepseekV4Model.forward`. The hc_head reduction + RMSNorm + LM head + are all deferred to `DeepseekV4MTP.compute_logits` so the wrapper's + forward output can be fed back in as `x` for the NEXT MTP draft + step (mtp_k > 1) without re-expanding from a `[N, dim]` post-reduction + state. + """ + assert ( + self.embed is not None + ), "MTPBlock requires .embed to be assigned by the wrapper" + e = self.enorm(self.embed(input_ids)) # [num_tokens, dim] + x = self.hnorm(x) # [num_tokens, hc, dim] + # Mix token-embed + hidden into a fresh residual. h_proj is FP8 + # (V4QuantConfig → ReplicatedLinear over `gemm_a8w8_blockscale_preshuffle`), + # which only supports 2D `[M, K]` input. Flatten the hc axis into the + # batch dim and reshape back after the projection. e_proj's input is + # already 2D [num_tokens, dim]; unsqueeze adds the hc axis so it + # broadcasts over h_proj_out's hc dim. + n_tok, hc, d = x.shape + h_proj_out = self.h_proj(x.reshape(n_tok * hc, d)).reshape(n_tok, hc, d) + x = self.e_proj(e).unsqueeze(-2) + h_proj_out # [num_tokens, hc, dim] + return super().forward(x, positions) # [num_tokens, hc, dim] + + +class DeepseekV4MTPModel(nn.Module): + """V4 MTP inner model: owns the MTP blocks. Each block's `embed` / `head` + are set externally by `DeepseekV4MTP.share_with_target` to point at the + already-loaded target instances; this module itself holds no embed/head. + """ + + def __init__(self, atom_config: Config, args: DeepseekV4Args) -> None: + super().__init__() + # ModelRunner reads `drafter.model.model.mtp_start_layer_idx` to bind + # the draft attention to KV slots `[n_layers, n_layers + n_mtp)`. + self.mtp_start_layer_idx = args.n_layers + # Real MTP blocks loaded via the standard load_model path with + # spec_decode=True. Each block's `.embed` / `.head` are set by + # share_with_target after load. + self.mtp = nn.ModuleList( + [ + MTPBlock(args.n_layers + i, args, prefix=f"mtp.{i}") + for i in range(args.n_mtp_layers) + ] + ) + + def forward( + self, + input_ids: torch.Tensor, # [num_tokens] int + positions: torch.Tensor, # [num_tokens] int + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head from target + spec_step_idx: int = 0, + ) -> torch.Tensor: # [num_tokens, hc, dim] pre-hc_head residual + """Returns the un-reduced mHC residual `[N, hc, dim]` from the + selected MTP block — same shape contract as `DeepseekV4Model.forward`. + EagleProposer feeds this back in as `hidden_states` for the NEXT + draft step (mtp_k > 1).""" + idx = spec_step_idx % len(self.mtp) + return self.mtp[idx](hidden_states, positions, input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, # [num_tokens, hc, dim] pre-hc_head residual + spec_step_idx: int = 0, + ) -> torch.Tensor: # [bs, vocab] + """Mirror `DeepseekV4ForCausalLM.compute_logits`: hc_head + RMSNorm + + LM head all here so MTPBlock.forward can return the un-reduced + residual stack. Each MTP block has its own `hc_head_fn/base/scale` + and own `norm` (the LM head is target-shared via `share_with_target`).""" + idx = spec_step_idx % len(self.mtp) + blk = self.mtp[idx] + x = blk.head.hc_head( + hidden_states, blk.hc_head_fn, blk.hc_head_scale, blk.hc_head_base + ) # [num_tokens, dim] + x = blk.norm(x) + return blk.head.get_logits(x) + + +class DeepseekV4MTP(nn.Module): + """Top-level V4 MTP wrapper. Owns its own MTP weights (loaded via the + standard `load_model` path with `spec_decode=True`) and shares only + `embed` / `head` with the target through `share_with_target`. + """ + + # Disk `mtp.{i}.*` -> wrapper param `model.mtp.{i}.*`. Prefix-anchored to + # avoid the `mtp.` substring colliding with anything inside MTPBlock subtree. + weights_mapper = WeightsMapper(orig_to_new_prefix={"mtp.": "model.mtp."}) + # Same on-disk -> internal name conventions as the target (V4 ckpt quirks). + weights_mapping = { + ".gate.bias": ".gate.e_score_correction_bias", + ".scale": ".weight_scale_inv", + } + # Same packed-module fusions as the target — MTPBlock subclasses Block, so + # the same attention / shared-experts fused param layouts apply. + packed_modules_mapping = { + "attn.wq_a": ("attn.wqkv_a", 0), + "attn.wkv": ("attn.wqkv_a", 1), + "compressor.wkv": ("compressor.wkv_gate", 0), + "compressor.wgate": ("compressor.wkv_gate", 1), + "shared_experts.w1": ("shared_experts.gate_up_proj", 0), + "shared_experts.w3": ("shared_experts.gate_up_proj", 1), + } + + def __init__(self, config: Config, prefix: str = "") -> None: + super().__init__() + self.atom_config = config + self.hf_config = config.hf_config + self.args = DeepseekV4Args.from_hf_config(self.hf_config) + self.args.quant_config = make_v4_quant_config(self.hf_config) + self.model = DeepseekV4MTPModel(atom_config=config, args=self.args) + + def remap_mtp_weight_name(self, name: str) -> str | None: + """Filter loader input to MTP-only weights. + + Called per ckpt entry AFTER `weights_mapper` rewrites `mtp.` -> + `model.mtp.`. Target-only entries (`embed.weight`, `layers.X.*`, + `head.weight`, `hc_head_*`, `norm.weight`) pass through unchanged + and have no matching wrapper param, which would otherwise generate + loud `dropped_ckpt_keys` warnings. Returning None drops them silently. + """ + return name if "mtp." in name else None + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + """FusedMoE expert param mapping for MTPBlock's MoE layer. Same + ckpt convention as the target (`ffn.experts.{e}.w{1,2,3}`). + """ + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.args.n_routed_experts, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # Hash MoE routing inside MTPBlock's MoE looks at this — same as + # `DeepseekV4ForCausalLM.forward`. + get_forward_context().context.input_ids = input_ids + return self.model(input_ids, positions, hidden_states, spec_step_idx) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> Optional[torch.Tensor]: + return self.model.compute_logits(hidden_states, spec_step_idx) + + def share_with_target(self, target_base: nn.Module, loaded: set[str]) -> None: + """Bind embed/head on each MTPBlock to the already-loaded target's + instances. MTPBlock requires `.embed` / `.head` set externally before + its first forward; `compute_logits` then reaches them via `mtp[0].head`. + """ + for blk in self.model.mtp: + blk.embed = target_base.model.embed + blk.head = target_base.model.head diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index d729701b94..5d86c62fd0 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -18,6 +18,7 @@ support_eagle_model_arch_dict = { "DeepSeekMTPModel": "atom.models.deepseek_mtp.DeepSeekMTP", + "DeepseekV4MTPModel": "atom.models.deepseek_v4_mtp.DeepseekV4MTP", "Qwen3NextMTPModel": "atom.models.qwen3_next_mtp.Qwen3NextMTP", "MiMoV2FlashMTPModel": "atom.models.mimo_v2_flash_mtp.MiMoV2FlashMTP", "Qwen3_5MTPModel": "atom.models.qwen3_5_mtp.Qwen3_5MTP", @@ -260,6 +261,15 @@ def load_model(self, target_model: nn.Module) -> None: # Resolve the base model (unwrap multimodal wrapper if present) target_base = getattr(target_model, "language_model", target_model) + # Model-specific share hook escape valve. Models whose embed/lm_head + # naming doesn't match the standard `model.embed_tokens` / + # `lm_head` convention (e.g. DeepSeek-V4 uses `model.embed` / + # `model.head`) implement `share_with_target(target_base)` to do + # their own setattr-rebinding and short-circuit the default path. + if hasattr(self.model, "share_with_target"): + self.model.share_with_target(target_base, loaded) + return + # Share embed_tokens with the target model if ( get_pp_group().world_size == 1 @@ -354,6 +364,13 @@ def propose( if draft_uses_mha: attn_metadata.slot_mapping = var["slot_mapping"].gpu[: len(input_ids)] + # Backends that expose flat per-seq kv_indices/kv_indptr (MLA, MHA) + # wire them through eagle's mid-step block; V4 has block_tables + + # context_lens instead (its v4_kv_indices_{swa,csa,hca} are per-token + # non-equivalent). Hoisted out of the loop so the value is bound for + # every iteration (used at i>=1 too, even though i==0 sets it). + has_flat_kv = "kv_indices" in var + for i in range(self.mtp_k): with record_function(f"draft[{i}/{self.mtp_k} bs={bs}]"): model_output = self.model( @@ -388,16 +405,17 @@ def propose( if i == 0: i0_max_seqlen_q = attn_metadata.max_seqlen_q attn_metadata.max_seqlen_q = 1 - kv_indptr = var["kv_indptr"].gpu[: bs + 1] - kv_indices = var["kv_indices"].gpu slot_mapping = var["slot_mapping"].gpu[ : bs * attn_metadata.max_seqlen_q ] cu_seqlens_q = var["cu_seqlens_q"].gpu[: bs + 1] - attn_metadata.kv_indptr = kv_indptr - attn_metadata.kv_indices = kv_indices attn_metadata.cu_seqlens_q = cu_seqlens_q attn_metadata.slot_mapping = slot_mapping + if has_flat_kv: + kv_indptr = var["kv_indptr"].gpu[: bs + 1] + kv_indices = var["kv_indices"].gpu + attn_metadata.kv_indptr = kv_indptr + attn_metadata.kv_indices = kv_indices if target_uses_mla: kv_last_page_lens = var["kv_last_page_lens"].gpu[:bs] attn_metadata.kv_last_page_lens = kv_last_page_lens @@ -410,7 +428,7 @@ def propose( "sparse_kv_indptr" ].gpu[: bs + 1] cu_seqlens_q[: bs + 1] = self.arrange_bs[: bs + 1] - if target_uses_mla: + if target_uses_mla and has_flat_kv: # MLA: block_size=1, kv_indptr tracks tokens kv_indptr[1 : bs + 1] -= torch.cumsum( num_reject_tokens, dim=0 @@ -436,7 +454,12 @@ def propose( ) for k, v in workinfos.items(): attn_metadata.__dict__[k] = v - slot_mapping[:] = kv_indices[kv_indptr[1 : bs + 1] - 1] + if has_flat_kv: + # MLA/MHA path: slot derived from flat kv_indices. + # V4 doesn't expose flat kv_indices and its kernels + # don't read attn_metadata.slot_mapping (state-ring + + # swa_write_indices instead), so the update is skipped. + slot_mapping[:] = kv_indices[kv_indptr[1 : bs + 1] - 1] input_ids = new_draft_ids positions += 1 diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py index 77b28fa490..cdac102e2d 100644 --- a/atom/utils/forward_context.py +++ b/atom/utils/forward_context.py @@ -336,6 +336,28 @@ class ForwardContext: ubatch_slices: Optional[list[Any]] = None + # Cached current_stream() captured at set_forward_context() time, so + # downstream code (V4 attention / MoE / metadata builder) doesn't have + # to query torch.cuda.current_stream() repeatedly during a forward — + # multiple call sites caching independent Stream handles was widening + # the hipStream handle pool and complicating reasoning about which + # logical stream each wait_stream() refers to. CG capture / TBO + # threads each call set_forward_context() inside their own stream + # context, so the cached value is correct for the captured graph or + # active thread. + main_stream: Optional[torch.cuda.Stream] = None + + # True only while the model forward runs inside a CUDAGraph capture + # block (model_runner.capture_model loop). Components that gate + # multi-stream side-launches (V4 main Compressor on alt_stream, + # indexer.compressor on compress_stream) check this flag: side-stream + # work is safe to emit inside a captured graph (graph records the + # fork-join edges and replay re-uses the same stream layout) but + # racy in eager mode where launches accumulate across layers and + # deadlock the hipStream queue. Replay does not re-execute Python + # forward, so it ignores the flag entirely. + in_hipgraph: bool = False + def __post_init__(self): if not hasattr(self, "no_compile_layers") or self.no_compile_layers is None: self.no_compile_layers = {} @@ -346,6 +368,10 @@ def __post_init__(self): _forward_context: Optional[ForwardContext] = ForwardContext() _forward_kv_cache_context: Optional[ForwardContext] = ForwardContext() +# Cached once at module import — CUDA availability does not change at +# runtime, so we don't pay torch.cuda.is_available() per set_forward_context(). +_CUDA_AVAILABLE: bool = torch.cuda.is_available() + # Thread-local storage for TBO dual-thread execution _forward_context_local = threading.local() @@ -373,6 +399,7 @@ def set_forward_context( num_tokens_across_dp: Optional[torch.Tensor] = None, spec_decode_metadata: Optional[SpecDecodeMetadata] = None, ubatch_slices: Optional[list[Any]] = None, + in_hipgraph: bool = False, ) -> None: global _forward_context dp_metadata: Optional[DPMetadata] = None @@ -392,6 +419,8 @@ def set_forward_context( dp_metadata=dp_metadata, spec_decode_metadata=spec_decode_metadata, ubatch_slices=ubatch_slices, + main_stream=(torch.cuda.current_stream() if _CUDA_AVAILABLE else None), + in_hipgraph=in_hipgraph, ) # _forward_context.attn_metadata = attn_metadata # _forward_context.no_compile_layers = atom_config.compilation_config.static_forward_context # _forward_context = ForwardContext(no_compile_layers=atom_config.compilation_config.static_forward_context, attn_metadata=attn_metadata) diff --git a/scripts/run_benchmark.sh b/scripts/run_benchmark.sh new file mode 100755 index 0000000000..8d9174669e --- /dev/null +++ b/scripts/run_benchmark.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Performance benchmark for ATOM serving +# Usage: bash run_benchmark.sh [MODEL] [PORT] [ISL] [OSL] [CONCURRENCY] [PROMPT_MULTIPLIER] [PROFILE] +# +# Examples: +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 1024 128 # normal +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 1024 64 2 1 # profile trace (conc*1 requests) +# bash run_benchmark.sh /data/DeepSeek-R1-0528 8000 1024 8192 32 + +set -euo pipefail + +MODEL="${1:-/data/DeepSeek-R1-0528}" +PORT="${2:-8000}" +ISL="${3:-1024}" +OSL="${4:-1024}" +CONCURRENCY="${5:-128}" +PROMPT_MULTIPLIER="${6:-10}" +PROFILE="${7:-0}" +_n=$(( $# < 7 ? $# : 7 )) +shift "$_n" 2>/dev/null || true +EXTRA_ARGS="$*" +NUM_PROMPTS=$(( CONCURRENCY * PROMPT_MULTIPLIER )) +RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO:-0.8}" +OUTPUT_DIR="/app/logs_claude" + +echo "========================================" +echo " ATOM Performance Benchmark" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: http://localhost:${PORT}" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Concurrency: $CONCURRENCY" +echo " Num Prompts: $NUM_PROMPTS (x${PROMPT_MULTIPLIER})" +echo " Range Ratio: $RANDOM_RANGE_RATIO" +echo " Profile: $PROFILE" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo "========================================" + +# Wait for server ready (dual check: HTTP + GPU VRAM) +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + VRAM_COUNT=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$VRAM_COUNT" -gt 0 ]; then + echo "Server is ready! (GPU VRAM loaded on $VRAM_COUNT GPUs)" + break + fi + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Server not ready after 120s" + exit 1 + fi + [ $((i % 10)) -eq 0 ] && echo " Waiting... (${i}s)" + sleep 1 +done + +LOG_FILE="${OUTPUT_DIR}/benchmark.log" +RESULT_FILE="${OUTPUT_DIR}/benchmark.json" + +# Write config header to log +{ +echo "========================================" +echo " ATOM Performance Benchmark" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: http://localhost:${PORT}" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Concurrency: $CONCURRENCY" +echo " Num Prompts: $NUM_PROMPTS (x${PROMPT_MULTIPLIER})" +echo " Range Ratio: $RANDOM_RANGE_RATIO" +echo " Profile: $PROFILE" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo " Date: $(date)" +echo "========================================" +} > "$LOG_FILE" + +python -m atom.benchmarks.benchmark_serving \ + --model="$MODEL" --backend=vllm --base-url="http://localhost:${PORT}" \ + --dataset-name=random \ + --random-input-len="$ISL" --random-output-len="$OSL" \ + --random-range-ratio="$RANDOM_RANGE_RATIO" \ + --max-concurrency="$CONCURRENCY" \ + --num-prompts="$NUM_PROMPTS" \ + --trust-remote-code \ + --num-warmups=$((CONCURRENCY * 2)) \ + --request-rate=inf --ignore-eos \ + --save-result \ + --result-filename="$RESULT_FILE" \ + --percentile-metrics="ttft,tpot,itl,e2el" \ + $([ "$PROFILE" = "1" ] && echo "--profile") \ + ${EXTRA_ARGS:-} \ + 2>&1 | tee -a "$LOG_FILE" diff --git a/scripts/run_benchmark_sweep.sh b/scripts/run_benchmark_sweep.sh new file mode 100755 index 0000000000..5f54809a45 --- /dev/null +++ b/scripts/run_benchmark_sweep.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Run benchmark across multiple concurrency levels +# Usage: bash run_benchmark_sweep.sh MODEL PORT ISL OSL [CONC_LIST] +# +# Examples: +# bash run_benchmark_sweep.sh /data/Kimi-K2.5-MXFP4 8000 1024 1024 +# bash run_benchmark_sweep.sh /data/Kimi-K2.5-MXFP4 8000 1024 1024 "4 8 16 32 64 128" +# bash run_benchmark_sweep.sh /data/DeepSeek-R1-0528 8000 1024 1024 "1 2 4 8 16 32 64 128 256" + +set -uo pipefail + +MODEL="${1:?Usage: $0 MODEL PORT ISL OSL [CONC_LIST]}" +PORT="${2:-8000}" +ISL="${3:-1024}" +OSL="${4:-1024}" +CONC_LIST="${5:-4 8 16 32 64 128}" + +echo "========================================" +echo " Benchmark Sweep" +echo "========================================" +echo " Model: $MODEL" +echo " ISL/OSL: ${ISL}/${OSL}" +echo " Conc: $CONC_LIST" +echo "========================================" + +for CONC in $CONC_LIST; do + echo "" + echo "================================================================" + echo "=== CONC=$CONC ===" + echo "================================================================" + bash /app/ATOM/scripts/run_benchmark.sh "$MODEL" "$PORT" "$ISL" "$OSL" "$CONC" +done + +echo "" +echo "=== Sweep complete ===" diff --git a/scripts/run_debug_agent.sh b/scripts/run_debug_agent.sh new file mode 100755 index 0000000000..67bd81e628 --- /dev/null +++ b/scripts/run_debug_agent.sh @@ -0,0 +1,82 @@ +#!/bin/bash +# Launch ATOM under rocm-debug-agent so a GPU memory fault dumps wave state + +# code-object disassembly. Intended for debugging async kernel races (Memory +# access fault, MEMORY_VIOLATION, ASSERT_TRAP, silent GPU livelocks). +# +# Two modes: +# +# Server mode (default): +# bash scripts/run_debug_agent.sh [MODEL] [TP] [PORT] [EXTRA_ARGS...] +# MODEL default /data/DeepSeek-V4-Pro +# TP default 8 +# PORT default 8000 +# EXTRA_ARGS forwarded to start_atom_server.sh +# +# Offline simple_inference mode: +# bash scripts/run_debug_agent.sh --simple [MODEL] [TP] [EXTRA_ARGS...] +# MODEL default /data/DeepSeek-V4-Pro +# TP default 8 +# (no PORT — simple_inference is offline) +# EXTRA_ARGS forwarded to start_simple_inference.sh +# Default log path: /app/logs_claude/simple_inference_debug_agent.log +# (override with LOG_FILE=...). +# +# REQUIRED in EXTRA_ARGS for either mode: --enforce-eager --level 0 +# (graph mode + Inductor are incompatible with the agent's no-caching +# allocator). The script does NOT inject these — pass them via EXTRA_ARGS so +# the model-specific launch matches your repro. +# +# Output: the agent prints wave dumps to stderr; both launchers redirect into +# /app/logs_claude/. Code objects (one per faulting kernel, +# ~4 MB each) land in /app/logs_claude/debug_run/. +# +# After fault triggers in server mode, the agent keeps running. Stop with: +# bash scripts/stop_atom_server.sh +# In simple mode the launcher exits when generation finishes (or faults). + +set -uo pipefail + +MODE="server" +if [ "${1:-}" = "--simple" ]; then + MODE="simple" + shift +fi + +MODEL="${1:-/data/DeepSeek-V4-Pro}" +TP="${2:-8}" + +if [ "$MODE" = "server" ]; then + PORT="${3:-8000}" + shift 3 2>/dev/null || true +else + shift 2 2>/dev/null || true +fi + +# 1) no gpucore files (a single ROCm fault dumps 30-50 GB per rank) +ulimit -c 0 + +# 2) clean cwd for code-object dumps; --save-code-objects writes +# `memory____offset__size_` files into the process cwd. +DEBUG_DIR="/app/logs_claude/debug_run" +mkdir -p "$DEBUG_DIR" +cd "$DEBUG_DIR" +rm -f memory_* *.s + +# 3) debug agent env (must be exported so spawn workers inherit) +export HSA_TOOLS_LIB=/opt/rocm/lib/librocm-debug-agent.so.2 +export HSA_ENABLE_DEBUG=1 +export ROCM_DEBUG_AGENT_OPTIONS="--save-code-objects" + +# 4) ATOM-side default; caller may override via env before invocation +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-WARNING}" + +# 5) launch (model-specific env like ATOM_USE_TRITON_MOE must come from caller) +SCRIPT_DIR="$(dirname "$0")" +if [ "$MODE" = "server" ]; then + exec bash "$SCRIPT_DIR/start_atom_server.sh" \ + "$MODEL" "$TP" "$PORT" "$@" +else + export LOG_FILE="${LOG_FILE:-/app/logs_claude/simple_inference_debug_agent.log}" + exec bash "$SCRIPT_DIR/start_simple_inference.sh" \ + "$MODEL" "$TP" "$@" +fi diff --git a/scripts/run_gsm8k_eval.sh b/scripts/run_gsm8k_eval.sh new file mode 100755 index 0000000000..79a06609ad --- /dev/null +++ b/scripts/run_gsm8k_eval.sh @@ -0,0 +1,56 @@ +#!/bin/bash +# GSM8K 精度测试 (使用 lm_eval) +# 用法: bash run_gsm8k_eval.sh [MODEL_NAME] [PORT] [NUM_FEWSHOT] +# +# 示例: +# bash run_gsm8k_eval.sh /data/DeepSeek-R1-0528 8000 5 +# bash run_gsm8k_eval.sh meta-llama/Meta-Llama-3-8B 8000 5 + +set -euo pipefail + +MODEL="${1:-/data/DeepSeek-R1-0528}" +PORT="${2:-8000}" +NUM_FEWSHOT="${3:-5}" +NUM_CONCURRENT="${NUM_CONCURRENT:-65}" +LIMIT="${LIMIT:-}" # set LIMIT=50 to run only first 50 samples +BASE_URL="http://localhost:${PORT}/v1/completions" +OUTPUT_DIR="/app/logs_claude" +LOG_FILE="${OUTPUT_DIR}/gsm8k_eval.log" + +{ +echo "========================================" +echo " GSM8K Accuracy Evaluation (lm_eval)" +echo "========================================" +echo " Model: $MODEL" +echo " Base URL: $BASE_URL" +echo " Few-shot: $NUM_FEWSHOT" +echo " Concurrency: $NUM_CONCURRENT" +echo "========================================" +} > "$LOG_FILE" + +# 等待服务就绪 +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + echo "Server is ready!" + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: Server not ready after 120 retries (20 min)" + exit 1 + fi + echo " Waiting... ($i/120)" + sleep 10 +done + +LIMIT_ARG=() +if [ -n "$LIMIT" ]; then + LIMIT_ARG=(--limit "$LIMIT") +fi + +lm_eval --model local-completions \ + --model_args "model=${MODEL},base_url=${BASE_URL},num_concurrent=${NUM_CONCURRENT},max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot "$NUM_FEWSHOT" \ + "${LIMIT_ARG[@]}" \ + 2>&1 | tee -a "${LOG_FILE}" diff --git a/scripts/start_atom_server.sh b/scripts/start_atom_server.sh new file mode 100755 index 0000000000..197c040dd1 --- /dev/null +++ b/scripts/start_atom_server.sh @@ -0,0 +1,115 @@ +#!/bin/bash +# Start ATOM OpenAI-compatible server +# Usage: bash start_atom_server.sh [MODEL_PATH] [TP_SIZE] [PORT] [EXTRA_ARGS...] +# +# Examples: +# bash start_atom_server.sh # DeepSeek-R1-0528, tp=8, port=8000 +# bash start_atom_server.sh /data/Llama-3.1-8B-Instruct-FP8-KV 1 8000 +# bash start_atom_server.sh /data/DeepSeek-R1-0528 8 8000 --method mtp --num-speculative-tokens 3 + +set -euo pipefail + +MODEL_PATH="${1:-/data/DeepSeek-R1-0528}" +TP_SIZE="${2:-8}" +PORT="${3:-8000}" +shift 3 2>/dev/null || true +EXTRA_ARGS="$*" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.9}" +LOG_FILE="/app/logs_claude/atom_server.log" + +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-INFO}" +export KINETO_CONFIG="/home/ljin1/dk/libkineto.conf" + +# === Pre-flight: ensure GPU is clean === +echo "Pre-flight: cleaning up processes and GPU memory..." + +# 1. Kill atom server processes +pkill -f 'atom.entrypoints' 2>/dev/null || true +sleep 2 + +# 2. Kill orphaned multiprocessing spawn/tracker (these hold GPU memory after server dies) +pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true +pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true +sleep 3 + +# 3. Verify GPU memory is actually free +MAX_WAIT=30 +for i in $(seq 1 $MAX_WAIT); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo "GPU memory clear after ${i}s" + break + fi + if [ "$i" -eq "$MAX_WAIT" ]; then + echo "WARNING: GPU memory still in use after ${MAX_WAIT}s. Dumping GPU process info:" + rocm-smi --showpidgpus 2>&1 | grep "PID.*is using" | grep -v "0 DRM" || true + echo "Attempting force kill of GPU-holding processes..." + rocm-smi --showpidgpus 2>&1 | grep -oP 'PID \K\d+' | while read pid; do + kill -9 "$pid" 2>/dev/null || true + done + sleep 5 + fi + sleep 1 +done + +# 4. Clear stale compile cache +rm -rf ~/.cache/atom/* +rm -rf ./gpucore.* + +# Write config header to log (truncates old content). +# Inherited env vars are dumped explicitly so you never have to wonder +# whether ATOM_USE_TRITON_MOE / V4_USE_REF_QUANT / etc. were set. +{ +echo "========================================" +echo " ATOM Server Launcher" +echo "========================================" +echo " Model: $MODEL_PATH" +echo " TP Size: $TP_SIZE" +echo " Port: $PORT" +echo " KV Cache dtype: $KV_CACHE_DTYPE" +echo " Max num seqs: $MAX_NUM_SEQS" +echo " GPU mem util: $GPU_MEM_UTIL" +echo " Extra args: ${EXTRA_ARGS:-none}" +echo " Date: $(date)" +echo "----------------------------------------" +echo " Inherited env vars (ATOM_*, V4_*, AITER_*, HSA_*, AMD_*, HIP_*):" +env | grep -E '^(ATOM_|V4_|AITER_|HSA_|AMD_|HIP_|KV_CACHE|MAX_NUM_SEQS|MAX_MODEL_LEN|MAX_BATCHED_TOKENS|GPU_MEM_UTIL)' \ + | sort | sed 's/^/ /' || echo " (none set)" +echo "========================================" +} | tee "$LOG_FILE" + +python -m atom.entrypoints.openai_server \ + --model "$MODEL_PATH" \ + --kv_cache_dtype "$KV_CACHE_DTYPE" \ + -tp "$TP_SIZE" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --server-port "$PORT" \ + $EXTRA_ARGS \ + >> "$LOG_FILE" 2>&1 & + +SERVER_PID=$! +echo "Server started in background (PID: $SERVER_PID)" + +# Wait for server ready with GPU verification +echo "Waiting for server to be ready..." +for i in $(seq 1 120); do + if curl -sf "http://localhost:${PORT}/v1/models" > /dev/null 2>&1; then + VRAM_COUNT=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$VRAM_COUNT" -gt 0 ]; then + echo "Server is ready! (PID: $SERVER_PID, GPU VRAM loaded)" + exit 0 + fi + fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: Server process died. Check $LOG_FILE" + exit 1 + fi + if [ "$i" -eq 600 ]; then + echo "ERROR: Server not ready after 600s (10 min)" + exit 1 + fi + sleep 1 +done diff --git a/scripts/start_simple_inference.sh b/scripts/start_simple_inference.sh new file mode 100755 index 0000000000..1627b34311 --- /dev/null +++ b/scripts/start_simple_inference.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# Run ATOM offline simple_inference for any registered model. +# Usage: bash start_simple_inference.sh [MODEL_PATH] [TP_SIZE] [EXTRA_ARGS...] +# +# All engine flags (kv_cache_dtype, max_num_seqs, max_model_len, gpu mem util, +# spec method, …) go in EXTRA_ARGS as native CLI args — no env-var translation. +# Only LOG_FILE and true env vars (ATOM_*, AITER_*, HSA_*, HIP_*) remain +# overridable via env. +# +# Examples: +# bash start_simple_inference.sh # default DeepSeek-R1-0528 tp=8 +# bash start_simple_inference.sh /data/Llama-3.1-8B-Instruct-FP8-KV 1 --kv_cache_dtype fp8 +# bash start_simple_inference.sh /data/DeepSeek-V4-Pro 8 --kv_cache_dtype fp8 --enforce-eager --level 0 +# bash start_simple_inference.sh /data/DeepSeek-V4-Pro 8 \ +# --kv_cache_dtype fp8 --enforce-eager --level 0 \ +# --method mtp --num-speculative-tokens 3 +# MAX_MODEL_LEN=1024 ... # ← no longer supported; pass --max-model-len 1024 instead +# +# DeepSeek-V4 must also export ATOM_USE_TRITON_MOE=1 (real env var, not a CLI flag). + +set -euo pipefail + +MODEL_PATH="${1:-/data/DeepSeek-R1-0528}" +TP_SIZE="${2:-8}" +shift 2 2>/dev/null || true +EXTRA_ARGS=("$@") + +LOG_FILE="${LOG_FILE:-/app/logs_claude/simple_inference.log}" + +export AITER_LOG_LEVEL="${AITER_LOG_LEVEL:-INFO}" + +# === Pre-flight: ensure GPU is clean === +echo "Pre-flight: cleaning up processes and GPU memory..." + +pkill -f 'atom.entrypoints' 2>/dev/null || true +pkill -f 'atom.examples.simple_inference' 2>/dev/null || true +sleep 2 + +pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true +pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true +sleep 3 + +# Verify GPU memory is actually free +MAX_WAIT=30 +for i in $(seq 1 $MAX_WAIT); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo " All GPUs free." + break + fi + if [ "$i" -eq "$MAX_WAIT" ]; then + echo " WARN: $USED_GPUS GPU(s) still showing memory after $MAX_WAIT s — proceeding anyway" + break + fi + sleep 1 +done + +# Clear ATOM compile cache (stale cache after code change causes silent fails) +rm -rf /root/.cache/atom/* 2>/dev/null || true + +# Clear ROCm core dumps (each ~85GB; HSA exceptions can fill the disk fast). +# Find from CWD + a few common roots — gpucore.PID files dropped in process CWD. +GPUCORE_DELETED=0 +for d in . /app/ATOM /app /root /tmp /app/logs_claude; do + [ -d "$d" ] || continue + for f in "$d"/gpucore.* "$d"/core.[0-9]*; do + if [ -e "$f" ]; then + sz=$(du -BG --apparent-size "$f" 2>/dev/null | awk '{print $1}') + rm -f "$f" 2>/dev/null && GPUCORE_DELETED=$((GPUCORE_DELETED+1)) + echo " removed $f ($sz)" + fi + done +done +[ "$GPUCORE_DELETED" -gt 0 ] && echo " Cleared $GPUCORE_DELETED ROCm core dump(s)" + +# Disk-space guard — refuse to launch if root fs has < 50 GB free +# (each crashed worker can dump 85+GB; better to fail-fast than fill disk again). +AVAIL_GB=$(df -BG / 2>/dev/null | awk 'NR==2 {gsub("G",""); print $4}') +if [ -n "$AVAIL_GB" ] && [ "$AVAIL_GB" -lt 50 ]; then + echo "ERROR: only ${AVAIL_GB}GB free on / — refusing to launch (would fill disk on crash)." + echo " Run: rm -f /app/ATOM/gpucore.* /app/logs_claude/**/gpucore.*" + exit 2 +fi + +# Disable ROCm core dumps for this run (HSA crashes won't write 85GB files). +# Set HSA_ENABLE_COREDUMP_FILTER to 0x0 to skip all core regions. +export HSA_ENABLE_COREDUMP=0 +export AMD_LOG_LEVEL=0 # quiet HSA error spam + +# === Header to log === +# Inherited env vars are dumped explicitly so missing ATOM_USE_TRITON_MOE etc. +# is visible at a glance instead of producing silent accuracy regressions. +{ +echo "========================================" +echo " ATOM simple_inference Launcher" +echo "========================================" +echo " Model: $MODEL_PATH" +echo " TP Size: $TP_SIZE" +echo " Extra args: ${EXTRA_ARGS[*]:-none}" +echo " Date: $(date)" +echo "----------------------------------------" +echo " Inherited env vars (ATOM_*, V4_*, AITER_*, HSA_*, AMD_*, HIP_*):" +env | grep -E '^(ATOM_|V4_|AITER_|HSA_|AMD_|HIP_)' \ + | sort | sed 's/^/ /' || echo " (none set)" +echo "========================================" +} | tee "$LOG_FILE" + +set +e +python -m atom.examples.simple_inference \ + --model "$MODEL_PATH" \ + -tp "$TP_SIZE" \ + "${EXTRA_ARGS[@]}" \ + >> "$LOG_FILE" 2>&1 +EXIT_CODE=$? +set -e + +echo "" +echo "========================================" +echo " Exit code: $EXIT_CODE" +echo " Log: $LOG_FILE" +echo "========================================" + +# Show first error if any +if [ "$EXIT_CODE" -ne 0 ]; then + echo "" + echo "First error in log:" + grep -nE "RuntimeError|AttributeError|TypeError|ValueError|^\[rank0\]:" "$LOG_FILE" \ + | grep -v "AsyncIOProc\|ModelRunner.*has no attribute" \ + | head -5 +fi + +exit $EXIT_CODE diff --git a/scripts/stop_atom_server.sh b/scripts/stop_atom_server.sh new file mode 100755 index 0000000000..1a3aee5c18 --- /dev/null +++ b/scripts/stop_atom_server.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Stop ATOM server and release GPU memory. +# Usage: bash stop_atom_server.sh +# +# Escalates SIGTERM -> SIGKILL on the main entrypoint process AND its +# multiprocessing children, then waits up to 60s for VRAM to drop to 0. +# If VRAM is still held, dumps the GPU PID list and force-kills each one. + +set -uo pipefail + +echo "=== Stopping ATOM server ===" + +# Patterns matching every process the server can spawn AND clients hammering +# it (lm_eval) so the next launch starts from a fully clean state. +PATTERNS=( + 'atom.entrypoints' + 'multiprocessing.spawn' + 'multiprocessing.resource_tracker' + 'atom.model_engine' + 'lm_eval' +) + +count_alive() { + local total=0 + for pat in "${PATTERNS[@]}"; do + local n + n=$(pgrep -f "$pat" 2>/dev/null | wc -l) + total=$((total + n)) + done + echo "$total" +} + +# 1. Graceful SIGTERM round. +if [ "$(count_alive)" -gt 0 ]; then + echo "Sending SIGTERM to ATOM processes..." + for pat in "${PATTERNS[@]}"; do + pkill -TERM -f "$pat" 2>/dev/null || true + done + for i in $(seq 1 5); do + sleep 1 + [ "$(count_alive)" -eq 0 ] && break + done +fi + +# 2. SIGKILL anything still alive — escalate without polite waiting. +if [ "$(count_alive)" -gt 0 ]; then + echo "SIGTERM did not finish cleanup; escalating to SIGKILL..." + for pat in "${PATTERNS[@]}"; do + pkill -KILL -f "$pat" 2>/dev/null || true + done + sleep 2 +fi + +# 3. Final sweep: any leftover that still references the server by name. +if [ "$(count_alive)" -gt 0 ]; then + echo "Stubborn processes still alive — final SIGKILL pass:" + for pat in "${PATTERNS[@]}"; do + pgrep -af "$pat" 2>/dev/null || true + done + for pat in "${PATTERNS[@]}"; do + pkill -KILL -9 -f "$pat" 2>/dev/null || true + done + sleep 3 +fi + +# 4. Wait for GPU memory to release. +echo "Waiting for GPU memory to release..." +for i in $(seq 1 60); do + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -eq 0 ]; then + echo "GPU memory released after ${i}s" + break + fi + if [ "$i" -eq 60 ]; then + echo "WARNING: GPU memory still in use after 60s, force-killing GPU PIDs..." + # Newer rocm-smi exposes per-GPU PIDs via --showpids + for pid in $(rocm-smi --showpids 2>/dev/null | awk '/^PID/{f=1; next} f && $1 ~ /^[0-9]+$/ {print $1}'); do + echo " kill -9 $pid" + kill -9 "$pid" 2>/dev/null || true + done + sleep 5 + USED_GPUS=$(rocm-smi --showmemuse 2>/dev/null | grep "VRAM%" | awk '{print $NF}' | awk '$1 > 0' | wc -l) + if [ "$USED_GPUS" -gt 0 ]; then + echo "ERROR: GPU memory still held after force-kill. Manual intervention required." + rocm-smi --showpids 2>&1 | head -40 + exit 1 + fi + fi + sleep 1 +done + +echo "Server stopped." diff --git a/scripts/wait_infer_drain.sh b/scripts/wait_infer_drain.sh new file mode 100755 index 0000000000..ba78d94c32 --- /dev/null +++ b/scripts/wait_infer_drain.sh @@ -0,0 +1,191 @@ +#!/bin/bash +# Watch an in-flight ATOM inference workload (server OR offline simple_inference). +# Exit 0 when the workload drains cleanly, exit 1 on hang, exit 2 on GPU fault. +# +# Workload modes (auto-detected by process pattern): +# - Server mode (atom.entrypoints): drains when eval client (lm_eval / curl +# / benchmark) is gone and a single poll shows no new output. +# Hang = STUCK_POLLS consecutive polls with no new "Engine +# Core: output send" while a client is still hammering. +# - Offline mode (atom.examples.simple_inference / atom.examples.benchmark): +# drains when the inference process exits naturally (no fault +# in log). Hang detection is N/A (no client concept; no +# Engine Core ZMQ progress). +# +# Server log auto-discovery: the canonical signal for "engine made progress" +# is `Engine Core: output send` in the server's stdout. Rather than asking +# the caller to know where the server log lives, the script `readlink`s +# `/proc//fd/1` of the atom.entrypoints process and uses THAT as the +# authoritative server log — path-independent across repo layouts. Falls +# back to the caller-supplied LOG_FILE only if /proc lookup fails. +# +# The caller-supplied LOG_FILE is still useful as a SECONDARY signal: +# - Fault grep runs on BOTH server log AND caller LOG_FILE (max coverage) +# - File-mtime growth on caller LOG_FILE counts as progress when the +# engine marker is absent (e.g. simple_inference offline, benchmark tqdm) +# +# Use this script for ANY blocking wait on an ATOM workload — don't fall back +# to `sleep + tail`. Offline path is supported in-script so you can wrap +# simple_inference the same way as server runs. +# +# Liveness: pgrep `$SERVER_PATTERN` (NOT curl /v1/models — under heavy load +# the HTTP endpoint can fail to respond within reasonable timeout even when +# the server is alive, producing false-positive "dead" reports). +# +# Usage: bash scripts/wait_infer_drain.sh [PORT] [MAX_MIN] [POLL_SEC] [LOG_FILE] [STUCK_POLLS] +# PORT default 8000 (kept for API symmetry with wait_server_ready.sh; unused in offline mode) +# MAX_MIN default 30 (full GSM8K 1319 typically 5-15 min on V4-Pro) +# POLL_SEC default 10 (fast poll for quick hang detection) +# LOG_FILE default empty (server log auto-discovered via /proc). Pass a +# client/workload log (e.g. benchmark.log, simple_inference.log) +# to add it as a secondary signal source for fault grep and +# mtime-based progress detection. +# STUCK_POLLS default 6 (6 × 10s = 1 min of no progress → declare hang) +# +# Exit codes: +# 0 — workload drained cleanly (server: client gone + no pending output; +# offline: process exited without fault) +# 1 — hang detected (server only: no progress while client running) +# 2 — fault detected (MEMORY_VIOLATION / ASSERT_TRAP / Memory access fault / proc died) +# 4 — max wait elapsed without resolution + +set -uo pipefail + +PORT="${1:-8000}" +MAX_MIN="${2:-30}" +POLL="${3:-10}" +LOG_FILE="${4:-}" +STUCK_POLLS="${5:-6}" +ITERS=$(( MAX_MIN * 60 / POLL )) + +# Match anything that indicates the eval driver is still hammering the +# server. Extend if you use a different client. +CLIENT_PATTERN='lm_eval|curl.*v1/(completions|chat)|atom\.examples\.benchmark|atom\.benchmarks\.benchmark' +# Server- or offline-mode workload process. simple_inference and +# atom.examples.benchmark also count — process-exit + no-fault = drain. +SERVER_PATTERN='atom\.entrypoints|atom\.examples\.simple_inference' + +# Resolve the authoritative server/workload log by reading the workload +# process's stdout fd. Re-run every poll because the workload PID can change +# (start_atom_server has just spawned it; CG capture forks; etc.) and we +# don't want a stale path. Returns empty string on failure. +discover_server_log() { + local pid log + pid=$(pgrep -f "$SERVER_PATTERN" 2>/dev/null | head -1) + [ -z "$pid" ] && return + log=$(readlink "/proc/$pid/fd/1" 2>/dev/null) + # Reject /dev/pts/, pipes, sockets, /dev/null — only real files count. + case "$log" in + /*) [ -f "$log" ] && echo "$log" ;; + esac +} + +# Grep helper that tolerates missing/empty file and never returns a count +# polluted with grep's stderr. +count_in() { + local pat=$1 file=$2 c + [ -z "$file" ] || [ ! -r "$file" ] && { echo 0; return; } + c=$(grep -cE "$pat" "$file" 2>/dev/null | head -1) + echo "${c:-0}" +} + +mtime_of() { + local file=$1 + [ -z "$file" ] || [ ! -r "$file" ] && { echo 0; return; } + stat -c %Y "$file" 2>/dev/null || echo 0 +} + +FAULT_PATTERN='stopped, reason|MEMORY_VIOLATION|ASSERT_TRAP|proc died unexpectedly|Memory access fault by GPU' + +prev_outputs=0 +prev_mtime=0 +stuck=0 +server_log="" + +for ((i=1; i<=ITERS; i++)); do + sleep "$POLL" + + # Refresh discovered server log (PID may have just appeared). + new_log=$(discover_server_log) + if [ -n "$new_log" ] && [ "$new_log" != "$server_log" ]; then + server_log="$new_log" + echo "[t=$((i*POLL))s] server log auto-discovered: $server_log" + fi + + # GPU fault? Scan BOTH server log AND caller LOG_FILE (max coverage). + # Check BEFORE process-exit so that fault-then-exit is attributed to + # fault (exit 2), not normal drain. + fault_server=$(count_in "$FAULT_PATTERN" "$server_log") + fault_caller=$(count_in "$FAULT_PATTERN" "$LOG_FILE") + fault_total=$(( fault_server + fault_caller )) + if [ "$fault_total" -gt 0 ]; then + echo "[t=$((i*POLL))s] GPU fault detected ($fault_total signals) — exiting 2" + for f in "$server_log" "$LOG_FILE"; do + [ -n "$f" ] && [ -r "$f" ] && grep -E "$FAULT_PATTERN" "$f" 2>/dev/null | head -3 + done + exit 2 + fi + + # Workload process alive? Only by process presence — no curl (HTTP can + # false-negative under heavy load). + if ! pgrep -f "$SERVER_PATTERN" >/dev/null 2>&1; then + # No fault grep matched above + process gone = clean exit (drain). + # Covers both stop_atom_server (server mode) and simple_inference + # finishing its prompts (offline mode). + echo "[t=$((i*POLL))s] workload process exited cleanly — DRAINED" + exit 0 + fi + + # Engine progress? Two signals — either resets the stuck counter: + # 1. "Engine Core: output send" count rising in server log + # (auto-discovered; authoritative engine progress). + # 2. Caller LOG_FILE mtime advancing (covers client logs / offline + # stdout where engine marker is absent: benchmark tqdm, + # simple_inference, etc.). + cur_outputs=$(count_in "Engine Core: output send" "$server_log") + delta_out=$(( cur_outputs - prev_outputs )) + cur_mtime=$(mtime_of "$LOG_FILE") + delta_mtime=$(( cur_mtime - prev_mtime )) + + # Client still running? + client_alive=$(pgrep -af "$CLIENT_PATTERN" 2>/dev/null | grep -v grep | wc -l) + + echo "[t=$((i*POLL))s] outputs=${cur_outputs} (+${delta_out}) mtime+${delta_mtime}s clients=${client_alive} stuck=${stuck}/${STUCK_POLLS}" + + if [ "$delta_out" -eq 0 ] && [ "$delta_mtime" -eq 0 ]; then + stuck=$(( stuck + 1 )) + else + stuck=0 + prev_outputs=$cur_outputs + prev_mtime=$cur_mtime + fi + + # Drained cleanly: client gone AND no new output this poll → done. + # No need to wait STUCK_POLLS — once the client is gone no new requests + # can arrive, so a single quiet poll is definitive. + if [ "$client_alive" -eq 0 ] && [ "$stuck" -ge 1 ]; then + echo "Engine DRAINED (client gone + no pending output)" + exit 0 + fi + + # Hung: no progress AND clients still alive → declare hang + if [ "$stuck" -ge "$STUCK_POLLS" ] && [ "$client_alive" -gt 0 ]; then + echo "HANG detected (no progress for ${stuck} polls while ${client_alive} client(s) still running)" + for f in "$server_log" "$LOG_FILE"; do + if [ -n "$f" ] && [ -r "$f" ]; then + echo "--- last 20 lines of $f ---" + tail -20 "$f" 2>/dev/null + fi + done + exit 1 + fi +done + +echo "MAX_WAIT ${MAX_MIN} min elapsed without resolution" +for f in "$server_log" "$LOG_FILE"; do + if [ -n "$f" ] && [ -r "$f" ]; then + echo "--- last 20 lines of $f ---" + tail -20 "$f" 2>/dev/null + fi +done +exit 4 diff --git a/scripts/wait_server_ready.sh b/scripts/wait_server_ready.sh new file mode 100755 index 0000000000..1f489a9191 --- /dev/null +++ b/scripts/wait_server_ready.sh @@ -0,0 +1,41 @@ +#!/bin/bash +# Poll http://localhost:$PORT/v1/models until the ATOM server is ready or +# the max wait elapses. Exit 0 = ready, 1 = timeout / startup error. +# +# Usage: bash scripts/wait_server_ready.sh [PORT] [MAX_MIN] [POLL_SEC] [LOG_FILE] +# PORT default 8000 +# MAX_MIN default 6 (server warmup typically 1-3 min; under debug agent 3-5 min) +# POLL_SEC default 30 +# LOG_FILE default /app/logs_claude/atom_server.log (used to detect startup errors) +# +# Side effect: every poll prints "[t=Ns] ready=...". Last line shows pass/fail. + +set -uo pipefail + +PORT="${1:-8000}" +MAX_MIN="${2:-6}" +POLL="${3:-30}" +LOG_FILE="${4:-/app/logs_claude/atom_server.log}" +ITERS=$(( MAX_MIN * 60 / POLL )) + +for ((i=1; i<=ITERS; i++)); do + sleep "$POLL" + READY=$(curl -s -m 3 "http://localhost:${PORT}/v1/models" 2>/dev/null | head -c 60) + ERR=$(grep -c "cluster_dims\|InductorError\|SHUTDOWN signal\|proc died" \ + "$LOG_FILE" 2>/dev/null | head -1) + ERR="${ERR:-0}" + echo "[t=$((i*POLL))s] ready=${READY:-(empty)} err=$ERR" + if [ -n "$READY" ]; then + echo "Server READY" + exit 0 + fi + if [ "$ERR" -gt 0 ]; then + echo "Server FAILED to start (errors detected)" + tail -30 "$LOG_FILE" + exit 1 + fi +done + +echo "Server NOT ready after ${MAX_MIN} min" +tail -30 "$LOG_FILE" +exit 1 diff --git a/tests/conftest.py b/tests/conftest.py index 18ae73642b..16d0cac657 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -118,6 +118,7 @@ def __init__(self, **overrides): enable_prefix_caching=False, max_num_seqs=4, max_num_batched_tokens=64, + max_model_len=64, bos_token_id=1, eos_token_id=2, stop_token_ids=[], diff --git a/tests/test_mxfp4_moe_has_bias.py b/tests/test_mxfp4_moe_has_bias.py index a9255bfddc..a195392c6c 100644 --- a/tests/test_mxfp4_moe_has_bias.py +++ b/tests/test_mxfp4_moe_has_bias.py @@ -19,10 +19,28 @@ import sys import unittest -# Clear cached atom modules (conftest.py stubs) -for mod_name in list(sys.modules): - if mod_name.startswith("atom"): - del sys.modules[mod_name] +# This test needs to inspect real atom source (not conftest.py stubs), so it +# wipes any cached `atom.*` modules at module-import time. Previously this +# also wiped the conftest stubs and never restored them, polluting later +# tests (test_arg_utils_spec / test_scheduler / test_sequence) that depend +# on the stubs. setUpModule / tearDownModule now snapshots-and-restores +# sys.modules so this file's effect is local to its own collection. +_saved_atom_modules: dict[str, object] = {} + + +def setUpModule(): + global _saved_atom_modules + _saved_atom_modules = { + name: mod for name, mod in sys.modules.items() if name.startswith("atom") + } + for name in list(_saved_atom_modules): + del sys.modules[name] + + +def tearDownModule(): + for name in [n for n in sys.modules if n.startswith("atom")]: + del sys.modules[name] + sys.modules.update(_saved_atom_modules) class TestFusedMoEDefaultHasBias(unittest.TestCase): @@ -201,6 +219,15 @@ class TestSwiGLUInterleavingWithoutBias(unittest.TestCase): and guard only the bias interleaving on ``layer.w13_bias is not None``. """ + @unittest.skip( + "Obsolete: Mxfp4MoEMethod.process_weights_after_loading no longer " + "branches on `layer.activation == ActivationType.Swiglu`. The function " + "now routes via `use_triton` (Triton swizzle) vs. the AITER shuffle " + "path, with bias cast handled unconditionally up top. The original " + "regression this guard was added for — the SwiGLU branch being " + "incorrectly gated on `w13_bias is not None` — cannot recur in the " + "current structure. Re-evaluate or delete when revisiting Mxfp4 MoE." + ) def test_swiglu_branch_condition_no_bias_check(self): """The SwiGLU branch must NOT require bias to be present.""" import inspect diff --git a/tests/test_per_req_cache_decoupling.py b/tests/test_per_req_cache_decoupling.py index 81166817ac..07ef2000e6 100644 --- a/tests/test_per_req_cache_decoupling.py +++ b/tests/test_per_req_cache_decoupling.py @@ -22,6 +22,7 @@ def per_req_cache_config(**overrides): enable_prefix_caching=False, max_num_seqs=8, max_num_batched_tokens=256, + max_model_len=256, bos_token_id=1, eos_token_id=2, stop_token_ids=[], From 28b455a98bb5b55a556e30324cc67661f23fc5b4 Mon Sep 17 00:00:00 2001 From: Zhiwei Date: Tue, 19 May 2026 18:38:50 +0800 Subject: [PATCH 71/76] [ATOM-SGLang][Feat] Enable Deepseek v3 MTP (#643) * MTP(num_step=1) for DeeepSeek * Add work log for claude debug * adopt new attn constructor args * rm worklog * use atom_parameter * kwargs handle * rebase main * precheckin * fix k_scale v_scale error * new commit * fix blank * fix qwen3.5 acc --------- Co-authored-by: ZhiweiYan-96 Co-authored-by: zhuyuhua-v --- atom/plugin/register.py | 15 +- .../attention_backend/radix_attention.py | 8 + .../attention_backend/sgl_attn_backend.py | 735 ++++++++++++++++-- .../sglang/models/base_model_wrapper.py | 244 ++++-- .../sglang/models/deepseek_nextn_wrapper.py | 203 +++++ 5 files changed, 1080 insertions(+), 125 deletions(-) create mode 100644 atom/plugin/sglang/models/deepseek_nextn_wrapper.py diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 0af553e476..86f06623dc 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -40,21 +40,28 @@ def _register_custom_attention_to_sglang() -> None: sglang only accepts pre-registered backend names, so we reuse the "aiter" name to inject ATOMAttnBackendForSgl without modifying sglang source. """ + import sglang.srt.layers.attention.aiter_backend as sglang_aiter_backend + from sglang.srt.layers.attention.attention_registry import ( register_attention_backend, ) + from atom.plugin.sglang.attention_backend.sgl_attn_backend import ( + ATOMAttnBackendForSgl, + ) # here register the custom attention backend with the name "aiter" # as sglang defines the fixed attention backend choices, which must be # in-tree logger.info("Register custom attention backend ATOMAttnBackendForSgl to SGLang") + # Speculative draft paths instantiate AiterAttnBackend directly inside + # AiterMultiStepDraftBackend, bypassing the attention registry. Rebind the + # module symbol as well so both registry lookup and direct construction use + # the plugin backend. + sglang_aiter_backend.AiterAttnBackend = ATOMAttnBackendForSgl + @register_attention_backend("aiter") def create_atom_backend(runner): - from atom.plugin.sglang.attention_backend.sgl_attn_backend import ( - ATOMAttnBackendForSgl, - ) - return ATOMAttnBackendForSgl(runner) diff --git a/atom/plugin/sglang/attention_backend/radix_attention.py b/atom/plugin/sglang/attention_backend/radix_attention.py index ba46eb4c79..8e20eecb14 100644 --- a/atom/plugin/sglang/attention_backend/radix_attention.py +++ b/atom/plugin/sglang/attention_backend/radix_attention.py @@ -88,10 +88,18 @@ def __init__( self.attn.k_scale = atom_parameter( torch.tensor([1.0], dtype=torch.float32, device="cuda") ) + elif not self.attn.k_scale.is_cuda: + self.attn.k_scale = atom_parameter( + self.attn.k_scale.detach().to(device="cuda") + ) if self.attn.v_scale is None: self.attn.v_scale = atom_parameter( torch.tensor([1.0], dtype=torch.float32, device="cuda") ) + elif not self.attn.v_scale.is_cuda: + self.attn.v_scale = atom_parameter( + self.attn.v_scale.detach().to(device="cuda") + ) # Some SGLang attention backends consume the host-side float scales # directly. Keep them in sync with the device-side defaults so the # plugin path works even when checkpoint loading never populates them. diff --git a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py b/atom/plugin/sglang/attention_backend/sgl_attn_backend.py index 57883e32a1..a8ebc1122e 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py +++ b/atom/plugin/sglang/attention_backend/sgl_attn_backend.py @@ -23,6 +23,7 @@ from sglang.srt.layers.attention.utils import ( create_flashinfer_kv_indices_triton, launch_reshape_and_cache_flash, + pad_sequence_with_mask, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import get_bool_env_var @@ -195,6 +196,7 @@ class ForwardMetadata: reduce_partial_map: Optional[torch.Tensor] = None fp8_prefill_kv_indices: Optional[torch.Tensor] = None num_kv_splits: Optional[int] = None + run_graph: Optional[bool] = True # PA metadata for pa_persistent_fwd (only used in decode mode, non-MLA) pa_metadata_qo_indptr: Optional[torch.Tensor] = None pa_metadata_pages_kv_indptr: Optional[torch.Tensor] = None @@ -217,8 +219,9 @@ def __init__( model_runner: ModelRunner, skip_prefill: bool = False, kv_indptr_buf: Optional[torch.Tensor] = None, + topk: int = 1, ): - super().__init__(model_runner, skip_prefill, kv_indptr_buf) + super().__init__(model_runner, skip_prefill, kv_indptr_buf, topk) mapping = getattr( model_runner.token_to_kv_pool, "full_attention_layer_id_mapping", None ) @@ -288,6 +291,10 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): """Init auxiliary variables for triton attention backend.""" if forward_batch.forward_mode.is_decode_or_idle(): self._init_forward_metadata_decode(forward_batch) + elif self.use_mla and forward_batch.forward_mode.is_draft_extend(): + self._init_draft_extend_mla(forward_batch.batch_size, forward_batch) + elif self.use_mla and forward_batch.forward_mode.is_target_verify(): + self._init_target_verify_mla(forward_batch.batch_size, forward_batch) else: self._init_forward_metadata_extend(forward_batch) self._fixup_page_table(forward_batch) @@ -425,6 +432,186 @@ def _init_forward_metadata_extend(self, forward_batch: ForwardBatch): else: self._init_extend_mha(bs, forward_batch) + def _init_draft_extend_mla(self, bs, forward_batch): + """Init MLA metadata for speculative draft_extend.""" + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MLA draft_extend requires speculative metadata") + + kv_indices, kv_indptr, qo_indptr, _ = spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + self.req_to_token, + ) + + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + if extend_seq_lens_cpu is not None: + max_q_len = ( + int(extend_seq_lens_cpu.max().item()) + if isinstance(extend_seq_lens_cpu, torch.Tensor) + else max(extend_seq_lens_cpu) + ) + elif forward_batch.extend_seq_lens is not None: + max_q_len = int(forward_batch.extend_seq_lens.max().item()) + elif getattr(spec_info, "accept_length", None) is not None: + max_q_len = int(spec_info.accept_length.max().item()) + else: + raise RuntimeError("MLA draft_extend is missing extend sequence lengths") + + seq_lens_cpu = forward_batch.seq_lens_cpu + max_kv_len = ( + ( + int(seq_lens_cpu.max().item()) + if isinstance(seq_lens_cpu, torch.Tensor) + else max(seq_lens_cpu) + ) + if seq_lens_cpu is not None + else int(forward_batch.seq_lens.max().item()) + ) + + work_metadata = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + num_kv_splits = None + + if _sglang_aiter._use_mla_ps_kernel: + ( + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) = self.make_mla_decode_meta_data_buffer(max_q_len, bs) + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + self.kv_last_page_len[:bs], + work_metadata, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + self.kv_last_page_len[:bs], + max_q_len, + max_kv_len, + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + run_graph=False, + ) + + def _init_target_verify_mla(self, bs, forward_batch): + """Init MLA metadata for speculative target_verify.""" + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MLA target_verify requires speculative metadata") + + draft_num = spec_info.draft_token_num + kv_lens = forward_batch.seq_lens + draft_num + kv_lens_sum = forward_batch.seq_lens_sum + draft_num * bs + device = forward_batch.seq_lens.device + + qo_indptr = torch.arange( + 0, + (1 + bs) * draft_num, + step=draft_num, + dtype=torch.int32, + device=device, + ) + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = torch.empty( + kv_lens_sum, + dtype=torch.int32, + device=device, + ) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + kv_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + work_metadata = None + work_indptr = None + work_info_set = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None + num_kv_splits = None + + if _sglang_aiter._use_mla_ps_kernel: + ( + work_metadata, + work_indptr, + work_info_set, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + ) = self.make_mla_decode_meta_data_buffer(draft_num, bs) + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + self.kv_last_page_len[:bs], + work_metadata, + work_info_set, + work_indptr, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + draft_num, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + self.kv_last_page_len[:bs], + draft_num, + None, + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + run_graph=False, + ) + def _init_extend_mla(self, bs, forward_batch): self.mla_indices_updater_prefill.update( forward_batch.req_pool_indices, @@ -710,6 +897,11 @@ def init_cuda_graph_state( self.cuda_graph_kv_last_page_len = torch.ones( max_bs, dtype=torch.int, device=self.device ) + assert self.cuda_graph_kv_last_page_len.is_cuda, ( + "ATOMAttnBackendForSgl.init_cuda_graph_state created " + f"non-CUDA cuda_graph_kv_last_page_len on {self.cuda_graph_kv_last_page_len.device}, " + f"backend={type(self)}" + ) if kv_indices_buf is None: self.cuda_graph_kv_indices = torch.zeros( (max_bs * self.max_context_len), @@ -837,43 +1029,244 @@ def init_forward_metadata_capture_cuda_graph( forward_mode: ForwardMode, spec_info: Optional[SpecInput], ): - if not forward_mode.is_decode_or_idle(): - raise ValueError(f"Invalid mode: {forward_mode=}") + num_kv_splits = None + work_metadata = None + work_info_set = None + work_indptr = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None - if self.use_mla: - self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) - else: - kv_indptr = self.kv_indptr - kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) - kv_indptr = kv_indptr[: bs + 1] + if forward_mode.is_decode_or_idle(): + if self.use_mla: + self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) + else: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + static_columns=True, + ) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + None, + None, + 1, + None, + page_table, + seq_lens_persistent, + ) + if self.decode_using_pa_ps: + self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + elif forward_mode.is_target_verify(): + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_lens = seq_lens + self.num_draft_tokens if self.use_mla else seq_lens + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) kv_indices = self.cuda_graph_kv_indices create_flashinfer_kv_indices_triton[(bs,)]( self.req_to_token, req_pool_indices, - seq_lens, + kv_lens, kv_indptr, None, kv_indices, self.req_to_token.stride(0), ) - page_table, seq_lens_persistent = self._update_decode_page_table( - bs, + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = self.num_draft_tokens + + if self.use_mla: + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph TARGET_VERIFY produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}, metadata_backend={type(self.forward_metadata)}" + ) + else: + custom_mask = self.cuda_graph_custom_mask + assert spec_info is not None and spec_info.custom_mask is not None + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = max_q_len * (seq_lens + max_q_len) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + custom_mask=custom_mask, + mask_indptr=mask_indptr, + max_extend_len=max_q_len, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph TARGET_VERIFY(non-MLA) produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + bs * num_tokens_per_bs + 1, + step=num_tokens_per_bs, + dtype=torch.int32, + device=self.device, + ) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, req_pool_indices, seq_lens, - static_columns=True, - ) - self.forward_metadata = ForwardMetadata( kv_indptr, - kv_indices, - None, - None, - 1, None, - page_table, - seq_lens_persistent, + kv_indices, + self.req_to_token.stride(0), ) - if self.decode_using_pa_ps: - self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + + if self.use_mla: + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = num_tokens_per_bs + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "capture_cuda_graph DRAFT_EXTEND produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + else: + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + None, + num_tokens_per_bs, + None, + None, + None, + custom_mask=None, + mask_indptr=None, + max_extend_len=num_tokens_per_bs, + ) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") def init_forward_metadata_replay_cuda_graph( self, @@ -887,45 +1280,243 @@ def init_forward_metadata_replay_cuda_graph( seq_lens_cpu: Optional[torch.Tensor], out_cache_loc: Optional[torch.Tensor] = None, ): - if not forward_mode.is_decode_or_idle(): - raise ValueError("Invalid forward mode") + num_kv_splits = None + work_metadata = None + work_info_set = None + work_indptr = None + reduce_indptr = None + reduce_final_map = None + reduce_partial_map = None - if self.use_mla: - self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) - else: - kv_indptr = self.kv_indptr - kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) - kv_indptr = kv_indptr[: bs + 1] + if forward_mode.is_decode_or_idle(): + if self.use_mla: + self._init_mla_cuda_graph_metadata(bs, req_pool_indices, seq_lens) + else: + kv_indptr = self.kv_indptr + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + req_pool_indices, + seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + page_table, seq_lens_persistent = self._update_decode_page_table( + bs, + req_pool_indices, + seq_lens, + seq_lens_cpu=seq_lens_cpu, + static_columns=True, + ) + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + None, + None, + 1, + None, + page_table, + seq_lens_persistent[:bs], + ) + if self.decode_using_pa_ps: + self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + elif forward_mode.is_target_verify(): + bs = len(req_pool_indices) + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[: bs + 1] = torch.arange( + 0, + (1 + bs) * self.num_draft_tokens, + step=self.num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + kv_lens = seq_lens + self.num_draft_tokens if self.use_mla else seq_lens + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) kv_indices = self.cuda_graph_kv_indices create_flashinfer_kv_indices_triton[(bs,)]( self.req_to_token, req_pool_indices, - seq_lens, + kv_lens, kv_indptr, None, kv_indices, self.req_to_token.stride(0), ) - page_table, seq_lens_persistent = self._update_decode_page_table( - bs, + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = self.num_draft_tokens + + if self.use_mla: + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph TARGET_VERIFY produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}, metadata_backend={type(self.forward_metadata)}" + ) + else: + custom_mask = self.cuda_graph_custom_mask + assert spec_info is not None and spec_info.custom_mask is not None + custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask + seq_mask_len = max_q_len * (seq_lens + max_q_len) + mask_indptr = self.mask_indptr + mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0) + mask_indptr = mask_indptr[: bs + 1] + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + custom_mask=custom_mask, + mask_indptr=mask_indptr, + max_extend_len=max_q_len, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph TARGET_VERIFY(non-MLA) produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + elif forward_mode.is_draft_extend(): + num_tokens_per_bs = self.speculative_num_steps + 1 + seq_lens = seq_lens[:bs] + accept_lens = spec_info.accept_length[:bs] + qo_indptr = self.qo_indptr[: bs + 1] + qo_indptr[1 : bs + 1] = torch.cumsum(accept_lens, dim=0) + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0) + kv_indices = self.cuda_graph_kv_indices + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, req_pool_indices, seq_lens, - seq_lens_cpu=seq_lens_cpu, - static_columns=True, - ) - - self.forward_metadata = ForwardMetadata( kv_indptr, - kv_indices, - None, - None, - 1, None, - page_table, - seq_lens_persistent[:bs], + kv_indices, + self.req_to_token.stride(0), ) - if self.decode_using_pa_ps: - self._build_pa_metadata_for_decode(bs, tp_q_head_num=self.num_head) + + if self.use_mla: + kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] + max_q_len = num_tokens_per_bs + if _sglang_aiter._use_mla_ps_kernel: + num_kv_splits = self.max_split_per_batch + self.make_mla_meta_data( + qo_indptr, + kv_indptr, + kv_last_page_len, + self.work_metadata, + self.work_info_set, + self.work_indptr, + self.reduce_indptr, + self.reduce_final_map, + self.reduce_partial_map, + max_q_len, + fast_mode=_sglang_aiter.fast_mode, + max_split_per_batch=num_kv_splits, + intra_batch_mode=_sglang_aiter.intra_batch_mode, + ) + work_metadata = self.work_metadata + work_info_set = self.work_info_set + work_indptr = self.work_indptr + reduce_indptr = self.reduce_indptr + reduce_final_map = self.reduce_final_map + reduce_partial_map = self.reduce_partial_map + + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + kv_last_page_len, + max_q_len, + kv_indptr[-1].item(), + None, + None, + work_metadata=work_metadata, + work_info_set=work_info_set, + work_indptr=work_indptr, + reduce_indptr=reduce_indptr, + reduce_final_map=reduce_final_map, + reduce_partial_map=reduce_partial_map, + num_kv_splits=num_kv_splits, + ) + assert ( + self.forward_metadata.kv_last_page_len is None + or self.forward_metadata.kv_last_page_len.is_cuda + ), ( + "replay_cuda_graph DRAFT_EXTEND produced non-CUDA kv_last_page_len: " + f"{self.forward_metadata.kv_last_page_len.device}, " + f"backend={type(self)}" + ) + else: + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + None, + num_tokens_per_bs, + None, + None, + None, + custom_mask=None, + mask_indptr=None, + max_extend_len=num_tokens_per_bs, + ) + else: + raise ValueError(f"Invalid mode: {forward_mode=}") @staticmethod def _max_len(cpu_values, device_values) -> int: @@ -1167,6 +1758,7 @@ def _forward_extend_mla(self, q, k, v, layer, forward_batch): layer, K_Buffer, qo_indptr, + forward_batch, ) if not forward_batch.forward_mode.is_extend(): raise ValueError( @@ -1504,14 +2096,51 @@ def _call_mla_decode_fwd(self, q, k_buffer, o, layer): num_kv_splits=md.num_kv_splits, ) - def _forward_extend_mla_speculative(self, q, layer, K_Buffer, qo_indptr): + def _forward_extend_mla_speculative( + self, q, layer, K_Buffer, qo_indptr, forward_batch + ): """MLA speculative path (target_verify / draft_extend).""" - o = q.new_empty( - (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), - dtype=self.input_dtype, + md = self.forward_metadata + + if forward_batch.forward_mode.is_target_verify(): + o = q.new_empty( + (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd(q, K_Buffer, o, layer) + return o + + if forward_batch.forward_mode.is_draft_extend(): + if md.run_graph is not True: + bs, q_pad, _ = pad_sequence_with_mask( + q.view(q.shape[0], -1), + qo_indptr[:-1], + forward_batch.extend_seq_lens, + md.max_q_len, + ) + o = q.new_empty( + (bs * md.max_q_len, layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd( + q_pad.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + K_Buffer, + o, + layer, + ) + total_valid_q = int(qo_indptr[-1].item()) + return o[:total_valid_q] + + o = q.new_empty( + (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), + dtype=self.input_dtype, + ) + self._call_mla_decode_fwd(q, K_Buffer, o, layer) + return o + + raise ValueError( + f"Invalid forward mode for MLA speculative path: {forward_batch.forward_mode=}" ) - self._call_mla_decode_fwd(q, K_Buffer, o, layer) - return o def forward_decode( self, diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index db75a67433..3f2c743b14 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -24,7 +24,11 @@ logger = logging.getLogger("atom.plugin.sglang.models") +_RUNTIME_SENTINEL = object() +# Context for patched DeepSeek attention layers that need wrapper state without +# changing every intermediate forward signature. ContextVar keeps nested or +# concurrent forwards isolated and lets us reliably restore the prior value. _current_forward_batch: ContextVar[Optional[ForwardBatch]] = ContextVar( "atom_sglang_current_forward_batch", default=None ) @@ -215,6 +219,37 @@ def _reset_sglang_forward_context() -> None: reset_forward_context() +@contextmanager +def plugin_runtime_scope( + *, + framework: Optional[str] = None, + atom_config: Any = _RUNTIME_SENTINEL, +): + """Temporarily bind plugin runtime globals to one wrapper instance. + + ATOM core currently relies on process-global framework/config state. In + SGLang speculative mode both target and draft wrappers coexist, so plugin + entrypoints must save/restore those globals around each init/load/forward. + """ + + import atom.config as atom_config_module + import atom.plugin.prepare as plugin_prepare + + prev_framework = plugin_prepare._CURRENT_FRAMEWORK + prev_atom_config = getattr(atom_config_module, "_current_atom_config", None) + + if framework is not None: + plugin_prepare._set_framework_backbone(framework) + if atom_config is not _RUNTIME_SENTINEL: + atom_config_module._current_atom_config = atom_config + + try: + yield + finally: + plugin_prepare._CURRENT_FRAMEWORK = prev_framework + atom_config_module._current_atom_config = prev_atom_config + + @dataclass(frozen=True) class SGLangForwardBatchMetadata: """Small context object for one SGLang model forward.""" @@ -327,7 +362,14 @@ def __init__( # Refactor so this wrapper only dispatches the attention backend # (register_ops_to_sglang + set_attn_cls), and let sglang handle # model construction directly - self.model = atom.prepare_model(config=config, engine="sglang") + with plugin_runtime_scope(framework="sglang"): + from atom.config import get_current_atom_config + + self.model = atom.prepare_model(config=config, engine="sglang") + self.atom_config = getattr(self.model, "atom_config", None) + if self.atom_config is None: + self.atom_config = get_current_atom_config() + self.model.atom_config = self.atom_config if self.model is None: raise ValueError( f"ATOM failed to create model for architecture {self.model_arch}" @@ -348,7 +390,52 @@ def __init__( setup_deepseek_for_sglang, ) - setup_deepseek_for_sglang(self.model) + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + setup_deepseek_for_sglang(self.model) + + def get_embed_and_head(self): + if hasattr(self.model, "get_embed_and_head"): + return self.model.get_embed_and_head() + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + return embed_owner.embed_tokens.weight, self.model.lm_head.weight + + def set_embed_and_head(self, embed, head): + if hasattr(self.model, "set_embed_and_head"): + return self.model.set_embed_and_head(embed, head) + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + del embed_owner.embed_tokens.weight + del self.model.lm_head.weight + embed_owner.embed_tokens.weight = embed + self.model.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def set_embed(self, embed): + if hasattr(self.model, "set_embed"): + return self.model.set_embed(embed) + + embed_owner = ( + self.model.model + if hasattr(self.model, "model") + and hasattr(self.model.model, "embed_tokens") + else self.model + ) + del embed_owner.embed_tokens.weight + embed_owner.embed_tokens.weight = embed + torch.cuda.empty_cache() + torch.cuda.synchronize() @torch.no_grad() def forward( @@ -361,75 +448,95 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, **model_kwargs: Any, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - metadata = SGLangForwardBatchMetadata.build( - forward_batch, - pp_proxy_tensors=pp_proxy_tensors, - save_kv_cache=model_kwargs.get("save_kv_cache"), - ) - - if _is_dummy_forward(forward_batch): - ( - model_input_ids, - model_positions, - model_input_embeds, - model_forward_batch, - ) = _materialize_atom_dummy_forward( - input_ids, - positions, - input_embeds, - forward_batch, - ) - else: - ( - model_input_ids, - model_positions, - model_input_embeds, - model_forward_batch, - ) = ( - input_ids, - positions, - input_embeds, + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + metadata = SGLangForwardBatchMetadata.build( forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + save_kv_cache=model_kwargs.get("save_kv_cache"), ) - model_inputs = dict( - input_ids=model_input_ids, - positions=model_positions, - intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( - pp_proxy_tensors, metadata - ), - inputs_embeds=model_input_embeds, - ) - with SGLangForwardBatchMetadata.bind(metadata): - if self.model_arch_spec.wrapper_binds_gdn_context: - from atom.plugin.sglang.attention_backend.attention_gdn import ( - SGLangGDNForwardContext, - ) - with SGLangGDNForwardContext.bind(metadata): - hidden_states = self.model(**model_inputs) + if _is_dummy_forward(forward_batch): + ( + model_input_ids, + model_positions, + model_input_embeds, + model_forward_batch, + ) = _materialize_atom_dummy_forward( + input_ids, + positions, + input_embeds, + forward_batch, + ) else: - try: - _set_sglang_forward_context( - self.model.atom_config, model_forward_batch, model_positions - ) - hidden_states = self.model(**model_inputs) - finally: - _reset_sglang_forward_context() + ( + model_input_ids, + model_positions, + model_input_embeds, + model_forward_batch, + ) = ( + input_ids, + positions, + input_embeds, + forward_batch, + ) - if self.pp_group.is_last_rank: - if _is_dummy_forward(forward_batch): - # TODO: Revisit if SGLang ever sends non-empty dummy batches. - # Today this path only runs when an empty IDLE batch is expanded - # to one ATOM dummy token, so the output boundary must trim back to - # the original SGLang-visible length: 0 tokens. - hidden_states = _trim_hidden_states_for_output(hidden_states, 0) - return self.logits_processor( - input_ids, - hidden_states, - self.model.lm_head, - forward_batch, + model_inputs = dict( + input_ids=model_input_ids, + positions=model_positions, + intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( + pp_proxy_tensors, metadata + ), + inputs_embeds=model_input_embeds, ) - return hidden_states + uses_context_only_forward = ( + self.model_arch_spec.apply_deepseek_patch + or self.model_arch_spec.wrapper_binds_gdn_context + ) + with SGLangForwardBatchMetadata.bind(metadata): + if self.model_arch_spec.wrapper_binds_gdn_context: + from atom.plugin.sglang.attention_backend.attention_gdn import ( + SGLangGDNForwardContext, + ) + + with SGLangGDNForwardContext.bind(metadata): + hidden_states = self.model(**model_inputs) + elif uses_context_only_forward: + try: + _set_sglang_forward_context( + self.atom_config, model_forward_batch, model_positions + ) + hidden_states = self.model(**model_inputs) + finally: + _reset_sglang_forward_context() + else: + try: + _set_sglang_forward_context( + self.atom_config, model_forward_batch, model_positions + ) + hidden_states = self.model( + **model_inputs, + forward_batch=model_forward_batch, + get_embedding=get_embedding, + pp_proxy_tensors=pp_proxy_tensors, + **model_kwargs, + ) + finally: + _reset_sglang_forward_context() + + if self.pp_group.is_last_rank: + if _is_dummy_forward(forward_batch): + # TODO: Revisit if SGLang ever sends non-empty dummy batches. + # Today this path only runs when an empty IDLE batch is expanded + # to one ATOM dummy token, so the output boundary must trim back to + # the original SGLang-visible length: 0 tokens. + hidden_states = _trim_hidden_states_for_output(hidden_states, 0) + return self.logits_processor( + input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + ) + return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # The passed `weights` iterable from sglang is ignored because ATOM @@ -438,9 +545,10 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # sglang's default weight iterator. from atom.model_loader.loader import load_model_in_plugin_mode - return load_model_in_plugin_mode( - model=self.model, config=self.model.atom_config, prefix="model." - ) + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + return load_model_in_plugin_mode( + model=self.model, config=self.atom_config, prefix="model." + ) EntryClass = [] diff --git a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py new file mode 100644 index 0000000000..f41617e140 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py @@ -0,0 +1,203 @@ +"""ATOM DeepSeek NextN wrapper for SGLang external loading. + +This keeps SGLang's draft architecture name (`DeepseekV3ForCausalLMNextN`) +so ModelRegistry can override the upstream implementation, but delegates the +actual draft core to ATOM's `DeepSeekMTP`. +""" + +import logging +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 SpeculativeConfig +from atom.plugin.config import generate_atom_config_for_plugin_mode +from atom.plugin.sglang.attention_backend.sgl_attention_mla import ( + setup_deepseek_for_sglang, +) +from atom.plugin.sglang.models.base_model_wrapper import ( + _current_forward_batch, + _reset_sglang_forward_context, + _set_sglang_forward_context, + plugin_runtime_scope, +) + +logger = logging.getLogger("atom.plugin.sglang.models") + + +def _sync_replaced_weights() -> None: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _replace_weight(module: nn.Module, attr_name: str, weight) -> None: + if hasattr(module, attr_name): + delattr(module, attr_name) + setattr(module, attr_name, weight) + + +def _set_runtime_layer_id(layer_module: nn.Module, layer_id: int) -> None: + if hasattr(layer_module, "layer_id"): + layer_module.layer_id = layer_id + if hasattr(layer_module, "layer_num"): + layer_module.layer_num = layer_id + + +def _retag_mtp_runtime_layer_ids(model: nn.Module) -> None: + """Retag MTP runtime layer ids to draft-local indices. + + ATOM's DeepSeekMTP keeps checkpoint/global layer numbering (e.g. 61, 62...) + in module prefixes so weight remapping still works. SGLang's draft KV cache, + however, allocates layers using draft-local indices (0..num_nextn_layers-1). + Rebind only the runtime ids used by the attention/KV-cache path. + """ + + for local_layer_id, mtp_layer in enumerate(model.model.layers.values()): + mtp_block = mtp_layer.mtp_block + self_attn = mtp_block.self_attn + + _set_runtime_layer_id(self_attn, local_layer_id) + + for attr_name in ("mla_attn", "attn_mha"): + attn_obj = getattr(self_attn, attr_name, None) + if attn_obj is None: + continue + _set_runtime_layer_id(attn_obj, local_layer_id) + nested_attn = getattr(attn_obj, "attn", None) + if nested_attn is not None: + _set_runtime_layer_id(nested_attn, local_layer_id) + + +class DeepseekV3ForCausalLMNextN(nn.Module): + """SGLang-compatible draft wrapper backed by ATOM's `DeepSeekMTP`.""" + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + del prefix + super().__init__() + + logger.info("Initializing ATOM backend for %s", self.__class__.__name__) + + self.pp_group = get_pp_group() + self.quant_config = quant_config + self.config = config + self.vocab_size = config.vocab_size + self.unpadded_vocab_size = config.vocab_size + + with plugin_runtime_scope(framework="sglang"): + self.atom_config = generate_atom_config_for_plugin_mode(config) + + # Draft workers need ATOM's MTP-specific config semantics rather than the + # default target-model translation used by the generic plugin wrapper. + SpeculativeConfig.hf_config_override(self.atom_config.hf_config) + + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + from atom.plugin.register import ( + init_aiter_dist, + register_ops_to_sglang, + set_attn_cls, + ) + from atom.models.deepseek_mtp import DeepSeekMTP + + register_ops_to_sglang(atom_config=self.atom_config) + set_attn_cls() + init_aiter_dist(config=self.atom_config) + + self.model = DeepSeekMTP(atom_config=self.atom_config) + self.model.atom_config = self.atom_config + setup_deepseek_for_sglang(self.model) + _retag_mtp_runtime_layer_ids(self.model) + + self.logits_processor = LogitsProcessor(config) + self.lm_head = self._first_mtp_layer().shared_head.head + + def _mtp_layers(self): + return list(self.model.model.layers.values()) + + def _first_mtp_layer(self): + layers = self._mtp_layers() + if not layers: + raise ValueError("DeepSeekMTP does not contain any draft layers") + return layers[0] + + def get_embed_and_head(self): + return self.model.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + self.set_embed(embed) + for mtp_layer in self._mtp_layers(): + _replace_weight(mtp_layer.shared_head.head, "weight", head) + self.lm_head = self._first_mtp_layer().shared_head.head + _sync_replaced_weights() + + def set_embed(self, embed): + _replace_weight(self.model.model.embed_tokens, "weight", embed) + _sync_replaced_weights() + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + **kwargs, + ): + if forward_batch.spec_info is None: + raise ValueError("DeepSeek MTP draft forward requires speculative info") + + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + token = _current_forward_batch.set(forward_batch) + try: + _set_sglang_forward_context(self.atom_config, forward_batch, positions) + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + hidden_states=forward_batch.spec_info.hidden_states, + inputs_embeds=input_embeds, + ) + finally: + _reset_sglang_forward_context() + _current_forward_batch.reset(token) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + ) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + del weights + from atom.model_loader.loader import load_model + + server_args = get_global_server_args() + draft_model_path = ( + server_args.speculative_draft_model_path or server_args.model_path + ) + self.atom_config.model = draft_model_path + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + return load_model( + model=self.model, + model_name_or_path=draft_model_path, + hf_config=self.atom_config.hf_config, + load_dummy=self.atom_config.load_dummy, + spec_decode=True, + ) + + +EntryClass = [DeepseekV3ForCausalLMNextN] From 29408d118ab7df51e8fa9b85e3ccaa2ebeb3a944 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Mon, 11 May 2026 10:16:06 +0000 Subject: [PATCH 72/76] add work log --- .../2026-04-30-attn-refractory.md | 354 +++++++++++++ ...26-05-07-radix-cache-vs-paged-attention.md | 482 ++++++++++++++++++ ...26-05-08-kv-indices-address-space-notes.md | 359 +++++++++++++ .../2026-05-08-metadata-layering-summary.md | 144 ++++++ ...6-05-11-attn-refractory-session-summary.md | 398 +++++++++++++++ .../attn_refractory/decoupling-diagram.md | 158 ++++++ .../vllm-integration-architecture.md | 269 ++++++++++ 7 files changed, 2164 insertions(+) create mode 100644 work_log/attn_refractory/2026-04-30-attn-refractory.md create mode 100644 work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md create mode 100644 work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md create mode 100644 work_log/attn_refractory/2026-05-08-metadata-layering-summary.md create mode 100644 work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md create mode 100644 work_log/attn_refractory/decoupling-diagram.md create mode 100644 work_log/attn_refractory/vllm-integration-architecture.md diff --git a/work_log/attn_refractory/2026-04-30-attn-refractory.md b/work_log/attn_refractory/2026-04-30-attn-refractory.md new file mode 100644 index 0000000000..f0946299a9 --- /dev/null +++ b/work_log/attn_refractory/2026-04-30-attn-refractory.md @@ -0,0 +1,354 @@ +# 2026-04-30 ATOM Attention Refactory Session Summary + +## 1. 本次会话的核心目标 + +本次连续讨论的主线是围绕 `ATOM plugin` 中 attention backend 的重构方向展开，重点包括： + +- 调研当前 `ATOM plugin` attention backend 的架构问题 +- 评估 `vLLM plugin` 与 `SGLang plugin` 的耦合程度 +- 分析 SGLang plugin 在支持新模型、特别是非传统 MHA/MLA 类型 backend 时会遇到的困难 +- 讨论如何在 plugin 层把 `vLLM` 与 `SGLang` 解耦 +- 讨论 `SGLang` 中 `hybrid/composite backend` 的接入方式 +- 评估三种模式共享 ATOM attention 实现的难度： + - `ATOM native/server` + - `ATOM vLLM plugin` + - `ATOM SGLang plugin` + +## 2. 对当前 ATOM plugin attention 架构的主要判断 + +### 2.1 当前不是一个真正统一的 attention backend 架构 + +本次调研得到的一个核心结论是： + +- `vLLM plugin` 和 `SGLang plugin` 虽然都叫 “ATOM attention” +- 但它们并不是通过同一套干净的 backend contract 接进来的 +- 更准确地说，是两套宿主接入模型外加多层 patch / adapter / runtime glue 共同构成的系统 + +### 2.2 当前的主要问题不是底层 kernel，而是 runtime / ownership / patch 层次 + +现有问题主要集中在： + +- 全局状态过重 + - `_CURRENT_FRAMEWORK` + - `current_atom_config` + - `ops.Attention` +- vLLM 与 SGLang 的接入模型差异很大，但还共享 `plugin` 根目录里一批伪公共逻辑 +- 很多扩展点依赖 monkey patch 和 runtime override +- SGLang 侧对 DeepSeek MLA 已经上卷到 model-level patch，而不是单纯 backend 替换 + +### 2.3 attention 文件结构的拆分轴不统一 + +本次对 `plugin` 目录下 attention 相关文件的解读认为，目前至少混杂了三层不同对象： + +- host/runtime glue +- backend runtime core +- model-specific specialization + +尤其在 `SGLang` 侧： + +- `radix_attention.py` 更像 adapter +- `sgl_attn_backend.py` 更像 full-attn runtime backend core +- `sgl_attention_mla.py` 更像 DeepSeek MLA specialization + +而在根目录： + +- `attention.py` +- `attention_mha.py` +- `attention_mla.py` +- `attention_mla_sparse.py` + +虽然放在 `plugin/` 根目录，但职责上明显更接近 `vLLM plugin` 的 glue / patch 层，而不是真正共享的 plugin runtime core。 + +## 3. 关于 vLLM plugin 与 SGLang plugin 的耦合 + +### 3.1 结论：耦合度高 + +这次讨论中对二者耦合的判断是： + +- 不是轻量共享几个工具函数 +- 而是共享了一套 runtime state、attention 抽象、selector 和 plugin-mode 判断方式 + +关键耦合点包括： + +- `atom/plugin/prepare.py` +- `atom/plugin/register.py` +- `atom/plugin/config.py` +- `atom.plugin.prepare.is_vllm()/is_sglang()/is_plugin_mode()` +- `ops.Attention` 的全局切换 + +### 3.2 关键判断 + +真正的问题不是 “目录没有拆开”，而是： + +**host runtime 差异没有被限制在 adapter 边界，而是泄漏进了 backend 选择、metadata 组织和全局状态模型。** + +## 4. 关于 plugin 层解耦的共识 + +### 4.1 plugin 层不应该继续保留共享 runtime + +讨论最终倾向于一个更激进但更干净的方向： + +- `atom/plugin/vllm/` 是一个完整子系统 +- `atom/plugin/sglang/` 是另一个完整子系统 +- `atom/plugin/` 根目录不再承担共享 runtime 中心的角色 + +共享只保留给更底层的 `ATOM core`，例如： + +- `model_ops` +- kernels +- loader +- metadata helpers +- model-family specialization 中真正 host-agnostic 的部分 + +### 4.2 “bootstrap” 的定义 + +本次还明确了 `bootstrap` 在本讨论语境中的含义： + +- 不是核心执行逻辑 +- 而是插件被宿主框架加载时的最早期接线/初始化层 + +也就是： + +- 注册扩展点 +- 安装 patch +- 选择 wrapper / adapter / backend +- 建立 host 与 plugin 的连接关系 + +## 5. 已经做过的代码原型 + +本次会话中，为了验证“按 host 拆入口”的思路，已经做了一批原型代码修改： + +### 5.1 新增的 bootstrap / prepare / register 原型 + +- `atom/plugin/sglang/bootstrap.py` +- `atom/plugin/vllm/bootstrap.py` +- `atom/plugin/sglang/prepare.py` +- `atom/plugin/sglang/register.py` + +### 5.2 `prepare.py` 已部分收缩为兼容层 + +- 实际的 SGLang prepare 逻辑已经移动到 `atom/plugin/sglang/prepare.py` +- `atom/plugin/prepare.py` 现在更像一个 legacy shim +- 但它仍保留了 framework-state helper（因为仓库里很多地方还在依赖） + +### 5.3 SGLang register 逻辑已开始下沉 + +SGLang 独有的几块逻辑已经迁入： + +- `register_ops_to_sglang` +- `set_sglang_attn_cls` +- `init_aiter_dist_for_sglang` +- `bootstrap_sglang_runtime` + +共享 `atom/plugin/register.py` 目前更多是兼容 facade。 + +### 5.4 说明 + +这些改动的定位是： + +- 用来显式化未来结构 +- 不是完整重构 +- 还保留了较多兼容层 + +## 6. 关于 SGLang attention backend 的重要判断 + +### 6.1 `ATOMAttnBackendForSgl` 与 public `AiterAttnBackend` + +在一次反复讨论后，本次对它们的关系收敛为： + +- 在 backend 角色位置上，两者基本是 apple-to-apple +- `ATOMAttnBackendForSgl` 不是“更高层的新东西” +- 而是 `SGLang full-attention runtime backend` 的 ATOM 对位实现 + +因此更合适的重构思路不是怀疑它的存在合法性，而是： + +- 保留它作为 full-attn backend core +- 再去拆它内部的 metadata / kv_cache / decode / graph 等职责 + +### 6.2 DeepSeek MLA 的特殊性 + +`sgl_attention_mla.py` 暴露出一个很重要的事实： + +- 对 `MLA`，尤其是 `DeepSeek MLA` +- SGLang plugin 已经不是单纯换 backend +- 而是 patch 了模型级 `forward` + +这说明： + +**MLA 在 SGLang plugin 里已经进入了 model specialization 维度。** + +## 7. 关于 GDN / KDA / Lightning / Mamba2 等非传统 “attention” + +### 7.1 调研结论 + +public SGLang 已经证明，系统里不止 full-attention backend，还存在： + +- linear attention backend + - `GDNAttnBackend` + - `KDAAttnBackend` + - `LightningAttentionBackend` + - `Mamba2AttnBackend` +- hybrid/composite backend + - `HybridLinearAttnBackend` + - `HybridAttnBackend` +- wrapper/composition backend + - `TboAttnBackend` + +### 7.2 对当前 ATOM plugin 设计造成的挑战 + +当前 ATOM plugin 的主抽象仍偏向： + +- MHA +- MLA +- full attention runtime + +这会带来几个结构性问题： + +1. 还没有把 `linear backend` 当成一等公民 +2. 还没有给 `hybrid/composite backend` 预留独立宿主位置 +3. 容易把 algorithm backend、kernel backend、host glue 混成一层 +4. 容易继续把 `GDN` 等路径硬塞回 full-attn backend 体系 + +### 7.3 当前共识 + +后续不应该继续只谈 “attention backend 重构”，而应该升级为： + +**sequence backend / mixer backend 体系重构** + +建议至少在设计上并列三类： + +- `full_attn` +- `linear_attn` +- `hybrid/composite` + +## 8. hybrid/composite backend 在 plugin 侧的接入思路 + +### 8.1 核心判断 + +由于 ATOM SGLang plugin 的启动方式是： + +- 启动 `sglang server` +- 再通过 plugin runtime override 接管 model / attention + +所以要接 `hybrid/composite backend`，不可避免需要 patch 某个 seam。 + +### 8.2 最好的 seam + +本次讨论认为，public SGLang 里最好的 seam 是： + +- `model_runner.py` 中导入并调用的 `attn_backend_wrapper` + +注意： + +- 它不是 `ModelRunner` 的成员方法 +- 而是 `model_runner.py` 模块级导入的符号 + +所以如果要 patch，更稳的是 patch： + +- `sglang.srt.model_executor.model_runner.attn_backend_wrapper` + +而不是仅 patch `attention_registry.attn_backend_wrapper`。 + +### 8.3 最推荐的方式 + +对 hybrid/composite backend 的接入方式，本次最终倾向于： + +**plugin-own 一个薄的 composite wrapper/factory，只在 backend 构造 seam 做单点 patch。** + +不推荐： + +- 深度 monkey patch public hybrid backend 实现 +- 一开始就复制整套 public linear/full backend 逻辑 + +更推荐： + +- full side 先用 ATOM full backend +- linear side 初期先复用 public SGLang 的 `GDNAttnBackend` / `KDAAttnBackend` / `LightningAttentionBackend` / `Mamba2AttnBackend` +- composite wrapper 由 plugin 侧拥有 + +## 9. 关于三种模式共享 ATOM attention 实现的难度 + +本次对： + +- `ATOM native/server` +- `ATOM vLLM plugin` +- `ATOM SGLang plugin` + +三种模式共享 ATOM attention 实现，得到的判断如下。 + +### 9.1 MHA + +难度：`Medium-High` + +原因： + +- `ATOM native` 与 `ATOM vLLM plugin` 已经在 `PagedAttentionImpl` 等层共享较多 +- 真正难的是 `SGLang plugin` 这一侧 +- 它不走 `PagedAttention` 的 ATOM 主路径，而走： + - `RadixAttention` + - `ATOMAttnBackendForSgl` + - `ForwardBatch` + +所以 MHA 共享的主要困难在于： + +**runtime orchestration 差异** + +### 9.2 MLA + +难度：`High` + +原因： + +- native/server 与 vLLM plugin 仍然较多共享 `MLAAttention` +- 但 SGLang plugin 下的 DeepSeek MLA 已经上卷到 model specialization +- 不只是 backend 不同，连 forward 组织方式都不同 + +所以 MLA 的困难在于： + +**runtime orchestration + model specialization 双重叠加** + +## 10. 已产出的文档与图 + +本次会话过程中，已额外产出下列文档/图，用于不同角度说明问题。 + +### 10.1 代码目录内 Markdown + +- `atom/plugin/decoupling-diagram.md` +- `atom/plugin/vllm-integration-architecture.md` +- `atom/plugin/sglang-attention-backend-survey.md` + +### 10.2 Canvas / 可视化分析 + +- `atom-attention-backend-architecture-review.canvas.tsx` +- `atom-plugin-coupling-risk-analysis.canvas.tsx` +- `atom-plugin-decoupling-diagram.canvas.tsx` +- `atom-vllm-plugin-architecture.canvas.tsx` +- `atom-attention-sharing-modes.canvas.tsx` + +### 10.3 这些产物对应的主题 + +- 当前 attention backend 架构缺陷 +- vLLM / SGLang plugin 耦合与风险 +- plugin 解耦方向与 bootstrap 理解 +- vLLM plugin 集成架构 +- public SGLang backend survey +- 三种模式共享 ATOM attention 的难度评估 + +## 11. 当前最值得继续推进的方向 + +如果延续本次会话的结论，后续工作最值得按下面顺序推进： + +1. 继续收缩 `atom/plugin/` 根目录的 shared runtime 语义 +2. 把 `full_attn / linear_attn / hybrid` 三层作为 plugin 后续结构设计的一等公民 +3. 在 `SGLang plugin` 侧先补出 `hybrid/composite` 组合层 +4. 初期复用 public SGLang linear backend，优先验证 runtime 结构是否成立 +5. 然后再评估哪些 linear backend 需要逐步替换成 ATOM 自己的实现 +6. 对 `MLA` 尽早拆开： + - generic MLA runtime + - DeepSeek specialization + +## 12. 一句话总结 + +本次会话最终把问题收敛为： + +**当前 ATOM plugin attention 的关键矛盾，不是底层 kernel 能不能共享，而是 vLLM / SGLang / native 三种模式在 runtime、metadata、model specialization 和 backend ownership 上没有对齐；后续重构应从“统一 attention backend”升级为“重新定义 plugin 的 sequence backend / host-owned runtime 结构”。** diff --git a/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md b/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md new file mode 100644 index 0000000000..f0f9d5bccc --- /dev/null +++ b/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md @@ -0,0 +1,482 @@ +# 2026-05-07 Radix Cache vs Paged Attention Notes + +> 预估阅读时间：15 分钟 +> 主题：梳理 `SGLang` 中 `radix cache` / `RadixAttention` 与 `ATOM` / `vLLM` 常见的 `paged attention` 之间的关系，解释为什么它们最终可以落到相同的 kernel，以及这件事对 `sglang plugin` attention backend 重构意味着什么。 + +## 1. 这篇笔记想回答什么 + +围绕 `sglang plugin attn backend` 的重构，最近反复出现了几个问题： + +1. `SGLang` 用的是 `RadixAttention`，`ATOM`/`vLLM plugin` 里常见的是 `PagedAttention`，两者是不是两套完全不同的注意力体系？ +2. 如果它们真的不同，为什么在 `DeepSeek MLA` 这类路径里，最终又会调用到相同的底层 kernel？ +3. 既然 `SGLang` 还多了一层 `radix cache` 前缀管理，为什么没有一个非常直观的结论说 "`SGLang` 一定比 `vLLM` 更快 / 更省显存"？ +4. 对当前重构来说，`radix` / `paged` 的差异到底应该被归类到： + - host runtime 差异 + - KV cache 管理差异 + - metadata lowering 差异 + - kernel 差异 + 的哪一层？ + +本文尝试把这几个问题放进同一个分析框架里。 + +## 2. 先给结论 + +### 2.1 最核心的一句话 + +`RadixAttention` 和 `PagedAttention` 主要不是在底层注意力数学上不同，而是在 **runtime 如何管理、共享、命中和定位历史 KV** 上不同。 +一旦 runtime 把历史 KV lowering 成底层 kernel 能吃的 `page_table`、`kv_indptr`、`kv_indices`、`qo_indptr` 之类的 metadata，kernel 就不再关心这些索引最初是来自 radix tree 还是来自 paged block table。 + +### 2.2 另一句更工程化的结论 + +对重构来说，不应该把 “`RadixAttention` vs `PagedAttention`” 当作 “能不能共享底层执行实现” 的直接判断依据。 +更应该把系统拆成三层： + +1. **host/runtime 层** + 例如 `ForwardBatch`、`RadixAttention`、`ModelRunner`、prefix cache、speculative 调度 +2. **execution metadata lowering 层** + 例如 `page_table`、`block_tables`、`kv_indptr`、`kv_indices`、`qo_indptr` +3. **kernel / execution core 层** + 例如 `pa_fwd_asm`、`pa_persistent_fwd`、`flash_attn_varlen_func`、`mla_decode_fwd`、`mla_prefill_fwd` + +`radix` 和 `paged` 的差异，主要在第 1、2 层； +兼容性和共享机会，主要出现在第 2、3 层。 + +## 3. 为什么名字容易让人误解 + +当前最容易造成误解的点是： + +- `PagedAttention` 这个名字听起来像“最终执行算法” +- `RadixAttention` 这个名字也听起来像“最终执行算法” + +但在工程上，它们都更接近 **host-facing attention runtime abstraction**，而不是最终 kernel 名字。 + +更具体一点： + +- `PagedAttention` 更强调 **按固定 page/block 管理 KV cache** +- `RadixAttention` 更强调 **按 prefix/radix tree 管理请求历史与前缀复用** + +这两者都不是在说： + +- `QK^T` 怎么算 +- softmax 怎么做 +- decode kernel 用哪套 asm/triton/flashinfer + +这些底层执行问题，是更下一层的 backend / kernel 决定的。 + +## 4. `RadixAttention` 到底是什么 + +### 4.1 `RadixAttention` 不是 kernel + +在 public `SGLang` 里，`RadixAttention` 本身并不直接定义底层 kernel 选择。 +它更像是： + +- SGLang 模型层里统一使用的 attention layer 壳 +- 它知道自己在 `ForwardBatch` 语义里运行 +- 它把真正的执行委托给 `forward_batch.attn_backend` + +因此，`RadixAttention` 更像一个 **layer/runtime 入口点**。 + +### 4.2 `radix` 的重点是 prefix 管理 + +`radix cache` 的本质，是一棵 prefix tree，用来做： + +- 前缀命中 +- 前缀复用 +- unfinished / finished request 的缓存插入 +- 基于 prefix 的 eviction / lifecycle 管理 + +它回答的问题更像： + +- 这个请求的历史前缀已经缓存到哪里了？ +- 哪一部分可以直接复用，不必重新 prefill？ +- 哪些 KV slot/page 仍然属于共享前缀？ +- 哪些历史片段可以回收？ + +所以 `radix` 优化的主要是 **prefix reuse**，而不是单次 attention kernel 的算力效率。 + +### 4.3 `radix cache` 在 SGLang 里不是完全“反 page”的 + +这一点很重要。public `SGLang` 的 `radix_cache` 本身就已经带有 page-aware 的逻辑。 + +例如它有 page 对齐： + +- `page_align_keys()` + +也有按 page 粒度匹配 prefix： + +- `_key_match_paged()` + +这意味着： + +- `radix cache` 解决的是前缀共享与命中问题 +- 但它并不排斥最终按 page/block 粒度来表达缓存 + +也就是说，**radix tree 的逻辑管理** 与 **paged/block 的物理表达** 并不冲突。 + +## 5. `PagedAttention` 到底是什么 + +### 5.1 `PagedAttention` 的重点不是 prefix tree + +`PagedAttention` 更强调： + +- KV cache 被组织成固定大小的 page/block +- 每个请求通过 `page_table` / `block_table` 找到自己历史对应的 page +- kernel 依照这些 page/block 索引读取历史 KV + +它回答的问题更像： + +- 当前请求历史有哪些 block/page？ +- 它们在物理内存中的位置是什么？ +- 这些 page 如何映射到 kernel 所需的 block table？ + +所以 `paged attention` 更偏： + +- **物理存储布局** +- **kernel 读取 KV 的地址组织方式** + +### 5.2 `PagedAttention` 天生不等于 prefix reuse + +`paged attention` 当然也可以配合 prefix cache 使用，但 prefix reuse 并不是它名字里最核心的概念。 +它的核心价值更在于： + +- 支持非连续物理布局 +- 让 decode kernel 可以按 page/block 高效读取 KV + +## 6. 为什么 `radix` 和 `paged` 最后可以兼容 + +### 6.1 因为它们主要解决的是不同层的问题 + +可以把它们看成： + +- `radix`: 逻辑历史管理器 +- `paged`: 物理页表表达方式 + +两者关系有点像： + +- `radix` 负责决定“哪些历史是共享的、哪些是命中的、请求当前逻辑历史是什么” +- `paged` 负责决定“这些历史最终以什么 page/block 索引形式交给 kernel” + +只要 runtime 最终能把逻辑历史翻译成 page/block 或 indptr/index 形式，底层 kernel 并不需要知道前缀最初是如何组织的。 + +### 6.2 兼容发生在 lowering 之后 + +这一点非常关键： + +上层看起来是： + +- `ForwardBatch` +- `RadixAttention` +- `req_to_token` +- prefix cache + +但 lower 到 backend 之后，会变成执行态 metadata，例如： + +- `page_table` +- `block_tables` +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `kv_last_page_len` + +到了这个层次，kernel 只看： + +- query 的 shape +- KV cache 的 shape +- 历史长度 +- 索引表 / indptr + +它不看： + +- 这套索引是不是来自 radix tree +- 是不是来自某个 prefix cache 命中 +- 是不是通过 `PagedAttention` 对象构造出来的 + +所以兼容性不是“名字兼容”，而是 **execution metadata 兼容**。 + +## 7. public `SGLang` 自己就在做类似 lowering + +这不是 `ATOM plugin` 独有的现象。public `SGLang` 本身就在这么做，只是不同 backend lower 到的 metadata 形态略有不同。 + +### 7.1 `aiter_backend` + +public `SGLang` 的 `aiter_backend` 会从： + +- `forward_batch.seq_lens` +- `forward_batch.req_pool_indices` +- `req_to_token` + +构造： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` + +也就是说，它会把 host runtime 的历史表示 lower 成 AITER kernel 所需的执行态索引。 + +### 7.2 `triton_backend` + +public `SGLang` 的 `triton_backend` 也做类似事情。 +它同样维护： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` + +然后再调用 Triton 的 decode/extend kernel。 + +### 7.3 `trtllm_mha_backend` / 部分 `flashinfer` 路径 + +这类 backend 更偏向： + +- 直接构造 `page_table` +- 或 `block_tables` + +再喂给 flashinfer / TRTLLM 风格的 kernel。 + +所以更准确地说，public `SGLang` 的共同模式不是： + +> “所有 backend 都把 radix runtime 翻成 paged attention” + +而是： + +> “所有 backend 都会把 radix/runtime 层的历史表示，lower 成各自 kernel 能吃的索引/页表 metadata” + +其中有的偏 `kv_indptr/kv_indices`，有的偏 `page_table/block_tables`。 + +## 8. 从 kernel 角度看，为什么可以兼容 + +### 8.1 kernel 真正关心什么 + +绝大多数 attention kernel 真正关心的是： + +1. 当前 query 的排列方式 +2. 每个请求可见的历史 KV 长度 +3. 如何从某个索引结构里拿到历史 KV 的物理地址 +4. page/block 大小、stride、layout +5. 是否需要 prefix/extend/speculative 的特殊 metadata + +它通常不关心： + +1. 这份索引最初是不是从 prefix tree 查出来的 +2. 命中前缀的逻辑是谁做的 +3. runtime 是 `RadixAttention` 还是 `PagedAttention` + +### 8.2 对 MHA kernel 的影响 + +MHA decode 常见地会落到： + +- `pa_fwd_asm` +- `pa_persistent_fwd` +- `flash_attn_varlen_func` + +这些 kernel 更偏 page/block 或 varlen index 风格。 +只要 runtime 最终给出： + +- `page_table` +- `context_lens` +- `block_tables` + +或者： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` + +它们都能算。 + +### 8.3 对 MLA kernel 的影响 + +`DeepSeek MLA` 更容易看出这种兼容性，因为 MLA 的底层算子路径更明确。 + +典型 kernel 包括： + +- `mla_decode_fwd` +- `mla_prefill_fwd` +- `concat_and_cache_mla` +- `fused_qk_rope_concat_and_cache_mla` + +这些 kernel 只要求： + +- 压缩后的 latent KV 表示 +- 以及相应的 `qo_indptr` / `kv_indptr` / `kv_indices` + +所以无论上层 host 是： + +- `ATOM native` +- `ATOM vLLM plugin` +- `ATOM SGLang plugin` + +只要最终 lower 成相同的 MLA execution metadata，就自然会落到同一套 MLA kernel。 + +## 9. `DeepSeek MLA` 为什么尤其容易“看起来收敛” + +### 9.1 因为 MLA 的 execution core 更专用 + +`DeepSeek MLA` 的真正执行对象不是： + +- “`RadixAttention` 版本的 MLA” +- “`PagedAttention` 版本的 MLA” + +而是： + +- 一组固定的 MLA latent/cache 表示 +- 一套固定的 rope + kv cache 写入方式 +- 一套固定的 prefill/decode kernel + +换句话说，MLA 的底层 core 更容易抽成“共享执行层”。 + +### 9.2 因为 host 差异主要体现在 metadata 准备与调度 + +对于 `DeepSeek MLA` 来说，上层 host 差异更多体现在： + +- `positions` 从哪里来 +- `forward_batch` 怎么组织 +- decode / extend / speculative / target_verify 怎么分流 +- `req_to_token` 怎么转换成 `kv_indices` + +一旦进入 `mla_decode_fwd` 或 `mla_prefill_fwd` 前，这些差异已经被消化掉了。 + +所以从外部观察，很容易得到一种印象： + +> “怎么上面完全不一样，下面却是同一套 kernel？” + +其实并不奇怪，因为两边只是上层 runtime 不同，而 MLA execution core 本来就在下层收敛。 + +## 10. 既然兼容，为什么 `SGLang` 没有天然更快或更省显存 + +这是另外一个很容易误判的问题。 + +### 10.1 `radix cache` 省的是前缀重复，不是所有 KV + +`radix cache` 能节省的，是重复前缀带来的重复 prefill 和重复 KV 占用。 + +它主要帮助的是： + +- 大量请求共享长前缀 +- 同 prompt 多分支采样 +- prefix-heavy workload +- 某些 speculative / chunked prefill 场景 + +如果 workload 不满足这些条件，它就不会天然表现为显著优势。 + +### 10.2 它优化的是 prefix reuse，不是单 token decode kernel + +radix 管理层并不会让 `mla_decode_fwd`、`pa_fwd_asm` 这种 kernel 自己更便宜。 +kernel 的 raw 速度主要还是受： + +- 头数 +- head dim +- KV layout +- quant +- page size +- GPU kernel 实现 + +这些因素决定。 + +所以常见现象是： + +- prefix-heavy workload 下，`SGLang` 的收益更明显 +- decode-heavy workload 下，收益没那么明显 +- 如果 `vLLM` 也有自己的 prefix caching/block caching，差距会进一步缩小 + +### 10.3 radix 自己也有管理成本 + +`radix cache` 不是免费层。它也有： + +- prefix match +- tree split / insert / eviction +- request -> KV slot 映射维护 +- speculative / unfinished request 的额外 bookkeeping + +如果命中率不高，这部分成本并不会自动转化成收益。 + +## 11. 对当前重构最有价值的判断 + +### 11.1 不要把 `RadixAttention` vs `PagedAttention` 当成是否共享 execution core 的边界 + +因为从 public `SGLang` 和当前 `ATOM plugin` 的代码路径看，二者最终都在做类似事情： + +- 从 host runtime 的历史表示出发 +- 构造底层 kernel 所需的 execution metadata + +所以： + +- runtime 不同，不代表 execution core 必须不同 +- 名字不同，不代表 kernel 层必须分叉 + +### 11.2 真正的边界更适合这样划 + +#### A. `sglang plugin runtime` 应拥有的部分 + +- `ForwardBatch` +- `RadixAttention` +- prefix cache / radix cache +- `ModelRunner` / `attn_backend_wrapper` 相关调度 +- speculative / graph / verify / chunked prefill 等宿主语义 + +#### B. execution metadata lowering 层 + +这层是最值得重新定义的边界。它负责： + +- 把 `req_to_token`、prefix 命中结果、cache state +- 转成 `page_table` / `kv_indptr` / `kv_indices` / `qo_indptr` + +这层既带有 host 语义，又已经开始接近共享执行核心，是未来最值得梳理清楚的一层。 + +#### C. shared execution core + +例如： + +- `mla_decode_fwd` +- `mla_prefill_fwd` +- `pa_fwd_asm` +- `pa_persistent_fwd` +- `flash_attn_varlen_func` + +以及与这些 kernel 紧耦合的那部分 ATOM core helper。 + +### 11.3 这对 `sglang plugin` 的重构意味着什么 + +当前 `sglang plugin` 最不该做的是： + +- 因为底层 kernel 能共享，就把 `sglang` 的 runtime seam 也强行对齐到 `vllM` +- 或者因为上层叫 `RadixAttention`，就认为它和 `PagedAttention` 完全不能共享底层执行实现 + +更合理的重构方向是： + +1. 保持 `sglang` 的 host/runtime 语义完整 +2. 识别 runtime lowering 层的清晰边界 +3. 尽量把 execution core 继续下沉回共享 `ATOM core` + +## 12. 一种对外更容易讲清楚的表述 + +如果后续要向别人解释，可以用下面这几句话： + +### 12.1 关于 `radix` vs `paged` + +`radix cache` 主要解决的是 prefix 命中与复用，`paged attention` 主要解决的是 page/block 形式的 KV 组织与读取。 +前者偏逻辑管理，后者偏物理表达；两者不是同一层次的问题。 + +### 12.2 关于为什么最终会调到同一个 kernel + +因为 attention kernel 只关心最终的执行态 metadata，比如 `page_table`、`kv_indptr`、`kv_indices`、`qo_indptr`。 +一旦 runtime 把历史 KV lowering 成这类形式，kernel 就不再关心这些索引最初是由 radix tree 还是 paged runtime 生成的。 + +### 12.3 关于为什么 `SGLang` 没有天然显著更快 + +因为 `radix` 带来的主要收益是 prefix reuse，而不是单 token decode kernel 的 raw speed。 +如果 workload 不是 prefix-heavy，或者 decode 占主要成本，那么 radix 带来的优势不会自然放大成“总是更快 / 更省显存”。 + +## 13. 最终总结 + +把这次探索收敛成一句话： + +**`RadixAttention` 与 `PagedAttention` 的主要差异，不在底层注意力数学，也不一定在最终 KV 的物理表示本身，而在 host runtime 如何管理、共享、命中并 lowering 历史 KV；一旦 lowering 完成，底层 kernel 完全可以是共享的。** + +对当前 `sglang plugin attn backend` 的重构来说，真正值得做的不是争论 “`radix` 和 `paged` 能不能统一成一个名字”，而是把系统拆清楚： + +- 哪些是 `sglang` 必须自有的 runtime / prefix cache / scheduling 语义 +- 哪些是 execution metadata lowering +- 哪些已经是可以共享的 `ATOM` execution core + +这比直接讨论 “`radix` 和 `paged` 谁更先进、谁应该替代谁” 更接近真正的工程问题。 diff --git a/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md b/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md new file mode 100644 index 0000000000..e5a63be0f1 --- /dev/null +++ b/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md @@ -0,0 +1,359 @@ +# 2026-05-08 KV Indices and Address Space Notes + +> 预估阅读时间：8-10 分钟 +> 主题：解释 `kv_indptr` / `kv_indices` 为什么不是纯 host metadata，也不是完全 storage-agnostic 的抽象，并用 `sglang` / `ATOM` / `vllm plugin` 的实际 KV cache 存储举例说明。 + +## 1. 想回答的问题 + +最近有一个很容易卡住重构讨论的问题： + +> `kv_indptr` / `kv_indices` 看起来和 KV cache 的存储方式深度绑定。 +> 既然如此，为什么还能把它们看成 `sglang` / `ATOM` / `vllm plugin` 之间可以共享的 metadata？ + +这个问题的关键在于区分两件事： + +1. **host/runtime 是否相同** +2. **kernel 最终消费的地址空间模型是否相同** + +一句话先说结论： + +**`kv_indptr` / `kv_indices` 不是完全 storage-agnostic 的抽象，它们绑定的是某种“可索引的 KV 地址空间”； +但它们可以是 host/runtime-agnostic 的，只要不同宿主最终都能把自己的 runtime state lowering 到同一种地址空间模型。** + +## 2. 三层视角 + +理解这个问题，最好先把 attention 相关代码拆成三层： + +### 2.1 host/runtime 层 + +这一层描述的是宿主框架如何组织请求和调度，例如： + +- `ATOM native` 的 `ScheduledBatch` +- `vllm plugin` 的 `query_start_loc` / `num_prefills` +- `sglang plugin` 的 `ForwardBatch` / `forward_mode` / `spec_info` + +这一层表达的是： + +- 哪些请求在 batch 中 +- 哪些是 decode，哪些是 prefill +- prefix / speculative / graph 的语义是什么 + +### 2.2 execution metadata lowering 层 + +这一层做的事情是： + +- 把 host/runtime 里的 batch 状态 +- 翻译成 kernel 真正能消费的索引结构 + +典型字段： + +- `block_tables` / `page_table` +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `kv_last_page_len` + +### 2.3 kernel / execution core 层 + +这一层只关心： + +- query tensor +- KV cache tensor +- index / indptr / page table +- persistent kernel workspace + +例如： + +- `mla_decode_fwd` +- `mla_prefill_fwd` +- `pa_fwd_asm` +- `pa_persistent_fwd` + +## 3. `sglang` 的实际 KV cache 存储 + +在 `sglang` 里，KV cache 本身不是存在 radix tree 里。 +radix tree 负责的是 prefix 复用与命中；真正的 KV 还是存在 memory pool 里。 + +`sglang` 自己的注释说得很清楚： + +- `ReqToTokenPool` 负责 request -> token location 映射 +- `TokenToKVPoolAllocator` 负责 KV cache index 管理 +- `KVCache` 真正持有物理 K/V tensor + +也就是说，`sglang` 的 runtime 世界里至少有两层： + +1. **逻辑视角**：某个 request 的第 `i` 个 token 属于哪里 +2. **物理视角**：这个 token 对应的 K/V 存在物理 pool 的哪个 slot + +### 3.1 `ReqToTokenPool` + +`ReqToTokenPool` 本质上是一张二维表： + +```text +req_to_token[req_id, token_pos] = physical_slot_id +``` + +它回答的是： + +- 某个 request 的第 `j` 个逻辑 token +- 对应到物理 KV pool 里的哪个 slot + +### 3.2 物理 KV pool + +物理 pool 里，K 和 V 都是按 slot 存的。 +写入 KV 时，最终就是按 `loc/indices` 写入： + +```text +k_cache[indices] = k +v_cache[indices] = v +``` + +所以从 kernel 角度看，最终它看到的是： + +- 一个可以索引的 KV 地址空间 +- 一组 slot index + +而不是看到“radix tree”本身。 + +## 4. `ATOM` / `vllm plugin` 的实际 KV cache 存储 + +`ATOM` / `vllm plugin` 更偏 paged/block 风格。 + +典型思路是： + +- 每个 request 有一张 `block_table` +- 每个 block 对应一个固定大小的 page +- kernel 按 `block_table` 或其展开后的 index 去读 KV + +所以它的上层直觉更像： + +```text +request -> block table -> page/block address -> physical KV +``` + +和 `sglang` 相比，最大的不同不是“有没有物理地址空间”，而是： + +- `sglang` 上层先经过 `req_to_token` +- `vllm/ATOM` 上层先经过 `block_table` + +但两边最终都能 lower 到一组可供 kernel 读取的地址索引。 + +## 5. 例子一：MLA decode 中的 `kv_indptr` / `kv_indices` + +这个例子里，可以把 `kv_indices` 理解为 **token-slot index**。 + +### 5.1 在 `sglang` 里 + +假设一个 request 当前长度是 10， +它在物理 KV pool 里的 slot 顺序不是连续的，而是： + +```text +[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] +``` + +那 `req_to_token[req_id, :10]` 就大致表示： + +```text +logical token position: 0 1 2 3 4 5 6 7 8 9 +physical slot id: 8 9 10 11 0 1 2 3 4 5 +``` + +如果这个 batch 只有这一个 request，那么 lowering 之后可以得到： + +```text +kv_indptr = [0, 10] +kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5] +``` + +这里： + +- `kv_indptr` 表示“第 0 个 request 的 KV 范围是 `[0, 10)`” +- `kv_indices` 表示“真正去物理 KV pool 里读这些 slot” + +### 5.2 在 `vllm/ATOM` 里 + +如果 page size 是 4，同样 10 个 token 对应的 block layout 可以写成： + +```text +block 2 -> slots [8, 9, 10, 11] +block 0 -> slots [0, 1, 2, 3] +block 1 -> slots [4, 5, 6, 7] +``` + +block table 大致就是： + +```text +[2, 0, 1] +``` + +展开前 10 个 token 后，对应的 slot 顺序仍然是： + +```text +[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] +``` + +于是 lower 之后，给 MLA decode kernel 的： + +```text +kv_indptr = [0, 10] +kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5] +``` + +和 `sglang` 这边是等价的。 + +### 5.3 这个例子说明什么 + +说明在 MLA decode 这个路径里： + +- `kv_indices` 绑定的是 **token-slot 地址空间** +- 它不是“完全与存储无关” +- 但它已经不关心 slot 列表最初是由 `req_to_token` 还是 `block_table` 推出来的 + +也就是说，**它和具体宿主无关，但和被选定的地址空间模型有关。** + +## 6. 例子二：MHA persistent decode 中的 `page_table` / `kv_indices` + +这个例子里，`kv_indices` 更接近 **page/block index**，而不是 token-slot index。 + +### 6.1 在 `sglang plugin` 里 + +仍然假设 page size = 4。 +如果某个 request 的 token-slot 顺序是： + +```text +[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] +``` + +那么每页的起始 slot 是： + +```text +[8, 0, 4] +``` + +除以 page size 后，对应 page id： + +```text +[2, 0, 1] +``` + +在 `sglang plugin` 里，这就是 `page_table` 的来源： +从 `req_to_token` 按页采样，再除以 `page_size`。 + +然后 `_build_pa_metadata_for_decode()` 再把它变成 `pa_persistent_fwd` 所需的 page-level `kv_indices`。 + +### 6.2 在 `ATOM` / `vllm` 里 + +这一层本来就是 paged/block 视角，所以 `block_table` 原生就已经是： + +```text +[2, 0, 1] +``` + +也就是说： + +- `sglang` 是 `req_to_token -> page_table` +- `vllm/ATOM` 是直接 `block_table -> page_table` + +但到了 persistent kernel 眼里，最后看到的是同一种 page/block 地址空间。 + +### 6.3 这个例子说明什么 + +说明在 MHA persistent decode 这个路径里： + +- `kv_indices` 绑定的是 **page/block 地址空间** +- 它仍然不是“完全 storage-agnostic” +- 但也不需要知道上层 host 是 `sglang` 还是 `vllm` + +只要两边都能 lower 到同一个 page/block 地址空间，kernel 就能共享。 + +## 7. 这两个例子真正说明的事 + +上面两个例子其实在说明同一个结论： + +### 7.1 `kv_indptr` / `kv_indices` 不是 host metadata + +它们不描述： + +- `ForwardBatch.forward_mode` +- `query_start_loc` +- `num_prefills` +- `run_graph` +- prefix hit 的高层语义 + +它们只描述： + +- 当前这次 kernel 调用 +- 要读哪些 KV 地址 +- 每个 request 的边界在哪里 + +所以它们更接近： + +- execution metadata +- kernel-facing metadata + +而不是 host/runtime metadata。 + +### 7.2 但它们也不是完全 storage-agnostic + +它们依赖： + +- 当前 KV cache 的地址空间模型 +- kernel 对该地址空间的解释方式 + +如果底层 cache 不是： + +- 可线性索引的 token slot +- 或可 flatten 的 page/block index + +那 `kv_indptr` / `kv_indices` 这套抽象就未必成立。 + +所以它们不是“无限通用”的。 + +## 8. 对重构的直接启示 + +这个背景知识对当前 `sglang attention` 重构最有价值的结论是： + +### 8.1 不要去统一 host metadata + +比如不要强行统一： + +- `ForwardBatch` +- `query_start_loc` +- `num_prefills` +- `run_graph` + +这些都属于宿主框架自己的 runtime 语言。 + +### 8.2 应该争取统一 execution metadata + +例如： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `page_table` +- `kv_last_page_len` +- `work_indptr` +- `reduce_indptr` + +以及围绕这些字段构造的更小 dataclass，例如当前 draft 里已经开始尝试的： + +- `MLADecodeKernelMetadata` +- `MHAAsmKernelMetadata` +- `MHAPersistentKernelMetadata` + +### 8.3 正确的共享边界 + +所以真正可共享的不是： + +- `sglang` 的 `ForwardMetadata` +- `vllm plugin` 那一整套 metadata family + +而是： + +- 它们再往下切出来的 +- kernel-facing execution metadata + +## 9. 一句话总结 + +**`kv_indptr` / `kv_indices` 不是“和存储彻底无关”的抽象，它们绑定的是某种被 kernel 接受的 KV 地址空间；但正因为它们已经脱离了宿主 batch/runtime 语义，`sglang` 和 `vllm/ATOM` 即使上层完全不同，也仍然可以在这层 metadata 上收敛并共享 kernel 调用逻辑。** diff --git a/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md b/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md new file mode 100644 index 0000000000..86bccce8e6 --- /dev/null +++ b/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md @@ -0,0 +1,144 @@ +# 2026-05-08 Metadata Layering Summary + +## 1. 目的 + +这份笔记简要总结 `ATOM native`、`vllm plugin`、`sglang plugin` 三边 metadata 的关系，重点回答： + +- `ATOM/atom/utils/forward_context.py` 中的 `AttentionMetaData` 是什么 +- `ATOM/atom/plugin/attention.py` 里的 metadata dataclass 在做什么 +- `ATOM/atom/plugin/sglang/attention_backend/sgl_attn_backend.py` 里的 `ForwardMetadata` 属于哪一层 +- 哪些 metadata 值得共享，哪些不值得强行统一 + +## 2. 三套 metadata 的定位 + +### 2.1 `AttentionMetaData` + +文件：`ATOM/atom/utils/forward_context.py` + +它是 `ATOM` attention 执行链路里的**通用外层容器**。 +里面既有： + +- 通用 attention 字段 + - `slot_mapping` + - `block_tables` + - `kv_indptr` + - `kv_indices` + - `work_meta_data` +- 也预留了 plugin 扩展入口 + - `plugin_metadata` + +所以它更像一个大容器，而不是某个 host 专属的数据模型。 + +### 2.2 `vllm plugin metadata` + +文件：`ATOM/atom/plugin/attention.py` + +`vllm plugin` 不是重写一套完全独立的 metadata，而是在 `AttentionMetaData` 外壳上，额外挂了一层 `plugin_metadata` payload。 + +代表类型包括： + +- `AiterFlashAttentionMetadataForPluginMode` +- `AiterMLACommonMetadataForPluginMode` +- `AiterMLADecodeMetadataForPluginMode` +- `AiterMLASparseMetadataForPluginMode` + +这些 dataclass 主要承载： + +- `vllm` 的 batch/runtime 语义 +- decode/prefill/extend phase 拆分 +- chunked prefill / DCP / sparse MLA 等 feature-specific 信息 + +一句话：**`vllm plugin` 是“外层复用 `AttentionMetaData`，内层自带一套 host-specific metadata family”。** + +### 2.3 `sglang plugin` 的 `ForwardMetadata` + +文件：`ATOM/atom/plugin/sglang/attention_backend/sgl_attn_backend.py` + +`ForwardMetadata` 不是 `AttentionMetaData.plugin_metadata` 那种 payload，而是 `sglang backend` 内部的**lowering 结果缓存对象**。 + +它更接近： + +- `ForwardBatch` +- `req_to_token` +- `token_to_kv_pool` +- graph/speculative path + +向 kernel metadata 的过渡层。 + +一句话：**`ForwardMetadata` 更像 backend-local lowering state，而不是 host-agnostic 的最终执行 metadata。** + +## 3. 三类字段 + +### 3.1 host/runtime-bound + +这类字段描述宿主框架如何组织 batch/runtime，而不是 kernel 如何计算。 + +例子： + +- `vllm plugin` + - `query_start_loc` + - `num_decodes` + - `num_prefills` + - `chunked_context` +- `sglang plugin` + - `run_graph` + - 各种 `ForwardBatch.forward_mode` 派生分支语义 + +### 3.2 execution/kernel-bound + +这类字段最接近 kernel 真正需要的执行信息。 + +例子： + +- `slot_mapping` +- `block_tables` +- `context_lens` +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `kv_last_page_len` +- `work_indptr` +- `reduce_indptr` + +### 3.3 hardware-sensitive + +这类字段和 page size、dtype、kernel 选择、并行配置相关。 + +例子： + +- `max_q_len` +- `max_kv_len` +- `num_kv_splits` +- `attn_out_dtype` +- `head_dim` + +## 4. 哪些值得共享 + +### 不建议强行共享的 + +- `AttentionMetaData` 整体 +- `vllm plugin` 整套 `plugin_metadata` +- `sglang plugin` 整体的 `ForwardMetadata` + +原因不是它们没价值，而是它们都带着很重的 host/runtime 语义。 + +### 值得重点共享的 + +应该共享的是更小的 **execution metadata view**，例如当前 draft 已经开始尝试的： + +- `MLADecodeKernelMetadata` +- `MHAAsmKernelMetadata` +- `MHAPersistentKernelMetadata` + +这些结构只保留某个 kernel call 真正需要的字段，更适合在： + +- `ATOM core` +- `vllm plugin` +- `sglang plugin` + +之间共享。 + +## 5. 一句话结论 + +现在不是三边共用“一套 metadata”，而是三边各自都有一层 host-owned metadata / lowering 体系。 +真正值得共享的，不是这些大 metadata 本身，而是从它们里面切出来的、更小的、kernel-facing execution metadata。 diff --git a/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md b/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md new file mode 100644 index 0000000000..3b2128f5bd --- /dev/null +++ b/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md @@ -0,0 +1,398 @@ +# 2026-05-11 ATOM Attention Refactory Session Summary + +> 主题：从 `sglang plugin` 设计者视角，梳理 `attention backend` 重构中的 runtime / metadata / KV cache / kernel reuse 边界，并记录几种 kernel 共用 draft 的结论。 + +## 1. 本次会话的主线 + +本次讨论的核心不是继续从 `vllm plugin` 出发看问题，而是明确切换到： + +- **`sglang plugin attn backend` 设计者视角** +- 关注 `sglang plugin` 自己的 runtime 语义 +- 判断哪些层应该继续由 `sglang plugin` 自有 +- 判断哪些层值得尽量与 `ATOM core` 共用 + +围绕这个目标，本次会话主要探索了 6 个问题： + +1. `RadixAttention` / `radix cache` 与 `PagedAttention` 的关系 +2. public `sglang` 是否也会把 radix runtime lowering 到 paged / index metadata +3. metadata 为什么会越长越多，以及哪些字段本质上属于 host/runtime +4. `kv_indptr` / `kv_indices` 为什么既绑定存储模型，又仍然能跨 host 共用 +5. `sglang plugin` 和 public `sglang` 在 KV cache 管理上的异同 +6. 在不大改文件结构、不做过度软件工程的前提下，怎么试探性地共用 kernel call + +## 2. 对 `sglang plugin` 重构最关键的几个判断 + +### 2.1 不是所有问题都该叫 “attention backend 复用” + +当前问题至少分成三层： + +1. **host/runtime 层** + - `ForwardBatch` + - `RadixAttention` + - speculative / graph / verify / extend 调度 +2. **execution metadata lowering 层** + - `page_table` + - `kv_indptr` + - `kv_indices` + - `qo_indptr` + - `kv_last_page_len` +3. **kernel / execution core 层** + - `flash_attn_varlen_func` + - `mla_decode_fwd` + - `mla_prefill_fwd` + - `pa_fwd_asm` + - `pa_persistent_fwd` + +后续任何“共用更多 code”的讨论，都应该先说明是在第几层说话。 + +### 2.2 `sglang plugin` 不应该为了复用去对齐到 `vllm` 的 host runtime + +`vllm plugin` 和 `sglang plugin` 上层接入模型不同： + +- `vllm plugin` 更像 layer-owned / metadata-builder-owned 路径 +- `sglang plugin` 更像 `ForwardBatch` 驱动的 backend runtime 路径 + +所以： + +- 不该共享 `vllm plugin` 那层 runtime facade +- 也不该强行统一大 metadata 对象 +- 真正值得共享的，是更靠近 kernel 的 execution metadata 和 kernel call + +### 2.3 `sglang plugin` 想共用 `ATOM` code，最现实的边界在 kernel call 一层 + +如果先不做大规模结构改造，最适合先共用的是： + +- `mla_decode_fwd` +- `pa_persistent_fwd` +- `pa_fwd_asm` +- `flash_attn_varlen_func` + +也就是： + +- 保留各自的 runtime lowering +- 先把“掉 kernel 前最后一跳”抽出来 + +## 3. `RadixAttention` / `radix cache` 和 `PagedAttention` 的结论 + +### 3.1 `RadixAttention` 不是最终 kernel + +`RadixAttention` 更像： + +- `sglang` 模型层统一使用的 attention layer 壳 +- 它知道自己跑在 `ForwardBatch` 语义里 +- 最终还是委托给 `forward_batch.attn_backend` + +所以它不是 “另一套底层注意力数学”，而是 **host-facing runtime abstraction**。 + +### 3.2 `PagedAttention` 和 `radix` 解决的问题不在同一层 + +- `radix cache` 更偏 prefix 命中 / 复用 / eviction +- `paged attention` 更偏 block/page 形式的物理组织和读取 + +可以理解成： + +- `radix` 管“历史怎么共享和命中” +- `paged` 管“命中的历史最后如何变成 page/block 索引给 kernel” + +### 3.3 public `sglang` 自己也在做类似 lowering + +这不是 `ATOM plugin` 才有的行为。 + +public `sglang` 的： + +- `aiter_backend` +- `triton_backend` +- `trtllm_mha_backend` +- 部分 `flashinfer` 路径 + +本身也会把： + +- `ForwardBatch` +- `req_to_token` +- `seq_lens` + +lower 成： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `page_table` / `block_tables` + +所以，`radix runtime -> paged/index metadata` 不是 plugin 特例，而是公共设计模式的一部分。 + +## 4. Metadata 分层结论 + +### 4.1 现在至少有三套 metadata family + +1. **`ATOM native`** + - `AttentionMetaData` +2. **`vllm plugin`** + - `AttentionMetaData` 外壳 + - `plugin_metadata` 内层 payload +3. **`sglang plugin`** + - `ForwardMetadata` + +### 4.2 `AttentionMetaData` 和 `plugin/attention.py` 的关系 + +`AttentionMetaData` 是外层通用容器； +`plugin/attention.py` 里的很多 dataclass，是 `vllm plugin` 为表达宿主 batch/runtime 语义而塞进 `plugin_metadata` 的 payload。 + +所以它们不是并列关系，而是： + +```text +AttentionMetaData + └─ plugin_metadata + ├─ flash-style plugin metadata + ├─ MLA plugin metadata + └─ sparse MLA plugin metadata +``` + +### 4.3 `ForwardMetadata` 的定位 + +`ForwardMetadata` 更像： + +- `sglang backend` 内部的 lowering state +- 它不是最终统一 metadata +- 更不是天然可共享的 host-agnostic object + +它里面混着： + +- runtime control 信息 +- execution metadata +- persistent kernel 预构造结果 + +## 5. 三类字段 + +为了判断能不能共享，本次把字段分成三类： + +### 5.1 host/runtime-bound + +典型例子： + +- `vllm` 的 `query_start_loc` +- `num_decodes` +- `num_prefills` +- `chunked_context` +- `sglang` 的 `run_graph` + +这些字段描述的是宿主框架如何组织 batch/runtime，而不是 kernel 怎么算。 + +### 5.2 execution/kernel-bound + +典型例子： + +- `slot_mapping` +- `block_tables` +- `page_table` +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `kv_last_page_len` +- `work_indptr` +- `reduce_indptr` + +这些字段才是 kernel 真正需要消费的执行信息。 + +### 5.3 hardware-sensitive + +典型例子： + +- `max_q_len` +- `max_kv_len` +- `num_kv_splits` +- `head_dim` +- `attn_out_dtype` + +这些字段不是 host 语义，但会影响： + +- page/block 解释 +- kernel 选路 +- workspace/split 策略 + +## 6. 关于 `kv_indptr` / `kv_indices` 的关键结论 + +这个问题本次单独做了较多背景解释，结论是： + +### 6.1 它们不是完全 storage-agnostic + +因为它们绑定的是某种 **可索引的 KV 地址空间**： + +- token slot index +- page/block index +- compacted KV index + +所以它们不能脱离底层存储模型独立存在。 + +### 6.2 但它们可以是 host/runtime-agnostic + +因为它们不表达： + +- `ForwardBatch.forward_mode` +- `query_start_loc` +- prefix 命中策略 +- graph/speculative 的宿主控制语义 + +它们只表达： + +- 这次 kernel 要去哪些 KV 地址读数据 +- 每个 request 的边界在哪里 + +### 6.3 两个例子 + +#### 例子 A：`sglang` + +- `req_to_token[req_id, token_pos] = slot_id` +- lowering 后得到： + - `kv_indptr = [0, 10]` + - `kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5]` + +#### 例子 B：`vllm/ATOM` + +- 上层是 `block_table` +- 展开后同样可得到： + - `kv_indptr = [0, 10]` + - `kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5]` + +这说明： + +- 上层来源不同 +- 但 lower 后的 execution metadata 可以收敛 + +## 7. public `sglang` 与 `sglang plugin` 在 KV cache 管理上的异同 + +### 7.1 相同点 + +- owner 都还是 public `sglang` + - `ReqToTokenPool` + - `TokenToKVPool` + - `KVCache` + - `radix_cache` +- request 到 slot 的映射机制没变 +- prefix / radix 复用体系没变 + +### 7.2 不同点 + +#### A. MHA 写 cache 路径不同 + +public `sglang` 更多沿用 pool 提供的标准写入接口： + +- `token_to_kv_pool.set_kv_buffer(...)` + +而 `sglang plugin` 在 MHA 路径里会绕过标准写法，自己做一次： + +- `set_kv_buffer_with_layout_shuffle(...)` +- 把同一块 public pool buffer 按更偏 `ATOM` kernel 的 SHUFFLE layout 写进去 + +#### B. MLA 路径更接近 public `sglang` + +MLA 上，plugin 更多沿用 public pool 的 contract： + +- `token_to_kv_pool.set_kv_buffer(...)` +- `get_key_buffer(...)` +- `get_value_buffer(...)` + +#### C. Lowering 目标不同 + +- public `sglang` 的 metadata 更通用 +- plugin 的 `ForwardMetadata` 更偏 `ATOM` kernel + - `page_table` + - `pa_metadata_*` + +### 7.3 一个更准确的说法 + +`sglang plugin` **没有改写 public `sglang` 的 pool owner / allocator / radix tree**， +但它在 MHA 路径上： + +- 改变了同一块 public pool buffer 的布局解释与写入约定 +- 并把 lowering 结果继续朝 `ATOM` kernel 所需的 metadata 形态推进 + +## 8. Kernel 共用的两种 draft + +本次会话里，实际做了两种方向的 draft。 + +### 8.1 方案 A：共享 kernel wrapper + +思路： + +- 新增共享小 metadata + - `MLADecodeKernelMetadata` + - `MHAAsmKernelMetadata` + - `MHAPersistentKernelMetadata` +- 新增共享 wrapper + - `run_mla_decode_kernel` + - `run_mha_asm_kernel` + - `run_mha_persistent_kernel` + +优点： + +- 边界清晰 +- 更像长期可演进的 execution API +- 不需要统一大 metadata + +缺点： + +- 还是引入了一层新的 shared contract + +### 8.2 方案 B：`sglang plugin` 主动适配 `ATOM core` + +思路： + +- 不改 `ATOM core` +- 在 `sglang plugin` 里把 `ForwardMetadata` 转成 `AttentionMetaData` +- 再伪造最小 `shim/self` +- 直接调用： + - `PagedAttentionImpl.paged_attention_asm` + - `PagedAttentionImpl.paged_attention_persistent_asm` + - `MLAAttention._forward_decode` + +优点： + +- 不改 `ATOM core` +- 更直观地证明 “plugin 主动复用 core” 是可行的 + +缺点： + +- 需要 shim +- 尤其 MLA 路径更脆弱 +- 更适合验证可行性，不适合长期保留 + +## 9. 当前最值得保留的判断 + +### 9.1 不应该强行统一大 metadata + +不建议直接统一： + +- `AttentionMetaData` +- `plugin_metadata` +- `ForwardMetadata` + +因为它们都带有较重的宿主 runtime 语义。 + +### 9.2 最值得共享的是 kernel-facing execution metadata + +真正最有价值的共享边界是： + +- `kv_indptr` +- `kv_indices` +- `qo_indptr` +- `page_table` +- `context_lens` +- `work_*` +- `reduce_*` + +以及基于它们的 kernel call / runner。 + +### 9.3 重构顺序建议 + +如果继续推进 `sglang plugin` 重构，更合理的顺序是： + +1. 保持 `sglang` runtime 自有 +2. 显式化 metadata lowering 边界 +3. 尽量共享 kernel call / execution helper +4. 最后再考虑是否继续上推到 runner 层 + +## 10. 一句话总结 + +本次会话最终收敛出的关键认识是： + +**`sglang plugin attention` 重构的关键，不是把 `sglang` runtime 改得更像 `vllm`，而是把 runtime、metadata lowering、kernel execution 这三层拆清楚；只要 lowering 之后能在 execution metadata 上对齐，就能在不统一宿主语义的前提下，共用更多 `ATOM` 的底层 attention 代码。** diff --git a/work_log/attn_refractory/decoupling-diagram.md b/work_log/attn_refractory/decoupling-diagram.md new file mode 100644 index 0000000000..b95c0031e5 --- /dev/null +++ b/work_log/attn_refractory/decoupling-diagram.md @@ -0,0 +1,158 @@ +# ATOM Plugin Decoupling Diagram + +下面这张图对比了当前入口架构和目标解耦架构。这里采用的是更激进、也更干净的方案： + +- `vLLM plugin` 和 `SGLang plugin` 在 `atom/plugin/` 层彻底分家 +- 不再保留共享的 plugin runtime +- 共享只发生在更底层的 `ATOM core` + +## Current + +```mermaid +flowchart TD + A[Shared Plugin Entry
prepare.py / register.py / config.py] + B[Global Runtime State
_CURRENT_FRAMEWORK
current_atom_config
ops.Attention] + C[vLLM Branch
platform / registry / patch / graph] + D[SGLang Branch
wrapper / registry hijack / RadixAttention / graph] + E[Mixed Plugin Core
attention selection / static_forward_context / loader glue] + + A --> B + B --> C + B --> D + C --> E + D --> E +``` + +### 当前设计的核心问题 + +- 入口层通过切换全局状态来区分 `vLLM` 和 `SGLang`，而不是通过物理隔离的子系统建立边界。 +- `prepare.py` / `register.py` / `config.py` 这些共享入口，本身就是耦合中心。 +- host 差异直接泄漏到 plugin 内部结构里，导致 backend 选择、attention 抽象和 runtime context 都知道上层框架。 + +## Target + +```mermaid +flowchart TD + A1[vLLM Plugin
bootstrap.py
register.py
config.py
runtime.py
attention glue
model adapters] + A2[SGLang Plugin
bootstrap.py
register.py
config.py
runtime.py
attention glue
model adapters] + + B[ATOM Core
model_ops
kernels
loader
metadata structs
model specialization] + + A1 --> B + A2 --> B +``` + +### 目标设计的关键点 + +- `vLLM` 和 `SGLang` 在 plugin 层是两个完整子系统，而不是一套共享 runtime 上的两个分支。 +- `atom/plugin/` 下不再存在共享的 `prepare_model()`、共享的 `register.py` 总入口、共享的 plugin config translator。 +- `vLLM plugin` 自己处理： + - platform registration + - model override + - graph / MLA / weight hook patch + - vLLM model adapter / attention glue +- `SGLang plugin` 自己处理： + - external model wrapper + - attention backend registration / override + - graph / speculative / wrapper patch + - SGLang model adapter / attention glue +- 真正共享的只保留在更底层的 `ATOM core`： + - `model_ops` + - kernels + - loader + - generic metadata structures + - model-family specialization + +## 最重要的切断点 + +```mermaid +flowchart LR + A[Cut 1
shared prepare.py] + B[Cut 2
shared register.py] + C[Cut 3
shared config.py] + D[Cut 4
framework global state] + E[Cut 5
ops.Attention global rebinding] + + A --> F[vllm/ and sglang/ own bootstrap] + B --> F + C --> F + D --> G[host-local runtime only] + E --> H[host-specific attention glue] +``` + +## 推荐目录形态 + +```text +atom/plugin/ + vllm/ + bootstrap.py + register.py + config.py + runtime.py + attention/ + models/ + graph/ + sglang/ + bootstrap.py + register.py + config.py + runtime.py + attention/ + models/ + graph/ +``` + +这里的重点不是“换目录”，而是： + +- `vllm/runtime.py` 和 `sglang/runtime.py` 各自独立 +- `vllm/config.py` 和 `sglang/config.py` 各自独立 +- `vllm/register.py` 和 `sglang/register.py` 各自独立 +- 不再有 plugin 级共享 runtime + +## 术语说明 + +### bootstrap + +`bootstrap` 在这里指的是“插件被宿主框架加载时，最先执行的接线初始化层”。 + +它负责的事情通常包括： + +- 注册 plugin 扩展点 +- 安装 patch / hook +- 构造 host adapter / wrapper +- 把宿主配置翻译到 ATOM 可用的格式 + +它不应该负责： + +- attention forward 本身 +- 长期 runtime state 管理 +- 每次请求都要走的核心执行逻辑 + +### register + +`register` 指的是把 plugin 的能力挂到宿主暴露的扩展点上，例如： + +- 注册 platform +- 注册 model class +- 注册 attention backend + +它更偏“声明接入点”，通常是 bootstrap 的一部分。 + +### adapter + +`adapter` 指的是宿主框架和 ATOM core 之间的桥接层。 +它的职责是做接口和参数语义翻译，而不是定义 attention / model 的核心语义。 + +### runtime + +`runtime` 指的是插件在宿主里长期存在、服务执行链路的那部分运行时逻辑。 +如果按这份设计推进，`runtime` 应该是 host-local 的： + +- `vLLM runtime` 只属于 `vLLM plugin` +- `SGLang runtime` 只属于 `SGLang plugin` + +而不是共享一个“Plugin runtime”。 + +## 一句话结论 + +真正的解耦不是把共享入口再包装一层，而是把 `atom/plugin` 从“共享状态机”改成“两个完全独立的 host plugin 子系统”，共享只留给更底层的 `ATOM core`。 diff --git a/work_log/attn_refractory/vllm-integration-architecture.md b/work_log/attn_refractory/vllm-integration-architecture.md new file mode 100644 index 0000000000..e8876268d0 --- /dev/null +++ b/work_log/attn_refractory/vllm-integration-architecture.md @@ -0,0 +1,269 @@ +# ATOM vLLM Plugin Integration Architecture + +## 1. 启动入口 + +vLLM plugin 的 setuptools entrypoint 定义在： + +- `pyproject.toml` + +关键位置： + +- `vllm.platform_plugins` + - `atom = "atom.plugin.vllm.register:register_platform"` +- `vllm.general_plugins` + - `atom_model_registry = "atom.plugin.vllm.register:register_model"` + +这意味着 vLLM 启动时会先进入： + +- `atom/plugin/vllm/register.py::register_platform()` +- `atom/plugin/vllm/register.py::register_model()` + +## 2. 关键代码文件与职责 + +### 2.1 `atom/plugin/vllm/register.py` + +这是 vLLM plugin 的启动接线层，主要负责： + +- 设置 plugin mode +- 返回自定义 platform +- 覆盖 vLLM 的 `ModelRegistry` +- 安装 MLA patch +- patch `Attention.process_weights_after_loading` +- 安装 graph capture patch + +重点符号： + +- `register_platform()` +- `register_model()` +- `_VLLM_MODEL_REGISTRY_OVERRIDES` + +## 2.2 `atom/plugin/vllm/platform.py` + +这是 vLLM 平台侧 attention backend 选择入口。 + +重点符号： + +- `ATOMPlatform.get_attn_backend_cls()` + +它根据 `attn_selector_config` 决定： + +- MHA -> `atom.model_ops.attentions.aiter_attention.AiterBackend` +- MLA -> `atom.model_ops.attentions.aiter_mla.AiterMLABackend` +- sparse MLA -> `atom.plugin.vllm.attention_backend.mla_sparse.AiterMLASparseBackend` + +这里是 vLLM runtime 真正选择 “哪个 ATOM attention backend class” 的入口。 + +## 2.3 `atom/plugin/vllm/model_wrapper.py` + +这是 vLLM model integration 的核心文件。 + +主要职责： + +- 把 `VllmConfig` 翻译成 `atom_config` +- 做 plugin 运行环境准备 +- 选择并实例化 ATOM 模型类 +- 把 vLLM forward context 中的 `positions` 传给 ATOM 路径 +- 对 sparse MLA indexer 做额外注册 + +重点符号： + +- `_prepare_env()` +- `ATOMModelBase.__init__()` +- `_register_indexer_caches_with_vllm()` +- `forward()` + +## 2.4 `atom/plugin/config.py` + +这是 vLLM 配置翻译的桥。 + +重点符号： + +- `generate_atom_config_for_vllm_plugin()` +- `_generate_atom_config_from_vllm_config()` + +它负责把： + +- `VllmConfig` +- `scheduler_config` +- `cache_config` +- `parallel_config` +- `quant_config` + +映射成 ATOM 自己的 `Config` 与 `PluginConfig`。 + +## 2.5 `atom/model_ops/paged_attention.py` + +这是 ATOM attention 在 vLLM plugin 下真正进入执行的关键对象。 + +重点符号： + +- `PagedAttention.__init__()` +- `PagedAttention.forward()` + +在 `is_vllm()` 下，它不会直接走 native/server 的那套 `impl` 初始化，而是： + +- 构造 vLLM 的 `Attention` / `MLAAttention` 外壳 +- 把额外 impl 参数塞进去 +- 注册到 `static_forward_context` +- forward 时通过 `unified_attention_with_output_base_for_plugin_mode()` 回调到 ATOM impl + +## 2.6 `atom/plugin/attention.py` + +这是 vLLM plugin mode 的核心 glue 层。 + +虽然文件在 `plugin/` 根目录，但从职责看它明显偏 vLLM。 + +主要职责： + +- 定义 `vllmAiterAttentionBackendMethods` +- 提供 backend decorator / metadata builder decorator +- 处理 plugin-mode metadata +- 处理 `static_forward_context` +- 为 vLLM 的 attention backend 补 OOT 接口行为 + +重点符号： + +- `vllmAiterAttentionBackendMethods` +- 各类 `*DecoratorForPluginMode` + +## 2.7 `atom/plugin/attention_mha.py` + +这是 vLLM plugin mode 下 MHA impl 的补丁层。 + +重点符号： + +- `PagedAttentionImplPluginModeMethods` + +它补的是： + +- rope/cache/update +- plugin-mode forward +- 与 vLLM KV cache/metadata 对齐的逻辑 + +## 2.8 `atom/plugin/attention_mla.py` + +这是 vLLM plugin mode 下 MLA impl 的补丁层。 + +重点符号： + +- `MLAAttentionImplPluginModeMethods` +- `_mla_plugin_mode_init` + +它补的是： + +- MLA plugin-mode metadata +- qk rope / kv cache / chunked prefill +- vLLM MLA 路径专有行为 + +## 2.9 `atom/plugin/vllm/mla_patch.py` + +这是 vLLM 原生 `MLAAttention` 的 monkey patch 入口。 + +重点符号： + +- `_patch_vllm_mla_attention_forward_impl()` +- `_patch_vllm_mla_attention_process_weights_after_loading()` +- `patch_vllm_mla_attention()` + +这层作用是： + +- 把 vLLM `MLAAttention.forward_impl()` 改接到 ATOM impl +- 把权重后处理改接到 ATOM 的 `impl.process_weights_after_loading()` + +## 2.10 `atom/plugin/vllm/graph_capture_patch.py` + +这是 vLLM graph capture 补丁入口。 + +重点符号： + +- `apply_graph_capture_patch()` + +它实际委托到共享实现： + +- `atom/plugin/graph_capture_patch.py` + +但调用点是 vLLM 自己的 plugin register 流程。 + +## 3. 当前 vLLM plugin 集成链路 + +```mermaid +flowchart TD + A[vLLM startup] + B[entrypoint
register_platform] + C[entrypoint
register_model] + D[ATOMPlatform.get_attn_backend_cls] + E[ATOMModelBase] + F[generate_atom_config_for_vllm_plugin] + G[_prepare_env
set_attn_cls + init_aiter_dist] + H[Instantiate ATOM model] + I[PagedAttention / MLAAttention wrapper] + J[plugin/attention*.py glue] + K[ATOM impl
PagedAttentionImpl / MLAAttention] + L[aiter kernels] + + A --> B + A --> C + B --> D + C --> E + E --> F + E --> G + E --> H + H --> I + I --> J + J --> K + K --> L +``` + +## 4. 最关键的架构判断 + +### 4.1 vLLM plugin 不是简单“替换 backend class” + +它实际上同时做了三件事： + +1. 替换 vLLM 的 model entry +2. 替换/选择 vLLM 的 attention backend class +3. 用 patch + decorator 的方式把 vLLM 的 layer/runtime 语义桥接到 ATOM impl + +### 4.2 真正共享得最好的是 native/server 与 vLLM plugin 的 attention impl + +尤其： + +- `PagedAttentionImpl` +- `MLAAttention` +- 低层 `aiter` kernel family + +也就是说，vLLM plugin 侧的关键复杂度主要不在 kernel，而在： + +- `ModelRegistry` override +- `VllmConfig -> atom_config` +- `static_forward_context` +- `MLAAttention` patch +- graph capture patch + +### 4.3 当前 `plugin/attention.py`、`attention_mha.py`、`attention_mla.py` + +虽然放在 `atom/plugin/` 根目录，但从 vLLM integration 的角度看，它们本质上更像： + +- `vLLM plugin impl glue` +- 而不是真正框架无关的 shared runtime + +## 5. 推荐你重点阅读的文件顺序 + +如果你是为了快速理解 vLLM plugin 集成架构，建议按这个顺序读： + +1. `pyproject.toml` +2. `atom/plugin/vllm/register.py` +3. `atom/plugin/vllm/platform.py` +4. `atom/plugin/vllm/model_wrapper.py` +5. `atom/plugin/config.py` +6. `atom/model_ops/paged_attention.py` +7. `atom/plugin/vllm/mla_patch.py` +8. `atom/plugin/attention.py` +9. `atom/plugin/attention_mha.py` +10. `atom/plugin/attention_mla.py` + +## 6. 一句话结论 + +vLLM plugin 的集成方式，本质上是： + +**用 entrypoint + platform + model wrapper 接入 vLLM，用 ATOM 自己的 attention impl 和 aiter kernels 提供核心执行，再用 patch/decorator 把 vLLM 的 layer/runtime 语义桥接进去。** From e1d06c729da7fecb4597449a4e606f07f88a93a6 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Tue, 12 May 2026 08:29:41 +0000 Subject: [PATCH 73/76] [ATOM-SGL][Attn refrac] Separate model-specific MLA from SGL full attention backend --- atom/model_ops/__init__.py | 4 +- atom/model_ops/attentions/aiter_attention.py | 4 +- atom/plugin/register.py | 2 +- .../full_attention/__init__.py | 8 + .../full_attention_backend.py} | 12 +- .../{ => full_attention}/radix_attention.py | 0 .../sglang/models/base_model_wrapper.py | 2 +- atom/plugin/sglang/models/deepseek_mla.py | 75 +++++++++ .../deepseek_mla_forward.py} | 154 +++--------------- tests/plugin/test_sglang_model_wrapper.py | 6 +- tests/plugin/test_sglang_register.py | 21 ++- 11 files changed, 137 insertions(+), 151 deletions(-) create mode 100644 atom/plugin/sglang/attention_backend/full_attention/__init__.py rename atom/plugin/sglang/attention_backend/{sgl_attn_backend.py => full_attention/full_attention_backend.py} (99%) rename atom/plugin/sglang/attention_backend/{ => full_attention}/radix_attention.py (100%) create mode 100644 atom/plugin/sglang/models/deepseek_mla.py rename atom/plugin/sglang/{attention_backend/sgl_attention_mla.py => models/deepseek_mla_forward.py} (86%) diff --git a/atom/model_ops/__init__.py b/atom/model_ops/__init__.py index 3883dd6944..4670a5cdb2 100644 --- a/atom/model_ops/__init__.py +++ b/atom/model_ops/__init__.py @@ -1,5 +1,7 @@ from .paged_attention import PagedAttention -from atom.plugin.sglang.attention_backend.radix_attention import RadixAttention +from atom.plugin.sglang.attention_backend.full_attention.radix_attention import ( + RadixAttention, +) # This global class is used to construct the attention op in model, # it can be assigned to different attention ops. diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index e4e7cc47ef..ba2bec9687 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -16,7 +16,9 @@ import atom.model_ops as ops from atom.model_ops.paged_attention import PagedAttention from atom.model_ops.attention_mha import PagedAttentionImpl -from atom.plugin.sglang.attention_backend.radix_attention import RadixAttention +from atom.plugin.sglang.attention_backend.full_attention.radix_attention import ( + RadixAttention, +) from atom.utils.forward_context import AttentionMetaData, Context from .backends import AttentionBackend, CommonAttentionBuilder diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 86f06623dc..043f14133a 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -45,7 +45,7 @@ def _register_custom_attention_to_sglang() -> None: from sglang.srt.layers.attention.attention_registry import ( register_attention_backend, ) - from atom.plugin.sglang.attention_backend.sgl_attn_backend import ( + from atom.plugin.sglang.attention_backend.full_attention.full_attention_backend import ( ATOMAttnBackendForSgl, ) diff --git a/atom/plugin/sglang/attention_backend/full_attention/__init__.py b/atom/plugin/sglang/attention_backend/full_attention/__init__.py new file mode 100644 index 0000000000..2b9f6f2726 --- /dev/null +++ b/atom/plugin/sglang/attention_backend/full_attention/__init__.py @@ -0,0 +1,8 @@ +from .radix_attention import RadixAttention +from .full_attention_backend import ATOMAttnBackendForSgl, ForwardMetadata + +__all__ = [ + "RadixAttention", + "ATOMAttnBackendForSgl", + "ForwardMetadata", +] diff --git a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py similarity index 99% rename from atom/plugin/sglang/attention_backend/sgl_attn_backend.py rename to atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py index a8ebc1122e..2b58f6fed5 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attn_backend.py +++ b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py @@ -1,10 +1,10 @@ from __future__ import annotations -# sglang-specific attention backend replacing sglang's built-in AiterAttnBackend. -# Shared by ALL models (DeepSeek, Qwen3, etc.) — handles KV cache writes, -# page-table fixup, pa_persistent_fwd decode path, and MLA prefill kernels. -# Sits at the lowest layer of the attention stack: sglang's RadixAttention -# delegates the actual kernel dispatch here. +# SGLang full-attention backend replacing sglang's built-in AiterAttnBackend. +# Shared by ALL full-attention models (DeepSeek, Qwen3, etc.) — handles KV +# cache writes, page-table fixup, pa_persistent_fwd decode path, and MLA +# prefill kernels. Sits at the lowest layer of the attention stack: +# sglang's RadixAttention delegates the actual kernel dispatch here. # # TODO: rewrite this file once sglang's attention flow is unified into ATOM's # attention layer — KV cache management and attention kernel dispatch will then @@ -47,7 +47,7 @@ except ImportError as e: raise ImportError( "Failed to import 'aiter', which provides AMD-specific attention kernels " - "required by sgl_attn_backend. Please ensure 'aiter' is installed and " + "required by full_attention_backend. Please ensure 'aiter' is installed and " f"available on your AMD system. Original import error: {e}" ) from e diff --git a/atom/plugin/sglang/attention_backend/radix_attention.py b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py similarity index 100% rename from atom/plugin/sglang/attention_backend/radix_attention.py rename to atom/plugin/sglang/attention_backend/full_attention/radix_attention.py diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index 3f2c743b14..dbace2e647 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -386,7 +386,7 @@ def __init__( # Apply ds model-specific sglang patches (attn dispatch, weight hooks, etc.) # TODO: will remove this after sglang supports atom attention backend if self.model_arch_spec.apply_deepseek_patch: - from atom.plugin.sglang.attention_backend.sgl_attention_mla import ( + from atom.plugin.sglang.models.deepseek_mla import ( setup_deepseek_for_sglang, ) diff --git a/atom/plugin/sglang/models/deepseek_mla.py b/atom/plugin/sglang/models/deepseek_mla.py new file mode 100644 index 0000000000..a8d4deb1a1 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_mla.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Model-level DeepSeek MLA patching for SGLang plugin mode. + +This module owns the monkey-patch entrypoints that adapt DeepSeek MLA models to +SGLang plugin mode. The heavy DeepSeek-specific forward and weight helpers live +in `atom.plugin.sglang.models.deepseek_mla_forward`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch + +from atom.plugin.sglang.models.deepseek_mla_forward import ( + forward_sgl_plugin_mode, + init_sgl_attrs, + process_mla_kv_b_proj_after_loading, +) + +if TYPE_CHECKING: + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + +def setup_deepseek_for_sglang(model) -> None: + """Patch a DeepSeek V2/V3 model for SGLang plugin mode.""" + config = model.config + + # Store atom_config for the OOT wrapper before install-time hooks run. + if not hasattr(model, "atom_config"): + from atom.config import get_current_atom_config + + model.atom_config = get_current_atom_config() + + kv_cache_dtype = model.atom_config.kv_cache_dtype + + # Initialise SGLang's MLA TP context before patching per-layer forwards. + from sglang.srt.configs.model_config import is_deepseek_nsa + from sglang.srt.layers.communicator import get_attn_tp_context + + get_attn_tp_context().init_context(config.q_lora_rank, is_deepseek_nsa(config)) + + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + for module in model.modules(): + if isinstance(module, DeepseekV2MLAAttention): + _patch_mla_attention_for_sglang(module, config, kv_cache_dtype) + + +def _patch_mla_attention_for_sglang( + attn: "DeepseekV2MLAAttention", + config: Any, + kv_cache_dtype: str = "bf16", +) -> None: + """Patch one DeepSeek MLA layer for SGLang plugin mode.""" + init_sgl_attrs(attn, config, kv_cache_dtype) + + def patched_forward( + positions: torch.Tensor, + hidden_states: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + from atom.plugin.sglang.models.base_model_wrapper import ( + get_current_forward_batch, + ) + + kwargs["forward_batch"] = get_current_forward_batch() + return forward_sgl_plugin_mode(attn, positions, hidden_states, **kwargs) + + attn.forward = patched_forward + attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading( + attn + ) diff --git a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py b/atom/plugin/sglang/models/deepseek_mla_forward.py similarity index 86% rename from atom/plugin/sglang/attention_backend/sgl_attention_mla.py rename to atom/plugin/sglang/models/deepseek_mla_forward.py index 1dae9349e3..25f1ef79aa 100644 --- a/atom/plugin/sglang/attention_backend/sgl_attention_mla.py +++ b/atom/plugin/sglang/models/deepseek_mla_forward.py @@ -1,17 +1,19 @@ -"""Sglang-specific MLA forward and weight processing for DeepseekV2/V3. +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. -DeepSeek MLA (Multi-Latent Attention) forward logic for sglang plugin mode: +"""Model-specific DeepSeek MLA helpers for SGLang plugin mode. + +DeepSeek MLA (Multi-Latent Attention) forward logic for SGLang plugin mode: absorbed BMM computation, MHA/MLA path dispatch (prefill -> MHA, decode -> MLA), -kv_b_proj weight splitting (w_kc/w_vc), and monkey-patch setup via -setup_deepseek_for_sglang(). +and kv_b_proj weight splitting (w_kc/w_vc). -This module is lazily imported from base_model_wrapper.py only when running in -sglang plugin mode (``is_sglang() == True``). Keeping all sglang-dependent -imports here avoids crashing when sglang is not installed. +This module lives under ``atom.plugin.sglang.models`` because the logic is +DeepSeek-model-specific rather than a generic SGLang attention backend. TODO: rewrite this file once sglang's attention flow is unified into ATOM's -attention layer — the MLA absorbed path and MHA dispatch will then be handled -natively by ATOM's attention ops, making this sglang-specific module unnecessary. +attention layer - the MLA absorbed path and MHA dispatch will then be handled +natively by ATOM's attention ops, making this sglang-specific module +unnecessary. """ from __future__ import annotations @@ -29,7 +31,6 @@ from atom.models.utils import maybe_prefix from atom.models.deepseek_v2 import _fuse_rmsnorm_quant -# sglang imports from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode @@ -59,7 +60,6 @@ from atom.models.deepseek_v2 import DeepseekV2MLAAttention -# bmm_fp8 custom-op wrapper (adapted from sglang forward_mla.py) if _is_cuda: from sgl_kernel import bmm_fp8 as _raw_bmm_fp8 from sglang.srt.utils.custom_op import register_custom_op @@ -90,7 +90,6 @@ def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): raise RuntimeError("bmm_fp8 requires CUDA (sgl_kernel)") -# NamedTuple for prepare → core data flow class SglPrepareResult(NamedTuple): q_pe: torch.Tensor k_pe: torch.Tensor @@ -182,7 +181,6 @@ def _prepare_weight_for_bmm( ) -# Init helpers def init_sgl_attrs( attn: DeepseekV2MLAAttention, config, @@ -216,7 +214,6 @@ def init_sgl_attrs( attn.attn_mha.attn.kv_b_proj = None -# Absorbed batched-matmul (shared by prepare and core) def mla_absorbed_bmm( attn: DeepseekV2MLAAttention, inp: torch.Tensor, @@ -225,12 +222,7 @@ def mla_absorbed_bmm( weight_scale_k: Optional[torch.Tensor], out_dim: int, ) -> torch.Tensor: - """Batched matmul for MLA absorbed weights (w_kc / w_vc). - - Handles deep_gemm, mxfp4, fp8-triton, fp8-cublas, and bf16 fallback paths. - inp: (num_tokens, num_heads, in_dim) — token-major - Returns: (num_tokens, num_heads, out_dim) — token-major - """ + """Batched matmul for MLA absorbed weights (w_kc / w_vc).""" effective_weight_scale = ( weight_scale_k if weight_scale_k is not None else weight_scale ) @@ -299,7 +291,6 @@ def mla_absorbed_bmm( ) return out.transpose(0, 1) - # CUDA fp8 path if weight.dtype == torch.float8_e4m3fn: val, scale = per_tensor_quant_mla_fp8( inp.transpose(0, 1), @@ -308,7 +299,6 @@ def mla_absorbed_bmm( out = bmm_fp8(val, weight, scale, effective_weight_scale, torch.bfloat16) return out.transpose(0, 1) - # bf16 fallback return torch.bmm(inp.transpose(0, 1), weight).transpose(0, 1) @@ -352,14 +342,13 @@ def mla_v_up_proj( ).flatten(1, 2) -# Forward: prepare → core def forward_sgl_prepare( attn: DeepseekV2MLAAttention, positions: torch.Tensor, hidden_states: torch.Tensor, **model_kwargs, ) -> SglPrepareResult: - """Prepare QKV for sglang MLA attention (adapted from sglang forward_absorb_prepare).""" + """Prepare QKV for sglang MLA attention.""" hidden_states_scale = None if isinstance(hidden_states, tuple): hidden_states, hidden_states_scale = hidden_states @@ -401,9 +390,6 @@ def forward_sgl_prepare( k_nope = latent_cache[..., : attn.kv_lora_rank] q_scale = None - # Reuse native ATOM gating for q/k RMSNorm fusion. Quant fusion is used - # when DeepSeek enables qknorm-quant; otherwise keep the non-quant fused - # path aligned with native ATOM before falling back to plain layernorm. if getattr(attn, "fuse_qknorm_quant", False): q, q_scale, q_lora, k_nope = _fuse_qk_rmsnorm_and_q_quant( attn, @@ -413,7 +399,6 @@ def forward_sgl_prepare( ) elif getattr(attn, "fuse_qknorm", False): q, k_nope = _fuse_qk_rmsnorm(attn, q, k_nope) - # Otherwise keep the original overlap path for unfused qk norm. elif attn.alt_stream is not None and get_is_capture_mode(): current_stream = torch.cuda.current_stream() attn.alt_stream.wait_stream(current_stream) @@ -425,11 +410,9 @@ def forward_sgl_prepare( q = attn.q_a_layernorm(q) k_nope = attn.kv_a_layernorm(k_nope) - if attn.use_nsa: - if q_lora is None: - q_lora = q + if attn.use_nsa and q_lora is None: + q_lora = q - # overlap q_b_proj and indexer during decode if ( attn.alt_stream is not None and get_is_capture_mode() @@ -506,7 +489,7 @@ def forward_sgl_core( attn: DeepseekV2MLAAttention, prepared: SglPrepareResult, ) -> torch.Tensor: - """Core MLA attention computation for sglang (adapted from sglang forward_absorb_core).""" + """Core MLA attention computation for sglang.""" save_kv_cache = True if attn.use_fused_qk_rope_concat_and_cache_mla: @@ -545,7 +528,6 @@ def forward_sgl_core( is_neox=attn.rotary_emb.is_neox_style, is_nope_first=True, ) - # Decode/speculative MLA consumes q plus packed MLA cache directly. k = None v = None save_kv_cache = False @@ -571,7 +553,6 @@ def forward_sgl_core( ) attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank) - # up-proj by w_vc attn_bmm_output = mla_v_up_proj( attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim ) @@ -580,15 +561,7 @@ def forward_sgl_core( def _dispatch_sgl_plugin_attn_path(forward_batch) -> str: - """Decide the attention algorithm for this batch based on forward_mode. - - Returns "mha" for extend/prefill (uses standard Q×K×V with flash_attn) - or "mla" for decode (uses absorbed weights + mla_decode_fwd). - - This is the per-batch *routing* decision, distinct from - ``_can_run_sgl_mha_now`` which is a *capability* gate checking whether - the model configuration supports the MHA path at all. - """ + """Decide the attention algorithm for this batch based on forward_mode.""" if forward_batch.forward_mode.is_extend_without_speculative(): return "mha" return "mla" @@ -635,12 +608,7 @@ def _set_mla_kv_buffer_for_mha( def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool: - """Check if the model configuration supports the MHA attention path. - - This is a *capability* gate — NSA models and MXFP4-quantised weights - (uint8) cannot use the MHA path. Distinct from - ``_dispatch_sgl_plugin_attn_path`` which routes each batch. - """ + """Check if the model configuration supports the MHA attention path.""" del forward_batch if attn.use_nsa: return False @@ -819,8 +787,6 @@ def prepare_qkv_latent( hidden_states, hidden_states_scale = hidden_states qkv_lora = attn.fused_qkv_a_proj(hidden_states, hidden_states_scale) - # Fallback: when communicator does not enable input_scattered gather, - # force qkv latent token dimension to align with positions. expected_tokens = 0 if hasattr(forward_batch, "positions") and forward_batch.positions is not None: expected_tokens = int(forward_batch.positions.shape[0]) @@ -843,7 +809,6 @@ def prepare_qkv_latent( return qkv_lora -# Top-level forward entry point def forward_sgl_plugin_mode( attn: DeepseekV2MLAAttention, positions: torch.Tensor, @@ -884,7 +849,6 @@ def forward_sgl_plugin_mode( raise ValueError(f"Unsupported plugin attention path: {attn_path}") -# Weight post-processing: decomposed into sub-functions def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor: """Read kv_b_proj weight, handling AWQ and fnuz dtypes.""" if hasattr(attn.kv_b_proj, "qweight"): @@ -901,8 +865,6 @@ def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor: else: w = attn.kv_b_proj.weight - # On ROCm, ATOM creates parameters with fnuz dtype but loads fn bytes. - # View-cast back to fn so the normalize path works correctly. if _is_fp8_fnuz and w.dtype == torch.float8_e4m3fnuz: w = w.view(torch.float8_e4m3fn) @@ -926,10 +888,7 @@ def _process_fp8_weight( w: torch.Tensor, weight_block_size: Optional[list[int]], ) -> tuple[torch.Tensor, bool, Optional[torch.Tensor]]: - """Process FP8 weights for kv_b_proj. - - Returns (w, use_deep_gemm_bmm, block_scale). - """ + """Process FP8 weights for kv_b_proj.""" from atom.model_ops.utils import normalize_e4m3fn_to_e4m3fnuz from sglang.srt.layers.quantization.fp8_utils import ( block_quant_dequant, @@ -1061,7 +1020,6 @@ def _split_and_assign_kc_vc( [attn.qk_nope_head_dim, attn.v_head_dim], dim=1 ) - # quark fp4 special path quant_method = getattr(attn.kv_b_proj, "quant_method", None) quant_config = getattr(quant_method, "quant_config", None) if ( @@ -1084,8 +1042,6 @@ def _split_and_assign_kc_vc( w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) w_vc = w_vc.contiguous().transpose(1, 2) - # Align bf16 kv_b_proj post-load handling with vLLM: split first, then - # quantize kc/vc independently for the fp8 BMM path. if w.dtype == torch.bfloat16 and (_is_hip or _is_cuda): w_kc, w_scale_k = dynamic_per_batched_tensor_quant(w_kc, dtype=dtypes.fp8) w_vc, w_scale_v = dynamic_per_batched_tensor_quant(w_vc, dtype=dtypes.fp8) @@ -1117,89 +1073,19 @@ def _split_and_assign_kc_vc( def process_mla_kv_b_proj_after_loading(attn: DeepseekV2MLAAttention) -> None: - """Process kv_b_proj weights after loading for sglang MLA mode. - - Orchestrates reading, quantization handling, and splitting of - kv_b_proj into absorbed w_kc / w_vc weights. - """ + """Process kv_b_proj weights after loading for sglang MLA mode.""" w = _read_kv_b_proj_weight(attn) weight_block_size = _get_weight_block_size(attn) use_deep_gemm_bmm = False block_scale = None - # fp8 path if w.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz): w, use_deep_gemm_bmm, block_scale = _process_fp8_weight( attn, w, weight_block_size ) - # int8 path if w.dtype == torch.int8: w = _process_int8_weight(attn, w, weight_block_size) - # split and assign kc/vc _split_and_assign_kc_vc(attn, w, use_deep_gemm_bmm, block_scale, weight_block_size) - - -# One-time model setup (called from base_model_wrapper.py) -def setup_deepseek_for_sglang(model) -> None: - """Patch a DeepseekV2/V3 model for sglang plugin mode. - - - Initialises sglang TP context - - Patches each MLAAttention.forward to dispatch to the sglang MLA path - - Registers process_weights_after_loading hooks - - Stores atom_config on the model - """ - config = model.config - - # Store atom_config (needed by load_weights in the OOT wrapper) - if not hasattr(model, "atom_config"): - from atom.config import get_current_atom_config - - model.atom_config = get_current_atom_config() - - kv_cache_dtype = model.atom_config.kv_cache_dtype - - # Initialise sglang TP context for MLA gather/scatter - from sglang.srt.configs.model_config import is_deepseek_nsa - from sglang.srt.layers.communicator import get_attn_tp_context - - get_attn_tp_context().init_context(config.q_lora_rank, is_deepseek_nsa(config)) - - # Patch each MLAAttention instance - from atom.models.deepseek_v2 import DeepseekV2MLAAttention - - for module in model.modules(): - if isinstance(module, DeepseekV2MLAAttention): - _patch_mla_attention_for_sglang(module, config, kv_cache_dtype) - - -def _patch_mla_attention_for_sglang(attn, config, kv_cache_dtype: str = "bf16") -> None: - """Patch a single DeepseekV2MLAAttention for sglang plugin mode. - - We patch attn.forward (rather than relying solely on ops.Attention = - RadixAttention) because MLA's absorbed-weight forward path replaces the - *entire* forward method — including RoPE, and absorbed - BMM — not just the attention backend. ops.Attention = RadixAttention - handles the backend layer (flash_attn / paged_attn dispatch) and is - already set via set_attn_cls(); this patch sits above that layer. - """ - init_sgl_attrs(attn, config, kv_cache_dtype) - - def patched_forward( - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs, - ) -> torch.Tensor: - from atom.plugin.sglang.models.base_model_wrapper import ( - get_current_forward_batch, - ) - - kwargs["forward_batch"] = get_current_forward_batch() - return forward_sgl_plugin_mode(attn, positions, hidden_states, **kwargs) - - attn.forward = patched_forward - attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading( - attn - ) diff --git a/tests/plugin/test_sglang_model_wrapper.py b/tests/plugin/test_sglang_model_wrapper.py index e4015ed9dc..20d0e07923 100644 --- a/tests/plugin/test_sglang_model_wrapper.py +++ b/tests/plugin/test_sglang_model_wrapper.py @@ -54,8 +54,7 @@ def _make_fake_modules(*, is_last_rank: bool, setup_hook=None) -> dict[str, Modu forward_batch_mod.ForwardBatch = object forward_batch_mod.PPProxyTensors = object - attn_backend_pkg = _package("atom.plugin.sglang.attention_backend") - mla_mod = ModuleType("atom.plugin.sglang.attention_backend.sgl_attention_mla") + mla_mod = ModuleType("atom.plugin.sglang.models.deepseek_mla") mla_mod.setup_deepseek_for_sglang = setup_hook or (lambda model: None) return { @@ -68,8 +67,7 @@ def _make_fake_modules(*, is_last_rank: bool, setup_hook=None) -> dict[str, Modu "sglang.srt.layers.quantization.base_config": quant_base_mod, "sglang.srt.model_executor": model_executor_pkg, "sglang.srt.model_executor.forward_batch_info": forward_batch_mod, - "atom.plugin.sglang.attention_backend": attn_backend_pkg, - "atom.plugin.sglang.attention_backend.sgl_attention_mla": mla_mod, + "atom.plugin.sglang.models.deepseek_mla": mla_mod, } diff --git a/tests/plugin/test_sglang_register.py b/tests/plugin/test_sglang_register.py index 562aaf8bc8..1adb20fb42 100644 --- a/tests/plugin/test_sglang_register.py +++ b/tests/plugin/test_sglang_register.py @@ -324,10 +324,13 @@ def __init__(self, runner): "atom.models.qwen3_moe": ModuleType("atom.models.qwen3_moe"), "atom.models.glm4_moe": ModuleType("atom.models.glm4_moe"), "atom.models.deepseek_v2": ModuleType("atom.models.deepseek_v2"), + "atom.models.minimax_m2": ModuleType("atom.models.minimax_m2"), + "atom.models.qwen3_next": ModuleType("atom.models.qwen3_next"), + "atom.models.qwen3_5": ModuleType("atom.models.qwen3_5"), "atom.config": ModuleType("atom.config"), "atom.plugin.prepare": fake_prepare_mod, - "atom.plugin.sglang.attention_backend.sgl_attn_backend": ModuleType( - "atom.plugin.sglang.attention_backend.sgl_attn_backend" + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend": ModuleType( + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend" ), } fake_modules["atom.models.qwen3"].Qwen3ForCausalLM = type( @@ -342,9 +345,21 @@ def __init__(self, runner): fake_modules["atom.models.deepseek_v2"].DeepseekV3ForCausalLM = type( "DeepseekV3ForCausalLM", (), {} ) + fake_modules["atom.models.minimax_m2"].MiniMaxM2ForCausalLM = type( + "MiniMaxM2ForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_next"].Qwen3NextForCausalLM = type( + "Qwen3NextForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_5"].Qwen3_5ForCausalLM = type( + "Qwen3_5ForCausalLM", (), {} + ) + fake_modules["atom.models.qwen3_5"].Qwen3_5MoeForCausalLM = type( + "Qwen3_5MoeForCausalLM", (), {} + ) fake_modules["atom.config"].Config = type("Config", (), {}) fake_modules[ - "atom.plugin.sglang.attention_backend.sgl_attn_backend" + "atom.plugin.sglang.attention_backend.full_attention.full_attention_backend" ].ATOMAttnBackendForSgl = _FakeBackend with patch.dict(sys.modules, fake_modules): From 35982548b938776f00a212f5984a5b7f46095815 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Tue, 12 May 2026 08:33:02 +0000 Subject: [PATCH 74/76] remove work log --- .../2026-04-30-attn-refractory.md | 354 ------------- ...26-05-07-radix-cache-vs-paged-attention.md | 482 ------------------ ...26-05-08-kv-indices-address-space-notes.md | 359 ------------- .../2026-05-08-metadata-layering-summary.md | 144 ------ ...6-05-11-attn-refractory-session-summary.md | 398 --------------- .../attn_refractory/decoupling-diagram.md | 158 ------ .../vllm-integration-architecture.md | 269 ---------- 7 files changed, 2164 deletions(-) delete mode 100644 work_log/attn_refractory/2026-04-30-attn-refractory.md delete mode 100644 work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md delete mode 100644 work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md delete mode 100644 work_log/attn_refractory/2026-05-08-metadata-layering-summary.md delete mode 100644 work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md delete mode 100644 work_log/attn_refractory/decoupling-diagram.md delete mode 100644 work_log/attn_refractory/vllm-integration-architecture.md diff --git a/work_log/attn_refractory/2026-04-30-attn-refractory.md b/work_log/attn_refractory/2026-04-30-attn-refractory.md deleted file mode 100644 index f0946299a9..0000000000 --- a/work_log/attn_refractory/2026-04-30-attn-refractory.md +++ /dev/null @@ -1,354 +0,0 @@ -# 2026-04-30 ATOM Attention Refactory Session Summary - -## 1. 本次会话的核心目标 - -本次连续讨论的主线是围绕 `ATOM plugin` 中 attention backend 的重构方向展开，重点包括： - -- 调研当前 `ATOM plugin` attention backend 的架构问题 -- 评估 `vLLM plugin` 与 `SGLang plugin` 的耦合程度 -- 分析 SGLang plugin 在支持新模型、特别是非传统 MHA/MLA 类型 backend 时会遇到的困难 -- 讨论如何在 plugin 层把 `vLLM` 与 `SGLang` 解耦 -- 讨论 `SGLang` 中 `hybrid/composite backend` 的接入方式 -- 评估三种模式共享 ATOM attention 实现的难度： - - `ATOM native/server` - - `ATOM vLLM plugin` - - `ATOM SGLang plugin` - -## 2. 对当前 ATOM plugin attention 架构的主要判断 - -### 2.1 当前不是一个真正统一的 attention backend 架构 - -本次调研得到的一个核心结论是： - -- `vLLM plugin` 和 `SGLang plugin` 虽然都叫 “ATOM attention” -- 但它们并不是通过同一套干净的 backend contract 接进来的 -- 更准确地说，是两套宿主接入模型外加多层 patch / adapter / runtime glue 共同构成的系统 - -### 2.2 当前的主要问题不是底层 kernel，而是 runtime / ownership / patch 层次 - -现有问题主要集中在： - -- 全局状态过重 - - `_CURRENT_FRAMEWORK` - - `current_atom_config` - - `ops.Attention` -- vLLM 与 SGLang 的接入模型差异很大，但还共享 `plugin` 根目录里一批伪公共逻辑 -- 很多扩展点依赖 monkey patch 和 runtime override -- SGLang 侧对 DeepSeek MLA 已经上卷到 model-level patch，而不是单纯 backend 替换 - -### 2.3 attention 文件结构的拆分轴不统一 - -本次对 `plugin` 目录下 attention 相关文件的解读认为，目前至少混杂了三层不同对象： - -- host/runtime glue -- backend runtime core -- model-specific specialization - -尤其在 `SGLang` 侧： - -- `radix_attention.py` 更像 adapter -- `sgl_attn_backend.py` 更像 full-attn runtime backend core -- `sgl_attention_mla.py` 更像 DeepSeek MLA specialization - -而在根目录： - -- `attention.py` -- `attention_mha.py` -- `attention_mla.py` -- `attention_mla_sparse.py` - -虽然放在 `plugin/` 根目录，但职责上明显更接近 `vLLM plugin` 的 glue / patch 层，而不是真正共享的 plugin runtime core。 - -## 3. 关于 vLLM plugin 与 SGLang plugin 的耦合 - -### 3.1 结论：耦合度高 - -这次讨论中对二者耦合的判断是： - -- 不是轻量共享几个工具函数 -- 而是共享了一套 runtime state、attention 抽象、selector 和 plugin-mode 判断方式 - -关键耦合点包括： - -- `atom/plugin/prepare.py` -- `atom/plugin/register.py` -- `atom/plugin/config.py` -- `atom.plugin.prepare.is_vllm()/is_sglang()/is_plugin_mode()` -- `ops.Attention` 的全局切换 - -### 3.2 关键判断 - -真正的问题不是 “目录没有拆开”，而是： - -**host runtime 差异没有被限制在 adapter 边界，而是泄漏进了 backend 选择、metadata 组织和全局状态模型。** - -## 4. 关于 plugin 层解耦的共识 - -### 4.1 plugin 层不应该继续保留共享 runtime - -讨论最终倾向于一个更激进但更干净的方向： - -- `atom/plugin/vllm/` 是一个完整子系统 -- `atom/plugin/sglang/` 是另一个完整子系统 -- `atom/plugin/` 根目录不再承担共享 runtime 中心的角色 - -共享只保留给更底层的 `ATOM core`，例如： - -- `model_ops` -- kernels -- loader -- metadata helpers -- model-family specialization 中真正 host-agnostic 的部分 - -### 4.2 “bootstrap” 的定义 - -本次还明确了 `bootstrap` 在本讨论语境中的含义： - -- 不是核心执行逻辑 -- 而是插件被宿主框架加载时的最早期接线/初始化层 - -也就是： - -- 注册扩展点 -- 安装 patch -- 选择 wrapper / adapter / backend -- 建立 host 与 plugin 的连接关系 - -## 5. 已经做过的代码原型 - -本次会话中，为了验证“按 host 拆入口”的思路，已经做了一批原型代码修改： - -### 5.1 新增的 bootstrap / prepare / register 原型 - -- `atom/plugin/sglang/bootstrap.py` -- `atom/plugin/vllm/bootstrap.py` -- `atom/plugin/sglang/prepare.py` -- `atom/plugin/sglang/register.py` - -### 5.2 `prepare.py` 已部分收缩为兼容层 - -- 实际的 SGLang prepare 逻辑已经移动到 `atom/plugin/sglang/prepare.py` -- `atom/plugin/prepare.py` 现在更像一个 legacy shim -- 但它仍保留了 framework-state helper（因为仓库里很多地方还在依赖） - -### 5.3 SGLang register 逻辑已开始下沉 - -SGLang 独有的几块逻辑已经迁入： - -- `register_ops_to_sglang` -- `set_sglang_attn_cls` -- `init_aiter_dist_for_sglang` -- `bootstrap_sglang_runtime` - -共享 `atom/plugin/register.py` 目前更多是兼容 facade。 - -### 5.4 说明 - -这些改动的定位是： - -- 用来显式化未来结构 -- 不是完整重构 -- 还保留了较多兼容层 - -## 6. 关于 SGLang attention backend 的重要判断 - -### 6.1 `ATOMAttnBackendForSgl` 与 public `AiterAttnBackend` - -在一次反复讨论后，本次对它们的关系收敛为： - -- 在 backend 角色位置上，两者基本是 apple-to-apple -- `ATOMAttnBackendForSgl` 不是“更高层的新东西” -- 而是 `SGLang full-attention runtime backend` 的 ATOM 对位实现 - -因此更合适的重构思路不是怀疑它的存在合法性，而是： - -- 保留它作为 full-attn backend core -- 再去拆它内部的 metadata / kv_cache / decode / graph 等职责 - -### 6.2 DeepSeek MLA 的特殊性 - -`sgl_attention_mla.py` 暴露出一个很重要的事实： - -- 对 `MLA`，尤其是 `DeepSeek MLA` -- SGLang plugin 已经不是单纯换 backend -- 而是 patch 了模型级 `forward` - -这说明： - -**MLA 在 SGLang plugin 里已经进入了 model specialization 维度。** - -## 7. 关于 GDN / KDA / Lightning / Mamba2 等非传统 “attention” - -### 7.1 调研结论 - -public SGLang 已经证明，系统里不止 full-attention backend，还存在： - -- linear attention backend - - `GDNAttnBackend` - - `KDAAttnBackend` - - `LightningAttentionBackend` - - `Mamba2AttnBackend` -- hybrid/composite backend - - `HybridLinearAttnBackend` - - `HybridAttnBackend` -- wrapper/composition backend - - `TboAttnBackend` - -### 7.2 对当前 ATOM plugin 设计造成的挑战 - -当前 ATOM plugin 的主抽象仍偏向： - -- MHA -- MLA -- full attention runtime - -这会带来几个结构性问题： - -1. 还没有把 `linear backend` 当成一等公民 -2. 还没有给 `hybrid/composite backend` 预留独立宿主位置 -3. 容易把 algorithm backend、kernel backend、host glue 混成一层 -4. 容易继续把 `GDN` 等路径硬塞回 full-attn backend 体系 - -### 7.3 当前共识 - -后续不应该继续只谈 “attention backend 重构”，而应该升级为： - -**sequence backend / mixer backend 体系重构** - -建议至少在设计上并列三类： - -- `full_attn` -- `linear_attn` -- `hybrid/composite` - -## 8. hybrid/composite backend 在 plugin 侧的接入思路 - -### 8.1 核心判断 - -由于 ATOM SGLang plugin 的启动方式是： - -- 启动 `sglang server` -- 再通过 plugin runtime override 接管 model / attention - -所以要接 `hybrid/composite backend`，不可避免需要 patch 某个 seam。 - -### 8.2 最好的 seam - -本次讨论认为，public SGLang 里最好的 seam 是： - -- `model_runner.py` 中导入并调用的 `attn_backend_wrapper` - -注意： - -- 它不是 `ModelRunner` 的成员方法 -- 而是 `model_runner.py` 模块级导入的符号 - -所以如果要 patch，更稳的是 patch： - -- `sglang.srt.model_executor.model_runner.attn_backend_wrapper` - -而不是仅 patch `attention_registry.attn_backend_wrapper`。 - -### 8.3 最推荐的方式 - -对 hybrid/composite backend 的接入方式，本次最终倾向于： - -**plugin-own 一个薄的 composite wrapper/factory，只在 backend 构造 seam 做单点 patch。** - -不推荐： - -- 深度 monkey patch public hybrid backend 实现 -- 一开始就复制整套 public linear/full backend 逻辑 - -更推荐： - -- full side 先用 ATOM full backend -- linear side 初期先复用 public SGLang 的 `GDNAttnBackend` / `KDAAttnBackend` / `LightningAttentionBackend` / `Mamba2AttnBackend` -- composite wrapper 由 plugin 侧拥有 - -## 9. 关于三种模式共享 ATOM attention 实现的难度 - -本次对： - -- `ATOM native/server` -- `ATOM vLLM plugin` -- `ATOM SGLang plugin` - -三种模式共享 ATOM attention 实现，得到的判断如下。 - -### 9.1 MHA - -难度：`Medium-High` - -原因： - -- `ATOM native` 与 `ATOM vLLM plugin` 已经在 `PagedAttentionImpl` 等层共享较多 -- 真正难的是 `SGLang plugin` 这一侧 -- 它不走 `PagedAttention` 的 ATOM 主路径，而走： - - `RadixAttention` - - `ATOMAttnBackendForSgl` - - `ForwardBatch` - -所以 MHA 共享的主要困难在于： - -**runtime orchestration 差异** - -### 9.2 MLA - -难度：`High` - -原因： - -- native/server 与 vLLM plugin 仍然较多共享 `MLAAttention` -- 但 SGLang plugin 下的 DeepSeek MLA 已经上卷到 model specialization -- 不只是 backend 不同，连 forward 组织方式都不同 - -所以 MLA 的困难在于： - -**runtime orchestration + model specialization 双重叠加** - -## 10. 已产出的文档与图 - -本次会话过程中，已额外产出下列文档/图，用于不同角度说明问题。 - -### 10.1 代码目录内 Markdown - -- `atom/plugin/decoupling-diagram.md` -- `atom/plugin/vllm-integration-architecture.md` -- `atom/plugin/sglang-attention-backend-survey.md` - -### 10.2 Canvas / 可视化分析 - -- `atom-attention-backend-architecture-review.canvas.tsx` -- `atom-plugin-coupling-risk-analysis.canvas.tsx` -- `atom-plugin-decoupling-diagram.canvas.tsx` -- `atom-vllm-plugin-architecture.canvas.tsx` -- `atom-attention-sharing-modes.canvas.tsx` - -### 10.3 这些产物对应的主题 - -- 当前 attention backend 架构缺陷 -- vLLM / SGLang plugin 耦合与风险 -- plugin 解耦方向与 bootstrap 理解 -- vLLM plugin 集成架构 -- public SGLang backend survey -- 三种模式共享 ATOM attention 的难度评估 - -## 11. 当前最值得继续推进的方向 - -如果延续本次会话的结论，后续工作最值得按下面顺序推进： - -1. 继续收缩 `atom/plugin/` 根目录的 shared runtime 语义 -2. 把 `full_attn / linear_attn / hybrid` 三层作为 plugin 后续结构设计的一等公民 -3. 在 `SGLang plugin` 侧先补出 `hybrid/composite` 组合层 -4. 初期复用 public SGLang linear backend，优先验证 runtime 结构是否成立 -5. 然后再评估哪些 linear backend 需要逐步替换成 ATOM 自己的实现 -6. 对 `MLA` 尽早拆开： - - generic MLA runtime - - DeepSeek specialization - -## 12. 一句话总结 - -本次会话最终把问题收敛为： - -**当前 ATOM plugin attention 的关键矛盾，不是底层 kernel 能不能共享，而是 vLLM / SGLang / native 三种模式在 runtime、metadata、model specialization 和 backend ownership 上没有对齐；后续重构应从“统一 attention backend”升级为“重新定义 plugin 的 sequence backend / host-owned runtime 结构”。** diff --git a/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md b/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md deleted file mode 100644 index f0f9d5bccc..0000000000 --- a/work_log/attn_refractory/2026-05-07-radix-cache-vs-paged-attention.md +++ /dev/null @@ -1,482 +0,0 @@ -# 2026-05-07 Radix Cache vs Paged Attention Notes - -> 预估阅读时间：15 分钟 -> 主题：梳理 `SGLang` 中 `radix cache` / `RadixAttention` 与 `ATOM` / `vLLM` 常见的 `paged attention` 之间的关系，解释为什么它们最终可以落到相同的 kernel，以及这件事对 `sglang plugin` attention backend 重构意味着什么。 - -## 1. 这篇笔记想回答什么 - -围绕 `sglang plugin attn backend` 的重构，最近反复出现了几个问题： - -1. `SGLang` 用的是 `RadixAttention`，`ATOM`/`vLLM plugin` 里常见的是 `PagedAttention`，两者是不是两套完全不同的注意力体系？ -2. 如果它们真的不同，为什么在 `DeepSeek MLA` 这类路径里，最终又会调用到相同的底层 kernel？ -3. 既然 `SGLang` 还多了一层 `radix cache` 前缀管理，为什么没有一个非常直观的结论说 "`SGLang` 一定比 `vLLM` 更快 / 更省显存"？ -4. 对当前重构来说，`radix` / `paged` 的差异到底应该被归类到： - - host runtime 差异 - - KV cache 管理差异 - - metadata lowering 差异 - - kernel 差异 - 的哪一层？ - -本文尝试把这几个问题放进同一个分析框架里。 - -## 2. 先给结论 - -### 2.1 最核心的一句话 - -`RadixAttention` 和 `PagedAttention` 主要不是在底层注意力数学上不同，而是在 **runtime 如何管理、共享、命中和定位历史 KV** 上不同。 -一旦 runtime 把历史 KV lowering 成底层 kernel 能吃的 `page_table`、`kv_indptr`、`kv_indices`、`qo_indptr` 之类的 metadata，kernel 就不再关心这些索引最初是来自 radix tree 还是来自 paged block table。 - -### 2.2 另一句更工程化的结论 - -对重构来说，不应该把 “`RadixAttention` vs `PagedAttention`” 当作 “能不能共享底层执行实现” 的直接判断依据。 -更应该把系统拆成三层： - -1. **host/runtime 层** - 例如 `ForwardBatch`、`RadixAttention`、`ModelRunner`、prefix cache、speculative 调度 -2. **execution metadata lowering 层** - 例如 `page_table`、`block_tables`、`kv_indptr`、`kv_indices`、`qo_indptr` -3. **kernel / execution core 层** - 例如 `pa_fwd_asm`、`pa_persistent_fwd`、`flash_attn_varlen_func`、`mla_decode_fwd`、`mla_prefill_fwd` - -`radix` 和 `paged` 的差异，主要在第 1、2 层； -兼容性和共享机会，主要出现在第 2、3 层。 - -## 3. 为什么名字容易让人误解 - -当前最容易造成误解的点是： - -- `PagedAttention` 这个名字听起来像“最终执行算法” -- `RadixAttention` 这个名字也听起来像“最终执行算法” - -但在工程上，它们都更接近 **host-facing attention runtime abstraction**，而不是最终 kernel 名字。 - -更具体一点： - -- `PagedAttention` 更强调 **按固定 page/block 管理 KV cache** -- `RadixAttention` 更强调 **按 prefix/radix tree 管理请求历史与前缀复用** - -这两者都不是在说： - -- `QK^T` 怎么算 -- softmax 怎么做 -- decode kernel 用哪套 asm/triton/flashinfer - -这些底层执行问题，是更下一层的 backend / kernel 决定的。 - -## 4. `RadixAttention` 到底是什么 - -### 4.1 `RadixAttention` 不是 kernel - -在 public `SGLang` 里，`RadixAttention` 本身并不直接定义底层 kernel 选择。 -它更像是： - -- SGLang 模型层里统一使用的 attention layer 壳 -- 它知道自己在 `ForwardBatch` 语义里运行 -- 它把真正的执行委托给 `forward_batch.attn_backend` - -因此，`RadixAttention` 更像一个 **layer/runtime 入口点**。 - -### 4.2 `radix` 的重点是 prefix 管理 - -`radix cache` 的本质，是一棵 prefix tree，用来做： - -- 前缀命中 -- 前缀复用 -- unfinished / finished request 的缓存插入 -- 基于 prefix 的 eviction / lifecycle 管理 - -它回答的问题更像： - -- 这个请求的历史前缀已经缓存到哪里了？ -- 哪一部分可以直接复用，不必重新 prefill？ -- 哪些 KV slot/page 仍然属于共享前缀？ -- 哪些历史片段可以回收？ - -所以 `radix` 优化的主要是 **prefix reuse**，而不是单次 attention kernel 的算力效率。 - -### 4.3 `radix cache` 在 SGLang 里不是完全“反 page”的 - -这一点很重要。public `SGLang` 的 `radix_cache` 本身就已经带有 page-aware 的逻辑。 - -例如它有 page 对齐： - -- `page_align_keys()` - -也有按 page 粒度匹配 prefix： - -- `_key_match_paged()` - -这意味着： - -- `radix cache` 解决的是前缀共享与命中问题 -- 但它并不排斥最终按 page/block 粒度来表达缓存 - -也就是说，**radix tree 的逻辑管理** 与 **paged/block 的物理表达** 并不冲突。 - -## 5. `PagedAttention` 到底是什么 - -### 5.1 `PagedAttention` 的重点不是 prefix tree - -`PagedAttention` 更强调： - -- KV cache 被组织成固定大小的 page/block -- 每个请求通过 `page_table` / `block_table` 找到自己历史对应的 page -- kernel 依照这些 page/block 索引读取历史 KV - -它回答的问题更像： - -- 当前请求历史有哪些 block/page？ -- 它们在物理内存中的位置是什么？ -- 这些 page 如何映射到 kernel 所需的 block table？ - -所以 `paged attention` 更偏： - -- **物理存储布局** -- **kernel 读取 KV 的地址组织方式** - -### 5.2 `PagedAttention` 天生不等于 prefix reuse - -`paged attention` 当然也可以配合 prefix cache 使用，但 prefix reuse 并不是它名字里最核心的概念。 -它的核心价值更在于： - -- 支持非连续物理布局 -- 让 decode kernel 可以按 page/block 高效读取 KV - -## 6. 为什么 `radix` 和 `paged` 最后可以兼容 - -### 6.1 因为它们主要解决的是不同层的问题 - -可以把它们看成： - -- `radix`: 逻辑历史管理器 -- `paged`: 物理页表表达方式 - -两者关系有点像： - -- `radix` 负责决定“哪些历史是共享的、哪些是命中的、请求当前逻辑历史是什么” -- `paged` 负责决定“这些历史最终以什么 page/block 索引形式交给 kernel” - -只要 runtime 最终能把逻辑历史翻译成 page/block 或 indptr/index 形式，底层 kernel 并不需要知道前缀最初是如何组织的。 - -### 6.2 兼容发生在 lowering 之后 - -这一点非常关键： - -上层看起来是： - -- `ForwardBatch` -- `RadixAttention` -- `req_to_token` -- prefix cache - -但 lower 到 backend 之后，会变成执行态 metadata，例如： - -- `page_table` -- `block_tables` -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `kv_last_page_len` - -到了这个层次，kernel 只看： - -- query 的 shape -- KV cache 的 shape -- 历史长度 -- 索引表 / indptr - -它不看： - -- 这套索引是不是来自 radix tree -- 是不是来自某个 prefix cache 命中 -- 是不是通过 `PagedAttention` 对象构造出来的 - -所以兼容性不是“名字兼容”，而是 **execution metadata 兼容**。 - -## 7. public `SGLang` 自己就在做类似 lowering - -这不是 `ATOM plugin` 独有的现象。public `SGLang` 本身就在这么做，只是不同 backend lower 到的 metadata 形态略有不同。 - -### 7.1 `aiter_backend` - -public `SGLang` 的 `aiter_backend` 会从： - -- `forward_batch.seq_lens` -- `forward_batch.req_pool_indices` -- `req_to_token` - -构造： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` - -也就是说，它会把 host runtime 的历史表示 lower 成 AITER kernel 所需的执行态索引。 - -### 7.2 `triton_backend` - -public `SGLang` 的 `triton_backend` 也做类似事情。 -它同样维护： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` - -然后再调用 Triton 的 decode/extend kernel。 - -### 7.3 `trtllm_mha_backend` / 部分 `flashinfer` 路径 - -这类 backend 更偏向： - -- 直接构造 `page_table` -- 或 `block_tables` - -再喂给 flashinfer / TRTLLM 风格的 kernel。 - -所以更准确地说，public `SGLang` 的共同模式不是： - -> “所有 backend 都把 radix runtime 翻成 paged attention” - -而是： - -> “所有 backend 都会把 radix/runtime 层的历史表示，lower 成各自 kernel 能吃的索引/页表 metadata” - -其中有的偏 `kv_indptr/kv_indices`，有的偏 `page_table/block_tables`。 - -## 8. 从 kernel 角度看，为什么可以兼容 - -### 8.1 kernel 真正关心什么 - -绝大多数 attention kernel 真正关心的是： - -1. 当前 query 的排列方式 -2. 每个请求可见的历史 KV 长度 -3. 如何从某个索引结构里拿到历史 KV 的物理地址 -4. page/block 大小、stride、layout -5. 是否需要 prefix/extend/speculative 的特殊 metadata - -它通常不关心： - -1. 这份索引最初是不是从 prefix tree 查出来的 -2. 命中前缀的逻辑是谁做的 -3. runtime 是 `RadixAttention` 还是 `PagedAttention` - -### 8.2 对 MHA kernel 的影响 - -MHA decode 常见地会落到： - -- `pa_fwd_asm` -- `pa_persistent_fwd` -- `flash_attn_varlen_func` - -这些 kernel 更偏 page/block 或 varlen index 风格。 -只要 runtime 最终给出： - -- `page_table` -- `context_lens` -- `block_tables` - -或者： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` - -它们都能算。 - -### 8.3 对 MLA kernel 的影响 - -`DeepSeek MLA` 更容易看出这种兼容性，因为 MLA 的底层算子路径更明确。 - -典型 kernel 包括： - -- `mla_decode_fwd` -- `mla_prefill_fwd` -- `concat_and_cache_mla` -- `fused_qk_rope_concat_and_cache_mla` - -这些 kernel 只要求： - -- 压缩后的 latent KV 表示 -- 以及相应的 `qo_indptr` / `kv_indptr` / `kv_indices` - -所以无论上层 host 是： - -- `ATOM native` -- `ATOM vLLM plugin` -- `ATOM SGLang plugin` - -只要最终 lower 成相同的 MLA execution metadata，就自然会落到同一套 MLA kernel。 - -## 9. `DeepSeek MLA` 为什么尤其容易“看起来收敛” - -### 9.1 因为 MLA 的 execution core 更专用 - -`DeepSeek MLA` 的真正执行对象不是： - -- “`RadixAttention` 版本的 MLA” -- “`PagedAttention` 版本的 MLA” - -而是： - -- 一组固定的 MLA latent/cache 表示 -- 一套固定的 rope + kv cache 写入方式 -- 一套固定的 prefill/decode kernel - -换句话说，MLA 的底层 core 更容易抽成“共享执行层”。 - -### 9.2 因为 host 差异主要体现在 metadata 准备与调度 - -对于 `DeepSeek MLA` 来说，上层 host 差异更多体现在： - -- `positions` 从哪里来 -- `forward_batch` 怎么组织 -- decode / extend / speculative / target_verify 怎么分流 -- `req_to_token` 怎么转换成 `kv_indices` - -一旦进入 `mla_decode_fwd` 或 `mla_prefill_fwd` 前，这些差异已经被消化掉了。 - -所以从外部观察，很容易得到一种印象： - -> “怎么上面完全不一样，下面却是同一套 kernel？” - -其实并不奇怪，因为两边只是上层 runtime 不同，而 MLA execution core 本来就在下层收敛。 - -## 10. 既然兼容，为什么 `SGLang` 没有天然更快或更省显存 - -这是另外一个很容易误判的问题。 - -### 10.1 `radix cache` 省的是前缀重复，不是所有 KV - -`radix cache` 能节省的，是重复前缀带来的重复 prefill 和重复 KV 占用。 - -它主要帮助的是： - -- 大量请求共享长前缀 -- 同 prompt 多分支采样 -- prefix-heavy workload -- 某些 speculative / chunked prefill 场景 - -如果 workload 不满足这些条件，它就不会天然表现为显著优势。 - -### 10.2 它优化的是 prefix reuse，不是单 token decode kernel - -radix 管理层并不会让 `mla_decode_fwd`、`pa_fwd_asm` 这种 kernel 自己更便宜。 -kernel 的 raw 速度主要还是受： - -- 头数 -- head dim -- KV layout -- quant -- page size -- GPU kernel 实现 - -这些因素决定。 - -所以常见现象是： - -- prefix-heavy workload 下，`SGLang` 的收益更明显 -- decode-heavy workload 下，收益没那么明显 -- 如果 `vLLM` 也有自己的 prefix caching/block caching，差距会进一步缩小 - -### 10.3 radix 自己也有管理成本 - -`radix cache` 不是免费层。它也有： - -- prefix match -- tree split / insert / eviction -- request -> KV slot 映射维护 -- speculative / unfinished request 的额外 bookkeeping - -如果命中率不高，这部分成本并不会自动转化成收益。 - -## 11. 对当前重构最有价值的判断 - -### 11.1 不要把 `RadixAttention` vs `PagedAttention` 当成是否共享 execution core 的边界 - -因为从 public `SGLang` 和当前 `ATOM plugin` 的代码路径看，二者最终都在做类似事情： - -- 从 host runtime 的历史表示出发 -- 构造底层 kernel 所需的 execution metadata - -所以： - -- runtime 不同，不代表 execution core 必须不同 -- 名字不同，不代表 kernel 层必须分叉 - -### 11.2 真正的边界更适合这样划 - -#### A. `sglang plugin runtime` 应拥有的部分 - -- `ForwardBatch` -- `RadixAttention` -- prefix cache / radix cache -- `ModelRunner` / `attn_backend_wrapper` 相关调度 -- speculative / graph / verify / chunked prefill 等宿主语义 - -#### B. execution metadata lowering 层 - -这层是最值得重新定义的边界。它负责： - -- 把 `req_to_token`、prefix 命中结果、cache state -- 转成 `page_table` / `kv_indptr` / `kv_indices` / `qo_indptr` - -这层既带有 host 语义，又已经开始接近共享执行核心，是未来最值得梳理清楚的一层。 - -#### C. shared execution core - -例如： - -- `mla_decode_fwd` -- `mla_prefill_fwd` -- `pa_fwd_asm` -- `pa_persistent_fwd` -- `flash_attn_varlen_func` - -以及与这些 kernel 紧耦合的那部分 ATOM core helper。 - -### 11.3 这对 `sglang plugin` 的重构意味着什么 - -当前 `sglang plugin` 最不该做的是： - -- 因为底层 kernel 能共享，就把 `sglang` 的 runtime seam 也强行对齐到 `vllM` -- 或者因为上层叫 `RadixAttention`，就认为它和 `PagedAttention` 完全不能共享底层执行实现 - -更合理的重构方向是： - -1. 保持 `sglang` 的 host/runtime 语义完整 -2. 识别 runtime lowering 层的清晰边界 -3. 尽量把 execution core 继续下沉回共享 `ATOM core` - -## 12. 一种对外更容易讲清楚的表述 - -如果后续要向别人解释，可以用下面这几句话： - -### 12.1 关于 `radix` vs `paged` - -`radix cache` 主要解决的是 prefix 命中与复用，`paged attention` 主要解决的是 page/block 形式的 KV 组织与读取。 -前者偏逻辑管理，后者偏物理表达；两者不是同一层次的问题。 - -### 12.2 关于为什么最终会调到同一个 kernel - -因为 attention kernel 只关心最终的执行态 metadata，比如 `page_table`、`kv_indptr`、`kv_indices`、`qo_indptr`。 -一旦 runtime 把历史 KV lowering 成这类形式，kernel 就不再关心这些索引最初是由 radix tree 还是 paged runtime 生成的。 - -### 12.3 关于为什么 `SGLang` 没有天然显著更快 - -因为 `radix` 带来的主要收益是 prefix reuse，而不是单 token decode kernel 的 raw speed。 -如果 workload 不是 prefix-heavy，或者 decode 占主要成本，那么 radix 带来的优势不会自然放大成“总是更快 / 更省显存”。 - -## 13. 最终总结 - -把这次探索收敛成一句话： - -**`RadixAttention` 与 `PagedAttention` 的主要差异，不在底层注意力数学，也不一定在最终 KV 的物理表示本身，而在 host runtime 如何管理、共享、命中并 lowering 历史 KV；一旦 lowering 完成，底层 kernel 完全可以是共享的。** - -对当前 `sglang plugin attn backend` 的重构来说，真正值得做的不是争论 “`radix` 和 `paged` 能不能统一成一个名字”，而是把系统拆清楚： - -- 哪些是 `sglang` 必须自有的 runtime / prefix cache / scheduling 语义 -- 哪些是 execution metadata lowering -- 哪些已经是可以共享的 `ATOM` execution core - -这比直接讨论 “`radix` 和 `paged` 谁更先进、谁应该替代谁” 更接近真正的工程问题。 diff --git a/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md b/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md deleted file mode 100644 index e5a63be0f1..0000000000 --- a/work_log/attn_refractory/2026-05-08-kv-indices-address-space-notes.md +++ /dev/null @@ -1,359 +0,0 @@ -# 2026-05-08 KV Indices and Address Space Notes - -> 预估阅读时间：8-10 分钟 -> 主题：解释 `kv_indptr` / `kv_indices` 为什么不是纯 host metadata，也不是完全 storage-agnostic 的抽象，并用 `sglang` / `ATOM` / `vllm plugin` 的实际 KV cache 存储举例说明。 - -## 1. 想回答的问题 - -最近有一个很容易卡住重构讨论的问题： - -> `kv_indptr` / `kv_indices` 看起来和 KV cache 的存储方式深度绑定。 -> 既然如此，为什么还能把它们看成 `sglang` / `ATOM` / `vllm plugin` 之间可以共享的 metadata？ - -这个问题的关键在于区分两件事： - -1. **host/runtime 是否相同** -2. **kernel 最终消费的地址空间模型是否相同** - -一句话先说结论： - -**`kv_indptr` / `kv_indices` 不是完全 storage-agnostic 的抽象，它们绑定的是某种“可索引的 KV 地址空间”； -但它们可以是 host/runtime-agnostic 的，只要不同宿主最终都能把自己的 runtime state lowering 到同一种地址空间模型。** - -## 2. 三层视角 - -理解这个问题，最好先把 attention 相关代码拆成三层： - -### 2.1 host/runtime 层 - -这一层描述的是宿主框架如何组织请求和调度，例如： - -- `ATOM native` 的 `ScheduledBatch` -- `vllm plugin` 的 `query_start_loc` / `num_prefills` -- `sglang plugin` 的 `ForwardBatch` / `forward_mode` / `spec_info` - -这一层表达的是： - -- 哪些请求在 batch 中 -- 哪些是 decode，哪些是 prefill -- prefix / speculative / graph 的语义是什么 - -### 2.2 execution metadata lowering 层 - -这一层做的事情是： - -- 把 host/runtime 里的 batch 状态 -- 翻译成 kernel 真正能消费的索引结构 - -典型字段： - -- `block_tables` / `page_table` -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `kv_last_page_len` - -### 2.3 kernel / execution core 层 - -这一层只关心： - -- query tensor -- KV cache tensor -- index / indptr / page table -- persistent kernel workspace - -例如： - -- `mla_decode_fwd` -- `mla_prefill_fwd` -- `pa_fwd_asm` -- `pa_persistent_fwd` - -## 3. `sglang` 的实际 KV cache 存储 - -在 `sglang` 里，KV cache 本身不是存在 radix tree 里。 -radix tree 负责的是 prefix 复用与命中；真正的 KV 还是存在 memory pool 里。 - -`sglang` 自己的注释说得很清楚： - -- `ReqToTokenPool` 负责 request -> token location 映射 -- `TokenToKVPoolAllocator` 负责 KV cache index 管理 -- `KVCache` 真正持有物理 K/V tensor - -也就是说，`sglang` 的 runtime 世界里至少有两层： - -1. **逻辑视角**：某个 request 的第 `i` 个 token 属于哪里 -2. **物理视角**：这个 token 对应的 K/V 存在物理 pool 的哪个 slot - -### 3.1 `ReqToTokenPool` - -`ReqToTokenPool` 本质上是一张二维表： - -```text -req_to_token[req_id, token_pos] = physical_slot_id -``` - -它回答的是： - -- 某个 request 的第 `j` 个逻辑 token -- 对应到物理 KV pool 里的哪个 slot - -### 3.2 物理 KV pool - -物理 pool 里，K 和 V 都是按 slot 存的。 -写入 KV 时，最终就是按 `loc/indices` 写入： - -```text -k_cache[indices] = k -v_cache[indices] = v -``` - -所以从 kernel 角度看，最终它看到的是： - -- 一个可以索引的 KV 地址空间 -- 一组 slot index - -而不是看到“radix tree”本身。 - -## 4. `ATOM` / `vllm plugin` 的实际 KV cache 存储 - -`ATOM` / `vllm plugin` 更偏 paged/block 风格。 - -典型思路是： - -- 每个 request 有一张 `block_table` -- 每个 block 对应一个固定大小的 page -- kernel 按 `block_table` 或其展开后的 index 去读 KV - -所以它的上层直觉更像： - -```text -request -> block table -> page/block address -> physical KV -``` - -和 `sglang` 相比，最大的不同不是“有没有物理地址空间”，而是： - -- `sglang` 上层先经过 `req_to_token` -- `vllm/ATOM` 上层先经过 `block_table` - -但两边最终都能 lower 到一组可供 kernel 读取的地址索引。 - -## 5. 例子一：MLA decode 中的 `kv_indptr` / `kv_indices` - -这个例子里，可以把 `kv_indices` 理解为 **token-slot index**。 - -### 5.1 在 `sglang` 里 - -假设一个 request 当前长度是 10， -它在物理 KV pool 里的 slot 顺序不是连续的，而是： - -```text -[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] -``` - -那 `req_to_token[req_id, :10]` 就大致表示： - -```text -logical token position: 0 1 2 3 4 5 6 7 8 9 -physical slot id: 8 9 10 11 0 1 2 3 4 5 -``` - -如果这个 batch 只有这一个 request，那么 lowering 之后可以得到： - -```text -kv_indptr = [0, 10] -kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5] -``` - -这里： - -- `kv_indptr` 表示“第 0 个 request 的 KV 范围是 `[0, 10)`” -- `kv_indices` 表示“真正去物理 KV pool 里读这些 slot” - -### 5.2 在 `vllm/ATOM` 里 - -如果 page size 是 4，同样 10 个 token 对应的 block layout 可以写成： - -```text -block 2 -> slots [8, 9, 10, 11] -block 0 -> slots [0, 1, 2, 3] -block 1 -> slots [4, 5, 6, 7] -``` - -block table 大致就是： - -```text -[2, 0, 1] -``` - -展开前 10 个 token 后，对应的 slot 顺序仍然是： - -```text -[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] -``` - -于是 lower 之后，给 MLA decode kernel 的： - -```text -kv_indptr = [0, 10] -kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5] -``` - -和 `sglang` 这边是等价的。 - -### 5.3 这个例子说明什么 - -说明在 MLA decode 这个路径里： - -- `kv_indices` 绑定的是 **token-slot 地址空间** -- 它不是“完全与存储无关” -- 但它已经不关心 slot 列表最初是由 `req_to_token` 还是 `block_table` 推出来的 - -也就是说，**它和具体宿主无关，但和被选定的地址空间模型有关。** - -## 6. 例子二：MHA persistent decode 中的 `page_table` / `kv_indices` - -这个例子里，`kv_indices` 更接近 **page/block index**，而不是 token-slot index。 - -### 6.1 在 `sglang plugin` 里 - -仍然假设 page size = 4。 -如果某个 request 的 token-slot 顺序是： - -```text -[8, 9, 10, 11, 0, 1, 2, 3, 4, 5] -``` - -那么每页的起始 slot 是： - -```text -[8, 0, 4] -``` - -除以 page size 后，对应 page id： - -```text -[2, 0, 1] -``` - -在 `sglang plugin` 里，这就是 `page_table` 的来源： -从 `req_to_token` 按页采样，再除以 `page_size`。 - -然后 `_build_pa_metadata_for_decode()` 再把它变成 `pa_persistent_fwd` 所需的 page-level `kv_indices`。 - -### 6.2 在 `ATOM` / `vllm` 里 - -这一层本来就是 paged/block 视角，所以 `block_table` 原生就已经是： - -```text -[2, 0, 1] -``` - -也就是说： - -- `sglang` 是 `req_to_token -> page_table` -- `vllm/ATOM` 是直接 `block_table -> page_table` - -但到了 persistent kernel 眼里，最后看到的是同一种 page/block 地址空间。 - -### 6.3 这个例子说明什么 - -说明在 MHA persistent decode 这个路径里： - -- `kv_indices` 绑定的是 **page/block 地址空间** -- 它仍然不是“完全 storage-agnostic” -- 但也不需要知道上层 host 是 `sglang` 还是 `vllm` - -只要两边都能 lower 到同一个 page/block 地址空间，kernel 就能共享。 - -## 7. 这两个例子真正说明的事 - -上面两个例子其实在说明同一个结论： - -### 7.1 `kv_indptr` / `kv_indices` 不是 host metadata - -它们不描述： - -- `ForwardBatch.forward_mode` -- `query_start_loc` -- `num_prefills` -- `run_graph` -- prefix hit 的高层语义 - -它们只描述： - -- 当前这次 kernel 调用 -- 要读哪些 KV 地址 -- 每个 request 的边界在哪里 - -所以它们更接近： - -- execution metadata -- kernel-facing metadata - -而不是 host/runtime metadata。 - -### 7.2 但它们也不是完全 storage-agnostic - -它们依赖： - -- 当前 KV cache 的地址空间模型 -- kernel 对该地址空间的解释方式 - -如果底层 cache 不是： - -- 可线性索引的 token slot -- 或可 flatten 的 page/block index - -那 `kv_indptr` / `kv_indices` 这套抽象就未必成立。 - -所以它们不是“无限通用”的。 - -## 8. 对重构的直接启示 - -这个背景知识对当前 `sglang attention` 重构最有价值的结论是： - -### 8.1 不要去统一 host metadata - -比如不要强行统一： - -- `ForwardBatch` -- `query_start_loc` -- `num_prefills` -- `run_graph` - -这些都属于宿主框架自己的 runtime 语言。 - -### 8.2 应该争取统一 execution metadata - -例如： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `page_table` -- `kv_last_page_len` -- `work_indptr` -- `reduce_indptr` - -以及围绕这些字段构造的更小 dataclass，例如当前 draft 里已经开始尝试的： - -- `MLADecodeKernelMetadata` -- `MHAAsmKernelMetadata` -- `MHAPersistentKernelMetadata` - -### 8.3 正确的共享边界 - -所以真正可共享的不是： - -- `sglang` 的 `ForwardMetadata` -- `vllm plugin` 那一整套 metadata family - -而是： - -- 它们再往下切出来的 -- kernel-facing execution metadata - -## 9. 一句话总结 - -**`kv_indptr` / `kv_indices` 不是“和存储彻底无关”的抽象，它们绑定的是某种被 kernel 接受的 KV 地址空间；但正因为它们已经脱离了宿主 batch/runtime 语义，`sglang` 和 `vllm/ATOM` 即使上层完全不同，也仍然可以在这层 metadata 上收敛并共享 kernel 调用逻辑。** diff --git a/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md b/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md deleted file mode 100644 index 86bccce8e6..0000000000 --- a/work_log/attn_refractory/2026-05-08-metadata-layering-summary.md +++ /dev/null @@ -1,144 +0,0 @@ -# 2026-05-08 Metadata Layering Summary - -## 1. 目的 - -这份笔记简要总结 `ATOM native`、`vllm plugin`、`sglang plugin` 三边 metadata 的关系，重点回答： - -- `ATOM/atom/utils/forward_context.py` 中的 `AttentionMetaData` 是什么 -- `ATOM/atom/plugin/attention.py` 里的 metadata dataclass 在做什么 -- `ATOM/atom/plugin/sglang/attention_backend/sgl_attn_backend.py` 里的 `ForwardMetadata` 属于哪一层 -- 哪些 metadata 值得共享，哪些不值得强行统一 - -## 2. 三套 metadata 的定位 - -### 2.1 `AttentionMetaData` - -文件：`ATOM/atom/utils/forward_context.py` - -它是 `ATOM` attention 执行链路里的**通用外层容器**。 -里面既有： - -- 通用 attention 字段 - - `slot_mapping` - - `block_tables` - - `kv_indptr` - - `kv_indices` - - `work_meta_data` -- 也预留了 plugin 扩展入口 - - `plugin_metadata` - -所以它更像一个大容器，而不是某个 host 专属的数据模型。 - -### 2.2 `vllm plugin metadata` - -文件：`ATOM/atom/plugin/attention.py` - -`vllm plugin` 不是重写一套完全独立的 metadata，而是在 `AttentionMetaData` 外壳上，额外挂了一层 `plugin_metadata` payload。 - -代表类型包括： - -- `AiterFlashAttentionMetadataForPluginMode` -- `AiterMLACommonMetadataForPluginMode` -- `AiterMLADecodeMetadataForPluginMode` -- `AiterMLASparseMetadataForPluginMode` - -这些 dataclass 主要承载： - -- `vllm` 的 batch/runtime 语义 -- decode/prefill/extend phase 拆分 -- chunked prefill / DCP / sparse MLA 等 feature-specific 信息 - -一句话：**`vllm plugin` 是“外层复用 `AttentionMetaData`，内层自带一套 host-specific metadata family”。** - -### 2.3 `sglang plugin` 的 `ForwardMetadata` - -文件：`ATOM/atom/plugin/sglang/attention_backend/sgl_attn_backend.py` - -`ForwardMetadata` 不是 `AttentionMetaData.plugin_metadata` 那种 payload，而是 `sglang backend` 内部的**lowering 结果缓存对象**。 - -它更接近： - -- `ForwardBatch` -- `req_to_token` -- `token_to_kv_pool` -- graph/speculative path - -向 kernel metadata 的过渡层。 - -一句话：**`ForwardMetadata` 更像 backend-local lowering state，而不是 host-agnostic 的最终执行 metadata。** - -## 3. 三类字段 - -### 3.1 host/runtime-bound - -这类字段描述宿主框架如何组织 batch/runtime，而不是 kernel 如何计算。 - -例子： - -- `vllm plugin` - - `query_start_loc` - - `num_decodes` - - `num_prefills` - - `chunked_context` -- `sglang plugin` - - `run_graph` - - 各种 `ForwardBatch.forward_mode` 派生分支语义 - -### 3.2 execution/kernel-bound - -这类字段最接近 kernel 真正需要的执行信息。 - -例子： - -- `slot_mapping` -- `block_tables` -- `context_lens` -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `kv_last_page_len` -- `work_indptr` -- `reduce_indptr` - -### 3.3 hardware-sensitive - -这类字段和 page size、dtype、kernel 选择、并行配置相关。 - -例子： - -- `max_q_len` -- `max_kv_len` -- `num_kv_splits` -- `attn_out_dtype` -- `head_dim` - -## 4. 哪些值得共享 - -### 不建议强行共享的 - -- `AttentionMetaData` 整体 -- `vllm plugin` 整套 `plugin_metadata` -- `sglang plugin` 整体的 `ForwardMetadata` - -原因不是它们没价值，而是它们都带着很重的 host/runtime 语义。 - -### 值得重点共享的 - -应该共享的是更小的 **execution metadata view**，例如当前 draft 已经开始尝试的： - -- `MLADecodeKernelMetadata` -- `MHAAsmKernelMetadata` -- `MHAPersistentKernelMetadata` - -这些结构只保留某个 kernel call 真正需要的字段，更适合在： - -- `ATOM core` -- `vllm plugin` -- `sglang plugin` - -之间共享。 - -## 5. 一句话结论 - -现在不是三边共用“一套 metadata”，而是三边各自都有一层 host-owned metadata / lowering 体系。 -真正值得共享的，不是这些大 metadata 本身，而是从它们里面切出来的、更小的、kernel-facing execution metadata。 diff --git a/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md b/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md deleted file mode 100644 index 3b2128f5bd..0000000000 --- a/work_log/attn_refractory/2026-05-11-attn-refractory-session-summary.md +++ /dev/null @@ -1,398 +0,0 @@ -# 2026-05-11 ATOM Attention Refactory Session Summary - -> 主题：从 `sglang plugin` 设计者视角，梳理 `attention backend` 重构中的 runtime / metadata / KV cache / kernel reuse 边界，并记录几种 kernel 共用 draft 的结论。 - -## 1. 本次会话的主线 - -本次讨论的核心不是继续从 `vllm plugin` 出发看问题，而是明确切换到： - -- **`sglang plugin attn backend` 设计者视角** -- 关注 `sglang plugin` 自己的 runtime 语义 -- 判断哪些层应该继续由 `sglang plugin` 自有 -- 判断哪些层值得尽量与 `ATOM core` 共用 - -围绕这个目标，本次会话主要探索了 6 个问题： - -1. `RadixAttention` / `radix cache` 与 `PagedAttention` 的关系 -2. public `sglang` 是否也会把 radix runtime lowering 到 paged / index metadata -3. metadata 为什么会越长越多，以及哪些字段本质上属于 host/runtime -4. `kv_indptr` / `kv_indices` 为什么既绑定存储模型，又仍然能跨 host 共用 -5. `sglang plugin` 和 public `sglang` 在 KV cache 管理上的异同 -6. 在不大改文件结构、不做过度软件工程的前提下，怎么试探性地共用 kernel call - -## 2. 对 `sglang plugin` 重构最关键的几个判断 - -### 2.1 不是所有问题都该叫 “attention backend 复用” - -当前问题至少分成三层： - -1. **host/runtime 层** - - `ForwardBatch` - - `RadixAttention` - - speculative / graph / verify / extend 调度 -2. **execution metadata lowering 层** - - `page_table` - - `kv_indptr` - - `kv_indices` - - `qo_indptr` - - `kv_last_page_len` -3. **kernel / execution core 层** - - `flash_attn_varlen_func` - - `mla_decode_fwd` - - `mla_prefill_fwd` - - `pa_fwd_asm` - - `pa_persistent_fwd` - -后续任何“共用更多 code”的讨论，都应该先说明是在第几层说话。 - -### 2.2 `sglang plugin` 不应该为了复用去对齐到 `vllm` 的 host runtime - -`vllm plugin` 和 `sglang plugin` 上层接入模型不同： - -- `vllm plugin` 更像 layer-owned / metadata-builder-owned 路径 -- `sglang plugin` 更像 `ForwardBatch` 驱动的 backend runtime 路径 - -所以： - -- 不该共享 `vllm plugin` 那层 runtime facade -- 也不该强行统一大 metadata 对象 -- 真正值得共享的，是更靠近 kernel 的 execution metadata 和 kernel call - -### 2.3 `sglang plugin` 想共用 `ATOM` code，最现实的边界在 kernel call 一层 - -如果先不做大规模结构改造，最适合先共用的是： - -- `mla_decode_fwd` -- `pa_persistent_fwd` -- `pa_fwd_asm` -- `flash_attn_varlen_func` - -也就是： - -- 保留各自的 runtime lowering -- 先把“掉 kernel 前最后一跳”抽出来 - -## 3. `RadixAttention` / `radix cache` 和 `PagedAttention` 的结论 - -### 3.1 `RadixAttention` 不是最终 kernel - -`RadixAttention` 更像： - -- `sglang` 模型层统一使用的 attention layer 壳 -- 它知道自己跑在 `ForwardBatch` 语义里 -- 最终还是委托给 `forward_batch.attn_backend` - -所以它不是 “另一套底层注意力数学”，而是 **host-facing runtime abstraction**。 - -### 3.2 `PagedAttention` 和 `radix` 解决的问题不在同一层 - -- `radix cache` 更偏 prefix 命中 / 复用 / eviction -- `paged attention` 更偏 block/page 形式的物理组织和读取 - -可以理解成： - -- `radix` 管“历史怎么共享和命中” -- `paged` 管“命中的历史最后如何变成 page/block 索引给 kernel” - -### 3.3 public `sglang` 自己也在做类似 lowering - -这不是 `ATOM plugin` 才有的行为。 - -public `sglang` 的： - -- `aiter_backend` -- `triton_backend` -- `trtllm_mha_backend` -- 部分 `flashinfer` 路径 - -本身也会把： - -- `ForwardBatch` -- `req_to_token` -- `seq_lens` - -lower 成： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `page_table` / `block_tables` - -所以，`radix runtime -> paged/index metadata` 不是 plugin 特例，而是公共设计模式的一部分。 - -## 4. Metadata 分层结论 - -### 4.1 现在至少有三套 metadata family - -1. **`ATOM native`** - - `AttentionMetaData` -2. **`vllm plugin`** - - `AttentionMetaData` 外壳 - - `plugin_metadata` 内层 payload -3. **`sglang plugin`** - - `ForwardMetadata` - -### 4.2 `AttentionMetaData` 和 `plugin/attention.py` 的关系 - -`AttentionMetaData` 是外层通用容器； -`plugin/attention.py` 里的很多 dataclass，是 `vllm plugin` 为表达宿主 batch/runtime 语义而塞进 `plugin_metadata` 的 payload。 - -所以它们不是并列关系，而是： - -```text -AttentionMetaData - └─ plugin_metadata - ├─ flash-style plugin metadata - ├─ MLA plugin metadata - └─ sparse MLA plugin metadata -``` - -### 4.3 `ForwardMetadata` 的定位 - -`ForwardMetadata` 更像： - -- `sglang backend` 内部的 lowering state -- 它不是最终统一 metadata -- 更不是天然可共享的 host-agnostic object - -它里面混着： - -- runtime control 信息 -- execution metadata -- persistent kernel 预构造结果 - -## 5. 三类字段 - -为了判断能不能共享，本次把字段分成三类： - -### 5.1 host/runtime-bound - -典型例子： - -- `vllm` 的 `query_start_loc` -- `num_decodes` -- `num_prefills` -- `chunked_context` -- `sglang` 的 `run_graph` - -这些字段描述的是宿主框架如何组织 batch/runtime，而不是 kernel 怎么算。 - -### 5.2 execution/kernel-bound - -典型例子： - -- `slot_mapping` -- `block_tables` -- `page_table` -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `kv_last_page_len` -- `work_indptr` -- `reduce_indptr` - -这些字段才是 kernel 真正需要消费的执行信息。 - -### 5.3 hardware-sensitive - -典型例子： - -- `max_q_len` -- `max_kv_len` -- `num_kv_splits` -- `head_dim` -- `attn_out_dtype` - -这些字段不是 host 语义，但会影响： - -- page/block 解释 -- kernel 选路 -- workspace/split 策略 - -## 6. 关于 `kv_indptr` / `kv_indices` 的关键结论 - -这个问题本次单独做了较多背景解释，结论是： - -### 6.1 它们不是完全 storage-agnostic - -因为它们绑定的是某种 **可索引的 KV 地址空间**： - -- token slot index -- page/block index -- compacted KV index - -所以它们不能脱离底层存储模型独立存在。 - -### 6.2 但它们可以是 host/runtime-agnostic - -因为它们不表达： - -- `ForwardBatch.forward_mode` -- `query_start_loc` -- prefix 命中策略 -- graph/speculative 的宿主控制语义 - -它们只表达： - -- 这次 kernel 要去哪些 KV 地址读数据 -- 每个 request 的边界在哪里 - -### 6.3 两个例子 - -#### 例子 A：`sglang` - -- `req_to_token[req_id, token_pos] = slot_id` -- lowering 后得到： - - `kv_indptr = [0, 10]` - - `kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5]` - -#### 例子 B：`vllm/ATOM` - -- 上层是 `block_table` -- 展开后同样可得到： - - `kv_indptr = [0, 10]` - - `kv_indices = [8, 9, 10, 11, 0, 1, 2, 3, 4, 5]` - -这说明： - -- 上层来源不同 -- 但 lower 后的 execution metadata 可以收敛 - -## 7. public `sglang` 与 `sglang plugin` 在 KV cache 管理上的异同 - -### 7.1 相同点 - -- owner 都还是 public `sglang` - - `ReqToTokenPool` - - `TokenToKVPool` - - `KVCache` - - `radix_cache` -- request 到 slot 的映射机制没变 -- prefix / radix 复用体系没变 - -### 7.2 不同点 - -#### A. MHA 写 cache 路径不同 - -public `sglang` 更多沿用 pool 提供的标准写入接口： - -- `token_to_kv_pool.set_kv_buffer(...)` - -而 `sglang plugin` 在 MHA 路径里会绕过标准写法，自己做一次： - -- `set_kv_buffer_with_layout_shuffle(...)` -- 把同一块 public pool buffer 按更偏 `ATOM` kernel 的 SHUFFLE layout 写进去 - -#### B. MLA 路径更接近 public `sglang` - -MLA 上，plugin 更多沿用 public pool 的 contract： - -- `token_to_kv_pool.set_kv_buffer(...)` -- `get_key_buffer(...)` -- `get_value_buffer(...)` - -#### C. Lowering 目标不同 - -- public `sglang` 的 metadata 更通用 -- plugin 的 `ForwardMetadata` 更偏 `ATOM` kernel - - `page_table` - - `pa_metadata_*` - -### 7.3 一个更准确的说法 - -`sglang plugin` **没有改写 public `sglang` 的 pool owner / allocator / radix tree**， -但它在 MHA 路径上： - -- 改变了同一块 public pool buffer 的布局解释与写入约定 -- 并把 lowering 结果继续朝 `ATOM` kernel 所需的 metadata 形态推进 - -## 8. Kernel 共用的两种 draft - -本次会话里，实际做了两种方向的 draft。 - -### 8.1 方案 A：共享 kernel wrapper - -思路： - -- 新增共享小 metadata - - `MLADecodeKernelMetadata` - - `MHAAsmKernelMetadata` - - `MHAPersistentKernelMetadata` -- 新增共享 wrapper - - `run_mla_decode_kernel` - - `run_mha_asm_kernel` - - `run_mha_persistent_kernel` - -优点： - -- 边界清晰 -- 更像长期可演进的 execution API -- 不需要统一大 metadata - -缺点： - -- 还是引入了一层新的 shared contract - -### 8.2 方案 B：`sglang plugin` 主动适配 `ATOM core` - -思路： - -- 不改 `ATOM core` -- 在 `sglang plugin` 里把 `ForwardMetadata` 转成 `AttentionMetaData` -- 再伪造最小 `shim/self` -- 直接调用： - - `PagedAttentionImpl.paged_attention_asm` - - `PagedAttentionImpl.paged_attention_persistent_asm` - - `MLAAttention._forward_decode` - -优点： - -- 不改 `ATOM core` -- 更直观地证明 “plugin 主动复用 core” 是可行的 - -缺点： - -- 需要 shim -- 尤其 MLA 路径更脆弱 -- 更适合验证可行性，不适合长期保留 - -## 9. 当前最值得保留的判断 - -### 9.1 不应该强行统一大 metadata - -不建议直接统一： - -- `AttentionMetaData` -- `plugin_metadata` -- `ForwardMetadata` - -因为它们都带有较重的宿主 runtime 语义。 - -### 9.2 最值得共享的是 kernel-facing execution metadata - -真正最有价值的共享边界是： - -- `kv_indptr` -- `kv_indices` -- `qo_indptr` -- `page_table` -- `context_lens` -- `work_*` -- `reduce_*` - -以及基于它们的 kernel call / runner。 - -### 9.3 重构顺序建议 - -如果继续推进 `sglang plugin` 重构，更合理的顺序是： - -1. 保持 `sglang` runtime 自有 -2. 显式化 metadata lowering 边界 -3. 尽量共享 kernel call / execution helper -4. 最后再考虑是否继续上推到 runner 层 - -## 10. 一句话总结 - -本次会话最终收敛出的关键认识是： - -**`sglang plugin attention` 重构的关键，不是把 `sglang` runtime 改得更像 `vllm`，而是把 runtime、metadata lowering、kernel execution 这三层拆清楚；只要 lowering 之后能在 execution metadata 上对齐，就能在不统一宿主语义的前提下，共用更多 `ATOM` 的底层 attention 代码。** diff --git a/work_log/attn_refractory/decoupling-diagram.md b/work_log/attn_refractory/decoupling-diagram.md deleted file mode 100644 index b95c0031e5..0000000000 --- a/work_log/attn_refractory/decoupling-diagram.md +++ /dev/null @@ -1,158 +0,0 @@ -# ATOM Plugin Decoupling Diagram - -下面这张图对比了当前入口架构和目标解耦架构。这里采用的是更激进、也更干净的方案： - -- `vLLM plugin` 和 `SGLang plugin` 在 `atom/plugin/` 层彻底分家 -- 不再保留共享的 plugin runtime -- 共享只发生在更底层的 `ATOM core` - -## Current - -```mermaid -flowchart TD - A[Shared Plugin Entry
prepare.py / register.py / config.py] - B[Global Runtime State
_CURRENT_FRAMEWORK
current_atom_config
ops.Attention] - C[vLLM Branch
platform / registry / patch / graph] - D[SGLang Branch
wrapper / registry hijack / RadixAttention / graph] - E[Mixed Plugin Core
attention selection / static_forward_context / loader glue] - - A --> B - B --> C - B --> D - C --> E - D --> E -``` - -### 当前设计的核心问题 - -- 入口层通过切换全局状态来区分 `vLLM` 和 `SGLang`，而不是通过物理隔离的子系统建立边界。 -- `prepare.py` / `register.py` / `config.py` 这些共享入口，本身就是耦合中心。 -- host 差异直接泄漏到 plugin 内部结构里，导致 backend 选择、attention 抽象和 runtime context 都知道上层框架。 - -## Target - -```mermaid -flowchart TD - A1[vLLM Plugin
bootstrap.py
register.py
config.py
runtime.py
attention glue
model adapters] - A2[SGLang Plugin
bootstrap.py
register.py
config.py
runtime.py
attention glue
model adapters] - - B[ATOM Core
model_ops
kernels
loader
metadata structs
model specialization] - - A1 --> B - A2 --> B -``` - -### 目标设计的关键点 - -- `vLLM` 和 `SGLang` 在 plugin 层是两个完整子系统，而不是一套共享 runtime 上的两个分支。 -- `atom/plugin/` 下不再存在共享的 `prepare_model()`、共享的 `register.py` 总入口、共享的 plugin config translator。 -- `vLLM plugin` 自己处理： - - platform registration - - model override - - graph / MLA / weight hook patch - - vLLM model adapter / attention glue -- `SGLang plugin` 自己处理： - - external model wrapper - - attention backend registration / override - - graph / speculative / wrapper patch - - SGLang model adapter / attention glue -- 真正共享的只保留在更底层的 `ATOM core`： - - `model_ops` - - kernels - - loader - - generic metadata structures - - model-family specialization - -## 最重要的切断点 - -```mermaid -flowchart LR - A[Cut 1
shared prepare.py] - B[Cut 2
shared register.py] - C[Cut 3
shared config.py] - D[Cut 4
framework global state] - E[Cut 5
ops.Attention global rebinding] - - A --> F[vllm/ and sglang/ own bootstrap] - B --> F - C --> F - D --> G[host-local runtime only] - E --> H[host-specific attention glue] -``` - -## 推荐目录形态 - -```text -atom/plugin/ - vllm/ - bootstrap.py - register.py - config.py - runtime.py - attention/ - models/ - graph/ - sglang/ - bootstrap.py - register.py - config.py - runtime.py - attention/ - models/ - graph/ -``` - -这里的重点不是“换目录”，而是： - -- `vllm/runtime.py` 和 `sglang/runtime.py` 各自独立 -- `vllm/config.py` 和 `sglang/config.py` 各自独立 -- `vllm/register.py` 和 `sglang/register.py` 各自独立 -- 不再有 plugin 级共享 runtime - -## 术语说明 - -### bootstrap - -`bootstrap` 在这里指的是“插件被宿主框架加载时，最先执行的接线初始化层”。 - -它负责的事情通常包括： - -- 注册 plugin 扩展点 -- 安装 patch / hook -- 构造 host adapter / wrapper -- 把宿主配置翻译到 ATOM 可用的格式 - -它不应该负责： - -- attention forward 本身 -- 长期 runtime state 管理 -- 每次请求都要走的核心执行逻辑 - -### register - -`register` 指的是把 plugin 的能力挂到宿主暴露的扩展点上，例如： - -- 注册 platform -- 注册 model class -- 注册 attention backend - -它更偏“声明接入点”，通常是 bootstrap 的一部分。 - -### adapter - -`adapter` 指的是宿主框架和 ATOM core 之间的桥接层。 -它的职责是做接口和参数语义翻译，而不是定义 attention / model 的核心语义。 - -### runtime - -`runtime` 指的是插件在宿主里长期存在、服务执行链路的那部分运行时逻辑。 -如果按这份设计推进，`runtime` 应该是 host-local 的： - -- `vLLM runtime` 只属于 `vLLM plugin` -- `SGLang runtime` 只属于 `SGLang plugin` - -而不是共享一个“Plugin runtime”。 - -## 一句话结论 - -真正的解耦不是把共享入口再包装一层，而是把 `atom/plugin` 从“共享状态机”改成“两个完全独立的 host plugin 子系统”，共享只留给更底层的 `ATOM core`。 diff --git a/work_log/attn_refractory/vllm-integration-architecture.md b/work_log/attn_refractory/vllm-integration-architecture.md deleted file mode 100644 index e8876268d0..0000000000 --- a/work_log/attn_refractory/vllm-integration-architecture.md +++ /dev/null @@ -1,269 +0,0 @@ -# ATOM vLLM Plugin Integration Architecture - -## 1. 启动入口 - -vLLM plugin 的 setuptools entrypoint 定义在： - -- `pyproject.toml` - -关键位置： - -- `vllm.platform_plugins` - - `atom = "atom.plugin.vllm.register:register_platform"` -- `vllm.general_plugins` - - `atom_model_registry = "atom.plugin.vllm.register:register_model"` - -这意味着 vLLM 启动时会先进入： - -- `atom/plugin/vllm/register.py::register_platform()` -- `atom/plugin/vllm/register.py::register_model()` - -## 2. 关键代码文件与职责 - -### 2.1 `atom/plugin/vllm/register.py` - -这是 vLLM plugin 的启动接线层，主要负责： - -- 设置 plugin mode -- 返回自定义 platform -- 覆盖 vLLM 的 `ModelRegistry` -- 安装 MLA patch -- patch `Attention.process_weights_after_loading` -- 安装 graph capture patch - -重点符号： - -- `register_platform()` -- `register_model()` -- `_VLLM_MODEL_REGISTRY_OVERRIDES` - -## 2.2 `atom/plugin/vllm/platform.py` - -这是 vLLM 平台侧 attention backend 选择入口。 - -重点符号： - -- `ATOMPlatform.get_attn_backend_cls()` - -它根据 `attn_selector_config` 决定： - -- MHA -> `atom.model_ops.attentions.aiter_attention.AiterBackend` -- MLA -> `atom.model_ops.attentions.aiter_mla.AiterMLABackend` -- sparse MLA -> `atom.plugin.vllm.attention_backend.mla_sparse.AiterMLASparseBackend` - -这里是 vLLM runtime 真正选择 “哪个 ATOM attention backend class” 的入口。 - -## 2.3 `atom/plugin/vllm/model_wrapper.py` - -这是 vLLM model integration 的核心文件。 - -主要职责： - -- 把 `VllmConfig` 翻译成 `atom_config` -- 做 plugin 运行环境准备 -- 选择并实例化 ATOM 模型类 -- 把 vLLM forward context 中的 `positions` 传给 ATOM 路径 -- 对 sparse MLA indexer 做额外注册 - -重点符号： - -- `_prepare_env()` -- `ATOMModelBase.__init__()` -- `_register_indexer_caches_with_vllm()` -- `forward()` - -## 2.4 `atom/plugin/config.py` - -这是 vLLM 配置翻译的桥。 - -重点符号： - -- `generate_atom_config_for_vllm_plugin()` -- `_generate_atom_config_from_vllm_config()` - -它负责把： - -- `VllmConfig` -- `scheduler_config` -- `cache_config` -- `parallel_config` -- `quant_config` - -映射成 ATOM 自己的 `Config` 与 `PluginConfig`。 - -## 2.5 `atom/model_ops/paged_attention.py` - -这是 ATOM attention 在 vLLM plugin 下真正进入执行的关键对象。 - -重点符号： - -- `PagedAttention.__init__()` -- `PagedAttention.forward()` - -在 `is_vllm()` 下，它不会直接走 native/server 的那套 `impl` 初始化，而是： - -- 构造 vLLM 的 `Attention` / `MLAAttention` 外壳 -- 把额外 impl 参数塞进去 -- 注册到 `static_forward_context` -- forward 时通过 `unified_attention_with_output_base_for_plugin_mode()` 回调到 ATOM impl - -## 2.6 `atom/plugin/attention.py` - -这是 vLLM plugin mode 的核心 glue 层。 - -虽然文件在 `plugin/` 根目录，但从职责看它明显偏 vLLM。 - -主要职责： - -- 定义 `vllmAiterAttentionBackendMethods` -- 提供 backend decorator / metadata builder decorator -- 处理 plugin-mode metadata -- 处理 `static_forward_context` -- 为 vLLM 的 attention backend 补 OOT 接口行为 - -重点符号： - -- `vllmAiterAttentionBackendMethods` -- 各类 `*DecoratorForPluginMode` - -## 2.7 `atom/plugin/attention_mha.py` - -这是 vLLM plugin mode 下 MHA impl 的补丁层。 - -重点符号： - -- `PagedAttentionImplPluginModeMethods` - -它补的是： - -- rope/cache/update -- plugin-mode forward -- 与 vLLM KV cache/metadata 对齐的逻辑 - -## 2.8 `atom/plugin/attention_mla.py` - -这是 vLLM plugin mode 下 MLA impl 的补丁层。 - -重点符号： - -- `MLAAttentionImplPluginModeMethods` -- `_mla_plugin_mode_init` - -它补的是： - -- MLA plugin-mode metadata -- qk rope / kv cache / chunked prefill -- vLLM MLA 路径专有行为 - -## 2.9 `atom/plugin/vllm/mla_patch.py` - -这是 vLLM 原生 `MLAAttention` 的 monkey patch 入口。 - -重点符号： - -- `_patch_vllm_mla_attention_forward_impl()` -- `_patch_vllm_mla_attention_process_weights_after_loading()` -- `patch_vllm_mla_attention()` - -这层作用是： - -- 把 vLLM `MLAAttention.forward_impl()` 改接到 ATOM impl -- 把权重后处理改接到 ATOM 的 `impl.process_weights_after_loading()` - -## 2.10 `atom/plugin/vllm/graph_capture_patch.py` - -这是 vLLM graph capture 补丁入口。 - -重点符号： - -- `apply_graph_capture_patch()` - -它实际委托到共享实现： - -- `atom/plugin/graph_capture_patch.py` - -但调用点是 vLLM 自己的 plugin register 流程。 - -## 3. 当前 vLLM plugin 集成链路 - -```mermaid -flowchart TD - A[vLLM startup] - B[entrypoint
register_platform] - C[entrypoint
register_model] - D[ATOMPlatform.get_attn_backend_cls] - E[ATOMModelBase] - F[generate_atom_config_for_vllm_plugin] - G[_prepare_env
set_attn_cls + init_aiter_dist] - H[Instantiate ATOM model] - I[PagedAttention / MLAAttention wrapper] - J[plugin/attention*.py glue] - K[ATOM impl
PagedAttentionImpl / MLAAttention] - L[aiter kernels] - - A --> B - A --> C - B --> D - C --> E - E --> F - E --> G - E --> H - H --> I - I --> J - J --> K - K --> L -``` - -## 4. 最关键的架构判断 - -### 4.1 vLLM plugin 不是简单“替换 backend class” - -它实际上同时做了三件事： - -1. 替换 vLLM 的 model entry -2. 替换/选择 vLLM 的 attention backend class -3. 用 patch + decorator 的方式把 vLLM 的 layer/runtime 语义桥接到 ATOM impl - -### 4.2 真正共享得最好的是 native/server 与 vLLM plugin 的 attention impl - -尤其： - -- `PagedAttentionImpl` -- `MLAAttention` -- 低层 `aiter` kernel family - -也就是说，vLLM plugin 侧的关键复杂度主要不在 kernel，而在： - -- `ModelRegistry` override -- `VllmConfig -> atom_config` -- `static_forward_context` -- `MLAAttention` patch -- graph capture patch - -### 4.3 当前 `plugin/attention.py`、`attention_mha.py`、`attention_mla.py` - -虽然放在 `atom/plugin/` 根目录，但从 vLLM integration 的角度看，它们本质上更像： - -- `vLLM plugin impl glue` -- 而不是真正框架无关的 shared runtime - -## 5. 推荐你重点阅读的文件顺序 - -如果你是为了快速理解 vLLM plugin 集成架构，建议按这个顺序读： - -1. `pyproject.toml` -2. `atom/plugin/vllm/register.py` -3. `atom/plugin/vllm/platform.py` -4. `atom/plugin/vllm/model_wrapper.py` -5. `atom/plugin/config.py` -6. `atom/model_ops/paged_attention.py` -7. `atom/plugin/vllm/mla_patch.py` -8. `atom/plugin/attention.py` -9. `atom/plugin/attention_mha.py` -10. `atom/plugin/attention_mla.py` - -## 6. 一句话结论 - -vLLM plugin 的集成方式，本质上是： - -**用 entrypoint + platform + model wrapper 接入 vLLM，用 ATOM 自己的 attention impl 和 aiter kernels 提供核心执行，再用 patch/decorator 把 vLLM 的 layer/runtime 语义桥接进去。** From a1f88ae4b3062297fb6aea4971cd08c8ccdc5b5e Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Wed, 13 May 2026 07:21:25 +0000 Subject: [PATCH 75/76] [ATOM-SGL][Attn refrac] Route DeepSeek MLA through an SGLang wrapper Move the SGLang DeepSeek MLA runtime entry from legacy forward glue into SGLangDeepseekMLAAttention while keeping RadixAttention and the full-attention backend as the host/backend layers. Shrink deepseek_mla_forward.py into a helper module and clarify absorbed vs non-absorbed path naming. --- atom/plugin/sglang/models/deepseek_mla.py | 28 +- .../sglang/models/deepseek_mla_attention.py | 332 ++++++++++ .../sglang/models/deepseek_mla_forward.py | 594 +----------------- 3 files changed, 358 insertions(+), 596 deletions(-) create mode 100644 atom/plugin/sglang/models/deepseek_mla_attention.py diff --git a/atom/plugin/sglang/models/deepseek_mla.py b/atom/plugin/sglang/models/deepseek_mla.py index a8d4deb1a1..a7165b7112 100644 --- a/atom/plugin/sglang/models/deepseek_mla.py +++ b/atom/plugin/sglang/models/deepseek_mla.py @@ -3,19 +3,19 @@ """Model-level DeepSeek MLA patching for SGLang plugin mode. -This module owns the monkey-patch entrypoints that adapt DeepSeek MLA models to -SGLang plugin mode. The heavy DeepSeek-specific forward and weight helpers live -in `atom.plugin.sglang.models.deepseek_mla_forward`. +This module owns the install-time hooks that adapt DeepSeek MLA models to +SGLang plugin mode. The heavy DeepSeek-specific runtime helpers live in +`atom.plugin.sglang.models.deepseek_mla_forward`. """ from __future__ import annotations from typing import TYPE_CHECKING, Any -import torch - +from atom.plugin.sglang.models.deepseek_mla_attention import ( + SGLangDeepseekMLAAttention, +) from atom.plugin.sglang.models.deepseek_mla_forward import ( - forward_sgl_plugin_mode, init_sgl_attrs, process_mla_kv_b_proj_after_loading, ) @@ -56,20 +56,8 @@ def _patch_mla_attention_for_sglang( ) -> None: """Patch one DeepSeek MLA layer for SGLang plugin mode.""" init_sgl_attrs(attn, config, kv_cache_dtype) - - def patched_forward( - positions: torch.Tensor, - hidden_states: torch.Tensor, - **kwargs: Any, - ) -> torch.Tensor: - from atom.plugin.sglang.models.base_model_wrapper import ( - get_current_forward_batch, - ) - - kwargs["forward_batch"] = get_current_forward_batch() - return forward_sgl_plugin_mode(attn, positions, hidden_states, **kwargs) - - attn.forward = patched_forward + if not isinstance(attn.mla_attn, SGLangDeepseekMLAAttention): + attn.mla_attn = SGLangDeepseekMLAAttention(attn, attn.mla_attn) attn.process_weights_after_loading = lambda: process_mla_kv_b_proj_after_loading( attn ) diff --git a/atom/plugin/sglang/models/deepseek_mla_attention.py b/atom/plugin/sglang/models/deepseek_mla_attention.py new file mode 100644 index 0000000000..bda5ef87d0 --- /dev/null +++ b/atom/plugin/sglang/models/deepseek_mla_attention.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek MLA wrapper for SGLang plugin mode. + +This adapter keeps the model-side entry at ``self.mla_attn(...)`` and owns the +SGLang-specific runtime dispatch for DeepSeek MLA. It is intentionally shaped +closer to the vLLM plugin path than the older model-side monkey-patched entry. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +if TYPE_CHECKING: + from atom.models.deepseek_v2 import DeepseekV2MLAAttention + + +class SGLangDeepseekMLAAttention(nn.Module): + """Enter SGLang DeepSeek MLA runtime through ``self.mla_attn(...)``.""" + + def __init__( + self, + owner_attn: "DeepseekV2MLAAttention", + base_attn: nn.Module, + ) -> None: + super().__init__() + self.owner_attn = owner_attn + self.base_attn = base_attn + + @property + def attn(self): + return getattr(self.base_attn, "attn", self.base_attn) + + def _get_forward_batch(self, kwargs: dict[str, Any]): + forward_batch = kwargs.get("forward_batch", None) + if forward_batch is None: + from atom.plugin.sglang.models.base_model_wrapper import ( + get_current_forward_batch, + ) + + forward_batch = get_current_forward_batch() + kwargs["forward_batch"] = forward_batch + if forward_batch is None: + raise RuntimeError( + "forward_batch is required for SGLang DeepSeek MLA wrapper" + ) + 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 _gather_runtime_inputs( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + input_scattered: bool, + ): + total_tokens = self._infer_total_tokens(forward_batch, q_input) + q_input = self._maybe_all_gather( + q_input, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + kv_c_normed = self._maybe_all_gather( + kv_c_normed, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + k_pe = self._maybe_all_gather( + k_pe, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + positions = self._maybe_all_gather( + positions, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + q_scale = self._maybe_all_gather( + q_scale, + total_tokens=total_tokens, + input_scattered=input_scattered, + ) + return q_input, kv_c_normed, k_pe, positions, q_scale + + def _project_q( + self, + q_input: torch.Tensor, + q_scale: torch.Tensor | None, + ) -> torch.Tensor: + attn = self.owner_attn + from atom.plugin.sglang.models.deepseek_mla_forward import _unwrap_linear_output + + if attn.q_lora_rank is not None: + q = ( + attn.q_b_proj(q_input, q_scale) + if q_scale is not None + else attn.q_b_proj(q_input) + ) + else: + q = ( + attn.q_proj(q_input, q_scale) + if q_scale is not None + else attn.q_proj(q_input) + ) + return _unwrap_linear_output(q).view( + -1, attn.num_local_heads, attn.qk_head_dim + ) + + def _forward_absorbed( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + ) -> torch.Tensor: + attn = self.owner_attn + 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 + + q = self._project_q(q_input, q_scale) + k_nope = kv_c_normed.unsqueeze(1) + k_pe = k_pe.unsqueeze(1) + q_nope, q_pe = q.split( + [attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1 + ) + q_nope_out = mla_absorbed_bmm( + attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank + ) + + if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla: + q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) + + if nsa_use_prefill_cp(forward_batch): + latent_cache = torch.cat([k_nope.squeeze(1), k_pe.squeeze(1)], dim=-1) + k_nope, k_pe = attn.rebuild_cp_kv_cache( + latent_cache, forward_batch, k_nope, k_pe + ) + + save_kv_cache = True + if attn.use_fused_qk_rope_concat_and_cache_mla: + mla_attn = _get_sglang_radix_attn(self.base_attn) + kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(mla_attn.layer_id) + q_out_dtype = ( + dtypes.fp8 + if attn.kv_cache_dtype == "fp8_e4m3" + else q_nope_out.dtype + ) + q = torch.empty( + ( + q_nope_out.shape[0], + attn.num_local_heads, + attn.kv_lora_rank + attn.qk_rope_head_dim, + ), + dtype=q_out_dtype, + device=q_nope_out.device, + ) + fused_qk_rope_concat_and_cache_mla( + q_nope_out, + q_pe, + k_nope, + k_pe, + kv_cache, + q, + forward_batch.out_cache_loc, + mla_attn.k_scale, + mla_attn.k_scale, + positions, + attn.rotary_emb.cos_cache, + attn.rotary_emb.sin_cache, + is_neox=attn.rotary_emb.is_neox_style, + is_nope_first=True, + ) + k = None + v = None + save_kv_cache = False + else: + q = torch.cat([q_nope_out, q_pe], dim=-1) + k = torch.cat([k_nope, k_pe], dim=-1) + v = k_nope + + attn_output = self.base_attn( + q, + k, + v, + forward_batch=forward_batch, + save_kv_cache=save_kv_cache, + ) + attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank) + attn_bmm_output = mla_v_up_proj( + attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim + ) + return attn.o_proj(attn_bmm_output) + + def _forward_non_absorbed( + self, + q_input: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + positions: torch.Tensor, + q_scale: torch.Tensor | None, + *, + forward_batch, + ) -> torch.Tensor: + attn = self.owner_attn + from atom.plugin.sglang.models.deepseek_mla_forward import ( + _set_mla_kv_buffer_for_non_absorbed, + _unwrap_linear_output, + ) + + q = self._project_q(q_input, q_scale) + _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) + + kv_a = kv_c_normed + k_pe = k_pe.unsqueeze(1) + if attn.rotary_emb is not None: + q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) + q[..., attn.qk_nope_head_dim :] = q_pe + + _set_mla_kv_buffer_for_non_absorbed(attn, kv_a, k_pe, forward_batch) + + kv = _unwrap_linear_output(attn.kv_b_proj(kv_a)).view( + -1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim + ) + k_nope = kv[..., : attn.qk_nope_head_dim] + v = kv[..., attn.qk_nope_head_dim :] + k = torch.cat( + [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)], + dim=-1, + ) + + attn_output = attn.attn_non_absorbed( + q, + k, + v, + forward_batch=forward_batch, + save_kv_cache=False, + ) + attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim) + return attn.o_proj(attn_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: + attn = self.owner_attn + forward_batch = self._get_forward_batch(kwargs) + + from atom.plugin.sglang.models.deepseek_mla_forward import ( + _can_run_non_absorbed_mla_now, + ) + 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): + q_input, kv_c_normed, k_pe, positions, q_scale = self._gather_runtime_inputs( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + input_scattered=attn_tp_context.input_scattered, + ) + + if forward_batch.forward_mode.is_extend_without_speculative(): + if _can_run_non_absorbed_mla_now(attn, forward_batch): + attn.current_sgl_plugin_attn_path = "non_absorbed" + return self._forward_non_absorbed( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + ) + attn.current_sgl_plugin_attn_path = "absorbed_fallback" + else: + attn.current_sgl_plugin_attn_path = "absorbed" + + return self._forward_absorbed( + q_input, + kv_c_normed, + k_pe, + positions, + q_scale, + forward_batch=forward_batch, + ) diff --git a/atom/plugin/sglang/models/deepseek_mla_forward.py b/atom/plugin/sglang/models/deepseek_mla_forward.py index 25f1ef79aa..f88aa762f8 100644 --- a/atom/plugin/sglang/models/deepseek_mla_forward.py +++ b/atom/plugin/sglang/models/deepseek_mla_forward.py @@ -1,38 +1,26 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. -"""Model-specific DeepSeek MLA helpers for SGLang plugin mode. +"""Helper functions for DeepSeek MLA in SGLang plugin mode. -DeepSeek MLA (Multi-Latent Attention) forward logic for SGLang plugin mode: -absorbed BMM computation, MHA/MLA path dispatch (prefill -> MHA, decode -> MLA), -and kv_b_proj weight splitting (w_kc/w_vc). - -This module lives under ``atom.plugin.sglang.models`` because the logic is -DeepSeek-model-specific rather than a generic SGLang attention backend. - -TODO: rewrite this file once sglang's attention flow is unified into ATOM's -attention layer - the MLA absorbed path and MHA dispatch will then be handled -natively by ATOM's attention ops, making this sglang-specific module -unnecessary. +This module now contains only the low-level helpers that are still shared by +the SGLang DeepSeek MLA wrapper and the install-time weight hooks: +absorbed BMM math, small utility helpers, non-absorbed cache staging, and +kv_b_proj post-load processing. """ from __future__ import annotations -from typing import TYPE_CHECKING, Any, NamedTuple, Optional +from typing import TYPE_CHECKING, Any, Optional import torch from aiter import dtypes -from aiter.dist.parallel_state import get_tensor_model_parallel_world_size, get_tp_group from atom.model_ops.base_attention import Attention from atom.model_ops.attention_mla import ( dynamic_per_batched_tensor_quant, - fused_qk_rope_concat_and_cache_mla, ) from atom.models.utils import maybe_prefix -from atom.models.deepseek_v2 import _fuse_rmsnorm_quant -from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context -from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.models.deepseek_common.utils import ( _use_aiter_gfx95, @@ -90,25 +78,6 @@ def bmm_fp8(A, B, A_scale, B_scale, dtype, out=None): raise RuntimeError("bmm_fp8 requires CUDA (sgl_kernel)") -class SglPrepareResult(NamedTuple): - q_pe: torch.Tensor - k_pe: torch.Tensor - q_nope_out: torch.Tensor - k_nope: torch.Tensor - forward_batch: Any - zero_allocator: Any - positions: torch.Tensor - topk_indices: Optional[torch.Tensor] - llama_4_scaling: Optional[Any] - - -class SglMhaPrepareResult(NamedTuple): - q: torch.Tensor - k: torch.Tensor - v: torch.Tensor - forward_batch: Any - - def _unwrap_linear_output(output: Any) -> torch.Tensor: """Normalize ATOM/public-SGLang linear outputs to a tensor.""" if isinstance(output, tuple): @@ -116,57 +85,6 @@ def _unwrap_linear_output(output: Any) -> torch.Tensor: return output -def _linear_quant_type_value(linear: Any) -> Optional[int]: - quant_type = getattr(linear, "quant_type", None) - return None if quant_type is None else getattr(quant_type, "value", quant_type) - - -def _fuse_qk_rmsnorm_and_q_quant( - attn: DeepseekV2MLAAttention, - q: torch.Tensor, - k_nope: torch.Tensor, - *, - output_unquantized_q: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: - """Fuse q/k RMSNorm and q quant using ATOM's DeepSeek-V2 path.""" - - (q_quantized, q_scale), q_normed, k_nope_normed, _ = _fuse_rmsnorm_quant( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - k_nope, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - None, - dtype_quant=attn.quant_dtype, - shuffle=False, - scale_shuffle_padding=False, - group_size=128, - quant_type=_linear_quant_type_value(attn.q_b_proj), - output_unquantized_inp1=output_unquantized_q, - transpose_scale=True, - ) - return q_quantized, q_scale, q_normed, k_nope_normed - - -def _fuse_qk_rmsnorm( - attn: DeepseekV2MLAAttention, - q: torch.Tensor, - k_nope: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fuse q/k RMSNorm without quantizing q.""" - from atom.models.deepseek_v2 import _fused_qk_rmsnorm - - return _fused_qk_rmsnorm( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - k_nope, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - ) - - def _prepare_weight_for_bmm( weight: torch.Tensor, in_dim: int, out_dim: int ) -> torch.Tensor: @@ -199,7 +117,7 @@ def init_sgl_attrs( attn.w_scale = None attn.w_scale_k = None attn.w_scale_v = None - attn.attn_mha = Attention( + attn.attn_non_absorbed = Attention( num_heads=attn.num_local_heads, head_dim=attn.qk_head_dim, scale=attn.scaling, @@ -208,10 +126,10 @@ def init_sgl_attrs( layer_num=attn.layer_num, use_mla=False, v_head_dim=attn.v_head_dim, - prefix=maybe_prefix(attn.prefix, "attn_mha"), + prefix=maybe_prefix(attn.prefix, "attn_non_absorbed"), ) - if hasattr(attn.attn_mha, "attn"): - attn.attn_mha.attn.kv_b_proj = None + if hasattr(attn.attn_non_absorbed, "attn"): + attn.attn_non_absorbed.attn.kv_b_proj = None def mla_absorbed_bmm( @@ -340,275 +258,31 @@ def mla_v_up_proj( return mla_absorbed_bmm( attn, inp, weight, weight_scale, weight_scale_k, out_dim ).flatten(1, 2) - - -def forward_sgl_prepare( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> SglPrepareResult: - """Prepare QKV for sglang MLA attention.""" - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - - forward_batch = model_kwargs.get("forward_batch", None) - zero_allocator = model_kwargs.get("zero_allocator", None) - llama_4_scaling = model_kwargs.get("llama_4_scaling", None) - q_lora = None - topk_indices = None - - if attn.q_lora_rank is not None: - q, latent_cache = ( - get_attn_tp_context() - .fetch_qkv_latent() - .split( - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - ) - - if ( - q.shape[0] != positions.shape[0] - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = torch.cat([q, latent_cache], dim=-1) - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] < positions.shape[0]: - raise RuntimeError( - f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {positions.shape[0]}" - ) - qkv_lora = qkv_lora[: positions.shape[0]] - q, latent_cache = torch.split( - qkv_lora, - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - - k_nope = latent_cache[..., : attn.kv_lora_rank] - q_scale = None - - if getattr(attn, "fuse_qknorm_quant", False): - q, q_scale, q_lora, k_nope = _fuse_qk_rmsnorm_and_q_quant( - attn, - q, - k_nope, - output_unquantized_q=attn.use_nsa, - ) - elif getattr(attn, "fuse_qknorm", False): - q, k_nope = _fuse_qk_rmsnorm(attn, q, k_nope) - elif attn.alt_stream is not None and get_is_capture_mode(): - current_stream = torch.cuda.current_stream() - attn.alt_stream.wait_stream(current_stream) - q = attn.q_a_layernorm(q) - with torch.cuda.stream(attn.alt_stream): - k_nope = attn.kv_a_layernorm(k_nope) - current_stream.wait_stream(attn.alt_stream) - else: - q = attn.q_a_layernorm(q) - k_nope = attn.kv_a_layernorm(k_nope) - - if attn.use_nsa and q_lora is None: - q_lora = q - - if ( - attn.alt_stream is not None - and get_is_capture_mode() - and forward_batch.forward_mode.is_decode_or_idle() - and q_lora is not None - ): - current_stream = torch.cuda.current_stream() - attn.alt_stream.wait_stream(current_stream) - with torch.cuda.stream(attn.alt_stream): - k_nope = k_nope.unsqueeze(1) - q = _unwrap_linear_output( - attn.q_b_proj(q, q_scale) - if q_scale is not None - else attn.q_b_proj(q) - ).view(-1, attn.num_local_heads, attn.qk_head_dim) - topk_indices = attn.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=attn.layer_num, - ) - current_stream.wait_stream(attn.alt_stream) - else: - k_nope = k_nope.unsqueeze(1) - q = _unwrap_linear_output( - attn.q_b_proj(q, q_scale) if q_scale is not None else attn.q_b_proj(q) - ).view(-1, attn.num_local_heads, attn.qk_head_dim) - if q_lora is not None: - topk_indices = attn.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=attn.layer_num, - ) - else: - q = _unwrap_linear_output(attn.q_proj(hidden_states)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - latent_cache = _unwrap_linear_output(attn.kv_a_proj_with_mqa(hidden_states)) - k_nope = latent_cache[..., : attn.kv_lora_rank] - k_nope = attn.kv_a_layernorm(k_nope).unsqueeze(1) - - q_nope, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) - k_pe = latent_cache[..., attn.kv_lora_rank :].unsqueeze(1) - - q_nope_out = mla_absorbed_bmm( - attn, q_nope, attn.w_kc, attn.w_scale, attn.w_scale_k, attn.kv_lora_rank - ) - - if attn.rotary_emb is not None and not attn.use_fused_qk_rope_concat_and_cache_mla: - q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) - - if nsa_use_prefill_cp(forward_batch): - k_nope, k_pe = attn.rebuild_cp_kv_cache( - latent_cache, forward_batch, k_nope, k_pe - ) - - return SglPrepareResult( - q_pe=q_pe, - k_pe=k_pe, - q_nope_out=q_nope_out, - k_nope=k_nope, - forward_batch=forward_batch, - zero_allocator=zero_allocator, - positions=positions, - topk_indices=topk_indices, - llama_4_scaling=llama_4_scaling, - ) - - -def forward_sgl_core( - attn: DeepseekV2MLAAttention, - prepared: SglPrepareResult, -) -> torch.Tensor: - """Core MLA attention computation for sglang.""" - save_kv_cache = True - - if attn.use_fused_qk_rope_concat_and_cache_mla: - mla_attn = _get_sglang_radix_attn(attn.mla_attn) - kv_cache = prepared.forward_batch.token_to_kv_pool.get_key_buffer( - mla_attn.layer_id - ) - q_out_dtype = ( - dtypes.fp8 - if attn.kv_cache_dtype == "fp8_e4m3" - else prepared.q_nope_out.dtype - ) - q = torch.empty( - ( - prepared.q_nope_out.shape[0], - attn.num_local_heads, - attn.kv_lora_rank + attn.qk_rope_head_dim, - ), - dtype=q_out_dtype, - device=prepared.q_nope_out.device, - ) - - fused_qk_rope_concat_and_cache_mla( - prepared.q_nope_out, - prepared.q_pe, - prepared.k_nope, - prepared.k_pe, - kv_cache, - q, - prepared.forward_batch.out_cache_loc, - mla_attn.k_scale, - mla_attn.k_scale, - prepared.positions, - attn.rotary_emb.cos_cache, - attn.rotary_emb.sin_cache, - is_neox=attn.rotary_emb.is_neox_style, - is_nope_first=True, - ) - k = None - v = None - save_kv_cache = False - else: - q = torch.cat([prepared.q_nope_out, prepared.q_pe], dim=-1) - k = torch.cat([prepared.k_nope, prepared.k_pe], dim=-1) - v = prepared.k_nope - - if prepared.llama_4_scaling is not None: - q = q * prepared.llama_4_scaling - - extra_kwargs = {} - if prepared.topk_indices is not None: - extra_kwargs["topk_indices"] = prepared.topk_indices - - attn_output = attn.mla_attn( - q, - k, - v, - forward_batch=prepared.forward_batch, - save_kv_cache=save_kv_cache, - **extra_kwargs, - ) - attn_output = attn_output.view(-1, attn.num_local_heads, attn.kv_lora_rank) - - attn_bmm_output = mla_v_up_proj( - attn, attn_output, attn.w_vc, attn.w_scale, attn.w_scale_v, attn.v_head_dim - ) - - return attn.o_proj(attn_bmm_output) - - -def _dispatch_sgl_plugin_attn_path(forward_batch) -> str: - """Decide the attention algorithm for this batch based on forward_mode.""" - if forward_batch.forward_mode.is_extend_without_speculative(): - return "mha" - return "mla" - - -def forward_sgl_plugin_mode_mla( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - prepared = forward_sgl_prepare(attn, positions, hidden_states, **model_kwargs) - from atom.utils.forward_context import get_forward_context - - if get_forward_context().context.is_dummy_run: - base_hidden_states = ( - hidden_states[0] if isinstance(hidden_states, tuple) else hidden_states - ) - dummy_output = base_hidden_states.new_empty( - (base_hidden_states.shape[0], base_hidden_states.shape[-1]) - ) - return dummy_output - return forward_sgl_core(attn, prepared) - - def _get_sglang_radix_attn(attn_module): return attn_module.attn if hasattr(attn_module, "attn") else attn_module -def _set_mla_kv_buffer_for_mha( +def _set_mla_kv_buffer_for_non_absorbed( attn: DeepseekV2MLAAttention, kv_a: torch.Tensor, k_pe: torch.Tensor, forward_batch, ) -> None: - attn_mha = _get_sglang_radix_attn(attn.attn_mha) + attn_non_absorbed = _get_sglang_radix_attn(attn.attn_non_absorbed) cache_k = torch.cat([kv_a.unsqueeze(1), k_pe], dim=-1) forward_batch.token_to_kv_pool.set_kv_buffer( - attn_mha, + attn_non_absorbed, forward_batch.out_cache_loc, cache_k, cache_k, ) -def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool: - """Check if the model configuration supports the MHA attention path.""" +def _can_run_non_absorbed_mla_now( + attn: DeepseekV2MLAAttention, + forward_batch, +) -> bool: + """Check if the non-absorbed MLA path is allowed for this batch.""" del forward_batch if attn.use_nsa: return False @@ -617,238 +291,6 @@ def _can_run_sgl_mha_now(attn: DeepseekV2MLAAttention, forward_batch) -> bool: return True -def forward_sgl_mha_prepare( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> SglMhaPrepareResult: - - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_mha_prepare") - - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - - attn_mha = _get_sglang_radix_attn(attn.attn_mha) - if getattr(attn_mha, "kv_b_proj", None) is None: - attn_mha.kv_b_proj = attn.kv_b_proj - - if attn.q_lora_rank is not None: - q, latent_cache = ( - get_attn_tp_context() - .fetch_qkv_latent() - .split( - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - ) - - if ( - q.shape[0] != positions.shape[0] - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = torch.cat([q, latent_cache], dim=-1) - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] < positions.shape[0]: - raise RuntimeError( - f"qkv_lora gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {positions.shape[0]}" - ) - qkv_lora = qkv_lora[: positions.shape[0]] - q, latent_cache = torch.split( - qkv_lora, - [attn.q_lora_rank, attn.kv_lora_rank + attn.qk_rope_head_dim], - dim=-1, - ) - - if _use_aiter_gfx95 and attn.q_b_proj.weight.dtype == torch.float8_e4m3fn: - (q, q_scale), _, _, _ = _fuse_rmsnorm_quant( - q, - attn.q_a_layernorm.weight, - attn.q_a_layernorm.eps, - None, - None, - None, - res1=None, - dtype_quant=torch.float8_e4m3fn, - group_size=128, - quant_type=_linear_quant_type_value(attn.q_b_proj), - output_unquantized_inp1=False, - transpose_scale=True, - ) - q = _unwrap_linear_output(attn.q_b_proj(q, q_scale)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - else: - q = attn.q_a_layernorm(q) - q = _unwrap_linear_output(attn.q_b_proj(q)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - else: - q = _unwrap_linear_output(attn.q_proj(hidden_states, hidden_states_scale)).view( - -1, attn.num_local_heads, attn.qk_head_dim - ) - latent_cache = _unwrap_linear_output( - attn.kv_a_proj_with_mqa(hidden_states, hidden_states_scale) - ) - - _, q_pe = q.split([attn.qk_nope_head_dim, attn.qk_rope_head_dim], dim=-1) - kv_a, _ = latent_cache.split([attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1) - latent_cache = latent_cache.unsqueeze(1) - - if _use_aiter_gfx95 and attn.kv_b_proj.weight.dtype == torch.float8_e4m3fn: - (kv_a_quanted, kv_a_quanted_scale), kv_a, _, _ = _fuse_rmsnorm_quant( - kv_a, - attn.kv_a_layernorm.weight, - attn.kv_a_layernorm.eps, - None, - None, - None, - res1=None, - dtype_quant=torch.float8_e4m3fn, - group_size=128, - quant_type=_linear_quant_type_value(attn.kv_b_proj), - output_unquantized_inp1=True, - transpose_scale=True, - ) - else: - kv_a_quanted = None - kv_a = attn.kv_a_layernorm(kv_a) - - k_pe = latent_cache[:, :, attn.kv_lora_rank :] - if attn.rotary_emb is not None: - q_pe, k_pe = attn.rotary_emb(positions, q_pe, k_pe) - q[..., attn.qk_nope_head_dim :] = q_pe - - _set_mla_kv_buffer_for_mha(attn, kv_a, k_pe, forward_batch) - - if kv_a_quanted is not None: - kv = _unwrap_linear_output(attn.kv_b_proj(kv_a_quanted, kv_a_quanted_scale)) - else: - kv = _unwrap_linear_output(attn.kv_b_proj(kv_a)) - kv = kv.view(-1, attn.num_local_heads, attn.qk_nope_head_dim + attn.v_head_dim) - k_nope = kv[..., : attn.qk_nope_head_dim] - v = kv[..., attn.qk_nope_head_dim :] - k = torch.cat( - [k_nope, k_pe.expand(-1, attn.num_local_heads, -1)], - dim=-1, - ) - return SglMhaPrepareResult(q=q, k=k, v=v, forward_batch=forward_batch) - - -def forward_sgl_mha_core( - attn: DeepseekV2MLAAttention, - prepared: SglMhaPrepareResult, -) -> torch.Tensor: - attn_output = attn.attn_mha( - prepared.q, - prepared.k, - prepared.v, - forward_batch=prepared.forward_batch, - save_kv_cache=False, - ) - attn_output = attn_output.reshape(-1, attn.num_local_heads * attn.v_head_dim) - return attn.o_proj(attn_output) - - -def forward_sgl_plugin_mode_mha( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode_mha") - if not _can_run_sgl_mha_now(attn, forward_batch): - attn.current_sgl_plugin_attn_path = "mla_fallback" - return forward_sgl_plugin_mode_mla( - attn, - positions, - hidden_states, - **model_kwargs, - ) - prepared = forward_sgl_mha_prepare(attn, positions, hidden_states, **model_kwargs) - return forward_sgl_mha_core(attn, prepared) - - -def prepare_qkv_latent( - attn: DeepseekV2MLAAttention, - hidden_states: torch.Tensor, - forward_batch, -) -> torch.Tensor: - """Prepare QKV latent tensor for the sglang communicator.""" - assert attn.q_lora_rank is not None - hidden_states_scale = None - if isinstance(hidden_states, tuple): - hidden_states, hidden_states_scale = hidden_states - qkv_lora = attn.fused_qkv_a_proj(hidden_states, hidden_states_scale) - - expected_tokens = 0 - if hasattr(forward_batch, "positions") and forward_batch.positions is not None: - expected_tokens = int(forward_batch.positions.shape[0]) - if expected_tokens <= 0: - expected_tokens = int(getattr(forward_batch, "seq_lens_sum", 0) or 0) - - if ( - expected_tokens > 0 - and qkv_lora.shape[0] != expected_tokens - and get_tensor_model_parallel_world_size() > 1 - ): - qkv_lora = get_tp_group().all_gather(qkv_lora, dim=0) - if qkv_lora.shape[0] > expected_tokens: - qkv_lora = qkv_lora[:expected_tokens] - elif qkv_lora.shape[0] < expected_tokens: - raise RuntimeError( - f"prepare_qkv_latent gather mismatch: got {qkv_lora.shape[0]}, " - f"expected {expected_tokens}" - ) - return qkv_lora - - -def forward_sgl_plugin_mode( - attn: DeepseekV2MLAAttention, - positions: torch.Tensor, - hidden_states: torch.Tensor, - **model_kwargs, -) -> torch.Tensor: - """Full MLA forward in sglang plugin mode.""" - forward_batch = model_kwargs.get("forward_batch", None) - if forward_batch is None: - raise RuntimeError("forward_batch is required in forward_sgl_plugin_mode") - - attn_tp_context = get_attn_tp_context() - with attn_tp_context.maybe_input_scattered(forward_batch): - if attn.q_lora_rank is not None: - attn_tp_context.set_attn_inputs( - AttentionInputs( - hidden_states, - forward_batch, - lambda hs, fb: prepare_qkv_latent(attn, hs, fb), - ) - ) - attn_path = _dispatch_sgl_plugin_attn_path(forward_batch) - attn.current_sgl_plugin_attn_path = attn_path - if attn_path == "mha": - return forward_sgl_plugin_mode_mha( - attn, - positions, - hidden_states, - **model_kwargs, - ) - if attn_path == "mla": - return forward_sgl_plugin_mode_mla( - attn, - positions, - hidden_states, - **model_kwargs, - ) - raise ValueError(f"Unsupported plugin attention path: {attn_path}") - - def _read_kv_b_proj_weight(attn: DeepseekV2MLAAttention) -> torch.Tensor: """Read kv_b_proj weight, handling AWQ and fnuz dtypes.""" if hasattr(attn.kv_b_proj, "qweight"): From 4aec32a65692acf82080fce9fd345b0f3794eee0 Mon Sep 17 00:00:00 2001 From: ZhiweiYan-96 Date: Mon, 18 May 2026 15:25:41 +0000 Subject: [PATCH 76/76] [ATOM SGL] runtime extraction --- .../sglang/attention_backend/attention_gdn.py | 2 +- .../full_attention/radix_attention.py | 2 +- .../sglang/models/base_model_wrapper.py | 453 +++--------------- .../sglang/models/deepseek_mla_attention.py | 2 +- .../sglang/models/deepseek_nextn_wrapper.py | 21 +- atom/plugin/sglang/models/qwen3_5.py | 4 +- atom/plugin/sglang/runtime/__init__.py | 20 + atom/plugin/sglang/runtime/context.py | 126 +++++ atom/plugin/sglang/runtime/forward_context.py | 234 +++++++++ atom/plugin/sglang/runtime/model_arch.py | 22 + 10 files changed, 481 insertions(+), 405 deletions(-) create mode 100644 atom/plugin/sglang/runtime/__init__.py create mode 100644 atom/plugin/sglang/runtime/context.py create mode 100644 atom/plugin/sglang/runtime/forward_context.py create mode 100644 atom/plugin/sglang/runtime/model_arch.py diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py index 2403efaa9e..f29a43b1b7 100644 --- a/atom/plugin/sglang/attention_backend/attention_gdn.py +++ b/atom/plugin/sglang/attention_backend/attention_gdn.py @@ -164,7 +164,7 @@ def _build_gdn_metadata( def build( cls, forward_batch_or_metadata: Any ) -> Optional["SGLangGDNForwardContext"]: - from atom.plugin.sglang.models.base_model_wrapper import ( + from atom.plugin.sglang.runtime import ( SGLangForwardBatchMetadata, ) diff --git a/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py index 8e20eecb14..613458ee37 100644 --- a/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py +++ b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py @@ -129,7 +129,7 @@ def forward_impl_plugin_mode( # for sglang, forward_batch is required forward_batch = kwargs.get("forward_batch", None) if forward_batch is None: - from atom.plugin.sglang.models.base_model_wrapper import ( + from atom.plugin.sglang.runtime import ( get_current_forward_batch, ) diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index dbace2e647..85bc6608e1 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -6,13 +6,8 @@ To add a new model, append its architecture class name to _MODEL_NAMES. """ -import copy - import logging -from contextlib import contextmanager -from contextvars import ContextVar -from dataclasses import dataclass -from typing import Any, ClassVar, Iterable, Optional, Tuple, Union +from typing import Any, Iterable, Optional, Tuple, Union import torch from torch import nn @@ -22,312 +17,26 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -logger = logging.getLogger("atom.plugin.sglang.models") - -_RUNTIME_SENTINEL = object() - -# Context for patched DeepSeek attention layers that need wrapper state without -# changing every intermediate forward signature. ContextVar keeps nested or -# concurrent forwards isolated and lets us reliably restore the prior value. -_current_forward_batch: ContextVar[Optional[ForwardBatch]] = ContextVar( - "atom_sglang_current_forward_batch", default=None +from atom.plugin.sglang.runtime import ( + MODEL_ARCH_SPECS, + SGLangForwardBatchMetadata, + SGLangPluginRuntime, + bind_current_forward_batch, + get_current_forward_batch, + get_model_arch_spec, + plugin_runtime_scope, ) +logger = logging.getLogger("atom.plugin.sglang.models") -def get_current_forward_batch(): - return _current_forward_batch.get() - - -def _is_dummy_forward(forward_batch: ForwardBatch) -> bool: - # SGLang's IDLE batch is the plugin-side equivalent of ATOM dummy run. - forward_mode = getattr(forward_batch, "forward_mode", None) - return bool( - forward_mode is not None - and hasattr(forward_mode, "is_idle") - and forward_mode.is_idle() - ) - - -def _pad_dummy_like( - tensor: Optional[torch.Tensor], - *, - length: int, - fill_value: int | float = 0, -) -> Optional[torch.Tensor]: - if tensor is None: - return None - shape = (length, *tensor.shape[1:]) - return torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device) - - -def _materialize_atom_dummy_forward( - input_ids: Optional[torch.Tensor], - positions: Optional[torch.Tensor], - input_embeds: Optional[torch.Tensor], - forward_batch: ForwardBatch, -) -> tuple[ - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - ForwardBatch, -]: - """Convert an empty SGLang IDLE batch into ATOM-style dummy forward inputs.""" - dummy_positions = positions.new_zeros((1,)) - dummy_input_ids = input_ids.new_zeros((1,)) - dummy_input_embeds = _pad_dummy_like(input_embeds, length=1, fill_value=0) - - model_forward_batch = copy.copy(forward_batch) - model_forward_batch.positions = dummy_positions - model_forward_batch.batch_size = 1 - model_forward_batch.seq_lens_sum = 1 - model_forward_batch.seq_lens = forward_batch.seq_lens.new_ones((1,)) - model_forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu.new_ones((1,)) - - return dummy_input_ids, dummy_positions, dummy_input_embeds, model_forward_batch - - -def _trim_hidden_states_for_output(hidden_states, num_tokens: int): - if torch.is_tensor(hidden_states): - return hidden_states[:num_tokens] - if isinstance(hidden_states, tuple): - return tuple( - tensor[:num_tokens] if torch.is_tensor(tensor) else tensor - for tensor in hidden_states - ) - return hidden_states - - -def _resolve_num_tokens_across_dp( - atom_config: Any, - forward_batch: ForwardBatch, - num_tokens: int, - is_dummy_run: bool, -) -> torch.Tensor: - """Resolve per-DP token counts for ATOM's CPU-side DPMetadata. - - Real SGLang dp-attention batches carry ``global_num_tokens_cpu`` from the - scheduler. That list is the source of truth for mixed prefill/decode/idle - batches, where token counts may look like [8, 1, 8, 8]. - - Some SGLang synthetic/static batches, especially CUDA graph capture batches, - only keep the global token buffer on GPU. ATOM's DPMetadata is CPU-side and - needs a CPU tensor before model forward, so avoid reading the GPU buffer back - to CPU. We only fallback when the batch advertises the same-shape DP buffer - layout (global_dp_buffer_len == local_num_tokens * dp_size), where the CPU - equivalent is exactly [local_num_tokens] * dp_size. - - IDLE batches are reported by SGLang as 0 tokens on the current rank, but - this wrapper materializes them as one local dummy token before entering - ATOM. Patch the current DP rank after resolving the distribution so - ``DPMetadata`` sees a local count that matches the actual ATOM input. - """ - global_num_tokens_cpu = getattr(forward_batch, "global_num_tokens_cpu", None) - if global_num_tokens_cpu is not None: - num_tokens_across_dp = torch.tensor( - global_num_tokens_cpu, dtype=torch.int32, device="cpu" - ) - else: - dp_size = atom_config.parallel_config.data_parallel_size - global_num_tokens_gpu = getattr(forward_batch, "global_num_tokens_gpu", None) - global_dp_buffer_len = getattr(forward_batch, "global_dp_buffer_len", None) - is_static_same_shape_batch = ( - global_num_tokens_gpu is not None - and global_dp_buffer_len == num_tokens * dp_size - ) - if not is_static_same_shape_batch: - raise RuntimeError( - "[SGL+ATOM] SGLang dp-attention requires " - "forward_batch.global_num_tokens_cpu unless the batch uses static " - "same-shape DP metadata." - ) - - # Static batches, such as CUDA graph capture batches, may only keep - # global token counts on GPU. Avoid GPU-to-CPU reads here and mirror - # their same-shape layout directly for ATOM's CPU DPMetadata. - num_tokens_across_dp = torch.full( - (dp_size,), num_tokens, dtype=torch.int32, device="cpu" - ) - - if is_dummy_run: - # SGLang reports idle ranks as 0 tokens, but ATOM materializes them - # as one local dummy token so collectives and DPMetadata stay aligned. - dp_rank = atom_config.parallel_config.data_parallel_rank - num_tokens_across_dp[dp_rank] = num_tokens - return num_tokens_across_dp - - -def _set_sglang_forward_context( - atom_config: Any, - forward_batch: ForwardBatch, - positions: torch.Tensor, -) -> None: - """Bridge SGLang batch metadata into ATOM's global forward context.""" - from atom.utils.forward_context import ( - AttentionMetaData, - Context, - set_forward_context, - ) - - forward_mode = forward_batch.forward_mode - # TODO: This max_seqlen_q is not the source of truth for prefill attention; - # SGLang plugin attention consumes forward_batch.attn_backend.forward_metadata - # directly. In this wrapper it is only needed by ATOM MoE padding: under - # dp-attention + TP (non-EP all_gather/reduce_scatter), decode/idle batches - # must use 1 so pad_for_all_gather keeps fixed-shape collectives aligned. - # Leaving it as 0 there can make active and dummy ranks send different - # shapes to DP all_gather and hang. - max_seqlen_q = 1 if forward_mode.is_decode_or_idle() else 0 - attn_metadata = AttentionMetaData(max_seqlen_q=max_seqlen_q) - batch_size = int(forward_batch.batch_size) - is_dummy_run = _is_dummy_forward(forward_batch) - is_prefill = forward_mode.is_prefill() - num_tokens = int(positions.shape[0]) - - enable_dp_attention = bool(atom_config.enable_dp_attention) - if enable_dp_attention: - # SGLang owns the cross-DP token distribution under dp-attention; ATOM - # uses it to derive graph_bs and fixed-size MoE gather/scatter buffers. - num_tokens_across_dp = _resolve_num_tokens_across_dp( - atom_config, forward_batch, num_tokens, is_dummy_run - ) - graph_bs = int(torch.max(num_tokens_across_dp).item()) - else: - # Without dp-attention, ATOM runs with local-rank shapes only. There is - # no cross-DP token distribution to pass into DPMetadata, so graph_bs - # follows the local prefill token count or decode batch size. - num_tokens_across_dp = None - graph_bs = num_tokens if is_prefill else batch_size - context = Context( - positions=positions, - is_prefill=is_prefill, - is_dummy_run=is_dummy_run, - batch_size=batch_size, - graph_bs=graph_bs, - ) - set_forward_context( - attn_metadata=attn_metadata, - atom_config=atom_config, - context=context, - num_tokens=num_tokens, - num_tokens_across_dp=num_tokens_across_dp, - ) - - -def _reset_sglang_forward_context() -> None: - from atom.utils.forward_context import reset_forward_context - - reset_forward_context() - - -@contextmanager -def plugin_runtime_scope( - *, - framework: Optional[str] = None, - atom_config: Any = _RUNTIME_SENTINEL, -): - """Temporarily bind plugin runtime globals to one wrapper instance. - - ATOM core currently relies on process-global framework/config state. In - SGLang speculative mode both target and draft wrappers coexist, so plugin - entrypoints must save/restore those globals around each init/load/forward. - """ - - import atom.config as atom_config_module - import atom.plugin.prepare as plugin_prepare - - prev_framework = plugin_prepare._CURRENT_FRAMEWORK - prev_atom_config = getattr(atom_config_module, "_current_atom_config", None) - - if framework is not None: - plugin_prepare._set_framework_backbone(framework) - if atom_config is not _RUNTIME_SENTINEL: - atom_config_module._current_atom_config = atom_config - - try: - yield - finally: - plugin_prepare._CURRENT_FRAMEWORK = prev_framework - atom_config_module._current_atom_config = prev_atom_config - - -@dataclass(frozen=True) -class SGLangForwardBatchMetadata: - """Small context object for one SGLang model forward.""" - - forward_batch: Optional[ForwardBatch] - pp_proxy_tensors: Optional[PPProxyTensors] = None - save_kv_cache: bool = True - _current: ClassVar[ContextVar[Optional["SGLangForwardBatchMetadata"]]] = ContextVar( - "atom_sglang_current_forward_batch_metadata", - default=None, - ) - - @classmethod - def current(cls) -> Optional["SGLangForwardBatchMetadata"]: - return cls._current.get() - - @classmethod - def build( - cls, - forward_batch: Optional[ - Union[ForwardBatch, "SGLangForwardBatchMetadata"] - ] = None, - *, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - save_kv_cache: Optional[bool] = None, - ) -> Optional["SGLangForwardBatchMetadata"]: - if isinstance(forward_batch, cls): - return forward_batch - if forward_batch is None and pp_proxy_tensors is None and save_kv_cache is None: - return cls.current() - return cls( - forward_batch=forward_batch, - pp_proxy_tensors=pp_proxy_tensors, - save_kv_cache=True if save_kv_cache is None else save_kv_cache, - ) - - @classmethod - @contextmanager - def bind(cls, metadata: Optional["SGLangForwardBatchMetadata"]): - meta_token = cls._current.set(metadata) - batch_token = _current_forward_batch.set( - None if metadata is None else metadata.forward_batch - ) - try: - yield metadata - finally: - _current_forward_batch.reset(batch_token) - cls._current.reset(meta_token) - - @staticmethod - def to_intermediate_tensors( - intermediate_tensors, - metadata: Optional["SGLangForwardBatchMetadata"], - ): - if intermediate_tensors is not None or metadata is None: - return intermediate_tensors - pp_proxy_tensors = metadata.pp_proxy_tensors - if pp_proxy_tensors is None: - return intermediate_tensors - tensors = getattr(pp_proxy_tensors, "tensors", None) - if tensors is None: - return intermediate_tensors - from atom.models.utils import IntermediateTensors - - return IntermediateTensors(dict(tensors)) - - -@dataclass(frozen=True) -class ModelArchSpec: - wrapper_binds_gdn_context: bool = False - apply_deepseek_patch: bool = False - - -_MODEL_ARCH_SPECS = { - "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True), - "Qwen3MoeForCausalLM": ModelArchSpec(), - "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True), -} +__all__ = [ + "EntryClass", + "SGLangForwardBatchMetadata", + "SGLangPluginRuntime", + "bind_current_forward_batch", + "get_current_forward_batch", + "plugin_runtime_scope", +] class _AtomCausalLMBaseForSglang(nn.Module): @@ -353,7 +62,7 @@ def __init__( self.vocab_size = config.vocab_size self.unpadded_vocab_size = config.vocab_size self.model_arch = getattr(config, "architectures", [""])[0] - self.model_arch_spec = _MODEL_ARCH_SPECS.get(self.model_arch, ModelArchSpec()) + self.model_arch_spec = get_model_arch_spec(self.model_arch) import atom @@ -449,94 +158,60 @@ def forward( **model_kwargs: Any, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): - metadata = SGLangForwardBatchMetadata.build( - forward_batch, - pp_proxy_tensors=pp_proxy_tensors, - save_kv_cache=model_kwargs.get("save_kv_cache"), - ) - - if _is_dummy_forward(forward_batch): - ( - model_input_ids, - model_positions, - model_input_embeds, - model_forward_batch, - ) = _materialize_atom_dummy_forward( - input_ids, - positions, - input_embeds, - forward_batch, + with SGLangPluginRuntime( + atom_config=self.atom_config, + forward_batch=forward_batch, + positions=positions, + input_ids=input_ids, + input_embeds=input_embeds, + set_forward_context=not self.model_arch_spec.wrapper_binds_gdn_context, + ) as runtime: + metadata = SGLangForwardBatchMetadata.build( + runtime.forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + save_kv_cache=model_kwargs.get("save_kv_cache"), ) - else: - ( - model_input_ids, - model_positions, - model_input_embeds, - model_forward_batch, - ) = ( - input_ids, - positions, - input_embeds, - forward_batch, + model_inputs = dict( + input_ids=runtime.input_ids, + positions=runtime.positions, + intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( + pp_proxy_tensors, metadata + ), + inputs_embeds=runtime.input_embeds, ) - - model_inputs = dict( - input_ids=model_input_ids, - positions=model_positions, - intermediate_tensors=SGLangForwardBatchMetadata.to_intermediate_tensors( - pp_proxy_tensors, metadata - ), - inputs_embeds=model_input_embeds, - ) - uses_context_only_forward = ( - self.model_arch_spec.apply_deepseek_patch - or self.model_arch_spec.wrapper_binds_gdn_context - ) - with SGLangForwardBatchMetadata.bind(metadata): - if self.model_arch_spec.wrapper_binds_gdn_context: - from atom.plugin.sglang.attention_backend.attention_gdn import ( - SGLangGDNForwardContext, - ) - - with SGLangGDNForwardContext.bind(metadata): - hidden_states = self.model(**model_inputs) - elif uses_context_only_forward: - try: - _set_sglang_forward_context( - self.atom_config, model_forward_batch, model_positions + uses_context_only_forward = ( + self.model_arch_spec.apply_deepseek_patch + or self.model_arch_spec.wrapper_binds_gdn_context + ) + with SGLangForwardBatchMetadata.bind(metadata): + if self.model_arch_spec.wrapper_binds_gdn_context: + from atom.plugin.sglang.attention_backend.attention_gdn import ( + SGLangGDNForwardContext, ) + + with SGLangGDNForwardContext.bind(metadata): + hidden_states = self.model(**model_inputs) + elif uses_context_only_forward: hidden_states = self.model(**model_inputs) - finally: - _reset_sglang_forward_context() - else: - try: - _set_sglang_forward_context( - self.atom_config, model_forward_batch, model_positions - ) + else: hidden_states = self.model( **model_inputs, - forward_batch=model_forward_batch, + forward_batch=runtime.forward_batch, get_embedding=get_embedding, pp_proxy_tensors=pp_proxy_tensors, **model_kwargs, ) - finally: - _reset_sglang_forward_context() - - if self.pp_group.is_last_rank: - if _is_dummy_forward(forward_batch): - # TODO: Revisit if SGLang ever sends non-empty dummy batches. - # Today this path only runs when an empty IDLE batch is expanded - # to one ATOM dummy token, so the output boundary must trim back to - # the original SGLang-visible length: 0 tokens. - hidden_states = _trim_hidden_states_for_output(hidden_states, 0) - return self.logits_processor( - input_ids, - hidden_states, - self.model.lm_head, - forward_batch, - ) - return hidden_states + + hidden_states = runtime.trim_output(hidden_states) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + ) + return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # The passed `weights` iterable from sglang is ignored because ATOM @@ -552,7 +227,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): EntryClass = [] -for _name in _MODEL_ARCH_SPECS: +for _name in MODEL_ARCH_SPECS: _cls = type(_name, (_AtomCausalLMBaseForSglang,), {}) globals()[_name] = _cls EntryClass.append(_cls) diff --git a/atom/plugin/sglang/models/deepseek_mla_attention.py b/atom/plugin/sglang/models/deepseek_mla_attention.py index bda5ef87d0..bd89870125 100644 --- a/atom/plugin/sglang/models/deepseek_mla_attention.py +++ b/atom/plugin/sglang/models/deepseek_mla_attention.py @@ -38,7 +38,7 @@ def attn(self): def _get_forward_batch(self, kwargs: dict[str, Any]): forward_batch = kwargs.get("forward_batch", None) if forward_batch is None: - from atom.plugin.sglang.models.base_model_wrapper import ( + from atom.plugin.sglang.runtime import ( get_current_forward_batch, ) diff --git a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py index f41617e140..730355aaff 100644 --- a/atom/plugin/sglang/models/deepseek_nextn_wrapper.py +++ b/atom/plugin/sglang/models/deepseek_nextn_wrapper.py @@ -19,13 +19,11 @@ from atom.config import SpeculativeConfig from atom.plugin.config import generate_atom_config_for_plugin_mode -from atom.plugin.sglang.attention_backend.sgl_attention_mla import ( +from atom.plugin.sglang.models.deepseek_mla import ( setup_deepseek_for_sglang, ) -from atom.plugin.sglang.models.base_model_wrapper import ( - _current_forward_batch, - _reset_sglang_forward_context, - _set_sglang_forward_context, +from atom.plugin.sglang.runtime import ( + SGLangPluginRuntime, plugin_runtime_scope, ) @@ -159,18 +157,19 @@ def forward( raise ValueError("DeepSeek MTP draft forward requires speculative info") with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): - token = _current_forward_batch.set(forward_batch) - try: - _set_sglang_forward_context(self.atom_config, forward_batch, positions) + with SGLangPluginRuntime( + atom_config=self.atom_config, + forward_batch=forward_batch, + positions=positions, + input_ids=input_ids, + input_embeds=input_embeds, + ): hidden_states = self.model( input_ids=input_ids, positions=positions, hidden_states=forward_batch.spec_info.hidden_states, inputs_embeds=input_embeds, ) - finally: - _reset_sglang_forward_context() - _current_forward_batch.reset(token) if self.pp_group.is_last_rank: return self.logits_processor( diff --git a/atom/plugin/sglang/models/qwen3_5.py b/atom/plugin/sglang/models/qwen3_5.py index 22294b8548..c0b6b0ca0b 100644 --- a/atom/plugin/sglang/models/qwen3_5.py +++ b/atom/plugin/sglang/models/qwen3_5.py @@ -37,7 +37,7 @@ from atom.plugin.sglang.attention_backend.attention_gdn import ( SGLangGDNForwardContext, ) -from atom.plugin.sglang.models.base_model_wrapper import ( +from atom.plugin.sglang.runtime import ( SGLangForwardBatchMetadata, ) @@ -448,5 +448,5 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # SGLang discovers these multimodal wrappers from this module's `EntryClass`. # They are not covered by `base_model_wrapper.py`, whose generated entries only -# handle the plain causal-LM architectures in `_MODEL_ARCH_SPECS`. +# handle the plain causal-LM architectures in `MODEL_ARCH_SPECS`. EntryClass = [Qwen3_5ForConditionalGeneration, Qwen3_5MoeForConditionalGeneration] diff --git a/atom/plugin/sglang/runtime/__init__.py b/atom/plugin/sglang/runtime/__init__.py new file mode 100644 index 0000000000..f63890418f --- /dev/null +++ b/atom/plugin/sglang/runtime/__init__.py @@ -0,0 +1,20 @@ +"""Runtime utilities for ATOM's SGLang plugin integration.""" + +from atom.plugin.sglang.runtime.context import ( + SGLangForwardBatchMetadata, + bind_current_forward_batch, + get_current_forward_batch, + plugin_runtime_scope, +) +from atom.plugin.sglang.runtime.forward_context import SGLangPluginRuntime +from atom.plugin.sglang.runtime.model_arch import MODEL_ARCH_SPECS, get_model_arch_spec + +__all__ = [ + "MODEL_ARCH_SPECS", + "SGLangForwardBatchMetadata", + "SGLangPluginRuntime", + "bind_current_forward_batch", + "get_current_forward_batch", + "get_model_arch_spec", + "plugin_runtime_scope", +] diff --git a/atom/plugin/sglang/runtime/context.py b/atom/plugin/sglang/runtime/context.py new file mode 100644 index 0000000000..1fffae92f2 --- /dev/null +++ b/atom/plugin/sglang/runtime/context.py @@ -0,0 +1,126 @@ +"""Runtime context helpers for ATOM's SGLang plugin path.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import ClassVar, Optional, Union + +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors + +_RUNTIME_SENTINEL = object() +_current_forward_batch: ContextVar[Optional[ForwardBatch]] = ContextVar( + "atom_sglang_current_forward_batch", default=None +) + + +def get_current_forward_batch(): + return _current_forward_batch.get() + + +@contextmanager +def bind_current_forward_batch(forward_batch: Optional[ForwardBatch]): + token = _current_forward_batch.set(forward_batch) + try: + yield + finally: + _current_forward_batch.reset(token) + + +@contextmanager +def plugin_runtime_scope( + *, + framework: Optional[str] = None, + atom_config=_RUNTIME_SENTINEL, +): + """Temporarily bind process-global ATOM plugin runtime state. + + SGLang target/draft wrappers can coexist during speculative decoding, while + ATOM core still reads process-global framework/config state in some paths. + Keep those globals scoped to one wrapper call and restore them afterwards. + """ + + import atom.config as atom_config_module + import atom.plugin.prepare as plugin_prepare + + prev_framework = plugin_prepare._CURRENT_FRAMEWORK + prev_atom_config = getattr(atom_config_module, "_current_atom_config", None) + + if framework is not None: + plugin_prepare._set_framework_backbone(framework) + if atom_config is not _RUNTIME_SENTINEL: + atom_config_module._current_atom_config = atom_config + + try: + yield + finally: + plugin_prepare._CURRENT_FRAMEWORK = prev_framework + atom_config_module._current_atom_config = prev_atom_config + + +@dataclass(frozen=True) +class SGLangForwardBatchMetadata: + """Small context object for one SGLang model forward.""" + + forward_batch: Optional[ForwardBatch] + pp_proxy_tensors: Optional[PPProxyTensors] = None + save_kv_cache: bool = True + _current: ClassVar[ContextVar[Optional["SGLangForwardBatchMetadata"]]] = ContextVar( + "atom_sglang_current_forward_batch_metadata", + default=None, + ) + + @classmethod + def current(cls) -> Optional["SGLangForwardBatchMetadata"]: + return cls._current.get() + + @classmethod + def build( + cls, + forward_batch: Optional[ + Union[ForwardBatch, "SGLangForwardBatchMetadata"] + ] = None, + *, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + save_kv_cache: Optional[bool] = None, + ) -> Optional["SGLangForwardBatchMetadata"]: + if isinstance(forward_batch, cls): + return forward_batch + if forward_batch is None and pp_proxy_tensors is None and save_kv_cache is None: + return cls.current() + return cls( + forward_batch=forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + save_kv_cache=True if save_kv_cache is None else save_kv_cache, + ) + + @classmethod + @contextmanager + def bind(cls, metadata: Optional["SGLangForwardBatchMetadata"]): + meta_token = cls._current.set(metadata) + batch_token = _current_forward_batch.set( + None if metadata is None else metadata.forward_batch + ) + try: + yield metadata + finally: + _current_forward_batch.reset(batch_token) + cls._current.reset(meta_token) + + @staticmethod + def to_intermediate_tensors( + intermediate_tensors, + metadata: Optional["SGLangForwardBatchMetadata"], + ): + if intermediate_tensors is not None or metadata is None: + return intermediate_tensors + pp_proxy_tensors = metadata.pp_proxy_tensors + if pp_proxy_tensors is None: + return intermediate_tensors + tensors = getattr(pp_proxy_tensors, "tensors", None) + if tensors is None: + return intermediate_tensors + from atom.models.utils import IntermediateTensors + + return IntermediateTensors(dict(tensors)) diff --git a/atom/plugin/sglang/runtime/forward_context.py b/atom/plugin/sglang/runtime/forward_context.py new file mode 100644 index 0000000000..05939cc7cb --- /dev/null +++ b/atom/plugin/sglang/runtime/forward_context.py @@ -0,0 +1,234 @@ +"""Scoped runtime adapter from SGLang batches to ATOM core.""" + +from __future__ import annotations + +import copy +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import Any, Optional + +import torch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +from atom.plugin.sglang.runtime.context import bind_current_forward_batch + + +def _is_dummy_forward(forward_batch: ForwardBatch) -> bool: + """Return whether an SGLang batch represents an empty/idle dummy run.""" + + forward_mode = getattr(forward_batch, "forward_mode", None) + return bool( + forward_mode is not None + and hasattr(forward_mode, "is_idle") + and forward_mode.is_idle() + ) + + +def _pad_dummy_like( + tensor: Optional[torch.Tensor], + *, + length: int, + fill_value: int | float = 0, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + shape = (length, *tensor.shape[1:]) + return torch.full(shape, fill_value, dtype=tensor.dtype, device=tensor.device) + + +def _materialize_atom_dummy_forward( + input_ids: Optional[torch.Tensor], + positions: Optional[torch.Tensor], + input_embeds: Optional[torch.Tensor], + forward_batch: ForwardBatch, +) -> tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + ForwardBatch, +]: + """Convert an empty SGLang IDLE batch into ATOM-style dummy inputs.""" + + if positions is None: + raise RuntimeError("SGLang dummy forward materialization requires positions") + if input_ids is None: + raise RuntimeError("SGLang dummy forward materialization requires input_ids") + + dummy_positions = positions.new_zeros((1,)) + dummy_input_ids = input_ids.new_zeros((1,)) + dummy_input_embeds = _pad_dummy_like(input_embeds, length=1, fill_value=0) + + model_forward_batch = copy.copy(forward_batch) + model_forward_batch.positions = dummy_positions + model_forward_batch.batch_size = 1 + model_forward_batch.seq_lens_sum = 1 + model_forward_batch.seq_lens = forward_batch.seq_lens.new_ones((1,)) + model_forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu.new_ones((1,)) + + return dummy_input_ids, dummy_positions, dummy_input_embeds, model_forward_batch + + +def _trim_hidden_states_for_output(hidden_states, num_tokens: int): + if torch.is_tensor(hidden_states): + return hidden_states[:num_tokens] + if isinstance(hidden_states, tuple): + return tuple( + tensor[:num_tokens] if torch.is_tensor(tensor) else tensor + for tensor in hidden_states + ) + return hidden_states + + +def _resolve_num_tokens_across_dp( + atom_config: Any, + forward_batch: ForwardBatch, + num_tokens: int, + is_dummy_run: bool, +) -> torch.Tensor: + """Resolve per-DP token counts for ATOM's CPU-side DPMetadata.""" + + global_num_tokens_cpu = getattr(forward_batch, "global_num_tokens_cpu", None) + if global_num_tokens_cpu is not None: + num_tokens_across_dp = torch.tensor( + global_num_tokens_cpu, dtype=torch.int32, device="cpu" + ) + else: + dp_size = atom_config.parallel_config.data_parallel_size + global_num_tokens_gpu = getattr(forward_batch, "global_num_tokens_gpu", None) + global_dp_buffer_len = getattr(forward_batch, "global_dp_buffer_len", None) + is_static_same_shape_batch = ( + global_num_tokens_gpu is not None + and global_dp_buffer_len == num_tokens * dp_size + ) + if not is_static_same_shape_batch: + raise RuntimeError( + "[SGL+ATOM] SGLang dp-attention requires " + "forward_batch.global_num_tokens_cpu unless the batch uses static " + "same-shape DP metadata." + ) + + # Static batches, such as CUDA graph capture batches, may only keep + # global token counts on GPU. Avoid GPU-to-CPU reads here and mirror + # their same-shape layout directly for ATOM's CPU DPMetadata. + num_tokens_across_dp = torch.full( + (dp_size,), num_tokens, dtype=torch.int32, device="cpu" + ) + + if is_dummy_run: + # SGLang reports idle ranks as 0 tokens, but ATOM materializes them + # as one local dummy token so collectives and DPMetadata stay aligned. + dp_rank = atom_config.parallel_config.data_parallel_rank + num_tokens_across_dp[dp_rank] = num_tokens + return num_tokens_across_dp + + +def _set_atom_forward_context( + atom_config: Any, + forward_batch: ForwardBatch, + positions: torch.Tensor, +) -> None: + """Bridge SGLang batch metadata into ATOM's global forward context.""" + + from atom.utils.forward_context import ( + AttentionMetaData, + Context, + set_forward_context, + ) + + forward_mode = forward_batch.forward_mode + # This value is only used by ATOM-side MoE padding in the SGLang wrapper. + max_seqlen_q = 1 if forward_mode.is_decode_or_idle() else 0 + attn_metadata = AttentionMetaData(max_seqlen_q=max_seqlen_q) + batch_size = int(forward_batch.batch_size) + is_dummy_run = _is_dummy_forward(forward_batch) + is_prefill = forward_mode.is_prefill() + num_tokens = int(positions.shape[0]) + + if bool(atom_config.enable_dp_attention): + num_tokens_across_dp = _resolve_num_tokens_across_dp( + atom_config, forward_batch, num_tokens, is_dummy_run + ) + graph_bs = int(torch.max(num_tokens_across_dp).item()) + else: + num_tokens_across_dp = None + graph_bs = num_tokens if is_prefill else batch_size + + context = Context( + positions=positions, + is_prefill=is_prefill, + is_dummy_run=is_dummy_run, + batch_size=batch_size, + graph_bs=graph_bs, + ) + set_forward_context( + attn_metadata=attn_metadata, + atom_config=atom_config, + context=context, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ) + + +def _reset_atom_forward_context() -> None: + from atom.utils.forward_context import reset_forward_context + + reset_forward_context() + + +@dataclass +class SGLangPluginRuntime: + """Scoped adapter for running ATOM model code under SGLang plugin runtime. + + The adapter owns the temporary translation from SGLang's ``ForwardBatch`` to + ATOM's process-local runtime state. Callers should use the normalized + ``input_ids``, ``positions``, ``input_embeds``, and ``forward_batch`` exposed + by this object while inside the context. + """ + + atom_config: Any + forward_batch: ForwardBatch + positions: torch.Tensor + input_ids: Optional[torch.Tensor] = None + input_embeds: Optional[torch.Tensor] = None + set_forward_context: bool = True + _original_forward_batch: ForwardBatch = field(init=False, repr=False) + _is_dummy_run: bool = field(init=False, default=False) + _exit_stack: ExitStack = field(init=False, repr=False) + + def __enter__(self) -> "SGLangPluginRuntime": + self._original_forward_batch = self.forward_batch + self._is_dummy_run = _is_dummy_forward(self.forward_batch) + + if self._is_dummy_run: + ( + self.input_ids, + self.positions, + self.input_embeds, + self.forward_batch, + ) = _materialize_atom_dummy_forward( + self.input_ids, + self.positions, + self.input_embeds, + self.forward_batch, + ) + + self._exit_stack = ExitStack() + self._exit_stack.enter_context(bind_current_forward_batch(self.forward_batch)) + if self.set_forward_context: + _set_atom_forward_context( + self.atom_config, + self.forward_batch, + self.positions, + ) + self._exit_stack.callback(_reset_atom_forward_context) + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self._exit_stack.close() + + def trim_output(self, hidden_states): + """Map ATOM-visible outputs back to SGLang-visible token count.""" + + if self._is_dummy_run: + return _trim_hidden_states_for_output(hidden_states, 0) + return hidden_states diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py new file mode 100644 index 0000000000..1ddd3772c9 --- /dev/null +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -0,0 +1,22 @@ +"""Runtime behavior flags for SGLang plugin model wrappers.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ModelArchSpec: + wrapper_binds_gdn_context: bool = False + apply_deepseek_patch: bool = False + + +MODEL_ARCH_SPECS = { + "DeepseekV3ForCausalLM": ModelArchSpec(apply_deepseek_patch=True), + "Qwen3MoeForCausalLM": ModelArchSpec(), + "Qwen3NextForCausalLM": ModelArchSpec(wrapper_binds_gdn_context=True), +} + + +def get_model_arch_spec(model_arch: str) -> ModelArchSpec: + return MODEL_ARCH_SPECS.get(model_arch, ModelArchSpec())