diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index efb3ff07a..b71a7c3a4 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1467,10 +1467,16 @@ class GroupCoordinator: # Bypass the function if we are using only 1 GPU. if self.world_size == 1: return input_ - # Broadcast. - torch.distributed.broadcast( - input_, src=self.ranks[src], group=self.device_group - ) + + # Always use pynccl to avoid capturing hip graph failure on torch + # version smaller than or equal to 2.11 + if is_hip() and self.pynccl_comm is not None and not self.pynccl_comm.disabled: + self.pynccl_comm.broadcast(input_, src=src) + else: + # Broadcast. + torch.distributed.broadcast( + input_, src=self.ranks[src], group=self.device_group + ) return input_ def broadcast_object(self, obj: Optional[Any] = None, src: int = 0): diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 1f9bc87d3..0d7c62d46 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -1018,32 +1018,25 @@ class AiterAttnBackend(AttentionBackend): k_descale, ): k_buffer = self.token_to_kv_pool.get_key_buffer(layer.layer_id) - q_mla = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) max_q_len = self.forward_metadata.max_q_len or 1 - if ( - prefer_mla_gluon_decode( - head_pad_mode=getattr(self, "head_pad_mode", "none"), - num_head=getattr(self, "num_head", layer.tp_q_head_num), - kv_cache_dtype=self.kv_cache_dtype, - ) - and max_q_len == 1 + if prefer_mla_gluon_decode( + head_pad_mode=getattr(self, "head_pad_mode", "none"), + num_head=getattr(self, "num_head", layer.tp_q_head_num), + kv_cache_dtype=self.kv_cache_dtype, ): - kv_scale = self._resolve_fp8_kv_scale_float(layer, k_descale) - min_kv_seq_len = self._resolve_mla_gluon_min_kv_seq_len(forward_batch) - gluon_out = mla_gluon_decode( - q=q_mla, + return mla_gluon_decode( + q=q, k_buffer=k_buffer, layer=layer, kv_indices=self.forward_metadata.kv_indices, kv_indptr=self.forward_metadata.kv_indptr, - seq_lens=forward_batch.seq_lens, sm_scale=layer.scaling, - kv_scale=kv_scale, - min_kv_seq_len=min_kv_seq_len, + kv_scale=self._resolve_fp8_kv_scale_float(layer, k_descale), + min_kv_seq_len=self._resolve_mla_gluon_min_kv_seq_len(forward_batch), + qlen=max_q_len, ) - if gluon_out is not None: - return gluon_out work_metadata = self.forward_metadata.work_metadata work_indptr = self.forward_metadata.work_indptr @@ -1054,7 +1047,7 @@ class AiterAttnBackend(AttentionBackend): num_kv_splits = self.forward_metadata.num_kv_splits return self._mla_decode_fwd_with_head_pad( - q_mla, + q, k_buffer.view(-1, 1, 1, layer.qk_head_dim), layer, qo_indptr=self.forward_metadata.qo_indptr, @@ -2471,6 +2464,25 @@ class AiterAttnBackend(AttentionBackend): ) return o elif forward_batch.forward_mode.is_target_verify(): + if prefer_mla_gluon_decode( + head_pad_mode=getattr(self, "head_pad_mode", "none"), + num_head=getattr(self, "num_head", layer.tp_q_head_num), + kv_cache_dtype=self.kv_cache_dtype, + ): + return mla_gluon_decode( + q=q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), + k_buffer=K_Buffer, + layer=layer, + kv_indices=self.forward_metadata.kv_indices, + kv_indptr=self.forward_metadata.kv_indptr, + sm_scale=layer.scaling, + kv_scale=self._resolve_fp8_kv_scale_float(layer, k_descale), + min_kv_seq_len=self._resolve_mla_gluon_min_kv_seq_len( + forward_batch + ), + qlen=self.forward_metadata.max_q_len or 1, + ) + work_metadata = self.forward_metadata.work_metadata work_indptr = self.forward_metadata.work_indptr work_info_set = self.forward_metadata.work_info_set diff --git a/python/sglang/srt/layers/attention/aiter_mla_gluon.py b/python/sglang/srt/layers/attention/aiter_mla_gluon.py index 50200c8f8..c2e9c524d 100644 --- a/python/sglang/srt/layers/attention/aiter_mla_gluon.py +++ b/python/sglang/srt/layers/attention/aiter_mla_gluon.py @@ -2,7 +2,7 @@ Uses aiter ``mla_gluon`` when import succeeds and Triton Gluon exposes ``cga_layout`` (needs Triton >= 3.7). Falls back to the caller (zero-pad + ``mla_decode_fwd``) when -Gluon is unavailable or fails at runtime. +Gluon is unavailable. Requires aiter ``main`` with ROCm/aiter #4480 (batch>1 ``bh16bn128``) and #4555 (decode CUDA graph KV splits). SGLang probes import + Triton API only; aiter version @@ -11,9 +11,9 @@ is not pinned at build time. from __future__ import annotations +import functools import inspect import logging -from dataclasses import dataclass from typing import TYPE_CHECKING, Optional import torch @@ -26,133 +26,49 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -_mla_gluon_fn = None -_mla_gluon_import_failed = False -_capability_cache: Optional[MlaGluonCapability] = None - -@dataclass(frozen=True) -class MlaGluonCapability: - """Runtime probe of aiter/Triton Gluon prerequisites for h12 + FP8 decode.""" - - enabled_by_env: bool - import_ok: bool - triton_version: str - triton_cga_layout_ok: bool - ready: bool - summary: str - - def missing_for_ready(self) -> list[str]: - missing = [] - if not self.enabled_by_env: - missing.append("SGLANG_AITER_MLA_GLUON=0") - if not self.import_ok: - missing.append("aiter.ops.triton.gluon.mla_gluon import") - if not self.triton_cga_layout_ok: - missing.append( - f"Triton Gluon cga_layout (have {self.triton_version or 'unknown'}, need >= 3.7)" - ) - return missing - - -def _triton_version() -> str: +@functools.lru_cache(maxsize=1) +def _gluon_fn(): + """aiter mla gluon entry point return None if disabled""" + if not envs.SGLANG_AITER_MLA_GLUON.get(): + logger.info("aiter mla gluon is disabled manually.") + return None try: import triton - - return getattr(triton, "__version__", "unknown") - except Exception: - return "missing" - - -def _triton_cga_layout_ok() -> bool: - try: import triton.experimental.gluon.language as gl - - return "cga_layout" in inspect.signature(gl.PaddedSharedLayout).parameters - except Exception: - return False - - -def _gluon_runtime_ok() -> bool: - return mla_gluon_available() and _triton_cga_layout_ok() - - -def _mla_gluon_enabled() -> bool: - return envs.SGLANG_AITER_MLA_GLUON.get() - - -def probe_mla_gluon_capability(*, force_refresh: bool = False) -> MlaGluonCapability: - global _capability_cache - if _capability_cache is not None and not force_refresh: - return _capability_cache - - enabled = _mla_gluon_enabled() - triton_ver = _triton_version() - import_ok = mla_gluon_available() if enabled else False - cga_ok = _triton_cga_layout_ok() - ready = enabled and import_ok and cga_ok - - if ready: - summary = f"Gluon MLA h12+fp8 ready (Triton={triton_ver})" - else: - cap = MlaGluonCapability( - enabled_by_env=enabled, - import_ok=import_ok, - triton_version=triton_ver, - triton_cga_layout_ok=cga_ok, - ready=False, - summary="", - ) - missing = cap.missing_for_ready() - summary = ( - "Gluon MLA h12+fp8 disabled; fallback to zero-pad mla_decode_fwd " - f"({'; '.join(missing)})" + from aiter.ops.triton.gluon.mla_gluon import mla_gluon + except ImportError as exc: + logger.info("aiter mla gluon import error message: %s", exc) + return None + # mla_gluon builds its shared layouts with cga_layout, added in Triton 3.7; + # older Triton fails at compile time with an opaque error instead. + if "cga_layout" not in inspect.signature(gl.PaddedSharedLayout).parameters: + logger.info( + "aiter mla gluon is disabled due to triton %s has no Gluon cga_layout (need >= 3.7)", + getattr(triton, "__version__", "unknown"), ) + return None + return mla_gluon - _capability_cache = MlaGluonCapability( - enabled_by_env=enabled, - import_ok=import_ok, - triton_version=triton_ver, - triton_cga_layout_ok=cga_ok, - ready=ready, - summary=summary, + +def log_mla_gluon_capability(log: logging.Logger | None = None) -> None: + """Report whether Gluon MLA decode is valid; the reason is logged by _gluon_fn.""" + ready = _gluon_fn() is not None + (log or logger).info( + "aiter mla gluon is %s", + "enabled" if ready else "disabled; falling back to zero-pad mla_decode_fwd", ) - return _capability_cache -def log_mla_gluon_capability(log: logging.Logger | None = None) -> MlaGluonCapability: - cap = probe_mla_gluon_capability() - (log or logger).info(cap.summary) - if not cap.ready: - for item in cap.missing_for_ready(): - (log or logger).info(" missing: %s", item) - return cap - - -def _in_cuda_graph_capture() -> bool: - try: - return bool(torch.cuda.is_current_stream_capturing()) - except Exception: - return False - - -def mla_gluon_available() -> bool: - if not _mla_gluon_enabled(): - return False - global _mla_gluon_fn, _mla_gluon_import_failed - if _mla_gluon_import_failed: - return False - if _mla_gluon_fn is not None: - return True - try: - from aiter.ops.triton.gluon.mla_gluon import mla_gluon as fn - - _mla_gluon_fn = fn - return True - except ImportError: - _mla_gluon_import_failed = True - logger.warning("mla_gluon import failed; Gluon MLA decode disabled.") - return False +def prefer_mla_gluon_decode( + *, head_pad_mode: str, num_head: int, kv_cache_dtype: torch.dtype +) -> bool: + return ( + head_pad_mode == "zero" + and num_head == 12 + and kv_cache_dtype == fp8_dtype + and _gluon_fn() is not None + ) def mla_gluon_decode( @@ -162,89 +78,46 @@ def mla_gluon_decode( layer: RadixAttention, kv_indices: torch.Tensor, kv_indptr: torch.Tensor, - seq_lens: torch.Tensor, sm_scale: float, + min_kv_seq_len: int, kv_scale: float = 1.0, - min_kv_seq_len: Optional[int] = None, + qlen: int = 1, ) -> Optional[torch.Tensor]: - """Run Gluon MLA decode for fused Q [B, H, 576] and MLA KV pool. - - Returns output [B, H, v_head_dim] on success, or None to fall back. - - ``min_kv_seq_len`` must be supplied by the caller during CUDA graph capture - (no GPU->CPU sync from ``seq_lens``). For eager decode, omit it to derive - from ``seq_lens`` when safe. + """Run Gluon MLA decode for fused Q [num_tokens, H, 576] and MLA KV pool. + Returns [num_tokens, H, v_head_dim], or None when Gluon is unavailable. """ - if not mla_gluon_available(): + mla_gluon = _gluon_fn() + if mla_gluon is None: return None - batch_size = q.shape[0] - + num_head = layer.tp_q_head_num kv_lora_rank = layer.v_head_dim qk_rope_head_dim = layer.qk_head_dim - kv_lora_rank + batch_size = q.shape[0] // qlen + q_nope, q_pe = torch.split(q, [kv_lora_rank, qk_rope_head_dim], dim=-1) + if qlen > 1: + # Splitting the leading dim is a stride change, so these stay views of + # the non-contiguous torch.split outputs. + q_nope = q_nope.view(batch_size, qlen, num_head, kv_lora_rank) + q_pe = q_pe.view(batch_size, qlen, num_head, qk_rope_head_dim) + o = q.new_empty((batch_size, qlen, num_head, kv_lora_rank)) + else: + o = q.new_empty((batch_size, num_head, kv_lora_rank)) - o = q.new_empty((batch_size, layer.tp_q_head_num, kv_lora_rank)) - - kv_c = k_buffer.view(-1, layer.qk_head_dim) - if min_kv_seq_len is None: - if _in_cuda_graph_capture(): - logger.warning( - "mla_gluon_decode: min_kv_seq_len missing during CUDA graph capture" - ) - min_kv_seq_len = 1 - elif seq_lens.numel(): - min_kv_seq_len = int(seq_lens.max().item()) - else: - min_kv_seq_len = 1 - - try: - _mla_gluon_fn( - q_nope, - q_pe, - kv_c, - o, - kv_indices, - kv_indptr, - sm_scale, - k_pe=None, - kv_pe_offset=kv_lora_rank, - use_2d_view=False, - kv_scale=kv_scale, - min_kv_seq_len=min_kv_seq_len, - ) - return o - except Exception as exc: - logger.warning( - "mla_gluon decode failed (num_head=%s, kv_dtype=%s, batch=%s): %s; " - "falling back to zero-pad mla_decode_fwd", - layer.tp_q_head_num, - k_buffer.dtype, - batch_size, - exc, - ) - return None - - -def prefer_mla_gluon_decode( - *, head_pad_mode: str, num_head: int, kv_cache_dtype: torch.dtype -) -> bool: - """Route Kimi-style h12 zero-pad MLA decode through Gluon when FP8 KV holds. - - ``head_pad_mode == "zero"`` selects the legacy ``mla_decode_fwd`` padding - topology (N heads padded to 16). Gluon is only validated for ``num_head == 12`` - today; other zero-pad head counts must stay on zero-pad + ``mla_decode_fwd``. - """ - if not _mla_gluon_enabled(): - return False - if head_pad_mode == "zero" and num_head == 12 and kv_cache_dtype == fp8_dtype: - return _gluon_runtime_ok() - return False - - -def reset_mla_gluon_state_for_test() -> None: - """Test helper: clear import/probe caches.""" - global _mla_gluon_fn, _mla_gluon_import_failed, _capability_cache - _mla_gluon_fn = None - _mla_gluon_import_failed = False - _capability_cache = None + mla_gluon( + q_nope, + q_pe, + k_buffer.view(-1, layer.qk_head_dim), + o, + kv_indices, + kv_indptr, + sm_scale, + k_pe=None, + kv_pe_offset=kv_lora_rank, + use_2d_view=False, + kv_scale=kv_scale, + min_kv_seq_len=min_kv_seq_len, + ) + # Hand back the caller's flat [num_tokens, H, v] layout either way. + return o.flatten(0, 1) if qlen > 1 else o diff --git a/test/registered/attention/test_aiter_gluon_h12_fp8.py b/test/registered/attention/test_aiter_gluon_h12_fp8.py new file mode 100644 index 000000000..aa759cb3b --- /dev/null +++ b/test/registered/attention/test_aiter_gluon_h12_fp8.py @@ -0,0 +1,321 @@ +"""Unit tests for the aiter Gluon MLA path: h12 + FP8 routing, MTP (qlen>1) +shaping, and dispatch against the zero-pad ``mla_decode_fwd`` fallback. + +Mocked throughout — no real aiter/Triton kernel is invoked. +""" + +import unittest +from unittest import mock + +import torch + +from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype +from sglang.srt.environ import envs +from sglang.srt.layers.attention import aiter_mla_gluon as mod +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +_GLUON_FN = "sglang.srt.layers.attention.aiter_mla_gluon._gluon_fn" + + +def _fake_gluon_modules(*, cga_layout: bool): + """sys.modules entries that make _gluon_fn() see a usable (or too-old) Triton. + + ``inspect.signature`` is called on ``gl.PaddedSharedLayout``, so that one has + to be a real callable rather than a Mock — the parameter list is exactly what + the Triton >= 3.7 probe reads. + """ + if cga_layout: + + def padded_shared_layout(*, cga_layout=None): + pass + + else: + + def padded_shared_layout(): + pass + + gl = mock.MagicMock() + gl.PaddedSharedLayout = padded_shared_layout + + # `import a.b.c as x` binds via getattr on the parent module, so the parents + # have to point at these exact objects -- a bare MagicMock parent would + # auto-create a different child and shadow them. + triton_gluon = mock.MagicMock(language=gl) + triton_experimental = mock.MagicMock(gluon=triton_gluon) + triton = mock.MagicMock(__version__="3.7.0", experimental=triton_experimental) + + fake_fn = mock.Mock(name="mla_gluon") + aiter_mla = mock.MagicMock(mla_gluon=fake_fn) + aiter_gluon = mock.MagicMock(mla_gluon=aiter_mla) + aiter_triton = mock.MagicMock(gluon=aiter_gluon) + aiter_ops = mock.MagicMock(triton=aiter_triton) + aiter = mock.MagicMock(ops=aiter_ops) + + return fake_fn, { + "triton": triton, + "triton.experimental": triton_experimental, + "triton.experimental.gluon": triton_gluon, + "triton.experimental.gluon.language": gl, + "aiter": aiter, + "aiter.ops": aiter_ops, + "aiter.ops.triton": aiter_triton, + "aiter.ops.triton.gluon": aiter_gluon, + "aiter.ops.triton.gluon.mla_gluon": aiter_mla, + } + + +class TestGluonAvailability(CustomTestCase): + """_gluon_fn() resolves the kernel, or None with the reason logged once.""" + + def setUp(self): + mod._gluon_fn.cache_clear() + + def tearDown(self): + mod._gluon_fn.cache_clear() + + def test_none_when_env_disabled(self): + with envs.SGLANG_AITER_MLA_GLUON.override(False): + self.assertIsNone(mod._gluon_fn()) + + def test_none_when_import_fails(self): + # A None entry in sys.modules makes `import ...` raise ImportError. + with mock.patch.dict("sys.modules", {"aiter.ops.triton.gluon.mla_gluon": None}): + self.assertIsNone(mod._gluon_fn()) + + def test_none_when_triton_lacks_cga_layout(self): + _fn, modules = _fake_gluon_modules(cga_layout=False) + with mock.patch.dict("sys.modules", modules): + self.assertIsNone(mod._gluon_fn()) + + def test_returns_kernel_when_ready(self): + fake_fn, modules = _fake_gluon_modules(cga_layout=True) + with mock.patch.dict("sys.modules", modules): + self.assertIs(mod._gluon_fn(), fake_fn) + + def test_result_is_cached(self): + fake_fn, modules = _fake_gluon_modules(cga_layout=True) + with mock.patch.dict("sys.modules", modules): + first = mod._gluon_fn() + # Second call must not re-probe: the fake modules are gone by now, so a + # re-probe would return None instead of the cached kernel. + self.assertIs(mod._gluon_fn(), first) + + +class TestPreferMlaGluonDecode(CustomTestCase): + """Only the validated h12 + zero-pad + FP8 topology may route to Gluon.""" + + def _prefer(self, **kwargs): + args = dict(head_pad_mode="zero", num_head=12, kv_cache_dtype=fp8_dtype) + args.update(kwargs) + return mod.prefer_mla_gluon_decode(**args) + + def test_true_for_h12_zero_pad_fp8(self): + with mock.patch(_GLUON_FN, return_value=mock.Mock()): + self.assertTrue(self._prefer()) + + def test_false_when_gluon_unavailable(self): + with mock.patch(_GLUON_FN, return_value=None): + self.assertFalse(self._prefer()) + + def test_false_for_other_head_counts(self): + with mock.patch(_GLUON_FN, return_value=mock.Mock()): + self.assertFalse(self._prefer(num_head=10)) + self.assertFalse(self._prefer(num_head=16)) + + def test_false_for_non_zero_pad_topology(self): + with mock.patch(_GLUON_FN, return_value=mock.Mock()): + self.assertFalse(self._prefer(head_pad_mode="repeat")) + self.assertFalse(self._prefer(head_pad_mode="none")) + + def test_false_for_non_fp8_kv(self): + with mock.patch(_GLUON_FN, return_value=mock.Mock()): + self.assertFalse(self._prefer(kv_cache_dtype=torch.bfloat16)) + + +def _layer(num_head=12, qk_head_dim=576, v_head_dim=512): + layer = mock.Mock() + layer.tp_q_head_num = num_head + layer.qk_head_dim = qk_head_dim + layer.v_head_dim = v_head_dim + layer.scaling = 0.125 + layer.logit_cap = 0.0 + layer.layer_id = 0 + return layer + + +class TestMlaGluonDecodeShapes(CustomTestCase): + """Plain decode stays 3-D; target-verify (qlen>1) goes in as 4-D MTP.""" + + def _call(self, *, num_tokens, qlen): + layer = _layer() + q = torch.zeros(num_tokens, 12, 576, dtype=torch.bfloat16) + captured = {} + + def fake_kernel(q_nope, q_pe, kv_c, o, *args, **kwargs): + captured["q_nope"] = q_nope.shape + captured["q_pe"] = q_pe.shape + captured["o"] = o.shape + + with mock.patch(_GLUON_FN, return_value=fake_kernel): + out = mod.mla_gluon_decode( + q=q, + k_buffer=torch.zeros(64, 576, dtype=torch.bfloat16), + layer=layer, + kv_indices=torch.zeros(64, dtype=torch.int32), + kv_indptr=torch.zeros(5, dtype=torch.int32), + sm_scale=layer.scaling, + min_kv_seq_len=128, + qlen=qlen, + ) + return out, captured + + def test_plain_decode_uses_3d(self): + out, cap = self._call(num_tokens=4, qlen=1) + self.assertEqual(cap["q_nope"], torch.Size([4, 12, 512])) + self.assertEqual(cap["q_pe"], torch.Size([4, 12, 64])) + self.assertEqual(cap["o"], torch.Size([4, 12, 512])) + self.assertEqual(out.shape, torch.Size([4, 12, 512])) + + def test_verify_uses_4d_mtp(self): + # 4 requests x 8 draft tokens, the DSPARK block-size-7 shape. + out, cap = self._call(num_tokens=32, qlen=8) + self.assertEqual(cap["q_nope"], torch.Size([4, 8, 12, 512])) + self.assertEqual(cap["q_pe"], torch.Size([4, 8, 12, 64])) + self.assertEqual(cap["o"], torch.Size([4, 8, 12, 512])) + # The caller's contract is the flat [num_tokens, H, v] layout. + self.assertEqual(out.shape, torch.Size([32, 12, 512])) + + def test_mtp_views_do_not_copy(self): + """The 4-D q must stay a view of the caller's tensor, not a copy.""" + layer = _layer() + q = torch.zeros(32, 12, 576, dtype=torch.bfloat16) + seen = {} + + def fake_kernel(q_nope, q_pe, *args, **kwargs): + seen["nope_ptr"] = q_nope.data_ptr() + seen["pe_ptr"] = q_pe.data_ptr() + + with mock.patch(_GLUON_FN, return_value=fake_kernel): + mod.mla_gluon_decode( + q=q, + k_buffer=torch.zeros(64, 576, dtype=torch.bfloat16), + layer=layer, + kv_indices=torch.zeros(64, dtype=torch.int32), + kv_indptr=torch.zeros(5, dtype=torch.int32), + sm_scale=layer.scaling, + min_kv_seq_len=128, + qlen=8, + ) + self.assertEqual(seen["nope_ptr"], q.data_ptr()) + self.assertEqual(seen["pe_ptr"], q[..., 512:].data_ptr()) + + def test_returns_none_when_gluon_unavailable(self): + with mock.patch(_GLUON_FN, return_value=None): + out = mod.mla_gluon_decode( + q=torch.zeros(4, 12, 576, dtype=torch.bfloat16), + k_buffer=torch.zeros(64, 576, dtype=torch.bfloat16), + layer=_layer(), + kv_indices=torch.zeros(64, dtype=torch.int32), + kv_indptr=torch.zeros(5, dtype=torch.int32), + sm_scale=0.125, + min_kv_seq_len=128, + ) + self.assertIsNone(out) + + +class TestForwardMlaDecodeDispatch(CustomTestCase): + """_forward_mla_decode picks Gluon or the zero-pad ASM path, never both.""" + + def _make_backend(self, max_q_len=1): + from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend + + be = AiterAttnBackend.__new__(AiterAttnBackend) + be.num_head = 12 + be.kv_cache_dtype = fp8_dtype + be.head_pad_mode = "zero" + be.num_head_padded = 16 + be.forward_metadata = mock.Mock( + max_q_len=max_q_len, + kv_indices=torch.zeros(4, dtype=torch.int32), + kv_indptr=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32), + kv_last_page_len=torch.ones(4, dtype=torch.int32), + qo_indptr=torch.arange(5, dtype=torch.int32), + work_metadata=None, + work_indptr=None, + work_info_set=None, + reduce_indptr=None, + reduce_final_map=None, + reduce_partial_map=None, + num_kv_splits=None, + ) + be.token_to_kv_pool = mock.Mock( + get_key_buffer=lambda _lid: torch.zeros(8, 576, dtype=fp8_dtype) + ) + be._resolve_fp8_kv_scale_float = mock.Mock(return_value=1.0) + be._resolve_mla_gluon_min_kv_seq_len = mock.Mock(return_value=128) + be._mla_decode_fwd_with_head_pad = mock.Mock( + return_value=torch.zeros(4, 12, 512) + ) + return be + + @mock.patch( + "sglang.srt.layers.attention.aiter_backend.prefer_mla_gluon_decode", + return_value=False, + ) + @mock.patch("sglang.srt.layers.attention.aiter_backend.mla_gluon_decode") + def test_uses_asm_when_gluon_not_preferred(self, mock_gluon, _prefer): + be = self._make_backend() + out = be._forward_mla_decode( + torch.zeros(4, 12, 576, dtype=torch.bfloat16), + _layer(), + mock.Mock(), + k_descale=1.0, + ) + mock_gluon.assert_not_called() + be._mla_decode_fwd_with_head_pad.assert_called_once() + self.assertIs(out, be._mla_decode_fwd_with_head_pad.return_value) + + @mock.patch( + "sglang.srt.layers.attention.aiter_backend.prefer_mla_gluon_decode", + return_value=True, + ) + @mock.patch("sglang.srt.layers.attention.aiter_backend.mla_gluon_decode") + def test_uses_gluon_output_when_preferred(self, mock_gluon, _prefer): + gluon_out = torch.ones(4, 12, 512) + mock_gluon.return_value = gluon_out + be = self._make_backend() + out = be._forward_mla_decode( + torch.zeros(4, 12, 576, dtype=torch.bfloat16), + _layer(), + mock.Mock(), + k_descale=1.0, + ) + mock_gluon.assert_called_once() + be._mla_decode_fwd_with_head_pad.assert_not_called() + self.assertIs(out, gluon_out) + + @mock.patch( + "sglang.srt.layers.attention.aiter_backend.prefer_mla_gluon_decode", + return_value=True, + ) + @mock.patch("sglang.srt.layers.attention.aiter_backend.mla_gluon_decode") + def test_passes_max_q_len_as_qlen(self, mock_gluon, _prefer): + """Target-verify must reach the kernel as qlen, not be silently dropped: + the ASM fallback cannot serve this topology above qSeqLen 4 at all.""" + mock_gluon.return_value = torch.ones(32, 12, 512) + be = self._make_backend(max_q_len=8) + be._forward_mla_decode( + torch.zeros(32, 12, 576, dtype=torch.bfloat16), + _layer(), + mock.Mock(), + k_descale=1.0, + ) + self.assertEqual(mock_gluon.call_args.kwargs["qlen"], 8) + be._mla_decode_fwd_with_head_pad.assert_not_called() + + +if __name__ == "__main__": + unittest.main()