fix: sync server_args.kv_cache_dtype when detecting FP8 KV cache

[zack041]

Motivation

fixes #18290 #12298

Source of error:
In sglang/srt/model_executor/model_runner.py

    def configure_kv_cache_dtype(self):
        if self.server_args.kv_cache_dtype == "auto":
            quant_config = getattr(self.model, "quant_config", None)
            kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None)
            if (
                isinstance(kv_cache_quant_algo, str)
                and kv_cache_quant_algo.upper() == "FP8"
            ):
                if _is_hip:
                    self.kv_cache_dtype = fp8_dtype
                    self.server_args.kv_cache_dtype = "FP8"
                else:
                    self.kv_cache_dtype = torch.float8_e4m3fn
                    self.server_args.kv_cache_dtype = "FP8"
            else:
                self.kv_cache_dtype = self.dtype

While configure_kv_cache_dtype correctly handles the update of self.kv_cache_dtype for modelopt case when server_args.kv_cache_dtype is "auto", it did not update itself. server_args.kv_cache_dtype is used in flashattention_backend.py to determine the data type for kv cache:

if self.kv_cache_dtype_str != "auto" and layer.head_dim <= 256:
            if layer.k_scale is not None:
                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
                k_descale = layer.k_scale.expand(descale_shape)
                v_descale = layer.v_scale.expand(descale_shape)
            q = q.to(self.kv_cache_dtype)
            q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
            k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None

self.kv_cache_dtype_str is assigned with server_args.kv_cache_dtype.

Modifications

To solve, the server_args.kv_cache_dtype is also updated to FP8 along with kv_cache_dtype.

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[b8zhong]

Thanks! I think this is the right solution, compared to the other ones. Could you maybe try to use a util (e.g mapping torch dtype to str for KV cache here)

[b8zhong]

/tag-and-rerun-ci

[zack041]

Thanks! I think this is the right solution, compared to the other ones. Could you maybe try to use a util (e.g mapping torch dtype to str for KV cache here)

Added a mapping dict!

[b8zhong]

Hi if you could add a test for this Modelopt model on SM90, that would be really great. @zack041 Thanks a lot!

[zack041]

Hi if you could add a test for this Modelopt model on SM90, that would be really great. @zack041 Thanks a lot!

Sure! I'll add the test in a follow-up PR.