sglang 0.2.14.post2__py3-none-any.whl → 0.3.0__py3-none-any.whl

sglang/api.py

@@ -78,6 +78,7 @@ def gen(
     choices: Optional[List[str]] = None,
     choices_method: Optional[ChoicesSamplingMethod] = None,
     regex: Optional[str] = None,
+    json_schema: Optional[str] = None,
 ):
     """Call the model to generate. See the meaning of the arguments in docs/en/sampling_params.md"""
 
@@ -114,6 +115,7 @@ def gen(
         return_text_in_logprobs,
         dtype,
         regex,
+        json_schema,
     )
 
 

sglang/bench_latency.py

@@ -11,26 +11,34 @@ python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct
 ## plot the results in series of lines:
 python -m sglang.bench_latency --result-filename out.jsonl --graph-sql="select run_name, batch_size, prefill_throughput from results"
 
-
 # Usage (correctness test):
 python -m sglang.bench_latency --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct
 
 ## Reference output (of the correctness test above, can be gpu dependent):
-prefill logits (first half) tensor([[-10.0312,  -9.5000,   0.8936,  ...,  -4.9414,  -3.2402,  -3.3633],
-        [-10.0312,  -9.5000,   0.8936,  ...,  -4.9414,  -3.2402,  -3.3633],
-        [ -9.1875, -10.2500,   2.7109,  ...,  -4.3359,  -4.0664,  -4.1328]],
-       device='cuda:0', dtype=torch.float16)
-prefill logits (final) tensor([[-8.3203, -7.1211,  3.3379,  ..., -4.9570, -4.1328, -3.4141],
-        [-8.9062, -9.0156,  4.1445,  ..., -4.9922, -4.4961, -4.0742],
-        [-9.6328, -9.0547,  4.0117,  ..., -5.3047, -4.7148, -4.4609]],
-       device='cuda:0', dtype=torch.float16)
-<s> The capital of France is.
+input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]]
+
+prefill logits (first half): tensor([[-10.0312,  -9.5000,   0.8931,  ...,  -4.9414,  -3.2422,  -3.3633],
+        [-10.0312,  -9.5000,   0.8931,  ...,  -4.9414,  -3.2422,  -3.3633],
+        [ -9.1875, -10.2500,   2.7129,  ...,  -4.3359,  -4.0664,  -4.1328]],
+       device='cuda:0')
+
+prefill logits (final): tensor([[-8.3125, -7.1172,  3.3457,  ..., -4.9570, -4.1328, -3.4141],
+        [-8.9141, -9.0156,  4.1445,  ..., -4.9922, -4.4961, -4.0781],
+        [-9.6328, -9.0547,  4.0195,  ..., -5.3047, -4.7148, -4.4570]],
+       device='cuda:0')
+
+========== Prompt 0 ==========
+<s> The capital of France is Paris.
 The capital of the United States is Washington, D.C.
 
-<s> The capital of the United Kindom is.
+
+========== Prompt 1 ==========
+<s> The capital of the United Kindom is London.
 The capital of the United Kingdom is London.
 The capital of the
