feat: support deterministic FA4 for GLM-4.7-Flash (#33945)

This commit is contained in:
Yuzhen Zhou
2026-08-12 16:57:34 +08:00
committed by GitHub
parent 1f008dc226
commit 2d76d537e5
6 changed files with 86 additions and 7 deletions
@@ -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")
+4 -3
View File
@@ -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()