[Kernel] Enable Helion backend for Kimi Delta-Attention (#32593)

Co-authored-by: Ethan Che <eche@meta.com>
This commit is contained in:
ethche
2026-08-14 22:22:46 -07:00
committed by GitHub
co-authored by Ethan Che
parent 3adbbec2fd
commit aeee1562e6
16 changed files with 4248 additions and 38 deletions
@@ -37,7 +37,14 @@ 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:
def _accepts(
backend: str,
*,
use_mla: bool,
unified: bool = True,
linear_decode: str | None = None,
linear_prefill: str | None = None,
) -> bool:
"""Run just `_handle_page_major_kv_layout` against a minimal stand-in.
ServerArgs' real constructor pulls in a model config; this exercises the
@@ -53,8 +60,8 @@ def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
"prefill_attention_backend": None,
"decode_attention_backend": None,
"linear_attn_backend": "triton",
"linear_attn_decode_backend": None,
"linear_attn_prefill_backend": None,
"linear_attn_decode_backend": linear_decode,
"linear_attn_prefill_backend": linear_prefill,
"mamba_backend": "triton",
}.items():
object.__setattr__(sa, name, value)
@@ -115,6 +122,17 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
f"{backend} has no dense-id remapping and must be rejected",
)
def test_helion_linear_attention_is_kda_only(self):
for unified in (True, False):
for phase in ("decode", "prefill"):
kwargs = {f"linear_{phase}": "helion"}
self.assertTrue(
_accepts("triton", use_mla=True, unified=unified, **kwargs)
)
self.assertFalse(
_accepts("triton", use_mla=False, unified=unified, **kwargs)
)
if __name__ == "__main__":
unittest.main()