-<s> Today is a sunny day and I like go for a walk in the park.
+
+========== Prompt 2 ==========
+<s> Today is a sunny day and I like to go for a walk in the park.
 I'm going to the park
 """
 
@@ -200,16 +208,16 @@ def extend(reqs, model_runner):
         tree_cache=None,
     )
     batch.prepare_for_extend(model_runner.model_config.vocab_size)
-    output = model_runner.forward(batch, ForwardMode.EXTEND)
-    next_token_ids = batch.sample(output.next_token_logits)
-    return next_token_ids, output.next_token_logits, batch
+    sample_output, logits_output = model_runner.forward(batch, ForwardMode.EXTEND)
+    next_token_ids = sample_output.batch_next_token_ids.tolist()
+    return next_token_ids, logits_output.next_token_logits, batch
 
 
 def decode(input_token_ids, batch, model_runner):
-    batch.prepare_for_decode(input_token_ids.cpu().numpy())
-    output = model_runner.forward(batch, ForwardMode.DECODE)
-    next_token_ids = batch.sample(output.next_token_logits)
-    return next_token_ids, output.next_token_logits
+    batch.prepare_for_decode(input_token_ids)
+    sample_output, logits_output = model_runner.forward(batch, ForwardMode.DECODE)
+    next_token_ids = sample_output.batch_next_token_ids.tolist()
+    return next_token_ids, logits_output.next_token_logits
 
 
 @torch.inference_mode()
@@ -225,12 +233,12 @@ def correctness_test(
 
     # Prepare inputs
     input_ids, reqs = prepare_inputs_for_correctness_test(bench_args, tokenizer)
-    rank_print(f"{input_ids=}")
+    rank_print(f"\n{input_ids=}\n")
 
     if bench_args.cut_len > 0:
         # Prefill
         next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
-        rank_print("prefill logits (first half)", next_token_logits)
+        rank_print(f"prefill logits (first half): {next_token_logits} \n")
 
     # Prepare extend inputs
     reqs = prepare_extend_inputs_for_correctness_test(
@@ -239,7 +247,7 @@ def correctness_test(
 
     # Extend
     next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
-    rank_print("prefill logits (final)", next_token_logits)
+    rank_print(f"prefill logits (final): {next_token_logits} \n")
 
     # Decode
     output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))]
@@ -250,7 +258,8 @@ def correctness_test(
 
     # Print
     for i in range(len(reqs)):
-        rank_print(tokenizer.decode(output_ids[i]))
+        rank_print(f"========== Prompt {i} ==========")
+        rank_print(tokenizer.decode(output_ids[i]), "\n")
 
 
 @torch.inference_mode()
@@ -292,6 +301,7 @@ def latency_test_run_once(
     measurement_results["prefill_throughput"] = throughput
 
     # Decode
+    decode_latencies = []
     for i in range(output_len):
         torch.cuda.synchronize()
         tic = time.time()
@@ -300,17 +310,18 @@ def latency_test_run_once(
         latency = time.time() - tic
         tot_latency += latency
         throughput = batch_size / latency
+        decode_latencies.append(latency)
         if i < 5:
             rank_print(
                 f"Decode.  latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s"
             )
-    avg_decode_latency = (tot_latency - prefill_latency) / output_len
-    avg_decode_throughput = batch_size / avg_decode_latency
+    med_decode_latency = np.median(decode_latencies)
+    med_decode_throughput = batch_size / med_decode_latency
     rank_print(
-        f"Decode.  avg latency: {avg_decode_latency:6.5f} s, avg throughput: {avg_decode_throughput:9.2f} token/s"
+        f"Decode.  median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s"
     )
-    measurement_results["avg_decode_latency"] = avg_decode_latency
-    measurement_results["avg_decode_throughput"] = avg_decode_throughput
+    measurement_results["median_decode_latency"] = med_decode_latency
+    measurement_results["median_decode_throughput"] = med_decode_throughput
 
     throughput = (input_len + output_len) * batch_size / tot_latency
     rank_print(

sglang/lang/backend/runtime_endpoint.py

@@ -4,7 +4,7 @@ from typing import List, Optional
 
 from sglang.global_config import global_config
 from sglang.lang.backend.base_backend import BaseBackend
-from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path
 from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
 from sglang.lang.interpreter import StreamExecutor
 from sglang.lang.ir import (
@@ -23,6 +23,7 @@ class RuntimeEndpoint(BaseBackend):
         base_url: str,
         api_key: Optional[str] = None,
         verify: Optional[str] = None,
+        chat_template_name: Optional[str] = None,
     ):
         super().__init__()
         self.support_concate_and_append = True
@@ -39,9 +40,12 @@ class RuntimeEndpoint(BaseBackend):
         self._assert_success(res)
         self.model_info = res.json()
 
-        self.chat_template = get_chat_template_by_model_path(
-            self.model_info["model_path"]
-        )
+        if chat_template_name:
+            self.chat_template = get_chat_template(chat_template_name)
+        else:
+            self.chat_template = get_chat_template_by_model_path(
+                self.model_info["model_path"]
+            )
 
     def get_model_name(self):
         return self.model_info["model_path"]

sglang/lang/interpreter.py

@@ -673,6 +673,7 @@ class StreamExecutor:
             "return_text_in_logprobs",
             "dtype",
             "regex",
+            "json_schema",
         ]:
             value = getattr(sampling_params, item, None)
             if value is not None:
@@ -854,6 +855,8 @@ class ProgramState:
         return self.stream_executor.get_meta_info(name)
 
     def __iadd__(self, other):
+        if other is None:
+            raise ValueError("Tried to append None to state.")
         self.stream_executor.submit(other)
         return self
 

sglang/lang/ir.py

@@ -30,6 +30,7 @@ class SglSamplingParams:
     logprob_start_len: Optional[int] = (None,)
     top_logprobs_num: Optional[int] = (None,)
     return_text_in_logprobs: Optional[bool] = (None,)
+    json_schema: Optional[str] = None
 
     # for constrained generation, not included in to_xxx_kwargs
     dtype: Optional[str] = None
@@ -51,6 +52,7 @@ class SglSamplingParams:
             self.logprob_start_len,
             self.top_logprobs_num,
             self.return_text_in_logprobs,
+            self.json_schema,
         )
 
     def to_openai_kwargs(self):
@@ -121,6 +123,7 @@ class SglSamplingParams:
             "presence_penalty": self.presence_penalty,
             "ignore_eos": self.ignore_eos,
             "regex": self.regex,
+            "json_schema": self.json_schema,
         }
 
 
@@ -425,6 +428,7 @@ class SglGen(SglExpr):
         return_text_in_logprobs: Optional[bool] = None,
         dtype: Optional[type] = None,
         regex: Optional[str] = None,
+        json_schema: Optional[str] = None,
     ):
         """Call the model to generate. See the meaning of the arguments in docs/en/sampling_params.md"""
         super().__init__()
@@ -446,6 +450,7 @@ class SglGen(SglExpr):
             return_text_in_logprobs=return_text_in_logprobs,
             dtype=dtype,
             regex=regex,
+            json_schema=json_schema,
         )
 
     def __repr__(self):

sglang/launch_server_llavavid.py

@@ -10,17 +10,17 @@ if __name__ == "__main__":
     args = parser.parse_args()
     server_args = ServerArgs.from_cli_args(args)
 
-    model_overide_args = {}
-    model_overide_args["mm_spatial_pool_stride"] = 2
-    model_overide_args["architectures"] = ["LlavaVidForCausalLM"]
-    model_overide_args["num_frames"] = 16
-    model_overide_args["model_type"] = "llavavid"
-    if model_overide_args["num_frames"] == 32:
-        model_overide_args["rope_scaling"] = {"factor": 2.0, "type": "linear"}
-        model_overide_args["max_sequence_length"] = 4096 * 2
-        model_overide_args["tokenizer_model_max_length"] = 4096 * 2
-        model_overide_args["model_max_length"] = 4096 * 2
+    model_override_args = {}
+    model_override_args["mm_spatial_pool_stride"] = 2
+    model_override_args["architectures"] = ["LlavaVidForCausalLM"]
+    model_override_args["num_frames"] = 16
+    model_override_args["model_type"] = "llavavid"
+    if model_override_args["num_frames"] == 32:
+        model_override_args["rope_scaling"] = {"factor": 2.0, "type": "linear"}
+        model_override_args["max_sequence_length"] = 4096 * 2
+        model_override_args["tokenizer_model_max_length"] = 4096 * 2
+        model_override_args["model_max_length"] = 4096 * 2
     if "34b" in args.model_path.lower():
-        model_overide_args["image_token_index"] = 64002
+        model_override_args["image_token_index"] = 64002
 
-    launch_server(server_args, model_overide_args, None)
+    launch_server(server_args, model_override_args, None)

sglang/srt/configs/__init__.py

@@ -0,0 +1,5 @@
+from sglang.srt.configs.exaone import ExaoneConfig
+
+__all__ = [
+    "ExaoneConfig",
+]

sglang/srt/configs/exaone.py

@@ -0,0 +1,195 @@
+# coding=utf-8
+# Copyright 2024 The LG AI Research EXAONE Lab. All rights reserved.
+# Copyright 2024 The LG CNS AI Engineering Team.
+# Copyright 2023-2024 SGLang Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" EXAONE model configuration """
+from typing import Any, Dict
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+EXAONE_PRETRAINED_CONFIG_ARCHIVE_MAP: Dict[str, Any] = {}
+
+
+# ruff: noqa: E501
+class ExaoneConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a :class:`~transformers.ExaoneModel`. It is used to
+    instantiate a EXAONE model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the Exaone
+
+    Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model
+    outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information.
+
+
+    Args:
+        vocab_size (:obj:`int`, `optional`, defaults to 102400):
+            Vocabulary size of the EXAONE model. Defines the number of different tokens that can be represented by the
+            :obj:`inputs_ids` passed when calling :class:`~transformers.ExaoneModel`. Vocabulary size of the model.
+            Defines the different tokens that can be represented by the `inputs_ids` passed to the forward method of
+            :class:`~transformers.EXAONEModel`.
+        max_position_embeddings (:obj:`int`, `optional`, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        hidden_size (:obj:`int`, `optional`, defaults to 2048):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_layers (:obj:`int`, `optional`, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (:obj:`int`, `optional`, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (:obj:`int`, `optional`):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        intermediate_size (:obj:`int`, `optional`, defaults to `hidden_size * 4`):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        rope_theta (:obj:`float`, `optional`, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (:obj:`Dict`, `optional`):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (:obj:`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (:obj:`float`, `optional`):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (:obj:`int`, `optional`):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (:obj:`float`, `optional`):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (:obj:`float`, `optional`):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (:obj:`float`, `optional`):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (:obj:`List[float]`, `optional`):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (:obj:`List[float]`, `optional`):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (:obj:`float`, `optional`):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (:obj:`float`, `optional`):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        embed_dropout (:obj:`float`, `optional`, defaults to 0.0):
+            The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_dropout (:obj:`float`, `optional`, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        layer_norm_epsilon (:obj:`float`, `optional`, defaults to 1e-5):
+            The epsilon used by the layer normalization layers.
+        initializer_range (:obj:`float`, `optional`, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if ``configs.is_decoder=True``.
+        bos_token_id (:obj:`int`, `optional`, defaults to 0):
+            Beginning of stream token id.
+        eos_token_id (:obj:`int`, `optional`, defaults to 2):
+            End of stream token id.
+        tie_word_embeddings (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to tie weight embeddings
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+
+        Example::
+
+            >>> from transformers import EXAONEModel, ExaoneConfig
+
+            >>> # Initializing a EXAONE configuration
+            >>> configuration = ExaoneConfig()
+
+            >>> # Initializing a model from configuration
+            >>> model = EXAONEModel(configuration)
+
+            >>> # Accessing the model configuration
+            >>> configuration = model.configs
+    """
+
+    model_type = "exaone"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {"num_hidden_layers": "num_layers"}
+
+    def __init__(
+        self,
+        vocab_size=102400,
+        max_position_embeddings=2048,
+        hidden_size=2048,
+        num_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        intermediate_size=None,
+        activation_function="silu",
+        rope_theta=10000.0,
+        rope_scaling=None,
+        embed_dropout=0.0,
+        attention_dropout=0.0,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        use_cache=True,
+        bos_token_id=0,
+        eos_token_id=2,
+        tie_word_embeddings=True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_layers
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        if intermediate_size:
+            self.intermediate_size = intermediate_size
+        else:
+            self.intermediate_size = hidden_size * 4
+        self.activation_function = activation_function
+        self.embed_dropout = embed_dropout
+        self.attention_dropout = attention_dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs
+        )

sglang/srt/constrained/fsm_cache.py

@@ -79,7 +79,7 @@ class FSMCache(BaseToolCache):
 
     def init_value(self, value):
         if self.json_schema_mode:
-            regex = build_regex_from_schema(value)
+            regex = build_regex_from_schema(value, whitespace_pattern=r"[\n\t ]*")
             return RegexGuide(regex, self.outlines_tokenizer), regex
         else:
             return RegexGuide(value, self.outlines_tokenizer)

sglang/srt/conversation.py

@@ -386,7 +386,16 @@ def generate_chat_conv(
     for message in request.messages:
         msg_role = message.role
         if msg_role == "system":
-            conv.system_message = message.content
+            if isinstance(message.content, str):
+                conv.system_message = message.content
+            elif isinstance(message.content, list):
+                if (
+                    len(message.content) != 1
+                    or getattr(message.content[0], "type", None) != "text"
+                ):
+                    raise ValueError("The system message should be a single text.")
+                else:
+                    conv.system_message = getattr(message.content[0], "text", "")
         elif msg_role == "user":
             # Handle the various types of Chat Request content types here.
             role = conv.roles[0]
@@ -414,7 +423,20 @@ def generate_chat_conv(
                         conv.append_image(content.image_url.url)
                 conv.append_message(conv.roles[0], real_content)
         elif msg_role == "assistant":
-            conv.append_message(conv.roles[1], message.content)
+            parsed_content = ""
+            if isinstance(message.content, str):
+                parsed_content = message.content
+            elif isinstance(message.content, list):
+                if (
+                    len(message.content) != 1
+                    or getattr(message.content[0], "type", None) != "text"
+                ):
+                    raise ValueError(
+                        "The assistant's response should be a single text."
+                    )
+                else:
+                    parsed_content = getattr(message.content[0], "text", "")
+            conv.append_message(conv.roles[1], parsed_content)
         else:
             raise ValueError(f"Unknown role: {msg_role}")
 

sglang/srt/hf_transformers_utils.py

@@ -15,6 +15,7 @@ limitations under the License.
 
 """Utilities for Huggingface Transformers."""
 
+import contextlib
 import functools
 import json
 import os
@@ -34,15 +35,20 @@ from transformers import (
 try:
     from vllm.transformers_utils.configs import ChatGLMConfig, DbrxConfig
 
+    from sglang.srt.configs import ExaoneConfig
+
     _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
         ChatGLMConfig.model_type: ChatGLMConfig,
         DbrxConfig.model_type: DbrxConfig,
+        ExaoneConfig.model_type: ExaoneConfig,
     }
 except ImportError:
     # We want this file to run without vllm dependency
     _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {}
 
-from sglang.srt.utils import is_multimodal_model
+for name, cls in _CONFIG_REGISTRY.items():
+    with contextlib.suppress(ValueError):
+        AutoConfig.register(name, cls)
 
 
 def download_from_hf(model_path: str):
@@ -52,17 +58,11 @@ def download_from_hf(model_path: str):
     return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
 
 
-def get_config_json(model_path: str):
-    with open(os.path.join(model_path, "config.json")) as f:
-        config = json.load(f)
-    return config
-
-
 def get_config(
     model: str,
     trust_remote_code: bool,
     revision: Optional[str] = None,
-    model_overide_args: Optional[dict] = None,
+    model_override_args: Optional[dict] = None,
 ):
     config = AutoConfig.from_pretrained(
         model, trust_remote_code=trust_remote_code, revision=revision
@@ -70,8 +70,8 @@ def get_config(
     if config.model_type in _CONFIG_REGISTRY:
         config_class = _CONFIG_REGISTRY[config.model_type]
         config = config_class.from_pretrained(model, revision=revision)
-    if model_overide_args:
-        config.update(model_overide_args)
+    if model_override_args:
+        config.update(model_override_args)
     return config
 
 
@@ -89,10 +89,10 @@ CONTEXT_LENGTH_KEYS = [
 
 
 def get_context_length(config):
-    """Get the context length of a model from a huggingface model config."""
+    """Get the context length of a model from a huggingface model configs."""
     rope_scaling = getattr(config, "rope_scaling", None)
     if rope_scaling:
-        rope_scaling_factor = config.rope_scaling["factor"]
+        rope_scaling_factor = config.rope_scaling.get("factor", 1)
         if "original_max_position_embeddings" in rope_scaling:
             rope_scaling_factor = 1
         if config.rope_scaling.get("rope_type", None) == "llama3":

sglang/srt/layers/extend_attention.py

@@ -127,8 +127,7 @@ def _fwd_kernel(
         )
         k = tl.load(K_Buffer + offs_buf_k, mask=mask_n[None, :], other=0.0)
 
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
+        qk = tl.dot(q.to(k.dtype), k)
         if BLOCK_DPE > 0:
             offs_kpe = (
                 offs_kv_loc[None, :] * stride_buf_kbs
@@ -140,7 +139,7 @@ def _fwd_kernel(
                 mask=mask_n[None, :],
                 other=0.0,
             )
-            qk += tl.dot(qpe, kpe)
+            qk += tl.dot(qpe.to(kpe.dtype), kpe)
         qk *= sm_scale
 
         if logit_cap > 0:
@@ -179,9 +178,7 @@ def _fwd_kernel(
         )
         k = tl.load(K_Extend + offs_k, mask=mask_n[None, :], other=0.0)
 
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k)
-
+        qk = tl.dot(q, k, out_dtype=tl.float32)
         if BLOCK_DPE > 0:
             offs_kpe = (
                 (cur_seq_extend_start_contiguous + start_n + offs_n[None, :])
@@ -276,9 +273,17 @@ def extend_attention_fwd(
     BLOCK_DV = Lv
 
     if CUDA_CAPABILITY[0] >= 9:
-        BLOCK_M, BLOCK_N = (128, 64)
+        if Lq <= 256:
+            BLOCK_M, BLOCK_N = (128, 64)
+        else:
+            BLOCK_M, BLOCK_N = (32, 64)
     elif CUDA_CAPABILITY[0] >= 8:
-        BLOCK_M, BLOCK_N = (128, 128) if Lq <= 128 else (64, 64)
+        if Lq <= 128:
+            BLOCK_M, BLOCK_N = (128, 128)
+        elif Lq <= 256:
+            BLOCK_M, BLOCK_N = (64, 64)
+        else:
+            BLOCK_M, BLOCK_N = (32, 64)
     else:
         BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)