diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py index 9ab592d1bc..06afa3a9f3 100644 --- a/atom/model_ops/attention_mha.py +++ b/atom/model_ops/attention_mha.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. +import logging from functools import cache import aiter @@ -29,6 +30,12 @@ from .attention_mla import MLAModules +logger = logging.getLogger("atom") + +# Once per process: the hint names the opt-in, not the shape, so one line says +# all a default user needs, and the decode loop pays one bool test after it. +_flydsl_off_hint_logged = False + @cache def use_pa_decode_bf16_asm() -> bool: @@ -181,16 +188,23 @@ def _can_use_prefill_sink_asm( return False return q.shape[1] % k.shape[1] == 0 - def _reject_per_token_scales_on_unified(self, k_scale, v_scale, phase: str): + @staticmethod + def _per_token_kv(k_scale, v_scale) -> bool: + return (k_scale is not None and k_scale.numel() > 1) or ( + v_scale is not None and v_scale.numel() > 1 + ) + + def _reject_per_token_scales_on_unified( + self, k_scale, v_scale, phase: str, flydsl_would_take: bool = False + ): """unified_attention carries one descale for the whole tensor. Feeding it a per-token cache is not an error downstream, just wrong numbers, so it has to be refused here. Both phases route into unified on - the same conditions, so both check. + the same conditions, so both check. ``flydsl_would_take``: the decode + is one ATOM_PA_FLYDSL=1 would offer to FlyDSL, so the error says so. """ - if (k_scale is not None and k_scale.numel() > 1) or ( - v_scale is not None and v_scale.numel() > 1 - ): + if self._per_token_kv(k_scale, v_scale): if phase == "decode" and not pa_decode_gluon_supports_arch(): reason = ( "set ATOM_PA_FLYDSL=1" @@ -202,9 +216,15 @@ def _reject_per_token_scales_on_unified(self, k_scale, v_scale, phase: str): f"cache on this GPU: {reason}, unified_attention cannot " "carry per-token KV scales, and there is no gluon fallback" ) + hint = ( + "; ATOM_PA_FLYDSL=1 lets aiter's FlyDSL pa_decode take it " + "(needs aiter with ROCm/aiter#7206)" + if flydsl_would_take + else "" + ) raise NotImplementedError( f"layer {self.layer_num} takes unified_attention for {phase}, " - "which cannot carry its per-token KV scales" + f"which cannot carry its per-token KV scales{hint}" ) def forward_impl( @@ -608,19 +628,61 @@ def _shuffle_v_for_flydsl(self, v_cache, k_cache): return v_cache return self._view_v_cache_for_pa_decode_bf16_asm(v_cache, k_cache) - def _flydsl_takes_decode(self, q, k_cache, v_cache, attn_metadata) -> bool: - """Without gluon decode (gfx1250): whether FlyDSL accepts this call. - - ATOM_USE_UNIFIED_ATTN keeps the call on unified_attention. - """ + def _pa_decode_kv_scales(self, k_scale, v_scale): + """The K/V scales ``paged_attention_triton`` hands ``run_pa_decode``.""" + if self.kv_cache_dtype == "bf16": + return None, None + if k_scale is not None and k_scale.numel() > 1: + return k_scale.unsqueeze(-1), v_scale.unsqueeze(-1) + return k_scale, v_scale + + def _flydsl_route_open(self, attn_metadata) -> bool: + """Whether this unified-bound decode is one ATOM_PA_FLYDSL offers to + FlyDSL: not forced to unified, paged layout, no window, no sinks, not + BF16 KV on a gluon GPU (ATOM's own routing in + ``_flydsl_pa_decode_num_seqs``), and on a gluon GPU past the gluon + envelope (inside it the gluon route is kept). aiter's predicate is not + asked here; the opt-in hint relies on this, so it never names a route + ATOM's routing would refuse.""" + gluon = pa_decode_gluon_supports_arch() if ( envs.ATOM_USE_UNIFIED_ATTN - or pa_decode_gluon_supports_arch() + or self.use_flash_layout or self.sliding_window > 0 + or self.sinks is not None + or (self.kv_cache_dtype == "bf16" and gluon) ): return False - num_seqs = attn_metadata.context_lens.shape[0] - return bool( + return not gluon or gluon_decode_over_limit( + attn_metadata.max_seqlen_q, self.num_heads, self.num_kv_heads + ) + + def _flydsl_takes_decode( + self, q, k_cache, v_cache, k_scale, v_scale, attn_metadata + ): + """Sequence count FlyDSL runs a decode routed to unified with, or a + falsy value to keep it on unified. + + Two cases reach unified with a call FlyDSL may take: a GPU without + gluon decode (gfx1250), and a gluon GPU whose query length x group is + past the gluon envelope (Qwen3.5 MTP-5: qlen 6, group 8). Inside that + envelope a gluon GPU keeps its route. The decision is made per call, + on the same tensors and metadata ``run_pa_decode`` then receives, and + the count is passed down as ``flydsl_seqs``, so aiter's predicate is + asked once per call and the two cannot disagree. + + With ATOM_PA_FLYDSL unset such a call stays on unified; + ``paged_attention_unified`` says what the flag would change. + + ATOM_USE_UNIFIED_ATTN keeps the call on unified_attention. + + gfx1250 with per-token FP8 KV past qlen 4 is decided by the same check; + it is not exercised by this path's tests or runs (out of scope here). + """ + if not envs.ATOM_PA_FLYDSL or not self._flydsl_route_open(attn_metadata): + return False + k_scale, v_scale = self._pa_decode_kv_scales(k_scale, v_scale) + return ( flydsl_decode_seqs( output=q, q=q, @@ -629,9 +691,6 @@ def _flydsl_takes_decode(self, q, k_cache, v_cache, attn_metadata) -> bool: block_tables=attn_metadata.block_tables, context_lens=attn_metadata.context_lens, max_seqlen_q=attn_metadata.max_seqlen_q, - max_context_partition_num=dense_decode_splits( - num_seqs, k_cache.shape[1] - ), context_partition_size=256, compute_type=( torch.bfloat16 @@ -639,11 +698,14 @@ def _flydsl_takes_decode(self, q, k_cache, v_cache, attn_metadata) -> bool: else aiter.dtypes.fp8 ), q_scale=None, + k_scale=k_scale, + v_scale=v_scale, alibi_slopes=None, sinks=self.sinks, sliding_window=self.sliding_window, ps=True, ) + or False ) @mark_trace(prefix="paged_attention_unified", torch_compile=False) @@ -657,11 +719,40 @@ def paged_attention_unified( of the reasons for it. """ attn_metadata = fwd_ctx.attn_metadata - if self._flydsl_takes_decode(q, k_cache, v_cache, attn_metadata): + flydsl_seqs = self._flydsl_takes_decode( + q, k_cache, v_cache, k_scale, v_scale, attn_metadata + ) + if flydsl_seqs: return self.paged_attention_triton( - q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx + q, + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + fwd_ctx, + flydsl_seqs=flydsl_seqs, + ) + # A default user gets the opt-in named once: in the per-token FP8 KV + # refusal (unified cannot carry those scales), else in one info line. + flydsl_would_take = not envs.ATOM_PA_FLYDSL and self._flydsl_route_open( + attn_metadata + ) + self._reject_per_token_scales_on_unified( + k_scale, v_scale, "decode", flydsl_would_take + ) + global _flydsl_off_hint_logged + if flydsl_would_take and not _flydsl_off_hint_logged: + _flydsl_off_hint_logged = True + logger.info( + "decode with max_seqlen_q=%d x %d q heads per kv head runs on " + "unified_attention because ATOM_PA_FLYDSL is unset; " + "ATOM_PA_FLYDSL=1 lets aiter's FlyDSL pa_decode take it " + "(needs aiter with ROCm/aiter#7206)", + attn_metadata.max_seqlen_q, + self.num_heads // self.num_kv_heads, ) - self._reject_per_token_scales_on_unified(k_scale, v_scale, "decode") if envs.ATOM_USE_UNIFIED_ATTN and self.kv_cache_dtype.startswith("fp8"): o = torch.empty(*q.shape, dtype=torch.bfloat16, device=q.device) @@ -696,13 +787,23 @@ def paged_attention_unified( @mark_trace(prefix="paged_attention_triton", torch_compile=False) def paged_attention_triton( - self, q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx: ForwardContext + self, + q, + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + fwd_ctx: ForwardContext, + flydsl_seqs: int | None = None, ): """Decode through ``run_pa_decode``: gluon, or opted-in FlyDSL. Callers must have cleared ``gluon_decode_over_limit`` in - ``_dispatch_decode`` or, on a GPU without gluon, established through - ``_flydsl_takes_decode`` that FlyDSL accepts the call. + ``_dispatch_decode`` or established through ``_flydsl_takes_decode`` + that FlyDSL accepts the call; the latter pass the sequence count it + returned as ``flydsl_seqs``, so ``run_pa_decode`` does not ask again. """ attn_metadata = fwd_ctx.attn_metadata @@ -723,9 +824,7 @@ def paged_attention_triton( max_context_partition_num = 1 context_partition_size = 128 - if k_scale is not None and k_scale.numel() > 1: - k_scale = k_scale.unsqueeze(-1) - v_scale = v_scale.unsqueeze(-1) + k_scale, v_scale = self._pa_decode_kv_scales(k_scale, v_scale) compute_type = ( torch.bfloat16 if self.kv_cache_dtype == "bf16" else aiter.dtypes.fp8 @@ -743,8 +842,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, + v_scale=v_scale, alibi_slopes=None, sinks=self.sinks, sliding_window=self.sliding_window, @@ -755,6 +854,7 @@ def paged_attention_triton( # max_context_length: the block-table bound is the same in graph # and eager, so both pick the same kernel specialization. work_plan=getattr(attn_metadata, "flydsl_work_plan", None), + flydsl_seqs=flydsl_seqs, ) return o @@ -1086,9 +1186,10 @@ def prefill_attention_triton( def _dispatch_decode(self, max_qlen: int = 1): """Select the coarse decode runner from model shape and configuration. - Cache-dependent dispatch happens later: on a GPU without gluon, - ``paged_attention_unified`` hands FlyDSL calls inside its supported - domain and keeps the rest on unified attention. The vLLM bridge shares + Cache-dependent dispatch happens later: on a GPU without gluon, or + past the gluon envelope on one with it, ``paged_attention_unified`` + hands FlyDSL calls inside its supported domain and keeps the rest on + unified attention. The vLLM bridge shares the constants and predicate, not this order -- its arms differ (no unified fallback, no sliding-window arm). Both paged kernels stop short of unified, at different points, and drafting diff --git a/atom/model_ops/base_attention.py b/atom/model_ops/base_attention.py index 2998d6f5e9..f9a3aff2f8 100644 --- a/atom/model_ops/base_attention.py +++ b/atom/model_ops/base_attention.py @@ -173,12 +173,9 @@ def run_pa_fwd_asm( ) -# Copies of aiter's own limits, pinned against its source by a test. Mirroring -# them lets the gate decline an unsupported call before entering aiter; a copy -# that drifts either over-rejects or stops protecting. -_FLYDSL_PA_MAX_PARTITIONS = 256 -_FLYDSL_PA_TILE = 256 -_FLYDSL_PA_BLOCK_SIZES = (16, 64, 128) +# The GPUs aiter's FlyDSL paged decode builds for, pinned against its source by +# a test. Only the MiniMax-M3 startup check reads it; per-call acceptance is +# aiter's can_use_flydsl_pa_decode. _FLYDSL_PA_ARCHS = ("gfx942", "gfx950", "gfx1250") _flydsl_pa_routed: set[tuple] = set() _flydsl_plan_refused: set[tuple] = set() @@ -252,6 +249,30 @@ def _flydsl_explicit_plan_api() -> bool: return ok +def _flydsl_pa_decode_can_use(): + """aiter's ``can_use_flydsl_pa_decode``; raise if this aiter lacks it. + + FlyDSL paged decode is opt-in (ATOM_PA_FLYDSL=1), and ATOM no longer keeps + its own copy of the kernel's acceptance rules, so an aiter without the + predicate cannot be routed safely. + """ + try: + mod = importlib.import_module("aiter.ops.flydsl.pa_decode") + except ImportError as exc: + raise RuntimeError( + "ATOM_PA_FLYDSL=1 needs aiter's FlyDSL paged decode, which this " + "aiter cannot import; unset ATOM_PA_FLYDSL or install a newer aiter" + ) from exc + can_use = getattr(mod, "can_use_flydsl_pa_decode", None) + if can_use is None: + raise RuntimeError( + "ATOM_PA_FLYDSL=1 needs aiter.ops.flydsl.pa_decode." + "can_use_flydsl_pa_decode (ROCm/aiter#7206), which this aiter " + "predates; unset ATOM_PA_FLYDSL or install a newer aiter" + ) + return can_use + + def _flydsl_pa_decode_num_seqs( *, output: torch.Tensor, @@ -261,20 +282,24 @@ def _flydsl_pa_decode_num_seqs( block_tables: torch.Tensor, context_lens: torch.Tensor, max_seqlen_q: int, - max_context_partition_num: int, context_partition_size: int, compute_type: torch.dtype, q_scale: torch.Tensor | None, + k_scale: torch.Tensor | None, + v_scale: torch.Tensor | None, alibi_slopes: torch.Tensor | None, sinks: torch.Tensor | None, sliding_window: int, ps: bool, + max_context_length: int | None = None, ) -> int | None: """Sequence count to run aiter's FlyDSL paged decode with, or None to decline. - Mirrors the kernel's own validation, so an unsupported call is rejected - before entering aiter. The caller uses gluon where available and raises - otherwise. Every clause below is a hard reject there, not a preference. + Whether the kernel takes the call is aiter's decision + (``can_use_flydsl_pa_decode``), asked with the exact arguments + ``run_pa_decode`` then passes to ``pa_decode``. ATOM decides only two + things here: which calls it routes to FlyDSL at all, and how many + sequences the call covers. Why a count rather than a bool: FlyDSL demands ``q.shape[0] == context_lens.shape[0] * query_length`` exactly, but ATOM @@ -288,66 +313,12 @@ def _flydsl_pa_decode_num_seqs( headroom over gluon lives, it runs ``num_spec + 1``, and FlyDSL tunes that shape (it has a query_length==4 MTP4 grid split). """ - import aiter - - if alibi_slopes is not None or sinks is not None or q_scale is not None: - return None - if sliding_window > 0 or not ps: + # ATOM's routing, not kernel limits: sliding-window, sink and + # non-persistent calls stay on gluon, and so does BF16 KV where gluon + # exists. + if sinks is not None or sliding_window > 0 or not ps: return None - # Native BF16 KV uses 8-element vectors and no scales; FP8 uses 16. - is_bf16_kv = compute_type is torch.bfloat16 and k_cache.dtype is torch.bfloat16 - is_fp8_kv = compute_type is aiter.dtypes.fp8 and k_cache.dtype is aiter.dtypes.fp8 - if not (is_fp8_kv or is_bf16_kv): - return None - # BF16 KV stays on gluon where gluon exists. - if is_bf16_kv and pa_decode_gluon_supports_arch(): - return None - x = 8 if is_bf16_kv else 16 - if k_cache.dim() != 5 or k_cache.shape[-1] != x: - return None - # block_size is dim -2 of the page-16 cache. --block-size 256/1024 reaches - # this site whenever use_triton_attn is set, which the published recipe does. - if k_cache.shape[-2] not in _FLYDSL_PA_BLOCK_SIZES: - return None - if not _flydsl_arch_supported(): - return None - if not _flydsl_explicit_plan_api(): - return None - if context_partition_size != _FLYDSL_PA_TILE: - return None - if not 1 <= max_context_partition_num <= _FLYDSL_PA_MAX_PARTITIONS: - return None - head_dim = q.shape[-1] - if not (head_dim == 64 or (head_dim % 128 == 0 and head_dim <= 1024)): - return None - # The cache encodes head_dim as num_hgroups * x; aiter rejects a q whose - # own head_dim disagrees, and nothing upstream forces the two to match. - if k_cache.shape[2] * x != head_dim: - return None - if q.dtype not in (torch.bfloat16, torch.float16): - return None - if output.dtype is not q.dtype or output.shape != q.shape: - return None - # aiter checks both head_dim axes (pa_decode.py:443,448); output is - # whatever the caller handed down, not necessarily contiguous. - if q.stride(2) != 1 or output.stride(2) != 1: - return None - if q.shape[-2] % k_cache.shape[1] != 0: - return None - if v_cache.dtype is not k_cache.dtype: - return None - if block_tables.dtype is not torch.int32 or context_lens.dtype is not torch.int32: - return None - # aiter requires all four contiguous (pa_decode.py:464-470). No in-tree - # path produces a non-contiguous one -- the cache views come from - # `.view()`, which would raise first -- but the SGLang bridge's pool is - # not this tree's to promise, and ATOM_PA_FLYDSL routes it too. - if not ( - k_cache.is_contiguous() - and v_cache.is_contiguous() - and block_tables.is_contiguous() - and context_lens.is_contiguous() - ): + if k_cache.dtype is torch.bfloat16 and pa_decode_gluon_supports_arch(): return None # The rectangle, recovered the way gluon recovers it. See the docstring. if max_seqlen_q < 1: @@ -357,6 +328,27 @@ def _flydsl_pa_decode_num_seqs( return None if num_seqs > block_tables.shape[0]: return None + if not _flydsl_pa_decode_can_use()( + output, + q, + k_cache, + v_cache, + context_lens[:num_seqs], + block_tables[:num_seqs], + max_seqlen_q, + context_partition_size=context_partition_size, + compute_type=compute_type, + query_scale=q_scale, + key_scale=k_scale, + value_scale=v_scale, + alibi_slopes=alibi_slopes, + sinks=sinks, + sliding_window=0, + max_context_length=_flydsl_context_bound( + max_context_length, block_tables, k_cache + ), + ): + return None return num_seqs @@ -602,7 +594,9 @@ def flydsl_decode_seqs(*, max_context_length=None, **kwargs): False when FlyDSL is off, None when the call is outside FlyDSL's domain. """ - flydsl_seqs = envs.ATOM_PA_FLYDSL and _flydsl_pa_decode_num_seqs(**kwargs) + flydsl_seqs = envs.ATOM_PA_FLYDSL and _flydsl_pa_decode_num_seqs( + max_context_length=max_context_length, **kwargs + ) # aiter's planner rejects more sequences than one plan holds. if flydsl_seqs and flydsl_seqs > _flydsl_plan_max_batch(): flydsl_seqs = None @@ -644,6 +638,7 @@ def run_pa_decode( plan_step_owner=None, plan_cache_owner=None, plan_partition_cap: int | None = None, + flydsl_seqs: int | None = None, ): """Run the AITER paged-attention decode kernel. @@ -652,32 +647,40 @@ def run_pa_decode( sites -- including the vLLM and SGLang bridges -- import it. gluon unless ``ATOM_PA_FLYDSL=1``, and then only where FlyDSL's - domain covers the call: ``_flydsl_pa_decode_num_seqs`` mirrors the kernel's - own validation. Unsupported shapes fall back on GPUs with gluon support; + domain covers the call, which aiter's ``can_use_flydsl_pa_decode`` decides + (see ``_flydsl_pa_decode_num_seqs``). Unsupported shapes fall back on GPUs + with gluon support; a GPU without gluon raises an explicit error. ``exp_sums`` / ``max_logits`` / ``temporary_output`` are gluon scratch. Omitted, ``pa_decode_gluon`` sizes and allocates its own, so a FlyDSL call, which uses its plan's cached scratch, allocates none. + + ``flydsl_seqs``: the sequence count the caller already got from + ``flydsl_decode_seqs`` for these exact arguments + (``PagedAttentionImpl._flydsl_takes_decode``). Given, the check is not + asked again, so a layer asks aiter's predicate once per call. """ - flydsl_seqs = flydsl_decode_seqs( - output=output, - q=q, - k_cache=k_cache, - v_cache=v_cache, - block_tables=block_tables, - context_lens=context_lens, - max_seqlen_q=max_seqlen_q, - max_context_partition_num=max_context_partition_num, - context_partition_size=context_partition_size, - compute_type=compute_type, - q_scale=q_scale, - alibi_slopes=alibi_slopes, - sinks=sinks, - sliding_window=sliding_window, - ps=ps, - max_context_length=max_context_length, - ) + if not flydsl_seqs: + flydsl_seqs = flydsl_decode_seqs( + output=output, + q=q, + k_cache=k_cache, + v_cache=v_cache, + block_tables=block_tables, + context_lens=context_lens, + max_seqlen_q=max_seqlen_q, + context_partition_size=context_partition_size, + compute_type=compute_type, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + alibi_slopes=alibi_slopes, + sinks=sinks, + sliding_window=sliding_window, + ps=ps, + max_context_length=max_context_length, + ) # Inside the guard, not before it: this runs 63 times per decode step and # attention is an eager region of the token graphs, so replay does not # elide it. A deployment that never enables FlyDSL should pay nothing here. diff --git a/atom/utils/envs.py b/atom/utils/envs.py index d0f7adcce7..22a756f614 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -911,7 +911,8 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: os.getenv("ATOM_V4_HCA_PERSIST_MIN_ROWS", "15") ), # Route the paged decode to aiter's FlyDSL kernel (with its work planner) - # instead of gluon. Needs aiter with #5809. + # instead of gluon. Needs aiter with #5809 (explicit work plan) and #7206 + # (can_use_flydsl_pa_decode, which decides per call whether FlyDSL takes it). "ATOM_PA_FLYDSL": lambda: os.getenv("ATOM_PA_FLYDSL", "0") == "1", # Qwen QSA: score decode-sized index calls (<= 256 rows) with the MFMA # kernel that stacks the index heads into M; 0 keeps the CUDA-core one. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index f730b1c6fc..e4b614dccd 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -145,7 +145,7 @@ cannot affect a server that does not use the feature. | **ATOM_USE_TRITON_MXFP4_BMM** | bool | 0 (false) | If set to `1`, use FP4 BMM in MLA attention module. | | **ATOM_USE_FLYDSL_GATHER_KV_B_PROJ** | bool | 1 (true) | Use the FlyDSL fused gather + `kv_b_proj` GEMM for MLA's cached-prefix path. Covers page_size-1 fp8 (e4m3) KV with an fp8 weight on gfx950 — i.e. Kimi-K3 / DeepSeek MLA under `--kv-cache-dtype fp8`. Unavailable imports or unsupported tensor configurations use Triton; kernel execution errors propagate. Set `0` to force Triton. | | **ATOM_USE_FLYDSL_FP8_PREFILL_ATTN** | bool | 0 (false) | Enable FlyDSL FP8 MLA prefill on gfx950, including fused K concatenation with QKV quantization and direct FP8 output from FlyDSL gather where supported. AITER capability checks are cached per layer. Missing FP8 attention support or unsupported device/model configurations select BF16 attention before quantization and emit a warning once per process. FP16 activations and synthetic RoPE-padding configurations also use the normal attention path. Cached K/V reuse descales `max(new_token_descale, 1e-6) * 2`; larger outliers saturate. Unsupported FP8 gather configurations use BF16 gather followed by fused K/V dynamic quantization when FP8 attention is supported. Kernel execution errors propagate without retry. Added 2026-09-10. | -| **ATOM_PA_FLYDSL** | bool | 0 (false) | Route the MHA paged decode to aiter's FlyDSL kernel, with its work planner, instead of gluon. On GPUs with gluon support, shapes outside FlyDSL's domain, calls over 4096 sequences and calls whose context bound is below the query window fall back to gluon. gfx1250 has no gluon fallback: native MiniMax-M3 requires this flag and validates the FlyDSL architecture, explicit-plan API and `max_num_batched_tokens * local_num_kv_heads` planner bound at model initialization. Dense calls FlyDSL declines stay on `unified_attention` when its cache scaling/layout is supported; per-token FP8 KV is rejected because unified attention cannot carry its scales. Applies to the vLLM bridge too. **Needs aiter with #5809**; upgrade ATOM and aiter together. Plans use aiter's default workgroup budget (2×CU) and partition ceiling (CU count), except at head_dim 256 (Qwen3.5), where they use 64 partitions and a 256-workgroup budget to keep aiter's head-256 reducer on its fast path (until aiter #5949). In native ATOM, captured dense batch sizes use plans built by the metadata builder (shared by target and EAGLE draft); other calls, and every vLLM bridge call, get a runtime plan. MiniMax-M3 sparse and the vLLM bridge keep their static split count. Sparse prefill keeps ATOM's plan/scratch tensors on its step metadata; when the local tail block is forced, all sparse layers also share one plan refresh. Sparse decode keeps process-lifetime plan/scratch tensors because captured graphs retain their addresses. A new row count still compiles aiter's planner once. `ATOM_PA_FLYDSL_PLAN` was removed and is ignored with a warning. | +| **ATOM_PA_FLYDSL** | bool | 0 (false) | Route the MHA paged decode to aiter's FlyDSL kernel, with its work planner, instead of gluon. On GPUs with gluon support, shapes outside FlyDSL's domain, calls over 4096 sequences and calls whose context bound is below the query window fall back to gluon. gfx1250 has no gluon fallback: native MiniMax-M3 requires this flag and validates the FlyDSL architecture, explicit-plan API and `max_num_batched_tokens * local_num_kv_heads` planner bound at model initialization. Dense calls FlyDSL declines stay on `unified_attention` when its cache scaling/layout is supported; per-token FP8 KV is rejected because unified attention cannot carry its scales. Applies to the vLLM bridge too. **Needs aiter with #5809 (explicit work plan) and #7206 (`can_use_flydsl_pa_decode`, aiter's per-call acceptance check; with an older aiter, setting the flag raises an error naming #7206)**; upgrade ATOM and aiter together. Plans use aiter's default workgroup budget (2×CU) and partition ceiling (CU count), except at head_dim 256 (Qwen3.5), where they use 64 partitions and a 256-workgroup budget to keep aiter's head-256 reducer on its fast path (until aiter #5949). In native ATOM, captured dense batch sizes use plans built by the metadata builder (shared by target and EAGLE draft); other calls, and every vLLM bridge call, get a runtime plan. MiniMax-M3 sparse and the vLLM bridge keep their static split count. Sparse prefill keeps ATOM's plan/scratch tensors on its step metadata; when the local tail block is forced, all sparse layers also share one plan refresh. Sparse decode keeps process-lifetime plan/scratch tensors because captured graphs retain their addresses. A new row count still compiles aiter's planner once. `ATOM_PA_FLYDSL_PLAN` was removed and is ignored with a warning. | | **ATOM_V4_HCA_PERSIST** | bool | 1 (true) | DeepSeek-V4 HCA (`compress_ratio` 128) fp8 decode through aiter's persistent V4-NM kernel (`aiter.mla.mla_decode_fwd_v4_nm_ps`): one launch that splits the KV on the GPU and merges in-kernel, instead of the decode ASM plus a host-built split plan. Taken only with fp8 KV, 128 local heads, gfx950, row-dense KV pools and `ATOM_V4_HCA_PERSIST_MIN_ROWS <= rows <= 32768`; everything else stays on the ASM path. An aiter without the kernel keeps the ASM path (logged once). One workspace per device (about 64 MiB) is allocated at model load, before KV sizing and CUDA-graph capture. Set `0` to force the ASM path. | | **ATOM_V4_HCA_PERSIST_MIN_ROWS** | int | 15 | Smallest decode call (q rows) that uses `ATOM_V4_HCA_PERSIST`; smaller calls stay on the ASM path. | | **ATOM_DCP_SPARSE_PREFILL_NONPS** | bool | 1 (true) | DSA models under DCP: run sparse prefill attention on the non-persistent fp8 gqa64 MLA kernel, which needs no work metadata. Applies with fp8 q/KV, 64 gathered heads (tp4/dcp4, tp8/dcp8), gfx950, `ATOM_MLA_PAGE_SIZE=1` and `--max-num-batched-tokens` <= 32768 (each prefill token is one kernel row; a larger step budget logs a warning and keeps the default kernel). Needs an aiter that ships the kernel (ROCm/aiter#6132); with an older aiter set `0`, or the first sparse prefill past `index_topk` raises. | diff --git a/tests/test_paged_attention_dispatch.py b/tests/test_paged_attention_dispatch.py index 10c0a69ad1..3fd69741e6 100644 --- a/tests/test_paged_attention_dispatch.py +++ b/tests/test_paged_attention_dispatch.py @@ -401,6 +401,348 @@ def test_persistent_asm_is_not_bounded_by_the_run_pa_fwd_envelope( ) +class TestFlyDSLPastGluon: + """Past the gluon envelope on a gluon GPU, FlyDSL takes the call per call. + + Qwen3.5 TP4 MTP-5: 8 q heads / 1 kv head per rank, head_dim 256, page-16 + fp8 KV with per-token scales, qlen 6 -> 48 rows. _dispatch_decode still + sends it to paged_attention_unified; the FlyDSL decision is made there on + the call's own tensors (the same check run_pa_decode applies) and passed + down as flydsl_seqs. + """ + + HEADS, KV_HEADS, HEAD_DIM, BLOCK = 8, 1, 256, 16 + + @staticmethod + def _meta(*shape, dtype=None): + import torch + + return torch.empty(shape, device="meta", dtype=dtype) + + def _layer(self, **over): + from atom.model_ops import attention_mha as mha + + layer = mha.PagedAttentionImpl.__new__(mha.PagedAttentionImpl) + layer.num_heads, layer.num_kv_heads = self.HEADS, self.KV_HEADS + layer.sliding_window, layer.kv_cache_dtype, layer.sinks = -1, "fp8", None + layer.use_flash_layout = False + layer.layer_num = 0 + for k, v in over.items(): + setattr(layer, k, v) + return layer + + def _call(self, qlen, seqs=4, block=None, kv_dtype=None): + import torch + + block = block or self.BLOCK + kv_dtype = kv_dtype or _FP8 + x = 8 if kv_dtype is torch.bfloat16 else 16 + m = self._meta + q = m(seqs * qlen, self.HEADS, self.HEAD_DIM, dtype=torch.bfloat16) + k_cache = m(64, self.KV_HEADS, self.HEAD_DIM // x, block, x, dtype=kv_dtype) + v_cache = m(64, self.KV_HEADS, self.HEAD_DIM, block, dtype=kv_dtype) + md = SimpleNamespace( + context_lens=m(seqs, dtype=torch.int32), + block_tables=m(seqs, 64, dtype=torch.int32), + max_seqlen_q=qlen, + ) + return q, k_cache, v_cache, md + + @staticmethod + def _flydsl_on(monkeypatch, on=True, aiter_accepts=True): + """FlyDSL opted in; aiter's predicate is faked (meta tensors), the rest + of ATOM's decision is the real one. Returns the predicate's calls.""" + from atom.model_ops import attention_mha as mha + from atom.model_ops import base_attention as ba + + monkeypatch.setattr(ba.envs, "ATOM_PA_FLYDSL", on) + monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", False) + calls = [] + + def can_use(*args, **kwargs): + calls.append((args, kwargs)) + return aiter_accepts + + monkeypatch.setattr(ba, "_flydsl_pa_decode_can_use", lambda: can_use) + return calls + + def test_dispatch_still_sends_past_gluon_to_unified(self, monkeypatch): + """No static gate in _dispatch_decode: qlen 6 routes as on main, and + never through the ASM arm (qlen 5 x group 3 is inside ASM's bound).""" + assert _route(monkeypatch, 6, num_heads=8) == "paged_attention_unified" + assert ( + _route(monkeypatch, 5, num_heads=3, use_triton_attn=False) + == "paged_attention_unified" + ) + + @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter FlyDSL pa_decode") + @pytest.mark.parametrize("qlen", [5, 6, 8]) + def test_qwen35_mtp5_takes_flydsl(self, monkeypatch, qlen): + import torch + + calls = self._flydsl_on(monkeypatch) + q, kc, vc, md = self._call(qlen) + ks = vs = torch.empty(64, self.KV_HEADS, self.BLOCK, device="meta") + assert self._layer()._flydsl_takes_decode(q, kc, vc, ks, vs, md) + # aiter is asked with what run_pa_decode will pass: per-token scales + # with their trailing singleton, the query length, no window. + ((args, kwargs),) = calls + assert args[1] is q and args[6] == qlen + assert kwargs["key_scale"].shape == (64, self.KV_HEADS, self.BLOCK, 1) + assert kwargs["value_scale"].shape == (64, self.KV_HEADS, self.BLOCK, 1) + assert kwargs["sliding_window"] == 0 + + @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter FlyDSL pa_decode") + @pytest.mark.parametrize("qlen", [1, 2, 4]) + def test_within_gluon_keeps_its_route(self, monkeypatch, qlen): + """MTP-3 (qlen 4) and draft steps (qlen 1) never divert to FlyDSL here.""" + calls = self._flydsl_on(monkeypatch) + q, kc, vc, md = self._call(qlen) + assert not self._layer()._flydsl_takes_decode(q, kc, vc, None, None, md) + assert not calls + + @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter FlyDSL pa_decode") + @pytest.mark.parametrize( + "case", + ["flydsl_off", "unified_env", "flash", "window", "aiter_declines", "bf16_kv"], + ) + def test_declined_calls_stay_on_unified(self, monkeypatch, case): + import torch + + from atom.model_ops import attention_mha as mha + + self._flydsl_on( + monkeypatch, + on=case != "flydsl_off", + aiter_accepts=case != "aiter_declines", + ) + if case == "unified_env": + monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", True) + layer = self._layer( + use_flash_layout=case == "flash", + sliding_window=128 if case == "window" else -1, + kv_cache_dtype="bf16" if case == "bf16_kv" else "fp8", + ) + q, kc, vc, md = self._call( + 6, kv_dtype=torch.bfloat16 if case == "bf16_kv" else None + ) + assert not layer._flydsl_takes_decode(q, kc, vc, None, None, md) + + def test_unified_passes_the_decision_down(self, monkeypatch): + from atom.model_ops import attention_mha as mha + + monkeypatch.setattr( + mha.PagedAttentionImpl, "_flydsl_takes_decode", lambda *a: 4 + ) + layer = self._layer() + layer.paged_attention_triton = lambda *a, **k: k + q, kc, vc, md = self._call(6) + out = mha.PagedAttentionImpl.paged_attention_unified( + layer, q, None, None, kc, vc, None, None, SimpleNamespace(attn_metadata=md) + ) + assert out == {"flydsl_seqs": 4} + + def test_the_gate_is_asked_once_per_call(self, monkeypatch): + """unified -> paged_attention_triton -> run_pa_decode asks ATOM's FlyDSL + gate (and so aiter's predicate) once, not once in _flydsl_takes_decode + and again in run_pa_decode. Stops at the plan, past the gate.""" + from atom.model_ops import attention_mha as mha + from atom.model_ops import base_attention as ba + + monkeypatch.setattr(ba.envs, "ATOM_PA_FLYDSL", True) + monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", False) + asked = [] + monkeypatch.setattr( + ba, "_flydsl_pa_decode_num_seqs", lambda **kw: asked.append(kw) or 4 + ) + monkeypatch.setattr(ba, "_flydsl_plan_max_batch", lambda: 4096) + monkeypatch.setattr(mha, "dense_decode_splits", lambda n, h: 8) + + class Reached(Exception): + pass + + def at_plan(**kw): + raise Reached(kw["num_seqs"]) + + monkeypatch.setattr(ba, "_flydsl_prepare_explicit_plan", at_plan) + layer = self._layer(scale=1.0) + q, kc, vc, md = self._call(6) + with pytest.raises(Reached) as hit: + mha.PagedAttentionImpl.paged_attention_unified( + *(layer, q, None, None, kc, vc, None, None), + SimpleNamespace(attn_metadata=md), + ) + assert hit.value.args == (4,) + assert len(asked) == 1 + + @pytest.mark.parametrize("per_token", [False, True]) + def test_flydsl_off_names_the_opt_in_once(self, monkeypatch, caplog, per_token): + """A default user past the gluon envelope is told about ATOM_PA_FLYDSL + once, and aiter's predicate is never asked. Per-token FP8 KV: unified + refuses those scales, so the refusal carries the hint and no info line + claims the call runs on unified.""" + import logging + + import torch + + from atom.model_ops import attention_mha as mha + + calls = self._flydsl_on(monkeypatch, on=False) + monkeypatch.setattr(mha, "_flydsl_off_hint_logged", False) + monkeypatch.setattr(mha, "unified_attention", lambda *a, **k: None) + layer = self._layer(scale=1.0, kv_scale=None) + q, kc, vc, md = self._call(6) + md.cu_seqlens_q = md.max_seqlen_k = None + scale = ( + torch.empty(64, self.KV_HEADS, self.BLOCK, device="meta") + if per_token + else None + ) + args = (layer, q, None, None, kc, vc, scale, scale) + ctx = SimpleNamespace(attn_metadata=md) + with caplog.at_level(logging.INFO, logger="atom"): + for _ in range(3): + if per_token: + with pytest.raises(NotImplementedError, match="ATOM_PA_FLYDSL=1"): + mha.PagedAttentionImpl.paged_attention_unified(*args, ctx) + else: + mha.PagedAttentionImpl.paged_attention_unified(*args, ctx) + hints = [r for r in caplog.records if "ATOM_PA_FLYDSL=1" in r.getMessage()] + assert len(hints) == (0 if per_token else 1) + if not per_token: + assert "max_seqlen_q=6 x 8" in hints[0].getMessage() + assert not calls + + @pytest.mark.parametrize("case", ["sinks", "bf16_kv"]) + def test_no_hint_where_atom_keeps_the_call_off_flydsl( + self, monkeypatch, caplog, case + ): + """Sinks, and BF16 KV on a gluon GPU, stay off FlyDSL even with + ATOM_PA_FLYDSL=1, so a default user past the envelope is not told + the flag would route them there.""" + import logging + + import torch + + from atom.model_ops import attention_mha as mha + + calls = self._flydsl_on(monkeypatch, on=False) + monkeypatch.setattr(mha, "_flydsl_off_hint_logged", False) + monkeypatch.setattr(mha, "unified_attention", lambda *a, **k: None) + layer = self._layer( + scale=1.0, + kv_scale=None, + sinks=torch.empty(self.HEADS, device="meta") if case == "sinks" else None, + kv_cache_dtype="bf16" if case == "bf16_kv" else "fp8", + ) + q, kc, vc, md = self._call( + 6, kv_dtype=torch.bfloat16 if case == "bf16_kv" else None + ) + md.cu_seqlens_q = md.max_seqlen_k = None + assert not layer._flydsl_route_open(md) + # The same layer with the condition lifted is open: not vacuous. + assert self._layer()._flydsl_route_open(md) + with caplog.at_level(logging.INFO, logger="atom"): + mha.PagedAttentionImpl.paged_attention_unified( + *(layer, q, None, None, kc, vc, None, None), + SimpleNamespace(attn_metadata=md), + ) + assert not [r for r in caplog.records if "ATOM_PA_FLYDSL" in r.getMessage()] + assert not calls + + def test_no_hint_inside_the_gluon_envelope(self, monkeypatch): + """MTP-3 (qlen 4) keeps gluon; the refusal names no opt-in.""" + import torch + + from atom.model_ops import attention_mha as mha + + self._flydsl_on(monkeypatch, on=False) + layer = self._layer() + q, kc, vc, md = self._call(4) + assert not layer._flydsl_route_open(md) + scale = torch.empty(64, self.KV_HEADS, self.BLOCK, device="meta") + with pytest.raises(NotImplementedError) as err: + mha.PagedAttentionImpl.paged_attention_unified( + *(layer, q, None, None, kc, vc, scale, scale), + SimpleNamespace(attn_metadata=md), + ) + assert "ATOM_PA_FLYDSL" not in str(err.value) + + +def _gpu_flydsl_arch(): + if not (_HAS_CUDA and _HAS_FLYDSL_PA): + return False + from aiter.jit.utils.chip_info import get_gfx_runtime + + return get_gfx_runtime() in ("gfx942", "gfx950") + + +@pytest.mark.skipif( + not _gpu_flydsl_arch(), reason="runs aiter FlyDSL pa_decode on gfx942/gfx950" +) +@pytest.mark.parametrize("qlen", [4, 6]) +def test_qwen35_mtp5_shape_is_accepted_by_real_aiter(monkeypatch, qlen): + """No monkeypatched probe: ATOM's gate accepts the Qwen3.5 MTP-5 call and + aiter's own pa_decode validation and kernel then take it, with per-token + fp8 KV scales and the count the gate returned passed as flydsl_seqs.""" + import torch + + from atom.model_ops import base_attention as ba + + monkeypatch.undo() # drop the autouse gfx950 / gluon pins + monkeypatch.setattr(ba.envs, "ATOM_PA_FLYDSL", True) + dev = torch.device("cuda", 0) + heads, kv_heads, head_dim, block, seqs, blocks_per_seq = 8, 1, 256, 16, 4, 32 + nblk = seqs * blocks_per_seq + q = torch.randn(seqs * qlen, heads, head_dim, device=dev, dtype=torch.bfloat16) + k_cache = ( + torch.randn(nblk, kv_heads, head_dim // 16, block, 16, device=dev) * 0.1 + ).to(_FP8) + v_cache = (torch.randn(nblk, kv_heads, head_dim, block, device=dev) * 0.1).to(_FP8) + k_scale = torch.rand(nblk, kv_heads, block, 1, device=dev) + 0.5 + v_scale = torch.rand(nblk, kv_heads, block, 1, device=dev) + 0.5 + block_tables = torch.arange(nblk, device=dev, dtype=torch.int32).view(seqs, -1) + context_lens = torch.tensor([7, 100, 300, 511], device=dev, dtype=torch.int32) + splits = ba.dense_decode_splits(seqs, kv_heads) + out = torch.empty_like(q) + common = { + "k_cache": k_cache, + "v_cache": v_cache, + "context_lens": context_lens, + "block_tables": block_tables, + "max_seqlen_q": qlen, + "context_partition_size": 256, + "compute_type": _FP8, + "q_scale": None, + } + assert ( + ba.flydsl_decode_seqs( + output=out, + q=q, + k_scale=k_scale, + v_scale=v_scale, + alibi_slopes=None, + sinks=None, + sliding_window=-1, + ps=True, + **common, + ) + == seqs + ) + ba.run_pa_decode( + output=out, + q=q, + softmax_scale=head_dim**-0.5, + k_scale=k_scale, + v_scale=v_scale, + max_context_partition_num=splits, + flydsl_seqs=seqs, + **common, + ) + torch.cuda.synchronize() + assert torch.isfinite(out).all() + + class _FakePlan: """Just the fields the op and the scratch helper read off a real plan.""" @@ -1488,298 +1830,162 @@ def no_flydsl(**kw): class TestFlyDSLCapabilityGate: - """`_flydsl_pa_decode_num_seqs` promises a fallback, not an exception. + """`_flydsl_pa_decode_num_seqs` asks aiter, and only aiter, whether the + kernel takes a call; ATOM adds its routing choices and the sequence count. - Every clause it checks is a hard reject inside aiter, so a clause that goes - missing does not degrade to gluon -- it raises from inside the kernel call - and kills the worker, while the server keeps answering /metrics and the - client warms up until it times out. This gate had no coverage at all, and - two of aiter's rejects were in fact missing from it. + A wrong answer is not a slow path: a call FlyDSL rejects raises from + inside the kernel call and kills the worker. So the arguments aiter is + asked about must be exactly the ones run_pa_decode then passes. """ @staticmethod - def _call(**over): - """A call that passes every clause, with one field overridable.""" + def _call(monkeypatch, accepts=True, gluon=True, **over): + """Run the gate on meta tensors with a recording fake of aiter's + predicate. Returns (result, the predicate's calls).""" import torch - from atom.model_ops.base_attention import _flydsl_pa_decode_num_seqs + from atom.model_ops import base_attention as ba + + calls = [] + + def can_use(*args, **kwargs): + calls.append((args, kwargs)) + return accepts + monkeypatch.setattr(ba, "_flydsl_pa_decode_can_use", lambda: can_use) + monkeypatch.setattr(ba, "pa_decode_gluon_supports_arch", lambda: gluon) q = over.pop("q", None) if q is None: q = torch.empty(16, 16, 128, device="meta", dtype=torch.bfloat16) kw = { - "output": ( - over.pop("output", None) if "output" in over else torch.empty_like(q) - ), + "output": torch.empty_like(q), "q": q, "k_cache": torch.empty(4, 1, 8, 128, 16, device="meta", dtype=_FP8), "v_cache": torch.empty(4, 1, 8, 128, 16, device="meta", dtype=_FP8), - "block_tables": torch.empty(4, 64, device="meta", dtype=torch.int32), - "context_lens": torch.empty(4, device="meta", dtype=torch.int32), + "block_tables": torch.empty(5, 64, device="meta", dtype=torch.int32), + "context_lens": torch.empty(5, device="meta", dtype=torch.int32), "max_seqlen_q": 4, - "max_context_partition_num": 32, "context_partition_size": 256, "compute_type": _FP8, "q_scale": None, + "k_scale": torch.empty(4, 1, 128, 1, device="meta"), + "v_scale": torch.empty(4, 1, 128, 1, device="meta"), "alibi_slopes": None, "sinks": None, "sliding_window": -1, "ps": True, } kw.update(over) - return _flydsl_pa_decode_num_seqs(**kw) + return ba._flydsl_pa_decode_num_seqs(**kw), calls, kw + + def test_aiter_is_asked_with_the_arguments_pa_decode_gets(self, monkeypatch): + """16 rows at query length 4 over 5 padded sequences: 4 sequences, and + aiter sees the sliced rectangle run_pa_decode passes to pa_decode.""" + seqs, calls, kw = self._call(monkeypatch, max_context_length=1000) + assert seqs == 4 + ((args, kwargs),) = calls + assert args[0] is kw["output"] and args[1] is kw["q"] + assert args[2] is kw["k_cache"] and args[3] is kw["v_cache"] + assert args[4].shape == (4,) and args[5].shape == (4, 64) + assert args[6] == 4 + assert kwargs == { + "context_partition_size": 256, + "compute_type": _FP8, + "query_scale": None, + "key_scale": kw["k_scale"], + "value_scale": kw["v_scale"], + "alibi_slopes": None, + "sinks": None, + "sliding_window": 0, + # The bound run_pa_decode hands pa_decode, clamped to the table. + "max_context_length": 1000, + } - def test_the_baseline_call_is_accepted(self, monkeypatch): - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True + def test_the_context_bound_matches_run_pa_decode(self, monkeypatch): + from atom.model_ops import base_attention as ba + + _, calls, kw = self._call(monkeypatch) + assert calls[0][1]["max_context_length"] == ba._flydsl_context_bound( + None, kw["block_tables"], kw["k_cache"] ) - assert self._call() == 4 + + def test_an_aiter_decline_falls_back(self, monkeypatch): + seqs, calls, _ = self._call(monkeypatch, accepts=False) + assert seqs is None and len(calls) == 1 @pytest.mark.parametrize( "field,value", - [ - ("sliding_window", 128), - ("ps", False), - ("compute_type", "bf16"), - ("context_partition_size", 128), - ("max_context_partition_num", 512), - ("q_scale", "present"), - ("alibi_slopes", "present"), - ("sinks", "present"), - ], + [("sliding_window", 128), ("ps", False), ("sinks", "present")], ) - def test_each_aiter_reject_falls_back(self, monkeypatch, field, value): + def test_atom_routing_keeps_these_on_gluon(self, monkeypatch, field, value): import torch - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - if value == "bf16": - value = torch.bfloat16 - elif value == "present": + if value == "present": value = torch.empty(1, device="meta") - assert self._call(**{field: value}) is None, f"{field}={value} must fall back" - - @pytest.mark.parametrize("block_size,ok", [(16, True), (128, True), (256, False)]) - def test_block_size_outside_aiters_whitelist_falls_back( - self, monkeypatch, block_size, ok - ): - """--block-size 256 reaches this site whenever use_triton_attn is set. - - aiter takes block_size from dim -2 and rejects anything but 16/64/128; - the guard used to check only dim -1, so a 256-block deployment routed - into FlyDSL and died there. - """ - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - k = torch.empty(4, 1, 8, block_size, 16, device="meta", dtype=_FP8) - assert (self._call(k_cache=k) == 4) is ok - - def test_unsupported_arch_falls_back(self, monkeypatch): - """aiter builds FlyDSL for gfx942/gfx950 only and raises on the rest.""" - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: False - ) - assert self._call() is None - - def test_row_count_must_divide_evenly(self, monkeypatch): - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - odd = torch.empty(15, 16, 128, device="meta", dtype=torch.bfloat16) - assert self._call(q=odd) is None - - @pytest.mark.parametrize("head_dim", [192, 96, 256]) - def test_head_dim_the_cache_does_not_encode_falls_back(self, monkeypatch, head_dim): - """Two clauses, and the second is the one that was missing. - - 192 and 96 fail ATOM's own whitelist. 256 passes it -- and used to be - asserted as accepted -- but the baseline cache is (.., 8, 128, 16), - whose num_hgroups encodes head_dim 128, and aiter rejects a q that - disagrees with its cache. Pinning that acceptance as the contract is - what would have kept the clause from ever being added. - """ - import torch + seqs, calls, _ = self._call(monkeypatch, **{field: value}) + assert seqs is None and not calls - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - q = torch.empty(16, 16, head_dim, device="meta", dtype=torch.bfloat16) - assert self._call(q=q) is None - - def test_head_dim_matching_the_cache_is_accepted(self, monkeypatch): - """The positive half: 256 is fine once the cache encodes 256.""" - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - q = torch.empty(16, 16, 256, device="meta", dtype=torch.bfloat16) - k = torch.empty(4, 1, 16, 128, 16, device="meta", dtype=_FP8) - v = torch.empty(4, 1, 16, 128, 16, device="meta", dtype=_FP8) - assert self._call(q=q, k_cache=k, v_cache=v) == 4 - - def test_zero_query_length_falls_back_instead_of_dividing(self, monkeypatch): - """`divmod(rows, 0)` is a ZeroDivisionError, not a fallback.""" - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - assert self._call(max_seqlen_q=0) is None - - def test_more_rows_than_sequences_falls_back(self, monkeypatch): - """The upper bound is what stops the per-sequence slices going short. - - `context_lens[:n]` with n past the buffer does not raise -- torch hands - back whatever is there -- so without this clause FlyDSL would be given - a rectangle that silently disagrees with the batch. - """ - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - q = torch.empty(32, 16, 128, device="meta", dtype=torch.bfloat16) - assert self._call(q=q) is None - - @pytest.mark.parametrize("shape", [(4, 1, 8, 128), (4, 1, 8, 128, 8)]) - def test_cache_layout_outside_page16_falls_back(self, monkeypatch, shape): - """A bf16 cache gives x == 8 on the last axis; aiter needs 16.""" + @pytest.mark.parametrize("gluon", [True, False]) + def test_bf16_kv_stays_on_gluon_where_gluon_exists(self, monkeypatch, gluon): import torch - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True + kv = torch.empty(4, 1, 16, 128, 8, device="meta", dtype=torch.bfloat16) + seqs, calls, _ = self._call( + monkeypatch, + gluon=gluon, + k_cache=kv, + v_cache=kv, + compute_type=torch.bfloat16, + k_scale=None, + v_scale=None, ) - k = torch.empty(*shape, device="meta", dtype=_FP8) - assert self._call(k_cache=k) is None + assert (seqs is None) == gluon + assert bool(calls) != gluon - @pytest.mark.parametrize("head_dim,hgroups", [(96, 6), (1152, 72)]) - def test_head_dim_outside_atoms_own_whitelist_falls_back( - self, monkeypatch, head_dim, hgroups + @pytest.mark.parametrize( + "rows,qlen,seqs", + [(15, 4, None), (16, 0, None), (24, 4, None), (16, 4, 4), (5, 1, 5)], + ) + def test_the_rectangle_is_recovered_or_declined( + self, monkeypatch, rows, qlen, seqs ): - """The whitelist clause alone, with the cache made to agree. - - The cases above fail the cache clause too, so deleting the whitelist - left them green. 96 is neither 64 nor a multiple of 128; 1152 is a - multiple but past aiter's 1024 ceiling. - """ + """Rows must split evenly over the query length, into no more + sequences than the padded lengths and block table hold.""" import torch - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - q = torch.empty(16, 16, head_dim, device="meta", dtype=torch.bfloat16) - k = torch.empty(4, 1, hgroups, 128, 16, device="meta", dtype=_FP8) - assert self._call(q=q, k_cache=k, v_cache=k) is None - - def test_a_q_dtype_the_kernel_has_no_arm_for_falls_back(self, monkeypatch): - """aiter picks its scale from q.dtype and raises on anything else.""" - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - q = torch.empty(16, 16, 128, device="meta", dtype=torch.float32) - assert self._call(q=q) is None - - def test_an_output_that_disagrees_with_q_falls_back(self, monkeypatch): - """aiter writes through the output pointer as if it were q's twin.""" - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - other_dtype = torch.empty(16, 16, 128, device="meta", dtype=torch.float16) - assert self._call(output=other_dtype) is None - other_shape = torch.empty(32, 16, 128, device="meta", dtype=torch.bfloat16) - assert self._call(output=other_shape) is None - - def test_a_non_contiguous_head_dim_axis_falls_back(self, monkeypatch): - """Both axes: aiter checks q and output separately (pa_decode:443,448).""" - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - skewed = torch.empty( - 16, 128, 16, device="meta", dtype=torch.bfloat16 - ).transpose(1, 2) - assert skewed.shape == (16, 16, 128) and skewed.stride(2) != 1 - contiguous = torch.empty(16, 16, 128, device="meta", dtype=torch.bfloat16) - assert self._call(q=skewed, output=contiguous) is None - assert self._call(output=skewed) is None - - def test_q_heads_that_do_not_divide_over_kv_heads_fall_back(self, monkeypatch): - """The kernel gives each kv head a whole group; 16 over 3 has no split.""" - import torch - - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - k = torch.empty(4, 3, 8, 128, 16, device="meta", dtype=_FP8) - assert self._call(k_cache=k, v_cache=k) is None + q = torch.empty(rows, 16, 128, device="meta", dtype=torch.bfloat16) + got, calls, _ = self._call(monkeypatch, q=q, output=q, max_seqlen_q=qlen) + assert got == seqs + assert bool(calls) == (seqs is not None) - def test_a_v_cache_of_another_dtype_falls_back(self, monkeypatch): - """One compute_type covers both caches; a split pair raises inside aiter.""" + def test_fewer_table_rows_than_sequences_falls_back(self, monkeypatch): import torch - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - v = torch.empty(4, 1, 8, 128, 16, device="meta", dtype=torch.bfloat16) - assert self._call(v_cache=v) is None + table = torch.empty(3, 64, device="meta", dtype=torch.int32) + seqs, calls, _ = self._call(monkeypatch, block_tables=table) + assert seqs is None and not calls - @pytest.mark.parametrize( - "field", ["k_cache", "v_cache", "block_tables", "context_lens"] - ) - def test_a_non_contiguous_argument_falls_back(self, monkeypatch, field): - """aiter demands all four contiguous (pa_decode.py:464-470). + def test_an_aiter_without_the_predicate_is_refused(self, monkeypatch): + """No fallback copy of aiter's rules: opting in on an old aiter is an + error that names the missing function.""" + import types - Nothing in this tree hands one over -- every cache view is a `.view()`, - which raises rather than returning a skewed tensor -- but the SGLang - bridge takes its pool from someone else and ATOM_PA_FLYDSL routes it - too, so the gate promises a fallback it could not keep. - """ - import torch + from atom.model_ops import base_attention as ba - monkeypatch.setattr( - "atom.model_ops.base_attention._flydsl_arch_supported", lambda: True - ) - skewed = { - "k_cache": lambda: torch.empty( - 4, 1, 8, 16, 128, device="meta", dtype=_FP8 - ).transpose(3, 4), - "v_cache": lambda: torch.empty( - 4, 1, 8, 16, 128, device="meta", dtype=_FP8 - ).transpose(3, 4), - "block_tables": lambda: torch.empty( - 64, 4, device="meta", dtype=torch.int32 - ).t(), - "context_lens": lambda: torch.empty(8, device="meta", dtype=torch.int32)[ - ::2 - ], - }[field]() - assert not skewed.is_contiguous() - assert self._call(**{field: skewed}) is None + old = types.ModuleType("aiter.ops.flydsl.pa_decode") + monkeypatch.setattr(ba.importlib, "import_module", lambda name: old) + with pytest.raises(RuntimeError, match="can_use_flydsl_pa_decode"): + ba._flydsl_pa_decode_can_use() - @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter #4332") - def test_the_arch_probe_binding_exists(self): - """The binding, not the module name. - - Every other test here monkeypatches `_flydsl_arch_supported`, so the - one thing none of them touches is whether it can be computed at all. - aiter moving `get_gfx_runtime` leaves the module importable, leaves all - of these green, and makes a guard that exists to prevent an exception - inside aiter raise one itself on the first decode. - """ - import importlib + @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter FlyDSL pa_decode") + def test_the_predicate_is_aiters(self): + """With an aiter that has it, the gate calls aiter's own function.""" + from atom.model_ops import base_attention as ba - mod = importlib.import_module("aiter.jit.utils.chip_info") - assert callable(getattr(mod, "get_gfx_runtime", None)) + mod = importlib.import_module("aiter.ops.flydsl.pa_decode") + if not hasattr(mod, "can_use_flydsl_pa_decode"): + pytest.skip("this aiter predates can_use_flydsl_pa_decode") + assert ba._flydsl_pa_decode_can_use() is mod.can_use_flydsl_pa_decode class TestFlyDSLConstantsMatchAiter: @@ -1797,17 +2003,6 @@ def _aiter_src(rel): return (pathlib.Path(aiter.__file__).parent / rel).read_text() - @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter #4332") - def test_block_size_whitelist(self): - from atom.model_ops.base_attention import _FLYDSL_PA_BLOCK_SIZES - - src = self._aiter_src("ops/flydsl/pa_decode.py") - assert "block_size not in (16, 64, 128)" in src, ( - "aiter's block_size whitelist moved; _FLYDSL_PA_BLOCK_SIZES is now " - "either over-rejecting or no longer protecting the call" - ) - assert _FLYDSL_PA_BLOCK_SIZES == (16, 64, 128) - @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter #4332") def test_supported_archs(self): from atom.model_ops.base_attention import _FLYDSL_PA_ARCHS @@ -1817,24 +2012,16 @@ def test_supported_archs(self): assert f'"{arch}"' in src, f"aiter no longer names {arch}" @pytest.mark.skipif(not _HAS_FLYDSL_PA, reason="needs aiter #4332") - def test_tile_size_and_batch_cap(self): - """The two limits the comment claimed were pinned and were not. - - _FLYDSL_PA_TILE is the only context_partition_size aiter accepts, and - _FLYDSL_PLAN_MAX_BATCH is where its planner stops. Drift either way is + def test_batch_cap(self): + """_FLYDSL_PLAN_MAX_BATCH is where aiter's planner stops. Drift is silent: too strict falls back more than it must, too loose hands aiter - a call it raises on. + a plan it raises on. """ from atom.model_ops.base_attention import ( - _FLYDSL_PA_TILE, _FLYDSL_PLAN_MAX_BATCH, _FLYDSL_PLAN_SCAN_LDS_BYTES, ) - src = self._aiter_src("ops/flydsl/pa_decode.py") - assert ( - f"context_partition_size={_FLYDSL_PA_TILE}" in src - ), "aiter's accepted partition size moved" plan_src = self._aiter_src("ops/flydsl/kernels/pa_decode_plan.py") assert ( str(_FLYDSL_PLAN_MAX_BATCH) in plan_src @@ -1923,8 +2110,9 @@ def test_unified_arm_reports_no_gluon_route(self, monkeypatch, env, accepts, err monkeypatch.setattr(mha, "dense_decode_splits", lambda n, h: 8) layer = mha.PagedAttentionImpl.__new__(mha.PagedAttentionImpl) layer.sliding_window, layer.kv_cache_dtype, layer.sinks = -1, "fp8", None + layer.use_flash_layout = False layer.layer_num = 0 - layer.paged_attention_triton = lambda *a: "flydsl" + layer.paged_attention_triton = lambda *a, **k: "flydsl" m = self._meta md = SimpleNamespace(context_lens=m(4), block_tables=m(4, 8), max_seqlen_q=1) call = lambda: mha.PagedAttentionImpl.paged_attention_unified( @@ -1945,13 +2133,23 @@ def test_unified_arm_reports_no_gluon_route(self, monkeypatch, env, accepts, err with pytest.raises(NotImplementedError, match=error): call() - def test_use_unified_attn_keeps_decode_off_flydsl(self, monkeypatch): + @pytest.mark.parametrize("unified", [False, True]) + def test_use_unified_attn_keeps_decode_off_flydsl(self, monkeypatch, unified): + """FlyDSL opted in and accepting: only ATOM_USE_UNIFIED_ATTN keeps the + call off it (the False case shows the call would otherwise be taken).""" mha, _ba = self._no_gluon(monkeypatch) - monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", True) + monkeypatch.setattr(mha.envs, "ATOM_PA_FLYDSL", True) + monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", unified) monkeypatch.setattr(mha, "flydsl_decode_seqs", lambda **kw: 4) layer = mha.PagedAttentionImpl.__new__(mha.PagedAttentionImpl) - layer.sliding_window = -1 - assert not layer._flydsl_takes_decode(None, None, None, None) + layer.sliding_window, layer.kv_cache_dtype, layer.sinks = -1, "fp8", None + layer.use_flash_layout = False + layer.num_heads, layer.num_kv_heads = 16, 1 + m = self._meta + md = SimpleNamespace(context_lens=m(4), block_tables=m(4, 8), max_seqlen_q=1) + kc, vc = m(8, 1, 8, 128, 16), m(8, 1, 8, 128, 16) + taken = layer._flydsl_takes_decode(m(4, 16, 128), kc, vc, None, None, md) + assert taken == (False if unified else 4) @pytest.mark.parametrize( "arch, takes_max_batch", [("gfx950", True), ("gfx1250", False)]