[Fix] Allow flashinfer_sparse_mla DSA backend for HiSparse on SM120 FP8 KV (#33075)
This commit is contained in:
@@ -163,7 +163,7 @@ python3 -m sglang.launch_server \
|
|||||||
--hisparse-config='{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10, "swap_in_block_size": 960}'
|
--hisparse-config='{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10, "swap_in_block_size": 960}'
|
||||||
```
|
```
|
||||||
|
|
||||||
> **Note**: For DSA models, `--kv-cache-dtype` defaults to `auto`, which resolves to `fp8_e4m3` on SM100+ (Blackwell) and `bfloat16` on older architectures. The DSA decode backend is automatically selected based on KV dtype (`bfloat16` → `flashmla_sparse`, `fp8_e4m3` → `flashmla_kv`). DSA backend flags apply only to DSA models; DeepSeek V4 uses its own `dsv4` attention backend.
|
> **Note**: For DSA models, `--kv-cache-dtype` defaults to `auto`, which resolves to `fp8_e4m3` on SM100+ (Blackwell) and `bfloat16` on older architectures. The DSA decode backend is automatically selected based on KV dtype (`bfloat16` → `flashmla_sparse`, `fp8_e4m3` → `flashmla_kv`), except for GLM DSA models on SM120/SM121 with `fp8_e4m3`, which use `flashinfer_sparse_mla`. DSA backend flags apply only to DSA models; DeepSeek V4 uses its own `dsv4` attention backend.
|
||||||
|
|
||||||
### Benchmark
|
### Benchmark
|
||||||
|
|
||||||
@@ -189,6 +189,7 @@ python3 -m sglang.bench_serving \
|
|||||||
- The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse.
|
- The prefill instance does not need `--enable-hisparse`; it is unaware of HiSparse.
|
||||||
- On the decode instance, `--enable-hisparse` and `--hisparse-config` are required for HiSparse.
|
- On the decode instance, `--enable-hisparse` and `--hisparse-config` are required for HiSparse.
|
||||||
- For DSA models, `--kv-cache-dtype bfloat16` uses `flashmla_sparse`, and `--kv-cache-dtype fp8_e4m3` uses `flashmla_kv`.
|
- For DSA models, `--kv-cache-dtype bfloat16` uses `flashmla_sparse`, and `--kv-cache-dtype fp8_e4m3` uses `flashmla_kv`.
|
||||||
|
- On SM120/SM121 (e.g. RTX PRO 6000, RTX 5090) with GLM DSA models and `--kv-cache-dtype fp8_e4m3`, both DSA backends resolve to `flashinfer_sparse_mla`, which is the only DSA kernel available on that architecture. HiSparse accepts it there; no extra flag is needed.
|
||||||
- For DeepSeek V4, DSA backend flags are not applicable. DeepSeek V4 uses the `dsv4` attention backend and `fp8_e4m3` KV cache by default.
|
- For DeepSeek V4, DSA backend flags are not applicable. DeepSeek V4 uses the `dsv4` attention backend and `fp8_e4m3` KV cache by default.
|
||||||
- `host_to_device_ratio` should be configured based on the host machine's available memory. For example:
|
- `host_to_device_ratio` should be configured based on the host machine's available memory. For example:
|
||||||
- **~1 TB** host memory → `host_to_device_ratio: 5`
|
- **~1 TB** host memory → `host_to_device_ratio: 5`
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = {
|
HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = {
|
||||||
"bfloat16": {"flashmla_sparse"},
|
"bfloat16": {"flashmla_sparse"},
|
||||||
"fp8_e4m3": {"flashmla_kv"},
|
"fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla"},
|
||||||
}
|
}
|
||||||
HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "aiter"}
|
HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "aiter"}
|
||||||
HISPARSE_KV_CACHE_DTYPES = ("bfloat16", "fp8_e4m3")
|
HISPARSE_KV_CACHE_DTYPES = ("bfloat16", "fp8_e4m3")
|
||||||
@@ -32,7 +32,7 @@ def _hisparse_allowed_backends(kv_cache_dtype: str) -> set[str]:
|
|||||||
if _is_hip():
|
if _is_hip():
|
||||||
return HISPARSE_ROCM_DSA_BACKENDS
|
return HISPARSE_ROCM_DSA_BACKENDS
|
||||||
return HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE.get(
|
return HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE.get(
|
||||||
kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"}
|
kv_cache_dtype, {"flashmla_sparse", "flashmla_kv", "flashinfer_sparse_mla"}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -56,9 +56,8 @@ def validate_hisparse_dsa_backend(
|
|||||||
f"HiSparse supports DSA {label} backend(s) {sorted(allowed_backends)} "
|
f"HiSparse supports DSA {label} backend(s) {sorted(allowed_backends)} "
|
||||||
f"on this platform with --kv-cache-dtype={kv_cache_dtype}, "
|
f"on this platform with --kv-cache-dtype={kv_cache_dtype}, "
|
||||||
f"but got --dsa-{label}-backend={backend}. "
|
f"but got --dsa-{label}-backend={backend}. "
|
||||||
f"Please use --dsa-{label}-backend="
|
f"Please use one of {sorted(allowed_backends)}, or omit the option "
|
||||||
f"{_hisparse_default_backend(kv_cache_dtype)} "
|
"to let SGLang pick a backend for this platform."
|
||||||
"or omit it."
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -661,6 +661,23 @@ class TestHiSparseDsaBackendPolicy(unittest.TestCase):
|
|||||||
self.assertEqual(resolved["dsa_prefill_backend"], "flashmla_kv")
|
self.assertEqual(resolved["dsa_prefill_backend"], "flashmla_kv")
|
||||||
self.assertEqual(resolved["dsa_decode_backend"], "flashmla_kv")
|
self.assertEqual(resolved["dsa_decode_backend"], "flashmla_kv")
|
||||||
|
|
||||||
|
@patch("sglang.srt.server_args.is_hip", return_value=False)
|
||||||
|
def test_hisparse_accepts_flashinfer_sparse_mla_on_cuda_fp8(self, _mock_is_hip):
|
||||||
|
"""SM120 GLM DSA resolves both DSA backends to flashinfer_sparse_mla, so
|
||||||
|
the fp8 hisparse allow-set must admit it or --enable-hisparse cannot
|
||||||
|
start there at all. The device/arch narrowing happens later, in
|
||||||
|
_validate_flashinfer_sparse_mla_backend."""
|
||||||
|
server_args = ServerArgs(
|
||||||
|
model_path="dummy",
|
||||||
|
enable_hisparse=True,
|
||||||
|
kv_cache_dtype="fp8_e4m3",
|
||||||
|
dsa_prefill_backend="flashinfer_sparse_mla",
|
||||||
|
dsa_decode_backend="flashinfer_sparse_mla",
|
||||||
|
)
|
||||||
|
|
||||||
|
server_args._validate_hisparse_dsa_backend("dsa_prefill_backend", "prefill")
|
||||||
|
server_args._validate_hisparse_dsa_backend("dsa_decode_backend", "decode")
|
||||||
|
|
||||||
@patch("sglang.srt.server_args.is_hip", return_value=True)
|
@patch("sglang.srt.server_args.is_hip", return_value=True)
|
||||||
def test_hisparse_defaults_to_tilelang_on_rocm(self, _mock_is_hip):
|
def test_hisparse_defaults_to_tilelang_on_rocm(self, _mock_is_hip):
|
||||||
resolved = self._resolve("bfloat16")
|
resolved = self._resolve("bfloat16")
|
||||||
|
|||||||
Reference in New Issue
Block a user