[unified-memory] Support fa3, the default MLA backend on pre-Blackwell hosts (#33046)

This commit is contained in:
Cheng Wan
2026-07-31 11:46:46 -07:00
committed by GitHub
parent 26486a957d
commit d3222bcc3a
9 changed files with 288 additions and 95 deletions
@@ -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):