[AMD] support qlen>1 for aiter gluon path for Kimi K3 (#37601)
This commit is contained in:
billishyahao
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user