# 传统方式 (无法用于 piecewise graph)
def forward(self, hidden_states, forward_batch):
    # ... 计算逻辑 ...
    return output  # 内部分配 output tensor

# _with_output 模式 (适用于 piecewise graph)
def forward(self, hidden_states, forward_batch):
    output = torch.empty_like(hidden_states)  # 预分配
    if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
        torch.ops.sglang.gdn_with_output(
            hidden_states,
            output,      # ← 传入预分配的输出 tensor
            self.layer_id,
        )
        return output
    else:
        return self._forward(hidden_states, forward_batch)

# 1. 实现实际执行函数
def gdn_with_output(hidden_states, output, layer_id):
    context = get_forward_context()
    attention_layer = context.attention_layers[layer_id]
    ret = attention_layer._forward(hidden_states, forward_batch)
    output.copy_(ret)
    return

# 2. 实现 fake 函数 (用于 torch.compile)
def gdn_with_output_fake(hidden_states, output, layer_id):
    return  # 空实现

# 3. 注册 custom op
direct_register_custom_op(
    op_name="gdn_with_output",
    op_func=gdn_with_output,
    mutates_args=["output"],  # 标记 output 为可变参数
    fake_impl=gdn_with_output_fake,
)

# backend.py:28-32
SPLIT_OPS = [
    "sglang.unified_attention_with_output",
    "sglang.inplace_all_reduce",
    "sglang.gdn_with_output",  # ← 新增
]

def _forward_input_proj(self, hidden_states: torch.Tensor):
    # 关键修改：在 piecewise graph 模式下禁用 dual stream
    if _is_npu or get_global_server_args().enable_piecewise_cuda_graph:
        DUAL_STREAM_TOKEN_THRESHOLD = 0  # ← 禁用 dual stream
    else:
        DUAL_STREAM_TOKEN_THRESHOLD = 1024  # ← 非 piecewise 模式启用

    seq_len, _ = hidden_states.shape
    if seq_len < DUAL_STREAM_TOKEN_THRESHOLD:
        # Dual stream 优化：使用两个 CUDA stream 并行执行
        current_stream = torch.cuda.current_stream()
        self.alt_stream.wait_stream(current_stream)
        projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
        with torch.cuda.stream(self.alt_stream):
            projected_states_ba, _ = self.in_proj_ba(hidden_states)
        current_stream.wait_stream(self.alt_stream)
    else:
        # 单 stream 执行
        projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
        projected_states_ba, _ = self.in_proj_ba(hidden_states)
    return projected_states_qkvz, projected_states_ba

def forward(
    self,
    hidden_states: torch.Tensor,
    forward_batch: ForwardBatch,
):
    # 预分配输出 tensor
    output = torch.empty_like(hidden_states)

    # 检测 piecewise graph 模式
    if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
        # 使用 custom op split 模式
        torch.ops.sglang.gdn_with_output(
            hidden_states,
            output,
            self.layer_id,
        )
        return output
    else:
        # 正常执行路径
        return self._forward(hidden_states, forward_batch)

def _forward(
    self,
    hidden_states: torch.Tensor,
    forward_batch: ForwardBatch,
):
    # 原始的完整计算逻辑
    seq_len, _ = hidden_states.shape
    is_cuda_graph = forward_batch.forward_mode.is_cuda_graph()

    projected_states_qkvz, projected_states_ba = self._forward_input_proj(
        hidden_states
    )

    # ... 其余计算逻辑保持不变 ...

def gdn_with_output(
    hidden_states: torch.Tensor,
    output: torch.Tensor,
    layer_id: int,
) -> None:
    # 1. 获取 forward context (由 piecewise context manager 设置)
    context = get_forward_context()
    forward_batch = context.forward_batch
    attention_layers = context.attention_layers

    # 2. 根据 layer_id 获取实际的 GDN layer 对象
    attention_layer = attention_layers[layer_id]

    # 3. 调用实际的 _forward 方法执行计算
    ret = attention_layer._forward(hidden_states, forward_batch)

    # 4. 将结果写入预分配的 output tensor
    assert (
        output.numel() == ret.numel()
    ), f"Output tensor element mismatch: {output.numel()} != {ret.numel()}"

    output.view(ret.shape).copy_(ret)
    return

def gdn_with_output_fake(
    hidden_states: torch.Tensor,
    output: torch.Tensor,
    layer_id: int,
) -> None:
    return  # 空实现

direct_register_custom_op(
    op_name="gdn_with_output",
    op_func=gdn_with_output,
    mutates_args=["output"],  # ← 标记 output 为可变参数
    fake_impl=gdn_with_output_fake,
)

# 原有代码处理标准 attention layer
if hasattr(layer, "self_attn"):
    if hasattr(layer.self_attn, "attn"):
        self.attention_layers.append(layer.self_attn.attn)
    elif hasattr(layer.self_attn, "attn_mqa"):
        self.attention_layers.append(layer.self_attn.attn_mqa)

