diff --git a/atom/config.py b/atom/config.py index 04f169bd23..52cf7652e6 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1292,6 +1292,13 @@ def compute_hash(self) -> str: ), ) ) + # Vocab-embedding replication (ATOM_REPLICATE_VOCAB_EMBED) changes both the + # embed weight shape ([vocab] vs [vocab/tp]) and the embed op (local + # F.embedding vs masked-embedding + all-reduce), so it alters the compiled + # graph and MUST be part of its key. Without this, toggling the flag — or + # deploying it on top of a cache built with the other setting — reuses a + # stale artifact and trips assert_size_stride at runtime. + factors.append(bool(envs.ATOM_REPLICATE_VOCAB_EMBED)) hash_str = hashlib.md5( str(factors).encode(), usedforsecurity=False diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ea846ebed3..3340e43a07 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -870,6 +870,9 @@ def _forward_prefill_mla( paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q paged_kv_indptr = attn_metadata.sparse_kv_indptr paged_kv_indices = self.sparse_kv_indices_buffer + # Sparse attention needs one last-page len per query token; the dense + # kv_last_page_lens (per-seq) would over-read -> illegal access. + kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens max_q_len = 1 if kv_c_and_k_pe_cache.numel() > 0: @@ -1047,8 +1050,14 @@ def _forward_decode( paged_kv_indices = self.sparse_kv_indices_buffer max_q_len = 1 else: + # Non-MTP sparse decode: KV is packed per token at + # page_size=1, so last_page_len is 1 for every seq. Use the + # all-1s sparse buffer, NOT the dense per-block + # kv_last_page_lens (which makes the asm kernel over-read + # past the written sparse-index region -> illegal access). paged_kv_indptr = attn_metadata.sparse_kv_indptr paged_kv_indices = self.sparse_kv_indices_buffer + paged_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens dp_size = get_dp_group().world_size use_persistent_mode = not (dp_size > 1) diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index b924ea0498..020457aa0a 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -8,6 +8,7 @@ import numpy as np import torch +import triton from atom.utils import envs from aiter import ( decode_update_mla_metadata_v1, @@ -20,6 +21,7 @@ from atom.utils import CpuGpuBuffer from atom.utils.block_convert import ( kv_indices_generate_triton, + mtp_prepare_decode_mla_kernel, ) from atom.utils.forward_context import AttentionMetaData, Context @@ -99,6 +101,11 @@ def get_impl_cls() -> Type["MLAAttention"]: class AiterMLAMetadataBuilder(CommonAttentionBuilder): + # EagleProposer folds the per-draft-step position/context bump into + # prepare_mtp_decode's fused kernel when this is set (matches the MHA + # backend). The fused kernel handles both sparse and dense MLA. + fuse_mtp_decode_position_update = True + def __init__(self, model_runner): if envs.ATOM_MLA_PAGE_SIZE > 1: self.block_size = envs.ATOM_MLA_PAGE_SIZE @@ -111,6 +118,10 @@ def __init__(self, model_runner): f"got --block-size {model_runner.block_size}" ) CommonAttentionBuilder.__init__(self, model_runner) + # Single-program block for the fused MTP-decode metadata kernel. Sized + # to the max batch (runtime bs <= max_bs) so one tl.cumsum spans the + # whole batch in a single launch. + self._mtp_fuse_block = triton.next_power_of_2(self.max_bs + 1) config = model_runner.config hf_config = config.hf_config # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. @@ -508,6 +519,7 @@ def set_mla_persistent_worker_buffers( max_q_len: int, only_update: bool = False, num_reject_tokens: torch.Tensor = None, + sparse_decode: bool = False, ): split_params = { "kv_granularity": max(self.block_size, 16), @@ -523,20 +535,41 @@ 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. + # This work metadata feeds sparse (DSA) attention when either: + # - max_q_len == 1: the plain single-token sparse decode, or + # - sparse_decode=True: the MTP draft (EagleProposer) whose single + # sparse block reuses these buffers but passes the target's original + # max_seqlen_qo (>1) through prepare_mtp_decode, so the max_q_len==1 + # test alone misses it. + # In both cases the KV is the per-token top-k selection, so the metadata + # must be built from sparse_kv_indptr; using the dense kv_indptr makes the + # asm kernel's kv_end run past sparse_kv_indptr[-1] into the stale region + # of the persistent sparse-index buffer once the context exceeds + # index_topk (dense >> sparse) -> illegal KV-cache access. + use_sparse_meta = self.is_sparse and (max_q_len == 1 or sparse_decode) kv_indptr_for_metadata = ( var["sparse_kv_indptr"].gpu[: bs + 1] - if self.is_sparse and max_q_len == 1 + if use_sparse_meta else var["kv_indptr"].gpu[: bs + 1] ) + # Sparse decode packs KV per query token at page_size=1, so every "page" + # is exactly one token -> last_page_len must be 1. The dense + # var["kv_last_page_lens"] holds the real last-BLOCK fill (1..block_size); + # feeding it here makes get_mla_metadata_v1 compute a per-seq KV extent of + # (sparse_count - 1 + dense_last_page_len), i.e. up to block_size-1 pages + # PAST the written sparse-index region -> stale-index over-read. Mirror + # kv_indptr_for_metadata (and the prefill/MTP-verify paths, which already + # use the all-1s sparse buffer). + kv_last_page_lens_for_metadata = ( + var["sparse_kv_last_page_lens"].gpu[:bs] + if use_sparse_meta + else var["kv_last_page_lens"].gpu[:bs] + ) if only_update: decode_update_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], kv_indptr_for_metadata, - var["kv_last_page_lens"].gpu[:bs], + kv_last_page_lens_for_metadata, self.padded_num_attention_heads, 1, # nhead_kv, True, @@ -557,7 +590,7 @@ def set_mla_persistent_worker_buffers( get_mla_metadata_v1( var["cu_seqlens_q"].gpu[: bs + 1], kv_indptr_for_metadata, - var["kv_last_page_lens"].gpu[:bs], + kv_last_page_lens_for_metadata, self.padded_num_attention_heads, 1, # nhead_kv, True, @@ -589,21 +622,53 @@ def prepare_mtp_decode( positions: torch.Tensor, # [total_tokens] int32 only_update: bool = False, num_reject_tokens: torch.Tensor = None, + *, + update_context_lens: bool = False, + positions_out: torch.Tensor | None = None, + last_token_indices: torch.Tensor | None = None, ): + """Per-draft-step MLA metadata update, fused into a single kernel. + + One ``_mtp_prepare_decode_mla_kernel`` launch performs, in place: + - ``kv_indptr += cu_seqlens_q`` (needed by kv_indices + slot_mapping), + - (sparse) per-seq ``min(kv_count, index_topk)`` cumsum -> + ``sparse_kv_indptr``, + - (fused position update) ``positions += 1`` when ``positions_out`` is + given, and ``context_lens += 1`` when ``update_context_lens`` is set. + + ``fuse_mtp_decode_position_update`` makes EagleProposer route the + per-step position/context bumps through here instead of launching them + as separate kernels. ``last_token_indices`` is accepted for signature + parity with the MHA backend but unused (MLA's ``positions`` is already + one entry per sequence at this point). + """ + del last_token_indices # MLA positions are already per-seq (1 per token) var = self.model_runner.forward_vars kv_indptr = var["kv_indptr"].gpu[: bs + 1] + cu_seqlens_q = var["cu_seqlens_q"].gpu[: bs + 1] if self.is_sparse: - # 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] + sparse_kv_indptr = None + + update_positions = positions_out is not None + context_lens = var["context_lens"].gpu[:bs] if update_context_lens else None + + mtp_prepare_decode_mla_kernel[(1,)]( + kv_indptr, + cu_seqlens_q, + sparse_kv_indptr if self.is_sparse else kv_indptr, + positions_out if update_positions else kv_indptr, + context_lens if update_context_lens else kv_indptr, + bs, + self.index_topk if self.is_sparse else 0, + positions_out.stride(0) if update_positions else 1, + IS_SPARSE=self.is_sparse, + UPDATE_POSITIONS=update_positions, + UPDATE_CONTEXT_LENS=update_context_lens, + BLOCK=self._mtp_fuse_block, + ) kv_indices_generate_triton( var["block_tables"].gpu[:bs], @@ -612,11 +677,33 @@ def prepare_mtp_decode( self.block_ratio, max_seqlen_k, ) - result = self.set_mla_persistent_worker_buffers( - bs, max_seqlen_q, only_update, num_reject_tokens - ) if self.is_sparse: + # The MTP draft's single sparse block reads sparse_kv_indptr, but it + # reuses the TARGET's work_info buffer, which was built dense. The + # incremental decode_update path cannot convert that dense work_info + # to sparse: it rebases each item's (dense) work_kv_len onto the new + # sparse seq_kv_end, driving kv_start negative and kv_end past the + # written sparse-index region -> illegal access. So do a FRESH sparse + # build (only_update=False) from sparse_kv_indptr instead. The draft + # emits exactly one query token per seq (cu_seqlens_q is an arange), + # so max_seqlen_qo must be 1 — passing the caller's max_seqlen_q (the + # target's verify width, e.g. 4) sets uni_seqlen_qo>1 while + # cu_seqlens_q says 1, which makes get_mla_metadata_v1 emit q ranges + # that run past the actual query rows. sparse_kv_indptr already + # reflects the reject-adjusted KV lengths, so num_reject_tokens is + # not needed here. + result = self.set_mla_persistent_worker_buffers( + bs, + 1, + only_update=False, + num_reject_tokens=None, + sparse_decode=True, + ) result["sparse_kv_indptr"] = sparse_kv_indptr + else: + result = self.set_mla_persistent_worker_buffers( + bs, max_seqlen_q, only_update, num_reject_tokens + ) return result def compute_block_bytes(self) -> int: @@ -814,9 +901,12 @@ def prepare_prefill(self, batch: ScheduledBatch): attn_metadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[ : sum_scheduled_tokens + 1 ] - attn_metadata.kv_last_page_lens = var["sparse_kv_last_page_lens"].gpu[ - :sum_scheduled_tokens - ] + # Sparse (DSA) attention: one last-page len per query token (all 1s, + # page_size=1). Lives only on sparse_kv_last_page_lens; kv_last_page_lens + # stays the dense per-seq buffer set by the has_cached block below. + attn_metadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:sum_scheduled_tokens] # Per-query req_id: token_id 0..sum_scheduled_tokens-1 maps to batch id. # Use counts (new tokens per batch), not context_lens (full seq len). @@ -835,7 +925,7 @@ def prepare_prefill(self, batch: ScheduledBatch): get_mla_metadata_v1( attn_metadata.sparse_cu_seqlens_q, attn_metadata.sparse_kv_indptr, - attn_metadata.kv_last_page_lens, + attn_metadata.sparse_kv_last_page_lens, self.padded_num_attention_heads, 1, # nhead_kv True, @@ -1191,6 +1281,15 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): ).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] + elif self.is_sparse: + # Non-MTP sparse decode (single token per seq): the sparse KV is + # packed at page_size=1, so last_page_len is 1 for every seq. Expose + # the all-1s buffer so _forward_decode passes it to mla_decode_fwd + # instead of the dense per-block kv_last_page_lens (which would make + # the kernel over-read past the written sparse-index region). + attn_metadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:bs] # 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: @@ -1325,13 +1424,19 @@ def _set_ubatch_mla_buffers(self, padded_bs, max_q_len, ubatch_idx): var = self.model_runner.forward_vars kv_indptr_for_mla = var[f"{p}kv_indptr"].gpu[: padded_bs + 1] + kv_last_page_lens_for_mla = var[f"{p}kv_last_page_lens"].gpu[:padded_bs] if self.is_sparse: kv_indptr_for_mla = var[f"{p}sparse_kv_indptr"].gpu[: padded_bs + 1] + # Sparse KV is packed per token at page_size=1 -> last_page_len is 1. + # The dense per-block buffer would over-read past the sparse indices + # (see set_mla_persistent_worker_buffers). The all-1s sparse buffer is + # batch-independent, so the shared (non-ubatch) copy is safe here. + kv_last_page_lens_for_mla = var["sparse_kv_last_page_lens"].gpu[:padded_bs] get_mla_metadata_v1( var[f"{p}cu_seqlens_q"].gpu[: padded_bs + 1], kv_indptr_for_mla, - var[f"{p}kv_last_page_lens"].gpu[:padded_bs], + kv_last_page_lens_for_mla, self.padded_num_attention_heads, 1, # nhead_kv True, @@ -1412,6 +1517,12 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData: 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] + elif self.is_sparse: + # Non-MTP sparse decode capture: all-1s per-token last-page lens, + # matching prepare_decode so _forward_decode reads the sparse buffer. + attn_matadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_lens" + ].gpu[:bs] positions = var["positions"].copy_to_gpu(sum_tokens) context = Context( positions=positions, is_prefill=False, batch_size=bs, graph_bs=bs @@ -1445,6 +1556,11 @@ def build_ubatch_metadata( if self.is_sparse else None ), + sparse_kv_last_page_lens=( + var["sparse_kv_last_page_lens"].gpu[:padded_bs] + if self.is_sparse + else None + ), work_meta_data=var[f"{p}work_meta_data"], work_info_set=var[f"{p}work_info_set"], work_indptr=var[f"{p}work_indptr"], @@ -1506,11 +1622,16 @@ def build_ubatch_prefill_metadata( if attn_metadata.slot_mapping is not None else 0 ) + # Sparse prefill: sparse_kv_last_page_lens is per query TOKEN, so slice it + # by the token slice. (The dense kv_last_page_lens is per-seq and is sliced + # by request in split_attn_metadata.) if ( - attn_metadata.kv_last_page_lens is not None - and attn_metadata.kv_last_page_lens.shape[0] == total_tokens + attn_metadata.sparse_kv_last_page_lens is not None + and attn_metadata.sparse_kv_last_page_lens.shape[0] == total_tokens ): - ub_attn.kv_last_page_lens = attn_metadata.kv_last_page_lens[ts] + ub_attn.sparse_kv_last_page_lens = attn_metadata.sparse_kv_last_page_lens[ + ts + ] if ( attn_metadata.sparse_kv_indptr is not None diff --git a/atom/model_ops/embed_head.py b/atom/model_ops/embed_head.py index 15a8505c97..749c952604 100644 --- a/atom/model_ops/embed_head.py +++ b/atom/model_ops/embed_head.py @@ -103,6 +103,37 @@ def masked_embedding( return _masked_embedding_launcher(x, weight, vocab_start_idx, vocab_end_idx) +def _replicated_embedding_fake(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.empty( + x.numel(), + weight.shape[1], + dtype=weight.dtype, + device=weight.device, + ) + + +@torch_compile_guard(gen_fake=_replicated_embedding_fake) +def replicated_embedding(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + # Keep the lookup opaque to torch.compile: inductor otherwise fuses the + # embedding gather into the surrounding graph, which corrupts the MTP draft + # rollout (acceptance collapses ~69%->45%) — the same reason + # VocabParallelEmbedding routes through the masked_embedding custom op. + # + # Route through the masked kernel with the full-table range [0, num_rows) so + # out-of-range ids never reach a raw gather. Under async scheduling + MTP + # spec-decode, input_ids can transiently carry the optimistic placeholder + # token -1 (an unresolved "assumed-accepted" draft/bonus slot, produced in + # gpu_model_runner and read back via prepare_next_token_ids_padded's backup + # before the deferred correction lands) — for BOTH the target and the shared + # draft embedding. A raw F.embedding(-1) reads the row before the table -> + # random illegal memory access. The masked load returns a zero vector for any + # out-of-range id: bit-identical to F.embedding for every valid token, and + # matching vLLM's VocabParallelEmbedding (which masks the same -1 to 0) so the + # unverified -1 slots — whose output is discarded/corrected by async + # spec-decode — see the same value native does. No accuracy change. + return _masked_embedding_launcher(x, weight, 0, weight.shape[0]) + + class VocabParallelEmbedding(nn.Module): def __init__( @@ -184,7 +215,7 @@ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor): - return F.embedding(x, self.weight) + return replicated_embedding(x, self.weight) class ParallelLMHead(VocabParallelEmbedding): diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 9ccae7d3e3..c030599353 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -1099,6 +1099,109 @@ def __call__( ) +# --------------------------------------------------------------------------- +# Fused dual RMSNorm + concat — single Triton launch. +# +# Two independent RMSNorms over the SAME hidden dim (different inputs, different +# weights, shared eps) whose bf16 results are concatenated on the last axis: +# out = cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1) # [M, 2H] +# One program per row does both norms and writes each into its half of the +# [M, 2H] output, so the concat is free — it eliminates the separate cat +# kernel's read+write of 2*M*H (halving traffic vs two norms + torch.cat) and +# folds three launches (enorm, hnorm, cat) into one. Used by the DeepSeek/GLM +# MTP eh_proj input (enorm(embed) ++ hnorm(prev_hidden)). +# --------------------------------------------------------------------------- + + +@triton.jit +def _fused_dual_rmsnorm_cat_kernel( + xe_ptr, + xh_ptr, + we_ptr, + wh_ptr, + out_ptr, + H, + stride_xe_m, + stride_xh_m, + stride_out_m, + eps, + BLOCK_H: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_H) + mask = cols < H + + # enorm half -> out[row, :H] + xe = tl.load(xe_ptr + row * stride_xe_m + cols, mask=mask, other=0.0).to(tl.float32) + we = tl.load(we_ptr + cols, mask=mask, other=0.0).to(tl.float32) + rstd_e = tl.rsqrt(tl.sum(xe * xe, axis=-1) / H + eps) + out_e = (xe * rstd_e * we).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out_m + cols, out_e, mask=mask) + + # hnorm half -> out[row, H:2H] + xh = tl.load(xh_ptr + row * stride_xh_m + cols, mask=mask, other=0.0).to(tl.float32) + wh = tl.load(wh_ptr + cols, mask=mask, other=0.0).to(tl.float32) + rstd_h = tl.rsqrt(tl.sum(xh * xh, axis=-1) / H + eps) + out_h = (xh * rstd_h * wh).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out_m + H + cols, out_h, mask=mask) + + +def _fused_dual_rmsnorm_cat_fake( + xe: torch.Tensor, + we: torch.Tensor, + xh: torch.Tensor, + wh: torch.Tensor, + eps: float, +) -> torch.Tensor: + M, H = xe.shape + return torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device) + + +@torch_compile_guard(gen_fake=_fused_dual_rmsnorm_cat_fake) +def fused_dual_rmsnorm_cat( + xe: torch.Tensor, + we: torch.Tensor, + xh: torch.Tensor, + wh: torch.Tensor, + eps: float, +) -> torch.Tensor: + """cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1) in one Triton launch. + + Args: + xe, xh: (M, H) bf16/fp16 inputs (same shape). + we, wh: (H,) RMSNorm weights (one per input). + eps: shared RMSNorm epsilon. + Returns: + (M, 2H) tensor: [:, :H] = rmsnorm(xe, we), [:, H:] = rmsnorm(xh, wh). + """ + assert xe.shape == xh.shape, f"shape mismatch {xe.shape} vs {xh.shape}" + assert xe.dim() == 2, f"expected 2-D inputs, got {xe.dim()}-D" + xe = xe.contiguous() + xh = xh.contiguous() + M, H = xe.shape + out = torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device) + if M == 0: + return out + BLOCK_H = triton.next_power_of_2(H) + num_warps = 8 if BLOCK_H >= 4096 else 4 + _fused_dual_rmsnorm_cat_kernel[(M,)]( + xe, + xh, + we, + wh, + out, + H, + xe.stride(0), + xh.stride(0), + out.stride(0), + eps, + BLOCK_H=BLOCK_H, + num_warps=num_warps, + num_stages=2, + ) + return out + + @torch_compile_guard() def layernorm2d_fwd_( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float, dim: int diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index dff057eb8e..ff6f993236 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -6,8 +6,12 @@ import torch.nn as nn from aiter.dist.communication_op import tensor_model_parallel_all_reduce 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.embed_head import ( + ParallelLMHead, + ReplicatedEmbedding, + VocabParallelEmbedding, +) +from atom.model_ops.layernorm import RMSNorm, fused_dual_rmsnorm_cat from atom.model_ops.linear import ReplicatedLinear from atom.model_ops.moe import FusedMoE from atom.models.utils import IntermediateTensors @@ -15,7 +19,11 @@ from atom.utils.decorators import support_torch_compile from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig -from .deepseek_v2 import DeepseekV2DecoderLayer, _can_fuse_indexer_wk_weights_proj +from .deepseek_v2 import ( + DeepseekV2DecoderLayer, + _can_fuse_indexer_wk_weights_proj, + use_replicated_vocab_embed, +) from .utils import ckpt_has_tensor_suffix, maybe_prefix @@ -87,13 +95,17 @@ def forward( 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) + # Fused enorm(inputs_embeds) ++ hnorm(previous_hidden_states) in a single + # Triton launch (folds the two RMSNorms + the torch.cat; enorm and hnorm + # share eps=rms_norm_eps). bf16-identical to the separate path. + eh_input = fused_dual_rmsnorm_cat( + inputs_embeds, + self.enorm.weight, + previous_hidden_states, + self.hnorm.weight, + self.enorm.eps, ) + hidden_states = self.eh_proj(eh_input) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None @@ -136,10 +148,18 @@ def __init__( ) } ) - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) + if use_replicated_vocab_embed(config): + # GLM-5.2 MTP: full table per rank, no post-embedding all-reduce. + # (Shared with the target's replicated embed by EagleProposer at load.) + self.embed_tokens = ReplicatedEmbedding( + config.vocab_size, + config.hidden_size, + ) + else: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) def forward( self, @@ -170,6 +190,24 @@ def compute_logits( logits = mtp_layer.shared_head.head(mtp_layer.shared_head(hidden_states)) return logits + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Greedy draft token via distributed argmax over the TP-sharded vocab — + avoids all-gathering the full [N, vocab] logits every draft step. + + Mirrors compute_logits() (same norm + shared head), but reduces each + rank's logit shard to (max_val, global_idx) and all-gathers only [N, 2] + instead of the O(vocab) logits. Token-identical to + compute_logits(...).argmax(-1). + """ + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + normed = mtp_layer.shared_head(hidden_states) + return mtp_layer.shared_head.head.compute_argmax_token(normed) + @support_torch_compile class DeepSeekMTP(nn.Module): @@ -259,6 +297,20 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Distributed greedy argmax for the MTP draft rollout (GLM-5.2). + + EagleProposer picks this over compute_logits().argmax(-1) when present + (``_draft_argmax_fused``), so the draft never all-gathers the full + [N, vocab] logits — it all-gathers only the packed [N, 2] per-rank + reductions. See DeepSeekMultiTokenPredictor.compute_draft_token. + """ + return self.model.compute_draft_token(hidden_states, spec_step_idx) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 4ec791444f..6a363ac8fa 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -60,7 +60,11 @@ triton_gather_kv_indices_sparse, ) from atom.model_ops.base_attention import Attention -from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.embed_head import ( + ParallelLMHead, + ReplicatedEmbedding, + VocabParallelEmbedding, +) from atom.model_ops.layernorm import LayerNorm, RMSNorm from atom.model_ops.linear import ( ColumnParallelLinear, @@ -2426,6 +2430,39 @@ def forward( return hidden_states, residual +def use_replicated_vocab_embed(config: PretrainedConfig) -> bool: + """Whether to hold the full vocab embedding on every TP rank (local lookup, + no post-embedding all-reduce) instead of a ``VocabParallelEmbedding`` shard. + + Enabled by default (gated by ``ATOM_REPLICATE_VOCAB_EMBED``) for GLM-5.2 + (``glm_moe_dsa``) — both the main model and its MTP draft — whose embedding is + independent of the still TP-sharded ``lm_head`` (``tie_word_embeddings=False``), + so the lookup is bit-identical to the sharded masked-embedding + all-reduce + path. + + Enabled under the **vLLM** plugin as well: its MTP proposer unconditionally + shares the *target* model's ``embed_tokens`` into the draft + (``llm_base_proposer._maybe_share_embeddings``), so replicating the main + model's table also removes the per-step all-reduce from the draft rollout — + and the plugin loader honours each param's ``weight_loader`` so every rank + loads the full (un-sharded) table. Left on the sharded path for the SGLang/RTP + plugins (their embedding lifecycle is not verified here) and whenever the + embedding is tied to the sharded head. + + The GLM-5.2 main model keeps ``model_type == "glm_moe_dsa"``; its MTP draft + config has ``model_type`` rewritten to ``"deepseek_mtp"`` (see + ``SpeculativeConfig.hf_config_override``) but still carries the GLM-only + ``index_share_for_mtp_iteration`` flag, so we detect either. + """ + if not envs.ATOM_REPLICATE_VOCAB_EMBED: + return False + if getattr(config, "tie_word_embeddings", False): + return False + return getattr(config, "model_type", None) == "glm_moe_dsa" or bool( + getattr(config, "index_share_for_mtp_iteration", False) + ) + + @support_torch_compile class DeepseekV2Model(nn.Module): def __init__( @@ -2446,10 +2483,22 @@ def __init__( self.is_v32 = hasattr(config, "index_topk") if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) + if use_replicated_vocab_embed(config): + # GLM-5.2: full table per rank, no post-embedding all-reduce. + self.embed_tokens = ReplicatedEmbedding( + config.vocab_size, + config.hidden_size, + ) + logger.info( + "vocab embedding: REPLICATED (full %d-row table per rank, " + "no post-embed all-reduce)", + config.vocab_size, + ) + else: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) else: self.embed_tokens = PPMissingLayer() diff --git a/atom/models/eagle3_llama.py b/atom/models/eagle3_llama.py index 0aac1f2f67..3f8380fb42 100644 --- a/atom/models/eagle3_llama.py +++ b/atom/models/eagle3_llama.py @@ -340,8 +340,8 @@ def __init__(self, atom_config: Config, prefix: str = "", layer_offset: int = 0) # full on every rank — a local lookup with no post-embedding all-reduce. # Bit-identical to the sharded path; on by default (trades memory for one # fewer collective per draft step). Falls back to the sharded embedding - # when ATOM_EAGLE_REPLICATE_EMBED=0. - if envs.ATOM_EAGLE_REPLICATE_EMBED: + # when ATOM_REPLICATE_VOCAB_EMBED=0. + if envs.ATOM_REPLICATE_VOCAB_EMBED: self.embed_tokens = ReplicatedEmbedding( config.vocab_size, config.hidden_size ) diff --git a/atom/plugin/vllm/attention/layer_sparse_mla.py b/atom/plugin/vllm/attention/layer_sparse_mla.py index 5f02dae5d4..704c27e4b3 100644 --- a/atom/plugin/vllm/attention/layer_sparse_mla.py +++ b/atom/plugin/vllm/attention/layer_sparse_mla.py @@ -25,6 +25,7 @@ from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits from atom.plugin.prepare import is_vllm +from atom.utils import envs from atom.utils.custom_register import direct_register_custom_op import triton @@ -35,6 +36,15 @@ logger = logging.getLogger("atom") +# Byte budget (MB) for the dense fp8_mqa_logits prefill logits buffer. Mirrors +# native ATOM (atom/models/deepseek_v4.py, issue #1376): the indexer logits +# matrix is [rows, total_seq_lens] fp32 and total_seq_lens (the sum of all +# co-scheduled prefill contexts) is unbounded by max_num_batched_tokens, so a +# burst of long-context requests can push a single allocation to tens of GiB +# and OOM the engine. Chunking along the Q-row dimension keeps the buffer within +# this budget. 0 disables chunking (always single-shot). +_SPARSE_INDEXER_LOGITS_BUDGET_MB = envs.ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB + @triton.jit def _convert_req_index_to_global_index_kernel( @@ -338,6 +348,8 @@ def sparse_attn_indexer_plugin_mode( if has_prefill: prefill_metadata = indexer_meta.prefill + assert topk_tokens == 2048, "top_k_per_row assumes size 2048" + budget_bytes = _SPARSE_INDEXER_LOGITS_BUDGET_MB * 1024 * 1024 for chunk in prefill_metadata.chunks: k_fp8 = torch.empty( [chunk.total_seq_lens, head_dim], @@ -359,32 +371,67 @@ def sparse_attn_indexer_plugin_mode( preshuffle=preshuffle_cache, ) - logits = fp8_mqa_logits( - Q=q_fp8[chunk.token_start : chunk.token_end], - KV=k_fp8, - kv_scales=k_scale, - weights=weights[chunk.token_start : chunk.token_end], - cu_starts=chunk.cu_seqlen_ks, - cu_ends=chunk.cu_seqlen_ke, - ) - num_rows = logits.shape[0] - assert topk_tokens == 2048, "top_k_per_row assumes size 2048" - topk_indices_prefill = topk_indices[ - chunk.token_start : chunk.token_end, :topk_tokens - ] - # Use top_k_per_row_prefill from vLLM to correctly handle row starts - # and ends. It also produces 0-based local indices, eliminating the - # need for conversion from global. - torch.ops._C.top_k_per_row_prefill( - logits, - chunk.cu_seqlen_ks, - chunk.cu_seqlen_ke, - topk_indices_prefill, - num_rows, - logits.stride(0), - logits.stride(1), - topk_tokens, - ) + # The dense fp8_mqa_logits buffer is [rows, total_committed] fp32. + # total_committed (the column dim = sum of co-scheduled prefill + # contexts) is unbounded by max_num_batched_tokens, so a burst of + # long-context requests can push a single allocation to tens of GiB + # (#1376). Chunk along the Q (row) dimension so [row_chunk, + # total_committed] fp32 stays within budget_bytes — row_chunk shrinks + # as total_committed grows. Each row chunk still scores the FULL KV, + # so every row's top-k is exact with no cross-chunk merge and the + # kernel's per-row column indices need no remapping. + total_committed = int(chunk.total_seq_lens) + total_rows = chunk.token_end - chunk.token_start + if ( + budget_bytes > 0 + and total_committed > 0 + and budget_bytes // (total_committed * 4) < total_rows + ): + # 4 bytes per fp32 logit; total_committed * 4 is one row's + # footprint. Round the budget-derived row count DOWN to a + # multiple of 128 (aligned to the kernel's row tiling); when the + # budget affords < 128 rows (extreme total_committed), fall back + # to a power-of-2 floor so it degrades 64/32/.../1 instead of + # collapsing straight to 1. + budget_rows = budget_bytes // (total_committed * 4) + if budget_rows >= 128: + row_chunk = (budget_rows // 128) * 128 + else: + row_chunk = 1 << (max(1, budget_rows).bit_length() - 1) + else: + row_chunk = total_rows + + for row_start in range(0, total_rows, row_chunk): + row_end = min(row_start + row_chunk, total_rows) + q_start = chunk.token_start + row_start + q_end = chunk.token_start + row_end + # cu_seqlen_ks/ke are per-row (length total_rows); slice 1:1 with + # this row chunk. KV stays full so per-row windows are unchanged. + row_ks = chunk.cu_seqlen_ks[row_start:row_end] + row_ke = chunk.cu_seqlen_ke[row_start:row_end] + logits = fp8_mqa_logits( + Q=q_fp8[q_start:q_end], + KV=k_fp8, + kv_scales=k_scale, + weights=weights[q_start:q_end], + cu_starts=row_ks, + cu_ends=row_ke, + ) + num_rows = logits.shape[0] + topk_indices_prefill = topk_indices[q_start:q_end, :topk_tokens] + # Use top_k_per_row_prefill from vLLM to correctly handle row + # starts and ends. It also produces 0-based local indices, + # eliminating the need for conversion from global. + torch.ops._C.top_k_per_row_prefill( + logits, + row_ks, + row_ke, + topk_indices_prefill, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) if has_decode: decode_metadata = indexer_meta.decode diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 1339638839..d3ef4cb0da 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/atom/plugin/vllm/attention/metadata.py @@ -1810,13 +1810,27 @@ def _resolve_indexer(layer_name): ) return shared_buffer - def build(self, common_prefix_len, common_attn_metadata, fast_build=False): + def build( + self, + common_prefix_len, + common_attn_metadata, + fast_build=False, + *, + _req_id_per_token=None, + _drafting=False, + ): num_tokens = common_attn_metadata.num_actual_tokens - starts = common_attn_metadata.query_start_loc_cpu.to(torch.int32) - seg_lengths = torch.diff(starts) - req_id_per_token = torch.repeat_interleave( - torch.arange(seg_lengths.shape[0], dtype=torch.int32), seg_lengths - ) + if _req_id_per_token is not None: + # Draft path (see build_for_drafting): req_id_per_token was already + # produced on-GPU, so the copy_ below is device-to-device and the + # per-step pageable H2D sync is avoided. + req_id_per_token = _req_id_per_token + else: + starts = common_attn_metadata.query_start_loc_cpu.to(torch.int32) + seg_lengths = torch.diff(starts) + req_id_per_token = torch.repeat_interleave( + torch.arange(seg_lengths.shape[0], dtype=torch.int32), seg_lengths + ) # Shrink-tail-only zeroing instead of three full-buffer fill_(0) every # step (the buffers are persistent across decode steps and zeros-init): # - req_id_per_token_buffer / paged_kv_indices: the kernel only reads @@ -1899,7 +1913,9 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): # adds no extra CPU work. num_reqs = common_attn_metadata.num_reqs decode_only = ( - int(common_attn_metadata.max_query_len) == 1 and num_tokens == num_reqs + not _drafting + and int(common_attn_metadata.max_query_len) == 1 + and num_tokens == num_reqs ) metadata_key = None if decode_only: @@ -1970,6 +1986,45 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): return attn_metadata + def build_for_drafting(self, common_attn_metadata, draft_index): + """Sync-free per-draft-step build for sparse MLA main attention. + + vLLM's EagleProposer rebuilds attention metadata for every MTP draft + step. The inherited build_for_drafting just calls build(), which + constructs ``req_id_per_token`` on the host and copies it into a GPU + buffer -- a pageable, synchronous H2D copy that stalls the draft + model's kernel dispatch right at the target->draft boundary (this is + the H2D sync visible between the main and draft models in the trace). + + During MTP/EAGLE drafting every request contributes exactly one decode + token, so ``req_id_per_token`` is simply ``arange(num_reqs)`` and can be + built on-GPU (device-to-device copy, no sync). We also pass + ``_drafting=True`` so build() skips the decode-only work-split + fingerprint (whose ``seq_lens.cpu()`` fallback would be another sync and + which MTP/spec batches invalidate anyway). Non-uniform batches (any + prefill or padding) fall back to the regular build(). + """ + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + if not ( + num_tokens == num_reqs and int(common_attn_metadata.max_query_len) == 1 + ): + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + fast_build=True, + ) + req_id_per_token = torch.arange( + num_tokens, dtype=torch.int32, device=self.device + ) + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + fast_build=True, + _req_id_per_token=req_id_per_token, + _drafting=True, + ) + class AiterMlaSparseIndexerMetadataBuilder(AttentionMetadataBuilder): _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH @@ -2037,8 +2092,20 @@ def __init__( dtype=torch.int32, device=self.device, ) + # ``arange_buffer`` is sliced to ``actual_expanded`` (the flattened + # multi-token decode count) in ``_build_indexer``. During MTP drafting a + # decode request carries up to ``1 + num_speculative_tokens`` tokens, so + # the expansion can reach ``num_decode_tokens`` (bounded by + # ``max_num_batched_tokens``) -- larger than ``max_num_seqs`` for a draft + # builder whose ``num_speculative_tokens`` is forced to 0. Guard with + # ``max_num_batched_tokens`` exactly as upstream vLLM's indexer does so + # the buffer always covers ``actual_expanded``. self.arange_buffer = torch.arange( - config.scheduler_config.max_num_seqs * (1 + self.num_speculative_tokens), + max( + config.scheduler_config.max_num_seqs + * (1 + self.num_speculative_tokens), + max_num_batched_tokens, + ), dtype=torch.int32, device=self.device, ) @@ -2155,14 +2222,31 @@ def _build_indexer( decode_metadata = None if num_decodes > 0: - torch.diff( - common_attn_metadata.query_start_loc[: num_decodes + 1], - out=self.decode_lens_buffer[:num_decodes], - ) - decode_lens = self.decode_lens_buffer[:num_decodes] + # Single-source the decode lengths from query_start_loc_cpu. + # + # `decode_lens` is used below as the `repeats` argument of + # torch.repeat_interleave, while `actual_expanded` (derived from + # decode_lens_cpu.sum()) is passed as its `output_size`. When + # output_size is given, repeat_interleave SKIPS the device sync and + # trusts that `output_size == repeats.sum()`; if that contract is + # violated it emits an internal gather index past the source rows and + # the underlying index_select over-reads -> random illegal memory + # access. + # + # Previously `decode_lens` was diffed from the GPU `query_start_loc` + # while `output_size` came from the CPU `query_start_loc_cpu`. Under + # async scheduling the GPU and CPU copies of query_start_loc can + # momentarily disagree, so `repeats` and `output_size` came from two + # different sources and could mismatch. Deriving `decode_lens` from + # the SAME CPU tensor (copied into the persistent GPU buffer) makes + # `decode_lens.sum() == actual_expanded` hold by construction. decode_lens_cpu = torch.diff( common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] ) + self.decode_lens_buffer[:num_decodes].copy_( + decode_lens_cpu, non_blocking=True + ) + decode_lens = self.decode_lens_buffer[:num_decodes] seq_lens = common_attn_metadata.seq_lens[:num_decodes] block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index 3d374014f5..804f6eb1fe 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -771,6 +771,26 @@ def compute_logits( logits = self.model.compute_logits(hidden_states) return logits + def get_top_tokens( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Local-argmax spec-decode hook for vLLM's MTP proposer. + + vLLM's ``LLMBaseProposer._greedy_sample`` calls this (when + ``use_local_argmax_reduction`` is enabled) in place of + ``compute_logits(...).argmax(-1)``. Bridge it to ATOM's distributed + greedy argmax (``compute_draft_token``): each rank reduces its logit + shard to ``(max_val, global_idx)`` and only ``[N, 2]`` is all-gathered + instead of the full ``[N, vocab]`` logits. Token-identical to + ``compute_logits(...).argmax(-1)``; returns ``[N]`` int64 token ids. + """ + if getattr(self, "_is_deepseek_v4_mtp", False): + hidden_states = _deepseek_v4_mtp_unflatten_hidden_states( + hidden_states, self.model + ) + return self.model.compute_draft_token(hidden_states) + class ATOMForCausalLM(ATOMModelBase, VllmModelForTextGeneration): ... diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index 7a921de241..ae1649861c 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -324,11 +324,21 @@ def load_model(self, target_model: nn.Module) -> None: self.model.share_with_target(target_base, loaded) return - # Share embed_tokens with the target model + # Share embed_tokens with the target model. Match on the *logical* vocab + # (num_embeddings) and hidden dim rather than the stored weight shape, so a + # replicated target embed ([vocab, hidden], ATOM_REPLICATE_VOCAB_EMBED) is + # still shared onto a TP-sharded draft embed ([vocab/tp, hidden]) — the + # draft then reuses the target's replicated table (no post-embed + # all-reduce). When both are sharded this is identical to the old check. + draft_embed = self.model.model.embed_tokens + target_embed = target_base.model.embed_tokens + draft_vocab = getattr(draft_embed, "num_embeddings", None) + target_vocab = getattr(target_embed, "num_embeddings", None) if ( get_pp_group().world_size == 1 - and self.model.model.embed_tokens.weight.shape - == target_base.model.embed_tokens.weight.shape + and draft_vocab is not None + and draft_vocab == target_vocab + and draft_embed.weight.shape[1] == target_embed.weight.shape[1] ): logger.info( "Assuming the EAGLE head shares the same vocab embedding" @@ -484,6 +494,18 @@ def propose( 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 + # Sparse (DSA) MLA decode packs KV per token at + # page_size=1, so it reads the all-1s + # sparse_kv_last_page_lens (NOT the dense per-block + # buffer, which makes the asm kernel over-read past + # the written sparse-index region -> illegal access). + # The draft reuses the target's attn_metadata but + # drops to max_seqlen_q=1, so it must re-point this to + # the per-seq all-1s slice itself. + if "sparse_kv_last_page_lens" in var: + attn_metadata.sparse_kv_last_page_lens = var[ + "sparse_kv_last_page_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] diff --git a/atom/utils/block_convert.py b/atom/utils/block_convert.py index 7ebfa41f2e..590ac8f69f 100644 --- a/atom/utils/block_convert.py +++ b/atom/utils/block_convert.py @@ -213,6 +213,64 @@ def kv_indices_generate_triton( return kv_indices_convert +@triton.jit +def mtp_prepare_decode_mla_kernel( + kv_indptr_ptr, # int32 [bs+1] in/out: += cu_seqlens_q + cu_seqlens_q_ptr, # int32 [bs+1] in + sparse_kv_indptr_ptr, # int32 [bs+1] out (IS_SPARSE only) + positions_ptr, # int64 [bs] in/out (UPDATE_POSITIONS only) + context_lens_ptr, # int32 [bs] in/out (UPDATE_CONTEXT_LENS only) + bs, + index_topk, + position_stride, + IS_SPARSE: tl.constexpr, + UPDATE_POSITIONS: tl.constexpr, + UPDATE_CONTEXT_LENS: tl.constexpr, + BLOCK: tl.constexpr, +): + """Single-program fused per-draft-step MTP-decode metadata for MLA. + + Collapses the elementwise chain that used to launch ~4 tiny kernels + (``kv_indptr += cu_seqlens_q``; ``counts = diff(kv_indptr)``; + ``clamp(index_topk)``; ``cumsum -> sparse_kv_indptr``) plus eagle's two + per-step bumps (``positions += 1``, ``context_lens += 1``) into one launch. + One program owns the whole batch so the sparse cumsum is a single in-register + ``tl.cumsum``; the caller guarantees ``bs + 1 <= BLOCK`` (bs <= max_bs). All + reads happen before any write, so the in-place ``kv_indptr`` update is + hazard-free. ``sparse_kv_indptr[0]`` is written on-device (no Python-scalar + H2D copy, which is what caused the original hot-path Host-Device sync). + """ + offs = tl.arange(0, BLOCK) + mask_p1 = offs < (bs + 1) # indptr arrays are length bs+1 + mask_bs = offs < bs # per-seq arrays are length bs + + # Load originals first (into registers) so the kv_indptr store below cannot + # clobber the shifted [i+1] read used for per-seq counts. + kvi = tl.load(kv_indptr_ptr + offs, mask=mask_p1, other=0) + cuq = tl.load(cu_seqlens_q_ptr + offs, mask=mask_p1, other=0) + new_kvi = kvi + cuq + + if IS_SPARSE: + kvi_next = tl.load(kv_indptr_ptr + offs + 1, mask=mask_bs, other=0) + cuq_next = tl.load(cu_seqlens_q_ptr + offs + 1, mask=mask_bs, other=0) + counts = tl.where(mask_bs, (kvi_next + cuq_next) - new_kvi, 0) + clamped = tl.minimum(counts, index_topk) + # inclusive scan: csum[i] == sparse_kv_indptr[i + 1] + csum = tl.cumsum(clamped, axis=0) + tl.store(sparse_kv_indptr_ptr + offs + 1, csum, mask=mask_bs) + tl.store(sparse_kv_indptr_ptr + offs, tl.zeros_like(csum), mask=(offs == 0)) + + tl.store(kv_indptr_ptr + offs, new_kvi, mask=mask_p1) + + if UPDATE_POSITIONS: + pos = tl.load(positions_ptr + offs * position_stride, mask=mask_bs, other=0) + tl.store(positions_ptr + offs * position_stride, pos + 1, mask=mask_bs) + + if UPDATE_CONTEXT_LENS: + ctx = tl.load(context_lens_ptr + offs, mask=mask_bs, other=0) + tl.store(context_lens_ptr + offs, ctx + 1, mask=mask_bs) + + if __name__ == "__main__": # Example usage and test diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c3725c7bb4..1ecdf35d7b 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -102,11 +102,18 @@ "ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION", "1") == "1" ), - # Replicate the EAGLE3 draft vocab embedding on every TP rank (full table per - # rank, local lookup) instead of sharding it — eliminates the post-embedding - # all-reduce. The draft embed is independent of the (sharded) lm_head. - "ATOM_EAGLE_REPLICATE_EMBED": lambda: ( - os.getenv("ATOM_EAGLE_REPLICATE_EMBED", "1") == "1" + # Replicate the vocab embedding on every TP rank (full table per rank, purely + # local lookup) instead of TP-sharding it — eliminates the post-embedding + # all-reduce. Applies to BOTH the main/target model and the speculative draft + # (EAGLE3 head + MTP draft steps), so the collective is dropped on every embed + # on the critical path. Only used where the embedding is independent of the + # still TP-sharded lm_head (EAGLE3 draft; GLM-5.2 target+MTP, whose + # tie_word_embeddings=False), so the lookup is bit-identical to the sharded + # masked-embedding + all-reduce path. Trades (tp-1)/tp of the embedding's + # memory per rank for one fewer collective per embed. Default on; set "0" to + # fall back to the sharded VocabParallelEmbedding. + "ATOM_REPLICATE_VOCAB_EMBED": lambda: ( + os.getenv("ATOM_REPLICATE_VOCAB_EMBED", "1") == "1" ), "ATOM_ENABLE_GDN_DECODE_LOSSY_FAST": lambda: ( os.getenv("ATOM_ENABLE_GDN_DECODE_LOSSY_FAST", "0").lower() == "1" diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py index ce2c45116e..b7705ff357 100644 --- a/atom/utils/forward_context.py +++ b/atom/utils/forward_context.py @@ -377,6 +377,10 @@ class AttentionMetaData: cu_seqlen_ks: Optional[torch.Tensor] = None cu_seqlen_ke: Optional[torch.Tensor] = None sparse_kv_indptr: Optional[torch.Tensor] = None + # Last-page lens for sparse (DSA) attention: all 1s, one per query token in + # prefill/MTP-verify and per seq in decode. Separate from kv_last_page_lens + # (the dense per-seq buffer) so the two never clobber each other. + sparse_kv_last_page_lens: Optional[torch.Tensor] = None work_meta_data: Optional[torch.Tensor] = None work_indptr: Optional[torch.Tensor] = None diff --git a/recipes/atom_vllm/GLM-5.md b/recipes/atom_vllm/GLM-5.md index dce2a88b6f..5697aa19f8 100644 --- a/recipes/atom_vllm/GLM-5.md +++ b/recipes/atom_vllm/GLM-5.md @@ -57,6 +57,28 @@ vllm serve amd/GLM-5.2-MXFP4 \ --no-enable-prefix-caching ``` +### GLM-5.2-MXFP4 MTP3 Recipe +```bash +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export AITER_USE_FLYDSL_MOE_SORTING=1 + +vllm serve amd/GLM-5.2-MXFP4 \ + --host localhost \ + --port 8004 \ + --async-scheduling \ + --load-format fastsafetensors \ + --trust-remote-code \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --kv-cache-dtype fp8 \ + --tensor-parallel-size 4 \ + --gpu-memory-utilization 0.9 \ + --max-num-batched-tokens 16384 \ + --additional-config '{"online_quant_config": {"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "*.mlp.gate", "model.layers.[0-9].mlp.*expert*", "model.layers.[1-6][0-9].mlp.*expert*", "model.layers.7[0-7].mlp.*expert*"]}}' \ + --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}' +``` + + + ## Step 3: Performance Benchmark Users can use the default vllm bench commands for performance benchmarking. ```bash