diff --git a/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py index 335e5aeba6..b4cae969bf 100644 --- a/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py +++ b/atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py @@ -11,12 +11,14 @@ # be handled by ATOM's native backend, making sglang-specific overrides # unnecessary. +import logging import math from typing import TYPE_CHECKING, Optional import torch import sglang.srt.layers.attention.aiter_backend as _sglang_aiter +from aiter.ops.triton.gluon.pa_decode_gluon import get_recommended_splits from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend from sglang.srt.layers.attention.utils import ( create_flashinfer_kv_indices_triton, @@ -26,6 +28,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import get_bool_env_var +from atom.model_ops.base_attention import run_pa_decode_gluon from atom.plugin.sglang.attention_backend.full_attention.kv_cache import ( set_kv_buffer_with_layout_shuffle as _set_kv_buffer_with_layout_shuffle, ) @@ -35,6 +38,9 @@ build_pa_metadata_for_decode as _build_pa_metadata_for_decode, build_pa_metadata_for_prefill as _build_pa_metadata_for_prefill, ) +from atom.utils import envs + +logger = logging.getLogger("atom.plugin.sglang.attention_backend.full_attention") if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -191,8 +197,14 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): self._init_draft_extend_v2_mla(forward_batch.batch_size, 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 (not self.use_mla) and forward_batch.forward_mode.is_draft_extend_v2(): + self._init_draft_extend_v2_mha(forward_batch.batch_size, forward_batch) + elif (not self.use_mla) and forward_batch.forward_mode.is_draft_extend(): + self._init_draft_extend_mha(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) + elif (not self.use_mla) and forward_batch.forward_mode.is_target_verify(): + self._init_target_verify_mha(forward_batch.batch_size, forward_batch) else: self._init_forward_metadata_extend(forward_batch) self._fixup_page_table(forward_batch) @@ -517,6 +529,66 @@ def _init_draft_extend_v2_mla(self, bs, forward_batch): run_graph=False, ) + def _init_draft_extend_v2_mha(self, bs, forward_batch): + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MHA draft_extend_v2 requires speculative metadata") + if forward_batch.extend_prefix_lens is None: + raise RuntimeError("MHA draft_extend_v2 requires extend_prefix_lens") + self.indices_updater_prefill.update( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + prefix_lens=forward_batch.extend_prefix_lens, + encoder_lens=forward_batch.encoder_lens, + spec_info=None, + ) + max_q_len = self._max_len( + forward_batch.extend_seq_lens_cpu, + forward_batch.extend_seq_lens, + ) + max_kv_len = self._max_len(forward_batch.seq_lens_cpu, forward_batch.seq_lens) + self.forward_metadata = ForwardMetadata( + self.indices_updater_prefill.kv_indptr, + self.indices_updater_prefill.kv_indices, + self.qo_indptr[: bs + 1], + None, + max_q_len, + max_kv_len, + None, + None, + custom_mask=None, + mask_indptr=None, + max_extend_len=max_q_len, + ) + def _init_draft_extend_mha(self, bs, forward_batch): + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MHA draft_extend requires speculative metadata") + kv_indices, kv_indptr, qo_indptr, custom_mask = ( + spec_info.generate_attn_arg_prefill( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + forward_batch.seq_lens_sum, + self.req_to_token, + ) + ) + kv_indices = kv_indices.to(torch.int64) + draft_max_extend_len = int(torch.max(spec_info.num_accept_tokens).item()) + self.forward_metadata = ForwardMetadata( + kv_indptr, + kv_indices, + qo_indptr, + None, + draft_max_extend_len, + None, + None, + None, + custom_mask=custom_mask, + mask_indptr=None, + max_extend_len=draft_max_extend_len, + ) + def _init_target_verify_mla(self, bs, forward_batch): """Init MLA metadata for speculative target_verify.""" spec_info = forward_batch.spec_info @@ -709,6 +781,55 @@ def _init_extend_mha(self, bs, forward_batch): forward_batch.seq_lens, ) + def _init_target_verify_mha(self, bs, forward_batch): + spec_info = forward_batch.spec_info + if spec_info is None: + raise RuntimeError("MHA target_verify requires speculative metadata") + + draft_num = spec_info.draft_token_num + qo_indptr = torch.arange( + 0, + (1 + bs) * draft_num, + step=draft_num, + dtype=torch.int32, + device=self.device, + ) + + kv_indptr = self.kv_indptr[: bs + 1] + kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0) + kv_indptr = kv_indptr[: bs + 1] + + kv_indices_len = int(kv_indptr[-1].item()) + kv_indices = torch.empty(kv_indices_len, dtype=torch.int64, device=self.device) + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, + kv_indices, + self.req_to_token.stride(0), + ) + + seq_mask_len = draft_num * (forward_batch.seq_lens + draft_num) + 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, + None, + draft_num, + None, + None, + None, + custom_mask=spec_info.custom_mask, + mask_indptr=mask_indptr, + max_extend_len=draft_num, + ) + def _fixup_page_table(self, forward_batch: ForwardBatch): """Post-process page_table for non-MLA extend mode.""" if ( @@ -764,6 +885,14 @@ def init_cuda_graph_state( else: self.cuda_graph_kv_indices = kv_indices_buf + draft_tokens = int(self.num_draft_tokens or 1) + custom_mask_size = max_bs * draft_tokens * (self.max_context_len + draft_tokens) + self.cuda_graph_custom_mask = torch.ones( + custom_mask_size, + dtype=torch.bool, + device=self.device, + ) + # Always use preshuffle layout for pa_fwd_asm self.page_table = torch.zeros( (max_bs, self.max_context_len // self.page_size), @@ -1584,8 +1713,9 @@ def _update_decode_page_table( def _should_use_native_dense_mha(self, layer) -> bool: sliding_window_size = getattr(layer, "sliding_window_size", None) is_minimax_m3 = bool(getattr(layer, "_atom_minimax_m3_dense_mha", False)) + is_eagle3 = bool(getattr(layer, "_atom_eagle3_dense_mha", False)) if ( - is_minimax_m3 + (is_minimax_m3 or is_eagle3) and not self.use_mla and not layer.is_cross_attention and layer.head_dim == 128 @@ -1740,25 +1870,16 @@ def forward_extend( 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( - layer.layer_id - ) - k_scale, v_scale = self._ensure_fp8_kv_scales( - layer, k_buffer, self.page_size - ) - self.set_kv_buffer_with_layout_shuffle( - cache_loc, - k, - v, - k_buffer, - v_buffer, - k_scale, - v_scale, - self.page_size, + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, cache_loc, k, v ) if self.use_mla: return self._forward_extend_mla(q, k, v, layer, forward_batch) + if forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend( + include_v2=True + ): + return self._forward_extend_mha_speculative(q, k, v, layer, forward_batch) if bool(getattr(layer, "_atom_minimax_m3_dense_mha", False)): # M3 dense decode benefits from the native ragged path, but batched # SGLang prefill is safer through the standard varlen extend path. @@ -1801,6 +1922,42 @@ def _forward_extend_native_dense_mha(self, q, layer, forward_batch): ) return o.view(-1, layer.tp_q_head_num * layer.head_dim) + def _forward_extend_mha_speculative(self, q, k, v, layer, forward_batch): + """Non-MLA EAGLE verify/draft-extend path. + + Mirrors /app/sglang's AiterAttnBackend: speculative MHA uses the + extend-attention kernel with custom masks, not plain flash_attn_varlen. + """ + + if k is None or v is None: + raise RuntimeError("MHA speculative extend requires explicit k/v tensors") + + if layer.qk_head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + self.extend_attention_fwd( + q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k.contiguous(), + v.contiguous(), + o.view(-1, layer.tp_q_head_num, layer.v_head_dim), + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id), + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id), + self.forward_metadata.qo_indptr, + self.forward_metadata.kv_indptr, + self.forward_metadata.kv_indices, + self.forward_metadata.custom_mask, + True, + self.forward_metadata.mask_indptr, + self.forward_metadata.max_extend_len, + 1.0, + 1.0, + layer.scaling, + logit_cap=layer.logit_cap, + ) + return o.view(-1, layer.tp_q_head_num * layer.v_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 @@ -2389,6 +2546,18 @@ def forward_decode( ) return self._forward_decode_native_dense_mha(q, layer, forward_batch) + if ( + envs.ATOM_FORCE_ATTN_TRITON + and not layer.is_cross_attention + and layer.qk_head_dim == layer.v_head_dim + and layer.head_dim == layer.qk_head_dim + ): + 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: @@ -2485,7 +2654,81 @@ 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) + k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer( + layer.layer_id + ) + if envs.ATOM_FORCE_ATTN_TRITON: + block_size = self.page_size + num_slots, num_kv_heads, head_size = k_cache.shape + num_blocks = num_slots // block_size + k_cache = k_cache[: num_blocks * block_size].view( + num_blocks, block_size, num_kv_heads, head_size + ) + v_cache = v_cache[: num_blocks * block_size].view( + num_blocks, block_size, num_kv_heads, layer.v_head_dim + ) + x = 16 // k_cache.element_size() + k_cache = ( + k_cache.view(num_blocks, block_size, num_kv_heads, head_size // x, x) + .permute(0, 2, 3, 1, 4) + .contiguous() + ) + v_cache = ( + v_cache.view(num_blocks, block_size // x, x, num_kv_heads, layer.v_head_dim) + .permute(0, 3, 1, 4, 2) + .contiguous() + ) + q_3d = q.view(-1, layer.tp_q_head_num, layer.head_dim) + out = torch.empty_like(q_3d, dtype=self.input_dtype) + num_seqs = self.forward_metadata.kv_lens.shape[0] + query_group_size = 1 * (layer.tp_q_head_num // layer.tp_k_head_num) + max_context_partition_num = get_recommended_splits( + num_seqs, layer.tp_k_head_num + ) + context_partition_size = 256 + intermediate_shape = ( + num_seqs, + layer.tp_k_head_num, + 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, + layer.head_dim, + dtype=q.dtype, + device=q.device, + ) + run_pa_decode_gluon( + output=out, + q=q_3d, + k_cache=k_cache, + v_cache=v_cache, + context_lens=self.forward_metadata.kv_lens, + block_tables=self.forward_metadata.page_table, + softmax_scale=layer.scaling, + max_seqlen_q=1, + max_context_partition_num=max_context_partition_num, + context_partition_size=context_partition_size, + compute_type=torch.bfloat16, + q_scale=None, + k_scale=None, + v_scale=None, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + alibi_slopes=None, + sinks=None, + sliding_window=-1, + ps=True, + ) + return out.reshape_as(q) + aiter_kv_str = self._get_aiter_paged_ragged_kv_cache_dtype() o = torch.empty_like(q, dtype=self.input_dtype) diff --git a/atom/plugin/sglang/attention_backend/full_attention/metadata.py b/atom/plugin/sglang/attention_backend/full_attention/metadata.py index b66feaa756..ce051a2e9f 100644 --- a/atom/plugin/sglang/attention_backend/full_attention/metadata.py +++ b/atom/plugin/sglang/attention_backend/full_attention/metadata.py @@ -29,6 +29,10 @@ class ForwardMetadata: fp8_prefill_kv_indices: Optional[torch.Tensor] = None num_kv_splits: Optional[int] = None run_graph: Optional[bool] = True + custom_mask: Optional[torch.Tensor] = None + mask_indptr: Optional[torch.Tensor] = None + max_extend_len: Optional[int] = None + swa_page_table: Optional[torch.Tensor] = None # 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 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 2e79ecd500..b3b1076ed7 100644 --- a/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py +++ b/atom/plugin/sglang/attention_backend/full_attention/radix_attention.py @@ -16,6 +16,7 @@ from atom.model_ops.base_attention import BaseAttention from atom.model_ops.layernorm import GemmaRMSNorm, fused_qk_norm from atom.model_ops.utils import atom_parameter +from atom.config import get_current_atom_config from atom.plugin.prepare import is_plugin_mode, is_sglang from atom.models.utils import maybe_prefix @@ -90,6 +91,34 @@ def __init__( v_head_dim=_v_head_dim, prefix=maybe_prefix(prefix, "attn"), ) + atom_config = get_current_atom_config() + hf_config = getattr(atom_config, "hf_config", None) + text_config = getattr(hf_config, "text_config", None) + model_type = getattr(hf_config, "model_type", "") + text_model_type = getattr(text_config, "model_type", "") + architectures = getattr(hf_config, "architectures", []) or [] + is_minimax_m3 = model_type == "minimax_m3_vl" or ( + text_model_type in {"minimax_m3", "minimax_m3_text", "minimax_m3_vl"} + ) + is_eagle3_llama = "LlamaForCausalLMEagle3" in architectures + if ( + is_minimax_m3 + and not use_mla + and rotary_emb is not None + and q_norm is not None + and k_norm is not None + and head_dim == 128 + and _v_head_dim == 128 + ): + self.attn._atom_minimax_m3_dense_mha = True + if ( + is_eagle3_llama + and not use_mla + and rotary_emb is not None + and head_dim == 128 + and _v_head_dim == 128 + ): + self.attn._atom_eagle3_dense_mha = True # sglang's RadixAttention expects k_scale/v_scale on device; # ensure they exist with identity scaling for non-quantised KV cache. # device="cuda" is safe here: this branch is guarded by is_sglang(), diff --git a/atom/plugin/sglang/attention_backend/minimax_m3_sparse.py b/atom/plugin/sglang/attention_backend/minimax_m3_sparse.py index baf9ae8df4..3e688015c7 100644 --- a/atom/plugin/sglang/attention_backend/minimax_m3_sparse.py +++ b/atom/plugin/sglang/attention_backend/minimax_m3_sparse.py @@ -51,12 +51,46 @@ def _slice_i32(tensor: torch.Tensor, batch_size: int) -> torch.Tensor: def _get_query_lens(forward_batch, batch_size: int) -> torch.Tensor: + if forward_batch.forward_mode.is_target_verify(): + spec_info = getattr(forward_batch, "spec_info", None) + if spec_info is None: + raise RuntimeError("MiniMax-M3 target_verify requires speculative metadata") + draft_num = int(spec_info.draft_token_num) + return torch.full( + (batch_size,), + draft_num, + dtype=torch.int32, + device=forward_batch.seq_lens.device, + ) query_lens = getattr(forward_batch, "extend_seq_lens", None) if query_lens is None: query_lens = getattr(forward_batch, "seq_lens") return _slice_i32(query_lens, batch_size) +def _get_seq_lens(forward_batch, batch_size: int) -> torch.Tensor: + seq_lens = _slice_i32(forward_batch.seq_lens, batch_size) + if forward_batch.forward_mode.is_target_verify(): + spec_info = getattr(forward_batch, "spec_info", None) + if spec_info is None: + raise RuntimeError("MiniMax-M3 target_verify requires speculative metadata") + seq_lens = seq_lens + int(spec_info.draft_token_num) + return seq_lens + + +def _get_max_seq_len(forward_batch, batch_size: int, seq_lens: torch.Tensor) -> int: + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + if seq_lens_cpu is not None and batch_size: + max_seq_len = int(seq_lens_cpu[:batch_size].max().item()) + if forward_batch.forward_mode.is_target_verify(): + spec_info = getattr(forward_batch, "spec_info", None) + if spec_info is None: + raise RuntimeError("MiniMax-M3 target_verify requires speculative metadata") + max_seq_len += int(spec_info.draft_token_num) + return max_seq_len + return int(seq_lens.max().item()) if batch_size else 0 + + def _get_prefix_lens( forward_batch, batch_size: int, @@ -607,21 +641,32 @@ def build_minimax_m3_block_table(forward_batch, page_size: int) -> torch.Tensor: if not forward_batch.forward_mode.is_decode_or_idle(): query_lens = _get_query_lens(forward_batch, batch_size) - seq_lens = _slice_i32(forward_batch.seq_lens, batch_size) + seq_lens = _get_seq_lens(forward_batch, batch_size) prefix_lens = _get_prefix_lens(forward_batch, batch_size, seq_lens, query_lens) out_cache_loc = forward_batch.out_cache_loc - offset = 0 - for req_idx in range(batch_size): - prefix_len = int(prefix_lens[req_idx].item()) - query_len = int(query_lens[req_idx].item()) - if query_len > 0: - token_table[req_idx, prefix_len : prefix_len + query_len] = ( - out_cache_loc[offset : offset + query_len] - ) - offset += query_len - - seq_lens = _slice_i32(forward_batch.seq_lens, batch_size) - if _is_stream_capturing() and forward_batch.forward_mode.is_decode_or_idle(): + if _is_stream_capturing(): + columns = torch.arange(token_table.shape[1], device=token_table.device) + rel_pos = columns.unsqueeze(0) - prefix_lens.unsqueeze(1) + mask = (rel_pos >= 0) & (rel_pos < query_lens.unsqueeze(1)) + query_offsets = torch.cumsum(query_lens, dim=0) - query_lens + src_idx = (query_offsets.unsqueeze(1) + rel_pos).clamp_min(0) + max_src_idx = max(int(out_cache_loc.numel()) - 1, 0) + src_idx = src_idx.clamp_max(max_src_idx).to(torch.long) + src_values = out_cache_loc.gather(0, src_idx.reshape(-1)).view_as(src_idx) + token_table = torch.where(mask, src_values, token_table) + else: + offset = 0 + for req_idx in range(batch_size): + prefix_len = int(prefix_lens[req_idx].item()) + query_len = int(query_lens[req_idx].item()) + if query_len > 0: + token_table[req_idx, prefix_len : prefix_len + query_len] = ( + out_cache_loc[offset : offset + query_len] + ) + offset += query_len + + seq_lens = _get_seq_lens(forward_batch, batch_size) + if _is_stream_capturing(): max_blocks = int(token_table.shape[1]) // page_size else: max_seq_len = int(seq_lens.max().item()) if batch_size else 0 @@ -639,11 +684,8 @@ def build_minimax_m3_forward_metadata( validate_minimax_m3_page_size(page_size) batch_size = _get_batch_size(forward_batch) - seq_lens = _slice_i32(forward_batch.seq_lens, batch_size) - if _is_stream_capturing() and forward_batch.forward_mode.is_decode_or_idle(): - max_seq_len = int(block_table.shape[1]) * page_size - else: - max_seq_len = int(seq_lens.max().item()) if batch_size else 0 + seq_lens = _get_seq_lens(forward_batch, batch_size) + max_seq_len = _get_max_seq_len(forward_batch, batch_size, seq_lens) if forward_batch.forward_mode.is_decode_or_idle(): return MiniMaxM3SGLangMetadata( @@ -655,16 +697,22 @@ def build_minimax_m3_forward_metadata( query_lens = _get_query_lens(forward_batch, batch_size) context_lens = _get_prefix_lens(forward_batch, batch_size, seq_lens, query_lens) - cu_seqlens_q = torch.empty( - batch_size + 1, dtype=torch.int32, device=seq_lens.device + cu_seqlens_q = torch.nn.functional.pad( + torch.cumsum(query_lens, dim=0, dtype=torch.int32), + (1, 0), ) - cu_seqlens_k = torch.empty( - batch_size + 1, dtype=torch.int32, device=seq_lens.device + cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seq_lens, dim=0, dtype=torch.int32), + (1, 0), ) - cu_seqlens_q[0] = 0 - cu_seqlens_k[0] = 0 - torch.cumsum(query_lens, dim=0, out=cu_seqlens_q[1:]) - torch.cumsum(seq_lens, dim=0, out=cu_seqlens_k[1:]) + + if _is_stream_capturing(): + if forward_batch.forward_mode.is_target_verify(): + max_query_len = int(getattr(forward_batch.spec_info, "draft_token_num", 1)) + else: + max_query_len = int(getattr(forward_batch, "extend_num_tokens", None) or 1) + else: + max_query_len = int(query_lens.max().item()) if batch_size else 0 return MiniMaxM3SGLangMetadata( is_decode=False, @@ -674,7 +722,7 @@ def build_minimax_m3_forward_metadata( cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, context_lens=context_lens, - max_query_len=int(query_lens.max().item()) if batch_size else 0, + max_query_len=max_query_len, ) @@ -825,7 +873,9 @@ def minimax_m3_sparse_attention_for_sglang( else: assert metadata.cu_seqlens_q is not None assert metadata.context_lens is not None - num_tokens = int(metadata.cu_seqlens_q[-1].item()) + num_tokens = int(q.shape[0]) if _is_stream_capturing() else int( + metadata.cu_seqlens_q[-1].item() + ) topk_idx = minimax_m3_index_topk( index_q[:num_tokens], index_cache, diff --git a/atom/plugin/sglang/models/base_model_wrapper.py b/atom/plugin/sglang/models/base_model_wrapper.py index 417b8be71b..f0f2317b2a 100644 --- a/atom/plugin/sglang/models/base_model_wrapper.py +++ b/atom/plugin/sglang/models/base_model_wrapper.py @@ -65,6 +65,11 @@ class _AtomCausalLMBaseForSglang(nn.Module): type that sglang expects. """ + # ATOM owns checkpoint quantization parsing, weight allocation, loading, and + # post-load processing for these external models. SGLang should not build + # its native quant_config from HF `quantization_config` first. + sglang_skip_quant_config = True + def __init__( self, config, @@ -81,6 +86,7 @@ def __init__( self.unpadded_vocab_size = config.vocab_size self.model_arch = getattr(config, "architectures", [""])[0] self.model_arch_spec = get_model_arch_spec(self.model_arch) + self.capture_aux_hidden_states = False with plugin_runtime_scope(framework="sglang"): from atom.config import get_current_atom_config @@ -102,7 +108,12 @@ def __init__( f"ATOM failed to create model for architecture {self.model_arch}" ) - if hasattr(self.model, "lm_head"): + if self.model_arch == "LlamaForCausalLMEagle3" and hasattr( + self.model, "compute_logits" + ): + self.logits_head = _ComputeLogitsHeadAdapter(self.model) + logits_head_handles_all_gather = True + elif hasattr(self.model, "lm_head"): self.logits_head = self.model.lm_head logits_head_handles_all_gather = False elif hasattr(self.model, "compute_logits"): @@ -123,6 +134,10 @@ def __init__( self.logits_processor = LogitsProcessor( config, skip_all_gather=plugin_skip_all_gather ) + self.load_lm_head_from_target = getattr( + self.model, "load_lm_head_from_target", False + ) + self.hot_token_id = getattr(self.model, "hot_token_id", None) # Apply model-specific install-time adapters (attn dispatch, weight hooks, etc.). if self.model_arch_spec.install_adapters is not None: @@ -150,46 +165,137 @@ def get_embed_and_head(self): if self.model_arch == "DeepseekV4ForCausalLM": return self.model.model.embed.weight, self.model.model.head.weight + embed_owner, head_owner = self._embed_and_head_owners() + return embed_owner.embed_tokens.weight, head_owner.lm_head.weight + + def _embed_and_head_owners(self): + if hasattr(self.model, "language_model"): + language_model = self.model.language_model + embed_owner = ( + language_model.model + if hasattr(language_model, "model") + and hasattr(language_model.model, "embed_tokens") + else language_model + ) + return embed_owner, language_model + 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 + return embed_owner, self.model 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) + if self.model_arch == "LlamaForCausalLMEagle3": + logger.info( + "Skip sharing target embed/lm_head for ATOM EAGLE3 draft; " + "the draft checkpoint owns independent weights." + ) + return None - embed_owner = ( - self.model.model - if hasattr(self.model, "model") - and hasattr(self.model.model, "embed_tokens") - else self.model - ) + embed_owner, head_owner = self._embed_and_head_owners() del embed_owner.embed_tokens.weight - del self.model.lm_head.weight + del head_owner.lm_head.weight embed_owner.embed_tokens.weight = embed - self.model.lm_head.weight = head + head_owner.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) + if self.model_arch == "LlamaForCausalLMEagle3": + logger.info( + "Skip sharing target embedding for ATOM EAGLE3 draft; " + "the draft checkpoint owns independent embedding." + ) + return None - embed_owner = ( - self.model.model - if hasattr(self.model, "model") - and hasattr(self.model.model, "embed_tokens") - else self.model - ) + embed_owner, _ = self._embed_and_head_owners() del embed_owner.embed_tokens.weight embed_owner.embed_tokens.weight = embed torch.cuda.empty_cache() torch.cuda.synchronize() + def set_eagle3_layers_to_capture(self, layer_ids: Optional[Iterable[int]] = None): + self.capture_aux_hidden_states = True + if layer_ids is None: + get_default_layers = getattr( + self.model, "get_eagle3_aux_hidden_state_layers", None + ) + if get_default_layers is None: + raise AttributeError( + f"ATOM model {type(self.model).__name__} does not define " + "get_eagle3_aux_hidden_state_layers" + ) + layer_ids = get_default_layers() + + layer_ids = tuple(int(layer_id) for layer_id in layer_ids) + if hasattr(self.model, "set_eagle3_layers_to_capture"): + return self.model.set_eagle3_layers_to_capture(layer_ids) + if hasattr(self.model, "set_aux_hidden_state_layers"): + return self.model.set_aux_hidden_state_layers(layer_ids) + raise AttributeError( + f"ATOM model {type(self.model).__name__} does not support " + "EAGLE3 auxiliary hidden-state capture" + ) + + def _split_aux_hidden_states(self, output): + if ( + isinstance(output, tuple) + and len(output) == 2 + and (torch.is_tensor(output[0]) or hasattr(output[0], "tensors")) + ): + return output[0], output[1] + return output, None + + def _trim_aux_hidden_states(self, runtime, aux_hidden_states): + if aux_hidden_states is None: + return None + if torch.is_tensor(aux_hidden_states): + return runtime.trim_output(aux_hidden_states) + if isinstance(aux_hidden_states, list): + return [ + runtime.trim_output(aux_hidden_state) + for aux_hidden_state in aux_hidden_states + ] + return aux_hidden_states + + def _forward_eagle3_draft_model(self, input_ids, positions, forward_batch): + spec_info = getattr(forward_batch, "spec_info", None) + hidden_states = getattr(spec_info, "hidden_states", None) + if hidden_states is None: + raise RuntimeError("EAGLE3 draft forward requires spec_info.hidden_states") + if hidden_states.shape[-1] != self.model.config.hidden_size: + hidden_states = self.model.combine_hidden_states(hidden_states) + return self.model( + input_ids=input_ids, positions=positions, hidden_states=hidden_states + ) + + def _load_eagle3_token_map(self, draft_path: str) -> None: + """Load SGLang EAGLE3 non-parameter token-map tensors in plugin scope.""" + from atom.model_loader.loader import safetensors_weights_iterator + + for name, weight_tensor in safetensors_weights_iterator(draft_path): + if "d2t" in name: + base = torch.arange( + weight_tensor.shape[0], + dtype=weight_tensor.dtype, + device=weight_tensor.device, + ) + self.hot_token_id = (weight_tensor + base).to(torch.int64) + setattr(self.model, "hot_token_id", self.hot_token_id) + logger.info( + "Loaded EAGLE3 draft token map %s with %d entries", + name, + self.hot_token_id.numel(), + ) + return + @torch.no_grad() def forward( self, @@ -244,7 +350,13 @@ def forward( ) with SGLangForwardBatchMetadata.bind(metadata): - if self.model_arch_spec.wrapper_binds_gdn_context: + if self.model_arch == "LlamaForCausalLMEagle3": + hidden_states = self._forward_eagle3_draft_model( + runtime.input_ids, + runtime.positions, + runtime.forward_batch, + ) + elif self.model_arch_spec.wrapper_binds_gdn_context: from atom.plugin.sglang.attention_backend.attention_gdn import ( SGLangGDNForwardContext, ) @@ -269,7 +381,13 @@ def forward( **self._filter_model_forward_kwargs(model_call_kwargs) ) + hidden_states, aux_hidden_states = self._split_aux_hidden_states( + hidden_states + ) hidden_states = runtime.trim_output(hidden_states) + aux_hidden_states = self._trim_aux_hidden_states( + runtime, aux_hidden_states + ) if self.pp_group.is_last_rank: if self.model_arch == "DeepseekV4ForCausalLM": @@ -278,6 +396,7 @@ def forward( hidden_states, self.logits_head, forward_batch, + aux_hidden_states=aux_hidden_states, hidden_states_before_norm=hidden_states, ) return self.logits_processor( @@ -285,6 +404,7 @@ def forward( hidden_states, self.logits_head, forward_batch, + aux_hidden_states=aux_hidden_states, ) return hidden_states @@ -293,6 +413,36 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # uses its own weight loading pipeline (handling AITER-specific quant # formats, kv_b_proj splitting, etc.) that is incompatible with # sglang's default weight iterator. + if self.model_arch == "LlamaForCausalLMEagle3": + from atom.model_loader.loader import load_model + + draft_path = None + try: + from sglang.srt.server_args import get_global_server_args + + server_args = get_global_server_args() + draft_path = getattr(server_args, "speculative_draft_model_path", None) + except Exception: + logger.exception("Failed to resolve SGLang EAGLE3 draft model path") + draft_path = draft_path or getattr(self.config, "_name_or_path", None) + draft_path = draft_path or getattr(self.config, "name_or_path", None) + if not draft_path: + raise RuntimeError("Cannot resolve EAGLE3 draft model path") + logger.info("Loading ATOM EAGLE3 draft weights from %s", draft_path) + self.atom_config.model = draft_path + self.atom_config.hf_config = self.config + with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): + result = load_model( + model=self.model, + model_name_or_path=draft_path, + hf_config=self.config, + load_dummy=self.atom_config.load_dummy, + prefix="", + is_plugin_mode=True, + ) + self._load_eagle3_token_map(draft_path) + return result + from atom.model_loader.loader import load_model_in_plugin_mode with plugin_runtime_scope(framework="sglang", atom_config=self.atom_config): diff --git a/atom/plugin/sglang/prepare.py b/atom/plugin/sglang/prepare.py index fd813628c1..939a60472c 100644 --- a/atom/plugin/sglang/prepare.py +++ b/atom/plugin/sglang/prepare.py @@ -47,8 +47,14 @@ def prepare_model(config: Any): set_attn_cls, ) - if model_arch not in _ATOM_SUPPORTED_MODELS: - supported_archs = list(_ATOM_SUPPORTED_MODELS.keys()) + supported_models = dict(_ATOM_SUPPORTED_MODELS) + if model_arch == "LlamaForCausalLMEagle3": + from atom.models.eagle3_llama import Eagle3LlamaModel + + supported_models[model_arch] = Eagle3LlamaModel + + if model_arch not in supported_models: + supported_archs = list(supported_models.keys()) raise ValueError( f"ATOM does not support the required model architecture: {model_arch}. " f"For now supported model architectures: {supported_archs}" @@ -57,8 +63,31 @@ def prepare_model(config: Any): from atom.plugin.config import generate_atom_config_for_plugin_mode atom_config = generate_atom_config_for_plugin_mode(config) + target_hf_config = getattr(atom_config, "hf_config", None) + target_hf_config = getattr(target_hf_config, "text_config", target_hf_config) + target_num_hidden_layers = int( + getattr(target_hf_config, "num_hidden_layers", 0) or 0 + ) + # `generate_atom_config_for_plugin_mode` is built from SGLang global + # ServerArgs, which always describe the target model. External draft + # runners (for example EAGLE3) pass their own HF config into this wrapper, + # so overwrite the model-local pieces here before constructing ATOM modules. + # Otherwise the draft modeling is built with target dimensions while + # SGLang allocates draft KV pools from the draft config. + atom_config.hf_config = config + model_path = getattr(config, "_name_or_path", None) or getattr( + config, "name_or_path", None + ) + if model_path: + atom_config.model = model_path + from atom.config import QuantizationConfig + + atom_config.quant_config = QuantizationConfig( + atom_config.hf_config, + online_quant_config=getattr(atom_config, "online_quant_config", None), + ) - model_cls = _ATOM_SUPPORTED_MODELS[model_arch] + model_cls = supported_models[model_arch] logger.info("ATOM model class for %s is %s", model_arch, model_cls) from atom.plugin.sglang.runtime import get_model_arch_spec @@ -83,8 +112,11 @@ def prepare_model(config: Any): apply_graph_capture_patch() init_params = inspect.signature(model_cls.__init__).parameters + init_kwargs = {} + if model_arch == "LlamaForCausalLMEagle3" and "layer_offset" in init_params: + init_kwargs["layer_offset"] = target_num_hidden_layers if "atom_config" in init_params: - model = model_cls(atom_config=atom_config) + model = model_cls(atom_config=atom_config, **init_kwargs) elif "config" in init_params: model = model_cls(config=atom_config) else: diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py new file mode 100644 index 0000000000..c6a222fde0 --- /dev/null +++ b/atom/plugin/sglang/register.py @@ -0,0 +1,89 @@ +import logging + +logger = logging.getLogger("atom.plugin.sglang.register") + + +def _is_atom_external_model_enabled() -> bool: + try: + from sglang.srt.environ import envs + + return envs.SGLANG_EXTERNAL_MODEL_PACKAGE.get() == "atom.plugin.sglang.models" + except Exception: + return False + + +def _hf_quant_method(model_config) -> str: + try: + quant_cfg = model_config._parse_quant_hf_config() + except Exception: + quant_cfg = None + if not quant_cfg: + return "" + return str(quant_cfg.get("quant_method", "")).lower() + + +def _install_model_config_quant_patch() -> None: + from sglang.srt.configs.model_config import ModelConfig + + if getattr(ModelConfig, "_atom_sglang_quant_patch", False): + return + + original_verify_quantization = ModelConfig._verify_quantization + + def verify_quantization_with_atom_external_bypass(self): + try: + return original_verify_quantization(self) + except ValueError as exc: + if ( + _is_atom_external_model_enabled() + and _hf_quant_method(self) == "mxfp8" + and "quantization is currently not supported in ROCm" in str(exc) + ): + logger.info( + "Skipping SGLang server-args quantization gate for ATOM " + "external MXFP8 model; ATOM owns quantized weight loading." + ) + self.quantization = None + return None + raise + + ModelConfig._verify_quantization = verify_quantization_with_atom_external_bypass + ModelConfig._atom_sglang_quant_patch = True + + +def _install_loader_quant_patch() -> None: + import sglang.srt.model_loader.loader as loader + + if getattr(loader, "_atom_sglang_quant_patch", False): + return + + original_get_quantization_config = loader._get_quantization_config + + def get_quantization_config_with_atom_external_bypass(model_config, load_config): + model_class, _ = loader.get_model_architecture(model_config) + if getattr(model_class, "sglang_skip_quant_config", False): + logger.info( + "Skipping SGLang native quant_config for external model %s; " + "the model wrapper owns quantized weight loading.", + model_class.__name__, + ) + return None + return original_get_quantization_config(model_config, load_config) + + loader._get_quantization_config = get_quantization_config_with_atom_external_bypass + loader._atom_sglang_quant_patch = True + + +def register_plugin() -> None: + """Install ATOM patches that must run before SGLang parses server args.""" + + _install_model_config_quant_patch() + _install_loader_quant_patch() + + try: + from atom.plugin.sglang.runtime import apply_load_config_patch + + apply_load_config_patch() + except Exception: + logger.exception("Failed to install ATOM SGLang load-config patch") + diff --git a/atom/plugin/sglang/runtime/model_arch.py b/atom/plugin/sglang/runtime/model_arch.py index f3620a57d5..1b1bc015d4 100644 --- a/atom/plugin/sglang/runtime/model_arch.py +++ b/atom/plugin/sglang/runtime/model_arch.py @@ -142,6 +142,9 @@ def _install_minimax_m3_adapters(model: Any) -> None: prepare_config=_prepare_minimax_m3_config, install_adapters=_install_minimax_m3_adapters, ), + "LlamaForCausalLMEagle3": SGLangModelAdapterSpec( + uses_context_only_forward=True, + ), } # Architectures whose SGLang EntryClass is generated by base_model_wrapper. @@ -159,6 +162,7 @@ def _install_minimax_m3_adapters(model: Any) -> None: "MiniMaxM2ForCausalLM", "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", + "LlamaForCausalLMEagle3", "DeepseekV4ForCausalLM", ) } diff --git a/pyproject.toml b/pyproject.toml index 127239ebf3..a41a4bb3c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,3 +42,6 @@ atom = "atom.plugin.vllm.register:register_platform" [project.entry-points."vllm.general_plugins"] atom_model_registry = "atom.plugin.vllm.register:register_model" + +[project.entry-points."sglang.srt.plugins"] +atom_sglang = "atom.plugin.sglang.register:register_plugin"