# 新增：处理 hybrid attention layer (Qwen3-Next)
elif hasattr(layer, "attn"):
    self.attention_layers.append(layer.attn)
elif hasattr(layer, "linear_attn"):
    self.attention_layers.append(layer.linear_attn)

class Qwen3NextDecoderLayer:
    self.attn = Qwen3Attention()        # 标准 attention
    self.linear_attn = Qwen3GatedDeltaNet()  # 线性 attention (GDN)
    # ...

# 修改前
o = torch.empty_like(v)  # ← 未初始化，包含垃圾值

# 修改后
o = torch.zeros_like(v)  # ← 初始化为 0

# 修改前
def mem_usage_bytes(self):
    return sum(get_tensor_size_bytes(t) for t in vars(self).values())
    # ↑ vars() 不包含 frozen dataclass 的字段

# 修改后
def mem_usage_bytes(self):
    return sum(
        get_tensor_size_bytes(getattr(self, f.name))
        for f in dataclasses.fields(self)  # ← 使用 dataclasses.fields()
    )

# 1. 设置 context (在 piecewise_cuda_graph_runner.py)
with set_forward_context(
    forward_batch,
    self.attention_layers,  # ← 包含所有 GDN layers
    self.quant_config,
    self.moe_layers,
):
    output = self.model_runner.model.forward(...)

# 2. 在 gdn_with_output 中访问 (在 qwen3_next.py)
context = get_forward_context()
attention_layer = context.attention_layers[layer_id]

# 3. Context 自动清理

@contextmanager
def set_forward_context(
    forward_batch: ForwardBatch,
    attention_layers: List[Any],
    quant_config: Any,
    moe_layers: List[Any],
):
    global _forward_context
    _forward_context = ForwardContext()
    _forward_context.set_forward_batch(forward_batch)
    _forward_context.set_attention_layers(attention_layers)
    # ...
    try:
        yield
    finally:
        _forward_context = None  # ← 自动清理

# unified_attention.py (已存在)
def forward(self, hidden_states, forward_batch):
    if is_extend() and get_forward_context() is not None:
        torch.ops.sglang.unified_attention_with_output(
            hidden_states,
            self.output_buffer,
            self.layer_id,
        )
        return self.output_buffer
    else:
        return self._forward(hidden_states, forward_batch)

# qwen3_next.py (本 PR 新增)
def forward(self, hidden_states, forward_batch):
    if is_extend() and get_forward_context() is not None:
        torch.ops.sglang.gdn_with_output(
            hidden_states,
            output,  # 预分配
            self.layer_id,
        )
        return output
    else:
        return self._forward(hidden_states, forward_batch)

class TestQwen3NextPiecewiseCudaGraph(CustomTestCase):

    @classmethod
    def setUpClass(cls):
        cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.process = popen_launch_server(
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            other_args=[
                "--tp", "4",
                "--enable-piecewise-cuda-graph",          # ← 启用 piecewise graph
                "--piecewise-cuda-graph-compiler", "eager",  # ← 使用 eager 编译器
            ],
        )

    def test_gsm8k(self):
        args = SimpleNamespace(
            num_shots=5,
            num_questions=200,
            max_new_tokens=512,
            parallel=128,
            # ...
        )
        metrics = run_eval(args)
        self.assertGreater(metrics["accuracy"], 0.93)  # ← 验证正确性

# 在代码中添加
from torch._dynamo.utils import lazy_format_graph_code
lazy_format_graph_code("after split", split_gm)  # 打印切分后的 graph

# ✓ 使用 _with_output 模式
def forward(self, x, ...):
    output = torch.empty_like(x)
    if get_forward_context() is not None:
        torch.ops.sglang.my_op_with_output(x, output, ...)
        return output
    return self._forward(x, ...)

# ✓ 通过 layer_id 查找实际 layer
def my_op_with_output(x, output, layer_id):
    layer = get_forward_context().attention_layers[layer_id]
    ret = layer._forward(x)
    output.copy_(ret)

# ✗ 不要在 custom op 中创建新 tensor
def my_op(x, layer_id):
    layer = get_forward_context().attention_layers[layer_id]
    return layer._forward(x)  # ← 返回新 tensor，内存地址不固定

# ✗ 不要硬编码 layer 引用
def my_op_with_output(x, output):
    global my_layer  # ← 避免全局变量
    ret = my_layer._forward(x)
    output.copy_(ret)

import logging
logging.basicConfig(level=logging.DEBUG)
# 查看 split 后的 subgraph 信息

context = get_forward_context()
assert context is not None, "Not in piecewise graph context"
assert context.attention_layers is not None
assert len(context.attention_layers) > layer_id