From ea2b4db9a49985c44486b2267ec4490d453c4168 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:01:27 +0000 Subject: [PATCH 01/13] fix: parse GLM tool calls inside reasoning --- lightllm/server/reasoning_parser.py | 66 ++++++- .../server/test_glm_reasoning_tool_calls.py | 172 ++++++++++++++++++ 2 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 unit_tests/server/test_glm_reasoning_tool_calls.py diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 1e7c2f1315..251a399bd1 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -756,6 +756,70 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False) ) +class Glm45Detector(Qwen3Detector): + """GLM tool calls implicitly end thinking, even without ````.""" + + tool_start_token = "" + tool_end_token = "" + + def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False): + super().__init__(stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) + self._in_tool_call = False + + def detect_and_parse(self, text: str) -> StreamingParseResult: + detector = type(self)(stream_reasoning=self.stream_reasoning, force_reasoning=self._in_reasoning) + result = detector.parse_streaming_increment(text) + remaining = detector.flush() + return StreamingParseResult( + reasoning_text=(result.reasoning_text + remaining.reasoning_text).strip(), + normal_text=result.normal_text + remaining.normal_text, + ) + + def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + self._buffer += new_text + reasoning_parts = [] + normal_parts = [] + markers = [self.think_start_token, self.think_end_token, self.tool_start_token] + + while self._buffer: + if self._in_tool_call: + end = self._buffer.find(self.tool_end_token) + if end == -1: + text, self._buffer = prefix_hold(self._buffer, [self.tool_end_token]) + normal_parts.append(text) + break + end += len(self.tool_end_token) + normal_parts.append(self._buffer[:end]) + self._buffer = self._buffer[end:] + self._in_tool_call = False + continue + + positions = [(self._buffer.find(marker), marker) for marker in markers] + positions = [(position, marker) for position, marker in positions if position != -1] + parts = reasoning_parts if self._in_reasoning else normal_parts + if not positions: + # A delimiter may share a chunk with reasoning or span chunks. + text, self._buffer = prefix_hold(self._buffer, markers) + parts.append(text) + break + + position, marker = min(positions) + parts.append(self._buffer[:position]) + if marker == self.tool_start_token: + self._in_reasoning = False + self._in_tool_call = True + # Preserve the opener and literal thinking tags in tool arguments. + normal_parts.append(marker) + else: + self._in_reasoning = marker == self.think_start_token + self._buffer = self._buffer[position + len(marker) :] + + return StreamingParseResult( + reasoning_text="".join(reasoning_parts), + normal_text="".join(normal_parts), + ) + + class KimiDetector(BaseReasoningFormatDetector): """ Detector for Kimi Thinking model. @@ -919,7 +983,7 @@ class ReasoningParser: DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = { "deepseek-r1": DeepSeekR1Detector, "deepseek-v3": Qwen3Detector, - "glm45": Qwen3Detector, + "glm45": Glm45Detector, "gpt-oss": GptOssDetector, "kimi": KimiDetector, "kimi_k2": DeepSeekR1Detector, diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py new file mode 100644 index 0000000000..9040910bae --- /dev/null +++ b/unit_tests/server/test_glm_reasoning_tool_calls.py @@ -0,0 +1,172 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from lightllm.server import api_http, api_openai, build_prompt, reasoning_parser +from lightllm.server.api_models import ChatCompletionRequest, Tool +from lightllm.server.core.objs import FinishStatus +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.function_call_parser import FunctionCallParser +from lightllm.server.httpserver.stop_seq_match import ReasoningStateMachine +from lightllm.server.reasoning_parser import ReasoningParser + + +TOOLS = [ + Tool.model_validate( + { + "type": "function", + "function": { + "name": "Write", + "parameters": {"type": "object", "properties": {"content": {"type": "string"}}}, + }, + } + ) +] +VALUE = "line 1\nliteral\nline 2" +TOOL_CALL = f"Writecontent{VALUE}" + + +@pytest.fixture(autouse=True) +def mock_token_ids(monkeypatch): + monkeypatch.setattr( + reasoning_parser, "get_token_id", {"": 98, "": 99, "": 77}.__getitem__ + ) + + +@pytest.mark.parametrize( + "prefix,force_reasoning", + [ + ("Need a file.", False), + ("Need a file.", True), + ("Need a file.", False), + ("Need a file.", True), + ], +) +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +def test_glm_reasoning_and_tools_agree_across_stream_boundaries(prefix, force_reasoning, chunk_size): + text = prefix + TOOL_CALL + TOOL_CALL + parser = ReasoningParser("glm45", force_reasoning=force_reasoning) + reasoning, content = parser.parse_non_stream(text) + assert reasoning == "Need a file." + assert content == TOOL_CALL + TOOL_CALL + + parser = ReasoningParser("glm45", force_reasoning=force_reasoning) + tool_parser = FunctionCallParser(TOOLS, "glm47") + reasoning_parts, content_parts, calls = [], [], {} + for start in range(0, len(text), chunk_size): + reasoning, content = parser.parse_stream_chunk(text[start : start + chunk_size]) + reasoning_parts.append(reasoning) + content, deltas = tool_parser.parse_stream_chunk(content) + content_parts.append(content) + for delta in deltas: + if delta.name is not None: + assert delta.tool_index not in calls + calls[delta.tool_index] = {"name": delta.name, "arguments": ""} + calls[delta.tool_index]["arguments"] += delta.parameters + + assert parser.flush() == ("", "") + assert "".join(reasoning_parts) == "Need a file." + assert "".join(content_parts) == "" + assert list(calls) == [0, 1] + assert all(call["name"] == "Write" for call in calls.values()) + assert all(json.loads(call["arguments"]) == {"content": VALUE} for call in calls.values()) + + +@pytest.mark.parametrize("suffix", [""]) +def test_truncated_or_similar_markers_remain_reasoning(suffix): + parser = ReasoningParser("glm45", force_reasoning=True) + reasoning, content = parser.parse_stream_chunk("Keep thinking." + suffix) + tail_reasoning, tail_content = parser.flush() + assert reasoning + tail_reasoning == "Keep thinking." + suffix + assert content + tail_content == "" + + +def test_non_thinking_glm_preserves_tool_payload(): + parser = ReasoningParser("glm45", force_reasoning=False) + assert parser.parse_non_stream(TOOL_CALL) == ("", TOOL_CALL) + assert parser.parse_stream_chunk(TOOL_CALL) == ("", TOOL_CALL) + + +def test_other_reasoning_parsers_keep_their_existing_boundary(): + parser = ReasoningParser("qwen3", force_reasoning=True) + reasoning, content = parser.parse_non_stream("Still thinking." + TOOL_CALL.replace(VALUE, "literal")) + assert reasoning.startswith("Still thinking.") + assert content == "" + + +def test_glm_tool_parsing_keeps_existing_reasoning_token_accounting(): + parser = ReasoningParser("glm45", force_reasoning=True) + for token_id in [1, 2, 77, 3, 99, 4]: + parser.update_reasoning_token_count(token_id) + assert parser.reasoning_tokens == 4 + + +def test_glm_tools_inside_reasoning_do_not_enable_string_stops(): + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [98], "": [99]}[text]) + state = ReasoningStateMachine("glm45", tokenizer, 1, []) + assert not any(state.handle(token_id) for token_id in [1, 77, 2, 3]) + assert state.in_reasoning + assert state.handle(99) is False + assert not state.in_reasoning + assert state.handle(4) is True + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("reasoning_end", ["missing", "before_tool", "after_tool"]) +def test_chat_completion_exposes_glm_tool_calls(monkeypatch, stream, reasoning_end): + args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47") + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="")) + segments = [("Need a file.", 1)] + if reasoning_end == "before_tool": + segments.append(("", 99)) + segments.extend([("", 77), (TOOL_CALL[len("") :], 2)]) + if reasoning_end == "after_tool": + segments.append(("", 99)) + + async def generate(*args, **kwargs): + for index, (text, token_id) in enumerate(segments): + finish = FinishStatus(FinishStatus.FINISHED_STOP) if index == len(segments) - 1 else FinishStatus() + yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish + + manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) + request = ChatCompletionRequest(model="glm", messages=[], tools=TOOLS, stream=stream) + + async def run(): + response = await api_openai.chat_completions_impl(request, None) + if not stream: + choice = response.choices[0] + return ( + choice.message.reasoning, + choice.message.content, + choice.message.tool_calls[0].function.arguments, + choice.finish_reason, + ) + + reasoning, content, arguments, finish_reason = "", "", "", None + async for chunk in response.body_iterator: + chunk = chunk.decode() if isinstance(chunk, bytes) else chunk + if chunk.strip() == "data: [DONE]": + continue + payload = json.loads(chunk.removeprefix("data: ")) + for choice in payload["choices"]: + delta = choice["delta"] + reasoning += delta.get("reasoning") or "" + content += delta.get("content") or "" + for call in delta.get("tool_calls") or []: + assert call["index"] == 0 + assert call["function"].get("name") in (None, "Write") + arguments += call["function"].get("arguments") or "" + finish_reason = choice.get("finish_reason") or finish_reason + return reasoning, content, arguments, finish_reason + + reasoning, content, arguments, finish_reason = asyncio.run(run()) + assert reasoning == "Need a file." + assert content == "" + assert json.loads(arguments) == {"content": VALUE} + assert finish_reason == "tool_calls" From e8c8bc883869a1c7e9b986044aee499019d44d20 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:52:11 +0000 Subject: [PATCH 02/13] fix: harden GLM reasoning and XML tool parsing --- lightllm/server/api_openai.py | 34 ++-- lightllm/server/function_call_parser.py | 47 +++--- lightllm/server/glm_xml_parser.py | 148 ++++++++++++++++++ lightllm/server/reasoning_parser.py | 27 +++- unit_tests/server/test_glm47_tool_parser.py | 81 +++++++++- .../server/test_glm_reasoning_tool_calls.py | 139 ++++++++++++++-- 6 files changed, 420 insertions(+), 56 deletions(-) create mode 100644 lightllm/server/glm_xml_parser.py diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index fa61a311f4..29adeffefc 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -180,13 +180,17 @@ def _get_history_tool_calls_cnt(request: ChatCompletionRequest) -> int: return idx -def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool: +def _is_force_thinking_mode(request: ChatCompletionRequest, prompt: Optional[str] = None) -> bool: """Whether this request uses forced thinking / reasoning (parser + template).""" from .build_prompt import tokenizer_supports_force_thinking reasoning_parser = get_env_start_args().reasoning_parser if not reasoning_parser: return False + if reasoning_parser == "glm45" and prompt is not None: + state = ReasoningParser.DetectorMap["glm45"].initial_reasoning_state(prompt) + if state is not None: + return state # Always-thinking models may have no thinking switch in their chat template. # Their answer grammar must wait for the reasoning delimiter nonetheless. if reasoning_parser in ["deepseek-r1", "qwen3-thinking", "gpt-oss", "minimax"]: @@ -201,9 +205,9 @@ def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool: return True # default -def _get_guided_reasoning_end(request: ChatCompletionRequest, tokenizer) -> List[int]: +def _get_guided_reasoning_end(request: ChatCompletionRequest, tokenizer, prompt: Optional[str] = None) -> List[int]: """Use the same thinking mode as the chat template for grammar activation.""" - if not _is_force_thinking_mode(request): + if not _is_force_thinking_mode(request, prompt): return [] parser_name = get_env_start_args().reasoning_parser detector = ReasoningParser.DetectorMap[parser_name.lower()]() @@ -217,13 +221,14 @@ def _process_reasoning_stream( reasoning_parser_dict: Dict[int, ReasoningParser], metadata: Dict[str, Any], request: ChatCompletionRequest, + force_reasoning: Optional[bool] = None, ) -> tuple[Optional[str], str]: """Process reasoning content and update its token usage.""" if index not in reasoning_parser_dict: reasoning_parser_dict[index] = ReasoningParser( get_env_start_args().reasoning_parser, request.stream_reasoning, - _is_force_thinking_mode(request), + _is_force_thinking_mode(request) if force_reasoning is None else force_reasoning, ) parser = reasoning_parser_dict[index] token_id = metadata.get("id") @@ -353,6 +358,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req multimodal_params = _build_multimodal_params(request) tools = _select_chat_template_tools(request) prompt = await build_prompt(request, tools) + request_enable_reasoning = _is_force_thinking_mode(request, prompt) sampling_params_dict = { "do_sample": request.do_sample, "presence_penalty": request.presence_penalty, @@ -387,7 +393,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req sampling_params_dict["stop_sequences"] = request.stop if get_env_start_args().reasoning_parser: # The prompt may already open reasoning; generated tokens may omit the opening marker. - sampling_params_dict["_initial_reasoning_state"] = int(_is_force_thinking_mode(request)) + sampling_params_dict["_initial_reasoning_state"] = int(request_enable_reasoning) # Chat 接口提取 reasoning，按请求设置 stop 范围；通用 generate 接口不区分阶段。 sampling_params_dict["stop_str_match_mode"] = ( StopStrMatchMode.ALL_TEXT if is_stop_matching_in_reasoning_enabled() else StopStrMatchMode.CONTENT_ONLY @@ -405,7 +411,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req if sampling_params_dict.get("guided_json") or sampling_params_dict.get("guided_grammar"): sampling_params_dict["guided_reasoning_end"] = _get_guided_reasoning_end( - request, g_objs.httpserver_manager.tokenizer + request, g_objs.httpserver_manager.tokenizer, prompt ) sampling_params = SamplingParams() @@ -419,7 +425,6 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req # Non-streaming case if not request.stream: reasoning_parser = get_env_start_args().reasoning_parser - request_enable_reasoning = _is_force_thinking_mode(request) if reasoning_parser else False final_output_dict = collections.defaultdict(list) count_output_tokens_dict = collections.defaultdict(lambda: 0) finish_reason_dict = {} @@ -592,8 +597,19 @@ async def stream_results() -> AsyncGenerator[bytes, None]: # Handle reasoning content if get_env_start_args().reasoning_parser: reasoning_text, delta = _process_reasoning_stream( - choice_index, delta, reasoning_parser_dict, metadata, request + choice_index, + delta, + reasoning_parser_dict, + metadata, + request, + request_enable_reasoning if get_env_start_args().reasoning_parser == "glm45" else None, ) + if current_finish_reason is not None and get_env_start_args().reasoning_parser == "glm45": + # Route a buffered GLM suffix through the tool parser too; + # truncating inside an XML argument must not leak content. + flush_reasoning, flush_text = reasoning_parser_dict[choice_index].flush() + reasoning_text = (reasoning_text or "") + (flush_reasoning or "") + delta = (delta or "") + (flush_text or "") if reasoning_text: if request.separate_reasoning: choice_data = ChatCompletionStreamResponseChoice( @@ -775,7 +791,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: if current_finish_reason is not None: # Flush any buffered reasoning content that was never released # (e.g., max_completion_tokens hit before was seen). - if get_env_start_args().reasoning_parser: + if get_env_start_args().reasoning_parser and get_env_start_args().reasoning_parser != "glm45": parser = reasoning_parser_dict.get(choice_index) if parser is not None: flush_reasoning, flush_text = parser.flush() diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index e5d7564081..19ad405c59 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -28,6 +28,7 @@ from pydantic import BaseModel, Field from .api_models import Tool +from .glm_xml_parser import GlmToolCallScanner, iter_glm_tool_call_spans, parse_glm_xml_value logger = logging.getLogger(__name__) ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"] @@ -1228,16 +1229,16 @@ def __init__(self): self.tool_call_separator = "\n" # Regex patterns for parsing GLM-4.7 tool calls - # Match complete tool call blocks - self.func_call_regex = re.compile(r".*?", re.DOTALL) # Extract function name and arguments from a tool call block # Function name can be followed by newline OR directly by # Pattern: function_name(\n|)... self.func_detail_regex = re.compile( - r"([^<\n]+?)(?:\n|(?=)|(?=))(.*?)", re.DOTALL + r"([^<\n]+?)(?:\n|(?=)|(?=))(.*?)\Z", re.DOTALL ) self.func_name_regex = re.compile(r"([^<\n]+?)(?:\n|(?=)|(?=))") self._streaming_tool_name: Optional[str] = None + self._tool_scanner = GlmToolCallScanner() + self._tool_scan_offset = len(self.bot_token) # Extract arg_key/arg_value pairs self.func_arg_regex = re.compile(r"(.*?)\s*(.*?)", re.DOTALL) @@ -1260,21 +1261,12 @@ def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict: args = {} properties = (tool.function.parameters or {}).get("properties", {}) + if not isinstance(properties, dict): + properties = {} matches = self.func_arg_regex.findall(arg_text) for key, value in matches: key = key.strip() - # XML string values are literal: file content and edit targets may - # contain meaningful whitespace or text that happens to be JSON. - if properties.get(key, {}).get("type") == "string": - args[key] = value - continue - # Try to parse value as JSON for complex types (arrays, objects, numbers, booleans) - try: - parsed_value = json.loads(value) - args[key] = parsed_value - except (json.JSONDecodeError, ValueError): - # Keep as string if not valid JSON - args[key] = value + args[key] = parse_glm_xml_value(value, properties.get(key)) return args def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: @@ -1293,11 +1285,13 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult tool_indices = self._get_tool_indices(tools) calls = [] + normal_parts = [] + cursor = 0 - # Find all ... blocks - match_result_list = self.func_call_regex.findall(text) - - for match_result in match_result_list: + for start, end in iter_glm_tool_call_spans(text): + normal_parts.append(text[cursor:start]) + cursor = end + match_result = text[start:end] try: # Extract function name and arguments func_detail = self.func_detail_regex.search(match_result) @@ -1327,7 +1321,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}, error: {str(e)}") continue - return StreamingParseResult(normal_text=normal_text, calls=calls) + # Keep ordinary text between and after calls. An unfinished envelope + # stays buffered in streaming mode and must not leak in non-streaming. + incomplete_start = text.find(self.bot_token, cursor) + normal_parts.append(text[cursor : incomplete_start if incomplete_start != -1 else len(text)]) + return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Buffer XML arguments while emitting tool deltas to keep the stream alive.""" @@ -1345,8 +1343,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami normal_text += self._buffer[:start] self._buffer = self._buffer[start:] - end = self._buffer.find(self.eot_token) - if end == -1: + consumed, complete = self._tool_scanner.scan(self._buffer[self._tool_scan_offset :]) + self._tool_scan_offset += consumed + if not complete: # Match the buffered Qwen3-Coder flow: announce the name once, # then send empty argument deltas until the XML call is complete. # Long Write/Bash arguments must not leave the HTTP stream idle. @@ -1361,7 +1360,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami calls.append(ToolCallItem(tool_index=self.current_tool_id, name=name, parameters="")) break - end += len(self.eot_token) + end = self._tool_scan_offset # Waiting for the closing tag avoids re-emitting partial JSON or # rewriting already sent arguments when another XML key arrives. result = self.detect_and_parse(self._buffer[:end], tools) @@ -1374,6 +1373,8 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami calls.append(call) self._streaming_tool_name = None self._buffer = self._buffer[end:] + self._tool_scanner = GlmToolCallScanner() + self._tool_scan_offset = len(self.bot_token) return StreamingParseResult(normal_text=normal_text, calls=calls) diff --git a/lightllm/server/glm_xml_parser.py b/lightllm/server/glm_xml_parser.py new file mode 100644 index 0000000000..fae339f3bb --- /dev/null +++ b/lightllm/server/glm_xml_parser.py @@ -0,0 +1,148 @@ +"""Shared boundaries and schema coercion for GLM's XML tool calls.""" + +import json +import math +from typing import Any, Iterator, Tuple + + +TOOL_START = "" +TOOL_END = "" +VALUE_START = "" +VALUE_END = "" +_TYPE_ALIASES = { + **dict.fromkeys(("str", "text", "varchar", "char", "enum"), "string"), + **dict.fromkeys(("int", "int32", "int64", "uint", "uint32", "uint64", "long", "short", "unsigned"), "integer"), + **dict.fromkeys(("float", "float32", "float64", "double"), "number"), + "bool": "boolean", + "dict": "object", + **dict.fromkeys(("arr", "list", "sequence"), "array"), +} + + +class GlmToolCallScanner: + """Scan after a tool opener, treating argument values as literal text. + + ``scan`` returns the consumed length and whether the outer call closed. + Callers retain the unconsumed suffix, which may be a partial delimiter, + and pass it with the next chunk. Only ```` closes a value; + thinking and tool-call tags inside that value have no structural meaning. + """ + + def __init__(self): + self.in_arg_value = False + + def scan(self, text: str) -> Tuple[int, bool]: + offset = 0 + while offset < len(text): + markers = (VALUE_END,) if self.in_arg_value else (VALUE_START, TOOL_END) + positions = [(text.find(marker, offset), marker) for marker in markers] + positions = [(position, marker) for position, marker in positions if position != -1] + if not positions: + max_hold = min(len(text) - offset, max(map(len, markers)) - 1) + for length in range(max_hold, 0, -1): + if any(marker.startswith(text[-length:]) for marker in markers): + return len(text) - length, False + return len(text), False + + position, marker = min(positions) + offset = position + len(marker) + if marker == TOOL_END: + return offset, True + self.in_arg_value = marker == VALUE_START + + return offset, False + + +def iter_glm_tool_call_spans(text: str) -> Iterator[Tuple[int, int]]: + start = text.find(TOOL_START) + while start != -1: + body_start = start + len(TOOL_START) + consumed, complete = GlmToolCallScanner().scan(text[body_start:]) + if not complete: + return + end = body_start + consumed + yield start, end + start = text.find(TOOL_START, end) + + +def _schema_types(schema: Any) -> set[str]: + if not isinstance(schema, dict): + return {"string"} + types = set() + declared = schema.get("type") + if isinstance(declared, str): + types.add(declared) + elif isinstance(declared, list): + types.update(item for item in declared if isinstance(item, str)) + for value in schema.get("enum", []) if isinstance(schema.get("enum"), list) else []: + if value is None: + types.add("null") + elif isinstance(value, bool): + types.add("boolean") + elif isinstance(value, int): + types.add("integer") + elif isinstance(value, float): + types.add("number") + elif isinstance(value, str): + types.add("string") + elif isinstance(value, list): + types.add("array") + elif isinstance(value, dict): + types.add("object") + for field in ("anyOf", "oneOf", "allOf"): + choices = schema.get(field) + if isinstance(choices, list): + for choice in choices: + types.update(_schema_types(choice)) + return types or {"string"} + + +def _finite_json(value: Any) -> bool: + if isinstance(value, float): + return math.isfinite(value) + if isinstance(value, dict): + return all(_finite_json(item) for item in value.values()) + if isinstance(value, list): + return all(_finite_json(item) for item in value) + return True + + +def parse_glm_xml_value(value: str, schema: Any) -> Any: + """Match vLLM's schema-based coercion without changing literal strings. + + Union types use null/integer/number/boolean/object/array/string priority. + Unspecified fields remain strings. Non-finite numbers remain raw text so + the resulting arguments are always valid JSON. + """ + types = {_TYPE_ALIASES.get(item.strip().lower(), item.strip().lower()) for item in _schema_types(schema)} + for candidate in ("null", "integer", "number", "boolean", "object", "array", "string"): + if candidate not in types: + continue + try: + if candidate == "null" and value.lower() == "null": + return None + if candidate == "integer": + return int(value) + if candidate == "number": + number = float(value) + if math.isfinite(number): + return number if number != int(number) else int(number) + if candidate == "boolean": + normalized = value.strip().lower() + if normalized in ("true", "1"): + return True + if normalized in ("false", "0"): + return False + if candidate in ("object", "array"): + parsed = json.loads(value) + if _finite_json(parsed): + return parsed + if candidate == "string": + return value + except (ValueError, TypeError): + continue + try: + parsed = json.loads(value) + return parsed if _finite_json(parsed) else value + except ValueError: + return value diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 251a399bd1..51502b0d2c 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -16,6 +16,7 @@ from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id +from lightllm.server.glm_xml_parser import GlmToolCallScanner @dataclass @@ -765,6 +766,20 @@ class Glm45Detector(Qwen3Detector): def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False): super().__init__(stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) self._in_tool_call = False + self._tool_scanner = GlmToolCallScanner() + + @classmethod + def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: + # GLM-5.3 always opens thinking, even if enable_thinking=False was + # requested. Inspect only the final assistant turn, never past turns. + turns = prompt.rsplit("<|assistant|>", 1) + if len(turns) != 2: + return None + tail = turns[-1] + if any(marker in tail for marker in ("<|system|>", "<|user|>", "<|observation|>")): + return None + start, end = tail.rfind(""), tail.rfind("") + return start > end if start != -1 or end != -1 else None def detect_and_parse(self, text: str) -> StreamingParseResult: detector = type(self)(stream_reasoning=self.stream_reasoning, force_reasoning=self._in_reasoning) @@ -783,14 +798,11 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: while self._buffer: if self._in_tool_call: - end = self._buffer.find(self.tool_end_token) - if end == -1: - text, self._buffer = prefix_hold(self._buffer, [self.tool_end_token]) - normal_parts.append(text) + consumed, complete = self._tool_scanner.scan(self._buffer) + normal_parts.append(self._buffer[:consumed]) + self._buffer = self._buffer[consumed:] + if not complete: break - end += len(self.tool_end_token) - normal_parts.append(self._buffer[:end]) - self._buffer = self._buffer[end:] self._in_tool_call = False continue @@ -808,6 +820,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: if marker == self.tool_start_token: self._in_reasoning = False self._in_tool_call = True + self._tool_scanner = GlmToolCallScanner() # Preserve the opener and literal thinking tags in tool arguments. normal_parts.append(marker) else: diff --git a/unit_tests/server/test_glm47_tool_parser.py b/unit_tests/server/test_glm47_tool_parser.py index 8f626d6343..eb57175ac6 100644 --- a/unit_tests/server/test_glm47_tool_parser.py +++ b/unit_tests/server/test_glm47_tool_parser.py @@ -84,6 +84,9 @@ def test_text_around_tool_calls_survives_chunk_boundaries(chunk_size): content, calls = collect_stream(text, chunk_size) assert content == "Before between after" assert list(calls) == [0, 1] + normal, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text) + assert normal == content + assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()] def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name(): @@ -95,7 +98,7 @@ def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name(): assert calls[0].name == "Write" assert calls[0].parameters == "" - value = ' line 1\n{"literal": true}\n' + "x" * 100_000 + value = ' line 1\n{"literal": true}\nRead\n' + "x" * 100_000 for chunk in (value[:13], value[13:27], value[27:]): normal, calls = parser.parse_stream_chunk(chunk) assert normal == "" @@ -125,3 +128,79 @@ def test_undefined_buffered_tool_does_not_consume_a_call_index(): assert len(calls) == 1 assert calls[0].name == "Read" assert calls[0].tool_index == 0 + + +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +def test_tool_delimiters_inside_values_do_not_close_or_create_calls(chunk_size): + value = "Read\nliteralliteral" + text = ( + f"Writecontent{value}" + "Readfile_path/tmp/a" + ) + normal, calls = collect_stream(text, chunk_size) + assert normal == "" + assert list(calls) == [0, 1] + assert [call["name"] for call in calls.values()] == ["Write", "Read"] + assert [json.loads(call["arguments"]) for call in calls.values()] == [{"content": value}, {"file_path": "/tmp/a"}] + normal, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text) + assert normal == "" + assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()] + + +@pytest.mark.parametrize( + "schema,value,expected", + [ + ({"type": ["string"]}, "123", "123"), + ({"anyOf": [{"type": "string"}]}, "false", "false"), + ({"oneOf": [{"type": "string"}]}, "null", "null"), + ({"allOf": [{"type": "string"}]}, "123", "123"), + ({"enum": ["123", "false"]}, "123", "123"), + ({"type": ["string", "null"]}, "null", None), + ({"anyOf": [{"type": "integer"}, {"type": "string"}]}, "0012", 12), + ({"type": "boolean"}, "False", False), + ({"type": "boolean"}, "1", True), + ({"type": "boolean"}, " 0 ", False), + ({"type": "integer"}, "0012", 12), + ({"type": "number"}, ".5", 0.5), + ({"type": "number"}, "1.0", 1), + ({"type": "array"}, '[1, "中", false]', [1, "中", False]), + ({"type": "object"}, '{"nested": {"ok": true}}', {"nested": {"ok": True}}), + ({}, "123", "123"), + (None, "false", "false"), + ({"type": "number"}, "NaN", "NaN"), + ({"type": "number"}, "Infinity", "Infinity"), + ({"type": "number"}, "1e999", "1e999"), + ({"type": "array"}, "[1e999]", "[1e999]"), + ({"type": "object"}, '{"x": NaN}', '{"x": NaN}'), + ], +) +def test_schema_coercion_preserves_declared_types_and_finite_json(schema, value, expected): + tools = [ + Tool.model_validate( + { + "type": "function", + "function": { + "name": "Probe", + "parameters": { + "type": "object", + "properties": {"value": schema} if schema is not None else {}, + }, + }, + } + ) + ] + text = f"Probevalue{value}" + parser = FunctionCallParser(tools, "glm47") + normal, calls = parser.parse_non_stream(text) + assert normal == "" + assert len(calls) == 1 + parsed = json.loads(calls[0].parameters)["value"] + assert parsed == expected + assert type(parsed) is type(expected) + json.dumps(parsed, allow_nan=False) + parameters = "" + for char in text: + normal, deltas = parser.parse_stream_chunk(char) + assert normal == "" + parameters += "".join(delta.parameters for delta in deltas) + assert parameters == calls[0].parameters diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py index 9040910bae..46352bbe52 100644 --- a/unit_tests/server/test_glm_reasoning_tool_calls.py +++ b/unit_tests/server/test_glm_reasoning_tool_calls.py @@ -25,7 +25,7 @@ } ) ] -VALUE = "line 1\nliteral\nline 2" +VALUE = ' 中文 "quoted" \\path\nliteralRead\n ' TOOL_CALL = f"Writecontent{VALUE}" @@ -115,40 +115,59 @@ def test_glm_tools_inside_reasoning_do_not_enable_string_stops(): @pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) @pytest.mark.parametrize("reasoning_end", ["missing", "before_tool", "after_tool"]) -def test_chat_completion_exposes_glm_tool_calls(monkeypatch, stream, reasoning_end): +@pytest.mark.parametrize("chunk_size", [1, 4096]) +def test_chat_completion_exposes_glm_tool_calls(monkeypatch, stream, separate_reasoning, reasoning_end, chunk_size): args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47") monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="")) + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="<|assistant|>")) segments = [("Need a file.", 1)] if reasoning_end == "before_tool": segments.append(("", 99)) - segments.extend([("", 77), (TOOL_CALL[len("") :], 2)]) + segments.extend([("", 77), (TOOL_CALL[len("") :], 2)] * 2) if reasoning_end == "after_tool": segments.append(("", 99)) + chunks = [ + (text[start : start + chunk_size], token_id if start + chunk_size >= len(text) else 1) + for text, token_id in segments + for start in range(0, len(text), chunk_size) + ] + async def generate(*args, **kwargs): - for index, (text, token_id) in enumerate(segments): - finish = FinishStatus(FinishStatus.FINISHED_STOP) if index == len(segments) - 1 else FinishStatus() + for index, (text, token_id) in enumerate(chunks): + finish = FinishStatus(FinishStatus.FINISHED_STOP) if index == len(chunks) - 1 else FinishStatus() yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) - request = ChatCompletionRequest(model="glm", messages=[], tools=TOOLS, stream=stream) + request = ChatCompletionRequest( + model="glm", + messages=[], + tools=TOOLS, + stream=stream, + separate_reasoning=separate_reasoning, + # GLM-5.3's always-on template ignores this toggle; its rendered + # generation prompt must determine the extraction state instead. + chat_template_kwargs={"enable_thinking": False}, + ) async def run(): response = await api_openai.chat_completions_impl(request, None) if not stream: choice = response.choices[0] + assert [call.function.name for call in choice.message.tool_calls] == ["Write", "Write"] + assert len({call.id for call in choice.message.tool_calls}) == 2 return ( choice.message.reasoning, choice.message.content, - choice.message.tool_calls[0].function.arguments, + [call.function.arguments for call in choice.message.tool_calls], choice.finish_reason, ) - reasoning, content, arguments, finish_reason = "", "", "", None + reasoning, content, calls, finish_reason = "", "", {}, None async for chunk in response.body_iterator: chunk = chunk.decode() if isinstance(chunk, bytes) else chunk if chunk.strip() == "data: [DONE]": @@ -159,14 +178,102 @@ async def run(): reasoning += delta.get("reasoning") or "" content += delta.get("content") or "" for call in delta.get("tool_calls") or []: - assert call["index"] == 0 - assert call["function"].get("name") in (None, "Write") - arguments += call["function"].get("arguments") or "" + index = call["index"] + if call["function"].get("name") is not None: + assert index not in calls + assert call["function"]["name"] == "Write" + assert call["id"] + calls[index] = {"id": call["id"], "arguments": ""} + assert index in calls + calls[index]["arguments"] += call["function"].get("arguments") or "" finish_reason = choice.get("finish_reason") or finish_reason - return reasoning, content, arguments, finish_reason + assert list(calls) == [0, 1] + assert len({call["id"] for call in calls.values()}) == 2 + return reasoning, content, [call["arguments"] for call in calls.values()], finish_reason reasoning, content, arguments, finish_reason = asyncio.run(run()) - assert reasoning == "Need a file." - assert content == "" - assert json.loads(arguments) == {"content": VALUE} + assert reasoning == ("Need a file." if separate_reasoning else "") + assert content == ("" if separate_reasoning else "Need a file.") + assert [json.loads(argument) for argument in arguments] == [{"content": VALUE}] * 2 assert finish_reason == "tool_calls" + + +@pytest.mark.parametrize( + "prompt,state", + [ + ("<|assistant|>", True), + ("<|assistant|>reasoning directive", True), + ("<|assistant|>", False), + ("<|assistant|>old<|user|>question<|assistant|>", True), + ("<|assistant|>old<|user|>question<|assistant|>", None), + ("<|assistant|>old<|observation|>tool result", None), + ("<|user|>quoted", None), + ], +) +def test_glm_initial_reasoning_state_comes_from_current_prompt_turn(prompt, state): + assert ReasoningParser.DetectorMap["glm45"].initial_reasoning_state(prompt) is state + + +def test_glm_prompt_override_does_not_change_other_models(monkeypatch): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="qwen3")) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) + request = ChatCompletionRequest(model="qwen", messages=[], chat_template_kwargs={"enable_thinking": False}) + assert api_openai._is_force_thinking_mode(request, "<|assistant|>") is False + + +@pytest.mark.parametrize("thinking", [False, True]) +def test_glm_grammar_activation_follows_rendered_prompt(monkeypatch, thinking): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="glm45")) + request = ChatCompletionRequest(model="glm", messages=[], chat_template_kwargs={"enable_thinking": not thinking}) + prompt = "<|assistant|>" + ("" if thinking else "") + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text]) + assert api_openai._is_force_thinking_mode(request, prompt) is thinking + assert api_openai._get_guided_reasoning_end(request, tokenizer, prompt) == ([99] if thinking else []) + + +@pytest.mark.parametrize("suffix", ["<", "Writecontentbody", { + "id": 2, + "prompt_tokens": 1, + }, FinishStatus() + yield 0, suffix, {"id": 3, "prompt_tokens": 1}, FinishStatus(FinishStatus.FINISHED_LENGTH) + + manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) + request = ChatCompletionRequest( + model="glm", messages=[], tools=TOOLS, stream=True, separate_reasoning=separate_reasoning + ) + + async def run(): + response = await api_openai.chat_completions_impl(request, None) + reasoning, content, names, arguments, finishes = [], [], [], [], [] + async for chunk in response.body_iterator: + chunk = chunk.decode() if isinstance(chunk, bytes) else chunk + if chunk.strip() == "data: [DONE]": + continue + for choice in json.loads(chunk.removeprefix("data: "))["choices"]: + delta = choice["delta"] + reasoning.append(delta.get("reasoning") or "") + content.append(delta.get("content") or "") + if choice.get("finish_reason"): + finishes.append(choice["finish_reason"]) + for call in delta.get("tool_calls") or []: + assert call["index"] == 0 + if call["function"].get("name"): + names.append(call["function"]["name"]) + arguments.append(call["function"].get("arguments") or "") + assert "".join(reasoning) == ("Need a file." if separate_reasoning else "") + assert "".join(content) == ("" if separate_reasoning else "Need a file.") + assert names == ["Write"] + assert "".join(arguments) == "" + assert finishes == ["length"] + + asyncio.run(run()) From 637c0796167cb191a15199c46cb0f75c3f966a92 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 21 Oct 2026 17:15:43 +0000 Subject: [PATCH 03/13] refactor: keep GLM tool parsing within existing detectors --- lightllm/server/api_openai.py | 24 +-- lightllm/server/function_call_parser.py | 156 +++++++++++++- lightllm/server/glm_xml_parser.py | 148 ------------- lightllm/server/reasoning_parser.py | 3 +- unit_tests/server/httpserver/test_grammar.py | 2 +- .../server/test_glm_reasoning_tool_calls.py | 196 ++++++++---------- 6 files changed, 249 insertions(+), 280 deletions(-) delete mode 100644 lightllm/server/glm_xml_parser.py diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 29adeffefc..91a357b25b 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -205,9 +205,13 @@ def _is_force_thinking_mode(request: ChatCompletionRequest, prompt: Optional[str return True # default -def _get_guided_reasoning_end(request: ChatCompletionRequest, tokenizer, prompt: Optional[str] = None) -> List[int]: +def _get_guided_reasoning_end( + request: ChatCompletionRequest, tokenizer, force_reasoning: Optional[bool] = None +) -> List[int]: """Use the same thinking mode as the chat template for grammar activation.""" - if not _is_force_thinking_mode(request, prompt): + if force_reasoning is None: + force_reasoning = _is_force_thinking_mode(request) + if not force_reasoning: return [] parser_name = get_env_start_args().reasoning_parser detector = ReasoningParser.DetectorMap[parser_name.lower()]() @@ -221,14 +225,14 @@ def _process_reasoning_stream( reasoning_parser_dict: Dict[int, ReasoningParser], metadata: Dict[str, Any], request: ChatCompletionRequest, - force_reasoning: Optional[bool] = None, + force_reasoning: bool, ) -> tuple[Optional[str], str]: """Process reasoning content and update its token usage.""" if index not in reasoning_parser_dict: reasoning_parser_dict[index] = ReasoningParser( get_env_start_args().reasoning_parser, request.stream_reasoning, - _is_force_thinking_mode(request) if force_reasoning is None else force_reasoning, + force_reasoning, ) parser = reasoning_parser_dict[index] token_id = metadata.get("id") @@ -411,7 +415,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req if sampling_params_dict.get("guided_json") or sampling_params_dict.get("guided_grammar"): sampling_params_dict["guided_reasoning_end"] = _get_guided_reasoning_end( - request, g_objs.httpserver_manager.tokenizer, prompt + request, g_objs.httpserver_manager.tokenizer, request_enable_reasoning ) sampling_params = SamplingParams() @@ -602,14 +606,8 @@ async def stream_results() -> AsyncGenerator[bytes, None]: reasoning_parser_dict, metadata, request, - request_enable_reasoning if get_env_start_args().reasoning_parser == "glm45" else None, + request_enable_reasoning, ) - if current_finish_reason is not None and get_env_start_args().reasoning_parser == "glm45": - # Route a buffered GLM suffix through the tool parser too; - # truncating inside an XML argument must not leak content. - flush_reasoning, flush_text = reasoning_parser_dict[choice_index].flush() - reasoning_text = (reasoning_text or "") + (flush_reasoning or "") - delta = (delta or "") + (flush_text or "") if reasoning_text: if request.separate_reasoning: choice_data = ChatCompletionStreamResponseChoice( @@ -791,7 +789,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: if current_finish_reason is not None: # Flush any buffered reasoning content that was never released # (e.g., max_completion_tokens hit before was seen). - if get_env_start_args().reasoning_parser and get_env_start_args().reasoning_parser != "glm45": + if get_env_start_args().reasoning_parser: parser = reasoning_parser_dict.get(choice_index) if parser is not None: flush_reasoning, flush_text = parser.flush() diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index 19ad405c59..24df428a46 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -13,6 +13,7 @@ import ast import json +import math import os import orjson import logging @@ -20,7 +21,7 @@ from abc import ABC, abstractmethod from json import JSONDecodeError, JSONDecoder from json.decoder import WHITESPACE -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any, Dict, Iterator, List, Optional, Tuple, Type import partial_json_parser from partial_json_parser.core.exceptions import MalformedJSON @@ -28,7 +29,6 @@ from pydantic import BaseModel, Field from .api_models import Tool -from .glm_xml_parser import GlmToolCallScanner, iter_glm_tool_call_spans, parse_glm_xml_value logger = logging.getLogger(__name__) ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"] @@ -1187,6 +1187,44 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult(normal_text=current_text) +class GlmToolCallScanner: + """Track a GLM tool envelope without interpreting tags inside argument values. + + The caller forwards all input until the outer closing tag. Only a possible + delimiter prefix is retained here for scanning the next chunk; those bytes + have already been forwarded and are excluded from the consumed count. + """ + + def __init__(self): + self._in_arg_value = False + self._pending = "" + + def scan(self, new_text: str) -> Tuple[int, bool]: + pending_length = len(self._pending) + text = self._pending + new_text + self._pending = "" + offset = 0 + while offset < len(text): + markers = ("",) if self._in_arg_value else ("", "") + positions = [(text.find(marker, offset), marker) for marker in markers] + positions = [(position, marker) for position, marker in positions if position != -1] + if not positions: + max_hold = min(len(text) - offset, max(map(len, markers)) - 1) + for length in range(max_hold, 0, -1): + if any(marker.startswith(text[-length:]) for marker in markers): + self._pending = text[-length:] + break + return len(new_text), False + + position, marker = min(positions) + offset = position + len(marker) + if marker == "": + return offset - pending_length, True + self._in_arg_value = marker == "" + + return len(new_text), False + + class Glm47Detector(BaseFormatDetector): """ Detector for GLM-4.7/GLM-4.7-Flash model function call format. @@ -1222,6 +1260,15 @@ class Glm47Detector(BaseFormatDetector): Reference: https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_parser.py """ + _TYPE_ALIASES = { + **dict.fromkeys(("str", "text", "varchar", "char", "enum"), "string"), + **dict.fromkeys(("int", "int32", "int64", "uint", "uint32", "uint64", "long", "short", "unsigned"), "integer"), + **dict.fromkeys(("float", "float32", "float64", "double"), "number"), + "bool": "boolean", + "dict": "object", + **dict.fromkeys(("arr", "list", "sequence"), "array"), + } + def __init__(self): super().__init__() self.bot_token = "" @@ -1246,6 +1293,91 @@ def has_tool_call(self, text: str) -> bool: """Check if the text contains a GLM-4.7 format tool call.""" return self.bot_token in text + @classmethod + def _schema_types(cls, schema: Any) -> set[str]: + if not isinstance(schema, dict): + return {"string"} + types = set() + declared = schema.get("type") + if isinstance(declared, str): + types.add(declared) + elif isinstance(declared, list): + types.update(item for item in declared if isinstance(item, str)) + for value in schema.get("enum", []) if isinstance(schema.get("enum"), list) else []: + if value is None: + types.add("null") + elif isinstance(value, bool): + types.add("boolean") + elif isinstance(value, int): + types.add("integer") + elif isinstance(value, float): + types.add("number") + elif isinstance(value, str): + types.add("string") + elif isinstance(value, list): + types.add("array") + elif isinstance(value, dict): + types.add("object") + for field in ("anyOf", "oneOf", "allOf"): + choices = schema.get(field) + if isinstance(choices, list): + for choice in choices: + types.update(cls._schema_types(choice)) + return types or {"string"} + + @classmethod + def _finite_json(cls, value: Any) -> bool: + if isinstance(value, float): + return math.isfinite(value) + if isinstance(value, dict): + return all(cls._finite_json(item) for item in value.values()) + if isinstance(value, list): + return all(cls._finite_json(item) for item in value) + return True + + @classmethod + def _convert_param_value(cls, value: str, schema: Any) -> Any: + """Match vLLM's schema-based coercion without changing literal strings. + + Union types use null/integer/number/boolean/object/array/string priority. + Unspecified fields remain strings. Non-finite numbers remain raw text so + the resulting arguments are always valid JSON. + """ + types = { + cls._TYPE_ALIASES.get(item.strip().lower(), item.strip().lower()) for item in cls._schema_types(schema) + } + for candidate in ("null", "integer", "number", "boolean", "object", "array", "string"): + if candidate not in types: + continue + try: + if candidate == "null" and value.lower() == "null": + return None + if candidate == "integer": + return int(value) + if candidate == "number": + number = float(value) + if math.isfinite(number): + return number if number != int(number) else int(number) + if candidate == "boolean": + normalized = value.strip().lower() + if normalized in ("true", "1"): + return True + if normalized in ("false", "0"): + return False + if candidate in ("object", "array"): + parsed = json.loads(value) + if cls._finite_json(parsed): + return parsed + if candidate == "string": + return value + except (ValueError, TypeError): + continue + try: + parsed = json.loads(value) + return parsed if cls._finite_json(parsed) else value + except ValueError: + return value + def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict: """ Parse XML-style arguments into a dictionary. @@ -1266,9 +1398,20 @@ def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict: matches = self.func_arg_regex.findall(arg_text) for key, value in matches: key = key.strip() - args[key] = parse_glm_xml_value(value, properties.get(key)) + args[key] = self._convert_param_value(value, properties.get(key)) return args + def _iter_tool_call_spans(self, text: str) -> Iterator[Tuple[int, int]]: + start = text.find(self.bot_token) + while start != -1: + body_start = start + len(self.bot_token) + consumed, complete = GlmToolCallScanner().scan(text[body_start:]) + if not complete: + return + end = body_start + consumed + yield start, end + start = text.find(self.bot_token, end) + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: """ One-time parsing: Detects and parses tool calls in the provided text. @@ -1277,18 +1420,15 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :param tools: List of available tools. :return: StreamingParseResult with normal_text and parsed calls. """ - idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text - if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=text, calls=[]) tool_indices = self._get_tool_indices(tools) calls = [] normal_parts = [] cursor = 0 - for start, end in iter_glm_tool_call_spans(text): + for start, end in self._iter_tool_call_spans(text): normal_parts.append(text[cursor:start]) cursor = end match_result = text[start:end] diff --git a/lightllm/server/glm_xml_parser.py b/lightllm/server/glm_xml_parser.py deleted file mode 100644 index fae339f3bb..0000000000 --- a/lightllm/server/glm_xml_parser.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Shared boundaries and schema coercion for GLM's XML tool calls.""" - -import json -import math -from typing import Any, Iterator, Tuple - - -TOOL_START = "" -TOOL_END = "" -VALUE_START = "" -VALUE_END = "" -_TYPE_ALIASES = { - **dict.fromkeys(("str", "text", "varchar", "char", "enum"), "string"), - **dict.fromkeys(("int", "int32", "int64", "uint", "uint32", "uint64", "long", "short", "unsigned"), "integer"), - **dict.fromkeys(("float", "float32", "float64", "double"), "number"), - "bool": "boolean", - "dict": "object", - **dict.fromkeys(("arr", "list", "sequence"), "array"), -} - - -class GlmToolCallScanner: - """Scan after a tool opener, treating argument values as literal text. - - ``scan`` returns the consumed length and whether the outer call closed. - Callers retain the unconsumed suffix, which may be a partial delimiter, - and pass it with the next chunk. Only ```` closes a value; - thinking and tool-call tags inside that value have no structural meaning. - """ - - def __init__(self): - self.in_arg_value = False - - def scan(self, text: str) -> Tuple[int, bool]: - offset = 0 - while offset < len(text): - markers = (VALUE_END,) if self.in_arg_value else (VALUE_START, TOOL_END) - positions = [(text.find(marker, offset), marker) for marker in markers] - positions = [(position, marker) for position, marker in positions if position != -1] - if not positions: - max_hold = min(len(text) - offset, max(map(len, markers)) - 1) - for length in range(max_hold, 0, -1): - if any(marker.startswith(text[-length:]) for marker in markers): - return len(text) - length, False - return len(text), False - - position, marker = min(positions) - offset = position + len(marker) - if marker == TOOL_END: - return offset, True - self.in_arg_value = marker == VALUE_START - - return offset, False - - -def iter_glm_tool_call_spans(text: str) -> Iterator[Tuple[int, int]]: - start = text.find(TOOL_START) - while start != -1: - body_start = start + len(TOOL_START) - consumed, complete = GlmToolCallScanner().scan(text[body_start:]) - if not complete: - return - end = body_start + consumed - yield start, end - start = text.find(TOOL_START, end) - - -def _schema_types(schema: Any) -> set[str]: - if not isinstance(schema, dict): - return {"string"} - types = set() - declared = schema.get("type") - if isinstance(declared, str): - types.add(declared) - elif isinstance(declared, list): - types.update(item for item in declared if isinstance(item, str)) - for value in schema.get("enum", []) if isinstance(schema.get("enum"), list) else []: - if value is None: - types.add("null") - elif isinstance(value, bool): - types.add("boolean") - elif isinstance(value, int): - types.add("integer") - elif isinstance(value, float): - types.add("number") - elif isinstance(value, str): - types.add("string") - elif isinstance(value, list): - types.add("array") - elif isinstance(value, dict): - types.add("object") - for field in ("anyOf", "oneOf", "allOf"): - choices = schema.get(field) - if isinstance(choices, list): - for choice in choices: - types.update(_schema_types(choice)) - return types or {"string"} - - -def _finite_json(value: Any) -> bool: - if isinstance(value, float): - return math.isfinite(value) - if isinstance(value, dict): - return all(_finite_json(item) for item in value.values()) - if isinstance(value, list): - return all(_finite_json(item) for item in value) - return True - - -def parse_glm_xml_value(value: str, schema: Any) -> Any: - """Match vLLM's schema-based coercion without changing literal strings. - - Union types use null/integer/number/boolean/object/array/string priority. - Unspecified fields remain strings. Non-finite numbers remain raw text so - the resulting arguments are always valid JSON. - """ - types = {_TYPE_ALIASES.get(item.strip().lower(), item.strip().lower()) for item in _schema_types(schema)} - for candidate in ("null", "integer", "number", "boolean", "object", "array", "string"): - if candidate not in types: - continue - try: - if candidate == "null" and value.lower() == "null": - return None - if candidate == "integer": - return int(value) - if candidate == "number": - number = float(value) - if math.isfinite(number): - return number if number != int(number) else int(number) - if candidate == "boolean": - normalized = value.strip().lower() - if normalized in ("true", "1"): - return True - if normalized in ("false", "0"): - return False - if candidate in ("object", "array"): - parsed = json.loads(value) - if _finite_json(parsed): - return parsed - if candidate == "string": - return value - except (ValueError, TypeError): - continue - try: - parsed = json.loads(value) - return parsed if _finite_json(parsed) else value - except ValueError: - return value diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 51502b0d2c..4076dfdc13 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -16,7 +16,7 @@ from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id -from lightllm.server.glm_xml_parser import GlmToolCallScanner +from lightllm.server.function_call_parser import GlmToolCallScanner @dataclass @@ -761,7 +761,6 @@ class Glm45Detector(Qwen3Detector): """GLM tool calls implicitly end thinking, even without ````.""" tool_start_token = "" - tool_end_token = "" def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False): super().__init__(stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) diff --git a/unit_tests/server/httpserver/test_grammar.py b/unit_tests/server/httpserver/test_grammar.py index 916c0f665c..bb29ca9b4a 100644 --- a/unit_tests/server/httpserver/test_grammar.py +++ b/unit_tests/server/httpserver/test_grammar.py @@ -44,7 +44,7 @@ def http_client(compiler, monkeypatch): monkeypatch.setattr(api_http.g_objs, "g_generate_stream_func", lightllm_generate_stream) monkeypatch.setattr(api_http.g_objs, "metric_client", SimpleNamespace(counter_inc=lambda *args: None)) monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="test")) - monkeypatch.setattr(api_openai, "_is_force_thinking_mode", lambda request: False) + monkeypatch.setattr(api_openai, "_is_force_thinking_mode", lambda request, prompt=None: False) yield httpx.AsyncClient(transport=httpx.ASGITransport(app=api_http.app), base_url="http://test") generate.assert_not_called() diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py index 46352bbe52..c9239ccc50 100644 --- a/unit_tests/server/test_glm_reasoning_tool_calls.py +++ b/unit_tests/server/test_glm_reasoning_tool_calls.py @@ -36,6 +36,70 @@ def mock_token_ids(monkeypatch): ) +@pytest.fixture +def complete_chat(monkeypatch): + args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47") + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="<|assistant|>")) + + def complete(chunks, finish_status=FinishStatus.FINISHED_STOP, **request_kwargs): + async def generate(*args, **kwargs): + for index, (text, token_id) in enumerate(chunks): + finish = FinishStatus(finish_status) if index == len(chunks) - 1 else FinishStatus() + yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish + + manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) + request = ChatCompletionRequest(model="glm", messages=[], tools=TOOLS, **request_kwargs) + + async def run(): + response = await api_openai.chat_completions_impl(request, None) + result = SimpleNamespace(reasoning="", content="", calls=[], finish_reason=None) + if not request.stream: + choice = response.choices[0] + result.reasoning = choice.message.reasoning or "" + result.content = choice.message.content or "" + result.calls = [ + {"id": call.id, "name": call.function.name, "arguments": call.function.arguments} + for call in choice.message.tool_calls or [] + ] + result.finish_reason = choice.finish_reason + return result + + calls = {} + done = False + async for chunk in response.body_iterator: + chunk = chunk.decode() if isinstance(chunk, bytes) else chunk + if chunk.strip() == "data: [DONE]": + done = True + continue + for choice in json.loads(chunk.removeprefix("data: "))["choices"]: + delta = choice["delta"] + result.reasoning += delta.get("reasoning") or "" + result.content += delta.get("content") or "" + for call in delta.get("tool_calls") or []: + index = call["index"] + name = call["function"].get("name") + if name is not None: + assert index not in calls # Send each name and ID once. + assert call["id"] + calls[index] = {"id": call["id"], "name": name, "arguments": ""} + assert index in calls + if call.get("id") is not None: + assert call["id"] == calls[index]["id"] + calls[index]["arguments"] += call["function"].get("arguments") or "" + result.finish_reason = choice.get("finish_reason") or result.finish_reason + assert done + assert list(calls) == list(range(len(calls))) + result.calls = list(calls.values()) + return result + + return asyncio.run(run()) + + return complete + + @pytest.mark.parametrize( "prefix,force_reasoning", [ @@ -118,84 +182,31 @@ def test_glm_tools_inside_reasoning_do_not_enable_string_stops(): @pytest.mark.parametrize("separate_reasoning", [False, True]) @pytest.mark.parametrize("reasoning_end", ["missing", "before_tool", "after_tool"]) @pytest.mark.parametrize("chunk_size", [1, 4096]) -def test_chat_completion_exposes_glm_tool_calls(monkeypatch, stream, separate_reasoning, reasoning_end, chunk_size): - args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47") - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) - monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="<|assistant|>")) +def test_chat_completion_exposes_glm_tool_calls(complete_chat, stream, separate_reasoning, reasoning_end, chunk_size): segments = [("Need a file.", 1)] if reasoning_end == "before_tool": segments.append(("", 99)) segments.extend([("", 77), (TOOL_CALL[len("") :], 2)] * 2) if reasoning_end == "after_tool": segments.append(("", 99)) - chunks = [ (text[start : start + chunk_size], token_id if start + chunk_size >= len(text) else 1) for text, token_id in segments for start in range(0, len(text), chunk_size) ] - - async def generate(*args, **kwargs): - for index, (text, token_id) in enumerate(chunks): - finish = FinishStatus(FinishStatus.FINISHED_STOP) if index == len(chunks) - 1 else FinishStatus() - yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish - - manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) - request = ChatCompletionRequest( - model="glm", - messages=[], - tools=TOOLS, + result = complete_chat( + chunks, stream=stream, separate_reasoning=separate_reasoning, - # GLM-5.3's always-on template ignores this toggle; its rendered - # generation prompt must determine the extraction state instead. + # GLM-5.3's template opens thinking even with this toggle disabled. chat_template_kwargs={"enable_thinking": False}, ) - - async def run(): - response = await api_openai.chat_completions_impl(request, None) - if not stream: - choice = response.choices[0] - assert [call.function.name for call in choice.message.tool_calls] == ["Write", "Write"] - assert len({call.id for call in choice.message.tool_calls}) == 2 - return ( - choice.message.reasoning, - choice.message.content, - [call.function.arguments for call in choice.message.tool_calls], - choice.finish_reason, - ) - - reasoning, content, calls, finish_reason = "", "", {}, None - async for chunk in response.body_iterator: - chunk = chunk.decode() if isinstance(chunk, bytes) else chunk - if chunk.strip() == "data: [DONE]": - continue - payload = json.loads(chunk.removeprefix("data: ")) - for choice in payload["choices"]: - delta = choice["delta"] - reasoning += delta.get("reasoning") or "" - content += delta.get("content") or "" - for call in delta.get("tool_calls") or []: - index = call["index"] - if call["function"].get("name") is not None: - assert index not in calls - assert call["function"]["name"] == "Write" - assert call["id"] - calls[index] = {"id": call["id"], "arguments": ""} - assert index in calls - calls[index]["arguments"] += call["function"].get("arguments") or "" - finish_reason = choice.get("finish_reason") or finish_reason - assert list(calls) == [0, 1] - assert len({call["id"] for call in calls.values()}) == 2 - return reasoning, content, [call["arguments"] for call in calls.values()], finish_reason - - reasoning, content, arguments, finish_reason = asyncio.run(run()) - assert reasoning == ("Need a file." if separate_reasoning else "") - assert content == ("" if separate_reasoning else "Need a file.") - assert [json.loads(argument) for argument in arguments] == [{"content": VALUE}] * 2 - assert finish_reason == "tool_calls" + assert result.reasoning == ("Need a file." if separate_reasoning else "") + assert result.content == ("" if separate_reasoning else "Need a file.") + assert [call["name"] for call in result.calls] == ["Write", "Write"] + assert len({call["id"] for call in result.calls}) == 2 + assert [json.loads(call["arguments"]) for call in result.calls] == [{"content": VALUE}] * 2 + assert result.finish_reason == "tool_calls" @pytest.mark.parametrize( @@ -227,53 +238,22 @@ def test_glm_grammar_activation_follows_rendered_prompt(monkeypatch, thinking): request = ChatCompletionRequest(model="glm", messages=[], chat_template_kwargs={"enable_thinking": not thinking}) prompt = "<|assistant|>" + ("" if thinking else "") tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text]) - assert api_openai._is_force_thinking_mode(request, prompt) is thinking - assert api_openai._get_guided_reasoning_end(request, tokenizer, prompt) == ([99] if thinking else []) + force_reasoning = api_openai._is_force_thinking_mode(request, prompt) + assert force_reasoning is thinking + assert api_openai._get_guided_reasoning_end(request, tokenizer, force_reasoning) == ([99] if thinking else []) @pytest.mark.parametrize("suffix", ["<", "Writecontentbody", { - "id": 2, - "prompt_tokens": 1, - }, FinishStatus() - yield 0, suffix, {"id": 3, "prompt_tokens": 1}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) - request = ChatCompletionRequest( - model="glm", messages=[], tools=TOOLS, stream=True, separate_reasoning=separate_reasoning +def test_truncated_tool_argument_suffix_does_not_leak_as_content(complete_chat, suffix, separate_reasoning): + result = complete_chat( + [("Need a file.", 1), ("Writecontentbody", 2), (suffix, 3)], + finish_status=FinishStatus.FINISHED_LENGTH, + stream=True, + separate_reasoning=separate_reasoning, ) - - async def run(): - response = await api_openai.chat_completions_impl(request, None) - reasoning, content, names, arguments, finishes = [], [], [], [], [] - async for chunk in response.body_iterator: - chunk = chunk.decode() if isinstance(chunk, bytes) else chunk - if chunk.strip() == "data: [DONE]": - continue - for choice in json.loads(chunk.removeprefix("data: "))["choices"]: - delta = choice["delta"] - reasoning.append(delta.get("reasoning") or "") - content.append(delta.get("content") or "") - if choice.get("finish_reason"): - finishes.append(choice["finish_reason"]) - for call in delta.get("tool_calls") or []: - assert call["index"] == 0 - if call["function"].get("name"): - names.append(call["function"]["name"]) - arguments.append(call["function"].get("arguments") or "") - assert "".join(reasoning) == ("Need a file." if separate_reasoning else "") - assert "".join(content) == ("" if separate_reasoning else "Need a file.") - assert names == ["Write"] - assert "".join(arguments) == "" - assert finishes == ["length"] - - asyncio.run(run()) + assert result.reasoning == ("Need a file." if separate_reasoning else "") + assert result.content == ("" if separate_reasoning else "Need a file.") + assert [call["name"] for call in result.calls] == ["Write"] + assert result.calls[0]["arguments"] == "" + assert result.finish_reason == "length" From 6c902a09ce7d04c1319bfe9d8bee30a745a51e04 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:59:47 +0000 Subject: [PATCH 04/13] refactor: separate thinking policy from reasoning state --- lightllm/server/api_openai.py | 83 +++++------------- lightllm/server/reasoning_parser.py | 43 +++++++--- unit_tests/server/httpserver/test_grammar.py | 2 +- .../server/test_glm_reasoning_tool_calls.py | 85 +++++++++++++++++-- 4 files changed, 132 insertions(+), 81 deletions(-) diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 91a357b25b..6c9224b875 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -180,17 +180,13 @@ def _get_history_tool_calls_cnt(request: ChatCompletionRequest) -> int: return idx -def _is_force_thinking_mode(request: ChatCompletionRequest, prompt: Optional[str] = None) -> bool: - """Whether this request uses forced thinking / reasoning (parser + template).""" +def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool: + """Resolve the request's thinking switch for rendering the chat template.""" from .build_prompt import tokenizer_supports_force_thinking reasoning_parser = get_env_start_args().reasoning_parser if not reasoning_parser: return False - if reasoning_parser == "glm45" and prompt is not None: - state = ReasoningParser.DetectorMap["glm45"].initial_reasoning_state(prompt) - if state is not None: - return state # Always-thinking models may have no thinking switch in their chat template. # Their answer grammar must wait for the reasoning delimiter nonetheless. if reasoning_parser in ["deepseek-r1", "qwen3-thinking", "gpt-oss", "minimax"]: @@ -206,12 +202,12 @@ def _is_force_thinking_mode(request: ChatCompletionRequest, prompt: Optional[str def _get_guided_reasoning_end( - request: ChatCompletionRequest, tokenizer, force_reasoning: Optional[bool] = None + request: ChatCompletionRequest, tokenizer, initial_reasoning_state: Optional[bool] = None ) -> List[int]: - """Use the same thinking mode as the chat template for grammar activation.""" - if force_reasoning is None: - force_reasoning = _is_force_thinking_mode(request) - if not force_reasoning: + """Delay answer constraints while generation starts inside reasoning.""" + if initial_reasoning_state is None: + initial_reasoning_state = _is_force_thinking_mode(request) + if not initial_reasoning_state: return [] parser_name = get_env_start_args().reasoning_parser detector = ReasoningParser.DetectorMap[parser_name.lower()]() @@ -362,7 +358,10 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req multimodal_params = _build_multimodal_params(request) tools = _select_chat_template_tools(request) prompt = await build_prompt(request, tools) - request_enable_reasoning = _is_force_thinking_mode(request, prompt) + reasoning_parser = get_env_start_args().reasoning_parser + initial_reasoning_state = ReasoningParser.get_initial_reasoning_state( + reasoning_parser, prompt, default=_is_force_thinking_mode(request) + ) sampling_params_dict = { "do_sample": request.do_sample, "presence_penalty": request.presence_penalty, @@ -395,9 +394,9 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req sampling_params_dict["max_new_tokens"] = request.max_tokens if request.stop is not None: sampling_params_dict["stop_sequences"] = request.stop - if get_env_start_args().reasoning_parser: + if reasoning_parser: # The prompt may already open reasoning; generated tokens may omit the opening marker. - sampling_params_dict["_initial_reasoning_state"] = int(request_enable_reasoning) + sampling_params_dict["_initial_reasoning_state"] = int(initial_reasoning_state) # Chat 接口提取 reasoning，按请求设置 stop 范围；通用 generate 接口不区分阶段。 sampling_params_dict["stop_str_match_mode"] = ( StopStrMatchMode.ALL_TEXT if is_stop_matching_in_reasoning_enabled() else StopStrMatchMode.CONTENT_ONLY @@ -415,7 +414,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req if sampling_params_dict.get("guided_json") or sampling_params_dict.get("guided_grammar"): sampling_params_dict["guided_reasoning_end"] = _get_guided_reasoning_end( - request, g_objs.httpserver_manager.tokenizer, request_enable_reasoning + request, g_objs.httpserver_manager.tokenizer, initial_reasoning_state ) sampling_params = SamplingParams() @@ -428,7 +427,6 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req # Non-streaming case if not request.stream: - reasoning_parser = get_env_start_args().reasoning_parser final_output_dict = collections.defaultdict(list) count_output_tokens_dict = collections.defaultdict(lambda: 0) finish_reason_dict = {} @@ -449,7 +447,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req parser = ReasoningParser( reasoning_parser, stream_reasoning=False, - force_reasoning=request_enable_reasoning, + force_reasoning=initial_reasoning_state, ) reasoning_parser_dict[sub_req_id] = parser token_id = metadata.get("id") @@ -599,15 +597,21 @@ async def stream_results() -> AsyncGenerator[bytes, None]: yield f"data: {_serialize_sse_chunk(first_chunk, _first_choice_nulls, _first_resp_nulls)}\n\n" # Handle reasoning content - if get_env_start_args().reasoning_parser: + if reasoning_parser: reasoning_text, delta = _process_reasoning_stream( choice_index, delta, reasoning_parser_dict, metadata, request, - request_enable_reasoning, + initial_reasoning_state, ) + if current_finish_reason is not None: + # Flush before tool parsing so buffered argument text follows + # the same path as every other chunk, including on truncation. + flush_reasoning, flush_text = reasoning_parser_dict[choice_index].flush() + reasoning_text = (reasoning_text or "") + (flush_reasoning or "") + delta = (delta or "") + (flush_text or "") if reasoning_text: if request.separate_reasoning: choice_data = ChatCompletionStreamResponseChoice( @@ -787,46 +791,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]: # Emit a per-choice final chunk with finish_reason (for tool_calls path # or when no delta was emitted alongside finish_reason). if current_finish_reason is not None: - # Flush any buffered reasoning content that was never released - # (e.g., max_completion_tokens hit before was seen). - if get_env_start_args().reasoning_parser: - parser = reasoning_parser_dict.get(choice_index) - if parser is not None: - flush_reasoning, flush_text = parser.flush() - if flush_reasoning: - if request.separate_reasoning: - flush_choice = ChatCompletionStreamResponseChoice( - index=choice_index, - delta=DeltaMessage(reasoning=flush_reasoning), - finish_reason=None, - ) - else: - # vLLM compat: emit buffered thinking as content - flush_choice = ChatCompletionStreamResponseChoice( - index=choice_index, - delta=DeltaMessage(content=flush_reasoning), - finish_reason=None, - ) - flush_chunk = ChatCompletionStreamResponse( - id=chat_completion_id, - created=created_time, - model=request.model, - choices=[flush_choice], - ) - yield f"data: {_serialize_sse_chunk(flush_chunk, _choice_nulls)}\n\n" - if flush_text: - flush_choice = ChatCompletionStreamResponseChoice( - index=choice_index, - delta=DeltaMessage(content=flush_text), - finish_reason=None, - ) - flush_chunk = ChatCompletionStreamResponse( - id=chat_completion_id, - created=created_time, - model=request.model, - choices=[flush_choice], - ) - yield f"data: {_serialize_sse_chunk(flush_chunk, _choice_nulls)}\n\n" if has_emitted_tool_calls[sub_req_id] and current_finish_reason == "stop": current_finish_reason = "tool_calls" final_choice = ChatCompletionStreamResponseChoice( @@ -842,7 +806,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) yield f"data: {_serialize_sse_chunk(final_chunk, _final_choice_nulls)}\n\n" - reasoning_parser = get_env_start_args().reasoning_parser completion_tokens_details = None if reasoning_parser: reasoning_tokens = sum(parser.reasoning_tokens for parser in reasoning_parser_dict.values()) diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 4076dfdc13..13a7f357f4 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -586,6 +586,11 @@ def __init__( class BaseReasoningFormatDetector: """Base class providing two sets of interfaces: one-time and streaming incremental.""" + @classmethod + def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: + """Infer the generation phase from the prompt, or defer to the request.""" + return None + def __init__( self, think_start_token: str, @@ -757,13 +762,13 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False) ) -class Glm45Detector(Qwen3Detector): - """GLM tool calls implicitly end thinking, even without ````.""" +class Glm45Detector(BaseReasoningFormatDetector): + """Route GLM tool envelopes to content, even before an explicit ````.""" tool_start_token = "" def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False): - super().__init__(stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) + super().__init__("", "", stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) self._in_tool_call = False self._tool_scanner = GlmToolCallScanner() @@ -771,14 +776,16 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False) def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: # GLM-5.3 always opens thinking, even if enable_thinking=False was # requested. Inspect only the final assistant turn, never past turns. - turns = prompt.rsplit("<|assistant|>", 1) - if len(turns) != 2: + _, assistant_marker, assistant_prompt = prompt.rpartition("<|assistant|>") + if not assistant_marker: + return None + if any(marker in assistant_prompt for marker in ("<|system|>", "<|user|>", "<|observation|>")): return None - tail = turns[-1] - if any(marker in tail for marker in ("<|system|>", "<|user|>", "<|observation|>")): + think_start = assistant_prompt.rfind("") + think_end = assistant_prompt.rfind("") + if think_start == think_end == -1: return None - start, end = tail.rfind(""), tail.rfind("") - return start > end if start != -1 or end != -1 else None + return think_start > think_end def detect_and_parse(self, text: str) -> StreamingParseResult: detector = type(self)(stream_reasoning=self.stream_reasoning, force_reasoning=self._in_reasoning) @@ -797,6 +804,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: while self._buffer: if self._in_tool_call: + # Tool arguments are literal, including any reasoning markers. consumed, complete = self._tool_scanner.scan(self._buffer) normal_parts.append(self._buffer[:consumed]) self._buffer = self._buffer[consumed:] @@ -807,15 +815,15 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: positions = [(self._buffer.find(marker), marker) for marker in markers] positions = [(position, marker) for position, marker in positions if position != -1] - parts = reasoning_parts if self._in_reasoning else normal_parts + output_parts = reasoning_parts if self._in_reasoning else normal_parts if not positions: # A delimiter may share a chunk with reasoning or span chunks. text, self._buffer = prefix_hold(self._buffer, markers) - parts.append(text) + output_parts.append(text) break position, marker = min(positions) - parts.append(self._buffer[:position]) + output_parts.append(self._buffer[:position]) if marker == self.tool_start_token: self._in_reasoning = False self._in_tool_call = True @@ -1009,6 +1017,17 @@ class ReasoningParser: "gemma4": Gemma4Detector, } + @classmethod + def get_initial_reasoning_state(cls, model_type: Optional[str], prompt: str, default: bool) -> bool: + """Resolve the starting phase once for parsing, stops and answer constraints.""" + if not model_type: + return False + detector_class = cls.DetectorMap.get(model_type.lower()) + if detector_class is None: + raise ValueError(f"Unsupported model type: {model_type}") + state = detector_class.initial_reasoning_state(prompt) + return default if state is None else state + def __init__( self, model_type: Optional[str] = None, diff --git a/unit_tests/server/httpserver/test_grammar.py b/unit_tests/server/httpserver/test_grammar.py index bb29ca9b4a..916c0f665c 100644 --- a/unit_tests/server/httpserver/test_grammar.py +++ b/unit_tests/server/httpserver/test_grammar.py @@ -44,7 +44,7 @@ def http_client(compiler, monkeypatch): monkeypatch.setattr(api_http.g_objs, "g_generate_stream_func", lightllm_generate_stream) monkeypatch.setattr(api_http.g_objs, "metric_client", SimpleNamespace(counter_inc=lambda *args: None)) monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="test")) - monkeypatch.setattr(api_openai, "_is_force_thinking_mode", lambda request, prompt=None: False) + monkeypatch.setattr(api_openai, "_is_force_thinking_mode", lambda request: False) yield httpx.AsyncClient(transport=httpx.ASGITransport(app=api_http.app), base_url="http://test") generate.assert_not_called() diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py index c9239ccc50..208c7fc7a5 100644 --- a/unit_tests/server/test_glm_reasoning_tool_calls.py +++ b/unit_tests/server/test_glm_reasoning_tool_calls.py @@ -41,21 +41,33 @@ def complete_chat(monkeypatch): args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47") monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="<|assistant|>")) - def complete(chunks, finish_status=FinishStatus.FINISHED_STOP, **request_kwargs): + def complete( + chunks, + finish_status=FinishStatus.FINISHED_STOP, + prompt="<|assistant|>", + parser_name="glm45", + compiler=None, + **request_kwargs, + ): + args.reasoning_parser = parser_name + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value=prompt)) + captured_params = [] + async def generate(*args, **kwargs): + captured_params.append(args[1]) for index, (text, token_id) in enumerate(chunks): finish = FinishStatus(finish_status) if index == len(chunks) - 1 else FinishStatus() yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish - manager = SimpleNamespace(tokenizer=SimpleNamespace(), generate=generate, output_grammar_compiler=None) + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text]) + manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler) monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) request = ChatCompletionRequest(model="glm", messages=[], tools=TOOLS, **request_kwargs) async def run(): response = await api_openai.chat_completions_impl(request, None) - result = SimpleNamespace(reasoning="", content="", calls=[], finish_reason=None) + result = SimpleNamespace(reasoning="", content="", calls=[], finish_reason=None, sampling_params=None) if not request.stream: choice = response.choices[0] result.reasoning = choice.message.reasoning or "" @@ -65,6 +77,7 @@ async def run(): for call in choice.message.tool_calls or [] ] result.finish_reason = choice.finish_reason + result.sampling_params = captured_params[0] return result calls = {} @@ -93,6 +106,7 @@ async def run(): assert done assert list(calls) == list(range(len(calls))) result.calls = list(calls.values()) + result.sampling_params = captured_params[0] return result return asyncio.run(run()) @@ -214,7 +228,9 @@ def test_chat_completion_exposes_glm_tool_calls(complete_chat, stream, separate_ [ ("<|assistant|>", True), ("<|assistant|>reasoning directive", True), + ("<|assistant|>", False), ("<|assistant|>", False), + ("<|assistant|>firstsecond", True), ("<|assistant|>old<|user|>question<|assistant|>", True), ("<|assistant|>old<|user|>question<|assistant|>", None), ("<|assistant|>old<|observation|>tool result", None), @@ -229,18 +245,71 @@ def test_glm_prompt_override_does_not_change_other_models(monkeypatch): monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="qwen3")) monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) request = ChatCompletionRequest(model="qwen", messages=[], chat_template_kwargs={"enable_thinking": False}) - assert api_openai._is_force_thinking_mode(request, "<|assistant|>") is False + default = api_openai._is_force_thinking_mode(request) + assert ReasoningParser.get_initial_reasoning_state("qwen3", "<|assistant|>", default) is False + + +@pytest.mark.parametrize("default", [False, True]) +@pytest.mark.parametrize("prompt", ["custom prompt", "<|assistant|>", "<|assistant|>old<|user|>question"]) +def test_glm_unrecognized_prompt_preserves_request_mode(prompt, default): + assert ReasoningParser.get_initial_reasoning_state("glm45", prompt, default) is default @pytest.mark.parametrize("thinking", [False, True]) def test_glm_grammar_activation_follows_rendered_prompt(monkeypatch, thinking): monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="glm45")) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) request = ChatCompletionRequest(model="glm", messages=[], chat_template_kwargs={"enable_thinking": not thinking}) prompt = "<|assistant|>" + ("" if thinking else "") tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text]) - force_reasoning = api_openai._is_force_thinking_mode(request, prompt) - assert force_reasoning is thinking - assert api_openai._get_guided_reasoning_end(request, tokenizer, force_reasoning) == ([99] if thinking else []) + default = api_openai._is_force_thinking_mode(request) + initial_reasoning_state = ReasoningParser.get_initial_reasoning_state("glm45", prompt, default) + assert default is not thinking + assert initial_reasoning_state is thinking + assert api_openai._get_guided_reasoning_end(request, tokenizer, initial_reasoning_state) == ( + [99] if thinking else [] + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("thinking", [False, True]) +def test_glm_prompt_phase_is_shared_by_output_stop_and_grammar(complete_chat, compiler, stream, thinking): + chunks = [("Need a file.", 1), ("", 99)] if thinking else [] + chunks.append(('{"ok": true}', 2)) + result = complete_chat( + chunks, + prompt="<|assistant|>" + ("" if thinking else ""), + compiler=compiler, + stream=stream, + separate_reasoning=True, + tool_choice="none", + stop=["STOP"], + response_format={"type": "json_object"}, + chat_template_kwargs={"enable_thinking": not thinking}, + ) + assert result.reasoning == ("Need a file." if thinking else "") + assert result.content == '{"ok": true}' + assert result.sampling_params._initial_reasoning_state == int(thinking) + assert result.sampling_params.constraint_out_params.guided_reasoning_end.to_list() == ([99] if thinking else []) + + +@pytest.mark.parametrize("parser_name", ["glm45", "qwen3"]) +@pytest.mark.parametrize("thinking", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +def test_stream_flush_preserves_truncated_text_in_each_phase(complete_chat, parser_name, thinking, separate_reasoning): + result = complete_chat( + [("Body", 1), ("<", 2)], + finish_status=FinishStatus.FINISHED_LENGTH, + parser_name=parser_name, + prompt="<|assistant|>" + ("" if thinking else ""), + stream=True, + separate_reasoning=separate_reasoning, + tool_choice="none", + chat_template_kwargs={"enable_thinking": thinking}, + ) + assert result.reasoning == ("Body<" if thinking and separate_reasoning else "") + assert result.content == ("" if thinking and separate_reasoning else "Body<") + assert result.finish_reason == "length" @pytest.mark.parametrize("suffix", ["<", " Date: Thu, 22 Oct 2026 15:13:05 +0000 Subject: [PATCH 05/13] refactor: simplify GLM parser state and argument handling --- lightllm/server/api_openai.py | 21 +- lightllm/server/function_call_parser.py | 257 ++++++++---------- lightllm/server/parser_utils.py | 62 +++++ lightllm/server/reasoning_parser.py | 75 ++--- .../structured_output/test_reasoning.py | 10 +- .../server/test_glm_reasoning_tool_calls.py | 6 +- 6 files changed, 217 insertions(+), 214 deletions(-) create mode 100644 lightllm/server/parser_utils.py diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 6c9224b875..fee264748a 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -201,15 +201,10 @@ def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool: return True # default -def _get_guided_reasoning_end( - request: ChatCompletionRequest, tokenizer, initial_reasoning_state: Optional[bool] = None -) -> List[int]: +def _get_guided_reasoning_end(tokenizer, parser_name: Optional[str], initial_reasoning_state: bool) -> List[int]: """Delay answer constraints while generation starts inside reasoning.""" - if initial_reasoning_state is None: - initial_reasoning_state = _is_force_thinking_mode(request) - if not initial_reasoning_state: + if not parser_name or not initial_reasoning_state: return [] - parser_name = get_env_start_args().reasoning_parser detector = ReasoningParser.DetectorMap[parser_name.lower()]() token_ids = tokenizer.encode(detector.grammar_start_token, add_special_tokens=False) return token_ids @@ -222,7 +217,7 @@ def _process_reasoning_stream( metadata: Dict[str, Any], request: ChatCompletionRequest, force_reasoning: bool, -) -> tuple[Optional[str], str]: +) -> tuple[str, str]: """Process reasoning content and update its token usage.""" if index not in reasoning_parser_dict: reasoning_parser_dict[index] = ReasoningParser( @@ -414,7 +409,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req if sampling_params_dict.get("guided_json") or sampling_params_dict.get("guided_grammar"): sampling_params_dict["guided_reasoning_end"] = _get_guided_reasoning_end( - request, g_objs.httpserver_manager.tokenizer, initial_reasoning_state + g_objs.httpserver_manager.tokenizer, reasoning_parser, initial_reasoning_state ) sampling_params = SamplingParams() @@ -604,14 +599,14 @@ async def stream_results() -> AsyncGenerator[bytes, None]: reasoning_parser_dict, metadata, request, - initial_reasoning_state, + force_reasoning=initial_reasoning_state, ) if current_finish_reason is not None: # Flush before tool parsing so buffered argument text follows # the same path as every other chunk, including on truncation. flush_reasoning, flush_text = reasoning_parser_dict[choice_index].flush() - reasoning_text = (reasoning_text or "") + (flush_reasoning or "") - delta = (delta or "") + (flush_text or "") + reasoning_text += flush_reasoning + delta += flush_text if reasoning_text: if request.separate_reasoning: choice_data = ChatCompletionStreamResponseChoice( @@ -627,7 +622,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls)}\n\n" else: - delta = reasoning_text + (delta or "") + delta = reasoning_text + delta if request.tool_choice != "none" and request.tools: # parse_increment => returns (normal_text, calls) diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index 24df428a46..e2b6d2ab84 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -29,6 +29,7 @@ from pydantic import BaseModel, Field from .api_models import Tool +from .parser_utils import GlmToolCallScanner logger = logging.getLogger(__name__) ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"] @@ -1187,87 +1188,40 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult(normal_text=current_text) -class GlmToolCallScanner: - """Track a GLM tool envelope without interpreting tags inside argument values. - - The caller forwards all input until the outer closing tag. Only a possible - delimiter prefix is retained here for scanning the next chunk; those bytes - have already been forwarded and are excluded from the consumed count. - """ - - def __init__(self): - self._in_arg_value = False - self._pending = "" - - def scan(self, new_text: str) -> Tuple[int, bool]: - pending_length = len(self._pending) - text = self._pending + new_text - self._pending = "" - offset = 0 - while offset < len(text): - markers = ("",) if self._in_arg_value else ("", "") - positions = [(text.find(marker, offset), marker) for marker in markers] - positions = [(position, marker) for position, marker in positions if position != -1] - if not positions: - max_hold = min(len(text) - offset, max(map(len, markers)) - 1) - for length in range(max_hold, 0, -1): - if any(marker.startswith(text[-length:]) for marker in markers): - self._pending = text[-length:] - break - return len(new_text), False - - position, marker = min(positions) - offset = position + len(marker) - if marker == "": - return offset - pending_length, True - self._in_arg_value = marker == "" - - return len(new_text), False - - class Glm47Detector(BaseFormatDetector): - """ - Detector for GLM-4.7/GLM-4.7-Flash model function call format. - - The GLM-4.7 format uses an XML-style envelope with arg_key/arg_value pairs - instead of JSON arguments. + """Parse GLM's XML-style tool calls. - Format Structure: - ``` - function_name - param1 - value1 - param2 - value2 - - ``` - - Example: - ``` - tool_brave_web_search_post - query - test search - count - 5 - - ``` - - Key Components: - - Tool Call Tags: `` and `` wrap each individual call - - Function Name: Appears on the first line after `` - - Arguments: Pairs of `name` and `value` - - Reference: https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_parser.py + Format: namekeyvalue + Arguments stay literal until the call is complete, then convert according + to the function's parameter schema. """ _TYPE_ALIASES = { - **dict.fromkeys(("str", "text", "varchar", "char", "enum"), "string"), - **dict.fromkeys(("int", "int32", "int64", "uint", "uint32", "uint64", "long", "short", "unsigned"), "integer"), - **dict.fromkeys(("float", "float32", "float64", "double"), "number"), + "str": "string", + "text": "string", + "varchar": "string", + "char": "string", + "enum": "string", + "int": "integer", + "int32": "integer", + "int64": "integer", + "uint": "integer", + "uint32": "integer", + "uint64": "integer", + "long": "integer", + "short": "integer", + "unsigned": "integer", + "float": "number", + "float32": "number", + "float64": "number", + "double": "number", "bool": "boolean", "dict": "object", - **dict.fromkeys(("arr", "list", "sequence"), "array"), + "arr": "array", + "list": "array", + "sequence": "array", } + _TYPE_PRIORITY = ("null", "integer", "number", "boolean", "object", "array", "string") def __init__(self): super().__init__() @@ -1275,19 +1229,11 @@ def __init__(self): self.eot_token = "" self.tool_call_separator = "\n" - # Regex patterns for parsing GLM-4.7 tool calls - # Extract function name and arguments from a tool call block - # Function name can be followed by newline OR directly by - # Pattern: function_name(\n|)... - self.func_detail_regex = re.compile( - r"([^<\n]+?)(?:\n|(?=)|(?=))(.*?)\Z", re.DOTALL - ) + # A function name ends at a newline, the first argument, or the closing tag. self.func_name_regex = re.compile(r"([^<\n]+?)(?:\n|(?=)|(?=))") - self._streaming_tool_name: Optional[str] = None + self.func_arg_regex = re.compile(r"(.*?)\s*(.*?)", re.DOTALL) self._tool_scanner = GlmToolCallScanner() self._tool_scan_offset = len(self.bot_token) - # Extract arg_key/arg_value pairs - self.func_arg_regex = re.compile(r"(.*?)\s*(.*?)", re.DOTALL) def has_tool_call(self, text: str) -> bool: """Check if the text contains a GLM-4.7 format tool call.""" @@ -1303,7 +1249,10 @@ def _schema_types(cls, schema: Any) -> set[str]: types.add(declared) elif isinstance(declared, list): types.update(item for item in declared if isinstance(item, str)) - for value in schema.get("enum", []) if isinstance(schema.get("enum"), list) else []: + enum_values = schema.get("enum", []) + if not isinstance(enum_values, list): + enum_values = [] + for value in enum_values: if value is None: types.add("null") elif isinstance(value, bool): @@ -1326,15 +1275,38 @@ def _schema_types(cls, schema: Any) -> set[str]: return types or {"string"} @classmethod - def _finite_json(cls, value: Any) -> bool: + def _is_finite_json(cls, value: Any) -> bool: if isinstance(value, float): return math.isfinite(value) if isinstance(value, dict): - return all(cls._finite_json(item) for item in value.values()) + return all(cls._is_finite_json(item) for item in value.values()) if isinstance(value, list): - return all(cls._finite_json(item) for item in value) + return all(cls._is_finite_json(item) for item in value) return True + @staticmethod + def _parse_typed_value(value: str, value_type: str) -> Any: + match value_type: + case "string": + return value + case "null": + if value.lower() == "null": + return None + case "integer": + return int(value) + case "number": + number = float(value) + return int(number) if number.is_integer() else number + case "boolean": + normalized = value.strip().lower() + if normalized in ("true", "1"): + return True + if normalized in ("false", "0"): + return False + case "object" | "array": + return json.loads(value) + raise ValueError(f"Cannot parse value as {value_type}") + @classmethod def _convert_param_value(cls, value: str, schema: Any) -> Any: """Match vLLM's schema-based coercion without changing literal strings. @@ -1343,38 +1315,23 @@ def _convert_param_value(cls, value: str, schema: Any) -> Any: Unspecified fields remain strings. Non-finite numbers remain raw text so the resulting arguments are always valid JSON. """ - types = { - cls._TYPE_ALIASES.get(item.strip().lower(), item.strip().lower()) for item in cls._schema_types(schema) - } - for candidate in ("null", "integer", "number", "boolean", "object", "array", "string"): - if candidate not in types: + schema_types = set() + for type_name in cls._schema_types(schema): + type_name = type_name.strip().lower() + schema_types.add(cls._TYPE_ALIASES.get(type_name, type_name)) + + for value_type in cls._TYPE_PRIORITY: + if value_type not in schema_types: continue try: - if candidate == "null" and value.lower() == "null": - return None - if candidate == "integer": - return int(value) - if candidate == "number": - number = float(value) - if math.isfinite(number): - return number if number != int(number) else int(number) - if candidate == "boolean": - normalized = value.strip().lower() - if normalized in ("true", "1"): - return True - if normalized in ("false", "0"): - return False - if candidate in ("object", "array"): - parsed = json.loads(value) - if cls._finite_json(parsed): - return parsed - if candidate == "string": - return value + parsed = cls._parse_typed_value(value, value_type) except (ValueError, TypeError): continue + if cls._is_finite_json(parsed): + return parsed try: parsed = json.loads(value) - return parsed if cls._finite_json(parsed) else value + return parsed if cls._is_finite_json(parsed) else value except ValueError: return value @@ -1405,10 +1362,10 @@ def _iter_tool_call_spans(self, text: str) -> Iterator[Tuple[int, int]]: start = text.find(self.bot_token) while start != -1: body_start = start + len(self.bot_token) - consumed, complete = GlmToolCallScanner().scan(text[body_start:]) - if not complete: + relative_end = GlmToolCallScanner().find_end(text[body_start:]) + if relative_end is None: return - end = body_start + consumed + end = body_start + relative_end yield start, end start = text.find(self.bot_token, end) @@ -1431,16 +1388,15 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult for start, end in self._iter_tool_call_spans(text): normal_parts.append(text[cursor:start]) cursor = end - match_result = text[start:end] + tool_text = text[start:end] try: - # Extract function name and arguments - func_detail = self.func_detail_regex.search(match_result) - if not func_detail: - logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}") + header = self.func_name_regex.match(tool_text) + if header is None: + logger.warning(f"Failed to parse GLM-4.7 tool call: {tool_text}") continue - func_name = func_detail.group(1).strip() - arg_text = func_detail.group(2) if func_detail.group(2) else "" + func_name = header.group(1).strip() + arg_text = tool_text[header.end() : -len(self.eot_token)] # Validate function name if func_name not in tool_indices: @@ -1458,15 +1414,33 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult ) ) except Exception as e: - logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}, error: {str(e)}") + logger.warning(f"Failed to parse GLM-4.7 tool call: {tool_text}, error: {str(e)}") continue # Keep ordinary text between and after calls. An unfinished envelope # stays buffered in streaming mode and must not leak in non-streaming. - incomplete_start = text.find(self.bot_token, cursor) - normal_parts.append(text[cursor : incomplete_start if incomplete_start != -1 else len(text)]) + content_end = text.find(self.bot_token, cursor) + if content_end == -1: + content_end = len(text) + normal_parts.append(text[cursor:content_end]) return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls) + def _stream_incomplete_call(self, tools: List[Tool]) -> Optional[ToolCallItem]: + """Announce the name once, then emit empty argument deltas while buffering.""" + if self.current_tool_name_sent: + return ToolCallItem(tool_index=self.current_tool_id, parameters="") + + header = self.func_name_regex.match(self._buffer) + if header is None: + return None + name = header.group(1).strip() + if name not in self._get_tool_indices(tools): + return None + + self.current_tool_id += 1 + self.current_tool_name_sent = True + return ToolCallItem(tool_index=self.current_tool_id, name=name, parameters="") + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Buffer XML arguments while emitting tool deltas to keep the stream alive.""" self._buffer += new_text @@ -1483,35 +1457,26 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami normal_text += self._buffer[:start] self._buffer = self._buffer[start:] - consumed, complete = self._tool_scanner.scan(self._buffer[self._tool_scan_offset :]) - self._tool_scan_offset += consumed - if not complete: - # Match the buffered Qwen3-Coder flow: announce the name once, - # then send empty argument deltas until the XML call is complete. - # Long Write/Bash arguments must not leave the HTTP stream idle. - name = None - if self._streaming_tool_name is None: - match = self.func_name_regex.match(self._buffer) - if match and match.group(1).strip() in self._get_tool_indices(tools): - name = match.group(1).strip() - self._streaming_tool_name = name - self.current_tool_id += 1 - if self._streaming_tool_name is not None: - calls.append(ToolCallItem(tool_index=self.current_tool_id, name=name, parameters="")) + relative_end = self._tool_scanner.find_end(self._buffer[self._tool_scan_offset :]) + if relative_end is None: + self._tool_scan_offset = len(self._buffer) + pending_call = self._stream_incomplete_call(tools) + if pending_call is not None: + calls.append(pending_call) break - end = self._tool_scan_offset + end = self._tool_scan_offset + relative_end # Waiting for the closing tag avoids re-emitting partial JSON or # rewriting already sent arguments when another XML key arrives. result = self.detect_and_parse(self._buffer[:end], tools) for call in result.calls: - if self._streaming_tool_name is None: - self.current_tool_id += 1 - else: + if self.current_tool_name_sent: call.name = None + else: + self.current_tool_id += 1 call.tool_index = self.current_tool_id calls.append(call) - self._streaming_tool_name = None + self.current_tool_name_sent = False self._buffer = self._buffer[end:] self._tool_scanner = GlmToolCallScanner() self._tool_scan_offset = len(self.bot_token) diff --git a/lightllm/server/parser_utils.py b/lightllm/server/parser_utils.py new file mode 100644 index 0000000000..8299357554 --- /dev/null +++ b/lightllm/server/parser_utils.py @@ -0,0 +1,62 @@ +# Copyright 2025 ModelTC Team +# Copyright 2023-2025 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. + +import re +from typing import Optional, Sequence, Tuple + + +def prefix_hold(text: str, markers: Sequence[str]) -> Tuple[str, str]: + """Split off the longest suffix that could become a complete marker.""" + max_length = min(len(text), max((len(marker) - 1 for marker in markers), default=0)) + for length in range(max_length, 0, -1): + suffix = text[-length:] + if any(len(suffix) < len(marker) and marker.startswith(suffix) for marker in markers): + return text[:-length], suffix + return text, "" + + +class GlmToolCallScanner: + """Find the outer GLM tool-call end while treating argument values as literal text.""" + + _MARKERS = ("", "", "") + _MARKER_PATTERN = re.compile(r"|") + + def __init__(self): + self._in_arg_value = False + self._pending_marker = "" + + def find_end(self, new_text: str) -> Optional[int]: + """Return the closing tag's end offset in this chunk, or None if incomplete. + + Callers may forward every character immediately. The retained marker + prefix is only for recognizing a delimiter split across chunks. + """ + pending_length = len(self._pending_marker) + text = self._pending_marker + new_text + self._pending_marker = "" + offset = 0 + + for match in self._MARKER_PATTERN.finditer(text): + offset = match.end() + marker = match.group() + if self._in_arg_value: + if marker == "": + self._in_arg_value = False + elif marker == "": + self._in_arg_value = True + elif marker == "": + return offset - pending_length + + _, self._pending_marker = prefix_hold(text[offset:], self._MARKERS) + return None diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 13a7f357f4..b0036ae288 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -16,7 +16,7 @@ from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id -from lightllm.server.function_call_parser import GlmToolCallScanner +from lightllm.server.parser_utils import GlmToolCallScanner, prefix_hold @dataclass @@ -37,28 +37,6 @@ class Token: end: int -def prefix_hold(text: str, tokens: List[str]) -> Tuple[str, str]: - """ - Holds back the longest suffix of `text` that could be a prefix of any token. - Returns (emit_now, keep_for_later). - """ - if not text: - return "", "" - max_hold = 0 - for tok in tokens: - if not tok: - continue - # Check for prefixes of tok in the suffix of text - L = min(len(tok) - 1, len(text)) - for k in range(L, 0, -1): - if tok.startswith(text[-k:]): - max_hold = max(max_hold, k) - break - if max_hold == 0: - return text, "" - return text[:-max_hold], text[-max_hold:] - - def iter_tokens(text: str, start_pos: int = 0) -> Iterator[Token]: """Iterate over structural tokens in left-to-right order.""" TOKENS = { @@ -586,11 +564,6 @@ def __init__( class BaseReasoningFormatDetector: """Base class providing two sets of interfaces: one-time and streaming incremental.""" - @classmethod - def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: - """Infer the generation phase from the prompt, or defer to the request.""" - return None - def __init__( self, think_start_token: str, @@ -606,6 +579,11 @@ def __init__( self._buffer = "" self.stripped_think_start = False + @classmethod + def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: + """Infer the generation phase from the prompt, or defer to the request.""" + return None + @property def grammar_start_token(self) -> str: """Generated delimiter after which the final answer's grammar applies.""" @@ -766,11 +744,12 @@ class Glm45Detector(BaseReasoningFormatDetector): """Route GLM tool envelopes to content, even before an explicit ````.""" tool_start_token = "" + _MARKERS = ("", "", tool_start_token) + _MARKER_PATTERN = re.compile("|".join(re.escape(marker) for marker in _MARKERS)) def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False): super().__init__("", "", stream_reasoning=stream_reasoning, force_reasoning=force_reasoning) - self._in_tool_call = False - self._tool_scanner = GlmToolCallScanner() + self._tool_scanner: Optional[GlmToolCallScanner] = None @classmethod def initial_reasoning_state(cls, prompt: str) -> Optional[bool]: @@ -800,39 +779,37 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: self._buffer += new_text reasoning_parts = [] normal_parts = [] - markers = [self.think_start_token, self.think_end_token, self.tool_start_token] while self._buffer: - if self._in_tool_call: + if self._tool_scanner is not None: # Tool arguments are literal, including any reasoning markers. - consumed, complete = self._tool_scanner.scan(self._buffer) - normal_parts.append(self._buffer[:consumed]) - self._buffer = self._buffer[consumed:] - if not complete: - break - self._in_tool_call = False + end = self._tool_scanner.find_end(self._buffer) + if end is None: + end = len(self._buffer) + else: + self._tool_scanner = None + normal_parts.append(self._buffer[:end]) + self._buffer = self._buffer[end:] continue - positions = [(self._buffer.find(marker), marker) for marker in markers] - positions = [(position, marker) for position, marker in positions if position != -1] + match = self._MARKER_PATTERN.search(self._buffer) output_parts = reasoning_parts if self._in_reasoning else normal_parts - if not positions: + if match is None: # A delimiter may share a chunk with reasoning or span chunks. - text, self._buffer = prefix_hold(self._buffer, markers) + text, self._buffer = prefix_hold(self._buffer, self._MARKERS) output_parts.append(text) break - position, marker = min(positions) - output_parts.append(self._buffer[:position]) + output_parts.append(self._buffer[: match.start()]) + marker = match.group() + self._buffer = self._buffer[match.end() :] if marker == self.tool_start_token: self._in_reasoning = False - self._in_tool_call = True self._tool_scanner = GlmToolCallScanner() # Preserve the opener and literal thinking tags in tool arguments. normal_parts.append(marker) else: self._in_reasoning = marker == self.think_start_token - self._buffer = self._buffer[position + len(marker) :] return StreamingParseResult( reasoning_text="".join(reasoning_parts), @@ -1070,17 +1047,17 @@ def update_reasoning_token_count(self, token_id: int) -> None: else: self.reasoning_tokens += 1 - def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]: + def parse_non_stream(self, full_text: str) -> Tuple[str, str]: """Non-streaming call: one-time parsing""" ret = self.detector.detect_and_parse(full_text) return ret.reasoning_text, ret.normal_text - def parse_stream_chunk(self, chunk_text: str) -> Tuple[Optional[str], Optional[str]]: + def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, str]: """Streaming call: incremental parsing""" ret = self.detector.parse_streaming_increment(chunk_text) return ret.reasoning_text, ret.normal_text - def flush(self) -> Tuple[Optional[str], Optional[str]]: + def flush(self) -> Tuple[str, str]: """Flush remaining buffered content when generation ends prematurely.""" ret = self.detector.flush() return ret.reasoning_text, ret.normal_text diff --git a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py index 8c265847db..e4f08421ad 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py +++ b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py @@ -49,7 +49,8 @@ def encode(text, add_special_tokens): chat_template_kwargs=template_kwargs, reasoning_effort=effort, ) - result = api_openai._get_guided_reasoning_end(request, SimpleNamespace(encode=encode)) + thinking = api_openai._is_force_thinking_mode(request) + result = api_openai._get_guided_reasoning_end(SimpleNamespace(encode=encode), parser, thinking) assert encoded == ([marker] if marker else []) assert result == ([7, 8] if marker else []) @@ -134,8 +135,11 @@ def test_fixed_reasoning_does_not_require_a_template_thinking_switch(monkeypatch monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False) tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) request = ChatCompletionRequest(model="test", messages=[]) - assert api_openai._is_force_thinking_mode(request) == (marker is not None) - assert api_openai._get_guided_reasoning_end(request, tokenizer) == (list(marker.encode()) if marker else []) + thinking = api_openai._is_force_thinking_mode(request) + assert thinking == (marker is not None) + assert api_openai._get_guided_reasoning_end(tokenizer, parser, thinking) == ( + list(marker.encode()) if marker else [] + ) @pytest.mark.parametrize("stream", [False, True]) diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py index 208c7fc7a5..92bc3e8496 100644 --- a/unit_tests/server/test_glm_reasoning_tool_calls.py +++ b/unit_tests/server/test_glm_reasoning_tool_calls.py @@ -54,8 +54,8 @@ def complete( monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value=prompt)) captured_params = [] - async def generate(*args, **kwargs): - captured_params.append(args[1]) + async def generate(_prompt, sampling_params, *_args, **_kwargs): + captured_params.append(sampling_params) for index, (text, token_id) in enumerate(chunks): finish = FinishStatus(finish_status) if index == len(chunks) - 1 else FinishStatus() yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish @@ -266,7 +266,7 @@ def test_glm_grammar_activation_follows_rendered_prompt(monkeypatch, thinking): initial_reasoning_state = ReasoningParser.get_initial_reasoning_state("glm45", prompt, default) assert default is not thinking assert initial_reasoning_state is thinking - assert api_openai._get_guided_reasoning_end(request, tokenizer, initial_reasoning_state) == ( + assert api_openai._get_guided_reasoning_end(tokenizer, "glm45", initial_reasoning_state) == ( [99] if thinking else [] ) From 710118e11598f2b1a68959957bf24535a12f792d Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:16:56 +0800 Subject: [PATCH 06/13] fix(glm5): align indexer, routing and mHC numerics Use the indexer's dedicated LayerNorm epsilon of 1e-6 and retain MoE correction bias in FP32 through checkpoint loading and TP/EP routing. For GLM mHC, normalize activations in FP32 before the mapping projection, using an IEEE FP32 Triton path independent of global TF32 settings. Keep pre coefficients in FP32 and round post/residual coefficients to the activation dtype before mixing. Preserve the existing fused FP32 post accumulation and the legacy default path for other model families. Add 11 CUDA regression cases covering actual GLM callers, pooled-indexer selection, TP/EP routing, projection order and coefficient rounding, including negative controls that reproduce the old numerical behavior. Validation: Python 3.10 syntax/bytecode, line-width and whitespace checks, independent source review, and CPU numerical counterexamples passed. CUDA kernel compilation, GPU accuracy tests and latency measurements remain pending because the requested GPU host was unreachable from the authoring environment. No bitwise equivalence across reduction backends or native versus fused post-mixing implementations is claimed. --- .../fused_moe/fused_moe_weight.py | 4 +- .../basemodel/triton_kernel/mhc/post.py | 8 +- .../basemodel/triton_kernel/mhc/pre_norm.py | 132 +++++++++-- lightllm/models/glm5_next/indexer.py | 4 +- .../layer_infer/transformer_layer_infer.py | 1 + .../basemodel/triton_kernel/mhc/test_mhc.py | 159 ++++++++++++- .../glm5_next/test_precision_alignment.py | 210 ++++++++++++++++++ 7 files changed, 496 insertions(+), 22 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_precision_alignment.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index fd3de5c9b4..f414dbeb44 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -306,9 +306,11 @@ def _create_weight(self): self.per_expert_scale = None # Create e_score_correction_bias if self.e_score_correction_bias_name: + # Routing compares close corrected scores in FP32. Preserve the + # checkpoint bias instead of rounding it to the expert weight dtype. self.e_score_correction_bias = torch.empty( (self.n_routed_experts,), - dtype=self.data_type_, + dtype=torch.float32, device=f"cuda:{self.device_id_}", ) self.e_score_correction_bias.load_ok = False diff --git a/lightllm/common/basemodel/triton_kernel/mhc/post.py b/lightllm/common/basemodel/triton_kernel/mhc/post.py index b4050e0765..31ecca43b9 100644 --- a/lightllm/common/basemodel/triton_kernel/mhc/post.py +++ b/lightllm/common/basemodel/triton_kernel/mhc/post.py @@ -13,8 +13,8 @@ def _hc_post_4stream_kernel( layer_output, # 输入 [T, H]，激活 dtype。 residual, # 输入 [T, 4 * H]，激活 dtype；当前子层 pre-mixing 前的残差。 - residual_mix, # 输入 [T, 4, 4]，fp32；最后两维为输入、输出 stream。 - post_mix, # 输入 [T, 4]，fp32。 + residual_mix, # 输入 [T, 4, 4]，fp32 或激活 dtype；最后两维为输入、输出 stream。 + post_mix, # 输入 [T, 4]，fp32 或激活 dtype。 output, # 输出 [T, 4 * H]，与 layer_output 相同 dtype。 hidden: tl.constexpr, layer_stride_m: tl.constexpr, @@ -93,8 +93,8 @@ def hc_post( Args: layer_output: [T, H]，连续的 attention / FFN 子层输出。 residual: [T, S * H]，连续的、该子层 pre-mixing 前保存的原始残差。 - residual_mix: [T, S, S]，fp32；来自 hc_pre_norm。 - post_mix: [T, S]，fp32；来自同一次 pre 调用。 + residual_mix: [T, S, S]，fp32 或激活 dtype；来自 hc_pre_norm。 + post_mix: [T, S]，与 residual_mix 相同 dtype；来自同一次 pre 调用。 streams: 残差 stream 数 S；当前 Triton 入口要求 S = 4。 Returns: diff --git a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py index 2710040673..d43a7ecaf8 100644 --- a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py +++ b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py @@ -14,6 +14,64 @@ LIGHTLLM_DISABLE_DEEPGEMM_MHC = os.getenv("LIGHTLLM_DISABLE_DEEPGEMM_MHC", "False").upper() in ["ON", "TRUE", "1"] +@triton.jit +def _hc_input_inv_rms_kernel( + x, + inv_rms, + FLATTENED_HIDDEN: tl.constexpr, + RMS_EPS: tl.constexpr, + BLOCK_K: tl.constexpr, +): + token = tl.program_id(0) + offsets = tl.arange(0, BLOCK_K) + values = tl.load(x + token * FLATTENED_HIDDEN + offsets, offsets < FLATTENED_HIDDEN, 0).to(tl.float32) + variance = tl.sum(values * values, axis=0) / FLATTENED_HIDDEN + tl.store(inv_rms + token, tl.rsqrt(variance + RMS_EPS)) + + +@triton.jit +def _hc_normalized_projection_kernel( + x, + fn, + inv_rms, + gemm_partial, + tokens, + FLATTENED_HIDDEN: tl.constexpr, + MIX_SIZE: tl.constexpr, + SPLIT_K_BLOCKS: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, +): + token_offsets = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + split = tl.program_id(1) + mix_offsets = tl.arange(0, BLOCK_M) + k_offsets = split * SPLIT_K_BLOCKS * BLOCK_K + tl.arange(0, BLOCK_K) + rms = tl.load(inv_rms + token_offsets, token_offsets < tokens, 0) + accumulator = tl.zeros((BLOCK_T, BLOCK_M), dtype=tl.float32) + for _ in range(SPLIT_K_BLOCKS): + values = tl.load( + x + token_offsets[:, None] * FLATTENED_HIDDEN + k_offsets[None, :], + (token_offsets[:, None] < tokens) & (k_offsets[None, :] < FLATTENED_HIDDEN), + 0, + ).to(tl.float32) + weights = tl.load( + fn + mix_offsets[None, :] * FLATTENED_HIDDEN + k_offsets[:, None], + (mix_offsets[None, :] < MIX_SIZE) & (k_offsets[:, None] < FLATTENED_HIDDEN), + 0, + ) + # GLM projects the FP32-normalized input. Moving the RMS factor after + # the dot, or rounding this tile to TF32/BF16, changes that computation. + normalized = values * rms[:, None] + accumulator = tl.dot(normalized, weights, accumulator, input_precision="ieee") + k_offsets += BLOCK_K + tl.store( + gemm_partial + (split * tokens + token_offsets[:, None]) * MIX_SIZE + mix_offsets[None, :], + accumulator, + (token_offsets[:, None] < tokens) & (mix_offsets[None, :] < MIX_SIZE), + ) + + @triton.jit def _hc_prepare_prenorm_kernel( gemm_partial, # 输入 [P, T, M]，fp32；投影的 split-K 部分和。 @@ -21,8 +79,8 @@ def _hc_prepare_prenorm_kernel( scale, # 输入 [3]，fp32；pre、post、residual 三组缩放系数。 base, # 输入 [M]，fp32；M = 2 * S + S * S。 pre, # 输出 [T, S]，fp32。 - post, # 输出 [T, S]，fp32。 - residual_mix, # 输出 [T, S, S]，fp32；最后两维为输入、输出 stream。 + post, # 输出 [T, S]，fp32 或激活 dtype。 + residual_mix, # 输出 [T, S, S]，fp32 或激活 dtype；最后两维为输入、输出 stream。 gemm_stride_s, gemm_stride_m: tl.constexpr, sqrsum_stride_s, @@ -36,6 +94,7 @@ def _hc_prepare_prenorm_kernel( POST_MULTIPLIER: tl.constexpr, SINKHORN_ITERS: tl.constexpr, N_SPLITS: tl.constexpr, + PROJECTION_IS_NORMALIZED: tl.constexpr, ): """归并分片并准备混合权重：[P, T, M] / [P, T] -> [T, S] / [T, S, S]。 @@ -55,11 +114,13 @@ def _hc_prepare_prenorm_kernel( pre_raw += tl.load(partial_base + stream_offsets) post_raw += tl.load(partial_base + STREAMS + stream_offsets) matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets) - sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) - inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) - pre_raw *= inv_rms - post_raw *= inv_rms - matrix_raw *= inv_rms + if not PROJECTION_IS_NORMALIZED: + sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m) + if not PROJECTION_IS_NORMALIZED: + inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS) + pre_raw *= inv_rms + post_raw *= inv_rms + matrix_raw *= inv_rms pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets)) @@ -141,6 +202,38 @@ def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.de return max(1, min(sms // max(grid_size, 1), k_blocks // 4)) +def _hc_normalized_projection(x: torch.Tensor, fn: torch.Tensor, rms_eps: float) -> torch.Tensor: + """FP32 normalize-before-project, without materializing a [T, S * H] FP32 copy.""" + tokens, flattened_hidden = x.shape + mix_size = fn.shape[0] + n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device) + inv_rms = torch.empty((tokens,), dtype=torch.float32, device=x.device) + gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device) + _hc_input_inv_rms_kernel[(tokens,)]( + x, + inv_rms, + FLATTENED_HIDDEN=flattened_hidden, + RMS_EPS=rms_eps, + BLOCK_K=triton.next_power_of_2(flattened_hidden), + num_warps=8, + ) + _hc_normalized_projection_kernel[(triton.cdiv(tokens, 16), n_splits)]( + x, + fn, + inv_rms, + gemm_partial, + tokens, + FLATTENED_HIDDEN=flattened_hidden, + MIX_SIZE=mix_size, + SPLIT_K_BLOCKS=triton.cdiv(triton.cdiv(flattened_hidden, 64), n_splits), + BLOCK_T=16, + BLOCK_M=triton.next_power_of_2(mix_size), + BLOCK_K=64, + num_warps=4, + ) + return gemm_partial + + def hc_pre_norm( x: torch.Tensor, fn: torch.Tensor, @@ -153,6 +246,8 @@ def hc_pre_norm( hc_eps: float, sinkhorn_iters: int, post_multiplier: float = 2.0, + *, + glm5_numerics: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """将 mHC pre-mixing 与紧随其后的子层 RMSNorm 融合。 @@ -170,15 +265,19 @@ def hc_pre_norm( hc_eps: pre sigmoid 后的偏移量及 Sinkhorn 的稳定项。 sinkhorn_iters: Sinkhorn 迭代次数。 post_multiplier: post sigmoid 的乘数，默认 2.0。 + glm5_numerics: GLM-5.3 的 FP32 先归一化后投影路径；post/residual 系数 + 舍入到激活 dtype，pre 系数仍保留 fp32。 Returns: - (layer_input, residual_mix, post_mix)：分别为 bf16 [T, H]、fp32 [T, S, S]、fp32 [T, S]。 + (layer_input, residual_mix, post_mix)：形状为 [T, H]、[T, S, S]、[T, S]。 + layer_input 为 bf16；两个混合系数默认 fp32，glm5_numerics=True 时为激活 dtype。 layer_input 已完成子层 RMSNorm；residual_mix[t, i, j] 对应输入 i 到输出 j。 x 保持不变，留给对应 hc_post 使用。 默认用 DeepGEMM 融合投影与平方和计算，Triton 完成 Sinkhorn、stream 合并和 RMSNorm。 合并结果先舍入到 bf16，再做 RMSNorm。 LIGHTLLM_DISABLE_DEEPGEMM_MHC=1 禁用 DeepGEMM，改用 PyTorch 计算投影与平方和。 + glm5_numerics=True 则使用独立的 IEEE FP32 Triton 投影，不依赖 DeepGEMM 或全局 TF32 设置。 """ assert x.ndim == 2 and x.shape[-1] % streams == 0 @@ -187,8 +286,13 @@ def hc_pre_norm( assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous() tokens, flattened_hidden = x.shape hidden = flattened_hidden // streams + assert fn.shape == ((2 + streams) * streams, flattened_hidden) - if LIGHTLLM_DISABLE_DEEPGEMM_MHC: + if glm5_numerics: + gemm_partial = _hc_normalized_projection(x, fn, rms_eps) + sqrsum_partial = None + n_splits = gemm_partial.shape[0] + elif LIGHTLLM_DISABLE_DEEPGEMM_MHC: # PyTorch produces one complete partition: [1, T, M] and [1, T]. x_fp32 = x.float() gemm_partial = (x_fp32 @ fn.T).unsqueeze(0) @@ -204,8 +308,9 @@ def hc_pre_norm( tf32_hc_prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits) pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device) - post = torch.empty_like(pre) - residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device) + mix_dtype = x.dtype if glm5_numerics else torch.float32 + post = torch.empty((tokens, streams), dtype=mix_dtype, device=x.device) + residual_mix = torch.empty((tokens, streams, streams), dtype=mix_dtype, device=x.device) _hc_prepare_prenorm_kernel[(tokens,)]( gemm_partial, sqrsum_partial, @@ -216,8 +321,8 @@ def hc_pre_norm( residual_mix, gemm_partial.stride(0), gemm_partial.stride(1), - sqrsum_partial.stride(0), - sqrsum_partial.stride(1), + sqrsum_partial.stride(0) if sqrsum_partial is not None else 0, + sqrsum_partial.stride(1) if sqrsum_partial is not None else 0, pre.stride(0), residual_mix.stride(0), FLATTENED_HIDDEN=flattened_hidden, @@ -227,6 +332,7 @@ def hc_pre_norm( POST_MULTIPLIER=post_multiplier, SINKHORN_ITERS=sinkhorn_iters, N_SPLITS=n_splits, + PROJECTION_IS_NORMALIZED=glm5_numerics, num_warps=1, ) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..ad36c875aa 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -15,7 +15,9 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.topk = network_config["index_topk"] self.heads = network_config["index_n_heads"] self.dim = network_config["index_head_dim"] - self.eps = network_config["rms_norm_eps"] + # The indexer's K LayerNorm has its own epsilon in the reference model; + # rms_norm_eps belongs to the surrounding RMSNorm layers. + self.eps = 1e-6 def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 620554bc06..81aec80665 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -428,6 +428,7 @@ def _hc_pre(self, streams, layer_weight, prefix, norm_weight): norm_eps=self.eps_, hc_eps=self.hc_eps, sinkhorn_iters=self.hc_sinkhorn_iters, + glm5_numerics=True, ) def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py index 03ae1f3c33..a71a1f7df6 100644 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 +import math import sys -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest import torch @@ -16,7 +17,19 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_iters, post_multiplier=2.0): +def _hc_pre_reference( + x, + fn, + scale, + base, + streams, + rms_eps, + hc_eps, + sinkhorn_iters, + post_multiplier=2.0, + *, + normalize_before_projection=False, +): """Unfused FP32 reference; returns the merged input before sublayer RMSNorm.""" assert x.ndim == 2 and x.shape[-1] % streams == 0 tokens, flattened_hidden = x.shape @@ -25,7 +38,10 @@ def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_ite x_fp32 = x.float() inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) - mixes = F.linear(x_fp32, fn) * inv_rms + if normalize_before_projection: + mixes = F.linear(x_fp32 * inv_rms, fn) + else: + mixes = F.linear(x_fp32, fn) * inv_rms pre_raw = mixes[:, :streams] post_raw = mixes[:, streams : 2 * streams] @@ -98,3 +114,140 @@ def test_mhc_reports_missing_deepgemm(monkeypatch, missing): with pytest.raises(ImportError, match="deep_gemm"): hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20) + + +@pytest.mark.parametrize("tokens,hidden,prefix", [(1, 4096, "attn"), (19, 780, "ffn"), (3073, 4096, "attn")]) +def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix, monkeypatch): + from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + streams = 4 + x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + merged, expected_mix, expected_post = _hc_pre_reference( + x, fn, scale, base, streams, 1e-5, 1e-6, 20, normalize_before_projection=True + ) + merged = merged.float() + expected_input = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5) * norm.float()).to(x.dtype) + expected_mix, expected_post = expected_mix.to(x.dtype), expected_post.to(x.dtype) + + # Call the actual model entry point, without constructing attention/MoE weights. + layer = SimpleNamespace(mhc_streams=streams, eps_=1e-5, hc_eps=1e-6, hc_sinkhorn_iters=20) + weights = SimpleNamespace( + **{ + f"hc_{prefix}_{name}": SimpleNamespace(weight=value) + for name, value in (("fn", fn), ("scale", scale), ("base", base)) + } + ) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", True) + actual_input, actual_mix, actual_post = Glm5NextTransformerLayerInfer._hc_pre( + layer, x, weights, prefix, SimpleNamespace(weight=norm) + ) + assert torch.backends.cuda.matmul.allow_tf32 + assert actual_mix.dtype == actual_post.dtype == x.dtype + torch.testing.assert_close(actual_input, expected_input, atol=0.004, rtol=0.008) + torch.testing.assert_close(actual_mix, expected_mix, atol=1e-5, rtol=0.008) + torch.testing.assert_close(actual_post, expected_post, atol=1e-5, rtol=0.008) + layer_out = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) + # The training fused post kernel consumes rounded coefficients, accumulates + # in FP32, and rounds only the final residual output. + torch.testing.assert_close( + hc_post(layer_out, x, actual_mix, actual_post, streams), + _hc_post_reference(layer_out, x, actual_mix, actual_post, streams), + atol=0.004, + rtol=0.008, + ) + + +@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (19, 780)]) +def test_glm_mhc_projection_retains_fp32_inputs(tokens, hidden, monkeypatch): + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) + fn = torch.randn(24, 4 * hidden, device="cuda") * 0.005 + x_fp32 = x.float() + expected = F.linear(x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5), fn) + + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", True) + partials = pre_norm._hc_normalized_projection(x, fn, 1e-5) + assert torch.backends.cuda.matmul.allow_tf32 + torch.testing.assert_close(partials.sum(0), expected, atol=2e-5, rtol=2e-5) + + +def test_glm_mhc_normalizes_before_cancelling_large_products(monkeypatch): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) + monkeypatch.setitem(sys.modules, "deep_gemm", None) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + # All four streams are identical. These finite, exactly representable inputs + # stress the operation order: the unnormalized dot overflows before its + # opposite products cancel; normalize-first has two finite, cancelling terms. + x = torch.full((1, 4 * 17), 2.0**32, device="cuda", dtype=torch.bfloat16) + fn = torch.zeros(24, x.shape[-1], device="cuda") + fn[:, 0], fn[:, 1] = 2.0**110, -(2.0**110) + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5) + expected = F.linear(x_fp32 * inv_rms, fn) + old_order = F.linear(x_fp32, fn) * inv_rms + assert torch.isfinite(expected).all() + assert not torch.isfinite(old_order).all() + torch.testing.assert_close(pre_norm._hc_normalized_projection(x, fn, 1e-5).sum(0), expected, atol=0, rtol=0) + actual = hc_pre_norm( + x, + fn, + torch.ones(3, device="cuda"), + torch.zeros(24, device="cuda"), + torch.ones(17, device="cuda", dtype=x.dtype), + 4, + 1e-5, + 1e-5, + 1e-6, + 20, + glm5_numerics=True, + ) + for value in actual: + assert torch.isfinite(value).all() + + +def test_glm_mhc_rounds_post_coefficients_before_mixing(): + x = torch.full((1, 4 * 17), 1 / 256, device="cuda", dtype=torch.bfloat16) + fn = torch.zeros(24, x.shape[-1], device="cuda") + scale = torch.zeros(3, device="cuda") + base = torch.zeros(24, device="cuda") + base[4:8] = math.log(1.003 / 0.997) + norm = torch.ones(17, device="cuda", dtype=x.dtype) + _, mix, post = hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20, glm5_numerics=True) + assert mix.dtype == post.dtype == x.dtype + torch.testing.assert_close(post, torch.ones_like(post), atol=0, rtol=0) + torch.testing.assert_close(mix, torch.full_like(mix, 0.25), atol=0, rtol=0) + layer_out = torch.ones(1, 17, device="cuda", dtype=x.dtype) + actual = hc_post(layer_out, x, mix, post, 4) + torch.testing.assert_close(actual, torch.ones_like(actual), atol=0, rtol=0) + _, unrounded_mix, unrounded_post = _hc_pre_reference( + x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True + ) + old_result = _hc_post_reference(layer_out, x, unrounded_mix, unrounded_post, 4) + assert (old_result != actual).all() + + +def test_glm_mhc_keeps_pre_coefficients_in_fp32(): + x = torch.zeros(1, 4, 17, device="cuda", dtype=torch.bfloat16) + x[:, 0], x[:, 1] = 1, -1 + x = x.flatten(1) + fn = torch.zeros(24, x.shape[-1], device="cuda") + scale = torch.zeros(3, device="cuda") + base = torch.zeros(24, device="cuda") + # Both coefficients round to BF16 0.5, which would erase this input entirely. + pre_values = torch.tensor([0.5005, 0.5], device="cuda") - 1e-6 + base[:2] = torch.logit(pre_values) + norm = torch.ones(17, device="cuda", dtype=x.dtype) + actual, _, _ = hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20, glm5_numerics=True) + merged, _, _ = _hc_pre_reference( + x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True + ) + merged = merged.float() + expected = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5)).to(x.dtype) + assert (expected > 0.1).all() + torch.testing.assert_close(actual, expected, atol=0.001, rtol=0.008) diff --git a/unit_tests/models/glm5_next/test_precision_alignment.py b/unit_tests/models/glm5_next/test_precision_alignment.py new file mode 100644 index 0000000000..e1b4cc9ab1 --- /dev/null +++ b/unit_tests/models/glm5_next/test_precision_alignment.py @@ -0,0 +1,210 @@ +import dataclasses +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton +from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for name, value in { + "GLOBAL_RANK": 0, + "GLOBAL_WORLD_SIZE": 1, + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_RANK_IN_NODE": 0, + "CURRENT_DEVICE_ID": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + for getter in (get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num): + getter.cache_clear() + yield + for getter in (get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num): + getter.cache_clear() + + +@pytest.fixture +def config(): + return { + "hidden_size": 128, + "n_embed": 128, + "intermediate_size": 128, + "moe_intermediate_size": 128, + "num_hidden_layers": 4, + "n_layer": 4, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-5, + "q_lora_rank": 128, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 0, + "v_head_dim": 128, + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 2048, + "index_kpool": 4, + "index_kpool_always_select_tail": True, + "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4}, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "mhc": False, + "n_routed_experts": 4, + "n_shared_experts": 1, + "first_k_dense_replace": 4, + "num_experts_per_tok": 2, + "norm_topk_prob": True, + "scoring_func": "sigmoid", + "n_group": 1, + "topk_group": 1, + } + + +def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkeypatch): + pytest.importorskip("deep_gemm") + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("FP8 indexer requires Hopper or newer") + # Exercise the real indexer and pooling kernels without requiring vLLM's + # optional top-k extension. The LayerNorm and projections are not stubbed. + monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False) + dim, heads, topk = config["index_head_dim"], config["index_n_heads"], config["index_topk"] + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + indexer = Glm5NextNsaInfer(3, config, 1) + columns = torch.arange(dim) + a = torch.where(columns % 2 == 0, -1.0, 1.0) + b = torch.where((columns // 2) % 2 == 0, -1.0, 1.0) + query = 2 * a + b + prefix = "model.language_model.layers.3.self_attn.indexer" + eye = torch.eye(dim, dtype=torch.bfloat16) + weight.load_hf_weights( + { + f"{prefix}.wk.weight": eye, + f"{prefix}.k_norm.weight": torch.ones(dim, dtype=torch.bfloat16), + f"{prefix}.k_norm.bias": torch.zeros(dim, dtype=torch.bfloat16), + f"{prefix}.wq_b.weight": eye.repeat(heads, 1), + f"{prefix}.weights_proj.weight": ((a + b) / dim).repeat(heads, 1), + f"{prefix}.index_kpool_compress_gate": torch.zeros(dim, dim, dtype=torch.bfloat16), + f"{prefix}.index_kpool_compress_ape": torch.zeros(4, dim), + } + ) + + # At 2052 tokens, 513 complete pools compete for 512 slots. The first + # 511 pools win under either epsilon. The last two have different K + # variances: eps=1e-6 selects A, while the global RMSNorm eps selects B. + hidden = torch.cat( + ( + (query * 2 ** -4).repeat(topk - 4, 1), + (a * 2 ** -10).repeat(4, 1), + (b * 2 ** -4).repeat(4, 1), + ) + ).to(device="cuda", dtype=torch.bfloat16) + tokens = hidden.shape[0] + q_lora = query.to(device="cuda", dtype=torch.bfloat16).repeat(tokens, 1) + projected_k = weight.wk_proj_.mm(hidden) + actual_k = weight.k_norm_(projected_k, eps=indexer.eps) + expected_k = F.layer_norm(projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=1e-6) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + wrong_k = F.layer_norm( + projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=config["rms_norm_eps"] + ) + assert not torch.equal(expected_k, wrong_k) + + storage = torch.zeros(tokens + 1, 1, 584, dtype=torch.bfloat16, device="cuda") + tail = torch.zeros(2, 4, 2 * dim, dtype=torch.bfloat16, device="cuda") + locations = torch.arange(tokens, dtype=torch.int32, device="cuda") + infer = SimpleNamespace( + is_prefill=True, + max_kv_seq_len=tokens, + max_q_seq_len=tokens, + b_req_idx=torch.zeros(1, dtype=torch.int32, device="cuda"), + b_seq_len=torch.tensor([tokens], dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), + req_manager=SimpleNamespace( + req_to_token_indexs=locations[None], + get_indexer_tail_buffer=lambda _: tail, + ), + mem_manager=SimpleNamespace(get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:]), + ) + state = SimpleNamespace(lengths=locations + 1, ks=torch.zeros_like(locations), ragged_mem_index=locations) + _, actual = indexer._get_indices(hidden, q_lora, infer, state, weight) + torch.testing.assert_close(actual[-1, :topk].sort().values, locations[:topk], atol=0, rtol=0) + assert (actual[-1, topk:] == -1).all() + + # Verify this fixture detects the old behavior, including through the FP8 + # indexer path, rather than only checking a constructor constant. + indexer.eps = config["rms_norm_eps"] + _, wrong = indexer._get_indices(hidden, q_lora, infer, state, weight) + wrong_selection = torch.cat((locations[: topk - 4], locations[topk:])) + torch.testing.assert_close(wrong[-1, :topk].sort().values, wrong_selection, atol=0, rtol=0) + + +@pytest.mark.parametrize("enable_ep_moe", [False, True], ids=["tp", "ep"]) +def test_checkpoint_correction_bias_preserves_fp32_routing(config, monkeypatch, enable_ep_moe): + args = StartArgs(tp=2, enable_ep_moe=enable_ep_moe, enable_fused_shared_experts=False) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + # Construct rank 1's real weights on this GPU; routing itself needs no + # distributed communication and must retain all experts' global biases. + for name, value in {"GLOBAL_RANK": 1, "GLOBAL_WORLD_SIZE": 2, "DP_WORLD_SIZE": 2, "CURRENT_RANK_IN_DP": 1}.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + get_env_start_args.cache_clear() + config.update(n_routed_experts=288, num_experts_per_tok=8, first_k_dense_replace=1) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + experts = weight.experts + expected_impl = FuseMoeDeepGEMM if enable_ep_moe else FuseMoeTriton + assert isinstance(experts.fuse_moe_impl, expected_impl) + source_bias = torch.zeros(config["n_routed_experts"], dtype=torch.float32) + source_bias[7:9] = torch.tensor([0.1002, 0.1001], dtype=torch.float32) + weight.load_hf_weights({"model.language_model.layers.3.mlp.gate.e_score_correction_bias": source_bias}) + assert experts.e_score_correction_bias.dtype == torch.float32 + assert experts.e_score_correction_bias.load_ok + torch.testing.assert_close(experts.e_score_correction_bias.cpu(), source_bias, atol=0, rtol=0) + assert experts.w13.weight.dtype == experts.w2.weight.dtype == torch.bfloat16 + + hidden = torch.zeros(1, config["hidden_size"], dtype=torch.bfloat16, device="cuda") + logits = torch.full((1, config["n_routed_experts"]), -10.0, dtype=torch.float32, device="cuda") + logits[:, :7] = 4.0 + logits[0, 7:9] = torch.tensor([0.0, 4e-5], dtype=torch.float32, device="cuda") + + def route(bias): + # This is the actual common selection entry used by TP expert GEMMs, + # EP low-latency dispatch, and EP select_experts_and_quant_input. + return experts.fuse_moe_impl._select_experts( + input_tensor=hidden, + router_logits=logits, + correction_bias=bias, + top_k=8, + renormalize=True, + use_grouped_topk=True, + topk_group=1, + num_expert_group=1, + scoring_func="sigmoid", + ) + + weights, ids, origin_ids = route(experts.e_score_correction_bias) + expected_ids = torch.arange(8, dtype=ids.dtype, device="cuda")[None] + torch.testing.assert_close(ids.sort(dim=-1).values, expected_ids, atol=0, rtol=0) + torch.testing.assert_close(origin_ids, ids, atol=0, rtol=0) + # Bias affects selection only; combination weights use the original + # sigmoid scores and then normalize over the selected experts. + expected_weights = logits.sigmoid().gather(1, ids.long()) + expected_weights /= expected_weights.sum(-1, keepdim=True) + torch.testing.assert_close(weights, expected_weights, atol=1e-7, rtol=1e-6) + + # Both biases round to the same BF16 value. Reproducing the old storage + # conversion therefore lets expert 8 displace expert 7 at the boundary. + _, rounded_ids, _ = route(experts.e_score_correction_bias.to(torch.bfloat16)) + rounded_expected = torch.tensor([list(range(7)) + [8]], dtype=rounded_ids.dtype, device="cuda") + torch.testing.assert_close(rounded_ids.sort(dim=-1).values, rounded_expected, atol=0, rtol=0) From 92d580e1c7704e0f4e3cb99fbb4f20a9e5114aa1 Mon Sep 17 00:00:00 2001 From: baishihao Date: Fri, 23 Oct 2026 00:55:52 +0800 Subject: [PATCH 07/13] fix(glm5): match training post rounding and close precision tests Round the GLM residual mix to the activation dtype before post expansion, matching the default compiled native training path. Preserve the existing fused path for other models. Account for two BF16 rounding steps after FP32 reduction differences, add a strict 3073-token FP32 projection check and another seed, and retain negative controls. Trim redundant test work and comments. Validation: all 24 mHC and GLM precision cases passed on H200 in bsh_dsv4; syntax and whitespace checks passed. --- .../fused_moe/fused_moe_weight.py | 3 +- .../basemodel/triton_kernel/mhc/post.py | 25 +++++++-- .../basemodel/triton_kernel/mhc/pre_norm.py | 6 +- lightllm/models/glm5_next/indexer.py | 3 +- .../layer_infer/transformer_layer_infer.py | 4 +- .../basemodel/triton_kernel/mhc/test_mhc.py | 56 +++++++++---------- .../glm5_next/test_precision_alignment.py | 20 ++----- 7 files changed, 60 insertions(+), 57 deletions(-) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index f414dbeb44..5652279702 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -306,8 +306,7 @@ def _create_weight(self): self.per_expert_scale = None # Create e_score_correction_bias if self.e_score_correction_bias_name: - # Routing compares close corrected scores in FP32. Preserve the - # checkpoint bias instead of rounding it to the expert weight dtype. + # Preserve FP32 correction scores at expert-selection boundaries. self.e_score_correction_bias = torch.empty( (self.n_routed_experts,), dtype=torch.float32, diff --git a/lightllm/common/basemodel/triton_kernel/mhc/post.py b/lightllm/common/basemodel/triton_kernel/mhc/post.py index 31ecca43b9..2b798018f7 100644 --- a/lightllm/common/basemodel/triton_kernel/mhc/post.py +++ b/lightllm/common/basemodel/triton_kernel/mhc/post.py @@ -23,6 +23,7 @@ def _hc_post_4stream_kernel( post_stride_m: tl.constexpr, out_stride_m: tl.constexpr, BLOCK_H: tl.constexpr, + GLM5_NUMERICS: tl.constexpr, ): """将子层输出和原始残差混合为四条新 stream：[T, H] + [T, 4, H] -> [T, 4, H]。 @@ -50,10 +51,10 @@ def _hc_post_4stream_kernel( post_base = post_mix + token * post_stride_m mix_base = residual_mix + token * mix_stride_m - accumulator_0 = layer_value * tl.load(post_base) - accumulator_1 = layer_value * tl.load(post_base + 1) - accumulator_2 = layer_value * tl.load(post_base + 2) - accumulator_3 = layer_value * tl.load(post_base + 3) + accumulator_0 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base) + accumulator_1 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 1) + accumulator_2 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 2) + accumulator_3 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 3) accumulator_0 += residual_0 * tl.load(mix_base) accumulator_0 += residual_1 * tl.load(mix_base + 4) @@ -72,6 +73,18 @@ def _hc_post_4stream_kernel( accumulator_3 += residual_2 * tl.load(mix_base + 11) accumulator_3 += residual_3 * tl.load(mix_base + 15) + if GLM5_NUMERICS: + # Megatron's compiled native post rounds the residual bmm to BF16, + # then fuses the post expansion and addition into the final BF16 output. + accumulator_0 = accumulator_0.to(residual.dtype.element_ty).to(tl.float32) + accumulator_1 = accumulator_1.to(residual.dtype.element_ty).to(tl.float32) + accumulator_2 = accumulator_2.to(residual.dtype.element_ty).to(tl.float32) + accumulator_3 = accumulator_3.to(residual.dtype.element_ty).to(tl.float32) + accumulator_0 += layer_value * tl.load(post_base) + accumulator_1 += layer_value * tl.load(post_base + 1) + accumulator_2 += layer_value * tl.load(post_base + 2) + accumulator_3 += layer_value * tl.load(post_base + 3) + output_base = output + token * out_stride_m + hidden_offsets tl.store(output_base, accumulator_0, mask=hidden_mask) tl.store(output_base + hidden, accumulator_1, mask=hidden_mask) @@ -85,6 +98,8 @@ def hc_post( residual_mix: torch.Tensor, post_mix: torch.Tensor, streams: int, + *, + glm5_numerics: bool = False, ) -> torch.Tensor: """用一次 Triton launch 将子层输出混回残差 streams。 @@ -96,6 +111,7 @@ def hc_post( residual_mix: [T, S, S]，fp32 或激活 dtype；来自 hc_pre_norm。 post_mix: [T, S]，与 residual_mix 相同 dtype；来自同一次 pre 调用。 streams: 残差 stream 数 S；当前 Triton 入口要求 S = 4。 + glm5_numerics: 匹配 GLM 训练默认 native post，在加子层输出前将残差混合结果舍入到激活 dtype。 Returns: 新残差 [T, S * H]，dtype 与 layer_output 相同；输入张量保持不变。 @@ -127,6 +143,7 @@ def hc_post( post_stride_m=post_mix.stride(0), out_stride_m=output.stride(0), BLOCK_H=block_h, + GLM5_NUMERICS=glm5_numerics, num_warps=8, ) return output diff --git a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py index d43a7ecaf8..3191ce0716 100644 --- a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py +++ b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py @@ -60,8 +60,7 @@ def _hc_normalized_projection_kernel( (mix_offsets[None, :] < MIX_SIZE) & (k_offsets[:, None] < FLATTENED_HIDDEN), 0, ) - # GLM projects the FP32-normalized input. Moving the RMS factor after - # the dot, or rounding this tile to TF32/BF16, changes that computation. + # Normalize in FP32 before the IEEE projection, independent of TF32 settings. normalized = values * rms[:, None] accumulator = tl.dot(normalized, weights, accumulator, input_precision="ieee") k_offsets += BLOCK_K @@ -265,8 +264,7 @@ def hc_pre_norm( hc_eps: pre sigmoid 后的偏移量及 Sinkhorn 的稳定项。 sinkhorn_iters: Sinkhorn 迭代次数。 post_multiplier: post sigmoid 的乘数，默认 2.0。 - glm5_numerics: GLM-5.3 的 FP32 先归一化后投影路径；post/residual 系数 - 舍入到激活 dtype，pre 系数仍保留 fp32。 + glm5_numerics: FP32 先归一化后投影；pre 系数保持 fp32，post/residual 舍入到激活 dtype。 Returns: (layer_input, residual_mix, post_mix)：形状为 [T, H]、[T, S, S]、[T, S]。 diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index ad36c875aa..30a9883881 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -15,8 +15,7 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.topk = network_config["index_topk"] self.heads = network_config["index_n_heads"] self.dim = network_config["index_head_dim"] - # The indexer's K LayerNorm has its own epsilon in the reference model; - # rms_norm_eps belongs to the surrounding RMSNorm layers. + # The indexer LayerNorm has a separate epsilon from the model RMSNorm. self.eps = 1e-6 def select_topk_indices(self, logits, lengths, indices): diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 81aec80665..2032e50c96 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -439,11 +439,11 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) else: layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight) - streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams, glm5_numerics=True) layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_) layer_output = self._ffn(layer_input, infer_state, layer_weight) - streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams) + streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams, glm5_numerics=True) # Only prefill autotuning truncates the model. Decode autotuning still # executes every layer and must keep all residual streams until the end. is_autotune_last_layer = ( diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py index a71a1f7df6..32d2b694ee 100644 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -60,11 +60,13 @@ def _hc_pre_reference( return layer_input, residual_mix, post -def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams): +def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams, *, round_residual=False): """Unfused FP32 reference; residual_mix maps input streams to output streams.""" tokens, hidden = layer_output.shape residual_3d = residual.view(tokens, streams, hidden) mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) + if round_residual: + mixed_residual = mixed_residual.to(layer_output.dtype).float() out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual return out.to(layer_output.dtype).reshape(tokens, streams * hidden) @@ -117,9 +119,11 @@ def test_mhc_reports_missing_deepgemm(monkeypatch, missing): @pytest.mark.parametrize("tokens,hidden,prefix", [(1, 4096, "attn"), (19, 780, "ffn"), (3073, 4096, "attn")]) -def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix, monkeypatch): +@pytest.mark.parametrize("seed", [1525, 2026]) +def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix, seed, monkeypatch): from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + torch.manual_seed(seed) monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) streams = 4 @@ -135,7 +139,6 @@ def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix expected_input = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5) * norm.float()).to(x.dtype) expected_mix, expected_post = expected_mix.to(x.dtype), expected_post.to(x.dtype) - # Call the actual model entry point, without constructing attention/MoE weights. layer = SimpleNamespace(mhc_streams=streams, eps_=1e-5, hc_eps=1e-6, hc_sinkhorn_iters=20) weights = SimpleNamespace( **{ @@ -149,21 +152,21 @@ def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix ) assert torch.backends.cuda.matmul.allow_tf32 assert actual_mix.dtype == actual_post.dtype == x.dtype - torch.testing.assert_close(actual_input, expected_input, atol=0.004, rtol=0.008) + # FP32 reduction differences can cross the merged BF16 midpoint; RMSNorm + # rounds to BF16 again. Check FP32 projection separately at strict tolerance. + torch.testing.assert_close(actual_input, expected_input, atol=0.004, rtol=2 * torch.finfo(x.dtype).eps) torch.testing.assert_close(actual_mix, expected_mix, atol=1e-5, rtol=0.008) torch.testing.assert_close(actual_post, expected_post, atol=1e-5, rtol=0.008) layer_out = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) - # The training fused post kernel consumes rounded coefficients, accumulates - # in FP32, and rounds only the final residual output. torch.testing.assert_close( - hc_post(layer_out, x, actual_mix, actual_post, streams), - _hc_post_reference(layer_out, x, actual_mix, actual_post, streams), + hc_post(layer_out, x, actual_mix, actual_post, streams, glm5_numerics=True), + _hc_post_reference(layer_out, x, actual_mix, actual_post, streams, round_residual=True), atol=0.004, rtol=0.008, ) -@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (19, 780)]) +@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (19, 780), (3073, 4096)]) def test_glm_mhc_projection_retains_fp32_inputs(tokens, hidden, monkeypatch): monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) @@ -181,9 +184,7 @@ def test_glm_mhc_normalizes_before_cancelling_large_products(monkeypatch): monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) monkeypatch.setitem(sys.modules, "deep_gemm", None) monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) - # All four streams are identical. These finite, exactly representable inputs - # stress the operation order: the unnormalized dot overflows before its - # opposite products cancel; normalize-first has two finite, cancelling terms. + # Project-first overflows; normalize-first has finite cancelling products. x = torch.full((1, 4 * 17), 2.0**32, device="cuda", dtype=torch.bfloat16) fn = torch.zeros(24, x.shape[-1], device="cuda") fn[:, 0], fn[:, 1] = 2.0**110, -(2.0**110) @@ -194,21 +195,6 @@ def test_glm_mhc_normalizes_before_cancelling_large_products(monkeypatch): assert torch.isfinite(expected).all() assert not torch.isfinite(old_order).all() torch.testing.assert_close(pre_norm._hc_normalized_projection(x, fn, 1e-5).sum(0), expected, atol=0, rtol=0) - actual = hc_pre_norm( - x, - fn, - torch.ones(3, device="cuda"), - torch.zeros(24, device="cuda"), - torch.ones(17, device="cuda", dtype=x.dtype), - 4, - 1e-5, - 1e-5, - 1e-6, - 20, - glm5_numerics=True, - ) - for value in actual: - assert torch.isfinite(value).all() def test_glm_mhc_rounds_post_coefficients_before_mixing(): @@ -223,7 +209,7 @@ def test_glm_mhc_rounds_post_coefficients_before_mixing(): torch.testing.assert_close(post, torch.ones_like(post), atol=0, rtol=0) torch.testing.assert_close(mix, torch.full_like(mix, 0.25), atol=0, rtol=0) layer_out = torch.ones(1, 17, device="cuda", dtype=x.dtype) - actual = hc_post(layer_out, x, mix, post, 4) + actual = hc_post(layer_out, x, mix, post, 4, glm5_numerics=True) torch.testing.assert_close(actual, torch.ones_like(actual), atol=0, rtol=0) _, unrounded_mix, unrounded_post = _hc_pre_reference( x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True @@ -232,6 +218,20 @@ def test_glm_mhc_rounds_post_coefficients_before_mixing(): assert (old_result != actual).all() +def test_glm_mhc_post_rounds_residual_like_training(): + x = torch.zeros(1, 4, 17, device="cuda", dtype=torch.bfloat16) + x[:, 0], x[:, 1] = 1, 1 / 256 + mix = torch.full((1, 4, 4), 0.25, device="cuda", dtype=x.dtype) + post = torch.ones(1, 4, device="cuda", dtype=x.dtype) + layer_out = torch.full((1, 17), 3 / 512, device="cuda", dtype=x.dtype) + # The native training bmm rounds 0.2509765625 to 0.25 before expansion. + mixed = torch.bmm(mix.transpose(1, 2), x) + expected = torch.addcmul(mixed.float(), layer_out[:, None].float(), post[:, :, None].float()).to(x.dtype) + actual = hc_post(layer_out, x.flatten(1), mix, post, 4, glm5_numerics=True) + torch.testing.assert_close(actual, expected.flatten(1), atol=0, rtol=0) + assert (hc_post(layer_out, x.flatten(1), mix, post, 4) != actual).all() + + def test_glm_mhc_keeps_pre_coefficients_in_fp32(): x = torch.zeros(1, 4, 17, device="cuda", dtype=torch.bfloat16) x[:, 0], x[:, 1] = 1, -1 diff --git a/unit_tests/models/glm5_next/test_precision_alignment.py b/unit_tests/models/glm5_next/test_precision_alignment.py index e1b4cc9ab1..746397c385 100644 --- a/unit_tests/models/glm5_next/test_precision_alignment.py +++ b/unit_tests/models/glm5_next/test_precision_alignment.py @@ -77,8 +77,6 @@ def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkey pytest.importorskip("deep_gemm") if torch.cuda.get_device_capability()[0] < 9: pytest.skip("FP8 indexer requires Hopper or newer") - # Exercise the real indexer and pooling kernels without requiring vLLM's - # optional top-k extension. The LayerNorm and projections are not stubbed. monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False) dim, heads, topk = config["index_head_dim"], config["index_n_heads"], config["index_topk"] weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) @@ -101,9 +99,7 @@ def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkey } ) - # At 2052 tokens, 513 complete pools compete for 512 slots. The first - # 511 pools win under either epsilon. The last two have different K - # variances: eps=1e-6 selects A, while the global RMSNorm eps selects B. + # 513 pools compete for 512 slots: epsilon changes the final selected pool. hidden = torch.cat( ( (query * 2 ** -4).repeat(topk - 4, 1), @@ -143,8 +139,7 @@ def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkey torch.testing.assert_close(actual[-1, :topk].sort().values, locations[:topk], atol=0, rtol=0) assert (actual[-1, topk:] == -1).all() - # Verify this fixture detects the old behavior, including through the FP8 - # indexer path, rather than only checking a constructor constant. + # Negative control: the previous epsilon selects the other boundary pool. indexer.eps = config["rms_norm_eps"] _, wrong = indexer._get_indices(hidden, q_lora, infer, state, weight) wrong_selection = torch.cat((locations[: topk - 4], locations[topk:])) @@ -155,8 +150,7 @@ def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkey def test_checkpoint_correction_bias_preserves_fp32_routing(config, monkeypatch, enable_ep_moe): args = StartArgs(tp=2, enable_ep_moe=enable_ep_moe, enable_fused_shared_experts=False) monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - # Construct rank 1's real weights on this GPU; routing itself needs no - # distributed communication and must retain all experts' global biases. + # Rank 1 must retain every expert's global correction bias. for name, value in {"GLOBAL_RANK": 1, "GLOBAL_WORLD_SIZE": 2, "DP_WORLD_SIZE": 2, "CURRENT_RANK_IN_DP": 1}.items(): monkeypatch.setenv("LIGHTLLM_" + name, str(value)) get_env_start_args.cache_clear() @@ -179,8 +173,6 @@ def test_checkpoint_correction_bias_preserves_fp32_routing(config, monkeypatch, logits[0, 7:9] = torch.tensor([0.0, 4e-5], dtype=torch.float32, device="cuda") def route(bias): - # This is the actual common selection entry used by TP expert GEMMs, - # EP low-latency dispatch, and EP select_experts_and_quant_input. return experts.fuse_moe_impl._select_experts( input_tensor=hidden, router_logits=logits, @@ -197,14 +189,12 @@ def route(bias): expected_ids = torch.arange(8, dtype=ids.dtype, device="cuda")[None] torch.testing.assert_close(ids.sort(dim=-1).values, expected_ids, atol=0, rtol=0) torch.testing.assert_close(origin_ids, ids, atol=0, rtol=0) - # Bias affects selection only; combination weights use the original - # sigmoid scores and then normalize over the selected experts. + # Correction bias affects selection; weights use the original sigmoid scores. expected_weights = logits.sigmoid().gather(1, ids.long()) expected_weights /= expected_weights.sum(-1, keepdim=True) torch.testing.assert_close(weights, expected_weights, atol=1e-7, rtol=1e-6) - # Both biases round to the same BF16 value. Reproducing the old storage - # conversion therefore lets expert 8 displace expert 7 at the boundary. + # Negative control: BF16 rounding makes expert 8 displace expert 7. _, rounded_ids, _ = route(experts.e_score_correction_bias.to(torch.bfloat16)) rounded_expected = torch.tensor([list(range(7)) + [8]], dtype=rounded_ids.dtype, device="cuda") torch.testing.assert_close(rounded_ids.sort(dim=-1).values, rounded_expected, atol=0, rtol=0) From aee915e10f5305896a718509399bd8890443d382 Mon Sep 17 00:00:00 2001 From: baishihao Date: Fri, 23 Oct 2026 21:27:29 +0800 Subject: [PATCH 08/13] fix(moe): honor UE8M0 scales in Triton expert activations --- .../fused_moe/grouped_fused_moe.py | 7 ++- .../test_ue8m0_fp8_quantization.py | 49 ++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index 90eeb5a4ca..64537ac606 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -25,6 +25,7 @@ from lightllm.utils.log_utils import init_logger from lightllm.utils.vllm_utils import vllm_ops from lightllm.utils.device_utils import triton_support_tensor_descriptor +from lightllm.utils.envs_utils import get_env_use_ue8m0_scales from .moe_silu_and_mul import silu_and_mul_fwd from .moe_sum_reduce import moe_sum_reduce from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 @@ -878,8 +879,12 @@ def grouped_matmul( # input 使用 per group quant 量化 _m, _k = token_inputs.shape assert _k % block_size_k == 0 + # Triton consumes FP32 activation scales even when weight scales are packed UE8M0. token_inputs, token_input_scale = per_token_group_quant_fp8( - token_inputs, block_size_k, dtype=expert_weights.dtype + token_inputs, + block_size_k, + dtype=expert_weights.dtype, + use_ue8m0_scales=get_env_use_ue8m0_scales() or use_packed_ue8m0, ) if reused_mblock_infos is None: diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py index cf87ac3ffb..2abf8ccd58 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -266,10 +266,56 @@ def test_weight_scale_loading_preserves_fused_views( assert len(packed_weight_layout) == (2 * (experts or 1) if packed else 0) +@pytest.mark.parametrize("use_ue8m0", [False, True]) +@pytest.mark.parametrize("shared_expert", [False, True]) +def test_tp_moe_activation_scale_rule(monkeypatch, use_ue8m0, shared_expert): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe as moe + from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import append_fused_shared_experts + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") + Quantcfg( + { + "n_layer": 1, + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0" if use_ue8m0 else "float32", + }, + } + ) + torch.manual_seed(20261023) + x = torch.randn((17, 512), device="cuda", dtype=torch.bfloat16) + w1, s1 = weight.weight_quant(torch.randn((3, 512, 512), device="cuda", dtype=torch.bfloat16) * 0.02) + w2, s2 = weight.weight_quant(torch.randn((3, 512, 256), device="cuda", dtype=torch.bfloat16) * 0.02) + ids = torch.tensor([0, 1], device="cuda", dtype=torch.int32).repeat(17, 1) + scores = torch.tensor([0.3, 0.7], device="cuda").repeat(17, 1) + if shared_expert: + scores, ids = append_fused_shared_experts(scores, ids, shared_expert_start_id=2, num_fused_shared_experts=1) + calls = [] + quantize = moe.per_token_group_quant_fp8 + + def check_activation(values, group_size, **kwargs): + q, scales = quantize(values, group_size, **kwargs) + assert group_size == 128 + assert scales.dtype == torch.float32 + expected = values.float().reshape(values.shape[0], -1, 128).abs().amax(-1).clamp_min(1e-10) / 448.0 + if use_ue8m0: + expected = torch.exp2(torch.ceil(torch.log2(expected.double()))).float() + torch.testing.assert_close(scales, expected, rtol=1e-6, atol=0) + calls.append(values.shape[1]) + return q, scales + + monkeypatch.setattr(moe, "per_token_group_quant_fp8", check_activation) + out = moe.fused_experts_impl(x, w1, w2, scores, ids, use_fp8_w8a8=True, w1_scale=s1, w2_scale=s2) + assert calls == [512, 256] # W13 and W2, including the optional shared-expert slot. + assert torch.isfinite(out).all() + + @pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("packed_env", [False, True]) @pytest.mark.parametrize("need_trans", [False, True]) @pytest.mark.parametrize("k", [512, 640]) -def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k): +def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k, packed_env): from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import grouped_matmul monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") @@ -287,6 +333,7 @@ def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, indices[1, :8] = torch.arange(9, 17, device="cuda") outputs = [] for sf in (scales, full.weight_scale): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1" if sf.dtype != torch.int32 or packed_env else "0") out = torch.empty((17, 256), device="cuda", dtype=torch.bfloat16) grouped_matmul( 17, From 3e7a5213cae0aeb666b07d87a8388cffbb38412e Mon Sep 17 00:00:00 2001 From: baishihao Date: Fri, 23 Oct 2026 23:46:51 +0800 Subject: [PATCH 09/13] fix(glm5): align eager SwiGLU, W2 routing and mHC post with training --- .../fused_moe/fused_moe_weight.py | 2 + .../meta_weights/fused_moe/impl/base_impl.py | 2 + .../fused_moe/impl/deepgemm_impl.py | 1 + .../fused_moe/impl/triton_impl.py | 1 + .../deepep_expanded_layout_kernels.py | 9 +- .../fused_moe/grouped_fused_moe.py | 39 ++- .../fused_moe/grouped_fused_moe_ep.py | 43 ++- .../fused_moe/moe_route_weight.py | 120 +++++++ .../fused_moe/moe_silu_and_mul.py | 11 +- .../moe_silu_and_mul_mix_quant_ep.py | 29 +- .../basemodel/triton_kernel/mhc/post.py | 15 +- .../layer_infer/transformer_layer_infer.py | 7 +- .../layer_weights/transformer_layer_weight.py | 1 + .../basemodel/triton_kernel/mhc/test_mhc.py | 21 +- .../fused_moe/test_deepgemm_activation.py | 18 +- .../test_glm53_training_alignment.py | 323 ++++++++++++++++++ unit_tests/models/glm5_next/test_layers.py | 16 +- 17 files changed, 621 insertions(+), 37 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py create mode 100644 unit_tests/common/fused_moe/test_glm53_training_alignment.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 5652279702..97aebbd3fc 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -35,6 +35,7 @@ def __init__( layer_num: int = 0, network_config: Dict[str, Any] = None, per_expert_scale_name: str = "", + glm5_numerics: bool = False, ) -> None: super().__init__(data_type=data_type) self.w1_weight_name = gate_proj_name @@ -67,6 +68,7 @@ def __init__( redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, + glm5_numerics=glm5_numerics, ) self.lock = threading.Lock() self._create_weight() diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index e242870d69..7a1c1219bb 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -22,7 +22,9 @@ def __init__( redundancy_expert_ids_tensor: torch.Tensor, routed_expert_counter_tensor: torch.Tensor, auto_update_redundancy_expert: bool, + glm5_numerics: bool = False, ): + self.glm5_numerics = glm5_numerics self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self.routed_scaling_factor = routed_scaling_factor diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 74e457426b..3f85309457 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -94,6 +94,7 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=self.glm5_numerics, ) return output diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index f748aea467..ccafa36292 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -87,6 +87,7 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=self.glm5_numerics, ) return input_tensor diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py index ab54ea5b79..21f3f6439c 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py @@ -104,7 +104,9 @@ def _ep_zero_padding_kernel( scale_offsets = tl.arange(0, BLOCK_SCALE_K) scale_mask = scale_offsets < scale_hidden_size scale_ptrs = ( - recv_x_scale + row_offsets[:, None] * recv_x_scale_stride_m + scale_offsets[None, :] * recv_x_scale_stride_k + recv_x_scale + + row_offsets[:, None] * recv_x_scale_stride_m + + scale_offsets[None, :] * recv_x_scale_stride_k ) tl.store(scale_ptrs, 0.0, mask=row_mask[:, None] & scale_mask[None, :]) tl.store(recv_topk_weights + row_offsets, 0.0, mask=row_mask) @@ -173,6 +175,7 @@ def _ep_gather_chunk_kernel( TOPK: tl.constexpr, BLOCK_D: tl.constexpr, NEED_HIDDEN_MASK: tl.constexpr, + APPLY_ROUTED_WEIGHT: tl.constexpr, ): hidden_block_id = tl.program_id(0) start_recv_token_id = tl.program_id(1) @@ -196,7 +199,7 @@ def _ep_gather_chunk_kernel( value = tl.load(chunk_ptrs, mask=hidden_mask, other=0.0) else: value = tl.load(chunk_ptrs) - weight = tl.load(weights + slot) + weight = tl.load(weights + slot) if APPLY_ROUTED_WEIGHT else 1.0 accumulator += value.to(tl.float32) * weight if NEED_HIDDEN_MASK: tl.store(output_ptrs, accumulator, mask=hidden_mask) @@ -211,6 +214,7 @@ def ep_gather_chunk( weights: torch.Tensor, # [num_expanded_tokens] recv_src_metadata: torch.Tensor, # [num_recv_tokens, topk + 2] output: torch.Tensor, # [num_recv_tokens, hidden_size] + apply_routed_weight: bool = True, ): """Accumulate one expanded W2-output chunk into dense receive-token rows. @@ -240,6 +244,7 @@ def ep_gather_chunk( output_stride_k=output.stride(1), hidden_size=hidden_size, TOPK=topk, + APPLY_ROUTED_WEIGHT=apply_routed_weight, BLOCK_D=block_d, # 常见整块 hidden size 保持原来的无 mask kernel；仅尾块不完整时启用边界保护。 NEED_HIDDEN_MASK=hidden_size % block_d != 0, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index 64537ac606..794ad76d94 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -27,6 +27,7 @@ from lightllm.utils.device_utils import triton_support_tensor_descriptor from lightllm.utils.envs_utils import get_env_use_ue8m0_scales from .moe_silu_and_mul import silu_and_mul_fwd +from .moe_route_weight import multiply_routing_weights, sorted_tp_routing_weights from .moe_sum_reduce import moe_sum_reduce from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 from lightllm.utils.torch_ops_utils import direct_register_custom_op @@ -392,7 +393,9 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t else: # Expert rows may be unordered, but grouped matmul reuses this same # mapping for up/down projections and writes back to original topk slots. - _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights) + _moe_align_fused_atomic_token( + expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights + ) return expert_to_token_index, expert_to_weight, expert_token_num @@ -581,7 +584,7 @@ def grouped_matmul_kernel( assert BLOCK_SIZE_K <= block_size_k token_scale_stride0 = token_stride_0 // block_size_k if TOKEN_INPUT_USE_TMA: - assert MUL_ROUTED_WEIGHT is True + assert topk_num == 1 a_scale_ptrs = token_scale_ptr + (token_start_index + tl.arange(0, BLOCK_SIZE_M)) * token_scale_stride0 else: a_scale_ptrs = token_scale_ptr + (a_m_index // topk_num) * token_scale_stride0 @@ -596,7 +599,7 @@ def grouped_matmul_kernel( else: # per token scale quant if TOKEN_INPUT_USE_TMA: - assert MUL_ROUTED_WEIGHT is True + assert topk_num == 1 a_scale_ptrs = token_scale_ptr + (token_start_index + tl.arange(0, BLOCK_SIZE_M))[:, None] else: a_scale_ptrs = token_scale_ptr + (a_m_index // topk_num)[:, None] @@ -797,6 +800,7 @@ def grouped_matmul( reused_mblock_infos=None, run_config: Optional[dict] = None, bias: Optional[torch.Tensor] = None, # per-expert bias [expert_num, N] + apply_routed_weight: bool = True, ): """ token_num_mul_topk_num is int equal token_num * topk_num, @@ -905,7 +909,8 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]): triton.set_allocator(alloc_fn) - # moe 分为 up 和 down 两次计算，当 mul_routed_weight 为 False 的时候为 up + # mul_routed_weight selects the W2 layout. GLM can separately disable + # its output weighting after applying p to the BF16 input before quantization. is_up_moe = not mul_routed_weight if is_up_moe: @@ -989,7 +994,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]): BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=GROUP_SIZE_M, - MUL_ROUTED_WEIGHT=mul_routed_weight, + MUL_ROUTED_WEIGHT=mul_routed_weight and apply_routed_weight, NEED_K_MASK=NEED_K_MASK, NEED_TRANS=NEED_TRANS, num_warps=num_warps, @@ -1024,6 +1029,7 @@ def fused_experts_impl( limit=None, alpha=None, clamp_up_add_one=True, + glm5_numerics=False, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -1103,9 +1109,18 @@ def fused_experts_impl( limit=limit, alpha=alpha, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, layout=layout, ) + if glm5_numerics: + p = curr_topk_weights.reshape(-1) + if triton_support_tensor_descriptor(): + p = sorted_tp_routing_weights( + curr_topk_weights, expert_to_weights, expert_to_token_num, reused_mblock_infos + ) + multiply_routing_weights(intermediate_cache2.view(-1, N // 2), p) + grouped_matmul( curr_topk_ids.numel(), intermediate_cache2.view(-1, N // 2), @@ -1123,6 +1138,7 @@ def fused_experts_impl( reused_mblock_infos=reused_mblock_infos, run_config=run_config, bias=w2_bias, + apply_routed_weight=not glm5_numerics, ) moe_sum_reduce( @@ -1150,6 +1166,7 @@ def inplace_fused_experts_impl( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, ) -> None: fused_experts_impl( hidden_states, @@ -1170,6 +1187,7 @@ def inplace_fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) @@ -1192,6 +1210,7 @@ def inplace_fused_experts_impl_fake( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, ) -> None: pass @@ -1223,7 +1242,8 @@ def outplace_fused_experts_impl( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, -) -> None: + glm5_numerics: bool = False, +) -> torch.Tensor: return fused_experts_impl( hidden_states, w1, @@ -1243,6 +1263,7 @@ def outplace_fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) @@ -1265,7 +1286,8 @@ def outplace_fused_experts_impl_fake( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, -) -> None: + glm5_numerics: bool = False, +) -> torch.Tensor: return torch.empty_like(hidden_states) @@ -1297,6 +1319,7 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, ): if inplace: torch.ops.lightllm.inplace_fused_experts_impl( @@ -1317,6 +1340,7 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) return hidden_states else: @@ -1338,4 +1362,5 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 61cd81556a..5047c67900 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -7,6 +7,7 @@ from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from .moe_route_weight import multiply_routing_weights, gather_low_latency_routing_weights from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( silu_and_mul_masked_post_quant_fwd, ) @@ -79,6 +80,8 @@ def masked_group_gemm( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, + routing_weights: Optional[torch.Tensor] = None, ): padded_m = recv_x[0].shape[1] E, N, _ = w1.shape @@ -99,6 +102,8 @@ def masked_group_gemm( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, + routing_weights=routing_weights, ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) @@ -221,6 +226,7 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) check_ep_expert_dtype(quant_method) @@ -249,6 +255,7 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) @@ -271,6 +278,7 @@ def fused_experts_impl( limit: Optional[float] = None, clamp_up_add_one: bool = True, block_size_k: int = 128, + glm5_numerics: bool = False, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -340,6 +348,7 @@ def fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) else: gather_out = torch.empty( @@ -361,6 +370,7 @@ def fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) _gemm_out_a, _silu_out = None, None del recv_x @@ -387,6 +397,9 @@ def fused_experts_impl( async_finish=False, return_recv_hook=False, ) + routing_weights = None + if glm5_numerics: + routing_weights = gather_low_latency_routing_weights(buffer, topk_idx, topk_weights, handle) # deepgemm gemm_out_b = masked_group_gemm( recv_x, @@ -397,13 +410,20 @@ def fused_experts_impl( w2, w2_scale, expected_m, + routing_weights=routing_weights, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( - gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False + gemm_out_b, + topk_idx, + torch.ones_like(topk_weights) if glm5_numerics else topk_weights, + handle, + async_finish=False, + return_recv_hook=False, ) return combined_x @@ -521,6 +541,7 @@ def chunked_expanded_moe_forward( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + glm5_numerics: bool = False, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" alignment = 128 @@ -588,7 +609,16 @@ def chunked_expanded_moe_forward( gemm_out_a, m_indices[chunk_start:chunk_end], ) - silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) + silu_and_mul_fwd( + gemm_out_a, + silu_out, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + glm5_numerics=glm5_numerics, + ) + if glm5_numerics: + multiply_routing_weights(silu_out, recv_topk_weights[chunk_start:chunk_end]) workspace_manager.free(gemm_out_a) del gemm_out_a @@ -624,7 +654,14 @@ def workspace_quant_alloc(shape, dtype, device): for quant_buffer in quant_buffers: workspace_manager.free(quant_buffer) - ep_gather_chunk(gemm_out_b, chunk_start, recv_topk_weights, recv_src_metadata, gather_out) + ep_gather_chunk( + gemm_out_b, + chunk_start, + recv_topk_weights, + recv_src_metadata, + gather_out, + apply_routed_weight=not glm5_numerics, + ) workspace_manager.free(gemm_out_b) finally: if is_autotune_warmup: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py new file mode 100644 index 0000000000..c10769a734 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py @@ -0,0 +1,120 @@ +"""Apply GLM routing probabilities at the BF16 W2 input boundary.""" + +import torch +import torch.distributed as dist +import triton +import triton.language as tl + + +@triton.jit +def _multiply_rows(X, P, ROWS: tl.constexpr, COLS: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(X + row * COLS + cols, cols < COLS, 0).to(tl.float32) + p = tl.load(P + row) + tl.store(X + row * COLS + cols, x * p, cols < COLS) + + +def multiply_routing_weights(x, weights): + """x is already rounded to BF16 by SwiGLU; round x.float()*p to BF16 again.""" + assert x.is_contiguous() and weights.is_contiguous() + assert x.numel() // x.shape[-1] == weights.numel() + _multiply_rows[(weights.numel(), triton.cdiv(x.shape[-1], 256))]( + x, weights, weights.numel(), x.shape[-1], 256, num_warps=4 + ) + + +@triton.jit +def _sort_tp_weights(Blocks, Counts, Weights, Out, STRIDE: tl.constexpr, BLOCK_M: tl.constexpr): + block = tl.program_id(0) + expert = tl.load(Blocks + block * 3) + if expert >= 0: + tile = tl.load(Blocks + block * 3 + 1) + start = tl.load(Blocks + block * 3 + 2) + lanes = tl.arange(0, BLOCK_M) + count = tl.load(Counts + expert) + p = tl.load(Weights + expert * STRIDE + tile * BLOCK_M + lanes, tile * BLOCK_M + lanes < count, 0) + tl.store(Out + start + lanes, p, tile * BLOCK_M + lanes < count) + + +def sorted_tp_routing_weights(topk_weights, expert_weights, counts, mblock_infos): + """Reorder p exactly like the TMA W13 output, including non-aligned expert tails.""" + out = torch.empty_like(topk_weights).view(-1) + blocks, block_m = mblock_infos + _sort_tp_weights[(blocks.shape[0],)]( + blocks, counts, expert_weights, out, expert_weights.stride(0), block_m, num_warps=4 + ) + return out + + +@triton.jit +def _gather_low_latency_weights( + Routes, + SrcInfo, + Layout, + Out, + NUM_RANKS: tl.constexpr, + MAX_TOKENS: tl.constexpr, + TOPK: tl.constexpr, + LOCAL_EXPERTS: tl.constexpr, + RANK: tl.constexpr, + BLOCK: tl.constexpr, + BLOCK_TOPK: tl.constexpr, +): + expert, src_rank = tl.program_id(0), tl.program_id(1) + layout = tl.load(Layout + expert * NUM_RANKS + src_rank) + count = (layout & 0xFFFFFFFF).to(tl.int32) + start = (layout >> 32).to(tl.int32) + k = tl.arange(0, BLOCK_TOPK) + for offset in range(0, count, BLOCK): + rows = offset + tl.arange(0, BLOCK) + valid = rows < count + dst = expert * (NUM_RANKS * MAX_TOKENS) + start + rows + src_token = tl.load(SrcInfo + dst, valid, 0) + ptr = Routes + ((src_rank * MAX_TOKENS + src_token[:, None]) * TOPK + k[None, :]) * 2 + ids = tl.load(ptr, valid[:, None] & (k[None, :] < TOPK), -1).to(tl.int32) + weights = tl.load(ptr + 1, valid[:, None] & (k[None, :] < TOPK), 0) + p = tl.sum(tl.where(ids == RANK * LOCAL_EXPERTS + expert, weights, 0), 1) + tl.store(Out + dst, p, valid) + + +def gather_low_latency_routing_weights(buffer, topk_idx, topk_weights, handle): + """Recover p using DeepEP's actual receive order, without CPU count synchronization. + + Legacy dispatch transports X but not p. All-gather only the small padded + FP32 (expert id, p) table, then use its src_info/layout_range handle. All + ranks participate, including empty/HOLD ranks, so CUDA graph replay has a + fixed collective shape. Expert ids are exactly representable in FP32. + """ + src_info, layout, max_tokens, _, num_experts = handle + ranks, topk = buffer.group_size, topk_idx.shape[1] + assert num_experts < 2**24 and topk_idx.shape[0] <= max_tokens + assert topk_idx.shape == topk_weights.shape and topk_weights.dtype == torch.float32 + assert src_info.shape == (num_experts // ranks, ranks * max_tokens) + assert layout.shape == (num_experts // ranks, ranks) + assert src_info.dtype == torch.int32 and layout.dtype == torch.int64 + routes = torch.zeros((max_tokens, topk, 2), dtype=torch.float32, device=topk_weights.device) + routes[:, :, 0].fill_(-1) + routes[: topk_idx.shape[0], :, 0].copy_(topk_idx) + routes[: topk_idx.shape[0], :, 1].copy_(topk_weights) + if ranks == 1: + all_routes = routes + else: + all_routes = torch.empty((ranks * max_tokens, topk, 2), dtype=routes.dtype, device=routes.device) + dist.all_gather_into_tensor(all_routes, routes, group=buffer.group) + out = torch.zeros(src_info.shape, dtype=torch.float32, device=topk_weights.device) + _gather_low_latency_weights[(num_experts // ranks, ranks)]( + all_routes, + src_info, + layout, + out, + ranks, + max_tokens, + topk, + num_experts // ranks, + buffer.rank, + 32, + triton.next_power_of_2(topk), + num_warps=4, + ) + return out diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index 8169cefc3c..f0771cb1a8 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py @@ -26,6 +26,7 @@ def _silu_and_mul_kernel_fast( USE_LIMIT_AND_ALPHA: tl.constexpr = False, CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, + GLM5_NUMERICS: tl.constexpr = False, ): stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) stride_output_m = tl.cast(stride_output_m, dtype=tl.int64) @@ -66,11 +67,14 @@ def _silu_and_mul_kernel_fast( other=other, ).to(tl.float32) + if GLM5_NUMERICS: + up = up.to(tl.float32) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) gate = 1 / (1 + tl.exp(-gate * alpha)) * gate - gate = gate.to(input_ptr.dtype.element_ty) + if not GLM5_NUMERICS: + gate = gate.to(input_ptr.dtype.element_ty) if CLAMP_UP_ADD_ONE: up += 1 tl.store(output_ptr + out_offsets, up * gate, mask=mask) @@ -83,7 +87,8 @@ def _silu_and_mul_kernel_fast( gate = 0.5 * gate * (1.0 + tanh_val) else: gate = gate / (1 + tl.exp(-gate)) - gate = gate.to(input_ptr.dtype.element_ty) + if not GLM5_NUMERICS: + gate = gate.to(input_ptr.dtype.element_ty) tl.store( output_ptr + out_offsets, @@ -121,6 +126,7 @@ def silu_and_mul_fwd( alpha=None, run_config=None, clamp_up_add_one=True, + glm5_numerics=False, ): assert input.stride(-1) == 1 assert output.is_contiguous() @@ -173,5 +179,6 @@ def silu_and_mul_fwd( USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA, CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + GLM5_NUMERICS=glm5_numerics, ) return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 8c94f088c2..69d371d4cd 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -22,6 +22,8 @@ def _silu_and_mul_post_quant_kernel( stride_output_scale_1, stride_output_scale_2, masked_m_ptr, + routing_weights, + ROUTING_STRIDE: tl.constexpr, size_n, fp8_max, fp8_min, @@ -33,6 +35,8 @@ def _silu_and_mul_post_quant_kernel( limit: tl.constexpr = None, CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_UE8M0_SCALE: tl.constexpr = False, + GLM5_NUMERICS: tl.constexpr = False, + WEIGHTED_INPUT: tl.constexpr = False, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -55,8 +59,12 @@ def _silu_and_mul_post_quant_kernel( ) for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): - gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) + gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to( + tl.float32 + ) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) + if GLM5_NUMERICS: + up = up.to(tl.float32) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) @@ -70,11 +78,15 @@ def _silu_and_mul_post_quant_kernel( gate = 0.5 * gate * (1.0 + tanh_val) else: gate = gate / (1 + tl.exp(-gate)) - gate = gate.to(input_ptr.dtype.element_ty) + if not GLM5_NUMERICS: + gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate - if USE_LIMIT_AND_ALPHA: + if USE_LIMIT_AND_ALPHA or GLM5_NUMERICS: # Match the BF16/FP16 activation stored before prefill's FP8 quantization. gate_up = gate_up.to(input_ptr.dtype.element_ty).to(tl.float32) + if WEIGHTED_INPUT: + p = tl.load(routing_weights + expert_id * ROUTING_STRIDE + token_index) + gate_up = (gate_up * p).to(input_ptr.dtype.element_ty).to(tl.float32) _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) output_s = _absmax / fp8_max if USE_UE8M0_SCALE: @@ -100,6 +112,8 @@ def silu_and_mul_masked_post_quant_fwd( alpha=None, limit=None, clamp_up_add_one=True, + glm5_numerics=False, + routing_weights=None, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -117,6 +131,10 @@ def silu_and_mul_masked_post_quant_fwd( assert input.shape[0] == masked_m.shape[0] assert input.shape[-1] % 2 == 0 + if routing_weights is not None: + assert routing_weights.shape == input.shape[:2] and routing_weights.is_contiguous() + assert routing_weights.dtype == torch.float32 + size_n = input.shape[-1] // 2 assert size_n % quant_group_size == 0 @@ -151,6 +169,8 @@ def silu_and_mul_masked_post_quant_fwd( output_scale, *output_scale.stride(), masked_m, + routing_weights, + routing_weights.stride(0) if routing_weights is not None else 0, size_n, fp8_max, fp8_min, @@ -162,6 +182,9 @@ def silu_and_mul_masked_post_quant_fwd( limit=limit, CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_UE8M0_SCALE=get_env_use_ue8m0_scales(), + GLM5_NUMERICS=glm5_numerics, + enable_fp_fusion=not glm5_numerics, + WEIGHTED_INPUT=routing_weights is not None, num_warps=num_warps, ) return diff --git a/lightllm/common/basemodel/triton_kernel/mhc/post.py b/lightllm/common/basemodel/triton_kernel/mhc/post.py index 2b798018f7..b420bb2edf 100644 --- a/lightllm/common/basemodel/triton_kernel/mhc/post.py +++ b/lightllm/common/basemodel/triton_kernel/mhc/post.py @@ -74,16 +74,16 @@ def _hc_post_4stream_kernel( accumulator_3 += residual_3 * tl.load(mix_base + 15) if GLM5_NUMERICS: - # Megatron's compiled native post rounds the residual bmm to BF16, - # then fuses the post expansion and addition into the final BF16 output. + # With TORCH_COMPILE_DISABLE=1, Megatron materializes both the + # residual bmm and the post product in the activation dtype. accumulator_0 = accumulator_0.to(residual.dtype.element_ty).to(tl.float32) accumulator_1 = accumulator_1.to(residual.dtype.element_ty).to(tl.float32) accumulator_2 = accumulator_2.to(residual.dtype.element_ty).to(tl.float32) accumulator_3 = accumulator_3.to(residual.dtype.element_ty).to(tl.float32) - accumulator_0 += layer_value * tl.load(post_base) - accumulator_1 += layer_value * tl.load(post_base + 1) - accumulator_2 += layer_value * tl.load(post_base + 2) - accumulator_3 += layer_value * tl.load(post_base + 3) + accumulator_0 += (layer_value * tl.load(post_base)).to(layer_output.dtype.element_ty).to(tl.float32) + accumulator_1 += (layer_value * tl.load(post_base + 1)).to(layer_output.dtype.element_ty).to(tl.float32) + accumulator_2 += (layer_value * tl.load(post_base + 2)).to(layer_output.dtype.element_ty).to(tl.float32) + accumulator_3 += (layer_value * tl.load(post_base + 3)).to(layer_output.dtype.element_ty).to(tl.float32) output_base = output + token * out_stride_m + hidden_offsets tl.store(output_base, accumulator_0, mask=hidden_mask) @@ -111,7 +111,7 @@ def hc_post( residual_mix: [T, S, S]，fp32 或激活 dtype；来自 hc_pre_norm。 post_mix: [T, S]，与 residual_mix 相同 dtype；来自同一次 pre 调用。 streams: 残差 stream 数 S；当前 Triton 入口要求 S = 4。 - glm5_numerics: 匹配 GLM 训练默认 native post，在加子层输出前将残差混合结果舍入到激活 dtype。 + glm5_numerics: 匹配 GLM 训练默认 native post，在加子层输出前将残差混合结果和 post 乘积分别舍入到激活 dtype。 Returns: 新残差 [T, S * H]，dtype 与 layer_output 相同；输入张量保持不变。 @@ -144,6 +144,7 @@ def hc_post( out_stride_m=output.stride(0), BLOCK_H=block_h, GLM5_NUMERICS=glm5_numerics, + enable_fp_fusion=not glm5_numerics, num_warps=8, ) return output diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 2032e50c96..f13d447066 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -58,7 +58,7 @@ def __init__(self, layer_num, network_config): self.q_lora_rank = network_config["q_lora_rank"] self.kv_lora_rank = network_config["kv_lora_rank"] self.v_head_dim = network_config["v_head_dim"] - self.softmax_scale = self.qk_nope_head_dim ** -0.5 + self.softmax_scale = self.qk_nope_head_dim**-0.5 self.indexer = Glm5NextNsaInfer( layer_idx=self.layer_num_, network_config=self.network_config_, @@ -94,6 +94,7 @@ def _ffn_tp(self, input, infer_state, layer_weight): limit=self.swiglu_limit, alpha=1.0, clamp_up_add_one=False, + glm5_numerics=True, ) return layer_weight.down_proj.mm(ffn1_out) @@ -434,7 +435,9 @@ def _hc_pre(self, streams, layer_weight, prefix, norm_weight): def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): streams = input_embeddings - layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) + layer_input, residual_mix, post_mix = self._hc_pre( + streams, layer_weight, "attn", layer_weight.att_norm_weight_ + ) if prefill: layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) else: diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index 32191809fe..b193ada4e7 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -212,6 +212,7 @@ def _init_moe(self): num_fused_shared_experts=self.num_fused_shared_experts, layer_num=self.layer_num_, network_config=self.network_config_, + glm5_numerics=True, ) def _init_kda(self): diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py index 32d2b694ee..d126c65ca8 100644 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -67,7 +67,10 @@ def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams, mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) if round_residual: mixed_residual = mixed_residual.to(layer_output.dtype).float() - out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + expanded = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + if round_residual: + expanded = expanded.to(layer_output.dtype).float() + out = expanded + mixed_residual return out.to(layer_output.dtype).reshape(tokens, streams * hidden) @@ -244,10 +247,20 @@ def test_glm_mhc_keeps_pre_coefficients_in_fp32(): base[:2] = torch.logit(pre_values) norm = torch.ones(17, device="cuda", dtype=x.dtype) actual, _, _ = hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20, glm5_numerics=True) - merged, _, _ = _hc_pre_reference( - x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True - ) + merged, _, _ = _hc_pre_reference(x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True) merged = merged.float() expected = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5)).to(x.dtype) assert (expected > 0.1).all() torch.testing.assert_close(actual, expected, atol=0.001, rtol=0.008) + + +def test_glm_mhc_rounds_post_product_before_residual_add(): + x = torch.full((1, 4 * 17), -1.015625, device="cuda", dtype=torch.bfloat16) + layer = torch.full((1, 17), 1.0078125, device="cuda", dtype=x.dtype) + mix = torch.eye(4, device="cuda", dtype=x.dtype).unsqueeze(0) + post = torch.full((1, 4), 1.0078125, device="cuda", dtype=x.dtype) + expected = _hc_post_reference(layer, x, mix, post, 4, round_residual=True) + actual = hc_post(layer, x, mix, post, 4, glm5_numerics=True) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert torch.count_nonzero(actual) == 0 + assert torch.count_nonzero(hc_post(layer, x, mix, post, 4)) > 0 diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py index e9ba2ab26e..d4ac0f3935 100644 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -11,10 +11,13 @@ ) +@pytest.mark.parametrize("glm5_numerics", [False, True]) @pytest.mark.parametrize("is_prefill", [False, True]) @pytest.mark.parametrize("use_ue8m0_scales", [False, True]) @pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) -def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one): +def test_grouped_gemm_activation_matches_reference( + monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one, glm5_numerics +): monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES") config = {"quant_method": "fp8", "weight_block_size": [128, 128]} @@ -30,7 +33,10 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ w1_scale = torch.ones(2, 2, 1, device="cuda") w2_scale = torch.ones(2, 1, 1, device="cuda") counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) - activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) + activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one, glm5_numerics=glm5_numerics) + p = torch.tensor([0.13, 0.71, 0.88, 1.0, 1.3], device="cuda") if glm5_numerics else torch.ones(5, device="cuda") + recv_p = torch.zeros(2, padded, device="cuda") + recv_p[0, :3], recv_p[1, :2] = p[:3], p[3:] recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) recv[0, :3] = x[:3] recv[1, :2] = x[3:] @@ -44,7 +50,7 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ num_recv_tokens_per_expert_list=[padded, padded], num_unaligned_recv_tokens_per_expert=counts, recv_x=(recv.view(2 * padded, dim), torch.ones(1, 2 * padded, device="cuda").T), - recv_topk_weights=torch.ones(2 * padded, device="cuda"), + recv_topk_weights=recv_p.view(-1), recv_src_metadata=metadata, w1=w1, w1_scale=w1_scale, @@ -65,6 +71,7 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ w2, w2_scale, expected_m=3, + routing_weights=recv_p if glm5_numerics else None, **activation_args, ) actual = torch.cat([output[0, :3], output[1, :2]]) @@ -74,8 +81,11 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ if limit is not None: gate = gate.clamp(max=limit) up = up.clamp(-limit, limit) + int(add_one) - gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() + gate = gate * torch.sigmoid(gate * (alpha or 1.0)) + if not glm5_numerics: + gate = gate.bfloat16().float() activation = (gate * up).bfloat16().float() + activation = (activation * p[:, None]).bfloat16().float() scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 if use_ue8m0_scales: scales = torch.exp2(torch.ceil(torch.log2(scales))) diff --git a/unit_tests/common/fused_moe/test_glm53_training_alignment.py b/unit_tests/common/fused_moe/test_glm53_training_alignment.py new file mode 100644 index 0000000000..6eb8c63176 --- /dev/null +++ b/unit_tests/common/fused_moe/test_glm53_training_alignment.py @@ -0,0 +1,323 @@ +"""GLM launch recipe: eager BF16 boundaries, FP32 SwiGLU and p before W2.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( + silu_and_mul_masked_post_quant_fwd, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_route_weight import ( + multiply_routing_weights, + gather_low_latency_routing_weights, + _gather_low_latency_weights, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def swiglu_reference(x): + gate, up = x.float().chunk(2, -1) + return (F.silu(gate.clamp(max=10)) * up.clamp(-10, 10)).to(x.dtype) + + +def quant_reference(x, ue8m0): + groups = x.float().reshape(*x.shape[:-1], -1, 128) + sf = groups.abs().amax(-1).clamp_min(1e-10) / 448 + if ue8m0: + sf = torch.exp2(torch.ceil(torch.log2(sf))) + q = (groups / sf.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) + return q.reshape_as(x), sf + + +@pytest.mark.parametrize("rows,width", [(1, 256), (19, 4096)]) +def test_swiglu_retains_fp32_product_and_bf16_weight_boundary(rows, width): + torch.manual_seed(53) + x = (torch.randn(rows, width, device="cuda") * 6).bfloat16() + out = torch.empty(rows, width // 2, device=x.device, dtype=x.dtype) + silu_and_mul_fwd(x, out, alpha=1.0, limit=10.0, clamp_up_add_one=False, glm5_numerics=True) + expected = swiglu_reference(x) + # CUDA exp implementations differ near BF16 midpoints. + torch.testing.assert_close(out, expected, atol=1e-6, rtol=0.008) + old = torch.empty_like(out) + silu_and_mul_fwd(x, old, alpha=1.0, limit=10.0, clamp_up_add_one=False) + assert torch.count_nonzero(old != expected) > torch.count_nonzero(out != expected) + p = torch.linspace(0.01, 1.7, rows, device=x.device) + before = out.clone() + multiply_routing_weights(out, p) + torch.testing.assert_close(out, (before.float() * p[:, None]).to(x.dtype), atol=0, rtol=0) + + +@pytest.mark.parametrize("ue8m0", [False, True]) +def test_masked_decode_swiglu_matches_training_before_quant(monkeypatch, ue8m0): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) + torch.manual_seed(53) + x = (torch.randn(3, 5, 4096, device="cuda") * 6).bfloat16() + p = torch.rand(3, 5, device=x.device) * 1.7 + p[1, 0] = 0 + counts = torch.tensor([0, 2, 5], device=x.device, dtype=torch.int32) + out = torch.full((3, 5, 2048), 1, device=x.device, dtype=torch.float8_e4m3fn) + sf = torch.full((3, 5, 16), 7.0, device=x.device) + silu_and_mul_masked_post_quant_fwd( + x, + out, + sf, + 128, + counts, + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + glm5_numerics=True, + routing_weights=p, + ) + weighted = (swiglu_reference(x).float() * p[..., None]).to(x.dtype) + q_ref, s_ref = quant_reference(weighted, ue8m0) + for expert, count in enumerate(counts.tolist()): + torch.testing.assert_close(sf[expert, :count], s_ref[expert, :count], atol=1e-6, rtol=1e-5) + actual = out[expert, :count].float() * sf[expert, :count].repeat_interleave(128, -1) + expected = q_ref[expert, :count].float() * s_ref[expert, :count].repeat_interleave(128, -1) + # A SwiGLU midpoint may flip one FP8 rounding decision. + err = (actual - expected).square().mean().sqrt() if count else actual.sum() + assert err < 0.01 + assert torch.all(out[expert, count:].float() == 1) + assert torch.all(sf[expert, count:] == 7) + + +@pytest.mark.parametrize("empty_local", [False, True]) +def test_low_latency_routing_uses_source_and_actual_receive_order(monkeypatch, empty_local): + from lightllm.common.basemodel.triton_kernel.fused_moe import moe_route_weight as route + + # Two ranks, unequal token counts; expert receive order differs by source rank. + max_tokens, ranks, local_experts, topk = 4, 2, 2, 2 + ids = torch.tensor([[[2, 3], [0, 2], [-1, -1], [-1, -1]], [[3, 2], [2, 1], [2, 3], [-1, -1]]], device="cuda") + p = torch.tensor( + [[[0.1, 0.9], [0.2, 0.8], [0, 0], [0, 0]], [[0.3, 0.7], [0.4, 0.6], [0.5, 0.5], [0, 0]]], device="cuda" + ) + if empty_local: + ids[1].fill_(-1) + p[1].zero_() + routes = torch.stack((ids.float(), p), -1).reshape(ranks * max_tokens, topk, 2) + src_info = torch.zeros((local_experts, ranks * max_tokens), device="cuda", dtype=torch.int32) + layout = torch.zeros((local_experts, ranks), device="cuda", dtype=torch.int64) + expected = torch.zeros_like(src_info, dtype=torch.float32) + rank = 1 + for expert in range(local_experts): + offset = 0 + for src_rank in (1, 0): + tokens = torch.where((ids[src_rank] == rank * local_experts + expert).any(-1))[0].flip(0) + count = tokens.numel() + src_info[expert, offset : offset + count] = tokens.int() + layout[expert, src_rank] = (offset << 32) | count + for row, token in enumerate(tokens.tolist()): + slot = (ids[src_rank, token] == rank * local_experts + expert).nonzero()[0, 0] + expected[expert, offset + row] = p[src_rank, token, slot] + offset += count + calls = [] + + def gather(out, inp, group): + calls.append(inp.clone()) + out.copy_(routes) + + monkeypatch.setattr(route.dist, "all_gather_into_tensor", gather) + buffer = SimpleNamespace(group_size=ranks, rank=rank, group=object()) + handle = (src_info, layout, max_tokens, 256, ranks * local_experts) + actual = gather_low_latency_routing_weights( + buffer, ids[1, : 0 if empty_local else 3], p[1, : 0 if empty_local else 3], handle + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert len(calls) == 1 + assert torch.all(calls[0][3, :, 0] == -1) + # The metadata mapping is graph replayable and reads changed inputs. + out = torch.zeros_like(expected) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + _gather_low_latency_weights[(local_experts, ranks)]( + routes, src_info, layout, out, ranks, max_tokens, topk, local_experts, rank, 32, 2 + ) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _gather_low_latency_weights[(local_experts, ranks)]( + routes, src_info, layout, out, ranks, max_tokens, topk, local_experts, rank, 32, 2 + ) + routes[:, :, 1].mul_(2) + graph.replay() + torch.testing.assert_close(out, expected * 2, atol=0, rtol=0) + + +@pytest.mark.parametrize("tokens", [2, 19, 257]) +@pytest.mark.parametrize("tma", [False, True]) +@pytest.mark.parametrize( + "quantized,ue8m0,packed", [(False, False, False), (True, False, False), (True, True, False), (True, True, True)] +) +def test_tp_moe_routes_before_w2_with_tma_sorting_and_shared_expert( + monkeypatch, tma, quantized, ue8m0, packed, tokens +): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe as tp + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0 and not packed))) + monkeypatch.setattr(tp, "triton_support_tensor_descriptor", lambda: tma) + torch.manual_seed(53) + dim, experts = 128, 4 + x = (torch.randn(tokens, dim, device="cuda") * 4).bfloat16() + ids = torch.stack( + ( + torch.arange(tokens, device=x.device) % 3, + (torch.arange(tokens, device=x.device) + 1) % 3, + torch.full((tokens,), 3, device=x.device), + ), + -1, + ).int() + p = torch.rand(tokens, 3, device=x.device) * 1.7 + p[:, -1] = 1 # Fused shared expert must remain unweighted. + eye = torch.eye(dim, device=x.device, dtype=x.dtype) + w1 = torch.cat((eye, eye)).repeat(experts, 1, 1) + w2 = eye.repeat(experts, 1, 1) + kwargs = {} + + def quantize(x): + # Ordinary scales use the installed SGL quantizer; its reciprocal + # multiplication can resolve FP8 midpoints differently from Torch's + # division. Keep that unchanged stage fixed while testing W2 placement. + if ue8m0: + return quant_reference(x, True) + from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 + + q, sf = per_token_group_quant_fp8(x.reshape(-1, dim), 128) + return q.reshape_as(x), sf.reshape(*x.shape[:-1], 1) + + if quantized: + w1, w2 = w1.to(torch.float8_e4m3fn), w2.to(torch.float8_e4m3fn) + kwargs = dict( + w1_scale=torch.ones(experts, 2, 1, device=x.device), w2_scale=torch.ones(experts, 1, 1, device=x.device) + ) + if packed: + # One K group: exponent 127 occupies byte 0, per output channel. + # Leave the env off to check packed B itself implies UE8M0 on A. + kwargs = dict( + w1_scale=torch.full((experts, 256, 1), 127, device=x.device, dtype=torch.int32), + w2_scale=torch.full((experts, 128, 1), 127, device=x.device, dtype=torch.int32), + ) + q, sf = quantize(x) + pre = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() + else: + pre = x + h = swiglu_reference(torch.cat((pre, pre), -1)) + h = (h[:, None, :].float() * p[..., None]).bfloat16() + if quantized: + q, sf = quantize(h) + h = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() + expected = h.float().sum(1).bfloat16() + actual = tp.fused_experts( + x.clone(), + w1, + w2, + p, + ids, + use_fp8_w8a8=quantized, + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + glm5_numerics=True, + **kwargs, + ) + torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.008) + legacy = tp.fused_experts( + x.clone(), + w1, + w2, + p, + ids, + use_fp8_w8a8=quantized, + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + **kwargs, + ) + assert torch.count_nonzero(legacy != expected) > torch.count_nonzero(actual != expected) + + +@pytest.mark.parametrize("ue8m0", [False, True]) +def test_ep_decode_dispatch_to_combine_weights_exactly_once(monkeypatch, ue8m0): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) + monkeypatch.setattr(ep, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 128) + torch.manual_seed(53) + tokens, dim, experts = 5, 128, 2 + x = (torch.randn(tokens, dim, device="cuda") * 3).bfloat16() + ids = torch.tensor([[0, 1]] * tokens, device=x.device, dtype=torch.int64) + p = torch.rand(tokens, 2, device=x.device) * 1.7 + orders = torch.tensor([[4, 1, 3, 0, 2], [2, 4, 1, 3, 0]], device=x.device) + eye = torch.eye(dim, device=x.device, dtype=x.dtype) + w1 = torch.cat((2 * eye, 3 * eye)).repeat(experts, 1, 1).to(torch.float8_e4m3fn) + w2 = eye.repeat(experts, 1, 1).to(torch.float8_e4m3fn) + seen = [] + + class Buffer: + group_size, rank = 1, 0 + + def low_latency_dispatch(self, hidden, topk_idx, max_tokens, num_experts, **kwargs): + assert kwargs["round_scale"] == ue8m0 + q, sf = quant_reference(hidden, ue8m0) + recv = torch.zeros(experts, max_tokens, dim, device=x.device, dtype=q.dtype) + scales = torch.zeros(experts, max_tokens, 1, device=x.device) + src = torch.zeros(experts, max_tokens, device=x.device, dtype=torch.int32) + for e in range(experts): + recv[e, :tokens] = q[orders[e]] + scales[e, :tokens] = sf[orders[e]] + src[e, :tokens] = orders[e].int() + counts = torch.full((experts,), tokens, device=x.device, dtype=torch.int32) + layout = torch.full((experts, 1), tokens, device=x.device, dtype=torch.int64) + return (recv, scales), counts, (src, layout, max_tokens, dim, experts), None, None + + def low_latency_combine(self, output, topk_idx, topk_weights, handle, **kwargs): + torch.testing.assert_close(topk_weights, torch.ones_like(p), atol=0, rtol=0) + seen.append(True) + result = torch.zeros(tokens, dim, device=x.device) + for e in range(experts): + result[orders[e]] += output[e, :tokens].float() + return result.bfloat16(), None, None + + actual = ep.fused_experts_impl( + x, + w1, + w2, + p, + ids, + experts, + Buffer(), + False, + use_fp8_w8a8=True, + use_fp8_all2all=True, + w1_scale=torch.ones(experts, 2, 1, device=x.device), + w2_scale=torch.ones(experts, 1, 1, device=x.device), + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + glm5_numerics=True, + ) + q, sf = quant_reference(x, ue8m0) + pre = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() + h = swiglu_reference(torch.cat(((2 * pre.float()).bfloat16(), (3 * pre.float()).bfloat16()), -1)) + weighted = (h[:, None, :].float() * p[..., None]).bfloat16() + q, sf = quant_reference(weighted, ue8m0) + expected = (q.float() * sf.repeat_interleave(128, -1)).bfloat16().float().sum(1).bfloat16() + if ue8m0: + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.008) + else: + # Arbitrary scales use reciprocal/division implementations with distinct + # FP8 midpoint ties. Bound the difference by one FP8 step per route, + # plus the final BF16 sum rounding; UE8M0 keeps the strict comparison. + q_float = q.float().abs().clamp_min(2.0**-6) + fp8_step = torch.exp2(torch.floor(torch.log2(q_float)) - 3) * sf.repeat_interleave(128, -1) + bound = fp8_step.sum(1) + expected.float().abs() * torch.finfo(torch.bfloat16).eps + assert torch.all((actual.float() - expected.float()).abs() <= bound + 1e-5) + nrmse = (actual.float() - expected.float()).square().mean().sqrt() / expected.float().square().mean().sqrt() + assert nrmse < 0.008 + assert seen == [True] diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 7b60c9a267..ee087b05e7 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -95,7 +95,16 @@ def test_kda_projections_need_no_sparse_attention_config(config): weight = Glm5NextTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) projections = { name: random_weight(size, 128) - for name, size in {"q": 512, "k": 512, "v": 512, "b": 4, "f_a": 128, "g_a": 128, "f_b": 512, "g_b": 512}.items() + for name, size in { + "q": 512, + "k": 512, + "v": 512, + "b": 4, + "f_a": 128, + "g_a": 128, + "f_b": 512, + "g_b": 512, + }.items() } weight.load_hf_weights( {f"model.language_model.layers.0.self_attn.{name}_proj.weight": value for name, value in projections.items()} @@ -226,10 +235,11 @@ def mlp_weights(name, gate_scale, up_scale): def reference_mlp(gate_scale, up_scale): gate = (x * gate_scale).float().clamp(max=10) up = (x * up_scale).float().clamp(-10, 10) - return (F.silu(gate).bfloat16().float() * up).bfloat16() + return (F.silu(gate) * up).bfloat16() expected = reference_mlp(1.5, 0.75) if mode != "dense": + assert weight.experts.fuse_moe_impl.glm5_numerics scores = F.linear(x.float(), router).sigmoid() probabilities, ids = scores.topk(2, dim=-1) probabilities /= probabilities.sum(-1, keepdim=True) @@ -270,7 +280,7 @@ def combined_output(hidden_states, **kwargs): monkeypatch.setattr(weight.experts, "experts", combined_output) actual = layer._ffn(x, SimpleNamespace(is_prefill=is_prefill), weight) - shared = (F.silu(x.float().clamp(max=10)).bfloat16().float() * x.float().clamp(-10, 10)).bfloat16() + shared = (F.silu(x.float().clamp(max=10)) * x.float().clamp(-10, 10)).bfloat16() torch.testing.assert_close(actual, shared + 17, atol=0.25, rtol=0.01) torch.testing.assert_close(x, original, atol=0, rtol=0) From 35f71e98d50016a064ae41a7a717433b8a9c1169 Mon Sep 17 00:00:00 2001 From: baishihao Date: Sat, 24 Oct 2026 00:09:41 +0800 Subject: [PATCH 10/13] fix(glm5): revert FFN numerical alignment for isolated indexer ablation --- .../fused_moe/fused_moe_weight.py | 2 - .../meta_weights/fused_moe/impl/base_impl.py | 2 - .../fused_moe/impl/deepgemm_impl.py | 1 - .../fused_moe/impl/triton_impl.py | 1 - .../deepep_expanded_layout_kernels.py | 9 +- .../fused_moe/grouped_fused_moe.py | 39 +-- .../fused_moe/grouped_fused_moe_ep.py | 43 +-- .../fused_moe/moe_route_weight.py | 120 ------- .../fused_moe/moe_silu_and_mul.py | 11 +- .../moe_silu_and_mul_mix_quant_ep.py | 29 +- .../layer_infer/transformer_layer_infer.py | 7 +- .../layer_weights/transformer_layer_weight.py | 1 - .../fused_moe/test_deepgemm_activation.py | 18 +- .../test_glm53_training_alignment.py | 323 ------------------ unit_tests/models/glm5_next/test_layers.py | 16 +- 15 files changed, 26 insertions(+), 596 deletions(-) delete mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py delete mode 100644 unit_tests/common/fused_moe/test_glm53_training_alignment.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 97aebbd3fc..5652279702 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -35,7 +35,6 @@ def __init__( layer_num: int = 0, network_config: Dict[str, Any] = None, per_expert_scale_name: str = "", - glm5_numerics: bool = False, ) -> None: super().__init__(data_type=data_type) self.w1_weight_name = gate_proj_name @@ -68,7 +67,6 @@ def __init__( redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, routed_expert_counter_tensor=self.routed_expert_counter_tensor, auto_update_redundancy_expert=self.auto_update_redundancy_expert, - glm5_numerics=glm5_numerics, ) self.lock = threading.Lock() self._create_weight() diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index 7a1c1219bb..e242870d69 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -22,9 +22,7 @@ def __init__( redundancy_expert_ids_tensor: torch.Tensor, routed_expert_counter_tensor: torch.Tensor, auto_update_redundancy_expert: bool, - glm5_numerics: bool = False, ): - self.glm5_numerics = glm5_numerics self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self.routed_scaling_factor = routed_scaling_factor diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 3f85309457..74e457426b 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -94,7 +94,6 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=self.glm5_numerics, ) return output diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index ccafa36292..f748aea467 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -87,7 +87,6 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=self.glm5_numerics, ) return input_tensor diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py index 21f3f6439c..ab54ea5b79 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py @@ -104,9 +104,7 @@ def _ep_zero_padding_kernel( scale_offsets = tl.arange(0, BLOCK_SCALE_K) scale_mask = scale_offsets < scale_hidden_size scale_ptrs = ( - recv_x_scale - + row_offsets[:, None] * recv_x_scale_stride_m - + scale_offsets[None, :] * recv_x_scale_stride_k + recv_x_scale + row_offsets[:, None] * recv_x_scale_stride_m + scale_offsets[None, :] * recv_x_scale_stride_k ) tl.store(scale_ptrs, 0.0, mask=row_mask[:, None] & scale_mask[None, :]) tl.store(recv_topk_weights + row_offsets, 0.0, mask=row_mask) @@ -175,7 +173,6 @@ def _ep_gather_chunk_kernel( TOPK: tl.constexpr, BLOCK_D: tl.constexpr, NEED_HIDDEN_MASK: tl.constexpr, - APPLY_ROUTED_WEIGHT: tl.constexpr, ): hidden_block_id = tl.program_id(0) start_recv_token_id = tl.program_id(1) @@ -199,7 +196,7 @@ def _ep_gather_chunk_kernel( value = tl.load(chunk_ptrs, mask=hidden_mask, other=0.0) else: value = tl.load(chunk_ptrs) - weight = tl.load(weights + slot) if APPLY_ROUTED_WEIGHT else 1.0 + weight = tl.load(weights + slot) accumulator += value.to(tl.float32) * weight if NEED_HIDDEN_MASK: tl.store(output_ptrs, accumulator, mask=hidden_mask) @@ -214,7 +211,6 @@ def ep_gather_chunk( weights: torch.Tensor, # [num_expanded_tokens] recv_src_metadata: torch.Tensor, # [num_recv_tokens, topk + 2] output: torch.Tensor, # [num_recv_tokens, hidden_size] - apply_routed_weight: bool = True, ): """Accumulate one expanded W2-output chunk into dense receive-token rows. @@ -244,7 +240,6 @@ def ep_gather_chunk( output_stride_k=output.stride(1), hidden_size=hidden_size, TOPK=topk, - APPLY_ROUTED_WEIGHT=apply_routed_weight, BLOCK_D=block_d, # 常见整块 hidden size 保持原来的无 mask kernel；仅尾块不完整时启用边界保护。 NEED_HIDDEN_MASK=hidden_size % block_d != 0, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index 794ad76d94..64537ac606 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -27,7 +27,6 @@ from lightllm.utils.device_utils import triton_support_tensor_descriptor from lightllm.utils.envs_utils import get_env_use_ue8m0_scales from .moe_silu_and_mul import silu_and_mul_fwd -from .moe_route_weight import multiply_routing_weights, sorted_tp_routing_weights from .moe_sum_reduce import moe_sum_reduce from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 from lightllm.utils.torch_ops_utils import direct_register_custom_op @@ -393,9 +392,7 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t else: # Expert rows may be unordered, but grouped matmul reuses this same # mapping for up/down projections and writes back to original topk slots. - _moe_align_fused_atomic_token( - expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights - ) + _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights) return expert_to_token_index, expert_to_weight, expert_token_num @@ -584,7 +581,7 @@ def grouped_matmul_kernel( assert BLOCK_SIZE_K <= block_size_k token_scale_stride0 = token_stride_0 // block_size_k if TOKEN_INPUT_USE_TMA: - assert topk_num == 1 + assert MUL_ROUTED_WEIGHT is True a_scale_ptrs = token_scale_ptr + (token_start_index + tl.arange(0, BLOCK_SIZE_M)) * token_scale_stride0 else: a_scale_ptrs = token_scale_ptr + (a_m_index // topk_num) * token_scale_stride0 @@ -599,7 +596,7 @@ def grouped_matmul_kernel( else: # per token scale quant if TOKEN_INPUT_USE_TMA: - assert topk_num == 1 + assert MUL_ROUTED_WEIGHT is True a_scale_ptrs = token_scale_ptr + (token_start_index + tl.arange(0, BLOCK_SIZE_M))[:, None] else: a_scale_ptrs = token_scale_ptr + (a_m_index // topk_num)[:, None] @@ -800,7 +797,6 @@ def grouped_matmul( reused_mblock_infos=None, run_config: Optional[dict] = None, bias: Optional[torch.Tensor] = None, # per-expert bias [expert_num, N] - apply_routed_weight: bool = True, ): """ token_num_mul_topk_num is int equal token_num * topk_num, @@ -909,8 +905,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]): triton.set_allocator(alloc_fn) - # mul_routed_weight selects the W2 layout. GLM can separately disable - # its output weighting after applying p to the BF16 input before quantization. + # moe 分为 up 和 down 两次计算，当 mul_routed_weight 为 False 的时候为 up is_up_moe = not mul_routed_weight if is_up_moe: @@ -994,7 +989,7 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]): BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=GROUP_SIZE_M, - MUL_ROUTED_WEIGHT=mul_routed_weight and apply_routed_weight, + MUL_ROUTED_WEIGHT=mul_routed_weight, NEED_K_MASK=NEED_K_MASK, NEED_TRANS=NEED_TRANS, num_warps=num_warps, @@ -1029,7 +1024,6 @@ def fused_experts_impl( limit=None, alpha=None, clamp_up_add_one=True, - glm5_numerics=False, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -1109,18 +1103,9 @@ def fused_experts_impl( limit=limit, alpha=alpha, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, layout=layout, ) - if glm5_numerics: - p = curr_topk_weights.reshape(-1) - if triton_support_tensor_descriptor(): - p = sorted_tp_routing_weights( - curr_topk_weights, expert_to_weights, expert_to_token_num, reused_mblock_infos - ) - multiply_routing_weights(intermediate_cache2.view(-1, N // 2), p) - grouped_matmul( curr_topk_ids.numel(), intermediate_cache2.view(-1, N // 2), @@ -1138,7 +1123,6 @@ def fused_experts_impl( reused_mblock_infos=reused_mblock_infos, run_config=run_config, bias=w2_bias, - apply_routed_weight=not glm5_numerics, ) moe_sum_reduce( @@ -1166,7 +1150,6 @@ def inplace_fused_experts_impl( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, ) -> None: fused_experts_impl( hidden_states, @@ -1187,7 +1170,6 @@ def inplace_fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) @@ -1210,7 +1192,6 @@ def inplace_fused_experts_impl_fake( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, ) -> None: pass @@ -1242,8 +1223,7 @@ def outplace_fused_experts_impl( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, -) -> torch.Tensor: +) -> None: return fused_experts_impl( hidden_states, w1, @@ -1263,7 +1243,6 @@ def outplace_fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) @@ -1286,8 +1265,7 @@ def outplace_fused_experts_impl_fake( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, -) -> torch.Tensor: +) -> None: return torch.empty_like(hidden_states) @@ -1319,7 +1297,6 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, ): if inplace: torch.ops.lightllm.inplace_fused_experts_impl( @@ -1340,7 +1317,6 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) return hidden_states else: @@ -1362,5 +1338,4 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 5047c67900..61cd81556a 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -7,7 +7,6 @@ from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from .moe_route_weight import multiply_routing_weights, gather_low_latency_routing_weights from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( silu_and_mul_masked_post_quant_fwd, ) @@ -80,8 +79,6 @@ def masked_group_gemm( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, - routing_weights: Optional[torch.Tensor] = None, ): padded_m = recv_x[0].shape[1] E, N, _ = w1.shape @@ -102,8 +99,6 @@ def masked_group_gemm( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, - routing_weights=routing_weights, ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) @@ -226,7 +221,6 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) check_ep_expert_dtype(quant_method) @@ -255,7 +249,6 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) @@ -278,7 +271,6 @@ def fused_experts_impl( limit: Optional[float] = None, clamp_up_add_one: bool = True, block_size_k: int = 128, - glm5_numerics: bool = False, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -348,7 +340,6 @@ def fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) else: gather_out = torch.empty( @@ -370,7 +361,6 @@ def fused_experts_impl( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) _gemm_out_a, _silu_out = None, None del recv_x @@ -397,9 +387,6 @@ def fused_experts_impl( async_finish=False, return_recv_hook=False, ) - routing_weights = None - if glm5_numerics: - routing_weights = gather_low_latency_routing_weights(buffer, topk_idx, topk_weights, handle) # deepgemm gemm_out_b = masked_group_gemm( recv_x, @@ -410,20 +397,13 @@ def fused_experts_impl( w2, w2_scale, expected_m, - routing_weights=routing_weights, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( - gemm_out_b, - topk_idx, - torch.ones_like(topk_weights) if glm5_numerics else topk_weights, - handle, - async_finish=False, - return_recv_hook=False, + gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False ) return combined_x @@ -541,7 +521,6 @@ def chunked_expanded_moe_forward( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, - glm5_numerics: bool = False, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" alignment = 128 @@ -609,16 +588,7 @@ def chunked_expanded_moe_forward( gemm_out_a, m_indices[chunk_start:chunk_end], ) - silu_and_mul_fwd( - gemm_out_a, - silu_out, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - glm5_numerics=glm5_numerics, - ) - if glm5_numerics: - multiply_routing_weights(silu_out, recv_topk_weights[chunk_start:chunk_end]) + silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) workspace_manager.free(gemm_out_a) del gemm_out_a @@ -654,14 +624,7 @@ def workspace_quant_alloc(shape, dtype, device): for quant_buffer in quant_buffers: workspace_manager.free(quant_buffer) - ep_gather_chunk( - gemm_out_b, - chunk_start, - recv_topk_weights, - recv_src_metadata, - gather_out, - apply_routed_weight=not glm5_numerics, - ) + ep_gather_chunk(gemm_out_b, chunk_start, recv_topk_weights, recv_src_metadata, gather_out) workspace_manager.free(gemm_out_b) finally: if is_autotune_warmup: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py deleted file mode 100644 index c10769a734..0000000000 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_route_weight.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Apply GLM routing probabilities at the BF16 W2 input boundary.""" - -import torch -import torch.distributed as dist -import triton -import triton.language as tl - - -@triton.jit -def _multiply_rows(X, P, ROWS: tl.constexpr, COLS: tl.constexpr, BLOCK: tl.constexpr): - row = tl.program_id(0) - cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - x = tl.load(X + row * COLS + cols, cols < COLS, 0).to(tl.float32) - p = tl.load(P + row) - tl.store(X + row * COLS + cols, x * p, cols < COLS) - - -def multiply_routing_weights(x, weights): - """x is already rounded to BF16 by SwiGLU; round x.float()*p to BF16 again.""" - assert x.is_contiguous() and weights.is_contiguous() - assert x.numel() // x.shape[-1] == weights.numel() - _multiply_rows[(weights.numel(), triton.cdiv(x.shape[-1], 256))]( - x, weights, weights.numel(), x.shape[-1], 256, num_warps=4 - ) - - -@triton.jit -def _sort_tp_weights(Blocks, Counts, Weights, Out, STRIDE: tl.constexpr, BLOCK_M: tl.constexpr): - block = tl.program_id(0) - expert = tl.load(Blocks + block * 3) - if expert >= 0: - tile = tl.load(Blocks + block * 3 + 1) - start = tl.load(Blocks + block * 3 + 2) - lanes = tl.arange(0, BLOCK_M) - count = tl.load(Counts + expert) - p = tl.load(Weights + expert * STRIDE + tile * BLOCK_M + lanes, tile * BLOCK_M + lanes < count, 0) - tl.store(Out + start + lanes, p, tile * BLOCK_M + lanes < count) - - -def sorted_tp_routing_weights(topk_weights, expert_weights, counts, mblock_infos): - """Reorder p exactly like the TMA W13 output, including non-aligned expert tails.""" - out = torch.empty_like(topk_weights).view(-1) - blocks, block_m = mblock_infos - _sort_tp_weights[(blocks.shape[0],)]( - blocks, counts, expert_weights, out, expert_weights.stride(0), block_m, num_warps=4 - ) - return out - - -@triton.jit -def _gather_low_latency_weights( - Routes, - SrcInfo, - Layout, - Out, - NUM_RANKS: tl.constexpr, - MAX_TOKENS: tl.constexpr, - TOPK: tl.constexpr, - LOCAL_EXPERTS: tl.constexpr, - RANK: tl.constexpr, - BLOCK: tl.constexpr, - BLOCK_TOPK: tl.constexpr, -): - expert, src_rank = tl.program_id(0), tl.program_id(1) - layout = tl.load(Layout + expert * NUM_RANKS + src_rank) - count = (layout & 0xFFFFFFFF).to(tl.int32) - start = (layout >> 32).to(tl.int32) - k = tl.arange(0, BLOCK_TOPK) - for offset in range(0, count, BLOCK): - rows = offset + tl.arange(0, BLOCK) - valid = rows < count - dst = expert * (NUM_RANKS * MAX_TOKENS) + start + rows - src_token = tl.load(SrcInfo + dst, valid, 0) - ptr = Routes + ((src_rank * MAX_TOKENS + src_token[:, None]) * TOPK + k[None, :]) * 2 - ids = tl.load(ptr, valid[:, None] & (k[None, :] < TOPK), -1).to(tl.int32) - weights = tl.load(ptr + 1, valid[:, None] & (k[None, :] < TOPK), 0) - p = tl.sum(tl.where(ids == RANK * LOCAL_EXPERTS + expert, weights, 0), 1) - tl.store(Out + dst, p, valid) - - -def gather_low_latency_routing_weights(buffer, topk_idx, topk_weights, handle): - """Recover p using DeepEP's actual receive order, without CPU count synchronization. - - Legacy dispatch transports X but not p. All-gather only the small padded - FP32 (expert id, p) table, then use its src_info/layout_range handle. All - ranks participate, including empty/HOLD ranks, so CUDA graph replay has a - fixed collective shape. Expert ids are exactly representable in FP32. - """ - src_info, layout, max_tokens, _, num_experts = handle - ranks, topk = buffer.group_size, topk_idx.shape[1] - assert num_experts < 2**24 and topk_idx.shape[0] <= max_tokens - assert topk_idx.shape == topk_weights.shape and topk_weights.dtype == torch.float32 - assert src_info.shape == (num_experts // ranks, ranks * max_tokens) - assert layout.shape == (num_experts // ranks, ranks) - assert src_info.dtype == torch.int32 and layout.dtype == torch.int64 - routes = torch.zeros((max_tokens, topk, 2), dtype=torch.float32, device=topk_weights.device) - routes[:, :, 0].fill_(-1) - routes[: topk_idx.shape[0], :, 0].copy_(topk_idx) - routes[: topk_idx.shape[0], :, 1].copy_(topk_weights) - if ranks == 1: - all_routes = routes - else: - all_routes = torch.empty((ranks * max_tokens, topk, 2), dtype=routes.dtype, device=routes.device) - dist.all_gather_into_tensor(all_routes, routes, group=buffer.group) - out = torch.zeros(src_info.shape, dtype=torch.float32, device=topk_weights.device) - _gather_low_latency_weights[(num_experts // ranks, ranks)]( - all_routes, - src_info, - layout, - out, - ranks, - max_tokens, - topk, - num_experts // ranks, - buffer.rank, - 32, - triton.next_power_of_2(topk), - num_warps=4, - ) - return out diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index f0771cb1a8..8169cefc3c 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py @@ -26,7 +26,6 @@ def _silu_and_mul_kernel_fast( USE_LIMIT_AND_ALPHA: tl.constexpr = False, CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, - GLM5_NUMERICS: tl.constexpr = False, ): stride_input_m = tl.cast(stride_input_m, dtype=tl.int64) stride_output_m = tl.cast(stride_output_m, dtype=tl.int64) @@ -67,14 +66,11 @@ def _silu_and_mul_kernel_fast( other=other, ).to(tl.float32) - if GLM5_NUMERICS: - up = up.to(tl.float32) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) gate = 1 / (1 + tl.exp(-gate * alpha)) * gate - if not GLM5_NUMERICS: - gate = gate.to(input_ptr.dtype.element_ty) + gate = gate.to(input_ptr.dtype.element_ty) if CLAMP_UP_ADD_ONE: up += 1 tl.store(output_ptr + out_offsets, up * gate, mask=mask) @@ -87,8 +83,7 @@ def _silu_and_mul_kernel_fast( gate = 0.5 * gate * (1.0 + tanh_val) else: gate = gate / (1 + tl.exp(-gate)) - if not GLM5_NUMERICS: - gate = gate.to(input_ptr.dtype.element_ty) + gate = gate.to(input_ptr.dtype.element_ty) tl.store( output_ptr + out_offsets, @@ -126,7 +121,6 @@ def silu_and_mul_fwd( alpha=None, run_config=None, clamp_up_add_one=True, - glm5_numerics=False, ): assert input.stride(-1) == 1 assert output.is_contiguous() @@ -179,6 +173,5 @@ def silu_and_mul_fwd( USE_LIMIT_AND_ALPHA=USE_LIMIT_AND_ALPHA, CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), - GLM5_NUMERICS=glm5_numerics, ) return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 69d371d4cd..8c94f088c2 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -22,8 +22,6 @@ def _silu_and_mul_post_quant_kernel( stride_output_scale_1, stride_output_scale_2, masked_m_ptr, - routing_weights, - ROUTING_STRIDE: tl.constexpr, size_n, fp8_max, fp8_min, @@ -35,8 +33,6 @@ def _silu_and_mul_post_quant_kernel( limit: tl.constexpr = None, CLAMP_UP_ADD_ONE: tl.constexpr = True, USE_UE8M0_SCALE: tl.constexpr = False, - GLM5_NUMERICS: tl.constexpr = False, - WEIGHTED_INPUT: tl.constexpr = False, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -59,12 +55,8 @@ def _silu_and_mul_post_quant_kernel( ) for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): - gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to( - tl.float32 - ) + gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) - if GLM5_NUMERICS: - up = up.to(tl.float32) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) @@ -78,15 +70,11 @@ def _silu_and_mul_post_quant_kernel( gate = 0.5 * gate * (1.0 + tanh_val) else: gate = gate / (1 + tl.exp(-gate)) - if not GLM5_NUMERICS: - gate = gate.to(input_ptr.dtype.element_ty) + gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate - if USE_LIMIT_AND_ALPHA or GLM5_NUMERICS: + if USE_LIMIT_AND_ALPHA: # Match the BF16/FP16 activation stored before prefill's FP8 quantization. gate_up = gate_up.to(input_ptr.dtype.element_ty).to(tl.float32) - if WEIGHTED_INPUT: - p = tl.load(routing_weights + expert_id * ROUTING_STRIDE + token_index) - gate_up = (gate_up * p).to(input_ptr.dtype.element_ty).to(tl.float32) _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) output_s = _absmax / fp8_max if USE_UE8M0_SCALE: @@ -112,8 +100,6 @@ def silu_and_mul_masked_post_quant_fwd( alpha=None, limit=None, clamp_up_add_one=True, - glm5_numerics=False, - routing_weights=None, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -131,10 +117,6 @@ def silu_and_mul_masked_post_quant_fwd( assert input.shape[0] == masked_m.shape[0] assert input.shape[-1] % 2 == 0 - if routing_weights is not None: - assert routing_weights.shape == input.shape[:2] and routing_weights.is_contiguous() - assert routing_weights.dtype == torch.float32 - size_n = input.shape[-1] // 2 assert size_n % quant_group_size == 0 @@ -169,8 +151,6 @@ def silu_and_mul_masked_post_quant_fwd( output_scale, *output_scale.stride(), masked_m, - routing_weights, - routing_weights.stride(0) if routing_weights is not None else 0, size_n, fp8_max, fp8_min, @@ -182,9 +162,6 @@ def silu_and_mul_masked_post_quant_fwd( limit=limit, CLAMP_UP_ADD_ONE=clamp_up_add_one, USE_UE8M0_SCALE=get_env_use_ue8m0_scales(), - GLM5_NUMERICS=glm5_numerics, - enable_fp_fusion=not glm5_numerics, - WEIGHTED_INPUT=routing_weights is not None, num_warps=num_warps, ) return diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index f13d447066..2032e50c96 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -58,7 +58,7 @@ def __init__(self, layer_num, network_config): self.q_lora_rank = network_config["q_lora_rank"] self.kv_lora_rank = network_config["kv_lora_rank"] self.v_head_dim = network_config["v_head_dim"] - self.softmax_scale = self.qk_nope_head_dim**-0.5 + self.softmax_scale = self.qk_nope_head_dim ** -0.5 self.indexer = Glm5NextNsaInfer( layer_idx=self.layer_num_, network_config=self.network_config_, @@ -94,7 +94,6 @@ def _ffn_tp(self, input, infer_state, layer_weight): limit=self.swiglu_limit, alpha=1.0, clamp_up_add_one=False, - glm5_numerics=True, ) return layer_weight.down_proj.mm(ffn1_out) @@ -435,9 +434,7 @@ def _hc_pre(self, streams, layer_weight, prefix, norm_weight): def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill): streams = input_embeddings - layer_input, residual_mix, post_mix = self._hc_pre( - streams, layer_weight, "attn", layer_weight.att_norm_weight_ - ) + layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "attn", layer_weight.att_norm_weight_) if prefill: layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight) else: diff --git a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py index b193ada4e7..32191809fe 100644 --- a/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/glm5_next/layer_weights/transformer_layer_weight.py @@ -212,7 +212,6 @@ def _init_moe(self): num_fused_shared_experts=self.num_fused_shared_experts, layer_num=self.layer_num_, network_config=self.network_config_, - glm5_numerics=True, ) def _init_kda(self): diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py index d4ac0f3935..e9ba2ab26e 100644 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -11,13 +11,10 @@ ) -@pytest.mark.parametrize("glm5_numerics", [False, True]) @pytest.mark.parametrize("is_prefill", [False, True]) @pytest.mark.parametrize("use_ue8m0_scales", [False, True]) @pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) -def test_grouped_gemm_activation_matches_reference( - monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one, glm5_numerics -): +def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one): monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES") config = {"quant_method": "fp8", "weight_block_size": [128, 128]} @@ -33,10 +30,7 @@ def test_grouped_gemm_activation_matches_reference( w1_scale = torch.ones(2, 2, 1, device="cuda") w2_scale = torch.ones(2, 1, 1, device="cuda") counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) - activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one, glm5_numerics=glm5_numerics) - p = torch.tensor([0.13, 0.71, 0.88, 1.0, 1.3], device="cuda") if glm5_numerics else torch.ones(5, device="cuda") - recv_p = torch.zeros(2, padded, device="cuda") - recv_p[0, :3], recv_p[1, :2] = p[:3], p[3:] + activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) recv[0, :3] = x[:3] recv[1, :2] = x[3:] @@ -50,7 +44,7 @@ def test_grouped_gemm_activation_matches_reference( num_recv_tokens_per_expert_list=[padded, padded], num_unaligned_recv_tokens_per_expert=counts, recv_x=(recv.view(2 * padded, dim), torch.ones(1, 2 * padded, device="cuda").T), - recv_topk_weights=recv_p.view(-1), + recv_topk_weights=torch.ones(2 * padded, device="cuda"), recv_src_metadata=metadata, w1=w1, w1_scale=w1_scale, @@ -71,7 +65,6 @@ def test_grouped_gemm_activation_matches_reference( w2, w2_scale, expected_m=3, - routing_weights=recv_p if glm5_numerics else None, **activation_args, ) actual = torch.cat([output[0, :3], output[1, :2]]) @@ -81,11 +74,8 @@ def test_grouped_gemm_activation_matches_reference( if limit is not None: gate = gate.clamp(max=limit) up = up.clamp(-limit, limit) + int(add_one) - gate = gate * torch.sigmoid(gate * (alpha or 1.0)) - if not glm5_numerics: - gate = gate.bfloat16().float() + gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() activation = (gate * up).bfloat16().float() - activation = (activation * p[:, None]).bfloat16().float() scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 if use_ue8m0_scales: scales = torch.exp2(torch.ceil(torch.log2(scales))) diff --git a/unit_tests/common/fused_moe/test_glm53_training_alignment.py b/unit_tests/common/fused_moe/test_glm53_training_alignment.py deleted file mode 100644 index 6eb8c63176..0000000000 --- a/unit_tests/common/fused_moe/test_glm53_training_alignment.py +++ /dev/null @@ -1,323 +0,0 @@ -"""GLM launch recipe: eager BF16 boundaries, FP32 SwiGLU and p before W2.""" - -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F - -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( - silu_and_mul_masked_post_quant_fwd, -) -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_route_weight import ( - multiply_routing_weights, - gather_low_latency_routing_weights, - _gather_low_latency_weights, -) - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def swiglu_reference(x): - gate, up = x.float().chunk(2, -1) - return (F.silu(gate.clamp(max=10)) * up.clamp(-10, 10)).to(x.dtype) - - -def quant_reference(x, ue8m0): - groups = x.float().reshape(*x.shape[:-1], -1, 128) - sf = groups.abs().amax(-1).clamp_min(1e-10) / 448 - if ue8m0: - sf = torch.exp2(torch.ceil(torch.log2(sf))) - q = (groups / sf.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) - return q.reshape_as(x), sf - - -@pytest.mark.parametrize("rows,width", [(1, 256), (19, 4096)]) -def test_swiglu_retains_fp32_product_and_bf16_weight_boundary(rows, width): - torch.manual_seed(53) - x = (torch.randn(rows, width, device="cuda") * 6).bfloat16() - out = torch.empty(rows, width // 2, device=x.device, dtype=x.dtype) - silu_and_mul_fwd(x, out, alpha=1.0, limit=10.0, clamp_up_add_one=False, glm5_numerics=True) - expected = swiglu_reference(x) - # CUDA exp implementations differ near BF16 midpoints. - torch.testing.assert_close(out, expected, atol=1e-6, rtol=0.008) - old = torch.empty_like(out) - silu_and_mul_fwd(x, old, alpha=1.0, limit=10.0, clamp_up_add_one=False) - assert torch.count_nonzero(old != expected) > torch.count_nonzero(out != expected) - p = torch.linspace(0.01, 1.7, rows, device=x.device) - before = out.clone() - multiply_routing_weights(out, p) - torch.testing.assert_close(out, (before.float() * p[:, None]).to(x.dtype), atol=0, rtol=0) - - -@pytest.mark.parametrize("ue8m0", [False, True]) -def test_masked_decode_swiglu_matches_training_before_quant(monkeypatch, ue8m0): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) - torch.manual_seed(53) - x = (torch.randn(3, 5, 4096, device="cuda") * 6).bfloat16() - p = torch.rand(3, 5, device=x.device) * 1.7 - p[1, 0] = 0 - counts = torch.tensor([0, 2, 5], device=x.device, dtype=torch.int32) - out = torch.full((3, 5, 2048), 1, device=x.device, dtype=torch.float8_e4m3fn) - sf = torch.full((3, 5, 16), 7.0, device=x.device) - silu_and_mul_masked_post_quant_fwd( - x, - out, - sf, - 128, - counts, - alpha=1.0, - limit=10.0, - clamp_up_add_one=False, - glm5_numerics=True, - routing_weights=p, - ) - weighted = (swiglu_reference(x).float() * p[..., None]).to(x.dtype) - q_ref, s_ref = quant_reference(weighted, ue8m0) - for expert, count in enumerate(counts.tolist()): - torch.testing.assert_close(sf[expert, :count], s_ref[expert, :count], atol=1e-6, rtol=1e-5) - actual = out[expert, :count].float() * sf[expert, :count].repeat_interleave(128, -1) - expected = q_ref[expert, :count].float() * s_ref[expert, :count].repeat_interleave(128, -1) - # A SwiGLU midpoint may flip one FP8 rounding decision. - err = (actual - expected).square().mean().sqrt() if count else actual.sum() - assert err < 0.01 - assert torch.all(out[expert, count:].float() == 1) - assert torch.all(sf[expert, count:] == 7) - - -@pytest.mark.parametrize("empty_local", [False, True]) -def test_low_latency_routing_uses_source_and_actual_receive_order(monkeypatch, empty_local): - from lightllm.common.basemodel.triton_kernel.fused_moe import moe_route_weight as route - - # Two ranks, unequal token counts; expert receive order differs by source rank. - max_tokens, ranks, local_experts, topk = 4, 2, 2, 2 - ids = torch.tensor([[[2, 3], [0, 2], [-1, -1], [-1, -1]], [[3, 2], [2, 1], [2, 3], [-1, -1]]], device="cuda") - p = torch.tensor( - [[[0.1, 0.9], [0.2, 0.8], [0, 0], [0, 0]], [[0.3, 0.7], [0.4, 0.6], [0.5, 0.5], [0, 0]]], device="cuda" - ) - if empty_local: - ids[1].fill_(-1) - p[1].zero_() - routes = torch.stack((ids.float(), p), -1).reshape(ranks * max_tokens, topk, 2) - src_info = torch.zeros((local_experts, ranks * max_tokens), device="cuda", dtype=torch.int32) - layout = torch.zeros((local_experts, ranks), device="cuda", dtype=torch.int64) - expected = torch.zeros_like(src_info, dtype=torch.float32) - rank = 1 - for expert in range(local_experts): - offset = 0 - for src_rank in (1, 0): - tokens = torch.where((ids[src_rank] == rank * local_experts + expert).any(-1))[0].flip(0) - count = tokens.numel() - src_info[expert, offset : offset + count] = tokens.int() - layout[expert, src_rank] = (offset << 32) | count - for row, token in enumerate(tokens.tolist()): - slot = (ids[src_rank, token] == rank * local_experts + expert).nonzero()[0, 0] - expected[expert, offset + row] = p[src_rank, token, slot] - offset += count - calls = [] - - def gather(out, inp, group): - calls.append(inp.clone()) - out.copy_(routes) - - monkeypatch.setattr(route.dist, "all_gather_into_tensor", gather) - buffer = SimpleNamespace(group_size=ranks, rank=rank, group=object()) - handle = (src_info, layout, max_tokens, 256, ranks * local_experts) - actual = gather_low_latency_routing_weights( - buffer, ids[1, : 0 if empty_local else 3], p[1, : 0 if empty_local else 3], handle - ) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - assert len(calls) == 1 - assert torch.all(calls[0][3, :, 0] == -1) - # The metadata mapping is graph replayable and reads changed inputs. - out = torch.zeros_like(expected) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(2): - _gather_low_latency_weights[(local_experts, ranks)]( - routes, src_info, layout, out, ranks, max_tokens, topk, local_experts, rank, 32, 2 - ) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - _gather_low_latency_weights[(local_experts, ranks)]( - routes, src_info, layout, out, ranks, max_tokens, topk, local_experts, rank, 32, 2 - ) - routes[:, :, 1].mul_(2) - graph.replay() - torch.testing.assert_close(out, expected * 2, atol=0, rtol=0) - - -@pytest.mark.parametrize("tokens", [2, 19, 257]) -@pytest.mark.parametrize("tma", [False, True]) -@pytest.mark.parametrize( - "quantized,ue8m0,packed", [(False, False, False), (True, False, False), (True, True, False), (True, True, True)] -) -def test_tp_moe_routes_before_w2_with_tma_sorting_and_shared_expert( - monkeypatch, tma, quantized, ue8m0, packed, tokens -): - from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe as tp - - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0 and not packed))) - monkeypatch.setattr(tp, "triton_support_tensor_descriptor", lambda: tma) - torch.manual_seed(53) - dim, experts = 128, 4 - x = (torch.randn(tokens, dim, device="cuda") * 4).bfloat16() - ids = torch.stack( - ( - torch.arange(tokens, device=x.device) % 3, - (torch.arange(tokens, device=x.device) + 1) % 3, - torch.full((tokens,), 3, device=x.device), - ), - -1, - ).int() - p = torch.rand(tokens, 3, device=x.device) * 1.7 - p[:, -1] = 1 # Fused shared expert must remain unweighted. - eye = torch.eye(dim, device=x.device, dtype=x.dtype) - w1 = torch.cat((eye, eye)).repeat(experts, 1, 1) - w2 = eye.repeat(experts, 1, 1) - kwargs = {} - - def quantize(x): - # Ordinary scales use the installed SGL quantizer; its reciprocal - # multiplication can resolve FP8 midpoints differently from Torch's - # division. Keep that unchanged stage fixed while testing W2 placement. - if ue8m0: - return quant_reference(x, True) - from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 - - q, sf = per_token_group_quant_fp8(x.reshape(-1, dim), 128) - return q.reshape_as(x), sf.reshape(*x.shape[:-1], 1) - - if quantized: - w1, w2 = w1.to(torch.float8_e4m3fn), w2.to(torch.float8_e4m3fn) - kwargs = dict( - w1_scale=torch.ones(experts, 2, 1, device=x.device), w2_scale=torch.ones(experts, 1, 1, device=x.device) - ) - if packed: - # One K group: exponent 127 occupies byte 0, per output channel. - # Leave the env off to check packed B itself implies UE8M0 on A. - kwargs = dict( - w1_scale=torch.full((experts, 256, 1), 127, device=x.device, dtype=torch.int32), - w2_scale=torch.full((experts, 128, 1), 127, device=x.device, dtype=torch.int32), - ) - q, sf = quantize(x) - pre = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() - else: - pre = x - h = swiglu_reference(torch.cat((pre, pre), -1)) - h = (h[:, None, :].float() * p[..., None]).bfloat16() - if quantized: - q, sf = quantize(h) - h = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() - expected = h.float().sum(1).bfloat16() - actual = tp.fused_experts( - x.clone(), - w1, - w2, - p, - ids, - use_fp8_w8a8=quantized, - alpha=1.0, - limit=10.0, - clamp_up_add_one=False, - glm5_numerics=True, - **kwargs, - ) - torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.008) - legacy = tp.fused_experts( - x.clone(), - w1, - w2, - p, - ids, - use_fp8_w8a8=quantized, - alpha=1.0, - limit=10.0, - clamp_up_add_one=False, - **kwargs, - ) - assert torch.count_nonzero(legacy != expected) > torch.count_nonzero(actual != expected) - - -@pytest.mark.parametrize("ue8m0", [False, True]) -def test_ep_decode_dispatch_to_combine_weights_exactly_once(monkeypatch, ue8m0): - from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep - - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) - monkeypatch.setattr(ep, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 128) - torch.manual_seed(53) - tokens, dim, experts = 5, 128, 2 - x = (torch.randn(tokens, dim, device="cuda") * 3).bfloat16() - ids = torch.tensor([[0, 1]] * tokens, device=x.device, dtype=torch.int64) - p = torch.rand(tokens, 2, device=x.device) * 1.7 - orders = torch.tensor([[4, 1, 3, 0, 2], [2, 4, 1, 3, 0]], device=x.device) - eye = torch.eye(dim, device=x.device, dtype=x.dtype) - w1 = torch.cat((2 * eye, 3 * eye)).repeat(experts, 1, 1).to(torch.float8_e4m3fn) - w2 = eye.repeat(experts, 1, 1).to(torch.float8_e4m3fn) - seen = [] - - class Buffer: - group_size, rank = 1, 0 - - def low_latency_dispatch(self, hidden, topk_idx, max_tokens, num_experts, **kwargs): - assert kwargs["round_scale"] == ue8m0 - q, sf = quant_reference(hidden, ue8m0) - recv = torch.zeros(experts, max_tokens, dim, device=x.device, dtype=q.dtype) - scales = torch.zeros(experts, max_tokens, 1, device=x.device) - src = torch.zeros(experts, max_tokens, device=x.device, dtype=torch.int32) - for e in range(experts): - recv[e, :tokens] = q[orders[e]] - scales[e, :tokens] = sf[orders[e]] - src[e, :tokens] = orders[e].int() - counts = torch.full((experts,), tokens, device=x.device, dtype=torch.int32) - layout = torch.full((experts, 1), tokens, device=x.device, dtype=torch.int64) - return (recv, scales), counts, (src, layout, max_tokens, dim, experts), None, None - - def low_latency_combine(self, output, topk_idx, topk_weights, handle, **kwargs): - torch.testing.assert_close(topk_weights, torch.ones_like(p), atol=0, rtol=0) - seen.append(True) - result = torch.zeros(tokens, dim, device=x.device) - for e in range(experts): - result[orders[e]] += output[e, :tokens].float() - return result.bfloat16(), None, None - - actual = ep.fused_experts_impl( - x, - w1, - w2, - p, - ids, - experts, - Buffer(), - False, - use_fp8_w8a8=True, - use_fp8_all2all=True, - w1_scale=torch.ones(experts, 2, 1, device=x.device), - w2_scale=torch.ones(experts, 1, 1, device=x.device), - alpha=1.0, - limit=10.0, - clamp_up_add_one=False, - glm5_numerics=True, - ) - q, sf = quant_reference(x, ue8m0) - pre = (q.float() * sf.repeat_interleave(128, -1)).bfloat16() - h = swiglu_reference(torch.cat(((2 * pre.float()).bfloat16(), (3 * pre.float()).bfloat16()), -1)) - weighted = (h[:, None, :].float() * p[..., None]).bfloat16() - q, sf = quant_reference(weighted, ue8m0) - expected = (q.float() * sf.repeat_interleave(128, -1)).bfloat16().float().sum(1).bfloat16() - if ue8m0: - torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.008) - else: - # Arbitrary scales use reciprocal/division implementations with distinct - # FP8 midpoint ties. Bound the difference by one FP8 step per route, - # plus the final BF16 sum rounding; UE8M0 keeps the strict comparison. - q_float = q.float().abs().clamp_min(2.0**-6) - fp8_step = torch.exp2(torch.floor(torch.log2(q_float)) - 3) * sf.repeat_interleave(128, -1) - bound = fp8_step.sum(1) + expected.float().abs() * torch.finfo(torch.bfloat16).eps - assert torch.all((actual.float() - expected.float()).abs() <= bound + 1e-5) - nrmse = (actual.float() - expected.float()).square().mean().sqrt() / expected.float().square().mean().sqrt() - assert nrmse < 0.008 - assert seen == [True] diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index ee087b05e7..7b60c9a267 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -95,16 +95,7 @@ def test_kda_projections_need_no_sparse_attention_config(config): weight = Glm5NextTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) projections = { name: random_weight(size, 128) - for name, size in { - "q": 512, - "k": 512, - "v": 512, - "b": 4, - "f_a": 128, - "g_a": 128, - "f_b": 512, - "g_b": 512, - }.items() + for name, size in {"q": 512, "k": 512, "v": 512, "b": 4, "f_a": 128, "g_a": 128, "f_b": 512, "g_b": 512}.items() } weight.load_hf_weights( {f"model.language_model.layers.0.self_attn.{name}_proj.weight": value for name, value in projections.items()} @@ -235,11 +226,10 @@ def mlp_weights(name, gate_scale, up_scale): def reference_mlp(gate_scale, up_scale): gate = (x * gate_scale).float().clamp(max=10) up = (x * up_scale).float().clamp(-10, 10) - return (F.silu(gate) * up).bfloat16() + return (F.silu(gate).bfloat16().float() * up).bfloat16() expected = reference_mlp(1.5, 0.75) if mode != "dense": - assert weight.experts.fuse_moe_impl.glm5_numerics scores = F.linear(x.float(), router).sigmoid() probabilities, ids = scores.topk(2, dim=-1) probabilities /= probabilities.sum(-1, keepdim=True) @@ -280,7 +270,7 @@ def combined_output(hidden_states, **kwargs): monkeypatch.setattr(weight.experts, "experts", combined_output) actual = layer._ffn(x, SimpleNamespace(is_prefill=is_prefill), weight) - shared = (F.silu(x.float().clamp(max=10)) * x.float().clamp(-10, 10)).bfloat16() + shared = (F.silu(x.float().clamp(max=10)).bfloat16().float() * x.float().clamp(-10, 10)).bfloat16() torch.testing.assert_close(actual, shared + 17, atol=0.25, rtol=0.01) torch.testing.assert_close(x, original, atol=0, rtol=0) From 6d219d3c43b26c160dee7c7952b89e3c78f2b20d Mon Sep 17 00:00:00 2001 From: baishihao Date: Sat, 24 Oct 2026 00:16:53 +0800 Subject: [PATCH 11/13] feat(glm5): add opt-in BF16 indexer precision ablation --- .../glm5_next_mem_manager.py | 4 + .../common/state_cache_manager/glm5_next.py | 10 +- lightllm/models/glm5_next/indexer.py | 50 +++++- .../glm5_next/triton_kernel/index_bf16.py | 76 +++++++++ .../models/glm5_next/triton_kernel/kpool.py | 48 ++++-- lightllm/utils/envs_utils.py | 5 + unit_tests/models/glm5_next/test_cache.py | 24 ++- .../models/glm5_next/test_indexer_bf16.py | 153 ++++++++++++++++++ unit_tests/models/glm5_next/test_kernels.py | 38 +++-- unit_tests/models/glm5_next/test_mtp.py | 7 +- unit_tests/models/glm5_next/test_pd_cache.py | 16 +- 11 files changed, 384 insertions(+), 47 deletions(-) create mode 100644 lightllm/models/glm5_next/triton_kernel/index_bf16.py create mode 100644 unit_tests/models/glm5_next/test_indexer_bf16.py diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py index 45a3be667c..1779fabf0c 100644 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -39,6 +39,10 @@ def get_att_input_params(self, layer_index): return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] def get_indexer_k_buffer(self, layer_index): + if self.linear_config.indexer_bf16: + return self._layer_buffer(layer_index)[ + :, :, self.mla_head_dim : self.mla_head_dim + self.linear_config.index_head_dim + ] return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] def write_to_shm(self, req_manager): diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 13195a74ac..4c833dd15d 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -1,6 +1,6 @@ import dataclasses -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_env_glm5_indexer_bf16 from lightllm.utils.torch_dtype_utils import get_torch_dtype from .linear_att import LinearAttCacheConfig @@ -12,6 +12,7 @@ class Glm5NextCacheConfig(LinearAttCacheConfig): index_kpool: int = 4 index_head_dim: int = 128 + indexer_bf16: bool = False INDEX_PADDING_BYTES = 144 @@ -26,8 +27,12 @@ def from_model_config(cls, config, args): "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) ] dtype = get_torch_dtype(args.data_type) + indexer_bf16 = get_env_glm5_indexer_bf16() + if indexer_bf16: + assert args.data_type in ("bf16", "bfloat16"), "BF16 indexer requires --data_type bfloat16" # Raw index keys and compression scores only live in request tails. - packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize + index_bytes = config["index_head_dim"] * dtype.itemsize if indexer_bf16 else cls.INDEX_PADDING_BYTES + packed_dim = config["kv_lora_rank"] + index_bytes // dtype.itemsize return cls( tp_world_size=tp, full_att_all_num_kv_heads=1, @@ -49,6 +54,7 @@ def from_model_config(cls, config, args): draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], + indexer_bf16=indexer_bf16, ) def get_cpu_cache_full_att_bytes(self): diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 30a9883881..652e848c1b 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -2,9 +2,18 @@ import triton from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops +from lightllm.utils.envs_utils import get_env_glm5_indexer_bf16 +from .triton_kernel.index_bf16 import bf16_mqa_logits from .triton_kernel.index_quant import hadamard_transform_quant_fp8 -from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk +from .triton_kernel.kpool import ( + compress_pools, + gather_pools, + gather_bf16_pools, + gather_paged_pools, + get_pool_ranges, + expand_topk, +) class Glm5NextNsaInfer: @@ -17,6 +26,7 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.dim = network_config["index_head_dim"] # The indexer LayerNorm has a separate epsilon from the model RMSNorm. self.eps = 1e-6 + self.indexer_bf16 = get_env_glm5_indexer_bf16() def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" @@ -39,6 +49,7 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig raw = torch.cat((k, gate), -1) tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + assert (packed_buffer.dtype == torch.bfloat16) == self.indexer_bf16, "Indexer/cache precision mismatch" compress_pools( raw=raw, tail=tail, @@ -55,15 +66,24 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig ) if infer_state.max_kv_seq_len <= self.topk: - return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True) + return expand_topk( + None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True + ) # The small indexer is replicated: no all-gather of query heads and # identical pool selection on every TP rank. q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.heads, self.dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim ** -0.5) - weights = layer_weight.weights_proj_.mm(hidden_states.float()) - weights = weights * (self.heads ** -0.5 * self.dim ** -0.5) * q_scale.squeeze(-1) max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 + if self.indexer_bf16: + # Keep the existing projections and pooling arithmetic: only remove + # the rotated FP8 Q/K representation and FP8 scoring from this A/B. + weights = layer_weight.weights_proj_.mm(hidden_states.float()) * self.heads**-0.5 + groups = self._get_bf16_indices(q, weights, packed_buffer, infer_state, att_state, max_pools) + return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + + q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim**-0.5) + weights = layer_weight.weights_proj_.mm(hidden_states.float()) + weights = weights * (self.heads**-0.5 * self.dim**-0.5) * q_scale.squeeze(-1) # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. # Decode has few queries, so paged MQA also splits K across SMs for parallelism. if infer_state.is_prefill: @@ -72,6 +92,26 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + def _get_bf16_indices(self, q, weights, packed_buffer, infer_state, att_state, max_pools): + # Decode uses one request entry per MTP position; its visible length is + # per row, so future verify pools and HOLD rows remain masked. + seq_lens = infer_state.b_seq_len if infer_state.is_prefill else att_state.lengths + keys = gather_bf16_pools( + packed_buffer, infer_state.req_manager.req_to_token_indexs, infer_state.b_req_idx, seq_lens, max_pools + ) + starts, ends, lengths = get_pool_ranges( + att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools + ) + groups = torch.empty((q.shape[0], self.topk // 4), dtype=torch.int32, device=q.device) + chunk_size = max(1, min(q.shape[0], 16 * 1024 * 1024 // max_pools)) + for start in range(0, q.shape[0], chunk_size): + end = min(start + chunk_size, q.shape[0]) + logits = bf16_mqa_logits( + q[start:end], keys, weights[start:end], starts[start:end], ends[start:end], max_pools + ) + self.select_topk_indices(logits, lengths[start:end], groups[start:end]) + return groups + def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): keys = gather_pools( packed_buffer, diff --git a/lightllm/models/glm5_next/triton_kernel/index_bf16.py b/lightllm/models/glm5_next/triton_kernel/index_bf16.py new file mode 100644 index 0000000000..d17ff0e50c --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/index_bf16.py @@ -0,0 +1,76 @@ +"""Unrotated BF16 Q/K with IEEE FP32 scoring for the indexer precision ablation.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _bf16_mqa_logits( + Q, + K, + Weights, + Starts, + Ends, + Out, + Q_STRIDE: tl.constexpr, + Q_HEAD_STRIDE: tl.constexpr, + W_STRIDE: tl.constexpr, + POOLS: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_P: tl.constexpr, +): + row = tl.program_id(1) + pools = tl.program_id(0) * BLOCK_P + tl.arange(0, BLOCK_P) + heads = tl.arange(0, BLOCK_H) + dims = tl.arange(0, 32) + start = tl.load(Starts + row).to(tl.int64) + length = tl.load(Ends + row).to(tl.int64) - start + valid = (pools < POOLS) & (pools < length) + # One query per program preserves request boundaries in ragged/MTP batches. + # IEEE FP32 avoids TF32 settings changing the precision experiment. + dots = tl.full((BLOCK_H, BLOCK_P), 0, tl.float32) + for block in range(tl.cdiv(DIM, 32)): + cols = block * 32 + dims + q = tl.load( + Q + row * Q_STRIDE + heads[:, None] * Q_HEAD_STRIDE + cols[None, :], + (heads[:, None] < HEADS) & (cols[None, :] < DIM), + 0, + ).to(tl.float32) + k = tl.load( + K + (start + pools[None, :]) * DIM + cols[:, None], + valid[None, :] & (cols[:, None] < DIM), + 0, + ).to(tl.float32) + dots = tl.dot(q, k, dots, input_precision="ieee") + weights = tl.load(Weights + row * W_STRIDE + heads, heads < HEADS, 0).to(tl.float32) + scores = tl.sum(tl.maximum(dots * (DIM**-0.5), 0) * weights[:, None], 0) + tl.store(Out + row * POOLS + pools, tl.where(valid, scores, -float("inf")), pools < POOLS) + + +def bf16_mqa_logits(q, keys, weights, starts, ends, max_pools): + """Return row-relative logits; weights already include heads**-0.5.""" + assert q.dtype == keys.dtype == torch.bfloat16 + assert keys.is_contiguous() and weights.stride(1) == 1 + logits = torch.empty((q.shape[0], max_pools), device=q.device, dtype=torch.float32) + _bf16_mqa_logits[(triton.cdiv(max_pools, 32), q.shape[0])]( + q, + keys, + weights, + starts, + ends, + logits, + q.stride(0), + q.stride(1), + weights.stride(0), + max_pools, + q.shape[1], + q.shape[2], + max(16, triton.next_power_of_2(q.shape[1])), + 32, + num_warps=4, + enable_fp_fusion=False, + ) + return logits diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index 0e126aea6a..0bf42af58b 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -45,6 +45,7 @@ def _compress_pools( SINGLE_QUERY: tl.constexpr, TAIL_SIZE: tl.constexpr, HAS_MTP_INDEX: tl.constexpr, + QUANTIZED: tl.constexpr, ): if SINGLE_QUERY: batch = tl.program_id(0) @@ -85,15 +86,18 @@ def _compress_pools( score = tl.exp(score - tl.max(score, 0)[None, :]) weights = score / tl.sum(score, 0)[None, :] key = tl.sum(raw * weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) - for step in tl.static_range(7): - key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) - key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) - scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) - key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) loc = tl.load(Ragged + start + length - 1).to(tl.int64) dest = Packed + loc * PACKED_STRIDE - tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) - tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + if QUANTIZED: + for step in tl.static_range(7): + key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) + key = (key * (128**-0.5)).to(tl.bfloat16).to(tl.float32) + scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) + key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) + tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) + tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + else: + tl.store(dest + cols, key.reshape(128)) @triton.jit @@ -156,6 +160,7 @@ def compress_pools( single_query, tail.shape[1], mtp_index is not None, + packed_buffer.dtype == torch.uint8, num_warps=4, ) # Complete all boundary pools before replacing the previous chunk's tail. @@ -226,6 +231,7 @@ def _gather_pools( PACKED_STRIDE: tl.constexpr, REQ_STRIDE: tl.constexpr, POOLS: tl.constexpr, + QUANTIZED: tl.constexpr, ): pool, batch = tl.program_id(0), tl.program_id(1) length = tl.load(SeqLen + batch) @@ -234,10 +240,13 @@ def _gather_pools( loc = tl.load(ReqTable + req * REQ_STRIDE + pool * 4 + 3, valid, 0).to(tl.int64) cols = tl.arange(0, 128) packed = tl.load(Packed + loc * PACKED_STRIDE + cols, valid, 0) - scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) row = batch.to(tl.int64) * POOLS + pool - tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) - tl.store(Scale + row, scale) + if QUANTIZED: + scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) + tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) + tl.store(Scale + row, scale) + else: + tl.store(K + row * 128 + cols, packed) def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): @@ -254,11 +263,30 @@ def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): packed_buffer.stride(0), req_table.stride(0), max_pools, + True, num_warps=4, ) return keys, scales +def gather_bf16_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): + keys = torch.empty((req_idx.numel() * max_pools, 128), device=packed_buffer.device, dtype=torch.bfloat16) + _gather_pools[(max_pools, req_idx.numel())]( + packed_buffer, + req_table, + req_idx, + seq_len, + keys, + None, + packed_buffer.stride(0), + req_table.stride(0), + max_pools, + False, + num_warps=4, + ) + return keys + + @triton.jit def _gather_paged_pools( Packed, diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 748f1ad2f7..ead9d04165 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -77,6 +77,11 @@ def get_env_use_ue8m0_scales() -> bool: return os.getenv("LIGHTLLM_USE_UE8M0_SCALES", "0").upper() in ["ON", "TRUE", "1"] +def get_env_glm5_indexer_bf16() -> bool: + """Precision ablation; set identically on all GLM5 prefill/decode workers.""" + return os.getenv("LIGHTLLM_GLM5_INDEXER_BF16", "0").upper() in ["ON", "TRUE", "1"] + + @lru_cache(maxsize=None) def is_stop_matching_in_reasoning_enabled() -> bool: """是否允许用户配置的 stop 在 reasoning 阶段参与匹配。""" diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index ae40702d75..d6646a5aea 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -14,7 +14,9 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_nope_cache_config_uses_native_mla_width(): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_nope_cache_config_uses_native_mla_width(monkeypatch, indexer_bf16): + monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", str(int(indexer_bf16))) config = Glm5NextCacheConfig.from_model_config( { "kv_lora_rank": 512, @@ -31,14 +33,16 @@ def test_nope_cache_config_uses_native_mla_width(): }, StartArgs(tp=4, data_type="bfloat16"), ) - assert config.full_att_head_dim == 584 - assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 + assert config.indexer_bf16 == indexer_bf16 + assert config.full_att_head_dim == (640 if indexer_bf16 else 584) + assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + (256 if indexer_bf16 else 144) @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, indexer_bf16): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -54,12 +58,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) + head_dim = 640 if indexer_bf16 else 584 config = Glm5NextCacheConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, + indexer_bf16=indexer_bf16, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, num_linear_k_heads=2, @@ -75,12 +81,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() - mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) + mem = Glm5NextMemManager(16, torch.bfloat16, 1, head_dim, full_layers, config) req = Glm5NextReqManager(3, 16, mem, config) att_kv = mem.get_att_input_params(3) assert att_kv.shape == (17, 1, 512) - assert att_kv.stride(0) == 584 + assert att_kv.stride(0) == head_dim index_bytes = mem.get_indexer_k_buffer(3) + assert index_bytes.dtype == (torch.bfloat16 if indexer_bf16 else torch.uint8) + assert index_bytes.shape[-1] == (128 if indexer_bf16 else 132) index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") @@ -133,7 +141,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t packed_bytes[:, 0].random_(0, 256) mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == 584 * 2 * full_layers + assert mem.get_cell_size() == head_dim * 2 * full_layers assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( diff --git a/unit_tests/models/glm5_next/test_indexer_bf16.py b/unit_tests/models/glm5_next/test_indexer_bf16.py new file mode 100644 index 0000000000..048a77d647 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_bf16.py @@ -0,0 +1,153 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next.triton_kernel.index_bf16 import bf16_mqa_logits +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("heads", [4, 32]) +@pytest.mark.parametrize("max_pools", [640, 8192, 262144]) +def test_bf16_scores_match_fp32_and_replay_changing_bounds(heads, max_pools): + torch.manual_seed(53) + q = torch.randn(6, heads, 128, device="cuda", dtype=torch.bfloat16) + keys = torch.randn(2 * max_pools, 128, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(6, heads, device="cuda") * heads**-0.5 + starts = torch.tensor([0, max_pools, 0, max_pools, 0, max_pools], device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 1, 511, 512, 639, max_pools], device="cuda", dtype=torch.int32) + ends = starts + lengths + bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) + # Explicit FP64 oracle avoids any global TF32 setting affecting the check. + for shift in range(2): + lengths.copy_(lengths.roll(shift)) + ends.copy_(starts + lengths) + graph.replay() + for row, (start, length) in enumerate(zip(starts.tolist(), lengths.tolist())): + reference = (q[row].double() @ keys[start : start + length].double().T * 128**-0.5).relu() + reference = (reference * weights[row, :, None].double()).sum(0).float() + torch.testing.assert_close(actual[row, :length], reference, atol=3e-6, rtol=2e-5) + assert torch.isneginf(actual[row, length:]).all() + + +@pytest.mark.parametrize("is_prefill", [True, False]) +@pytest.mark.parametrize("cuda_graph", [False, True]) +def test_bf16_indexer_forward_causal_pools_mtp_and_hold(monkeypatch, is_prefill, cuda_graph): + monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", "1") + monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False) + + def unexpected_fp8(*args, **kwargs): + raise AssertionError("BF16 indexer must not quantize Q or call the FP8 scorer") + + monkeypatch.setattr("lightllm.models.glm5_next.indexer.hadamard_transform_quant_fp8", unexpected_fp8) + torch.manual_seed(57) + heads, dim = 4, 128 + prefixes, counts, requests = [2300, 2049, 7], [5, 3, 3], [2, 0, 1] + lengths = [prefix + count for prefix, count in zip(prefixes, counts)] + table = torch.randperm(4 * 2500, device="cuda", dtype=torch.int32).view(4, 2500) + storage = torch.zeros(4 * 2500, 1, 640, device="cuda", dtype=torch.bfloat16) + packed = storage[:, :, 512:] + tail = torch.zeros(4, 8, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, dim, device="cuda") + tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + sources = [torch.randn(length, 256, device="cuda", dtype=torch.bfloat16) for length in lengths] + for req, prefix, source in zip(requests, prefixes, sources): + compress_pools( + source[:prefix], + tail, + packed, + ape, + tensor(range(1, prefix + 1)), + tensor([0] * prefix), + table[req], + tensor([req]), + tensor([0, prefix]), + tensor([prefix]), + prefix, + ) + hidden = torch.cat([source[prefix:] for source, prefix in zip(sources, prefixes)]) + q = torch.randn(hidden.shape[0], heads, dim, device="cuda", dtype=torch.bfloat16) + visible, req_rows, mtp_rows, ragged, starts, cu = [], [], [], [], [], [0] + for req, prefix, count, length in zip(requests, prefixes, counts, lengths): + visible.extend(range(prefix + 1, length + 1)) + req_rows.extend([req] * count) + mtp_rows.extend(range(count)) + starts.extend([len(ragged)] * count) + ragged.extend(table[req, :length].tolist()) + cu.append(cu[-1] + count) + if not is_prefill: + # Each verify position is an independent row with its own visible pool + # length and ragged start; graph padding uses the HOLD request slot. + hidden = torch.cat((hidden, torch.zeros(1, 256, device="cuda", dtype=torch.bfloat16))) + q = torch.cat((q, torch.zeros(1, heads, dim, device="cuda", dtype=torch.bfloat16))) + req_rows.append(3) + mtp_rows.append(0) + visible.append(0) + starts.append(len(ragged)) + infer = SimpleNamespace( + is_prefill=is_prefill, + max_kv_seq_len=max(lengths), + max_q_seq_len=max(counts) if is_prefill else 1, + b_req_idx=tensor(requests if is_prefill else req_rows), + b_seq_len=tensor(lengths if is_prefill else visible), + b1_cu_q_seq_len=tensor(cu if is_prefill else range(hidden.shape[0] + 1)), + b_mtp_index=tensor(mtp_rows), + mem_manager=SimpleNamespace(get_indexer_k_buffer=lambda _: packed), + req_manager=SimpleNamespace(req_to_token_indexs=table, get_indexer_tail_buffer=lambda _: tail), + ) + att = SimpleNamespace(lengths=tensor(visible), ks=tensor(starts), ragged_mem_index=tensor(ragged)) + weight = SimpleNamespace( + wk_proj_=SimpleNamespace(mm=lambda x: x[:, :dim].contiguous()), + k_norm_=lambda x, **kwargs: x, + index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x[:, dim:].contiguous()), + index_kpool_compress_ape=SimpleNamespace(weight=ape), + wq_b_proj_=SimpleNamespace(mm=lambda x: x), + weights_proj_=SimpleNamespace(mm=lambda x: x[:, :heads].contiguous() * 0.1), + ) + indexer = Glm5NextNsaInfer(3, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim}, 8) + run = lambda: indexer._get_indices(hidden, q.flatten(1), infer, att, weight) + before = tail.clone() + run() + tail.copy_(before) + if cuda_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual, relative = run() + tail.copy_(before) + graph.replay() + else: + actual, relative = run() + row = 0 + for req, prefix, count, source in zip(requests, prefixes, counts, sources): + for step in range(count): + length = prefix + step + 1 + pools = length // 4 + values = source[: pools * 4].view(pools, 4, 256) + pooled = (values[:, :, :dim].float() * (values[:, :, dim:].float() + ape).softmax(1)).sum(1).bfloat16() + scores = (q[row].double() @ pooled.double().T * dim**-0.5).relu() + weights = hidden[row, :heads].float() * 0.1 * heads**-0.5 + scores = (scores * weights[:, None].double()).sum(0) + selected = relative[row][relative[row] >= 0] + # Dense requests include all tokens; sparse requests include full + # selected pools followed by the uncompressed partial tail. + if length <= 2048: + assert sorted(selected.tolist()) == list(range(length)) + else: + groups = selected[:2048:4] // 4 + assert groups.unique().numel() == 512 + expected = scores.topk(512).values.sort().values + got = scores[groups.long()].sort().values + torch.testing.assert_close(got, expected, atol=2e-5, rtol=2e-5) + assert selected[2048:].tolist() == list(range(pools * 4, length)) + assert torch.equal(actual[row, : selected.numel()], table[req, selected.long()]) + row += 1 + if not is_prefill: + assert (actual[-1] == -1).all() + assert not tail[3].any() diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 1a5e32567b..8a8ec45afd 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -182,7 +182,8 @@ def set_inputs(counts, phase): @pytest.mark.parametrize("cuda_graph", [False, True]) -def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph, indexer_bf16): # Unaligned chunks, reordered requests, and pools spanning old/new tails. seqs = [9, 7] ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] @@ -191,8 +192,10 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): table[1, :9], table[3, :7] = ragged[:9], ragged[9:] source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda") - packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = packed_storage.view(torch.uint8)[:, :, -132:] + head_dim = 640 if indexer_bf16 else 584 + packed_storage = torch.zeros(40, 1, head_dim, device="cuda", dtype=torch.bfloat16) + get_packed = lambda storage: storage[:, :, 512:] if indexer_bf16 else storage.view(torch.uint8)[:, :, -132:] + packed = get_packed(packed_storage) tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) unchanged = tail[[0, 2, 4]].clone() tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) @@ -237,13 +240,22 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): assert torch.equal(tail[[0, 2, 4]], unchanged) # Model the full KV copy performed by cache offload/load or request move. packed_storage = packed_storage.clone() - packed = packed_storage.view(torch.uint8)[:, :, -132:] - keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) + packed = get_packed(packed_storage) + seq_lens = torch.tensor(seqs, device="cuda", dtype=torch.int32) + if indexer_bf16: + from lightllm.models.glm5_next.triton_kernel.kpool import gather_bf16_pools + + keys = gather_bf16_pools(packed, table, req_idx, seq_lens, 3) + else: + keys, scales = gather_pools(packed, table, req_idx, seq_lens, 3) for batch, start in enumerate([0, 9]): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) + if indexer_bf16: + torch.testing.assert_close(keys[batch * 3 + group], expected, atol=0, rtol=0) + continue + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -300,7 +312,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -392,7 +404,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2 ** 31 // (584 * 2) + 1 + loc = 2**31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -405,7 +417,7 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @@ -482,7 +494,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) out, final = chunk_gated_delta_rule( q, k, @@ -637,10 +649,10 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) + scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 00d4e326f0..e863ab91a7 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -72,13 +72,16 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste @pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) @pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) -def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix, indexer_bf16): torch.manual_seed(53) width, capacity = mtp_step + 1, 128 raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) + packed = torch.zeros( + capacity, 128 if indexer_bf16 else 132, device="cuda", dtype=torch.bfloat16 if indexer_bf16 else torch.uint8 + ) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) def prefill(values, destination, tail): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..c561ea5da9 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -14,12 +14,13 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.fixture(params=[0, 2]) +@pytest.fixture(params=[(mtp, bf16) for mtp in (0, 2) for bf16 in (False, True)]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step = request.param + mtp_step, indexer_bf16 = request.param + head_dim = 640 if indexer_bf16 else 584 args = StartArgs( data_type="bfloat16", linear_att_hash_page_size=4, @@ -45,7 +46,8 @@ def create(tp): full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, + indexer_bf16=indexer_bf16, global_linear_k_heads=8, global_linear_v_heads=8, num_linear_k_heads=8 // tp, @@ -62,7 +64,7 @@ def create(tp): ) mems = [] for _ in range(tp): - mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) + mem = Glm5NextMemManager(32, torch.bfloat16, 1, head_dim, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) mem.write_to_shm(req) assert mem.big_page_buffers is not None @@ -73,7 +75,7 @@ def create(tp): return create -@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) +@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1), (8, 1)]) @pytest.mark.parametrize("remainder", range(4)) def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder): source, dest = make_mems(prefill_tp), make_mems(decode_tp) @@ -83,7 +85,7 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") + packed = torch.randint(0, 256, (layers, length, 1, source[0].head_dim * 2), dtype=torch.uint8, device="cuda") conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") # Non-live slots deliberately contain stale data. The sequence length, not @@ -170,4 +172,4 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems): with pytest.raises(AssertionError, match="smaller than global linear att state"): helper.assert_page_size() shape = mem.get_paged_kv_move_buffer_shape(2, 2048) - assert shape == (2, 2048, mem.layer_num, 1, 584) + assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim) From bb6833d2d0d5b88a9e1ab82c5dcd6e93c55b5003 Mon Sep 17 00:00:00 2001 From: baishihao Date: Sat, 24 Oct 2026 00:34:44 +0800 Subject: [PATCH 12/13] feat(fp8): decouple activation scale format from checkpoint weights --- .../fused_moe/impl/deepgemm_impl.py | 4 +- .../fused_moe/grouped_fused_moe.py | 8 +- .../fused_moe/grouped_fused_moe_ep.py | 14 +- .../moe_silu_and_mul_mix_quant_ep.py | 8 +- lightllm/common/quantization/deepgemm.py | 9 +- lightllm/common/quantization/w8a8.py | 7 +- lightllm/utils/envs_utils.py | 19 +++ .../fused_moe/test_deepgemm_activation.py | 28 +++- .../test_activation_scale_override.py | 149 ++++++++++++++++++ .../test_ue8m0_fp8_quantization.py | 36 +++-- 10 files changed, 251 insertions(+), 31 deletions(-) create mode 100644 unit_tests/common/quantization/test_activation_scale_override.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 74e457426b..3aa9b8b32f 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -6,7 +6,7 @@ from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_env_use_ue8m0_scales, + get_env_use_ue8m0_activation_scales, ) from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( fused_experts, @@ -130,7 +130,7 @@ def low_latency_dispatch( num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, num_experts=self.total_expert_num_contain_redundancy, use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), + round_scale=get_env_use_ue8m0_activation_scales(), async_finish=False, return_recv_hook=True, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index 64537ac606..d2a56c188c 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -25,7 +25,7 @@ from lightllm.utils.log_utils import init_logger from lightllm.utils.vllm_utils import vllm_ops from lightllm.utils.device_utils import triton_support_tensor_descriptor -from lightllm.utils.envs_utils import get_env_use_ue8m0_scales +from lightllm.utils.envs_utils import get_env_use_ue8m0_scales, get_env_use_ue8m0_activation_scales from .moe_silu_and_mul import silu_and_mul_fwd from .moe_sum_reduce import moe_sum_reduce from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 @@ -392,7 +392,9 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t else: # Expert rows may be unordered, but grouped matmul reuses this same # mapping for up/down projections and writes back to original topk slots. - _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights) + _moe_align_fused_atomic_token( + expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights + ) return expert_to_token_index, expert_to_weight, expert_token_num @@ -884,7 +886,7 @@ def grouped_matmul( token_inputs, block_size_k, dtype=expert_weights.dtype, - use_ue8m0_scales=get_env_use_ue8m0_scales() or use_packed_ue8m0, + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(get_env_use_ue8m0_scales() or use_packed_ue8m0), ) if reused_mblock_infos is None: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 61cd81556a..da37626171 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -22,7 +22,7 @@ from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_env_use_ue8m0_scales, + get_env_use_ue8m0_activation_scales, ) from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.device_utils import is_sm100_gpu @@ -158,7 +158,7 @@ def mega_moe_impl( qinput_tensor = per_token_cast_to_fp8( hidden_states, - use_ue8m0=True, + use_ue8m0=get_env_use_ue8m0_activation_scales(default=True, require_ue8m0=True), gran_k=quant_method.block_size, use_packed_ue8m0=True, ) @@ -192,7 +192,7 @@ def quantize_fused_experts_input( return per_token_cast_to_fp8( hidden_states, - use_ue8m0=True, + use_ue8m0=get_env_use_ue8m0_activation_scales(default=True, require_ue8m0=True), gran_k=quant_method.block_size, use_packed_ue8m0=True, ) @@ -204,7 +204,7 @@ def quantize_fused_experts_input( hidden_states, block_size_k, dtype=w13.weight.dtype, - use_ue8m0_scales=get_env_use_ue8m0_scales(), + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(), ) @@ -288,7 +288,7 @@ def fused_experts_impl( combined_x = None if is_prefill: qinput_tensor, input_scale = per_token_group_quant_fp8( - hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales() + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_activation_scales() ) allocate_on_comm_stream = previous_event is not None # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: @@ -383,7 +383,7 @@ def fused_experts_impl( num_max_dispatch_tokens_per_rank, num_experts, use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), + round_scale=get_env_use_ue8m0_activation_scales(), async_finish=False, return_recv_hook=False, ) @@ -608,7 +608,7 @@ def workspace_quant_alloc(shape, dtype, device): column_major_scales=True, scale_tma_aligned=True, alloc_func=workspace_quant_alloc, - use_ue8m0_scales=get_env_use_ue8m0_scales(), + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(), ) workspace_manager.free(silu_out) del silu_out diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 8c94f088c2..9a631c8ea7 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -4,7 +4,7 @@ import triton.language as tl from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu -from lightllm.utils.envs_utils import get_env_use_ue8m0_scales +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales @triton.jit @@ -55,7 +55,9 @@ def _silu_and_mul_post_quant_kernel( ) for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): - gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) + gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to( + tl.float32 + ) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) @@ -161,7 +163,7 @@ def silu_and_mul_masked_post_quant_fwd( alpha=alpha, limit=limit, CLAMP_UP_ADD_ONE=clamp_up_add_one, - USE_UE8M0_SCALE=get_env_use_ue8m0_scales(), + USE_UE8M0_SCALE=get_env_use_ue8m0_activation_scales(), num_warps=num_warps, ) return diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index c953afd27a..b0221f3304 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -6,6 +6,7 @@ from lightllm.common.quantization.registry import QUANTMETHODS from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -109,7 +110,9 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, - use_ue8m0_scales=self.use_ue8m0_scales, + use_ue8m0_scales=get_env_use_ue8m0_activation_scales( + self.use_ue8m0_scales, require_ue8m0=self.use_packed_ue8m0 + ), use_packed_ue8m0=self.use_packed_ue8m0, ) @@ -187,7 +190,9 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): else: num_groups, n, k = weight.shape packed_weight = torch.empty((num_groups, n, k // 2), device=weight.device, dtype=torch.int8) - weight_scale = torch.empty((num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32) + weight_scale = torch.empty( + (num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32 + ) for i in range(num_groups): packed_weight[i], weight_scale[i] = per_token_cast_to_fp4( weight[i], use_ue8m0=True, gran_k=self.block_size diff --git a/lightllm/common/quantization/w8a8.py b/lightllm/common/quantization/w8a8.py index 472db0d543..92075b8282 100644 --- a/lightllm/common/quantization/w8a8.py +++ b/lightllm/common/quantization/w8a8.py @@ -13,6 +13,7 @@ from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_gemm_kernel import w8a8_block_fp8_matmul from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops, cutlass_scaled_mm from lightllm.utils.sgl_utils import HAS_SGL_KERNEL, sgl_ops +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales _HAS_SGL_FP8 = HAS_SGL_KERNEL and sgl_ops is not None and hasattr(sgl_ops, "fp8_scaled_mm") @@ -447,7 +448,11 @@ def _dynamic_quant_input( alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty if input_scale is None: qinput_tensor, input_scale = per_token_group_quant_fp8( - input_tensor, self.block_size, dtype=qweight.dtype, alloc_func=alloc_func + input_tensor, + self.block_size, + dtype=qweight.dtype, + alloc_func=alloc_func, + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(default=False), ) if out is None: out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index ead9d04165..ba8f690b36 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -77,6 +77,25 @@ def get_env_use_ue8m0_scales() -> bool: return os.getenv("LIGHTLLM_USE_UE8M0_SCALES", "0").upper() in ["ON", "TRUE", "1"] +def get_env_use_ue8m0_activation_scales(default=None, *, require_ue8m0=False) -> bool: + """Select block-FP8 activation scales without changing checkpoint weights. + + LIGHTLLM_ACTIVATION_SCALE_FORMAT: auto keeps the caller's existing policy; + fp32 uses continuous scales; ue8m0 rounds scales upward to powers of two. + Indexer Q/K quantization has its own precision switch. + """ + scale_format = os.getenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "auto").strip().lower() + if scale_format == "auto": + use_ue8m0 = get_env_use_ue8m0_scales() if default is None else default + elif scale_format in ("fp32", "ue8m0"): + use_ue8m0 = scale_format == "ue8m0" + else: + raise ValueError("LIGHTLLM_ACTIVATION_SCALE_FORMAT must be auto, fp32, or ue8m0") + if require_ue8m0 and not use_ue8m0: + raise ValueError("This packed DeepGEMM backend requires UE8M0 activation scales; fp32 is supported on Hopper") + return use_ue8m0 + + def get_env_glm5_indexer_bf16() -> bool: """Precision ablation; set identically on all GLM5 prefill/decode workers.""" return os.getenv("LIGHTLLM_GLM5_INDEXER_BF16", "0").upper() in ["ON", "TRUE", "1"] diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py index e9ba2ab26e..2df337f6b2 100644 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -14,7 +14,12 @@ @pytest.mark.parametrize("is_prefill", [False, True]) @pytest.mark.parametrize("use_ue8m0_scales", [False, True]) @pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) -def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one): +@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"]) +def test_grouped_gemm_activation_matches_reference( + monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one, activation_format +): + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format) + activation_ue8m0 = use_ue8m0_scales if activation_format == "auto" else activation_format == "ue8m0" monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES") config = {"quant_method": "fp8", "weight_block_size": [128, 128]} @@ -31,6 +36,24 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ w2_scale = torch.ones(2, 1, 1, device="cuda") counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) + captured_scales = [] + quantize = ep.per_token_group_quant_fp8 + masked_quantize = ep.silu_and_mul_masked_post_quant_fwd + + def capture_contiguous(*args, **kwargs): + assert kwargs["use_ue8m0_scales"] == activation_ue8m0 + result = quantize(*args, **kwargs) + count = (3, 2)[len(captured_scales)] + captured_scales.append(result[1][:count].clone()) + return result + + def capture_masked(*args, **kwargs): + result = masked_quantize(*args, **kwargs) + captured_scales.append(torch.cat([args[2][0, :3], args[2][1, :2]]).clone()) + return result + + monkeypatch.setattr(ep, "per_token_group_quant_fp8", capture_contiguous) + monkeypatch.setattr(ep, "silu_and_mul_masked_post_quant_fwd", capture_masked) recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) recv[0, :3] = x[:3] recv[1, :2] = x[3:] @@ -77,7 +100,8 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() activation = (gate * up).bfloat16().float() scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 - if use_ue8m0_scales: + if activation_ue8m0: scales = torch.exp2(torch.ceil(torch.log2(scales))) + torch.testing.assert_close(torch.cat(captured_scales), scales, atol=1e-6, rtol=1e-5) expected = ((activation / scales).to(torch.float8_e4m3fn).float() * scales).bfloat16() torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.02) diff --git a/unit_tests/common/quantization/test_activation_scale_override.py b/unit_tests/common/quantization/test_activation_scale_override.py new file mode 100644 index 0000000000..21fc9e522a --- /dev/null +++ b/unit_tests/common/quantization/test_activation_scale_override.py @@ -0,0 +1,149 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.quantization import Quantcfg, deepgemm +from lightllm.common.quantization.quantize_method import WeightPack +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales, get_env_use_ue8m0_scales + + +@pytest.mark.parametrize("scale_format", [None, "auto", "fp32", "ue8m0", " FP32 "]) +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("default", [None, False, True]) +def test_activation_override_preserves_weight_policy(monkeypatch, scale_format, legacy, default): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(legacy))) + if scale_format is None: + monkeypatch.delenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", raising=False) + else: + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) + expected = legacy if default is None else default + if scale_format is not None and scale_format.strip().lower() != "auto": + expected = scale_format.strip().lower() == "ue8m0" + assert get_env_use_ue8m0_activation_scales(default) == expected + assert get_env_use_ue8m0_scales() == legacy + + +def test_invalid_or_unsupported_activation_format_fails_explicitly(monkeypatch): + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "typo") + with pytest.raises(ValueError, match="LIGHTLLM_ACTIVATION_SCALE_FORMAT"): + get_env_use_ue8m0_activation_scales() + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "fp32") + with pytest.raises(ValueError, match="requires UE8M0"): + get_env_use_ue8m0_activation_scales(True, require_ue8m0=True) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not deepgemm.HAS_DEEPGEMM, reason="CUDA + DeepGEMM required") +@pytest.mark.parametrize("weight_ue8m0", [False, True]) +def test_dense_activation_override_reuses_exact_weight_pack(monkeypatch, weight_ue8m0): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Continuous FP32 activation scales target Hopper") + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") + config = { + "n_layer": 1, + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0" if weight_ue8m0 else "float32", + }, + } + method = Quantcfg(config).get_quant_method(0, "q_proj") + pack, _ = method.create_weight([128], 512, torch.bfloat16, 0) + torch.manual_seed(54) + method.load_weight(torch.randn(128, 512, device="cuda", dtype=torch.bfloat16), pack) + before_weight, before_scale = pack.weight.view(torch.uint8).clone(), pack.weight_scale.clone() + x = torch.randn(17, 512, device="cuda", dtype=torch.bfloat16) + quantize = deepgemm.per_token_group_quant_fp8 + calls = [] + + def capture(values, group_size, **kwargs): + result = quantize(values, group_size, **kwargs) + expected = values.float().reshape(17, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448 + if kwargs["use_ue8m0_scales"]: + expected = expected.log2().ceil().exp2() + assert group_size == 128 and result[1].dtype == torch.float32 + torch.testing.assert_close(result[1], expected, rtol=1e-6, atol=0) + calls.append(kwargs["use_ue8m0_scales"]) + return result + + monkeypatch.setattr(deepgemm, "per_token_group_quant_fp8", capture) + for scale_format in ("auto", "fp32", "ue8m0"): + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) + result = method.apply(x, pack, use_custom_tensor_mananger=False) + assert torch.isfinite(result).all() + assert torch.equal(pack.weight.view(torch.uint8), before_weight) + assert torch.equal(pack.weight_scale, before_scale) + assert method.use_ue8m0_scales == weight_ue8m0 + assert calls == [weight_ue8m0, False, True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"]) +def test_fallback_block_quantization_honors_override(monkeypatch, scale_format): + from lightllm.common.quantization.w8a8 import FP8w8a8B128QuantizationMethod + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) + method = object.__new__(FP8w8a8B128QuantizationMethod) + method.block_size = 128 + x = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16) + pack = WeightPack( + weight=torch.ones(128, 512, device="cuda", dtype=torch.float8_e4m3fn), + weight_scale=torch.full((1, 4), 0.125, device="cuda"), + ) + result = method._dynamic_quant_input(x, pack, use_custom_tensor_mananger=False) + expected = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448 + if scale_format == "ue8m0": + expected = expected.log2().ceil().exp2() + torch.testing.assert_close(result[2], expected, rtol=1e-6, atol=0) + assert torch.all(pack.weight_scale == 0.125) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"]) +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("overlap", [False, True]) +def test_deepep_dispatch_uses_activation_override(monkeypatch, scale_format, legacy, overlap): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(legacy))) + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) + expected = legacy if scale_format == "auto" else scale_format == "ue8m0" + hidden = torch.zeros(3, 128, device="cuda", dtype=torch.bfloat16) + ids = torch.zeros(3, 1, device="cuda", dtype=torch.int64) + weights = torch.ones(3, 1, device="cuda") + calls = [] + + def dispatch(*args, **kwargs): + calls.append(kwargs["round_scale"]) + return object(), object(), object(), None, None + + buffer = SimpleNamespace( + group_size=1, + low_latency_dispatch=dispatch, + low_latency_combine=lambda *args, **kwargs: (hidden, None, None), + ) + monkeypatch.setattr(ep.dist_group_manager, "ep_low_latency_buffer", buffer, raising=False) + if overlap: + impl = object.__new__(FuseMoeDeepGEMM) + impl.quant_method = SimpleNamespace(method_name="fp8w8a8-b128-deepgemm") + impl.total_expert_num_contain_redundancy = 1 + impl._select_experts = lambda **kwargs: (weights, ids, ids) + impl.low_latency_dispatch(hidden, hidden, None, False, 1, True, 1, 1, "sigmoid") + else: + monkeypatch.setattr(ep, "masked_group_gemm", lambda *args, **kwargs: hidden) + ep.fused_experts_impl( + hidden, + torch.zeros(1, 256, 128, device="cuda", dtype=torch.float8_e4m3fn), + torch.zeros(1, 128, 128, device="cuda", dtype=torch.float8_e4m3fn), + weights, + ids, + buffer=buffer, + is_prefill=False, + num_experts=1, + use_fp8_w8a8=True, + use_fp8_all2all=True, + block_size_k=128, + ) + assert calls == [expected] diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py index 2abf8ccd58..f8565027b5 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -61,14 +61,14 @@ def record_call(*args, **kwargs): assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) if use_ue8m0_scales: actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(actual, torch.full_like(actual, 2.0**-8), rtol=0, atol=0) @pytest.mark.parametrize("use_packed_ue8m0", [False, True]) def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") + powers = torch.tensor([2.0**exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) below = torch.nextafter(powers, torch.zeros_like(powers)) magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 @@ -134,8 +134,10 @@ def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): x[0].zero_() if rows > 1: x[1].fill_(1e-12) - x[-1, -group_size:].fill_(2.0 ** 16) - q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) + x[-1, -group_size:].fill_(2.0**16) + q, packed_scales = activation.per_token_group_quant_fp8( + x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True + ) assert packed_scales.dtype == torch.int32 assert packed_scales.shape == (rows, (groups + 3) // 4) @@ -194,9 +196,9 @@ def check_gemm_inputs(a, b, out): seen.append(a[1].dtype) assert b[1].dtype == (torch.int32 if sm100 else torch.float32) scales = _unpack_scales(a[1], 4) if sm100 else a[1] - torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(scales, torch.full_like(scales, 2.0**-8), rtol=0, atol=0) weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] - torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0**-8), rtol=0, atol=0) out.zero_() # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. @@ -268,11 +270,14 @@ def test_weight_scale_loading_preserves_fused_views( @pytest.mark.parametrize("use_ue8m0", [False, True]) @pytest.mark.parametrize("shared_expert", [False, True]) -def test_tp_moe_activation_scale_rule(monkeypatch, use_ue8m0, shared_expert): +@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"]) +def test_tp_moe_activation_scale_rule(monkeypatch, use_ue8m0, shared_expert, activation_format): from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe as moe from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import append_fused_shared_experts monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format) + activation_ue8m0 = use_ue8m0 if activation_format == "auto" else activation_format == "ue8m0" Quantcfg( { "n_layer": 1, @@ -299,7 +304,7 @@ def check_activation(values, group_size, **kwargs): assert group_size == 128 assert scales.dtype == torch.float32 expected = values.float().reshape(values.shape[0], -1, 128).abs().amax(-1).clamp_min(1e-10) / 448.0 - if use_ue8m0: + if activation_ue8m0: expected = torch.exp2(torch.ceil(torch.log2(expected.double()))).float() torch.testing.assert_close(scales, expected, rtol=1e-6, atol=0) calls.append(values.shape[1]) @@ -315,10 +320,14 @@ def check_activation(values, group_size, **kwargs): @pytest.mark.parametrize("packed_env", [False, True]) @pytest.mark.parametrize("need_trans", [False, True]) @pytest.mark.parametrize("k", [512, 640]) -def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k, packed_env): +@pytest.mark.parametrize("activation_format", ["auto", "fp32"]) +def test_tp_moe_consumes_packed_weight_scales( + monkeypatch, packed_weight_layout, need_trans, k, packed_env, activation_format +): from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import grouped_matmul monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format) monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: True) method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") full, _ = method.create_moe_weight([256], k, torch.bfloat16, 0, 2) @@ -489,11 +498,15 @@ def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout @pytest.mark.parametrize("is_prefill", [False, True]) @pytest.mark.parametrize("sm100", [False, True]) @pytest.mark.parametrize("scale_source", ["none", "env", "config"]) -def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_prefill, sm100, scale_source): +@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"]) +def test_ep_moe_rounds_scales_and_preserves_logical_block_size( + monkeypatch, is_prefill, sm100, scale_source, activation_format +): from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep from lightllm.utils.envs_utils import get_env_use_ue8m0_scales monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False) + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format) assert not get_env_use_ue8m0_scales() config = {"n_layer": 1} if scale_source == "env": @@ -501,6 +514,7 @@ def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_p elif scale_source == "config": config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"} use_ue8m0_scales = scale_source != "none" + activation_ue8m0 = use_ue8m0_scales if activation_format == "auto" else activation_format == "ue8m0" monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") assert get_env_use_ue8m0_scales() == use_ue8m0_scales @@ -511,7 +525,7 @@ def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_p assert q.shape == x.shape and sf.shape == (3, 4) assert sf.dtype == torch.float32 expected_scales = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448 - if use_ue8m0_scales: + if activation_ue8m0: expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales))) expected_q = (x.float().reshape(3, 4, 128) / expected_scales.unsqueeze(-1)).to(q.dtype).reshape_as(q) torch.testing.assert_close(q.float(), expected_q.float(), atol=0, rtol=0) From a139d3f5817d75de6a1378d41cb81c429a32498f Mon Sep 17 00:00:00 2001 From: baishihao Date: Tue, 27 Oct 2026 14:14:16 +0800 Subject: [PATCH 13/13] refactor(glm5): trim precision fixes to FP8 and prefill combine --- .../fused_moe/fused_moe_weight.py | 2 + .../fused_moe/impl/deepgemm_impl.py | 2 + .../fused_moe/impl/triton_impl.py | 2 + .../fused_moe/grouped_fused_moe.py | 4 +- .../fused_moe/grouped_fused_moe_ep.py | 11 ++ .../moe_silu_and_mul_mix_quant_ep.py | 4 +- .../glm5_next_mem_manager.py | 4 - lightllm/common/quantization/deepgemm.py | 4 +- .../common/state_cache_manager/glm5_next.py | 10 +- lightllm/models/glm5_next/indexer.py | 50 +----- .../layer_infer/transformer_layer_infer.py | 10 +- .../glm5_next/triton_kernel/index_bf16.py | 76 --------- .../glm5_next/triton_kernel/index_quant.py | 9 +- .../models/glm5_next/triton_kernel/kpool.py | 53 ++---- lightllm/utils/envs_utils.py | 7 +- .../basemodel/triton_kernel/mhc/test_mhc.py | 4 +- .../test_ue8m0_fp8_quantization.py | 14 +- unit_tests/models/glm5_next/test_cache.py | 24 +-- .../models/glm5_next/test_indexer_bf16.py | 153 ------------------ .../models/glm5_next/test_indexer_scale.py | 30 ++++ unit_tests/models/glm5_next/test_kernels.py | 40 ++--- unit_tests/models/glm5_next/test_mtp.py | 7 +- unit_tests/models/glm5_next/test_pd_cache.py | 14 +- .../glm5_next/test_precision_alignment.py | 4 +- .../models/glm5_next/test_shared_combine.py | 148 +++++++++++++++++ 25 files changed, 275 insertions(+), 411 deletions(-) delete mode 100644 lightllm/models/glm5_next/triton_kernel/index_bf16.py delete mode 100644 unit_tests/models/glm5_next/test_indexer_bf16.py create mode 100644 unit_tests/models/glm5_next/test_indexer_scale.py create mode 100644 unit_tests/models/glm5_next/test_shared_combine.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 5652279702..ef970735d1 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -140,6 +140,7 @@ def experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_output: Optional[torch.Tensor] = None, ) -> torch.Tensor: # Captures MoE topk expert ids for routed-experts metadata when enabled. moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) @@ -162,6 +163,7 @@ def experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + **({"shared_expert_output": shared_expert_output} if shared_expert_output is not None else {}), ) def low_latency_dispatch( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 3aa9b8b32f..c22a5c595a 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -80,6 +80,7 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_output: Optional[torch.Tensor] = None, ): output = fused_experts( hidden_states=input_tensor, @@ -94,6 +95,7 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_output=shared_expert_output, ) return output diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index f748aea467..222e209374 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -111,6 +111,7 @@ def __call__( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_output: Optional[torch.Tensor] = None, ): topk_weights, topk_ids, origin_topk_ids = self._select_experts( input_tensor=input_tensor, @@ -140,5 +141,6 @@ def __call__( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + **({"shared_expert_output": shared_expert_output} if shared_expert_output is not None else {}), ) return output diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py index d2a56c188c..19df68e041 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py @@ -392,9 +392,7 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t else: # Expert rows may be unordered, but grouped matmul reuses this same # mapping for up/down projections and writes back to original topk slots. - _moe_align_fused_atomic_token( - expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights - ) + _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights) return expert_to_token_index, expert_to_weight, expert_token_num diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index da37626171..fda0ab09f9 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -221,10 +221,15 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_output: Optional[torch.Tensor] = None, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + if shared_expert_output is not None and not is_prefill: + raise ValueError("Shared combine bias requires normal prefill combine") check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): + if shared_expert_output is not None: + raise NotImplementedError("Shared combine bias requires DeepEP normal combine") if limit is not None: raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU") return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) @@ -249,6 +254,7 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_output=shared_expert_output, ) @@ -271,7 +277,11 @@ def fused_experts_impl( limit: Optional[float] = None, clamp_up_add_one: bool = True, block_size_k: int = 128, + shared_expert_output: Optional[torch.Tensor] = None, ): + # Shared output joins normal combine before its final BF16 cast. + if shared_expert_output is not None and not is_prefill: + raise ValueError("Shared combine bias requires normal prefill combine") # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" assert topk_weights.shape == topk_idx.shape, "topk shape mismatch" @@ -372,6 +382,7 @@ def fused_experts_impl( topk_weights=None, previous_event=previous_event, allocate_on_comm_stream=allocate_on_comm_stream, + **({"bias": shared_expert_output} if shared_expert_output is not None else {}), ) else: # low latency dispatch diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 9a631c8ea7..ef058407dd 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -55,9 +55,7 @@ def _silu_and_mul_post_quant_kernel( ) for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): - gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to( - tl.float32 - ) + gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py index 1779fabf0c..45a3be667c 100644 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -39,10 +39,6 @@ def get_att_input_params(self, layer_index): return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] def get_indexer_k_buffer(self, layer_index): - if self.linear_config.indexer_bf16: - return self._layer_buffer(layer_index)[ - :, :, self.mla_head_dim : self.mla_head_dim + self.linear_config.index_head_dim - ] return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] def write_to_shm(self, req_manager): diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index b0221f3304..5aca76bd5d 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -190,9 +190,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): else: num_groups, n, k = weight.shape packed_weight = torch.empty((num_groups, n, k // 2), device=weight.device, dtype=torch.int8) - weight_scale = torch.empty( - (num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32 - ) + weight_scale = torch.empty((num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32) for i in range(num_groups): packed_weight[i], weight_scale[i] = per_token_cast_to_fp4( weight[i], use_ue8m0=True, gran_k=self.block_size diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 4c833dd15d..13195a74ac 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -1,6 +1,6 @@ import dataclasses -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_env_glm5_indexer_bf16 +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args from lightllm.utils.torch_dtype_utils import get_torch_dtype from .linear_att import LinearAttCacheConfig @@ -12,7 +12,6 @@ class Glm5NextCacheConfig(LinearAttCacheConfig): index_kpool: int = 4 index_head_dim: int = 128 - indexer_bf16: bool = False INDEX_PADDING_BYTES = 144 @@ -27,12 +26,8 @@ def from_model_config(cls, config, args): "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) ] dtype = get_torch_dtype(args.data_type) - indexer_bf16 = get_env_glm5_indexer_bf16() - if indexer_bf16: - assert args.data_type in ("bf16", "bfloat16"), "BF16 indexer requires --data_type bfloat16" # Raw index keys and compression scores only live in request tails. - index_bytes = config["index_head_dim"] * dtype.itemsize if indexer_bf16 else cls.INDEX_PADDING_BYTES - packed_dim = config["kv_lora_rank"] + index_bytes // dtype.itemsize + packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize return cls( tp_world_size=tp, full_att_all_num_kv_heads=1, @@ -54,7 +49,6 @@ def from_model_config(cls, config, args): draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], - indexer_bf16=indexer_bf16, ) def get_cpu_cache_full_att_bytes(self): diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 652e848c1b..30a9883881 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -2,18 +2,9 @@ import triton from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops -from lightllm.utils.envs_utils import get_env_glm5_indexer_bf16 -from .triton_kernel.index_bf16 import bf16_mqa_logits from .triton_kernel.index_quant import hadamard_transform_quant_fp8 -from .triton_kernel.kpool import ( - compress_pools, - gather_pools, - gather_bf16_pools, - gather_paged_pools, - get_pool_ranges, - expand_topk, -) +from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk class Glm5NextNsaInfer: @@ -26,7 +17,6 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.dim = network_config["index_head_dim"] # The indexer LayerNorm has a separate epsilon from the model RMSNorm. self.eps = 1e-6 - self.indexer_bf16 = get_env_glm5_indexer_bf16() def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" @@ -49,7 +39,6 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig raw = torch.cat((k, gate), -1) tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) - assert (packed_buffer.dtype == torch.bfloat16) == self.indexer_bf16, "Indexer/cache precision mismatch" compress_pools( raw=raw, tail=tail, @@ -66,24 +55,15 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig ) if infer_state.max_kv_seq_len <= self.topk: - return expand_topk( - None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True - ) + return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True) # The small indexer is replicated: no all-gather of query heads and # identical pool selection on every TP rank. q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.heads, self.dim) - max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 - if self.indexer_bf16: - # Keep the existing projections and pooling arithmetic: only remove - # the rotated FP8 Q/K representation and FP8 scoring from this A/B. - weights = layer_weight.weights_proj_.mm(hidden_states.float()) * self.heads**-0.5 - groups = self._get_bf16_indices(q, weights, packed_buffer, infer_state, att_state, max_pools) - return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) - - q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim**-0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim ** -0.5) weights = layer_weight.weights_proj_.mm(hidden_states.float()) - weights = weights * (self.heads**-0.5 * self.dim**-0.5) * q_scale.squeeze(-1) + weights = weights * (self.heads ** -0.5 * self.dim ** -0.5) * q_scale.squeeze(-1) + max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. # Decode has few queries, so paged MQA also splits K across SMs for parallelism. if infer_state.is_prefill: @@ -92,26 +72,6 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) - def _get_bf16_indices(self, q, weights, packed_buffer, infer_state, att_state, max_pools): - # Decode uses one request entry per MTP position; its visible length is - # per row, so future verify pools and HOLD rows remain masked. - seq_lens = infer_state.b_seq_len if infer_state.is_prefill else att_state.lengths - keys = gather_bf16_pools( - packed_buffer, infer_state.req_manager.req_to_token_indexs, infer_state.b_req_idx, seq_lens, max_pools - ) - starts, ends, lengths = get_pool_ranges( - att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools - ) - groups = torch.empty((q.shape[0], self.topk // 4), dtype=torch.int32, device=q.device) - chunk_size = max(1, min(q.shape[0], 16 * 1024 * 1024 // max_pools)) - for start in range(0, q.shape[0], chunk_size): - end = min(start + chunk_size, q.shape[0]) - logits = bf16_mqa_logits( - q[start:end], keys, weights[start:end], starts[start:end], ends[start:end], max_pools - ) - self.select_topk_indices(logits, lengths[start:end], groups[start:end]) - return groups - def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): keys = gather_pools( packed_buffer, diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 2032e50c96..e748701967 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -17,7 +17,7 @@ ) from lightllm.common.triton_utils.autotuner import Autotuner from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor @@ -37,6 +37,7 @@ def __init__(self, layer_num, network_config): self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20) self.swiglu_limit = network_config["swiglu_limit"] self.enable_ep_moe = get_env_start_args().enable_ep_moe + self.ep_shared_combine_bias = self.enable_ep_moe and enable_env_vars("LIGHTLLM_EP_SHARED_COMBINE_BIAS") self.is_moe = ( network_config["n_routed_experts"] is not None and layer_num >= network_config["first_k_dense_replace"] @@ -101,10 +102,12 @@ def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor: hidden_states = input.view(-1, self.embed_dim_) num_tokens, hidden_dim = hidden_states.shape - # if fused_shared_experts is not enabled, compute shared_output + shared_output = None if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: shared_output = self._ffn_tp(hidden_states, infer_state, layer_weight) + combine_shared = self.ep_shared_combine_bias and infer_state.is_prefill and shared_output is not None + moe_gate_dtype = layer_weight.moe_gate.data_type_ router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype)) output = layer_weight.experts.experts( @@ -120,9 +123,10 @@ def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor: alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False, + **({"shared_expert_output": shared_output} if combine_shared else {}), ) - if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0: + if shared_output is not None and not combine_shared: output.add_(shared_output) return output.view(num_tokens, hidden_dim) diff --git a/lightllm/models/glm5_next/triton_kernel/index_bf16.py b/lightllm/models/glm5_next/triton_kernel/index_bf16.py deleted file mode 100644 index d17ff0e50c..0000000000 --- a/lightllm/models/glm5_next/triton_kernel/index_bf16.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Unrotated BF16 Q/K with IEEE FP32 scoring for the indexer precision ablation.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _bf16_mqa_logits( - Q, - K, - Weights, - Starts, - Ends, - Out, - Q_STRIDE: tl.constexpr, - Q_HEAD_STRIDE: tl.constexpr, - W_STRIDE: tl.constexpr, - POOLS: tl.constexpr, - HEADS: tl.constexpr, - DIM: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_P: tl.constexpr, -): - row = tl.program_id(1) - pools = tl.program_id(0) * BLOCK_P + tl.arange(0, BLOCK_P) - heads = tl.arange(0, BLOCK_H) - dims = tl.arange(0, 32) - start = tl.load(Starts + row).to(tl.int64) - length = tl.load(Ends + row).to(tl.int64) - start - valid = (pools < POOLS) & (pools < length) - # One query per program preserves request boundaries in ragged/MTP batches. - # IEEE FP32 avoids TF32 settings changing the precision experiment. - dots = tl.full((BLOCK_H, BLOCK_P), 0, tl.float32) - for block in range(tl.cdiv(DIM, 32)): - cols = block * 32 + dims - q = tl.load( - Q + row * Q_STRIDE + heads[:, None] * Q_HEAD_STRIDE + cols[None, :], - (heads[:, None] < HEADS) & (cols[None, :] < DIM), - 0, - ).to(tl.float32) - k = tl.load( - K + (start + pools[None, :]) * DIM + cols[:, None], - valid[None, :] & (cols[:, None] < DIM), - 0, - ).to(tl.float32) - dots = tl.dot(q, k, dots, input_precision="ieee") - weights = tl.load(Weights + row * W_STRIDE + heads, heads < HEADS, 0).to(tl.float32) - scores = tl.sum(tl.maximum(dots * (DIM**-0.5), 0) * weights[:, None], 0) - tl.store(Out + row * POOLS + pools, tl.where(valid, scores, -float("inf")), pools < POOLS) - - -def bf16_mqa_logits(q, keys, weights, starts, ends, max_pools): - """Return row-relative logits; weights already include heads**-0.5.""" - assert q.dtype == keys.dtype == torch.bfloat16 - assert keys.is_contiguous() and weights.stride(1) == 1 - logits = torch.empty((q.shape[0], max_pools), device=q.device, dtype=torch.float32) - _bf16_mqa_logits[(triton.cdiv(max_pools, 32), q.shape[0])]( - q, - keys, - weights, - starts, - ends, - logits, - q.stride(0), - q.stride(1), - weights.stride(0), - max_pools, - q.shape[1], - q.shape[2], - max(16, triton.next_power_of_2(q.shape[1])), - 32, - num_warps=4, - enable_fp_fusion=False, - ) - return logits diff --git a/lightllm/models/glm5_next/triton_kernel/index_quant.py b/lightllm/models/glm5_next/triton_kernel/index_quant.py index 67485d9b52..25c09fd908 100644 --- a/lightllm/models/glm5_next/triton_kernel/index_quant.py +++ b/lightllm/models/glm5_next/triton_kernel/index_quant.py @@ -2,6 +2,7 @@ import triton import triton.language as tl from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales @triton.jit @@ -13,6 +14,7 @@ def _hadamard_transform_quant_fp8_kernel( scale: tl.constexpr, BLOCK_R: tl.constexpr, BLOCK_N: tl.constexpr, + USE_UE8M0_SCALES: tl.constexpr, ): pid = tl.program_id(0) rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) @@ -32,7 +34,9 @@ def _hadamard_transform_quant_fp8_kernel( # Match the unfused path's bf16 Hadamard output before FP8 quantization. x = (x * scale).to(tl.bfloat16).to(tl.float32) absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) - quant_scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0)))) + quant_scale = absmax * (1.0 / 448.0) + if USE_UE8M0_SCALES: + quant_scale = tl.exp2(tl.ceil(tl.log2(quant_scale))) y = tl.minimum(tl.maximum(x / quant_scale[:, None], -448.0), 448.0) tl.store(Y + offsets, y, mask=row_mask[:, None]) @@ -40,7 +44,7 @@ def _hadamard_transform_quant_fp8_kernel( def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]: - """Fuse Hadamard-128 with the following ue8m0 FP8 quantization.""" + """Fuse Hadamard-128 and FP8 quantization with FP32 scale storage.""" assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors" assert x.dtype == torch.bfloat16, "Hadamard transform expects bfloat16 input" @@ -64,6 +68,7 @@ def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[t scale, BLOCK_R=block_r, BLOCK_N=128, + USE_UE8M0_SCALES=get_env_use_ue8m0_activation_scales(default=True), num_warps=2, ) return output.view(original_shape), output_scale diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index 0bf42af58b..c40ab42abd 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -4,6 +4,7 @@ from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage from lightllm.utils.device_utils import get_device_sm_count +from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales @triton.jit @@ -45,7 +46,7 @@ def _compress_pools( SINGLE_QUERY: tl.constexpr, TAIL_SIZE: tl.constexpr, HAS_MTP_INDEX: tl.constexpr, - QUANTIZED: tl.constexpr, + USE_UE8M0_SCALES: tl.constexpr, ): if SINGLE_QUERY: batch = tl.program_id(0) @@ -86,18 +87,17 @@ def _compress_pools( score = tl.exp(score - tl.max(score, 0)[None, :]) weights = score / tl.sum(score, 0)[None, :] key = tl.sum(raw * weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) + for step in tl.static_range(7): + key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) + key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32) + scale = tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0 + if USE_UE8M0_SCALES: + scale = tl.exp2(tl.ceil(tl.log2(scale))) + key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) loc = tl.load(Ragged + start + length - 1).to(tl.int64) dest = Packed + loc * PACKED_STRIDE - if QUANTIZED: - for step in tl.static_range(7): - key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128) - key = (key * (128**-0.5)).to(tl.bfloat16).to(tl.float32) - scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0))) - key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) - tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) - tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) - else: - tl.store(dest + cols, key.reshape(128)) + tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) + tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) @triton.jit @@ -160,7 +160,7 @@ def compress_pools( single_query, tail.shape[1], mtp_index is not None, - packed_buffer.dtype == torch.uint8, + USE_UE8M0_SCALES=get_env_use_ue8m0_activation_scales(default=True), num_warps=4, ) # Complete all boundary pools before replacing the previous chunk's tail. @@ -231,7 +231,6 @@ def _gather_pools( PACKED_STRIDE: tl.constexpr, REQ_STRIDE: tl.constexpr, POOLS: tl.constexpr, - QUANTIZED: tl.constexpr, ): pool, batch = tl.program_id(0), tl.program_id(1) length = tl.load(SeqLen + batch) @@ -240,13 +239,10 @@ def _gather_pools( loc = tl.load(ReqTable + req * REQ_STRIDE + pool * 4 + 3, valid, 0).to(tl.int64) cols = tl.arange(0, 128) packed = tl.load(Packed + loc * PACKED_STRIDE + cols, valid, 0) + scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) row = batch.to(tl.int64) * POOLS + pool - if QUANTIZED: - scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) - tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) - tl.store(Scale + row, scale) - else: - tl.store(K + row * 128 + cols, packed) + tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) + tl.store(Scale + row, scale) def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): @@ -263,30 +259,11 @@ def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): packed_buffer.stride(0), req_table.stride(0), max_pools, - True, num_warps=4, ) return keys, scales -def gather_bf16_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): - keys = torch.empty((req_idx.numel() * max_pools, 128), device=packed_buffer.device, dtype=torch.bfloat16) - _gather_pools[(max_pools, req_idx.numel())]( - packed_buffer, - req_table, - req_idx, - seq_len, - keys, - None, - packed_buffer.stride(0), - req_table.stride(0), - max_pools, - False, - num_warps=4, - ) - return keys - - @triton.jit def _gather_paged_pools( Packed, diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index ba8f690b36..7ad9d68cce 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -82,7 +82,7 @@ def get_env_use_ue8m0_activation_scales(default=None, *, require_ue8m0=False) -> LIGHTLLM_ACTIVATION_SCALE_FORMAT: auto keeps the caller's existing policy; fp32 uses continuous scales; ue8m0 rounds scales upward to powers of two. - Indexer Q/K quantization has its own precision switch. + GLM5 indexer Q/K use this policy too; their payload and cache remain FP8. """ scale_format = os.getenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "auto").strip().lower() if scale_format == "auto": @@ -96,11 +96,6 @@ def get_env_use_ue8m0_activation_scales(default=None, *, require_ue8m0=False) -> return use_ue8m0 -def get_env_glm5_indexer_bf16() -> bool: - """Precision ablation; set identically on all GLM5 prefill/decode workers.""" - return os.getenv("LIGHTLLM_GLM5_INDEXER_BF16", "0").upper() in ["ON", "TRUE", "1"] - - @lru_cache(maxsize=None) def is_stop_matching_in_reasoning_enabled() -> bool: """是否允许用户配置的 stop 在 reasoning 阶段参与匹配。""" diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py index d126c65ca8..2fea058ea6 100644 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -188,9 +188,9 @@ def test_glm_mhc_normalizes_before_cancelling_large_products(monkeypatch): monkeypatch.setitem(sys.modules, "deep_gemm", None) monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) # Project-first overflows; normalize-first has finite cancelling products. - x = torch.full((1, 4 * 17), 2.0**32, device="cuda", dtype=torch.bfloat16) + x = torch.full((1, 4 * 17), 2.0 ** 32, device="cuda", dtype=torch.bfloat16) fn = torch.zeros(24, x.shape[-1], device="cuda") - fn[:, 0], fn[:, 1] = 2.0**110, -(2.0**110) + fn[:, 0], fn[:, 1] = 2.0 ** 110, -(2.0 ** 110) x_fp32 = x.float() inv_rms = torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5) expected = F.linear(x_fp32 * inv_rms, fn) diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py index f8565027b5..0bdfbcd63f 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -61,14 +61,14 @@ def record_call(*args, **kwargs): assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) if use_ue8m0_scales: actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual, torch.full_like(actual, 2.0**-8), rtol=0, atol=0) + torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) @pytest.mark.parametrize("use_packed_ue8m0", [False, True]) def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - powers = torch.tensor([2.0**exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") + powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) below = torch.nextafter(powers, torch.zeros_like(powers)) magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 @@ -134,10 +134,8 @@ def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): x[0].zero_() if rows > 1: x[1].fill_(1e-12) - x[-1, -group_size:].fill_(2.0**16) - q, packed_scales = activation.per_token_group_quant_fp8( - x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True - ) + x[-1, -group_size:].fill_(2.0 ** 16) + q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) assert packed_scales.dtype == torch.int32 assert packed_scales.shape == (rows, (groups + 3) // 4) @@ -196,9 +194,9 @@ def check_gemm_inputs(a, b, out): seen.append(a[1].dtype) assert b[1].dtype == (torch.int32 if sm100 else torch.float32) scales = _unpack_scales(a[1], 4) if sm100 else a[1] - torch.testing.assert_close(scales, torch.full_like(scales, 2.0**-8), rtol=0, atol=0) + torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] - torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0**-8), rtol=0, atol=0) + torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) out.zero_() # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index d6646a5aea..ae40702d75 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -14,9 +14,7 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("indexer_bf16", [False, True]) -def test_nope_cache_config_uses_native_mla_width(monkeypatch, indexer_bf16): - monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", str(int(indexer_bf16))) +def test_nope_cache_config_uses_native_mla_width(): config = Glm5NextCacheConfig.from_model_config( { "kv_lora_rank": 512, @@ -33,16 +31,14 @@ def test_nope_cache_config_uses_native_mla_width(monkeypatch, indexer_bf16): }, StartArgs(tp=4, data_type="bfloat16"), ) - assert config.indexer_bf16 == indexer_bf16 - assert config.full_att_head_dim == (640 if indexer_bf16 else 584) - assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + (256 if indexer_bf16 else 144) + assert config.full_att_head_dim == 584 + assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -@pytest.mark.parametrize("indexer_bf16", [False, True]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, indexer_bf16): +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -58,14 +54,12 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - head_dim = 640 if indexer_bf16 else 584 config = Glm5NextCacheConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=head_dim, - indexer_bf16=indexer_bf16, + full_att_head_dim=584, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, num_linear_k_heads=2, @@ -81,14 +75,12 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() - mem = Glm5NextMemManager(16, torch.bfloat16, 1, head_dim, full_layers, config) + mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) req = Glm5NextReqManager(3, 16, mem, config) att_kv = mem.get_att_input_params(3) assert att_kv.shape == (17, 1, 512) - assert att_kv.stride(0) == head_dim + assert att_kv.stride(0) == 584 index_bytes = mem.get_indexer_k_buffer(3) - assert index_bytes.dtype == (torch.bfloat16 if indexer_bf16 else torch.uint8) - assert index_bytes.shape[-1] == (128 if indexer_bf16 else 132) index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") @@ -141,7 +133,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t packed_bytes[:, 0].random_(0, 256) mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == head_dim * 2 * full_layers + assert mem.get_cell_size() == 584 * 2 * full_layers assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( diff --git a/unit_tests/models/glm5_next/test_indexer_bf16.py b/unit_tests/models/glm5_next/test_indexer_bf16.py deleted file mode 100644 index 048a77d647..0000000000 --- a/unit_tests/models/glm5_next/test_indexer_bf16.py +++ /dev/null @@ -1,153 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer -from lightllm.models.glm5_next.triton_kernel.index_bf16 import bf16_mqa_logits -from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("heads", [4, 32]) -@pytest.mark.parametrize("max_pools", [640, 8192, 262144]) -def test_bf16_scores_match_fp32_and_replay_changing_bounds(heads, max_pools): - torch.manual_seed(53) - q = torch.randn(6, heads, 128, device="cuda", dtype=torch.bfloat16) - keys = torch.randn(2 * max_pools, 128, device="cuda", dtype=torch.bfloat16) - weights = torch.randn(6, heads, device="cuda") * heads**-0.5 - starts = torch.tensor([0, max_pools, 0, max_pools, 0, max_pools], device="cuda", dtype=torch.int32) - lengths = torch.tensor([0, 1, 511, 512, 639, max_pools], device="cuda", dtype=torch.int32) - ends = starts + lengths - bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) - # Explicit FP64 oracle avoids any global TF32 setting affecting the check. - for shift in range(2): - lengths.copy_(lengths.roll(shift)) - ends.copy_(starts + lengths) - graph.replay() - for row, (start, length) in enumerate(zip(starts.tolist(), lengths.tolist())): - reference = (q[row].double() @ keys[start : start + length].double().T * 128**-0.5).relu() - reference = (reference * weights[row, :, None].double()).sum(0).float() - torch.testing.assert_close(actual[row, :length], reference, atol=3e-6, rtol=2e-5) - assert torch.isneginf(actual[row, length:]).all() - - -@pytest.mark.parametrize("is_prefill", [True, False]) -@pytest.mark.parametrize("cuda_graph", [False, True]) -def test_bf16_indexer_forward_causal_pools_mtp_and_hold(monkeypatch, is_prefill, cuda_graph): - monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", "1") - monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False) - - def unexpected_fp8(*args, **kwargs): - raise AssertionError("BF16 indexer must not quantize Q or call the FP8 scorer") - - monkeypatch.setattr("lightllm.models.glm5_next.indexer.hadamard_transform_quant_fp8", unexpected_fp8) - torch.manual_seed(57) - heads, dim = 4, 128 - prefixes, counts, requests = [2300, 2049, 7], [5, 3, 3], [2, 0, 1] - lengths = [prefix + count for prefix, count in zip(prefixes, counts)] - table = torch.randperm(4 * 2500, device="cuda", dtype=torch.int32).view(4, 2500) - storage = torch.zeros(4 * 2500, 1, 640, device="cuda", dtype=torch.bfloat16) - packed = storage[:, :, 512:] - tail = torch.zeros(4, 8, 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, dim, device="cuda") - tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - sources = [torch.randn(length, 256, device="cuda", dtype=torch.bfloat16) for length in lengths] - for req, prefix, source in zip(requests, prefixes, sources): - compress_pools( - source[:prefix], - tail, - packed, - ape, - tensor(range(1, prefix + 1)), - tensor([0] * prefix), - table[req], - tensor([req]), - tensor([0, prefix]), - tensor([prefix]), - prefix, - ) - hidden = torch.cat([source[prefix:] for source, prefix in zip(sources, prefixes)]) - q = torch.randn(hidden.shape[0], heads, dim, device="cuda", dtype=torch.bfloat16) - visible, req_rows, mtp_rows, ragged, starts, cu = [], [], [], [], [], [0] - for req, prefix, count, length in zip(requests, prefixes, counts, lengths): - visible.extend(range(prefix + 1, length + 1)) - req_rows.extend([req] * count) - mtp_rows.extend(range(count)) - starts.extend([len(ragged)] * count) - ragged.extend(table[req, :length].tolist()) - cu.append(cu[-1] + count) - if not is_prefill: - # Each verify position is an independent row with its own visible pool - # length and ragged start; graph padding uses the HOLD request slot. - hidden = torch.cat((hidden, torch.zeros(1, 256, device="cuda", dtype=torch.bfloat16))) - q = torch.cat((q, torch.zeros(1, heads, dim, device="cuda", dtype=torch.bfloat16))) - req_rows.append(3) - mtp_rows.append(0) - visible.append(0) - starts.append(len(ragged)) - infer = SimpleNamespace( - is_prefill=is_prefill, - max_kv_seq_len=max(lengths), - max_q_seq_len=max(counts) if is_prefill else 1, - b_req_idx=tensor(requests if is_prefill else req_rows), - b_seq_len=tensor(lengths if is_prefill else visible), - b1_cu_q_seq_len=tensor(cu if is_prefill else range(hidden.shape[0] + 1)), - b_mtp_index=tensor(mtp_rows), - mem_manager=SimpleNamespace(get_indexer_k_buffer=lambda _: packed), - req_manager=SimpleNamespace(req_to_token_indexs=table, get_indexer_tail_buffer=lambda _: tail), - ) - att = SimpleNamespace(lengths=tensor(visible), ks=tensor(starts), ragged_mem_index=tensor(ragged)) - weight = SimpleNamespace( - wk_proj_=SimpleNamespace(mm=lambda x: x[:, :dim].contiguous()), - k_norm_=lambda x, **kwargs: x, - index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x[:, dim:].contiguous()), - index_kpool_compress_ape=SimpleNamespace(weight=ape), - wq_b_proj_=SimpleNamespace(mm=lambda x: x), - weights_proj_=SimpleNamespace(mm=lambda x: x[:, :heads].contiguous() * 0.1), - ) - indexer = Glm5NextNsaInfer(3, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim}, 8) - run = lambda: indexer._get_indices(hidden, q.flatten(1), infer, att, weight) - before = tail.clone() - run() - tail.copy_(before) - if cuda_graph: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual, relative = run() - tail.copy_(before) - graph.replay() - else: - actual, relative = run() - row = 0 - for req, prefix, count, source in zip(requests, prefixes, counts, sources): - for step in range(count): - length = prefix + step + 1 - pools = length // 4 - values = source[: pools * 4].view(pools, 4, 256) - pooled = (values[:, :, :dim].float() * (values[:, :, dim:].float() + ape).softmax(1)).sum(1).bfloat16() - scores = (q[row].double() @ pooled.double().T * dim**-0.5).relu() - weights = hidden[row, :heads].float() * 0.1 * heads**-0.5 - scores = (scores * weights[:, None].double()).sum(0) - selected = relative[row][relative[row] >= 0] - # Dense requests include all tokens; sparse requests include full - # selected pools followed by the uncompressed partial tail. - if length <= 2048: - assert sorted(selected.tolist()) == list(range(length)) - else: - groups = selected[:2048:4] // 4 - assert groups.unique().numel() == 512 - expected = scores.topk(512).values.sort().values - got = scores[groups.long()].sort().values - torch.testing.assert_close(got, expected, atol=2e-5, rtol=2e-5) - assert selected[2048:].tolist() == list(range(pools * 4, length)) - assert torch.equal(actual[row, : selected.numel()], table[req, selected.long()]) - row += 1 - if not is_prefill: - assert (actual[-1] == -1).all() - assert not tail[3].any() diff --git a/unit_tests/models/glm5_next/test_indexer_scale.py b/unit_tests/models/glm5_next/test_indexer_scale.py new file mode 100644 index 0000000000..a6bcd30cd5 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_scale.py @@ -0,0 +1,30 @@ +import pytest +import torch + +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"]) +@pytest.mark.parametrize("rows", [0, 1, 33]) +def test_indexer_scale_matches_hadamard_reference(monkeypatch, scale_format, rows): + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) + # FP32 matrix products are exact for these power-of-two input fractions. + x = (torch.randint(-64, 64, (rows, 128), device="cuda").float() / 64).bfloat16() + actual, scales = hadamard_transform_quant_fp8(x, 128 ** -0.5) + assert actual.dtype == torch.float8_e4m3fn and scales.dtype == torch.float32 + assert actual.shape == x.shape and scales.shape == (rows, 1) + if not rows: + return + matrix = torch.ones((1, 1), device="cuda") + for _ in range(7): + matrix = torch.cat((torch.cat((matrix, matrix), 1), torch.cat((matrix, -matrix), 1)), 0) + transformed = ((x.float() @ matrix) * 128 ** -0.5).bfloat16().float() + expected_scale = transformed.abs().amax(-1, keepdim=True).clamp_min(1e-4) / 448 + if scale_format != "fp32": + expected_scale = expected_scale.log2().ceil().exp2() + torch.testing.assert_close(scales, expected_scale, rtol=1e-6, atol=0) + expected = (transformed / expected_scale).clamp(-448, 448).to(torch.float8_e4m3fn) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 8a8ec45afd..c87afd3f2a 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -182,8 +182,9 @@ def set_inputs(counts, phase): @pytest.mark.parametrize("cuda_graph", [False, True]) -@pytest.mark.parametrize("indexer_bf16", [False, True]) -def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph, indexer_bf16): +@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"]) +def test_kpool_chunk_boundaries_and_fragmented_token_kv(monkeypatch, cuda_graph, scale_format): + monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format) # Unaligned chunks, reordered requests, and pools spanning old/new tails. seqs = [9, 7] ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] @@ -192,10 +193,8 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph, indexer_bf16 table[1, :9], table[3, :7] = ragged[:9], ragged[9:] source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda") - head_dim = 640 if indexer_bf16 else 584 - packed_storage = torch.zeros(40, 1, head_dim, device="cuda", dtype=torch.bfloat16) - get_packed = lambda storage: storage[:, :, 512:] if indexer_bf16 else storage.view(torch.uint8)[:, :, -132:] - packed = get_packed(packed_storage) + packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = packed_storage.view(torch.uint8)[:, :, -132:] tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) unchanged = tail[[0, 2, 4]].clone() tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) @@ -240,22 +239,13 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph, indexer_bf16 assert torch.equal(tail[[0, 2, 4]], unchanged) # Model the full KV copy performed by cache offload/load or request move. packed_storage = packed_storage.clone() - packed = get_packed(packed_storage) - seq_lens = torch.tensor(seqs, device="cuda", dtype=torch.int32) - if indexer_bf16: - from lightllm.models.glm5_next.triton_kernel.kpool import gather_bf16_pools - - keys = gather_bf16_pools(packed, table, req_idx, seq_lens, 3) - else: - keys, scales = gather_pools(packed, table, req_idx, seq_lens, 3) + packed = packed_storage.view(torch.uint8)[:, :, -132:] + keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) for batch, start in enumerate([0, 9]): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - if indexer_bf16: - torch.testing.assert_close(keys[batch * 3 + group], expected, atol=0, rtol=0) - continue - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -312,7 +302,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -404,7 +394,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2**31 // (584 * 2) + 1 + loc = 2 ** 31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -417,7 +407,7 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @@ -494,7 +484,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) out, final = chunk_gated_delta_rule( q, k, @@ -649,10 +639,10 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) + scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index e863ab91a7..00d4e326f0 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -72,16 +72,13 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste @pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) @pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) -@pytest.mark.parametrize("indexer_bf16", [False, True]) -def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix, indexer_bf16): +def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): torch.manual_seed(53) width, capacity = mtp_step + 1, 128 raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros( - capacity, 128 if indexer_bf16 else 132, device="cuda", dtype=torch.bfloat16 if indexer_bf16 else torch.uint8 - ) + packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) def prefill(values, destination, tail): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index c561ea5da9..39ee6c039c 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -14,13 +14,12 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.fixture(params=[(mtp, bf16) for mtp in (0, 2) for bf16 in (False, True)]) +@pytest.fixture(params=[0, 2]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step, indexer_bf16 = request.param - head_dim = 640 if indexer_bf16 else 584 + mtp_step = request.param args = StartArgs( data_type="bfloat16", linear_att_hash_page_size=4, @@ -46,8 +45,7 @@ def create(tp): full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=head_dim, - indexer_bf16=indexer_bf16, + full_att_head_dim=584, global_linear_k_heads=8, global_linear_v_heads=8, num_linear_k_heads=8 // tp, @@ -64,7 +62,7 @@ def create(tp): ) mems = [] for _ in range(tp): - mem = Glm5NextMemManager(32, torch.bfloat16, 1, head_dim, 2 + int(mtp_step > 0), config) + mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) mem.write_to_shm(req) assert mem.big_page_buffers is not None @@ -85,7 +83,7 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - packed = torch.randint(0, 256, (layers, length, 1, source[0].head_dim * 2), dtype=torch.uint8, device="cuda") + packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") # Non-live slots deliberately contain stale data. The sequence length, not @@ -172,4 +170,4 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems): with pytest.raises(AssertionError, match="smaller than global linear att state"): helper.assert_page_size() shape = mem.get_paged_kv_move_buffer_shape(2, 2048) - assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim) + assert shape == (2, 2048, mem.layer_num, 1, 584) diff --git a/unit_tests/models/glm5_next/test_precision_alignment.py b/unit_tests/models/glm5_next/test_precision_alignment.py index 746397c385..309e56bf05 100644 --- a/unit_tests/models/glm5_next/test_precision_alignment.py +++ b/unit_tests/models/glm5_next/test_precision_alignment.py @@ -113,9 +113,7 @@ def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkey actual_k = weight.k_norm_(projected_k, eps=indexer.eps) expected_k = F.layer_norm(projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=1e-6) torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) - wrong_k = F.layer_norm( - projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=config["rms_norm_eps"] - ) + wrong_k = F.layer_norm(projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=config["rms_norm_eps"]) assert not torch.equal(expected_k, wrong_k) storage = torch.zeros(tokens + 1, 1, 584, dtype=torch.bfloat16, device="cuda") diff --git a/unit_tests/models/glm5_next/test_shared_combine.py b/unit_tests/models/glm5_next/test_shared_combine.py new file mode 100644 index 0000000000..67c11f657e --- /dev/null +++ b/unit_tests/models/glm5_next/test_shared_combine.py @@ -0,0 +1,148 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as weight_module +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import deepgemm_impl +from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep + + +@pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("shared_mode", ["absent", "separate", "native_fused"]) +def test_shared_output_is_added_once(prefill, enabled, shared_mode): + hidden = torch.zeros(2, 128, dtype=torch.bfloat16) + shared = torch.full_like(hidden, 0.25) + calls = [] + + def experts(input_tensor, **kwargs): + calls.append(kwargs) + routed = torch.ones_like(input_tensor) + return routed if "shared_expert_output" not in kwargs else routed + kwargs["shared_expert_output"] + + layer = SimpleNamespace( + embed_dim_=128, + n_shared_experts=None if shared_mode == "absent" else 1, + ep_shared_combine_bias=enabled, + swiglu_limit=7.0, + num_experts_per_tok=2, + norm_topk_prob=True, + n_group=1, + topk_group=1, + _ffn_tp=lambda *args: shared, + ) + weights = SimpleNamespace( + num_fused_shared_experts=int(shared_mode == "native_fused"), + moe_gate=SimpleNamespace(data_type_=torch.float32, mm=lambda x: torch.zeros(2, 4)), + experts=SimpleNamespace(experts=experts), + ) + result = Glm5NextTransformerLayerInfer._moe_ffn(layer, hidden, SimpleNamespace(is_prefill=prefill), weights) + expected = torch.ones_like(hidden) + (shared if shared_mode == "separate" else 0) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + if enabled and prefill and shared_mode == "separate": + assert calls[0]["shared_expert_output"] is shared + else: + assert "shared_expert_output" not in calls[0] + + +@pytest.mark.parametrize("with_shared", [False, True]) +def test_shared_bias_reaches_normal_combine_backend(monkeypatch, with_shared): + hidden = torch.zeros(2, 128, dtype=torch.bfloat16) + shared = torch.ones_like(hidden) if with_shared else None + ids = torch.zeros(2, 1, dtype=torch.int32) + impl = object.__new__(deepgemm_impl.FuseMoeDeepGEMM) + impl._select_experts = lambda **kwargs: (torch.ones(2, 1), ids, ids) + impl.total_expert_num_contain_redundancy = 8 + impl.quant_method = SimpleNamespace(block_size=128) + pack = SimpleNamespace(weight=hidden, weight_scale=None) + calls = [] + + def capture(**kwargs): + calls.append(kwargs) + return hidden + + monkeypatch.setattr(weight_module, "get_moe_capture_callback", lambda *args: None) + monkeypatch.setattr(ep, "check_ep_expert_dtype", lambda method: None) + monkeypatch.setattr(ep, "use_sm100_mega_moe", lambda method: False) + monkeypatch.setattr(ep, "fused_experts_impl", capture) + monkeypatch.setattr(ep.dist_group_manager, "ep_buffer", object(), raising=False) + weight = SimpleNamespace( + layer_num_=3, + fuse_moe_impl=impl, + w13=pack, + w2=pack, + e_score_correction_bias=None, + scoring_func="sigmoid", + per_expert_scale=None, + ) + result = weight_module.FusedMoeWeight.experts( + weight, + hidden, + hidden, + top_k=1, + renormalize=True, + use_grouped_topk=False, + topk_group=1, + num_expert_group=1, + is_prefill=True, + shared_expert_output=shared, + ) + assert result is hidden and calls[0]["shared_expert_output"] is shared + + +def test_unsupported_shared_bias_fails_before_dispatch(monkeypatch): + x = torch.zeros(1, 128, dtype=torch.bfloat16) + args = dict( + hidden_states=x, + w13=None, + w2=None, + topk_weights=None, + topk_idx=None, + num_experts=8, + quant_method=None, + shared_expert_output=x, + ) + with pytest.raises(ValueError, match="normal prefill combine"): + ep.fused_experts(**args, is_prefill=False) + monkeypatch.setattr(ep, "check_ep_expert_dtype", lambda method: None) + monkeypatch.setattr(ep, "use_sm100_mega_moe", lambda method: True) + with pytest.raises(NotImplementedError, match="DeepEP normal combine"): + ep.fused_experts(**args, is_prefill=True) + + +@pytest.mark.parametrize("with_shared", [False, True]) +def test_normal_combine_preserves_default_call_contract(monkeypatch, with_shared): + x = torch.zeros(2, 128, dtype=torch.bfloat16) + bias = torch.ones_like(x) if with_shared else None + calls = [] + handle = SimpleNamespace(num_recv_tokens_per_expert_list=[0], recv_src_metadata=torch.empty(0, 3)) + + def combine(*args, **kwargs): + calls.append(kwargs) + return x, None, None + + buffer = SimpleNamespace(dispatch=lambda *args, **kwargs: ((x, x), None, None, handle, None), combine=combine) + monkeypatch.setattr(ep, "per_token_group_quant_fp8", lambda *args, **kwargs: (x, x)) + monkeypatch.setattr(ep, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 256) + monkeypatch.setattr(ep.Autotuner, "is_kernel_autotune_warmup", lambda *args: False) + ids = torch.zeros(2, 1, dtype=torch.int64) + result = ep.fused_experts_impl( + x, + torch.empty(1, 256, 128), + torch.empty(1, 128, 128), + torch.ones(2, 1), + ids, + num_experts=1, + buffer=buffer, + is_prefill=True, + use_fp8_w8a8=True, + use_fp8_all2all=True, + shared_expert_output=bias, + ) + assert result is x + if with_shared: + assert calls[0]["bias"] is bias + else: + assert "bias" not in calls[0]