[Kernel] Enable Helion backend for Kimi Delta-Attention (#32593)
Co-authored-by: Ethan Che <eche@meta.com>
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user