diff --git a/docs/CN/source/tutorial/function_calling.rst b/docs/CN/source/tutorial/function_calling.rst index e6eab454ec..afa588572f 100644 --- a/docs/CN/source/tutorial/function_calling.rst +++ b/docs/CN/source/tutorial/function_calling.rst @@ -53,8 +53,6 @@ Mistral **格式**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`` -启用 chat logprobs 时暂不支持 ``mistral`` 工具调用解析器。 - DeepSeek-V3 ~~~~~~~~~~~ diff --git a/docs/EN/source/tutorial/function_calling.rst b/docs/EN/source/tutorial/function_calling.rst index 99db50b0a8..d038eb9022 100644 --- a/docs/EN/source/tutorial/function_calling.rst +++ b/docs/EN/source/tutorial/function_calling.rst @@ -53,8 +53,6 @@ Mistral **Format**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`` -The ``mistral`` tool-call parser is not supported when chat logprobs are enabled. - DeepSeek-V3 ~~~~~~~~~~~ diff --git a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py index 504a1a6ffb..b7ed312469 100644 --- a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py +++ b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py @@ -77,7 +77,10 @@ def _output_logprobs_partial( rows, candidates, sampled_ids, - higher_counts, + ranks, + partial_maxes, + partial_sums, + partial_posinf_counts, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -91,10 +94,17 @@ def _output_logprobs_partial( row_logits = logits + row * stride_row ids = tile * BLOCK + tl.arange(0, BLOCK) scores = tl.load(row_logits + ids, ids < vocab_size, -float("inf")).to(tl.float32) - sampled = tl.load(row_logits + tl.load(sampled_ids + row)).to(tl.float32) - # 只数严格更高的 logit。同分不占名次，采样 token 仍算在前 20 里。 - higher = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0) - tl.store(higher_counts + index * tiles + tile, higher) + sampled = tl.load(row_logits + tl.load(sampled_ids + row)) + rank = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0) + tl.store(ranks + index * tiles + tile, rank) + posinf_count = tl.sum(((ids < vocab_size) & (scores == float("inf"))).to(tl.int32), 0) + finite_scores = tl.where((scores != float("inf")) & (scores != -float("inf")), scores, -float("inf")) + tile_max = tl.max(finite_scores, 0) + tile_sum = tl.sum(tl.exp(finite_scores - tile_max), 0) + tile_sum = tl.where(tile_max > -float("inf"), tile_sum, 0.0) + tl.store(partial_maxes + index * tiles + tile, tile_max) + tl.store(partial_sums + index * tiles + tile, tile_sum) + tl.store(partial_posinf_counts + index * tiles + tile, posinf_count) # tl.topk 只能比较整数。高 32 位放可排序的分数，低 32 位放反转后的 token id。 # 取最大的键时，分数高的在前；分数相同则 id 更小的在前。词表外的位置写成最小整数。 bits = scores.to(tl.int32, bitcast=True) @@ -115,7 +125,11 @@ def _output_logprobs_finish( candidates, token_ids, logprobs, - higher_counts, + partial_ranks, + sampled_ranks, + partial_maxes, + partial_sums, + partial_posinf_counts, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -127,7 +141,6 @@ def _output_logprobs_finish( row = tl.load(rows + index * 3).to(tl.int64) row_logits = logits + row * stride_row sampled = tl.load(row_logits + tl.load(sampled_ids + row)) - sampled_logprob = tl.load(sampled_logprobs + row) width = tl.load(rows + index * 3 + 1) offset = tl.load(rows + index * 3 + 2) columns = tl.arange(0, TOPK) @@ -137,16 +150,26 @@ def _output_logprobs_finish( # 键的低 32 位是反转后的 token id，这里还原。结果仍按分数从高到低排列。 ids = 0xFFFFFFFF - (best & 0xFFFFFFFF) scores = tl.load(row_logits + ids, (columns < width) & (ids < vocab_size), -float("inf")) - stored_ids = tl.where(tl.abs(scores) < float("inf"), ids, -1).to(tl.int64) - # 候选和采样 token 的 logsumexp 相同，用采样 logprob 平移即可，不必再做整词表 softmax。 - row_logprobs = scores - sampled + sampled_logprob + stored_ids = tl.where(scores == scores, ids, -1).to(tl.int64) + rank_tiles = tl.arange(0, TILE_BLOCK) + tile_maxes = tl.load(partial_maxes + index * tiles + rank_tiles, rank_tiles < tiles, -float("inf")) + tile_sums = tl.load(partial_sums + index * tiles + rank_tiles, rank_tiles < tiles, 0.0) + posinf_counts = tl.load(partial_posinf_counts + index * tiles + rank_tiles, rank_tiles < tiles, 0) + row_max = tl.max(tile_maxes, 0) + row_sum = tl.sum(tl.where(tile_maxes > -float("inf"), tile_sums * tl.exp(tile_maxes - row_max), 0.0), 0) + posinf_count = tl.sum(posinf_counts, 0) + finite_logprobs = scores - (row_max + tl.log(row_sum)) + uniform_logprob = -tl.log(float(vocab_size)) + row_logprobs = tl.where( + posinf_count > 0, + tl.where(scores == float("inf"), -tl.log(posinf_count.to(tl.float32)), -float("inf")), + tl.where(row_max > -float("inf"), finite_logprobs, uniform_logprob), + ) tl.store(token_ids + offset + columns, stored_ids, columns < width) tl.store(logprobs + offset + columns, row_logprobs, columns < width) - # 候选平移用的是原始采样 logprob。至少 20 个 token 严格更高时，再把采样分数写成 -9999。 - higher_tiles = tl.arange(0, TILE_BLOCK) - higher = tl.sum(tl.load(higher_counts + index * tiles + higher_tiles, higher_tiles < tiles, 0), 0) - outside = (higher >= 20) | ~(tl.abs(sampled) < float("inf")) - tl.store(sampled_logprobs + row, -9999.0, mask=outside) + rank = 1 + tl.sum(tl.load(partial_ranks + index * tiles + rank_tiles, rank_tiles < tiles, 0), 0) + rank = tl.where(sampled > -float("inf"), rank, vocab_size + 1) + tl.store(sampled_ranks + index, rank) def capture_output_logprobs( @@ -156,12 +179,12 @@ def capture_output_logprobs( rows: torch.Tensor, max_topk: int, output_count: int, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """返回每行的 top-k token id 和 logprob。 rows 每项是 (logits 行号, K, 写入起点)。写入起点按 token 个数累加， - 第 r 行的候选是两个数组的 [起点, 起点 + K)。至少 20 个 token 的 logit - 严格高于采样 token 时，把 sampled_logprobs 里对应位置写成 -9999。K=0 也判断。 + 第 r 行的候选是两个数组的 [起点, 起点 + K)。另返回每行采样 token 的 + 1-based rank（同分共享排名），不修改内部采样 logprob。K=0 仍计算排名。 """ block = 1024 # 先按词表分段留下局部 top-k，再逐行合并成最终候选。 @@ -169,13 +192,20 @@ def capture_output_logprobs( # Triton 的 topk 块长必须是 2 的幂。请求的 K 最大 20，这里向上取整；真正写入仍按每行的 K 截断。 topk = triton.next_power_of_2(max(2, max_topk)) candidates = logits.new_empty((rows.shape[0], tiles, topk), dtype=torch.int64) - higher_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) + partial_ranks = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) + sampled_ranks = logits.new_empty((rows.shape[0],), dtype=torch.int32) + partial_maxes = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32) + partial_sums = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32) + partial_posinf_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) _output_logprobs_partial[(rows.shape[0], tiles)]( logits=logits, rows=rows, candidates=candidates, sampled_ids=sampled_ids, - higher_counts=higher_counts, + ranks=partial_ranks, + partial_maxes=partial_maxes, + partial_sums=partial_sums, + partial_posinf_counts=partial_posinf_counts, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, @@ -193,7 +223,11 @@ def capture_output_logprobs( candidates=candidates, token_ids=token_ids, logprobs=logprobs, - higher_counts=higher_counts, + partial_ranks=partial_ranks, + sampled_ranks=sampled_ranks, + partial_maxes=partial_maxes, + partial_sums=partial_sums, + partial_posinf_counts=partial_posinf_counts, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, @@ -201,4 +235,4 @@ def capture_output_logprobs( TILE_BLOCK=triton.next_power_of_2(tiles), num_warps=4, ) - return token_ids, logprobs + return token_ids, logprobs, sampled_ranks diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 2104a40da6..b7dde16480 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -12,6 +12,7 @@ from lightllm.server.reasoning_parser import ReasoningParser from .chat_logprob_buffer import ChatLogprobBuffer +from .text_spans import TextSpan from .function_call_parser import TOOLS_TAG_LIST, FunctionCallParser, ToolCallItem from .build_prompt import build_prompt, init_tokenizer @@ -228,7 +229,8 @@ def _process_reasoning_stream( reasoning_parser_dict: Dict[int, ReasoningParser], metadata: Dict[str, Any], request: ChatCompletionRequest, -) -> tuple[Optional[str], str]: + with_spans: bool = False, +): """Process reasoning content and update its token usage.""" if index not in reasoning_parser_dict: reasoning_parser_dict[index] = ReasoningParser( @@ -240,10 +242,12 @@ def _process_reasoning_stream( token_id = metadata.get("id") if token_id is not None: parser.update_reasoning_token_count(int(token_id)) - return parser.parse_stream_chunk(delta) + return parser.parse_stream_chunk_with_spans(delta) if with_spans else parser.parse_stream_chunk(delta) -def _process_tools_stream(index: int, delta: str, parser_dict: Dict, request: ChatCompletionRequest): +def _process_tools_stream( + index: int, delta: str, parser_dict: Dict, request: ChatCompletionRequest, with_spans: bool = False +): from .api_http import g_objs if index not in parser_dict: @@ -256,8 +260,7 @@ def _process_tools_stream(index: int, delta: str, parser_dict: Dict, request: Ch parser = parser_dict[index] # parse_increment => returns (normal_text, calls) - normal_text, calls = parser.parse_stream_chunk(delta) - return normal_text, calls + return parser.parse_stream_chunk_with_spans(delta) if with_spans else parser.parse_stream_chunk(delta) def _split_tool_argument_delta(arguments: Optional[str]) -> List[str]: @@ -424,31 +427,21 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, **sampling_params_dict) await sampling_params.verify_cpu_async(g_objs.httpserver_manager.output_grammar_compiler) - logprob_model_type = get_env_start_args().reasoning_parser - logprob_force_reasoning = _is_force_thinking_mode(request) if logprob_model_type else False - - def _new_logprob_buffer() -> ChatLogprobBuffer: - return ChatLogprobBuffer( - tokenizer=g_objs.httpserver_manager.tokenizer, - model_type=logprob_model_type, - force_reasoning=logprob_force_reasoning, - separate_reasoning=bool(request.separate_reasoning), - tools=request.tools if request.tool_choice != "none" else None, - tool_call_parser=getattr(g_objs.args, "tool_call_parser", None) or "llama3", - ) - - logprob_buffers = None - if request.logprobs: - # 构造函数会检查思维链和工具调用的分隔符能不能编成稳定的 token 序列，不合法就抛 ValueError。 - # defaultdict 要等第一条输出 token 才创建 buffer。流式响应那时已经返回，这个错误进不了接口的 400 处理。 - # 这里先构造一次，只为在返回前完成检查；这个对象随即丢掉，真正使用的仍是下面按子请求创建的 buffer。 - # generate 放在这次检查之后，参数不合法时请求不会进入推理。 - _new_logprob_buffer() - logprob_buffers = collections.defaultdict(_new_logprob_buffer) - results_generator = g_objs.httpserver_manager.generate( prompt, sampling_params, multimodal_params, request=raw_request ) + logprob_buffers = ( + collections.defaultdict( + lambda: ChatLogprobBuffer( + tokenizer=g_objs.httpserver_manager.tokenizer, + separate_reasoning=bool(request.separate_reasoning), + skip_special_tokens=bool(sampling_params.skip_special_tokens), + spaces_between_special_tokens=bool(sampling_params.add_spaces_between_special_tokens), + ) + ) + if request.logprobs + else None + ) # Non-streaming case if not request.stream: @@ -499,13 +492,14 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: sub_req_id = sub_ids[i] finish_reason = finish_reason_dict[sub_req_id] text = "".join(final_output_dict[sub_req_id]) + output_spans = [TextSpan(0, len(text))] # Handle reasoning content reasoning_text = None if reasoning_parser: try: parser = reasoning_parser_dict[sub_req_id] - reasoning_text, text = parser.parse_non_stream(text) + reasoning_text, text, output_spans = parser.parse_non_stream_with_spans(text) except Exception as e: logger.error(f"Reasoning parsing error: {e}") return create_error_response( @@ -516,8 +510,12 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: text = (reasoning_text or "") + (text or "") reasoning_text = None + if logprob_buffers is not None: + logprob_buffers[sub_req_id].assign(output_spans, len(text)) + # Handle tool_calls parsing tool_calls = None + content_spans = None tool_choice = request.tool_choice tools = request.tools if tool_choice != "none" and any([i in text for i in TOOLS_TAG_LIST]): @@ -525,7 +523,7 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: # 为 tool_call_parser 提供默认值 tool_parser = getattr(g_objs.args, "tool_call_parser", None) or "llama3" parser = FunctionCallParser(tools, tool_parser) - text, call_info_list = parser.parse_non_stream(text) + text, call_info_list, content_spans = parser.parse_non_stream_with_spans(text) tool_calls = [] history_tool_calls_cnt = _get_history_tool_calls_cnt(request) for call_info in call_info_list: @@ -548,7 +546,7 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: finish_reason = "tool_calls" choice_logprobs = None if logprob_buffers is not None: - content_logprobs, reasoning_logprobs = logprob_buffers[sub_req_id].message_logprobs() + content_logprobs, reasoning_logprobs = logprob_buffers[sub_req_id].message_logprobs(content_spans) choice_logprobs = ChatLogprobs(content=content_logprobs, reasoning_content=reasoning_logprobs) chat_message = ChatMessage( role="assistant", @@ -612,7 +610,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: logprob_buffer = None if logprob_buffers is not None: logprob_buffer = logprob_buffers[sub_req_id] - logprob_buffer.record(text=delta, metadata=metadata, finish_status=finish_status) + raw_span = logprob_buffer.record(text=delta, metadata=metadata, finish_status=finish_status) current_finish_reason = finish_status.get_finish_reason() if current_finish_reason == "error" and completion_tokens == 1: raise GenerationError("Generation failed before producing output") @@ -636,9 +634,12 @@ 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 + reasoning_text, delta, output_spans = _process_reasoning_stream( + choice_index, delta, reasoning_parser_dict, metadata, request, with_spans=True ) + if logprob_buffer is not None: + content_length = len(delta or "") + (0 if request.separate_reasoning else len(reasoning_text or "")) + logprob_buffer.assign(output_spans, content_length) if reasoning_text: if request.separate_reasoning: choice_logprobs = logprob_buffer.reasoning_logprobs() if logprob_buffer is not None else None @@ -656,11 +657,13 @@ async def stream_results() -> AsyncGenerator[bytes, None]: yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls, logprobs=choice_logprobs)}\n\n" else: delta = reasoning_text + (delta or "") + elif logprob_buffer is not None: + logprob_buffer.assign([raw_span], len(delta)) if request.tool_choice != "none" and request.tools: # parse_increment => returns (normal_text, calls) - normal_text, calls = _process_tools_stream( - index=choice_index, delta=delta, parser_dict=parser_dict, request=request + normal_text, calls, content_spans = _process_tools_stream( + index=choice_index, delta=delta, parser_dict=parser_dict, request=request, with_spans=True ) # 1) if there's normal_text, output it as normal content @@ -676,8 +679,13 @@ async def stream_results() -> AsyncGenerator[bytes, None]: choices=[choice_data], model=request.model, ) - choice_logprobs = logprob_buffer.content_logprobs() if logprob_buffer is not None else None + choice_logprobs = ( + logprob_buffer.content_logprobs(content_spans) if logprob_buffer is not None else None + ) yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls, logprobs=choice_logprobs)}\n\n" + elif logprob_buffer is not None: + if current_finish_reason is None or not parser_dict[choice_index].detector._buffer: + logprob_buffer.discard_content(content_spans) # 2) if we found calls, we output them as separate chunk(s) history_tool_calls_cnt = _get_history_tool_calls_cnt(request) @@ -828,7 +836,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]: 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() + flush_reasoning, flush_text, output_spans = parser.flush_with_spans() + if logprob_buffer is not None: + content_length = len(flush_text or "") + ( + 0 if request.separate_reasoning else len(flush_reasoning or "") + ) + flush_spans = logprob_buffer.assign(output_spans, content_length) if flush_reasoning: if request.separate_reasoning: flush_logprobs = ( @@ -841,7 +854,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) else: flush_logprobs = ( - logprob_buffer.content_logprobs() if logprob_buffer is not None else None + logprob_buffer.content_logprobs(flush_spans) if logprob_buffer is not None else None ) # vLLM compat: emit buffered thinking as content flush_choice = ChatCompletionStreamResponseChoice( @@ -869,9 +882,30 @@ async def stream_results() -> AsyncGenerator[bytes, None]: model=request.model, choices=[flush_choice], ) - flush_logprobs = logprob_buffer.content_logprobs() if logprob_buffer is not None else None + flush_logprobs = ( + logprob_buffer.content_logprobs(flush_spans) if logprob_buffer is not None else None + ) payload = _serialize_sse_chunk(flush_chunk, _choice_nulls, logprobs=flush_logprobs) yield f"data: {payload}\n\n" + if request.tool_choice != "none" and request.tools and choice_index in parser_dict: + flush_text, _, output_spans = parser_dict[choice_index].flush_with_spans() + 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], + ) + flush_logprobs = ( + logprob_buffer.content_logprobs(output_spans) if logprob_buffer is not None else None + ) + payload = _serialize_sse_chunk(flush_chunk, _choice_nulls, logprobs=flush_logprobs) + yield f"data: {payload}\n\n" if has_emitted_tool_calls[sub_req_id] and current_finish_reason == "stop": current_finish_reason = "tool_calls" final_choice = ChatCompletionStreamResponseChoice( diff --git a/lightllm/server/chat_logprob_buffer.py b/lightllm/server/chat_logprob_buffer.py index 2f41449b9c..fb77e139a2 100644 --- a/lightllm/server/chat_logprob_buffer.py +++ b/lightllm/server/chat_logprob_buffer.py @@ -1,265 +1,121 @@ -from __future__ import annotations - -from typing import Any, NamedTuple +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs.logprob_utils import chat_token_logprob -from lightllm.server.function_call_parser import FunctionCallParser -from lightllm.server.reasoning_parser import ReasoningParser +from lightllm.server.core.objs.logprob_utils import ChatTokenDecoder +from lightllm.server.text_spans import TextSpan + + +@dataclass +class PendingLogprob: + start: int + end: int + token: Dict[str, Any] class ChatLogprobBuffer: - """按 token 序列识别分隔符，把可见文本的 logprob 放进正文或思维链。 + """Attribute generated tokens to the visible source ranges returned by parsers. - 分隔符可以是一个或多个 token。凑齐才切换通道，分隔符本身不输出。 - 空增量的分数先累加，加到后面第一个非空文本上。生成结束时还没拼上的空增量丢掉。 + A token spanning two channels belongs to the first visible channel. Empty + detokenization increments stay pending until the character they complete is + classified; recovering raw bytes is independent of retaining that token. """ def __init__( self, tokenizer: Any, - model_type: str | None, - force_reasoning: bool, separate_reasoning: bool, - tools: list | None = None, - tool_call_parser: str | None = None, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, ): self.tokenizer = tokenizer self._separate_reasoning = separate_reasoning - self._channel = "content" # 当前输出通道：content、reasoning，或 commentary 的 drop。 - self._in_tool = False # 是否处在工具调用开始和结束分隔符之间。 - self._pending_logprob = 0.0 # 尚未拼成字符的空增量分数，等下一个非空文本一起输出。 - self._reasoning: list[dict[str, Any]] = [] # 分开思考时的思维链 logprob。 - self._content: list[dict[str, Any]] = [] # 正文 logprob；不分开思考时思维链也放这里。 - self._sequences: list[tuple[tuple[int, ...], str]] = [] # 分隔符 token id 序列和切换动作。 - self._held: list[_Held] = [] # 可能是分隔符前缀、尚未判定的 token。 - self._mode = "body" # gpt-oss 解析位置：body 正文、header 角色头、channel 通道名。 - self._channel_name = "" # <|channel|> 和 <|message|> 之间拼出来的通道名。 - if model_type: - self._init_reasoning(tokenizer=tokenizer, model_type=model_type, force_reasoning=force_reasoning) - if tools: - self._init_tools(tokenizer=tokenizer, tool_call_parser=tool_call_parser) - - def record(self, text: str, metadata: dict[str, Any], finish_status: FinishStatus) -> None: + self._hold: List[PendingLogprob] = [] + self._reasoning: List[PendingLogprob] = [] + self._content: List[PendingLogprob] = [] + self._source_length = 0 + self._content_length = 0 + self._decoder = ChatTokenDecoder(tokenizer, skip_special_tokens, spaces_between_special_tokens) + + def record(self, text: str, metadata: Dict[str, Any], finish_status: FinishStatus) -> TextSpan: + start = self._source_length + self._source_length += len(text) token_id, score = metadata.get("id"), metadata.get("logprob") - if token_id is not None and score is not None: - self._held.append(_Held(token_id=token_id, text=text, score=score, metadata=metadata)) - self._resolve_held(closing=False) - if finish_status.is_finished(): - self._resolve_held(closing=True) - self._pending_logprob = 0.0 - - def reasoning_logprobs(self) -> dict[str, Any] | None: - tokens = self._drain(items=self._reasoning) - return {"content": None, "reasoning_content": tokens} if tokens else None - - def content_logprobs(self) -> dict[str, Any] | None: - tokens = self._drain(items=self._content) - return {"content": tokens, "reasoning_content": None} if tokens else None - - def message_logprobs(self) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]] | None]: - return self._drain(items=self._content) or None, self._drain(items=self._reasoning) or None - - def _init_reasoning(self, tokenizer: Any, model_type: str, force_reasoning: bool) -> None: - key = model_type.lower() - detector_class = ReasoningParser.DetectorMap.get(key) - if detector_class is None: - raise ValueError(f"logprobs does not support reasoning parser {model_type!r}") - if key == "gpt-oss": - for text, action in ( - ("<|start|>", "harmony_start"), - ("<|channel|>", "harmony_channel"), - ("<|message|>", "harmony_message"), - ("<|end|>", "harmony_end"), - ("<|return|>", "harmony_end"), - ("<|call|>", "harmony_end"), - ): - # <|call|> 不是每个词表都有，缺了就跳过。 - self._add_sequence(tokenizer=tokenizer, text=text, action=action, required=text != "<|call|>") - if force_reasoning: - self._channel = "reasoning" - return - detector = detector_class(stream_reasoning=True, force_reasoning=force_reasoning) - self._add_sequence(tokenizer=tokenizer, text=detector.think_start_token, action="reasoning") - self._add_sequence(tokenizer=tokenizer, text=detector.think_end_token, action="content") - # 以解析器为准。Gemma 4 会忽略请求里的 force_reasoning，生成不是从已打开的思考通道开始。 - if detector._in_reasoning: - self._channel = "reasoning" - - def _init_tools(self, tokenizer: Any, tool_call_parser: str | None) -> None: - if not tool_call_parser: - raise ValueError("logprobs requires tool_call_parser when tools are set") - if tool_call_parser == "mistral": - raise ValueError("logprobs does not support mistral tool_call_parser") - parser = FunctionCallParser(tools=[], tool_call_parser=tool_call_parser) - self._add_sequence(tokenizer=tokenizer, text=parser.detector.bot_token, action="tool_on") - if parser.detector.eot_token: - self._add_sequence(tokenizer=tokenizer, text=parser.detector.eot_token, action="tool_off") - - def _add_sequence(self, tokenizer: Any, text: str, action: str, required: bool = True) -> None: - token_ids = _delimiter_ids(tokenizer=tokenizer, text=text) - if token_ids is None: - if required: - raise ValueError(f"logprobs requires {text!r} to encode to the same token sequence") - return - sequence = tuple(token_ids) - for existing, existing_action in self._sequences: - if existing == sequence and existing_action != action: - raise ValueError(f"logprobs delimiter {text!r} maps to conflicting actions") - if any(existing == sequence for existing, _ in self._sequences): - return - self._sequences.append((sequence, action)) - - def _resolve_held(self, closing: bool) -> None: - # 还是更长分隔符的前缀时先留着。生成结束则按已经凑齐的短分隔符结算，否则当正文。 - while self._held: - token_ids = tuple(item.token_id for item in self._held) - exact, waiting, completed = self._lookup(token_ids=token_ids) - if waiting and not closing: - return - if exact is not None: - self._apply(action=exact) - self._held.clear() - return - if completed is not None: - action, size = completed - self._apply(action=action) - del self._held[:size] - continue - self._release(item=self._held.pop(0)) - - def _lookup(self, token_ids: tuple[int, ...]) -> tuple[str | None, bool, tuple[str, int] | None]: - exact = None - waiting = False - completed = None - for sequence, action in self._sequences: - if sequence == token_ids: - exact = action - elif sequence[: len(token_ids)] == token_ids: - waiting = True - elif token_ids[: len(sequence)] == sequence and (completed is None or len(sequence) > completed[1]): - completed = (action, len(sequence)) - return exact, waiting, completed - - def _apply(self, action: str) -> None: - match action: - case "reasoning": - # 思维链开始。这组 token 是分隔符，不输出。 - self._channel = "reasoning" - case "content": - # 思维链结束，后面的文本回到正文。 - self._channel = "content" - case "tool_on": - # 工具调用开始，参数不进入正文。 - self._in_tool = True - case "tool_off": - # 工具调用结束，后面的文本回到当前通道。 - self._in_tool = False - case "harmony_start": - # <|start|> 到下一个结构标记之间是角色头。 - self._mode = "header" - case "harmony_channel": - # <|channel|> 后面开始收集通道名。 - self._mode = "channel" - self._channel_name = "" - case "harmony_message": - # <|message|>：analysis 进入思维链，commentary 整段丢弃，其余进入正文。 - name = self._channel_name.strip().lower() - if name.startswith("analysis"): - self._channel = "reasoning" - elif name.startswith("commentary"): - self._channel = "drop" - else: - self._channel = "content" - self._mode = "body" - self._channel_name = "" - case "harmony_end": - # 这段消息结束。通道保持不变，下一次 <|message|> 再切换。 - self._mode = "body" - - def _release(self, item: _Held) -> None: - token_id, text, score, metadata = item - match self._mode: - case "channel": - # 通道名本身，如 analysis，不输出。 - self._channel_name += text - case "header": - # 头部里的角色名，如 assistant，不输出。 - return - case _: - if self._in_tool or self._channel == "drop": - # 工具参数或 commentary 正文不输出。 - return - self._save(text=text, score=score, metadata=metadata, token_id=token_id) - - def _save(self, text: str, score: float, metadata: dict[str, Any], token_id: int) -> None: - if not text: - # 未完成的 UTF-8 字节并进后面的字符。特殊 token 的空文本直接丢掉。 - if not self._is_special(metadata=metadata, token_id=token_id): - self._pending_logprob += score - return - logprob = self._pending_logprob + score - self._pending_logprob = 0.0 - token = { - "token": text, - "logprob": logprob, - "bytes": None if "\ufffd" in text else list(text.encode("utf-8")), - "top_logprobs": [ - _candidate_logprob(tokenizer=self.tokenizer, token_id=candidate, logprob=value) + if "logprob_prompt_ids" in metadata: + self._decoder.set_prompt(metadata["logprob_prompt_ids"]) + special = metadata.get("special", token_id in getattr(self.tokenizer, "all_special_ids", ())) + if token_id is not None and score is not None and (text or not special): + rank = metadata.get("logprobs", {}).get(token_id, {}).get("rank") + token = self._decoder.token_logprob(token_id, -9999.0 if rank is not None and rank > 20 else score) + token["top_logprobs"] = [ + self._decoder.token_logprob(candidate, value) for candidate, value in metadata.get("output_top_logprobs", ()) - ], - } - # 分开思考时思维链单独列出；否则和正文共用 content。 - if self._channel == "reasoning" and self._separate_reasoning: - self._reasoning.append(token) - else: - self._content.append(token) + ] + self._hold.append(PendingLogprob(start, self._source_length, token)) + if token_id is not None: + self._decoder.advance(token_id) + return TextSpan(start, self._source_length) - def _is_special(self, metadata: dict[str, Any], token_id: int) -> bool: - if "special" in metadata: - return bool(metadata["special"]) - return token_id in getattr(self.tokenizer, "all_special_ids", ()) + @staticmethod + def _overlaps(item, span): + if item.start == item.end: + return span.start <= item.start < span.end + return item.start < span.end and item.end > span.start + + def assign(self, spans: List[TextSpan], content_length: int) -> List[TextSpan]: + spans = [span for span in spans if span.start < span.end] + visible = [TextSpan(self._content_length, self._content_length + content_length)] if content_length else [] + content_items = [] + content_offset = self._content_length + for span in spans: + selected = [item for item in self._hold if self._overlaps(item, span)] + self._hold = [item for item in self._hold if not self._overlaps(item, span)] + if span.channel == "reasoning" and self._separate_reasoning: + self._reasoning.extend(selected) + else: + for item in selected: + content_items.append( + PendingLogprob( + content_offset + max(item.start - span.start, 0), + content_offset + min(item.end, span.end) - span.start, + item.token, + ) + ) + content_offset += span.end - span.start + if spans: + consumed = max(span.end for span in spans) + self._hold = [item for item in self._hold if item.end > consumed] + self._content_length += content_length + self._content.extend(content_items) + return visible + + def _take_content(self, spans=None): + if spans is None: + return self._drain(self._content) + selected = [item for item in self._content if any(self._overlaps(item, span) for span in spans)] + if spans: + consumed = max(span.end for span in spans) + self._content = [item for item in self._content if item.end > consumed and item not in selected] + return [item.token for item in selected] + + def discard_content(self, spans: List[TextSpan]) -> None: + self._take_content(spans) @staticmethod - def _drain(items: list[dict[str, Any]]) -> list[dict[str, Any]]: - tokens = list(items) + def _drain(items): + tokens = [item.token for item in items] items.clear() return tokens + def reasoning_logprobs(self) -> Optional[Dict[str, Any]]: + tokens = self._drain(self._reasoning) + return {"content": None, "reasoning_content": tokens} if tokens else None -class _Held(NamedTuple): - token_id: int - text: str - score: float - metadata: dict[str, Any] - - -def _candidate_logprob(tokenizer: Any, token_id: int, logprob: float) -> dict[str, Any]: - return chat_token_logprob(tokenizer=tokenizer, token_id=token_id, logprob=logprob) - - -def _delimiter_ids(tokenizer: Any, text: str) -> list[int] | None: - if not text: - return None - encode = getattr(tokenizer, "encode", None) - if encode is not None: - token_ids = list(encode(text, add_special_tokens=False)) - if token_ids and tokenizer.decode(token_ids, skip_special_tokens=False) == text: - return token_ids - return _longest_token_ids(tokenizer=tokenizer, text=text) - + def content_logprobs(self, spans=None) -> Optional[Dict[str, Any]]: + tokens = self._take_content(spans) + return {"content": tokens, "reasoning_content": None} if tokens else None -def _longest_token_ids(tokenizer: Any, text: str) -> list[int] | None: - vocab = tokenizer.get_vocab() - pieces = sorted((piece for piece in vocab if piece), key=len, reverse=True) - token_ids = [] - index = 0 - while index < len(text): - piece = next((item for item in pieces if text.startswith(item, index)), None) - if piece is None: - return None - token_ids.append(vocab[piece]) - index += len(piece) - if tokenizer.decode(token_ids, skip_special_tokens=False) != text: - return None - return token_ids + def message_logprobs( + self, content_spans=None + ) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[Dict[str, Any]]]]: + return self._take_content(content_spans) or None, self._drain(self._reasoning) or None diff --git a/lightllm/server/core/objs/logprob_utils.py b/lightllm/server/core/objs/logprob_utils.py index 055fb9e213..75f0348fc1 100644 --- a/lightllm/server/core/objs/logprob_utils.py +++ b/lightllm/server/core/objs/logprob_utils.py @@ -1,3 +1,6 @@ +import json +import re + from typing import Any, Dict @@ -20,13 +23,69 @@ def logprob_info(tokenizer, token_id: int, logprob: float, rank: int): ) -def chat_token_logprob(tokenizer: Any, token_id: int, logprob: float) -> Dict[str, Any]: +def _decoder_types(decoder): + if decoder is None: + return set() + config = json.loads(decoder.__getstate__()) + + def types(node): + result = {node.get("type")} + for child in node.get("decoders", []): + result.update(types(child)) + return result + + return types(config) + + +def chat_token_logprob( + tokenizer: Any, token_id: int, logprob: float, prefix_ids=(), pending_ids=(), decode_options=None +) -> Dict[str, Any]: token = tokenizer.decode([token_id], skip_special_tokens=False) + standalone_token = token + special = token_id in getattr(tokenizer, "all_special_ids", ()) + if (prefix_ids or pending_ids) and "\ufffd" not in standalone_token and not special: + options = decode_options or {"skip_special_tokens": False} + prefix = tokenizer.decode([*prefix_ids, *pending_ids], **options) + combined = tokenizer.decode([*prefix_ids, *pending_ids, token_id], **options) + if combined.startswith(prefix): + token = combined[len(prefix) :] raw_bytes = list(token.encode("utf-8")) if "\ufffd" not in token else None decoder = getattr(getattr(tokenizer, "backend_tokenizer", None), "decoder", None) - added_tokens = getattr(tokenizer, "added_tokens_decoder", {}) - if type(decoder).__name__ == "ByteLevel" and token_id not in added_tokens: + decoder_types = _decoder_types(decoder) + if token_id not in getattr(tokenizer, "added_tokens_decoder", {}): piece = tokenizer.convert_ids_to_tokens(token_id) - if piece and all(char in _BYTE_LEVEL_BYTES for char in piece): + if "ByteLevel" in decoder_types and piece and all(char in _BYTE_LEVEL_BYTES for char in piece): raw_bytes = [_BYTE_LEVEL_BYTES[char] for char in piece] + token = standalone_token + elif "ByteFallback" in decoder_types and piece and re.fullmatch(r"<0x[0-9a-fA-F]{2}>", piece): + raw_bytes = [int(piece[3:5], 16)] + token = standalone_token return {"token": token, "logprob": logprob, "bytes": raw_bytes} + + +class ChatTokenDecoder: + def __init__(self, tokenizer, skip_special_tokens=True, spaces_between_special_tokens=True): + self.tokenizer = tokenizer + self.prefix_ids = [] + self.pending_ids = [] + self.options = { + "skip_special_tokens": skip_special_tokens, + "spaces_between_special_tokens": spaces_between_special_tokens, + } + + def set_prompt(self, token_ids): + from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START + + self.prefix_ids = [token_id for token_id in token_ids if token_id < MULTIMODAL_TOKEN_ID_START] + self.pending_ids = [] + + def token_logprob(self, token_id, score): + return chat_token_logprob(self.tokenizer, token_id, score, self.prefix_ids, self.pending_ids, self.options) + + def advance(self, token_id): + prefix = self.tokenizer.decode(self.prefix_ids, **self.options) + self.pending_ids.append(token_id) + combined = self.tokenizer.decode(self.prefix_ids + self.pending_ids, **self.options) + if len(combined) > len(prefix) and not combined.endswith("\ufffd"): + self.prefix_ids = self.pending_ids + self.pending_ids = [] diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 29f5e29d98..dc074256fb 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -349,7 +349,8 @@ def _create_output_topk_logprobs_shm_array(self, topk: int) -> None: self.shm_output_topk_logprobs = ShmArray( name=f"shm_output_topk_logprobs_{self.index_in_shm_mem}", shape=(self.sample_params.max_new_tokens + 128, 2, topk), - dtype=np.float32, + # Token IDs must stay exact beyond float32's 2^24 integer limit. + dtype=np.float64, ) self.shm_output_topk_logprobs.create_shm() # PD decode 可能直接复用 prefill 已采样的第一个 token，先把 token id 标成无效。 @@ -360,7 +361,7 @@ def _link_output_topk_logprobs_shm_array(self, topk: int) -> None: self.shm_output_topk_logprobs = ShmArray( name=f"shm_output_topk_logprobs_{self.index_in_shm_mem}", shape=(self.sample_params.max_new_tokens + 128, 2, topk), - dtype=np.float32, + dtype=np.float64, ) self.shm_output_topk_logprobs.link_shm() diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index 7140d3209c..c4959ae03f 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 .text_spans import TextSpan, extract_text, slice_spans logger = logging.getLogger(__name__) ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"] @@ -55,8 +56,9 @@ class ToolCallItem(BaseModel): class StreamingParseResult: """Result of streaming incremental parsing.""" - def __init__(self, normal_text: str = "", calls: Optional[List[ToolCallItem]] = None): + def __init__(self, normal_text: str = "", calls: Optional[List[ToolCallItem]] = None, spans=None): self.normal_text = normal_text + self.spans = spans or [] self.calls = calls or [] @@ -117,6 +119,7 @@ def __init__(self): # Streaming state management # Buffer for accumulating incomplete patterns that arrive across multiple streaming chunks self._buffer = "" + self._source_length = 0 # Stores complete tool call info (name and arguments) for each tool being parsed. # Used by serving layer for completion handling when streaming ends. # Format: [{"name": str, "arguments": dict}, ...] @@ -227,11 +230,11 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami ): # Only clear buffer if we're sure no tool call is starting if not self._ends_with_partial_token(self._buffer, self.bot_token): - normal_text = self._buffer + normal_text, normal_spans = extract_text( + self._buffer, self._source_length - len(self._buffer), remove=(self.eot_token,) + ) self._buffer = "" - if self.eot_token in normal_text: - normal_text = normal_text.replace(self.eot_token, "") - return StreamingParseResult(normal_text=normal_text) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans) else: # Might be partial bot_token, keep buffering return StreamingParseResult() @@ -405,6 +408,7 @@ def __init__(self): self.eot_token = "" self.tool_call_separator = "\n" self._normal_text_buffer = "" # Buffer for handling partial end tokens + self._normal_text_spans = [] def has_tool_call(self, text: str) -> bool: """Check if the text contains a Qwen 2.5 format tool call.""" @@ -419,9 +423,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) # Find all \n...\n blocks pattern = rf"{re.escape(self.bot_token)}(.*?){re.escape(self.eot_token)}" @@ -434,7 +438,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult except json.JSONDecodeError as e: logger.warning(f"Failed to parse JSON part: {match_result}, JSON parse error: {str(e)}") continue - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -446,11 +450,19 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami # Handle partial end tokens that are streamed character by character if result.normal_text: self._normal_text_buffer += result.normal_text + self._normal_text_spans.extend(result.spans) # Check if buffer contains complete end token (without leading newline) end_token_without_newline = self.eot_token # "" if end_token_without_newline in self._normal_text_buffer: - cleaned_text = self._normal_text_buffer.replace(end_token_without_newline, "") + source = self._normal_text_buffer + cleaned_text, relative_spans = extract_text(source, remove=(end_token_without_newline,)) + result.spans = [ + part + for span in relative_spans + for part in slice_spans(self._normal_text_spans, span.start, span.end) + ] + self._normal_text_spans = [] self._normal_text_buffer = "" result.normal_text = cleaned_text else: @@ -459,12 +471,17 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if partial_match_len: # Keep potential partial match in buffer, return the rest + end = len(self._normal_text_buffer) - partial_match_len + result.spans = slice_spans(self._normal_text_spans, 0, end) + self._normal_text_spans = slice_spans(self._normal_text_spans, end, len(self._normal_text_buffer)) result.normal_text = self._normal_text_buffer[:-partial_match_len] self._normal_text_buffer = self._normal_text_buffer[-partial_match_len:] else: # No partial match, return all buffered text result.normal_text = self._normal_text_buffer + result.spans = self._normal_text_spans self._normal_text_buffer = "" + self._normal_text_spans = [] return result @@ -507,16 +524,16 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) # Extract the JSON array part from [TOOL_CALLS] [...] # Use bracket counting to properly handle nested brackets in JSON content json_array_str = self._extract_json_array(text) if not json_array_str: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) calls = [] try: @@ -528,7 +545,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult except json.JSONDecodeError as e: logger.warning(f"Failed to parse JSON part: {json_array_str}, JSON parse error: {str(e)}") - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def _extract_json_array(self, text: str) -> str: """ @@ -601,7 +618,7 @@ def has_tool_call(self, text: str) -> bool: def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: """Parse function calls from text, handling multiple JSON objects.""" if "<|python_tag|>" not in text and not text.startswith("{"): - return StreamingParseResult(normal_text=text, calls=[]) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))], calls=[]) if "<|python_tag|>" in text: normal_text, action_text = text.split("<|python_tag|>", maxsplit=1) @@ -631,8 +648,10 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # Only process if we found valid JSON objects calls = self.parse_base_json(all_actions, tools) if all_actions else [] # Use safe_idx to avoid idx containing the last part of an invalid JSON object - trailing_text = action_text[safe_idx:].strip() if safe_idx < action_text_len else "" - return StreamingParseResult(normal_text=normal_text + trailing_text, calls=calls) + action_offset = len(text) - len(action_text) + trailing_text, trailing_spans = extract_text(action_text[safe_idx:], action_offset + safe_idx, strip=True) + spans = ([TextSpan(0, len(normal_text))] if normal_text else []) + trailing_spans + return StreamingParseResult(normal_text=normal_text + trailing_text, spans=spans, calls=calls) class KimiK2Detector(BaseFormatDetector): @@ -684,7 +703,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ if self.bot_token not in text: - return StreamingParseResult(normal_text=text, calls=[]) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))], calls=[]) try: # there are two possible captures - between tags, or between a # tag and end-of-string so the result of @@ -715,12 +734,12 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult ) content = text[: text.find(self.bot_token)] - return StreamingParseResult(normal_text=content, calls=tool_calls) + return StreamingParseResult(normal_text=content, spans=[TextSpan(0, len(content))], calls=tool_calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -734,10 +753,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, self.tool_call_end_token]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, self.tool_call_end_token] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -791,7 +810,6 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami parsed_args_diff = argument_diff.split("<|tool_call_end|>", 1)[0] if parsed_args_diff: - calls.append( ToolCallItem( tool_index=self.current_tool_id, @@ -829,7 +847,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class DeepSeekV31Detector(BaseFormatDetector): @@ -883,9 +903,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) calls = [] try: @@ -898,11 +918,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # construct match_result for parse_base_json match_result = {"name": func_name, "parameters": func_args} calls.extend(self.parse_base_json(match_result, tools)) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -916,10 +936,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, "<｜tool▁call▁end｜>"]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, "<｜tool▁call▁end｜>"] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -1009,7 +1029,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class DeepSeekV3Detector(BaseFormatDetector): @@ -1064,9 +1086,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) calls = [] try: @@ -1079,11 +1101,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # construct match_result for parse_base_json match_result = {"name": func_name, "parameters": func_args} calls.extend(self.parse_base_json(match_result, tools)) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -1097,10 +1119,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, "```", "<｜tool▁call▁end｜>"]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, "```", "<｜tool▁call▁end｜>"] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -1192,7 +1214,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class Glm47Detector(BaseFormatDetector): @@ -1295,10 +1319,10 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: StreamingParseResult with normal_text and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) tool_indices = self._get_tool_indices(tools) calls = [] @@ -1336,22 +1360,27 @@ 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) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, 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.""" self._buffer += new_text normal_text = "" + normal_spans = [] calls: List[ToolCallItem] = [] while self._buffer: start = self._buffer.find(self.bot_token) if start == -1: keep = self._ends_with_partial_token(self._buffer, self.bot_token) end = len(self._buffer) - keep + offset = self._source_length - len(self._buffer) + normal_spans.extend([TextSpan(offset, offset + end)] if end else []) normal_text += self._buffer[:end] self._buffer = self._buffer[end:] break + offset = self._source_length - len(self._buffer) + normal_spans.extend([TextSpan(offset, offset + start)] if start else []) normal_text += self._buffer[:start] self._buffer = self._buffer[start:] end = self._buffer.find(self.eot_token) @@ -1384,7 +1413,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._streaming_tool_name = None self._buffer = self._buffer[end:] - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) class DeepSeekV32Detector(BaseFormatDetector): @@ -1462,9 +1491,9 @@ def _dsml_params_to_json(self, params: List[tuple]) -> str: def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: """One-time parsing for DSML format tool calls.""" idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) tool_indices = self._get_tool_indices(tools) calls = [] @@ -1486,7 +1515,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult ) ) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Streaming incremental parsing for DSML format tool calls.""" @@ -1502,10 +1531,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult() self._buffer = "" - for e_token in [self.eot_token, self.invoke_end_token]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, self.invoke_end_token] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) # Mark that we're inside a function_calls block if self.has_tool_call(current_text): @@ -1849,10 +1878,10 @@ def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if " StreamingParseResult item.tool_index = len(calls) calls.append(item) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Streaming incremental parsing for Qwen3-Coder XML tool calls.""" @@ -1878,6 +1907,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._tool_indices = self._get_tool_indices(tools) normal_text = "" + normal_spans = [] calls: List[ToolCallItem] = [] while True: @@ -1887,13 +1917,19 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if tool_call_start == -1: partial_len = self._ends_with_partial_token(current_text, self.bot_token) if partial_len: - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if current_text: - normal_text += current_text.replace(self.eot_token, "") + visible, spans = extract_text( + current_text, self._source_length - len(current_text), remove=(self.eot_token,) + ) + normal_text += visible + normal_spans.extend(spans) self._buffer = "" - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if tool_call_start > 0: + offset = self._source_length - len(current_text) + normal_spans.append(TextSpan(offset, offset + tool_call_start)) normal_text += current_text[:tool_call_start] self._buffer = current_text[tool_call_start:] current_text = self._buffer @@ -1903,7 +1939,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not self.current_tool_name_sent: function_name = self._get_incomplete_function_name(current_text) if function_name is None: - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if self.current_tool_id == -1: self.current_tool_id = 0 @@ -1925,7 +1961,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami parameters="", ) ) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) complete_block = current_text[: eot_pos + len(self.eot_token)] func_matches = self.function_regex.findall(complete_block) @@ -1994,6 +2030,44 @@ def __init__(self, tools: List[Tool], tool_call_parser: str): self.detector = detector self.tools = tools + self._emitted_calls = False + + def parse_non_stream_with_spans(self, full_text: str): + if not self.tools: + return full_text, [], [TextSpan(0, len(full_text))] + result = self.detector.detect_and_parse(full_text, self.tools) + if not result.calls: + return full_text, [], [TextSpan(0, len(full_text))] + return result.normal_text, result.calls, result.spans + + def parse_stream_chunk_with_spans(self, chunk_text: str): + self.detector._source_length += len(chunk_text) + if not self.tools: + end = self.detector._source_length + return chunk_text, [], [TextSpan(end - len(chunk_text), end)] + result = self.detector.parse_streaming_increment(chunk_text, self.tools) + if result.calls: + self._emitted_calls = True + return result.normal_text, result.calls, result.spans + + def flush_with_spans(self): + """Return buffered input as visible content when generation ends early.""" + if not self.tools: + return "", [], [] + buffered = self.detector._buffer + if self._emitted_calls: + if self.detector.eot_token: + buffered = buffered.replace(self.detector.eot_token, "") + if self.detector.tool_call_separator and buffered.startswith(self.detector.tool_call_separator): + buffered = buffered[len(self.detector.tool_call_separator) :] + else: + buffered = getattr(self.detector, "_normal_text_buffer", "") + buffered + self.detector._normal_text_buffer = "" + self.detector._buffer = "" + if not buffered: + return "", [], [] + end = self.detector._source_length + return buffered, [], [TextSpan(end - len(buffered), end)] def has_tool_call(self, text: str) -> bool: """ @@ -2022,14 +2096,8 @@ def parse_non_stream(self, full_text: str) -> Tuple[str, list[ToolCallItem]]: - The remaining text after parsing that was not consumed by the detector (can be treated as normal text) - A list of tool calls parsed from the text """ - if not self.tools: - return full_text, [] - parsed_result = self.detector.detect_and_parse(full_text, self.tools) - tool_call_list = parsed_result.calls - if tool_call_list: - return parsed_result.normal_text, tool_call_list - else: - return full_text, [] + normal, calls, _ = self.parse_non_stream_with_spans(full_text) + return normal, calls def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: """ @@ -2043,16 +2111,5 @@ def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: - The normal text that should be displayed to the user - A list of tool calls parsed from the chunk """ - if not self.tools: - return chunk_text, [] - final_normal_text = "" - final_calls = [] - - sp_result = self.detector.parse_streaming_increment(chunk_text, self.tools) - if sp_result.normal_text: - final_normal_text = sp_result.normal_text - if sp_result.calls: - final_calls.extend(sp_result.calls) - final_normal_text = sp_result.normal_text - - return final_normal_text, final_calls + normal, calls, _ = self.parse_stream_chunk_with_spans(chunk_text) + return normal, calls diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 95dd057212..677b236acb 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -344,7 +344,6 @@ async def _generate( # 用于等待 pd_master 下发的交换信息 pd_event: asyncio.Event = None, ) -> AsyncGenerator[Tuple[int, str, dict, FinishStatus], None]: - start_time = time.time() # 在请求登记、编码和 PD 握手前获取语法；缓存未命中时先完成编译，避免后续流程等待或校验失败。 compiled_grammar = await sampling_params.get_compiled_grammar(self.output_grammar_compiler) @@ -756,7 +755,6 @@ async def transfer_to_next_module( self, group_req_objs: Optional[GroupReqObjs] = None, ): - if self.pd_mode.is_P_or_NORMAL(): if not self.args.disable_vision: self.send_to_visual.send_pyobj(group_req_objs.to_group_req_index(), protocol=pickle.HIGHEST_PROTOCOL) @@ -799,6 +797,7 @@ async def _wait_to_token_package( req_status: "ReqStatus", request: Request, ): + from lightllm.server.detokenization.decode_req import LIGHTLLM_DECODE_PREFIX_LENGTH event = req_status.event unfinished_count = sampling_params.best_of @@ -838,6 +837,10 @@ async def _wait_to_token_package( for sub_req_id, out_str, metadata, finish_status in req_status.out_token_info_list: # pd master 节点需要这个做统计信息， 所以放在元数据中返回给 pd master 节点 metadata["prompt_tokens"] = prompt_tokens + if sampling_params.output_top_logprobs >= 0 and metadata.get("count_output_tokens") == 1: + metadata["logprob_prompt_ids"] = prompt_ids[ + max(len(prompt_ids) - LIGHTLLM_DECODE_PREFIX_LENGTH, 0) : + ] gpu_prompt_cache_len = metadata.pop("prompt_cache_len", 0) cpu_prompt_cache_len = metadata.pop("cpu_prompt_cache_len", 0) diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 52e27992ec..fd787c6eea 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -104,7 +104,6 @@ class PDUpKVStatus: pd_kv_trans_params: bytes # pd kv 传输建立连接所使用的元数据对象 def __post_init__(self): - if not isinstance(self.group_request_id, int): error_info = "group_request_id only can be int" logger.error(error_info) @@ -173,8 +172,10 @@ class PDChunckedTransTask: decode_num_pages: Optional[int] decode_page_reg_desc: Optional[bytes] - first_gen_token_id: Optional[int] - first_gen_token_logprob: Optional[float] + first_gen_token_id: Optional[int] = None + first_gen_token_logprob: Optional[float] = None + first_gen_token_rank: Optional[int] = None + first_gen_token_output_top_logprobs: Optional[tuple[list[int], list[float]]] = None write_stage: Optional[str] = None @@ -273,6 +274,8 @@ def createRetObj(self) -> "PDChunckedTransTaskRet": error_info=self.error_info, first_gen_token_id=self.first_gen_token_id, first_gen_token_logprob=self.first_gen_token_logprob, + first_gen_token_rank=self.first_gen_token_rank, + first_gen_token_output_top_logprobs=self.first_gen_token_output_top_logprobs, ) return ret @@ -302,6 +305,8 @@ class PDChunckedTransTaskRet: error_info: str = None first_gen_token_id: Optional[int] = None first_gen_token_logprob: Optional[float] = None + first_gen_token_rank: Optional[int] = None + first_gen_token_output_top_logprobs: Optional[tuple[list[int], list[float]]] = None def get_key(self) -> str: return f"{self.request_id}_{self.start_kv_index}_{self.end_kv_index}" diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 9570ce0536..e6c1535211 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -12,10 +12,11 @@ # limitations under the License. import re -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id +from lightllm.server.text_spans import TextSpan, extract_text, slice_spans @dataclass @@ -25,6 +26,7 @@ class Event: event_type: str content: str raw_text: str = None # Original text including structural markers + spans: List[TextSpan] = field(default_factory=list) @dataclass @@ -165,7 +167,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if pos == len(tokens) - 1: # Last token emit, hold = prefix_hold(text[token.start : token.end], self.guard_tokens) if emit: - events.append(Event("normal", emit)) + events.append(Event("normal", emit, spans=[TextSpan(token.start, token.start + len(emit))])) return events, hold else: # Check if this might be commentary filler between blocks @@ -176,7 +178,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: content = text[token.start : token.end] # Skip standalone structural tokens that shouldn't be emitted as normal text if not self._is_standalone_structural_token(content): - events.append(Event("normal", content)) + events.append(Event("normal", content, spans=[TextSpan(token.start, token.end)])) pos += 1 elif token.type in ("START", "CHANNEL"): @@ -206,7 +208,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: # Unexpected token - only emit as text if it's not a standalone structural token content = text[token.start : token.end] if not self._is_standalone_structural_token(content): - events.append(Event("normal", content)) + events.append(Event("normal", content, spans=[TextSpan(token.start, token.end)])) pos += 1 return events, "" @@ -248,7 +250,7 @@ def _parse_partial_analysis(self, text: str, tokens: List[Token], start_pos: int # Return partial reasoning content and preserve the channel structure for next parse remaining_text = text[tokens[start_pos].start : content_start] - return Event("reasoning", content), remaining_text + return Event("reasoning", content, spans=[TextSpan(content_start, len(text), "reasoning")]), remaining_text def _extract_channel_type(self, header_text: str) -> Optional[str]: """Extract channel type from header, ignoring other attributes like to=... or <|constrain|>...""" @@ -298,7 +300,10 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if end_token_pos is None: return None # Incomplete content = text[content_start : tokens[end_token_pos].start] - return Event("normal", content), end_token_pos + 1 + return ( + Event("normal", content, spans=[TextSpan(content_start, tokens[end_token_pos].start)]), + end_token_pos + 1, + ) # Standard channel block processing - message_pos is already found above pos = channel_pos + 1 # Skip CHANNEL token @@ -334,7 +339,7 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if channel_type == "final": # Final blocks can end at end of input without requiring <|return|> content = text[content_start:] - return Event("normal", content), end_pos + return Event("normal", content, spans=[TextSpan(content_start, len(text))]), end_pos return None # Analysis and commentary need proper end tokens end_token = tokens[end_pos] @@ -345,15 +350,24 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if end_token.type == "CALL": # Built-in tools (browser, python) use analysis channel with <|call|> raw_text = text[tokens[start_pos].start : end_token.end] - return Event("tool_call", content.strip(), raw_text), end_pos + 1 + return ( + Event("tool_call", content.strip(), raw_text, [TextSpan(tokens[start_pos].start, end_token.end)]), + end_pos + 1, + ) else: - return Event("reasoning", content), end_pos + 1 + return ( + Event("reasoning", content, spans=[TextSpan(content_start, end_token.start, "reasoning")]), + end_pos + 1, + ) elif channel_type == "commentary": if end_token.type == "CALL": raw_text = text[tokens[start_pos].start : end_token.end] - return Event("tool_call", content.strip(), raw_text), end_pos + 1 + return ( + Event("tool_call", content.strip(), raw_text, [TextSpan(tokens[start_pos].start, end_token.end)]), + end_pos + 1, + ) else: - return Event("normal", content), end_pos + 1 + return Event("normal", content, spans=[TextSpan(content_start, end_token.start)]), end_pos + 1 elif channel_type == "final": # For final blocks, include any trailing TEXT immediately after <|return|> final_content = content @@ -361,8 +375,18 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option next_token = tokens[end_pos + 1] if next_token.type == "TEXT": final_content += text[next_token.start : next_token.end] - return Event("normal", final_content), end_pos + 2 - return Event("normal", final_content), end_pos + 1 + return ( + Event( + "normal", + final_content, + spans=[ + TextSpan(content_start, end_token.start), + TextSpan(next_token.start, next_token.end), + ], + ), + end_pos + 2, + ) + return Event("normal", final_content, spans=[TextSpan(content_start, end_token.start)]), end_pos + 1 return None, end_pos + 1 @@ -439,11 +463,45 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if m: channel, reasoning, final = m.groups() if channel.lower() == "analysis" and reasoning.strip(): - events.append(Event("reasoning", reasoning.strip())) + events.append( + Event( + "reasoning", + reasoning.strip(), + spans=[ + TextSpan( + m.start(2) + len(reasoning) - len(reasoning.lstrip()), + m.end(2) - len(reasoning) + len(reasoning.rstrip()), + "reasoning", + ) + ], + ) + ) elif channel.lower() == "commentary" and reasoning.strip(): - events.append(Event("normal", reasoning.strip())) + events.append( + Event( + "normal", + reasoning.strip(), + spans=[ + TextSpan( + m.start(2) + len(reasoning) - len(reasoning.lstrip()), + m.end(2) - len(reasoning) + len(reasoning.rstrip()), + ) + ], + ) + ) if final.strip(): - events.append(Event("normal", final.strip())) + events.append( + Event( + "normal", + final.strip(), + spans=[ + TextSpan( + m.start(m.lastindex) + len(final) - len(final.lstrip()), + m.end(m.lastindex) - len(final) + len(final.rstrip()), + ) + ], + ) + ) return events, "" # If assistantfinal appears to be incomplete (e.g., 'assistantfin'), hold entire buffer @@ -456,7 +514,18 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if m: final = m.group(1) if final.strip(): - events.append(Event("normal", final.strip())) + events.append( + Event( + "normal", + final.strip(), + spans=[ + TextSpan( + m.start(m.lastindex) + len(final) - len(final.lstrip()), + m.end(m.lastindex) - len(final) + len(final.rstrip()), + ) + ], + ) + ) return events, "" m = self.patterns["analysis_only"].match(text) @@ -465,7 +534,9 @@ def parse(self, text: str) -> Tuple[List[Event], str]: emit, hold = prefix_hold(content, ["assistantfinal"]) if channel.lower() == "analysis" and emit: # Stream reasoning content as-is based on structural markers only. - events.append(Event("reasoning", emit)) + events.append( + Event("reasoning", emit, spans=[TextSpan(m.start(2), m.start(2) + len(emit), "reasoning")]) + ) # Keep the channel header in the remaining buffer to continue parsing # subsequent chunks in the text fallback format. Preserve any held # prefix that may complete into "assistantfinal". @@ -476,7 +547,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: elif channel.lower() == "commentary" and emit: # For commentary, stream as normal text. Preserve spaces unless holding. content_out = emit if hold else emit.strip() - events.append(Event("normal", content_out)) + events.append(Event("normal", content_out, spans=extract_text(emit, m.start(2), strip=not hold)[1])) if hold: return events, text[: m.start(2)] + hold else: @@ -486,7 +557,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: emit, hold = prefix_hold(text, ["analysis", "commentary", "assistantfinal"]) if emit: - events.append(Event("normal", emit)) + events.append(Event("normal", emit, spans=[TextSpan(0, len(emit))])) return events, hold @@ -494,12 +565,20 @@ class HarmonyParser: """Facade for parsing Harmony format, switching between strategies.""" def __init__(self): + self.track_spans = False + self.last_spans = [] + self._source_offsets = [] + self._source_length = 0 self.strategy = None self._buffer = "" self._should_filter_commentary = False # Track if we should filter commentary in next chunks self._partial_commentary = "" # Track partial commentary being built across chunks def parse(self, chunk: str) -> List[Event]: + if self.track_spans: + self._source_offsets.extend(range(self._source_length, self._source_length + len(chunk))) + self._source_length += len(chunk) + self.last_spans = [] self._buffer += chunk if self.strategy is None: @@ -524,6 +603,18 @@ def parse(self, chunk: str) -> List[Event]: # Check if we should start filtering commentary (after <|call|> token or tool_call event) buffer_has_call_token = self._buffer.rstrip().endswith("<|call|>") + if self.track_spans: + source = self._buffer + offsets = self._source_offsets + start = source.rfind(remaining) + if start >= 0: + self._source_offsets = offsets[start : start + len(remaining)] + else: + prefix = 0 + while prefix < min(len(source), len(remaining)) and source[prefix] == remaining[prefix]: + prefix += 1 + suffix = len(remaining) - prefix + self._source_offsets = offsets[:prefix] + (offsets[-suffix:] if suffix else []) self._buffer = remaining # Filter events for streaming case @@ -565,6 +656,20 @@ def parse(self, chunk: str) -> List[Event]: elif buffer_has_call_token: self._should_filter_commentary = True # Filter commentary after <|call|> token + if self.track_spans: + spans = event.spans + channel = "reasoning" if event.event_type == "reasoning" else "content" + for span in spans: + for pos in range(span.start, span.end): + absolute = offsets[pos] + if ( + self.last_spans + and self.last_spans[-1].end == absolute + and self.last_spans[-1].channel == channel + ): + self.last_spans[-1].end += 1 + else: + self.last_spans.append(TextSpan(absolute, absolute + 1, channel)) filtered_events.append(event) return filtered_events @@ -577,7 +682,9 @@ def __init__( self, normal_text: Optional[str] = None, reasoning_text: Optional[str] = None, + spans=None, ): + self.spans = spans or [] self.normal_text = normal_text or "" self.reasoning_text = reasoning_text or "" @@ -598,6 +705,7 @@ def __init__( self.stream_reasoning = stream_reasoning self._buffer = "" + self._source_length = 0 self.stripped_think_start = False @property @@ -613,21 +721,29 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: in_reasoning = self._in_reasoning or self.think_start_token in text if not in_reasoning: - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) # The text is considered to be in a reasoning block. - processed_text = text.replace(self.think_start_token, "").strip() + processed_text, spans = extract_text(text, strip=True, remove=(self.think_start_token,)) if self.think_end_token not in processed_text: # Assume reasoning was truncated before `` token - return StreamingParseResult(reasoning_text=processed_text) + for span in spans: + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=processed_text, spans=spans) # Extract reasoning content splits = processed_text.split(self.think_end_token, maxsplit=1) reasoning_text = splits[0] normal_text = splits[1].strip() - - return StreamingParseResult(normal_text=normal_text, reasoning_text=reasoning_text) + content_start = len(reasoning_text) + len(self.think_end_token) + len(splits[1]) - len(splits[1].lstrip()) + reasoning_spans = slice_spans(spans, 0, len(reasoning_text)) + for span in reasoning_spans: + span.channel = "reasoning" + content_spans = slice_spans(spans, content_start, content_start + len(normal_text)) + return StreamingParseResult( + normal_text=normal_text, reasoning_text=reasoning_text, spans=reasoning_spans + content_spans + ) def flush(self) -> StreamingParseResult: """ @@ -640,10 +756,12 @@ def flush(self) -> StreamingParseResult: return StreamingParseResult() remaining = self._buffer self._buffer = "" + span = TextSpan(self._source_length - len(remaining), self._source_length) if self._in_reasoning: - return StreamingParseResult(reasoning_text=remaining) + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=remaining, spans=[span]) else: - return StreamingParseResult(normal_text=remaining) + return StreamingParseResult(normal_text=remaining, spans=[span]) def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: """ @@ -658,6 +776,8 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: """ self._buffer += new_text current_text = self._buffer + offset = self._source_length - len(current_text) + spans = [TextSpan(offset, self._source_length)] if current_text else [] # If the current text is a prefix of the think token, keep buffering # until we can confirm or refute the tag. @@ -669,7 +789,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: # Strip `` token if present if not self.stripped_think_start and self.think_start_token in current_text: - current_text = current_text.replace(self.think_start_token, "") + current_text, spans = extract_text(current_text, offset, remove=(self.think_start_token,)) self.stripped_think_start = True self._in_reasoning = True @@ -683,18 +803,27 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: self._in_reasoning = False normal_text = current_text[end_idx + len(self.think_end_token) :] - return StreamingParseResult(normal_text=normal_text, reasoning_text=reasoning_text.rstrip()) + reasoning_text = reasoning_text.rstrip() + reasoning_spans = slice_spans(spans, 0, len(reasoning_text)) + for span in reasoning_spans: + span.channel = "reasoning" + content_spans = slice_spans(spans, end_idx + len(self.think_end_token), len(current_text)) + return StreamingParseResult( + normal_text=normal_text, reasoning_text=reasoning_text, spans=reasoning_spans + content_spans + ) # Always stream reasoning content immediately. # stream_reasoning flag is ignored for streaming responses. if self._in_reasoning: self._buffer = "" - return StreamingParseResult(reasoning_text=current_text) + for span in spans: + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=current_text, spans=spans) # If we're not in a reasoning block return as normal text if not self._in_reasoning: self._buffer = "" - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=current_text, spans=spans) return StreamingParseResult() @@ -794,9 +923,12 @@ def grammar_start_token(self) -> str: return "<|channel|>final<|message|>" def detect_and_parse(self, text: str) -> StreamingParseResult: + self.parser.track_spans = True events = self.parser.parse(text) + spans = list(self.parser.last_spans) # Flush the buffer for one-shot parsing events += self.parser.parse("") + spans.extend(self.parser.last_spans) reasoning_text = "".join([e.content for e in events if e.event_type == "reasoning"]) normal_parts = [] @@ -812,10 +944,13 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: return StreamingParseResult( normal_text=normal_text, reasoning_text=reasoning_text, + spans=spans, ) def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + self.parser.track_spans = True events = self.parser.parse(new_text) + spans = self.parser.last_spans reasoning_text = "".join([e.content for e in events if e.event_type == "reasoning"]) normal_parts = [] @@ -830,6 +965,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: return StreamingParseResult( normal_text=normal_text, reasoning_text=reasoning_text, + spans=spans, ) @@ -849,13 +985,14 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False) self.is_first_chunk = False def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + spans = [TextSpan(self._source_length - len(new_text), self._source_length)] if new_text else [] if not self.is_first_chunk: self.is_first_chunk = True new_text = self.think_start_token + new_text - return StreamingParseResult(normal_text=new_text) + return StreamingParseResult(normal_text=new_text, spans=spans) def detect_and_parse(self, text: str) -> StreamingParseResult: - return StreamingParseResult(normal_text=self.think_start_token + text) + return StreamingParseResult(normal_text=self.think_start_token + text, spans=[TextSpan(0, len(text))]) class NanoV3Detector(BaseReasoningFormatDetector): @@ -978,15 +1115,30 @@ def update_reasoning_token_count(self, token_id: int) -> None: def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]: """Non-streaming call: one-time parsing""" - ret = self.detector.detect_and_parse(full_text) - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.parse_non_stream_with_spans(full_text) + return reasoning, content + + def parse_non_stream_with_spans(self, full_text: str): + result = self.detector.detect_and_parse(full_text) + spans = sorted(result.spans, key=lambda span: span.channel != "reasoning") + return result.reasoning_text, result.normal_text, spans def parse_stream_chunk(self, chunk_text: str) -> Tuple[Optional[str], Optional[str]]: """Streaming call: incremental parsing""" - ret = self.detector.parse_streaming_increment(chunk_text) - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.parse_stream_chunk_with_spans(chunk_text) + return reasoning, content + + def parse_stream_chunk_with_spans(self, chunk_text: str): + self.detector._source_length += len(chunk_text) + result = self.detector.parse_streaming_increment(chunk_text) + spans = sorted(result.spans, key=lambda span: span.channel != "reasoning") + return result.reasoning_text, result.normal_text, spans + + def flush_with_spans(self): + result = self.detector.flush() + return result.reasoning_text, result.normal_text, result.spans def flush(self) -> Tuple[Optional[str], Optional[str]]: """Flush remaining buffered content when generation ends prematurely.""" - ret = self.detector.flush() - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.flush_with_spans() + return reasoning, content diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 0026c7cbf0..aea1e7ffb5 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1009,13 +1009,20 @@ def handle( eos_ids: List[int], is_master_in_dp: bool, extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None, - pd_prefill_chunked_handle_func: Optional[Callable[[InferReq, int, float, int], None]] = None, + pd_prefill_chunked_handle_func: Optional[Callable[[InferReq, int, float, int, int, object], None]] = None, output_top_logprobs=None, ): # pd_prefill_chunked_handle_func 主要是为了处理 pd prefill 模式下 # 分块 prefill 后，形成对应的pd 分块传输处理。 if pd_prefill_chunked_handle_func is not None: - pd_prefill_chunked_handle_func(self.req_obj, next_token_id, next_token_logprob, self.output_len) + pd_prefill_chunked_handle_func( + self.req_obj, + next_token_id, + next_token_logprob, + self.output_len, + next_token_rank, + output_top_logprobs, + ) if self.output_len <= 0: return diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 208fc46fbc..6ab4920fc4 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -603,9 +603,12 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): InferReqUpdatePack(req_obj=req, output_len=req.cur_output_len).handle( next_token_id=obj.first_gen_token_id, next_token_logprob=obj.first_gen_token_logprob, - next_token_rank=-1, + next_token_rank=( + obj.first_gen_token_rank if obj.first_gen_token_rank is not None else -1 + ), eos_ids=self.eos_id, is_master_in_dp=self.is_master_in_dp, + output_top_logprobs=obj.first_gen_token_output_top_logprobs, ) return @@ -939,7 +942,9 @@ def _post_handle( pack: InferReqUpdatePack = pack output_top_logprobs = None if topk_row is not None: - output_top_logprobs = (topk_row.token_ids.tolist(), topk_row.logprobs.tolist()) + next_token_rank = int(topk_row.sampled_rank.item()) + if topk_row.token_ids.numel(): + output_top_logprobs = (topk_row.token_ids.tolist(), topk_row.logprobs.tolist()) pack.handle( next_token_id=next_token_id, next_token_logprob=next_token_logprob, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index e3e01da612..d48e092e26 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -17,11 +17,13 @@ class TopkLogprobRow: """一行 top-k 候选。token_ids 与 logprobs 等长，按分数从高到低对齐。 batch_index 是该行在本次采样 logits 中的下标。返回列表与 logits 行对齐，未请求 logprobs 的位置是 None。 + K=0 的行只有 sampled_rank，候选数组为空。 """ batch_index: int token_ids: torch.Tensor logprobs: torch.Tensor + sampled_rank: torch.Tensor TopkLogprobRows = List[Optional[TopkLogprobRow]] @@ -143,14 +145,13 @@ def _collect_topk_logprobs_rows( ) -> Tuple[List[Tuple[int, int]], int]: """Rows that asked for candidate topk logprobs, as (batch_index, topk), plus the maximum topk. - topk == 0 still runs the fixed top-20 cutoff on the sampled logprob and writes no candidates. - batch_index is the position in the original reqs and logits. max_topk is 0 when the list is empty. + topk == 0 still needs the sampled token rank. batch_index is the position + in the original reqs and logits. max_topk is 0 when the list is empty. """ rows = [] max_topk = 0 for batch_index, req_obj in enumerate(reqs): topk = req_obj.sampling_param.shm_param.output_top_logprobs - # topk == 0 也要进来：没有候选列表，但仍要把落在前 20 名之外的采样 logprob 写成 -9999。 if topk >= 0 and (has_output is None or has_output[batch_index]): rows.append((batch_index, topk)) max_topk = max(max_topk, topk) @@ -172,7 +173,7 @@ def _capture_topk_logprobs( rows = g_pin_mem_manager.alloc_pin_tensor("output_logprob_rows", 3 * len(capture_rows), torch.int32).view(-1, 3) rows.numpy()[:] = capture_rows - token_ids, logprobs = capture_output_logprobs( + token_ids, logprobs, ranks = capture_output_logprobs( logits=logits, sampled_ids=next_token_ids, sampled_logprobs=next_token_logprobs, @@ -182,16 +183,16 @@ def _capture_topk_logprobs( ) token_ids_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_topk_token_ids", gpu_tensor=token_ids) logprobs_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_topk_logprobs", gpu_tensor=logprobs) + ranks_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_sampled_ranks", gpu_tensor=ranks) topk_rows: TopkLogprobRows = [None] * logits.shape[0] offset = 0 - for batch_index, topk in topk_logprobs_rows: - # topk == 0 只参与上面的采样分数改写，不返回候选。 - if topk > 0: - topk_rows[batch_index] = TopkLogprobRow( - batch_index=batch_index, - token_ids=token_ids_cpu[offset : offset + topk], - logprobs=logprobs_cpu[offset : offset + topk], - ) + for index, (batch_index, topk) in enumerate(topk_logprobs_rows): + topk_rows[batch_index] = TopkLogprobRow( + batch_index=batch_index, + token_ids=token_ids_cpu[offset : offset + topk], + logprobs=logprobs_cpu[offset : offset + topk], + sampled_rank=ranks_cpu[index : index + 1], + ) offset += topk return next_token_logprobs, topk_rows diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index ca4b51e9b5..480e28fba9 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -333,6 +333,10 @@ def accept_peer_task_loop( local_trans_task.shared_time_out_mark.update() local_trans_task.first_gen_token_id = remote_trans_task.first_gen_token_id local_trans_task.first_gen_token_logprob = remote_trans_task.first_gen_token_logprob + local_trans_task.first_gen_token_rank = remote_trans_task.first_gen_token_rank + local_trans_task.first_gen_token_output_top_logprobs = ( + remote_trans_task.first_gen_token_output_top_logprobs + ) self.ready_page_task_queue.put(local_trans_task) if self.args.pd_detail_log: logger.debug(f"recv WRITE done from prefill: {remote_trans_task.to_str()}") diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py index da4ce0d073..1d7ec7d9a1 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py @@ -66,7 +66,13 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: return ans_list def _prefill_chuncked_handle_func( - self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int + self, + req_obj: InferReq, + next_token_id: int, + next_token_prob: float, + output_len: int, + next_token_rank: int, + output_top_logprobs=None, ): """ 在每一步chuncked prefill 后，尝试生成chuncked 传输任务，发个 kv_move_manager 进行处理。 @@ -105,6 +111,8 @@ def _prefill_chuncked_handle_func( ) trans_task_list[-1].first_gen_token_id = next_token_id trans_task_list[-1].first_gen_token_logprob = next_token_prob + trans_task_list[-1].first_gen_token_rank = next_token_rank + trans_task_list[-1].first_gen_token_output_top_logprobs = output_top_logprobs if self.is_master_in_dp: for trans_task in trans_task_list: @@ -155,6 +163,8 @@ def _create_pd_trans_task( decode_page_reg_desc=pd_decode_node_info.page_reg_desc, first_gen_token_id=None, first_gen_token_logprob=None, + first_gen_token_rank=None, + first_gen_token_output_top_logprobs=None, page_kind=page_kind, req_idx=req_idx, ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py index 5f4465a730..673a17365e 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py @@ -24,10 +24,22 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: return PDChunkedPrefillForPrefillNode._filter_not_ready_reqs(self, req_ids) def _prefill_chuncked_handle_func( - self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int + self, + req_obj: InferReq, + next_token_id: int, + next_token_prob: float, + output_len: int, + next_token_rank: int, + output_top_logprobs=None, ): return PDChunkedPrefillForPrefillNode._prefill_chuncked_handle_func( - self, req_obj=req_obj, next_token_id=next_token_id, next_token_prob=next_token_prob, output_len=output_len + self, + req_obj=req_obj, + next_token_id=next_token_id, + next_token_prob=next_token_prob, + output_len=output_len, + next_token_rank=next_token_rank, + output_top_logprobs=output_top_logprobs, ) def _create_pd_trans_task( diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index b3c1dcc609..dfeb322d51 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -446,8 +446,6 @@ def success_loop(self): self.transporter.release_xfer_handle(trans_task.xfer_handle) ret = trans_task.createRetObj() - ret.first_gen_token_id = None - ret.first_gen_token_logprob = None self.task_out_queue.put(ret) if self.args.pd_detail_log: diff --git a/lightllm/server/text_spans.py b/lightllm/server/text_spans.py new file mode 100644 index 0000000000..212fbf160c --- /dev/null +++ b/lightllm/server/text_spans.py @@ -0,0 +1,40 @@ +from dataclasses import dataclass +from typing import List + + +@dataclass +class TextSpan: + start: int + end: int + channel: str = "content" + + +def slice_spans(spans: List[TextSpan], start: int, end: int) -> List[TextSpan]: + selected = [] + offset = 0 + for span in spans: + size = span.end - span.start + left, right = max(start - offset, 0), min(end - offset, size) + if left < right: + selected.append(TextSpan(span.start + left, span.start + right, span.channel)) + offset += size + return selected + + +def extract_text(text: str, offset=0, channel="content", strip=False, remove=()): + spans = [TextSpan(offset, offset + len(text), channel)] if text else [] + for marker in remove: + if not marker: + continue + start = text.find(marker) + while start >= 0: + end = start + len(marker) + spans = slice_spans(spans, 0, start) + slice_spans(spans, end, len(text)) + text = text[:start] + text[end:] + start = text.find(marker, start) + if strip: + start = len(text) - len(text.lstrip()) + end = len(text.rstrip()) + spans = slice_spans(spans, start, end) + text = text.strip() + return text, spans diff --git a/unit_tests/README.md b/unit_tests/README.md new file mode 100644 index 0000000000..a6ad0a6682 --- /dev/null +++ b/unit_tests/README.md @@ -0,0 +1,26 @@ +# Unit tests + +Keep tests that can catch an incorrect result, broken public contract, resource +lifetime bug, or a previously observed failure. Test the production entry point +with small inputs; mock external dependencies only as needed. + +- Add a regression case to an existing test when it exercises the same behavior. + Parameterize inputs instead of copying setup and assertions into another test. +- For kernels, compare outputs and mutated state against an independent reference. + Include relevant dtype, layout, padding, and boundary cases without taking the + Cartesian product of unrelated options. +- Do not test class inheritance, private method inventories, or dataclass field + declarations. Exercise the behavior those implementations support. +- A CPU reimplementation tested only against its own expected output does not + validate the CUDA path. Run the real implementation or explicitly skip when its + prerequisites are unavailable. +- Keep timing loops and performance reports out of unit tests. Benchmarks and + service-level tests belong under `test/`; ordinary unit tests need assertions. +- When removing a test, check whether another test covers the same failure mode. + Similar setup alone does not make cache ownership, transfer boundaries, or + asynchronous failure tests redundant. + +Run affected tests with `python -m pytest ` in an environment with their +required dependencies. For changed Python files, run +`pre-commit run --files `. Report skipped hardware/provider checks separately +from passing tests. diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py index badc2adf0c..4a99ef89bc 100644 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -10,6 +10,7 @@ from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace from lightllm.utils.envs_utils import get_env_start_args @@ -79,6 +80,11 @@ class TestModel(SimpleNamespace): from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend backend = DeepseekV4FlashMlaFp8SparseAttBackend(model=model) + model.dsv4_workspace = DeepseekV4Workspace(model) + model.dsv4_workspace.c128_compress_scratch = torch.empty( + (1, model.batch_max_tokens, model.config["head_dim"]), dtype=torch.float32, device="cuda" + ) + manager.dsv4_workspace = model.dsv4_workspace yield manager, window, backend torch.cuda.synchronize() get_env_start_args.cache_clear() @@ -88,6 +94,7 @@ def make_infer(manager, window): infer = DeepseekV4InferStateInfo() infer.att_state_manager = manager infer.att_state_role = "main" + infer.dsv4_workspace = manager.dsv4_workspace infer.first_logical_layer_id = window.config.logical_layer_ids[0] return infer diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py index 26d4425ac7..c943a78bee 100644 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -2,12 +2,6 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - from typing import TYPE_CHECKING import json @@ -16,7 +10,6 @@ import pytest import torch -from lightllm.common.attention_state.full_att import MLAFullAttentionState, PackedMLAFullAttentionState from lightllm.utils.envs_utils import get_env_start_args from unit_tests.common.att_resources.test_capacity import allocated_bytes from unit_tests.server.router.model_infer.test_cpu_resource_layout import CASES, make_runtime, runtime @@ -32,48 +25,6 @@ ] -@pytest.mark.parametrize("kind,encoding,model_cls", CASES) -def test_model_storage_write_and_copy_preserve_all_fields( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - kind: str, - encoding: str, - model_cls: type[TpPartBaseModel], -): - _, (resources,) = make_runtime( - monkeypatch=monkeypatch, - tmp_path=tmp_path, - kind=kind, - encoding=encoding, - model_cls=model_cls, - tp=1, - ) - (state,) = resources.get_full_attention_states("main") - config = state.config - for storage_buffer in full_state_test_buffers(state).values(): - storage_buffer.zero_() - ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") - is_mla = isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) - dim = config.latent_dim if is_mla else config.head_dim - kv = torch.randn( - size=(2, 1 if is_mla else 2 * config.head_num, dim), - device="cuda", - dtype=torch.bfloat16, - ) - state.write_full_att_kv(config.logical_layer_groups[-1][0], ids, kv) - if len(config.logical_layer_groups) > 1: - assert not read_kv_storage(state)[0].any() - resources.copy_full_att_tokens(ids, torch.tensor([1, 5], dtype=torch.int32, device="cuda")) - for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): - if "token" in field.axes: - assert torch.equal( - field.tensor[:, ids].contiguous().view(torch.uint8), - field.tensor[:, [1, 5]].contiguous().view(torch.uint8), - ), name - assert state.capacity == resources.full_att_allocator.size == 16 - assert read_kv_storage(state).shape[1] == 20 # Includes the whole HOLD page. - - @pytest.mark.parametrize("mode", ["normal", "prefill"]) @pytest.mark.parametrize("kind,encoding,model_cls", CASES[:8]) def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index f66767b9f4..c416d12cdc 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -122,6 +122,8 @@ def _task(peer: PDAgentMetadata, kind: str): decode_page_reg_desc=peer.page_reg_desc, first_gen_token_id=42, first_gen_token_logprob=-0.5, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([42, 7], [-0.5, -0.6]), src_page_index=1 if kind == "kv" else 2, page_kind=kind, req_idx=1, diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py deleted file mode 100644 index ca97cc7b70..0000000000 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ /dev/null @@ -1,37 +0,0 @@ -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -import torch - -from lightllm.common.attention_state.full_att import ( - GroupQuantizedAttentionConfig, - GroupQuantizedFullAttentionState, -) - - -def test_quantized_kv_buffers_include_the_full_hold_page(): - allocations = [] - for bits in (8, 4): - state = GroupQuantizedFullAttentionState( - GroupQuantizedAttentionConfig( - logical_layer_groups=((0,), (1,), (2,)), - global_head_num=2, - head_dim=16, - scale_dtype=torch.float16, - bits=bits, - ), - page_size=4, - device="cpu", - ) - state.allocate_storage(32) - allocations.extend(tuple(buffer.shape) for buffer in full_state_test_buffers(state).values()) - - assert allocations == [ - (3, 36, 2, 16), - (3, 36, 2, 16), - (3, 36, 2, 2), - (3, 36, 2, 2), - (3, 36, 2, 8), - (3, 36, 2, 8), - (3, 36, 2, 2), - (3, 36, 2, 2), - ] diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py deleted file mode 100644 index 67c2fe8c2e..0000000000 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from typing import TYPE_CHECKING -import pytest -import torch - -from unit_tests.common.att_resources.model_helpers import main_linear_state -from unit_tests.common.att_resources.test_model_resources import make_resources - -if TYPE_CHECKING: - from easydict import EasyDict - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("glm", [False, True]) -def test_main_and_draft_transfers_copy_pd_and_dp_windows( - glm: bool, - pd_args: EasyDict, - monkeypatch: pytest.MonkeyPatch, -): - pd_args.pd_kv_page_size = 128 - source, _ = make_resources( - glm=glm, - tp=1, - rank=0, - mtp=2, - draft_layers=2, - ) - target, _ = make_resources( - glm=glm, - tp=1, - rank=0, - mtp=2, - draft_layers=2, - ) - for index, role in enumerate(("main", "draft0")): - state = source.get_full_attention_states(role)[0] - for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): - if "token" in field.axes: - field.tensor.fill_(11 + index) - if glm: - getattr(source, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 0].fill_(7 + index) - getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 2].fill_(19) - main_linear_state(source).conv_state.fill_(3) - main_linear_state(source).ssm_state.fill_(5) - page = torch.empty(source.get_pd_operator().get_page_bytes(), dtype=torch.uint8, device="cuda") - for kind, getter in ( - ("kv", "get_pd_full_att_page_layouts"), - ("att_state", "get_pd_bounded_att_page_layouts"), - ): - source.get_pd_operator().write_page( - pd_page_buffers=page.unsqueeze(0), - pd_page_index=0, - kind=kind, - layouts=getattr(source.get_pd_operator(), getter)(), - **( - {"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} - if kind == "kv" - else {"req_idx": 0, "state_token_len": 8} - ), - ) - target.get_pd_operator().read_page( - pd_page_buffers=page.unsqueeze(0), - pd_page_index=0, - kind=kind, - layouts=getattr(target.get_pd_operator(), getter)(), - **( - {"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} - if kind == "kv" - else {"req_idx": 1, "state_token_len": 8} - ), - ) - src_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - dst_indices = torch.tensor([10, 11], dtype=torch.int64, device="cuda") - - def reject_host_list(tensor: torch.Tensor) -> list[int]: - pytest.fail("DP manager and States must preserve Tensor indices") - - with monkeypatch.context() as patch: - patch.setattr(torch.Tensor, "tolist", reject_host_list) - target.get_dp_shared_operator().copy_from_dp( - remote_att_manager=source, - src_token_indices=src_indices, - dst_token_indices=dst_indices, - src_big_page_ids=[], - dst_big_page_ids=[], - src_req_idx=0, - dst_req_idx=2, - ) - for index, role in enumerate(("main", "draft0")): - state = target.get_full_attention_states(role)[0] - for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): - if "token" in field.axes: - assert (field.tensor[:, 8:12] == 11 + index).all() - if glm: - tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer - expected_tail = torch.zeros_like(tail[:, 1]) - for position in range(4, 8): - expected_tail[:, position % tail.shape[2]].fill_(7 + index) - assert torch.equal(tail[:, 1], expected_tail) # PD restores only the prefill logical window. - assert (tail[:, 2] == 7 + index).all() # DP 直接复制相同 MTP 容量的 ring。 - - slot = source.bounded_att_big_page_allocator.alloc(1)[0] - source.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_slot_ids=[slot], - ) - torch.cuda.synchronize() - cpu = torch.empty((1, source.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) - source.get_cpu_cache_operator().write_req( - token_indices=torch.arange(8, dtype=torch.long, device="cuda"), - cpu_page_buffers=cpu, - cpu_page_indices=[0], - layouts=source.get_cpu_cache_operator().get_page_layouts(), - big_page_slot_ids=[slot], - ) - torch.cuda.synchronize() - dst_slot = target.bounded_att_big_page_allocator.alloc(1)[0] - target.get_cpu_cache_operator().read_req( - cpu_page_buffers=cpu, - cpu_page_indices=[0], - token_indices=torch.arange( - start=8, - end=16, - dtype=torch.long, - device="cuda", - ), - state_token_len=8, - layouts=target.get_cpu_cache_operator().get_page_layouts(), - big_page_slot_ids=[dst_slot], - req_idx=1, - ) - for index, role in enumerate(("main", "draft0")): - state = target.get_full_attention_states(role)[0] - for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): - if "token" in field.axes: - assert (field.tensor[:, 8:16] == 11 + index).all() - if glm: - # CPU 组边界恢复不修改无重叠窗口，保留先前 PD/DP 的数据。 - tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer - for position in range(4, 8): - assert (tail[:, 1, position % tail.shape[2]] == 7 + index).all() - assert (tail[:, 2] == 7 + index).all() diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py index 1371cfe47e..c44cc68b5a 100644 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -316,6 +316,9 @@ def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( pytest.skip("two CUDA devices required for IPC peer mapping") monkeypatch.setattr(shared_att_resources_utils, "kv_trans_use_p2p", lambda: p2p) monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", mp.reductions.reduce_tensor.__code__) + # Some containers lack pidfd_getfd, which recent torch releases require for + # expandable-segment CUDA IPC handles. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") source_device = int(p2p) with torch.cuda.device(source_device): resources, full = make_resources( diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py index 6eb8fa1be8..5e838a4931 100644 --- a/unit_tests/common/att_resources/test_state_creation.py +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -12,12 +12,6 @@ from lightllm.common.attention_state.full_att import ( GQAAttentionConfig, MLAConfig, - FullAttentionState, - GQAFullAttentionState, - GroupQuantizedFullAttentionState, - StaticFP8FullAttentionState, - MLAFullAttentionState, - PackedMLAFullAttentionState, PackedMLALogicalConfig, create_gqa_state, create_mla_state, @@ -216,20 +210,3 @@ def test_create_rejects_encoding_from_another_logical_family(): config=MLAConfig(logical_layer_groups=((0,),), latent_dim=576, latent_dtype=torch.bfloat16), device="meta", ) - - -@pytest.mark.parametrize( - "state_class", - [ - GQAFullAttentionState, - GroupQuantizedFullAttentionState, - StaticFP8FullAttentionState, - MLAFullAttentionState, - PackedMLAFullAttentionState, - ], -) -def test_full_attention_implementations_are_independent(state_class: type[FullAttentionState]): - # 具体方案只共享抽象契约，不继承其他具体方案的存储或操作实现。 - assert state_class.__bases__ == (FullAttentionState,) - assert not state_class.__abstractmethods__ - assert FullAttentionState.__abstractmethods__ <= state_class.__dict__.keys() diff --git a/unit_tests/common/att_resources/test_validate_model_cache.py b/unit_tests/common/att_resources/test_validate_model_cache.py index 42508cd0ee..5e31d2a0d5 100644 --- a/unit_tests/common/att_resources/test_validate_model_cache.py +++ b/unit_tests/common/att_resources/test_validate_model_cache.py @@ -1,3 +1,5 @@ +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path from types import SimpleNamespace import pytest @@ -5,7 +7,13 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import create_compressed_att_state from lightllm.models.deepseek_v4.config import DSV4AttentionConfig -from test.att_resources import validate_model_cache + +spec = spec_from_file_location( + "validate_model_cache", Path(__file__).parents[3] / "test/att_resources/validate_model_cache.py" +) +assert spec is not None and spec.loader is not None +validate_model_cache = module_from_spec(spec) +spec.loader.exec_module(validate_model_cache) @pytest.mark.parametrize("draft_count,bad_role", [(0, None), (2, None), (2, "main"), (2, "draft1")]) diff --git a/unit_tests/common/basemodel/test_att_sink_weight.py b/unit_tests/common/basemodel/test_att_sink_weight.py deleted file mode 100644 index 996fda7c86..0000000000 --- a/unit_tests/common/basemodel/test_att_sink_weight.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "heads,tp_size,rank,padding_multiple,data_type,start,end,padded_heads", - [ - (128, 8, 0, 64, torch.float32, 0, 16, 64), - (128, 8, 7, 64, torch.float32, 112, 128, 64), - (128, 2, 1, 64, torch.float32, 64, 128, 64), - (256, 2, 1, 64, torch.float32, 128, 256, 128), - (12, 4, 2, None, torch.bfloat16, 6, 9, 3), - (4, 8, 3, None, torch.bfloat16, 1, 2, 1), - ], -) -def test_load_tp_sink_weight(monkeypatch, heads, tp_size, rank, padding_multiple, data_type, start, end, padded_heads): - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_size)) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - kwargs = {} if padding_multiple is None else {"head_padding_multiple": padding_multiple} - sink = TpAttSinkWeight(all_q_head_num=heads, weight_name="attn_sink", data_type=data_type, **kwargs) - assert sink.weight.shape == (padded_heads,) - assert not sink.verify_load() - storage_ptr = sink.weight.data_ptr() - checkpoint = torch.arange(heads, dtype=torch.float64) + 0.5 - - sink.load_hf_weights({"attn_sink": checkpoint}) - - assert sink.verify_load() - assert sink.weight.data_ptr() == storage_ptr - assert sink.weight.dtype == data_type - assert sink.weight.is_cuda - torch.testing.assert_close(sink.weight[: end - start], checkpoint[start:end].to(device="cuda", dtype=data_type)) - assert torch.count_nonzero(sink.weight[end - start :]) == 0 diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py deleted file mode 100644 index 52b1c1cf2d..0000000000 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ /dev/null @@ -1,130 +0,0 @@ -from contextlib import contextmanager -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import cuda_graph as cuda_graph_module -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune - - -@pytest.fixture -def capture_env(monkeypatch): - env = SimpleNamespace(capturing=False, capture_count=0, graphs=[]) - env.args = SimpleNamespace( - enable_decode_microbatch_overlap=False, - enable_tpsp_mix_mode=False, - enable_torch_memory_saver=False, - ) - monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: env.args) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") - monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: None) - - class FakeGraph: - def __init__(self): - self.replays = 0 - env.graphs.append(self) - - def replay(self): - self.replays += 1 - - @contextmanager - def capture(graph, **kwargs): - assert not Autotuner.is_autotune_warmup() - env.capture_count += 1 - env.capturing = True - try: - yield - finally: - env.capturing = False - - monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) - monkeypatch.setattr(torch.cuda, "graph", capture) - return env - - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("level", [0, 1, 2, 3]) -def test_decode_attention_tunes_before_capture_only(capture_env, tmp_path, monkeypatch, overlap, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - capture_env.args.enable_decode_microbatch_overlap = overlap - graph = CudaGraph(1, 2, 1, max_batch_size=2) - benchmarks = [] - - def make_kernel(kernel_type): - @autotune( - kernel_name=kernel_type.value, - kernel_type=kernel_type, - configs_gen_func=lambda: [{"block": 1}, {"block": 2}], - static_key_func=lambda: {}, - run_key_func=lambda size: size, - ) - def kernel(size, run_config=None): - return size - - def bench(size, run_config): - assert kernel_type == AutotuneKernelType.DECODE_ATTENTION - assert Autotuner.is_kernel_autotune_warmup(kernel_type) - assert not capture_env.capturing - benchmarks.append(run_config) - return 1.0 / run_config["block"] - - cache_dir = tmp_path / kernel_type.value - cache_dir.mkdir() - kernel._cache_dir = str(cache_dir) - monkeypatch.setattr(kernel, "_bench", bench) - return kernel, cache_dir / "configs.json" - - general, general_cache = make_kernel(AutotuneKernelType.GENERAL) - attention, attention_cache = make_kernel(AutotuneKernelType.DECODE_ATTENTION) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] - outputs = [object() for _ in states] - forward_calls = [] - - def decode_func(*infer_states): - forward_calls.append(capture_env.capturing) - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) == (not capture_env.capturing) - for original, state in zip(states, infer_states): - assert (state is original) == capture_env.capturing - assert not hasattr(state, "temporary_buffer") - state.temporary_buffer = object() - general(state.input_ids.shape[0]) - attention(state.input_ids.shape[0]) - return tuple(outputs) if overlap else outputs[0] - - result = graph.capture_decode(decode_func, *states) - assert result == (tuple(outputs) if overlap else outputs[0]) - assert forward_calls == [False, True] - assert not Autotuner.is_autotune_warmup() - assert capture_env.capture_count == 1 - assert capture_env.graphs[0].replays == 1 - expected_benchmarks = 0 - if level in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]: - expected_benchmarks = 2 - assert len(benchmarks) == expected_benchmarks - assert attention_cache.exists() == (expected_benchmarks > 0) - assert not general_cache.exists() - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_warmup_failure_restores_phase_without_capturing(capture_env, overlap): - capture_env.args.enable_decode_microbatch_overlap = overlap - graph = CudaGraph(1, 2, 1, max_batch_size=2) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] - - def decode_func(*infer_states): - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - raise RuntimeError("decode warmup failed") - - with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): - with pytest.raises(RuntimeError, match="decode warmup failed"): - graph.capture_decode(decode_func, *states) - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert not Autotuner.is_autotune_warmup() - assert capture_env.capture_count == 0 - assert graph.graph == {} diff --git a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py deleted file mode 100644 index bed4f3967f..0000000000 --- a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py +++ /dev/null @@ -1,81 +0,0 @@ -from types import SimpleNamespace - -import torch - -import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph - - -class _GraphInferState: - def __init__(self, hidden_collector): - self.hidden_collector = hidden_collector - self.input_ids = torch.empty(4, dtype=torch.int64) - self.copied_from = None - self.replayed_with = None - - def copy_for_prefill_cuda_graph(self, new_infer_state): - self.copied_from = new_infer_state - - def prefill_replay(self, new_infer_state): - self.replayed_with = new_infer_state - - -class _IdentityPreInfer: - @staticmethod - def _tpsp_allgather(input, infer_state): - del infer_state - return input - - -def _create_layer_collector(monkeypatch): - monkeypatch.setattr( - hidden_collector_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_draft_model_dir=["/models/draft"]), - ) - monkeypatch.setattr( - hidden_collector_module.PretrainedConfig, - "get_config_dict", - lambda _: ({"target_layer_ids": [0]}, {}), - ) - model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) - return LayerHiddenCollector(model=model) - - -def test_replay_restores_captured_hidden_state_into_request_collector(monkeypatch): - graph_collector = _create_layer_collector(monkeypatch) - graph_collector.add(layer_index=0, hidden=torch.randn(2, 3)) - request_collector = graph_collector.new_instance() - graph_infer_state = _GraphInferState(hidden_collector=graph_collector) - request_infer_state = SimpleNamespace( - input_ids=torch.empty(4, dtype=torch.int64), - hidden_collector=request_collector, - ) - graph_output = torch.randn(2, 3) - prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} - - outputs = prefill_graph._replay(input_tensors=[], infer_state=request_infer_state) - - assert len(outputs) == 1 - assert outputs[0] is graph_output - assert request_infer_state.hidden_collector is request_collector - assert request_collector.layer_hiddens is not graph_collector.layer_hiddens - assert request_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] - assert graph_infer_state.copied_from is request_infer_state - assert graph_infer_state.replayed_with is request_infer_state - - -def test_first_replay_replaces_capture_collector_with_runtime_instance(monkeypatch): - graph_collector = _create_layer_collector(monkeypatch) - graph_collector.add(layer_index=0, hidden=torch.randn(2, 3)) - graph_infer_state = _GraphInferState(hidden_collector=graph_collector) - graph_output = torch.randn(2, 3) - prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} - - prefill_graph._replay(input_tensors=[], infer_state=graph_infer_state) - - assert graph_infer_state.hidden_collector is not graph_collector - assert graph_infer_state.hidden_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py deleted file mode 100644 index df357539e1..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ /dev/null @@ -1,295 +0,0 @@ -import weakref -import collections -import inspect -import json -import math -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( - int4kv_flash_decoding_stage1 as stage1_module, - ppl_int4kv_flash_decoding as decode_module, -) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -def make_inputs(device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2): - # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 - low = torch.randint(0, 15, (token_count, 4, head_dim // 2), device=device) - high = torch.randint(0, 15, low.shape, device=device) - kv = (low | (high << 4)).to(torch.int8) - scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.1 + 0.02).to(dtype) - return dict( - q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), - k=kv[:, :2], - k_scale=scale[:, :2], - v=kv[:, 2:], - v_scale=scale[:, 2:], - Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), - B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), - B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), - max_kv_seq_len=3, - mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), - mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=dtype, device=device), - block_seq=256, - ) - - -def rebuild(inputs): - args, kwargs = stage1_module.rebuild_inputs(**inputs) - return inspect.signature(stage1_module.int4kv_flash_decode_stage1.fn).bind(*args, **kwargs).arguments - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) -def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_lookup_uses_actual_length(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: - assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected - - -@pytest.mark.parametrize("token_count", [3, 128]) -def test_rebuild_preserves_packed_cache_and_scales(monkeypatch, token_count): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") - inputs = make_inputs(token_count=token_count) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - rebuilt = rebuild(inputs) - torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) - assert rebuilt["B_req_idx"].tolist() == [0, 1] - assert rebuilt["B_Seqlen"].tolist() == [17, 17] - assert rebuilt["max_kv_seq_len"] == 17 - assert rebuilt["block_seq"] == inputs["block_seq"] - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - for name, value in snapshots.items(): - torch.testing.assert_close(inputs[name], value) - - -@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) -def test_rebuild_rejects_inconsistent_token_capacity(name): - inputs = make_inputs() - inputs[name] = inputs[name][:-1] - with pytest.raises(AssertionError, match="same number of tokens"): - rebuild(inputs) - - -def test_rebuild_rejects_empty_cache(): - with pytest.raises(ValueError, match="non-empty KV cache"): - rebuild(make_inputs(token_count=0)) - - -def test_static_key_separates_quant_groups_without_head_or_block_counts(): - kernel = stage1_module.int4kv_flash_decode_stage1 - inputs = make_inputs() - key = kernel._static_key(**inputs) - inputs["q"] = inputs["q"][:, :4] - for name in ["k", "k_scale", "v", "v_scale"]: - inputs[name] = inputs[name][:, :1] - inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) - assert kernel._static_key(**inputs) == key - inputs["k_scale"] = inputs["k_scale"][:, :, :2] - assert kernel._static_key(**inputs) != key - - -def test_decode_passes_saved_length_to_stage1(monkeypatch): - inputs = make_inputs() - infer_state = InferStateInfo() - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 8193 - infer_state.b_req_idx = inputs["B_req_idx"] - infer_state.b_seq_len = inputs["B_Seqlen"] - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int4kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 - calls = [] - monkeypatch.setattr(stage1_module, "int4kv_flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) - stage2 = inspect.getmodule(flash_decode_stage2) - monkeypatch.setattr(stage2, "flash_decode_stage2", lambda *args: None) - state.decode_att( - inputs["q"], - (inputs["k"], inputs["k_scale"]), - (inputs["v"], inputs["v_scale"]), - alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), - ) - assert calls[0]["max_kv_seq_len"] == 8193 - assert calls[0]["block_seq"] == 256 - - -def reference_attention(inputs): - def dequant(packed, scale): - packed = packed.to(torch.uint8) - unpacked = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2).float() - 7 - group_size = unpacked.shape[-1] // scale.shape[-1] - # 与算子一致，以 scale 的 dtype 产生解量化 K/V，再用 FP32 SDPA 验证 attention 本身。 - return (unpacked * scale.float().repeat_interleave(group_size, dim=-1)).to(scale.dtype).float() - - k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) - outputs = [] - for row, length in enumerate(inputs["B_Seqlen"].tolist()): - indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() - with sdpa_kernel(SDPBackend.MATH): - outputs.append( - F.scaled_dot_product_attention( - inputs["q"][row].float().unsqueeze(1), - k[indices].transpose(0, 1), - v[indices].transpose(0, 1), - enable_gqa=True, - ).squeeze(1) - ) - return torch.stack(outputs) - - -def reduce_output(inputs): - output = torch.empty_like(inputs["q"]) - flash_decode_stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"]) - return output - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) -@pytest.mark.parametrize( - "head_dim,group_size,dtype", [(64, 8, torch.bfloat16), (128, 32, torch.float16), (256, 8, torch.bfloat16)] -) -def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len)) - inputs["B_Seqlen"][1] -= 1 - reference = reference_attention(inputs) - # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 - for block_n in [16, 32, 64, 128]: - stage1_module.int4kv_flash_decode_stage1.fn( - **inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2} - ) - output = reduce_output(inputs) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8193, 16384]) -def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - kernel = stage1_module.int4kv_flash_decode_stage1 - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - assert kernel.mutates_args == [] - inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) - inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - reference = reference_attention(inputs) - rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench - rebuild_count, timings = 0, [] - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - rebuild_count += 1 - return rebuild_func(*args, **kwargs) - - def checked_bench(*args, **kwargs): - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] - assert rebuilt["Req_to_tokens"].shape == (2, kv_len) - assert rebuilt["max_kv_seq_len"] == kv_len - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - elapsed = benchmark(*args, **kwargs) - timings.append(elapsed) - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) - run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) - assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) - assert len(list(tmp_path.glob("*.json"))) == 1 - assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) - torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) - - # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 - - output = torch.empty_like(inputs["q"]) - - def decode(): - kernel(**inputs) - flash_decode_stage2( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] - ) - - decode() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - decode() - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 - for name, snapshot in snapshots.items(): - if name not in ["mid_out", "mid_out_logsumexp"]: - torch.testing.assert_close(inputs[name], snapshot) - - # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 - inputs["B_Seqlen"].fill_(kv_len) - inputs["max_kv_seq_len"] = kv_len - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py deleted file mode 100644 index f3cb8de463..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ /dev/null @@ -1,188 +0,0 @@ -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import ( - flash_decode_stage1, -) - - -def create_tensors( - batch_size=4, - num_heads=4, - kv_head_num=1, - seq_len=256, - max_len_in_batch=8192, - max_batch_group_size=4, - kv_len=None, - req_to_tokens_len=None, -): - head_dim = 128 - block_seq = 256 - quant_group_size = 8 - - test_dtype = torch.bfloat16 - - kv_len = max_len_in_batch if kv_len is None else kv_len - req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len - - kv_shape = (batch_size * kv_len, kv_head_num, head_dim) - kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size) - - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") - k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view( - batch_size, req_to_tokens_len - ) - B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda") - mid_out = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda" - ) - mid_out_logsumexp = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda" - ) - - return { - "q": q, - "k": k, - "k_scale": k_scale, - "v": v, - "v_scale": v_scale, - "Req_to_tokens": Req_to_tokens, - "B_req_idx": B_req_idx, - "b_shared_seq_len": b_shared_seq_len, - "b_mark_shared_group": b_mark_shared_group, - "max_len_in_batch": max_len_in_batch, - "mid_out": mid_out, - "mid_out_logsumexp": mid_out_logsumexp, - "block_seq": block_seq, - "max_batch_group_size": max_batch_group_size, - } - - -@pytest.fixture -def setup_tensors(): - return create_tensors() - - -def test_flash_decode_stage1_execution(setup_tensors): - flash_decode_stage1( - q=setup_tensors["q"], - k=setup_tensors["k"], - k_scale=setup_tensors["k_scale"], - v=setup_tensors["v"], - v_scale=setup_tensors["v_scale"], - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - b_shared_seq_len=setup_tensors["b_shared_seq_len"], - b_mark_shared_group=setup_tensors["b_mark_shared_group"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=setup_tensors["mid_out"], - mid_out_logsumexp=setup_tensors["mid_out_logsumexp"], - block_seq=setup_tensors["block_seq"], - max_batch_group_size=setup_tensors["max_batch_group_size"], - ) - - q = setup_tensors["q"] - k = setup_tensors["k"] - v = setup_tensors["v"] - true_mid_out = torch.zeros_like(setup_tensors["mid_out"]) - true_mid_out_logsumexp = torch.zeros_like(setup_tensors["mid_out_logsumexp"]) - new_q = q - new_k = k.to(q.dtype) - new_v = v.to(q.dtype) - - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( - flash_decode_stage1 as gqa_flash_decode_stage1, - ) - - gqa_flash_decode_stage1( - q=new_q, - k=new_k, - v=new_v, - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=setup_tensors["b_shared_seq_len"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=true_mid_out, - mid_out_logsumexp=true_mid_out_logsumexp, - block_seq=setup_tensors["block_seq"], - ) - print(setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0]) - assert torch.allclose( - setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2 - ), "Mid output does not match expected values" - assert torch.allclose( - setup_tensors["mid_out_logsumexp"], true_mid_out_logsumexp, atol=1e-2 - ), "LogSumExp output does not match expected values" - - -def autotune_and_benchmark(): - import triton - - batch_sizes = [8, 16, 32, 64] - seq_lens = [1024, 2048, 4096] - - results = [] - for batch in batch_sizes: - for seq in seq_lens: - # Clear GPU cache to reduce CUDA Graph capture failures. - torch.cuda.empty_cache() - - setup_tensors = create_tensors( - batch_size=batch, - num_heads=32, - kv_head_num=8, - seq_len=seq, - max_len_in_batch=8192, - max_batch_group_size=8, - kv_len=seq, - req_to_tokens_len=seq, - ) - - def fn_triton(st=setup_tensors): - return flash_decode_stage1( - q=st["q"], - k=st["k"], - k_scale=st["k_scale"], - v=st["v"], - v_scale=st["v_scale"], - Req_to_tokens=st["Req_to_tokens"], - B_req_idx=st["B_req_idx"], - b_shared_seq_len=st["b_shared_seq_len"], - b_mark_shared_group=st["b_mark_shared_group"], - max_len_in_batch=st["max_len_in_batch"], - mid_out=st["mid_out"], - mid_out_logsumexp=st["mid_out_logsumexp"], - block_seq=st["block_seq"], - max_batch_group_size=st["max_batch_group_size"], - ) - - ms_triton = triton.testing.do_bench_cudagraph(fn_triton, rep=100) - - results.append( - { - "batch_size": batch, - "seq_len": seq, - "triton_ms": ms_triton, - } - ) - print(results[-1]) - - del setup_tensors - - print(f"\n{'='*80}") - print("SUMMARY - Performance Comparison") - print(f"{'='*80}") - print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}") - print(f"{'-'*80}") - for r in results: - print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}") - print(f"{'='*80}") - - -if __name__ == "__main__": - autotune_and_benchmark() diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py deleted file mode 100644 index 985a33289f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ /dev/null @@ -1,192 +0,0 @@ -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, -) - - -def create_tensors( - shared_seq_len, - batch_size=4, - seq_len=256, - max_len_in_batch=8192, - max_batch_group_size=4, - kv_len=None, - req_to_tokens_len=None, -): - num_heads = 32 - kv_head_num = 8 - head_dim = 128 - block_seq = 256 - quant_group_size = 8 - - test_dtype = torch.bfloat16 - - kv_len = max_len_in_batch if kv_len is None else kv_len - req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len - - kv_shape = (batch_size * kv_len, kv_head_num, head_dim) - kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size) - - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") - k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view( - batch_size, req_to_tokens_len - ) - B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda") - b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda") - mid_out = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda" - ) - mid_out_logsumexp = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda" - ) - - return { - "q": q, - "k": k, - "k_scale": k_scale, - "v": v, - "v_scale": v_scale, - "Req_to_tokens": Req_to_tokens, - "B_req_idx": B_req_idx, - "b_seq_len": b_seq_len, - "b_shared_seq_len": b_shared_seq_len, - "b_mark_shared_group": b_mark_shared_group, - "max_len_in_batch": max_len_in_batch, - "mid_out": mid_out, - "mid_out_logsumexp": mid_out_logsumexp, - "block_seq": block_seq, - "max_batch_group_size": max_batch_group_size, - "head_dim": head_dim, - } - - -@pytest.mark.parametrize("shared_seq_len", [0, 47, 77, 128, 200, 255]) -def test_flash_decode_stage2_execution(shared_seq_len): - setup_tensors = create_tensors(shared_seq_len) - - flash_decode_stage2( - q=setup_tensors["q"], - k=setup_tensors["k"], - k_scale=setup_tensors["k_scale"], - v=setup_tensors["v"], - v_scale=setup_tensors["v_scale"], - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=setup_tensors["b_seq_len"], - b_shared_seq_len=setup_tensors["b_shared_seq_len"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=setup_tensors["mid_out"], - mid_out_logsumexp=setup_tensors["mid_out_logsumexp"], - block_seq=setup_tensors["block_seq"], - ) - seq_block_idx = (setup_tensors["b_shared_seq_len"][0].item() + setup_tensors["block_seq"] - 1) // setup_tensors[ - "block_seq" - ] - mid_out = setup_tensors["mid_out"][:, :, seq_block_idx:, :] - mid_out_logsumexp = setup_tensors["mid_out_logsumexp"][:, :, seq_block_idx:] - - q = setup_tensors["q"] - k = setup_tensors["k"] - v = setup_tensors["v"] - true_mid_out = torch.zeros_like(mid_out) - true_mid_out_logsumexp = torch.zeros_like(mid_out_logsumexp) - new_q = q - new_k = k.to(q.dtype) - new_v = v.to(q.dtype) - - b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"] - req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :] - - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( - flash_decode_stage1 as gqa_flash_decode_stage1, - ) - - gqa_flash_decode_stage1( - q=new_q, - k=new_k, - v=new_v, - Req_to_tokens=req_to_tokens, - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=b_seq_len, - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=true_mid_out, - mid_out_logsumexp=true_mid_out_logsumexp, - block_seq=setup_tensors["block_seq"], - ) - print(f"\nshared_seq_len={shared_seq_len}") - print(f"mid_out: {mid_out[0:4, 0, 0, 0]}") - print(f"true_mid_out: {true_mid_out[0:4, 0, 0, 0]}") - abs_diff = (mid_out - true_mid_out).abs() - max_diff = abs_diff.max() - max_diff_idx = abs_diff.argmax() - max_diff_idx_unraveled = torch.unravel_index(max_diff_idx, abs_diff.shape) - mid_out_value = mid_out[max_diff_idx_unraveled] - true_mid_out_value = true_mid_out[max_diff_idx_unraveled] - print(f"max abs diff: {max_diff}, mid_out value: {mid_out_value}, " f"true_mid_out value: {true_mid_out_value}") - - assert torch.allclose( - mid_out[0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2 - ), f"Mid output does not match expected values for shared_seq_len={shared_seq_len}" - assert torch.allclose( - mid_out_logsumexp, true_mid_out_logsumexp, atol=1e-2 - ), f"LogSumExp output does not match expected values for shared_seq_len={shared_seq_len}" - - -if __name__ == "__main__": - # 可选：对 Triton diverse stage2 做 cudagraph bench（仅本仓库内核，无外部 CUDA 扩展）。 - - import triton - - batch_sizes = [8, 16, 32, 64] - seq_lens = [32, 64, 128, 256] - - results = [] - for batch in batch_sizes: - for seq in seq_lens: - torch.cuda.empty_cache() - - setup_tensors = create_tensors( - shared_seq_len=0, - batch_size=batch, - seq_len=seq, - max_len_in_batch=8192, - kv_len=seq, - req_to_tokens_len=seq, - ) - st = setup_tensors - - def bench_stage2(): - flash_decode_stage2( - q=st["q"], - k=st["k"], - k_scale=st["k_scale"], - v=st["v"], - v_scale=st["v_scale"], - Req_to_tokens=st["Req_to_tokens"], - B_req_idx=st["B_req_idx"], - B_Seqlen=st["b_seq_len"], - b_shared_seq_len=st["b_shared_seq_len"], - max_len_in_batch=st["max_len_in_batch"], - mid_out=st["mid_out"], - mid_out_logsumexp=st["mid_out_logsumexp"], - block_seq=st["block_seq"], - ) - - ms = triton.testing.do_bench_cudagraph(bench_stage2, rep=100) - results.append({"batch_size": batch, "seq_len": seq, "flash_decode_stage2_ms": ms}) - print(results[-1]) - del setup_tensors - - print(f"\n{'='*80}") - print(f"{'batch_size':<10} {'seq_len':<10} {'flash_decode_stage2_ms':<22}") - print(f"{'-'*80}") - for r in results: - print(f"{r['batch_size']:<10} {r['seq_len']:<10} {r['flash_decode_stage2_ms']:<22.4f}") - print(f"{'='*80}") diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py deleted file mode 100644 index ecab784529..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py +++ /dev/null @@ -1,340 +0,0 @@ -"""Unit test for ``context_attention_fwd_neo``. - -Torch reference expresses the *semantics* of the attention, not the kernel's -internal block structure — it has no notion of BLOCK_N / BLOCK_M. For each -batch element we gather K/V for the whole request (prompt + new tokens) via -``req_to_token_indexs`` and apply:: - - allow[m, k] = (k <= q_pos[m]) OR (k < image_end[m]) - -i.e. normal queries are causal, and image-token queries can only see future -tokens from the same image span. If the Triton kernel disagrees with this -reference, the kernel is wrong. - -Run directly for quick debugging: - - python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_context_attention_fwd_neo.py - -or via pytest: - - pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_context_attention_fwd_neo.py -x -s -""" - -import math -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo - - -def torch_reference_context_attention_neo( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b_req_idx: torch.Tensor, - b_start_loc: torch.Tensor, - b_seq_len: torch.Tensor, - b_prompt_cache_len: torch.Tensor, - req_to_token_indexs: torch.Tensor, - b_image_token_end: torch.Tensor, -) -> torch.Tensor: - device = q.device - dtype = q.dtype - _, Hq, D = q.shape - Hk = k.shape[1] - kv_group = Hq // Hk - scale = 1.0 / math.sqrt(D) - - out = torch.empty_like(q) - - for b in range(b_seq_len.shape[0]): - req_idx = int(b_req_idx[b].item()) - total = int(b_seq_len[b].item()) - prompt = int(b_prompt_cache_len[b].item()) - new = total - prompt - if new <= 0: - continue - - q_start = int(b_start_loc[b].item()) - q_blk = q[q_start : q_start + new] # [M, Hq, D] - q_image_end = b_image_token_end[q_start : q_start + new].to(torch.int64) - - token_locs = req_to_token_indexs[req_idx, :total].to(torch.int64) - k_blk = k[token_locs] # [total, Hk, D] - v_blk = v[token_locs] - - q_pos = torch.arange(prompt, total, device=device, dtype=torch.int64) # [M] - k_pos = torch.arange(0, total, device=device, dtype=torch.int64) # [total] - causal = k_pos[None, :] <= q_pos[:, None] - allow = causal | (k_pos[None, :] < q_image_end[:, None]) - - out_blk = torch.empty_like(q_blk) - for h in range(Hq): - h_k = h // kv_group - q_h = q_blk[:, h, :].to(torch.float32) - k_h = k_blk[:, h_k, :].to(torch.float32) - v_h = v_blk[:, h_k, :].to(torch.float32) - - scores = (q_h @ k_h.transpose(0, 1)) * scale - scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8)) - probs = torch.softmax(scores, dim=-1) - out_h = (probs @ v_h).to(dtype) - out_blk[:, h, :] = out_h - - out[q_start : q_start + new] = out_blk - - return out - - -def _build_inputs( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - device: str, - max_new: int = 256, - max_prompt: int = 512, - image_prob: float = 0.7, - num_image_spans_max: int = 3, - image_span_len_max: int = 24, - kv_pool_slack: int = 4096, - seed: int = 0, -): - g = torch.Generator(device="cpu").manual_seed(seed) - - new_lens = torch.randint(low=1, high=max_new + 1, size=(batch,), generator=g) - prompt_lens = torch.randint(low=0, high=max_prompt + 1, size=(batch,), generator=g) - total_lens = new_lens + prompt_lens - - sum_new = int(new_lens.sum().item()) - sum_total = int(total_lens.sum().item()) - max_total_len = int(total_lens.max().item()) - max_new_len = int(new_lens.max().item()) - - b_start_loc = torch.zeros(batch, dtype=torch.int32) - cur = 0 - for i in range(batch): - b_start_loc[i] = cur - cur += int(new_lens[i].item()) - - # Permute so batch idx != request idx: exercises the Req_to_tokens indexing. - b_req_idx = torch.randperm(batch, generator=g).to(torch.int32) - - # Global KV pool with scattered, non-contiguous slot assignment per request. - base = 1024 - kv_pool_size = base + sum_total + kv_pool_slack - pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base - - req_to_token_indexs = torch.zeros((batch, max_total_len), dtype=torch.int32) - p = 0 - for r_logical, req_id in enumerate(b_req_idx.tolist()): - L = int(total_lens[r_logical].item()) - req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32) - p += L - - b_image_token_end = torch.zeros(sum_new, dtype=torch.int32) - for i in range(batch): - M = int(new_lens[i].item()) - P = int(prompt_lens[i].item()) - start_pack = int(b_start_loc[i].item()) - if M < 2: - continue - if torch.rand((), generator=g).item() > image_prob: - continue - n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item()) - cursor = 0 - for _ in range(n_spans): - remaining = M - cursor - if remaining <= 0: - break - gap = int(torch.randint(0, remaining, (1,), generator=g).item()) - s_rel = cursor + gap - max_span_len = min(image_span_len_max, M - s_rel) - if max_span_len <= 0: - break - span_len = int(torch.randint(1, max_span_len + 1, (1,), generator=g).item()) - e_rel = s_rel + span_len - image_end = P + e_rel - b_image_token_end[start_pack + s_rel : start_pack + e_rel] = image_end - cursor = e_rel - - b_seq_len = total_lens.to(torch.int32) - b_prompt_cache_len = prompt_lens.to(torch.int32) - - q = torch.randn((sum_new, Hq, D), dtype=dtype, device=device) - k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - o = torch.empty_like(q) - - return dict( - q=q, - k=k, - v=v, - o=o, - b_req_idx=b_req_idx.to(device), - b_start_loc=b_start_loc.to(device), - b_seq_len=b_seq_len.to(device), - b_prompt_cache_len=b_prompt_cache_len.to(device), - max_new_len=max_new_len, - req_to_token_indexs=req_to_token_indexs.to(device), - b_image_token_end=b_image_token_end.to(device), - new_lens=new_lens, - prompt_lens=prompt_lens, - ) - - -def _report_per_batch_error(out_triton, out_ref, new_lens, b_start_loc, image_token_end, tag=""): - print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):") - for i in range(new_lens.shape[0]): - s = int(b_start_loc[i].item()) - m = int(new_lens[i].item()) - if m == 0: - continue - a = out_triton[s : s + m].float() - b = out_ref[s : s + m].float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - n_img = int((image_token_end[s : s + m] > 0).sum().item()) - print( - f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | " - f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}" - ) - - -def _run_case( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - seed: int, - max_new: int, - max_prompt: int, - atol: float = 5e-2, - rtol: float = 5e-2, - cos_threshold: float = 0.99, - verbose: bool = True, -): - assert Hq % Hk == 0 - device = "cuda" - - inputs = _build_inputs( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - device=device, - max_new=max_new, - max_prompt=max_prompt, - seed=seed, - ) - - context_attention_fwd_neo( - inputs["q"], - inputs["k"], - inputs["v"], - inputs["o"], - inputs["b_req_idx"], - inputs["b_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["max_new_len"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - out_triton = inputs["o"] - - out_ref = torch_reference_context_attention_neo( - inputs["q"], - inputs["k"], - inputs["v"], - inputs["b_req_idx"], - inputs["b_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - - a = out_triton.float() - b = out_ref.float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - - n_image = int((inputs["b_image_token_end"] > 0).sum().item()) - n_tokens = int(inputs["b_image_token_end"].numel()) - if verbose: - print( - f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} " - f"seed={seed} image_tokens={n_image}/{n_tokens}" - ) - print( - f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} " - f"(allclose atol={atol}, rtol={rtol}? " - f"{torch.allclose(a, b, atol=atol, rtol=rtol)})" - ) - _report_per_batch_error( - out_triton, - out_ref, - inputs["new_lens"], - inputs["b_start_loc"], - inputs["b_image_token_end"], - tag=f"seed={seed}", - ) - - return abs_err, rel_err, cos - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") -@pytest.mark.parametrize( - "batch,Hq,Hk,D,dtype,seed,max_new,max_prompt", - [ - (4, 8, 2, 128, torch.bfloat16, 0, 128, 256), - (4, 8, 2, 128, torch.bfloat16, 1, 256, 512), - (8, 16, 4, 128, torch.bfloat16, 2, 256, 512), - (16, 28, 4, 128, torch.bfloat16, 3, 128, 256), - (4, 8, 2, 128, torch.float16, 4, 256, 512), - (4, 8, 8, 64, torch.bfloat16, 5, 128, 256), - (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024), - ], -) -def test_context_attention_fwd_neo(batch, Hq, Hk, D, dtype, seed, max_new, max_prompt): - abs_err, rel_err, cos = _run_case( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - seed=seed, - max_new=max_new, - max_prompt=max_prompt, - verbose=True, - ) - assert cos > 0.99, f"cosine similarity too low: {cos}" - assert rel_err < 5e-2, f"max relative error too large: {rel_err}" - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("No CUDA available.") - raise SystemExit(0) - - torch.manual_seed(0) - - cases = [ - dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_new=128, max_prompt=256), - dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_new=256, max_prompt=512), - dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_new=128, max_prompt=256), - dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_new=256, max_prompt=512), - ] - - for cfg in cases: - _run_case(**cfg, verbose=True) diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py deleted file mode 100644 index ed0c6e369f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py +++ /dev/null @@ -1,42 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv -from lightllm.utils.log_utils import init_logger -import torch.nn.functional as F - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, seqlen, heads, nope_head, rope_head, copy_len", - [ - (a, b, c, d, e, f) - for a in [1, 16, 32, 128, 512] - for b in [1024, 2048] - for c in [1] - for d in [512] - for e in [64] - for f in [10, 20, 100, 1024] - ], -) -def test_destindex_copy_kv(batch, seqlen, heads, nope_head, rope_head, copy_len): - B, N_CTX, H, NOPE_HEAD, ROPE_HEAD, COPY_LEN = batch, seqlen, heads, nope_head, rope_head, copy_len - dtype = torch.bfloat16 - dest_loc = torch.randperm(COPY_LEN).cuda() - kv = torch.randn((len(dest_loc), H, NOPE_HEAD + ROPE_HEAD), dtype=dtype).cuda() - O_nope = torch.zeros((B * N_CTX, H, NOPE_HEAD), dtype=dtype).cuda() - O_rope = torch.zeros((B * N_CTX, H, ROPE_HEAD), dtype=dtype).cuda() - - kv_nope = kv[:, :, :NOPE_HEAD] - kv_rope = kv[:, :, NOPE_HEAD:] - destindex_copy_kv(kv_nope, kv_rope, dest_loc, O_nope, O_rope) - - assert torch.allclose(O_nope[dest_loc], kv_nope, atol=1e-2, rtol=0) - assert torch.allclose(O_rope[dest_loc], kv_rope, atol=1e-2, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py deleted file mode 100644 index 83537ec708..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py +++ /dev/null @@ -1,62 +0,0 @@ -import torch -import pytest -import numpy as np -from typing import Tuple -from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv, dequantize_int4kv -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_quanted_and_dequant(): - """Test quantization followed by dequantization.""" - batch_size = 1 - seq_len = 8 - head_num = 4 - k_head_num = 2 - v_head_num = 2 - assert k_head_num + v_head_num == head_num - head_dim = 64 - quant_group_size = 8 - - # Create original data - original_kv = torch.randn(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).clamp_(-1, 1).cuda() - dest_loc = torch.arange(batch_size * seq_len, dtype=torch.int64).cuda() - - # Quantize - group_count = head_dim // quant_group_size - kv_buffer = torch.zeros(batch_size * seq_len, head_num, head_dim // 2, dtype=torch.int8).cuda() - kv_scale_buffer = torch.zeros(batch_size * seq_len, head_num, group_count, dtype=torch.float32).cuda() - destindex_copy_int4kv(original_kv, dest_loc, kv_buffer, kv_scale_buffer, quant_group_size) - - # Dequantize - req_to_token_indexs = torch.arange(seq_len, dtype=torch.int64).view(1, -1).cuda() - b_seq_len = torch.tensor([seq_len], dtype=torch.int32).cuda() - b_req_idx = torch.tensor([0], dtype=torch.int32).cuda() - b_kv_start_loc = torch.tensor([0], dtype=torch.int32).cuda() - - recovered_kv = torch.zeros(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).cuda() - - dequantize_int4kv( - k=kv_buffer[:, 0:k_head_num, :], - k_scale=kv_scale_buffer[:, 0:k_head_num, :], - v=kv_buffer[:, k_head_num:, :], - v_scale=kv_scale_buffer[:, k_head_num:, :], - req_to_token_indexs=req_to_token_indexs, - b_seq_len=b_seq_len, - b_req_idx=b_req_idx, - b_kv_start_loc=b_kv_start_loc, - k_out=recovered_kv[:, :k_head_num, :], - v_out=recovered_kv[:, k_head_num:, :], - max_len_in_batch=seq_len, - quant_group_size=quant_group_size, - ) - - logger.info("Round-trip test completed!") - assert torch.allclose(recovered_kv, original_kv, atol=2 / 14, rtol=0) - cos = torch.nn.CosineSimilarity(0) - assert cos(recovered_kv.flatten().float(), original_kv.flatten().float()) > 0.99 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py deleted file mode 100644 index 149c9894ae..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( - dequantize_int8kv, - destindex_copy_quantize_kv, -) -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def torch_dequant(kv, kv_scale, b_req_idx, b_seq_len, req_to_token_indexs, odtype, group_quant_size): - batch = b_req_idx.shape[0] - tmp_out = [] - for i in range(batch): - req_idx = b_req_idx[i] - seq_len = b_seq_len[i] - kv_loc = req_to_token_indexs[req_idx, :seq_len] - head_num = kv.shape[1] - cur_kv = kv[kv_loc, :, :].reshape(seq_len, head_num, -1, group_quant_size).to(odtype) - cur_scale = kv_scale[kv_loc, :, :].reshape(seq_len, head_num, -1, 1) - out = cur_kv * cur_scale - tmp_out.append(out.reshape(seq_len, head_num, -1)) - return torch.cat(tmp_out, dim=0) - - -@pytest.mark.parametrize( - "B, H, N_CTX, D_HEAD, group_quant_size", - [ - (b, H, N_CTX, D_HEAD, group_quant_size) - for b in [1, 2, 4] - for H in [1, 8] - for N_CTX in [3, 10, 1024] - for D_HEAD in [64, 128] - for group_quant_size in [8, 16] - ], -) -def test_dequantize_int8kv(B, H, N_CTX, D_HEAD, group_quant_size): - dtype = torch.bfloat16 - kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=torch.int8, device="cuda").random_(-10, 10) - kv_scale = torch.randn((B * N_CTX, 2 * H, D_HEAD // group_quant_size), dtype=dtype, device="cuda") - out = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda") - req_to_token_indexs = torch.empty((B, N_CTX), dtype=torch.int32, device="cuda") - max_input_len = N_CTX - b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") - b_seq_len.fill_(N_CTX) - b_req_idx = torch.arange(0, B, dtype=torch.int32, device="cuda") - req_to_token_indexs.view(-1)[:] = torch.arange(0, B * N_CTX, dtype=torch.int32, device="cuda") - b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len - - k = kv[:, :H, :] - v = kv[:, H:, :] - k_scale = kv_scale[:, :H, :] - v_scale = kv_scale[:, H:, :] - - ground_out = torch_dequant( - kv=kv, - kv_scale=kv_scale, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - req_to_token_indexs=req_to_token_indexs, - odtype=out.dtype, - group_quant_size=group_quant_size, - ) - dequantize_int8kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=req_to_token_indexs, - b_seq_len=b_seq_len, - b_req_idx=b_req_idx, - b_kv_start_loc=b_kv_start_loc, - k_out=out[:, :H, :], - v_out=out[:, H:, :], - max_len_in_batch=max_input_len, - quant_group_size=group_quant_size, - ) - assert torch.allclose(out, ground_out, atol=1e-2, rtol=0) - cos = torch.nn.CosineSimilarity(0) - assert cos(out.flatten().float(), ground_out.flatten().float()) > 0.99 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index f7c742db01..1c984a06a4 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -1,971 +1,107 @@ -from typing import Optional - import pytest import torch +import torch.nn.functional as F from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update -def causal_conv1d_ref( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - mtp_step: int, - bias: Optional[torch.Tensor] = None, - activation: Optional[str] = None, - conv_state_indices: Optional[torch.Tensor] = None, - num_accepted_tokens: Optional[torch.Tensor] = None, - query_start_loc: Optional[torch.Tensor] = None, - pad_slot_id: int = -1, -) -> torch.Tensor: - """ - Pure-PyTorch reference for causal_conv1d_update. - - The kernel's algorithm per token iteration: - acc = bias - for j in range(width): - if j == 0: acc += col0 * weight[:, 0] - elif j < width-1: acc += col{j} * weight[:, j] - else: acc += x[t] * weight[:, width-1] - shift: col0=col1, col1=col2, ..., col_{w-3}=col_{w-2}, col_{w-2}=x[t] - """ - assert conv_state_indices is not None - batch = conv_state_indices.size(0) - dim = x.size(1) - assert x.size(0) % batch == 0 - seqlen = x.size(0) // batch - _, width = weight.shape - _, _, state_len = conv_state.size() - assert state_len == width - 1 + mtp_step +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - out = x.clone() - conv_state = conv_state.clone() - for b in range(batch): - slot = conv_state_indices[b].item() - if slot == pad_slot_id: +def reference(x, state, weight, bias, activation, slots, accepted, starts): + output, next_state = x.clone(), state.clone() + width = weight.shape[1] + for i, slot in enumerate(slots.tolist()): + start, end = starts[i : i + 2].tolist() + if slot == -1 or start == end: continue - - if query_start_loc is not None: - start = query_start_loc[b].item() - end = query_start_loc[b + 1].item() - local_seqlen = end - start - else: - start = b * seqlen - local_seqlen = seqlen - - accepted = 1 - if num_accepted_tokens is not None: - accepted = num_accepted_tokens[b].item() - offset = accepted - 1 - - state_3d = conv_state[slot] - - # STEP 1: initial columns from conv_state starting at offset - cols = [state_3d[:, offset + k].clone() for k in range(width - 1)] - - # STEP 2: update conv_state in the reference too - new_state = [] - for k in range(width - 2): - new_state.append(state_3d[:, offset + k + 1].clone()) - x_chunk = x[start : start + local_seqlen, :] - for t_ in range(local_seqlen): - new_state.append(x_chunk[t_, :].clone()) - while len(new_state) < state_len: - new_state.append(torch.zeros(dim, device=x.device, dtype=x.dtype)) - stacked = torch.stack(new_state, dim=1) - write_len = min(state_len, stacked.size(1)) - conv_state[slot, :, :write_len] = stacked[:, :write_len] - - # STEP 3-5: compute output - for t_ in range(local_seqlen): - acc = torch.zeros(dim, device=x.device, dtype=torch.float32) - if bias is not None: - acc += bias.float() - - for k in range(width): - if k == 0: - v = cols[0].float() - elif k == width - 1: - v = x_chunk[t_, :].float() - else: - v = cols[k].float() - w = weight[:, k].float() - acc += v * w - - if activation in ("silu", "swish"): - acc = acc / (1 + torch.exp(-acc)) - - out[start + t_, :] = acc.to(x.dtype) - - # shift cols: col0=col1, col1=col2, ..., col_{w-2}=x[t] - for k in range(width - 2): - cols[k] = cols[k + 1].clone() - if width >= 2: - cols[width - 2] = x_chunk[t_, :].clone() - - return out - - -def make_tensors( - batch: int, - dim: int, - width: int, - mtp_step: int, - has_bias: bool, - num_slots: int = None, - dtype: torch.dtype = torch.float16, - seed: int = 42, - device: str = "cuda", -): - torch.manual_seed(seed) - if num_slots is None: - num_slots = batch - - state_len = width - 1 + mtp_step - x = torch.randn(batch * (mtp_step + 1), dim, device=device, dtype=dtype) - weight = torch.randn(dim, width, device=device, dtype=dtype) - bias = torch.randn(dim, device=device, dtype=dtype) if has_bias else None - conv_state = torch.randn(num_slots, dim, state_len, device=device, dtype=dtype) - conv_state.add_(0.5) - - conv_state_indices = torch.arange(batch, device=device, dtype=torch.int32) % num_slots - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - query_start_loc = torch.arange(batch + 1, device=device, dtype=torch.int32) * (mtp_step + 1) - - return x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -@pytest.mark.parametrize("has_bias", [True, False]) -@pytest.mark.parametrize("activation", [None, "silu"]) -def test_single_request(width, dim, mtp_step, has_bias, activation): - """Single request, no pad slots, basic functionality.""" - batch = 1 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=has_bias, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch: width={width}, dim={dim}, mtp_step={mtp_step}, " - f"bias={has_bias}, activation={activation}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("batch", [2, 4]) -@pytest.mark.parametrize("width", [3, 4, 5]) -@pytest.mark.parametrize("dim", [64, 256]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_multi_request(batch, width, dim, mtp_step): - """Multiple requests.""" - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch: batch={batch}, width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [3, 4]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1]) -def test_pad_slots(width, dim, mtp_step): - """Some slots are padded, should produce same output as reference.""" - batch = 4 - num_slots = 4 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - num_slots=num_slots, - dtype=torch.float16, - ) - - conv_state_indices[2:] = -1 - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - pad_slot_id=-1, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - pad_slot_id=-1, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch with pad slots: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - for b in [2, 3]: - if conv_state_indices[b] == -1: - start = query_start_loc[b].item() - end = query_start_loc[b + 1].item() - for t_ in range(start, end): - assert torch.allclose( - x_orig[t_], out_triton[t_], rtol=1e-6, atol=1e-6 - ), f"Padded entry should not be modified at batch={b}, token={t_ - start}" - - -@pytest.mark.parametrize("width", [2, 3, 4, 5]) -@pytest.mark.parametrize("dim", [64, 256]) -@pytest.mark.parametrize("mtp_step", [1, 2]) -def test_num_accepted_tokens_varied(width, dim, mtp_step): - """Varying num_accepted_tokens across batch.""" - batch = 3 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - - num_accepted_tokens[0] = 1 - num_accepted_tokens[1] = mtp_step + 1 - num_accepted_tokens[2] = max(1, mtp_step // 2 + 1) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch with varied accepted tokens: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1]) -@pytest.mark.parametrize("activation", [None, "silu"]) -def test_conv_state_update_correctness(width, dim, mtp_step, activation): - """Verify that conv_state is updated correctly after the forward pass.""" - batch = 1 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_before = conv_state.clone() - - causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - slot = conv_state_indices[0].item() - - expected = torch.zeros_like(conv_state_before[slot]) - offset = num_accepted_tokens[0].item() - 1 - - for k in range(width - 2): - src = offset + k + 1 - if src < state_len: - expected[:, k] = conv_state_before[slot, :, src] - - for t_ in range(seqlen): - idx = (width - 2) + t_ - if idx < state_len: - expected[:, idx] = x_orig[t_, :] - - actual = conv_state[slot] - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(actual, expected, rtol=rtol, atol=atol), ( - f"Conv state mismatch: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(actual - expected).max().item():.6f}" - ) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_dtype_consistency(dtype): - """Test that the function runs correctly with different dtypes.""" - width = 4 - dim = 128 - mtp_step = 1 - - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=2, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=dtype, - ) - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = (1e-2, 1e-2) if dtype == torch.float16 else (2e-2, 2e-2) - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch for dtype {dtype}: " f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) + offset = accepted[i].item() - 1 + history = state[slot, :, offset : offset + width - 1] + window = torch.cat((history, x[start:end].T), dim=1) + values = ( + F.conv1d( + window.float().unsqueeze(0), + weight.float().unsqueeze(1), + bias=None if bias is None else bias.float(), + groups=x.shape[1], + ) + .squeeze(0) + .T + ) + output[start:end] = F.silu(values) if activation else values + saved = window[:, 1:] + next_state[slot].zero_() + next_state[slot, :, : saved.shape[1]] = saved + return output, next_state @pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("activation", [None, "silu"]) -@pytest.mark.parametrize("has_bias", [True, False]) -def test_known_values(width, activation, has_bias): - """Deterministic known values to verify numerically.""" - batch = 1 - dim = 4 - mtp_step = 1 - device = "cuda" - - torch.manual_seed(999) - x = torch.randn(batch * (mtp_step + 1), dim, device=device, dtype=torch.float16) - weight = torch.randn(dim, width, device=device, dtype=torch.float16) - bias = torch.randn(dim, device=device, dtype=torch.float16) if has_bias else None - state_len = width - 1 + mtp_step - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float16) - conv_state.add_(0.5) - - conv_state_indices = torch.zeros(batch, device=device, dtype=torch.int32) - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.arange(batch + 1, device=device, dtype=torch.int32) * (mtp_step + 1) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - out_triton = causal_conv1d_update( +@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) +@pytest.mark.parametrize( + "dtype,dim,has_bias,activation", + [ + (torch.float16, 64, False, None), + (torch.float16, 128, True, "silu"), + (torch.bfloat16, 65, False, "silu"), + (torch.bfloat16, 256, True, None), + ], +) +def test_output_and_state_match_convolution(width, mtp_step, dtype, dim, has_bias, activation): + torch.manual_seed(42) + lengths = [mtp_step + 1, max(1, mtp_step), 0, mtp_step + 1] + starts = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32, device="cuda") + slots = torch.tensor([2, 0, 3, -1], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, mtp_step + 1, 1, 1], dtype=torch.int32, device="cuda") + x = torch.randn(sum(lengths), dim, dtype=dtype, device="cuda") + state = torch.randn(5, dim, width - 1 + mtp_step, dtype=dtype, device="cuda") + weight = torch.randn(dim, width, dtype=dtype, device="cuda") + bias = torch.randn(dim, dtype=dtype, device="cuda") if has_bias else None + expected, next_state = reference(x, state, weight, bias, activation, slots, accepted, starts) + + output = causal_conv1d_update( x, - conv_state, + state, weight, mtp_step=mtp_step, bias=bias, activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Known value mismatch: width={width}, activation={activation}, bias={has_bias}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -# ============================================================================= -# Edge-case tests for kernel simplification verification -# ============================================================================= - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) -def test_single_step_kernel_vs_pytorch_conv(width, mtp_step): - """ - Verify the triton kernel against a direct torch.nn.functional.conv1d - reference (transposed to per-channel depthwise conv). This test ensures - the causal_conv1d kernel computes the mathematically correct causal conv. - """ - batch = 1 - dim = 32 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - - torch.manual_seed(123) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - conv_state_indices = torch.zeros(batch, device=device, dtype=torch.int32) - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ).float() - - # PyTorch causal conv1d reference: pad input with state on the left, - # then run conv1d and slice output - history = conv_state_ref[0, :, : width - 1] - padded_input = torch.cat( - [ - history.T.contiguous(), - x_orig, - ], - dim=0, - ) - - w_3d = weight.unsqueeze(1).float() - pytorch_out = ( - torch.nn.functional.conv1d( - padded_input.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - padding=0, - stride=1, - ) - .squeeze(0) - .T[:seqlen, :] + conv_state_indices=slots, + num_accepted_tokens=accepted, + query_start_loc=starts, ) - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, pytorch_out, rtol=rtol, atol=atol), ( - f"Conv1d mismatch against torch.nn.functional.conv1d: " - f"width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - pytorch_out).max().item():.6f}" - ) + assert output.data_ptr() == x.data_ptr() + tolerance = 2e-2 if dtype == torch.bfloat16 else 1e-2 + torch.testing.assert_close(output, expected, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(state, next_state, rtol=0, atol=0) + torch.testing.assert_close(output[starts[-2] :], expected[starts[-2] :], rtol=0, atol=0) @pytest.mark.parametrize("width", [2, 3, 4]) -@pytest.mark.parametrize("num_steps", [2, 3, 5]) -def test_multi_step_decode_sliding_window(width, num_steps): - """ - Simulate multiple consecutive decode steps. Each step produces one output - token and shifts a sliding window of conv_state. The outputs over multi-step - must match a single forward pass that sees the full sequence at once. - """ - batch = 1 - dim = 32 - mtp_step = 0 - state_len = width - 1 + mtp_step - device = "cuda" - +@pytest.mark.parametrize("mtp_step", [0, 1, 2, 3]) +def test_repeated_decode_preserves_accepted_history(width, mtp_step): torch.manual_seed(77) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - init_state = torch.zeros(1, dim, state_len, device=device, dtype=torch.float32) - - all_tokens = torch.randn(num_steps, dim, device=device, dtype=torch.float32) - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - nat = torch.ones(batch, device=device, dtype=torch.int32) - qsl1 = torch.tensor([0, 1], device=device, dtype=torch.int32) - - conv_state_step = init_state.clone().half() - step_outputs = [] - - for step in range(num_steps): - x_step = all_tokens[step : step + 1, :].clone().half() - step_outputs.append( - causal_conv1d_update( - x_step, - conv_state_step, - weight.half(), - mtp_step=0, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl1, - ).float() - ) - - step_outputs = torch.cat(step_outputs, dim=0) - - # Build full forward: preload state, then conv1d over full sequence - full_state = init_state.clone().float() - padded = torch.cat( - [ - full_state[0, :, : width - 1].T.contiguous(), - all_tokens, - ], - dim=0, - ) - - w_3d = weight.unsqueeze(1).float() - expected = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:num_steps, :] - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(step_outputs, expected, rtol=rtol, atol=atol), ( - f"Multi-step mismatch: width={width}, num_steps={num_steps}\n" - f"max diff={torch.abs(step_outputs - expected).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [2, 3, 4]) -@pytest.mark.parametrize("mtp_step", [1, 2, 3]) -def test_mtp_decode_multi_token_per_step(width, mtp_step): - """ - MTP decode: each step processes (mtp_step+1) tokens. After each step, - some tokens are accepted (num_accepted_tokens varies). The conv_state - sliding window must correctly preserve history across steps. - """ - batch = 1 - dim = 32 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - num_steps = 3 - - torch.manual_seed(1234) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - conv_state_ref = conv_state.clone() - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - # Phase 1: one-step MTP decode triton - x_full = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - out_triton = causal_conv1d_update( - x_full.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=torch.ones(batch, device=device, dtype=torch.int32), - query_start_loc=qsl, - ).float() - - # Reference: torch conv1d with state preload, then silu - history = conv_state_ref[0, :, : width - 1].float() - padded = torch.cat([history.T.contiguous(), x_full], dim=0) - w_3d = weight.unsqueeze(1).float() - ref_out = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:seqlen, :] - ) - ref_out = ref_out / (1 + torch.exp(-ref_out)) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, ref_out, rtol=rtol, atol=atol), ( - f"MTP decode mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - ) - - # Phase 2: multi-step MTP decode with varying acceptance - conv_state_step = conv_state.clone().half() - - for step in range(num_steps): - x_step = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - x_step_half = x_step.clone().half() - - accepted = torch.randint(1, seqlen + 1, (batch,), device=device, dtype=torch.int32) - nat_step = accepted.clone() - qsl_step = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - x_orig_step = x_step_half.clone() - conv_state_before_step = conv_state_step.clone().half() - - step_out = causal_conv1d_update( - x_step_half, - conv_state_step, - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=nat_step, - query_start_loc=qsl_step, - ).float() - - ref_step = causal_conv1d_ref( - x_orig_step, - conv_state_before_step, - weight.half(), + dim, length = 32, mtp_step + 1 + weight = torch.randn(dim, width, dtype=torch.float16, device="cuda") + bias = torch.randn(dim, dtype=torch.float16, device="cuda") + state = torch.randn(1, dim, width - 1 + mtp_step, dtype=torch.float16, device="cuda") + expected_state = state.clone() + slots = torch.tensor([0], dtype=torch.int32, device="cuda") + starts = torch.tensor([0, length], dtype=torch.int32, device="cuda") + for step in range(5): + accepted = torch.tensor([1 + step % length], dtype=torch.int32, device="cuda") + x = torch.randn(length, dim, dtype=torch.float16, device="cuda") + expected, expected_state = reference(x, expected_state, weight, bias, "silu", slots, accepted, starts) + output = causal_conv1d_update( + x, + state, + weight, mtp_step=mtp_step, - bias=bias.half(), + bias=bias, activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=nat_step, - query_start_loc=qsl_step, - ).float() - - assert torch.allclose(step_out, ref_step, rtol=rtol, atol=atol), ( - f"Multi-step MTP step={step} mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(step_out - ref_step).max().item():.6f}" + conv_state_indices=slots, + num_accepted_tokens=accepted, + query_start_loc=starts, ) - - -@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) -def test_kernel_width_2_correctness(mtp_step): - """ - KERNEL_WIDTH=2 is a degenerate case: restore_conv_state_len = 0, - meaning no history is preserved in the state update (the mask is always - False). Verify the kernel still computes correct outputs and the single - history token (col0) is properly read from the state offset. - """ - width = 2 - batch = 1 - dim = 64 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - - torch.manual_seed(555) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - # Test with multiple offsets (varying num_accepted_tokens) - for nac_val in range(1, min(seqlen + 1, mtp_step + 2)): - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - x_half = x.clone().half() - nat = torch.full((batch,), nac_val, device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x_half, - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl, - ).float() - - # Reference: causal conv with history from state at offset nac_val-1 - history = conv_state[0, :, nac_val - 1 : state_len].float() - history_pad = history[:, : width - 1] - padded = torch.cat([history_pad.T.contiguous(), x], dim=0) - w_3d = weight.unsqueeze(1).float() - ref_out = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:seqlen, :] - ) - - assert torch.allclose(out_triton, ref_out, rtol=1e-2, atol=1e-2), ( - f"Width=2 mismatch: mtp_step={mtp_step}, nac={nac_val}\n" - f"max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - ) - - -def test_kernel_width_2_no_state_history(): - """ - KERNEL_WIDTH=2 with num_accepted_tokens=1 (offset=0): the restore mask - is always False, so the state update is fully overwritten by x tokens - (no history preserved). Verify this yields correct output. - """ - batch, dim, width, mtp_step = 1, 16, 2, 2 - state_len = width - 1 + mtp_step - seqlen = mtp_step + 1 - device = "cuda" - - torch.manual_seed(42) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.zeros(1, dim, state_len, device=device, dtype=torch.float32) - conv_state[0, :, 0] = 99.0 # Place a known history value at offset 0 - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - nat = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl, - ).float() - - # Manual reference: for each token t, acc = bias + col0*w0 + x[t]*w1 - # col0 for token 0 is state[0] = 99.0. Then col0 shifts to x[t]. - ref_out = torch.zeros(seqlen, dim, device=device, dtype=torch.float32) - col0 = conv_state_ref[0, :, 0].clone() - for t in range(seqlen): - acc = bias.clone() - acc += col0 * weight[:, 0] - acc += x_orig[t] * weight[:, 1] - ref_out[t] = acc - col0 = x_orig[t].clone() - - assert torch.allclose( - out_triton, ref_out, rtol=1e-2, atol=1e-2 - ), f"Width=2 no-history: max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_output_overwrites_x_inplace(width, mtp_step): - """ - Verify the kernel indeed overwrites x in-place (out == x). Also - verify that the output values are different from the original input. - """ - batch = 1 - dim = 64 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - x_before = x.clone() - - out = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - assert x.data_ptr() == out.data_ptr(), "out must be the same tensor as x" - assert torch.allclose(x, out, rtol=0, atol=0), "x must equal out after call" - assert not torch.allclose( - x, x_before, rtol=1e-6, atol=1e-6 - ), "x must be overwritten (different from original input)" - - -@pytest.mark.parametrize("width", [3, 4, 5]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_bias_none_gives_same_shape(width, mtp_step): - """bias=None should produce the same output shape as bias specified.""" - batch = 1 - dim = 32 - x, conv_state, weight, _, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=False, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=None, - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=None, - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - assert out_triton.shape == out_ref.shape - assert out_triton.shape == x_orig.shape - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"No-bias mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--tb=short"]) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py deleted file mode 100644 index 8d92d6cc8d..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ /dev/null @@ -1,287 +0,0 @@ -import collections -import functools -import importlib -import inspect -import json -from itertools import accumulate - -import pytest -import torch -import triton - -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda -from lightllm.common.kernel_config import KernelConfigs -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel - - -KERNELS = [ - ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), - ("_chunk_kda_scaled_dot_kkt_sub_intra", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra"), - ("recompute_w_u_fwd", "recompute_w_u_fwd_kernel"), - ("chunk_gla_fwd_o_gk", "chunk_gla_fwd_kernel_o"), -] - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.CLOSE_AUTOTUNE) - - -def wrapper_inputs(kernel, tokens=65): - q = torch.empty(1, tokens, 2, 128, dtype=torch.bfloat16) - v = torch.empty(1, tokens, 2, 80, dtype=q.dtype) - g = torch.empty_like(q, dtype=torch.float32) - chunk_indices = torch.tensor([[0, 0], *[[1, i] for i in range(triton.cdiv(tokens - 3, 64))]]) - values = dict( - q=q, - k=q, - v=v, - g=g, - gk=g, - beta=torch.empty(1, tokens, 2), - Akk=torch.empty(1, tokens, 2, 64), - Aqk=torch.empty(1, tokens, 2, 64), - R=torch.empty(1, tokens, 2, 64, dtype=q.dtype), - h=torch.empty(1, len(chunk_indices), 2, 128, 80, dtype=q.dtype), - o=torch.empty_like(v), - scale=128 ** -0.5, - cu_seqlens=torch.tensor([0, 3, tokens]), - chunk_indices=chunk_indices, - ) - return {name: values[name] for name in inspect.signature(kernel.fn).parameters if name in values} - - -@pytest.mark.parametrize("wrapper_name,jit_name", KERNELS) -def test_lightllm_tuning_selects_persists_and_reuses_launch_config(monkeypatch, tmp_path, wrapper_name, jit_name): - kernel = getattr(kda, wrapper_name) - assert isinstance(kernel, Autotuner) - inputs = wrapper_inputs(kernel) - candidates = [kernel.configs_gen_func()[0], kernel.configs_gen_func()[-1]] - launches, benchmarks = [], [] - - class LaunchRecorder: - def __getitem__(self, grid): - def launch(**kwargs): - launches.append((grid, kwargs)) - - return launch - - def bench(*args, run_config, **kwargs): - benchmarks.append(run_config) - return 1.0 if run_config == candidates[-1] else 2.0 - - monkeypatch.setattr(kda, jit_name, LaunchRecorder()) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.kernel_config.get_current_device_name", lambda: "test-device") - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - monkeypatch.setattr(kernel, "configs_gen_func", lambda: candidates) - monkeypatch.setattr(kernel, "_bench", bench) - - with Autotuner.autotune_warmup(): - kernel(**inputs) - assert benchmarks == candidates - for name, value in candidates[-1].items(): - assert launches[-1][1][name] == value - - cache_file = tmp_path / KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) - assert json.loads(cache_file.read_text()) == {"65": candidates[-1]} - - # A new wrapper invocation must reload the same selected config from disk. - kernel.cached_configs.clear() - kernel.fast_match_configs.clear() - kernel.warmuped_configs_set.clear() - benchmarks.clear() - kernel(**inputs) - assert benchmarks == [] - for name, value in candidates[-1].items(): - assert launches[-1][1][name] == value - - # Packed B=1 still needs different run keys as the token count grows. - assert kernel._run_key(**inputs) == 65 - assert kernel._run_key(**wrapper_inputs(kernel, tokens=129)) == 129 - assert kernel._static_key(**inputs) == kernel._static_key(**wrapper_inputs(kernel, tokens=129)) - # Request boundaries are runtime data; all calls use the packed path. - single_request = dict(inputs, cu_seqlens=torch.tensor([0, 65])) - assert kernel._static_key(**inputs) == kernel._static_key(**single_request) - - -def recurrent_reference(q, k, v, g, beta, initial_state, sequences): - q, k, v, g, beta = [x.float().cpu() for x in (q, k, v, g, beta)] - q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() / q.shape[-1] ** 0.5 - k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() - output = torch.empty_like(v) - final = initial_state.float().cpu().clone() - for n, (batch, start, end) in enumerate(sequences): - state = final[n] - for t in range(start, end): - kt = k[batch, t] - state = state * g[batch, t].exp()[..., None] - delta = beta[batch, t, :, None] * (v[batch, t] - torch.einsum("hk,hkv->hv", kt, state)) - state = state + kt[..., None] * delta[:, None, :] - output[batch, t] = torch.einsum("hk,hkv->hv", q[batch, t], state) - final[n] = state - return output, final - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("config_index", [None, 0, -1], ids=["default", "first-candidate", "last-candidate"]) -@pytest.mark.parametrize( - "seq_lens,key_dim,value_dim", - [((65, 65), 64, 96), ((3, 65, 129), 128, 80)], - ids=["equal-length-packed", "ragged-packed"], -) -def test_kda_configs_match_token_recurrence(monkeypatch, config_index, seq_lens, key_dim, value_dim): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - if config_index is not None: - for wrapper_name, _ in KERNELS: - kernel = getattr(kda, wrapper_name) - config = kernel.configs_gen_func()[config_index] - monkeypatch.setattr(kda, wrapper_name, functools.partial(kernel, run_config=config)) - - # Include partial final chunks, distinct K/V sizes, and V tiles with masked columns. - batch, tokens, heads = 1, sum(seq_lens), 2 - boundaries = [0, *accumulate(seq_lens)] - sequences = [(0, start, end) for start, end in zip(boundaries, boundaries[1:])] - cu_seqlens = torch.tensor(boundaries, device="cuda", dtype=torch.int32) - q, k = [torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)] - v = torch.randn(batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16) - g = -torch.rand(batch, tokens, heads, key_dim, device="cuda") * 0.1 - beta = torch.rand(batch, tokens, heads, device="cuda") - initial = torch.randn(len(sequences), heads, key_dim, value_dim, device="cuda") * 0.1 - initial_before = initial.clone() - expected, expected_final = recurrent_reference(q, k, v, g, beta, initial, sequences) - - actual, final = kda.chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=initial, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) - torch.testing.assert_close(actual.float().cpu(), expected, atol=4e-3, rtol=3e-2) - torch.testing.assert_close(final.cpu(), expected_final, atol=8e-3, rtol=3e-2) - torch.testing.assert_close(initial, initial_before, atol=0, rtol=0) - - -@pytest.mark.parametrize("fused_gate", [False, True]) -@pytest.mark.parametrize("invalid_input", ["batch-dimension", "missing-boundaries"]) -def test_kda_requires_packed_inputs(fused_gate, invalid_input): - batch = 2 if invalid_input == "batch-dimension" else 1 - q = torch.empty(batch, 3, 2, 128) - inputs = dict( - q=q, - k=q, - v=q, - beta=torch.empty(batch, 3, 2), - cu_seqlens=None if invalid_input == "missing-boundaries" else torch.tensor([0, 3, 6]), - ) - if fused_gate: - kernel = kda.chunk_kda_with_fused_gate - inputs.update(raw_g=q, A_log=torch.empty(2), g_bias=None) - else: - kernel = kda.chunk_kda - inputs.update(g=q) - expected_message = "packed q" if invalid_input == "batch-dimension" else "cu_seqlens is required" - with pytest.raises(AssertionError, match=expected_message): - kernel(**inputs) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "heads,input_dtype,output_dtype,precision", - [ - (2, torch.float32, torch.bfloat16, "ieee"), - (32, torch.bfloat16, torch.bfloat16, "ieee"), - (64, torch.float16, torch.float16, "ieee"), - (64, torch.float32, torch.float32, "ieee"), - (64, torch.float32, torch.bfloat16, "tf32"), - (64, torch.float32, torch.bfloat16, "tf32x3"), - ], -) -def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): - if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: - pytest.skip("TF32 requires SM80+") - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) - lengths = (1, 16, 1519) - cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) - rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) - matrix = torch.randn(1, sum(lengths), heads, 64, device="cuda", dtype=input_dtype) * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, cu_seqlens, output_dtype) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, cu_seqlens, output_dtype, split_diagonal=True) - assert len(launches) == 1 - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) -def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - rows = torch.arange(tokens, device="cuda") % 64 - matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, output_dtype=torch.bfloat16) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, output_dtype=torch.bfloat16, split_diagonal=True) - assert len(launches) == 1 - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "lengths,expect_split", - [ - ((512,) * 16, True), - ((64,) * 128, True), - ((16,) * 16, True), - ((128,) * 128, True), - ((1,) * 256, False), - ((32,), False), - ], - ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], -) -def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) - rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) - matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16, split_diagonal=True) - assert len(launches) == int(expect_split) - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py deleted file mode 100644 index 2f8544cf75..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ /dev/null @@ -1,374 +0,0 @@ -import collections -import inspect -import json -import math - -import pytest -import torch -import torch.nn.functional as F - -from lightllm.common.basemodel.triton_kernel.linear_att import ( - mtp_fused_recurrent as kernel_module, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import ( - Autotuner, - AutotuneKernelType, - AutotuneLevel, -) - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch, tmp_path): - torch.manual_seed(42) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr( - autotuner_module, - "get_triton_autotune_level", - lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE, - ) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - - -def make_inputs( - lengths=(3, 2, 0), - mtp_size=3, - device="cpu", - head_k_dim=64, - head_v_dim=64, - heads=2, - value_heads=4, - dtype=torch.bfloat16, - state_dtype=torch.float32, - separate_write=False, - padding=0, -): - num_seqs = len(lengths) - num_tokens = sum(lengths) + padding - # 模拟生产环境从 mixed QKV 和 gate 投影切分的非连续 token 视图。 - mixed = torch.randn( - num_tokens, - heads * head_k_dim * 2 + value_heads * head_v_dim, - dtype=dtype, - device=device, - ) - q, k, v = mixed.split([heads * head_k_dim, heads * head_k_dim, value_heads * head_v_dim], dim=-1) - gates = torch.randn(num_tokens, value_heads * 3, dtype=dtype, device=device) - indices = torch.arange(1, num_seqs * mtp_size + 1, dtype=torch.int32, device=device).view(num_seqs, mtp_size) - return dict( - q=q.view(1, num_tokens, heads, head_k_dim), - k=k.view(1, num_tokens, heads, head_k_dim), - v=v.view(1, num_tokens, value_heads, head_v_dim), - initial_state=torch.randn( - 2 * num_seqs * mtp_size + 5, - value_heads, - head_k_dim, - head_v_dim, - dtype=state_dtype, - device=device, - ) - * 0.25, - cu_seqlens=torch.tensor( - [0, *torch.tensor(lengths).cumsum(0).tolist()], - dtype=torch.int64, - device=device, - ), - ssm_state_indices=indices, - ssm_state_write_indices=(indices + num_seqs * mtp_size if separate_write else indices), - num_accepted_tokens=torch.tensor( - [i % mtp_size + 1 for i in range(num_seqs)], - dtype=torch.int32, - device=device, - ), - A_log=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, - dt_bias=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, - a_raw=gates[:, :value_heads], - b_raw=gates[:, value_heads : 2 * value_heads], - ) - - -def rebuild(inputs): - args, kwargs = kernel_module.rebuild_inputs(**inputs) - return inspect.signature(kernel_module.mtp_fused_recurrent_gated_delta_rule.fn).bind(*args, **kwargs).arguments - - -def reference(inputs): - state = inputs["initial_state"].clone() - output = torch.zeros_like(inputs["v"], dtype=torch.float32) - q, k, v = [inputs[name][0].float() for name in ["q", "k", "v"]] - group_size = v.shape[1] // q.shape[1] - q = (q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) - k = (k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) - q *= q.shape[-1] ** -0.5 - cumulative = inputs["cu_seqlens"].tolist() - for row, (start, end) in enumerate(zip(cumulative, cumulative[1:])): - if start == end: - continue - read_index = inputs["ssm_state_indices"][row, inputs["num_accepted_tokens"][row] - 1] - h = inputs["initial_state"][read_index].float().clone() - for token in range(start, end): - g = -inputs["A_log"].float().exp() * F.softplus(inputs["a_raw"][token].float() + inputs["dt_bias"].float()) - h *= g.exp()[:, None, None] - delta = (v[token] - (h * k[token, :, :, None]).sum(1)) * inputs["b_raw"][token].float().sigmoid()[:, None] - h += k[token, :, :, None] * delta[:, None, :] - output[0, token] = (h * q[token, :, :, None]).sum(1) - state[inputs["ssm_state_write_indices"][row, token - start]] = h.to(state.dtype) - return output, state - - -@pytest.mark.parametrize("lengths,mtp_size,padding", [((3, 3), 3, 0), ((0,) * 7, 3, 7), ((2, 1, 0), 4, 2)]) -def test_rebuild_uses_independent_state_slots(lengths, mtp_size, padding): - inputs = make_inputs(lengths=lengths, mtp_size=mtp_size, padding=padding) - snapshots = {name: value.clone() for name, value in inputs.items()} - rebuilt = rebuild(inputs) - tokens = inputs["q"].shape[1] - active_seqs = (tokens + mtp_size - 1) // mtp_size - assert rebuilt["cu_seqlens"].tolist() == [min(i * mtp_size, tokens) for i in range(len(lengths) + 1)] - assert rebuilt["initial_state"] is inputs["initial_state"] - assert rebuilt["num_accepted_tokens"].tolist() == [1] * len(lengths) - assert rebuilt["ssm_state_write_indices"][:active_seqs].flatten().tolist() == list(range(active_seqs * mtp_size)) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - assert kernel._run_key(**inputs) == kernel._run_key(**rebuilt) - for name in ["q", "k", "v", "A_log", "dt_bias", "a_raw", "b_raw"]: - assert rebuilt[name] is inputs[name] - for name, snapshot in snapshots.items(): - torch.testing.assert_close(inputs[name], snapshot) - - -def test_rebuild_checks_active_state_capacity(): - inputs = make_inputs(lengths=(0,) * 7, padding=7) - # 7 个有效 token 只需 7 个槽；空序列及末组未使用的索引取余后也应在池内。 - inputs["initial_state"] = inputs["initial_state"][:7] - rebuilt = rebuild(inputs) - assert rebuilt["initial_state"] is inputs["initial_state"] - for name in ["ssm_state_indices", "ssm_state_write_indices"]: - indices = rebuilt[name] - assert torch.all((indices >= 0) & (indices < 7)) - assert indices.flatten()[:7].tolist() == list(range(7)) - assert indices[2].tolist() == [6, 0, 1] - # 槽数不足时仍拒绝调优，不能靠取余让有效请求的状态槽相互重叠。 - inputs["initial_state"] = inputs["initial_state"][:6] - with pytest.raises(AssertionError, match="Not enough SSM state slots"): - rebuild(inputs) - - -def test_keys_use_mtp_workload_and_state_dtype(monkeypatch): - # 历史上下文长度与递推计算量无关，不应读取 full attention 的调优长度环境变量。 - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - inputs = make_inputs() - key = kernel._static_key(**inputs) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel._run_key(**inputs) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 0))) == kernel._run_key(**inputs) - assert kernel._run_key(**make_inputs(lengths=(3, 2))) == kernel._run_key(**inputs) - assert kernel._run_key(**make_inputs(mtp_size=4)) == 16 - # 不同序列数共享 token 桶，11/12 个 token 共用一桶，第 13 个 token 进入下一桶。 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 2, 0))) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 0))) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 1))) == 24 - # head 数不再编码进 run key，但不同模型/TP 的 head 规模应使用不同静态配置文件。 - more_heads = make_inputs(heads=4, value_heads=8) - assert kernel._run_key(**more_heads) == kernel._run_key(**inputs) - assert kernel._static_key(**more_heads) != key - assert kernel._static_key(**make_inputs(state_dtype=torch.bfloat16)) != key - assert kernel._static_key(**make_inputs(mtp_size=4)) != key - assert kernel._static_key(**make_inputs(padding=1)) == key - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "lengths,mtp_size,head_k_dim,head_v_dim,heads,value_heads,dtype,state_dtype,separate_write", - [ - ((1, 0, 1), 1, 64, 64, 2, 8, torch.bfloat16, torch.bfloat16, False), - ((4, 2, 0), 4, 128, 128, 4, 8, torch.bfloat16, torch.float32, True), - ((2, 3, 1), 3, 64, 80, 2, 2, torch.float16, torch.float32, False), - ((3, 3), 3, 128, 128, 16, 16, torch.bfloat16, torch.float32, False), - ((3, 1), 3, 256, 128, 2, 4, torch.bfloat16, torch.bfloat16, True), - ], -) -def test_all_candidates_match_fp32( - lengths, - mtp_size, - head_k_dim, - head_v_dim, - heads, - value_heads, - dtype, - state_dtype, - separate_write, -): - inputs = make_inputs( - lengths, - mtp_size, - "cuda", - head_k_dim, - head_v_dim, - heads, - value_heads, - dtype, - state_dtype, - separate_write, - ) - output_ref, state_ref = reference(inputs) - state_before = inputs["initial_state"].clone() - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - for config in kernel_module.get_test_configs(): - inputs["initial_state"].copy_(state_before) - output, state = kernel(**inputs, run_config=config) - assert state is inputs["initial_state"] - torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(state.float(), state_ref.float(), atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("dynamic", [False, True]) -def test_autotune_reuses_state_cache_and_graph(monkeypatch, tmp_path, dynamic): - inputs = make_inputs( - lengths=(3, 2, 0) if dynamic else (3, 3), - device="cuda", - padding=2 if dynamic else 0, - ) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - state_before = inputs["initial_state"].clone() - output_ref, state_ref = reference(inputs) - benchmark = kernel._bench - timings = [] - num_configs = len(kernel_module.get_test_configs()) - post_tuning_reference = None - - def checked_bench(*args, **kwargs): - nonlocal post_tuning_reference - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["cu_seqlens"].tolist() == ([0, 3, 6, 7] if dynamic else [0, 3, 6]) - assert rebuilt["initial_state"] is inputs["initial_state"] - elapsed = benchmark(*args, **kwargs) - timings.append(elapsed) - # benchmark 原地更新池内前 num_tokens 个槽，未使用的槽应保持不变。 - assert torch.isfinite(rebuilt["initial_state"]).all() - num_tokens = inputs["q"].shape[1] - assert not torch.equal(inputs["initial_state"][:num_tokens], state_before[:num_tokens]) - torch.testing.assert_close(inputs["initial_state"][num_tokens:], state_before[num_tokens:], atol=0, rtol=0) - if len(timings) == num_configs: - # 最后一次 benchmark 后的状态作为正式执行的初值，不再假定调优前后状态隔离。 - post_tuning_reference = reference(inputs) - return elapsed - - monkeypatch.setattr(kernel, "_bench", checked_bench) - # 普通 warmup 不搜索；decode 调优会原地更新状态，随后按原始请求布局正式执行。 - with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): - kernel(**inputs) - assert not timings - inputs["initial_state"].copy_(state_before) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - output, _ = kernel(**inputs) - assert len(timings) == num_configs and all(math.isfinite(t) for t in timings) - tuned_output_ref, tuned_state_ref = post_tuning_reference - valid_tokens = int(inputs["cu_seqlens"][-1]) - torch.testing.assert_close( - output[:, :valid_tokens].float(), - tuned_output_ref[:, :valid_tokens], - atol=2e-3, - rtol=2e-2, - ) - torch.testing.assert_close(inputs["initial_state"], tuned_state_ref, atol=2e-3, rtol=2e-2) - config_path = next(tmp_path.glob("*.json")) - assert str(kernel._run_key(**inputs)) in json.loads(config_path.read_text()) - - # 关闭已有配置预热后，文件重载和 FORCE 复用都只能正式执行一次，即使在 decode warmup 阶段。 - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr( - autotuner_module, - "get_triton_autotune_level", - lambda: AutotuneLevel.FORCE_AUTOTUNE, - ) - inputs["initial_state"].copy_(state_before) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert len(timings) == num_configs - torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) - assert not kernel.warmuped_configs_set - # 预热结束后恢复测试请求的初始状态，正常调用及 Graph 回放每次只推进一次。 - inputs["initial_state"].copy_(state_before) - kernel(**inputs) - torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) - - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_output, graph_state = kernel(**inputs) - inputs["initial_state"].copy_(state_before) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close( - graph_output[:, :valid_tokens].float(), - output_ref[:, :valid_tokens], - atol=2e-3, - rtol=2e-2, - ) - torch.testing.assert_close(graph_state, state_ref, atol=2e-3, rtol=2e-2) - # 同一 Graph 更新变长分组和接受位置，确保回放依赖原始输入而非调优构造的元数据。 - inputs["cu_seqlens"].copy_(torch.tensor([0, 1, 3, 3] if dynamic else [0, 2, 3], device="cuda")) - inputs["num_accepted_tokens"].fill_(2) - inputs["initial_state"].copy_(state_before) - new_output_ref, new_state_ref = reference(inputs) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(graph_output[:, :3].float(), new_output_ref[:, :3], atol=2e-3, rtol=2e-2) - torch.testing.assert_close(graph_state, new_state_ref, atol=2e-3, rtol=2e-2) - assert len(timings) == num_configs - - # 改变序列数和 token 数但仍在同一桶内，decode warmup 应复用已有配置，不再搜索。 - same_bucket_inputs = make_inputs(lengths=(3,), device="cuda") - assert kernel._static_key(**same_bucket_inputs) == kernel._static_key(**inputs) - assert kernel._run_key(**same_bucket_inputs) == kernel._run_key(**inputs) - same_output_ref, same_state_ref = reference(same_bucket_inputs) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - same_output, same_state = kernel(**same_bucket_inputs) - assert len(timings) == num_configs - torch.testing.assert_close(same_output.float(), same_output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(same_state, same_state_ref, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("level", [0, 1, 2]) -def test_first_decode_loads_history_without_advancing_state_twice(monkeypatch, tmp_path, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - inputs = make_inputs(lengths=(3, 3), device="cuda") - # 模拟两个刚完成 prefill 的请求，从各自的 canonical 状态槽开始执行首次 decode。 - inputs["ssm_state_indices"] = torch.arange(6, dtype=torch.int32, device="cuda").view(2, 3) - inputs["ssm_state_write_indices"] = inputs["ssm_state_indices"] - inputs["num_accepted_tokens"].fill_(1) - output_ref, state_ref = reference(inputs) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - config = {"BV": 8, "num_warps": 1, "num_stages": 1} - filename = autotuner_module.KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) - cache_file = tmp_path / filename - cache_file.write_text( - json.dumps( - { - str(kernel._run_key(**inputs)): config, - "24": {"BV": 16, "num_warps": 2, "num_stages": 1}, - } - ) - ) - calls = [] - fn = kernel.fn - - def count_calls(*args, **kwargs): - calls.append(kwargs.get("run_config")) - return fn(*args, **kwargs) - - monkeypatch.setattr(kernel, "fn", count_calls) - assert not Autotuner.is_autotune_warmup() - output, state = kernel(**inputs) - assert calls == [config] - assert state is inputs["initial_state"] - assert not kernel.warmuped_configs_set - torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(state, state_ref, atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..5e03665d43 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -101,9 +101,7 @@ def test_mtp_verify_path(mtp_step): b_raw, ) - assert torch.equal( - o_old, o_new - ), f"output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" + torch.testing.assert_close(o_old, o_new, rtol=torch.finfo(o_old.dtype).eps, atol=2e-3) if not torch.equal(fs_old, fs_new): assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( f"final_state mismatch at mtp_step={mtp_step}, " @@ -153,9 +151,7 @@ def test_variable_length_cu_seqlens(seed): b_raw, ) - assert torch.equal( - o_old, o_new - ), f"seed={seed}: output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" + torch.testing.assert_close(o_old, o_new, rtol=torch.finfo(o_old.dtype).eps, atol=2e-3) if not torch.equal(fs_old, fs_new): assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( f"seed={seed}: final_state mismatch, " diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 786376a362..e6c7d7761c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -120,35 +120,6 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state torch.testing.assert_close(ssm_snapshot(ssm_backend, 0)[0].float(), expected_state, atol=128, rtol=0.01) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_flush_export_updates_checkpoint_and_clears_history(dtype): - state = torch.randn(1, 2, 1, 32, 64, device="cuda", dtype=dtype) * 0.01 - backend = ReplaySSMExecutor( - state, - capacity=4, - verify_width=1, - activation_dtype=torch.bfloat16, - ) - backend.raw_keys.normal_(0, 0.1) - backend.raw_values.normal_(0, 0.1) - backend.gates.fill_(-0.1) - backend.betas.fill_(0.5) - backend.cursors[0] = 2 - expected = ssm_snapshot(backend, 0).clone() - target = ssm_snapshot(backend, 0) - reqs = torch.tensor([0], dtype=torch.int32, device="cuda") - backend.fold_accepted_prefix(reqs) - torch.cuda.synchronize() - - assert backend.cursors[0].item() == 0 - torch.testing.assert_close(target, expected, atol=0, rtol=0) - torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) - - empty_history_target = ssm_snapshot(backend, 0) - torch.cuda.synchronize() - torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) - - @pytest.mark.parametrize("capacity", [16, 32, 64]) @pytest.mark.parametrize("width", [1, 3, 5, 16]) @pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) @@ -263,254 +234,6 @@ def step(): torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) -@pytest.mark.parametrize( - "state_dtype,activation_dtype", - [(torch.float32, torch.float32), (torch.float32, torch.bfloat16), (torch.bfloat16, torch.bfloat16)], -) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) -@pytest.mark.parametrize("width", [1, 4]) -@pytest.mark.parametrize("gate_shift", [-3, -10]) -def test_fold_matches_reference_across_precisions( - state_dtype, - activation_dtype, - compute_mode, - dims, - width, - gate_shift, -): - test_replay_acceptance_flush_and_graph( - width, - dims, - 8, - state_dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode=compute_mode, - activation_dtype=activation_dtype, - gate_shift=gate_shift, - ) - - -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 3), (8, 4), (8, 5)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_short_replay_history_matches_independent_reference(capacity, width, dtype): - test_replay_acceptance_flush_and_graph(width, (128, 128), capacity, dtype) - - -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 4), (8, 5), (16, 4), (16, 16), (32, 4)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_precomputed_projections_match_independent_reference(capacity, width, dtype): - test_replay_acceptance_flush_and_graph( - width, - (128, 128), - capacity, - dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode="precompute", - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_production_gdn_geometry_matches_independent_reference(dtype): - test_replay_acceptance_flush_and_graph( - 4, - (128, 128), - 8, - dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode="precompute", - head_geometry=(4, 12), - ) - - -@pytest.mark.parametrize("warps", [1, 4]) -@pytest.mark.parametrize("decay", [-3, -10]) -@pytest.mark.parametrize( - "state_dtype,dtype", - [(torch.float32, torch.float32), (torch.bfloat16, torch.bfloat16), (torch.float32, torch.bfloat16)], -) -@pytest.mark.parametrize("bv", [8, 32]) -@pytest.mark.parametrize("width", [1, 4]) -def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, dtype, width, bv): - torch.manual_seed(35) - state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) - config = {"BV": bv, "num_warps": warps, "num_stages": 3} - ssm_backends = [ - ReplaySSMExecutor( - state.clone(), - 16 if width == 1 else 8, - width, - dtype, - num_key_heads=2, - compute_mode=mode, - ) - for mode in ("inline", "precompute") - ] - reqs = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - tokens = 2 if width == 1 else 7 - cu = torch.tensor([0, 1, 2] if width == 1 else [0, 4, 7], dtype=torch.int32, device="cuda") - q = torch.randn((tokens, 2, 128), device="cuda", dtype=dtype) - k = torch.randn_like(q) - v = torch.randn((tokens, 4, 128), device="cuda", dtype=dtype) - a = torch.randn((tokens, 4), device="cuda", dtype=dtype) + decay - b = torch.randn_like(a) - alog = torch.zeros(4, device="cuda") - accepted = torch.zeros(3, dtype=torch.int32, device="cuda") - for step in range(64): - for tensor in (q, k, v, b): - tensor.normal_() - accepted[0] = step % max(1, width - 1) - accepted[1] = step % width - outputs = [] - step_lengths = (cu[1:] - cu[:-1]).tolist() - if any( - ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity - for req, length in zip(reqs.tolist(), step_lengths) - ): - overflow = [ - req - for req, length in zip(reqs.tolist(), step_lengths) - if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity - ] - overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(overflow_reqs) - for ssm_backend in ssm_backends: - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu, run_config=config)) - accept_replay_window(ssm_backend, reqs, accepted, cu) - if dtype == torch.float32: - torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) - else: - torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=0.01) - error = torch.linalg.vector_norm(outputs[0].float() - outputs[1].float()) - scale = torch.linalg.vector_norm(outputs[0].float()) - assert error < scale * (torch.finfo(dtype).eps / 2) - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(reqs) - torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) - - -def test_bf16_commit_rounds_before_the_next_token(): - # 前缀写满后先压回 SSM，下一 token 从长度 0 重写。两次 commit 必须得到同一块 BF16 状态。 - state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) - ssm_backends = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] - reqs = torch.tensor([0], device="cuda", dtype=torch.int32) - for ssm_backend in ssm_backends: - ssm_backend.raw_keys.zero_() - ssm_backend.raw_values.zero_() - ssm_backend.gates.fill_(-0.01) - ssm_backend.betas.zero_() - ssm_backend.cursors[0] = 16 - ssm_backend.fold_accepted_prefix(reqs) - q = torch.ones((1, 1, 32), device="cuda") - k = torch.zeros_like(q) - v = torch.zeros((1, 1, 64), device="cuda") - a = torch.full((1, 1), -3.0, device="cuda") - b = torch.zeros_like(a) - alog = torch.zeros(1, device="cuda") - bias = torch.zeros_like(alog) - outputs = [] - cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - assert ssm_backend.cursors[0].item() == 0 - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)) - torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) - torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) - - -def test_large_bf16_state_precompute_matches_inline(): - # 输出投影用 tf32x3。重建缓存的 k/d 在 bf16 激活下是 fp16，输出允许一个 fp16 单位。 - # 折回 SSM 仍从原始 k/v 在 fp32 中计算，两条路径的 checkpoint 必须一致。 - heads = 16 - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - reqs = torch.tensor([0], dtype=torch.int32, device="cuda") - alog = torch.zeros(heads, device="cuda") - bias = torch.zeros_like(alog) - - def check(capacity, width, steps): - torch.manual_seed(7) - state = (torch.randn((1, 2, heads, 128, 128), device="cuda") * 10).bfloat16() - backends = [ - ReplaySSMExecutor( - state.clone(), - capacity, - width, - torch.bfloat16, - num_key_heads=heads, - compute_mode=mode, - ) - for mode in ("inline", "precompute") - ] - cu = torch.tensor([0, width], dtype=torch.int32, device="cuda") - accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") - for step in range(steps): - torch.manual_seed(100 + step) - q = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) - a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 - b = torch.randn_like(a) - outputs = [ - replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for backend in backends - ] - # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 - torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) - for backend in backends: - accept_replay_window(backend, reqs, accepted, cu) - overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] - if overflow.numel(): - backend.fold_accepted_prefix(overflow) - for backend in backends: - backend.fold_accepted_prefix(reqs) - torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) - - check(4, 1, 16) - check(16, 2, 4) - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capacity", [4, 8]) -def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): - # 放不下时，GDN 和 KDA 都在 replay kernel 内折回。cursor 等两层都算完才改。 - torch.manual_seed(11 + capacity + int(kda)) - hv, kh, dim = 2, 2, 32 - state = torch.randn(2, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda) - twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda) - assert backend.gates.shape[-1] == (dim if kda else capacity) - reqs = torch.tensor([0], device="cuda", dtype=torch.int32) - cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - alog = torch.zeros(hv, device="cuda") - bias = torch.zeros(hv * dim if kda else hv, device="cuda") - for step in range(capacity + 1): - q = torch.randn(1, kh, dim, device="cuda", dtype=dtype) - k = torch.randn_like(q) - v = torch.randn(1, hv, dim, device="cuda", dtype=dtype) - a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() - b = torch.randn(1, hv, device="cuda", dtype=dtype) - if step == capacity: - fold_prefix(twin, reqs, config) - twin.cursors.index_fill_(0, reqs.long(), 0) - assert int(backend.cursors[0]) == capacity - outs, twin_outs = [], [] - for layer in range(2): - outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) - twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) - if step == capacity: - assert int(backend.cursors[0]) == capacity - for actual, expected in zip(outs, twin_outs): - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - accept_replay_window(backend, reqs) - accept_replay_window(twin, reqs) - assert int(backend.cursors[0]) == 1 - assert int(twin.cursors[0]) == 1 - # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 - state_atol = 0 if dtype == torch.bfloat16 else 1e-6 - torch.testing.assert_close(backend.ssm_state, twin.ssm_state, atol=state_atol, rtol=0) - - @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("bv", [32, 128]) @@ -581,173 +304,12 @@ def step(owner): torch.testing.assert_close(backend.ssm_state[:, batch], initial[:, batch], rtol=0, atol=0) -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): - state = torch.randn(1, 3, 2, 32, 32, device="cuda") - backend = ReplaySSMExecutor(state, 8, 4, torch.float32, num_key_heads=2, kda=kda) - backend.cursors[0] = 3 - reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) - cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) - q = torch.randn(1, 2, 32, device="cuda") - k = torch.randn_like(q) - v = torch.randn(1, 2, 32, device="cuda") - a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") - b = torch.zeros(1, 2, device="cuda") - bias = torch.zeros(2 * 32 if kda else 2, device="cuda") - replay_layer( - backend, - 0, - q, - k, - v, - a, - b, - bias, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) - accept_replay_window(backend, reqs, accepted, cu) - assert backend.cursors.tolist() == [3, 0, 0] - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("capture_metadata", [False, True]) -def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_mode, capture_metadata): - torch.manual_seed(31) - state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) - ssm_backend = ReplaySSMExecutor( - state, - 8, - 4, - dtype, - num_key_heads=2, - kda=kda, - compute_mode=compute_mode, - ) - expected = state.float().clone() - pending = [0, 0] - reqs = torch.full((3,), 3, dtype=torch.int32, device="cuda") - cu = torch.tensor([0, 4, 8, 8], dtype=torch.int32, device="cuda") - accepted = torch.zeros(4, dtype=torch.int32, device="cuda") - q = torch.randn(8, 2, 128, dtype=dtype, device="cuda") - k = torch.randn_like(q) - v = torch.randn(8, 4, 128, dtype=dtype, device="cuda") - a = torch.full((8, 4 * 128 if kda else 4), -3.0, dtype=dtype, device="cuda") - b = torch.randn(8, 4, dtype=dtype, device="cuda") - alog = torch.zeros(4, device="cuda") - bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - - def forward(): - return [ - replay_layer( - ssm_backend, - layer, - q, - k, - v, - a, - b, - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - for layer in range(2) - ] - - def step(): - positions = ssm_backend.cursors.index_select(0, reqs.long()) - out = forward() - accept_replay_window(ssm_backend, reqs, accepted, cu) - return positions, out - - step() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - positions = torch.zeros_like(reqs) - with torch.cuda.graph(graph): - if capture_metadata: - positions, outputs = step() - else: - outputs = forward() - lengths = [(3, 4), (2, 1), (3, 2), (1, 4), (4, 3), (0, 2), (2, 0), (4, 4)] - for iteration in range(24): - ids = [0, 1] if iteration < 3 or iteration % 2 == 0 else [1, 0] - widths = lengths[iteration % len(lengths)] - counts = ( - list(widths) if iteration < 3 else [min(width, 1 + (iteration + i) % 4) for i, width in enumerate(widths)] - ) - reqs.copy_(reqs.new_tensor(ids + [3])) - cu.copy_(cu.new_tensor([0, widths[0], sum(widths), 8])) - accepted.fill_(-1) - for req, count in zip(ids, counts): - accepted[req] = count - 1 - for tensor in (q, k, v, b): - tensor.normal_() - before = state.clone() - for req, width in zip(ids, widths): - if pending[req] + width > 8: - ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) - expected[:, req] = expected[:, req].to(dtype).float() - pending[req] = 0 - old_cursors = [pending[req] for req in ids] + [0] - if not capture_metadata: - positions.copy_(ssm_backend.cursors.index_select(0, reqs.long())) - graph.replay() - if not capture_metadata: - accept_replay_window(ssm_backend, reqs, accepted, cu) - assert positions.tolist() == old_cursors - for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): - for layer in range(2): - current = expected[layer, req].clone() - for j in range(width): - t = (0 if seq == 0 else widths[0]) + j - out, current = reference(q[t], k[t], v[t], a[t], b[t], alog, bias, current, kda=kda) - torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) - if j + 1 == count: - expected[layer, req].copy_(current) - pending[req] += count - assert ssm_backend.cursors.tolist() == pending + [0, 0] - for output in outputs: - assert torch.count_nonzero(output[sum(widths) :]).item() == 0 - torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) - if iteration == 2: - assert pending[0] == 8 - torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) - ssm_backend.fold_accepted_prefix(reqs[:2]) - torch.testing.assert_close( - state.float(), - expected.to(dtype).float(), - atol=2e-5 if dtype == torch.float32 else 0.002, - rtol=2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps, - ) - - -def test_snapshot_without_history_copies_active_state(): - state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) - ssm_backend = ReplaySSMExecutor(state, 4, 3) - snapshot = ssm_snapshot(ssm_backend, 1) - torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) - assert ssm_backend.cursors[1].item() == 0 - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("run_config", get_configs(), ids=lambda c: f"bv{c['BV']}-w{c['num_warps']}-s{c['num_stages']}") -def test_autotune_candidates_match_recurrence(dtype, run_config): - test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) - - @pytest.mark.parametrize("save_big", [False, True]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_args, save_big, dtype): + from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig - from unit_tests.common.test_linear_state import make_state + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState torch.manual_seed(42) state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps @@ -761,7 +323,11 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg conv_dtype=torch.bfloat16, ssm_dtype=dtype, ) - state, _, layout = make_state(mode="replay", config=config) + state = ReplayGDNState(config, 4, mtp_step=2, device="cuda", ssm_replay_config={"cache_len": 4}) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) state.mtp_state_index[1] = 1 state.conv_state.normal_() @@ -810,119 +376,6 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg assert not state.conv_state[:, 2].any() -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("gate_lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) -@pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) -def test_kda_replay_reference_fold_and_graph( - dtype, capacity, width, compute_mode, gate_lower_bound, gate_shift, activation_dtype -): - test_replay_acceptance_flush_and_graph( - width, - (128, 128), - capacity, - dtype, - run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, - compute_mode=compute_mode, - gate_shift=gate_shift, - activation_dtype=activation_dtype, - head_geometry=(2, 2), - kda=True, - gate_lower_bound=gate_lower_bound, - ) - - -@pytest.mark.parametrize("config", get_configs()) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -def test_kda_tuning_candidates_match_reference(config, compute_mode): - test_replay_acceptance_flush_and_graph( - 3, - (128, 128), - 4, - torch.bfloat16, - run_config=config, - compute_mode=compute_mode, - head_geometry=(2, 2), - kda=True, - gate_shift=-3, - ) - - -@pytest.mark.parametrize("kda", [False, True]) -@pytest.mark.parametrize( - "bucket,live,folding,width,separate", - [ - (12, 12, 12, 4, True), - (13, 13, 13, 4, False), - (255, 255, 127, 4, True), - (256, 256, 128, 4, False), - (257, 257, 128, 4, False), - (256, 256, 128, 3, True), - (256, 12, 12, 4, True), - (256, 13, 4, 4, True), - ], -) -def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): - torch.manual_seed(72) - initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 - backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) - reqs = torch.cat( - ( - torch.arange(live, device="cuda", dtype=torch.int32), - torch.full((bucket - live,), live, device="cuda", dtype=torch.int32), - ) - ) - indices = torch.zeros(live + 1, device="cuda", dtype=torch.int32) - expected = initial[0, :live].clone() - alog = torch.zeros(1, device="cuda") - bias = torch.zeros(32 if kda else 1, device="cuda") - for step, query_width in enumerate((4, 4, width)): - if step == 2: - backend.cursors.index_fill_(0, reqs[folding:live].long(), 0) - expected[folding:] = initial[0, folding:live] - cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) - q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 - k, v = torch.randn_like(q), torch.randn_like(q) - a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") - b = torch.randn(live, query_width, 1, device="cuda") - if (backend.cursors[:live] + query_width > backend.capacity).any(): - backend.fold_accepted_prefix(reqs[:live]) - actual = replay_layer( - backend, - 0, - q.reshape(-1, 1, 32), - k.reshape(-1, 1, 32), - v.reshape(-1, 1, 32), - a.flatten(0, 1), - b.flatten(0, 1), - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ).reshape(live, query_width, 1, 32) - for token in range(query_width): - output, expected = reference( - q[:, token], - k[:, token], - v[:, token], - a[:, token], - b[:, token], - alog, - bias.expand(live, 1, 32) if kda else bias, - expected, - kda=kda, - ) - torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) - lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) - advance_cursor(backend, reqs[:live], cu[: live + 1]) - commit_ssm_prefix(indices, reqs[:live], lengths, lengths, backend.cursors) - backend.fold_accepted_prefix(reqs[:live]) - torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) - assert torch.equal(backend.ssm_state[0, live], initial[0, live]) - - @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py deleted file mode 100644 index a94411fbd5..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ /dev/null @@ -1,70 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning -from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( - attach_replay_history as ReplaySSMExecutor, -) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("width", [1, 5]) -def test_rebuild_uses_one_layer_scratch(width): - state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) - tokens = 3 * width - q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") - gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") - a_log = torch.zeros(2, device="cuda") - cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - reqs = torch.arange(3, device="cuda", dtype=torch.int32) - _, kwargs = tuning.rebuild_inputs(ssm_backend, 0, q, q, q, gates, gates, a_log, a_log, reqs, cu) - scratch = kwargs["state"] - assert scratch is not ssm_backend - assert scratch.ssm_state.shape[0] == 1 - assert scratch.capacity == ssm_backend.capacity - assert torch.count_nonzero(ssm_backend.ssm_state) == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("width", [1, 5]) -def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): - from lightllm.common.triton_utils import autotuner as autotuner_module - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType - - state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) - for tensor in vars(ssm_backend).values(): - if isinstance(tensor, torch.Tensor): - tensor.zero_() - tokens = 3 * width - q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") - gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") - a_log = torch.zeros(2, device="cuda") - bias = torch.zeros(2, device="cuda") - cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - reqs = torch.arange(3, device="cuda", dtype=torch.int32) - before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} - - def oom(*args, **kwargs): - raise torch.cuda.OutOfMemoryError("scratch") - - tuner = tuning.replay - monkeypatch.setattr(tuner, "rebuild_input_func", oom) - monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(tuner, "cached_configs", {}) - monkeypatch.setattr(tuner, "fast_match_configs", {}) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) - monkeypatch.setattr( - tuner, "_bench", lambda *args, **kwargs: pytest.fail("Failed rebuild must not reach benchmarks") - ) - with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu) - assert not list(tmp_path.iterdir()) - config = {"BV": 32, "num_warps": 1, "num_stages": 3} - for name, expected in before.items(): - torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) - - tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) - torch.cuda.synchronize() diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py deleted file mode 100644 index b4e1c1ac12..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py +++ /dev/null @@ -1,105 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - - -import pytest -import torch -import triton - -from lightllm.common.basemodel.triton_kernel.mhc import head, post, pre_norm -from lightllm.common.basemodel.triton_kernel.mhc.configs import post_projection_configs -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel -from unit_tests.common.basemodel.triton_kernel.mhc.test_mhc import _hc_post_reference, _hc_pre_reference -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward - - -@pytest.fixture(autouse=True) -def tuning_environment(monkeypatch): - torch.manual_seed(1525) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.CLOSE_AUTOTUNE) - if torch.cuda.is_available(): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("tokens,hidden", [(3, 4096), (16, 6144)]) -def test_all_pre_and_post_projection_configs_match_reference(tokens, hidden): - x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) - fn = torch.randn(24, 4 * hidden, device="cuda") * 0.005 - scale, base = torch.randn(3, device="cuda"), torch.randn(24, device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=x.dtype) - expected = _hc_pre_reference(x, fn, scale, base, 4, 1e-6, 1e-6, 20) - - def check(actual, reference): - torch.testing.assert_close(actual[0], rmsnorm_forward(reference[0], norm, 1e-6), atol=0.04, rtol=0.03) - for a, b in zip(actual[1:], reference[1:]): - torch.testing.assert_close(a, b, atol=2e-4, rtol=2e-3) - - for config in pre_norm._hc_prenorm_projection.configs_gen_func(): - gp, sp = pre_norm._hc_prenorm_projection(x, fn, run_config=config) - check(pre_norm._finish_pre_norm(x, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0), expected) - gp, sp = pre_norm._hc_prenorm_projection(x, fn) - for config in pre_norm._finish_pre_norm.configs_gen_func(): - check( - pre_norm._finish_pre_norm(x, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0, run_config=config), - expected, - ) - - layer_output = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) - residual_mix = torch.randn(tokens, 4, 4, device="cuda").softmax(-1) - post_mix = torch.rand(tokens, 4, device="cuda") * 2 - expected_residual = _hc_post_reference(layer_output, x, residual_mix, post_mix, 4) - expected = _hc_pre_reference(expected_residual, fn, scale, base, 4, 1e-6, 1e-6, 20) - before = x.clone() - for config in post_projection_configs(): - updated, gp, sp = pre_norm._hc_post_prenorm_projection( - layer_output, x, residual_mix, post_mix, fn, run_config=config - ) - torch.testing.assert_close(updated, expected_residual, atol=0.04, rtol=0.02) - check(pre_norm._finish_pre_norm(updated, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0), expected) - torch.testing.assert_close(x, before, atol=0, rtol=0) - for config in pre_norm._hc_post_pre_norm_decode.configs_gen_func(): - actual = pre_norm._hc_post_pre_norm_decode( - layer_output, - x, - residual_mix, - post_mix, - fn, - scale, - base, - norm, - 4, - 1e-6, - 1e-6, - 1e-6, - 20, - 2.0, - run_config=config, - ) - torch.testing.assert_close(actual[1], expected_residual, atol=0.04, rtol=0.02) - check((actual[0], actual[2], actual[3]), expected) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("hidden", [4096, 6144]) -def test_all_post_and_head_configs_match_reference(hidden): - tokens = 19 - x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) - y = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) - mix, post_mix = torch.randn(tokens, 4, 4, device="cuda").softmax(-1), torch.rand(tokens, 4, device="cuda") * 2 - expected_post = _hc_post_reference(y, x, mix, post_mix, 4) - for config in post.hc_post.configs_gen_func(): - actual = post.hc_post(y, x, mix, post_mix, 4, run_config=config) - torch.testing.assert_close(actual, expected_post, atol=0.04, rtol=0.02) - fn, scale, base = ( - torch.randn(4, 4 * hidden, device="cuda") * 0.005, - torch.randn(1, device="cuda"), - torch.randn(4, device="cuda"), - ) - inv_rms = torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) - gates = torch.sigmoid((x.float() @ fn.T) * inv_rms * scale + base) + 1e-6 - expected_head = (x.view(tokens, 4, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() - for config in head.hc_head.configs_gen_func(): - actual = head.hc_head(x, fn, scale, base, 4, 1e-6, 1e-6, run_config=config) - torch.testing.assert_close(actual, expected_head, atol=0.04, rtol=0.02) diff --git a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py deleted file mode 100644 index a5ac9708d2..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py +++ /dev/null @@ -1,112 +0,0 @@ -import torch -import time -import pytest -import numpy as np -import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.mla_att.decode_att.gqa_flash_decoding import ( - gqa_token_decode_attention_flash_decoding, -) -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, seqlen, heads, nope_head, rope_head", - [ - (a, b, c, d, e) - for a in [1, 16, 32, 128, 512] - for b in [16, 32, 512, 2048] - for c in [16] - for d in [512] - for e in [64] - ], -) -def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): - Z, N_CTX, H, D_HEAD, ROPE_HEAD = batch, seqlen, heads, nope_head, rope_head - dtype = torch.bfloat16 - sm_scale = 1.0 / ((D_HEAD + ROPE_HEAD) ** 0.5) - q_nope = torch.randn((Z, H, D_HEAD), dtype=dtype, device="cuda") - q_rope = torch.randn((Z, H, ROPE_HEAD), dtype=dtype, device="cuda") - - kv = torch.randn((Z * N_CTX, 1, D_HEAD + ROPE_HEAD), dtype=dtype, device="cuda") - - max_input_len = Z * N_CTX - req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX) - b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX - b_start_loc = torch.arange(Z).cuda().int() * N_CTX - b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda() - kv_starts = torch.cat([b_start_loc, b_start_loc[-1:] + b_seq_len[-1:]], dim=0) - - o = torch.zeros((Z, H, D_HEAD), dtype=dtype, device="cuda") - o1 = torch.zeros((Z, H, D_HEAD), dtype=dtype, device="cuda") - - infer_state = Deepseek2InferStateInfo() - infer_state.batch_size = Z - infer_state.max_len_in_batch = N_CTX - infer_state.total_token_num = Z * N_CTX - infer_state.req_manager = type("Object", (), {})() - infer_state.req_manager.req_to_token_indexs = req_to_token_indexs - infer_state.b_req_idx = b_req_idx - infer_state.b_seq_len = b_seq_len - infer_state.kv_starts = kv_starts - - kv_nope = kv[:, :, :D_HEAD] - kv_rope = kv[:, :, D_HEAD:] - gqa_token_decode_attention_flash_decoding( - q_nope, - q_rope, - kv_nope, - kv_rope, - infer_state, - sm_scale, - o, - ) - - batch_size = Z - head_dim_ckv = D_HEAD - head_dim_kpe = ROPE_HEAD - num_heads = H - page_size = 1 - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0) - q_indptr = torch.arange(batch_size + 1).to(0).int() - kv_indptr = infer_state.kv_starts - kv_indices = torch.arange(Z * N_CTX).cuda().int() - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - kv_indices[start : start + sl] = req_to_token_indexs[b][:sl] - kv_lens = b_seq_len - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - workspace_buffer, - use_cuda_graph=True, - qo_indptr=q_indptr, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_len_arr=kv_lens, - ) - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, # causal - sm_scale, - q_nope.dtype, - kv.dtype, - ) - wrapper.run(q_nope, q_rope, kv_nope, kv_rope, out=o1, return_lse=False) - - cos_sim1 = F.cosine_similarity(o, o1).mean() - assert cos_sim1 == 1.0 diff --git a/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py b/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py deleted file mode 100644 index 7e9a55a775..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py +++ /dev/null @@ -1,31 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy -from lightllm.common.basemodel.triton_kernel.add_in_place import add_in_place -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "dim1, dim2, alpha", - [ - (dim1, dim2, alpha) - for dim1 in range(1, 1024, 100) - for dim2 in range(1, 1024, 100) - for alpha in [0.1, 0.3, 0.5, 0.7, 0.1] - ], -) -def test_add_in_place(dim1, dim2, alpha): - input = torch.rand((dim1, dim2), device="cuda") - other = torch.rand((dim1, dim2), device="cuda") - - output = input + other * alpha - add_in_place(input, other, alpha=alpha) - rlt = torch.allclose(input, output, atol=1e-5, rtol=0) - assert rlt - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py b/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py deleted file mode 100644 index c6768053e4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py +++ /dev/null @@ -1,104 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.apply_constraint_mask import apply_constraint_mask - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Constraint mask kernel requires CUDA") - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("vocab_size,padding", [(1, 0), (31, 5), (32, 0), (33, 31), (257, 7), (8195, 43)]) -@pytest.mark.parametrize("mtp", [False, True]) -def test_request_indexed_mask_matches_reference(dtype, vocab_size, padding, mtp): - generator = torch.Generator().manual_seed(2026) - masks = torch.randint( - -(1 << 31), - (1 << 31) - 1, - (6, 4, (vocab_size + 31) // 32), - dtype=torch.int32, - generator=generator, - pin_memory=True, - ) - enabled = torch.tensor([False, True, False, False, True, False], pin_memory=True) - req_ids = torch.tensor([4, 1, 4, 2, 1], dtype=torch.int32) - positions = torch.tensor([2, 0, 1, 3, 3], dtype=torch.int32) if mtp else torch.zeros(5, dtype=torch.int32) - # Exercise a padded row stride as well as unused vocabulary columns. - width = vocab_size + padding - storage = torch.randn(5, width + 11, dtype=dtype, device="cuda") - logits = storage[:, :width] - expected = logits.cpu().clone() - token_ids = torch.arange(vocab_size) - for prediction, (req_id, position) in enumerate(zip(req_ids, positions)): - if enabled[req_id]: - words = masks[req_id, position, token_ids // 32] - allowed = ((words >> (token_ids % 32)) & 1).bool() - expected[prediction, :vocab_size].masked_fill_(~allowed, float("-inf")) - expected[:, vocab_size:] = float("-inf") - untouched_padding = storage[:, width:].clone() - - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids.cuda(), - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=vocab_size, - b_mtp_index=positions.cuda() if mtp else None, - ) - - torch.testing.assert_close(logits.cpu(), expected, rtol=0, atol=0) - torch.testing.assert_close(storage[:, width:], untouched_padding, rtol=0, atol=0) - - -def test_pinned_masks_can_be_reused_after_sampling_completion(): - masks = torch.zeros(3, 2, 2, dtype=torch.int32, pin_memory=True) - enabled = torch.tensor([True, True, False], pin_memory=True) - masks[0, 0, 0] = -(1 << 31) - masks[1, 1, 1] = 1 - req_ids = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - positions = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - logits = torch.zeros(2, 33, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - # This represents the existing sampling/post_handle completion boundary. - done = torch.cuda.Event() - with torch.cuda.stream(stream): - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids, - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=33, - b_mtp_index=positions, - ) - done.record() - done.synchronize() - assert torch.isfinite(logits[0]).nonzero().flatten().tolist() == [32] - assert torch.isfinite(logits[1]).nonzero().flatten().tolist() == [31] - - # Reuse one slot for an ordinary request, and update the other's prefix. - enabled[1] = False - masks[0].fill_(-1) - with torch.cuda.stream(stream): - logits.fill_(1.25) - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids, - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=33, - b_mtp_index=positions, - ) - done.record() - done.synchronize() - torch.testing.assert_close(logits, torch.full_like(logits, 1.25), rtol=0, atol=0) - - -def test_empty_batch_needs_no_kernel_launch(): - apply_constraint_mask( - logits=torch.empty(0, 33, device="cuda"), - b_req_idx=torch.empty(0, dtype=torch.int32, device="cuda"), - req_to_bitmask=torch.empty(1, 1, 2, dtype=torch.int32, device="cuda"), - req_to_bitmask_enabled=torch.zeros(1, dtype=torch.bool, device="cuda"), - vocab_size=33, - ) diff --git a/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py b/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py deleted file mode 100644 index 3b2f159f62..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import ( - apply_invalid_token_ids, -) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_apply_invalid_token_ids(dtype): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for Triton kernels.") - - batch_size = 4 - vocab_size = 32 - logits = torch.randn((batch_size, vocab_size), device="cuda", dtype=dtype) - expected = logits.clone() - - invalid_token_ids_per_batch = [ - [1, 3, 5], - [], - [0, 2, 31], - [7], - ] - - flat_ids = [] - cu_invalid_token_num = [0] - invalid_token_num_start = 0 - for ids in invalid_token_ids_per_batch: - flat_ids.extend(ids) - invalid_token_num_start += len(ids) - cu_invalid_token_num.append(invalid_token_num_start) - - invalid_token_ids = torch.tensor(flat_ids, device="cuda", dtype=torch.int32) - cu_invalid_token_num = torch.tensor(cu_invalid_token_num, device="cuda", dtype=torch.int32) - - for batch_idx, ids in enumerate(invalid_token_ids_per_batch): - if ids: - expected[batch_idx, ids] = float("-inf") - - apply_invalid_token_ids( - Logits=logits, - invalid_token_ids=invalid_token_ids, - cu_invalid_token_num=cu_invalid_token_num, - ) - assert torch.equal(logits, expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py b/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py deleted file mode 100644 index 0afcd5558a..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py +++ /dev/null @@ -1,28 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.atomic_event import wait_value, add_value -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_add_in_place(): - input = torch.zeros((1,), device="cuda", dtype=torch.int32) - wait_value(input, 0) - add_value(input) - wait_value(input, 1) - add_value(input) - wait_value(input, 2) - add_value(input) - wait_value(input, 3) - assert input.item() == 3, "最终值应为 3" - - -# @pytest.mark.timeout(2) -# def test_wait_timeout(): -# input = torch.zeros((1,), device="cuda", dtype=torch.int32) -# wait_value(input, 4) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py deleted file mode 100644 index 0c1c2df08a..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py +++ /dev/null @@ -1,48 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as int8kv_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel import diverse_utils - - -@pytest.mark.parametrize( - "node_ids,expected_markers", - [ - ([10, 10, 10, 10, 20, 20], [0, 0, 3, 1, 0, 2]), - ([10, 10, 20, 10, 10], [0, 2, 1, 0, 2]), - ([10, 10, -1, -1], [0, 2, 1, 1]), - ([-1, -1, -1, -1], [1, 1, 1, 1]), - ], -) -def test_build_diverse_shared_group_markers(monkeypatch, node_ids, expected_markers): - monkeypatch.setattr(diverse_utils, "get_diverse_max_batch_shared_group_size", lambda: 3) - b_shared_radix_node_id = torch.tensor(node_ids, dtype=torch.int64, device="cuda") - - markers = diverse_utils.build_diverse_shared_group_markers( - b_shared_radix_node_id=b_shared_radix_node_id, - ) - - assert markers.cpu().tolist() == expected_markers - - -def test_int8kv_decode_state_rebuilds_diverse_metadata(monkeypatch): - monkeypatch.setattr(int8kv_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: True) - monkeypatch.setattr(diverse_utils, "get_diverse_max_batch_shared_group_size", lambda: 3) - infer_state = InferStateInfo() - infer_state.max_kv_seq_len = 8192 - infer_state.b_shared_seq_len = torch.tensor([8, 8, 5, 0], dtype=torch.int32, device="cuda") - infer_state.b_shared_radix_node_id = torch.tensor([10, 10, 20, -1], dtype=torch.int64, device="cuda") - state = Int8kvTritonDecodeAttState(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer_state)) - - state.init_state() - - assert state.b_mark_shared_group.cpu().tolist() == [0, 2, 1, 1] - assert state.b_shared_seq_len.cpu().tolist() == [8, 8, 0, 0] diff --git a/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py b/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py deleted file mode 100644 index cd136ae713..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py +++ /dev/null @@ -1,360 +0,0 @@ -import json - -import torch -import pytest -import triton -import numpy as np - -from lightllm.common.basemodel.batch_objs import ModelInput -from lightllm.common.basemodel.mtp_manager import MtpManager -from lightllm.common.basemodel.triton_kernel import dynamic_mtp_utils -from lightllm.common.basemodel.triton_kernel.dynamic_mtp_utils import ( - _fwd_kernel_cumprod_scores, - sample_dynamic_mtp_row_mask, -) -from lightllm.utils.envs_utils import get_env_start_args - - -@pytest.fixture(autouse=True) -def _reset_mtp_manager(): - MtpManager._instance = None - yield - MtpManager._instance = None - - -def test_compact_dynamic_mtp_model_input(monkeypatch): - monkeypatch.setenv( - "LIGHTLLM_START_ARGS", - json.dumps( - { - "diverse_mode": False, - "llm_kv_type": "fp16", - "mtp_mode": "eagle3", - "mtp_dynamic_verify": True, - "mtp_step": 3, - "llm_decode_att_backend": "triton", - } - ), - ) - monkeypatch.setenv("LIGHTLLM_MAX_BATCH_SHARED_GROUP_SIZE", "4") - get_env_start_args.cache_clear() - - model_input = ModelInput( - batch_size=12, - total_token_num=54, - max_q_seq_len=1, - max_kv_seq_len=6, - input_ids=torch.arange(12, dtype=torch.int64, device="cuda") + 1000, - b_req_idx=torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6], dtype=torch.int32, device="cuda"), - b_position_delta=torch.arange(12, dtype=torch.int32, device="cuda") + 200, - b_shared_seq_len=torch.tensor([0, 0, 0, 0, 7, 7, 7, 7, 9, 9, 9, 9], dtype=torch.int32, device="cuda"), - b_shared_radix_node_id=torch.tensor( - [10, 10, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12], dtype=torch.int64, device="cuda" - ), - is_prefill=False, - multimodal_params=[{"row": i, "images": [], "audios": []} for i in range(12)], - mtp_draft_input_hiddens=(torch.arange(12 * 5, dtype=torch.float32, device="cuda").reshape(12, 5) + 0.5), - ) - req_to_next_token_scores = torch.tensor( - [ - [1.0, 0.95, 0.90, 0.10, 0.0, 0.0], - [1.0, 0.20, 0.80, 0.80, 0.0, 0.0], - [1.0, 0.99, 0.99, 0.99, 0.0, 0.0], - ], - dtype=torch.float32, - device="cuda", - ) - - compacted_input, selected_row_mask = dynamic_mtp_utils.prepare_dynamic_mtp_model_input( - model_input=model_input, - req_num=3, - dynamic_batch_size=8, - req_to_next_token_scores=req_to_next_token_scores, - ) - torch.cuda.synchronize() - - expected_selected_mask = torch.tensor([1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32) - expected_selected_rows = torch.where(expected_selected_mask == 1)[0] - - assert torch.equal(selected_row_mask.cpu(), expected_selected_mask) - assert compacted_input.batch_size == 8 - assert compacted_input.max_q_seq_len == 1 - assert compacted_input.multimodal_params == [{"images": [], "audios": []}] * 8 - - assert torch.equal( - compacted_input.input_ids.cpu(), torch.arange(12, dtype=torch.int64)[expected_selected_rows] + 1000 - ) - assert torch.equal(compacted_input.b_req_idx.cpu(), torch.tensor([0, 0, 0, 1, 2, 2, 2, 2], dtype=torch.int32)) - assert torch.equal(compacted_input.b_mtp_index.cpu(), torch.tensor([0, 1, 2, 0, 0, 1, 2, 3], dtype=torch.int32)) - assert torch.equal(compacted_input.b_seq_len.cpu(), torch.tensor([3, 4, 5, 3, 3, 4, 5, 6], dtype=torch.int32)) - assert torch.equal( - compacted_input.b_position_delta.cpu(), - torch.tensor([200, 201, 202, 204, 208, 209, 210, 211], dtype=torch.int32), - ) - assert torch.equal( - compacted_input.b_shared_seq_len.cpu(), torch.tensor([0, 0, 0, 7, 9, 9, 9, 9], dtype=torch.int32) - ) - assert torch.equal( - compacted_input.b_shared_radix_node_id.cpu(), - torch.tensor([10, 10, 10, 11, 12, 12, 12, 12], dtype=torch.int64), - ) - expected_hiddens = (torch.arange(12 * 5, dtype=torch.float32).reshape(12, 5) + 0.5)[expected_selected_rows] - assert torch.equal(compacted_input.mtp_draft_input_hiddens.cpu(), expected_hiddens) - - -def test_compaction_preserves_shared_radix_metadata(): - model_input = ModelInput( - batch_size=5, - total_token_num=36, - max_q_seq_len=1, - max_kv_seq_len=6, - input_ids=None, - b_req_idx=torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4, 5, 6, 7], dtype=torch.int32, device="cuda"), - b_position_delta=torch.zeros(5, dtype=torch.int32, device="cuda"), - b_shared_seq_len=torch.full((5,), 7, dtype=torch.int32, device="cuda"), - b_shared_radix_node_id=torch.full((5,), 10, dtype=torch.int64, device="cuda"), - is_prefill=False, - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - ) - selected_row_mask = torch.ones((5,), dtype=torch.int32, device="cuda") - - compacted_input = dynamic_mtp_utils._compact_decode_model_input( - model_input=model_input, - selected_row_mask=selected_row_mask, - dynamic_batch_size=5, - ) - torch.cuda.synchronize() - - assert torch.equal(compacted_input.b_req_idx.cpu(), torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32)) - assert torch.equal(compacted_input.b_mtp_index.cpu(), torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32)) - assert torch.equal(compacted_input.b_shared_seq_len.cpu(), torch.full((5,), 7, dtype=torch.int32)) - assert torch.equal(compacted_input.b_shared_radix_node_id.cpu(), torch.full((5,), 10, dtype=torch.int64)) - - -def _reference_cumprod_scores(req_to_next_token_scores, b_req_idx, max_draft_step: int) -> torch.Tensor: - scores = req_to_next_token_scores.clone() - req_num = b_req_idx.shape[0] // (max_draft_step + 1) - for req_i in range(req_num): - req_idx = int(b_req_idx[req_i * (max_draft_step + 1)].item()) - row = scores[req_idx, : max_draft_step + 1].clone() - row[0] = 1.0 - row[1:] = torch.clamp(row[1:], min=0.01, max=0.99) - scores[req_idx, : max_draft_step + 1] = torch.cumprod(row, dim=0) - return scores - - -def _flat_cumprod_scores( - b_req_idx: torch.Tensor, - req_to_next_token_scores: torch.Tensor, - max_draft_step: int, -) -> torch.Tensor: - scores = _reference_cumprod_scores(req_to_next_token_scores, b_req_idx, max_draft_step) - req_num = b_req_idx.shape[0] // (max_draft_step + 1) - all_num = req_num * (max_draft_step + 1) - flat_scores = [] - for offset in range(all_num): - req_idx = int(b_req_idx[offset].item()) - mtp_index = offset % (max_draft_step + 1) - flat_scores.append(scores[req_idx, mtp_index]) - return torch.stack(flat_scores) - - -def _assert_topk_mask(select: torch.Tensor, flat_scores: torch.Tensor, dynamic_batch_size: int) -> None: - k = dynamic_batch_size - assert int(select.sum().item()) == k - selected_scores = flat_scores[select.bool()] - unselected_scores = flat_scores[(select == 0).bool()] - if unselected_scores.numel() > 0: - assert selected_scores.min() >= unselected_scores.max() - 1e-5 - - -def _make_batch_scores(req_num: int, max_draft_step: int, rows): - max_req = req_num - scores = torch.zeros((max_req + 1, 16), dtype=torch.float32, device="cuda") - for req_idx, row in enumerate(rows): - scores[req_idx, : max_draft_step + 1] = torch.tensor(row, dtype=torch.float32, device="cuda") - b_req_idx = torch.arange(req_num, dtype=torch.int32, device="cuda").repeat_interleave(max_draft_step + 1) - return scores, b_req_idx - - -@pytest.mark.parametrize("max_draft_step", [1, 3]) -def test_cumprod_scores_kernel(max_draft_step: int): - req_num = 2 - scores, b_req_idx = _make_batch_scores( - req_num, - max_draft_step, - rows=[ - [1.0] + [0.5] * max_draft_step, - [1.0] + [0.2] * max_draft_step, - ], - ) - scores_clone = scores.clone() - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores_clone, - req_to_next_token_scores_stride=scores_clone.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - expected = _reference_cumprod_scores(scores, b_req_idx, max_draft_step) - assert torch.allclose( - scores_clone[:, : max_draft_step + 1], expected[:, : max_draft_step + 1], rtol=1e-5, atol=1e-5 - ) - - -def test_cumprod_scores_clamps_invalid_values(): - max_draft_step = 2 - req_num = 1 - scores, b_req_idx = _make_batch_scores(req_num, max_draft_step, rows=[[1.0, 0.0, 1.5]]) - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores, - req_to_next_token_scores_stride=scores.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - row = scores[0, : max_draft_step + 1] - # Index 0 is the always-accepted target sample; clamping applies only to drafts. - assert row[0].item() == pytest.approx(1.0) - assert row[1].item() == pytest.approx(0.01, rel=1e-4) - assert row[2].item() == pytest.approx(0.01 * 0.99, rel=1e-4) - - -def test_cumprod_scores_clamps_boundary_values(): - max_draft_step = 3 - req_num = 1 - scores, b_req_idx = _make_batch_scores(req_num, max_draft_step, rows=[[1.0, 0.995, 0.005, 0.5]]) - raw_scores = scores.clone() - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores, - req_to_next_token_scores_stride=scores.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - expected = _reference_cumprod_scores(raw_scores, b_req_idx, max_draft_step) - row = scores[0, : max_draft_step + 1] - assert torch.allclose(row, expected[0, : max_draft_step + 1], rtol=1e-5, atol=1e-5) - # Draft scores 0.995 and 0.005 clamp to 0.99 and 0.01. - assert row[0].item() == pytest.approx(1.0) - assert row[1].item() == pytest.approx(0.99, rel=1e-4) - assert row[2].item() == pytest.approx(0.99 * 0.01, rel=1e-4) - assert row[3].item() == pytest.approx(0.99 * 0.01 * 0.5, rel=1e-4) - - -def test_sample_select_count(): - max_draft_step = 3 - req_num = 3 - scores, b_req_idx = _make_batch_scores( - req_num, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - all_num = req_num * (max_draft_step + 1) - for dynamic_batch_size in [3, 8, all_num]: - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=dynamic_batch_size, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - assert select.dtype == torch.int32 - assert select.shape[0] == all_num - assert int(select.sum().item()) == dynamic_batch_size - assert torch.all((select == 0) | (select == 1)) - - -def test_sample_accepts_numpy_scalar_dynamic_batch_size(): - max_draft_step = 3 - scores, b_req_idx = _make_batch_scores( - 3, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=np.int64(8), - b_req_idx=b_req_idx, - req_to_next_token_scores=scores, - max_draft_step=np.int64(max_draft_step), - ) - assert int(select.sum().item()) == 8 - - -def test_sample_topk_by_cumprod_score(): - max_draft_step = 3 - scores, b_req_idx = _make_batch_scores( - 3, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - for dynamic_batch_size in [1, 4, 8, 12]: - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=dynamic_batch_size, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, dynamic_batch_size) - - -def test_sample_picks_highest_cumprod_rows(): - max_draft_step = 1 - scores, b_req_idx = _make_batch_scores( - 2, - max_draft_step, - rows=[ - [1.0, 0.9], - [1.0, 0.1], - ], - ) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=2, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, 2) - # top-2 scores are both 0.99 at mtp_index==0 (req0 and req1 main rows) - assert select[0].item() == 1 - assert select[2].item() == 1 - - -def test_sample_single_request(): - max_draft_step = 2 - scores, b_req_idx = _make_batch_scores(1, max_draft_step, rows=[[1.0, 0.5, 0.25]]) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=2, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, 2) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py b/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py deleted file mode 100644 index c1ef686f8e..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py +++ /dev/null @@ -1,82 +0,0 @@ -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params - - -def _reference_dynamic_spec_fa3_decode_params(b_req_idx, b_seq_len, b_mark_mtp_shared_group, hold_req_id): - batch_size = b_req_idx.shape[0] - b_req_idx_cpu = b_req_idx.cpu() - b_seq_len_cpu = b_seq_len.cpu() - b_mark_cpu = b_mark_mtp_shared_group.cpu() - valid_idx = torch.nonzero(b_mark_cpu > 0, as_tuple=False).flatten() - valid_size = valid_idx.numel() - - b_q_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - b_kv_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - b_att_req_idx = torch.full((batch_size,), hold_req_id, dtype=torch.int32) - b_att_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - - b_q_seq_len[:valid_size] = b_mark_cpu[valid_idx] - b_kv_seq_len[:valid_size] = b_seq_len_cpu[valid_idx] - b_att_req_idx[:valid_size] = b_req_idx_cpu[valid_idx] - b_att_seq_len[:valid_size] = b_seq_len_cpu[valid_idx] - return b_q_seq_len, b_kv_seq_len, b_att_req_idx, b_att_seq_len - - -@pytest.mark.parametrize("batch_size", [1, 7, 256, 257, 777, 1025, 1537]) -def test_build_dynamic_spec_fa3_decode_params(batch_size): - hold_req_id = -1 - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") + 100 - b_seq_len = (torch.arange(batch_size, dtype=torch.int32, device="cuda") % 97) + 1 - b_mark_mtp_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") - - mark_positions = [0, 3, 17, 255, 256, 511, batch_size - 1] - for pos in mark_positions: - if 0 <= pos < batch_size: - b_mark_mtp_shared_group[pos] = pos % 5 + 1 - - actual = build_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=batch_size, - hold_req_id=hold_req_id, - ) - expected = _reference_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - hold_req_id=hold_req_id, - ) - - for actual_tensor, expected_tensor in zip(actual, expected, strict=True): - assert torch.equal(actual_tensor.cpu(), expected_tensor) - - -def test_build_dynamic_spec_fa3_decode_params_all_padding(): - batch_size = 513 - hold_req_id = -1 - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_seq_len = torch.arange(batch_size, dtype=torch.int32, device="cuda") + 1 - b_mark_mtp_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") - - actual = build_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=batch_size, - hold_req_id=hold_req_id, - ) - expected = _reference_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - hold_req_id=hold_req_id, - ) - - for actual_tensor, expected_tensor in zip(actual, expected, strict=True): - assert torch.equal(actual_tensor.cpu(), expected_tensor) diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py deleted file mode 100644 index 41bc217b94..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.gen_decode_params import gen_decode_params - - -def test_gen_decode_params_basic(): - b_seq_len = torch.ones((9,), dtype=torch.int64, device="cuda") * 8192 - ( - b_q_seq_len, - b1_cu_q_seq_len, - b_kv_seq_len, - b1_cu_kv_seq_len, - position_ids, - ) = gen_decode_params(b_seq_len) - - true_b_q_seq_len = torch.ones_like(b_seq_len) - - assert torch.equal(b_q_seq_len, true_b_q_seq_len) - assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(b_kv_seq_len, b_seq_len) - assert torch.equal(b1_cu_kv_seq_len, torch.nn.functional.pad(torch.cumsum(b_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(position_ids, b_seq_len - 1) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py deleted file mode 100644 index 2702ca9453..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py +++ /dev/null @@ -1,44 +0,0 @@ -import torch -import pytest -import numpy as np -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.gen_mtp_prefill_params import gen_mtp_new_input_ids - - -def test_gen_mtp_new_input_ids_empty_batch(): - input_ids = torch.empty((0,), dtype=torch.int64, device="cuda") - b_next_token_ids = torch.empty((0,), dtype=torch.int64, device="cuda") - b_seq_len = torch.empty((0,), dtype=torch.int32, device="cuda") - - new_input_ids = gen_mtp_new_input_ids( - input_ids=input_ids, - b_next_token_ids=b_next_token_ids, - b_seq_len=b_seq_len, - ) - - assert new_input_ids.is_cuda - assert new_input_ids.dtype == torch.int64 - assert new_input_ids.shape == (0,) - - -def test_gen_mtp_new_input_ids_0(): - input_ids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9]).long().cuda() - b_next_token_ids = torch.tensor([10, 11, 12]).long().cuda() - b_seq_len = torch.tensor([3, 3, 3]).int().cuda() - expected_output = torch.tensor([2, 3, 10, 5, 6, 11, 8, 9, 12]).long().cuda() - new_input_ids = gen_mtp_new_input_ids(input_ids, b_next_token_ids, b_seq_len) - assert torch.equal(new_input_ids, expected_output) - - -def test_gen_mtp_new_input_ids_1(): - input_ids = torch.tensor([1, 2, 3, 4, 5, 6]).long().cuda() - b_next_token_ids = torch.tensor([10, 11, 12]).long().cuda() - b_seq_len = torch.tensor([3, 3, 3]).int().cuda() - b_ready_cache_len = torch.tensor([1, 1, 1]).int().cuda() - expected_output = torch.tensor([2, 10, 4, 11, 6, 12]).long().cuda() - new_input_ids = gen_mtp_new_input_ids(input_ids, b_next_token_ids, b_seq_len, b_ready_cache_len=b_ready_cache_len) - assert torch.equal(new_input_ids, expected_output) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py b/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py deleted file mode 100644 index ba168f4a2f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py +++ /dev/null @@ -1,87 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke, gen_same_req_mark - - -def test_gen_nsa_ks_ke_basic(): - """Test basic functionality of gen_nsa_ks_ke with simple inputs.""" - # Setup test data - b_seq_len = torch.tensor( - [ - 10, - ], - dtype=torch.int32, - device="cuda", - ) - b_q_seq_len = torch.tensor( - [ - 5, - ], - dtype=torch.int32, - device="cuda", - ) - b_req_idx = torch.tensor( - [ - 1, - ], - dtype=torch.int32, - device="cuda", - ) - q_token_num = b_q_seq_len.sum().item() - - req_to_token_index = torch.arange(0, 1000).cuda().view(100, -1) - ragged_mem_index = torch.empty_like(req_to_token_index.view(-1)) - - ks, ke, lengths = gen_nsa_ks_ke( - b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index - ) - - assert torch.equal(ks, torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32, device="cuda")) - assert torch.equal(ke, torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32, device="cuda")) - assert torch.equal(lengths, ke - ks) - assert torch.equal( - ragged_mem_index[0:10], torch.tensor([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=torch.int32, device="cuda") - ) - return - - -def test_gen_nsa_ks_ke_batch(): - b_seq_len = torch.tensor([10, 11], dtype=torch.int32, device="cuda") - b_q_seq_len = torch.tensor([1, 1], dtype=torch.int32, device="cuda") - b_req_idx = torch.tensor([1, 1], dtype=torch.int32, device="cuda") - q_token_num = b_q_seq_len.sum().item() - - req_to_token_index = torch.arange(0, 1000).cuda().view(10, -1) - ragged_mem_index = torch.empty_like(req_to_token_index.view(-1)) - - ks, ke, lengths = gen_nsa_ks_ke( - b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index - ) - - assert torch.equal( - ks, - torch.tensor( - [ - 0, - 0, - ], - dtype=torch.int32, - device="cuda", - ), - ) - assert torch.equal(ke, torch.tensor([10, 11], dtype=torch.int32, device="cuda")) - assert torch.equal(lengths, ke - ks) - assert torch.equal(ragged_mem_index[0:11], torch.arange(100, 100 + 11).cuda()) - - -def test_gen_same_req_mark(): - b_req_idx = torch.tensor([0, 0, 1, 1, 1, 2], dtype=torch.int32, device="cuda") - expected_same_req_mark = torch.tensor([0, 2, 0, 0, 3, 1], dtype=torch.int32, device="cuda") - - same_req_mark = gen_same_req_mark(b_req_idx) - - assert torch.equal(same_req_mark, expected_same_req_mark) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py deleted file mode 100644 index e981b0bb73..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py +++ /dev/null @@ -1,43 +0,0 @@ -import torch -import pytest -import numpy as np -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_prefill_params - - -def test_gen_prefill_params_basic(): - - b_ready_cache_len = torch.ones((9,), dtype=torch.int64, device="cuda") - b_seq_len = torch.ones((9,), dtype=torch.int64, device="cuda") * 8192 - b_seq_len[0:5] = 4567 - - input_token_num = (b_seq_len - b_ready_cache_len).sum().item() - true_b_q_seq_len = b_seq_len - b_ready_cache_len - - ( - b_q_seq_len, - b1_cu_q_seq_len, - b_kv_seq_len, - b1_cu_kv_seq_len, - position_ids, - ) = gen_prefill_params(input_token_num, b_ready_cache_len, b_seq_len) - - assert torch.equal(b_q_seq_len, true_b_q_seq_len) - assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(b_kv_seq_len, b_seq_len) - assert torch.equal(b1_cu_kv_seq_len, torch.nn.functional.pad(torch.cumsum(b_seq_len, dim=0), (1, 0), value=0)) - - b_ready_cache_len_numpy = b_ready_cache_len.cpu().numpy() - b_seq_len_numpy = b_seq_len.cpu().numpy() - true_position_ids = torch.from_numpy( - np.concatenate( - [np.arange(b_ready_cache_len_numpy[i], b_seq_len_numpy[i]) for i in range(len(b_seq_len_numpy))], - axis=0, - ) - ).cuda() - - assert torch.equal(position_ids, true_position_ids) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py deleted file mode 100644 index e9d0193279..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.gen_sampling_params import token_id_counter, gen_sampling_params -from lightllm.common.basemodel.triton_kernel.gen_sampling_params import update_req_to_token_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_token_id_counter(): - test_prompt_ids = torch.randint(0, 100, size=(7853,), device="cuda") - test_token_id_counter = torch.zeros((100,), dtype=torch.int32, device="cuda") - token_id_counter(prompt_ids=test_prompt_ids, out_token_id_counter=test_token_id_counter) - - true_out = torch.bincount(test_prompt_ids, minlength=test_token_id_counter.shape[0]) - assert torch.equal(test_token_id_counter, true_out) - - test_prompt_ids = torch.randint(0, 50 * 10000, size=(128 * 1024,), device="cuda") - test_token_id_counter = torch.zeros((50 * 10000,), device="cuda", dtype=torch.int32) - - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - for _ in range(100): - token_id_counter(prompt_ids=test_prompt_ids, out_token_id_counter=test_token_id_counter) - end_event.record() - end_event.synchronize() - logger.info(f"test_token_id_count cost time: {start_event.elapsed_time(end_event)} ms") - - -def test_gen_sampling_params(): - - # Mocking ReqSamplingParamsManager for testing - class MockReqSamplingParamsManager: - def __init__(self, batch_size, vocab_size): - self.req_to_presence_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_frequency_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_repetition_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_temperature = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_exponential_decay_length_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - - batch_size = 1083 - vocab_size = 100 - req_sampling_params_manager = MockReqSamplingParamsManager(batch_size, vocab_size) - b_req_idx = torch.arange(batch_size, device="cuda").flip(dims=[0]) - - ( - b_presence_penalty, - b_frequency_penalty, - b_repetition_penalty, - b_temperature, - b_exponential_decay_length_penalty, - ) = gen_sampling_params(b_req_idx, req_sampling_params_manager) - - assert torch.equal(b_presence_penalty, req_sampling_params_manager.req_to_presence_penalty[b_req_idx]) - assert torch.equal(b_frequency_penalty, req_sampling_params_manager.req_to_frequency_penalty[b_req_idx]) - assert torch.equal(b_repetition_penalty, req_sampling_params_manager.req_to_repetition_penalty[b_req_idx]) - assert torch.equal(b_temperature, req_sampling_params_manager.req_to_temperature[b_req_idx]) - assert torch.equal( - b_exponential_decay_length_penalty, - req_sampling_params_manager.req_to_exponential_decay_length_penalty[b_req_idx], - ) - - -def test_update_req_to_token_id_counter(): - b_req_idx = torch.tensor([0, 1, 3, 2], dtype=torch.int32, device="cuda") - next_token_ids = torch.tensor([0, 1, 1, 0], dtype=torch.int32, device="cuda") - req_to_out_token_id_counter = torch.zeros((4, 4), dtype=torch.int32, device="cuda") - update_req_to_token_id_counter(b_req_idx, next_token_ids, req_to_out_token_id_counter) - expected_output = torch.tensor( - [ - [1, 0, 0, 0], - [0, 1, 0, 0], - [1, 0, 0, 0], - [0, 1, 0, 0], - ], - dtype=torch.int32, - device="cuda", - ) - - assert torch.equal(req_to_out_token_id_counter, expected_output) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py b/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py deleted file mode 100644 index 9ecfa008c4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py +++ /dev/null @@ -1,107 +0,0 @@ -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel import mtp_utils - - -@pytest.mark.parametrize( - "max_group_size,req_indexes,expected_markers", - [ - (2, [7, 7, 7, 11, 11, 20], [0, 2, 1, 0, 2, 1]), - (8, [7, 7, 7, 11, 11, 11, 20, 20, 20], [0, 0, 3, 0, 0, 3, 0, 0, 3]), - (3, [7, 7, 7, 7, 7, 7], [0, 0, 3, 0, 0, 3]), - (3, [7, 7, -1, -1, 11, 11], [0, 2, 1, 1, 0, 2]), - ], -) -def test_build_mtp_shared_group_markers(monkeypatch, max_group_size, req_indexes, expected_markers): - monkeypatch.setattr(mtp_utils, "get_diverse_max_batch_shared_group_size", lambda: max_group_size) - b_req_idx = torch.tensor(req_indexes, dtype=torch.int32, device="cuda") - - markers = mtp_utils.build_mtp_shared_group_markers(b_req_idx, hold_req_id=-1) - - assert markers.cpu().tolist() == expected_markers - - -@pytest.mark.parametrize("pinned", [False, True]) -def test_mtp_verify_scatter_and_start_locations(pinned): - req_to_next_token_ids = torch.tensor( - [[1, 2, -2, -1, -1], [1, 2, 0, -1, -1], [1, 3, 4, 4, 5]], - dtype=torch.int64, - device="cpu" if pinned else "cuda", - pin_memory=pinned, - ) - b_req_idx = torch.tensor([0, 0, 2, 2, 2], dtype=torch.int32, device="cuda") - b_mtp_index = torch.tensor([0, 1, 0, 1, 2], dtype=torch.int32, device="cuda") - b_req_mtp_start_loc = mtp_utils.gen_b_req_mtp_start_loc(b_mtp_index, num_reqs=2) - new_next_token_ids = torch.tensor([1, 4, 2, 4, 13], dtype=torch.int64, device="cuda") - draft_token_ids = torch.tensor( - [[2, 3], [8, 9]], - dtype=torch.int64, - device="cuda", - ) - req_to_next_token_scores = torch.full((3, 5), -1.0, dtype=torch.float32, device="cuda") - schedule_scores = torch.tensor( - [[0.9, 0.8], [0.7, 0.6]], - dtype=torch.float32, - device="cuda", - ) - - mtp_accept_len, accepted_index = mtp_utils.mtp_verify( - req_to_next_token_ids, b_req_mtp_start_loc, new_next_token_ids, b_req_idx - ) - assert mtp_accept_len.is_cuda and accepted_index.is_cuda - mtp_utils.mtp_scatter_next_token_ids( - req_to_next_token_ids=req_to_next_token_ids, - b_req_mtp_start_loc=b_req_mtp_start_loc, - target_next_token_ids=new_next_token_ids, - draft_token_ids=draft_token_ids, - b_req_idx=b_req_idx, - mtp_accept_len=mtp_accept_len, - req_to_next_token_scores=req_to_next_token_scores, - schedule_scores=schedule_scores, - ) - torch.cuda.synchronize() - - assert torch.equal(b_req_mtp_start_loc.cpu(), torch.tensor([0, 2], dtype=torch.int64)) - assert torch.equal(mtp_accept_len.cpu(), torch.tensor([1, 1], dtype=torch.int32)) - assert torch.equal(accepted_index.cpu(), torch.tensor([1, 0, 1, 0, 0], dtype=torch.int32)) - assert torch.equal( - req_to_next_token_ids.cpu(), - torch.tensor( - [[1, 2, 3, 1, 1], [1, 2, 0, -1, -1], [2, 8, 9, 1, 1]], - dtype=torch.int64, - ), - ) - torch.testing.assert_close( - req_to_next_token_scores.cpu(), - torch.tensor( - [[1.0, 0.9, 0.8, 0.0, 0.0], [-1.0, -1.0, -1.0, -1.0, -1.0], [1.0, 0.7, 0.6, 0.0, 0.0]], - dtype=torch.float32, - ), - ) - - -@pytest.mark.parametrize("pinned", [False, True]) -def test_mtp_scatter_handles_zero_draft_step(pinned): - req_to_next_token_ids = torch.full( - (1, 4), -1, dtype=torch.int64, device="cpu" if pinned else "cuda", pin_memory=pinned - ) - req_to_next_token_scores = torch.full((1, 4), -1.0, dtype=torch.float32, device="cuda") - - mtp_utils.mtp_scatter_next_token_ids( - req_to_next_token_ids=req_to_next_token_ids, - b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - target_next_token_ids=torch.tensor([42], dtype=torch.int64, device="cuda"), - draft_token_ids=torch.empty((1, 0), dtype=torch.int64, device="cuda"), - b_req_idx=torch.tensor([0], dtype=torch.int32, device="cuda"), - mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), - req_to_next_token_scores=req_to_next_token_scores, - schedule_scores=torch.empty((1, 0), dtype=torch.float32, device="cuda"), - ) - torch.cuda.synchronize() - - assert req_to_next_token_ids.cpu().tolist() == [[42, 1, 1, 1]] - assert req_to_next_token_scores.cpu().tolist() == [[1.0, 0.0, 0.0, 0.0]] diff --git a/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py b/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py deleted file mode 100755 index ae3bcf4aed..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py +++ /dev/null @@ -1,54 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.multimodal_emb import mark_multimodal_obj, multimodal_emb -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_mark_multimodal_obj(): - virtual_token_base = 2 ** 62 - obj_start_ids = torch.tensor( - [virtual_token_base, virtual_token_base + 3, virtual_token_base + 99], device="cuda", dtype=torch.int64 - ) - obj_token_lens = torch.tensor([1, 3, 2], device="cuda", dtype=torch.int64) - input_ids = torch.tensor([virtual_token_base, virtual_token_base + 6, 9, 333], device="cuda", dtype=torch.int64) - - mark_obj = mark_multimodal_obj( - obj_start_token_ids=obj_start_ids, obj_token_lens=obj_token_lens, input_ids=input_ids - ) - - assert torch.equal(mark_obj, torch.tensor([1, 0, 0], device="cuda")) - - -def test_multimodal_emb(): - S, D = 1024 * 10, 128 * 64 - vob_size = 3200 - image_size = 10 - image_token_size = 512 - virtual_token_base = 2 ** 62 - - text_weight = torch.randn((vob_size, D), device="cuda", dtype=torch.float16) - img_weight = torch.randn((image_size * image_token_size, 1, D), device="cuda", dtype=torch.float16) - img_token_lens = torch.full((image_size,), image_token_size, device="cuda", dtype=torch.long) - img_start_token_ids = ( - torch.arange(0, image_size * image_token_size, image_token_size, dtype=torch.int64, device="cuda") - + virtual_token_base - ) - img_start_locs = torch.arange(0, image_size * image_token_size, image_token_size).cuda().long() - - prompt_ids = torch.arange(0, S, 1).cuda().long() - prompt_ids[0 : image_size * image_token_size] = virtual_token_base + torch.arange( - 0, image_size * image_token_size, 1, dtype=torch.int64, device="cuda" - ) - - out = torch.zeros((S, D), dtype=torch.float16, device="cuda") - multimodal_emb( - out, prompt_ids, text_weight, img_weight, img_token_lens, img_start_token_ids, img_start_locs, 0, vob_size, 1 - ) - assert torch.allclose(out[0 : image_size * image_token_size, :].view(-1), img_weight.view(-1), atol=0.0001) - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py b/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py deleted file mode 100644 index 1e7d5080ca..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py +++ /dev/null @@ -1,64 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.norm.qk_norm import ( - qk_rmsnorm_fused_forward, - qk_rmsnorm_forward, -) - - -def torch_qk_rmsnorm(q, w_q, eps=1e-6): - input_dtype = q.dtype - head_dim = w_q.shape[0] - q_fp32 = q.to(torch.float32) - q_fp32 = q_fp32.view(-1, head_dim) - variance = q_fp32.pow(2).mean(dim=-1, keepdim=True) - q_fp32 = q_fp32 * torch.rsqrt(variance + eps) - return (q_fp32 * w_q.to(input_dtype)).view_as(q) - - -def test_qk_rmsnorm_fused_matches_reference(): - """Compare fused QK RMSNorm with separate reference RMSNorm kernels.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for qk_rmsnorm_fused test") - - torch.manual_seed(0) - - # 模拟配置: Batch=2, Seq=128, Head_Dim=128 - # Q: 16 Heads, K: 4 Heads (GQA 场景) - B, S, D = 2, 128, 128 - H_Q = 16 - H_K = 4 - - q = torch.randn((B * S, H_Q * D), device="cuda", dtype=torch.bfloat16) - k = torch.randn((B * S, H_K * D), device="cuda", dtype=torch.bfloat16) - - w_q = torch.ones((D,), device="cuda", dtype=torch.bfloat16) - w_k = torch.ones((D,), device="cuda", dtype=torch.bfloat16) - - # 复制一份做对比（reference 会在新 tensor 上计算） - q_ref = q.clone() - k_ref = k.clone() - - # fused kernel in-place 计算 - q_out, k_out = qk_rmsnorm_fused_forward(q, k, w_q, w_k, eps=1e-6) - - # reference: 分别对 Q / K 做 RMSNorm - q_ref_out = qk_rmsnorm_forward(q_ref, w_q, eps=1e-6) - k_ref_out = qk_rmsnorm_forward(k_ref, w_k, eps=1e-6) - # q_ref_out = torch_qk_rmsnorm(q_ref, w_q, eps=1e-6) - # k_ref_out = torch_qk_rmsnorm(k_ref, w_k, eps=1e-6) - - # fused 是 in-place 的，返回的 q_out/k_out 应该与 q/k 引用一致 - assert q_out.data_ptr() == q.data_ptr() - assert k_out.data_ptr() == k.data_ptr() - - # 误差容忍度: 由于 bfloat16 计算，设定一个合理的 atol - q_max_diff = (q_out - q_ref_out).abs().max().item() - k_max_diff = (k_out - k_ref_out).abs().max().item() - - print(f"Q max diff: {q_max_diff}") - print(f"K max diff: {k_max_diff}") - - assert q_max_diff < 1e-5, f"Q max diff too large: {q_max_diff}" - assert k_max_diff < 1e-5, f"K max diff too large: {k_max_diff}" diff --git a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py b/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py deleted file mode 100644 index 16131ef935..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py +++ /dev/null @@ -1,151 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import expert_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_redundancy_topk_ids_repair(): - ep_expert_num = 4 - global_rank = 0 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 0, - ], - dtype=torch.int64, - device="cuda", - ) - - expert_id_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - expert_counter=expert_id_counter, - enable_counter=True, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - assert torch.equal( - expert_id_counter, torch.tensor([1, 2, 1, 2, 0, 1, 0, 2, 0, 1, 1, 1], dtype=torch.int64, device="cuda") - ) - - ep_expert_num = 4 - global_rank = 1 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 5, - ], - dtype=torch.int64, - device="cuda", - ) - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - - -def test_expert_id_counter(): - token_num = 256 - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - - ans_expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - ids, counts = torch.unique(tok_ids.view(-1), return_counts=True) - ans_expert_counter[ids] = counts - - assert torch.equal(expert_counter, ans_expert_counter) - - # test speed - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for _ in range(100): - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - graph.replay() - - start_event = torch.cuda.Event(enable_timing=True) - start_event.record() - graph.replay() - end_event = torch.cuda.Event(enable_timing=True) - end_event.record() - torch.cuda.synchronize() - logger.info(f"expert_id_counter time cost: {start_event.elapsed_time(end_event)} ms") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py deleted file mode 100644 index 15a4447466..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py +++ /dev/null @@ -1,50 +0,0 @@ -import torch -import pytest -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, max_seq_len", - [(a, b) for a in [1, 16, 32, 128, 512] for b in [16, 32, 512, 2048]], -) -def test_repack_kv_index(batch, max_seq_len): - def repack_kv_ref(req_to_token_indexs, b_req_idx, b_seq_len, b_start_loc, output): - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - output[start : start + sl] = req_to_token_indexs[b][:sl] - - BATCH, MAX_SEQ_LEN = batch, max_seq_len - rand_idx = torch.randperm(2 * MAX_SEQ_LEN * BATCH).cuda().int() - b_req_idx = torch.randperm(BATCH).cuda().int() - b_seq_len = torch.randint(1, MAX_SEQ_LEN, (BATCH,)).cuda().int() - req_to_token_indexs = torch.zeros((2 * BATCH, 2 * MAX_SEQ_LEN)).cuda().int() - b_start_loc = ( - torch.cat([torch.zeros([1], device=b_seq_len.device, dtype=b_seq_len.dtype), b_seq_len[0:-1].cumsum(0)]) - .cuda() - .int() - ) - - output = torch.zeros((b_seq_len.sum(),)).cuda().int() - ref = torch.zeros((b_seq_len.sum(),)).cuda().int() - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - req_to_token_indexs[b][:sl] = rand_idx[start : start + sl] - - repack_kv_ref(req_to_token_indexs, b_req_idx, b_seq_len, b_start_loc, ref) - repack_kv_index( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=b_req_idx, - b_token_len=b_seq_len, - b_page_start_loc=b_start_loc, - max_token_len=MAX_SEQ_LEN, - out_page_indices=output, - ) - assert torch.allclose(output.float(), ref.float()) diff --git a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py b/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py deleted file mode 100644 index 80cc75ee33..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py +++ /dev/null @@ -1,137 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.routing_capture import scatter_routing_topk_to_cpu - - -def _skip_without_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for Triton kernels.") - - -@pytest.mark.parametrize("dtype,dtype_id,max_value", [(torch.uint8, 1, 255), (torch.int16, 2, 1024)]) -def test_scatter_routing_topk_to_pinned_cpu(dtype, dtype_id, max_value): - _skip_without_cuda() - - num_tokens = 5 - num_moe_layers = 3 - topk = 4 - kv_cache_size = 32 - moe_layer_index = 1 - - topk_ids = torch.randint(0, max_value, (num_tokens, topk), dtype=torch.int64, device="cuda") - mem_indexes = torch.tensor([17, 3, 29, 8, 11], dtype=torch.int32, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=dtype, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=moe_layer_index, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=dtype_id, - ) - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu().long(), moe_layer_index, :] = topk_ids.cpu().to(dtype) - assert torch.equal(routing_buffer, expected) - - -def test_scatter_routing_topk_respects_layer_index(): - _skip_without_cuda() - - num_tokens = 3 - num_moe_layers = 2 - topk = 2 - kv_cache_size = 16 - - topk_ids = torch.arange(num_tokens * topk, dtype=torch.int64, device="cuda").view(num_tokens, topk) - mem_indexes = torch.tensor([10, 4, 13], dtype=torch.int64, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=torch.int16, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=1, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=2, - ) - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu(), 1, :] = topk_ids.cpu().to(torch.int16) - assert torch.equal(routing_buffer, expected) - - -def test_scatter_routing_topk_is_cuda_graph_capturable(): - _skip_without_cuda() - - num_tokens = 4 - num_moe_layers = 2 - topk = 3 - kv_cache_size = 16 - - topk_ids = torch.arange(num_tokens * topk, dtype=torch.int64, device="cuda").view(num_tokens, topk) - mem_indexes = torch.tensor([2, 4, 6, 8], dtype=torch.int32, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=torch.uint8, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=0, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=1, - ) - torch.cuda.synchronize() - routing_buffer.zero_() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=0, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=1, - ) - - graph.replay() - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu(), 0, :] = topk_ids.cpu().to(torch.uint8) - assert torch.equal(routing_buffer, expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py b/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py deleted file mode 100644 index 713a1c20c4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py +++ /dev/null @@ -1,146 +0,0 @@ -from collections import Counter -from types import SimpleNamespace - -import pytest -import torch -from triton.runtime.jit import JITFunction - -from lightllm.common.basemodel.triton_kernel import gather_token_id as token_ops -from lightllm.common.basemodel.triton_kernel import gen_sampling_params as sampling_ops -from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import ( - apply_penalty_gpu_cache, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _observe(monkeypatch, kernels): - misses = Counter() - hashes = {kernel: set() for kernel in kernels} - original_compile = JITFunction._do_compile - - def compile_kernel(self, *args, **kwargs): - if self in hashes: - misses[self] += 1 - return original_compile(self, *args, **kwargs) - - monkeypatch.setattr(JITFunction, "_do_compile", compile_kernel) - for kernel in kernels: - original_run = kernel.run - - def run(*args, _kernel=kernel, _original=original_run, **kwargs): - compiled = _original(*args, **kwargs) - hashes[_kernel].add(compiled.hash) - return compiled - - monkeypatch.setattr(kernel, "run", run) - return misses, hashes - - -def _assert_reused(misses, hashes): - assert all(len(values) == 1 for values in hashes.values()) - assert all(misses[kernel] <= 1 for kernel in hashes) - - -@pytest.mark.parametrize("has_mask", [False, True]) -@pytest.mark.parametrize("width", [1, 3, 16, 17]) -def test_token_helpers_reuse_runtime_counts_with_fixed_strides(monkeypatch, has_mask, width): - # MTP width and vocabulary are fixed for each service. Preserve their stride - # specializations while varying the request count across integer classes. - kernels = [ - token_ops._fwd_kernel_scatter, - token_ops._fwd_kernel_gather, - token_ops._fwd_kernel_gather_prefill_decode_mixed, - sampling_ops._token_id_counter_update_kernel, - ] - misses, hashes = _observe(monkeypatch, kernels) - table = torch.zeros((128, width), dtype=torch.int64, device="cuda") - counters = torch.zeros((128, width), dtype=torch.int32, device="cuda") - for batch in [1, 3, 16, 17, 33, 65]: - requests = torch.arange(batch, dtype=torch.int32, device="cuda") - columns = requests % width - values = requests.long() + 100 - enabled = requests % 3 != 1 - mask = enabled if has_mask else None - table.zero_() - token_ops.scatter_token(values, table, requests, columns, mask) - actual = token_ops.gather_token(table, requests, columns) - expected = torch.where(enabled, values, 0) if has_mask else values - assert torch.equal(actual, expected) - - input_ids = torch.full((batch * 2,), -1, dtype=torch.int64, device="cuda") - is_decode = requests % 2 == 0 - starts = requests * 2 - expected_input = input_ids.clone() - expected_input[starts[is_decode]] = expected[is_decode] - token_ops.gather_token_prefill_decode_mixed(input_ids, table, requests, columns, is_decode, starts) - assert torch.equal(input_ids, expected_input) - - counters.zero_() - owners = requests % 7 - sampling_ops.update_req_to_token_id_counter(owners, columns, counters, mask) - included = enabled if has_mask else torch.ones_like(enabled) - flat = owners[included].long() * width + columns[included] - expected_counts = torch.bincount(flat, minlength=counters.numel()).reshape_as(counters) - assert torch.equal(counters, expected_counts) - _assert_reused(misses, hashes) - - -@pytest.mark.parametrize("vocab", [1, 3, 16, 17]) -def test_sampling_parameter_and_prompt_counts_reuse_compilation(monkeypatch, vocab): - kernels = [sampling_ops._gen_sampling_params_kernel, sampling_ops._token_id_counter_kernel] - misses, hashes = _observe(monkeypatch, kernels) - fields = [ - "presence_penalty", - "frequency_penalty", - "repetition_penalty", - "temperature", - "exponential_decay_length_penalty", - ] - manager = SimpleNamespace( - **{ - f"req_to_{field}": torch.arange(128, device="cuda", dtype=torch.float32) + index - for index, field in enumerate(fields) - } - ) - for count in [1, 3, 16, 17, 65]: - requests = torch.arange(count, device="cuda", dtype=torch.int32).flip(0) - actual = sampling_ops.gen_sampling_params(requests, manager) - for value, field in zip(actual, fields): - assert torch.equal(value, getattr(manager, f"req_to_{field}")[requests]) - tokens = torch.arange(count, device="cuda") % vocab - counters = torch.zeros(vocab, dtype=torch.int32, device="cuda") - sampling_ops.token_id_counter(tokens, counters) - assert torch.equal(counters, torch.bincount(tokens, minlength=vocab)) - _assert_reused(misses, hashes) - - -def test_penalties_match_reference_across_shapes(): - # Preserve the specialized stride/EOS variants selected by performance - # measurements; their outputs must still agree across shapes. - for batch, vocab, eos_count in [(1, 1, 0), (3, 3, 1), (16, 16, 3), (17, 17, 17), (33, 3, 3)]: - requests = torch.arange(batch, dtype=torch.int32, device="cuda") - counters = (torch.arange(128 * vocab, device="cuda") % 3).int().reshape(128, vocab) - manager = SimpleNamespace( - vocab_size=vocab, - req_to_out_token_id_counter=counters, - req_to_presence_penalty=torch.full((128,), 0.5, device="cuda"), - req_to_frequency_penalty=torch.full((128,), 0.25, device="cuda"), - req_to_repetition_penalty=torch.full((128,), 2.0, device="cuda"), - req_to_exponential_decay_length_penalty=torch.full((128,), 2.0, device="cuda"), - ) - logits = ((torch.arange(batch * vocab, device="cuda") % 9 - 4).float() / 4).reshape(batch, vocab) - counts = counters[requests] - repeated = torch.where(logits > 0, logits / 2, logits * 2) - expected = torch.where(counts > 0, repeated, logits) - counts * 0.25 - expected -= torch.where(counts > 0, 0.5, 0.0) - length = requests % 3 - mask_eos = requests % 4 == 2 - eos_ids = torch.arange(eos_count, device="cuda") % vocab - scales = torch.exp2(length.float()) - 1 - for eos in eos_ids.tolist(): - value = expected[:, eos] - expected[:, eos] = torch.where(mask_eos, -10000000.0, value + value.abs() * scales) - apply_penalty_gpu_cache(logits, requests, length, mask_eos, eos_ids, manager) - torch.testing.assert_close(logits, expected, atol=0, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py b/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py deleted file mode 100644 index 0dab062934..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py +++ /dev/null @@ -1,75 +0,0 @@ -import pytest -import torch -from types import SimpleNamespace - -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( - select_kv_index_from_req, - select_kv_index_from_req_prefill, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def test_select_kv_index_from_req_for_decode(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 6], dtype=torch.int32, device="cuda"), - ) - - assert mem_indexes.cpu().tolist() == [22, 15] - - -def test_select_kv_index_from_req_for_prefill(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req_prefill( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([4, 7], dtype=torch.int32, device="cuda"), - b_ready_cache_len=torch.tensor([2, 3], dtype=torch.int32, device="cuda"), - b_start_loc=torch.tensor([0, 2], dtype=torch.int32, device="cuda"), - max_q_seq_len=4, - token_num=6, - ) - - assert mem_indexes.cpu().tolist() == [22, 23, 13, 14, 15, 16] - - -def test_model_selects_reserved_indexes_when_execution_starts(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(page_size=4) - model.req_manager = SimpleNamespace( - req_to_token_indexs=torch.tensor( - [[10, 11, 12, 13], [20, 21, 22, 23]], - dtype=torch.int32, - device="cuda", - ) - ) - model_input = SimpleNamespace( - is_prefill=False, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4], dtype=torch.int32, device="cuda"), - ) - - mem_indexes = model._select_mem_indexes(model_input) - - assert mem_indexes.cpu().tolist() == [22, 13] diff --git a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py b/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py deleted file mode 100644 index 772644c28f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py +++ /dev/null @@ -1,36 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "token_num, hidden_dim, sp_world_size", - [ - (token_num, hidden_dim, sp_world_size) - for token_num in range(3, 6) - for hidden_dim in [257, 2048] - for sp_world_size in range(2, 5) - ], -) -def test_sp_pad_copy(token_num, hidden_dim, sp_world_size): - in_tensor = torch.randn((token_num, hidden_dim), dtype=torch.float16, device="cuda") - if token_num % sp_world_size: - # TPSP pads the model input before splitting it between ranks. - with pytest.raises(AssertionError, match="in_token_num % sp_world_size != 0"): - sp_pad_copy(in_tensor=in_tensor, sp_rank_id=0, sp_world_size=sp_world_size) - padding = in_tensor[: sp_world_size - token_num % sp_world_size] - in_tensor = torch.cat([in_tensor, padding]) - out_tensors = [ - sp_pad_copy(in_tensor=in_tensor, sp_rank_id=rank_id, sp_world_size=sp_world_size) - for rank_id in range(sp_world_size) - ] - out_tensor = torch.cat(out_tensors, dim=0) - assert torch.equal(in_tensor, out_tensor) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py index b39478bd12..ed4175596a 100644 --- a/unit_tests/common/dsv4_test_utils.py +++ b/unit_tests/common/dsv4_test_utils.py @@ -245,10 +245,6 @@ def free(self, indices: torch.Tensor): def get_cell_size(self): return sum(state.get_token_bytes() for state in self.full_states.values()) - def get_att_input_params(self, layer: int): - index = self.cache_config.layer_to_compressed_kv_idx.get(layer) - return self.kv_buffer[index] if index is not None else None - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): for state in self.full_states.values(): state.copy_full_att_tokens(src_mem_index, dst_mem_index) @@ -282,12 +278,6 @@ def alloc_paged_kv_move_buffer(self, count: int, page_size: int): ) return self.kv_move_buffer - def get_pd_kv_move_buffer_size(self): - return self.manager.get_pd_operator().get_page_bytes() * get_env_start_args().pd_kv_page_num - - def get_paged_kv_move_buffer_shape(self, count: int, page_size: int): - return (count, self.manager.get_pd_operator().get_page_bytes()) - def write_mem_to_page_kv_move_buffer( self, indices: torch.Tensor, diff --git a/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py b/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py deleted file mode 100644 index 8967ded513..0000000000 --- a/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py +++ /dev/null @@ -1,61 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( - _ep_compact_metadata_kernel, - _ep_gather_chunk_kernel, - ep_compact_metadata, - ep_gather_chunk, -) - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("aligned_chunks", [False, True]) -def test_ep_runtime_counts_reuse_compilation(monkeypatch, aligned_chunks): - torch.manual_seed(51) - # Production chunks are 128-aligned. Generic unaligned calls remain valid, - # but need not share their boundary specialization with other chunks. - chunks = [(0, 128), (128, 128), (32768, 128)] if aligned_chunks else [(0, 1), (1, 3), (3, 16)] - weight_count = max(start + count for start, count in chunks) + 1 - hashes = {"gather": {}, "compact": set()} - bounds = None - for name, kernel in [("gather", _ep_gather_chunk_kernel), ("compact", _ep_compact_metadata_kernel)]: - original = kernel.run - - def observe(*args, _name=name, _original=original, **kwargs): - compiled = _original(*args, **kwargs) - if _name == "gather": - hashes[_name].setdefault(bounds, set()).add(compiled.hash) - else: - hashes[_name].add(compiled.hash) - return compiled - - monkeypatch.setattr(kernel, "run", observe) - for tokens in [1, 3, 16, 17, 33, 129]: - metadata = torch.randint(-1, 32, (tokens, 10), dtype=torch.int32, device="cuda") - weights = torch.randn(weight_count, dtype=torch.float32, device="cuda") - for start, count in chunks: - bounds = (start, count) - metadata[:, 2:] = torch.randint(start, start + count, (tokens, 8), device="cuda", dtype=torch.int32) - metadata[:, 2] = start - 1 - metadata[:, 3] = start - metadata[:, 4] = start + count - 1 - metadata[:, 5] = start + count - chunk = torch.randn((count, 129), dtype=torch.bfloat16, device="cuda") - output = torch.zeros((tokens, 129), dtype=chunk.dtype, device="cuda") - slots = metadata[:, 2:].long() - valid = (slots >= start) & (slots < start + count) - selected = chunk[(slots - start).clamp(0, count - 1)].float() - reference = (selected * weights[slots.clamp(0, weight_count - 1), None] * valid[..., None]).sum(1) - ep_gather_chunk(chunk, start, weights, metadata, output) - torch.testing.assert_close(output.float(), reference, atol=0.04, rtol=0.01) - original_header = metadata[:, :2].clone() - ep_compact_metadata(metadata) - assert torch.equal(metadata[:, :2], original_header) - assert torch.equal(metadata[:, 2], torch.arange(tokens, device="cuda", dtype=torch.int32)) - assert (metadata[:, 3:] == -1).all() - assert len(hashes["compact"]) == 1 - assert all(len(values) == 1 for values in hashes["gather"].values()), hashes - if aligned_chunks: - assert len(set.union(*hashes["gather"].values())) == 1 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe.py b/unit_tests/common/fused_moe/test_grouped_fused_moe.py index 4ecb7ed81c..8bfa703886 100644 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe.py +++ b/unit_tests/common/fused_moe/test_grouped_fused_moe.py @@ -5,9 +5,7 @@ from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import ( _moe_align_fused_atomic_token, fused_experts_impl, - moe_align, moe_align_fused, - moe_align1, moe_align2, grouped_matmul, ) @@ -23,60 +21,6 @@ torch.cuda.manual_seed_all(seed) -def test_moe_align(): - expert_num = 5 - token_num = 3 - topk = 3 - topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4]], dtype=torch.int32, device="cuda") - out = torch.zeros((expert_num, token_num * topk), dtype=torch.int32, device="cuda") - out.fill_(0) - moe_align(topk_ids, out) - true = torch.tensor( - [ - [1, 0, 0, 1, 0, 0, 0, 0, 0], - [0, 1, 0, 0, 0, 1, 0, 1, 0], - [0, 0, 1, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 1, 0, 1, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 1], - ], - dtype=torch.int32, - device="cuda", - ) - assert torch.equal(out, true) - - -def test_moe_align1(): - experts_info = torch.tensor( - [[1, 0, 0, 0], [0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0]], - dtype=torch.int32, - device="cuda", - ) - topk_weights = torch.tensor([[0.3, 0.7], [0.2, 0.8]], dtype=torch.float32, device="cuda") - experts_token_num = torch.zeros((4,), dtype=torch.int32, device="cuda") - experts_weights = torch.zeros(experts_info.shape, dtype=torch.float32, device="cuda") - - moe_align1(experts_info, topk_weights, experts_weights, experts_token_num, 2) - - true_experts_token_num = torch.tensor([1, 2, 1, 0], device="cuda", dtype=torch.int32) - true_experts_info = torch.tensor( - [[0, 0, 0, 0], [1, 2, 1, 0], [3, 0, 0, 1], [0, 0, 0, 0]], device="cuda:0", dtype=torch.int32 - ) - true_experts_weights = torch.tensor( - [ - [0.3000, 0.0000, 0.0000, 0.0000], - [0.7000, 0.2000, 0.0000, 0.0000], - [0.8000, 0.0000, 0.0000, 0.0000], - [0.0000, 0.0000, 0.0000, 0.0000], - ], - device="cuda", - dtype=torch.float32, - ) - - assert torch.allclose(true_experts_weights, experts_weights) - assert torch.equal(experts_token_num, true_experts_token_num) - assert torch.equal(experts_info, true_experts_info) - - def _check_moe_align_fused(topk_ids, topk_weights, expert_num, ordered=True): expert_to_token_index = torch.empty((expert_num, topk_ids.numel()), dtype=torch.int32, device="cuda") expert_to_weight = torch.empty((expert_num, topk_ids.numel()), dtype=torch.float32, device="cuda") diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py b/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py deleted file mode 100644 index 35b6ce7104..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py +++ /dev/null @@ -1,11 +0,0 @@ -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common import get_ep_num_sms - - -def test_get_ep_num_sms_selects_overlap_quota(monkeypatch): - import lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common as common - - monkeypatch.setattr(common.dist_group_manager, "ep_num_sms", 20, raising=False) - monkeypatch.setattr(common.dist_group_manager, "ep_non_overlap_num_sms", 64, raising=False) - - assert get_ep_num_sms(overlap_with_compute=True) == 20 - assert get_ep_num_sms(overlap_with_compute=False) == 64 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py b/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py deleted file mode 100644 index 769002517d..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py +++ /dev/null @@ -1,51 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import moe_align, moe_align1, grouped_matmul -from lightllm.utils.log_utils import init_logger - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - -logger = init_logger(__name__) - - -@pytest.mark.parametrize("token_num", [200, 256, 8 * 1024]) -def test_moe_align1(token_num): - expert_num = 160 - topk_num = 6 - print(token_num) - - def get_one(): - rnd_logics = torch.randn(token_num, expert_num, device="cuda") - topk_values, topk_ids = torch.topk(rnd_logics, topk_num, dim=1) - - experts_info = torch.zeros((expert_num, token_num * topk_num), dtype=torch.int32, device="cuda") - experts_info.fill_(0) - moe_align(topk_ids, experts_info) - - topk_weights = torch.randn((token_num, topk_num), dtype=torch.float32, device="cuda") - experts_token_num = torch.zeros((expert_num,), dtype=torch.int32, device="cuda") - experts_weights = torch.zeros(experts_info.shape, dtype=torch.float32, device="cuda") - return experts_info, topk_weights, experts_weights, experts_token_num - - test_datas = [get_one() for _ in range(100)] - - moe_align1(*test_datas[0], topk_num) - - torch.cuda.synchronize() - start = time.time() - - for i in range(60): - moe_align1(*test_datas[i + 1], topk_num) - torch.cuda.synchronize() - - print(f"token_num: {token_num} cost time: {time.time() - start} s") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py deleted file mode 100644 index b569ec8429..0000000000 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ /dev/null @@ -1,196 +0,0 @@ -import torch -import pytest - - -def is_fp8_native_supported(): - """检查是否为 H100/B200 等原生支持 FP8 的硬件 (SM90+)""" - if not torch.cuda.is_available(): - return False - major, _ = torch.cuda.get_device_capability() - return major >= 9 - - -if not is_fp8_native_supported(): - pytest.skip(reason="not support fp8 test in this gpu card", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( - silu_and_mul_masked_post_quant_fwd, - silu_and_mul_psum_post_quant_fwd, -) -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "expert_num, token_num, hidden_dim", - [ - ( - expert_num, - token_num, - hidden_dim, - ) - for expert_num in range(3, 6) - for hidden_dim in [256, 128 * 4, 2048] - for token_num in range(1, 7, 2) - ], -) -def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): - quant_group_size = 128 - in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") - out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") - out_scale_tensor = torch.randn( - (expert_num, token_num, hidden_dim // 2 // quant_group_size), dtype=torch.float32, device="cuda" - ) - - true_out_tensor_mid = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=in_tensor.dtype, device="cuda") - - masked_m = torch.full((expert_num,), token_num, dtype=torch.int32, device="cuda") - - silu_and_mul_fwd(in_tensor.view(-1, hidden_dim), true_out_tensor_mid.view(-1, hidden_dim // 2)) - true_out_tensor, true_out_scale_tensor = per_token_group_quant_fp8( - true_out_tensor_mid.view(-1, hidden_dim // 2), - quant_group_size, - alloc_func=torch.empty, - ) - - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) - - true_out_tensor = true_out_tensor.view(out_tensor.shape) - true_out_scale_tensor = true_out_scale_tensor.view(out_scale_tensor.shape) - - hidden_dim_scale_count = hidden_dim // 2 // quant_group_size - for expert_id, expert_token_num in enumerate(masked_m.cpu().numpy()): - true_scale = true_out_scale_tensor[expert_id, :expert_token_num, :hidden_dim_scale_count] - out_scale = out_scale_tensor[expert_id, :expert_token_num, :hidden_dim_scale_count] - assert torch.allclose( - true_scale, - out_scale, - atol=1e-3, - rtol=1e-2, - ) - - true_out = true_out_tensor[expert_id, :expert_token_num, :].to(torch.float32) - out = out_tensor[expert_id, :expert_token_num, :].to(torch.float32) - true_dequant = true_out * true_scale.repeat_interleave(quant_group_size, dim=-1) - out_dequant = out * out_scale.repeat_interleave(quant_group_size, dim=-1) - assert torch.allclose(true_dequant, out_dequant, atol=1e-1, rtol=1e-1) - return - - -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -def test_silu_and_mul_masked_skips_padded_tokens(monkeypatch, use_ue8m0_scales): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(use_ue8m0_scales))) - expert_num = 3 - token_num = 4 - hidden_dim = 256 - quant_group_size = 128 - masked_m = torch.tensor([0, 2, token_num], dtype=torch.int32, device="cuda") - - in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") - out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") - out_scale_tensor = torch.empty( - (expert_num, token_num, hidden_dim // 2 // quant_group_size), dtype=torch.float32, device="cuda" - ) - out_tensor.fill_(1.0) - out_scale_tensor.fill_(7.0) - - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) - torch.cuda.synchronize() - - for expert_id, expert_token_num in enumerate(masked_m.cpu().tolist()): - assert torch.equal( - out_tensor[expert_id, expert_token_num:, :], - torch.ones_like(out_tensor[expert_id, expert_token_num:, :]), - ) - assert torch.equal( - out_scale_tensor[expert_id, expert_token_num:, :], - torch.full_like(out_scale_tensor[expert_id, expert_token_num:, :], 7.0), - ) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -@pytest.mark.parametrize("alpha,limit,add_one", [(1.0, 10.0, False), (1.702, 7.0, True)]) -def test_clamped_swiglu_quantization_matches_reference(monkeypatch, dtype, use_ue8m0_scales, alpha, limit, add_one): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(use_ue8m0_scales))) - torch.manual_seed(53) - x = (torch.randn(3, 5, 4096, device="cuda") * 16).to(dtype) - x[1, 0].zero_() - masked_m = torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda") - out = torch.full((3, 5, 2048), 1.0, dtype=torch.float8_e4m3fn, device="cuda") - scales = torch.full((3, 5, 16), 7.0, dtype=torch.float32, device="cuda") - - silu_and_mul_masked_post_quant_fwd( - x, - out, - scales, - 128, - masked_m, - alpha=alpha, - limit=limit, - clamp_up_add_one=add_one, - ) - - gate, up = x.float().chunk(2, dim=-1) - gate = gate.clamp(max=limit) - gate = (gate * torch.sigmoid(alpha * gate)).to(dtype).float() - activation = (gate * (up.clamp(-limit, limit) + int(add_one))).to(dtype).float() - groups = activation.reshape(3, 5, 16, 128) - expected_scales = groups.abs().amax(-1).clamp(min=1e-10) / 448 - if use_ue8m0_scales: - expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales))) - expected_q = (groups / expected_scales.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) - expected = (expected_q.float() * expected_scales.unsqueeze(-1)).reshape(3, 5, 2048) - actual = out.float() * scales.repeat_interleave(128, dim=-1) - for expert, count in enumerate(masked_m.tolist()): - torch.testing.assert_close(scales[expert, :count], expected_scales[expert, :count], atol=1e-6, rtol=1e-5) - torch.testing.assert_close(actual[expert, :count], expected[expert, :count], atol=0.05, rtol=0.02) - assert torch.all(out[expert, count:].float() == 1) - assert torch.all(scales[expert, count:] == 7) - - -def test_silu_and_mul_psum_skips_alignment_padding(): - hidden = 128 - alignment = 8 - # expert0 使用 3 行，expert1 从对齐后的第 8 行开始再使用 2 行。 - expert_token_psum = torch.tensor([3, 10], device="cuda", dtype=torch.int32) - rows = 16 - gate = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) - up = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) - packed = torch.cat([gate, up], dim=1).contiguous() - output = torch.zeros((rows, hidden), device="cuda", dtype=torch.float8_e4m3fn) - scales = torch.zeros((rows, hidden // 128), device="cuda", dtype=torch.float32) - - silu_and_mul_psum_post_quant_fwd( - input=packed, - output=output, - output_scale=scales, - expert_token_psum=expert_token_psum, - expert_alignment=alignment, - quant_group_size=128, - ) - - def reference(row): - # Match the kernel: SiLU in fp32, then round the gate back to bf16 before multiplying. - activated_gate = torch.nn.functional.silu(gate[row].float()).to(torch.bfloat16) - activated = activated_gate.float() * up[row].float() - scale = activated.abs().amax().clamp(min=1e-10) / torch.finfo(torch.float8_e4m3fn).max - quantized = (activated / scale).clamp(-448, 448).to(torch.float8_e4m3fn) - return quantized, scale - - for row in (0, 1, 2, 8, 9): - expected_q, expected_scale = reference(row) - torch.testing.assert_close(scales[row], expected_scale.view(1), atol=1e-4, rtol=1e-3) - actual = output[row].float() * scales[row] - expected = expected_q.float() * expected_scale - # FP8 E4M3 bins are coarse, so compare the dequantized activation within one bin. - torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.15) - # expert0 与 expert1 之间的对齐 padding 不应被写成有效值。 - assert torch.count_nonzero(output[3:8].float()) == 0 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_softmax_topk.py b/unit_tests/common/fused_moe/test_softmax_topk.py index 7c3e483df8..a4050f236b 100755 --- a/unit_tests/common/fused_moe/test_softmax_topk.py +++ b/unit_tests/common/fused_moe/test_softmax_topk.py @@ -1,118 +1,20 @@ -import torch -import time import pytest -import numpy as np -from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def benchmark(M, N, K, renorm, runs): - try: - import sgl_kernel as sgl_ops - except Exception as e: - pytest.skip(f"no sgl_kernel error: {str(e)}", allow_module_level=True) - - gating = torch.randn(M, N, device="cuda", dtype=torch.float32) - torch.cuda.synchronize() - - # 1. SGL kernel - sgl_vals = torch.empty((M, K), dtype=torch.float32, device="cuda") - sgl_ids = torch.empty((M, K), dtype=torch.int32, device="cuda") - # Warm-up - sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating) - torch.cuda.synchronize() - start = torch.cuda.Event(True) - end = torch.cuda.Event(True) - start.record() - for _ in range(runs): - sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating) - if renorm: - sgl_vals.div_(sgl_vals.sum(-1, keepdim=True).clamp_min(1e-8)) - - end.record() - torch.cuda.synchronize() - t_sgl = start.elapsed_time(end) / runs - - # 2. Triton kernel - t0 = torch.cuda.Event(True) - t1 = torch.cuda.Event(True) - # Warm-up - softmax_topk(gating, K) - torch.cuda.synchronize() - t0.record() - for _ in range(runs): - triton_vals, triton_ids = softmax_topk(gating, K, renorm) - t1.record() - torch.cuda.synchronize() - t_triton = t0.elapsed_time(t1) / runs - - # 3. Native PyTorch - _ = torch.softmax(gating, dim=-1) - _, _ = torch.topk(_, K, dim=-1) - torch.cuda.synchronize() - - start, end = torch.cuda.Event(True), torch.cuda.Event(True) - start.record() - for _ in range(runs): - probs = torch.softmax(gating, dim=-1) - torch_vals, torch_ids = torch.topk(probs, K, dim=-1) - if renorm: - torch_vals.div_(torch_vals.sum(-1, keepdim=True).clamp_min(1e-8)) - end.record() - torch.cuda.synchronize() - t_torch = start.elapsed_time(end) / runs - - # Compare indices and weights - # Count mismatches of ordered indices - diff_sgl_triton_ids = (sgl_ids != triton_ids).sum().item() - diff_torch_triton_ids = (torch_ids != triton_ids).sum().item() - # Max absolute difference of weights aligned by position - max_err_triton_torch = (triton_vals - torch_vals).abs().max().item() - max_err_triton_torch_sgl = (sgl_vals - torch_vals).abs().max().item() - max_err_triton_sgl = (triton_vals - sgl_vals).abs().max().item() - - assert diff_sgl_triton_ids == 0, f"Mismatch SGL vs Triton ids: {diff_sgl_triton_ids}" - assert diff_torch_triton_ids == 0, f"Mismatch Torch vs Triton ids: {diff_torch_triton_ids}" - assert max_err_triton_torch < 1e-3, f"Max err Triton vs Torch: {max_err_triton_torch}" - assert max_err_triton_torch_sgl < 1e-3, f"Max err Triton vs SGL: {max_err_triton_torch_sgl}" - assert max_err_triton_sgl < 1e-3, f"Max err Torch vs SGL: {max_err_triton_sgl}" +import torch - results = { - "time_sgl": t_sgl, - "time_triton": t_triton, - "time_torch": t_torch, - "mismatch_sgl_triton_ids": diff_sgl_triton_ids, - "mismatch_torch_triton_ids": diff_torch_triton_ids, - "max_err_triton_torch": max_err_triton_torch, - "max_err_triton_sgl": max_err_triton_sgl, - "max_err_triton_torch_sgl": max_err_triton_torch_sgl, - "sgl_ids": sgl_ids, - "triton_ids": triton_ids, - "torch_ids": torch_ids, - "sgl_vals": sgl_vals, - "triton_vals": triton_vals, - "torch_vals": torch_vals, - } - return results +from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -def test_softmax_topk(): - M, N, K = 8192, 1024, 8 - res = benchmark(M, N, K, False, 1000) - print(f"SGL time: {res['time_sgl']:.6f}ms") - print(f"Triton time: {res['time_triton']:.6f}ms") - print(f"PyTorch time: {res['time_torch']:.6f}ms") - print("Mismatch SGL vs Triton ids:", res["mismatch_sgl_triton_ids"]) - print("Mismatch Torch vs Triton ids:", res["mismatch_torch_triton_ids"]) - print("Max err Triton vs Torch :", res["max_err_triton_torch"]) - print("Max err Triton vs SGL :", res["max_err_triton_sgl"]) - print("Max err Torch vs SGL :", res["max_err_triton_torch_sgl"]) - benchmark(M, N, K, True, 10) - benchmark(M, 256, 5, True, 10) - benchmark(M, 127, 5, True, 10) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("experts,topk", [(1024, 8), (256, 5), (127, 5)]) +@pytest.mark.parametrize("renorm", [False, True]) +def test_softmax_topk(experts, topk, renorm): + torch.manual_seed(42) + gating = torch.randn(32, experts, device="cuda") + expected_weights, expected_ids = torch.softmax(gating, dim=-1).topk(topk, dim=-1) + if renorm: + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + weights, ids = softmax_topk(gating, topk, renorm) -if __name__ == "__main__": - pytest.main() + torch.testing.assert_close(ids.long(), expected_ids, rtol=0, atol=0) + torch.testing.assert_close(weights, expected_weights, rtol=0, atol=1e-3) diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index f5bd92e08a..06de22959d 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -86,36 +86,6 @@ def cache(monkeypatch, tmp_path, request): get_env_start_args.cache_clear() -def test_attention_input_params_select_token_buffers(cache): - manager, requests = cache - for layer_index, slot in ((0, 0), (1, 1), (2, None)): - swa, compressed_kv = requests.get_att_input_params(layer_index) - assert swa.data_ptr() == requests.req_to_swa.flatten(1, 2)[layer_index].data_ptr() - if slot is None: - assert compressed_kv is None - else: - token_buffer = manager.get_att_input_params(layer_index) - assert token_buffer is manager.kv_buffer[slot] - assert token_buffer.shape[:2] == (manager.size + manager.page_size, 1) - assert compressed_kv.untyped_storage().data_ptr() == token_buffer.data_ptr() - - -def test_big_page_pool_reserves_cpu_cache_staging_slots(cache): - manager, _ = cache - pool = manager.big_page_buffers - args = get_env_start_args() - page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - capacity = (manager.size + page_tokens - 1) // page_tokens - assert pool.size == capacity + 2 and pool.keep_num == 2 - assert manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID == capacity - assert manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID == capacity + 1 - slots = pool.alloc_state_cache(capacity) - assert slots == list(range(capacity)) - assert pool.alloc_one_state_cache() is None - pool.free_state_cache(slots) - assert pool.get_free_cache_num() == capacity - - def test_copy_mem_to_mem_preserves_complete_block_indices(cache): manager, _ = cache # 压缩块内的 KV/scale/indexer 不可拆分；允许按非连续、乱序的完整块复制。 @@ -160,47 +130,6 @@ def test_cache_io_copies_individual_token_slots(cache): assert buffer[destination[2:]].eq(0).all() -def test_request_runtime_is_independent_of_compressed_kv_capacity(cache, monkeypatch): - original, requests = cache - runtime = requests.window - before = tuple( - tensor.data_ptr() - for tensor in ( - runtime.swa_buffer, - runtime.c4_buffer, - runtime.indexer_buffer, - runtime.c128_buffer, - ) - ) - # 扩展压缩 KV 容量不能重新创建或移动按请求数分配的运行态。 - for state in original.full_states.values(): - state.allocate_storage(token_capacity=16384) - assert state.storage.shape[1] == 16384 + state.page_size - assert before == tuple( - tensor.data_ptr() - for tensor in ( - runtime.swa_buffer, - runtime.c4_buffer, - runtime.indexer_buffer, - runtime.c128_buffer, - ) - ) - assert runtime.request_capacity == requests.max_request_num + 1 - - -@pytest.mark.parametrize("page_size", [256, 1024, 4096]) -def test_pd_memory_reservation_matches_allocated_bytes(cache, monkeypatch, page_size): - manager, _ = cache - args = get_env_start_args() - monkeypatch.setattr(args, "run_mode", "decode") - monkeypatch.setattr(args, "pd_kv_page_num", 3, raising=False) - monkeypatch.setattr(args, "pd_kv_page_size", page_size, raising=False) - reserved = manager.get_pd_kv_move_buffer_size() - buffer = manager.alloc_paged_kv_move_buffer(3, page_size) - assert reserved == buffer.nbytes - assert buffer.shape == manager.get_paged_kv_move_buffer_shape(3, page_size) - - def _packed_swa(buffer): return buffer.view(buffer.shape[0], -1, 74880) if buffer.shape[-1] == 585 else buffer @@ -373,173 +302,6 @@ def _reserve(manager, requests, length): return req_idx, slots -@pytest.mark.parametrize("ratio", [4, 128]) -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_noncontiguous_pages_obey_logical_compression_boundaries(cache, ratio): - manager, requests = cache - page_size = manager.page_size - req_idx, allocated = _reserve(manager, requests, 3 * page_size) - slots = allocated.view(3, page_size)[torch.tensor([2, 0, 1], device="cuda")].reshape(-1) - requests.req_to_token_indexs[req_idx, : slots.numel()] = slots - positions = torch.tensor( - [0, 2, 3, 126, 127, 254, 255, page_size - 1, page_size, 2 * page_size - 1, 2 * page_size], device="cuda" - ) - indexes = torch.empty((len(positions), slots.numel() // ratio), dtype=torch.int32, device="cuda") - lengths = torch.empty_like(positions, dtype=torch.int32) - build_compressed_kv_indices( - torch.full_like(positions, req_idx), positions, requests.req_to_token_indexs, ratio, indexes, lengths - ) - group_slots = slots[ratio - 1 :: ratio] // ratio - for row, position in enumerate(positions.tolist()): - count = (position + 1) // ratio - torch.testing.assert_close(indexes[row, :count], group_slots[:count]) - assert indexes[row, count:].eq(-1).all() - assert lengths[row].item() == max(1, count) - manager.free(allocated) - assert manager.allocator.can_use_mem_size == manager.size - - -def test_indexer_writer_only_writes_closed_groups(cache): - manager, requests = cache - _, slots = _reserve(manager, requests, 512) - slots = slots.view(2, 256).flip(0).reshape(-1) - keys = torch.randn((259, 128), dtype=torch.bfloat16, device="cuda") - positions = torch.arange(259, dtype=torch.int32, device="cuda") - _compressed_kv_pools(manager)[1].zero_() - from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k import destindex_copy_indexer_k - - destindex_copy_indexer_k(keys, slots[:259], positions, manager.full_states[0].get_layer_view(0).indexer_k, 64) - closed = torch.arange(3, 259, 4, device="cuda") - actual = _read_packed_rows(_compressed_kv_pools(manager)[1], 0, slots[closed] // 4) - # Reference quantization is independent of the cache writer and stays in the test. - selected_keys = keys[closed].float() - scale = selected_keys.abs().amax(dim=-1, keepdim=True).clamp_min(1e-4) / 448.0 - quantized = (selected_keys / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) - expected = torch.cat((quantized.view(torch.uint8), scale.view(torch.uint8)), dim=-1) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - # Positions 256..258 cannot publish their reserved next compressed slot. - assert _read_packed_rows(_compressed_kv_pools(manager)[1], 0, slots[256:257] // 4).eq(0).all() - - -def test_speculative_retry_reuses_swa_without_freeing_token_pages(cache): - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - - def cpu(data): - return torch.tensor(data, dtype=torch.int32) - - _prepare_prefill(requests, cpu([req_idx]), cpu([0]), cpu([254])) - for sequences in ([255, 256, 257, 258], [255, 256, 257, 258], [256, 257, 258, 259]): - swa_slots = _swa_slots(requests, req_idx, torch.tensor(sequences, device="cuda") - 1) - assert swa_slots.unique().numel() == len(sequences) - assert manager.allocator.can_use_mem_size == manager.size - 512 - requests.free([req_idx], slots) - assert manager.allocator.can_use_mem_size == manager.size - - -@pytest.mark.parametrize("model_kind", ["target", "mtp", "dspark"]) -def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): - from lightllm.common.basemodel.batch_objs import ModelInput - from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel - - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - _prepare_swa(requests, req_idx, 0, 256) - model_cls = {"target": DeepseekV4TpPartModel, "mtp": DeepseekV4MTPModel, "dspark": DeepseekV4DSparkModel}[ - model_kind - ] - model = model_cls.__new__(model_cls) - model.req_manager = requests - model.att_state_manager = requests.memory.manager - model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] - model.tp_rank_ = 0 - model.mtp_previous_draft_models = [] - model.__dict__["att_state_role"] = "main" - model_input = ModelInput( - batch_size=1, - total_token_num=257, - max_q_seq_len=1, - max_kv_seq_len=257, - b_req_idx=torch.tensor([req_idx], dtype=torch.int32), - b_seq_len=torch.tensor([257], dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_position_delta=torch.zeros(1, dtype=torch.int32), - b_shared_seq_len=torch.zeros(1, dtype=torch.int32), - b_shared_radix_node_id=torch.full((1,), -1, dtype=torch.int64), - multimodal_params=[{"images": [], "audios": []}], - ) - model_input.to_cuda() - for _ in range(2): - torch.testing.assert_close(model._select_mem_indexes(model_input), slots[256:257]) - assert manager.allocator.can_use_mem_size == manager.size - 512 - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("cache", [(256, 0), (256, 5)], indirect=True) -def test_prefill_chunk_keeps_all_new_swa_rows(cache): - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 2048) - _prepare_prefill(requests, torch.tensor([req_idx]), torch.tensor([0]), torch.tensor([2048])) - positions = torch.arange(2048, device="cuda") - swa_slots = _swa_slots(requests, req_idx, positions) - assert swa_slots.unique().numel() == 2048 - packed = torch.randint(0, 256, (2048, 584), dtype=torch.uint8, device="cuda") - _write_swa(requests, 0, req_idx, positions, packed) - # Checkpoints can still read any boundary of the complete chunk here. - torch.testing.assert_close(_read_packed_rows(_swa_buffer(requests, req_idx), 0, swa_slots), packed) - _finish_prefill(requests, [req_idx]) - tail = positions[-requests.sliding_window :] - torch.testing.assert_close(_read_swa(requests, 0, req_idx, tail), packed[tail]) - requests.free([req_idx], slots) - - -def test_prefill_checkpoint_reads_precede_window_retention(cache): - """The chunk-end checkpoint and late draft KV must survive a long prefill.""" - manager, requests = cache - req, slots = _reserve(manager, requests, 2048) - _prepare_swa(requests, req, 0, 2048) - positions = torch.arange(2048, device="cuda") - packed = torch.randint(0, 256, (2048, 584), dtype=torch.uint8, device="cuda") - _write_swa(requests, 0, req, positions, packed) - _state_buffer(requests, "c4").uniform_() - _state_buffer(requests, "c4_indexer").uniform_() - live_tail = requests.req_to_c4_state.clone() - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - requests.save_state(_checkpoint_req(req, 2048), index, states) - window, _, _ = states.get_state_cache(index) - torch.testing.assert_close(window[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) - for boundary in range(256, 2049, 256): - torch.testing.assert_close( - _read_swa(requests, 0, req, positions[boundary - requests.sliding_window : boundary]), - packed[boundary - requests.sliding_window : boundary], - rtol=0, - atol=0, - ) - - # A DSpark context commit writes KV after the target model has returned. - _write_swa(requests, 2, req, positions, packed.flip(0)) - _finish_prefill(requests, [req]) - # Reuse the scratch: only fixed request windows and CPU snapshots remain. - requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model)[0].zero_() - requests.save_state(_checkpoint_req(req, 2048), index, states) - saved_swa = states.get_state_cache(index)[0] - torch.testing.assert_close(saved_swa[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) - tail = positions[-requests.sliding_window :] - for layer, expected in ((0, packed), (2, packed.flip(0))): - torch.testing.assert_close( - _read_swa(requests, layer, req, tail), - expected[tail], - rtol=0, - atol=0, - ) - torch.testing.assert_close(requests.req_to_c4_state, live_tail) - requests.free([req], slots) - states.free_state_cache([index]) - - def test_non_aligned_chunks_preserve_window_and_checkpoint_history(cache): manager, requests = cache req, slots = _reserve(manager, requests, 3072) @@ -564,51 +326,6 @@ def test_non_aligned_chunks_preserve_window_and_checkpoint_history(cache): requests.free([req], slots) -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_compressed_kv_copy_preserves_packed_data_and_scale_regions(cache): - manager, requests = cache - source, src = _reserve(manager, requests, 2 * manager.page_size) - destination, dst = _reserve(manager, requests, 2 * manager.page_size) - requests.req_to_swa.random_(0, 256) - swa_state = requests.req_to_swa.clone() - src = src.view(2, manager.page_size).flip(0).reshape(-1) - for pool in _compressed_kv_pools(manager): - pool.random_(0, 256) - manager.operator.copy_mem_to_mem(src_mem_index=src, dst_mem_index=dst) - for pool in _compressed_kv_pools(manager): - torch.testing.assert_close(pool[:, src[::256].long() // 256], pool[:, dst[::256].long() // 256], rtol=0, atol=0) - torch.testing.assert_close(requests.req_to_swa, swa_state) - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_hold_page_is_separate_from_request_rings(cache): - manager, requests = cache - hold = requests.req_to_token_indexs[requests.HOLD_REQUEST_ID] - torch.testing.assert_close( - hold[: manager.page_size], - torch.arange(manager.size, manager.size + manager.page_size, dtype=torch.int32, device="cuda"), - ) - assert ( - _swa_blocks(requests, requests.HOLD_REQUEST_ID, torch.arange(32, device="cuda")) - .eq(requests.HOLD_SWA_SLOT // 128) - .all() - ) - assert (hold // 4).lt(_compressed_kv_pools(manager)[0].shape[1] * 64).all() - assert (hold // 128).lt(_compressed_kv_pools(manager)[2].shape[1] * 2).all() - free_tokens = manager.allocator.can_use_mem_size - reqs = [requests.alloc(), requests.alloc()] - assert set(_request_pages(requests, reqs[0])).isdisjoint(_request_pages(requests, reqs[1])) - assert ( - _swa_blocks(requests, torch.tensor(reqs, device="cuda")[:, None], torch.arange(32, device="cuda")[None, :]) - .lt(requests.HOLD_SWA_SLOT // 128) - .all() - ) - assert manager.allocator.can_use_mem_size == free_tokens - for req in reqs: - requests.free_req(req) - - @pytest.mark.parametrize("length", [256, 512, 2048]) def test_shared_compressed_kv_restores_private_continuation(cache, length): manager, requests = cache @@ -714,1348 +431,3 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): prefill.free([source], src) requests.free([destination], dst) assert manager.allocator.can_use_mem_size == manager.size - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") -@pytest.mark.parametrize("length", [257, 2051]) -def test_pd_transport_completion_covers_peer_gpu_state(cache, length): - import ctypes - - transport, _ = cache - with torch.cuda.device(1): - peer = DeepseekV4MemoryManager(8192, torch.bfloat16, 1, 512, 3, transport.cache_config) - requests = DeepseekV4ReqManager(2, 4096, None, peer.cache_config, mtp_step=get_env_start_args().mtp_step) - requests.bind_mem_manager(peer) - source, src = _reserve(peer, requests, length) - destination, dst = _reserve(peer, requests, length) - _randomize_compressed_kv(peer) - requests.req_to_swa.flatten(1, 2).random_(0, 256) - for buffer in ( - _state_buffer(requests, "c4"), - _state_buffer(requests, "c4_indexer"), - _state_buffer(requests, "c128"), - ): - buffer.uniform_() - for name in ( - "req_to_swa", - "req_to_c4_state", - "req_to_c4_indexer_state", - "req_to_c128_state", - ): - setattr(peer, name, getattr(requests, name)) - torch.cuda.synchronize() - peer_stream = torch.cuda.Stream() - consumer_stream = torch.cuda.Stream() - copy_stream = torch.cuda.Stream() - # CUDA IPC enables peer access in production. A same-process test needs - # to establish that access before Triton dereferences the peer pointer. - cudart = ctypes.CDLL(f"libcudart.so.{torch.version.cuda.split('.')[0]}") - assert cudart.cudaDeviceEnablePeerAccess(1, 0) in (0, 704) - cudart.cudaGetLastError() - transport.state_token_len = peer.state_token_len = length - transport.alloc_paged_kv_move_buffer(1, transport.page_size) - mem_managers = [transport, peer] - for start in range(0, length, transport.page_size): - transport.write_mem_to_page_kv_move_buffer( - src[start : start + transport.page_size].tolist(), 0, 1, mem_managers, 1 - ) - torch.cuda.current_stream().synchronize() # The network consumes the completed packed page. - transport.read_page_kv_move_buffer_to_mem( - dst[start : start + transport.page_size].tolist(), 0, 1, mem_managers, 1 - ) - torch.cuda.current_stream().synchronize() # Return the page to the pool only after completion. - with torch.cuda.stream(peer_stream), torch.cuda.stream(copy_stream): - transport.write_mem_to_page_kv_move_buffer([], 0, 1, mem_managers, 1, "att_state", source) - with torch.cuda.device(1): - torch.cuda._sleep(20_000_000) - transport.read_page_kv_move_buffer_to_mem([], 0, 1, mem_managers, 1, "att_state", destination) - completed = torch.cuda.Event() - completed.record(copy_stream) - completed.synchronize() - # 在独立流验证逻辑窗口；无效 MTP 槽位不属于 PD 数据。 - with torch.cuda.stream(consumer_stream): - positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda:1") - for layer in range(requests.cache_config.layer_num): - torch.testing.assert_close( - _read_swa(requests, layer, destination, positions, prefill=False), - _read_swa(requests, layer, source, positions, prefill=False), - rtol=0, - atol=0, - ) - for buffer, width in ( - (requests.req_to_c4_state, 8), - (requests.req_to_c4_indexer_state, 8), - (requests.req_to_c128_state, 128), - ): - positions = torch.arange(max(0, length - width), length, device="cuda:1") % buffer.shape[2] - torch.testing.assert_close(buffer[:, destination, positions], buffer[:, source, positions], rtol=0, atol=0) - for buffer in peer.kv_buffer: - torch.testing.assert_close(buffer[src.long()], buffer[dst.long()], rtol=0, atol=0) - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize("ratio,layer,is_indexer", [(4, 0, False), (128, 1, False), (4, 0, True)]) -@pytest.mark.parametrize("length", [512, 2048]) -@pytest.mark.parametrize("cache", [(256, 0), (256, 1), (256, 5)], indirect=True) -def test_compressor_checkpoint_continuation_matches_full_prefill(cache, ratio, layer, is_indexer, length): - from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - - manager, requests = cache - workspace = DeepseekV4Workspace(SimpleNamespace(tp_world_size_=1, batch_max_tokens=length)) - workspace.c128_compress_scratch = torch.empty((1, length, 512), dtype=torch.float32, device="cuda") - source, src = _reserve(manager, requests, length) - destination, dst = _reserve(manager, requests, length) - _prepare_swa(requests, source, 0, length) - head_dim = requests.cache_config.indexer_head_dim if is_indexer else requests.cache_config.head_dim - scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), dtype=torch.float32, device="cuda") - cos = torch.ones((length, 32), dtype=torch.float32, device="cuda") - sin = torch.zeros_like(cos) - weight = torch.ones(head_dim, dtype=torch.float32, device="cuda") - - def compress(req_idx, slots, start, end, decode=False): - positions = torch.arange(start, end, dtype=torch.int32, device="cuda") - width = end - start - state = SimpleNamespace( - dsv4_workspace=workspace, - microbatch_index=0, - mem_manager=manager, - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - mem_index=slots[start:end], - position_ids=positions, - is_prefill=not decode, - _dsv4_token_to_batch_idx=None if decode else torch.zeros(width, dtype=torch.int32, device="cuda"), - b_req_idx=torch.full((width if decode else 1,), req_idx, dtype=torch.int32, device="cuda"), - b_mtp_index=torch.arange(width, dtype=torch.int32, device="cuda") - if decode - else torch.zeros(1, dtype=torch.int32, device="cuda"), - b_seq_len=positions + 1 if decode else torch.tensor([end], dtype=torch.int32, device="cuda"), - b_ready_cache_len=None if decode else torch.tensor([start], dtype=torch.int32, device="cuda"), - b_q_start_loc=None if decode else torch.zeros(1, dtype=torch.int32, device="cuda"), - dsv4_swa_write_slots=_swa_slots(requests, req_idx, positions), - dsv4_swa_runtime_write_slots=torch.where( - positions >= end - requests.swa_capacity_tokens, - _swa_slots(requests, req_idx, positions, prefill=False), - -1, - ) - if not decode - else None, - ) - output = fused_compress( - kv_score=scores[start:end], - infer_state=state, - layer_idx=layer, - norm_weight=weight, - eps=1e-6, - head_dim=head_dim, - qk_rope_head_dim=64, - compress_ratio=ratio, - cos_table=cos, - sin_table=sin, - is_in_indexer=is_indexer, - ) - if is_indexer and req_idx == destination: - closed_groups = (positions + 1) % ratio == 0 - torch.testing.assert_close(output[closed_groups], reference[start:end][closed_groups], rtol=0, atol=0) - return output - - reference = compress(source, src, 0, length) - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - # A short final chunk must preserve the previous rows needed by the checkpoint tail. - checkpoint_len = length // 2 - for start, end in ((0, checkpoint_len - 2), (checkpoint_len - 2, checkpoint_len)): - _prepare_swa(requests, source, start, end) - compress(source, src, start, end) - requests.save_state(_checkpoint_req(source, checkpoint_len), index, states) - requests.restore_state(SimpleNamespace(req_idx=destination, cur_kv_len=checkpoint_len), states, index) - verify_width = requests.c4_state_tokens - 8 + 1 - for start in range(checkpoint_len, length - verify_width, 256): - end = min(start + 256, length - verify_width) - _prepare_swa(requests, destination, start, end) - compress(destination, dst, start, end) - _finish_prefill(requests, [destination]) - # Rejected speculative rows are written again in the same held token page. - compress(destination, dst, length - verify_width, length, decode=True) - compress(destination, dst, length - verify_width, length, decode=True) - if not is_indexer: - pool = _compressed_kv_pools(manager)[0] if ratio == 4 else _compressed_kv_pools(manager)[2] - torch.testing.assert_close( - _read_packed_rows(pool, 0, src[checkpoint_len + ratio - 1 : length : ratio].long() // ratio), - _read_packed_rows(pool, 0, dst[checkpoint_len + ratio - 1 : length : ratio].long() // ratio), - rtol=0, - atol=0, - ) - requests.free([source, destination], torch.cat([src, dst])) - states.free_state_cache([index]) - - -def test_swa_index_cuda_graph_reads_updated_private_pages(cache): - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - - manager, requests = cache - source, src = _reserve(manager, requests, 512) - destination, dst = _reserve(manager, requests, 512) - _prepare_prefill(requests, torch.tensor([source, destination]), torch.zeros(2), torch.full((2,), 512)) - req_ids = torch.tensor([source, requests.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda") - positions = torch.tensor([255, 1], dtype=torch.int32, device="cuda") - indexes = torch.empty((2, 128), dtype=torch.int32, device="cuda") - lengths = torch.empty(2, dtype=torch.int32, device="cuda") - write_slots = torch.empty_like(lengths) - - def build(): - build_swa_indices( - req_ids, - positions, - requests, - indexes, - lengths, - write_slots, - prefill_state=requests._test_prefill_states[source], - ) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - build() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - build() - req_ids[0] = destination - positions[0] = 256 - graph.replay() - expected = _swa_slots(requests, destination, torch.arange(256, 128, -1, dtype=torch.int32, device="cuda")) - torch.testing.assert_close(indexes[0], expected) - assert write_slots[0].item() == expected[0].item() - assert write_slots[1].item() == requests._test_prefill_states[source].HOLD_SWA_SLOT + 1 - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize( - "dspark,cache", - [(False, (256, 0)), (False, (256, 1)), (False, (256, 5)), (True, (256, 1)), (True, (256, 5))], - indirect=["cache"], -) -def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, dspark): - from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - from lightllm.models.deepseek_v4.triton_kernel.build_dspark_swa_indices import build_dspark_swa_indices - - torch.manual_seed(1593) - manager, requests = cache - source, src = _reserve(manager, requests, 4096) - destination, dst = _reserve(manager, requests, 4096) - _prepare_swa(requests, destination, 0, 512) - # The reference retains every position in the other forward workspace. - _prepare_swa(requests, source, 0, 4096, microbatch_index=1) - values = torch.randn((4096, 512), dtype=torch.bfloat16, device="cuda") * 0.1 - packed = torch.zeros((4096, 584), dtype=torch.uint8, device="cuda") - packed[:, :448] = values[:, :448].to(torch.float8_e4m3fn).view(torch.uint8) - packed[:, 448:576] = values[:, 448:].contiguous().view(torch.uint8) - packed[:, 576:583] = 127 - positions = torch.arange(4096, device="cuda") - _write_swa_rows(requests, 2, source, positions, packed) - _write_swa(requests, 2, destination, positions[:512], packed[:512]) - _finish_prefill(requests, [destination]) - - mtp_step = requests.c4_state_tokens - 8 - token_num = mtp_step if dspark else mtp_step + 1 - width = 256 if dspark else 128 - req_idx = torch.full((token_num,), destination, dtype=torch.int32, device="cuda") - query_positions = torch.arange(512, 512 + token_num, dtype=torch.int32, device="cuda") - indices = torch.empty((token_num, width), dtype=torch.int32, device="cuda") - lengths = torch.empty(token_num, dtype=torch.int32, device="cuda") - write_slots = torch.empty_like(lengths) - - def build(): - if dspark: - build_dspark_swa_indices( - token_req_idx=req_idx, - position_ids=query_positions, - swa_state=requests.swa_state, - swa_indices=indices, - swa_lengths=lengths, - swa_write_slots=write_slots, - block_size=token_num, - ) - else: - build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) - - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend( - type( - "Model", - (), - { - "config": dict( - num_attention_heads=64, - compress_ratios=[0, 0, 0], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - sliding_window=128, - ), - "layers_infer": [SimpleNamespace(layer_num_=2)], - "tp_rank_": 0, - "tp_world_size_": 1, - "req_manager": requests, - }, - )() - ) - state_info = SimpleNamespace(mem_manager=manager) - state = backend.create_att_decode_state(state_info) - state.init_state() - params = dict( - compress_ratio=0, - head_dim_v=512, - softmax_scale=512 ** -0.5, - swa_indices=indices.unsqueeze(1), - swa_lengths=lengths, - attn_sink=torch.zeros(64, device="cuda"), - ) - control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) - q = torch.randn((token_num, 64, 512), dtype=torch.bfloat16, device="cuda") * 0.1 - packed_kv = requests.get_att_input_params(2) - - def forward(): - build() - return state.decode_att(q, packed_kv, None, att_control=control) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - forward() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - actual = forward() - - reference_state = backend.create_att_decode_state(state_info) - reference_state.init_state() - reference_kv = backend.get_att_input_params( - SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - is_prefill=True, - prefill_att_state=requests._test_prefill_states[source], - ), - 2, - ) - committed_end = 512 - # Both sides of page boundaries, repeated retries, and several complete ring wraps. - for start in (512, 513, 638, 639, 640, 640, 641, 767, 768, 1023, 1279, 1280, 1919, 1920, 2047): - # Advance intervening committed KV in small steps. A large parallel - # prefill would instead borrow pages before writing the whole chunk. - for gap_start in range(committed_end, start, 128): - gap = positions[gap_start : min(start, gap_start + 128)] - _write_swa_rows(requests, 2, source, gap, packed[gap]) - _write_swa_rows(requests, 2, destination, gap, packed[gap]) - query_positions.copy_(torch.arange(start, start + token_num, dtype=torch.int32)) - current = packed[start : start + token_num].clone() - current[:, :448] = ( - (values[start : start + token_num, :448] * (1 + start % 7)).to(torch.float8_e4m3fn).view(torch.uint8) - ) - logical_positions = query_positions.long() - _write_swa_rows(requests, 2, source, logical_positions, current) - _write_swa_rows(requests, 2, destination, logical_positions, current) - history = torch.arange(max(0, start - 128), start, device="cuda") - graph.replay() - if dspark: - visible = torch.cat((history.flip(0), logical_positions)) - reference_indices = _swa_slots(requests, source, visible).expand(token_num, -1).contiguous() - reference_lengths = torch.full_like(lengths, visible.numel()) - else: - visible = logical_positions[:, None] - torch.arange(128, device="cuda") - reference_indices = _swa_slots(requests, source, visible) - reference_lengths = torch.full_like(lengths, 128) - reference_indices = torch.nn.functional.pad( - reference_indices, (0, width - reference_indices.shape[1]), value=-1 - ) - reference_indices = reference_indices.to(torch.int32) - reference_params = dict(params, swa_indices=reference_indices.unsqueeze(1), swa_lengths=reference_lengths) - expected = reference_state.decode_att( - q, reference_kv, None, att_control=attention.AttControl(nsa_decode=True, nsa_decode_dict=reference_params) - ) - torch.testing.assert_close(actual, expected, rtol=0, atol=0, msg=f"SWA mismatch at position {start}") - torch.testing.assert_close(write_slots, _swa_slots(requests, destination, logical_positions).to(torch.int32)) - # Accept one position; speculative future rows stay dirty until the next write. - committed_end = start + 1 - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize( - "cache,hash_page_size", [(256, 256), (256, 512), (512, 512), (768, 768), (1024, 1024)], indirect=["cache"] -) -@pytest.mark.parametrize("hit_extra_page", [False, True]) -def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash_page_size, hit_extra_page): - from sortedcontainers import SortedDict - from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context, CacheTier - from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache - from lightllm.utils.kv_cache_utils import compute_token_list_hash - - manager, requests = cache - small = requests.create_small_page_cache_manager(2) - block_num = 2048 // hash_page_size - big_page_tokens = block_num * hash_page_size - radix = HybridAttPagedRadixCache(manager.size, 0, hash_page_size, block_num, manager.manager) - args = get_env_start_args() - args.linear_att_hash_page_size = hash_page_size - args.linear_att_page_block_num = block_num - tail_len = big_page_tokens + hash_page_size - hit_len = tail_len if hit_extra_page else big_page_tokens - for name, value in { - "req_manager": requests, - "att_state_manager": manager.manager, - "radix_cache": radix, - "args": args, - "is_hybrid_att_model": True, - }.items(): - monkeypatch.setattr(g_infer_context, name, value, raising=False) - - def request(req_idx, total_len): - req = InferReq.__new__(InferReq) - req.req_idx, req.args = req_idx, args - req.cur_kv_len, req.hold_kv_len, req.cur_output_len = 0, 0, 0 - req.shared_kv_node = None - req.tail_small_page_buffer_id = None - req.hybrid_len_to_big_page_id = SortedDict() - req.hybrid_cache_len = (total_len - 1) // hash_page_size * hash_page_size - req.mm_input_spans = [] - req.cache_tiers = {CacheTier.GPU} - req.sampling_param = SimpleNamespace(disable_prompt_cache=False) - req.prompt_selected_logprobs = SimpleNamespace(copy_capture_slots_if_needed=lambda **kwargs: None) - tokens = list(range(total_len)) - hashes = compute_token_list_hash(tokens, hash_page_size) - req.shm_req = SimpleNamespace( - input_len=total_len, - shm_prompt_ids=SimpleNamespace(arr=tokens), - hybrid_token_hash_list=SimpleNamespace(size=len(hashes), get_all=lambda: hashes), - shm_cur_kv_len=0, - prompt_cache_len=0, - ) - req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att - req.req_manager = requests - req.args = args - if not hasattr(args, "req_token_index_mode"): - args.req_token_index_mode = "pin_mem_page_cache" - from lightllm.server.router.model_infer.req_token_index import create_req_token_index - - req.token_index = create_req_token_index(req, torch.cuda.current_stream()) - return req - - source, held = _reserve(manager, requests, tail_len + 1) - req = request(source, tail_len + 1) - req.token_index.append_token_indexes(held.cpu()) - req.hold_kv_len = held.numel() - for end in (big_page_tokens, tail_len): - assert req.get_chuncked_input_token_len() == end - _prepare_swa(requests, source, req.cur_kv_len, end) - for pool in (requests.req_to_swa.flatten(1, 2), *_compressed_kv_pools(manager)): - pool.random_(0, 256) - _state_buffer(requests, "c4").uniform_() - _state_buffer(requests, "c4_indexer").uniform_() - g_infer_context.save_hybrid_state_to_cache(torch.tensor([source], device="cuda"), [req]) - req.cur_kv_len = end - _prepare_swa(requests, source, req.cur_kv_len, tail_len + 1) - req.cur_kv_len = tail_len + 1 - freed = [] - g_infer_context.free_a_req_mem(freed, req) - requests.memory.allocator.free(torch.cat(freed)) - requests.free_req(source) - assert radix.get_tree_total_tokens_num() == tail_len - - forks = [request(requests.alloc(), hit_len + 1) for _ in range(2)] - for fork in forks: - from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex - - fork._hybrid_match_radix_cache() - ReqTokenIndex.submit_pending_uploads([fork.token_index]) - assert fork.cur_kv_len == hit_len - assert fork.hold_kv_len == hit_len - assert fork.shared_kv_node.node_prefix_total_len == big_page_tokens - if hit_extra_page: - assert requests.req_to_token_indexs[fork.req_idx, big_page_tokens].item() != held[big_page_tokens].item() - restored = requests.req_to_token_indexs[fork.req_idx, :hit_len] - for pool, ratio in zip(_compressed_kv_pools(manager), (4, 4, 128)): - torch.testing.assert_close( - _read_packed_rows(pool, 0, held[ratio - 1 : hit_len : ratio].long() // ratio), - _read_packed_rows(pool, 0, restored[ratio - 1 : hit_len : ratio].long() // ratio), - ) - assert set(_request_pages(requests, forks[0].req_idx)).isdisjoint(_request_pages(requests, forks[1].req_idx)) - # Pause can publish a reusable prefix; abort with GPU placement disabled only dereferences it. - forks[1].cache_tiers.clear() - for fork in forks: - freed = [] - g_infer_context.free_a_req_mem(freed, fork) - requests.memory.allocator.free(torch.cat(freed)) - requests.free_req(fork.req_idx) - radix.free_radix_cache_to_get_enough_token(manager.size) - assert manager.allocator.can_use_mem_size == manager.size - assert ( - manager.big_page_buffers.get_free_cache_num() - == manager.big_page_buffers.size - manager.big_page_buffers.keep_num - ) - assert small.get_free_cache_num() == small.size - - -@pytest.mark.parametrize("page_cap", [3, 130]) -def test_c4_indexer_page_table_uses_each_query_request_and_causal_length(cache, page_cap): - from lightllm.models.deepseek_v4.triton_kernel.build_c4_indexer_page_table import build_c4_indexer_page_table - - manager, requests = cache - first, first_slots = _reserve(manager, requests, 512) - second, second_slots = _reserve(manager, requests, 768) - req_ids = [second, first, second, requests.HOLD_REQUEST_ID, first] - lengths = [1, 0, 129, 64, 65] - token_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda") - c4_lengths = torch.tensor(lengths, dtype=torch.int32, device="cuda") - storage = torch.full((len(req_ids) * page_cap + 7,), -17, dtype=torch.int32, device="cuda") - out = storage[:-7].view(len(req_ids), page_cap) - actual = build_c4_indexer_page_table( - token_req_idx, - c4_lengths, - page_cap * 64, - requests.req_to_token_indexs, - requests.HOLD_REQUEST_ID, - out=out, - ) - expected = torch.zeros_like(out) - request_slots = {first: first_slots, second: second_slots} - for row, (req_idx, length) in enumerate(zip(req_ids, lengths)): - if req_idx == requests.HOLD_REQUEST_ID: - continue - compressed_slots = request_slots[req_idx][3 : length * 4 : 4].long() // 4 - pages = compressed_slots[::64] // 64 - expected[row, : pages.numel()] = pages.to(torch.int32) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert actual.data_ptr() == out.data_ptr() - assert storage[-7:].eq(-17).all() - empty = build_c4_indexer_page_table( - token_req_idx[:0], - c4_lengths[:0], - page_cap * 64, - requests.req_to_token_indexs, - requests.HOLD_REQUEST_ID, - ) - assert empty.shape == (0, page_cap) - requests.free([first, second], torch.cat((first_slots, second_slots))) - - -@pytest.mark.parametrize("mode", ["decode", "prefill"]) -def test_c4_indexer_metadata_scoring_and_graph_replay(cache, monkeypatch, mode): - from lightllm.models.deepseek_v4 import infer_struct as infer_struct_module - from lightllm.models.deepseek_v4.layer_infer import indexer as indexer_module - - torch.manual_seed(1593) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 4096) - heads, tokens = 64, 5 - config = dict( - compress_ratios=[4, 128, 0], - num_attention_heads=heads, - index_topk=512, - index_head_dim=128, - qk_rope_head_dim=64, - index_n_heads=heads, - sliding_window=128, - head_dim=512, - rms_norm_eps=1e-6, - ) - indexer = indexer_module.DeepseekV4IndexInfer(0, config, 1) - state = infer_struct_module.DeepseekV4InferStateInfo() - state.mem_manager, state.req_manager = manager, requests - state.att_state_manager = requests.memory.manager - state.att_state_role = "main" - state.first_logical_layer_id = 0 - state.first_c4_logical_layer_id = 0 - state.is_prefill = mode != "decode" - state.microbatch_index, state.max_kv_seq_len = 0, 4096 - state.b_req_idx = torch.tensor([req_idx] if state.is_prefill else [req_idx] * tokens, device="cuda") - state.b_seq_len = torch.full_like(state.b_req_idx, 4096) - state.position_ids = torch.arange(3068, 3068 + tokens, device="cuda") - state.dsv4_sparse_req_idx = torch.full_like(state.position_ids, req_idx) - state._dsv4_token_to_batch_idx = torch.zeros(tokens, dtype=torch.int64, device="cuda") - # Exercise chunk scheduling without requiring a long prompt or large logits allocation. - monkeypatch.setattr(infer_struct_module, "C4_PREFILL_LOGITS_BUDGET_BYTES", 2 * 1024 * 4) - keys = (torch.randn((1024, 128), device="cuda") * 0.1).to(torch.float8_e4m3fn) - scales = torch.ones((1024, 1), dtype=torch.float32, device="cuda") - packed = torch.cat((keys.view(torch.uint8), scales.view(torch.uint8)), dim=1) - c4_slots = slots[3::4].long() // 4 - _write_packed_rows(_compressed_kv_pools(manager)[1], 0, c4_slots, packed) - q = (torch.randn((tokens, heads, 128), device="cuda") * 0.125).to(torch.float8_e4m3fn) - weights = torch.rand((tokens, heads), device="cuda") - - def select(): - return indexer._c4_indices(state, q, weights) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - select() - state.c4_indexer_metadata = None - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - captured, captured_lengths = select() - # Change both causal lengths and Q inputs; graph metadata must be recomputed. - state.position_ids.add_(256) - state.b_seq_len.fill_(3584) - q.copy_((q.float() * 1.25).to(q.dtype)) - graph.replay() - actual = captured.clone() - actual_lengths = captured_lengths.clone() - state.c4_indexer_metadata = None - eager, eager_lengths = select() - torch.testing.assert_close(actual.sort(dim=-1).values, eager.sort(dim=-1).values, rtol=0, atol=0) - torch.testing.assert_close(actual_lengths, eager_lengths, rtol=0, atol=0) - torch.testing.assert_close(actual_lengths, torch.full_like(actual_lengths, 512)) - inverse_slots = torch.full((_compressed_kv_pools(manager)[1].shape[1] * 64,), -1, device="cuda", dtype=torch.long) - inverse_slots[c4_slots] = torch.arange(1024, device="cuda") - logical_indices = inverse_slots[actual[:, 0].long()] - scores = ((q.float() @ keys.float().T).relu() * weights[:, :, None]).sum(1) - valid_lengths = (state.position_ids + 1) // 4 - for row in range(tokens): - selected = logical_indices[row] - assert ((selected >= 0) & (selected < valid_lengths[row])).all() - selected_scores = scores[row, selected].sort(descending=True).values - reference_scores = scores[row, : valid_lengths[row]].topk(512).values - torch.testing.assert_close(selected_scores, reference_scores, rtol=2e-3, atol=2e-3) - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("owned_layer_ids", [(0, 1, 2), (0, 2), (2,)]) -def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill, owned_layer_ids): - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - - manager, requests = cache - first, first_slots = _reserve(manager, requests, 512) - second, second_slots = _reserve(manager, requests, 512) - backend = _test_backend(requests) - model = backend.model - model.layers_infer = [layer for layer in model.layers_infer if layer.layer_num_ in owned_layer_ids] - has_c128_layer = 1 in owned_layer_ids - model.graph = SimpleNamespace(can_run=lambda *args: True, graph_max_len_in_batch=16384) - model.dsv4_workspace = DeepseekV4Workspace(model) - rope = torch.zeros((1024, 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - - def make_state(lengths, microbatch_index): - state = DeepseekV4InferStateInfo() - state.is_prefill, state.microbatch_index = is_prefill, microbatch_index - state.req_manager, state.mem_manager = requests, manager - state.att_state_manager = requests.memory.manager - state.att_state_role = "main" - state.first_logical_layer_id = 0 - state.first_c4_logical_layer_id = 0 - state.batch_size = 2 - state.max_kv_seq_len = max(lengths) - state.input_ids = torch.zeros(2, dtype=torch.int64, device="cuda") - state.b_req_idx = torch.tensor([first, second], dtype=torch.int32, device="cuda") - state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") - if is_prefill: - state.b_ready_cache_len = state.b_seq_len - 1 - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) - state.init_some_extra_state(model) - return state - - captured_state = make_state([130, 258], 0) - other_state = make_state([130, 258], 1) - swa_before = captured_state.dsv4_swa_indices.clone() - other_state.dsv4_swa_indices.fill_(-7) - torch.testing.assert_close(captured_state.dsv4_swa_indices, swa_before) - if has_c128_layer: - assert captured_state.dsv4_c128_indices.shape == (2, 64 if is_prefill else 128) - else: - assert captured_state.dsv4_c128_indices is None - assert captured_state.dsv4_c128_lengths is None - captured_state.init_c4_indices(512) - other_state.init_c4_indices(512) - captured_state.dsv4_c4_indices.fill_(11) - other_state.dsv4_c4_indices.fill_(22) - assert captured_state.dsv4_c4_indices.eq(11).all() - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - swa_output = captured_state.dsv4_swa_indices.clone() - if has_c128_layer: - c128_output = captured_state.dsv4_c128_indices.clone() - updated = make_state([258, 386], 0) - assert updated.dsv4_c4_indices is None - if is_prefill: - captured_state.copy_for_prefill_cuda_graph(updated) - else: - captured_state.copy_for_cuda_graph(updated) - graph.replay() - torch.testing.assert_close(swa_output, updated.dsv4_swa_indices) - if has_c128_layer: - torch.testing.assert_close(c128_output, updated.dsv4_c128_indices) - assert captured_state.dsv4_c4_indices.eq(11).all() - requests.free([first, second], torch.cat((first_slots, second_slots))) - - -@pytest.mark.parametrize("microbatch_index", [0, 1]) -@pytest.mark.parametrize("mode", ["proposal", "commit", "prefill"]) -def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, mode): - from unittest.mock import Mock - - import lightllm.models.deepseek_v4.triton_kernel.build_swa_indices as swa_kernel - import lightllm.models.deepseek_v4_dspark.infer_struct as dspark_state - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - from lightllm.models.deepseek_v4_dspark.infer_struct import DeepseekV4DSparkInferStateInfo - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend - - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - args = get_env_start_args() - monkeypatch.setitem(args, "mtp_mode", "dspark") - monkeypatch.setitem(args, "enable_decode_microbatch_overlap", True) - monkeypatch.setitem(args, "enable_prefill_microbatch_overlap", False) - config = dict( - sliding_window=128, - index_topk=512, - num_attention_heads=64, - compress_ratios=[4, 128, 0], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - ) - model = type("TestModel", (), {})() - vars(model).update( - config=config, - batch_max_tokens=3, - max_req_num=requests.max_request_num, - max_seq_length=4096, - graph=None, - has_vision=False, - args=args, - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - tp_rank_=0, - tp_world_size_=1, - att_state_role="main", - layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], - ) - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) - model.block_size = 3 - model.dsv4_workspace = DeepseekV4Workspace(model) - rope = torch.zeros((512, 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - state = DeepseekV4DSparkInferStateInfo() - state.is_prefill, state.microbatch_index = mode == "prefill", microbatch_index - state.is_draft_model = True - state.req_manager, state.mem_manager = requests, manager - state.att_state_manager = requests.memory.manager - state.att_state_role = "main" - state.first_logical_layer_id = 0 - state.first_c4_logical_layer_id = 0 - if state.is_prefill: - state.batch_size = 1 - state.input_ids = torch.ones(3, dtype=torch.int64, device="cuda") - state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") - state.b_seq_len = torch.tensor([260], dtype=torch.int32, device="cuda") - state.b_ready_cache_len = torch.tensor([257], dtype=torch.int32, device="cuda") - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) - state.b_req_idx = torch.full((3,), req_idx, dtype=torch.int32, device="cuda") - state.b_seq_len = torch.arange(258, 261, dtype=torch.int32, device="cuda") - if mode == "commit": - state.mtp_draft_input_hiddens = torch.zeros((3, 1), device="cuda") - state.max_kv_seq_len = 260 - causal_builder = Mock(wraps=swa_kernel.build_swa_indices) - draft_builder = Mock(wraps=dspark_state.build_dspark_swa_indices) - monkeypatch.setattr(swa_kernel, "build_swa_indices", causal_builder) - monkeypatch.setattr(dspark_state, "build_dspark_swa_indices", draft_builder) - state.init_some_extra_state(model) - proposal = mode == "proposal" - assert causal_builder.call_count == int(not proposal) - assert draft_builder.call_count == int(proposal) - assert state.dsv4_swa_indices.shape == (3, 256 if proposal else 128) - torch.testing.assert_close( - state.dsv4_swa_lengths, torch.full_like(state.dsv4_swa_lengths, 131 if proposal else 128) - ) - expected_slots = _get_swa_token_indices( - requests, req_idx, state.position_ids, prefill_state=state.prefill_att_state if state.is_prefill else None - ).to(torch.int32) - torch.testing.assert_close(state.dsv4_swa_write_slots, expected_slots) - requests.free([req_idx], slots) - - -def test_empty_dspark_commit_is_local_but_proposal_keeps_common_padding_path(monkeypatch): - from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel - - model = DeepseekV4DSparkModel.__new__(DeepseekV4DSparkModel) - model_input = SimpleNamespace(batch_size=0, mtp_draft_input_hiddens=torch.empty((0, 8), device="cuda")) - # A KV-only empty commit is local; a full proposal still takes the common - # model path and participates in DP collectives through HOLD padding. - output = model._decode(model_input) - assert output.logits.shape == (0, 1) - assert output.logits.device == model_input.mtp_draft_input_hiddens.device - - proposal_calls = [] - monkeypatch.setattr(DeepseekV4TpPartModel, "_decode", lambda self, inputs: proposal_calls.append(inputs)) - model_input.mtp_draft_input_hiddens = None - model._decode(model_input) - assert proposal_calls == [model_input] - - -@pytest.mark.parametrize( - "length,prefix", [(768, 0), (768, 256), (2816, 0), (2816, 2304), (4096, 0), (4096, 256), (4096, 2304)] -) -@pytest.mark.parametrize("exhaust_small_pool", [False, True]) -@pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) -def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( - cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode -): - """Offload after the runtime ring is overwritten, then resume through the common loader.""" - from collections import deque - from lightllm.server.multi_level_kv_cache import CacheTier - from lightllm.server.router.model_infer.infer_batch import InferReq - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - from unit_tests.common.dsv4_test_utils import DeepseekV4StateCacheManager - - manager, requests = cache - source, src = _reserve(manager, requests, length) - requests.create_small_page_cache_manager(4) - page_tokens = get_env_start_args().cpu_cache_token_page_size - ends = list(range(page_tokens, length + 1, page_tokens)) - if length % page_tokens: - ends.append(length) - source_checkpoints = {} - tail = None - expected_checkpoints = [] - for end in ends: - # Populate a runtime state at this endpoint without generating new tokens. - _prepare_swa(requests, source, end, end) - value = 11 + end // 256 - requests.req_to_swa.flatten(1, 2).fill_(value) - _state_buffer(requests, "c4").fill_(value) - _state_buffer(requests, "c4_indexer").fill_(value) - state_pool = requests.small_page_buffers if end % page_tokens else manager.big_page_buffers - slot = state_pool.alloc_one_state_cache() - requests.save_state(_checkpoint_req(source, end), slot, state_pool) - torch.cuda.synchronize() - expected_checkpoints.append(state_pool.buffer[slot].clone()) - if end % page_tokens: - tail = slot - else: - source_checkpoints[end] = slot - offload_before = manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID].clone() - # The live ring no longer contains the prompt's continuation checkpoints. - requests.req_to_swa.flatten(1, 2).fill_(99) - _state_buffer(requests, "c4").fill_(88) - _state_buffer(requests, "c4_indexer").fill_(77) - _randomize_compressed_kv(manager) - pages = torch.full( - (len(ends) + 2, 1, 1, 1, manager.cache_config.get_cpu_cache_big_page_bytes()), - 221, - dtype=torch.uint8, - pin_memory=True, - ) - published, released = [], [] - client = SimpleNamespace( - cpu_kv_cache_tensor=pages, - lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), - allocate_pages=lambda hashes, **kwargs: (list(range(1, len(hashes) + 1)), [False] * len(hashes)), - update_pages_status_to_ready=lambda **kwargs: published.append(kwargs), - deref_pages=lambda page_list: released.extend(page_list), - ) - req = SimpleNamespace( - req_idx=source, - req_id=42, - cur_kv_len=length, - hold_kv_len=src.numel(), - shared_kv_node=None, - hybrid_len_to_big_page_id=source_checkpoints, - hybrid_cache_len=length, - tail_small_page_buffer_id=tail, - cpu_cache_offload_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, - cache_tiers=(CacheTier.CPU,), - shm_req=SimpleNamespace( - group_req_id=42, - request_id=42, - input_len=length + 1, - token_hash_list=SimpleNamespace(get_all=lambda: list(range(len(ends)))), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - attach_req_token_index(req, requests, src.cpu(), mode=token_index_mode) - # 已命中的大页来自 radix，其余完整大页来自当前请求。 - shared_page_ids = [source_checkpoints.pop(end) for end in range(page_tokens, prefix + 1, page_tokens)] - producer = torch.cuda.current_stream() - context = common.g_infer_context - monkeypatch.setattr(context, "req_manager", requests) - monkeypatch.setattr(context, "att_state_manager", manager.manager) - monkeypatch.setattr( - context, - "radix_cache", - SimpleNamespace( - get_big_page_ids_by_node=lambda node: shared_page_ids.copy(), - free_one_small_page_buffer=lambda: None, - free_radix_cache_to_get_enough_token=lambda need: None, - ), - ) - monkeypatch.setattr(context, "is_hybrid_att_model", True) - monkeypatch.setattr(context, "get_overlap_stream", lambda: producer) - monkeypatch.setattr(context, "get_can_alloc_token_num", lambda: manager.allocator.can_use_mem_size) - monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace( - page_size=manager.page_size, - cpu_cache_token_page_size=page_tokens, - linear_att_hash_page_size=256, - disable_linear_att_small_page_cpu_cache=False, - ) - module.backend = SimpleNamespace( - is_master_in_dp=True, - dp_world_size=1, - is_hybrid_att_model=True, - model=SimpleNamespace(att_state_manager=manager.manager, req_manager=requests), - ) - module.cpu_cache_client = client - module.offload_tasks = deque() - module.load_tasks = deque() - module.page_index_buffer = torch.empty(8, dtype=torch.int32, device="cuda") - module.page_ready_buffer = torch.empty(8, dtype=torch.bool, device="cuda") - module.gloo_group = module.filter_group = module.init_sync_group = None - module.need_sync_compute_stream = lambda: False - assert module.offload_finished_reqs_to_cpu_cache([req]) == [] - assert req.cpu_cache_offload_task_status.is_running() - torch.cuda.synchronize() - module.update_cpu_cache_offload_task_states() - assert module.offload_finished_reqs_to_cpu_cache([req]) == [req] - assert published == [ - dict( - page_list=list(range(1, len(ends) + 1)), - deref=True, - disk_offload_enable=False, - token_num_in_page_list=length, - ) - ] - flat = pages.view(pages.shape[0], -1) - for i, expected in enumerate(expected_checkpoints, 1): - torch.testing.assert_close( - flat[i, manager.cache_config.get_cpu_cache_full_att_bytes() :], expected, rtol=0, atol=0 - ) - if tail is not None: - torch.testing.assert_close( - manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID], - offload_before, - rtol=0, - atol=0, - ) - - destination = requests.alloc() - dst = torch.empty(0, dtype=torch.int32, device="cuda") - if prefix: - dst = manager.alloc(prefix).cuda() - requests.req_to_token_indexs[destination, :prefix] = dst - manager.operator.copy_mem_to_mem(src[:prefix], dst) - restored = SimpleNamespace( - req_idx=destination, - cur_kv_len=prefix, - hold_kv_len=prefix, - mm_input_spans=[], - hybrid_cache_len=length, - hybrid_len_to_big_page_id={}, - tail_small_page_buffer_id=None, - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=length + 1, - disk_prompt_cache_len=0, - cpu_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: list(range(1, len(ends) + 1))), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - attach_req_token_index(restored, requests, dst.cpu(), mode=token_index_mode) - if prefix >= page_tokens: - slot = manager.big_page_buffers.alloc_one_state_cache() - manager.big_page_buffers.buffer[slot].copy_(expected_checkpoints[0]) - restored.hybrid_len_to_big_page_id[page_tokens] = slot - - busy_small_slots = [] - if exhaust_small_pool: - busy_small_slots = requests.small_page_buffers.alloc_state_cache( - requests.small_page_buffers.get_free_cache_num() - ) - module.load_cpu_cache_to_reqs([restored]) - torch.cuda.synchronize() - module.update_cpu_cache_load_task_states() - assert restored.cur_kv_len == restored.hold_kv_len == length - assert released == list(range(1, len(ends) + 1)) - if length % page_tokens: - torch.testing.assert_close( - manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID], - expected_checkpoints[-1], - rtol=0, - atol=0, - ) - # CPU 碎页直接经 LOAD 大页恢复运行态，不重复占用普通小页。 - assert restored.tail_small_page_buffer_id is None - dst = requests.req_to_token_indexs[destination, :length] - for pool in _compressed_kv_pools(manager): - torch.testing.assert_close( - pool[:, (src[::256] // 256).long()], pool[:, (dst[::256] // 256).long()], rtol=0, atol=0 - ) - check = DeepseekV4StateCacheManager(1, manager.cache_config) - _prepare_swa(requests, destination, length, length) - requests.save_state(_checkpoint_req(destination, length), 0, check) - torch.cuda.synchronize() - torch.testing.assert_close(check.buffer[0], expected_checkpoints[-1], rtol=0, atol=0) - requests.free([source, destination], torch.cat((src, dst))) - manager.big_page_buffers.free_state_cache( - shared_page_ids + list(source_checkpoints.values()) + list(restored.hybrid_len_to_big_page_id.values()) - ) - requests.small_page_buffers.free_state_cache( - [x for x in (tail, restored.tail_small_page_buffer_id) if x is not None] + busy_small_slots - ) - assert manager.allocator.can_use_mem_size == manager.size - assert ( - manager.big_page_buffers.get_free_cache_num() - == manager.big_page_buffers.size - manager.big_page_buffers.keep_num - ) - - -def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): - from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - - manager, requests = cache - req, held = _reserve(manager, requests, 2048) - reference = torch.zeros((1, 17, 74880), dtype=torch.uint8, device="cuda") - kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") - weight = torch.randn(512, dtype=torch.bfloat16, device="cuda") - angles = torch.randn((2048, 32), device="cuda") - freqs = torch.complex(angles.cos(), angles.sin()) - layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs - layer_weight = SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight)) - for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): - _prepare_swa(requests, req, start, end) - positions = torch.arange(start, end, device="cuda", dtype=torch.int32) - indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") - lengths = torch.empty(positions.numel(), dtype=torch.int32, device="cuda") - writes = torch.empty_like(lengths) - runtime = torch.empty_like(lengths) - build_swa_indices( - torch.full_like(positions, req), - positions, - requests, - indices, - lengths, - writes, - prefill_state=requests._test_prefill_states[req], - swa_runtime_write_slots=runtime, - token_to_batch_idx=torch.zeros_like(positions), - b_seq_len=torch.tensor([end], device="cuda"), - ) - slots = _swa_slots(requests, req, positions) - state = SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - is_prefill=True, - prefill_att_state=requests._test_prefill_states[req], - dsv4_swa_write_slots=slots, - dsv4_swa_runtime_write_slots=runtime, - position_ids=positions, - ) - layer._norm_rope_and_cache_kv(kv[start:end], state, layer_weight) - fused_kv_norm_rope_to_swa(kv[start:end], weight, 1e-6, freqs, positions, positions, reference[0], 128) - visible = torch.arange(max(0, start - 127), end, device="cuda") - # 第八个 scale 字节仅用于对齐；workspace 复用时无需保持其初始值。 - torch.testing.assert_close( - _read_swa(requests, 2, req, visible)[:, :-1], - _read_packed_rows(reference, 0, visible)[:, :-1], - rtol=0, - atol=0, - ) - logical = positions[:, None] - torch.arange(128, device="cuda")[None, :] - valid = logical >= 0 - expected_slots = _swa_slots(requests, req, logical.clamp_min(0)) - torch.testing.assert_close(indices[valid], expected_slots[valid].to(torch.int32)) - _finish_prefill(requests, [req]) - tail = torch.arange(max(0, end - requests.sliding_window), end, device="cuda") - torch.testing.assert_close( - _read_swa(requests, 2, req, tail)[:, :-1], - _read_packed_rows(reference, 0, tail)[:, :-1], - rtol=0, - atol=0, - ) - requests.free([req], held) - - -def test_cpu_cache_preserves_already_published_pages(cache): - manager, requests = cache - req, slots = _reserve(manager, requests, 2048) - _randomize_compressed_kv(manager) - layout = manager.cache_config - pages = torch.full((3, layout.get_cpu_cache_big_page_bytes()), 221, dtype=torch.uint8, pin_memory=True) - page_ids = [1] - readies = [True] - state_id = manager.big_page_buffers.alloc_one_state_cache() - manager.big_page_buffers.buffer[state_id].fill_(7) - checkpoint_ids = [state_id] - manager.manager.get_cpu_cache_operator().write_req( - token_indices=slots, - cpu_page_buffers=pages, - cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), - cpu_page_ready_flags=readies, - big_page_slot_ids=checkpoint_ids, - ) - torch.cuda.synchronize() - assert pages.eq(221).all() - readies[0] = False - manager.manager.get_cpu_cache_operator().write_req( - token_indices=slots, - cpu_page_buffers=pages, - cpu_page_indices=page_ids, - layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), - cpu_page_ready_flags=readies, - big_page_slot_ids=checkpoint_ids, - ) - torch.cuda.synchronize() - offset = 0 - for state in manager.manager.get_full_attention_states(): - expected = state.storage[0, slots.long()].cpu().reshape(-1) - torch.testing.assert_close(pages[1, offset : offset + expected.numel()], expected, rtol=0, atol=0) - offset += expected.numel() - assert pages[1, layout.get_cpu_cache_full_att_bytes() :].eq(7).all() - assert pages[0].eq(221).all() and pages[2].eq(221).all() - manager.big_page_buffers.free_state_cache([state_id]) - requests.free([req], slots) - - -@pytest.mark.parametrize("prefix_len", range(4)) -@pytest.mark.parametrize("cached_first_image", [False, True]) -def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_len, cached_first_image): - from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - from lightllm.server.multimodal_params import MultimodalParams - from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans - - tokenizer = DeepSeekV4Tokenizer( - SimpleNamespace(convert_tokens_to_ids=lambda token: 99), "/unused", {"vision_n_layers": 1} - ) - params = MultimodalParams(images=[{"type": "url", "data": "unused"} for _ in range(3)]) - token_base, canonical_length = 2 ** 62, 13 - for image in params.images: - image.token_id, image.token_num, image.start_index_in_embed_cache = token_base, canonical_length, 0 - input_ids = tokenizer.encode([7] * prefix_len + [99, 99, 8, 99, 8], params) - spans = get_mm_input_spans(input_ids, params.to_dict()) - expected_left, expected_right = [0] * len(input_ids), [0] * len(input_ids) - for start, end in spans: - image_start = input_ids.index(token_base + 3, start, end) - expected_left[image_start:end] = range(end - image_start) - expected_right[image_start:end] = range(end - image_start - 1, -1, -1) - - ready = spans[0][1] if cached_first_image else 0 - manager, requests = cache - req_idx, slots = _reserve(manager, requests, len(input_ids)) - backend = _test_backend(requests, token_capacity=128) - model = backend.model - model.has_vision = True - model.config["vision_max_n_token"] = 384 - model.dsv4_workspace = None - rope = torch.zeros((len(input_ids), 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - state = DeepseekV4InferStateInfo() - state.is_prefill = True - state.req_manager, state.mem_manager = requests, manager - state.att_state_manager = requests.memory.manager - state.att_state_role = "main" - state.first_logical_layer_id = 0 - state.first_c4_logical_layer_id = 0 - state.input_ids = torch.tensor(input_ids[ready:], dtype=torch.int64, device="cuda") - state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") - state.b_seq_len = torch.tensor([len(input_ids)], dtype=torch.int32, device="cuda") - state.b_ready_cache_len = torch.tensor([ready], dtype=torch.int32, device="cuda") - state.max_kv_seq_len = len(input_ids) - state.multimodal_params = [params.to_dict()] - state.prefill_att_state = backend.create_att_prefill_state(state) - state.init_some_extra_state(model) - assert state.dsv4_image_left.cpu().tolist() == expected_left[ready:] - assert state.dsv4_image_right.cpu().tolist() == expected_right[ready:] - - text_embeddings = torch.randn((100, 16), dtype=torch.bfloat16, device="cuda") - image_embeddings = torch.randn((canonical_length, 1, 16), dtype=torch.bfloat16, device="cuda") - out = torch.zeros((len(input_ids) - ready, 16), dtype=torch.bfloat16, device="cuda") - multimodal_emb( - out, - state.input_ids, - text_embeddings, - image_embeddings, - torch.tensor([canonical_length], dtype=torch.int64, device="cuda"), - torch.tensor([token_base], dtype=torch.int64, device="cuda"), - torch.tensor([0], dtype=torch.int64, device="cuda"), - 0, - 100, - 1, - ) - expected = torch.stack( - [ - text_embeddings[token] if token < 100 else image_embeddings[token - token_base, 0] - for token in input_ids[ready:] - ] - ) - torch.testing.assert_close(out, expected, rtol=0, atol=0) - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("padding_tokens", [0, 3, 5]) -def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - - manager, requests = cache - backend = _test_backend(requests) - model = backend.model - model.has_vision = False - rope = torch.zeros((8, 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - - model.dsv4_workspace = DeepseekV4Workspace(model) - real_tokens = 5 - padding_tokens - req_ids, lengths = ([0], [real_tokens]) if real_tokens else ([], []) - if padding_tokens: - req_ids.append(requests.HOLD_REQUEST_ID) - lengths.append(padding_tokens) - infer_state = DeepseekV4InferStateInfo() - infer_state.is_prefill = True - infer_state.req_manager = requests - infer_state.att_state_manager = requests.memory.manager - infer_state.att_state_role = "main" - infer_state.first_logical_layer_id = 0 - infer_state.first_c4_logical_layer_id = 0 - infer_state.mem_manager = manager - infer_state.input_ids = torch.ones(5, dtype=torch.int64, device="cuda") - infer_state.b_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda") - infer_state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") - infer_state.b_ready_cache_len = torch.zeros_like(infer_state.b_seq_len) - infer_state.max_kv_seq_len = max(lengths) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) - infer_state.init_some_extra_state(model) - - previous_debug_mode = torch.cuda.get_sync_debug_mode() - torch.cuda.set_sync_debug_mode("error") - try: - infer_state.init_some_extra_state(model) - finally: - torch.cuda.set_sync_debug_mode(previous_debug_mode) - assert infer_state.dsv4_sparse_req_idx.shape == (5,) - - -@pytest.mark.parametrize("window", [64, 128, 256, 512]) -@pytest.mark.parametrize("mtp_step", [0, 5]) -def test_configured_window_survives_chunks_and_checkpoint_restore(cache, window, mtp_step): - from dataclasses import replace - - manager, _ = cache - config = replace(manager.cache_config, sliding_window=window) - requests = DeepseekV4ReqManager(2, 4096, manager, config, mtp_step=mtp_step) - source, destination = requests.alloc(), requests.alloc() - packed = torch.randint(0, 256, (2048, config.kv_bytes), dtype=torch.uint8, device="cuda") - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - for start, end in ((0, 257), (257, 512), (512, 513), (513, 2048)): - _prepare_swa(requests, source, start, end) - positions = torch.arange(start, end, device="cuda") - _write_swa(requests, 0, source, positions, packed[start:end]) - if end % 256 == 0: - requests.save_state(_checkpoint_req(source, end), index, states) - requests.restore_state(SimpleNamespace(req_idx=destination, cur_kv_len=end), states, index) - positions = torch.arange(max(0, end - window), end, device="cuda") - actual = _read_swa(requests, 0, destination, positions) - torch.testing.assert_close(actual, packed[positions], rtol=0, atol=0) - _finish_prefill(requests, [source]) - positions = torch.arange(max(0, end - window), end, device="cuda") - actual = _read_swa(requests, 0, source, positions) - torch.testing.assert_close(actual, packed[positions], rtol=0, atol=0) - requests.free_req(source) - requests.free_req(destination) - states.free_state_cache([index]) - - -@pytest.mark.parametrize("mtp_step", [0, 3]) -def test_full_batch_unaligned_prefill_workspace_is_bounded_and_private(cache, mtp_step): - manager, _ = cache - request_count, chunk_tokens, ready = 8, 64, 382 - requests = DeepseekV4ReqManager(request_count, 4096, manager, manager.cache_config, mtp_step=mtp_step) - reqs = [requests.alloc() for _ in range(request_count)] - page_token_num = manager.cache_config.swa_page_token_num - first = (ready - requests.sliding_window + 1) // page_token_num * page_token_num - end = ready + chunk_tokens - required = request_count * ((end + page_token_num - 1) // page_token_num * page_token_num - first) - requests._test_backend = _test_backend(requests, request_count * chunk_tokens) - requests._test_prefill_states = {} - workspace, capacity = requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model) - assert capacity >= required - assert workspace.untyped_storage().data_ptr() != requests.req_to_swa.untyped_storage().data_ptr() - workspace[:, capacity : 2 * capacity].fill_(173) - old_positions = torch.arange(ready - requests.sliding_window + 1, ready, device="cuda") - new_positions = torch.arange(ready, end, device="cuda") - for req in reqs: - _write_swa_rows( - requests, - 0, - req, - old_positions, - torch.full((old_positions.numel(), 584), req + 1, dtype=torch.uint8, device="cuda"), - prefill=False, - ) - _prepare_prefill( - requests, torch.tensor(reqs), torch.full((request_count,), ready), torch.full((request_count,), end) - ) - for req in reqs: - _write_swa( - requests, - 0, - req, - new_positions, - torch.full((chunk_tokens, 584), req + 17, dtype=torch.uint8, device="cuda"), - ) - for req in reqs: - positions = torch.arange(ready - requests.sliding_window + 1, end, device="cuda") - slots = _swa_slots(requests, req, positions) - assert slots.unique().numel() == positions.numel() - actual = _read_packed_rows(workspace, 0, slots) - expected = torch.where(positions[:, None] < ready, req + 1, req + 17).expand_as(actual).to(torch.uint8) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert workspace[:, capacity : 2 * capacity].eq(173).all() - _finish_prefill(requests, reqs) - for req in reqs: - positions = torch.arange(end - requests.sliding_window, end, device="cuda") - actual = _read_swa(requests, 0, req, positions) - expected = torch.where(positions[:, None] < ready, req + 1, req + 17).expand_as(actual).to(torch.uint8) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - requests.free_req(req) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py index ba349372ca..dc3bc15ecb 100644 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -95,6 +95,9 @@ def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): def test_shared_attention_state_roundtrip_preserves_worker_resources( runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization ): + # Torch 2.11 uses pidfd_getfd for expandable segments, which is unavailable in + # some containers. Disable expandable segments before allocating the exported state. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") state, resources, _ = make_state(cls=cls, mode=mode) state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) diff --git a/unit_tests/common/test_mem_manager_ipc.py b/unit_tests/common/test_mem_manager_ipc.py index 4008cd6686..523b61f82f 100644 --- a/unit_tests/common/test_mem_manager_ipc.py +++ b/unit_tests/common/test_mem_manager_ipc.py @@ -18,6 +18,10 @@ def set_value(self, value): def _make_manager(monkeypatch, size, device="cpu"): + if device == "cuda": + # The spawn receiver reopens CUDA IPC handles. Recent torch releases use + # pidfd_getfd for expandable segments, which some containers do not expose. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") monkeypatch.setattr(allocator_module, "SharedInt", _Counter) monkeypatch.setattr(allocator_module, "get_current_rank_in_node", lambda: 0) manager = AttentionStateManager(request_capacity=2) diff --git a/unit_tests/common/test_pinned_memory_non_blocking_copy.py b/unit_tests/common/test_pinned_memory_non_blocking_copy.py deleted file mode 100644 index 594157560b..0000000000 --- a/unit_tests/common/test_pinned_memory_non_blocking_copy.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest -import torch -from torch.profiler import ProfilerActivity, profile - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def _profile_request_table_copy(non_blocking: bool): - token_num = 1024 * 1024 - mem_indexes = torch.arange(token_num, dtype=torch.int32, pin_memory=True) - req_to_token_indexes = torch.empty((1, token_num), dtype=torch.int32, device="cuda") - - # 先完成分配等准备操作，避免其 CUDA runtime 调用混入待验证区间。 - torch.cuda.synchronize() - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - acc_events=True, - ) as prof: - req_to_token_indexes[0].copy_(mem_indexes, non_blocking=non_blocking) - if non_blocking: - # 异步路径需要在读取结果前显式等待；这里会记录为 cudaDeviceSynchronize， - # 与 copy_ 内部可能产生的 cudaStreamSynchronize 可以明确区分。 - torch.cuda.synchronize() - - runtime_calls = {event.key for event in prof.key_averages()} - assert torch.equal(req_to_token_indexes[0].cpu(), mem_indexes) - return runtime_calls - - -def test_pinned_cpu_to_gpu_non_blocking_copy_avoids_internal_stream_sync(): - blocking_calls = _profile_request_table_copy(non_blocking=False) - non_blocking_calls = _profile_request_table_copy(non_blocking=True) - - # 两条路径底层都提交异步 H2D memcpy；同步路径随后在 copy_ 内部等待当前 stream。 - assert "cudaMemcpyAsync" in blocking_calls - assert "cudaMemcpyAsync" in non_blocking_calls - assert "cudaStreamSynchronize" in blocking_calls - assert "cudaStreamSynchronize" not in non_blocking_calls diff --git a/unit_tests/common/test_pinned_memory_stream.py b/unit_tests/common/test_pinned_memory_stream.py index f87ee97b51..064d2cb894 100644 --- a/unit_tests/common/test_pinned_memory_stream.py +++ b/unit_tests/common/test_pinned_memory_stream.py @@ -50,24 +50,6 @@ def create_pool(*args, **kwargs): submit(lambda: None).result(timeout=10) -def test_empty_batch_and_fallback_do_not_create_release_pool(monkeypatch): - copier = cuda_wrapper.PinnedH2DBatchCopier - monkeypatch.setattr(copier, "_release_pool", None) - monkeypatch.setattr(copier, "_native_batch_available", False) - monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(cuda_stream=0)) - - def unexpected_pool(*args, **kwargs): - pytest.fail("Empty batches and copy_ fallback must not create the source release pool") - - monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", unexpected_pool) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([], []) - source = torch.arange(8, dtype=torch.int32) - target = torch.empty_like(source) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - assert copier._release_pool is None - assert torch.equal(target, source) - - def test_release_pool_is_initialized_once_for_concurrent_callers(monkeypatch): copier = cuda_wrapper.PinnedH2DBatchCopier monkeypatch.setattr(copier, "_release_pool", None) @@ -98,82 +80,6 @@ def get_pool(): assert creations == [{"max_workers": 1, "thread_name_prefix": "lightllm-pinned-release"}] -@requires_cuda -@pytest.mark.parametrize("native_batch", [False, True]) -def test_default_stream_copy_orders_reads_and_reuses_upload_stream(monkeypatch, native_batch): - runtime = pytest.importorskip("cuda.bindings.runtime") - copier = cuda_wrapper.PinnedH2DBatchCopier - monkeypatch.setattr(copier, "_native_batch_available", native_batch) - stream = torch.cuda.default_stream() - source = torch.arange(8192, dtype=torch.int32, pin_memory=True) - target = torch.empty(8192, dtype=torch.int32, device="cuda") - native_copy = runtime.cudaMemcpyBatchAsync - submitted_streams, observed = [], [] - - def record_batch(*args): - submitted_streams.append(args[-1]) - return native_copy(*args) - - monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_batch) - with torch.cuda.stream(stream): - for _ in range(2): - torch.cuda._sleep(20_000_000) - target.fill_(-1) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - assert torch.cuda.current_stream() == stream - observed.append(target.clone()) - for result in observed: - assert torch.equal(result.cpu(), source) - if native_batch: - assert len(submitted_streams) == 2 - assert submitted_streams[0] == submitted_streams[1] != 0 - else: - assert submitted_streams == [] - - -@requires_cuda -def test_raw_batch_records_one_completion_event_without_small_copy(monkeypatch, release_futures): - runtime = pytest.importorskip("cuda.bindings.runtime") - source = torch.arange(2048, dtype=torch.int32, pin_memory=True) - another_source = torch.arange(2048, dtype=torch.int32, pin_memory=True).add_(10000) - sources = [source[128:384], source[1000:1256], another_source[17:273]] - target = torch.full((1024,), -1, dtype=torch.int32, device="cuda") - targets = [target[2:258], target[300:556], target[600:856]] - expected = torch.full((1024,), -1, dtype=torch.int32) - expected[2:258], expected[300:556], expected[600:856] = sources - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - calls = [] - native_copy, event_record = runtime.cudaMemcpyBatchAsync, torch.cuda.Event.record - - def record_native(*args): - calls.append(("batch", args[3], args[-1])) - return native_copy(*args) - - def record_event(event, on_stream=None): - calls.append(("event", 1, on_stream.cuda_stream)) - return event_record(event, on_stream) - - def unexpected_copy(*args, **kwargs): - pytest.fail("Raw batch source protection must not submit a small copy_ or enter per-tensor fallback") - - def unexpected_sync(*args, **kwargs): - pytest.fail("Normal raw batches must not synchronize a CUDA stream on the submitting thread") - - monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_native) - monkeypatch.setattr(torch.cuda.Event, "record", record_event) - monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) - with torch.cuda.stream(stream), monkeypatch.context() as no_sync: - no_sync.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) - no_sync.setattr(torch.cuda, "synchronize", unexpected_sync) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async(targets, sources) - assert len(release_futures) == 1 - release_futures[0].result(timeout=10) - assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert calls == [("batch", 3, stream.cuda_stream), ("event", 1, stream.cuda_stream)] - assert torch.equal(target.cpu(), expected) - - @requires_cuda def test_background_wait_holds_tensor_references_without_blocking_submitter(monkeypatch, release_futures): pytest.importorskip("cuda.bindings.runtime") @@ -254,74 +160,6 @@ def test_async_copy_keeps_freed_source_storage_out_of_reuse(monkeypatch, native_ assert torch.equal(target.cpu(), expected) -@requires_cuda -@pytest.mark.parametrize("failure", ["record", "pool", "submit"]) -def test_event_setup_failure_propagates_without_sync_or_fallback(monkeypatch, failure): - pytest.importorskip("cuda.bindings.runtime") - source = torch.arange(1024, dtype=torch.int32, pin_memory=True) - target = torch.empty_like(source, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - synchronize = torch.cuda.Stream.synchronize - - def fail_setup(*args, **kwargs): - raise RuntimeError(f"test {failure} failure") - - def unexpected_sync(*args, **kwargs): - pytest.fail("Source protection setup errors must propagate without synchronous cleanup") - - def unexpected_copy(*args, **kwargs): - pytest.fail("copy_ fallback must not hide source protection failure after raw batch submission") - - if failure == "record": - monkeypatch.setattr(torch.cuda.Event, "record", fail_setup) - elif failure == "pool": - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_release_pool", None) - monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", fail_setup) - else: - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier._get_release_pool(), "submit", fail_setup) - monkeypatch.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) - monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) - try: - with torch.cuda.stream(stream): - torch.cuda._sleep(20_000_000) - with pytest.raises(RuntimeError, match=f"test {failure} failure"): - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - finally: - # 测试捕获注入异常后仍要继续运行，自行持有源并等待在途传输结束。 - synchronize(stream) - assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert torch.equal(target.cpu(), source) - - -@requires_cuda -def test_pytorch_fallback_needs_no_completion_event_or_pool_task(monkeypatch, release_futures): - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", False) - source = torch.arange(1024, dtype=torch.int32, pin_memory=True) - target = torch.empty_like(source, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - copy_sizes = [] - copy = torch.Tensor.copy_ - - def unexpected_event(*args, **kwargs): - pytest.fail("PyTorch copy_ already protects source storage and must not create an extra completion event") - - def record_copy(dst, src, *args, **kwargs): - copy_sizes.append(src.numel() * src.element_size()) - return copy(dst, src, *args, **kwargs) - - monkeypatch.setattr(torch.cuda, "Event", unexpected_event) - monkeypatch.setattr(torch.Tensor, "copy_", record_copy) - with torch.cuda.stream(stream): - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - stream.synchronize() - assert not cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert release_futures == [] - assert copy_sizes == [source.numel() * source.element_size()] - assert torch.equal(target.cpu(), source) - - @requires_cuda def test_source_release_pool_is_reused_across_batches(release_futures): pytest.importorskip("cuda.bindings.runtime") diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 4671576018..dd0da891c1 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -191,70 +191,3 @@ def _make_scheduler_req(manager, req_idx, current, total): req.kv_cache_alloc_need = MethodType(InferReq.kv_cache_alloc_need, req) req.alloc_kv_mem = MethodType(InferReq.alloc_kv_mem, req) return req - - -@pytest.mark.parametrize( - "hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause", - [(4, 8, 3, [4], False), (8, 0, 3, [], False), (4, 4, 6, [4], True)], -) -def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( - monkeypatch, hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause -): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, 4, available, batch_max_tokens) - reqs = [_make_scheduler_req(manager, i, 3, lambda: 6) for i in range(2)] - for req in reqs: - req.alloc_kv_mem(hold_kv_len) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: reqs - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0, 1]) - - assert prefill_reqs == [reqs[0]] - assert decode_reqs == [] - assert prefill_tokens == 3 - assert allocations == expected_alloc_sizes - assert reqs[0].hold_kv_len == 8 - assert not reqs[0].wait_pause - assert reqs[1].hold_kv_len == hold_kv_len - assert reqs[1].wait_pause is second_req_wait_pause - torch.cuda.synchronize() - assert torch.equal(manager.req_to_token_indexs[0, :8].cpu(), reqs[0].token_index.get_cpu_token_indexes(0, 8)) - - -def test_decode_scheduler_publishes_batched_request_table(monkeypatch): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, 4, 12) - req = _make_scheduler_req(manager, 0, 4, lambda: 5) - req.alloc_kv_mem(4) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: [req] - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) - - assert prefill_reqs == [] - assert decode_reqs == [req] - assert prefill_tokens == 0 - assert req.hold_kv_len == 16 - assert allocations == [12] - torch.cuda.synchronize() - assert manager.req_to_token_indexs[0, :16].tolist() == list(range(16)) - - -@pytest.mark.parametrize("page_size, expected_alloc_token_num", [(1, 9), (3, 9), (4, 12), (7, 14), (8, 8)]) -def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_steps( - monkeypatch, page_size, expected_alloc_token_num -): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, page_size, 32) - req = _make_scheduler_req(manager, 0, page_size, lambda: req.cur_kv_len + 1) - req.alloc_kv_mem(page_size) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: [req] - - for _ in range(expected_alloc_token_num): - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) - assert prefill_reqs == [] - assert decode_reqs == [req] - assert prefill_tokens == 0 - req.cur_kv_len += 1 - - assert req.hold_kv_len == page_size + expected_alloc_token_num - assert allocations == [expected_alloc_token_num] diff --git a/unit_tests/models/deepseek2/test_rope_repeat.py b/unit_tests/models/deepseek2/test_rope_repeat.py deleted file mode 100644 index 2f578e03eb..0000000000 --- a/unit_tests/models/deepseek2/test_rope_repeat.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch -import time -import pytest -from lightllm.models.deepseek2.triton_kernel.repeat_rope import repeat_rope - - -def test_torch_cat(): - source = torch.randn((100, 1, 1077), device="cuda") - dest = torch.randn((100, 7, 1077), device="cuda") - - repeat_rope(dest, source) - torch.equal(dest[:, 0, :], source) - torch.equal(dest[:, -1, :], source) - - source = torch.randn((100, 1, 128), device="cuda") - dest = torch.randn((100, 64, 128), device="cuda") - - repeat_rope(dest, source) - torch.equal(dest[:, 0, :], source) - torch.equal(dest[:, -1, :], source) - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/deepseek_v4/test_kernels.py b/unit_tests/models/deepseek_v4/test_kernels.py index bf29e0e70b..d5afb9b219 100644 --- a/unit_tests/models/deepseek_v4/test_kernels.py +++ b/unit_tests/models/deepseek_v4/test_kernels.py @@ -3,10 +3,31 @@ from lightllm.models.deepseek_v4.triton_kernel import norm_rope, topk_softplus_sqrt, indexer_topk from lightllm.models.deepseek_v4.triton_kernel.rotary_emb import inverse_rotary_emb +from lightllm.utils.vllm_utils import vllm_ops pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.fixture(autouse=True) +def topk_provider_fallback(monkeypatch): + if hasattr(vllm_ops, "top_k_per_row_decode"): + return + + def top_k_per_row_decode(scores, indent, lengths, out, num_rows, row_stride, stride, topk): + available = min(row_stride, scores.shape[1]) + for row in range(scores.shape[0]): + row_length = lengths[row].clamp(max=available) + count = min(available, topk) + columns = torch.arange(available, device=scores.device) + valid_scores = torch.where(columns < row_length, scores[row, :available], -torch.inf) + indices = torch.topk(valid_scores, count).indices + mask = columns[:count] < row_length + out[row, :count] = torch.where(mask, indices, -1) + out[row, count:] = -1 + + monkeypatch.setattr(vllm_ops, "top_k_per_row_decode", top_k_per_row_decode, raising=False) + + def freqs_and_positions(batch, dtype=torch.int64): torch.manual_seed(13) angles = torch.randn(128, 32, device="cuda") diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 8394d18b48..7f6e8d5183 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -266,6 +266,7 @@ def test_output_projection_matches_bf16_reference(monkeypatch, tp_size, rank): layer.prefix, layer.data_type_ = "layers.0", torch.bfloat16 layer.hidden, layer.n_heads, layer.head_dim = 128, 8, 64 layer.q_lora_rank, layer.o_lora_rank, layer.o_groups = 128, 128, 2 + layer.num_fused_shared_experts = 0 layer.get_quant_method = lambda name: None layer._init_qkvo() layer.experts = SimpleNamespace(quant_method=SimpleNamespace(weight_scale_suffix=None)) diff --git a/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py b/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py deleted file mode 100644 index 0338c98575..0000000000 --- a/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py +++ /dev/null @@ -1,200 +0,0 @@ -import pytest -import torch -import triton -from triton.runtime.jit import JITFunction - -from lightllm.models.glm5_next.triton_kernel.index_quant import ( - hadamard_transform_quant_fp8, -) -from lightllm.models.glm5_next.triton_kernel.kpool import ( - _compress_pools, - _gather_pools, - _get_pool_ranges, - compress_pools, - gather_pools, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _check_pool_ranges(batch, max_pools, single_query): - counts = [1] * batch if single_query else [0 if i % 7 == 3 else i % 3 + 1 for i in range(batch)] - cu = [0] - visible, owners = [], [] - for owner, count in enumerate(counts): - visible.extend(range(4, count + 4)) - owners.extend([owner] * count) - cu.append(len(visible)) - lengths = torch.tensor(visible, device="cuda", dtype=torch.int32) - cu_q_lens = torch.tensor(cu, device="cuda", dtype=torch.int32) - outputs = [torch.empty_like(lengths) for _ in range(3)] - tokens = lengths.numel() - blocks = triton.cdiv(tokens, 256) if single_query else tokens // 256 + batch - compiled = _get_pool_ranges[(blocks,)]( - lengths, cu_q_lens, *outputs, batch, tokens, max_pools, single_query, 256, num_warps=4 - ) - starts = torch.tensor(owners, device="cuda", dtype=torch.int32) * max_pools - expected = (starts, starts + lengths // 4, lengths // 4) - for actual, reference in zip(outputs, expected): - assert torch.equal(actual, reference) - return compiled.hash - - -@pytest.mark.parametrize("single_query", [False, True]) -def test_pool_ranges_reuse_compilation_across_runtime_sizes(single_query): - # The indexer rounds capacity up to 128 pools. Changing that capacity keeps - # its default alignment specialization while batch/token counts are dynamic. - hashes = { - _check_pool_ranges(batch, max_pools, single_query) - for batch, max_pools in [(1, 128), (3, 256), (16, 384), (17, 640), (127, 768), (128, 1024), (129, 1152)] - } - assert len(hashes) == 1 - - -@pytest.mark.parametrize("single_query", [False, True]) -@pytest.mark.parametrize("batch,max_pools", [(1, 1), (3, 3), (16, 16), (129, 1025)]) -def test_pool_ranges_nonstandard_capacities_preserve_correctness(batch, max_pools, single_query): - _check_pool_ranges(batch, max_pools, single_query) - - -def test_pool_compression_reuses_compilation_within_static_query_layouts(monkeypatch): - torch.manual_seed(1525) - tail = torch.zeros((130, 4, 256), device="cuda", dtype=torch.bfloat16) - ape = torch.zeros((4, 128), device="cuda") - hashes = {} - layout = None - misses = [] - original_run = _compress_pools.run - original_compile = JITFunction._do_compile - - def observe_compile(self, *args, **kwargs): - if self is _compress_pools: - misses.append(layout) - return original_compile(self, *args, **kwargs) - - def observe(*args, **kwargs): - compiled = original_run(*args, **kwargs) - hashes.setdefault(layout, set()).add(compiled.hash) - return compiled - - monkeypatch.setattr(_compress_pools, "run", observe) - monkeypatch.setattr(JITFunction, "_do_compile", observe_compile) - for batch, single_query in [(1, False), (3, True), (33, False), (1, True), (17, True), (1, True)]: - values = torch.randint(-16, 17, (batch, 4, 256), device="cuda").bfloat16() - values[:, :, 128:] = 0 - tail[:batch].zero_() - packed = torch.zeros((batch * 4, 1, 720), device="cuda", dtype=torch.uint8) - requests = torch.arange(batch, device="cuda", dtype=torch.int32) - ragged = torch.arange(batch * 4, device="cuda", dtype=torch.int32) - seq_lens = torch.full((batch,), 4, device="cuda", dtype=torch.int32) - if single_query: - tail[:batch, :3] = values[:, :3] - raw = values[:, 3].contiguous() - lengths = seq_lens - mem_index = requests * 4 + 3 - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) - else: - raw = values.reshape(batch * 4, 256) - lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32).repeat(batch) - mem_index = ragged - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 - # Query mode and address strides stay constexpr for the faster kernel. - layout = (single_query, raw.stride(0)) - seen_layout = layout in hashes - misses_before = len(misses) - compress_pools(raw, tail, packed, ape, lengths, mem_index, requests, cu, seq_lens, 1 if single_query else 4) - if seen_layout: - assert len(misses) == misses_before - keys, scales = gather_pools(packed, ragged.view(batch, 4), requests, seq_lens, 1) - expected = values[:, :, :128].float().mean(1).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected, scale=128 ** -0.5) - assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) - assert torch.equal(scales, expected_scale.flatten()) - assert torch.equal(tail[:batch], values) - assert len(hashes) == 3 - assert all(len(variants) == 1 for variants in hashes.values()) - - -def test_gather_pools_reuses_compilation_across_pool_counts(): - torch.manual_seed(1525) - packed = torch.zeros((16, 1, 720), device="cuda", dtype=torch.uint8) - source_keys = torch.randint(0, 256, (16, 128), device="cuda", dtype=torch.uint8) - source_scales = torch.rand(16, device="cuda") + 0.1 - packed[:, 0, :128] = source_keys - packed[:, 0, 128:132] = source_scales.view(torch.uint8).view(16, 4) - locations = torch.randint(0, 16, (6, 129), device="cuda", dtype=torch.int32) - table = torch.zeros((6, 129 * 4), device="cuda", dtype=torch.int32) - table[:, 3::4] = locations - requests = torch.tensor([2, 0, 5], device="cuda", dtype=torch.int32) - hashes = set() - for pools in [1, 3, 16, 33, 64, 129]: - lengths = torch.tensor([pools * 4, min(3, pools) * 4, 2], device="cuda", dtype=torch.int32) - keys = torch.empty((3, pools, 128), device="cuda", dtype=torch.float8_e4m3fn) - scales = torch.empty((3, pools), device="cuda") - compiled = _gather_pools[(triton.cdiv(pools, 32), 3)]( - packed, - table, - requests, - lengths, - keys, - scales, - packed.stride(0), - table.stride(0), - pools, - POOLS_PER_BLOCK=32, - num_warps=4, - ) - hashes.add(compiled.hash) - for row, (request, count) in enumerate([(2, pools), (0, min(3, pools)), (5, 0)]): - indices = locations[request, :count].long() - assert torch.equal(keys.view(torch.uint8)[row, :count], source_keys[indices]) - assert torch.equal(scales[row, :count], source_scales[indices]) - assert not keys.view(torch.uint8)[row, count:].any() - assert (scales[row, count:] == 1).all() - assert len(hashes) == 1 - - -def test_pool_compression_reuses_compilation_across_request_counts(monkeypatch): - torch.manual_seed(1525) - tail = torch.zeros((130, 4, 256), device="cuda", dtype=torch.bfloat16) - # Exact uniform weights isolate indexing from softmax/FP8 rounding differences. - ape = torch.zeros((4, 128), device="cuda") - hashes = set() - original_run = _compress_pools.run - - def observe(*args, **kwargs): - compiled = original_run(*args, **kwargs) - hashes.add(compiled.hash) - return compiled - - monkeypatch.setattr(_compress_pools, "run", observe) - for batch in [1, 3, 16, 17, 127, 128, 129]: - raw = torch.randint(-16, 17, (batch * 4, 256), device="cuda").bfloat16() - raw[:, 128:] = 0 - packed = torch.zeros((batch * 4, 1, 720), device="cuda", dtype=torch.uint8) - requests = torch.arange(batch, device="cuda", dtype=torch.int32) - lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32).repeat(batch) - ragged = torch.arange(batch * 4, device="cuda", dtype=torch.int32) - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 - seq_lens = torch.full((batch,), 4, device="cuda", dtype=torch.int32) - compress_pools( - raw, - tail, - packed, - ape, - lengths, - ragged, - requests, - cu, - seq_lens, - max_q_seq_len=4, - ) - keys, scales = gather_pools(packed, ragged.view(batch, 4), requests, seq_lens, 1) - values = raw.view(batch, 4, 256) - expected = (values[:, :, :128].float() * (values[:, :, 128:].float() + ape).softmax(1)).sum(1).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected, scale=128 ** -0.5) - assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) - assert torch.equal(scales, expected_scale.flatten()) - assert torch.equal(tail[:batch], values) - assert len(hashes) == 1 diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 1aca878b2d..4af14b0149 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -324,7 +324,7 @@ def test_mtp_pre_layer_applies_main_norm_before_hidden_fusion(config, monkeypatc main_norm_weight_=norms[0], enorm_weight_=norms[1], hnorm_weight_=norms[2], - eh_proj_weight_=SimpleNamespace(mm=lambda x: x @ projection), + eh_proj_weight_=SimpleNamespace(mm=lambda *, input_tensor: input_tensor @ projection), ) state = SimpleNamespace(mtp_draft_input_hiddens=hidden.clone()) monkeypatch.setattr(LlamaMultimodalPreLayerInfer, method, lambda self, ids, state, weight: embeddings.clone()) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 4d15d6bf1c..15dcb4a499 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -304,7 +304,11 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): + from lightllm.common.attention_state import pd_operator + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) + caplog.set_level(30, logger=pd_operator.logger.name) + monkeypatch.setattr(pd_operator.logger, "handlers", [caplog.handler]) mem = make_mems(1)[0] full_bytes = max( layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_full_att_page_layouts().values() diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py index 6a25a8d4ea..e4e38812c4 100644 --- a/unit_tests/models/glm5_next/test_visual.py +++ b/unit_tests/models/glm5_next/test_visual.py @@ -197,20 +197,14 @@ def test_target_and_mtp_prefill_embed_images_before_hidden_fusion(monkeypatch): expanded = Glm5NextPreLayerInfer(config).context_forward(ids, state, weight) torch.testing.assert_close(expanded.reshape(4, 4, hidden), expected[:, None].expand(-1, 4, -1)) - def norm(input, eps, out): - out.copy_(input * torch.rsqrt(input.square().mean(-1, keepdim=True) + eps)) - return out - state.mtp_draft_input_hiddens = torch.randn(4, hidden, device="cuda") - main_norm_weight = torch.randn(hidden, device="cuda") old_hidden = state.mtp_draft_input_hiddens.clone() projection = torch.randn(2 * hidden, hidden, device="cuda") * 0.1 - weight.enorm_weight_ = weight.hnorm_weight_ = norm - weight.main_norm_weight_ = lambda input, eps, out: out.copy_( - torch.nn.functional.rms_norm(input, (hidden,), main_norm_weight, eps) + weight.enorm_weight_ = weight.hnorm_weight_ = weight.main_norm_weight_ = SimpleNamespace( + weight=torch.ones(hidden, device="cuda") ) - weight.eh_proj_weight_ = SimpleNamespace(mm=lambda x: x @ projection) + weight.eh_proj_weight_ = SimpleNamespace(mm=lambda *, input_tensor: input_tensor @ projection) actual = Glm5NextMTPPreLayerInfer(config).context_forward(ids, state, weight) - old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), main_norm_weight, 1e-5) + old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), weight.main_norm_weight_.weight, 1e-5) normalized = [x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) for x in (expected, old_hidden)] torch.testing.assert_close(actual, torch.cat(normalized, -1) @ projection) diff --git a/unit_tests/models/phi3/test_rotary_emb.py b/unit_tests/models/phi3/test_rotary_emb.py deleted file mode 100644 index 3127f110ad..0000000000 --- a/unit_tests/models/phi3/test_rotary_emb.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest -import torch - -from lightllm.models.phi3.triton_kernel.rotary_emb import rotary_emb_fwd - - -def ref_rotary(x, cos, sin, head_dim): - rot = head_dim // 2 - x = x.float() - x1, x2 = x[..., :rot], x[..., rot:head_dim] - c, s = cos.float()[:, None, :], sin.float()[:, None, :] - out = x.clone() - out[..., :rot] = x1 * c - x2 * s - out[..., rot:head_dim] = x1 * s + x2 * c - return out - - -@pytest.mark.parametrize( - "shape", - [ - # (total_len, q_heads, k_heads, head_dim); 96 is Phi-3-mini's head_dim, - # whose rot_dim 48 is not a power of two, so BLOCK_DMODEL has padding - # lanes that every load must mask off - (33, 32, 32, 96), - (128, 32, 8, 80), - (64, 16, 16, 128), - ], -) -@pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.5]) -def test_rotary_matches_reference(shape, partial_rotary_factor): - torch.manual_seed(0) - device = "cuda" - T, HQ, HK, D = shape - rot_dim = int(D * partial_rotary_factor) // 2 - - q = torch.randn(T, HQ, D, device=device, dtype=torch.float16) - k = torch.randn(T, HK, D, device=device, dtype=torch.float16) - pos = torch.arange(T, device=device, dtype=torch.float32) - inv = 1.0 / (10000 ** (torch.arange(0, rot_dim, device=device, dtype=torch.float32) / rot_dim)) - angles = pos[:, None] * inv[None, :] - cos = angles.cos().to(torch.float16).contiguous() - sin = angles.sin().to(torch.float16).contiguous() - - q_ref = ref_rotary(q, cos, sin, int(D * partial_rotary_factor)) - k_ref = ref_rotary(k, cos, sin, int(D * partial_rotary_factor)) - - rotary_emb_fwd(q, k, cos, sin, partial_rotary_factor=partial_rotary_factor) - - assert torch.allclose(q.float(), q_ref, atol=5e-3, rtol=1e-2) - assert torch.allclose(k.float(), k_ref, atol=5e-3, rtol=1e-2) diff --git a/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py b/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py deleted file mode 100644 index 7226b05239..0000000000 --- a/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -import pytest - -from lightllm.models.qwen2_vl.triton_kernel.get_mrope_position_ids import get_mrope_position_triton - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="需要 CUDA 支持") -def test_get_mrope_position_triton(): - """ - 测试 get_mrope_position_triton 函数的正确性 - """ - b_image_start_idx = torch.tensor([0, 0, 4], dtype=torch.int32, device="cuda") - b_image_thwd = torch.tensor([[1, 2, 2, -2], [1, 2, 2, -2], [1, 2, 2, -2]], dtype=torch.int32, device="cuda") - b_image_nums = torch.tensor([1, 2], dtype=torch.int32, device="cuda") - b_image_start_num = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - b_image_len = torch.tensor([4, 4, 4], dtype=torch.int32, device="cuda") - position_ids = ( - torch.tensor([0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=torch.int32, device="cuda") - .unsqueeze(0) - .expand(3, -1) - .contiguous() - ) - b_ready_cache_len = torch.tensor([0, 0], dtype=torch.int32, device="cuda") - b_q_seq_len = torch.tensor([7, 13], dtype=torch.int32, device="cuda") - b_start_loc = torch.tensor([0, 7], dtype=torch.int32, device="cuda") - get_mrope_position_triton( - b_image_start_idx, - b_image_thwd, - b_image_nums, - b_image_start_num, - b_image_len, - position_ids, - b_ready_cache_len, - b_q_seq_len, - b_start_loc, - ) - # 预期的输出结果 - expected_output = torch.tensor( - [ - [0, 0, 0, 0, 2, 3, 4, 0, 0, 0, 0, 2, 2, 2, 2, 4, 5, 6, 7, 8], - [0, 0, 1, 1, 2, 3, 4, 0, 0, 1, 1, 2, 2, 3, 3, 4, 5, 6, 7, 8], - [0, 1, 0, 1, 2, 3, 4, 0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 6, 7, 8], - ], - dtype=torch.int32, - device="cuda", - ) - - assert torch.equal(position_ids, expected_output), "position_ids 输出与预期结果不一致" diff --git a/unit_tests/models/qwen2_vl/test_infer_struct.py b/unit_tests/models/qwen2_vl/test_infer_struct.py deleted file mode 100644 index b239e64067..0000000000 --- a/unit_tests/models/qwen2_vl/test_infer_struct.py +++ /dev/null @@ -1,68 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.models.qwen2_vl.infer_struct import Qwen2VLInferStateInfo - - -def _patch_base_position_ids(monkeypatch): - def init_positions(self, model): - self.position_ids = torch.tensor([9, 19], dtype=torch.int32) - self.b_q_seq_len = torch.ones(2, dtype=torch.int32) - - monkeypatch.setattr(InferStateInfo, "init_some_extra_state", init_positions) - - -def _make_model(): - return SimpleNamespace( - config={"rope_scaling": {}}, - _cos_cached=torch.arange(32, dtype=torch.float32).view(32, 1), - _sin_cached=torch.arange(32, dtype=torch.float32).view(32, 1), - ) - - -def test_normal_prompt_prefill_builds_multimodal_positions(monkeypatch): - _patch_base_position_ids(monkeypatch) - expected_position_ids = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) - monkeypatch.setattr( - Qwen2VLInferStateInfo, - "get_mrope_position", - lambda self, multimodal_params: expected_position_ids, - ) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = True - infer_state.b_position_delta = None - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - infer_state.init_some_extra_state(_make_model()) - - assert torch.equal(infer_state.position_ids, expected_position_ids) - - -def test_normal_decode_applies_position_delta(monkeypatch): - _patch_base_position_ids(monkeypatch) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = False - infer_state.b_position_delta = torch.tensor([3, 5], dtype=torch.int32) - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - infer_state.init_some_extra_state(_make_model()) - - expected_position_ids = torch.tensor([[12, 24]] * 3, dtype=torch.int32) - assert torch.equal(infer_state.position_ids, expected_position_ids) - - -def test_prefill_rejects_position_delta(monkeypatch): - _patch_base_position_ids(monkeypatch) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = True - infer_state.b_position_delta = torch.tensor([3, 5], dtype=torch.int32) - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - with pytest.raises(AssertionError, match="prefill must not provide b_position_delta"): - infer_state.init_some_extra_state(_make_model()) diff --git a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py deleted file mode 100644 index 34e41578a6..0000000000 --- a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py +++ /dev/null @@ -1,55 +0,0 @@ -import math -import torch -import pytest - -from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton - - -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb_vision(tensor: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - orig_dtype = tensor.dtype - tensor = tensor.float() - cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() - sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() - output = (tensor * cos) + (rotate_half(tensor) * sin) - output = output.to(orig_dtype) - return output - - -@pytest.mark.parametrize( - "shape", - [ - (1296, 64, 80), - (1024, 2, 192), - (1024, 1, 256), - (1024, 3, 160), - ], -) -def test_triton_matches_reference(shape): - L, H, D = shape - assert D % 2 == 0 - - torch.manual_seed(0) - - freqs = torch.randn(L, D // 2, device="cuda", dtype=torch.bfloat16) - cos = freqs.cos() - sin = freqs.sin() - - tensor = torch.randn(L, H, D, device="cuda", dtype=torch.bfloat16) - - ref = apply_rotary_pos_emb_vision(tensor.unsqueeze(0), cos, sin).squeeze(0) - out = apply_rotary_pos_emb_triton(tensor, cos, sin) - - assert out.dtype == tensor.dtype, "输出 dtype 应与输入一致" - assert out.shape == tensor.shape, "输出形状应与输入一致" - assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致" - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/qwen3next/test_shared_expert_gate.py b/unit_tests/models/qwen3next/test_shared_expert_gate.py deleted file mode 100644 index 44d46e4471..0000000000 --- a/unit_tests/models/qwen3next/test_shared_expert_gate.py +++ /dev/null @@ -1,22 +0,0 @@ -import pytest -import torch - -from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ - - -@pytest.mark.parametrize("gate_width", [1, 6144]) -def test_sigmoid_mul_matches_unfused_dtype_rounding(gate_width): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required") - - torch.manual_seed(2026) - x = torch.randn((17, 6144), device="cuda", dtype=torch.bfloat16) * 2 - gate = torch.randn((17, gate_width), device="cuda", dtype=torch.bfloat16) * 3 - - expected = x.clone() - expected.mul_(gate.clone().sigmoid_()) - - actual = x.clone() - sigmoid_mul_(actual, gate) - - torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/models/test_multimodal_tokenizer_positions.py b/unit_tests/models/test_multimodal_tokenizer_positions.py deleted file mode 100644 index 9425e02587..0000000000 --- a/unit_tests/models/test_multimodal_tokenizer_positions.py +++ /dev/null @@ -1,98 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.models.gemma3.model import Gemma3Tokenizer -from lightllm.models.internvl.model import InternvlTokenizer -from lightllm.models.llava.model import LlavaTokenizer -from lightllm.models.qwen3_omni_moe_thinker.model import QWen3OmniTokenizer -from lightllm.models.qwen_vl.model import QWenVLTokenizer -from lightllm.models.tarsier2.model import Tarsier2Tokenizer -from lightllm.server.multimodal_params import MultimodalParams -from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans - - -def make_params(images=(), audios=()): - params = MultimodalParams( - images=[{"type": "url", "data": "unused"} for _ in images], - audios=[{"type": "url", "data": "unused"} for _ in audios], - ) - for items, specs in ((params.images, images), (params.audios, audios)): - for item, (token_id, token_num) in zip(items, specs): - item.token_id, item.token_num = token_id, token_num - return params - - -def make_tokenizer(tokenizer_class, origin_ids): - tokenizer = tokenizer_class.__new__(tokenizer_class) - tokenizer.tokenizer = Mock() - tokenizer.tokenizer.encode.return_value = origin_ids - tokenizer.image_start_id, tokenizer.image_end_id, tokenizer.image_token_id = 10, 11, 12 - tokenizer.audio_start_id, tokenizer.audio_end_id, tokenizer.audio_token_id = 20, 21, 22 - return tokenizer - - -@pytest.mark.parametrize( - "tokenizer_class", [LlavaTokenizer, Gemma3Tokenizer, QWenVLTokenizer, Tarsier2Tokenizer, InternvlTokenizer] -) -def test_images_record_positions_even_when_embeddings_are_reused(tokenizer_class): - tokenizer = make_tokenizer(tokenizer_class, [1, 10, 11, 2, 10, 11, 3]) - tokenizer.tokenizer.side_effect = [SimpleNamespace(input_ids=[token]) for token in (1, 2, 3)] - tokenizer.tokenizer.tokenizer = tokenizer.tokenizer - tokenizer.image_token = tokenizer.boi_token = "" - tokenizer.image_length = tokenizer.mm_tokens_per_image = 2 - tokenizer.boi_token_index, tokenizer.eoi_token_index = 10, 11 - tokenizer.image_start_tag, tokenizer.image_end_tag = "", "" - tokenizer.skip_start = False - params = make_params(images=[(100, 2), (100, 2)]) - - prompt = "abc" if tokenizer_class is QWenVLTokenizer else "abc" - kwargs = {"add_special_tokens": False} if tokenizer_class is InternvlTokenizer else {} - input_ids = tokenizer.encode(prompt, params, **kwargs) - - if tokenizer_class is LlavaTokenizer: - assert input_ids == [1, 100, 101, 2, 100, 101, 3] - expected_spans = [(1, 3), (4, 6)] - else: - assert input_ids == [1, 10, 100, 101, 11, 2, 10, 100, 101, 11, 3] - expected_spans = [(2, 4), (7, 9)] - assert get_mm_input_spans(input_ids, params.to_dict()) == expected_spans - - -@pytest.mark.parametrize("tokenizer_class", [InternvlTokenizer, QWen3OmniTokenizer]) -@pytest.mark.parametrize( - "origin_ids,image_starts,audio_starts,expected_ids", - [ - ( - [1, 20, 21, 2, 10, 11, 3, 20, 21, 4, 10, 11, 5], - [7, 19], - [2, 13], - [1, 20, 200, 201, 21, 2, 10, 100, 101, 102, 11, 3, 20, 300, 301, 302, 21, 4, 10, 100, 101, 102, 11, 5], - ), - ( - [1, 10, 11, 2, 20, 21, 3, 10, 11, 4, 20, 21, 5], - [2, 13], - [8, 19], - [1, 10, 100, 101, 102, 11, 2, 20, 200, 201, 21, 3, 10, 100, 101, 102, 11, 4, 20, 300, 301, 302, 21, 5], - ), - ], -) -def test_positions_refer_to_final_input_after_both_modalities_expand( - tokenizer_class, origin_ids, image_starts, audio_starts, expected_ids -): - tokenizer = make_tokenizer(tokenizer_class, origin_ids) - params = make_params(images=[(100, 3), (100, 3)], audios=[(200, 2), (300, 3)]) - - input_ids = tokenizer.encode("mixed image/audio input", params, add_special_tokens=False) - - assert input_ids == expected_ids - assert [image.start_idx for image in params.images] == image_starts - assert [audio.start_idx for audio in params.audios] == audio_starts - spans = get_mm_input_spans(input_ids, params.to_dict()) - assert len(spans) == 4 - for item in params.images + params.audios: - assert (item.start_idx, item.start_idx + item.token_num) in spans - assert input_ids[item.start_idx : item.start_idx + item.token_num] == list( - range(item.token_id, item.token_id + item.token_num) - ) diff --git a/unit_tests/models/test_neo_registry.py b/unit_tests/models/test_neo_registry.py deleted file mode 100644 index ab3efee352..0000000000 --- a/unit_tests/models/test_neo_registry.py +++ /dev/null @@ -1,30 +0,0 @@ -import sys -from types import ModuleType - -import pytest - -from lightllm import models -from lightllm.models.registry import ModelRegistry - - -@pytest.mark.parametrize( - "llm_type,module_name,class_name", - [ - ("qwen3", "lightllm.models.neo_chat.model", "NeoTpPartModel"), - ("qwen3_moe", "lightllm.models.neo_chat_moe.model", "NeoTpMOEPartModel"), - ], -) -def test_neo_models_resolve_through_lazy_registry(monkeypatch, llm_type, module_name, class_name): - config = {"model_type": "neo_chat", "llm_config": {"model_type": llm_type}} - registration = ModelRegistry.get_model_config(config) - assert registration.is_multimodal - assert registration.model_class == f"{module_name}:{class_name}" - - implementation = ModuleType(module_name) - model_class = type(class_name, (), {}) - setattr(implementation, class_name, model_class) - monkeypatch.setitem(sys.modules, module_name, implementation) - - assert models.get_model_class(config) is model_class - assert getattr(models, class_name) is model_class - assert models.get_model_class(config) is model_class diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index e020c6709b..0e229938a4 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -78,72 +78,6 @@ def candidates(*args): assert prompt.logits_token_ids is None -def test_vocab_parallel_topk_returns_candidates_on_single_rank(monkeypatch) -> None: - _mock_vocab_parallel_pack(monkeypatch) - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 1 - head.alloc_tensor = lambda shape, dtype, **kwargs: torch.empty(shape, dtype=dtype) - local_logits = torch.tensor( - [ - [1.0, 8.0], - [7.0, 2.0], - [3.0, 6.0], - [5.0, 4.0], - ] - ) - head._project_local_logits = lambda *args: local_logits - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=4, tp_vocab_start_id=0)) - - logits, token_ids = head._vocab_parallel_topk(None, 2, weight, SimpleNamespace(dist_group=None), 2) - - expected_logits, expected_token_ids = torch.topk(local_logits, k=2, dim=0, sorted=False) - assert logits.dtype == torch.float32 and token_ids.dtype == torch.int64 - assert token_ids.is_contiguous() - torch.testing.assert_close(logits, expected_logits.permute(1, 0)) - torch.testing.assert_close(token_ids, expected_token_ids.permute(1, 0)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("vocab_size, token_num, top_k", [(257, 15, 512), (257, 65, 64), (248320, 64, 64)]) -def test_single_rank_vocab_topk_graph_preserves_output_contract(vocab_size, token_num, top_k) -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 1 - local_logits = torch.randn((vocab_size, token_num), dtype=torch.bfloat16, device="cuda") - head._project_local_logits = lambda *args: local_logits - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=vocab_size, tp_vocab_start_id=0)) - state = SimpleNamespace(dist_group=None) - - for _ in range(3): - head._vocab_parallel_topk(None, token_num, weight, state, top_k) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - values, token_ids = head._vocab_parallel_topk(None, token_num, weight, state, top_k) - - local_logits.normal_() - graph.replay() - torch.cuda.synchronize() - candidate_count = min(top_k, vocab_size) - expected_values = torch.topk(local_logits.T, k=candidate_count, dim=1, sorted=True).values.float() - assert values.shape == token_ids.shape == (token_num, candidate_count) - assert values.dtype == torch.float32 and token_ids.dtype == torch.int64 - assert token_ids.is_contiguous() - assert torch.all((token_ids >= 0) & (token_ids < vocab_size)) - assert torch.all(token_ids.sort(dim=1).values.diff(dim=1) > 0) - torch.testing.assert_close(values, local_logits.T.gather(1, token_ids).float(), rtol=0, atol=0) - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - -def test_vocab_parallel_topk_requires_enough_local_vocabulary() -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 2 - head._project_local_logits = lambda *args: torch.ones(2, 1) - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=6, tp_vocab_start_id=0)) - - with pytest.raises(AssertionError, match="local vocabulary size 2 must be at least the candidate count 3"): - head._vocab_parallel_topk(None, 1, weight, SimpleNamespace(dist_group=None), 3) - - def test_vocab_parallel_topk_merges_rank_candidates(monkeypatch) -> None: _mock_vocab_parallel_pack(monkeypatch) head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) @@ -210,32 +144,6 @@ def test_vocab_parallel_topk_pack_round_trip_preserves_token_id_bits(transposed) torch.testing.assert_close(token_ids, expected_token_ids) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("shape", [(129, 15), (129, 16), (129, 17), (62080, 256)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) -def test_local_vocab_topk_preserves_values_and_token_ids(shape, dtype) -> None: - logits = torch.randn(shape, dtype=dtype, device="cuda") - values, token_ids = local_vocab_topk(logits, top_k=64) - expected_values = torch.topk(logits, k=64, dim=0, sorted=True).values.T - - assert values.shape == token_ids.shape == (shape[1], 64) - assert values.dtype == dtype and token_ids.dtype == torch.int64 - assert torch.all((token_ids >= 0) & (token_ids < shape[0])) - assert torch.all(token_ids.sort(dim=1).values.diff(dim=1) > 0) - torch.testing.assert_close(values, logits.T.gather(1, token_ids), rtol=0, atol=0) - # Equal-valued boundary candidates can have different valid IDs/order. - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_local_vocab_topk_accepts_noncontiguous_logits() -> None: - logits = torch.randn((65, 129), dtype=torch.bfloat16, device="cuda").T - values, token_ids = local_vocab_topk(logits, top_k=8) - expected_values = torch.topk(logits, k=8, dim=0, sorted=True).values.T - torch.testing.assert_close(values, logits.T.gather(1, token_ids), rtol=0, atol=0) - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("top_k", [8, 64, 512]) def test_local_vocab_topk_graph_replay_reads_updated_logits(top_k) -> None: @@ -286,17 +194,6 @@ def test_prefill_prompt_logits_do_not_replace_candidate_output(draft) -> None: assert state.prompt_logics is prompt -def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: - model = TpPartBaseModel.__new__(TpPartBaseModel) - logits = torch.randn(3, 2) - ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) - state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None) - output = model._create_model_output(PostLayerOutput(logits, ids), state) - unpadded = model._create_unpad_decode_model_output(output, 2) - torch.testing.assert_close(unpadded.logits, logits[:2]) - torch.testing.assert_close(unpadded.logits_token_ids, ids[:2]) - - def test_gemma_softcap_preserves_candidate_ids(monkeypatch) -> None: ids = torch.tensor([[10, 20]]) logits = torch.tensor([[1.0, 4.0]]) @@ -307,12 +204,3 @@ def test_gemma_softcap_preserves_candidate_ids(monkeypatch) -> None: output = head.token_forward(None, state, None) torch.testing.assert_close(output.logits, torch.tanh(logits / 2) * 2) assert output.logits_token_ids is ids - - -def test_overlap_keeps_each_microbatch_mapping() -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - outputs = [PostLayerOutput(torch.ones(1, 2), torch.tensor([[10, 20]])), PostLayerOutput(torch.ones(1, 3))] - head.token_forward = lambda hidden, state, layer_weight: outputs[state] - output0, output1 = head.overlap_tpsp_token_forward(None, None, 0, 1, None) - assert output0 is outputs[0] - assert output1 is outputs[1] diff --git a/unit_tests/models/test_qwen_vl_tokenizer.py b/unit_tests/models/test_qwen_vl_tokenizer.py deleted file mode 100644 index b1d8d87fb0..0000000000 --- a/unit_tests/models/test_qwen_vl_tokenizer.py +++ /dev/null @@ -1,69 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3_vl.model import QWen3VLTokenizer -from lightllm.utils.error_utils import InvalidRequestError - - -@pytest.fixture(params=[QWen2VLTokenizer, QWen3VLTokenizer]) -def tokenizer(request): - instance = request.param.__new__(request.param) - instance.tokenizer = Mock() - instance.image_start_id = 10 - instance.image_end_id = 11 - instance.image_token_id = 12 - return instance - - -def test_image_expansion(tokenizer): - images = [SimpleNamespace(token_id=100, token_num=2), SimpleNamespace(token_id=200, token_num=3)] - prompt = [1, 10, 12, 11, 2, 10, 11, 3] - assert tokenizer.encode(prompt, SimpleNamespace(images=images)) == [ - 1, - 10, - 100, - 101, - 11, - 2, - 10, - 200, - 201, - 202, - 11, - 3, - ] - assert [image.start_idx for image in images] == [2, 7] - assert prompt == [1, 10, 12, 11, 2, 10, 11, 3] - assert tokenizer.encode(prompt, None) == [1, 10, 11, 2, 10, 11, 3] - - -@pytest.mark.parametrize("prompt,image_count", [([10, 12, 11], 0), ([10, 11, 10, 11], 1), ([1], 1)]) -def test_image_count_mismatch(tokenizer, prompt, image_count): - images = [SimpleNamespace(token_id=100, token_num=2) for _ in range(image_count)] - with pytest.raises(InvalidRequestError, match="invalid image tag num"): - tokenizer.encode(prompt, SimpleNamespace(images=images)) - assert all(not hasattr(image, "start_idx") for image in images) - - -@pytest.mark.parametrize("prompt", [[10], [10, 12], [10, 1, 11], [10, 10, 11]]) -@pytest.mark.parametrize("params", [None, SimpleNamespace(images=[])]) -def test_invalid_image_sequence(tokenizer, prompt, params): - with pytest.raises(InvalidRequestError, match="image"): - tokenizer.encode(prompt, params) - - -def test_text_and_encode_options(tokenizer): - tokenizer.tokenizer.encode.return_value = [1, 12, 2] - assert tokenizer.encode("text", SimpleNamespace(images=[]), add_special_tokens=False) == [1, 12, 2] - tokenizer.tokenizer.encode.assert_called_once_with("text", add_special_tokens=False) - - -def test_unpaired_image_end_tokens(tokenizer): - assert tokenizer.encode([1, 11, 2], SimpleNamespace(images=[])) == [1, 11, 2] - images = [SimpleNamespace(token_id=100, token_num=2), SimpleNamespace(token_id=200, token_num=1)] - prompt = [11, 1, 10, 12, 11, 11, 2, 10, 11, 11] - assert tokenizer.encode(prompt, SimpleNamespace(images=images)) == [11, 1, 10, 100, 101, 11, 11, 2, 10, 200, 11, 11] - assert tokenizer.encode(prompt, None) == [11, 1, 10, 11, 11, 2, 10, 11, 11] diff --git a/unit_tests/server/core/objs/test_sampling_params.py b/unit_tests/server/core/objs/test_sampling_params.py index b78608bf55..f8027b0fec 100644 --- a/unit_tests/server/core/objs/test_sampling_params.py +++ b/unit_tests/server/core/objs/test_sampling_params.py @@ -37,48 +37,6 @@ }""" -@pytest.mark.parametrize( - "sequence, expected", - [ - ([1, 2, 3], [1, 2, 3]), - ([1] * (STOP_SEQUENCE_MAX_LENGTH), [1] * STOP_SEQUENCE_MAX_LENGTH), - ], -) -def test_stop_sequence_initialization(sequence, expected): - seq = StopSequence() - seq.initialize(sequence) - assert seq.size == len(expected) - assert seq.to_list() == expected - - -def test_stop_sequence_initialization_too_many(): - seq = StopSequence() - with pytest.raises(AssertionError): - seq.initialize([1] * (STOP_SEQUENCE_MAX_LENGTH + 1)) - - -@pytest.mark.parametrize( - "stop_sequences, expected_ids, expected_strings, expected_original", - [ - (["stop1", "stop2"], [], ["stop1", "stop2"], ["stop1", "stop2"]), - ("END", [], ["END"], ["END"]), - ([[1, 2], [3, 4]], [[1, 2], [3, 4]], [], [[1, 2], [3, 4]]), - (["END", [1, 2], "结束"], [[1, 2]], ["END", "结束"], ["END", [1, 2], "结束"]), - (["", [], "END", [1]], [[1]], ["END"], ["END", [1]]), - ([], [], [], []), - (None, [], [], []), - ], -) -def test_stop_sequence_groups_initialization(stop_sequences, expected_ids, expected_strings, expected_original): - groups = StopSequenceGroups() - groups.initialize(stop_sequences) - assert groups.size == len(expected_original) - assert groups.to_token_ids_list() == expected_ids - assert groups.to_stop_str_list() == expected_strings - assert groups.has_stop_strs() is bool(expected_strings) - assert groups.to_origin_list() == expected_original - - @pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) def test_string_stops_are_not_encoded_and_round_trip_with_token_ids(mode): tokenizer = Mock() @@ -101,53 +59,6 @@ def test_string_stops_are_not_encoded_and_round_trip_with_token_ids(mode): tokenizer.encode.assert_not_called() -def test_regular_constraint_initialization(): - constraint = RegularConstraint() - constraint.initialize("[a-zA-Z]+") - assert constraint.length == len("[a-zA-Z]+") - assert constraint.to_str() == "[a-zA-Z]+" - - with pytest.raises(AssertionError): - constraint.initialize("a" * (REGULAR_CONSTRAINT_MAX_LENGTH + 1)) - - -def test_guided_grammar_initialization(): - grammar = GuidedGrammar() - grammar.initialize(grammar_str, None) - assert grammar.to_str() == grammar_str - - with pytest.raises(AssertionError): - grammar.initialize("a" * (GRAMMAR_CONSTRAINT_MAX_LENGTH + 1), None) - - -def test_guided_json_schema_initialization(): - schema = GuidedJsonSchema() - schema.initialize(schema_str, None) - assert schema.to_str() == schema_str - - with pytest.raises(AssertionError): - schema.initialize("a" * (JSON_SCHEMA_MAX_LENGTH + 1), None) - - -def test_exponential_decay_length_penalty_initialization(): - penalty = ExponentialDecayLengthPenalty() - penalty.initialize((5, 1.5)) - assert penalty.to_tuple() == (5, 1.5) - - with pytest.raises(AssertionError): - penalty.initialize((5, 0.5)) - - -def test_node_uuid_initialization(): - node_id = 12345678901234567890 - node_uuid = NodeUUId() - node_uuid.initialize(node_id) - - assert node_uuid.node_id_high == (node_id >> 64) & 0xFFFFFFFFFFFFFFFF - assert node_uuid.node_id_low == node_id & 0xFFFFFFFFFFFFFFFF - assert node_uuid.get() == node_id - - def test_sampling_params_initialization(): params = SamplingParams() pd_master_node_id = 12345678901234567890 @@ -238,15 +149,6 @@ def test_stop_matching_policy_round_trip(kwargs, expected): assert py_restored.stop_str_match_mode.get_mode() == expected -@pytest.mark.parametrize("mode", [-1, 0, 3, 2 ** 32, "CONTENT_ONLY", None, 0.0]) -def test_invalid_stop_str_match_mode(mode): - params = SamplingParams() - with pytest.raises((ValueError, TypeError)): - params.init(None, stop_str_match_mode=mode) - with pytest.raises((ValueError, TypeError)): - PySamplingParams(stop_str_match_mode=mode) - - def test_stop_str_match_mode_updates_shared_buffer(): params = SamplingParams() params.init(None, include_stop_str_in_output=True, _initial_reasoning_state=1) @@ -264,12 +166,6 @@ def test_stop_str_match_mode_updates_shared_buffer(): assert reader.stop_str_match_mode.is_content_only() -def test_infer_high_priority_is_internal_and_ignores_external_value(): - params = SamplingParams() - params.init(None, infer_high_priority=-3) - assert params.infer_high_priority == 0 - - # Mock tokenizer for testing class MockTokenizer: def encode(self, text, add_special_tokens=False): diff --git a/unit_tests/server/core/objs/test_shm_req_manager.py b/unit_tests/server/core/objs/test_shm_req_manager.py index c51e3e6b95..3b59f7d7ca 100644 --- a/unit_tests/server/core/objs/test_shm_req_manager.py +++ b/unit_tests/server/core/objs/test_shm_req_manager.py @@ -1,7 +1,6 @@ import os import asyncio import pytest -import time import uuid import numpy as np from unittest.mock import MagicMock @@ -96,18 +95,6 @@ def test_alloc_req_index_no_available(shm_req_manager): assert index is None # No indices should be available -def test_performance_alloc_release(shm_req_manager): - start_time = time.time() - - for _ in range(100): - index = shm_req_manager.alloc_req_index() - if index is not None: - shm_req_manager.release_req_index(index) - - duration = time.time() - start_time - print(f"Time taken for 100 alloc/release cycles: {duration:.6f} seconds") - - @pytest.fixture def batch_req_manager(monkeypatch): service_name = f"test_batch_req_manager_{uuid.uuid4().hex}" diff --git a/unit_tests/server/core/objs/test_stop_sequence_output.py b/unit_tests/server/core/objs/test_stop_sequence_output.py index 4e0d514c63..d935550dc7 100644 --- a/unit_tests/server/core/objs/test_stop_sequence_output.py +++ b/unit_tests/server/core/objs/test_stop_sequence_output.py @@ -210,170 +210,6 @@ async def stop_request(req_id): assert not buffer.ready_token_list -@pytest.mark.parametrize("include", [False, True]) -def test_only_possible_stop_prefix_delays_whole_token(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(1, "hello", 1)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("hello", 1)] - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(2, "worldE", 2)) - ready_tokens = _drain_ready_tokens(buffer) - assert not ready_tokens - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(3, "X", 3)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("worldE", 2), ("X", 3)] - assert buffer.state is buffer.State.SCANNING - - -@pytest.mark.parametrize("include", [False, True]) -def test_releasing_one_token_preserves_next_token_before_stop(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - _handle_output(buffer, *_output(1, "aE", "first")) - assert not buffer.ready_token_list - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(2, "bE", "second")) - released_tokens = _drain_ready_tokens(buffer) - assert released_tokens == [("aE", "first")] - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(3, "NDextra", "third")) - stopped_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert stopped_tokens == [("bE" if include else "b", "second"), ("ND" if include else "", "third")] - assert released_tokens == [("aE", "first")] - - -@pytest.mark.parametrize("include", [False, True]) -def test_regex_partial_match_keeps_longest_tail_and_releases_safe_tokens(include): - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, _params(["abcd", "bcx"], include), []) - _handle_output(buffer, *_output(1, "safe", "first")) - assert _drain_ready_tokens(buffer) == [("safe", "first")] - _handle_output(buffer, *_output(2, "a", "second")) - assert not buffer.ready_token_list - _handle_output(buffer, *_output(3, "bc", "third")) - assert not buffer.ready_token_list - assert buffer._pending_text == "abc" - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(4, "xab", "fourth")) - # 后一个 stop 已完整命中，即使文本尾部仍是前一个 stop 的前缀，也应立即结束。 - assert _drain_ready_tokens(buffer) == [ - ("a", "second"), - ("bc" if include else "", "third"), - ("x" if include else "", "fourth"), - ] - assert buffer.state is buffer.State.STOPPED - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_unicode_stop_matching_does_not_normalize_text(include, mode): - params = _params("é", include) - params.stop_str_match_mode.set_mode(mode) - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, params, []) - _handle_output(buffer, *_output(1, "e", "first")) - assert _drain_ready_tokens(buffer) == [("e", "first")] - _handle_output(buffer, *_output(2, "\u0301", "last", finish=FinishStatus.FINISHED_LENGTH)) - assert buffer.ready_token_list[-1][3].is_finished_length() - assert _drain_ready_tokens(buffer) == [("\u0301", "last")] - assert buffer.state is buffer.State.STOPPED - manager.stop_request.assert_not_awaited() - - -def test_none_token_resets_text_matching_without_simulated_flag(): - buffer = StopStrMatchBuffer(_http_manager(), _params(), []) - _handle_output(buffer, *_output(10, "helloE", "first")) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(None, "N", "marker")) - marker_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(11, "D", "last")) - last_tokens = _drain_ready_tokens(buffer) - assert first_tokens + marker_tokens + last_tokens == [("helloE", "first"), ("N", "marker"), ("D", "last")] - _handle_output(buffer, *_output(10, "E", "new first")) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(11, "ND", "new last")) - assert buffer.ready_token_list[-1][3].is_stopped() - - -@pytest.mark.parametrize( - "token_id,finish", - [(None, FinishStatus.NO_FINISH), (1001, FinishStatus.FINISHED_ERROR), (1001, FinishStatus.FINISHED_ABORTED)], -) -def test_non_real_token_does_not_leave_reasoning(token_id, finish): - buffer = _stop_buffer() - _handle_output(buffer, *_output(token_id, "", "marker", finish=finish)) - marker_tokens = _drain_ready_tokens(buffer) - assert marker_tokens == [("", "marker")] - assert buffer._reasoning_state.in_reasoning - if finish == FinishStatus.NO_FINISH: - _handle_output(buffer, *_output(10, "END", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("END", "thought")] - else: - assert buffer.state is buffer.State.STOPPED - - -@pytest.mark.parametrize("include", [False, True]) -def test_unmatched_tail_flushes(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - _handle_output(buffer, *_output(1, "helloE", 1)) - assert not buffer.ready_token_list - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(2, "N", 2, finish=FinishStatus.FINISHED_LENGTH)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("helloE", 1), ("N", 2)] - assert buffer.state is buffer.State.STOPPED - - -def test_no_string_stops_emit_immediately(): - buffer = StopStrMatchBuffer(_http_manager(), _params([[10, 11]]), []) - _handle_output(buffer, *_output(1, "hello", 1)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("hello", 1)] - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("alphabet", ["ab结束", "ab结束🙂.*[]\\\n"]) -def test_string_stop_matches_reference_for_random_token_boundaries(include, alphabet): - async def run(): - rng = random.Random(20261011) - for _ in range(300): - stops = ["".join(rng.choices(alphabet, k=rng.randint(1, 5))) for _ in range(rng.randint(1, 4))] - chunks = ["".join(rng.choices(alphabet, k=rng.randint(0, 4))) for _ in range(20)] - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, _params(stops, include), []) - text = "" - expected_text = None - expected_count = len(chunks) - outputs = [] - for index, chunk in enumerate(chunks): - text += chunk - if expected_text is None: - hits = [(text.find(stop), -len(stop), stop) for stop in stops if stop in text] - if hits: - position, _, stop = min(hits) - expected_text = text[: position + (len(stop) if include else 0)] - expected_count = index + 1 - finish = FinishStatus.FINISHED_LENGTH if index == len(chunks) - 1 else FinishStatus.NO_FINISH - await buffer.handle(*_output(index, chunk, index, finish)) - outputs.extend(_drain_ready_tokens(buffer)) - assert "".join(chunk for chunk, _ in outputs) == (text if expected_text is None else expected_text) - assert [payload for _, payload in outputs] == list(range(expected_count)) - assert buffer.state is buffer.State.STOPPED - assert manager.stop_request.await_count == int(expected_text is not None and expected_count < len(chunks)) - - asyncio.run(run()) - - @pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) def test_token_id_stop_is_wrapped_without_text_buffering_even_in_reasoning(mode): params = _params([[10, 11]]) @@ -402,205 +238,6 @@ async def run(): manager.stop_request.assert_awaited_once_with(80) -@pytest.mark.parametrize( - "stop,chunks,expected_text", - [ - ("END", [(1, "E"), (2, "E"), (3, "ND")], "E"), - ], -) -def test_waiting_tracks_overlapping_stop_prefixes(stop, chunks, expected_text): - buffer = StopStrMatchBuffer(_http_manager(), _params(stop), []) - outputs = [] - for index, (token_id, text) in enumerate(chunks): - _handle_output(buffer, *_output(token_id, text, index)) - ready_tokens = _drain_ready_tokens(buffer) - outputs.extend(ready_tokens) - if index < len(chunks) - 1: - assert buffer.state is buffer.State.WAITING - else: - assert buffer.state is buffer.State.STOPPED - assert buffer.state is buffer.State.STOPPED - assert "".join(text for text, _ in outputs) == expected_text - assert [payload for _, payload in outputs] == list(range(len(chunks))) - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_token_id_stop_does_not_enter_text_waiting(mode): - params = _params([[10, 11]]) - params.stop_str_match_mode.set_mode(mode) - buffer = StopStrMatchBuffer(_http_manager(), params, []) - _handle_output(buffer, *_output(10, "first", 1)) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(11, "last", 2, finish=FinishStatus.FINISHED_LENGTH)) - last_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert first_tokens + last_tokens == [("first", 1), ("last", 2)] - - -@pytest.mark.parametrize("stop_in_reasoning,stop_at", [(False, 3), (True, 0)]) -@pytest.mark.parametrize("hidden_delimiter", [False, True]) -def test_reasoning_stop_matches_only_selected_scope(stop_in_reasoning, stop_at, hidden_delimiter): - buffer = _stop_buffer(stop_in_reasoning) - chunks = [(10, "END"), (1001, "" if hidden_delimiter else ""), (11, "answerE"), (12, "NDextra")] - results = [] - for index, (token_id, text) in enumerate(chunks): - _handle_output(buffer, *_output(token_id, text, None)) - ready_tokens = _drain_ready_tokens(buffer) - results.extend(ready_tokens) - assert (buffer.state is buffer.State.STOPPED) is (index == stop_at) - if buffer.state is buffer.State.STOPPED: - break - assert "".join(text for text, _ in results) == ( - "END" + ("" if hidden_delimiter else "") + "answer" if stop_at == 3 else "" - ) - - -@pytest.mark.parametrize("include", [False, True]) -def test_reasoning_eos_preserves_ignored_stop_text(include): - buffer = _stop_buffer(include=include) - _handle_output(buffer, *_output(10, "reasoning END", None)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("reasoning END", None)] - - -@pytest.mark.parametrize("include", [False, True]) -def test_scoped_match_filters_content_and_keeps_reasoning(include): - buffer = _stop_buffer(include=include) - tokens = [] - for token_id, text in [(10, "END"), (1001, ""), (11, "answerENDextra")]: - _handle_output(buffer, *_output(token_id, text, None)) - tokens.extend(_drain_ready_tokens(buffer)) - assert "".join(text for text, _ in tokens) == ("ENDanswer" + ("END" if include else "")) - - -def test_stop_cannot_cross_reasoning_content_boundary(): - buffer = _stop_buffer() - tokens = [] - for token_id, text in [(1, "E"), (1001, ""), (2, "ND")]: - _handle_output(buffer, *_output(token_id, text, None)) - assert buffer.state is buffer.State.SCANNING - tokens.extend(_drain_ready_tokens(buffer)) - assert "".join(text for text, _ in tokens) == "END" - - -def test_entering_reasoning_releases_and_resets_pending_string_stop(): - buffer = _stop_buffer(force=0) - _handle_output(buffer, *_output(10, "helloE", "first")) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(1000, "", "start")) - start_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - assert first_tokens + start_tokens == [("helloE", "first"), ("", "start")] - _handle_output(buffer, *_output(11, "ND", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("ND", "thought")] - _handle_output(buffer, *_output(1001, "", "end")) - assert not buffer._reasoning_state.in_reasoning - _handle_output(buffer, *_output(10, "E", "answer first")) - _handle_output(buffer, *_output(11, "ND", "answer last")) - assert buffer.ready_token_list[-1][3].is_stopped() - - -def test_text_buffer_does_not_match_token_id_stops_in_any_phase(): - buffer = _stop_buffer(stop=[[10, 11]]) - for token_id in [10, 11, 1001, 10, 11]: - _handle_output(buffer, *_output(token_id, "", None)) - assert buffer.state is buffer.State.SCANNING - assert all(not output[3].is_finished() for output in buffer.ready_token_list) - - -def test_pd_first_token_updates_reasoning_state(): - buffer = _stop_buffer() - _handle_output(buffer, *_output(1001, "", None)) - assert not buffer._reasoning_state.in_reasoning - _handle_output(buffer, *_output(10, "END", None)) - assert buffer.ready_token_list[-1][3].is_stopped() - - -def test_prompt_start_marker_prevents_stop_inside_reasoning(): - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - params = _params(force=0) - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - buffer = StopStrMatchBuffer(manager, params, [1, 1000]) - _handle_output(buffer, *_output(10, "END", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("END", "thought")] - _handle_output(buffer, *_output(1001, "", "end")) - assert _drain_ready_tokens(buffer) == [("", "end")] - _handle_output(buffer, *_output(11, "answerEND", "answer")) - answer_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert answer_tokens == [("answer", "answer")] - - -def test_end_marker_outside_reasoning_can_match_stop(): - buffer = _stop_buffer(stop="", force=0) - payload = object() - _handle_output(buffer, *_output(1001, "", payload)) - ready_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert ready_tokens == [("", payload)] - - -def test_pd_capacity_marker_preserves_reasoning_state(): - buffer = _stop_buffer() - _handle_output(buffer, *_output(None, "", None)) - # 模拟 token 的 不结束思考，后续 END 仍属于 reasoning。 - _handle_output(buffer, *_output(10, "END", None)) - assert buffer._reasoning_state.in_reasoning - assert buffer.state is buffer.State.SCANNING - assert _drain_ready_tokens(buffer) == [("", None), ("END", None)] - - -@pytest.mark.parametrize( - "finish", - [ - FinishStatus.FINISHED_LENGTH, - FinishStatus.FINISHED_STOP, - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - FinishStatus.FINISHED_PD_DECODE_CAPACITY, - ], -) -def test_detokenization_forwards_raw_tokens_before_release(monkeypatch, finish): - tokens = list(range(10, 30)) - params = _params() - params.stop_str_match_mode.set_mode(StopStrMatchMode.ALL_TEXT) - decode = DecodeReq(_req(tokens, finish, params)) - simulated = finish in ( - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - FinishStatus.FINISHED_PD_DECODE_CAPACITY, - ) - chunks = {tokens[-2]: "helloE", tokens[-1]: "" if simulated else "NDextra"} - monkeypatch.setattr( - "lightllm.server.detokenization.manager.decode_token", - lambda tokenizer, req, token_id, eos: chunks.get(token_id, ""), - ) - manager = object.__new__(DeTokenizationManager) - manager.req_id_to_out = {80: decode} - manager.tokenizer = _ReasoningTokenizer() - manager.eos_id = [] - manager.all_special_ids = set() - manager.pub_to_httpserver = SimpleNamespace(send_pyobj=lambda *a, **kw: None) - manager.shm_req_manager = SimpleNamespace(put_back_req_obj=lambda req: None) - results = [] - for _ in range(100): - manager.gen_token_out() - if decode.out_queue_is_full() and decode.need_detoken(): - assert not decode.req.can_released_mark - while not decode.req.out_tokens_queue.is_empty(): - results.append(decode.req.out_tokens_queue.pop()) - if decode.req.can_released_mark: - break - assert decode.req.can_released_mark - assert "".join(text for text, *_ in results) == ("helloE" if simulated else "helloENDextra") - assert [count for *_, count in results] == list(range(1, len(tokens) + 1)) - - @pytest.mark.parametrize("include", [False, True]) @pytest.mark.parametrize("empty_tokens", [2, 35]) def test_http_stop_drains_detokenization_before_release(monkeypatch, include, empty_tokens): @@ -647,565 +284,6 @@ async def raw_results(): asyncio.run(run()) -def test_http_uses_request_finish_index(): - async def run(): - manager = _http_manager() - manager.pd_mode = SimpleNamespace(is_normal=lambda: False) - req = _req([10, 11, 12], FinishStatus.FINISHED_ABORTED) - req.out_tokens_queue.push("hello", 2, False, 1) - first = await _read_http_tokens(manager, req) - req.out_tokens_queue.push("", 4, False, 3) - last = await _read_http_tokens(manager, req) - assert first[0][1] == "hello" - assert not first[0][3].is_finished() - assert last[0][2]["id"] is None - assert last[0][3].status == FinishStatus.FINISHED_ABORTED - req.merge_final_token_metadata.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize("include", [False, True]) -def test_http_stop_aborts_only_matching_choice_and_drains_late_tokens(include): - async def run(): - manager = _http_manager() - params = _params(include=include) - params.n = params.best_of = 2 - req = _req([10, 11, 12, 13], params=params) - sibling = _req([50, 51], FinishStatus.FINISHED_LENGTH, params, request_id=81) - drained = [] - - async def raw_results(): - for target, index, text in [ - (req, 0, "helloE"), - (sibling, 0, "otherE"), - (req, 1, "NDextra"), - (req, 2, "late"), - (sibling, 1, "X"), - (req, 3, ""), - ]: - if target is req and index >= 2: - assert req.is_aborted - assert not sibling.is_aborted - if target is req and index == 3: - req.finish_status.set_status(FinishStatus.FINISHED_ABORTED) - req.finish_token_index = 5 - req.can_released_mark = True - target.out_tokens_queue.push(text, 2 + index, False, index + 1) - drained.append((target.request_id, index)) - for result in await _read_http_tokens(manager, target): - yield result - - results = [ - result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req, sibling]) - ] - assert len(drained) == 6 - assert req.out_tokens_queue.is_empty() - assert "".join(r[1] for r in results if r[0] == 80) == "hello" + ("END" if include else "") - assert "".join(r[1] for r in results if r[0] == 81) == "otherEX" - assert [r[2]["id"] for r in results if r[0] == 80] == [10, 11] - assert {r[0]: r[3].get_finish_reason() for r in results if r[3].is_finished()} == {80: "stop", 81: "length"} - req.merge_final_token_metadata.assert_awaited_once() - sibling.merge_final_token_metadata.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize( - "finish", - [ - FinishStatus.FINISHED_STOP, - FinishStatus.FINISHED_LENGTH, - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - ], -) -def test_http_finish_flushes_unmatched_prefix_and_preserves_final_metadata(finish): - async def run(): - manager = _http_manager() - params = _params() - req = _req([10, 11], finish, params) - req.merge_final_token_metadata.side_effect = lambda metadata, *args, **kwargs: metadata.update( - prompt_token_ids=[1, 2] - ) - - async def raw_results(): - for index, text in enumerate(["helloE", ""]): - req.out_tokens_queue.push(text, 2 + index, False, index + 1) - for result in await _read_http_tokens(manager, req): - yield result - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - assert "".join(r[1] for r in results) == "helloE" - assert results[-1][2]["prompt_token_ids"] == [1, 2] - assert results[-1][3].status == finish - assert results[-1][2]["id"] == (None if finish == FinishStatus.FINISHED_ABORTED else 11) - assert not req.is_aborted - - asyncio.run(run()) - - -@pytest.mark.parametrize("stop_in_reasoning", [False, True]) -@pytest.mark.parametrize("token_stop", [False, True]) -def test_http_scopes_stop_matching_and_preserves_hidden_reasoning_marker(stop_in_reasoning, token_stop): - async def run(): - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - params = _params([[10, 11]] if token_stop else "END", force=1) - params.stop_str_match_mode.set_mode( - StopStrMatchMode.ALL_TEXT if stop_in_reasoning else StopStrMatchMode.CONTENT_ONLY - ) - tokens = [10, 11, 1001, 10, 11] - req = _req(tokens, params=params) - infer_req = infer_batch.InferReq.__new__(infer_batch.InferReq) - infer_req.stop_sequences = params.stop_sequences.to_token_ids_list() - infer_req.shm_req = req - infer_req.sampling_param = SimpleNamespace(shm_param=params) - infer_req.finish_status = req.finish_status - - async def raw_results(): - for index, text in enumerate(["E", "ND", "", "answerE", "NDextra"]): - infer_req.update_finish_status([], output_len=index + 1) - if req.finish_status.is_finished(): - req.finish_token_index = req.input_len + index - req.out_tokens_queue.push(text, 2 + index, False, index + 1) - for result in await _read_http_tokens(manager, req): - yield result - if req.finish_status.is_finished(): - break - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - expected = "END" if token_stop else ("" if stop_in_reasoning else "ENDanswer") - assert "".join(r[1] for r in results) == expected - assert [r[2]["id"] for r in results] == tokens[: 2 if token_stop or stop_in_reasoning else 5] - assert results[-1][3].is_stopped() - - asyncio.run(run()) - - -def test_stop_groups_can_separate_explicit_token_ids_from_strings(): - params = _params(["END", [10, 11]]) - assert params.stop_sequences.to_token_ids_list() == [[10, 11]] - assert params.stop_sequences.to_stop_str_list() == ["END"] - assert params.to_dict()["stop_sequences"] == ["END", [10, 11]] - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_string_stop_is_matched_only_after_detokenization(mode): - params = _params() - params.stop_str_match_mode.set_mode(mode) - tokens = [ord(c) for c in "END"] - shm_req = _req(tokens, params=params) - req = infer_batch.InferReq.__new__(infer_batch.InferReq) - req.stop_sequences = params.stop_sequences.to_token_ids_list() - req.shm_req = shm_req - req.sampling_param = SimpleNamespace(shm_param=params) - req.finish_status = FinishStatus() - req.update_finish_status([], output_len=len(tokens)) - assert not req.finish_status.is_finished() - - decode = DecodeReq(shm_req) - assert decode.need_detoken() - assert not decode.can_set_release_mark() - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -@pytest.mark.parametrize("reasoning", [False, True]) -@pytest.mark.parametrize( - "stop_ids, output_ids, matched", - [ - ([10], [10], True), - ([10, 11], [10, 11], True), - ([10, 11], [10], False), - ([10, 11], [10, 12, 11], False), - ([10, 10, 11], [10, 10, 10, 11], True), - ], -) -def test_inference_token_stop_is_independent_of_string_mode_and_reasoning( - mode, reasoning, stop_ids, output_ids, matched -): - req = infer_batch.InferReq.__new__(infer_batch.InferReq) - params = _params(["END", stop_ids], force=int(reasoning)) - params.stop_str_match_mode.set_mode(mode) - req.stop_sequences = params.stop_sequences.to_token_ids_list() - req.shm_req = SimpleNamespace(input_len=1, shm_prompt_ids=SimpleNamespace(arr=np.array([1, *output_ids]))) - req.sampling_param = SimpleNamespace(shm_param=params) - req.finish_status = FinishStatus() - req.update_finish_status([], output_len=len(output_ids)) - assert req.finish_status.is_stopped() is matched - - # EOS and the output length limit remain independent of user stop sequences. - req.finish_status = FinishStatus() - req.update_finish_status([output_ids[-1]], output_len=len(output_ids)) - assert req.finish_status.is_stopped() is True - req.finish_status = FinishStatus() - req.stop_sequences = [] - req.sampling_param.shm_param.max_new_tokens = 1 - req.update_finish_status([], output_len=1) - assert req.finish_status.is_finished_length() is True - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("last_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_LENGTH]) -@pytest.mark.parametrize( - "stops, output_ids, match_index", - [ - ([[10]], [10, 20], 0), - ([[10, 11]], [10, 11, 20], 1), - ([[10, 10, 11]], [10, 10, 10, 11, 20], 3), - ([[10, 11], [11, 12]], [11, 12, 20], 1), - ([[10, 11]], [10, 12, 11], None), - ], -) -def test_token_id_buffer_preserves_tokens_and_only_corrects_matched_finish( - include, last_status, stops, output_ids, match_index -): - async def run(): - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopTokenIdsMatchBuffer(manager, _params(stops, include)) - results = [] - raw_outputs = [] - for index, token_id in enumerate(output_ids): - status = last_status if index == match_index or index == len(output_ids) - 1 else FinishStatus.NO_FINISH - output = _output(token_id, "" if token_id == 10 else str(token_id), index, status) - raw_outputs.append(output) - await buffer.handle(*output) - while buffer.ready_token_list: - results.append(buffer.ready_token_list.popleft()) - expected = raw_outputs if match_index is None else raw_outputs[: match_index + 1] - assert [r[:3] for r in results] == [r[:3] for r in expected] - assert all(result[2] is raw[2] for result, raw in zip(results, expected)) - if match_index is None: - assert results[-1][3].status == last_status - manager.stop_request.assert_not_awaited() - else: - assert results[-1][3].is_stopped() - assert raw_outputs[match_index][3].status == last_status - if last_status == FinishStatus.NO_FINISH: - manager.stop_request.assert_awaited_once_with(80) - else: - manager.stop_request.assert_not_awaited() - - asyncio.run(run()) - - -@pytest.mark.parametrize("last_status", [FinishStatus.FINISHED_ERROR, FinishStatus.FINISHED_LENGTH]) -def test_token_id_stream_preserves_error_and_simulated_final_output(last_status): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - - async def raw_results(): - yield _output(10, "first", 0) - yield _output(11 if last_status == FinishStatus.FINISHED_ERROR else None, "final", 1, last_status) - yield _output(11, "late", 2) - - async def run(): - return [result async for result in match_stop_token_ids_stream(manager, raw_results(), params)] - - results = asyncio.run(run()) - assert [r[1] for r in results] == ["first", "final"] - assert results[-1][3].status == last_status - manager.stop_request.assert_not_awaited() - - -def test_token_id_stream_keeps_interleaved_choices_independent(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - params.n = params.best_of = 2 - - async def raw_results(): - for req_id, token_id in [(80, 10), (81, 11), (80, 11), (80, 20), (81, 10), (81, 11), (81, 20)]: - yield req_id, str(token_id), {"id": token_id}, FinishStatus() - - async def run(): - return [result async for result in match_stop_token_ids_stream(manager, raw_results(), params)] - - results = asyncio.run(run()) - assert [r[2]["id"] for r in results if r[0] == 80] == [10, 11] - assert [r[2]["id"] for r in results if r[0] == 81] == [11, 10, 11] - assert [r[0] for r in results if r[3].is_stopped()] == [80, 81] - assert [call.args for call in manager.stop_request.await_args_list] == [(80,), (81,)] - - -def test_token_id_stream_does_not_match_prompt_ids(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - - async def raw_results(): - yield 80, "answer", {"id": 11, "prompt_token_ids": [10]}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - async def run(): - return [result async for result in wrap_stop_match_stream(manager, raw_results(), params, [10])] - - results = asyncio.run(run()) - assert results[0][3].is_finished_length() - manager.stop_request.assert_not_awaited() - - -def test_token_id_matching_runs_before_string_matching(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params(["END", [10, 11]]) - - async def raw_results(): - yield _output(10, "helloE", 0) - yield _output(11, "N", 1) - yield _output(12, "Dextra", 2, FinishStatus.FINISHED_LENGTH) - - async def run(): - # prompt 的 ID 即使是 stop 前缀，也不进入 token ID 匹配历史。 - return [result async for result in wrap_stop_match_stream(manager, raw_results(), params, [10])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results) == "helloEN" - assert [r[2]["id"] for r in results] == [10, 11] - assert results[-1][3].is_stopped() - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("with_stop", [None, "END", [[10, 11]], ["END", [10, 11]]]) -def test_stop_stream_close_closes_raw_generator(with_stop): - params = _params(stop=with_stop) - closed = [] - - async def raw_results(): - try: - yield 80, "hello", {"id": 10}, FinishStatus() - pytest.fail("Closing the wrapper must not consume another raw token") - finally: - closed.append(True) - - async def run(): - manager = SimpleNamespace( - args=SimpleNamespace(enable_rl=False, enable_return_routed_experts=False, reasoning_parser=None), - tokenizer=_ReasoningTokenizer(), - stop_request=AsyncMock(), - ) - raw_generator = raw_results() - output = wrap_stop_match_stream(manager, raw_generator, params, "prompt") - assert (output is raw_generator) is (not with_stop) - result = await output.__anext__() - assert result[1] == "hello" - await output.aclose() - - asyncio.run(run()) - assert closed == [True] - - -@pytest.mark.parametrize("node_finished", [False, True]) -def test_stop_stream_calls_manager_without_mutating_raw_finish(node_finished): - params = _params() - raw_finish = FinishStatus(FinishStatus.FINISHED_LENGTH if node_finished else FinishStatus.NO_FINISH) - manager = _http_manager() - manager.stop_request = AsyncMock() - - async def raw_results(): - yield 80, "helloENDextra", {"id": 10}, raw_finish - yield 80, "late", {"id": 11}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, "prompt")] - - results = asyncio.run(run()) - assert results[0][1] == "hello" - assert results[0][3].is_stopped() - assert raw_finish.status == (FinishStatus.FINISHED_LENGTH if node_finished else FinishStatus.NO_FINISH) - if node_finished: - manager.stop_request.assert_not_awaited() - else: - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("include", [False, True]) -def test_http_stream_discards_late_tokens_in_same_full_queue_batch(include): - async def run(): - manager = _http_manager() - params = _params(include=include) - req = _req(list(range(10, 20)), params=params) - index = 0 - while not req.out_tokens_queue.is_full(): - chunks = ["helloE", "N", "Dextra"] - req.out_tokens_queue.push(chunks[index] if index < len(chunks) else "late", 2 + index, False, index + 1) - index += 1 - - async def raw_results(): - for result in await _read_http_tokens(manager, req): - yield result - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - assert req.out_tokens_queue.is_empty() - assert req.is_aborted - assert "".join(r[1] for r in results) == "hello" + ("END" if include else "") - assert [r[2]["id"] for r in results] == [10, 11, 12] - assert results[-1][3].is_stopped() - - asyncio.run(run()) - - -def test_stop_stream_uses_prompt_tail_for_all_choices_and_ignores_metadata(): - params = _params(force=0) - params.n = params.best_of = 2 - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - manager.stop_request = AsyncMock() - manager.tokenizer.encode = Mock(wraps=manager.tokenizer.encode) - - async def raw_results(): - yield 80, "END", {"id": 10, "prompt_token_ids": [1, 1001]}, FinishStatus() - yield 81, "END", {"id": 10}, FinishStatus() - for req_id in (80, 81): - yield req_id, "", {"id": 1001}, FinishStatus() - yield req_id, "answerENDextra", {"id": 11}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [1] * 4096 + [1000])] - - results = asyncio.run(run()) - for req_id in (80, 81): - assert "".join(result[1] for result in results if result[0] == req_id) == "ENDanswer" - assert [call.args for call in manager.stop_request.await_args_list] == [(80,), (81,)] - assert [call.args for call in manager.tokenizer.encode.call_args_list] == [("",), ("",)] * 2 - assert results[0][2]["prompt_token_ids"] == [1, 1001] - - -@pytest.mark.parametrize("prompt", ["prompt", "x" * 4096 + "prompt", [1, 2], [1] * 4096 + [2]]) -@pytest.mark.parametrize("prompt_metadata", [{}, {"prompt_token_ids": None}, {"prompt_token_ids": [1000]}]) -@pytest.mark.parametrize( - "mode, parser", - [ - (StopStrMatchMode.ALL_TEXT, "qwen3"), - (StopStrMatchMode.CONTENT_ONLY, None), - (StopStrMatchMode.CONTENT_ONLY, "qwen3"), - ], -) -def test_stop_stream_encodes_prompt_tail_independent_of_mode_and_metadata(prompt, prompt_metadata, mode, parser): - params = _params() - params.stop_str_match_mode.set_mode(mode) - manager = _http_manager() - manager.args.reasoning_parser = parser - manager.stop_request = AsyncMock() - manager.tokenizer.encode = Mock(wraps=manager.tokenizer.encode) - - async def raw_results(): - yield 80, "helloENDextra", {"id": 10, **prompt_metadata}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, prompt)] - - results = asyncio.run(run()) - assert results[0][1] == "hello" - assert results[0][3].is_stopped() - prompt_calls = [ - call for call in manager.tokenizer.encode.call_args_list if call.args[0] not in ("", "") - ] - assert len(prompt_calls) == int(isinstance(prompt, str)) - if isinstance(prompt, str): - assert prompt_calls[0].args == (prompt[-2048:],) - assert prompt_calls[0].kwargs == {"add_special_tokens": False} - - -def test_stop_stream_keeps_reasoning_state_independent_for_interleaved_choices(): - params = _params(force=0) - params.n = params.best_of = 2 - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - manager.stop_request = AsyncMock() - - async def raw_results(): - for req_id, token_id, text in [ - (80, 1000, ""), - (81, 10, "END"), - (80, 10, "END"), - (81, 11, "late"), - (80, 1001, ""), - (80, 12, "answerENDextra"), - ]: - yield req_id, text, {"id": token_id}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results if r[0] == 80) == "ENDanswer" - assert "".join(r[1] for r in results if r[0] == 81) == "" - assert [call.args for call in manager.stop_request.await_args_list] == [(81,), (80,)] - assert params._initial_reasoning_state == 0 - - -@pytest.mark.parametrize("first_finished", [False, True]) -@pytest.mark.parametrize( - "late_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_LENGTH, FinishStatus.FINISHED_ERROR] -) -def test_stopped_choice_ignores_late_output_before_and_after_sibling_starts(first_finished, late_status): - params = _params() - params.n = params.best_of = 2 - manager = _http_manager() - manager.stop_request = AsyncMock() - closed = [] - exhausted = [] - - async def raw_results(): - try: - finish = FinishStatus(FinishStatus.FINISHED_LENGTH if first_finished else FinishStatus.NO_FINISH) - yield 80, "first" if first_finished else "firstENDextra", {"id": 10}, finish - yield 80, "lateEND", {"id": 11}, FinishStatus(late_status) - yield 81, "secondE", {"id": 20}, FinishStatus() - yield 80, "laterEND", {"id": 12}, FinishStatus(late_status) - yield 81, "NDextra", {"id": 21}, FinishStatus(FinishStatus.FINISHED_LENGTH) - yield 80, "lastEND", {"id": 13}, FinishStatus(late_status) - yield 81, "lastEND", {"id": 22}, FinishStatus(late_status) - exhausted.append(True) - finally: - closed.append(True) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert [result[0] for result in results] == [80, 81, 81] - assert "".join(result[1] for result in results) == "firstsecond" - assert results[0][3].status == (FinishStatus.FINISHED_LENGTH if first_finished else FinishStatus.FINISHED_STOP) - assert results[-1][3].is_stopped() - if first_finished: - manager.stop_request.assert_not_awaited() - else: - manager.stop_request.assert_awaited_once_with(80) - assert closed == [True] - assert exhausted == [True] - - -def test_application_stop_keeps_hit_metadata_and_discards_backend_finish(): - params = _params() - manager = _http_manager() - manager.stop_request = AsyncMock() - - async def raw_results(): - yield 80, "helloE", {"id": 10, "logprob": -0.1, "count_output_tokens": 1}, FinishStatus() - yield 80, "NDextra", {"id": 11, "logprob": -0.2, "count_output_tokens": 2}, FinishStatus() - yield 80, "late", {"id": 12, "logprob": -0.3, "count_output_tokens": 3}, FinishStatus( - FinishStatus.FINISHED_LENGTH - ) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results) == "hello" - assert results[-1][3].is_stopped() - assert results[-1][2] == { - "id": 11, - "logprob": -0.2, - "count_output_tokens": 2, - } - manager.stop_request.assert_awaited_once_with(80) - - @pytest.mark.parametrize("include", [False, True]) @pytest.mark.parametrize("finished", [False, True]) def test_ready_queue_preserves_unconsumed_tokens_after_stop(include, finished): @@ -1232,32 +310,3 @@ def test_ready_queue_preserves_unconsumed_tokens_after_stop(include, finished): ("N" if finished else "ND" if include else "", payloads[3]), ] assert not ready_queue - - -@pytest.mark.parametrize("matched,finished", [(False, True), (True, False), (True, True)]) -def test_stopped_buffer_ignores_later_tokens(matched, finished): - buffer = _stop_buffer(stop_in_reasoning=True) - finish = FinishStatus.FINISHED_LENGTH if finished else FinishStatus.NO_FINISH - _handle_output(buffer, *_output(10, "helloEND" if matched else "helloE", "first", finish=finish)) - first_finish = buffer.ready_token_list[-1][3] - first_tokens = _drain_ready_tokens(buffer) - assert first_finish.is_stopped() is matched - assert "".join(text for text, _ in first_tokens) == ("hello" if matched else "helloE") - assert buffer.state is buffer.State.STOPPED - for finished in [False, True]: - finish = FinishStatus.FINISHED_LENGTH if finished else FinishStatus.NO_FINISH - _handle_output(buffer, *_output(11, "laterEND", "later", finish=finish)) - ready_tokens = _drain_ready_tokens(buffer) - assert not ready_tokens - assert "".join(text for text, _ in first_tokens) == ("hello" if matched else "helloE") - - -def test_simulated_finish_releases_prefix_and_preserves_payloads(): - buffer = _stop_buffer(stop_in_reasoning=True) - first_payload, finish_payload = object(), object() - _handle_output(buffer, *_output(10, "helloE", first_payload)) - assert not buffer.ready_token_list - _handle_output(buffer, *_output(None, "", finish_payload, finish=FinishStatus.FINISHED_LENGTH)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("helloE", first_payload), ("", finish_payload)] - assert buffer.state is buffer.State.STOPPED diff --git a/unit_tests/server/embed_cache/test_multimodal_token_id_range.py b/unit_tests/server/embed_cache/test_multimodal_token_id_range.py deleted file mode 100644 index b47a1f912f..0000000000 --- a/unit_tests/server/embed_cache/test_multimodal_token_id_range.py +++ /dev/null @@ -1,61 +0,0 @@ -from types import SimpleNamespace - -from lightllm.server.embed_cache.impl import naive_memory_cache - - -def test_multimodal_token_id_range_selection_uses_time_seed(monkeypatch): - seed_values = [] - selected_index = 1234 - - monkeypatch.setattr(naive_memory_cache.time, "time", lambda: 123.5) - monkeypatch.setattr(naive_memory_cache.random, "seed", seed_values.append) - monkeypatch.setattr(naive_memory_cache.random, "randint", lambda start, end: selected_index) - - cache = naive_memory_cache.InMemoryCache.__new__(naive_memory_cache.InMemoryCache) - range_index = cache._set_prefill_node_random_token_id_range() - range_size = ((2 ** 63 - 1) - 100000000) // 10000 - - assert seed_values == [123.5] - assert range_index == selected_index - assert cache.token_id_range_start == 100000000 + selected_index * range_size - assert cache.token_id_range_end == cache.token_id_range_start + range_size - - -def test_cache_manager_selects_and_logs_probabilistic_range(monkeypatch): - node_uuid = 0x12345678123456781234567812345678 - expected_index = 4321 - range_size = ((2 ** 63 - 1) - 100000000) // 10000 - expected_start = 100000000 + expected_index * range_size - expected_end = expected_start + range_size - warning_messages = [] - - def set_expected_range(cache): - cache.token_id_range_start = expected_start - cache.token_id_range_end = expected_end - return expected_index - - monkeypatch.setattr(naive_memory_cache, "CpuEmbedCacheClient", lambda **kwargs: object()) - monkeypatch.setattr( - naive_memory_cache.InMemoryCache, - "_set_prefill_node_random_token_id_range", - set_expected_range, - ) - monkeypatch.setattr(naive_memory_cache.logger, "warning", warning_messages.append) - - cache = naive_memory_cache.InMemoryCache( - SimpleNamespace( - cache_capacity=200, - config_server_host=None, - config_server_port=None, - pd_node_id=node_uuid, - run_mode="prefill", - ) - ) - cache._check_and_set_new_id_range(1) - - assert cache.token_id_range_start == expected_start - assert cache.token_id_range_end == expected_end - assert len(warning_messages) == 1 - assert f"slot={expected_index}/9999" in warning_messages[0] - assert "does not guarantee global uniqueness" in warning_messages[0] - assert "0.0100% probability" in warning_messages[0] diff --git a/unit_tests/server/httpserver/test_grammar.py b/unit_tests/server/httpserver/test_grammar.py deleted file mode 100644 index 916c0f665c..0000000000 --- a/unit_tests/server/httpserver/test_grammar.py +++ /dev/null @@ -1,340 +0,0 @@ -import asyncio -import json -import pickle -from threading import Event -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest -import xgrammar as xgr - -from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs.sampling_params import SamplingParams -from unit_tests.server.grammar_helpers import configure, reset_args_cache - - -@pytest.fixture(autouse=True) -def start_args(monkeypatch): - configure(monkeypatch) - - -def params(compiler=None, **kwargs): - result = SamplingParams() - result.init(None, **kwargs) - asyncio.run(result.verify_cpu_async(compiler)) - return result - - -@pytest.fixture -def http_client(compiler, monkeypatch): - from lightllm.server import api_http, api_openai - from lightllm.server.api_lightllm import lightllm_generate, lightllm_generate_stream - - configure(monkeypatch, run_mode="pd_master", disable_delay_response_start=True) - generate = MagicMock(side_effect=AssertionError("Invalid request reached generation")) - manager = SimpleNamespace( - tokenizer=None, - output_grammar_compiler=compiler, - generate=generate, - args=SimpleNamespace(enable_return_routed_experts=False), - ) - monkeypatch.setattr(api_http.g_objs, "httpserver_manager", manager) - monkeypatch.setattr(api_http.g_objs, "g_generate_func", lightllm_generate) - 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) - yield httpx.AsyncClient(transport=httpx.ASGITransport(app=api_http.app), base_url="http://test") - generate.assert_not_called() - - -def assert_bad_request(client, path, body): - async def send(): - async with client: - return await client.post(path, json=body) - - response = asyncio.run(send()) - assert response.status_code == 400, response.text - assert response.json()["error"]["type"] == "BadRequestError" - return response.text - - -@pytest.mark.parametrize("failure", ["syntax", "exception", "timeout"]) -@pytest.mark.parametrize("early_headers", [False, True]) -def test_compilation_failure_returns_400_before_generation(compiler, http_client, monkeypatch, failure, early_headers): - configure(monkeypatch, disable_delay_response_start=early_headers) - if failure == "exception": - monkeypatch.setattr(compiler, "_compile", MagicMock(side_effect=RuntimeError("injected compiler error"))) - elif failure == "timeout": - compiler.timeout = 0 - grammar = "[" if failure == "syntax" else "ab" - body = {"inputs": "test", "parameters": {"regular_constraint": grammar}} - error = assert_bad_request(http_client, "/generate_stream", body) - assert ("timed out" if failure == "timeout" else "Failed to compile") in error - - -@pytest.mark.parametrize("path", ["/generate", "/generate_stream"]) -@pytest.mark.parametrize("early_headers", [False, True]) -@pytest.mark.parametrize("invalid_params", [{"max_new_tokens": 0}, {"do_sample": True, "top_p": 0}]) -def test_invalid_sampling_params_return_400(http_client, monkeypatch, path, early_headers, invalid_params): - configure(monkeypatch, disable_delay_response_start=early_headers) - assert_bad_request(http_client, path, {"inputs": "test", "parameters": invalid_params}) - - -@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions"]) -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("failure", ["syntax", "exception", "timeout"]) -@pytest.mark.parametrize("early_headers", [False, True]) -def test_openai_grammar_failures_return_400(compiler, http_client, monkeypatch, path, stream, failure, early_headers): - configure(monkeypatch, disable_delay_response_start=early_headers) - if failure == "exception": - monkeypatch.setattr(compiler, "_compile", MagicMock(side_effect=RuntimeError("injected compiler error"))) - elif failure == "timeout": - compiler.timeout = 0 - schema = {"type": 42} if failure == "syntax" else {"type": "object"} - body = { - "model": "test", - "stream": stream, - "response_format": {"type": "json_schema", "json_schema": {"name": "test", "schema": schema}}, - } - body.update({"messages": []} if "chat" in path else {"prompt": "test"}) - assert_bad_request(http_client, path, body) - - -@pytest.mark.parametrize( - "constraint,text", - [ - ({"guided_grammar": 'root ::= "ab"'}, b"ab"), - ({"regular_constraint": "ab"}, b"ab"), - ({"guided_json": {"const": "ok"}}, b'"ok"'), - ], -) -def test_copied_and_pickled_params_get_artifact_from_compiler_cache(compiler, monkeypatch, constraint, text): - sampling_params = params(compiler=compiler, **constraint) - # Choices and PD continuation transport scalar parameters, then get the - # immutable artifact from the compiler's completed-result cache. - copied = pickle.loads(pickle.dumps(SamplingParams.from_buffer_copy(sampling_params))) - with monkeypatch.context() as patcher: - patcher.setattr(compiler._executor, "submit", MagicMock(side_effect=AssertionError("Cache miss"))) - patcher.setattr(xgr, "GrammarCompiler", MagicMock(side_effect=AssertionError("Backend recompiled"))) - payload = asyncio.run(copied.get_compiled_grammar(compiler)) - assert payload is asyncio.run(sampling_params.get_compiled_grammar(compiler)) - grammar = compiler.output_constraint_state_factory.get_grammar(payload) - matcher = xgr.GrammarMatcher(grammar) - assert all(matcher.accept_token(token) for token in text) - assert matcher.accept_token(256) and matcher.is_terminated() - assert "compiled_grammar" not in sampling_params.to_dict() - assert not hasattr(copied, "compiled_grammar") - copied.max_new_tokens = 2 - assert sampling_params.max_new_tokens != copied.max_new_tokens - - -def test_completed_artifact_is_returned_without_executor_work(compiler, monkeypatch): - async def run(): - payload = await compiler.compile("regex", "ab") - monkeypatch.setattr( - compiler._executor, "submit", MagicMock(side_effect=AssertionError("Cache hit reached executor")) - ) - assert await compiler.compile("regex", "ab") is payload - - asyncio.run(run()) - - -def test_completed_artifact_cache_evicts_least_recently_used(compiler, monkeypatch): - compiler._cache_size = 2 - compile_grammar = MagicMock(wraps=compiler._compile) - monkeypatch.setattr(compiler, "_compile", compile_grammar) - - async def run(): - for grammar in ("a", "b", "a", "c", "a", "b"): - await compiler.compile("regex", grammar) - - asyncio.run(run()) - assert [call.args[1] for call in compile_grammar.call_args_list] == ["a", "b", "c", "b"] - - -def test_failed_compilation_is_retried_before_caching_success(compiler, monkeypatch): - compile_grammar = MagicMock(side_effect=[RuntimeError("temporary failure"), b"compiled artifact"]) - monkeypatch.setattr(compiler, "_compile", compile_grammar) - - async def run(): - with pytest.raises(ValueError, match="temporary failure"): - await compiler.compile("regex", "ab") - payload = await compiler.compile("regex", "ab") - assert payload == b"compiled artifact" - assert await compiler.compile("regex", "ab") is payload - - asyncio.run(run()) - assert compile_grammar.call_count == 2 - - -def test_no_grammar_does_not_use_compiler_and_none_mode_rejects_constraints(): - compiler = MagicMock(side_effect=AssertionError("Ordinary request compiled")) - params(compiler=compiler) - compiler.compile.assert_not_called() - with pytest.raises(ValueError, match="output_constraint_mode xgrammar"): - params(regular_constraint="ab") - - -def test_getter_uses_current_constraints_after_reinitialization(compiler): - sampling_params = SamplingParams() - - async def run(): - sampling_params.init(None, regular_constraint="ab") - await sampling_params.verify_cpu_async(compiler) - first = await sampling_params.get_compiled_grammar(compiler) - sampling_params.init(None, regular_constraint="cd") - assert await sampling_params.get_compiled_grammar(compiler) != first - sampling_params.init(None) - assert await sampling_params.get_compiled_grammar(None) is None - - asyncio.run(run()) - - -@pytest.mark.parametrize("stop_waiter", ["cancel", "timeout"]) -def test_cold_requests_compile_independently_and_cache_only_success(compiler, monkeypatch, stop_waiter): - release = Event() - compile_grammar = compiler._compile - - async def run(): - loop = asyncio.get_running_loop() - started, finished = asyncio.Event(), asyncio.Event() - - def slow_first_compile(*args): - loop.call_soon_threadsafe(started.set) - try: - assert release.wait(5) - return b"late artifact" - finally: - loop.call_soon_threadsafe(finished.set) - - monkeypatch.setattr(compiler, "_compile", slow_first_compile) - compiler.timeout = 0.02 if stop_waiter == "timeout" else 5 - first = asyncio.create_task(compiler.compile("regex", "ab")) - try: - await asyncio.wait_for(started.wait(), timeout=5) - # The same grammar can finish on another worker while the first job - # is blocked. Neither a cache hit nor shared pending work is involved. - second_compile = MagicMock(wraps=compile_grammar) - monkeypatch.setattr(compiler, "_compile", second_compile) - compiler.timeout = 5 - payload = await compiler.compile("regex", "ab") - second_compile.assert_called_once_with("regex", "ab") - - if stop_waiter == "cancel": - first.cancel() - with pytest.raises(asyncio.CancelledError): - await first - else: - with pytest.raises(ValueError, match="timed out"): - await first - finally: - release.set() - - await asyncio.wait_for(finished.wait(), timeout=5) - # The abandoned worker's eventual result must not replace a cached success. - assert await compiler.compile("regex", "ab") is payload - - try: - asyncio.run(run()) - finally: - release.set() - - -@pytest.mark.parametrize("explicit_eos", [None, [255]]) -def test_pd_master_compiler_resolves_eos_before_model_loading(monkeypatch, tmp_path, explicit_eos): - from lightllm.server.core.objs import StartArgs - from lightllm.server.httpserver.grammar import create_output_grammar_compiler - - (tmp_path / "config.json").write_text('{"model_type":"llama","vocab_size":257,"eos_token_id":256}') - args = StartArgs(run_mode="pd_master", model_dir=str(tmp_path), eos_id=explicit_eos) - stop_ids = explicit_eos if explicit_eos is not None else [256] - info = xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=stop_ids) - tokenizer = object() - with patch.object(xgr.TokenizerInfo, "from_huggingface", return_value=info) as create_info: - compiler = create_output_grammar_compiler(args, tokenizer) - try: - create_info.assert_called_once_with(tokenizer, vocab_size=257, stop_token_ids=stop_ids) - payload = asyncio.run(compiler.compile("regex", "a")) - matcher = xgr.GrammarMatcher(xgr.CompiledGrammar.deserialize_json(payload.decode(), info)) - assert matcher.accept_token(ord("a")) - assert matcher.accept_token(stop_ids[0]) and matcher.is_terminated() - finally: - compiler.shutdown() - - -def test_dictionary_schema_is_serialized_during_initialization(): - params = SamplingParams() - params.init(None, guided_json={"type": "object"}) - assert json.loads(params.constraint_out_params.guided_json.to_str()) == {"type": "object"} - - -@pytest.mark.parametrize( - "invalid_params", [{"max_new_tokens": 0}, {"do_sample": True, "top_p": 0}, {"best_of": 2, "n": 1}] -) -def test_parameter_validation_precedes_grammar_compilation(invalid_params): - params = SamplingParams() - compiler = SimpleNamespace(compile=AsyncMock()) - with pytest.raises(ValueError): - params.init(None, regular_constraint="ab", **invalid_params) - asyncio.run(params.verify_cpu_async(compiler)) - compiler.compile.assert_not_called() - assert not hasattr(params, "compiled_grammar") - - -@pytest.mark.parametrize( - "constraint", [{"guided_grammar": "bad grammar"}, {"guided_json": "not json"}, {"regular_constraint": "["}] -) -def test_invalid_grammar_is_rejected_by_output_constraint_validation(compiler, constraint): - params = SamplingParams() - params.init(None, **constraint) - with pytest.raises(ValueError): - asyncio.run(params.verify_cpu_async(compiler)) - assert not hasattr(params, "compiled_grammar") - - -@pytest.mark.parametrize("status", [FinishStatus.FINISHED_ERROR, FinishStatus.FINISHED_STOP]) -def test_http_preserves_finish_status_and_inflight_token(status): - import numpy as np - - from lightllm.server.core.objs import FinishStatus, Req - from lightllm.server.httpserver.manager import HttpServerManager, ReqStatus - from lightllm.server.pd_io_struct import NodeRole - - req = Req() - req.request_id = req.group_req_id = 1 - req.input_len = 1 - req.shm_prompt_ids = SimpleNamespace(arr=np.array([ord("P"), ord("!")])) - req.shm_logprobs = SimpleNamespace(arr=np.array([(0.0,), (-0.2,)], dtype=[("logprob", np.float32)])) - req.finish_status.set_status(status) - req.finish_token_index = 1 - req.out_tokens_queue.push("!", 1, False, 1) - req.get_output_logprobs_metadata = MagicMock(return_value={"token": "!"}) - req.merge_final_token_metadata = AsyncMock() - - async def collect(): - manager = HttpServerManager.__new__(HttpServerManager) - manager.args = SimpleNamespace(use_reward_model=False, enable_return_routed_experts=False) - manager.tokenizer = None - manager.is_multinode_tp_slave = False - manager.pd_mode = NodeRole.NORMAL - wakeups = asyncio.Queue() - wakeups.put_nowait(None) - manager.zmq_recv_socket = SimpleNamespace(recv_pyobj=wakeups.get) - manager.recycle_resource_loop = AsyncMock() - status = ReqStatus(1, None, [req], 0) - manager.req_id_to_out_inf = {1: status} - task = asyncio.create_task(manager.handle_loop()) - try: - await asyncio.wait_for(status.event.wait(), timeout=2) - return status.out_token_info_list[0] - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - _, text, metadata, finish = asyncio.run(collect()) - assert text == "!" and finish.status == status - assert metadata["id"] == ord("!") and metadata["logprob"] == pytest.approx(-0.2) - assert metadata["logprobs"] == {"token": "!"} - req.merge_final_token_metadata.assert_awaited_once() diff --git a/unit_tests/server/httpserver/test_pd_master_cached_tokens.py b/unit_tests/server/httpserver/test_pd_master_cached_tokens.py index f71534f9be..cb815fb469 100644 --- a/unit_tests/server/httpserver/test_pd_master_cached_tokens.py +++ b/unit_tests/server/httpserver/test_pd_master_cached_tokens.py @@ -9,7 +9,7 @@ from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager, logger from lightllm.server.httpserver_for_pd_master.pd_selector import PDSelectionExtraInfo from lightllm.server.pd_io_struct import PD_Client_Obj @@ -622,6 +622,7 @@ def test_pd_segment_summary_is_recorded_on_exit(monkeypatch, caplog, exit_mode, params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4, group_request_id=80) request = SimpleNamespace(headers={"X-Request-Id": "summary-test"}, is_disconnected=lambda: asyncio.sleep(0, False)) caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") + monkeypatch.setattr(logger, "handlers", [caplog.handler]) async def raw_stream(*_args): if exit_mode == "empty": @@ -684,6 +685,7 @@ def test_pd_stopped_choice_records_segment_summary(monkeypatch, caplog, late_sta params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4) request = SimpleNamespace(headers={}, is_disconnected=lambda: asyncio.sleep(0, False)) caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") + monkeypatch.setattr(logger, "handlers", [caplog.handler]) async def raw_stream(_p_node, _d_node, _prompt, sampling_params, *_args): metadata = {"id": 1, "prompt_tokens": 10, "prompt_cache_len": 3} diff --git a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py deleted file mode 100644 index 4238e199f5..0000000000 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ /dev/null @@ -1,120 +0,0 @@ -import time -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager -from lightllm.utils.kv_cache_utils import CpuKVCacheMeta - - -def test_disk_roundtrip_preserves_raw_pages_before_mm_boundary_trim(cpu_cache_factory, tmp_path, monkeypatch): - light_mem = pytest.importorskip("light_mem") - from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker - - client = cpu_cache_factory(CpuKVCacheMeta(3, 4096)) - pages, _ = client.allocate_pages([11, 22, 33], False) - for index in pages: - client.cpu_kv_cache_tensor[index].fill_(index + 1) - client.update_pages_status_to_ready(pages, disk_offload_enable=True, token_num_in_page_list=768) - expected = client.cpu_kv_cache_tensor.clone() - worker = DiskCacheWorker.__new__(DiskCacheWorker) - worker.cpu_cache_client = client - worker.max_concurrent_write_tasks = 8 - worker.service = light_mem.PyLocalCacheService( - worker._prepare_tensor(client.cpu_kv_cache_tensor), - str(tmp_path / "cache"), - storage_size=16 * 1024 * 1024, - num_shard=1, - num_worker=2, - ) - for payloads in worker._gather_offload_payloads(): - worker._persist_pages_to_disk(payloads) - deadline = time.monotonic() + 10 - while not all(worker.service.query([11, 22, 33])): - assert time.monotonic() < deadline, "Disk write did not finish" - time.sleep(0.01) - - # Evict CPU copies and overwrite their bytes before restoring from disk. - replaced, _ = client.allocate_pages([44, 55, 66], False) - client.recycle_pages(replaced) - client.cpu_kv_cache_tensor.fill_(165) - manager = MultiLevelKVCacheManager.__new__(MultiLevelKVCacheManager) - manager.args = SimpleNamespace(enable_disk_cache=True) - manager.cpu_cache_client = client - manager.disk_cache_worker = worker - loaded, disk_pages = manager._disk_cache_match([11, 22, 33], []) - assert len(loaded) == disk_pages == 3 - assert torch.equal(client.cpu_kv_cache_tensor[loaded], expected) - - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - - copied_pages = [] - req = SimpleNamespace( - req_idx=0, - cur_kv_len=0, - hold_kv_len=0, - mm_input_spans=[(600, 800)], - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=1024, - disk_prompt_cache_len=768, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: loaded), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [256, 512, 768]), - ), - ) - - def allocate(alloc_token_num): - if alloc_token_num == 0: - return None - start = req.cur_kv_len - req.hold_kv_len += alloc_token_num - return torch.arange(start, start + alloc_token_num, dtype=torch.int32) - - req.args = SimpleNamespace(page_size=1) - req.token_index = SimpleNamespace( - get_gpu_token_indexes=lambda start, end: torch.empty(end - start, dtype=torch.int32) - ) - req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) - req.alloc_kv_mem = allocate - - from collections import deque - - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.load_tasks = deque() - module.args = SimpleNamespace(page_size=1, cpu_cache_token_page_size=256) - module.cpu_cache_client = client - module.init_sync_group = None - module.need_sync_compute_stream = lambda: False - module.backend = SimpleNamespace( - is_master_in_dp=True, - dp_world_size=1, - is_hybrid_att_model=False, - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs=torch.arange(1024).view(1, -1), - alloc_token_indexes=allocate, - wait_token_indexes=lambda: None, - ), - att_state_manager=SimpleNamespace( - get_cpu_cache_operator=lambda: SimpleNamespace( - get_page_layouts=lambda: {}, - read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: copied_pages.extend( - cpu_page_indices - ), - ), - full_att_allocator=None, - ), - ), - ) - monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 1024) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - module.load_cpu_cache_to_reqs([req]) - for task in module.load_tasks: - task.event.synchronize() - module.update_cpu_cache_load_task_states() - assert copied_pages == loaded[:2] - assert req.cur_kv_len == req.hold_kv_len == req.shm_req.disk_prompt_cache_len == 512 - assert req.shm_req.cpu_prompt_cache_len == 0 - assert all(client.page_items.linked_items[page].ref_count == 0 for page in loaded) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py index 8f99c6add3..98042570c1 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py @@ -40,6 +40,8 @@ def make_task(page_kind): decode_page_reg_desc=b"decode pages", first_gen_token_id=42, first_gen_token_logprob=-0.5, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([42, 7], [-0.5, -0.6]), page_kind=page_kind, req_idx=17, src_page_index=3, @@ -77,6 +79,17 @@ def test_copy_for_notify_preserves_local_task(page_kind, with_mark): assert task.shared_time_out_mark.time_mark == time_mark +def test_task_return_preserves_first_token_metadata(): + task = make_task("kv") + + result = task.createRetObj() + + assert result.first_gen_token_id == 42 + assert result.first_gen_token_logprob == -0.5 + assert result.first_gen_token_rank == 21 + assert result.first_gen_token_output_top_logprobs == ([42, 7], [-0.5, -0.6]) + + _UNPICKLE_WITHOUT_DECODE = """ import importlib.abc import pickle diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py index d2adf388f4..3c7ad49f4a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py @@ -29,7 +29,11 @@ def test_hybrid_pd_uses_the_shared_att_state_task_without_model_flags(monkeypatc def create(**kwargs): task = SimpleNamespace( - page_kind=kwargs.get("page_kind", "kv"), first_gen_token_id=None, first_gen_token_logprob=None + page_kind=kwargs.get("page_kind", "kv"), + first_gen_token_id=None, + first_gen_token_logprob=None, + first_gen_token_rank=None, + first_gen_token_output_top_logprobs=None, ) tasks.append(task) if "group" in kwargs: @@ -38,8 +42,17 @@ def create(**kwargs): backend._create_pd_trans_task = create if role == "prefill": - backend._prefill_chuncked_handle_func(req, next_token_id=17, next_token_prob=-0.2, output_len=1) + backend._prefill_chuncked_handle_func( + req, + next_token_id=17, + next_token_prob=-0.2, + output_len=1, + next_token_rank=21, + output_top_logprobs=([17, 16], [-0.2, -0.3]), + ) assert tasks[-1].first_gen_token_id == 17 + assert tasks[-1].first_gen_token_rank == 21 + assert tasks[-1].first_gen_token_output_top_logprobs == ([17, 16], [-0.2, -0.3]) else: backend.radix_cache = None backend.model = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py deleted file mode 100644 index a760f44a35..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -import pytest -import torch - -from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.req_manager import ReqManager -from lightllm.common.req_manager.base import _ReqIndexPool -from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps -from lightllm.utils.torch_memory_saver_utils import MemoryTag - - -@pytest.mark.parametrize( - ("tags", "should_init_hold_indexes"), - [ - ([MemoryTag.WEIGHT], False), - ([MemoryTag.GRAPH], False), - ([], False), - ([MemoryTag.KV_CACHE], True), - (None, True), - ], -) -def test_resume_memory_initializes_hold_indexes_only_for_kv_cache( - tags, should_init_hold_indexes, kv_allocator_environment -): - allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) - req_manager = ReqManager.__new__(ReqManager) - req_manager.req_list = _ReqIndexPool(2) - req_manager.HOLD_REQUEST_ID = 2 - req_manager.page_size = allocator.page_size - req_manager._cpu_token_index_blocks = np.empty((3, 2), dtype=np.int32) - req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) - resumed = [] - - def resume(tag): - assert req_manager.req_to_token_indexs.eq(-1).all() - resumed.append(tag) - - model = SimpleNamespace( - torch_memory_saver=SimpleNamespace(resume=resume), - req_manager=req_manager, - att_state_manager=SimpleNamespace(full_att_allocator=allocator), - ) - backend = SimpleNamespace(model=model, logger=Mock()) - ops = RlBackendOps(backend) - ops._clear_cuda_cache = lambda: None - - success, message = ops.resume_memory_occupation(tags) - - assert success, message - assert resumed == list(RlBackendOps.MEMORY_TAG_ORDER if tags is None else tags) - expected = list(allocator.hold_token_indices) * 2 if should_init_hold_indexes else [-1] * 8 - assert req_manager.req_to_token_indexs[req_manager.HOLD_REQUEST_ID].tolist() == expected - assert req_manager.req_to_token_indexs[:2].eq(-1).all() - backend.logger.error.assert_not_called() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py deleted file mode 100644 index 63e78e4e5d..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py +++ /dev/null @@ -1,129 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import DSparkProposer -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager - - -def test_dspark_prefill_uses_a_shallow_copy_for_target_hidden(): - forwarded_inputs = [] - draft_model = SimpleNamespace(forward=forwarded_inputs.append) - proposer = DSparkProposer( - backend=SimpleNamespace(draft_models=[draft_model]), - enable_dynmaic_mtp=False, - ) - model_input = SimpleNamespace( - is_prefill=True, - b_position_delta=None, - b_req_idx=torch.tensor([3, 5], dtype=torch.int32), - input_ids=torch.arange(5, dtype=torch.int64), - mtp_draft_input_hiddens=None, - ) - target_hidden = torch.empty((5, 8)) - - proposer.fill_draft_model_kv_state( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=torch.tensor([11, 13], dtype=torch.int64), - ) - - assert len(forwarded_inputs) == 1 - assert forwarded_inputs[0] is not model_input - assert forwarded_inputs[0].mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - -def test_dspark_commits_verify_kv_and_builds_parallel_block(monkeypatch): - block_size = 3 - draft_token_ids = torch.tensor([30, 31, 32, 40, 41, 42], dtype=torch.int64) - confidence_logits = torch.tensor( - [ - [-10.0, 0.0, 10.0], - [10.0, 0.0, -10.0], - ] - ) - block_output = SimpleNamespace( - mtp_collector=SimpleNamespace( - draft_token_ids=draft_token_ids, - confidence_logits=confidence_logits, - ) - ) - forwarded_inputs = [] - - def forward(model_input): - forwarded_inputs.append(model_input) - return block_output if len(forwarded_inputs) == 2 else SimpleNamespace() - - draft_model = SimpleNamespace( - block_size=block_size, - mask_token_id=99, - forward=forward, - ) - proposer = DSparkProposer( - backend=SimpleNamespace(draft_models=[draft_model]), - enable_dynmaic_mtp=True, - ) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) - monkeypatch.setattr( - g_pin_mem_manager, - "get_const_gpu_tensor", - lambda *, shape, fill_value, dtype, **_: torch.full(shape, fill_value, dtype=dtype), - ) - monkeypatch.setattr( - g_pin_mem_manager, - "async_copy_from_gpu_tensor", - lambda *, gpu_tensor, **_: gpu_tensor.clone(), - ) - - target_hidden = torch.empty((5, 8)) - model_input = SimpleNamespace( - is_prefill=False, - batch_size=5, - total_token_num=40, - max_q_seq_len=1, - max_kv_seq_len=9, - input_ids=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int64), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32), - b_mtp_index=torch.arange(5, dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 8, 9], dtype=torch.int32), - b_position_delta=torch.tensor([1, 1, 1, 2, 2], dtype=torch.int32), - b_shared_seq_len=torch.tensor([3, 3, 3, 6, 6], dtype=torch.int32), - b_shared_radix_node_id=torch.tensor([17, 17, 17, 19, 19], dtype=torch.int64), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - mtp_draft_input_hiddens=None, - ) - - proposal = proposer.propose_next( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32), - ) - - assert len(forwarded_inputs) == 2 - verify_draft_input, block_draft_input = forwarded_inputs - assert verify_draft_input is not model_input - assert verify_draft_input.mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - assert torch.equal(block_draft_input.input_ids, torch.tensor([11, 99, 99, 21, 99, 99])) - assert torch.equal(block_draft_input.b_req_idx, torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_mtp_index, torch.zeros(6, dtype=torch.int32)) - assert torch.equal(block_draft_input.b_seq_len, torch.tensor([6, 7, 8, 10, 11, 12], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_position_delta, torch.tensor([1, 1, 1, 2, 2, 2], dtype=torch.int32)) - assert block_draft_input.mtp_draft_input_hiddens is None - - assert torch.equal(proposal.token_ids, torch.tensor([[30, 31], [40, 41]], dtype=torch.int64)) - torch.testing.assert_close( - proposal.schedule_scores, - torch.tensor([[0.01, 0.5], [0.99, 0.5]], dtype=torch.float32), - ) - assert torch.equal(proposal.schedule_scores_cpu, proposal.schedule_scores) - assert proposal.schedule_scores_cpu is not proposal.schedule_scores diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py deleted file mode 100644 index 71de619b9a..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py +++ /dev/null @@ -1,93 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_no_att import ( - EagleNoAttProposer, -) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_no_att_recurrently_proposes_from_accepted_tails(): - device = "cuda" - draft_calls = [] - draft_outputs = [ - SimpleNamespace( - token_ids=torch.tensor([21, 24], device=device), - token_probs=torch.tensor([0.8, 0.7], device=device), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[102.0, 103.0], [108.0, 109.0]], device=device)), - ), - SimpleNamespace( - token_ids=torch.tensor([31, 34], device=device), - token_probs=torch.tensor([0.6, 0.5], device=device), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[202.0, 203.0], [208.0, 209.0]], device=device)), - ), - ] - - def forward(model_input): - draft_calls.append( - { - "batch_size": model_input.batch_size, - "input_ids": model_input.input_ids.cpu(), - "draft_hidden": model_input.mtp_draft_input_hiddens.cpu(), - "b_req_idx": model_input.b_req_idx.cpu(), - "b_seq_len": model_input.b_seq_len.cpu(), - } - ) - return draft_outputs[len(draft_calls) - 1] - - backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward)], - _gen_argmax_token_ids=lambda output: output.token_ids, - _gen_argmax_token_ids_and_prob=lambda output: ( - output.token_ids, - output.token_probs, - ), - ) - proposer = EagleNoAttProposer(backend=backend, enable_dynmaic_mtp=True) - target_model_input = SimpleNamespace( - batch_size=5, - input_ids=torch.tensor([10, 11, 12, 13, 14], device=device), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32, device=device), - b_mtp_index=torch.tensor([0, 1, 2, 0, 1], dtype=torch.int32, device=device), - b_seq_len=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int32, device=device), - b_position_delta=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device=device), - b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6], dtype=torch.int32, device=device), - b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90], dtype=torch.int64, device=device), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - ) - target_model_output = SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=torch.arange(10, dtype=torch.float32, device=device).reshape(5, 2)) - ) - - proposal = proposer.propose_next( - target_model_input=target_model_input, - target_model_output=target_model_output, - target_next_token_ids=target_model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32, device=device), - ) - - torch.testing.assert_close(proposal.token_ids, torch.tensor([[21, 31], [24, 34]], device=device)) - torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.8, 0.6], [0.7, 0.5]], device=device)) - assert len(draft_calls) == 2 - assert draft_calls[0]["batch_size"] == 2 - torch.testing.assert_close(draft_calls[0]["input_ids"], torch.tensor([11, 14])) - torch.testing.assert_close(draft_calls[0]["draft_hidden"], torch.tensor([[2.0, 3.0], [8.0, 9.0]])) - torch.testing.assert_close(draft_calls[0]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[0]["b_seq_len"], torch.tensor([11, 21], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([21, 24])) - torch.testing.assert_close( - draft_calls[1]["draft_hidden"], - torch.tensor([[102.0, 103.0], [108.0, 109.0]]), - ) - assert target_model_input.batch_size == 5 - torch.testing.assert_close(target_model_input.input_ids.cpu(), torch.tensor([10, 11, 12, 13, 14])) - - -def test_eagle_no_att_fill_hook_is_noop(): - proposer = EagleNoAttProposer(backend=SimpleNamespace(draft_models=[]), enable_dynmaic_mtp=False) - - proposer.fill_draft_model_kv_state(None, None, None) diff --git a/unit_tests/server/router/model_infer/structured_output/test_sampling.py b/unit_tests/server/router/model_infer/structured_output/test_sampling.py index bad4430c71..e26c60d9a2 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_sampling.py +++ b/unit_tests/server/router/model_infer/structured_output/test_sampling.py @@ -96,8 +96,6 @@ def test_sample_returns_requested_output_logprobs(compiler, sampling_manager, mo reference = logits.log_softmax(dim=-1) sampled_id = int(token_ids[0]) expected = reference[0, sampled_id] - if int((logits[0] > logits[0, sampled_id]).sum()) >= 20: - expected = expected.new_tensor(-9999.0) torch.testing.assert_close(logprobs[0], expected) if topk: assert len(output_topk_logprobs) == 1 @@ -106,7 +104,8 @@ def test_sample_returns_requested_output_logprobs(compiler, sampling_manager, mo assert row.token_ids.tolist() == list(range(req.vocab_size - 1, req.vocab_size - 1 - topk, -1)) torch.testing.assert_close(row.logprobs, reference[0, row.token_ids].cpu()) else: - assert output_topk_logprobs == [None] + assert output_topk_logprobs[0].token_ids.numel() == 0 + assert output_topk_logprobs[0].sampled_rank.item() == req.vocab_size - sampled_id def test_sample_candidate_logprobs_follow_mtp_row_selection(compiler, sampling_manager, monkeypatch): @@ -148,8 +147,10 @@ def test_sample_candidate_logprobs_follow_mtp_row_selection(compiler, sampling_m reference = logits.log_softmax(dim=-1).cpu() assert middle["next_token_logprob"] == pytest.approx(reference[2, 10].item()) assert first["output_top_logprobs"][1] == pytest.approx(reference[0, [256, 255, 254]].tolist()) - assert last["output_top_logprobs"][1] == pytest.approx(reference[3, [256, 255]].tolist()) - assert logprobs_cpu[1].item() == -9999.0 + assert last["output_top_logprobs"][1] == pytest.approx(reference[3, [256, 255]].tolist(), rel=1e-5) + assert logprobs_cpu[1].item() == pytest.approx(reference[1, 0].item()) + assert first["next_token_rank"] == 1 + assert last["next_token_rank"] == 1 @pytest.mark.parametrize("topk", [0, 1, 3, 20]) @@ -170,8 +171,6 @@ def test_sample_reuses_sampled_normalization(compiler, sampling_manager, monkeyp topk_by_index = output_topk_logprobs for row in range(4): expected = reference[row, ids[row]] - if row in rows and int((logits[row] > logits[row, ids[row]]).sum()) >= 20: - expected = expected.new_tensor(-9999.0) torch.testing.assert_close(logprobs[row], expected) if topk and row in rows: top_ids, top_scores = topk_by_index[row].token_ids, topk_by_index[row].logprobs @@ -179,7 +178,7 @@ def test_sample_reuses_sampled_normalization(compiler, sampling_manager, monkeyp torch.testing.assert_close(top_scores, reference[row, top_ids].cpu()) -def test_sample_preserves_sampled_score_and_keeps_ties(compiler, sampling_manager, monkeypatch): +def test_sample_rank_keeps_ties_and_includes_zero_width_candidates(compiler, sampling_manager, monkeypatch): reqs = [sampling_req(compiler, sampling_manager, req_idx=i, regular_constraint=None) for i in range(4)] for req, topk in zip(reqs, [0, 0, 3, -1]): req.sampling_param.shm_param.top_k = req.vocab_size @@ -190,12 +189,10 @@ def test_sample_preserves_sampled_score_and_keeps_ties(compiler, sampling_manage ids = torch.full((4,), 256, dtype=torch.int64, device="cuda") monkeypatch.setattr(sampling, "_random_sample", lambda *args: ids) _, logprobs, output_topk_logprobs = sampling.sample(logits, reqs, [256]) - assert [None if row is None else row.batch_index for row in output_topk_logprobs] == [None, None, 2, None] + assert [None if row is None else row.batch_index for row in output_topk_logprobs] == [0, 1, 2, None] + assert [row.sampled_rank.item() for row in output_topk_logprobs[:3]] == [20, 21, 21] reference = logits.log_softmax(dim=-1) - expected = reference[:, 256].clone() - expected[1] = -9999.0 - expected[2] = -9999.0 - torch.testing.assert_close(logprobs, expected) + torch.testing.assert_close(logprobs, reference[:, 256]) candidate_ids, candidate_logprobs = output_topk_logprobs[2].token_ids, output_topk_logprobs[2].logprobs assert len(set(candidate_ids.tolist())) == 3 and (candidate_ids < 20).all() torch.testing.assert_close(candidate_logprobs, reference[2, candidate_ids].cpu()) @@ -442,16 +439,59 @@ def copy_snapshot_and_wait(manager): completed.synchronize() +@pytest.mark.parametrize("topk", [0, 1, 3, 20]) +@pytest.mark.parametrize("sampled_rank", [2, 20, 21]) +def test_sample_rank_is_independent_of_candidate_count(compiler, sampling_manager, monkeypatch, topk, sampled_rank): + req = sampling_req(compiler, sampling_manager, regular_constraint=None) + req.sampling_param.shm_param.top_k = req.vocab_size + req.sampling_param.shm_param.output_top_logprobs = topk + logits = torch.linspace(3.0, 0.1, req.vocab_size, device="cuda").view(1, -1) + ids = torch.tensor([sampled_rank - 1], device="cuda") + monkeypatch.setattr(sampling, "_random_sample", lambda *args: ids) + _, logprobs, rows = sampling.sample(logits, [req], [256]) + torch.testing.assert_close(logprobs[0], logits.log_softmax(-1)[0, ids[0]]) + assert rows[0].sampled_rank.item() == sampled_rank + assert rows[0].token_ids.numel() == topk + + @pytest.mark.parametrize("vocab_size", [30, 2051, 32769]) -def test_capture_across_vocab_tiles_preserves_sampled_scores(vocab_size): +def test_capture_rank_across_vocab_tiles_preserves_sampled_scores(vocab_size): logits = torch.linspace(3.0, 0.1, vocab_size, device="cuda").repeat(4, 1) ids = torch.tensor([1, 19, 20, vocab_size - 1], device="cuda") logprobs = logits.log_softmax(-1).gather(1, ids[:, None]).squeeze(1) original = logprobs.clone() rows = torch.tensor([[0, 0, 0], [1, 1, 0], [2, 3, 1], [3, 20, 4]], device="cuda", dtype=torch.int32) - candidate_ids, scores = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 20, 24) - torch.testing.assert_close(logprobs[:2], original[:2], rtol=0, atol=0) - assert logprobs[2].item() == -9999.0 and logprobs[3].item() == -9999.0 + candidate_ids, scores, ranks = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 20, 24) + torch.testing.assert_close(logprobs, original, rtol=0, atol=0) + assert ranks.tolist() == [2, 20, 21, vocab_size] for row, width, offset in rows.tolist(): assert candidate_ids[offset : offset + width].tolist() == list(range(width)) torch.testing.assert_close(scores[offset : offset + width], logits.log_softmax(-1)[row, :width]) + + +@pytest.mark.parametrize( + "logits,sampled_id,expected_ids,expected_rank", + [ + ([0.0, -torch.inf, 1.0], 1, [2, 0, 1], 4), + ([0.0, torch.inf, 1.0], 1, [1, 2, 0], 1), + ([-torch.inf, -torch.inf, -torch.inf], 0, [0, 1, 2], 4), + ], +) +def test_capture_preserves_nonfinite_candidates(logits, sampled_id, expected_ids, expected_rank): + logits = torch.tensor([logits], device="cuda") + ids = torch.tensor([sampled_id], device="cuda") + logprobs = torch.tensor([[0.0]], device="cuda") + rows = torch.tensor([[0, 3, 0]], device="cuda", dtype=torch.int32) + + candidate_ids, scores, ranks = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 3, 3) + + assert candidate_ids.tolist() == expected_ids + assert ranks.tolist() == [expected_rank] + assert not torch.isnan(scores).any() + if logits[0, sampled_id] == -torch.inf and not torch.isneginf(logits).all(): + torch.testing.assert_close(scores[:2], logits.log_softmax(-1)[0, [2, 0]]) + assert scores[2].item() == -torch.inf + elif logits[0, sampled_id] == torch.inf: + assert scores.tolist() == [0.0, -torch.inf, -torch.inf] + else: + torch.testing.assert_close(scores, torch.full_like(scores, -torch.log(torch.tensor(3.0)))) diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py index d76269b0d6..9a76597ce3 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_state.py +++ b/unit_tests/server/router/model_infer/structured_output/test_state.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest +import numpy as np import torch import xgrammar as xgr @@ -308,6 +309,9 @@ def test_pd_first_token_updates_constraint_before_sampling( req.pd_task_success_num = 0 req.shm_req.shm_prompt_ids.arr.append(0) req.shm_req.shm_logprobs = SimpleNamespace(arr=[None] * (req.shm_req.input_len + 1)) + req.sampling_param.shm_param.output_top_logprobs = 2 + req.shm_req.shm_output_topk_logprobs = SimpleNamespace(arr=np.full((2, 2, 2), -1, dtype=np.float64)) + req.shm_req.output_topk_logprobs_row = lambda output_index: req.shm_req.shm_output_topk_logprobs.arr[output_index] req.set_next_gen_token_id = MethodType(infer_batch.InferReq.set_next_gen_token_id, req) req.update_finish_status = lambda **kwargs: None sampling_params = make_mask_buffers([req, good]) @@ -316,6 +320,8 @@ def test_pd_first_token_updates_constraint_before_sampling( has_error=False, first_gen_token_id=first_token, first_gen_token_logprob=0.0, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([first_token, 7], [-0.1, -0.2]), ) backend = SimpleNamespace( shm_pd_trans_io_buffer=SimpleNamespace(read_obj=lambda: [transfer], sub_state=lambda: None), @@ -333,9 +339,14 @@ def test_pd_first_token_updates_constraint_before_sampling( assert req.cur_output_len == 1 and req.pd_task_success_num == 1 assert req.shm_req.shm_prompt_ids.arr == [ord("P")] + list(history) + [first_token] + assert req.shm_req.shm_logprobs.arr[req.shm_req.input_len] == (0.0, 21) assert sampling_params.req_to_next_token_ids[req.req_idx, 0].item() == first_token if is_master_in_dp: assert req.shm_req.shm_cur_output_len == 1 + assert req.shm_req.shm_output_topk_logprobs.arr[0].tolist() == [ + [first_token, 7], + [-0.1, -0.2], + ] state = req.output_constraint if first_token == ord("!"): assert req.finish_status.status == FinishStatus.FINISHED_ERROR diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py deleted file mode 100644 index c243c07a49..0000000000 --- a/unit_tests/server/router/test_kv_capacity.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.attention_state.allocators import full_att_kv_allocator - - -pytestmark = pytest.mark.usefixtures("kv_allocator_environment") - - -@pytest.mark.parametrize( - "tp,dp,nnodes,selected_ranks", - [(4, 1, 1, [0]), (8, 4, 1, [0, 2, 4, 6]), (8, 4, 2, [0, 2]), (8, 8, 2, [0, 1, 2, 3]), (8, 1, 2, [0])], -) -@pytest.mark.parametrize("reader_first", [False, True]) -def test_router_reads_each_local_dp_groups_allocator( - monkeypatch: pytest.MonkeyPatch, - tp: int, - dp: int, - nnodes: int, - selected_ranks: list[int], - reader_first: bool, -): - monkeypatch.setattr( - full_att_kv_allocator, "get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=dp, nnodes=nnodes) - ) - reader = full_att_kv_allocator.KVCapacityReader() if reader_first else None - allocators = [] - for rank in range(tp // nnodes): - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", str(rank)) - allocator = FullAttKVAllocator(16 + 4 * rank, 4, pin_memory=False) - assert allocator._shared_can_use_token_num.get_value() == allocator.size - allocator.alloc(4) - allocators.append(allocator) - reader = reader or full_att_kv_allocator.KVCapacityReader() - assert len(reader.counters) == len(selected_ranks) - for dp_rank, rank in enumerate(selected_ranks): - allocator = allocators[rank] - initial = allocator.size - 4 - assert reader.get_available_tokens(dp_rank) == initial - tokens = allocator.alloc(4) - assert reader.get_available_tokens(dp_rank) == initial - 4 - allocator.free(tokens) - assert reader.get_available_tokens(dp_rank) == initial - allocator.free_all() - assert reader.get_available_tokens(dp_rank) == allocator.size - - -def test_failed_page_operations_leave_shared_capacity_unchanged(): - allocator = FullAttKVAllocator(16, 4, pin_memory=False) - counter = allocator._shared_can_use_token_num - tokens = allocator.alloc(16) - with pytest.raises(AssertionError): - allocator.alloc(4) - with pytest.raises(AssertionError): - allocator.free(tokens[:3]) - assert counter.get_value() == allocator.can_use_mem_size == 0 - allocator.free(tokens) - assert counter.get_value() == allocator.can_use_mem_size == 16 diff --git a/unit_tests/server/router/test_profiler_service.py b/unit_tests/server/router/test_profiler_service.py deleted file mode 100644 index 020c495cb8..0000000000 --- a/unit_tests/server/router/test_profiler_service.py +++ /dev/null @@ -1,53 +0,0 @@ -import socket -import types - -import pytest -import rpyc - -from lightllm.server.router.profiler_service import RouterProfilerCmdQueue, start_router_profiler_server - - -def _free_port(): - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - - -@pytest.fixture -def profiler_server(monkeypatch): - port = _free_port() - from lightllm.utils import shm_port_args - - monkeypatch.setattr(shm_port_args, "get_shm_port_args", lambda: types.SimpleNamespace(router_profiler_port=port)) - queue = RouterProfilerCmdQueue() - server, _ = start_router_profiler_server(types.SimpleNamespace(enable_profiling="torch_profiler"), queue) - yield server, port, queue - server.close() - - -def test_profiler_server_binds_loopback_only(profiler_server): - server, _port, _queue = profiler_server - assert server.listener.getsockname()[0] == "127.0.0.1" - - -def test_profiler_server_accepts_known_cmds(profiler_server): - _server, port, queue = profiler_server - conn = rpyc.connect("127.0.0.1", port) - conn.root.profiler_cmd("start") - assert queue.pop() == "start" - conn.root.profiler_cmd("stop") - assert queue.pop() == "stop" - conn.close() - - -def test_profiler_server_rejects_unknown_cmd(profiler_server): - _server, port, queue = profiler_server - conn = rpyc.connect("127.0.0.1", port) - with pytest.raises(Exception): - conn.root.profiler_cmd("drop_tables") - assert queue.pop() is None - conn.close() - - -def test_disabled_profiling_starts_nothing(): - assert start_router_profiler_server(types.SimpleNamespace(enable_profiling=None), None) == (None, None) diff --git a/unit_tests/server/test_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py deleted file mode 100644 index 949021cd34..0000000000 --- a/unit_tests/server/test_api_start_deepep.py +++ /dev/null @@ -1,45 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): - for name in ( - "_set_envs_and_config", - "auto_set_max_req_total_len", - "auto_set_fused_shared_experts", - "set_unique_server_name", - "auto_set_page_size", - ): - monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) - - def stop(args): - raise RuntimeError("startup validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", stop) - (tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"})) - args = StartArgs( - run_mode="decode", - model_dir=str(tmp_path), - enable_ep_moe=True, - running_max_req_size=256, - mtp_mode="vanilla_no_att", - mtp_step=2, - max_req_total_len=4096, - batch_max_tokens=4096, - chunked_prefill_size=2048, - eos_id=[0], - data_type="float16", - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - router_token_ratio=0.0, - ) - - with pytest.raises(RuntimeError, match="startup validation passed"): - _launch_subprocesses(args) - - assert args.batch_max_tokens == 768 diff --git a/unit_tests/server/test_application_stop.py b/unit_tests/server/test_application_stop.py index 0603c29c67..6825028434 100644 --- a/unit_tests/server/test_application_stop.py +++ b/unit_tests/server/test_application_stop.py @@ -113,6 +113,54 @@ async def raw_results(*args, **kwargs): asyncio.run(run()) +@pytest.mark.parametrize("output_top_logprobs", [-1, 0, 2]) +def test_tgi_num_beam_keeps_true_scores(output_top_logprobs): + async def run(): + async def raw_results( + prompt, + sampling_params, + multimodal, + request, + pd_upload_websocket, + pd_event, + ): + assert sampling_params.output_top_logprobs == output_top_logprobs + for sub_id, scores in enumerate(((-2.302585, -3.912023), (-2.995732, -6.907755))): + text = "A" if sub_id == 0 else "B" + for index, score in enumerate(scores): + cumlogprob = sum(scores[: index + 1]) / (index + 1) + yield ( + sub_id, + text, + { + "prompt_tokens": 1, + "logprobs": {}, + "logprob": score, + "cumlogprob": cumlogprob, + }, + FinishStatus(FinishStatus.FINISHED_LENGTH if index else FinishStatus.NO_FINISH), + ) + + manager = _manager( + SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), + SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), + raw_results, + ) + request = SimpleNamespace( + json=AsyncMock( + return_value={ + "inputs": "prompt", + "num_beam": 2, + "parameters": {"output_top_logprobs": output_top_logprobs}, + } + ) + ) + response = await tgi_generate_impl(request, manager) + assert json.loads(response.body)["generated_text"] == "AA" + + asyncio.run(run()) + + @pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) @pytest.mark.parametrize("slave", [False, True]) def test_http_generate_only_wraps_client_output(mode, slave): diff --git a/unit_tests/server/test_chat_output_logprobs.py b/unit_tests/server/test_chat_output_logprobs.py index 18d4e64b11..f7d4704921 100644 --- a/unit_tests/server/test_chat_output_logprobs.py +++ b/unit_tests/server/test_chat_output_logprobs.py @@ -7,7 +7,7 @@ import pytest from pydantic import ValidationError -from tokenizers import Tokenizer, decoders, models, pre_tokenizers +from tokenizers import Tokenizer, decoders, models from transformers import PreTrainedTokenizerFast from lightllm.server import api_http, api_openai, build_prompt as prompt_builder, reasoning_parser @@ -107,9 +107,8 @@ async def generate(prompt, params, *args, **kwargs): assert [entry["token"] for entry in reasoning] == (["think"] if separate_reasoning else []) expected_content = "你 yes" if separate_reasoning else "think你 yes" assert bytes(byte for entry in content for byte in entry["bytes"]).decode() == expected_content - assert [entry["token"] for entry in content] == (["你", " yes"] if separate_reasoning else ["think", "你", " yes"]) - assert [entry["logprob"] for entry in content] == ([-0.75, -0.25] if separate_reasoning else [-0.25, -0.75, -0.25]) - assert [entry["logprob"] for entry in reasoning] == ([-0.25] if separate_reasoning else []) + assert len(content) == (4 if separate_reasoning else 5) + assert all(entry["logprob"] == -0.25 for entry in content + reasoning) expected = [{"token": " yes", "logprob": -0.125, "bytes": [32, 121, 101, 115]}] if topk > 0 else [] assert all(entry["top_logprobs"] == expected for entry in content + reasoning) @@ -129,13 +128,14 @@ def test_candidate_shm_is_reused_and_resized(monkeypatch): assert worker.shm_output_topk_logprobs.shm.size == worker.shm_output_topk_logprobs.arr.nbytes assert worker.output_topk_logprobs_row(output_index=0)[0].tolist() == [-1, -1] worker.shm_logprobs.arr[3] = (-0.25, 1) - worker.output_topk_logprobs_row(output_index=0)[:] = ([9, 8], [-0.5, -1.0]) + worker.output_topk_logprobs_row(output_index=0)[:] = ([16_777_217, 8], [-0.5, -1.0]) worker.output_topk_logprobs_row(output_index=1)[:] = ([7, 6], [-0.25, -1.5]) worker.detach_shm_arrays() owner.detach_shm_arrays() owner.link_logprobs_shm_array() assert owner.shm_logprobs.arr[3].tolist() == (-0.25, 1) + assert owner.output_topk_logprobs_row(output_index=0)[0].tolist() == [16_777_217, 8] assert owner.output_topk_logprobs_row(output_index=1).tolist() == [[7, 6], [-0.25, -1.5]] owner.detach_shm_arrays() @@ -166,11 +166,26 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p params.output_top_logprobs = 2 params.group_request_id = 0 finish = FinishStatus() - d_first = (0, "first", {"count_output_tokens": 1, "node_mode": "decode", "output_top_logprobs": []}, finish) + d_first = ( + 0, + "first", + { + "count_output_tokens": 1, + "node_mode": "decode", + "output_top_logprobs": [(1, -0.5)], + "logprob_prompt_ids": prompt_ids, + }, + finish, + ) p_first = ( 0, "first", - {"count_output_tokens": 1, "node_mode": "prefill", "output_top_logprobs": [(1, -0.5)]}, + { + "count_output_tokens": 1, + "node_mode": "prefill", + "output_top_logprobs": [(1, -0.5)], + "logprob_prompt_ids": prompt_ids, + }, finish, ) d_next = (0, "second", {"count_output_tokens": 2}, FinishStatus(end_status)) @@ -198,7 +213,8 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p try: _, text, metadata, _ = await stream.__anext__() assert text == "first" - assert metadata["output_top_logprobs"] == ([(1, -0.5)] if needs_prefill else []) + assert metadata["logprob_prompt_ids"] == prompt_ids + assert metadata["output_top_logprobs"] == [(1, -0.5)] _, text, _, status = await stream.__anext__() assert text == "second" and status.status == end_status finally: @@ -207,7 +223,23 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p @pytest.fixture def chat_response(monkeypatch): - async def run(tokenizer, pieces, *, stream, separate_reasoning=True, model_type=None, tools=None, topk=0, rank=2): + async def run( + tokenizer, + pieces, + *, + stream, + separate_reasoning=True, + model_type=None, + tools=None, + topk=0, + rank=2, + tool_parser="qwen25", + scores=None, + first_metadata=None, + candidates=None, + logprobs=True, + return_choices=False, + ): async def build_prompt(*args): return "prompt" @@ -215,10 +247,11 @@ async def generate(prompt, params, *args, **kwargs): for index, (token_id, delta) in enumerate(pieces): yield 0, delta, { "id": token_id, - "logprob": -2.401099, + **(first_metadata or {} if index == 0 else {}), + "logprob": scores[index] if scores is not None else -2.401099, "logprobs": {token_id: {"rank": rank}}, "prompt_tokens": 10, - "output_top_logprobs": [(0, -0.125)] * topk, + "output_top_logprobs": candidates if candidates is not None else [(0, -0.125)] * topk, }, FinishStatus(FinishStatus.FINISHED_LENGTH if index == len(pieces) - 1 else FinishStatus.NO_FINISH) monkeypatch.setattr(api_openai, "build_prompt", build_prompt) @@ -235,13 +268,13 @@ async def generate(prompt, params, *args, **kwargs): ), raising=False, ) - monkeypatch.setattr(api_http.g_objs, "args", SimpleNamespace(tool_call_parser="qwen25"), raising=False) + monkeypatch.setattr(api_http.g_objs, "args", SimpleNamespace(tool_call_parser=tool_parser), raising=False) request = ChatCompletionRequest( messages=[], stream=stream, separate_reasoning=separate_reasoning, - logprobs=True, - top_logprobs=topk, + logprobs=logprobs, + top_logprobs=topk if logprobs else None, tools=tools, ) response = await api_openai.chat_completions_impl(request, SimpleNamespace()) @@ -255,6 +288,8 @@ async def generate(prompt, params, *args, **kwargs): choices.extend(json.loads(line[6:])["choices"]) else: choices = response.model_dump()["choices"] + if return_choices: + return choices messages = [choice["delta" if stream else "message"] for choice in choices] logprobs = [choice["logprobs"] for choice in choices if choice["logprobs"]] return ( @@ -278,22 +313,26 @@ def make_tokenizer(pieces, decoder=None): @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("topk", [0, 1, 3, 20]) @pytest.mark.parametrize("rank", [2, 20, 21]) -async def test_chat_keeps_the_sampled_logprob(chat_response, stream, topk, rank): +async def test_chat_sampled_score_uses_fixed_rank_cutoff(chat_response, stream, topk, rank): tokenizer = make_tokenizer(["answer"]) _, _, content, _ = await chat_response(tokenizer, [(0, "answer")], stream=stream, topk=topk, rank=rank) - assert content[0]["logprob"] == -2.401099 + assert content[0]["logprob"] == (-9999.0 if rank > 20 else -2.401099) assert len(content[0]["top_logprobs"]) == topk @pytest.mark.anyio @pytest.mark.parametrize("stream", [False, True]) -async def test_empty_increments_fuse_into_the_following_text(chat_response, stream): - tokenizer = make_tokenizer(["<0xE4>", "<0xBD>", "<0xA0>"]) +@pytest.mark.parametrize("sequence", [False, True]) +async def test_byte_fallback_keeps_every_generated_token(chat_response, stream, sequence): + decoder = decoders.ByteFallback() + if sequence: + decoder = decoders.Sequence([decoder, decoders.Fuse()]) + tokenizer = make_tokenizer(["<0xE4>", "<0xBD>", "<0xA0>"], decoder) text, _, content, _ = await chat_response(tokenizer, [(0, ""), (1, ""), (2, "你")], stream=stream) assert text == "你" - assert [entry["token"] for entry in content] == ["你"] - assert content[0]["logprob"] == pytest.approx(-2.401099 * 3) - assert content[0]["bytes"] == list("你".encode()) + assert len(content) == 3 + assert [entry["bytes"] for entry in content] == [[0xE4], [0xBD], [0xA0]] + assert bytes(byte for entry in content for byte in entry["bytes"]).decode() == text @pytest.mark.anyio @@ -349,6 +388,28 @@ async def test_tool_logprobs_only_contain_message_content(chat_response, stream, assert reasoning == [] +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_mistral_tool_parser_keeps_message_scores(chat_response, stream): + pieces = [ + "Checking.", + "[TOOL_CALLS] [", + '{"name":"weather","arguments":{"city":"Paris"}}', + "]", + ] + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, _ = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + tools=tools, + tool_parser="mistral", + ) + assert text == "Checking." + assert [entry["token"] for entry in content] == ["Checking."] + + @pytest.mark.anyio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("separate_reasoning", [False, True]) @@ -371,10 +432,9 @@ async def test_harmony_headers_are_not_visible_logprobs(chat_response, stream, s if split_header: pieces[5:6] = ["<|sta", "rt|>"] tokenizer = make_tokenizer(list(dict.fromkeys(pieces)), decoders.Fuse()) - generated = [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces] text, reasoning_text, content, reasoning = await chat_response( tokenizer, - generated, + [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces], stream=stream, separate_reasoning=separate_reasoning, model_type="gpt-oss", @@ -416,153 +476,473 @@ async def test_harmony_repeated_header_words_are_real_content(chat_response, str @pytest.mark.anyio -async def test_tool_delimiter_fused_into_another_token_is_rejected(chat_response): - pieces = ["Checking.", '{"name":"weather","arguments":{"city":"Paris"}}', ""] +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_harmony_commentary_content_keeps_its_scores(chat_response, stream, separate_reasoning): + pieces = [ + "<|channel|>", + "commentary to=functions.weather", + "<|message|>", + "Checking.", + "<|end|>", + "<|start|>", + "assistant", + "<|channel|>", + "final", + "<|message|>", + "Answer", + "<|return|>", + ] + ids = {piece: index for index, piece in enumerate(dict.fromkeys(pieces))} + tokenizer = make_tokenizer(list(ids), decoders.Fuse()) + text, _, content, reasoning = await chat_response( + tokenizer, + [(ids[piece], piece) for piece in pieces], + stream=stream, + separate_reasoning=separate_reasoning, + model_type="gpt-oss", + ) + assert text == "Checking.Answer" + assert [entry["token"] for entry in content] == ["Checking.", "Answer"] + assert reasoning == [] + + +@pytest.mark.anyio +@pytest.mark.parametrize("valid_tool", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_nonstream_tool_logprobs_follow_actual_extraction(chat_response, valid_tool, stream): + pieces = ["Checking.", '{"name":"weather","arguments":{"city":"Paris"}}'] + if valid_tool: + pieces += [""] tokenizer = make_tokenizer(pieces, decoders.Fuse()) tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] - with pytest.raises(ValueError, match="token sequence"): - await chat_response(tokenizer, list(enumerate(pieces)), stream=False, tools=tools) + text, _, content, _ = await chat_response(tokenizer, list(enumerate(pieces)), stream=stream, tools=tools) + assert text == ("Checking." if valid_tool else "".join(pieces)) + assert [entry["token"] for entry in content] == (pieces[:1] if valid_tool else pieces) -def test_empty_increments_fuse_and_trailing_ones_are_dropped(): +def test_unknown_token_bytes_do_not_drop_empty_increments(): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: "") - buffer = ChatLogprobBuffer(tokenizer, None, False, True) + tokenizer = SimpleNamespace( + decode=lambda ids, **kwargs: "\ufffd" if len(ids) == 1 else "你", + convert_ids_to_tokens=lambda token_id: "unknown-byte", + ) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) for token_id, text in enumerate(["", "", "你"]): - buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + span = buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + buffer.assign([span], len(text)) content, _ = buffer.message_logprobs() - assert [token["token"] for token in content] == ["你"] - assert content[0]["logprob"] == -0.75 - assert content[0]["bytes"] == list("你".encode()) - - buffer.record("", {"id": 3, "logprob": -0.5}, FinishStatus(FinishStatus.FINISHED_STOP)) - assert buffer.message_logprobs() == (None, None) + assert len(content) == 3 + assert all(token["bytes"] is None for token in content) -def test_reasoning_delimiter_must_be_one_token(): +def test_combined_reasoning_and_content_token_is_only_returned_once(monkeypatch): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = make_tokenizer(["", "ThinkingAnswer"], decoders.Fuse()) - with pytest.raises(ValueError, match="token sequence"): - ChatLogprobBuffer(tokenizer, "qwen3", False, True) + pieces = ["", "ThinkingAnswer"] + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + monkeypatch.setattr(reasoning_parser, "get_token_id", tokenizer.convert_tokens_to_ids) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=False) + for token_id, text in enumerate(pieces): + buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + _, content, spans = parser.parse_stream_chunk_with_spans(text) + buffer.assign(spans, len(content)) + content, reasoning = buffer.message_logprobs() + assert content is None + assert [token["token"] for token in reasoning] == pieces[1:] -def test_multitoken_reasoning_delimiter_includes_newline(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer +def test_reasoning_span_api_preserves_literal_delimiters_and_flush(monkeypatch): + monkeypatch.setattr(reasoning_parser, "get_token_id", lambda token: 0) + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=False) + text = "literal " + reasoning, content, spans = parser.parse_stream_chunk_with_spans(text) + assert not reasoning and content == text + assert "".join(text[span.start : span.end] for span in spans) == text + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=True) + assert parser.parse_stream_chunk_with_spans("{"name":"weather","arguments":{"city":"', + "Paris", + '"}};', + "Paris", + ] + tokenizer = make_tokenizer(list(dict.fromkeys(pieces)), decoders.Fuse()) + scores = [-0.1 * (index + 1) for index in range(len(pieces))] + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, _ = await chat_response( + tokenizer, + [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces], + stream=stream, + tools=tools, + tool_parser="llama3", + scores=scores, + ) + assert text == prefix + "Paris" + assert [item["token"] for item in content] == ([prefix] if prefix else []) + ["Paris"] + assert [item["logprob"] for item in content] == ([scores[0]] if prefix else []) + [scores[-1]] - pieces = ["<|channel>", "thought", "\n", "secret", "", "answer"] + +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_nonstream_reasoning_uses_message_trimming(chat_response, stream, separate_reasoning): + pieces = ["Thinking", "", "\n\n", "Answer"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gemma4", False, True) - for piece in pieces: - buffer.record( - piece, - {"id": tokenizer.convert_tokens_to_ids(piece), "logprob": -0.2}, - FinishStatus(FinishStatus.FINISHED_STOP if piece == pieces[-1] else FinishStatus.NO_FINISH), + text, reasoning_text, content, reasoning = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + separate_reasoning=separate_reasoning, + model_type="qwen3", + scores=[-0.1, -0.2, -0.3, -0.4], + ) + expected = (["Thinking"] if not separate_reasoning else []) + (["\n\n"] if stream else []) + ["Answer"] + assert text == "".join(expected) + assert [item["token"] for item in content] == expected + assert [item["logprob"] for item in content] == ( + ([-0.1] if not separate_reasoning else []) + ([-0.3] if stream else []) + [-0.4] + ) + assert bytes(byte for item in content for byte in item["bytes"]).decode() == text + assert [item["token"] for item in reasoning] == (["Thinking"] if separate_reasoning else []) + + +async def first_token_metadata(prompt_ids, output_index=1): + import asyncio + from unittest.mock import Mock + from lightllm.server.httpserver.manager import HttpServerManager + + manager = HttpServerManager.__new__(HttpServerManager) + manager.is_multinode_tp_slave = False + manager.first_time_costs = Mock() + manager.latest_success_infer_time_mark = Mock() + status = SimpleNamespace( + event=asyncio.Event(), + lock=asyncio.Lock(), + has_timed_out_waiting_for_inference=lambda: False, + out_token_info_list=[(0, "", {"count_output_tokens": output_index}, FinishStatus())], + ) + status.event.set() + params = SamplingParams() + params.best_of = 1 + params.output_top_logprobs = 1 + generator = manager._wait_to_token_package(0, prompt_ids, 0, params, status, None) + try: + _, _, metadata, _ = await generator.__anext__() + return metadata + finally: + await generator.aclose() + + +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("sequence", [False, True]) +async def test_prompt_generation_sentencepiece_boundary(chat_response, stream, sequence): + import numpy as np + from lightllm.server.detokenization.decode import decode_token + from lightllm.server.detokenization.decode_req import DecodeReq + from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START + + decoder = ( + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] ) - content, reasoning = buffer.message_logprobs() - assert [token["token"] for token in reasoning] == ["secret"] - assert [token["token"] for token in content] == ["answer"] + if sequence + else decoders.Metaspace() + ) + tokenizer = make_tokenizer(["▁Hello", "▁world", "▁there"], decoder) + prompt_ids = [0] * 6 + [MULTIMODAL_TOKEN_ID_START] + req = SimpleNamespace( + request_id=0, + group_req_id=0, + input_len=len(prompt_ids), + shm_prompt_ids=SimpleNamespace(arr=np.array(prompt_ids + [1])), + sample_params=SimpleNamespace( + print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=True + ), + ) + decode_req = DecodeReq(req) + decode_req.output_ids.append(1) + delta = decode_token(tokenizer, decode_req, 1, []) + metadata = await first_token_metadata(prompt_ids) + assert metadata["logprob_prompt_ids"] == prompt_ids[-5:] + assert "logprob_prompt_ids" not in await first_token_metadata(prompt_ids, output_index=2) + text, _, content, _ = await chat_response( + tokenizer, + [(1, delta)], + stream=stream, + first_metadata=metadata, + topk=1, + candidates=[(2, -0.2)], + ) + assert text == delta == " world" + assert content[0]["token"] == " world" + assert bytes(content[0]["bytes"]).decode() == text + candidate = content[0]["top_logprobs"][0] + assert candidate["token"] == " there" + assert bytes(candidate["bytes"]).decode() == " there" -def test_gemma4_answer_without_thought_header_stays_in_content(): +@pytest.mark.anyio +@pytest.mark.parametrize("parser_name", ["glm47", "qwen3_coder"]) +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +async def test_stream_tool_body_uses_generated_position(parser_name, chunk_size): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer + from lightllm.server.function_call_parser import FunctionCallParser + from lightllm.server.api_models import Tool - pieces = ["<|channel>", "thought", "\n", "", "2"] + call = ( + "weathercityParis" + if parser_name == "glm47" + else "Paris" + ) + pieces = [call, "Paris"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gemma4", True, True) - buffer.record("2", {"id": tokenizer.convert_tokens_to_ids("2"), "logprob": -0.2}, FinishStatus()) - content, reasoning = buffer.message_logprobs() - assert reasoning is None - assert [token["token"] for token in content] == ["2"] + tools = [ + Tool( + type="function", + function={"name": "weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, + ) + ] + parser = FunctionCallParser(tools, parser_name) + raw = "".join(pieces) + normal_text, all_calls, positions = "", [], [] + for start in range(0, len(raw), chunk_size): + normal, calls, spans = parser.parse_stream_chunk_with_spans(raw[start : start + chunk_size]) + normal_text += normal + all_calls.extend(calls) + positions.extend(pos for span in spans for pos in range(span.start, span.end)) + assert normal_text == "Paris" and all_calls + assert positions == list(range(len(call), len(raw))) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) + parser = FunctionCallParser(tools, parser_name) + for index, piece in enumerate(pieces): + raw_span = buffer.record(piece, {"id": index, "logprob": -0.1 * (index + 1)}, FinishStatus()) + buffer.assign([raw_span], len(piece)) + _, _, spans = parser.parse_stream_chunk_with_spans(piece) + content = buffer.content_logprobs(spans)["content"] + assert [(item["token"], item["logprob"]) for item in content] == [("Paris", -0.2)] -def test_incomplete_delimiter_prefix_is_visible_text(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer +@pytest.mark.anyio +@pytest.mark.parametrize("text", ['{"name":"weather","arguments":{"city":"Paris"}}', '{"name":123}']) +@pytest.mark.parametrize("tool_parser", [None, "llama3"]) +async def test_plain_json_with_tools_keeps_scores(chat_response, text, tool_parser): + tokenizer = make_tokenizer([text], decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + for enabled in (False, True): + choices = await chat_response( + tokenizer, + [(0, text)], + stream=False, + tools=tools, + tool_parser=tool_parser, + scores=[-0.6], + logprobs=enabled, + return_choices=True, + ) + assert choices[0]["message"]["content"] == text + assert choices[0]["message"]["tool_calls"] is None + if enabled: + assert [(item["token"], item["logprob"]) for item in choices[0]["logprobs"]["content"]] == [(text, -0.6)] + else: + assert choices[0]["logprobs"] is None + - pieces = ["", "", "nope"] +@pytest.mark.anyio +@pytest.mark.parametrize("tool_parser", ["glm47", "qwen3_coder"]) +async def test_suppressed_tool_whitespace_has_no_scores(chat_response, tool_parser): + call = ( + "weathercityParis" + if tool_parser == "glm47" + else "Paris" + ) + pieces = [call, "\n", "Done"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "qwen3", False, True) - for piece in ["", "ok"], decoders.Fuse()) - buffer = ChatLogprobBuffer( - tokenizer, - None, - False, - True, - tools=[{"type": "function", "function": {"name": "ok"}}], - tool_call_parser="llama3", + tokenizer = make_tokenizer( + ["▁Hello", "<0xE4>", "<0xBD>", "<0xA0>", "▁world"], + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] + ), ) - assert buffer._sequences + req = SimpleNamespace( + request_id=0, + group_req_id=0, + input_len=1, + shm_prompt_ids=SimpleNamespace(arr=np.array([0, 1, 2, 3])), + sample_params=SimpleNamespace( + print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=True + ), + ) + decode_req = DecodeReq(req) + decoder = ChatTokenDecoder(tokenizer) + decoder.set_prompt([0]) + deltas = [] + for token_id in (1, 2, 3): + decode_req.output_ids.append(token_id) + deltas.append(decode_token(tokenizer, decode_req, token_id, [])) + assert decoder.token_logprob(token_id, -0.1)["bytes"] == [[0xE4], [0xBD], [0xA0]][token_id - 1] + decoder.advance(token_id) + candidate = decoder.token_logprob(4, -0.2) + assert candidate == {"token": " world", "logprob": -0.2, "bytes": list(b" world")} + assert deltas == ["", "", "你"] -def test_byte_level_candidate_preserves_raw_bytes(): - from lightllm.server.chat_logprob_buffer import _candidate_logprob +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_special_candidates_keep_identity(chat_response, monkeypatch, stream): + from lightllm.server.core.objs import sampling_params as params_module - backend = Tokenizer(models.BPE({"Ã": 0}, [])) - backend.pre_tokenizer = pre_tokenizers.ByteLevel() - backend.decoder = decoders.ByteLevel() - tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) - assert _candidate_logprob(tokenizer, 0, -0.5)["bytes"] == [195] + tokenizer = make_tokenizer(["prefix", "answer", "", ""], decoders.Fuse()) + tokenizer.add_special_tokens({"eos_token": "", "pad_token": ""}) + monkeypatch.setattr(params_module, "SKIP_SPECIAL_TOKENS", True) + _, _, content, _ = await chat_response( + tokenizer, + [(1, "answer")], + stream=stream, + topk=2, + first_metadata={"logprob_prompt_ids": [0]}, + candidates=[(2, -0.2), (3, -0.3)], + ) + assert content[0]["top_logprobs"] == [ + {"token": "", "logprob": -0.2, "bytes": list(b"")}, + {"token": "", "logprob": -0.3, "bytes": list(b"")}, + ] -def test_harmony_commentary_attributes_are_hidden_logprobs(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_pending_byte_candidates_in_chat(chat_response, stream): + tokenizer = make_tokenizer( + ["▁Hello", "<0xE4>", "<0xBD>", "<0xA0>", "▁world"], + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] + ), + ) + text, _, content, _ = await chat_response( + tokenizer, + [(1, ""), (2, ""), (3, "你"), (4, " world")], + stream=stream, + first_metadata={"logprob_prompt_ids": [0]}, + topk=1, + candidates=[(4, -0.2)], + ) + assert text == "你 world" + assert len(content) == 4 + assert bytes(byte for item in content for byte in item["bytes"]).decode() == text + assert all( + item["top_logprobs"] == [{"token": " world", "logprob": -0.2, "bytes": list(b" world")}] for item in content + ) + +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_reasoning_then_tool_uses_api_spans(chat_response, stream, separate_reasoning): pieces = [ - "<|start|>", - "assistant", - "<|channel|>", - "commentary to=functions.weather", - "<|message|>", - '{"city":"Paris"}', - "<|end|>", - "<|start|>", - "assistant", - "<|channel|>", - "final", - "<|message|>", - "Answer", - "<|return|>", + "Thinking", + "", + "Checking.", + "", + '{"name":"weather","arguments":{"city":"Paris"}}', + "", ] - unique = list(dict.fromkeys(pieces)) - ids = {piece: index for index, piece in enumerate(unique)} - tokenizer = make_tokenizer(unique, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gpt-oss", False, True) - for index, piece in enumerate(pieces): - buffer.record( - piece, - {"id": ids[piece], "logprob": -0.1}, - FinishStatus(FinishStatus.FINISHED_STOP if index == len(pieces) - 1 else FinishStatus.NO_FINISH), - ) - content, reasoning = buffer.message_logprobs() - assert [token["token"] for token in content] == ["Answer"] - assert reasoning is None - + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, reasoning = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + separate_reasoning=separate_reasoning, + model_type="qwen3", + tools=tools, + ) + expected = ["Checking."] if separate_reasoning else ["Thinking", "Checking."] + assert text == "".join(expected) + assert [item["token"] for item in content] == expected + assert [item["token"] for item in reasoning] == (["Thinking"] if separate_reasoning else []) -def test_mistral_tool_parser_is_rejected(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = make_tokenizer(["[TOOL_CALLS] [", "]"], decoders.Fuse()) - with pytest.raises(ValueError, match="mistral"): - ChatLogprobBuffer( - tokenizer, - None, - False, - True, - tools=[{"type": "function", "function": {"name": "ok"}}], - tool_call_parser="mistral", - ) +@pytest.mark.anyio +async def test_reasoning_flush_does_not_include_hidden_tool_scores(chat_response): + call = "weathercityParis" + pieces = [call, "\n" - f'<{DSML}invoke name="write_file">\n' - f'<{DSML}parameter name="path" string="true">src/app.py\n' - f'<{DSML}parameter name="content" string="true">print("hi")\n\n' - f"\n" - f'<{DSML}invoke name="get_weather">\n' - f'<{DSML}parameter name="location" string="true">Hangzhou\n' - f'<{DSML}parameter name="days" string="false">3\n' - f"\n" - f"" - ) - - -def collect_stream(parser_name, text, chunk_size): - parser = FunctionCallParser(TOOLS, parser_name) - calls = {} - for start in range(0, len(text), chunk_size): - _, deltas = parser.parse_stream_chunk(text[start : start + chunk_size]) - for delta in deltas: - if delta.name is not None: - calls[delta.tool_index] = {"name": delta.name, "arguments": ""} - calls[delta.tool_index]["arguments"] += delta.parameters - return calls - - -@pytest.mark.parametrize("parser_name,wrapper", [("deepseekv32", "function_calls"), ("deepseekv4", "tool_calls")]) -@pytest.mark.parametrize("chunk_size", [1, 5, 16]) -def test_streamed_arguments_concatenate_to_valid_json(parser_name, wrapper, chunk_size): - calls = collect_stream(parser_name, dsml_text(wrapper), chunk_size) - assert [call["name"] for call in calls.values()] == ["write_file", "get_weather"] - assert json.loads(calls[0]["arguments"]) == {"path": "src/app.py", "content": 'print("hi")\n'} - assert json.loads(calls[1]["arguments"]) == {"location": "Hangzhou", "days": 3} diff --git a/unit_tests/server/test_deepseek_tool_parser_stream.py b/unit_tests/server/test_deepseek_tool_parser_stream.py deleted file mode 100644 index e18eb9a005..0000000000 --- a/unit_tests/server/test_deepseek_tool_parser_stream.py +++ /dev/null @@ -1,143 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser - -TOOLS = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "write_file", - "parameters": {"type": "object", "properties": {"content": {"type": "string"}}}, - }, - } - ), - Tool.model_validate( - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"location": {"type": "string"}}}, - }, - } - ), -] - -DSML = "｜DSML｜" - - -def _dsml_invoke(name, param, value): - return ( - f'<{DSML}invoke name="{name}">\n' - f'<{DSML}parameter name="{param}" string="true">{value}\n' - f"" - ) - - -def _stream(parser, deltas): - names = [] - arguments = {} - for delta in deltas: - _, calls = parser.parse_stream_chunk(delta) - for call in calls: - if call.name: - names.append(call.name) - arguments[call.tool_index] = arguments.get(call.tool_index, "") + (call.parameters or "") - return names, arguments - - -@pytest.mark.parametrize("tool_call_parser,block", [("deepseekv32", "function_calls"), ("deepseekv4", "tool_calls")]) -def test_dsml_streams_long_invokes_in_small_deltas(tool_call_parser, block): - content = "print(1)\n" * 500 - deltas = [f"<{DSML}{block}>", "\n", f"<{DSML}invoke", ' name="write_file">\n', f"<{DSML}parameter"] - deltas += [' name="content" string="true">'] - deltas += [content[i : i + 4] for i in range(0, len(content), 4)] - deltas += [f"", "\n", f"", "\n"] - deltas += [f"<{DSML}invoke", ' name="get_weather">\n', f"<{DSML}parameter", ' name="location" string="true">'] - deltas += ["Hang", "zhou", f"", "\n", f"", "\n", f""] - - parser = FunctionCallParser(TOOLS, tool_call_parser) - names, arguments = _stream(parser, deltas) - - assert names == ["write_file", "get_weather"] - assert json.loads(arguments[0]) == {"content": content} - assert json.loads(arguments[1]) == {"location": "Hangzhou"} - - -def test_dsml_emits_buffered_invoke_on_next_delta(): - # Both invokes arrive in one delta; the second one is still in the buffer - # and must be emitted by the next delta, which has no invoke end tag. - text = f"<{DSML}function_calls>\n" - text += _dsml_invoke("write_file", "content", "x = 1") + "\n" - text += _dsml_invoke("get_weather", "location", "Paris") + "\n" - - parser = FunctionCallParser(TOOLS, "deepseekv32") - names, arguments = _stream(parser, [text, f""]) - - assert names == ["write_file", "get_weather"] - assert json.loads(arguments[0]) == {"content": "x = 1"} - assert json.loads(arguments[1]) == {"location": "Paris"} - - -def _v3_call(name, payload): - return "<｜tool▁call▁begin｜>function<｜tool▁sep｜>" + name + "\n```json\n" + payload + "\n```<｜tool▁call▁end｜>" - - -def test_deepseekv3_streams_long_arguments_in_small_deltas(): - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - call = _v3_call("write_file", payload) - deltas = ["<｜tool▁calls▁begin｜>"] + [call[i : i + 4] for i in range(0, len(call), 4)] - - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == arguments - - -def test_deepseekv3_holds_arguments_until_the_json_fence_closes(): - payload = json.dumps({"content": "x" * 200}) - head = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>write_file\n```json\n" - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, [head] + [payload[i : i + 4] for i in range(0, len(payload), 4)]) - assert names == [] - assert streamed == {} - - # The closing fence is split across deltas. The name is emitted on the delta that - # completes it, and the arguments follow on the next delta, same as one regex search. - names, streamed = _stream(parser, ["\n", "``", "`<｜tool▁call▁end｜>", "<｜tool▁calls▁end｜>"]) - assert names == ["write_file"] - assert json.loads(streamed[0]) == {"content": "x" * 200} - - -def test_deepseekv3_keeps_a_fence_inside_the_argument(): - content = "line1\n```\nline2\n```python\ncode" - payload = json.dumps({"content": content}) - text = "<｜tool▁calls▁begin｜>" + _v3_call("write_file", payload) - deltas = [text[: len("<｜tool▁calls▁begin｜>")]] + [ - text[i : i + 4] for i in range(len("<｜tool▁calls▁begin｜>"), len(text), 4) - ] - - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == {"content": content} - - -def test_deepseekv31_streams_long_arguments_in_small_deltas(): - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - deltas = ["<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>", "write_file", "<｜tool▁sep｜>"] - deltas += [payload[i : i + 4] for i in range(0, len(payload), 4)] - deltas += ["<｜tool▁call▁end｜>", "<｜tool▁calls▁end｜>"] - - parser = FunctionCallParser(TOOLS, "deepseekv31") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == arguments diff --git a/unit_tests/server/test_function_call_common_prefix.py b/unit_tests/server/test_function_call_common_prefix.py deleted file mode 100644 index 498f5d45a2..0000000000 --- a/unit_tests/server/test_function_call_common_prefix.py +++ /dev/null @@ -1,67 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser, _find_common_prefix - - -def _reference_common_prefix(s1, s2): - prefix = "" - for a, b in zip(s1, s2): - if a != b: - break - prefix += a - return prefix - - -@pytest.mark.parametrize( - "s1,s2", - [ - ("", ""), - ("", "abc"), - ("abc", ""), - ("abc", "abc"), - ("abc", "abcd"), - ("abcd", "abc"), - ("abc", "abd"), - ("xbc", "abc"), - ('{"a": "x"}', '{"a": "xy"}'), - ('{"a": "é😀"}', '{"a": "é😀b"}'), - ("a" * 1000 + "b", "a" * 1000 + "c"), - ], -) -def test_find_common_prefix_matches_reference(s1, s2): - assert _find_common_prefix(s1, s2) == _reference_common_prefix(s1, s2) - - -def test_qwen25_streams_long_arguments_in_small_deltas(): - tools = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "write_file", - "parameters": {"type": "object", "properties": {"content": {"type": "string"}}}, - }, - } - ) - ] - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - deltas = ["", '\n{"name": "write_file", "arguments": '] - deltas += [payload[i : i + 4] for i in range(0, len(payload), 4)] - deltas += ["}\n", ""] - - parser = FunctionCallParser(tools, "qwen25") - streamed = "" - names = [] - for delta in deltas: - _, calls = parser.parse_stream_chunk(delta) - for call in calls: - if call.name: - names.append(call.name) - streamed += call.parameters or "" - - assert names == ["write_file"] - assert json.loads(streamed) == arguments diff --git a/unit_tests/server/test_glm47_tool_parser.py b/unit_tests/server/test_glm47_tool_parser.py deleted file mode 100644 index 8f626d6343..0000000000 --- a/unit_tests/server/test_glm47_tool_parser.py +++ /dev/null @@ -1,127 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser - - -TOOLS = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "Read", - "parameters": {"type": "object", "properties": {"file_path": {"type": "string"}}}, - }, - } - ), - Tool.model_validate( - { - "type": "function", - "function": { - "name": "Write", - "parameters": { - "type": "object", - "properties": {"content": {"type": "string"}, "overwrite": {"type": "boolean"}}, - }, - }, - } - ), -] - - -def collect_stream(text, chunk_size): - parser = FunctionCallParser(TOOLS, "glm47") - content = "" - calls = {} - for start in range(0, len(text), chunk_size): - normal, deltas = parser.parse_stream_chunk(text[start : start + chunk_size]) - content += normal - for delta in deltas: - if delta.name is not None: - assert delta.tool_index not in calls, "a tool call must have exactly one name event" - calls[delta.tool_index] = {"name": delta.name, "arguments": ""} - calls[delta.tool_index]["arguments"] += delta.parameters - return content, calls - - -@pytest.mark.parametrize("separator", ["", "\n"]) -@pytest.mark.parametrize("chunk_size", [1, 7, 39, 4096]) -def test_repeated_tool_calls_have_separate_indices(separator, chunk_size): - text = ( - f'Write{separator}content{{"ready": true}}' - "overwritefalse" - f"Read{separator}file_path/tmp/a" - f"Read{separator}file_path/tmp/b" - ) - _, calls = collect_stream(text, chunk_size) - assert list(calls) == [0, 1, 2] - assert [call["name"] for call in calls.values()] == ["Write", "Read", "Read"] - assert [json.loads(call["arguments"]) for call in calls.values()] == [ - {"content": '{"ready": true}', "overwrite": False}, - {"file_path": "/tmp/a"}, - {"file_path": "/tmp/b"}, - ] - _, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text) - assert [call.tool_index for call in nonstream] == [0, 1, 2] - assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()] - - -@pytest.mark.parametrize("value", ['{"count": 1}', "123", "false", " return 1\n", '"quoted"']) -def test_string_arguments_preserve_type_and_whitespace(value): - text = f"Writecontent{value}" - _, calls = collect_stream(text, 1) - assert json.loads(calls[0]["arguments"])["content"] == value - - -@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) -def test_text_around_tool_calls_survives_chunk_boundaries(chunk_size): - text = ( - "Before Readfile_path/tmp/a" - " between Readfile_path/tmp/b after" - ) - content, calls = collect_stream(text, chunk_size) - assert content == "Before between after" - assert list(calls) == [0, 1] - - -def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name(): - parser = FunctionCallParser(TOOLS, "glm47") - normal, calls = parser.parse_stream_chunk("Writecontent") - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name == "Write" - assert calls[0].parameters == "" - - value = ' line 1\n{"literal": true}\n' + "x" * 100_000 - for chunk in (value[:13], value[13:27], value[27:]): - normal, calls = parser.parse_stream_chunk(chunk) - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert calls[0].parameters == "" - - normal, calls = parser.parse_stream_chunk("") - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"content": value} - - _, calls = parser.parse_stream_chunk("Write\n") - assert len(calls) == 1 - assert calls[0].tool_index == 1 - assert calls[0].name == "Write" - - -def test_undefined_buffered_tool_does_not_consume_a_call_index(): - parser = FunctionCallParser(TOOLS, "glm47") - _, calls = parser.parse_stream_chunk("unknown\ncontentx") - assert calls == [] - _, calls = parser.parse_stream_chunk("Read\n") - assert len(calls) == 1 - assert calls[0].name == "Read" - assert calls[0].tool_index == 0 diff --git a/unit_tests/server/test_hypercorn_config.py b/unit_tests/server/test_hypercorn_config.py deleted file mode 100644 index 46ee8b7e7a..0000000000 --- a/unit_tests/server/test_hypercorn_config.py +++ /dev/null @@ -1,26 +0,0 @@ -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _hypercorn_config_args - - -def test_hypercorn_config_defaults_to_none(): - args = make_argument_parser().parse_args([]) - - assert args.hypercorn_config is None - - -def test_hypercorn_config_is_parsed(): - args = make_argument_parser().parse_args(["--hypercorn_config", "hypercorn.toml"]) - - assert args.hypercorn_config == "hypercorn.toml" - - -def test_hypercorn_args_use_default_keep_alive_without_config(): - args = make_argument_parser().parse_args([]) - - assert _hypercorn_config_args(args) == ["--keep-alive", "10"] - - -def test_hypercorn_args_do_not_set_keep_alive_with_config(): - args = make_argument_parser().parse_args(["--hypercorn_config", "hypercorn.toml"]) - - assert _hypercorn_config_args(args) == ["--config", "hypercorn.toml"] diff --git a/unit_tests/server/test_multimodal_chunk_args.py b/unit_tests/server/test_multimodal_chunk_args.py deleted file mode 100644 index c799235904..0000000000 --- a/unit_tests/server/test_multimodal_chunk_args.py +++ /dev/null @@ -1,15 +0,0 @@ -import argparse - -from lightllm.server.api_cli import add_cli_args -from lightllm.server.core.objs.start_args_type import StartArgs - - -def test_multimodal_chunk_flag(): - parser = add_cli_args(argparse.ArgumentParser()) - assert parser.parse_args(["--disable_chunked_mm_input"]).disable_chunked_mm_input is True - - -def test_multimodal_chunking_keeps_existing_defaults(): - parser = add_cli_args(argparse.ArgumentParser()) - assert parser.parse_args([]).disable_chunked_mm_input is False - assert StartArgs().disable_chunked_mm_input is False diff --git a/unit_tests/server/test_pd_cache_aware.py b/unit_tests/server/test_pd_cache_aware.py index 1b111122fb..f71105c0d7 100644 --- a/unit_tests/server/test_pd_cache_aware.py +++ b/unit_tests/server/test_pd_cache_aware.py @@ -1,105 +1,6 @@ -import asyncio -from types import SimpleNamespace -from unittest.mock import patch - import pytest -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( - create_prompt_cache_tree_client, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware import ( - BalanceRelThresholdController, - CacheAwareConfig, - CacheAwarePolicy, -) from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import PromptCacheTree -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_selector import ( - LoadBalancedCacheAwareSelector, - PDSelectionExtraInfo, -) - -pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") - - -def _worker(address: str, dispatched_prompt_chars: int = 0, dispatched_req_num: int = 0): - return SimpleNamespace( - client_ip_port=address, - connection_key=address, - dispatched_prompt_chars=dispatched_prompt_chars, - dispatched_req_num=dispatched_req_num, - ) - - -def _local_policy(config): - return CacheAwarePolicy(prompt_cache_tree=create_prompt_cache_tree_client(config), config=config) - - -def _insert_prompt(policy, prompt, worker): - sample_stride = policy.config.sample_stride - asyncio.run(policy.prompt_cache_tree.insert_sampled(prompt[::sample_stride], worker.connection_key, sample_stride)) - - -def _prefix_match(policy, prompt): - sample_stride = policy.config.sample_stride - return asyncio.run( - policy.prompt_cache_tree.prefix_match_sampled(prompt[::sample_stride], len(prompt), sample_stride) - ) - - -def test_balance_threshold_controller_adjusts_threshold_each_window(): - config = CacheAwareConfig(cache_hit_rate_window_size=3, balance_rel_threshold_step=0.1) - controller = BalanceRelThresholdController() - - for cache_hit_rate, expected_threshold in ( - (0.5, 1.5), - (0.6, 1.6), - (0.4, 1.5), - (0.4, 1.5), - ): - for _ in range(config.cache_hit_rate_window_size): - controller.append(cache_hit_rate) - controller.update_config(config) - assert config.balance_rel_threshold == pytest.approx(expected_threshold) - - -@pytest.mark.parametrize( - ("initial_threshold", "first_hit_rate", "second_hit_rate", "expected_threshold"), - ( - (1.95, 0.5, 0.6, 2.0), - (1.05, 0.5, 0.4, 1.0), - ), -) -def test_balance_threshold_controller_limits_threshold_range( - initial_threshold, - first_hit_rate, - second_hit_rate, - expected_threshold, -): - config = CacheAwareConfig( - balance_rel_threshold=initial_threshold, - cache_hit_rate_window_size=1, - balance_rel_threshold_step=0.1, - ) - controller = BalanceRelThresholdController() - - controller.append(first_hit_rate) - controller.update_config(config) - controller.append(second_hit_rate) - controller.update_config(config) - - assert config.balance_rel_threshold == pytest.approx(expected_threshold) - - -def test_cache_aware_updates_threshold_from_inference_cache_hit_rate(): - policy = _local_policy(CacheAwareConfig(cache_hit_rate_window_size=2)) - - for _ in range(2): - policy.record_prompt_cache_hit_rate(0.25) - assert policy.config.balance_rel_threshold == 1.5 - for _ in range(2): - policy.record_prompt_cache_hit_rate(0.75) - - assert policy.config.balance_rel_threshold == pytest.approx(1.55) @pytest.mark.parametrize("operation", ["insert", "prefix_match"]) @@ -179,143 +80,3 @@ def test_prompt_cache_tree_iterative_insert_keeps_lru_eviction_order(): result = tree.prefix_match_sampled("ab", input_char_count=2, sample_stride=1) assert result.prefill_node == "prefill-ab" assert result.matched_char_count == 1 - - -def test_cache_aware_selector_returns_selected_nodes_and_estimated_cache_hit_rate(): - selector = LoadBalancedCacheAwareSelector(pd_manager=None) - p_node = _worker("10.0.0.1:8000") - d_node = SimpleNamespace( - client_ip_port="10.0.0.2:8000", - run_status=SimpleNamespace(total_token_usage_rate=0.0), - ) - selector.update_nodes([p_node], [d_node]) - prompt = "x" * 1025 - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=123.0, - ): - _insert_prompt(selector.policy, prompt, p_node) - - selected_p_node, selected_d_node, selection_extra_info = asyncio.run(selector.select_p_d_node(prompt, None, None)) - - assert selected_p_node is p_node - assert selected_d_node is d_node - assert selection_extra_info == PDSelectionExtraInfo( - estimated_cache_hit_rate=pytest.approx(1.0), - cache_last_insert_time=123.0, - ) - - -def test_cache_aware_reinsert_refreshes_matched_node_time(): - policy = _local_policy(CacheAwareConfig()) - worker = _worker("10.0.0.1:8000") - prompt = "shared prefix " * 100 - - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=100.0, - ): - _insert_prompt(policy, prompt, worker) - first_result = _prefix_match(policy, prompt) - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=160.0, - ): - _insert_prompt(policy, prompt, worker) - second_result = _prefix_match(policy, prompt) - - assert first_result.last_insert_time == 100.0 - assert second_result.last_insert_time == 160.0 - - -def test_cache_aware_inserts_prompt_only_when_explicitly_recorded(): - policy = _local_policy(CacheAwareConfig()) - worker = _worker("10.0.0.1:8000") - prompt = "new prompt" - - selected_worker, cache_info = asyncio.run(policy.select_worker([worker], prompt)) - assert selected_worker is worker - assert cache_info == PDSelectionExtraInfo() - assert _prefix_match(policy, prompt).prefill_node is None - - _insert_prompt(policy, prompt, worker) - - assert _prefix_match(policy, prompt).prefill_node == worker.connection_key - - -def test_cache_aware_keeps_cache_worker_when_inflight_load_is_balanced(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=110, dispatched_req_num=2) - least_loaded_worker = _worker("10.0.0.2:8000", dispatched_prompt_chars=100, dispatched_req_num=2) - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, cache_info = asyncio.run(policy.select_worker([cache_worker, least_loaded_worker], prompt)) - - assert selected_worker is cache_worker - # 前缀树每 512 个字符抽样一次，命中率保持原有的保守估算方式。 - assert cache_info.estimated_cache_hit_rate == pytest.approx(1025 / len(prompt)) - - -def test_cache_aware_uses_least_loaded_worker_when_cache_worker_is_overloaded(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=2000, dispatched_req_num=2) - least_loaded_worker = _worker("10.0.0.2:8000", dispatched_prompt_chars=100, dispatched_req_num=2) - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, cache_info = asyncio.run(policy.select_worker([cache_worker, least_loaded_worker], prompt)) - - assert selected_worker is least_loaded_worker - assert cache_info.estimated_cache_hit_rate == 0.0 - - -def test_cache_aware_keeps_cache_worker_when_both_workers_are_idle(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000") - other_worker = _worker("10.0.0.2:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([other_worker, cache_worker], prompt)) - - assert selected_worker is cache_worker - - -def test_cache_aware_forces_idle_worker_over_busy_cache_worker(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=100, dispatched_req_num=1) - idle_worker = _worker("10.0.0.2:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([cache_worker, idle_worker], prompt)) - - assert selected_worker is idle_worker - - -def test_cache_aware_randomly_selects_an_idle_worker_without_cache_match(): - policy = _local_policy(CacheAwareConfig()) - first = _worker("10.0.0.1:8000") - second = _worker("10.0.0.2:8000") - - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware.random.choice", return_value=second - ) as choice: - selected_worker, _ = asyncio.run(policy.select_worker([first, second], "new prompt")) - - assert selected_worker is second - choice.assert_called_once_with([first, second]) - - -def test_cache_aware_matches_cache_only_within_idle_workers(): - policy = _local_policy(CacheAwareConfig()) - busy_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=100, dispatched_req_num=1) - cache_idle_worker = _worker("10.0.0.2:8000") - other_idle_worker = _worker("10.0.0.3:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_idle_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([other_idle_worker, busy_worker, cache_idle_worker], prompt)) - - assert selected_worker is cache_idle_worker diff --git a/unit_tests/server/test_pd_master_launcher_integration.py b/unit_tests/server/test_pd_master_launcher_integration.py deleted file mode 100644 index 8833be9202..0000000000 --- a/unit_tests/server/test_pd_master_launcher_integration.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Exercise real PD master startup, supervision, and cleanup.""" - -import json -from pathlib import Path -import time -from urllib.error import HTTPError - -import pytest - -from unit_tests.server.pd_master_test_utils import assert_processes_exit, running_pd_master - - -@pytest.mark.parametrize( - "failed_process", - [ - "api_server", - "pd_node_stats_server", - "prompt_cache_tree_server", - "req_id_generator_server", - "metric_manager", - "haproxy", - ], -) -def test_real_launcher_discovers_workers_before_pd_readiness_and_fails_as_a_group(tmp_path, failed_process): - with running_pd_master(tmp_path) as master: - assert len({worker["node_id"] for worker in master.workers}) == 2 - assert len({worker["port"] for worker in master.workers}) == 2 - public_port = int(master.public_url.rsplit(":", 1)[1]) - for endpoint in master.workers: - info = master.get_json("/get_server_info", "http://" + endpoint["host_ip_port"]) - assert info["port"] == public_port - assert info["pd_node_id"] == endpoint["node_id"] - assert info["httpserver_workers"] == len(master.workers) - assert set(info["pd_master_multi_ports"]) == {worker["port"] for worker in master.workers} - assert info["config_server_host"] is None - assert info["config_server_port"] is None - assert "pd_master_ipc_name" not in info - assert "metric_port" not in info - - deadline = time.monotonic() + 15 - while True: - try: - master.get_json("/readiness") - except HTTPError as readiness: - assert readiness.code == 503 - assert json.load(readiness) == {"status": "not ready"} - break - assert time.monotonic() < deadline, "Public readiness checks did not converge" - time.sleep(0.1) - with master.opener.open(master.public_url + "/metrics", timeout=3) as response: - assert b"lightllm_request_count" in response.read() - - runtime = master.haproxy_runtime_dir() - assert runtime.is_dir() - descendants = master.descendants() - if failed_process == "haproxy": - candidates = [ - process - for process in descendants - if "haproxy" in process.name() and process.ppid() == master.process.pid - ] - else: - candidates = [process for process in descendants if "::" + failed_process in process.name()] - assert candidates, [(process.pid, process.name()) for process in descendants] - candidates[0].kill() - assert master.process.wait(timeout=20) != 0 - assert_processes_exit(descendants) - assert not runtime.exists() - - -@pytest.mark.parametrize( - "failure_stage", - ["backend_start", "proxy_start"], -) -def test_partial_startup_failure_stops_every_started_process_and_removes_runtime(tmp_path, failure_stage): - launcher_script = Path(__file__).with_name("pd_master_fault_launcher.py") - observations = tmp_path / "cleanup.json" - with running_pd_master( - tmp_path, - launcher_script=launcher_script, - wait_for_startup=False, - extra_args=("--failure-stage", failure_stage, "--cleanup-observations", str(observations)), - ) as master: - descendants = {} - deadline = time.monotonic() + 60 - while master.process.poll() is None: - descendants.update({process.pid: process for process in master.descendants()}) - assert time.monotonic() < deadline, master.log_path.read_text() - time.sleep(0.05) - assert master.process.returncode != 0 - assert "injected" in master.log_path.read_text() or failure_stage == "backend_start" - assert_processes_exit(descendants.values()) - cleanup = json.loads(observations.read_text()) - assert all(not Path(runtime).exists() for runtime in cleanup["runtime_paths"]) - assert cleanup["runtime_before_stop"] == [True] - assert cleanup["runtime_after_stop"] == [True] - assert cleanup["backends_stopped"] is True diff --git a/unit_tests/server/test_prefill_queue_start_args.py b/unit_tests/server/test_prefill_queue_start_args.py deleted file mode 100644 index 6a802fd1ae..0000000000 --- a/unit_tests/server/test_prefill_queue_start_args.py +++ /dev/null @@ -1,27 +0,0 @@ -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _auto_set_prefill_queue_strategy -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize( - "run_mode, automatic_strategy", - [("normal", "default"), ("prefill", "hrrn"), ("decode", "default"), ("visual_only", "default")], -) -@pytest.mark.parametrize("configured_strategy", [None, "auto", "default", "promote_shortest", "hrrn"]) -def test_auto_set_prefill_queue_strategy(run_mode, automatic_strategy, configured_strategy): - cli_args = ["--run_mode", run_mode] - if configured_strategy is not None: - cli_args.extend(["--prefill_queue_strategy", configured_strategy]) - args = StartArgs(**vars(make_argument_parser().parse_args(cli_args))) - assert args.prefill_queue_strategy == (configured_strategy or "auto") - expected = automatic_strategy if configured_strategy in (None, "auto") else configured_strategy - _auto_set_prefill_queue_strategy(args) - assert args.prefill_queue_strategy == expected - - -def test_start_args_default_prefill_queue_strategy_matches_cli(): - assert StartArgs().prefill_queue_strategy == make_argument_parser().parse_args([]).prefill_queue_strategy == "auto" - cli_option = next(action for action in make_argument_parser()._actions if action.dest == "prefill_queue_strategy") - assert set(cli_option.choices) == set(StartArgs.__dataclass_fields__["prefill_queue_strategy"].metadata["choices"]) diff --git a/unit_tests/server/test_prompt_cache_tree_service.py b/unit_tests/server/test_prompt_cache_tree_service.py deleted file mode 100644 index b85cd67902..0000000000 --- a/unit_tests/server/test_prompt_cache_tree_service.py +++ /dev/null @@ -1,228 +0,0 @@ -"""Prompt cache tree local and RPyC component behavior.""" - -import asyncio -import os -import signal -import socket -import threading -import uuid -from unittest.mock import Mock, patch - -import pytest -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import PDNodeStatsState -from lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware import ( - CacheAwareConfig, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import ( - PromptCacheMatchResult, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( - LocalPromptCacheTreeClient, - RemotePromptCacheTreeClient, - create_prompt_cache_tree_client, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_server import PromptCacheTreeState -from lightllm.utils.net_utils import get_abstract_unix_socket_address - -SAMPLE_STRIDE = CacheAwareConfig().sample_stride - - -def _key(prompt): - return prompt[::SAMPLE_STRIDE] - - -def _remote_tree(ipc_name): - with patch.object( - prompt_cache_tree_client, - "get_pd_master_unix_socket_address", - return_value=get_abstract_unix_socket_address("prompt_cache_tree", ipc_name), - ): - return RemotePromptCacheTreeClient() - - -def test_local_client_owns_an_independent_tree(monkeypatch): - monkeypatch.setattr( - prompt_cache_tree_client, - "get_env_start_args", - lambda: Mock(httpserver_workers=1), - ) - config = CacheAwareConfig() - first = create_prompt_cache_tree_client(config) - second = create_prompt_cache_tree_client(config) - assert isinstance(first, LocalPromptCacheTreeClient) - assert isinstance(second, LocalPromptCacheTreeClient) - prompt = "shared prefix" * 200 - - async def run(): - await first.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = await first.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - missing = await second.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match.prefill_node == "prefill" - assert missing.prefill_node is None - - asyncio.run(run()) - - -def test_remote_client_uses_symmetric_sampled_operations(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - tree = _remote_tree("test-prompt-cache") - prompt = "0123456789abcdef" * 128 - result = PromptCacheMatchResult(None, 0, len(prompt), None) - tree._prefix_match_sampled = Mock(return_value=result) - tree._insert_sampled = Mock() - - async def run(): - assert await tree.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) is result - await tree.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - - asyncio.run(run()) - tree._prefix_match_sampled.assert_called_once_with(_key(prompt), len(prompt), SAMPLE_STRIDE) - tree._insert_sampled.assert_called_once_with(_key(prompt), "prefill", SAMPLE_STRIDE) - - -def test_cancelled_tree_operation_waits_for_the_client_lock_and_still_executes(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - tree = _remote_tree("test-prompt-cache-cancellation") - tree._insert_sampled = Mock() - - async def run(): - await tree.async_lock.acquire() - task = asyncio.create_task(tree.insert_sampled("key", "prefill", SAMPLE_STRIDE)) - await asyncio.sleep(0) - task.cancel() - await asyncio.sleep(0) - assert not task.done() - tree.async_lock.release() - with pytest.raises(asyncio.CancelledError): - await task - - asyncio.run(run()) - tree._insert_sampled.assert_called_once_with("key", "prefill", SAMPLE_STRIDE) - - -def test_tree_rpc_failure_terminates_the_worker(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - terminate = Mock() - monkeypatch.setattr(prompt_cache_tree_client.os, "kill", terminate) - tree = _remote_tree("test-prompt-cache-failure") - tree._insert_sampled = Mock(side_effect=ConnectionError("connection lost")) - - with pytest.raises(ConnectionError, match="connection lost"): - asyncio.run(tree.insert_sampled("key", "prefill", SAMPLE_STRIDE)) - - terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) - - -def test_tree_state_does_not_create_node_statistics(): - tree_state = PromptCacheTreeState() - stats_state = PDNodeStatsState() - prompt = "prefix" * 200 - tree_state.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = tree_state.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match[0] == "prefill" - assert not stats_state.stats - - -@pytest.fixture -def rpc_tree_cluster(): - state = PromptCacheTreeState() - ipc_name = uuid.uuid4().hex - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("prompt_cache_tree", ipc_name), - protocol_config={"allow_pickle": True}, - ) - thread = threading.Thread(target=server.start, daemon=True) - thread.start() - clients = [_remote_tree(ipc_name) for _ in range(4)] - try: - yield state, clients - finally: - for client in clients: - connection = client.connection - if connection is not None: - connection.close() - server.close() - thread.join(timeout=5) - assert not thread.is_alive() - - -def test_real_rpc_clients_share_prompt_cache(rpc_tree_cluster): - _, clients = rpc_tree_cluster - prompt = "shared prefix" * 200 - - async def run(): - await clients[0].insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = await clients[1].prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match.prefill_node == "prefill" - assert match.matched_char_count > len(prompt) // 2 - transport = clients[0].connection._channel.stream.sock - assert transport.family == socket.AF_UNIX - - asyncio.run(run()) - - -def test_concurrent_prompts_remain_separate(rpc_tree_cluster): - _, clients = rpc_tree_cluster - prompts = ["a" * 2048, "b" * 2048] - - async def run(): - await asyncio.gather( - *( - clients[index].insert_sampled(_key(prompt), f"prefill-{index}", SAMPLE_STRIDE) - for index, prompt in enumerate(prompts) - ) - ) - matches = await asyncio.gather( - *( - clients[index + 2].prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - for index, prompt in enumerate(prompts) - ) - ) - assert [match.prefill_node for match in matches] == [ - "prefill-0", - "prefill-1", - ] - - asyncio.run(run()) - - -def test_tree_and_statistics_servers_use_independent_sockets_and_state(): - ipc_name = uuid.uuid4().hex - tree_address = get_abstract_unix_socket_address("prompt_cache_tree", ipc_name) - stats_address = get_abstract_unix_socket_address("pd_node_stats", ipc_name) - tree_state, stats_state = PromptCacheTreeState(), PDNodeStatsState() - tree_server = ThreadedServer( - tree_state, - socket_path=tree_address, - protocol_config={"allow_pickle": True}, - ) - stats_server = ThreadedServer( - stats_state, - socket_path=stats_address, - protocol_config={"allow_pickle": True}, - ) - tree_thread = threading.Thread(target=tree_server.start, daemon=True) - stats_thread = threading.Thread(target=stats_server.start, daemon=True) - tree_thread.start() - stats_thread.start() - tree = _remote_tree(ipc_name) - try: - prompt = "prefix" * 200 - asyncio.run(tree.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE)) - assert not stats_state.stats - tree_server.close() - tree_thread.join(timeout=5) - assert not tree_thread.is_alive() - assert stats_thread.is_alive() - assert stats_state.get_dispatched_info(("prefill",)) == ((0, 0),) - finally: - if tree.connection is not None: - tree.connection.close() - tree_server.close() - stats_server.close() - tree_thread.join(timeout=5) - stats_thread.join(timeout=5) diff --git a/unit_tests/server/test_qwen3_coder_stream_fc.py b/unit_tests/server/test_qwen3_coder_stream_fc.py deleted file mode 100644 index bd8ff6f08c..0000000000 --- a/unit_tests/server/test_qwen3_coder_stream_fc.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Coverage for buffered Qwen3-Coder streaming keepalive deltas.""" - -import json -from types import SimpleNamespace - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_models import Function, Tool -from lightllm.server.function_call_parser import FunctionCallParser, Qwen3CoderDetector - - -def _tool() -> Tool: - return Tool( - type="function", - function=Function( - name="write_file", - description="", - parameters={ - "type": "object", - "properties": {"content": {"type": "string"}}, - }, - ), - ) - - -def test_qwen3_coder_selects_buffered_parser(): - assert FunctionCallParser.ToolCallParserEnum["qwen3_coder"] is Qwen3CoderDetector - assert "qwen3_coder_legacy" not in FunctionCallParser.ToolCallParserEnum - - args = make_argument_parser().parse_args(["--tool_call_parser", "qwen3_coder"]) - assert args.tool_call_parser == "qwen3_coder" - - -def test_qwen3_coder_emits_one_tool_delta_for_every_decode_after_the_name(): - parser = FunctionCallParser([_tool()], "qwen3_coder") - - normal_text, calls = parser.parse_stream_chunk("\n\n\n") - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name == "write_file" - assert calls[0].parameters == "" - - streamed_calls = list(calls) - for chunk in ("a", "b", "c" * 100_000): - normal_text, calls = parser.parse_stream_chunk(chunk) - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert calls[0].parameters == "" - streamed_calls.extend(calls) - - normal_text, calls = parser.parse_stream_chunk("\n\n\n") - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"content": "ab" + "c" * 100_000} - streamed_calls.extend(calls) - - assert [call.name for call in streamed_calls].count("write_file") == 1 - assert json.loads("".join(call.parameters for call in streamed_calls)) == {"content": "ab" + "c" * 100_000} - - -def test_qwen3_coder_accepts_an_undefined_tool_when_name_check_is_disabled(monkeypatch): - import lightllm.server.function_call_parser as parser_module - - monkeypatch.setattr(parser_module, "ENABLE_TOOL_NAME_CHECK", False) - parser = FunctionCallParser([_tool()], "qwen3_coder") - - _, calls = parser.parse_stream_chunk("\n\n") - assert len(calls) == 1 - assert calls[0].name == "unknown" - assert calls[0].parameters == "" - - _, calls = parser.parse_stream_chunk("x\n\n") - assert len(calls) == 1 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"value": "x"} - - -def test_qwen3_coder_rejects_an_undefined_tool_when_name_check_is_enabled(monkeypatch): - import lightllm.server.function_call_parser as parser_module - - monkeypatch.setattr(parser_module, "ENABLE_TOOL_NAME_CHECK", True) - parser = FunctionCallParser([_tool()], "qwen3_coder") - - _, calls = parser.parse_stream_chunk("\n\n") - assert calls == [] - - _, calls = parser.parse_stream_chunk("x\n\n") - assert calls == [] - - -def test_api_tool_parser_keeps_its_two_item_return_contract(): - from lightllm.server.api_openai import _process_tools_stream - - request = SimpleNamespace(tools=[_tool()]) - parser_dict = {0: FunctionCallParser(request.tools, "qwen3_coder")} - - result = _process_tools_stream( - 0, - "\n\n\n", - parser_dict, - request, - ) - - assert len(result) == 2 - normal_text, calls = result - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].name == "write_file" - - -def test_empty_arguments_are_preserved_in_the_sse_payload(): - from lightllm.server.api_models import ( - ChatCompletionStreamResponse, - ChatCompletionStreamResponseChoice, - DeltaMessage, - FunctionResponse, - ToolCall, - ) - from lightllm.server.api_openai import _serialize_sse_chunk - - chunk = ChatCompletionStreamResponse( - id="chatcmpl-test", - created=0, - model="test", - choices=[ - ChatCompletionStreamResponseChoice( - index=0, - delta=DeltaMessage( - tool_calls=[ - ToolCall( - index=0, - function=FunctionResponse(arguments=""), - ) - ] - ), - finish_reason=None, - ) - ], - ) - - payload = json.loads(_serialize_sse_chunk(chunk, ("logprobs", "token_ids", "finish_reason"))) - function_delta = payload["choices"][0]["delta"]["tool_calls"][0]["function"] - assert function_delta["arguments"] == "" diff --git a/unit_tests/server/test_req_id_generator_performance.py b/unit_tests/server/test_req_id_generator_performance.py deleted file mode 100644 index da0ab4fcaa..0000000000 --- a/unit_tests/server/test_req_id_generator_performance.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Opt-in microbenchmark for request ID generation through RPyC. - -Run with:: - - LIGHTLLM_RUN_REQ_ID_PERF=1 pytest -q -s \ - unit_tests/server/test_req_id_generator_performance.py -""" - -import os -import statistics -import threading -import time -import uuid -from concurrent.futures import ThreadPoolExecutor -from unittest.mock import patch - -import pytest -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver import req_id_generator_client -from lightllm.server.httpserver.req_id_generator_client import ReqIDGenerator -from lightllm.server.httpserver.req_id_generator_server import ( - CLIENT_REQUEST_ID_RANGE_SIZE, - MAX_BEST_OF, - StandaloneReqIDGeneratorState, -) -from lightllm.utils.net_utils import get_abstract_unix_socket_address - -RUN_PERFORMANCE_TEST = os.getenv("LIGHTLLM_RUN_REQ_ID_PERF") == "1" - - -class _LocalIDGenerator: - def __init__(self): - self.current_id = MAX_BEST_OF - self.lock = threading.Lock() - - def generate_id(self): - with self.lock: - request_id = self.current_id - self.current_id += MAX_BEST_OF - return request_id - - -def _client(ipc_name): - socket_address = get_abstract_unix_socket_address("req_id_generator", ipc_name) - with patch.object( - req_id_generator_client, - "get_req_id_generator_socket_address", - return_value=socket_address, - ): - return ReqIDGenerator() - - -def _measure_calls(generate_id, count: int): - latencies_ns = [] - request_ids = [] - start = time.perf_counter_ns() - for _ in range(count): - call_start = time.perf_counter_ns() - request_ids.append(generate_id()) - latencies_ns.append(time.perf_counter_ns() - call_start) - elapsed_ns = time.perf_counter_ns() - start - return request_ids, latencies_ns, elapsed_ns - - -@pytest.mark.skipif( - not RUN_PERFORMANCE_TEST, - reason="set LIGHTLLM_RUN_REQ_ID_PERF=1 to run the request-ID microbenchmark", -) -def test_req_id_generator_rpc_performance(): - iterations = int(os.getenv("LIGHTLLM_REQ_ID_PERF_ITERATIONS", "10000")) - client_count = 4 - assert iterations >= client_count - - ipc_name = uuid.uuid4().hex - state = StandaloneReqIDGeneratorState() - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("req_id_generator", ipc_name), - protocol_config={"allow_pickle": True}, - ) - server_thread = threading.Thread(target=server.start, daemon=True) - server_thread.start() - clients = [_client(ipc_name) for _ in range(client_count)] - try: - for _ in range(100): - clients[0].generate_id() - - local_generator = _LocalIDGenerator() - _, _, local_elapsed_ns = _measure_calls(local_generator.generate_id, iterations) - request_ids, rpc_latencies_ns, rpc_elapsed_ns = _measure_calls(clients[0].generate_id, iterations) - - calls_per_client = iterations // client_count - concurrent_call_count = calls_per_client * client_count - concurrent_start_ns = time.perf_counter_ns() - with ThreadPoolExecutor(max_workers=client_count) as executor: - futures = [executor.submit(_measure_calls, client.generate_id, calls_per_client) for client in clients] - concurrent_ids = [request_id for future in futures for request_id in future.result()[0]] - concurrent_elapsed_ns = time.perf_counter_ns() - concurrent_start_ns - - assert len(set(request_ids)) == len(request_ids) - assert len(set(concurrent_ids)) == len(concurrent_ids) - assert set(request_ids).isdisjoint(concurrent_ids) - assert all(request_id % MAX_BEST_OF == 0 for request_id in request_ids) - - local_ns_per_call = local_elapsed_ns / iterations - rpc_ns_per_call = rpc_elapsed_ns / iterations - rpc_calls_per_second = iterations * 1e9 / rpc_elapsed_ns - concurrent_calls_per_second = concurrent_call_count * 1e9 / concurrent_elapsed_ns - rpc_latency_us = sorted(rpc_latencies_ns) - rpc_p50_us = statistics.median(rpc_latency_us) / 1e3 - rpc_p95_us = rpc_latency_us[int(len(rpc_latency_us) * 0.95)] / 1e3 - print( - "request ID performance: " - f"local={local_ns_per_call:.0f} ns/call, " - f"ranged_client={rpc_calls_per_second:.0f} calls/s, " - f"client_p50={rpc_p50_us:.1f} us, client_p95={rpc_p95_us:.1f} us, " - f"client_overhead={rpc_ns_per_call / local_ns_per_call:.1f}x, " - f"ranged_{client_count}_clients={concurrent_calls_per_second:.0f} calls/s, " - f"range_size={CLIENT_REQUEST_ID_RANGE_SIZE}" - ) - - minimum_calls_per_second = float(os.getenv("LIGHTLLM_REQ_ID_MIN_CALLS_PER_SECOND", "0")) - if minimum_calls_per_second: - assert rpc_calls_per_second >= minimum_calls_per_second - finally: - for client in clients: - client.connection.close() - server.close() - server_thread.join(timeout=5) - assert not server_thread.is_alive() diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py deleted file mode 100644 index aa1c6d6273..0000000000 --- a/unit_tests/server/test_tokenizer.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Tokenizer entry points must not eagerly load unrelated model implementations.""" - -import json -import os -from pathlib import Path -import subprocess -import sys -from types import ModuleType, SimpleNamespace -from unittest.mock import Mock - -import pytest - - -ROOT = Path(__file__).resolve().parents[2] - - -@pytest.mark.parametrize("model_type", ["deepseek_v32", "deepseek_v3"]) -def test_deepseek_bytelevel_checkpoint_preserves_whitespace(tmp_path: Path, model_type: str) -> None: - from tokenizers import Tokenizer, decoders, pre_tokenizers, trainers - from tokenizers.models import BPE - from transformers import PreTrainedTokenizerFast - from lightllm.server.tokenizer import get_tokenizer - - text = "Question: What is 2 + 3?\nAnswer: #### 5" - backend = Tokenizer(BPE(unk_token="")) - backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) - backend.decoder = decoders.ByteLevel() - backend.train_from_iterator( - [text], - trainers.BpeTrainer( - vocab_size=300, - initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), - special_tokens=[""], - ), - ) - fast = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="") - fast.save_pretrained(tmp_path) - # 真实 DeepSeek checkpoint 声明 LlamaTokenizerFast，但实际编码和解码均为 ByteLevel。 - config_path = tmp_path / "tokenizer_config.json" - config = json.loads(config_path.read_text()) - config.update(tokenizer_class="LlamaTokenizerFast", legacy=True) - config_path.write_text(json.dumps(config)) - (tmp_path / "config.json").write_text( - json.dumps( - { - "model_type": model_type, - "architectures": ["DeepseekV3ForCausalLM"], - } - ) - ) - - tokenizer = get_tokenizer(str(tmp_path)) - expected_ids = backend.encode(text).ids - assert tokenizer.encode(text, add_special_tokens=False) == expected_ids - assert tokenizer.decode(expected_ids, skip_special_tokens=True) == text - - -@pytest.mark.parametrize("model_type", ["llama", "deepseek_v32", "deepseek_v4"]) -def test_tokenizer_imports_are_isolated(model_type): - result = subprocess.run( - [ - sys.executable, - "-c", - """ -import importlib.abc -import json -from pathlib import Path -import sys -import tempfile -from tokenizers import Tokenizer -from tokenizers.models import WordLevel -from transformers import AutoConfig, PreTrainedTokenizerFast - -model_type = sys.argv[1] -class RejectUnrelatedModels(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname.startswith("lightllm.models."): - allowed = {"lightllm.models.registry", "lightllm.models.draft_registry", "lightllm.models.builtin"} - if model_type == "deepseek_v32": - allowed.update({ - "lightllm.models.deepseek3_2", - "lightllm.models.deepseek3_2.tokenizer", - "lightllm.models.deepseek3_2.encoding_dsv32", - }) - if model_type == "deepseek_v4": - allowed.update({"lightllm.models.deepseek_v4", "lightllm.models.deepseek_v4.tokenizer"}) - if fullname not in allowed: - raise AssertionError("Unexpected model import: " + fullname) -sys.meta_path.insert(0, RejectUnrelatedModels()) -from lightllm.server import tokenizer as module -assert not any(name.startswith("lightllm.models.") for name in sys.modules) - -with tempfile.TemporaryDirectory() as directory: - fast = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]"))) - fast.save_pretrained(directory) - Path(directory, "config.json").write_text(json.dumps({ - "model_type": model_type, "architectures": ["LlamaForCausalLM"], - })) - tokenizer = module.get_tokenizer(directory) - assert tokenizer.encode("hello") == [0] - if model_type == "deepseek_v32": - assert AutoConfig.for_model("deepseek_v32").model_type == "deepseek_v32" - assert tokenizer.__class__.__module__ == "lightllm.models.deepseek3_2.tokenizer" - prompt = tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}]) - assert "hello" in prompt - thinking_prompt = tokenizer.apply_chat_template( - messages=[{"role": "user", "content": "hello"}], thinking=True - ) - assert "" in thinking_prompt - tool_prompt = tokenizer.apply_chat_template( - messages=[{"role": "user", "content": "hello"}], - tools=[{"name": "echo", "description": "Echo text", "parameters": {"type": "object"}}], - ) - assert "" in tool_prompt and "echo" in tool_prompt - assert "lightllm.models.deepseek3_2.model" not in sys.modules -""", - model_type, - ], - cwd=ROOT, - env={**os.environ, "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1"}, - capture_output=True, - text=True, - timeout=60, - ) - assert result.returncode == 0, result.stdout + result.stderr - - -@pytest.mark.parametrize( - "config,module_name,class_name", - [ - ( - {"model_type": "llava", "architectures": ["TarsierForConditionalGeneration"]}, - "tarsier2.model", - "Tarsier2Tokenizer", - ), - ({"model_type": "llava"}, "llava.model", "LlavaTokenizer"), - ({"model_type": "internlmxcomposer2"}, "llava.model", "LlavaTokenizer"), - ({"model_type": "qwen", "visual": {}}, "qwen_vl.model", "QWenVLTokenizer"), - ({"model_type": "qwen2_vl", "vision_config": {}}, "qwen2_vl.model", "QWen2VLTokenizer"), - ({"model_type": "qwen2_5_vl", "vision_config": {}}, "qwen2_vl.model", "QWen2VLTokenizer"), - ({"model_type": "qwen3_vl", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), - ({"model_type": "qwen3_vl_moe", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), - ({"model_type": "qwen3_5", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"model_type": "sensenova_loop", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"model_type": "qwen3_5_moe", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"thinker_config": {}}, "qwen3_omni_moe_thinker.model", "QWen3OmniTokenizer"), - ({"model_type": "internvl_chat"}, "internvl.model", "InternvlTokenizer"), - ({"model_type": "gemma3"}, "gemma3.model", "Gemma3Tokenizer"), - ({"model_type": "gemma4"}, "gemma4.tokenizer", "Gemma4Tokenizer"), - ({"model_type": "gemma4", "vision_config": {}}, "gemma4.tokenizer", "Gemma4Tokenizer"), - ({"model_type": "neo_chat"}, "neo_chat_moe.model", "NeoChatTokenizer"), - ], -) -def test_selected_multimodal_tokenizer_is_loaded(monkeypatch, config, module_name, class_name): - import transformers - from lightllm.server import tokenizer as module - - config = {"architectures": ["Example"], **config} - base_tokenizer = object() - monkeypatch.setattr(module.PretrainedConfig, "get_config_dict", lambda *a, **kw: (config, {})) - monkeypatch.setattr(module.AutoTokenizer, "from_pretrained", lambda *a, **kw: base_tokenizer) - processor = SimpleNamespace(image_processor=object()) - monkeypatch.setattr(transformers.AutoProcessor, "from_pretrained", lambda *a, **kw: processor) - monkeypatch.setattr(transformers.AutoImageProcessor, "from_pretrained", lambda *a, **kw: processor.image_processor) - vision = ModuleType("lightllm.models.qwen2_vl.vision_process") - vision.Qwen2VLImageProcessor = SimpleNamespace(from_pretrained=lambda *a, **kw: processor.image_processor) - monkeypatch.setitem(sys.modules, vision.__name__, vision) - implementation = ModuleType("lightllm.models." + module_name) - constructor = Mock() - setattr(implementation, class_name, constructor) - monkeypatch.setitem(sys.modules, implementation.__name__, implementation) - - assert module.get_tokenizer("local-checkpoint") is constructor.return_value - constructor.assert_called_once() - args, kwargs = constructor.call_args - assert (args[0] if args else kwargs["tokenizer"]) is base_tokenizer - assert (args[1] if len(args) > 1 else kwargs["model_cfg"]) is config.get("thinker_config", config) diff --git a/unit_tests/server/test_vocab_sampling_start_args.py b/unit_tests/server/test_vocab_sampling_start_args.py deleted file mode 100644 index f76dd0f7ec..0000000000 --- a/unit_tests/server/test_vocab_sampling_start_args.py +++ /dev/null @@ -1,57 +0,0 @@ -import pytest - -from lightllm.server.api_start import _launch_subprocesses, pd_master_start -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) -@pytest.mark.parametrize( - "constraint_args", - [ - {}, - {"output_constraint_mode": "xgrammar"}, - ], -) -def test_target_candidates_reject_constraints_before_model_loading(monkeypatch, run_mode, constraint_args): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - - def unexpected_model_loading(args): - pytest.fail("Incompatible candidate sampling must be rejected before loading model configuration") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", unexpected_model_loading) - args = StartArgs(run_mode=run_mode, target_vocab_topk_sampling=2, **constraint_args) - - with pytest.raises(AssertionError, match="Disable --target_vocab_topk_sampling when using output constraints"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode", "pd_master"]) -@pytest.mark.parametrize( - "sampling_args", - [ - {}, - {"target_vocab_topk_sampling": 2, "output_constraint_mode": "none"}, - {"draft_vocab_topk_sampling": 2}, - {"diverse_mode": True, "output_constraint_mode": "xgrammar"}, - {"use_reward_model": True, "output_constraint_mode": "xgrammar"}, - { - "dp": 2, - "mtp_mode": "vanilla_no_att", - "enable_prefill_microbatch_overlap": True, - "enable_decode_microbatch_overlap": True, - "output_constraint_mode": "xgrammar", - }, - ], -) -def test_compatible_sampling_options_reach_model_configuration(monkeypatch, run_mode, sampling_args): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - - def stop_at_model_configuration(args): - raise RuntimeError("model configuration reached") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", stop_at_model_configuration) - start = pd_master_start if run_mode == "pd_master" else _launch_subprocesses - - with pytest.raises(RuntimeError, match="model configuration reached"): - start(StartArgs(run_mode=run_mode, **sampling_args)) diff --git a/unit_tests/test_format_out.py b/unit_tests/test_format_out.py deleted file mode 100644 index ec42be9ef7..0000000000 --- a/unit_tests/test_format_out.py +++ /dev/null @@ -1,47 +0,0 @@ -import json -from unittest.mock import Mock - -import pytest -import xgrammar as xgr -from pydantic import BaseModel - -from format_out.impl import ChatSession, SamplingParams -from lightllm.server.core.objs.sampling_params import SamplingParams as ServerSamplingParams - - -class HelperResult(BaseModel): - answer: str - - -@pytest.mark.parametrize("mode", ["plain", "regex", "grammar", "json", "inherited_json"]) -def test_helper_payload_initializes_server_constraints(monkeypatch, mode): - payloads = [] - - def post(url, headers, data): - payload = json.loads(data) - server_params = ServerSamplingParams() - server_params.init(None, **payload["parameters"]) - assert server_params.constraint_out_params.has_constraint() == (mode != "plain") - assert all(value is not None for value in payload["parameters"].values()) - payloads.append(payload) - return Mock(status_code=200, json=lambda: {"generated_text": ['{"answer":"OK"}']}) - - monkeypatch.setattr("format_out.impl.requests.post", post) - params = SamplingParams(do_sample=False, guided_json={} if mode == "inherited_json" else None) - session = ChatSession("Answer briefly.", params, disable_log=True) - if mode == "regex": - session.generate(regex="OK") - assert payloads[0]["parameters"]["regular_constraint"] == "OK" - elif mode == "grammar": - session.generate(grammar=xgr.Grammar.from_ebnf('root ::= "OK"')) - assert payloads[0]["parameters"]["guided_grammar"] - elif mode == "json": - output = session.gen_json_object(HelperResult, prefix_regex=r"\s{0,20}") - assert HelperResult.model_validate_json(output).answer == "OK" - else: - session.generate() - if mode == "inherited_json": - assert payloads[0]["parameters"]["guided_json"] == {} - assert params.regular_constraint is None - assert params.guided_grammar is None - assert params.guided_json == ({} if mode == "inherited_json" else None) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py deleted file mode 100644 index 1f4efc4baa..0000000000 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ /dev/null @@ -1,155 +0,0 @@ -import sys -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -import torch - - -REPO_ROOT = Path(__file__).resolve().parents[1] -STATIC_BENCHMARK_DIR = REPO_ROOT / "test" / "benchmark" / "static_inference" -for path in (REPO_ROOT, STATIC_BENCHMARK_DIR): - if str(path) not in sys.path: - sys.path.append(str(path)) - -import static_benchmark -from static_benchmark import ( - BenchmarkCase, - StaticBenchmarkExecutor, - decode_profile_batch_divisor, - resolve_batch_max_prefill_cases, -) - - -class FakeKVAllocator: - def __init__(self, page_size: int): - self.page_size = page_size - self.next_index = 0 - self.alloc_sizes = [] - - def alloc(self, size: int) -> torch.Tensor: - assert size % self.page_size == 0 - self.alloc_sizes.append(size) - indexes = torch.arange(self.next_index, self.next_index + size, dtype=torch.int32) - self.next_index += size - return indexes - - -def make_executor(page_size: int = 4): - allocator = FakeKVAllocator(page_size=page_size) - req_manager = SimpleNamespace( - req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), - ) - req_manager.page_size = page_size - - model = SimpleNamespace(att_state_manager=SimpleNamespace(full_att_allocator=allocator), req_manager=req_manager) - args = SimpleNamespace(mtp_mode=None, mtp_step=0) - executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) - return executor, allocator, req_manager - - -def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): - for start in range(0, indexes.numel(), page_size): - page = indexes[start : start + page_size] - assert torch.all(page // page_size == page[0] // page_size) - - -def test_prefill_allocates_and_reuses_complete_pages_per_request(): - executor, allocator, req_manager = make_executor(page_size=4) - req_idx = torch.tensor([0, 1], dtype=torch.int32) - - model_input = executor._make_prefill_input( - token_chunk=np.arange(6, dtype=np.int64).reshape(2, 3), - req_idx=req_idx, - ready_cache_len=0, - ) - - assert allocator.alloc_sizes == [4, 4] - assert model_input.b_is_decode_req.tolist() == [False, False] - assert executor._hold_kv_len_by_req == {0: 4, 1: 4} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :4], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :4], page_size=4) - - # 两个请求都还在已有页内继续 decode，不应产生新的分配。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([4, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2], dtype=torch.int64), - max_kv_seq_len=4, - ) - assert allocator.alloc_sizes == [4, 4] - - # 跨过页边界后，每个请求分别追加一整页。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([5, 5], dtype=torch.int32), - input_ids=torch.tensor([3, 4], dtype=torch.int64), - max_kv_seq_len=5, - ) - assert allocator.alloc_sizes == [4, 4, 4, 4] - assert executor._hold_kv_len_by_req == {0: 8, 1: 8} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) - - -def test_mtp_expanded_rows_allocate_once_per_request(): - executor, allocator, req_manager = make_executor(page_size=4) - executor.args.mtp_mode = "eagle_with_att" - executor.args.mtp_step = 3 - - model_input = executor._make_decode_input( - batch_size=4, - req_idx=torch.tensor([0, 0, 1, 1], dtype=torch.int32), - mtp_index=torch.tensor([0, 1, 0, 1], dtype=torch.int32), - seq_len=torch.tensor([3, 4, 3, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2, 3, 4], dtype=torch.int64), - max_kv_seq_len=4, - extra_kv_len=2, - ) - - assert model_input.batch_size == 4 - assert allocator.alloc_sizes == [8, 8] - assert executor._hold_kv_len_by_req == {0: 8, 1: 8} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) - - -def test_profile_capacity_uses_page_aligned_per_request_cost(): - args = SimpleNamespace( - mtp_mode=None, - mtp_step=0, - page_size=4, - max_batch_size=0, - batch_max_tokens=64, - ) - prefill_case = BenchmarkCase( - name="prefill", - stage="prefill", - batch_size=8, - context_len=5, - output_len=0, - prefill_uncached_len=5, - prefill_batch_size_by_batch_max_tokens=8, - ) - - resolved = resolve_batch_max_prefill_cases(args, [prefill_case], profiled_max_total_token_num=24) - - assert resolved[0].batch_size == 3 - assert resolved[0].profiled_batch_divisor == 8 - - decode_case = BenchmarkCase( - name="decode", - stage="decode", - batch_size=1, - context_len=6, - output_len=2, - ) - assert decode_profile_batch_divisor(args, decode_case) == 20 - - args.mtp_mode = "eagle_with_att" - args.mtp_step = 3 - assert decode_profile_batch_divisor(args, decode_case) == 24 diff --git a/unit_tests/utils/test_atexit_exit_modes.py b/unit_tests/utils/test_atexit_exit_modes.py deleted file mode 100644 index 0883f0efb7..0000000000 --- a/unit_tests/utils/test_atexit_exit_modes.py +++ /dev/null @@ -1,188 +0,0 @@ -"""用独立进程验证 atexit；所有信号只发送给本测试创建的进程。""" - -import selectors -import signal -import subprocess -import sys -import textwrap - -import pytest - - -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="Exit codes and signals below target Linux") - -PROBE_SCRIPT = textwrap.dedent( - """ - import atexit - import multiprocessing as mp - import os - import resource - import signal - import sys - - - def record(marker_path, event): - with open(marker_path, "a", encoding="utf-8") as output: - output.write(event + "\\n") - - - def cleanup(marker_path, interrupt=False): - record(marker_path, "started") - if interrupt: - print("READY", flush=True) - while True: - signal.pause() - record(marker_path, "finished") - - - def worker(marker_path): - atexit.register(cleanup, marker_path) - record(marker_path, "worker_registered") - - - def main(): - # SIGABRT/SIGSEGV 实验不生成 core dump。 - resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) - mode, marker_path, signal_number = sys.argv[1:] - signal_number = int(signal_number) - signal.signal(signal.SIGINT, signal.default_int_handler) - - if mode.startswith("multiprocessing_"): - context = mp.get_context(mode.removeprefix("multiprocessing_")) - process = context.Process(target=worker, args=(marker_path,)) - process.start() - process.join(timeout=10) - if process.is_alive(): - process.kill() - process.join(timeout=10) - raise RuntimeError("multiprocessing worker failed to exit") - assert process.exitcode == 0, process.exitcode - return - - atexit.register(cleanup, marker_path, mode == "interrupt_cleanup") - - if mode in ("normal_return", "interrupt_cleanup"): - return - if mode == "sys_exit_0": - sys.exit(0) - if mode == "sys_exit_7": - sys.exit(7) - if mode == "unhandled_exception": - raise RuntimeError("intentional test exception") - if mode == "keyboard_interrupt": - raise KeyboardInterrupt - if mode == "os_exit_0": - os._exit(0) - if mode == "os_exit_7": - os._exit(7) - if mode == "abort": - os.abort() - if mode == "exec_replace": - os.execv(sys.executable, [sys.executable, "-I", "-c", "pass"]) - - if mode == "signal_sys_exit": - signal.signal(signal_number, lambda signum, frame: sys.exit(0)) - elif mode == "signal_os_exit": - signal.signal(signal_number, lambda signum, frame: os._exit(7)) - elif mode == "signal_os_default" or (mode == "signal_default" and signal_number != signal.SIGINT): - if signal_number != signal.SIGKILL: - signal.signal(signal_number, signal.SIG_DFL) - elif mode != "signal_default": - raise ValueError(mode) - - # 父进程收到 READY 后才发信号，确保回调和信号处理方式已注册。 - print("READY", flush=True) - while True: - signal.pause() - - - if __name__ == "__main__": - main() - """ -) - - -@pytest.fixture -def probe_script(tmp_path): - script_path = tmp_path / "atexit_probe.py" - script_path.write_text(PROBE_SCRIPT, encoding="utf-8") - return script_path - - -def run_probe(probe_script, tmp_path, mode, signum=0): - marker_path = tmp_path / "callback_events.txt" - with subprocess.Popen( - [sys.executable, "-I", str(probe_script), mode, str(marker_path), str(int(signum))], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - start_new_session=True, - ) as process: - try: - if signum: - with selectors.DefaultSelector() as selector: - selector.register(process.stdout, selectors.EVENT_READ) - assert selector.select(timeout=10), "Probe did not become ready" - assert process.stdout.readline().strip() == "READY", "Probe exited before registering its callback" - process.send_signal(signum) - stdout, stderr = process.communicate(timeout=30) - finally: - if process.poll() is None: - process.kill() - process.communicate(timeout=10) - - events = marker_path.read_text(encoding="utf-8").splitlines() if marker_path.exists() else [] - print(f"{mode} signal={int(signum)} returncode={process.returncode} events={events}") - return process.returncode, events, stdout + stderr - - -@pytest.mark.parametrize( - "mode,signum,expected_returncode,callback_completed", - [ - pytest.param("normal_return", 0, 0, True, id="normal-return"), - pytest.param("sys_exit_0", 0, 0, True, id="sys-exit-0"), - pytest.param("sys_exit_7", 0, 7, True, id="sys-exit-nonzero"), - pytest.param("unhandled_exception", 0, 1, True, id="unhandled-exception"), - pytest.param("keyboard_interrupt", 0, -signal.SIGINT, True, id="unhandled-keyboard-interrupt"), - pytest.param("signal_default", signal.SIGINT, -signal.SIGINT, True, id="sigint-python-default"), - pytest.param("signal_os_default", signal.SIGINT, -signal.SIGINT, False, id="sigint-os-default"), - pytest.param("signal_default", signal.SIGTERM, -signal.SIGTERM, False, id="sigterm-default"), - pytest.param("signal_default", signal.SIGHUP, -signal.SIGHUP, False, id="sighup-default"), - pytest.param("signal_default", signal.SIGQUIT, -signal.SIGQUIT, False, id="sigquit-default"), - pytest.param("signal_default", signal.SIGSEGV, -signal.SIGSEGV, False, id="sigsegv-default"), - pytest.param("signal_default", signal.SIGKILL, -signal.SIGKILL, False, id="sigkill"), - pytest.param("signal_sys_exit", signal.SIGINT, 0, True, id="sigint-handler-sys-exit"), - pytest.param("signal_sys_exit", signal.SIGTERM, 0, True, id="sigterm-handler-sys-exit"), - pytest.param("signal_sys_exit", signal.SIGHUP, 0, True, id="sighup-handler-sys-exit"), - pytest.param("signal_os_exit", signal.SIGTERM, 7, False, id="sigterm-handler-os-exit"), - pytest.param("os_exit_0", 0, 0, False, id="os-exit-0"), - pytest.param("os_exit_7", 0, 7, False, id="os-exit-nonzero"), - pytest.param("abort", 0, -signal.SIGABRT, False, id="abort"), - pytest.param("exec_replace", 0, 0, False, id="exec-replaces-interpreter"), - ], -) -def test_atexit_on_process_exit(probe_script, tmp_path, mode, signum, expected_returncode, callback_completed): - returncode, events, output = run_probe(probe_script, tmp_path, mode, signum) - - assert returncode == expected_returncode, output - assert events == (["started", "finished"] if callback_completed else []), output - - -@pytest.mark.skipif( - sys.version_info[:2] != (3, 10), reason="This experiment records Python 3.10 multiprocessing behavior" -) -@pytest.mark.parametrize("start_method,callback_completed", [("spawn", True), ("fork", False)]) -def test_atexit_in_multiprocessing_worker(probe_script, tmp_path, start_method, callback_completed): - returncode, events, output = run_probe(probe_script, tmp_path, f"multiprocessing_{start_method}") - - assert returncode == 0, output - # 本机 Python 3.10：spawn_main() 使用 sys.exit()，fork 的 _launch() 使用 os._exit()。 - expected_events = ["worker_registered"] + (["started", "finished"] if callback_completed else []) - assert events == expected_events, output - - -@pytest.mark.parametrize("signum", [signal.SIGINT, signal.SIGKILL]) -def test_signal_can_interrupt_atexit_callback(probe_script, tmp_path, signum): - _returncode, events, output = run_probe(probe_script, tmp_path, "interrupt_cleanup", signum) - - assert events == ["started"], output diff --git a/unit_tests/utils/test_custom_kernel_utils.py b/unit_tests/utils/test_custom_kernel_utils.py deleted file mode 100644 index 4f9d63388d..0000000000 --- a/unit_tests/utils/test_custom_kernel_utils.py +++ /dev/null @@ -1,35 +0,0 @@ -import torch -import time -import pytest -from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch, torch_cat_3 - - -def test_torch_cat(): - a = torch.tensor([[[1, 2], [3, 4]]], device="cuda") - b = torch.tensor([[[5, 6], [7, 8]]], device="cuda") - c = torch_cat_3([a, b], dim=0) - torch.equal(torch.cat((a, b), dim=0), c) - - d = torch_cat_3([a, b], dim=1) - torch.equal(torch.cat((a, b), dim=1), d) - - e = torch_cat_3([a, b], dim=-1) - torch.equal(torch.cat((a, b), dim=-1), e) - - empty = torch.empty((0, 2), device="cuda") - torch_cat_3([a, empty, b], dim=0) - return - - -def test_pad2dim_tensor_to_new_batch(): - input_tensor = torch.tensor([[1.0, 2.0]]) - padded_tensor = pad2dim_tensor_to_new_batch(input=input_tensor, new_batch_size=3) - assert torch.equal(padded_tensor, torch.tensor([[1.0, 2.0], [1.0, 2.0], [1.0, 2.0]])) - - empty_input = torch.empty((0, 2)) - padded_empty_input = pad2dim_tensor_to_new_batch(input=empty_input, new_batch_size=2) - assert torch.equal(padded_empty_input, torch.zeros((2, 2))) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py index 75f48cb174..18d0072a7a 100644 --- a/unit_tests/utils/test_envs_utils.py +++ b/unit_tests/utils/test_envs_utils.py @@ -5,120 +5,9 @@ get_deepep_num_max_dispatch_tokens_per_rank, get_deepep_num_max_dispatch_tokens_per_rank_decode, get_deepep_num_max_dispatch_tokens_per_rank_prefill, - is_stop_matching_in_reasoning_enabled, - get_pd_cache_high_priority_max_age_seconds, - get_pd_cache_high_priority_min_prompt_tokens, - get_pd_node_busy_retry_timeout_seconds, - get_pd_node_continuation_resource_wait_timeout_seconds, - get_pd_node_resource_wait_timeout_seconds, ) -def test_stop_in_reasoning_defaults_to_false(monkeypatch): - monkeypatch.delenv("LIGHTLLM_STOP_IN_REASONING", raising=False) - is_stop_matching_in_reasoning_enabled.cache_clear() - assert is_stop_matching_in_reasoning_enabled() is False - is_stop_matching_in_reasoning_enabled.cache_clear() - - -def test_stop_in_reasoning_reads_environment(monkeypatch): - for value, expected in [("1", True), ("true", True), ("ON", True), ("0", False), ("false", False)]: - monkeypatch.setenv("LIGHTLLM_STOP_IN_REASONING", value) - is_stop_matching_in_reasoning_enabled.cache_clear() - assert is_stop_matching_in_reasoning_enabled() is expected - is_stop_matching_in_reasoning_enabled.cache_clear() - - -def test_pd_cache_high_priority_max_age_defaults_to_180_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS", raising=False) - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - assert get_pd_cache_high_priority_max_age_seconds() == 180 - - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - -def test_pd_cache_high_priority_max_age_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS", "30") - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - assert get_pd_cache_high_priority_max_age_seconds() == 30 - - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - -def test_pd_cache_high_priority_min_prompt_tokens_defaults_to_2048(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS", raising=False) - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - assert get_pd_cache_high_priority_min_prompt_tokens() == 2048 - - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - -def test_pd_cache_high_priority_min_prompt_tokens_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS", "2048") - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - assert get_pd_cache_high_priority_min_prompt_tokens() == 2048 - - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - -def test_pd_node_resource_wait_timeout_defaults_to_20_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_RESOURCE_WAIT_TIMEOUT_SECONDS", raising=False) - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_resource_wait_timeout_seconds() == 20 - - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_resource_wait_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_RESOURCE_WAIT_TIMEOUT_SECONDS", "30") - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_resource_wait_timeout_seconds() == 30 - - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_continuation_resource_wait_timeout_defaults_to_60_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_CONTINUATION_RESOURCE_WAIT_TIMEOUT_SECONDS", raising=False) - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_continuation_resource_wait_timeout_seconds() == 60 - - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_continuation_resource_wait_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_CONTINUATION_RESOURCE_WAIT_TIMEOUT_SECONDS", "90") - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_continuation_resource_wait_timeout_seconds() == 90 - - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_busy_retry_timeout_defaults_to_zero(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", raising=False) - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - assert get_pd_node_busy_retry_timeout_seconds() == 0 - - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - -def test_pd_node_busy_retry_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", "60") - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - assert get_pd_node_busy_retry_timeout_seconds() == 60 - - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - def test_deepep_decode_buffer_capacity_covers_mtp_rows_and_safety_margin(monkeypatch): monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) monkeypatch.setattr( diff --git a/unit_tests/utils/test_log_utils.py b/unit_tests/utils/test_log_utils.py deleted file mode 100644 index 3534656835..0000000000 --- a/unit_tests/utils/test_log_utils.py +++ /dev/null @@ -1,42 +0,0 @@ -import logging - -import pytest - -from lightllm.utils import log_utils - - -def _format_message(message: str) -> str: - formatter = log_utils.NewLineFormatter("%(node_role)s%(message)s") - record = logging.LogRecord("test", logging.INFO, __file__, 1, message, (), None) - return formatter.format(record) - - -@pytest.mark.parametrize( - ("run_mode", "expected_marker"), - [ - ("prefill", "P"), - ("decode", "D"), - ("normal", ""), - ("pd_master", ""), - ], -) -def test_log_node_role_marker_is_limited_to_pd_workers(monkeypatch, run_mode, expected_marker): - monkeypatch.delenv("LIGHTLLM_LOG_NODE_ROLE", raising=False) - monkeypatch.setattr(log_utils, "_LOG_NODE_ROLE", "") - - log_utils.set_log_node_role(run_mode) - - formatted_marker = f"[{expected_marker}] " if expected_marker else "" - assert _format_message("ready") == f"{formatted_marker}ready" - if expected_marker: - assert log_utils.os.environ["LIGHTLLM_LOG_NODE_ROLE"] == expected_marker - else: - assert "LIGHTLLM_LOG_NODE_ROLE" not in log_utils.os.environ - - -def test_log_node_role_marker_is_applied_to_every_line(monkeypatch): - monkeypatch.delenv("LIGHTLLM_LOG_NODE_ROLE", raising=False) - monkeypatch.setattr(log_utils, "_LOG_NODE_ROLE", "") - log_utils.set_log_node_role("prefill") - - assert _format_message("first\nsecond") == "[P] first\r\n[P] second" diff --git a/unit_tests/utils/test_multimodal_chunk_utils.py b/unit_tests/utils/test_multimodal_chunk_utils.py deleted file mode 100644 index d553a62994..0000000000 --- a/unit_tests/utils/test_multimodal_chunk_utils.py +++ /dev/null @@ -1,68 +0,0 @@ -import random -import unittest - -from lightllm.utils.multimodal_chunk_utils import ( - adjust_mm_chunk_end, - get_mm_input_spans, -) - - -class MultimodalChunkTests(unittest.TestCase): - def test_embedding_position_locates_a_qwen_image(self): - params = {"images": [{"start_idx": 3, "token_id": 100, "token_num": 4}], "audios": []} - input_ids = [7, 8, 9, 100, 101, 102, 103] - self.assertEqual(get_mm_input_spans(input_ids, params), [(3, 7)]) - - def test_adjacent_images_audio_and_reused_embeddings_remain_separate(self): - params = { - "images": [ - {"start_idx": 1, "token_id": 100, "token_num": 3}, - {"start_idx": 6, "token_id": 100, "token_num": 3}, - ], - "audios": [{"start_idx": 4, "token_id": 200, "token_num": 2}], - } - input_ids = [7, 100, 101, 102, 200, 201, 100, 101, 102] - self.assertCountEqual(get_mm_input_spans(input_ids, params), [(1, 4), (4, 6), (6, 9)]) - - def test_text_before_and_after_image_can_still_be_chunked(self): - spans = [(1800, 2100)] - chunks = [] - start = 0 - while start < 4000: - end = adjust_mm_chunk_end(start, min(start + 2048, 4000), spans) - chunks.append((start, end)) - start = end - self.assertEqual(chunks, [(0, 1800), (1800, 3848), (3848, 4000)]) - - def test_small_chunk_target_can_yield_to_a_complete_item(self): - spans = [(100, 800)] - self.assertEqual(adjust_mm_chunk_end(0, 256, spans), 100) - self.assertEqual(adjust_mm_chunk_end(100, 356, spans), 800) - - def test_checkpoint_boundary_can_be_crossed_without_splitting_an_item(self): - spans = [(1900, 2200)] - self.assertEqual(adjust_mm_chunk_end(0, 2048, spans), 1900) - self.assertEqual(adjust_mm_chunk_end(1900, 2048, spans), 2200) - - def test_random_chunk_partitions_always_make_progress_within_budget(self): - rng = random.Random(0) - for _ in range(100): - spans = [] - position = 0 - for _ in range(8): - start = position + rng.randrange(0, 50) - position = start + rng.randrange(1, 129) - spans.append((start, position)) - total = position + 50 - target = rng.randrange(1, 129) - position = 0 - while position < total: - end = adjust_mm_chunk_end(position, min(position + target, total), spans) - self.assertLess(position, end) - self.assertLessEqual(end - position, 128) - self.assertFalse(any(start < end < stop for start, stop in spans)) - position = end - - -if __name__ == "__main__": - unittest.main() diff --git a/unit_tests/utils/test_startup_process_lifecycle.py b/unit_tests/utils/test_startup_process_lifecycle.py deleted file mode 100644 index 9b9d38fb60..0000000000 --- a/unit_tests/utils/test_startup_process_lifecycle.py +++ /dev/null @@ -1,135 +0,0 @@ -import ctypes -import importlib -import multiprocessing as mp -import os -import signal -import sys -import time -import traceback -from types import SimpleNamespace - -import psutil -import pytest - -from lightllm.utils import process_check, start_utils - - -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="requires Linux subreaper and fork") - - -def _kill_and_reap_children(): - for child in psutil.Process().children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - while True: - try: - os.waitpid(-1, 0) - except ChildProcessError: - break - - -def _run_metrics_lifecycle(failure_mode, result_writer): - """Own all launcher descendants, including orphans, in an isolated process.""" - result = None - try: - libc = ctypes.CDLL(None, use_errno=True) - if libc.prctl(36, 1, 0, 0, 0) != 0: # PR_SET_CHILD_SUBREAPER - raise OSError(ctypes.get_errno(), "could not become a child subreaper") - - mp.set_start_method("fork", force=True) - ctx = mp.get_context("fork") - events, event_writer = ctx.Pipe(duplex=False) - metrics = importlib.import_module("lightllm.server.metrics.manager") - envs_utils = importlib.import_module("lightllm.utils.envs_utils") - - # Keep the real metrics entry point and RPyC server, without monitor/shared-memory setup. - metrics.Monitor = lambda args: SimpleNamespace() - metrics.get_unique_server_name = lambda: "metrics-lifecycle-regression" - metrics.get_shm_port_args = lambda: SimpleNamespace(metric_port=0) - envs_utils.get_env_start_args = lambda: SimpleNamespace(enable_mps=False) - # Speed up only the monitor's polling, without changing the shared time module. - process_check.time = SimpleNamespace(sleep=lambda _: time.sleep(0.01)) - - def failing_submodule(pipe_writer): - event_writer.send(("failing", os.getpid())) - pipe_writer.send("injected initialization failure") - while True: - time.sleep(100) - - def launcher_main(): - manager = start_utils.SubmoduleManager() - metric_processes = manager.start_submodule_processes( - start_funcs=[metrics.start_metric_manager], - start_args=[(SimpleNamespace(metric_gateway=None, push_interval=1),)], - ) - event_writer.send(("metrics", metric_processes[0].pid)) - if failure_mode == "parent-exit": - # Bypass multiprocessing joins to exercise the metrics parent monitor itself. - os._exit(2) - manager.start_submodule_processes(start_funcs=[failing_submodule], start_args=[()]) - - launcher = ctx.Process(target=launcher_main) - launcher.start() - reported = {} - expected = {"metrics", "failing"} if failure_mode == "startup-failure" else {"metrics"} - deadline = time.monotonic() + 8 - while set(reported) != expected: - assert time.monotonic() < deadline, "launcher did not start the expected submodules" - if events.poll(0.05): - name, pid = events.recv() - reported[name] = pid - - launcher.join(timeout=8) - assert not launcher.is_alive(), "launcher hung while joining an earlier metrics process" - assert launcher.exitcode == (1 if failure_mode == "startup-failure" else 2) - - if failure_mode == "parent-exit": - # The real monitor must kill metrics after its launcher's immediate exit. - deadline = time.monotonic() + 5 - while True: - reaped_pid, status = os.waitpid(reported["metrics"], os.WNOHANG) - if reaped_pid: - assert os.WIFSIGNALED(status) - assert os.WTERMSIG(status) == signal.SIGKILL - break - assert time.monotonic() < deadline, "metrics did not detect its dead launcher" - time.sleep(0.01) - - assert all(not psutil.pid_exists(pid) for pid in reported.values()), "submodule leaked after launcher exit" - result = "ok" - except BaseException: - result = traceback.format_exc() - finally: - _kill_and_reap_children() - result_writer.send(result) - - -@pytest.mark.parametrize("failure_mode", ["startup-failure", "parent-exit"]) -def test_metrics_cleanup_on_launcher_failure(failure_mode): - # Import before forking so framework imports happen only once per pytest run. - importlib.import_module("lightllm.server.metrics.manager") - ctx = mp.get_context("fork") - result_reader, result_writer = ctx.Pipe(duplex=False) - controller = ctx.Process(target=_run_metrics_lifecycle, args=(failure_mode, result_writer)) - controller.start() - try: - controller.join(timeout=20) - assert not controller.is_alive(), "metrics lifecycle regression controller timed out" - assert controller.exitcode == 0 - assert result_reader.poll(timeout=1), "metrics lifecycle controller returned no result" - result = result_reader.recv() - assert result == "ok", result - finally: - if controller.is_alive(): - for child in psutil.Process(controller.pid).children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - controller.kill() - controller.join(timeout=5) - controller.close() - result_reader.close() - result_writer.close()