diff --git a/atom/plugin/vllm/attention/backend.py b/atom/plugin/vllm/attention/backend.py index 3734c9c89..adaa4e315 100644 --- a/atom/plugin/vllm/attention/backend.py +++ b/atom/plugin/vllm/attention/backend.py @@ -1,4 +1,5 @@ import torch +from vllm.v1.attention.backend import MultipleOf from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend from atom.model_ops.minimax_m3.sparse_attn import SPARSE_BLOCK_SIZE @@ -34,7 +35,19 @@ def get_name() -> str: @staticmethod def get_supported_kernel_block_sizes(): - return [16] + # The AITER asm_pa kernel only ships a bf16/bf16 paged-attention variant + # for kernel block size 16, but the Triton paged-attention path reads + # block_size from the cache shape at runtime and handles any multiple of + # 16. AttentionForVllmMHA routes to the Triton path whenever the bf16 KV + # cache block size is not 16 (see layer_mha.use_triton_attn), so this + # backend genuinely supports any MultipleOf(16). Declaring it as such + # lets vLLM pick the kv-manager block size (e.g. 128) as the common + # kernel block size, so a layer using this backend (the Eagle3 draft) + # can share the uniform-type KV cache group with the block-128 + # sparse/dense layers instead of forcing a singleton group or a + # "No common block size" failure. This matches native vLLM, whose draft + # attention also supports the model block size. + return [MultipleOf(16)] @classmethod def supports_block_size(cls, block_size: int | None) -> bool: diff --git a/atom/plugin/vllm/attention/layer_mha.py b/atom/plugin/vllm/attention/layer_mha.py index b79217160..b10f61e26 100644 --- a/atom/plugin/vllm/attention/layer_mha.py +++ b/atom/plugin/vllm/attention/layer_mha.py @@ -240,8 +240,19 @@ def rope_cache( attn_metadata = attention_metadata slot_mapping = attn_metadata.slot_mapping[: q.shape[0]] - use_triton_attn = self.sliding_window != -1 or self.head_dim != 128 - # use_triton_attn = True + # The AITER asm paged-attention kernel only has a bf16/bf16 variant for + # kernel block size 16. When the KV cache uses a different block size + # (e.g. the Eagle3 draft sharing the model's block 128 so it can join + # vLLM's uniform-type per-layer-tensor KV cache group instead of + # collapsing to a singleton group), asm has no kernel and the worker + # crashes ("cannot get heuristic kernel"). The Triton path (insert + + # decode) is block-size agnostic, so route bf16 caches with a non-16 + # block size through Triton. + use_triton_attn = ( + self.sliding_window != -1 + or self.head_dim != 128 + or (not self.kv_cache_dtype.startswith("fp8") and block_size != 16) + ) self.use_triton_attn = use_triton_attn if ( @@ -311,7 +322,11 @@ def rope_cache( v_scale=v_scale, ) elif use_triton_attn and self.rotary_emb is not None: - k_scale = v_scale = self.per_tensor_scale + # `per_tensor_scale` is only populated for fp8 KV caches (see + # forward_impl). For a bf16 cache (e.g. the Eagle3 draft routed here + # because its block size != 16) it is absent and unused, since + # apply_scale below is False for non-fp8 dtypes. + k_scale = v_scale = getattr(self, "per_tensor_scale", None) self.per_token_quant = False q, k, _k_cache, _v_cache = fused_qk_rope_reshape_and_cache( q, @@ -423,7 +438,9 @@ def paged_attention_triton( query_group_size, ) compute_type = ( - torch.bfloat16 if self.kv_cache_dtype == "bf16" else aiter.dtypes.fp8 + aiter.dtypes.fp8 + if self.kv_cache_dtype.startswith("fp8") + else torch.bfloat16 ) exp_sums = torch.empty(intermediate_shape, dtype=torch.float32, device=q.device) max_logits = torch.empty( @@ -457,8 +474,8 @@ def paged_attention_triton( context_partition_size=context_partition_size, compute_type=compute_type, q_scale=None, - k_scale=None if self.kv_cache_dtype == "bf16" else k_scale, - v_scale=None if self.kv_cache_dtype == "bf16" else v_scale, + k_scale=k_scale if self.kv_cache_dtype.startswith("fp8") else None, + v_scale=v_scale if self.kv_cache_dtype.startswith("fp8") else None, exp_sums=exp_sums, max_logits=max_logits, temporary_output=temporary_output, @@ -908,6 +925,14 @@ def get_kv_cache_spec(self, vllm_config): assert self.attn_type == AttentionType.DECODER block_size = vllm_config.cache_config.block_size + # `self.sliding_window` uses -1 (not None) as the "no sliding window" + # sentinel. Only emit a SlidingWindowSpec for a *real* window (> 0); + # otherwise emit FullAttentionSpec. + # + # As Eagle3 draft has no sliding window, it must be a FullAttentionSpec. + # MLAAttentionSpec (M3's sparse indexer cache) subclasses FullAttentionSpec, + # so a spec set of {full/sparse target, full draft} stays uniform-type and + # vLLM allocates a separate KV tensor per layer. if self.sliding_window is not None and self.sliding_window > 0: return SlidingWindowSpec( block_size=block_size, diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 98e1debe8..9972650bf 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/atom/plugin/vllm/attention/metadata.py @@ -435,11 +435,13 @@ class MinimaxM3SparseMetadata: class MinimaxM3SparseAttentionMetadataBuilder(AttentionMetadataBuilder): - # Only uniform single-token decode is safe to capture. Prefill/mixed batches - # still use build(), where variable query lengths and CPU-side max reduction - # are allowed. The decode kernels consume per-step seq_lens/block_table from - # vLLM's fixed metadata buffers and keep their grids shape-constant. - _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + # Uniform decode batches are safe to capture, including spec-decode verify + # (query_len == num_spec + 1): the decode index-topk and sparse-attn kernels + # thread MAX_Q with per-token causality (causal_len = seq_len - MAX_Q + tok + + # 1) and their grids depend only on shape constants, so a captured (batch, + # query_len) shape is fixed. Prefill/mixed batches still use build(), where + # variable query lengths and CPU-side max reduction are allowed. + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH reorder_batch_threshold = 1 def __init__( diff --git a/atom/plugin/vllm/model_wrapper.py b/atom/plugin/vllm/model_wrapper.py index 87cba672c..6439e205f 100644 --- a/atom/plugin/vllm/model_wrapper.py +++ b/atom/plugin/vllm/model_wrapper.py @@ -541,49 +541,94 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # 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"}) + # Attribute names under which ATOM models nest their inner backbone. Walked + # in order at each level to build the inner-model chain. To support a model + # that nests under a new name, add it here. + _INNER_MODEL_ATTRS = ("model", "language_model") + + def _inner_model_chain(self, outer: nn.Module) -> list[nn.Module]: + """`outer` followed by each nested backbone, deepest last. + + ATOM models nest their language backbone at different names and depths: + - flat EAGLE3 draft (Eagle3LlamaModel): depth 0, attrs on `outer` + - MTP draft / text-only target: depth 1, under `.model` + - VL target (MiniMax-M3 ...TextOnly): depth 2, + `.language_model` then `.model` + Walking `_INNER_MODEL_ATTRS` at each level yields a single chain that + covers all of them, so a wanted attribute can be found wherever it lives. + """ + chain: list[nn.Module] = [] + node: nn.Module | None = outer + seen: set[int] = set() + while node is not None and id(node) not in seen: + seen.add(id(node)) + chain.append(node) + nxt = None + for name in self._INNER_MODEL_ATTRS: + cand = getattr(node, name, None) + if isinstance(cand, nn.Module) and id(cand) not in seen: + nxt = cand + break + node = nxt + return chain 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. + vLLM reads embed_tokens / embedding / layers at ``wrapper.model.`` + and the head at ``wrapper.lm_head``. ATOM nests these under one or more + backbone levels (see `_inner_model_chain`), so mirror the first holder + found in the chain onto the paths vLLM reads. + + Draft vs target differ in one way: + - Draft/MTP: vLLM *replaces* embed_tokens / lm_head with shared target + weights, so register real submodules (`setattr`) and install a + `__setattr__` hook that propagates those writes down to the inner + module the forward path reads from. + - Target: vLLM only *reads* these attrs, so mirror them as plain aliases + (`object.__setattr__`) that stay invisible to the module tree and the + weight loader — no ownership change, no sync hook. """ 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 + chain = self._inner_model_chain(model) + inner = chain[1] if len(chain) > 1 else None + is_draft = self.is_spec_draft_model + put = setattr if is_draft else object.__setattr__ + + def first_holder(attr: str) -> nn.Module | None: + for node in chain: + if hasattr(node, attr): + return node + return None - # ATOM DeepSeek-V4 names these shared modules `embed` / `head`, while - # vLLM's generic MTP proposer expects `embedding` / `lm_head`. - if not hasattr(model, "embedding") and hasattr(inner, "embed"): - model.embedding = inner.embed - if not hasattr(model, "lm_head") and hasattr(inner, "head"): - model.lm_head = inner.head + # ATOM DeepSeek-V4 names these shared modules `embed` / `head` on its + # immediate backbone child, while vLLM's generic MTP proposer expects + # `embedding` / `lm_head` on the outer model. + if inner is not None: + if not hasattr(model, "embedding") and hasattr(inner, "embed"): + put(model, "embedding", inner.embed) + if not hasattr(model, "lm_head") and hasattr(inner, "head"): + put(model, "lm_head", inner.head) - # (1) Mirror: make attrs visible on the outer model for vLLM discovery. + # (1) Mirror backbone attrs onto the outer model, and the head onto self. 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. - # Create the one-off subclass only once + if not hasattr(model, attr): + holder = first_holder(attr) + if holder is not None and holder is not model: + put(model, attr, getattr(holder, attr)) + + if not hasattr(self, "lm_head"): + holder = first_holder("lm_head") + if holder is not None: + put(self, "lm_head", holder.lm_head) + + # (2) Draft only: propagate vLLM's later writes on the outer model down + # to the inner module the forward path reads from. Create the one-off + # subclass only once, and only when there is an inner level to sync. + if not is_draft: + return + if inner is None: + return if getattr(model, "_atom_vllm_shared_attr_sync_patched", False): return shared = self._WEIGHT_SHARED_ATTRS