[unified-memory] Let Kimi-Linear use the paged MLA attention backends (#32972)

test/registered/unit/server_args/test_page_major_backend_allowlist.py

@@ -0,0 +1,111 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""`--enable-page-major-kv-layout` full-attention backend allowlist.
+
+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.
+
+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.
+
+    python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
+"""
+
+import unittest
+
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
+    """Run just `_handle_page_major_kv_layout` against a minimal stand-in.
+
+    ServerArgs' real constructor pulls in a model config; this exercises the
+    single handler under test with the fields it reads.
+    """
+    sa = ServerArgs.__new__(ServerArgs)
+    for name, value in {
+        "enable_unified_memory": unified,
+        # The unified pool sets this itself; without it the flag must be explicit
+        # or the handler returns before reaching the allowlist.
+        "enable_page_major_kv_layout": not unified,
+        "attention_backend": backend,
+        "prefill_attention_backend": None,
+        "decode_attention_backend": None,
+        "linear_attn_backend": "triton",
+        "linear_attn_decode_backend": None,
+        "linear_attn_prefill_backend": None,
+        "mamba_backend": "triton",
+    }.items():
+        object.__setattr__(sa, name, value)
+    sa.use_mla_backend = lambda: use_mla
+    sa._resolved_attention_backends = lambda: [backend]
+    try:
+        ServerArgs._handle_page_major_kv_layout(sa)
+        return True
+    except AssertionError:
+        return False
+
+
+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")
+    # No dense-id remapping: must stay rejected until they get one.
+    UNWIRED_BACKENDS = ("fa3", "flashmla", "cutlass_mla", "trtllm_mha", "aiter")
+
+    def test_triton_always_allowed(self):
+        for use_mla in (True, False):
+            self.assertTrue(_accepts("triton", use_mla=use_mla))
+
+    def test_dense_mla_backends_allowed_under_unified_mla(self):
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertTrue(
+                _accepts(backend, use_mla=True),
+                f"{backend} should be allowed with the unified-memory MLA pool",
+            )
+
+    def test_dense_mla_backends_rejected_for_mha(self):
+        """The dense-view exception is MLA-only -- MHA sub-pools stay strided."""
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertFalse(
+                _accepts(backend, use_mla=False),
+                f"{backend} must stay rejected for a non-MLA model",
+            )
+
+    def test_dense_mla_backends_rejected_without_unified_memory(self):
+        """Plain --enable-page-major-kv-layout (no unified pool) keeps the
+        strided views, so only Triton can read them."""
+        for backend in self.DENSE_MLA_BACKENDS:
+            self.assertFalse(
+                _accepts(backend, use_mla=True, unified=False),
+                f"{backend} must stay rejected without --enable-unified-memory",
+            )
+
+    def test_unwired_backends_always_rejected(self):
+        for backend in self.UNWIRED_BACKENDS:
+            for use_mla in (True, False):
+                self.assertFalse(
+                    _accepts(backend, use_mla=use_mla),
+                    f"{backend} has no dense-id remapping and must be rejected",
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()