[unified-memory] Support fa3, the default MLA backend on pre-Blackwell hosts (#33046)
This commit is contained in:
Cheng Wan
@@ -1,28 +1,26 @@
|
||||
"""Kimi-Linear (MLA full attention + KDA linear attention) served from the
|
||||
unified memory pool.
|
||||
|
||||
`--enable-unified-memory` replaces the statically-partitioned hybrid pools with
|
||||
one byte buffer split dynamically between the full-attention KV sub-pool and the
|
||||
per-request KDA state sub-pool. For an MLA model the full side is exposed as
|
||||
DENSE per-layer views (`build_dense_mla_views`) and every loc the kernels see is
|
||||
a translated virtual id, so the whole read/write path differs from the static
|
||||
pool: `translate_kv_loc_dense` for kv_indices and the cuda-graph write loc,
|
||||
`HybridLinearKVPool._full_translate` for the model-level MLA entry points, and
|
||||
page-envelope relocation on allocator compaction.
|
||||
Under `--enable-unified-memory` the MLA full side is exposed as DENSE per-layer
|
||||
views and every loc the kernels see is a translated virtual id, so the whole
|
||||
read/write path differs from the static pool. The unit tests pin that pool in
|
||||
isolation; this is the end-to-end guard. `test_prefix_cache_branching` carries
|
||||
most of the weight: a radix hit replays virtual locs whose physical pages may
|
||||
have moved under compaction.
|
||||
|
||||
None of that is covered by the CPU/GPU unit tests, which pin the pool in
|
||||
isolation. This is the end-to-end guard: accuracy must match the static-pool
|
||||
baseline, and the prefix-cache branching case must still hit, since a radix hit
|
||||
replays virtual locs whose physical pages may have moved under compaction.
|
||||
No `--attention-backend` is pinned on purpose -- the test runs whatever the host
|
||||
resolves to (`fa3` on this suite's H100 runner, also the H200 default). Both
|
||||
defects found in review on #32972 were reachable only under a resolved default,
|
||||
which a pinned test hides by construction.
|
||||
|
||||
Reference numbers on 2x H200 TP2, GSM8K 400 examples (2026-07-30):
|
||||
static pools 0.915, `--enable-unified-memory` 0.900 (1 sigma ~= 0.015) -- both with
|
||||
the attention backend pinned to triton, as this test runs it. For reference the
|
||||
paged MLA kernels land in the same band on a single B300 TP1, GSM8K 200, unified
|
||||
(2026-07-31): 0.915 with flashinfer prefill+decode, 0.900 with trtllm_mla. Those
|
||||
are not exercised here (see the comment on `other_args`).
|
||||
Nightly-only: it needs a second full 48B server launch, which is too much to add
|
||||
to per-PR CI on top of the existing Kimi-Linear e2e coverage.
|
||||
Reference GSM8K, all with `--enable-unified-memory`:
|
||||
- 2x H200 TP2, resolved default (fa3): 0.917 @400, vs 0.915 static (1 sigma
|
||||
~= 0.015). This file as written scores 0.920 @200.
|
||||
- 2x H200 TP2, `--attention-backend triton`: 0.900 @400.
|
||||
- 1x B300 TP1: 0.915 flashinfer, 0.900 trtllm_mla, @200.
|
||||
|
||||
Nightly-only: a second full 48B server launch is too much for per-PR CI on top
|
||||
of the existing Kimi-Linear e2e coverage.
|
||||
|
||||
python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory.py -v
|
||||
"""
|
||||
@@ -45,7 +43,7 @@ class TestKimiLinearUnifiedMemory(
|
||||
model = KIMI_LINEAR_MODEL
|
||||
cache_chunk_size = 64
|
||||
# Same bar as the static-pool Kimi-Linear e2e test: unified memory must not
|
||||
# cost accuracy (measured 0.900 vs 0.915 static, see the module docstring).
|
||||
# cost accuracy (measured 0.917 vs 0.915 static, see the module docstring).
|
||||
gsm8k_score_threshold = 0.88
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
@@ -54,22 +52,6 @@ class TestKimiLinearUnifiedMemory(
|
||||
"--chunked-prefill-size",
|
||||
"2048",
|
||||
"--enable-unified-memory",
|
||||
# Pinned because the resolved default is not portable: on pre-Blackwell
|
||||
# (this suite's runner is H100) an unspecified backend resolves to `fa3`,
|
||||
# which cannot read the dense views at all, so the un-pinned form fails at
|
||||
# startup with the page-major allowlist assertion. Unified memory on such a
|
||||
# host currently REQUIRES an explicit compatible --attention-backend; that
|
||||
# is a real usability gap, tracked separately, not something this test can
|
||||
# paper over.
|
||||
#
|
||||
# Consequence to keep in mind: pinning triton means this test does NOT
|
||||
# cover the paged MLA backends (trtllm_mla / flashinfer / cutedsl_mla /
|
||||
# tokenspeed_mla), which is where dense-id translation bugs live -- a
|
||||
# captured flashinfer decode reading untranslated virtual ids scored GSM8K
|
||||
# 0.000 on a healthy server. Those paths are covered by the unit tests plus
|
||||
# manual B300 runs; an sm100-gated case here would close the gap.
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -20,9 +20,12 @@ block table filled with DENSE page ids:
|
||||
|
||||
dense_page(virtual_page) = v2p[virtual_page] * layer_num
|
||||
|
||||
`create_flashmla_kv_indices_triton` does that in-kernel via `v2p_ptr` / `PAGE_MULT`
|
||||
(trtllm_mla / cutedsl_mla / tokenspeed_mla), and the flashinfer_mla updaters do it
|
||||
by post-gathering `translate_kv_loc_dense` over the token-level kv_indices.
|
||||
Three backend families reach that same formula by different routes:
|
||||
- `create_flashmla_kv_indices_triton` in-kernel via `v2p_ptr` / `PAGE_MULT`
|
||||
(trtllm_mla / cutedsl_mla / tokenspeed_mla);
|
||||
- the flashinfer_mla updaters, post-gathering `translate_kv_loc_dense` over the
|
||||
token-level kv_indices;
|
||||
- `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
|
||||
|
||||
Covered here:
|
||||
- kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
|
||||
@@ -30,7 +33,10 @@ Covered here:
|
||||
ragged sequence lengths and a non-identity v2p permutation;
|
||||
- padded block-table lanes never index the v2p table out of bounds;
|
||||
- the token-level dense translate the flashinfer updaters apply agrees with the
|
||||
page-level block table the trtllm path builds.
|
||||
page-level block table the trtllm path builds;
|
||||
- fa3's fused metadata kernels agree with the same reference, on both the
|
||||
page_size == 1 fast path (which is what Kimi-Linear takes: fa3 imposes no
|
||||
page-size constraint) and the general path.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py -v
|
||||
"""
|
||||
@@ -199,6 +205,105 @@ class TestDenseBlockTable(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
|
||||
class TestFa3MetadataDenseBlockTable(unittest.TestCase):
|
||||
"""fa3 folds the unified remap into `normal_decode_set_metadata`, the fused
|
||||
gather that writes its captured-decode page table, so the kernel itself has
|
||||
to get the mapping right. Two kernels back it: a page_size == 1 / no-SWA fast
|
||||
path (what Kimi-Linear takes, since fa3 imposes no page-size constraint) and
|
||||
a general one.
|
||||
"""
|
||||
|
||||
def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048):
|
||||
from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata
|
||||
|
||||
maker = TestDenseBlockTable._make_batch
|
||||
rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx)
|
||||
v2p_arg = v2p_full if v2p else None
|
||||
|
||||
max_pages = (max_ctx + page_size - 1) // page_size
|
||||
page_table = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
|
||||
cache_seqlens = torch.zeros((bs,), dtype=torch.int32, device=_DEV)
|
||||
cu_seqlens_k = torch.zeros((bs + 1,), dtype=torch.int32, device=_DEV)
|
||||
strided = torch.arange(0, max_ctx, page_size, device=_DEV)
|
||||
max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
|
||||
|
||||
normal_decode_set_metadata(
|
||||
cache_seqlens,
|
||||
cu_seqlens_k,
|
||||
page_table,
|
||||
rt,
|
||||
rpi,
|
||||
strided,
|
||||
max_seq_pages,
|
||||
sl.to(torch.int64),
|
||||
0,
|
||||
page_size,
|
||||
v2p_page_table=v2p_arg,
|
||||
kernel_page_multiplier=mult,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult)
|
||||
return page_table, want, sl
|
||||
|
||||
def _assert_live_prefix(self, got, want, sl, page_size):
|
||||
"""The kernel contract only (re)writes each row's live page prefix; the
|
||||
tail keeps stale values that consumers bound by cache_seqlens."""
|
||||
for r in range(got.shape[0]):
|
||||
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
|
||||
self.assertTrue(
|
||||
torch.equal(got[r, :n_pages].long(), want[r, :n_pages]),
|
||||
f"row {r} (page_size={page_size}):\n"
|
||||
f"got ={got[r, :n_pages]}\nwant={want[r, :n_pages]}",
|
||||
)
|
||||
|
||||
def test_identity_when_hooks_absent(self):
|
||||
"""Static pool: no v2p, multiplier 1 -> byte-identical to pre-change."""
|
||||
for page_size in (1, 64):
|
||||
got, want, sl = self._run(page_size, v2p=False, mult=1)
|
||||
self._assert_live_prefix(got, want, sl, page_size)
|
||||
|
||||
def test_dense_mapping_ps1_fast_path(self):
|
||||
got, want, sl = self._run(1, v2p=True, mult=_LAYERS)
|
||||
self._assert_live_prefix(got, want, sl, 1)
|
||||
|
||||
def test_dense_mapping_general_path(self):
|
||||
got, want, sl = self._run(64, v2p=True, mult=_LAYERS)
|
||||
self._assert_live_prefix(got, want, sl, 64)
|
||||
|
||||
def test_single_full_attention_layer(self):
|
||||
"""multiplier 1 with a real v2p: the gather alone is the translation."""
|
||||
for page_size in (1, 64):
|
||||
got, want, sl = self._run(page_size, v2p=True, mult=1)
|
||||
self._assert_live_prefix(got, want, sl, page_size)
|
||||
virtual = _reference(
|
||||
*TestDenseBlockTable._make_batch(self, page_size)[:3],
|
||||
page_size,
|
||||
v2p=None,
|
||||
mult=1,
|
||||
)
|
||||
self.assertFalse(
|
||||
torch.equal(want, virtual),
|
||||
"test batch degenerated: v2p is the identity on the pages used",
|
||||
)
|
||||
|
||||
def test_agrees_with_flashmla_block_table(self):
|
||||
"""fa3 and trtllm_mla build the same table two different ways; a
|
||||
disagreement means one family is addressing the wrong pages."""
|
||||
for page_size in (1, 64):
|
||||
got, _, sl = self._run(page_size, v2p=True, mult=_LAYERS)
|
||||
rt, rpi, sl2, v2p = TestDenseBlockTable._make_batch(self, page_size)
|
||||
other = _fill_block_table(
|
||||
rt, rpi, sl2, page_size, v2p=v2p, mult=_LAYERS
|
||||
).long()
|
||||
for r in range(got.shape[0]):
|
||||
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
|
||||
self.assertTrue(
|
||||
torch.equal(got[r, :n_pages].long(), other[r, :n_pages]),
|
||||
f"fa3 and flashmla block tables disagree (row {r}, ps={page_size})",
|
||||
)
|
||||
|
||||
|
||||
class TestUnifiedMLAHookDetection(unittest.TestCase):
|
||||
"""`unified_mla_hooks` decides whether the paged MLA backends translate at
|
||||
all. Getting the predicate wrong is silent: the block table and KV write loc
|
||||
@@ -207,7 +312,7 @@ class TestUnifiedMLAHookDetection(unittest.TestCase):
|
||||
|
||||
@staticmethod
|
||||
def _probe(**attrs):
|
||||
from sglang.srt.layers.attention.flashinfer_mla_backend import (
|
||||
from sglang.srt.layers.attention.unified_mem_hooks import (
|
||||
unified_mla_hooks,
|
||||
)
|
||||
|
||||
|
||||
@@ -16,12 +16,15 @@
|
||||
The page-major envelope K/V views are strided, which only the Triton attention
|
||||
kernels read. The one exception is the unified-memory MLA pool: it exposes each
|
||||
layer as a DENSE contiguous view (`build_dense_mla_views`), so the paged MLA
|
||||
backends (`trtllm_mla` and its `cutedsl_mla` / `tokenspeed_mla` subclasses, plus
|
||||
`flashinfer`'s MLA backend) can read it directly once their kv_indices / block
|
||||
tables are remapped to dense ids.
|
||||
backends can read it directly once their kv_indices / block tables are remapped
|
||||
to dense ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
|
||||
`cutedsl_mla` / `tokenspeed_mla` subclasses.
|
||||
|
||||
Pinned here so the exception cannot silently widen to a backend that has no
|
||||
dense-id remapping (`fa3`, `flashmla`, ...) or leak into the MHA path.
|
||||
dense-id remapping (`flashmla`, `cutlass_mla`, ...) or leak into the MHA path.
|
||||
`fa3` matters most: it is the resolved default on pre-Blackwell hosts, so it is
|
||||
the one entry whose absence used to make `--enable-unified-memory` fail to boot
|
||||
under its own default configuration.
|
||||
|
||||
python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
|
||||
"""
|
||||
@@ -66,9 +69,15 @@ def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
|
||||
|
||||
class TestPageMajorBackendAllowlist(unittest.TestCase):
|
||||
# Wired for the dense per-layer MLA views (see the module docstring).
|
||||
DENSE_MLA_BACKENDS = ("trtllm_mla", "flashinfer", "cutedsl_mla", "tokenspeed_mla")
|
||||
DENSE_MLA_BACKENDS = (
|
||||
"fa3",
|
||||
"trtllm_mla",
|
||||
"flashinfer",
|
||||
"cutedsl_mla",
|
||||
"tokenspeed_mla",
|
||||
)
|
||||
# No dense-id remapping: must stay rejected until they get one.
|
||||
UNWIRED_BACKENDS = ("fa3", "flashmla", "cutlass_mla", "trtllm_mha", "aiter")
|
||||
UNWIRED_BACKENDS = ("flashmla", "cutlass_mla", "trtllm_mha", "aiter")
|
||||
|
||||
def test_triton_always_allowed(self):
|
||||
for use_mla in (True, False):
|
||||
|
||||
Reference in New Issue
Block a user