[vLLM-ATOM] Support EAGLE3 spec decoding for MiniMax-M3

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:

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,

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__(

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.<attr>``
+        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