sglang 0.3.3.post1__py3-none-any.whl → 0.3.4__py3-none-any.whl

sglang/bench_latency.py

@@ -232,17 +232,18 @@ def extend(reqs, model_runner):
     model_worker_batch = batch.get_model_worker_batch()
     forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
     logits_output = model_runner.forward(forward_batch)
-    next_token_ids = model_runner.sample(logits_output, forward_batch).tolist()
+    next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits, batch
 
 
 @torch.inference_mode()
 def decode(input_token_ids, batch, model_runner):
-    batch.prepare_for_decode(input_token_ids)
+    batch.output_ids = input_token_ids
+    batch.prepare_for_decode()
     model_worker_batch = batch.get_model_worker_batch()
     forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
     logits_output = model_runner.forward(forward_batch)
-    next_token_ids = model_runner.sample(logits_output, forward_batch).tolist()
+    next_token_ids = model_runner.sample(logits_output, forward_batch)
     return next_token_ids, logits_output.next_token_logits
 
 
@@ -252,6 +253,7 @@ def correctness_test(
     bench_args,
     tp_rank,
 ):
+    configure_logger(server_args, prefix=f" TP{tp_rank}")
     rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
 
     # Load the model
@@ -279,8 +281,9 @@ def correctness_test(
     output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))]
     for _ in range(bench_args.output_len[0] - 1):
         next_token_ids, _ = decode(next_token_ids, batch, model_runner)
+        next_token_ids_list = next_token_ids.tolist()
         for i in range(len(reqs)):
-            output_ids[i].append(next_token_ids[i])
+            output_ids[i].append(next_token_ids_list[i])
 
     # Print
     for i in range(len(reqs)):
@@ -288,8 +291,15 @@ def correctness_test(
         rank_print(tokenizer.decode(output_ids[i]), "\n")
 
 
+def synchronize(device):
+    if device == "cuda":
+        torch.cuda.synchronize()
+    elif device == "xpu":
+        torch.xpu.synchronize()
+
+
 def latency_test_run_once(
-    run_name, model_runner, rank_print, reqs, batch_size, input_len, output_len
+    run_name, model_runner, rank_print, reqs, batch_size, input_len, output_len, device
 ):
     max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len)
     if batch_size > max_batch_size:
