diff --git a/atom/config.py b/atom/config.py index 89905f8edf..2e955c7280 100644 --- a/atom/config.py +++ b/atom/config.py @@ -2318,17 +2318,21 @@ def __post_init__(self): ) # Spec-decode + DCP arch gating. Any speculative method (mtp / # eagle3 / dspark) runs a q>1 verify pass, and DCP decode with q>1 - # uses the round-robin CP (cprr) MLA kernel, which is persistent-only - # and ships only on gfx950; on gfx942 the non-persistent fallback - # ignores the cprr masking and silently produces WRONG output. + # needs a persistent MLA kernel that applies the causal mask on the + # GLOBAL positions of the round-robin KV shard: gfx950's cprr kernel, + # or gfx1250's PS1 one, which takes g_kv_indptr. On gfx942 the + # non-persistent fallback ignores the cprr masking and silently + # produces WRONG output. if self.speculative_config is not None: from aiter.jit.utils.chip_info import get_gfx - gfx = get_gfx() - assert gfx == "gfx950", ( - f"Speculative decode + DCP is only supported on gfx950 (needs " - f"the persistent cprr MLA kernel); got {gfx}. Disable DCP or " - f"speculative decode on this GPU." + from atom.distributed.dcp_utils import dcp_persistent_supported + + assert dcp_persistent_supported(), ( + f"Speculative decode + DCP needs the persistent round-robin CP " + f"MLA decode (gfx950, or gfx1250 with " + f"AITER_MLA_DECODE_PS1_FLYDSL=1); got {get_gfx()}. Disable DCP " + f"or speculative decode on this GPU." ) # TBO slices its ubatch work buffers by request, while the # sparse DCP indexer indexes them by query token. The two agree diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 1eeb571e6a..e18df5ad59 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -34,6 +34,7 @@ concat_and_cache_mla_seg = None fused_qk_rope_concat_and_cache_mla_seg = None from aiter.dist.parallel_state import get_dp_group, get_tensor_model_parallel_rank +from aiter.jit.utils.chip_info import get_gfx from aiter.mla import mla_decode_fwd, mla_prefill_fwd from aiter.ops.triton.attention.mla import ( mla_decode_fwd as triton_shuffle_mla_decode_fwd, @@ -291,6 +292,11 @@ def mla_kernel_num_heads(num_heads: int) -> int: # width sets below, which only hold widths that have their own kernel. _MLA_DCP_MAX_KERNEL_HEADS = 128 +# gfx1250's persistent decode is aiter's PS1 kernel. It takes a multi-token +# (q > 1) query only at these gathered widths, and aiter does not fold the +# others onto it, so a wider verify or draft-block decode is padded up to one. +_MLA_DCP_KERNEL_WIDTHS_PS1_MULTI_TOKEN = (16, 96) + _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT = (16, 32, 128) _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT_FP8 = (16, 128) _MLA_DCP_SPARSE_PREFILL_WIDTHS = (16, 128) @@ -311,6 +317,7 @@ def mla_dcp_kernel_num_heads( *, kv_cache_dtype: str, persistent: bool, + multi_token: bool = False, ) -> int: """Width to gather the query heads to for a DCP decode. @@ -324,9 +331,19 @@ def mla_dcp_kernel_num_heads( ``kv_cache_dtype`` only selects the non-persistent set: gqa=64 is excluded for both dtypes there (fp8 aborts on it, bf16 silently miscomputes it), and fp8 lacks a gqa=32 kernel on top of that while bf16 does not. + + ``multi_token``: the decode takes more than one query per sequence (a + speculative verify pass, or a draft block). gfx1250's persistent kernel + serves that only at ``_MLA_DCP_KERNEL_WIDTHS_PS1_MULTI_TOKEN``. """ gathered = mla_kernel_num_heads(max(num_heads * dcp_world_size, min_kernel_heads)) - if persistent: + widest = _MLA_DCP_MAX_KERNEL_HEADS + if persistent and multi_token and get_gfx() == "gfx1250": + widest = _MLA_DCP_KERNEL_WIDTHS_PS1_MULTI_TOKEN[-1] + for width in _MLA_DCP_KERNEL_WIDTHS_PS1_MULTI_TOKEN: + if width >= gathered: + return width + elif persistent: if gathered <= _MLA_DCP_MAX_KERNEL_HEADS: return gathered else: @@ -343,7 +360,7 @@ def mla_dcp_kernel_num_heads( _dcp_kernel_width_warned = True logger.warning( f"DCP decode gathers {gathered} query heads, past the widest MLA " - f"kernel aiter dispatches ({_MLA_DCP_MAX_KERNEL_HEADS}); it serves " + f"kernel aiter dispatches ({widest}); it serves " "neither a kernel nor a fold that wide and will abort in " "mla_decode_fwd. Lower decode_context_parallel_size or raise tp." ) @@ -914,6 +931,9 @@ def _configure_dcp_decode_head_padding(self, dcp_world_size: int) -> None: self, "sparse_dcp_metadata_rebuild", False ), ), + # A speculative run decodes several queries per sequence: the + # target's verify pass and the draft's block pass alike. + multi_token=get_current_atom_config().speculative_config is not None, ) self.dcp_head_pad = ( self.dcp_kernel_num_heads - self.num_heads * dcp_world_size diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 477957445a..19ed7f77af 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -464,6 +464,7 @@ def __init__(self, model_runner): dcp_persistent, sparse_metadata_rebuild=self.sparse_dcp_metadata_rebuild, ), + multi_token=getattr(model_runner, "num_spec_tokens", 0) > 0, ) else: self.persistent_num_heads = self.padded_num_attention_heads