diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index 354ec9a32b2f..8951bcdc9633 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -78,6 +78,7 @@ def get_model_config( topk = config.num_experts_per_tok intermediate_size = config.intermediate_size elif architecture in [ + "BerryLMForCausalLM", "Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", diff --git a/docs/docs/advanced_features/separate_reasoning.mdx b/docs/docs/advanced_features/separate_reasoning.mdx index 86412b8a276c..0a9d9e828095 100644 --- a/docs/docs/advanced_features/separate_reasoning.mdx +++ b/docs/docs/advanced_features/separate_reasoning.mdx @@ -41,6 +41,12 @@ SGLang supports parsing reasoning content out from "normal" content for reasonin `deepseek-v3` Including [DeepSeek‑V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2-Exp). Supports `thinking` parameter + + BerryLM-OS + `` … `` + `berrylm` + Supports `enable_thinking`; a `` may follow the reasoning without `` + [Standard Qwen3 models](https://huggingface.co/collections/Qwen/qwen3-67dd247413f0e2e4f653967f) `` … `` diff --git a/docs/docs/advanced_features/tool_parser.mdx b/docs/docs/advanced_features/tool_parser.mdx index 27b935b3ad2f..f30d402216de 100644 --- a/docs/docs/advanced_features/tool_parser.mdx +++ b/docs/docs/advanced_features/tool_parser.mdx @@ -82,6 +82,11 @@ This guide demonstrates how to use SGLang’s [Function calling](https://platfor Qwen series (e.g. `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-VL-30B-A3B-Thinking`) except Qwen3-Coder + + `berrylm` + BerryLM-OS (e.g. `rwb-ai/BerryLM-OS`), XML `` format + + `qwen3_coder` Qwen3-Coder (e.g. `Qwen/Qwen3-Coder-30B-A3B-Instruct`) diff --git a/docs/docs/supported-models/generative_models.mdx b/docs/docs/supported-models/generative_models.mdx index 6811d2954c72..8e7f3d1c9ddb 100644 --- a/docs/docs/supported-models/generative_models.mdx +++ b/docs/docs/supported-models/generative_models.mdx @@ -48,6 +48,11 @@ in the GitHub search bar. `moonshotai/Kimi-K2-Instruct` Moonshot AI's 1 trillion parameter MoE model (32B active) with 128K–256K context; state-of-the-art agentic intelligence with stable long-horizon agency across 200–300 sequential tool calls. Features MLA attention and native INT4 quantization. See Reasoning Parser docs + + **BerryLM-OS** (18B-A3B) + `rwb-ai/BerryLM-OS` + Hybrid linear/full-attention MoE (18.3B total, 3.2B active) with a thinking mode: per-channel KDA forget gates on the linear-attention layers and Gated Block AttnRes depth mixing; 180k vocabulary, 256k context. + **Kimi Linear** (48B-A3B) `moonshotai/Kimi-Linear-48B-A3B-Instruct` diff --git a/python/sglang/kernels/ops/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py index 3b877ec6751f..f756b889b5cc 100644 --- a/python/sglang/kernels/ops/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -1139,6 +1139,7 @@ def chunk_kda_fwd( dt_bias: Optional[torch.Tensor] = None, lower_bound: Optional[float] = None, output_intermediate_states: bool = False, + fused_intra: Optional[bool] = None, track_state: Optional[torch.Tensor] = None, track_chunk_idx: Optional[torch.Tensor] = None, beta_is_raw: bool = False, @@ -1197,6 +1198,9 @@ def chunk_kda_fwd( _H_pr = q.shape[-2] _B = q.shape[0] _small_grid = _B * _NT_pr * _H_pr <= 256 + # Layers whose per-channel decays exceed the +-126 (log2) clamp of the fused diagonal + # factorization opt out (RadixLinearAttention.kda_fused_intra = False). + _fused_intra = _small_grid if fused_intra is None else bool(fused_intra) w, u, _, kg, Aqk, _ = chunk_kda_fwd_intra( q=q, k=k, @@ -1208,8 +1212,8 @@ def chunk_kda_fwd( chunk_size=chunk_size, chunk_indices=chunk_indices, safe_gate=lower_bound is not None, - fuse_diagonal=_small_grid, - fuse_recompute=_small_grid, + fuse_diagonal=_fused_intra, + fuse_recompute=_fused_intra, ) h, v_new = chunk_gated_delta_rule_fwd_h( @@ -1265,6 +1269,7 @@ def chunk_kda( dt_bias: Optional[torch.Tensor] = None, lower_bound: Optional[float] = None, output_intermediate_states: bool = False, + fused_intra: Optional[bool] = None, track_state: Optional[torch.Tensor] = None, track_chunk_idx: Optional[torch.Tensor] = None, beta_is_raw: bool = False, @@ -1292,6 +1297,7 @@ def chunk_kda( dt_bias=dt_bias, lower_bound=lower_bound, output_intermediate_states=output_intermediate_states, + fused_intra=fused_intra, track_state=track_state, track_chunk_idx=track_chunk_idx, beta_is_raw=beta_is_raw, diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 06239f17ee1d..2198c5348951 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -1,6 +1,7 @@ from sglang.srt.configs.afmoe import AfmoeConfig from sglang.srt.configs.bailing_hybrid import BailingHybridConfig, BailingMoeV3VLConfig from sglang.srt.configs.bailing_moe_v2 import BailingMM2Config +from sglang.srt.configs.berrylm import BerryLMConfig from sglang.srt.configs.chatglm import ChatGLMConfig from sglang.srt.configs.cohere2_moe import Cohere2MoeConfig from sglang.srt.configs.cosmos3 import ( @@ -95,6 +96,7 @@ "BailingMM2Config", "BailingMoeV3VLConfig", "ExaoneConfig", + "BerryLMConfig", "ChatGLMConfig", "Cosmos3Config", "Cosmos3EdgeConfig", diff --git a/python/sglang/srt/configs/berrylm.py b/python/sglang/srt/configs/berrylm.py new file mode 100644 index 000000000000..93c95ae96189 --- /dev/null +++ b/python/sglang/srt/configs/berrylm.py @@ -0,0 +1,191 @@ +# Copyright 2025-2026 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. +# ============================================================================== +"""BerryLM-OS model configuration (text-only hybrid MoE decoder: gated delta-net +linear attention with a per-channel KDA forget gate, full attention every +``full_attention_interval`` layers, sparse MoE + shared expert, and a Gated +Block AttnRes mixer on the residual stream).""" + +from transformers import PretrainedConfig + +from sglang.srt.configs.linear_attn_model_registry import ( + LinearAttnModelSpec, + register_linear_attn_model, +) +from sglang.srt.configs.mamba_utils import ( + KimiLinearCacheParams, + KimiLinearStateShape, + mamba2_state_dtype, +) + + +class BerryLMConfig(PretrainedConfig): + model_type = "berrylm" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=180224, + hidden_size=2048, + num_hidden_layers=40, + num_attention_heads=16, + num_key_value_heads=2, + hidden_act="silu", + max_position_embeddings=262144, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_parameters=None, + rope_scaling=None, + partial_rotary_factor=0.25, + attention_bias=False, + attention_dropout=0.0, + attn_output_gate=True, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts_per_tok=8, + num_experts=128, + norm_topk_prob=True, + output_router_logits=False, + router_aux_loss_coef=0.001, + layer_types=None, + full_attention_interval=4, + attn_res_block_size=8, + attn_res_gated=True, + attn_res_eps=1e-6, + kda_gate_bottleneck=128, + kda_safe_gate=False, + kda_gate_lower_bound=None, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.attn_output_gate = attn_output_gate + self.head_dim = head_dim + self.partial_rotary_factor = partial_rotary_factor + # transformers v5 stores RoPE settings in ``rope_parameters``; keep both names populated. + rope = dict(rope_parameters or rope_scaling or {}) + rope.setdefault("rope_type", "default") + rope.setdefault("rope_theta", 10000000.0) + rope.setdefault("partial_rotary_factor", partial_rotary_factor) + self.rope_parameters = rope + self.rope_scaling = rope + self.rope_theta = rope["rope_theta"] + + self.full_attention_interval = full_attention_interval + self.layer_types = layer_types + if self.layer_types is None: + self.layer_types = [ + ( + "linear_attention" + if bool((i + 1) % full_attention_interval) + else "full_attention" + ) + for i in range(self.num_hidden_layers) + ] + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError("layer_types must have num_hidden_layers entries") + + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + + self.attn_res_block_size = int(attn_res_block_size) + self.attn_res_gated = bool(attn_res_gated) + self.attn_res_eps = float(attn_res_eps) + self.kda_gate_bottleneck = int(kda_gate_bottleneck) + self.kda_safe_gate = bool(kda_safe_gate) + self.kda_gate_lower_bound = kda_gate_lower_bound + + @property + def layers_block_type(self): + # SGLang HybridLayerType values: "attention" / "linear_attention". + return [ + "attention" if t == "full_attention" else "linear_attention" + for t in self.layer_types + ] + + @property + def linear_layer_ids(self): + return [ + i for i, t in enumerate(self.layers_block_type) if t == "linear_attention" + ] + + @property + def full_attention_layer_ids(self): + return [i for i, t in enumerate(self.layers_block_type) if t == "attention"] + + @property + def mamba2_cache_params(self) -> KimiLinearCacheParams: + # KDA backend state layout (same as Kimi Linear): one fused q|k|v conv window + # [kernel-1, 2*H*K + HV*V] and a [HV, V, K] recurrent state per linear layer. + from sglang.srt.runtime_context import get_parallel + + if self.linear_key_head_dim != self.linear_value_head_dim: + raise ValueError( + "BerryLM KDA cache needs linear_key_head_dim == linear_value_head_dim" + ) + shape = KimiLinearStateShape.create( + tp_world_size=get_parallel().attn_tp_size, + num_heads=self.linear_num_value_heads, + head_dim=self.linear_value_head_dim, + num_k_heads=self.linear_num_key_heads, + head_k_dim=self.linear_key_head_dim, + conv_kernel_size=self.linear_conv_kernel_dim, + ) + return KimiLinearCacheParams( + shape=shape, layers=self.linear_layer_ids, dtype=mamba2_state_dtype(self) + ) + + +# Hybrid plumbing (mamba pool sizing, linear-attention backend, radix-cache args) is +# driven by the linear-attention model registry: BerryLM shares the KDA backend with +# Kimi Linear (same fused q|k|v conv state layout; grouped value heads are handled in +# KDAAttnBackend.forward_extend). +register_linear_attn_model( + LinearAttnModelSpec( + config_class=BerryLMConfig, + backend_class_name="sglang.srt.layers.attention.linear.kda_backend.KDAAttnBackend", + arch_names=["BerryLMForCausalLM"], + uses_mamba_radix_cache=True, + support_mamba_cache=True, + support_mamba_cache_extra_buffer=True, + ) +) diff --git a/python/sglang/srt/function_call/berrylm_detector.py b/python/sglang/srt/function_call/berrylm_detector.py new file mode 100644 index 000000000000..76b3b65ec1c8 --- /dev/null +++ b/python/sglang/srt/function_call/berrylm_detector.py @@ -0,0 +1,494 @@ +import json +import logging +import re +from typing import Any, List, Optional + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import ( + StreamingParseResult, + ToolCallItem, + _GetInfoFunc, +) +from sglang.srt.function_call.utils import ( + get_schema_properties, + infer_type_from_json_schema, + safe_literal_eval, +) + +logger = logging.getLogger(__name__) + + +class BerryLMFunctionCallDetector(BaseFormatDetector): + """Detector for the BerryLM XML tool-call format:: + + + + + value + + + + """ + + def __init__(self): + super().__init__() + + # Sentinel tokens + self.tool_call_start_token: str = "" + self.tool_call_end_token: str = "" + self.tool_call_prefix: str = "(.*?)", re.DOTALL) + self.tool_call_function_regex = re.compile( + r"||(?=)|$)", + re.DOTALL, + ) + + # Streaming State + # Base class already initializes _buffer, we just use it directly + # No need to check with hasattr - we control the lifecycle through inheritance + + # Index pointing to the next character to be processed in buffer + self.parsed_pos: int = 0 + # Parameter count inside the current tool being processed, used to determine whether to add comma + self.current_tool_param_count: int = 0 + # Flag indicating whether current tool has already sent '{' + self.json_started: bool = False + + # [FIX] New state flag: mark whether inside tool_call structure block + self.is_inside_tool_call: bool = False + + # Initialize attributes that were missing in the original PR + self.current_func_name: Optional[str] = None + + def has_tool_call(self, text: str) -> bool: + return self.tool_call_start_token in text + + def _get_arguments_config( + self, func_name: str, tools: Optional[list[Tool]] + ) -> dict: + """Extract argument configuration for a function.""" + if tools is None: + return {} + for config in tools: + try: + config_type = config.type + config_function = config.function + config_function_name = config_function.name + except AttributeError: + continue + + if config_type == "function" and config_function_name == func_name: + try: + params = config_function.parameters + except AttributeError: + return {} + + if isinstance(params, dict): + properties = get_schema_properties(params) + if properties or "properties" in params: + return properties + return params + else: + return {} + logger.warning(f"Tool '{func_name}' is not defined in the tools list.") + return {} + + def _get_param_type(self, param_schema: Any) -> str: + """Infer the parser conversion type from a JSON schema parameter.""" + inferred_type = infer_type_from_json_schema(param_schema) + if inferred_type is None: + return "string" + return str(inferred_type).strip().lower() + + def _convert_param_value( + self, param_value: str, param_name: str, param_config: dict, func_name: str + ) -> Any: + """Convert parameter value based on its type in the schema.""" + # Handle null value for any type + if param_value.lower() == "null": + return None + + if param_name not in param_config: + if param_config != {}: + logger.warning( + f"Parsed parameter '{param_name}' is not defined in the tool " + f"parameters for tool '{func_name}', directly returning the string value." + ) + return param_value + + param_type = self._get_param_type(param_config[param_name]) + if param_type in ["string", "str", "text", "varchar", "char", "enum"]: + return param_value + elif ( + param_type.startswith("int") + or param_type.startswith("uint") + or param_type.startswith("long") + or param_type.startswith("short") + or param_type.startswith("unsigned") + ): + try: + param_value = int(param_value) + except Exception: + logger.warning( + f"Parsed value '{param_value}' of parameter '{param_name}' is not an integer in tool " + f"'{func_name}', degenerating to string." + ) + return param_value + elif param_type.startswith("num") or param_type.startswith("float"): + try: + maybe_convert = ( + False if "." in param_value or "e" in param_value.lower() else True + ) + param_value: float = float(param_value) + if maybe_convert and param_value.is_integer(): + param_value = int(param_value) + except Exception: + logger.warning( + f"Parsed value '{param_value}' of parameter '{param_name}' is not a float in tool " + f"'{func_name}', degenerating to string." + ) + return param_value + elif param_type in ["boolean", "bool", "binary"]: + param_value = param_value.lower() + if param_value not in ["true", "false"]: + logger.warning( + f"Parsed value '{param_value}' of parameter '{param_name}' is not a boolean (`true` of `false`) in tool '{func_name}', degenerating to false." + ) + return param_value == "true" + else: + if ( + param_type in ["object", "array", "arr"] + or param_type.startswith("dict") + or param_type.startswith("list") + ): + try: + param_value = json.loads(param_value) + return param_value + except Exception: + logger.warning( + f"Parsed value '{param_value}' of parameter '{param_name}' cannot be parsed with json.loads in tool " + f"'{func_name}', will try other methods to parse it." + ) + try: + param_value = safe_literal_eval(param_value) + except Exception: + logger.warning( + f"Parsed value '{param_value}' of parameter '{param_name}' cannot be converted via Python `ast.literal_eval()` in tool '{func_name}', degenerating to string." + ) + return param_value + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """One-shot parsing for non-streaming scenarios.""" + if self.tool_call_start_token not in text: + return StreamingParseResult(normal_text=text) + + calls = [] + try: + # Simple cleanup of the text to find tool calls + # Note: This is a simplified regex approach consistent with vLLM + raw_tool_calls = self.tool_call_regex.findall(text) + if not raw_tool_calls: + # Fallback: maybe the whole text is inside the tag or tags are stripped + if self.tool_call_prefix in text: + raw_tool_calls = [text] + + tool_idx = 0 + for tool_content in raw_tool_calls: + # Find function calls + funcs = self.tool_call_function_regex.findall(tool_content) + for func_match in funcs: + func_body = func_match[0] or func_match[1] + if ">" not in func_body: + continue + + name_end = func_body.index(">") + func_name = func_body[:name_end] + params_str = func_body[name_end + 1 :] + + param_config = self._get_arguments_config(func_name, tools) + parsed_params = {} + + for p_match in self.tool_call_parameter_regex.findall(params_str): + if ">" not in p_match: + continue + p_idx = p_match.index(">") + p_name = p_match[:p_idx] + p_val = p_match[p_idx + 1 :] + # Remove prefixing and trailing \n + if p_val.startswith("\n"): + p_val = p_val[1:] + if p_val.endswith("\n"): + p_val = p_val[:-1] + + parsed_params[p_name] = self._convert_param_value( + p_val, p_name, param_config, func_name + ) + + calls.append( + ToolCallItem( + tool_index=tool_idx, + name=func_name, + parameters=json.dumps(parsed_params, ensure_ascii=False), + ) + ) + tool_idx += 1 + + # Determine normal text (text before the first tool call) + start_idx = text.find(self.tool_call_start_token) + if start_idx == -1: + start_idx = text.find(self.tool_call_prefix) + normal_text = text[:start_idx] if start_idx > 0 else "" + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + except Exception as e: + logger.error(f"Error in detect_and_parse: {e}") + return StreamingParseResult(normal_text=text) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Robust cursor-based streaming parser. + """ + self._buffer += new_text + + # Guard against empty buffer + if not self._buffer: + return StreamingParseResult() + + calls = [] + normal_text_chunks = [] + + while True: + # Working text slice + current_slice = self._buffer[self.parsed_pos :] + + # Optimization: If almost empty, wait for more + if not current_slice: + break + + # ------------------------------------------------------- + # 1. Priority detection: check if it's the start of Tool Call + # ------------------------------------------------------- + if current_slice.startswith(self.tool_call_start_token): + self.parsed_pos += len(self.tool_call_start_token) + self.is_inside_tool_call = True + continue + + # ------------------------------------------------------- + # 2. Function Name: + # ------------------------------------------------------- + if current_slice.startswith(self.tool_call_prefix): + end_angle = current_slice.find(">") + if end_angle != -1: + func_name = current_slice[len(self.tool_call_prefix) : end_angle] + + self.current_tool_id += 1 + self.current_tool_name_sent = True + self.current_tool_param_count = 0 + self.json_started = False + self.current_func_name = func_name + + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, + name=func_name, + parameters="", + ) + ) + + self.parsed_pos += end_angle + 1 + continue + else: + # Incomplete tag + break + + # ------------------------------------------------------- + # 3. Parameter: value... + # ------------------------------------------------------- + if current_slice.startswith(self.parameter_prefix): + name_end = current_slice.find(">") + if name_end != -1: + value_start_idx = name_end + 1 + rest_of_slice = current_slice[value_start_idx:] + + # A parameter can end in multiple ways: + # 1. [Normal] Encounter + # 2. [Abnormal] Encounter next + # So we need to find the smallest one as the parameter end position. + cand_end_param = rest_of_slice.find(self.parameter_end_token) + cand_next_param = rest_of_slice.find(self.parameter_prefix) + cand_end_func = rest_of_slice.find(self.function_end_token) + + candidates = [] + if cand_end_param != -1: + candidates.append( + (cand_end_param, len(self.parameter_end_token)) + ) + if cand_next_param != -1: + candidates.append((cand_next_param, 0)) + if cand_end_func != -1: + candidates.append((cand_end_func, 0)) + + if candidates: + best_cand = min(candidates, key=lambda x: x[0]) + end_pos = best_cand[0] + end_token_len = best_cand[1] + + param_name = current_slice[ + len(self.parameter_prefix) : name_end + ] + raw_value = rest_of_slice[:end_pos] + + # Cleanup value + if raw_value.startswith("\n"): + raw_value = raw_value[1:] + if raw_value.endswith("\n"): + raw_value = raw_value[:-1] + + # JSON Construction + if not self.json_started: + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, parameters="{" + ) + ) + self.json_started = True + + param_config = self._get_arguments_config( + self.current_func_name, tools + ) + converted_val = self._convert_param_value( + raw_value, param_name, param_config, self.current_func_name + ) + + # Construct JSON fragment: "key": value + # Note: We must be careful with json.dumps to ensure valid JSON streaming + json_key_val = f"{json.dumps(param_name)}: {json.dumps(converted_val, ensure_ascii=False)}" + + if self.current_tool_param_count > 0: + fragment = f", {json_key_val}" + else: + fragment = json_key_val + + calls.append( + ToolCallItem( + tool_index=self.current_tool_id, parameters=fragment + ) + ) + self.current_tool_param_count += 1 + + # Advance cursor + total_len = (name_end + 1) + end_pos + end_token_len + self.parsed_pos += total_len + continue + + # Incomplete parameter tag or value + break + + # ------------------------------------------------------- + # 4. Function End: + # ------------------------------------------------------- + if current_slice.startswith(self.function_end_token): + if not self.json_started: + calls.append( + ToolCallItem(tool_index=self.current_tool_id, parameters="{") + ) + self.json_started = True + + calls.append( + ToolCallItem(tool_index=self.current_tool_id, parameters="}") + ) + self.parsed_pos += len(self.function_end_token) + self.current_func_name = None + continue + + # ------------------------------------------------------- + # 5. Tool Call End: + # ------------------------------------------------------- + if current_slice.startswith(self.tool_call_end_token): + self.parsed_pos += len(self.tool_call_end_token) + self.is_inside_tool_call = False # [FIX] Exit tool call region + continue + + # ------------------------------------------------------- + # 6. Handling content / whitespace / normal text + # ------------------------------------------------------- + # If current position is not the start of a tag (i.e., doesn't start with <), it might be plain text, + # or a newline between two tags. + # But we need to be careful not to output truncated tags like " 0: + self._buffer = self._buffer[self.parsed_pos :] + self.parsed_pos = 0 + + normal_text = "".join(normal_text_chunks) if normal_text_chunks else "" + return StreamingParseResult(calls=calls, normal_text=normal_text) + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError + + def get_structural_tag_name(self) -> Optional[str]: + return None diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 3184ac7dc3e3..269386e099a7 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -13,6 +13,7 @@ from sglang.srt.environ import ToolStrictLevel, envs from sglang.srt.function_call.apertus2509_detector import Apertus2509Detector from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.berrylm_detector import BerryLMFunctionCallDetector from sglang.srt.function_call.cohere_command4_detector import CohereCommand4Detector from sglang.srt.function_call.core_types import ToolCallItem from sglang.srt.function_call.deepseekv3_detector import DeepSeekV3Detector @@ -75,6 +76,7 @@ class FunctionCallParser: ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = { "apertus2509": Apertus2509Detector, "cohere_command4": CohereCommand4Detector, + "berrylm": BerryLMFunctionCallDetector, "deepseekv3": DeepSeekV3Detector, "deepseekv31": DeepSeekV31Detector, "deepseekv32": DeepSeekV32Detector, diff --git a/python/sglang/srt/function_call/parser_names.py b/python/sglang/srt/function_call/parser_names.py index 3da09d6b59d9..0b4807ba955c 100644 --- a/python/sglang/srt/function_call/parser_names.py +++ b/python/sglang/srt/function_call/parser_names.py @@ -10,6 +10,7 @@ TOOL_CALL_PARSER_NAMES = [ "apertus2509", "cohere_command4", + "berrylm", "deepseekv3", "deepseekv31", "deepseekv32", diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 29fdbbed6b93..756dc91ee7d5 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -876,6 +876,13 @@ def forward_extend( if gate_was_flat: a = a.unflatten(-1, (-1, layer.head_k_dim)) + if layer.num_v_heads != layer.num_q_heads: + # Grouped value heads (BerryLM: 16 q/k heads over 32 value heads). The + # chunk KDA kernels take H == HV, so repeat q/k per value head; the + # decode / verify kernels index the groups themselves. + group = layer.num_v_heads // layer.num_q_heads + q = q.repeat_interleave(group, dim=2) + k = k.repeat_interleave(group, dim=2) track_ssm = self.forward_metadata.has_mamba_track_mask track_chunk_idx = self.forward_metadata.track_chunk_idx h_track_buf = None @@ -919,6 +926,7 @@ def forward_extend( lower_bound=layer.lower_bound, beta_is_raw=gate_was_flat, extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + fused_intra=getattr(layer, "kda_fused_intra", None), extend_prefix_lens=forward_batch.extend_prefix_lens, layer_id=layer.layer_id, # draft_extend_v2 must stay rollback-able, so kernels that commit state diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index e0e807aa476d..3a7d3c038f8e 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -249,6 +249,7 @@ def extend( lower_bound=lower_bound, beta_is_raw=beta_is_raw, output_intermediate_states=return_intermediate_states, + fused_intra=kwargs.get("fused_intra"), track_state=kwargs.get("track_state"), track_chunk_idx=kwargs.get("track_chunk_idx"), ) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000000..9767eaf33a61 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200.json new file mode 100644 index 000000000000..50d2854421c3 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=128,N=512,device_name=NVIDIA_H200.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/models/berrylm.py b/python/sglang/srt/models/berrylm.py new file mode 100644 index 000000000000..2d238e360171 --- /dev/null +++ b/python/sglang/srt/models/berrylm.py @@ -0,0 +1,1290 @@ +# Copyright 2025-2026 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. +# ============================================================================== +"""Inference-only BerryLM-OS model compatible with HuggingFace weights. + +BerryLM is a text-only hybrid MoE decoder: three gated delta-net linear-attention +layers per full-attention layer, sparse MoE + shared expert on every layer, and two +additions on top of the usual hybrid: + +* a per-channel **KDA forget gate** on the linear layers + (``g = -exp(A_log)[h] * softplus(f_up(f_down(x))[h, k] + dt_bias[h, k])``) on the + gated-delta-net projection layout (fused q|k|v conv, ``in_proj_z`` output gate) — served + by the KDA linear-attention backend (registered in ``configs/berrylm.py``); +* a **Gated Block AttnRes** mixer on the residual stream: before layer *i* the + stream ``x`` becomes ``x + tanh(gate_i) * (mix_i - x)`` where ``mix_i`` is a + softmax mixture over depth of the streams committed at block boundaries and + ``x`` itself. Token-local, so it does not touch any cache. +""" + +import logging +from typing import Iterable, List, Optional, Set, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated +from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add +from sglang.srt.configs.berrylm import BerryLMConfig +from sglang.srt.distributed import tensor_model_parallel_all_reduce +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layer_boundary import declare_attn, declare_ffn, make_stages +from sglang.srt.layers.layer_boundary.residual import batch as residual_batch +from sglang.srt.layers.layernorm import GemmaRMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_skip_post_experts_all_reduce, +) +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + filter_moe_weight_param_global_expert, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from sglang.srt.runtime_context import get_exec, get_parallel, get_stream +from sglang.srt.utils import add_prefix, is_cuda, make_pp_layers, set_weight_attrs +from sglang.srt.utils.custom_op import register_custom_op + +logger = logging.getLogger(__name__) +_is_cuda = is_cuda() + + +# --------------------------------------------------------------------------- +# Gated Block AttnRes +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Fused Gated Block AttnRes mixer (Triton, forward only): one program per token, +# one pass over the committed streams for the norms / logits and one for the +# mixture; numerics follow the torch chain (fp32 accumulation, the blend rounded +# in the stream dtype op by op). The torch chain stays as the off-GPU fallback. +# --------------------------------------------------------------------------- + +try: + import triton + import triton.language as tl + from triton.language.extra import libdevice + + HAS_TRITON = True +except ImportError: # pragma: no cover - CPU-only environments + HAS_TRITON = False + + +if HAS_TRITON: + + @triton.jit + def _gated_attnres_kernel( + stream_ptr, + blocks_ptr, + query_ptr, + gate_ptr, + out_ptr, + n_blocks, + D, + stride_st, # stream / out: row stride + stride_bn, # block_streams: source stride + stride_bt, # block_streams: row stride + eps, + HAS_GATE: tl.constexpr, + BLOCK_D: tl.constexpr, + ): + t = tl.program_id(0) + offs = tl.arange(0, BLOCK_D) + mask = offs < D + q = tl.load(query_ptr + offs, mask=mask, other=0.0).to(tl.float32) + inv_d = 1.0 / D + + # pass 1: logits of the committed sources and of the current stream, running max for a stable softmax + x_cur = tl.load(stream_ptr + t * stride_st + offs, mask=mask, other=0.0).to( + tl.float32 + ) + rstd = tl.rsqrt(tl.sum(x_cur * x_cur, axis=0) * inv_d + eps) + logit_cur = tl.sum(x_cur * q, axis=0) * rstd + m = logit_cur + for i in range(n_blocks): + x = tl.load( + blocks_ptr + i * stride_bn + t * stride_bt + offs, mask=mask, other=0.0 + ).to(tl.float32) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) * inv_d + eps) + logit = tl.sum(x * q, axis=0) * rstd + m = tl.maximum(m, logit) + + # pass 2: the softmax denominator (recomputing a logit is cheaper than keeping n rows live) + w_cur = tl.exp(logit_cur - m) + denom = w_cur + for i in range(n_blocks): + x = tl.load( + blocks_ptr + i * stride_bn + t * stride_bt + offs, mask=mask, other=0.0 + ).to(tl.float32) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) * inv_d + eps) + denom += tl.exp(tl.sum(x * q, axis=0) * rstd - m) + # pass 3: the mixture with normalized weights, committed sources first, the current stream last (the + # reference's order of summation) + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + for i in range(n_blocks): + x = tl.load( + blocks_ptr + i * stride_bn + t * stride_bt + offs, mask=mask, other=0.0 + ).to(tl.float32) + rstd = tl.rsqrt(tl.sum(x * x, axis=0) * inv_d + eps) + p = tl.exp(tl.sum(x * q, axis=0) * rstd - m) / denom + acc += p * x + mixed = acc + (w_cur / denom) * x_cur + + dt: tl.constexpr = out_ptr.dtype.element_ty + if HAS_GATE: + # The reference blends in the stream dtype op by op (`stream + tanh(gate) * (mixed - stream)` on bf16 + # tensors rounds after every op); mirror that rounding sequence so the results match ulp for ulp. + g = tl.load(gate_ptr).to(tl.float32) + scale = libdevice.tanh(g).to(dt).to(tl.float32) + mixed_r = mixed.to(dt).to(tl.float32) + diff = (mixed_r - x_cur).to(dt).to(tl.float32) + prod = (scale * diff).to(dt).to(tl.float32) + out = x_cur + prod + else: + out = mixed + tl.store(out_ptr + t * stride_st + offs, out.to(dt), mask=mask) + + +def gated_attnres( + stream: torch.Tensor, + block_streams: torch.Tensor, + pseudo_query: torch.Tensor, + gate: Optional[torch.Tensor], + eps: float, +) -> torch.Tensor: + """Fused mixer. ``stream`` [T, D], ``block_streams`` [n, T, D] (n >= 0), ``pseudo_query`` [D], ``gate`` 0-d or + None. Falls back to the torch chain off-GPU or for rows wider than 8192.""" + T, D = stream.shape + if ( + not HAS_TRITON + or stream.device.type not in ("cuda", "xpu") + or D > 8192 + or T == 0 + ): + return _gated_attnres_torch(stream, block_streams, pseudo_query, gate, eps) + if stream.stride(-1) != 1 or block_streams.stride(-1) != 1: + stream, block_streams = stream.contiguous(), block_streams.contiguous() + n = block_streams.shape[0] + out = torch.empty_like(stream) + pq = pseudo_query if pseudo_query.dtype == torch.float32 else pseudo_query.float() + has_gate = gate is not None + _gated_attnres_kernel[(T,)]( + stream, + block_streams, + pq, + gate if has_gate else stream, # any pointer when unused + out, + n, + D, + stream.stride(0), + block_streams.stride(0) if n else 0, + block_streams.stride(1) if n else 0, + eps, + HAS_GATE=has_gate, + BLOCK_D=triton.next_power_of_2(D), + num_warps=8 if D > 1024 else 4, + ) + return out + + +def _gated_attnres_torch( + stream: torch.Tensor, + block_streams: torch.Tensor, + pseudo_query: torch.Tensor, + gate: Optional[torch.Tensor], + eps: float, +) -> torch.Tensor: + """The torch chain of the mixer (fallback off-GPU): ``stream`` [T, D], ``block_streams`` [n, T, D].""" + stacked = torch.cat([block_streams, stream.unsqueeze(0)], dim=0).float() + keys = F.rms_norm(stacked, (stream.shape[-1],), None, eps) + logits = torch.einsum("ntd,d->nt", keys, pseudo_query.float()) + mixed = (logits.softmax(dim=0).unsqueeze(-1) * stacked).sum(dim=0).to(stream.dtype) + if gate is None: + return mixed + scale = torch.tanh(gate.float()).to(stream.dtype) + return stream + scale * (mixed - stream) + + +def _gated_block_attn_res( + pseudo_query: torch.Tensor, + stream: torch.Tensor, + completed: List[torch.Tensor], + gate: Optional[torch.Tensor], + eps: float, +) -> torch.Tensor: + hidden = stream.shape[-1] + stream2d = stream.reshape(-1, hidden) + if completed: + block_streams = torch.stack([s.reshape(-1, hidden) for s in completed], dim=0) + else: + block_streams = stream2d.new_empty((0, *stream2d.shape)) + return gated_attnres(stream2d, block_streams, pseudo_query, gate, eps).reshape( + stream.shape + ) + + +@register_custom_op(mutates_args=["output"]) +def berrylm_gated_attnres_with_output( + pseudo_query: torch.Tensor, + stream: torch.Tensor, + completed: List[torch.Tensor], + gate: Optional[torch.Tensor], + eps: float, + output: torch.Tensor, +) -> None: + # Opaque op: the multi-source fp32 stack/rms_norm/softmax chain must not be + # fused by a graph compiler (it stays a plain kernel sequence inside graphs). + output.copy_(_gated_block_attn_res(pseudo_query, stream, completed, gate, eps)) + + +class BerryLMAttnRes(nn.Module): + """Per-layer mixer: ``pseudo_query`` [hidden] (zero == uniform mean over the + sources) and a scalar ``gate`` (zero == exact identity). Replicated on every + TP rank. Checkpoint keys: ``model.layers.{i}.attn_res.{pseudo_query,gate}``.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6, gated: bool = True): + super().__init__() + self.eps = eps + self.pseudo_query = nn.Parameter(torch.zeros(hidden_size)) + self.gate = nn.Parameter(torch.zeros(())) if gated else None + + def forward( + self, stream: torch.Tensor, completed: List[torch.Tensor] + ) -> torch.Tensor: + output = torch.empty_like(stream) + berrylm_gated_attnres_with_output( + self.pseudo_query, stream, list(completed), self.gate, self.eps, output + ) + return output + + +# --------------------------------------------------------------------------- +# Linear attention with the KDA forget gate +# --------------------------------------------------------------------------- + + +class BerryLMKDA(nn.Module): + """Gated delta net (stock GDN projection layout) with a per-channel KDA gate. + + Projections: ``in_proj_qkv`` [2K+V, D] (fused in the checkpoint), ``in_proj_z`` and + ``in_proj_b`` (beta logits) — one merged GEMM at run time —, ``f_down_proj`` [bottleneck, D] / ``f_up_proj`` + [HV*K, bottleneck] (raw forget gate), depthwise ``conv1d`` over q|k|v, + ``A_log`` [HV], ``dt_bias`` [HV*K], SiLU-gated output RMSNorm on ``z``, + ``out_proj``. The KDA backend applies the gate activation (prefill in + ``chunk_kda``, decode in the fused recurrent kernel), like Kimi Linear. + """ + + def __init__( + self, + config: BerryLMConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_tp_size + self.hidden_size = config.hidden_size + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_size = config.linear_conv_kernel_dim + self.gate_bottleneck = config.kda_gate_bottleneck + self.layer_id = layer_id + self.activation = config.hidden_act + self.layer_norm_epsilon = config.rms_norm_eps + self.lower_bound = ( + float(config.kda_gate_lower_bound) + if getattr(config, "kda_safe_gate", False) + else None + ) + if self.num_v_heads % self.num_k_heads != 0: + raise ValueError( + "linear_num_value_heads must be a multiple of linear_num_key_heads" + ) + if self.num_k_heads % self.attn_tp_size or self.num_v_heads % self.attn_tp_size: + raise ValueError( + "BerryLM linear attention needs the key and value head counts " + f"divisible by the attention TP size ({self.attn_tp_size})" + ) + self.local_num_k_heads = self.num_k_heads // self.attn_tp_size + self.local_num_v_heads = self.num_v_heads // self.attn_tp_size + + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = ColumnParallelLinear( + input_size=self.conv_kernel_size, + output_size=self.conv_dim, + bias=False, + quant_config=None, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("conv1d", prefix), + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + + # Checkpoint ships in_proj_qkv / in_proj_z / in_proj_b separately; q|k|v|z run as one GEMM + # (every partition is sharded by head, so the merged layer's TP split is the natural one). + # The beta projection (one row per value head) stays a separate bf16 layer: smaller than a + # quantization block, and a merged layer takes a single quantization scheme. + self.in_proj_qkvz = MergedColumnParallelLinear( + input_size=self.hidden_size, + output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim], + bias=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("in_proj_qkvz", prefix), + ) + self.in_proj_b = ColumnParallelLinear( + input_size=self.hidden_size, + output_size=self.num_v_heads, + bias=False, + quant_config=None, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("in_proj_b", prefix), + ) + # Low-rank forget gate: D -> bottleneck (replicated) -> HV*K (sharded by head). Never + # quantized: the decay gate is the precision-critical part of the recurrence. + self.f_down_proj = ReplicatedLinear( + self.hidden_size, + self.gate_bottleneck, + bias=False, + quant_config=None, + prefix=add_prefix("f_down_proj", prefix), + ) + self.f_up_proj = ColumnParallelLinear( + input_size=self.gate_bottleneck, + output_size=self.num_v_heads * self.head_k_dim, + bias=False, + quant_config=None, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("f_up_proj", prefix), + ) + + query_key_settings = (self.key_dim, 0, False) + value_settings = (self.value_dim, 0, False) + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": mamba_v2_sharded_weight_loader( + [query_key_settings, query_key_settings, value_settings], + self.attn_tp_size, + self.attn_tp_rank, + ) + }, + ) + self.A_log = nn.Parameter( + torch.empty(self.local_num_v_heads, dtype=torch.float32) + ) + self.dt_bias = nn.Parameter( + torch.empty(self.local_num_v_heads * self.head_k_dim, dtype=torch.float32) + ) + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + self.attn = RadixLinearAttention( + layer_id=layer_id, + num_q_heads=self.local_num_k_heads, + num_k_heads=self.local_num_k_heads, + num_v_heads=self.local_num_v_heads, + head_q_dim=self.head_k_dim, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + conv_weights=conv_weights, + bias=self.conv1d.bias, + activation=self.activation, + A_log=self.A_log, + dt_bias=self.dt_bias, + ) + self.attn.lower_bound = self.lower_bound + # BerryLM's per-channel log-decays reach -300 and beyond; the fused "small-grid" intra-chunk + # kernel factorizes exp2(g_i - g_j) with factors clamped to +-126 and collapses such blocks. + # Route prefill through the difference-form (token-parallel) intra-chunk kernels instead. + self.attn.kda_fused_intra = False + act = getattr(config, "kda_output_gate_act", "silu") + act = "swish" if act in ("silu", "swish") else act + self.norm = RMSNormGated( + self.head_v_dim, + eps=self.layer_norm_epsilon, + group_size=None, + norm_before_gate=True, + device=torch.get_device_module().current_device(), + dtype=torch.get_default_dtype(), + activation=act, + ) + self.out_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + reduce_results=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("out_proj", prefix), + ) + + def forward(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + num_tokens = hidden_states.shape[0] + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + qkv_size = (2 * self.key_dim + self.value_dim) // self.attn_tp_size + z_size = self.value_dim // self.attn_tp_size + mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) + z = z.reshape(num_tokens, -1, self.head_v_dim) + beta, _ = self.in_proj_b(hidden_states) + beta = beta.float() + raw_g = self.f_up_proj(self.f_down_proj(hidden_states)[0])[0] + # Prefill passes raw gates to chunk KDA; decode and target-verify kernels + # apply the activation internally (same contract as Kimi Linear). + if ( + not forward_batch.forward_mode.is_decode() + and not forward_batch.forward_mode.is_target_verify() + ): + raw_g = raw_g.unflatten(-1, (-1, self.head_k_dim)).unsqueeze(0) + beta = beta.sigmoid() + beta = beta.unsqueeze(0) + core_attn_out = self.attn(forward_batch, mixed_qkv=mixed_qkv, a=raw_g, b=beta) + core_attn_out = self.norm( + core_attn_out.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim) + ) + core_attn_out = core_attn_out.reshape(num_tokens, -1) + output, _ = self.out_proj(core_attn_out) + return output + + +# --------------------------------------------------------------------------- +# Full attention (gated, partial RoPE, QK Gemma-RMSNorm) +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Sparse MoE block (routed experts + a sigmoid-gated shared expert). The standard +# TP / EP / DeepEP paths only; kept in this file so the model has no dependency on +# another model file. +# --------------------------------------------------------------------------- +class BerryLMMLP(nn.Module): + """Dense SwiGLU MLP (the shared expert).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class BerryLMSparseMoeBlock(nn.Module): + def __init__( + self, + layer_id: int, + config: BerryLMConfig, + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_parallel().tp_size + self.layer_id = layer_id + self.alt_stream = alt_stream + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + self.num_experts = config.num_experts + a2a = get_moe_a2a_backend() + self._a2a_is_ep = a2a.is_deepep() or a2a.is_deepep_v2() or a2a.is_mori() + + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + layer_id=layer_id, + ) + # The shared-expert add reads hidden_states after the experts -> no inplace MoE. + self.experts = get_moe_impl_class(quant_config)( + layer_id=self.layer_id, + top_k=config.num_experts_per_tok, + num_experts=config.num_experts + get_exec().moe.ep_num_redundant_experts, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + prefix=add_prefix("experts", prefix), + routing_method_type=RoutingMethodType.RenormalizeNaive, + inplace=config.shared_expert_intermediate_size <= 0, + ) + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + if config.shared_expert_intermediate_size > 0: + self.shared_expert = BerryLMMLP( + hidden_size=config.hidden_size, + intermediate_size=config.shared_expert_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_expert", prefix), + **( + dict(tp_rank=0, tp_size=1) + if ( + self._a2a_is_ep + or a2a.is_flashinfer() + or a2a.is_flashinfer_megamoe() + ) + else {} + ), + ) + else: + self.shared_expert = None + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + + if self._a2a_is_ep: + # TODO: tp < ep is not supported yet + self.ep_size = get_parallel().moe_ep_size + self.num_experts = ( + config.num_experts + get_exec().moe.ep_num_redundant_experts + ) + self.top_k = config.num_experts_per_tok + + def get_moe_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + and filter_moe_weight_param_global_expert( + name, x, self.experts.num_local_experts + ) + ] + + def _forward_shared_experts( + self, hidden_states: torch.Tensor, apply_gate: bool = True + ): + if self.shared_expert is None: + return None + shared_output = self.shared_expert(hidden_states) + if apply_gate: + shared_output = ( + F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output + ) + return shared_output + + def _forward_router_experts(self, hidden_states: torch.Tensor): + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + return self.experts(hidden_states, topk_output) + + def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + shared_output = None + if hidden_states.shape[0] > 0: + router_logits, _ = self.gate(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.moe_num_token_non_padded(), + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + else: + topk_output = self.topk.empty_topk_output(hidden_states.device) + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + ) + if shared_output is not None: + final_hidden_states.add_(shared_output) + return final_hidden_states + + def forward_normal_dual_stream( + self, hidden_states: torch.Tensor, use_fused_gate: bool = False + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + # Shared expert on the alternate stream, routed experts on the current one + # (CUDA-graph capture path). + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + shared_output = ( + self._forward_shared_experts( + hidden_states.clone(), apply_gate=not use_fused_gate + ) + if self.shared_expert is not None + else None + ) + with torch.cuda.stream(self.alt_stream): + router_output = self._forward_router_experts(hidden_states) + current_stream.wait_stream(self.alt_stream) + return router_output, shared_output + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + ) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + if self._a2a_is_ep: + return self._forward_deepep(hidden_states, forward_batch) + + # sigmoid(gate) * shared + routed in one fused kernel + use_fused_gate = self.shared_expert_gate is not None + + if hidden_states.shape[0] == 0: + # M=0 guard for idle DP ranks: still call self.experts() so the rank takes + # part in the all-to-all collective. + shared_output = None + topk_output = self.topk.empty_topk_output(hidden_states.device) + final_hidden_states = self.experts(hidden_states, topk_output) + elif ( + self.alt_stream is not None + and get_is_capture_mode() + and not torch.compiler.is_compiling() + ): + final_hidden_states, shared_output = self.forward_normal_dual_stream( + hidden_states, use_fused_gate=use_fused_gate + ) + else: + shared_output = self._forward_shared_experts( + hidden_states, apply_gate=not use_fused_gate + ) + final_hidden_states = self._forward_router_experts(hidden_states) + + if shared_output is not None: + if use_fused_gate: + fused_gate_sigmoid_mul_add( + hidden_states, + self.shared_expert_gate.weight.squeeze(), + shared_output, + final_hidden_states, + ) + else: + final_hidden_states += shared_output + if ( + self.tp_size > 1 + and not should_skip_post_experts_all_reduce(is_tp_path=True) + and not get_moe_a2a_backend().is_flashinfer() + ): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class BerryLMAttention(nn.Module): + def __init__( + self, + config: BerryLMConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_tp_size + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % self.attn_tp_size == 0 + self.num_heads = self.total_num_heads // self.attn_tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= self.attn_tp_size: + assert self.total_num_kv_heads % self.attn_tp_size == 0 + else: + assert self.attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) + self.head_dim = config.head_dim or (self.hidden_size // self.num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.layer_id = layer_id + self.alt_stream = alt_stream + self.attn_output_gate = getattr(config, "attn_output_gate", True) + + rope = ( + getattr(config, "rope_parameters", None) + or getattr(config, "rope_scaling", None) + or {} + ) + self.rotary_emb = get_rope( + head_size=self.head_dim, + rotary_dim=self.head_dim, + max_position=getattr(config, "max_position_embeddings", 8192), + rope_scaling=rope, + base=rope.get("rope_theta", 10000), + partial_rotary_factor=rope.get("partial_rotary_factor", 1.0), + is_neox_style=True, + dtype=torch.get_default_dtype(), + ) + attn_quant_config = ( + None + if quant_config and quant_config.get_name() == "modelopt_fp4" + else quant_config + ) + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads * (1 + self.attn_output_gate), + self.total_num_kv_heads, + bias=False, + quant_config=attn_quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=attn_quant_config, + reduce_results=False, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + prefix=add_prefix("o_proj", prefix), + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + prefix=add_prefix("attn", prefix), + ) + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if self.alt_stream is not None and get_is_capture_mode(): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + q_by_head = self.q_norm(q.reshape(-1, self.head_dim)) + with torch.cuda.stream(self.alt_stream): + k_by_head = self.k_norm(k.reshape(-1, self.head_dim)) + current_stream.wait_stream(self.alt_stream) + else: + q_by_head = self.q_norm(q.reshape(-1, self.head_dim)) + k_by_head = self.k_norm(k.reshape(-1, self.head_dim)) + return q_by_head.view(q.shape), k_by_head.view(k.shape) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + if self.attn_output_gate: + q_gate, k, v = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + orig_shape = q_gate.shape[:-1] + q_gate = q_gate.view(*orig_shape, self.num_heads, -1) + q, gate = torch.chunk(q_gate, 2, dim=-1) + q = q.reshape(*orig_shape, -1) + gate = gate.reshape(*orig_shape, -1) + else: + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + gate = None + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + if gate is not None: + attn_output = attn_output * torch.sigmoid(gate) + output, _ = self.o_proj(attn_output) + return output + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + + +class BerryLMDecoderLayer(nn.Module): + """One decoder layer: KDA linear attention or full attention, then sparse MoE. + The AttnRes mixer is owned by the layer (``attn_res``) but applied in the model + loop, on the residual stream, before ``input_layernorm``.""" + + def __init__( + self, + config: BerryLMConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + self.config = config + self.layer_id = layer_id + self.layer_type = config.layers_block_type[layer_id] + if self.layer_type == "linear_attention": + if not getattr(config, "kda_insert", True): + raise NotImplementedError( + "BerryLM without the KDA gate is not supported" + ) + linear_attn_quant_config = ( + None + if quant_config and quant_config.get_name() == "modelopt_fp4" + else quant_config + ) + self.linear_attn = BerryLMKDA( + config, + layer_id, + linear_attn_quant_config, + prefix=add_prefix("linear_attn", prefix), + ) + else: + self.self_attn = BerryLMAttention( + config, + layer_id, + quant_config, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + ) + self.mlp = BerryLMSparseMoeBlock( + layer_id=layer_id, + config=config, + quant_config=quant_config, + alt_stream=alt_stream, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + # every layer is sparse (MoE FFN), so are its neighbours + self.attn_boundary, self.ffn_boundary = make_stages( + (declare_attn(), self.input_layernorm), + ( + declare_ffn(sparse=True, next_layer_sparse=True), + self.post_attention_layernorm, + ), + previous=declare_ffn(sparse=True, next_layer_sparse=True) + if layer_id != 0 + else None, + terminal=layer_id == config.num_hidden_layers - 1, + ) + if config.attn_res_block_size > 0: + self.attn_res = BerryLMAttnRes( + config.hidden_size, + eps=getattr(config, "attn_res_eps", 1e-6), + gated=getattr(config, "attn_res_gated", True), + ) + else: + self.attn_res = None + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + hidden_states = self.attn_boundary.prepare(hidden_states, forward_batch) + if not forward_batch.forward_mode.is_idle(): + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states, forward_batch) + else: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = self.attn_boundary.finish(hidden_states, forward_batch) + hidden_states = self.ffn_boundary.prepare(hidden_states, forward_batch) + with self.ffn_boundary.exit(forward_batch) as ffn_exit: + hidden_states = self.mlp(hidden_states, forward_batch) + return ffn_exit.finish(hidden_states) + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + + +class BerryLMModel(nn.Module): + def __init__( + self, + config: BerryLMConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.attn_res_block_size = int(config.attn_res_block_size) + self.pp_group = get_parallel().pp_group + if self.pp_group.world_size > 1: + raise NotImplementedError( + "BerryLM does not support pipeline parallelism yet" + ) + if self.attn_res_block_size > 0 and is_dp_attention_enabled(): + # The mixer needs every rank to hold the full residual stream of every + # token; DP-attention scatter modes break that invariant. + raise NotImplementedError( + "BerryLM AttnRes does not support DP attention yet" + ) + alt_stream = get_stream("alt") if _is_cuda else None + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers, self._start_layer, self._end_layer = make_pp_layers( + config.num_hidden_layers, + lambda idx, prefix: BerryLMDecoderLayer( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=alt_stream, + ), + prefix=add_prefix("layers", prefix), + ) + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual_batch.start(forward_batch) + block = self.attn_res_block_size + # Residual streams committed at block boundaries (embeddings first). + completed: List[torch.Tensor] = [] + for layer_idx in range(self._start_layer, self._end_layer): + layer = self.layers[layer_idx] + with get_global_expert_distribution_recorder().with_current_layer( + layer_idx + ): + if block: + # The mixer reads the complete residual stream: fold the + # previous layer's output into it, mix, and hand the mixed + # stream to the layer as an already-written residual. + stream = residual_batch.fold(hidden_states, forward_batch) + if layer_idx % block == 0: + completed.append(stream) + hidden_states = residual_batch.set_written( + layer.attn_res(stream, completed), forward_batch + ) + hidden_states = layer( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + return residual_batch.final_norm( + hidden_states, forward_batch, self.norm, skip_empty=True + ) + + +class BerryLMForCausalLM(nn.Module): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__( + self, + config: BerryLMConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = BerryLMModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + # transformers-v5 fused expert tensors: experts.gate_up_proj [E, 2I, D], + # experts.down_proj [E, D, I]. + fused_expert_params_mapping = [ + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), + ] + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + num_experts = self.config.num_experts + is_fused_expert = False + + def load_fused_expert_weights(name, params_dict, loaded_weight, shard_id): + param = params_dict[name] + weight_loader = param.weight_loader + for expert_id in range(num_experts): + weight_loader( + param, loaded_weight[expert_id], name, shard_id, expert_id + ) + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "mtp" in name or "visual" in name: + continue + if "language_model" in name: + name = name.replace("model.language_model.", "model.") + # the checkpoint calls the KDA mixer `kda`; `linear_attn` is the layer attribute SGLang + # looks up for hybrid linear-attention layers (model_runner_components/layer_setup.py) + name = name.replace(".kda.", ".linear_attn.") + if "experts.gate_up_proj" in name or "experts.down_proj" in name: + is_fused_expert = True + expert_params_mapping = fused_expert_params_mapping + if ".linear_attn.in_proj_qkv." in name or ".linear_attn.in_proj_z." in name: + # in_proj_qkv [2K+V, D] -> partitions 0..2 (split here: the v1 loader takes one index), + # in_proj_z -> 3 of the merged GEMM; in_proj_b is its own layer (default loader) + key_dim = ( + self.config.linear_num_key_heads * self.config.linear_key_head_dim + ) + value_dim = ( + self.config.linear_num_value_heads + * self.config.linear_value_head_dim + ) + if ".in_proj_qkv." in name: + name = name.replace(".in_proj_qkv.", ".in_proj_qkvz.") + param = params_dict[name] + # block-quantized checkpoints ship `weight_scale_inv` with one row per + # 128-row block: split it in the same proportions as the weight + scale = loaded_weight.shape[0] / (2 * key_dim + value_dim) + sizes = [ + int(key_dim * scale), + int(key_dim * scale), + int(value_dim * scale), + ] + for sid, part in enumerate(loaded_weight.split(sizes, dim=0)): + param.weight_loader(param, part, sid) + else: + name = name.replace(".in_proj_z.", ".in_proj_qkvz.") + param = params_dict[name] + param.weight_loader(param, loaded_weight, 3) + loaded_params.add(name) + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # ".linear_attn." tensors are never stacked here (f_up_proj contains "up_proj"). + if ( + weight_name not in name + or "mlp.experts" in name + or ".linear_attn." in name + ): + continue + name = name.replace(weight_name, param_name) + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for ( + param_name, + weight_name, + expert_id, + shard_id, + ) in expert_params_mapping: + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + if is_fused_expert: + if "experts.gate_up_proj" in name: + w1, w3 = loaded_weight.chunk(2, dim=-2) + load_fused_expert_weights( + name_mapped, params_dict, w1, "w1" + ) + load_fused_expert_weights( + name_mapped, params_dict, w3, "w3" + ) + else: + load_fused_expert_weights( + name_mapped, params_dict, loaded_weight, shard_id + ) + else: + if ( + name_mapped.endswith(ignore_suffixes) + and name_mapped not in params_dict + ): + continue + param = params_dict[name_mapped] + param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + ) + name = name_mapped + break + else: + if is_expert_weight: + continue + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + logger.warning(f"Parameter {name} not found in params_dict") + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None, + ) + + +EntryClass = BerryLMForCausalLM diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 4dec02573a15..c247d977979f 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -2239,6 +2239,48 @@ def finish(self) -> StreamingParseResult: return StreamingParseResult(normal_text=normal, reasoning_text=reasoning) +class BerryLMDetector(BaseReasoningFormatDetector): + """ + Detector for BerryLM models. Assumes reasoning format: + ()*(.*) + + ``enable_thinking`` in the request switches between thinking mode and normal mode: + - enable_thinking=True: "reasoning contentThe answer is 42." + - enable_thinking=False: "The answer is 42." (no thinking tokens) + + Args: + stream_reasoning (bool): If False, accumulates reasoning content until the end tag. + If True, streams reasoning content as it arrives. + """ + + def __init__( + self, + stream_reasoning: bool = True, + force_reasoning: bool = False, + continue_final_message: bool = False, + previous_content: str = "", + force_nonempty_content: bool = False, + ): + think_excluded_tokens = [ + "", + "", + "<|im_end|>", + "<|endoftext|>", + ] + super().__init__( + "", + "", + think_excluded_tokens=think_excluded_tokens, + force_reasoning=force_reasoning, + stream_reasoning=stream_reasoning, + # the model may open ```` without closing ````: implicit reasoning end + tool_start_token="", + continue_final_message=continue_final_message, + previous_content=previous_content, + force_nonempty_content=force_nonempty_content, + ) + + class ReasoningParser: """ Parser that handles both streaming and non-streaming scenarios for extracting @@ -2256,6 +2298,7 @@ class ReasoningParser: DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = { "apertus2509": Apertus2509Detector, + "berrylm": BerryLMDetector, "deepseek-r1": DeepSeekR1Detector, "deepseek-v3": _DeepSeekV3Detector, "deepseek-v4": DeepSeekV4Detector, diff --git a/python/sglang/srt/parser/reasoning_parser_names.py b/python/sglang/srt/parser/reasoning_parser_names.py index c06157212976..7c9c72d05bf4 100644 --- a/python/sglang/srt/parser/reasoning_parser_names.py +++ b/python/sglang/srt/parser/reasoning_parser_names.py @@ -5,6 +5,7 @@ REASONING_PARSER_NAMES = [ "apertus2509", + "berrylm", "deepseek-r1", "deepseek-v3", "deepseek-v4", diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 5c90ddfd93df..771d78feae34 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -26,6 +26,7 @@ BailingHybridConfig, BailingMM2Config, BailingMoeV3VLConfig, + BerryLMConfig, ChatGLMConfig, Cosmos3Config, Cosmos3EdgeConfig, @@ -119,6 +120,7 @@ BailingHybridConfig, BailingMM2Config, BailingMoeV3VLConfig, + BerryLMConfig, ChatGLMConfig, DbrxConfig, ExaoneConfig, diff --git a/test/registered/e2e/models/test_berrylm.py b/test/registered/e2e/models/test_berrylm.py new file mode 100644 index 000000000000..ee6910bbf3e3 --- /dev/null +++ b/test/registered/e2e/models/test_berrylm.py @@ -0,0 +1,43 @@ +"""BerryLM-OS: hybrid MoE with KDA gated delta-net linear attention and block +attention residuals (1 GPU, GSM8K accuracy through the default server fixture).""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase + +# 1-gpu-large: the runner class of the other ~18B-parameter model tests. +register_cuda_ci( + est_time=600, + disabled="rwb-ai/BerryLM-OS is not public yet", + stage="extra-a", + runner_config="1-gpu-large", +) + +BERRYLM_MODEL = "rwb-ai/BerryLM-OS" + + +class TestBerryLM(GSM8KMixin, DefaultServerBase): + model = BERRYLM_MODEL + # Thinking is on by default: scored through the chat API with the reasoning + # parser, as the other reasoning-model tests do (the completion backend's + # 512-token cap cuts the reasoning off). + gsm8k_backend = "sgl_eval" + gsm8k_thinking = True + gsm8k_num_examples = 200 + gsm8k_num_threads = 32 + gsm8k_max_tokens = 16384 + gsm8k_score_threshold = 0.90 # 0.95 on the release checkpoint + other_args = [ + "--tp-size", + "1", + "--reasoning-parser", + "berrylm", + "--tool-call-parser", + "berrylm", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/function_call/test_berrylm_detector.py b/test/registered/unit/function_call/test_berrylm_detector.py new file mode 100644 index 000000000000..3eed29f3561f --- /dev/null +++ b/test/registered/unit/function_call/test_berrylm_detector.py @@ -0,0 +1,253 @@ +"""Unit tests of the BerryLM parsers: the XML tool-call detector +(``--tool-call-parser berrylm``) and the ```` reasoning detector +(``--reasoning-parser berrylm``). CPU only, no model.""" + +import json +import unittest + +from sglang.srt.entrypoints.openai.protocol import Function, Tool +from sglang.srt.function_call.berrylm_detector import BerryLMFunctionCallDetector +from sglang.srt.parser.reasoning_parser import BerryLMDetector, ReasoningParser +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(5, "base-a-test-cpu") + +TOOL_CALL_START = "" +TOOL_CALL_END = "" + +WEATHER_CALL = ( + f"{TOOL_CALL_START}\n\n" + "Tokyo\n" + "3\n" + f"\n{TOOL_CALL_END}" +) + + +def _make_tool(name="get_weather"): + return Tool( + type="function", + function=Function( + name=name, + description=f"{name} tool", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string", "description": "City"}, + "days": {"type": "integer", "description": "Forecast days"}, + "path": {"type": "string"}, + "content": {"type": "string"}, + }, + "required": ["city"], + }, + ), + ) + + +def _collect_streaming_tool_calls(detector, chunks, tools): + """Run streaming chunks through a detector and collect assembled tool calls.""" + tool_calls = [] + all_normal_text = "" + for chunk in chunks: + result = detector.parse_streaming_increment(chunk, tools) + all_normal_text += result.normal_text + for tc_chunk in result.calls: + if tc_chunk.tool_index is not None: + while len(tool_calls) <= tc_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tc_chunk.tool_index] + if tc_chunk.name: + tc["name"] = tc_chunk.name + if tc_chunk.parameters: + tc["parameters"] += tc_chunk.parameters + return tool_calls, all_normal_text + + +class TestBerryLMFunctionCallDetector(unittest.TestCase): + def setUp(self): + self.detector = BerryLMFunctionCallDetector() + self.tools = [_make_tool(), _make_tool("write_file")] + + def test_has_tool_call(self): + self.assertTrue(self.detector.has_tool_call(WEATHER_CALL)) + self.assertFalse(self.detector.has_tool_call("Just an answer.")) + + def test_single_tool_call_typed_arguments(self): + result = self.detector.detect_and_parse(WEATHER_CALL, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + # ``days`` is an integer in the tool schema: converted from the XML text + self.assertEqual( + json.loads(result.calls[0].parameters), {"city": "Tokyo", "days": 3} + ) + self.assertEqual(result.normal_text.strip(), "") + + def test_multiple_tool_calls(self): + text = WEATHER_CALL + ( + f"\n{TOOL_CALL_START}\n\n" + "Osaka\n" + f"\n{TOOL_CALL_END}" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual([c.name for c in result.calls], ["get_weather", "get_weather"]) + self.assertEqual( + json.loads(result.calls[0].parameters), {"city": "Tokyo", "days": 3} + ) + self.assertEqual(json.loads(result.calls[1].parameters), {"city": "Osaka"}) + + def test_normal_text_before_tool_call(self): + result = self.detector.detect_and_parse( + "Checking the forecast.\n" + WEATHER_CALL, self.tools + ) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.normal_text.strip(), "Checking the forecast.") + + def test_no_tool_call(self): + result = self.detector.detect_and_parse("Just an answer.", self.tools) + self.assertEqual(len(result.calls), 0) + self.assertEqual(result.normal_text, "Just an answer.") + + def test_multiline_value(self): + text = ( + f"{TOOL_CALL_START}\n\n" + "/tmp/x.py\n" + "\nline 1\nline 2\n\n" + f"\n{TOOL_CALL_END}" + ) + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(result.calls[0].name, "write_file") + self.assertEqual( + json.loads(result.calls[0].parameters), + {"path": "/tmp/x.py", "content": "line 1\nline 2"}, + ) + + def test_streaming_single_tool_call(self): + chunks = [ + f"{TOOL_CALL_START}\n", + "\n", + "Tok", + "yo\n", + "3\n", + "\n", + f"{TOOL_CALL_END}", + ] + calls, normal = _collect_streaming_tool_calls( + BerryLMFunctionCallDetector(), chunks, self.tools + ) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["name"], "get_weather") + self.assertEqual( + json.loads(calls[0]["parameters"]), {"city": "Tokyo", "days": 3} + ) + self.assertEqual(normal.strip(), "") + + def test_streaming_multiple_tool_calls(self): + chunks = [ + f"{TOOL_CALL_START}\n\n", + "Tokyo\n\n", + f"{TOOL_CALL_END}\n{TOOL_CALL_START}\n", + "\nOsaka\n", + f"\n{TOOL_CALL_END}", + ] + calls, _ = _collect_streaming_tool_calls( + BerryLMFunctionCallDetector(), chunks, self.tools + ) + self.assertEqual([c["name"] for c in calls], ["get_weather", "get_weather"]) + self.assertEqual(json.loads(calls[0]["parameters"]), {"city": "Tokyo"}) + self.assertEqual(json.loads(calls[1]["parameters"]), {"city": "Osaka"}) + + def test_streaming_plain_text_passes_through(self): + det = BerryLMFunctionCallDetector() + normal = "" + for chunk in ["Just ", "an answer", " with a < sign."]: + normal += det.parse_streaming_increment(chunk, self.tools).normal_text + self.assertEqual(normal, "Just an answer with a < sign.") + + def test_registered_in_function_call_parser(self): + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + self.assertIs( + FunctionCallParser.ToolCallParserEnum["berrylm"], + BerryLMFunctionCallDetector, + ) + + +class TestBerryLMReasoningDetector(unittest.TestCase): + def test_registered(self): + parser = ReasoningParser(model_type="berrylm") + self.assertIsInstance(parser.detector, BerryLMDetector) + + def test_think_then_content(self): + det = BerryLMDetector() + result = det.detect_and_parse( + "Plan the answer.The answer is 42." + ) + self.assertEqual(result.reasoning_text, "Plan the answer.") + self.assertEqual(result.normal_text, "The answer is 42.") + + def test_template_opened_think(self): + """The chat template ends the prompt with ````, so the generated + text starts inside the reasoning block; the server passes + ``force_reasoning=True`` for such templates (template manager).""" + det = BerryLMDetector(force_reasoning=True) + result = det.detect_and_parse("Plan the answer.The answer is 42.") + self.assertEqual(result.reasoning_text, "Plan the answer.") + self.assertEqual(result.normal_text, "The answer is 42.") + + def test_no_forced_reasoning_is_plain_content(self): + """Without ``force_reasoning`` (thinking off: the template closes an + empty think block in the prompt) text without ```` is content.""" + det = BerryLMDetector() + result = det.detect_and_parse("The answer is 42.") + self.assertEqual(result.reasoning_text, "") + self.assertEqual(result.normal_text, "The answer is 42.") + + def test_tool_call_without_think_end(self): + """The model may open ```` without closing ````: the + tool call must reach the tool parser, never the reasoning text.""" + det = BerryLMDetector() + result = det.detect_and_parse("I need the forecast." + WEATHER_CALL) + self.assertNotIn(TOOL_CALL_START, result.reasoning_text) + self.assertIn("I need the forecast.", result.reasoning_text) + self.assertNotIn("", result.reasoning_text) + self.assertTrue(result.normal_text.startswith(TOOL_CALL_START)) + + def _run_streaming(self, chunks, **kwargs): + det = BerryLMDetector(**kwargs) + reasoning, normal = "", "" + for chunk in chunks: + r = det.parse_streaming_increment(chunk) + reasoning += r.reasoning_text + normal += r.normal_text + return reasoning, normal + + def test_streaming_think_then_content(self): + reasoning, normal = self._run_streaming( + ["", "Plan the ", "answer.", "", "The answer", " is 42."] + ) + self.assertEqual(reasoning, "Plan the answer.") + self.assertEqual(normal, "The answer is 42.") + + def test_streaming_template_opened_think(self): + reasoning, normal = self._run_streaming( + ["Plan the ", "answer.", "", "The answer", " is 42."], + force_reasoning=True, + ) + self.assertEqual(reasoning, "Plan the answer.") + self.assertEqual(normal, "The answer is 42.") + + def test_streaming_tool_call_without_think_end(self): + reasoning, normal = self._run_streaming( + [ + "", + "I need the forecast.", + TOOL_CALL_START, + "\n", + ] + ) + self.assertNotIn(TOOL_CALL_START, reasoning) + self.assertIn("I need the forecast.", reasoning) + self.assertIn(TOOL_CALL_START, normal) + + +if __name__ == "__main__": + unittest.main()