feat: support deterministic FA4 for GLM-4.7-Flash (#33945)
This commit is contained in:
@@ -1448,9 +1448,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--bf16-gemm-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects <code>cutedsl</code> on SM100/SM103 (Blackwell), otherwise uses cuBLAS via `torch.nn.functional.linear`), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10X; dispatches between the CuTe DSL kernel and cuBLAS).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the backend for unquantized BF16 GEMM operations. Options: <code>auto</code> (default; selects <code>cutedsl</code> on SM10x GPUs, except deterministic inference selects <code>torch</code>; otherwise uses cuBLAS via <code>torch.nn.functional.linear</code>), <code>cutedsl</code> (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the CuTe DSL kernel and cuBLAS), <code>torch</code> (always uses cuBLAS via <code>torch.nn.functional.linear</code>).</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`auto`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>cutedsl</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>cutedsl</code>, <code>torch</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--disable-flashinfer-autotune`</td>
|
||||
|
||||
@@ -2050,6 +2050,7 @@ def _deterministic_is_deepseek_model(view: Any) -> bool:
|
||||
"MistralLarge3ForCausalLM",
|
||||
"PixtralForConditionalGeneration",
|
||||
"GlmMoeDsaForCausalLM",
|
||||
"Glm4MoeLiteForCausalLM",
|
||||
]
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@@ -90,11 +90,18 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
|
||||
|
||||
backend_str = server_args.bf16_gemm_backend
|
||||
if backend_str == "auto" and is_sm100_supported():
|
||||
backend_str = "cutedsl"
|
||||
backend_str = (
|
||||
"torch" if server_args.enable_deterministic_inference else "cutedsl"
|
||||
)
|
||||
|
||||
backend = Bf16GemmBackend(backend_str)
|
||||
|
||||
if backend.is_cutedsl():
|
||||
if server_args.enable_deterministic_inference:
|
||||
raise ValueError(
|
||||
"--bf16-gemm-backend cutedsl is batch-size dependent and cannot "
|
||||
"be combined with --enable-deterministic-inference"
|
||||
)
|
||||
if not is_sm100_supported():
|
||||
raise ValueError("--bf16-gemm-backend cutedsl requires an SM10x GPU")
|
||||
|
||||
|
||||
@@ -1754,7 +1754,7 @@ class ServerArgs:
|
||||
bf16_gemm_backend: A[
|
||||
str,
|
||||
Arg(
|
||||
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the CuTe DSL kernel and cuBLAS), 'torch' (always uses cuBLAS via torch.nn.functional.linear).",
|
||||
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, except deterministic inference selects 'torch'; otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the CuTe DSL kernel and cuBLAS), 'torch' (always uses cuBLAS via torch.nn.functional.linear).",
|
||||
cli_name="--bf16-gemm-backend",
|
||||
choices=BF16_GEMM_BACKEND_CHOICES,
|
||||
),
|
||||
@@ -8082,6 +8082,7 @@ class ServerArgs:
|
||||
"MistralLarge3ForCausalLM",
|
||||
"PixtralForConditionalGeneration",
|
||||
"GlmMoeDsaForCausalLM",
|
||||
"Glm4MoeLiteForCausalLM",
|
||||
]
|
||||
except Exception:
|
||||
pass
|
||||
@@ -8096,11 +8097,11 @@ class ServerArgs:
|
||||
not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND
|
||||
):
|
||||
raise ValueError(
|
||||
f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
|
||||
f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with absorbed-MLA models. But you're using {attention_backend}."
|
||||
)
|
||||
if attention_backend == "fa4" and not is_sm100_or_sm110_supported():
|
||||
raise ValueError(
|
||||
"Deterministic inference with DeepSeek models on the fa4 "
|
||||
"Deterministic inference with absorbed-MLA models on the fa4 "
|
||||
"attention backend requires SM100/SM110: it runs "
|
||||
"absorbed MLA, whose qv argument flash_attn.cute only "
|
||||
"implements on those archs."
|
||||
|
||||
@@ -31,7 +31,9 @@ class TestDeterministicBase(CustomTestCase):
|
||||
def setUpClass(cls):
|
||||
cls.model = cls.get_model()
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
if "--attention-backend" not in cls.get_server_args():
|
||||
# Identity, not a probe for --attention-backend: a subclass that
|
||||
# deliberately leaves the backend unspecified is a real test case.
|
||||
if cls is TestDeterministicBase:
|
||||
raise unittest.SkipTest("Skip the base test class")
|
||||
|
||||
cls.process = popen_launch_server(
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Deterministic inference for GLM-4.7-Flash (Glm4MoeLiteForCausalLM) on Blackwell.
|
||||
|
||||
Two entry paths matter. The fa4 backend is what this model needs for
|
||||
deterministic absorbed MLA, and the unspecified-backend path has to reach
|
||||
triton through the absorbed-MLA arch probe instead of the flashinfer default
|
||||
that every non-absorbed model gets on Blackwell.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_glm4_moe_lite_deterministic.TestGlm4MoeLiteFa4Deterministic
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_deterministic_utils import (
|
||||
COMMON_SERVER_ARGS,
|
||||
TestDeterministicBase,
|
||||
)
|
||||
from sglang.test.test_utils import DEFAULT_URL_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=900, stage="nightly", runner_config="4-gpu-gb300")
|
||||
|
||||
GLM_MODEL = "zai-org/GLM-4.7-Flash"
|
||||
|
||||
# COMMON_SERVER_ARGS is shared module state; copy it. Extending it in place
|
||||
# would leak the fa4 flag into the auto class below and silently make that
|
||||
# test a second fa4 test.
|
||||
SERVER_ARGS = COMMON_SERVER_ARGS + [
|
||||
"--chunked-prefill-size",
|
||||
"2048",
|
||||
"--max-prefill-tokens",
|
||||
"2048",
|
||||
"--mem-fraction-static",
|
||||
"0.8",
|
||||
]
|
||||
|
||||
|
||||
class TestGlm4MoeLiteFa4Deterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return GLM_MODEL
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return SERVER_ARGS + ["--attention-backend", "fa4"]
|
||||
|
||||
|
||||
class TestGlm4MoeLiteAutoBackendDeterministic(TestDeterministicBase):
|
||||
@classmethod
|
||||
def get_model(cls):
|
||||
return GLM_MODEL
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return SERVER_ARGS
|
||||
|
||||
def test_auto_backend_resolves_to_triton(self):
|
||||
# Guards the arch probe itself: if Glm4MoeLiteForCausalLM stopped
|
||||
# counting as absorbed MLA the fill would hand back flashinfer, which
|
||||
# deterministic inference then rejects at launch.
|
||||
info = requests.get(DEFAULT_URL_FOR_TEST + "/server_info").json()
|
||||
self.assertEqual(info["attention_backend"], "triton")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user