@@ -312,10 +322,10 @@ def latency_test_run_once(
     tot_latency = 0
 
     # Prefill
-    torch.cuda.synchronize()
+    synchronize(device)
     tic = time.time()
     next_token_ids, _, batch = extend(reqs, model_runner)
-    torch.cuda.synchronize()
+    synchronize(device)
     prefill_latency = time.time() - tic
     tot_latency += prefill_latency
     throughput = input_len * batch_size / prefill_latency
@@ -328,10 +338,10 @@ def latency_test_run_once(
     # Decode
     decode_latencies = []
     for i in range(output_len - 1):
-        torch.cuda.synchronize()
+        synchronize(device)
         tic = time.time()
         next_token_ids, _ = decode(next_token_ids, batch, model_runner)
-        torch.cuda.synchronize()
+        synchronize(device)
         latency = time.time() - tic
         tot_latency += latency
         throughput = batch_size / latency
@@ -387,6 +397,7 @@ def latency_test(
         bench_args.batch_size[0],
         bench_args.input_len[0],
         8,  # shorter decoding to speed up the warmup
+        server_args.device,
     )
     rank_print("Benchmark ...")
 
@@ -397,7 +408,14 @@ def latency_test(
     ):
         reqs = prepare_synthetic_inputs_for_latency_test(bs, il)
         ret = latency_test_run_once(
-            bench_args.run_name, model_runner, rank_print, reqs, bs, il, ol
+            bench_args.run_name,
+            model_runner,
+            rank_print,
+            reqs,
+            bs,
+            il,
+            ol,
+            server_args.device,
         )
         if ret is not None:
             result_list.append(ret)

sglang/bench_server_latency.py

@@ -6,6 +6,8 @@ It accepts arguments similar to those of launch_server.py.
 Usage:
 
 python3 -m sglang.bench_server_latency --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8
+
+python3 -m sglang.bench_server_latency --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8
 """
 
 import argparse
@@ -32,6 +34,8 @@ class BenchArgs:
     input_len: Tuple[int] = (1024,)
     output_len: Tuple[int] = (16,)
     result_filename: str = "result.jsonl"
+    base_url: str = ""
+    skip_warmup: bool = False
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
@@ -48,6 +52,8 @@ class BenchArgs:
         parser.add_argument(
             "--result-filename", type=str, default=BenchArgs.result_filename
         )
+        parser.add_argument("--base-url", type=str, default=BenchArgs.base_url)
+        parser.add_argument("--skip-warmup", action="store_true")
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
@@ -139,17 +145,21 @@ def run_one_case(
 
 
 def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
-    proc, base_url = launch_server_process(server_args)
+    if bench_args.base_url:
+        proc, base_url = None, bench_args.base_url
+    else:
+        proc, base_url = launch_server_process(server_args)
 
     # warmup
-    run_one_case(
-        base_url,
-        batch_size=16,
-        input_len=1024,
-        output_len=16,
-        run_name="",
-        result_filename="",
-    )
+    if not bench_args.skip_warmup:
+        run_one_case(
+            base_url,
+            batch_size=16,
+            input_len=1024,
+            output_len=16,
+            run_name="",
+            result_filename="",
+        )
 
     # benchmark
     try:
@@ -165,7 +175,8 @@ def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs):
                 bench_args.result_filename,
             )
     finally:
-        kill_child_process(proc.pid)
+        if proc:
+            kill_child_process(proc.pid)
 
     print(f"\nResults are saved to {bench_args.result_filename}")
 

sglang/bench_serving.py

@@ -222,6 +222,85 @@ async def async_request_openai_completions(
     return output
 
 
+async def async_request_sglang_generate(
+    request_func_input: RequestFuncInput,
+    pbar: Optional[tqdm] = None,
+) -> RequestFuncOutput:
+    api_url = request_func_input.api_url
+    prompt = request_func_input.prompt
+
+    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
+        payload = {
+            "text": prompt,
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": request_func_input.output_len,
+                "ignore_eos": not args.disable_ignore_eos,
+            },
+            "stream": not args.disable_stream,
+            **request_func_input.extra_request_body,
+        }
+        headers = {}
+
+        output = RequestFuncOutput()
+        output.prompt_len = request_func_input.prompt_len
+
+        generated_text = ""
+        ttft = 0.0
+        st = time.perf_counter()
+        most_recent_timestamp = st
+        try:
+            async with session.post(
+                url=api_url, json=payload, headers=headers
+            ) as response:
+                if response.status == 200:
+                    async for chunk_bytes in response.content:
+                        chunk_bytes = chunk_bytes.strip()
+                        if not chunk_bytes:
+                            continue
+                        # print(chunk_bytes)
+
+                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
+                        latency = time.perf_counter() - st
+                        if chunk == "[DONE]":
+                            pass
+                        else:
+                            data = json.loads(chunk)
+
+                            # NOTE: Some completion API might have a last
+                            # usage summary response without a token so we
+                            # want to check a token was generated
+                            if data["text"]:
+                                timestamp = time.perf_counter()
+                                # First token
+                                if ttft == 0.0:
+                                    ttft = time.perf_counter() - st
+                                    output.ttft = ttft
+
+                                # Decoding phase
+                                else:
+                                    output.itl.append(timestamp - most_recent_timestamp)
+
+                                most_recent_timestamp = timestamp
+                                generated_text = data["text"]
+
+                    output.generated_text = generated_text
+                    output.success = True
+                    output.latency = latency
+                    output.output_len = request_func_input.output_len
+                else:
+                    output.error = response.reason or ""
+                    output.success = False
+        except Exception:
+            output.success = False
+            exc_info = sys.exc_info()
+            output.error = "".join(traceback.format_exception(*exc_info))
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 async def async_request_gserver(
     request_func_input: RequestFuncInput,
     pbar: Optional[tqdm] = None,
@@ -264,7 +343,9 @@ def get_tokenizer(
 
 
 ASYNC_REQUEST_FUNCS = {
-    "sglang": async_request_openai_completions,
+    "sglang": async_request_sglang_generate,
+    "sglang-native": async_request_sglang_generate,
+    "sglang-oai": async_request_openai_completions,
     "vllm": async_request_openai_completions,
     "lmdeploy": async_request_openai_completions,
     "trt": async_request_trt_llm,
@@ -387,6 +468,8 @@ def sample_sharegpt_requests(
             continue
         filtered_dataset.append((prompt, prompt_len, output_len))
 
+    print(f"#Input tokens: {np.sum([x[1] for x in filtered_dataset])}")
+    print(f"#Output tokens: {np.sum([x[2] for x in filtered_dataset])}")
     return filtered_dataset
 
 
@@ -587,6 +670,8 @@ async def benchmark(
     else:
         print("Initial test run completed. Starting main benchmark run...")
 
+    time.sleep(1.5)
+
     pbar = None if disable_tqdm else tqdm(total=len(input_requests))
 
     benchmark_start_time = time.perf_counter()
@@ -782,24 +867,33 @@ def run_benchmark(args_: argparse.Namespace):
     if args.port is None:
         args.port = {
             "sglang": 30000,
+            "sglang-native": 30000,
+            "sglang-oai": 30000,
             "lmdeploy": 23333,
             "vllm": 8000,
             "trt": 8000,
             "gserver": 9988,
         }.get(args.backend, 30000)
 
-    api_url = (
-        f"{args.base_url}/v1/completions"
-        if args.base_url
-        else f"http://{args.host}:{args.port}/v1/completions"
-    )
     model_url = (
         f"{args.base_url}/v1/models"
         if args.base_url
         else f"http://{args.host}:{args.port}/v1/models"
     )
 
-    if args.backend == "trt":
+    if args.backend in ["sglang", "sglang-native"]:
+        api_url = (
+            f"{args.base_url}/generate"
+            if args.base_url
+            else f"http://{args.host}:{args.port}/generate"
+        )
+    elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]:
+        api_url = (
+            f"{args.base_url}/v1/completions"
+            if args.base_url
+            else f"http://{args.host}:{args.port}/v1/completions"
+        )
+    elif args.backend == "trt":
         api_url = (
             f"{args.base_url}/v2/models/ensemble/generate_stream"
             if args.base_url

sglang/global_config.py

@@ -19,7 +19,6 @@ class GlobalConfig:
         self.new_token_ratio_decay = 0.001
 
         # Runtime constants: others
-        self.num_continue_decode_steps = 10
         self.retract_decode_steps = 20
         self.flashinfer_workspace_size = os.environ.get(
             "FLASHINFER_WORKSPACE_SIZE", 384 * 1024 * 1024

sglang/srt/layers/attention/__init__.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 
+import torch
 from torch import nn
 
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -18,13 +19,13 @@ class AttentionBackend(ABC):
         raise NotImplementedError()
 
     def init_forward_metadata_capture_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
+        self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
     ):
         """Init the metadata for a forward pass for capturing a cuda graph."""
         raise NotImplementedError()
 
     def init_forward_metadata_replay_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
+        self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
     ):
         """Init the metadata for a forward pass for replying a cuda graph."""
         raise NotImplementedError()
@@ -33,17 +34,38 @@ class AttentionBackend(ABC):
         """Get the fill value for padded seq lens. Typically, it is 0 or 1."""
         raise NotImplementedError()
 
-    def forward(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: nn.Module,
+        forward_batch: ForwardBatch,
+    ):
         """Run forward on an attention layer."""
         if forward_batch.forward_mode.is_decode():
             return self.forward_decode(q, k, v, layer, forward_batch)
         else:
             return self.forward_extend(q, k, v, layer, forward_batch)
 
-    def forward_decode(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: nn.Module,
+        forward_batch: ForwardBatch,
+    ):
         """Run a forward for decode."""
         raise NotImplementedError()
 
-    def forward_extend(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: nn.Module,
+        forward_batch: ForwardBatch,
+    ):
         """Run a forward for extend."""
         raise NotImplementedError()

sglang/srt/layers/attention/double_sparsity_backend.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+from sglang.srt.layers.attention import AttentionBackend
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class DoubleSparseAttnBackend(AttentionBackend):
+    def __init__(self, model_runner: ModelRunner):
+        # Lazy import to avoid the initialization of cuda context
+        from sglang.srt.layers.attention.triton_ops.double_sparsity_attention import (
+            flash_decode_attention_fwd,
+            flash_decode_sparse_attention_fwd,
+        )
+        from sglang.srt.layers.attention.triton_ops.extend_attention import (
+            extend_attention_fwd,
+        )
+
+        super().__init__()
+
+        self.decode_attention_fwd = flash_decode_attention_fwd
+        self.decode_sparse_attention_fwd = flash_decode_sparse_attention_fwd
+        self.extend_attention_fwd = extend_attention_fwd
+        self.num_head = model_runner.model_config.num_attention_heads
+        self.head_dim = model_runner.model_config.hidden_size // self.num_head
+        self.heavy_token_num = model_runner.server_args.ds_heavy_token_num
+
+        self.sorted_channels = model_runner.sorted_channels
+        self.sparse_decode_thresold = (
+            model_runner.server_args.ds_sparse_decode_threshold
+        )
+        self.att_out_approx: torch.Tensor = None
+        self.mid_out: torch.Tensor = None
+        self.mid_o_logexpsum: torch.Tensor = None
+
+        # TODO: Change the hard-coded block_seq_num
+        self.BLOCK_SEQ = 128
+
+        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
+            self.reduce_dtype = torch.float32
+        else:
+            self.reduce_dtype = torch.float16
+
+        self.forward_metadata = None
+
+        self.cuda_graph_max_seq_len = model_runner.model_config.context_len
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init auxiliary variables for triton attention backend."""
+
+        if forward_batch.forward_mode.is_decode():
+            start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
+            start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)
+
+            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
+            attn_logits = torch.empty(
+                (self.num_head, total_num_tokens),
+                dtype=self.reduce_dtype,
+                device="cuda",
+            )
+
+            max_seq_len = torch.max(forward_batch.seq_lens).item()
+            min_seq_len = torch.min(forward_batch.seq_lens).item()
+            max_extend_len = None
+            # NOTE: Align sequence order with req_to_token order
+            ds_req_to_token = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices
+            ]
+
+            bsz = forward_batch.seq_lens.shape[0]
+
+            att_out_approx = torch.empty(
+                [self.num_head, bsz, max_seq_len],
+                dtype=self.reduce_dtype,
+                device="cuda",
+            )
+
+            block_seq_num = (
+                self.heavy_token_num + self.BLOCK_SEQ - 1
+            ) // self.BLOCK_SEQ
+
+            mid_out = torch.empty(
+                [bsz, self.num_head, block_seq_num, self.head_dim],
+                dtype=torch.float32,
+                device="cuda",
+            )
+            mid_o_logexpsum = torch.empty(
+                [bsz, self.num_head, block_seq_num], dtype=torch.float32, device="cuda"
+            )
+            self.att_out_approx = att_out_approx
+            self.mid_out = mid_out
+            self.mid_o_logexpsum = mid_o_logexpsum
+
+        else:
+            start_loc = attn_logits = max_seq_len = min_seq_len = None
+            prefix_lens = forward_batch.extend_prefix_lens
+            max_extend_len = torch.max(forward_batch.seq_lens - prefix_lens).item()
+            ds_req_to_token = None
+
+        self.forward_metadata = (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        )
+
+    def init_cuda_graph_state(self, max_bs: int):
+        # TODO(Andy): Support CUDA graph for double sparse attention
+        raise ValueError(
+            "Double sparse attention does not support CUDA graph for now. Please --disable-cuda-graph"
+        )
+        self.cuda_graph_max_total_num_tokens = max_bs * self.cuda_graph_max_seq_len
+
+        self.cuda_graph_start_loc = torch.zeros(
+            (max_bs,), dtype=torch.int32, device="cuda"
+        )
+        self.cuda_graph_attn_logits = torch.empty(
+            (
+                self.num_head,
+                self.cuda_graph_max_total_num_tokens,
+            ),
+            dtype=self.reduce_dtype,
+            device="cuda",
+        )
+
+    def init_forward_metadata_capture_cuda_graph(
+        self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
+    ):
+        self.forward_metadata = (
+            self.cuda_graph_start_loc,
+            self.cuda_graph_attn_logits,
+            self.cuda_graph_max_seq_len,
+            None,
+        )
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor
+    ):
+        self.cuda_graph_start_loc.zero_()
+        self.cuda_graph_start_loc[1:bs] = torch.cumsum(seq_lens[: bs - 1], dim=0)
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 1
+
+    def forward_extend(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        k_label = torch.gather(
+            k,
+            2,
+            self.sorted_channels[layer.layer_id]
+            .unsqueeze(0)
+            .expand(k.shape[0], -1, -1),
+        )
+
+        forward_batch.token_to_kv_pool.set_kv_buffer(
+            layer.layer_id, forward_batch.out_cache_loc, k, v, k_label
+        )
+
+        (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        ) = self.forward_metadata
+        self.extend_attention_fwd(
+            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            k.contiguous(),
+            v.contiguous(),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+            forward_batch.req_to_token_pool.req_to_token,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.extend_seq_lens,
+            forward_batch.extend_start_loc,
+            max_extend_len,
+            layer.scaling,
+            layer.logit_cap,
+        )
+        return o
+
+    def forward_decode(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        # During torch.compile, there is a bug in rotary_emb that causes the
+        # output value to have a 3D tensor shape. This reshapes the output correctly.
+        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+
+        # TODO: reuse the buffer across layers
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        # TODO: Add min seqlen
+        (
+            start_loc,
+            attn_logits,
+            max_seq_len,
+            min_seq_len,
+            max_extend_len,
+            ds_req_to_token,
+        ) = self.forward_metadata
+
+        k_label = torch.gather(
+            k,
+            2,
+            self.sorted_channels[layer.layer_id]
+            .unsqueeze(0)
+            .expand(k.shape[0], -1, -1),
+        )
+
+        forward_batch.token_to_kv_pool.set_kv_buffer(
+            layer.layer_id, forward_batch.out_cache_loc, k, v, k_label
+        )
+
+        # NOTE(Andy) shouldn't be used when max_len_in_batch < heavy_token_num
+        #            and set a minimum value for sparse_decode
+        if (
+            min_seq_len < self.heavy_token_num
+            or max_seq_len < self.sparse_decode_thresold
+        ):
+            self.decode_attention_fwd(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+                o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+                forward_batch.req_to_token_pool.req_to_token,
+                forward_batch.req_pool_indices,
+                start_loc,
+                forward_batch.seq_lens,
+                attn_logits,
+                max_seq_len,
+                layer.scaling,
+                layer.logit_cap,
+            )
+        else:
+            # TODO(Andy): indexing with torch.gather or torch.index_select or customized kernel
+            q_label = torch.gather(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                2,
+                self.sorted_channels[layer.layer_id]
+                .unsqueeze(0)
+                .expand(q.shape[0], -1, -1),
+            )
+            self.decode_sparse_attention_fwd(
+                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
+                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
+                o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+                q_label,
+                forward_batch.token_to_kv_pool.get_label_buffer(layer.layer_id),
+                ds_req_to_token,
+                forward_batch.seq_lens,
+                max_seq_len,
+                layer.scaling,
+                layer.logit_cap,
+                self.heavy_token_num,
+                self.att_out_approx,
+                self.mid_out,
+                self.mid_o_logexpsum,
+                self.BLOCK_SEQ,
+            )
+
+        return o