From fed12eed7536bdc0149787a99f7848f2b6dea972 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Sat, 10 Oct 2026 21:21:29 -0700 Subject: [PATCH 1/7] refactor: consolidate attention eager regions and isolate Inkling policy --- .../srt/layers/attention/graph_utils.py | 240 +++++ .../srt/layers/attention/hpc_ops_backend.py | 125 ++- python/sglang/srt/layers/radix_attention.py | 788 ++-------------- .../srt/layers/radix_linear_attention.py | 118 +-- .../runner/prefill_cuda_graph_runner.py | 5 + .../runner_utils/prefill_graph.py | 35 + .../attention_forward_methods/forward_mla.py | 150 ++- python/sglang/srt/models/inkling.py | 75 +- .../srt/models/inkling_common/kernels/comm.py | 24 +- .../breakable/test_breakable_cuda_graph.py | 26 +- .../unit/layers/test_radix_attention.py | 871 +++++------------- .../test_radix_attention_padded_extend.py | 2 +- .../layers/test_radix_linear_attention.py | 105 ++- .../models/test_inkling_graph_fusion_gates.py | 71 ++ 14 files changed, 996 insertions(+), 1639 deletions(-) create mode 100644 python/sglang/srt/layers/attention/graph_utils.py create mode 100644 python/sglang/srt/model_executor/runner_utils/prefill_graph.py create mode 100644 test/registered/unit/models/test_inkling_graph_fusion_gates.py diff --git a/python/sglang/srt/layers/attention/graph_utils.py b/python/sglang/srt/layers/attention/graph_utils.py new file mode 100644 index 000000000000..d1bd0746704e --- /dev/null +++ b/python/sglang/srt/layers/attention/graph_utils.py @@ -0,0 +1,240 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Token slicing and temporary batch state for graph attention calls.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.model_executor.forward_context import get_attn_backend + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +def _zero_padded_tokens(output, actual_tokens): + if actual_tokens is not None: + output[actual_tokens:].zero_() + + +@contextmanager +def attention_input_scope(forward_batch, output, query_tokens, kv_tokens, kwargs): + """Narrow per-token operands and restore the caller's batch even on failure.""" + kwargs = dict(kwargs) + for name in ( + "q_rope", + "topk_indices", + "rel_bias", + "q_descale", + "idx_q", + "mxfp8_norm_rope_positions", + "mxfp8_norm_rope_temp_scale", + ): + if kwargs.get(name) is not None: + kwargs[name] = kwargs[name][:query_tokens] + for name in ("k_rope", "k_descale", "v_descale", "idx_k", "idx_v"): + if kwargs.get(name) is not None: + kwargs[name] = kwargs[name][:kv_tokens] + if kwargs.get("aux_tensors") is not None: + kwargs["aux_tensors"] = [ + tensor[:query_tokens] for tensor in kwargs["aux_tensors"] + ] + cache_loc, positions = forward_batch.out_cache_loc, forward_batch.positions + previous_output = forward_batch._attn_output + forward_batch.out_cache_loc = cache_loc[:query_tokens] + if positions is not None: + forward_batch.positions = positions[:query_tokens] + forward_batch._attn_output = output[:query_tokens] + try: + yield kwargs + finally: + forward_batch.out_cache_loc, forward_batch.positions = cache_loc, positions + forward_batch._attn_output = previous_output + + +def allocate_attention_outputs(layer, query, value, index_query): + """Allocate graph-owned buffers before entering the eager attention region.""" + rows = query.shape[0] + sparse = index_query is not None + dtype = query.dtype if sparse or value is None else value.dtype + if dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + dtype = torch.bfloat16 + shape = ( + (rows, layer.tp_q_head_num * layer.v_head_dim) + if sparse or layer.qk_head_dim != layer.v_head_dim + else query.shape + ) + output = torch.empty(shape, dtype=dtype, device=query.device) + index_output = ( + query.new_empty((rows, index_query.shape[1] * index_query.shape[2])) + if sparse + else None + ) + return output, index_output + + +def _new_attention_output( + layer: RadixAttention, q: torch.Tensor, v: Optional[torch.Tensor] +) -> torch.Tensor: + """An uninitialized buffer for the layer's attention output over q's rows. + + Its dtype follows v (the model dtype) when available: qk-norm may emit q in + a different dtype without changing the dtype the backend writes. FP8 q/v + (e.g. mxfp8 KV-cache attention) still produce a bf16 attention output; + sizing the buffer off an fp8 dtype would silently cast-copy the result to + fp8. + """ + out_dtype = v.dtype if v is not None else q.dtype + if out_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + out_dtype = torch.bfloat16 + if layer.qk_head_dim != layer.v_head_dim: + return q.new_empty( + (q.shape[0], layer.tp_q_head_num * layer.v_head_dim), dtype=out_dtype + ) + return torch.empty_like(q, dtype=out_dtype) + + +# Per-row attention inputs beside Q, K and V that the tc-piecewise op schema +# cannot carry. The padded paths narrow them with the rows they belong to. +_PER_ROW_EXTRA_KWARGS = ( + "rel_bias", + "q_descale", + "k_descale", + "v_descale", + "mxfp8_norm_rope_positions", + "mxfp8_norm_rope_temp_scale", +) +# The per-row inputs that follow K and V rather than the queries. +_KEY_ROW_KWARGS = frozenset({"k_rope", "k_descale", "v_descale", "idx_k", "idx_v"}) + + +def padded_extend_real_tokens( + q: torch.Tensor, forward_batch: ForwardBatch +) -> Optional[int]: + """The number of real rows of an extend batch that MLP sync padded to a multiple of + attention TP, whose attention metadata covers only those rows; None for a + batch without such padding. Target verify plans its padded rows itself.""" + mode = forward_batch.forward_mode + if not mode.is_extend() or mode.is_target_verify(): + return None + real_num_tokens = forward_batch.global_num_token_non_padded_cpu + if real_num_tokens is None or not 0 < real_num_tokens < q.shape[0]: + return None + return real_num_tokens + + +def _attention_on_real_rows( + layer: RadixAttention, + real_num_tokens: int, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + forward_batch: ForwardBatch, + save_kv_cache: bool, + *, + key_value_num_tokens: Optional[int] = None, + **kwargs, +): + """Run the backend on a padded batch's real rows, with the cache locations + and positions narrowed for the call, and return its output padded back + with zero rows, as the prefill graph path does. K and V, and the per-row + inputs that follow them, are narrowed with the queries unless the caller + gives their own extent (K and V that hold a cached prefix), the layer is a + cross-attention (K and V are the encoder's), or they have a different row + count. A sparse indexer's q, k and v go with the rows they + index. A backend that writes into forward_batch._attn_output fills the + padded output's real rows in place.""" + num_tokens = q.shape[0] + + def queries(t): + return t[:real_num_tokens] if t is not None else None + + keys_own_rows = key_value_num_tokens is not None or layer.is_cross_attention + + def keys(t): + if t is None or keys_own_rows or t.shape[0] != num_tokens: + return t + return t[:real_num_tokens] + + for name in ( + "q_rope", + "k_rope", + "topk_indices", + "idx_q", + "idx_k", + "idx_v", + *_PER_ROW_EXTRA_KWARGS, + ): + if kwargs.get(name) is not None: + narrow = keys if name in _KEY_ROW_KWARGS else queries + kwargs[name] = narrow(kwargs[name]) + if kwargs.get("aux_tensors") is not None: + kwargs["aux_tensors"] = [queries(t) for t in kwargs["aux_tensors"]] + # A backend that returns more than the output (its LSE, or a sparse + # indexer's output beside it) returns a tuple, padded as it comes. + output = ( + None + if kwargs.get("return_lse") + or forward_batch.mha_return_lse + or kwargs.get("idx_q") is not None + else _new_attention_output(layer, q, v) + ) + out_cache_loc = forward_batch.out_cache_loc + positions = forward_batch.positions + attn_output = forward_batch._attn_output + forward_batch.out_cache_loc = out_cache_loc[:real_num_tokens] + if positions is not None: + forward_batch.positions = positions[:real_num_tokens] + if output is not None: + forward_batch._attn_output = output[:real_num_tokens] + try: + ret = get_attn_backend().forward( + queries(q), + keys(k), + keys(v), + layer, + forward_batch, + save_kv_cache, + **kwargs, + ) + finally: + forward_batch.out_cache_loc = out_cache_loc + forward_batch.positions = positions + forward_batch._attn_output = attn_output + + def padded(t): + if not isinstance(t, torch.Tensor) or t.shape[0] != real_num_tokens: + return t + full = t.new_empty((num_tokens, *t.shape[1:])) + full[:real_num_tokens].copy_(t) + full[real_num_tokens:].zero_() + return full + + if isinstance(ret, tuple): + return tuple(padded(t) for t in ret) + if ( + output is None + or not isinstance(ret, torch.Tensor) + or ret.shape[0] != real_num_tokens + or ret.numel() != output[:real_num_tokens].numel() + ): + return padded(ret) + if ret.data_ptr() != output.data_ptr(): + output[:real_num_tokens].view(ret.shape).copy_(ret) + output[real_num_tokens:].zero_() + return output.view(num_tokens, *ret.shape[1:]) diff --git a/python/sglang/srt/layers/attention/hpc_ops_backend.py b/python/sglang/srt/layers/attention/hpc_ops_backend.py index bf6aa687ce6f..5c882e66f393 100644 --- a/python/sglang/srt/layers/attention/hpc_ops_backend.py +++ b/python/sglang/srt/layers/attention/hpc_ops_backend.py @@ -40,20 +40,21 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( build_trtllm_mha_page_table, ) -from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.model_executor.forward_context import get_attn_backend +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, +) from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( eager_on_graph, is_in_breakable_cuda_graph, ) -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, +from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + is_in_full_prefill_graph, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -393,7 +394,7 @@ def fused_qk_rope_store_kv_fp8( head_dim]. The subsequent ``self.attn(...)`` call must pass ``save_kv_cache=False`` (K/V are already written here). - Under a captured prefill graph (breakable / tc_piecewise) this routes + Under a captured prefill graph (breakable / full) this routes through a graph-splitting op — like attention itself — so the Python Q-scale hand-off between this op and attention stays alive at replay. @@ -410,29 +411,22 @@ def fused_qk_rope_store_kv_fp8( device=qkv.device, ) - if is_extend and get_tc_piecewise_forward_context() is not None: - # Captured prefill graph: run through the splitting op (eager at - # capture AND at replay, keeping the metadata hand-off alive). - if is_in_breakable_cuda_graph(): - breakable_hpc_ops_fp8_rope_store_kv( - qkv, - cos_sin_cache, - out_q, - layer.layer_id, - qk_norm_policy, - q_norm_weight=q_norm_weight, - k_norm_weight=k_norm_weight, - ) - else: - hpc_ops_fp8_rope_store_kv( - qkv, - cos_sin_cache, - out_q, - layer.layer_id, - qk_norm_policy, - q_norm_weight=q_norm_weight, - k_norm_weight=k_norm_weight, - ) + if is_extend and ( + is_in_full_prefill_graph() + or ( + is_in_breakable_cuda_graph() + and forward_batch.forward_mode.is_extend_without_speculative() + ) + ): + self._eager_rope_store_kv( + qkv, + cos_sin_cache, + out_q, + layer, + qk_norm_policy, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + ) else: self._run_fp8_rope_store_kv( layer=layer, @@ -657,44 +651,37 @@ def forward_decode( ) return o.view(-1, layer.tp_q_head_num * layer.head_dim) + @eager_on_graph + def _eager_rope_store_kv( + self, + qkv: torch.Tensor, + cos_sin_cache: torch.Tensor, + out_q: torch.Tensor, + layer, + qk_norm_policy: int, + *, + q_norm_weight: Optional[torch.Tensor] = None, + k_norm_weight: Optional[torch.Tensor] = None, + ) -> None: + """Graph-splitting wrapper for the fused QKNorm+RoPE+FP8-quant+StoreKV op. -@register_custom_op(mutates_args=["out_q"]) -@register_split_op() -def hpc_ops_fp8_rope_store_kv( - qkv: torch.Tensor, - cos_sin_cache: torch.Tensor, - out_q: torch.Tensor, - layer_id: int, - qk_norm_policy: int, - *, - q_norm_weight: Optional[torch.Tensor] = None, - k_norm_weight: Optional[torch.Tensor] = None, -) -> None: - """Graph-splitting wrapper for the fused QKNorm+RoPE+FP8-quant+StoreKV op. - - Like ``unified_attention_with_output``, this runs eagerly between captured - prefill-graph segments (at capture and at every replay), so the Python - hand-off of the dynamic Q scales to the following attention op stays - alive. ``out_q`` is preallocated by the captured segment and mutated in - place, which is what stitches the surrounding graph segments together. - """ - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layer = context.attention_layers[layer_id] - real_num_tokens = forward_batch.global_num_token_non_padded_cpu - - backend = get_attn_backend() - backend._run_fp8_rope_store_kv( - layer=attention_layer, - forward_batch=forward_batch, - qkv=qkv[:real_num_tokens], - cos_sin_cache=cos_sin_cache, - q_norm_weight=q_norm_weight, - k_norm_weight=k_norm_weight, - qk_norm_policy=qk_norm_policy, - is_extend=True, - out_q=out_q[:real_num_tokens], - ) - - -breakable_hpc_ops_fp8_rope_store_kv = eager_on_graph(True)(hpc_ops_fp8_rope_store_kv) + Like ``RadixAttention._eager_attention``, this runs eagerly between captured + prefill-graph segments (at capture and at every replay), so the Python + hand-off of the dynamic Q scales to the following attention op stays + alive. ``out_q`` is preallocated by the captured segment and mutated in + place, which is what stitches the surrounding graph segments together. + """ + forward_batch = get_forward_batch() + real_num_tokens = forward_batch.global_num_token_non_padded_cpu + + get_attn_backend()._run_fp8_rope_store_kv( + layer=layer, + forward_batch=forward_batch, + qkv=qkv[:real_num_tokens], + cos_sin_cache=cos_sin_cache, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + qk_norm_policy=qk_norm_policy, + is_extend=True, + out_q=out_q[:real_num_tokens], + ) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 136e4ff0f4f0..ef5b75d4611d 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -15,66 +15,31 @@ from __future__ import annotations -from contextlib import contextmanager -from contextvars import ContextVar from enum import Enum from typing import TYPE_CHECKING, Optional import torch from torch import nn -from sglang.srt.compilation.compilation_config import register_split_op -from sglang.srt.model_executor.forward_context import get_attn_backend +from sglang.srt.layers.attention.graph_utils import ( + _attention_on_real_rows, + _zero_padded_tokens, + allocate_attention_outputs, + attention_input_scope, + padded_extend_real_tokens, +) +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, +) from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( eager_on_graph, is_in_breakable_cuda_graph, ) -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, +from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + get_prefill_raw_num_tokens, + is_in_full_prefill_graph, ) -from sglang.srt.utils.common import is_hip -from sglang.srt.utils.custom_op import register_custom_op - -_is_hip = is_hip() - -# When set, RadixAttention.forward runs the attention backend eagerly instead of -# routing through the tc-piecewise split op. A caller already inside a -# breakable-CUDA-graph eager break (e.g. Inkling wrapping norm+attn+sconv in one eager -# region for multi-seq-correct short-conv metadata) sets this so the attn does not -# start a nested break (which would assert on the ended segment). Default off. -_force_eager_attn: ContextVar[bool] = ContextVar("_force_eager_attn", default=False) - - -@contextmanager -def force_eager_attention(): - token = _force_eager_attn.set(True) - try: - yield - finally: - _force_eager_attn.reset(token) - - -def _zero_padded_pcg_tail(buf: torch.Tensor, context) -> None: - """Zero the padded tail ``buf`` leaves as torch.empty garbage under PCG - replay, so NaN/Inf cannot reach residual / MoE routing / allreduce.""" - pcg_static_tokens = context.num_tokens - actual_tokens = context.raw_num_tokens - if ( - pcg_static_tokens is not None - and actual_tokens is not None - and pcg_static_tokens > actual_tokens - ): - first_dim = buf.shape[0] - elems_per_token = buf.numel() // first_dim - buf.view(first_dim, elems_per_token)[actual_tokens:].zero_() - - -def _zero_skipped_attn_outputs(*bufs: Optional[torch.Tensor]) -> None: - """Zero outputs when an idle DP rank skips attention work.""" - for buf in bufs: - if buf is not None: - buf.zero_() - if TYPE_CHECKING: from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -82,24 +47,14 @@ def _zero_skipped_attn_outputs(*bufs: Optional[torch.Tensor]) -> None: class AttentionType(Enum): - """ - Attention type. - Use string to be compatible with `torch.compile`. - """ + """String-valued attention types, compatible with torch.compile.""" - # Decoder attention between previous layer Q/K/V DECODER = "decoder" - # Decoder bidirectional attention between image tokens DECODER_BIDIRECTIONAL = "decoder_bidirectional" - # Encoder attention between previous layer Q/K/V ENCODER_ONLY = "encoder_only" class RadixAttention(nn.Module): - """ - The attention layer implementation. - """ - def __init__( self, num_heads: int, @@ -137,9 +92,6 @@ def __init__( self.v_scale = None self.k_scale_float = None self.v_scale_float = None - # MiniMax-M3 fp8 attention-GEMM scales (fp8 attn-GEMM mode): main q and - # lightning-indexer q/k/v. No checkpoint loader populates them yet; - # None means unit scale. self.q_scale_float = None self.idx_q_scale_float = None self.idx_k_scale_float = None @@ -167,658 +119,110 @@ def forward( **kwargs, ): if k is not None: - # For cross-layer sharing, kv can be None assert v is not None + k = k.view( + -1, + self.tp_k_head_num, + self.v_head_dim if "k_rope" in kwargs else self.qk_head_dim, + ) if "k_rope" not in kwargs: - k = k.view(-1, self.tp_k_head_num, self.qk_head_dim) v = v.view(-1, self.tp_v_head_num, self.v_head_dim) - else: - k = k.view(-1, self.tp_k_head_num, self.v_head_dim) - - context = get_tc_piecewise_forward_context() - if ( + breakable_cg = is_in_breakable_cuda_graph() + full_cg = is_in_full_prefill_graph() + use_eager_attention = ( self.use_prefill_attention_wrapper and forward_batch.forward_mode.is_extend() - and context is not None - # ``_force_eager_attn`` is only set inside Inkling's eager - # norm+attn+sconv region, never during tc-piecewise capture. Reading - # the ContextVar under the fullgraph torch.compile trace is - # untraceable ("Unsupported method call: ContextVar.get"), so - # short-circuit it while compiling -- force-eager is always off there. - and (torch.compiler.is_compiling() or not _force_eager_attn.get()) - ): - if kwargs.get("idx_q") is not None: - idx_q = kwargs["idx_q"] - idx_k = kwargs["idx_k"] - idx_v = kwargs.get("idx_v") - attn_out = q.new_empty( - (q.shape[0], self.tp_q_head_num * self.v_head_dim) - ) - idx_out = q.new_empty((q.shape[0], idx_q.shape[1] * idx_q.shape[2])) - op = ( - breakable_unified_sparse_attention_with_output - if is_in_breakable_cuda_graph() - else unified_sparse_attention_with_output - ) - op( - q, - k, - v, - attn_out, - idx_out, - idx_q, - idx_k, - save_kv_cache, - self.layer_id, - idx_v=idx_v, - ) - return idx_out, attn_out - output = _new_attention_output(self, q, v) - if any( - key in kwargs - for key in ( - "score_mod", - "aux_tensors", - "rel_bias", - "return_lse", - "q_descale", - "k_descale", - "v_descale", - "mxfp8_norm_rope_positions", - ) - ): - # A score_mod callable, aux_tensors, rel_bias, mxfp8 descale - # tensors, or the mxfp8 deferred norm/RoPE operands can't cross - # the unified_attention_with_output custom-op schema; route this - # backend's extend attention through the plain eager path. - if is_in_breakable_cuda_graph(): - lse = breakable_attention_with_output_extra_kwargs( - q, k, v, output, save_kv_cache, self.layer_id, kwargs - ) - else: - lse = attention_with_output_extra_kwargs( - q, k, v, output, save_kv_cache, self.layer_id, kwargs - ) - if kwargs.get("return_lse") or forward_batch.mha_return_lse: - assert lse is not None - return output.view(-1, self.tp_q_head_num, self.v_head_dim), lse - return output - # Chunked-prefix MHA needs LSE to merge independently normalized - # suffix and cached-prefix attention states. - return_lse = bool(forward_batch.mha_return_lse) - mha_companion_layers = context.mha_companion_layers - use_mha_companion = ( - mha_companion_layers is not None - and mha_companion_layers[self.layer_id] is self + and (breakable_cg or full_cg) + ) + if use_eager_attention: + is_sparse = kwargs.get("idx_q") is not None + output, idx_output = allocate_attention_outputs( + self, q, v, kwargs.get("idx_q") ) - if is_in_breakable_cuda_graph(): - op = ( - breakable_unified_attention_with_output_and_lse - if return_lse - else breakable_unified_attention_with_output - ) - else: - op = ( - unified_attention_with_output_and_lse - if return_lse - else unified_attention_with_output - ) - lse = op( + lse = self._eager_attention( q, k, v, output, save_kv_cache, - self.layer_id, - use_mha_companion=use_mha_companion, - key_value_num_tokens=key_value_num_tokens, + key_value_num_tokens, + idx_output, **kwargs, ) - if return_lse: + if is_sparse: + return idx_output, output + if kwargs.get("return_lse") or forward_batch.mha_return_lse: return output.view(-1, self.tp_q_head_num, self.v_head_dim), lse return output - else: - real_num_tokens = padded_extend_real_tokens(q, forward_batch) - if real_num_tokens is not None: - return _attention_on_real_rows( - self, - real_num_tokens, - q, - k, - v, - forward_batch, - save_kv_cache, - key_value_num_tokens=key_value_num_tokens, - **kwargs, - ) - return get_attn_backend().forward( + real_num_tokens = padded_extend_real_tokens(q, forward_batch) + if real_num_tokens is not None: + return _attention_on_real_rows( + self, + real_num_tokens, q, k, v, - self, forward_batch, save_kv_cache, + key_value_num_tokens=key_value_num_tokens, **kwargs, ) - - -def _new_attention_output( - layer: RadixAttention, q: torch.Tensor, v: Optional[torch.Tensor] -) -> torch.Tensor: - """An uninitialized buffer for the layer's attention output over q's rows. - - Its dtype follows v (the model dtype) when available: qk-norm may emit q in - a different dtype without changing the dtype the backend writes. FP8 q/v - (e.g. mxfp8 KV-cache attention) still produce a bf16 attention output; - sizing the buffer off an fp8 dtype would silently cast-copy the result to - fp8. - """ - out_dtype = v.dtype if v is not None else q.dtype - if out_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - out_dtype = torch.bfloat16 - if layer.qk_head_dim != layer.v_head_dim: - return q.new_empty( - (q.shape[0], layer.tp_q_head_num * layer.v_head_dim), dtype=out_dtype + return get_attn_backend().forward( + q, k, v, self, forward_batch, save_kv_cache, **kwargs ) - return torch.empty_like(q, dtype=out_dtype) - - -# Per-row attention inputs beside Q, K and V that the tc-piecewise op schema -# cannot carry. The padded paths narrow them with the rows they belong to. -_PER_ROW_EXTRA_KWARGS = ( - "rel_bias", - "q_descale", - "k_descale", - "v_descale", - "mxfp8_norm_rope_positions", - "mxfp8_norm_rope_temp_scale", -) -# The per-row inputs that follow K and V rather than the queries. -_KEY_ROW_KWARGS = frozenset({"k_rope", "k_descale", "v_descale", "idx_k", "idx_v"}) - - -def padded_extend_real_tokens( - q: torch.Tensor, forward_batch: ForwardBatch -) -> Optional[int]: - """The number of real rows of an extend batch that MLP sync padded to a multiple of - attention TP, whose attention metadata covers only those rows; None for a - batch without such padding. Target verify plans its padded rows itself.""" - mode = forward_batch.forward_mode - if not mode.is_extend() or mode.is_target_verify(): - return None - real_num_tokens = forward_batch.global_num_token_non_padded_cpu - if real_num_tokens is None or not 0 < real_num_tokens < q.shape[0]: - return None - return real_num_tokens - - -def _attention_on_real_rows( - layer: RadixAttention, - real_num_tokens: int, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - forward_batch: ForwardBatch, - save_kv_cache: bool, - *, - key_value_num_tokens: Optional[int] = None, - **kwargs, -): - """Run the backend on a padded batch's real rows, with the cache locations - and positions narrowed for the call, and return its output padded back - with zero rows, as the prefill graph path does. K and V, and the per-row - inputs that follow them, are narrowed with the queries unless the caller - gives their own extent (K and V that hold a cached prefix), the layer is a - cross-attention (K and V are the encoder's), or they have a different row - count. A sparse indexer's q, k and v go with the rows they - index. A backend that writes into forward_batch._attn_output fills the - padded output's real rows in place.""" - num_tokens = q.shape[0] - - def queries(t): - return t[:real_num_tokens] if t is not None else None - - keys_own_rows = key_value_num_tokens is not None or layer.is_cross_attention - - def keys(t): - if t is None or keys_own_rows or t.shape[0] != num_tokens: - return t - return t[:real_num_tokens] - - for name in ( - "q_rope", - "k_rope", - "topk_indices", - "idx_q", - "idx_k", - "idx_v", - *_PER_ROW_EXTRA_KWARGS, - ): - if kwargs.get(name) is not None: - narrow = keys if name in _KEY_ROW_KWARGS else queries - kwargs[name] = narrow(kwargs[name]) - if kwargs.get("aux_tensors") is not None: - kwargs["aux_tensors"] = [queries(t) for t in kwargs["aux_tensors"]] - # A backend that returns more than the output (its LSE, or a sparse - # indexer's output beside it) returns a tuple, padded as it comes. - output = ( - None - if kwargs.get("return_lse") - or forward_batch.mha_return_lse - or kwargs.get("idx_q") is not None - else _new_attention_output(layer, q, v) - ) - out_cache_loc = forward_batch.out_cache_loc - positions = forward_batch.positions - attn_output = forward_batch._attn_output - forward_batch.out_cache_loc = out_cache_loc[:real_num_tokens] - if positions is not None: - forward_batch.positions = positions[:real_num_tokens] - if output is not None: - forward_batch._attn_output = output[:real_num_tokens] - try: - ret = get_attn_backend().forward( - queries(q), - keys(k), - keys(v), - layer, - forward_batch, - save_kv_cache, - **kwargs, - ) - finally: - forward_batch.out_cache_loc = out_cache_loc - forward_batch.positions = positions - forward_batch._attn_output = attn_output - - def padded(t): - if not isinstance(t, torch.Tensor) or t.shape[0] != real_num_tokens: - return t - full = t.new_empty((num_tokens, *t.shape[1:])) - full[:real_num_tokens].copy_(t) - full[real_num_tokens:].zero_() - return full - - if isinstance(ret, tuple): - return tuple(padded(t) for t in ret) - if ( - output is None - or not isinstance(ret, torch.Tensor) - or ret.shape[0] != real_num_tokens - or ret.numel() != output[:real_num_tokens].numel() - ): - return padded(ret) - if ret.data_ptr() != output.data_ptr(): - output[:real_num_tokens].view(ret.shape).copy_(ret) - output[real_num_tokens:].zero_() - return output.view(num_tokens, *ret.shape[1:]) - - -def _unified_attention_with_output_impl( - query: torch.Tensor, - key: Optional[torch.Tensor], - value: Optional[torch.Tensor], - output: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - use_mha_companion: bool, - return_lse: bool, - *, - key_value_num_tokens: Optional[int] = None, - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - sinks: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - # MLA / TRT-LLM / NSA paths pass these through RadixAttention.forward(**kwargs); - # they must appear in the schema when --cuda-graph-backend-prefill=tc_piecewise is on. - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = None, - llama_4_scaling: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layers = context.attention_layers - attention_layer = attention_layers[layer_id] - real_query_num_tokens = forward_batch.global_num_token_non_padded_cpu - # Ordinary PCG attention pads Q/K/V to the same token bucket. Prefix MHA - # instead supplies a fixed-capacity K/V chunk whose extent is independent - # of the suffix queries, so its caller must preserve that separate extent. - if key_value_num_tokens is None: - key_value_num_tokens = real_query_num_tokens - - if real_query_num_tokens == 0: - _zero_skipped_attn_outputs(output) - if return_lse: - # unified_attention_with_output_and_lse asserts a tensor comes back. - # Match _unified_attention_with_output_and_lse_fake's meta shape and - # the padded LSE the normal path returns below (padded row count, - # i.e. query before narrowing). - return query.new_zeros( - (query.shape[0], query.shape[1]), dtype=torch.float32 - ) - return None - - query = query[:real_query_num_tokens] - if key is not None: - key = key[:key_value_num_tokens] - if value is not None: - value = value[:key_value_num_tokens] - # DeepSeek MLA has two RadixAttention instances per layer (attn_mqa and - # attn_mha) that share the same layer_id. Preserve the calling instance's - # identity through the custom-op boundary; save_kv_cache is not an identity - # signal because absorbed MLA can also disable a redundant cache store. - if use_mha_companion: - assert context.mha_companion_layers is not None - attention_layer = context.mha_companion_layers[layer_id] - assert attention_layer is not None - - kwargs = {} - if q_rope is not None: - kwargs["q_rope"] = q_rope[:real_query_num_tokens] - if k_rope is not None: - kwargs["k_rope"] = k_rope[:key_value_num_tokens] - if sinks is not None: - kwargs["sinks"] = sinks - if attn_sink is not None: - kwargs["attn_sink"] = attn_sink - if cos_sin_cache is not None: - kwargs["cos_sin_cache"] = cos_sin_cache - if is_neox is not None: - kwargs["is_neox"] = is_neox - if llama_4_scaling is not None: - kwargs["llama_4_scaling"] = llama_4_scaling - if topk_indices is not None: - kwargs["topk_indices"] = topk_indices[:real_query_num_tokens] - - original_out_cache_loc = forward_batch.out_cache_loc - original_positions = forward_batch.positions - # Keep the original ForwardBatch object and only narrow cache locations for - # this backend call so model/backend state is still written to the same batch. - forward_batch.out_cache_loc = original_out_cache_loc[:real_query_num_tokens] - if original_positions is not None: - forward_batch.positions = original_positions[:real_query_num_tokens] - - # Store pre-allocated output for FA backend to write directly into. - # Must slice to real_query_num_tokens to match the narrowed query shape — - # the FA kernel validates out.size(0) == q.size(0). - forward_batch._attn_output = output[:real_query_num_tokens] - - ret = get_attn_backend().forward( - query, - key, - value, - attention_layer, - forward_batch, - save_kv_cache, - **kwargs, - ) - forward_batch.out_cache_loc = original_out_cache_loc - forward_batch.positions = original_positions - - lse = None - if return_lse: - assert isinstance(ret, tuple) - ret, lse, *_ = ret - else: - assert isinstance(ret, torch.Tensor) - - if ret.data_ptr() != output.data_ptr(): - output[:real_query_num_tokens].view(ret.shape).copy_(ret) - - # During PCG replay the attention backend writes only the narrowed - # real-token slice (output[:real_query_num_tokens]) and leaves padded positions - # as uninitialized torch.empty garbage. Zero them so garbage (NaN/Inf) does - # not propagate through residual connections, MoE routing, and allreduce. - # This affects every backend that varlen-writes under PCG, not just ROCm. - # Use context.raw_num_tokens (pre-padding count from PCG runner) instead of - # forward_batch.extend_num_tokens, which is None for TARGET_VERIFY batches. - _zero_padded_pcg_tail(output, context) - if lse is not None and lse.shape[0] != output.shape[0]: - padded_lse = lse.new_zeros((output.shape[0], *lse.shape[1:])) - padded_lse[:real_query_num_tokens].copy_(lse) - lse = padded_lse - return lse - - -@register_custom_op(mutates_args=["output"]) -@register_split_op() -def unified_attention_with_output( - query: torch.Tensor, - key: Optional[torch.Tensor], - value: Optional[torch.Tensor], - output: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - *, - use_mha_companion: bool = False, - key_value_num_tokens: Optional[int] = None, - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - sinks: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = None, - llama_4_scaling: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, -) -> None: - _unified_attention_with_output_impl( - query, - key, - value, - output, - save_kv_cache, - layer_id, - use_mha_companion, - False, - key_value_num_tokens=key_value_num_tokens, - q_rope=q_rope, - k_rope=k_rope, - sinks=sinks, - attn_sink=attn_sink, - cos_sin_cache=cos_sin_cache, - is_neox=is_neox, - llama_4_scaling=llama_4_scaling, - topk_indices=topk_indices, - ) - - -def _unified_attention_with_output_and_lse_fake( - query: torch.Tensor, *args, **kwargs -) -> torch.Tensor: - return query.new_empty((query.shape[0], query.shape[1]), dtype=torch.float32) - - -@register_custom_op( - mutates_args=["output"], fake_impl=_unified_attention_with_output_and_lse_fake -) -@register_split_op() -def unified_attention_with_output_and_lse( - query: torch.Tensor, - key: Optional[torch.Tensor], - value: Optional[torch.Tensor], - output: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - *, - use_mha_companion: bool = False, - key_value_num_tokens: Optional[int] = None, - q_rope: Optional[torch.Tensor] = None, - k_rope: Optional[torch.Tensor] = None, - sinks: Optional[torch.Tensor] = None, - attn_sink: Optional[torch.Tensor] = None, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = None, - llama_4_scaling: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, -) -> torch.Tensor: - lse = _unified_attention_with_output_impl( - query, - key, - value, + @eager_on_graph + def _eager_attention( + self, + q, + k, + v, output, - save_kv_cache, - layer_id, - use_mha_companion, - True, - key_value_num_tokens=key_value_num_tokens, - q_rope=q_rope, - k_rope=k_rope, - sinks=sinks, - attn_sink=attn_sink, - cos_sin_cache=cos_sin_cache, - is_neox=is_neox, - llama_4_scaling=llama_4_scaling, - topk_indices=topk_indices, - ) - assert lse is not None - return lse - - -@register_custom_op(mutates_args=["attn_out", "idx_out"]) -@register_split_op() -def unified_sparse_attention_with_output( - query: torch.Tensor, - key: Optional[torch.Tensor], - value: Optional[torch.Tensor], - attn_out: torch.Tensor, - idx_out: torch.Tensor, - idx_q: torch.Tensor, - idx_k: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - *, - idx_v: Optional[torch.Tensor] = None, -) -> None: - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layer = context.attention_layers[layer_id] - real_num_tokens = forward_batch.global_num_token_non_padded_cpu - - if real_num_tokens == 0: - _zero_skipped_attn_outputs(attn_out, idx_out) - return - - query = query[:real_num_tokens] - if key is not None: - key = key[:real_num_tokens] - if value is not None: - value = value[:real_num_tokens] - idx_q = idx_q[:real_num_tokens] - idx_k = idx_k[:real_num_tokens] - if idx_v is not None: - idx_v = idx_v[:real_num_tokens] - - original_out_cache_loc = forward_batch.out_cache_loc - forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens] - - ret_idx, ret_out = get_attn_backend().forward( - query, - key, - value, - attention_layer, - forward_batch, - save_kv_cache, - idx_q=idx_q, - idx_k=idx_k, - idx_v=idx_v, - ) - forward_batch.out_cache_loc = original_out_cache_loc - - attn_out[:real_num_tokens].view(ret_out.shape).copy_(ret_out) - # disable_value layers return ret_idx=None; the guard keeps idx_out's - # untouched real-token slice safe (model returns before index_o_proj). - if ret_idx is not None: - idx_out[:real_num_tokens].view(ret_idx.shape).copy_(ret_idx) - - for buf in (attn_out, idx_out): - _zero_padded_pcg_tail(buf, context) - return - - -breakable_unified_attention_with_output = eager_on_graph(True)( - unified_attention_with_output -) -breakable_unified_attention_with_output_and_lse = eager_on_graph(True)( - unified_attention_with_output_and_lse -) -breakable_unified_sparse_attention_with_output = eager_on_graph(True)( - unified_sparse_attention_with_output -) - - -def attention_with_output_extra_kwargs( - query: torch.Tensor, - key: Optional[torch.Tensor], - value: Optional[torch.Tensor], - output: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - extra_kwargs: dict, -) -> Optional[torch.Tensor]: - """Breakable/tc_piecewise attention for backends whose forward needs kwargs - that cannot cross the ``unified_attention_with_output`` custom-op schema -- - a ``score_mod`` callable and/or ``aux_tensors`` (e.g. Inkling's relative-bias - fa4 attention), or the per-token mxfp8 deferred norm/RoPE operands. Plain - (not a custom op) so the callable passes through; still - runs eagerly between graph segments under BCG via the wrapper below. Mirrors - the real-token narrowing + padded-output write of - ``unified_attention_with_output``, and narrows per-token ``aux_tensors`` too. - """ - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layer = context.attention_layers[layer_id] - real_num_tokens = forward_batch.global_num_token_non_padded_cpu - - if real_num_tokens == 0: - _zero_skipped_attn_outputs(output) - return - - query = query[:real_num_tokens] - if key is not None: - key = key[:real_num_tokens] - if value is not None: - value = value[:real_num_tokens] - - kwargs = dict(extra_kwargs) - aux_tensors = kwargs.get("aux_tensors") - if aux_tensors is not None: - kwargs["aux_tensors"] = [t[:real_num_tokens] for t in aux_tensors] - for per_token_key in _PER_ROW_EXTRA_KWARGS: - t = kwargs.get(per_token_key) - if t is not None: - kwargs[per_token_key] = t[:real_num_tokens] - - original_out_cache_loc = forward_batch.out_cache_loc - forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens] - forward_batch._attn_output = output[:real_num_tokens] - - ret = get_attn_backend().forward( - query, key, value, attention_layer, forward_batch, save_kv_cache, **kwargs - ) - forward_batch.out_cache_loc = original_out_cache_loc - - return_lse = bool(kwargs.get("return_lse") or forward_batch.mha_return_lse) - if return_lse: - assert isinstance(ret, tuple) - ret, lse, *_ = ret - else: - assert isinstance(ret, torch.Tensor) + save_kv_cache: bool = True, + key_value_num_tokens: Optional[int] = None, + idx_output=None, + **kwargs, + ): + """Run real tokens, retaining padded outputs for the next graph segment.""" + forward_batch = get_forward_batch() + rows = output.shape[0] + n = forward_batch.global_num_token_non_padded_cpu + kv_n = n if key_value_num_tokens is None else key_value_num_tokens + return_lse = bool(kwargs.get("return_lse") or forward_batch.mha_return_lse) lse = None - - if ret.data_ptr() != output.data_ptr(): - output[:real_num_tokens].view(ret.shape).copy_(ret) - - if _is_hip: - _zero_padded_pcg_tail(output, context) - if lse is not None and lse.shape[0] != output.shape[0]: - padded_lse = lse.new_zeros((output.shape[0], *lse.shape[1:])) - padded_lse[:real_num_tokens].copy_(lse) - lse = padded_lse - return lse - - -breakable_attention_with_output_extra_kwargs = eager_on_graph(True)( - attention_with_output_extra_kwargs -) + if n == 0: + output.zero_() + if idx_output is not None: + idx_output.zero_() + if return_lse: + lse = q.new_zeros((rows, q.shape[1]), dtype=torch.float32) + else: + with attention_input_scope( + forward_batch, output, n, kv_n, kwargs + ) as kwargs: + result = get_attn_backend().forward( + q[:n], + k[:kv_n] if k is not None else None, + v[:kv_n] if v is not None else None, + self, + forward_batch, + save_kv_cache, + **kwargs, + ) + if idx_output is not None: + idx_result, result = result + if idx_result is not None: + idx_output[:n].view(idx_result.shape).copy_(idx_result) + elif return_lse: + result, lse, *_ = result + if result.data_ptr() != output.data_ptr(): + output[:n].view(result.shape).copy_(result) + raw = get_prefill_raw_num_tokens() + _zero_padded_tokens(output, raw) + if idx_output is not None: + _zero_padded_tokens(idx_output, raw) + if lse is not None and lse.shape[0] != rows: + padded = lse.new_zeros((rows, *lse.shape[1:])) + padded[:n].copy_(lse) + lse = padded + return lse diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 3f0333772cc8..2ecafa55a18d 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -20,16 +20,17 @@ import torch from torch import nn -from sglang.srt.compilation.compilation_config import register_split_op -from sglang.srt.model_executor.forward_context import get_attn_backend +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, +) from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( eager_on_graph, is_in_breakable_cuda_graph, ) -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, +from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + is_in_full_prefill_graph, ) -from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -85,7 +86,13 @@ def forward( b: torch.Tensor, ) -> torch.Tensor: is_extend = forward_batch.forward_mode.is_extend() - if is_extend and get_tc_piecewise_forward_context() is not None: + if is_extend and ( + is_in_full_prefill_graph() + or ( + is_in_breakable_cuda_graph() + and forward_batch.forward_mode.is_extend_without_speculative() + ) + ): # Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim) seq_len = mixed_qkv.shape[0] output = torch.empty( @@ -93,22 +100,12 @@ def forward( dtype=mixed_qkv.dtype, device=mixed_qkv.device, ) - if is_in_breakable_cuda_graph() and not _linear_extend_in_graph(): - bcg_unified_linear_attention_with_output( - mixed_qkv, - a, - b, - output, - self.layer_id, + if is_in_breakable_cuda_graph() and _linear_extend_in_graph(): + _linear_attention_with_output_impl( + mixed_qkv, a, b, output, self, forward_batch ) else: - unified_linear_attention_with_output( - mixed_qkv, - a, - b, - output, - self.layer_id, - ) + self._eager_linear_attention(mixed_qkv, a, b, output) return output # Target verify rebuilds query_start_loc from the physical padded input, @@ -117,7 +114,7 @@ def forward( is_extend and not forward_batch.forward_mode.is_target_verify() ) real_num_tokens = ( - getattr(forward_batch, "global_num_token_non_padded_cpu", None) + forward_batch.global_num_token_non_padded_cpu if should_trim_padded_extend else None ) @@ -147,19 +144,13 @@ def forward( b=b, ) + def _capture_stub_linear_attention(self, mixed_qkv, a, b, output): + output.zero_() -def _linear_extend_in_graph() -> bool: - """Whether the backend runs this breakable prefill capture's linear extend - on static tables inside the graph (no eager break at the layer).""" - from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( - HybridLinearAttnBackend, - ) - - backend = get_attn_backend() - return ( - isinstance(backend, HybridLinearAttnBackend) - and backend.linear_extend_in_graph() - ) + @eager_on_graph(capture_stub=_capture_stub_linear_attention) + def _eager_linear_attention(self, mixed_qkv, a, b, output): + forward_batch = get_forward_batch() + _linear_attention_with_output_impl(mixed_qkv, a, b, output, self, forward_batch) def _linear_attention_with_output_impl( @@ -206,58 +197,15 @@ def _linear_attention_with_output_impl( output[:, real_num_tokens:].zero_() -def _unified_linear_attention_with_output_impl( - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - output: torch.Tensor, - layer_id: int, -) -> None: - """Eager implementation kept separate for backend-independent tests.""" - context = get_tc_piecewise_forward_context() - forward_batch = context.forward_batch - attention_layers = context.attention_layers - attention_layer = attention_layers[layer_id] - _linear_attention_with_output_impl( - mixed_qkv=mixed_qkv, - a=a, - b=b, - output=output, - attention_layer=attention_layer, - forward_batch=forward_batch, +def _linear_extend_in_graph() -> bool: + """Whether the backend runs this breakable prefill capture's linear extend + on static tables inside the graph (no eager break at the layer).""" + from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, ) - return - -@register_custom_op(mutates_args=["output"]) -@register_split_op() -def unified_linear_attention_with_output( - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - output: torch.Tensor, - layer_id: int, -) -> None: - """Custom op wrapper for linear attention computation only.""" - _unified_linear_attention_with_output_impl( - mixed_qkv=mixed_qkv, - a=a, - b=b, - output=output, - layer_id=layer_id, + backend = get_attn_backend() + return ( + isinstance(backend, HybridLinearAttnBackend) + and backend.linear_extend_in_graph() ) - - -def _linear_attention_capture_stub( - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - output: torch.Tensor, - layer_id: int, -) -> None: - output.zero_() - - -bcg_unified_linear_attention_with_output = eager_on_graph( - True, capture_stub=_linear_attention_capture_stub -)(unified_linear_attention_with_output) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 7a11965ca46d..4d2ef4d1b682 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -128,6 +128,7 @@ from sglang.srt.model_executor.runner_utils.pool import ( get_or_create_global_graph_capture_stream, ) +from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope from sglang.srt.model_loader.utils import resolve_language_model from sglang.srt.runtime_context import ( get_exec, @@ -780,6 +781,9 @@ def _prefill_forward_context( ): with ( set_forward_batch(forward_batch), + prefill_graph_scope( + full_graph=self._is_full_backend, raw_num_tokens=raw_num_tokens + ), forward_context( ForwardContext(attn_backend=self.model_runner.attn_backend) ), @@ -913,6 +917,7 @@ def run_dummy_multimodal_deepstack_forward( with ( set_forward_batch(fb), + prefill_graph_scope(full_graph=self._is_full_backend), forward_context( ForwardContext(attn_backend=self.model_runner.attn_backend) ), diff --git a/python/sglang/srt/model_executor/runner_utils/prefill_graph.py b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py new file mode 100644 index 000000000000..dfad7e50ad24 --- /dev/null +++ b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py @@ -0,0 +1,35 @@ +"""Runner-owned prefill graph policy, independent of the attention backend.""" + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + + +@dataclass(frozen=True) +class _PrefillGraphState: + full_graph: bool + raw_num_tokens: int | None + + +_current: ContextVar[_PrefillGraphState | None] = ContextVar( + "prefill_graph", default=None +) + + +def is_in_full_prefill_graph() -> bool: + state = _current.get() + return state is not None and state.full_graph + + +def get_prefill_raw_num_tokens() -> int | None: + state = _current.get() + return state.raw_num_tokens if state is not None else None + + +@contextmanager +def prefill_graph_scope(*, full_graph: bool, raw_num_tokens: int | None = None): + reset_token = _current.set(_PrefillGraphState(full_graph, raw_num_tokens)) + try: + yield + finally: + _current.reset(reset_token) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index a784b6de11c7..e53502706a13 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -11,11 +11,11 @@ per_tensor_quant_mla_fp8, per_token_group_quant_mla_deep_gemm_masked_fp8, ) -from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.dsa.utils import is_graph_dsa_split_op_surface from sglang.srt.layers.attention.dsa_backend import prepare_kv_for_attention +from sglang.srt.layers.attention.graph_utils import padded_extend_real_tokens from sglang.srt.layers.dcp import ( all_gather_kv_cache_for_mla_extend, all_gather_q_for_mla_decode, @@ -24,10 +24,6 @@ ) from sglang.srt.layers.layer_boundary import get_attn_tp_context from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run -from sglang.srt.layers.radix_attention import ( - padded_extend_real_tokens, - unified_attention_with_output, -) from sglang.srt.lora.deepseek_mla_correction import ( apply_q_correction as apply_kv_b_lora_q_correction, ) @@ -48,10 +44,6 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( is_in_breakable_cuda_graph, ) -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, - is_in_tc_piecewise_cuda_graph, -) from sglang.srt.models.deepseek_common.utils import ( FORWARD_ABSORB_CORE_ATTENTION_BACKENDS, _is_cpu, @@ -65,13 +57,16 @@ maybe_capture_indexer_topk, ) from sglang.srt.utils import BumpAllocator -from sglang.srt.utils.custom_op import register_custom_op logger = logging.getLogger(__name__) _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() _ENABLE_DSA_Q8KV8_BORN_FP8_Q = envs.SGLANG_ENABLE_DSA_Q8KV8_BORN_FP8_Q.get() _ENABLE_DSA_Q8KV8_QPREP_OVERLAP = envs.SGLANG_ENABLE_DSA_Q8KV8_QPREP_OVERLAP.get() +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + is_in_full_prefill_graph, +) + if TYPE_CHECKING: from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA @@ -162,9 +157,8 @@ def _can_fuse_bmm_into_attention( ) -> bool: if getattr(self, "_kimi_split_gguf_kv_b", False): return False - # Shared activation surface with the DSA indexer graph dispatch - # (in piecewise/breakable graph + non-speculative extend). Like the indexer - # dispatch, this fusion is on by default on that surface. + # Like the DSA indexer eager region, this fusion is enabled for + # non-speculative CUDA BCG prefill. if not is_graph_dsa_split_op_surface(forward_batch): return False if not self.use_dsa: @@ -256,11 +250,7 @@ def _q8kv8_born_fp8_q_backend( return None # Graph/compile surfaces run their own dispatch; the python-side # stash handshake is eager-only. - if is_graph_dsa_split_op_surface(forward_batch): - return None - if get_tc_piecewise_forward_context() is not None: - return None - if is_in_breakable_cuda_graph(): + if is_in_breakable_cuda_graph() or is_in_full_prefill_graph(): return None if get_is_capture_mode(): return None @@ -614,8 +604,6 @@ def forward_absorb_prepare( absorbed_bmm_concat_cast_q_fp8( q_fp8, q_nope, self.w_kc, q_pe, self.num_local_heads ) - # The attention call consumes the batch's real rows, a prefix of - # the rows written here. born_q_backend.q8kv8_stash_born_q( padded_extend_real_tokens(q_nope, forward_batch) or num_tokens, self.attn_mqa.layer_id, @@ -705,12 +693,7 @@ def forward_absorb_core( llama_4_scaling=llama_4_scaling, ) if fusion_plan is not None: - bmm_attention_fn = ( - bcg_mla_bmm_then_unified_attention - if is_in_breakable_cuda_graph() - else mla_bmm_then_unified_attention - ) - bmm_attention_fn( + self._eager_bmm_attention( fusion_plan.q_nope_t, self.w_kc, fusion_plan.q_nope_out_buf, @@ -718,7 +701,6 @@ def forward_absorb_core( k_nope, fusion_plan.attn_output_buf, save_kv_cache, - self.layer_id, q_pe, k_pe, cos_sin_cache=extra_args.get("cos_sin_cache"), @@ -890,26 +872,18 @@ def forward_absorb_core( ) attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) else: - if is_in_tc_piecewise_cuda_graph(): - # torch dynamo requires out= op was called where output tensor was non-contiguous - attn_bmm_output = ( - torch.bmm(attn_output.transpose(0, 1), self.w_vc) - .transpose(0, 1) - .flatten(1, 2) - ) - else: - attn_bmm_output = torch.empty( - (attn_output.shape[0], self.num_local_heads * self.v_head_dim), - dtype=attn_output.dtype, - device=attn_output.device, - ) - torch.bmm( - attn_output.transpose(0, 1), - self.w_vc, - out=attn_bmm_output.view( - -1, self.num_local_heads, self.v_head_dim - ).transpose(0, 1), - ) + attn_bmm_output = torch.empty( + (attn_output.shape[0], self.num_local_heads * self.v_head_dim), + dtype=attn_output.dtype, + device=attn_output.device, + ) + torch.bmm( + attn_output.transpose(0, 1), + self.w_vc, + out=attn_bmm_output.view( + -1, self.num_local_heads, self.v_head_dim + ).transpose(0, 1), + ) if _SGLANG_EXPERIMENTAL_LORA_OPTI: from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import ( kv_b_lora_v_apply, @@ -962,53 +936,35 @@ def _fuse_rope_for_trtllm_mla( and get_attn_backend().data_type == torch.float8_e4m3fn ) - -# Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output` -# into one eager split op under both PCG and BCG. Without this, the bf16 -# fallback BMM is captured alone in its own single-kernel CUDA graph submodule, -# paying per-submodule host overhead with no fusion benefit. -# -# `q_nope_out_view` aliases `q_nope_out_buf` (transposed). The op writes -# `q_nope_out_buf` via `torch.bmm(..., out=...)` and then reads through -# `q_nope_out_view`, so the alias's storage is mutated too. Declare it in -# `mutates_args` to keep the schema honest. -@register_custom_op( - mutates_args=["q_nope_out_buf", "q_nope_out_view", "attn_output_buf"] -) -@register_split_op() -def mla_bmm_then_unified_attention( - q_nope_t: torch.Tensor, - w_kc: torch.Tensor, - q_nope_out_buf: torch.Tensor, - q_nope_out_view: torch.Tensor, - k_nope: torch.Tensor, - attn_output_buf: torch.Tensor, - save_kv_cache: bool, - layer_id: int, - q_pe: torch.Tensor, - k_pe: torch.Tensor, - cos_sin_cache: Optional[torch.Tensor] = None, - is_neox: Optional[bool] = None, - llama_4_scaling: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, -) -> None: - torch.bmm(q_nope_t, w_kc, out=q_nope_out_buf) - unified_attention_with_output( - q_nope_out_view, - k_nope, - k_nope, - attn_output_buf, - save_kv_cache, - layer_id, - q_rope=q_pe, - k_rope=k_pe, - cos_sin_cache=cos_sin_cache, - is_neox=is_neox, - llama_4_scaling=llama_4_scaling, - topk_indices=topk_indices, - ) - - -bcg_mla_bmm_then_unified_attention = eager_on_graph(True)( - mla_bmm_then_unified_attention -) + @eager_on_graph + def _eager_bmm_attention( + self, + q_nope_t: torch.Tensor, + w_kc: torch.Tensor, + q_nope_out_buf: torch.Tensor, + q_nope_out_view: torch.Tensor, + k_nope: torch.Tensor, + attn_output_buf: torch.Tensor, + save_kv_cache: bool, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = None, + llama_4_scaling: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + ) -> None: + torch.bmm(q_nope_t, w_kc, out=q_nope_out_buf) + self.attn_mqa._eager_attention.__wrapped__( + self.attn_mqa, + q_nope_out_view, + k_nope, + k_nope, + attn_output_buf, + save_kv_cache, + q_rope=q_pe, + k_rope=k_pe, + cos_sin_cache=cos_sin_cache, + is_neox=is_neox, + llama_4_scaling=llama_4_scaling, + topk_indices=topk_indices, + ) diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py index 63944d61975d..0e04d3aa345d 100644 --- a/python/sglang/srt/models/inkling.py +++ b/python/sglang/srt/models/inkling.py @@ -4,6 +4,7 @@ import logging import re from array import array +from contextlib import contextmanager from typing import Iterable, Optional, Set, Tuple import torch @@ -21,7 +22,6 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import get_moe_runner_backend from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.radix_attention import force_eager_attention from sglang.srt.layers.utils import get_layer_id from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -41,9 +41,7 @@ eager_on_graph, is_in_breakable_cuda_graph, ) -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, -) +from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.inkling_common.attn import ( InklingAttention, @@ -143,6 +141,16 @@ def _is_unsupported_mm_weight_name(name: str) -> bool: ) +@contextmanager +def force_eager_attention(layer): + previous = layer.use_prefill_attention_wrapper + layer.use_prefill_attention_wrapper = False + try: + yield + finally: + layer.use_prefill_attention_wrapper = previous + + class InklingDecoderLayer(nn.Module): def __init__( self, @@ -250,13 +258,11 @@ def __init__( # Under BCG the short-conv metadata (cu_seqlens/seq_idx) is baked at bs=1 # during capture, which is wrong for multi-seq prefill. Running every # sconv (and the attn whose k/v_sconv it wraps) eagerly makes them re-read - # the LIVE per-seq metadata at replay. `_breakable_attn_group` groups the + # the LIVE per-seq metadata at replay. `_eager_attn_group` groups the # prior layer's (deferred) mlp_sconv + attn_norm + attn + attn_sconv into # ONE eager break; only mlp_norm + MoE stay captured. Outside a capture - # these wrappers just run inline. `_breakable_mlp_sconv` runs the final + # these wrappers just run inline. `_eager_mlp_sconv` runs the final # layer's deferred mlp_sconv after the layer loop. - self._breakable_attn_group = eager_on_graph(True)(self._attn_group_impl) - self._breakable_mlp_sconv = eager_on_graph(True)(self._mlp_sconv_impl) def _attn_block( self, @@ -352,11 +358,13 @@ def _attn_block( # then restore the full buffer on forward_batch (shared across layers/replays). orig_out_cache_loc = forward_batch.out_cache_loc forward_batch.out_cache_loc = orig_out_cache_loc[: hs.shape[0]] - with force_eager_attention(): - hs = self.attn( - hs, positions, forward_batch, log_scaling_tau=log_scaling_tau - ) - forward_batch.out_cache_loc = orig_out_cache_loc + try: + with force_eager_attention(self.attn.attn): + hs = self.attn( + hs, positions, forward_batch, log_scaling_tau=log_scaling_tau + ) + finally: + forward_batch.out_cache_loc = orig_out_cache_loc else: hs = self.attn( hs, @@ -378,7 +386,8 @@ def _attn_block( hs = all_gather_hidden(hs, self.attn_tp_group) return hs, res - def _attn_group_impl( + @eager_on_graph + def _eager_attn_group( self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], @@ -390,9 +399,8 @@ def _attn_group_impl( ) -> None: """Eager break: run `_attn_block` on the REAL (non-padded) tokens with the LIVE forward_batch and write the result into the padded output buffers. Mutates - attn_out / residual_out and returns None (the eager_on_graph copy-back is - per-tensor, not per-tuple, so outputs must be pre-allocated buffers).""" - forward_batch = get_tc_piecewise_forward_context().forward_batch + attn_out / residual_out and returns None (the following graph segment reads these stable output buffers).""" + forward_batch = get_forward_batch() n = forward_batch.global_num_token_non_padded_cpu # log_scaling_tau is per-token, so narrow it to match the real tokens too. hs, res = self._attn_block( @@ -408,7 +416,8 @@ def _attn_group_impl( if attn_out.shape[0] != n: torch._foreach_zero_((attn_out[n:], residual_out[n:])) - def _mlp_sconv_impl( + @eager_on_graph + def _eager_mlp_sconv( self, hidden_states: torch.Tensor, positions: torch.Tensor, @@ -416,7 +425,7 @@ def _mlp_sconv_impl( ) -> None: """Eager break for the final layer's deferred mlp_sconv: run on the real tokens with the live forward_batch, write the padded output buffer.""" - forward_batch = get_tc_piecewise_forward_context().forward_batch + forward_batch = get_forward_batch() n = forward_batch.global_num_token_non_padded_cpu y = self.mlp_sconv(hidden_states[:n], positions[:n], forward_batch) if self.scattered_sconv: @@ -454,30 +463,24 @@ def forward( if forward_batch.forward_mode.is_idle(): return hidden_states, residual - # The eager group reads the LIVE forward_batch from the tc_piecewise context - # (the only hook evaluated at BCG replay time — the break's captured args are - # frozen at capture-bucket shapes). Only the prefill BCG runner installs that - # context; the decode breakable backend sets is_in_breakable_cuda_graph() but - # NOT the context, so gate on both and otherwise fall through to the inline - # path below (which uses the passed forward_batch — correct for decode). + # The eager group reads the live prepared batch from the runner scope. if ( is_in_breakable_cuda_graph() - and get_tc_piecewise_forward_context() is not None + and forward_batch.forward_mode.is_extend_without_speculative() ): # BCG prefill path: the AR fusion is decode-only, so partials never # reach (or leave) this branch. assert not prev_mlp_partial and not fuse_ar_sconv and not fuse_attn_ar # BCG: {prev mlp_sconv, attn_norm, attn, attn_sconv} run eagerly (one - # break under capture); mlp_norm + MoE stay captured. (The live - # forward_batch inside the break is read from the shared tc_piecewise - # context, which the prefill BCG runner populates at capture and replay.) + # break under capture); mlp_norm + MoE stay captured. The eager method + # reads the live prepared forward_batch from the runner scope. # Under scattered sconv the group's INPUT can be the previous layer's # [T, H/P] MoE shard while its OUTPUT is post-all-gather [T, H], so # size the output buffers explicitly (residual is always [T, H]). out_shape = (hidden_states.shape[0], self.attn_norm.weight.shape[0]) attn_out = hidden_states.new_empty(out_shape) residual_out = hidden_states.new_empty(out_shape) - self._breakable_attn_group( + self._eager_attn_group( hidden_states, residual, positions, @@ -900,13 +903,11 @@ def forward( if self._dflash_layers_to_capture else hidden_states ) - # Same gate as the per-layer group: the eager break needs the tc_piecewise - # context (installed only by the prefill BCG runner) to read the live - # forward_batch at replay; else run inline with the passed forward_batch. + # Match the per-layer prefill eager region. scattered = self.layers[-1].scattered_sconv if ( is_in_breakable_cuda_graph() - and get_tc_piecewise_forward_context() is not None + and forward_batch.forward_mode.is_extend_without_speculative() ): # Under scattered sconv the input is the last MoE's [T, H/P] # shard; the break's output buffer is post-all-gather [T, H]. @@ -916,8 +917,10 @@ def forward( else hidden_states.shape ) mlp_sconv_out = hidden_states.new_empty(out_shape) - self.layers[-1]._breakable_mlp_sconv( - hidden_states, positions, mlp_sconv_out + self.layers[-1]._eager_mlp_sconv( + hidden_states, + positions, + mlp_sconv_out, ) hidden_states = mlp_sconv_out else: diff --git a/python/sglang/srt/models/inkling_common/kernels/comm.py b/python/sglang/srt/models/inkling_common/kernels/comm.py index fced9b720ed0..1c094fe6da8e 100644 --- a/python/sglang/srt/models/inkling_common/kernels/comm.py +++ b/python/sglang/srt/models/inkling_common/kernels/comm.py @@ -7,6 +7,9 @@ import torch from sglang.srt.environ import envs +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + is_in_breakable_cuda_graph, +) from sglang.srt.runtime_context import get_exec from sglang.srt.utils import is_cuda @@ -710,10 +713,9 @@ def scattered_ar_sconv_fusable( if not (fm.is_extend() or fm.is_decode()): return False # Prefill scope: the BCG runner's eager-break sites are not wired (its - # baked flags would disagree with the break bodies) and tc_piecewise's - # FX pieces can't carry the cross-layer producer/consumer contract, so - # both fall back to the unfused chain. The FULL prefill CUDA-graph - # backend (context.full_graph -- the whole model captured uniformly) IS + # baked flags would disagree with the break bodies), so it falls back + # to the unfused chain. The FULL prefill CUDA-graph + # backend (the whole model captured uniformly) IS # supported: the kernel is capture-safe (barrier epochs advance across # replays; validated capture+replay) and all its metadata (qsl/si/ # cache_mask/safe_idx/track rows) is recomputed in-graph from the @@ -721,12 +723,8 @@ def scattered_ar_sconv_fusable( # rows only write pad rows of the OUT region (the eager tail slices # [:raw]), and sentinel request slots have qlen == 0 so the in-kernel # cache update/track skip them. - from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, - ) - tc_ctx = get_tc_piecewise_forward_context() - if tc_ctx is not None and not tc_ctx.full_graph: + if is_in_breakable_cuda_graph() and fm.is_extend_without_speculative(): return False comm = group.torch_symm_mem_comm if ( @@ -1050,15 +1048,11 @@ def fullwidth_ar_sconv_fusable( return False if num_tokens < _INKLING_AR_FW_MIN_TOKENS: return False - # Same prefill-runner scope as the scattered gate: BCG / tc_piecewise + # Same prefill-runner scope as the scattered gate: BCG # pieces can't carry the cross-layer producer contract; the FULL prefill # CUDA-graph backend is supported (capture-safe kernel, in-graph metadata). - from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, - ) - tc_ctx = get_tc_piecewise_forward_context() - if tc_ctx is not None and not tc_ctx.full_graph: + if is_in_breakable_cuda_graph(): return False comm = group.torch_symm_mem_comm if ( diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index b44147dc05c2..7283e09b2977 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py @@ -241,22 +241,23 @@ def scale(src): ) def test_attention_narrows_padded_positions(self): - from sglang.srt.layers.radix_attention import unified_attention_with_output + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.runner_utils.forward_batch import ( + set_forward_batch, + ) + from sglang.srt.model_executor.runner_utils.prefill_graph import ( + prefill_graph_scope, + ) num_tokens = 3 padded_num_tokens = 5 forward_batch = SimpleNamespace( global_num_token_non_padded_cpu=num_tokens, + mha_return_lse=False, + _attn_output=None, out_cache_loc=torch.arange(padded_num_tokens, device=self.device), positions=torch.arange(padded_num_tokens, device=self.device), ) - context = SimpleNamespace( - forward_batch=forward_batch, - attention_layers=[object()], - mha_companion_layers=None, - num_tokens=padded_num_tokens, - raw_num_tokens=num_tokens, - ) observed = {} def attention_forward(query, key, value, layer, batch, save_kv_cache): @@ -266,22 +267,19 @@ def attention_forward(query, key, value, layer, batch, save_kv_cache): output = torch.full((padded_num_tokens, 2), float("nan"), device=self.device) with ( - patch( - "sglang.srt.layers.radix_attention.get_tc_piecewise_forward_context", - return_value=context, - ), + set_forward_batch(forward_batch), + prefill_graph_scope(full_graph=False, raw_num_tokens=num_tokens), patch( "sglang.srt.layers.radix_attention.get_attn_backend", return_value=SimpleNamespace(forward=attention_forward), ), ): - unified_attention_with_output( + RadixAttention(1, 2, 1.0, 1, 0)._eager_attention( torch.zeros((padded_num_tokens, 2), device=self.device), torch.zeros((padded_num_tokens, 1, 2), device=self.device), torch.zeros((padded_num_tokens, 1, 2), device=self.device), output, True, - 0, ) expected = torch.arange(num_tokens, device=self.device) diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index a632cd95de74..1e290a4c09f9 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/test/registered/unit/layers/test_radix_attention.py @@ -1,643 +1,274 @@ -"""CPU unit tests for the graph-safe ``RadixAttention`` interface.""" +"""CPU contracts shared by the full and breakable prefill attention paths.""" -import unittest -from contextlib import ExitStack +import sys +from contextlib import contextmanager from types import SimpleNamespace from unittest.mock import patch +import pytest import torch -import sglang.srt.layers.radix_attention as radix_attention_module from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context -from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( - get_tc_piecewise_forward_context, - set_tc_piecewise_forward_context, -) +from sglang.srt.model_executor.runner_utils.forward_batch import set_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=11, suite="base-a-test-cpu") +register_cpu_ci(est_time=5, suite="base-a-test-cpu") -class _RecordingAttentionBackend: - def __init__(self, *, return_lse=True): - self.calls = [] - self.return_lse = return_lse - def forward( - self, - query, - key, - value, - attention_layer, - forward_batch, - save_kv_cache, - **kwargs, +@contextmanager +def attention_scope(batch, backend, *, full_graph=False, raw_num_tokens=None): + with ( + forward_context(ForwardContext(backend)), + set_forward_batch(batch), + prefill_graph_scope(full_graph=full_graph, raw_num_tokens=raw_num_tokens), ): - self.calls.append( - SimpleNamespace( - query=query, - key=key, - value=value, - attention_layer=attention_layer, - forward_batch=forward_batch, - output=forward_batch._attn_output, - out_cache_loc=forward_batch.out_cache_loc.clone(), - save_kv_cache=save_kv_cache, - kwargs=kwargs, - ) - ) - output = torch.full_like(query, 3) - lse = torch.full((query.shape[0], query.shape[1]), 7, dtype=torch.float32) - return (output, lse) if self.return_lse else output - - -class TestRadixAttentionGraphInterface(CustomTestCase): - @staticmethod - def _new_layer() -> RadixAttention: - layer = RadixAttention( - num_heads=2, - head_dim=3, - scaling=1.0, - num_kv_heads=2, - layer_id=0, - ) - return layer - - @staticmethod - def _new_impl_context( - attention_layers, - *, - mha_companion_layers=None, - num_tokens=4, - real_num_tokens=2, + yield + + +def make_batch(tokens=2, rows=4, return_lse=False): + return SimpleNamespace( + forward_mode=SimpleNamespace( + is_extend=lambda: True, is_target_verify=lambda: False + ), + global_num_token_non_padded_cpu=tokens, + out_cache_loc=torch.arange(rows), + positions=torch.arange(rows), + _attn_output=None, + mha_return_lse=return_lse, + ) + + +@pytest.mark.parametrize("breakable", [False, True]) +@pytest.mark.parametrize("return_lse", [False, True]) +@pytest.mark.parametrize("tokens", [0, 2, 4]) +def test_dense_outputs_padding_and_lse(breakable, return_lse, tokens): + layer = RadixAttention(2, 3, 1.0, 2, 7) + batch = make_batch(tokens=tokens, return_lse=return_lse) + original = batch.out_cache_loc, batch.positions + calls = [] + + def attention(q, k, v, actual_layer, live_batch, save_kv_cache, **kwargs): + assert actual_layer is layer + assert live_batch is batch + assert q.shape[0] == tokens + assert k.shape[0] == v.shape[0] == 3 + assert batch.positions.shape[0] == batch.out_cache_loc.shape[0] == tokens + calls.append(q) + out = torch.full_like(q, 3) + if return_lse: + return out, torch.full((tokens, 2), 7.0) + return out + + q = torch.zeros(4, 2, 3) + with ( + attention_scope( + batch, + SimpleNamespace(forward=attention), + full_graph=not breakable, + raw_num_tokens=tokens, + ), + patch( + "sglang.srt.layers.radix_attention.is_in_breakable_cuda_graph", + return_value=breakable, + ), ): - forward_batch = SimpleNamespace( - global_num_token_non_padded_cpu=real_num_tokens, - out_cache_loc=torch.arange(num_tokens, dtype=torch.int64), - positions=torch.arange(num_tokens, dtype=torch.int64), - _attn_output=None, - mha_return_lse=False, - ) - return SimpleNamespace( - forward_batch=forward_batch, - attention_layers=attention_layers, - mha_companion_layers=mha_companion_layers, - num_tokens=None, - raw_num_tokens=None, - ) - - def test_forward_dispatches_all_graph_and_lse_variants(self): - layer = self._new_layer() - query = torch.zeros((4, 2, 3)) - key = torch.zeros_like(query) - value = torch.zeros_like(query) - - op_names = { - (False, False): "unified_attention_with_output", - (False, True): "unified_attention_with_output_and_lse", - (True, False): "breakable_unified_attention_with_output", - (True, True): "breakable_unified_attention_with_output_and_lse", - } - - for breakable in (False, True): - for return_lse in (False, True): - with self.subTest(breakable=breakable, return_lse=return_lse): - forward_batch = SimpleNamespace( - forward_mode=ForwardMode.EXTEND, - mha_return_lse=return_lse, - ) - calls = [] - - def output_only(*args, **kwargs): - args[3].fill_(5) - calls.append(kwargs) - - def output_and_lse(*args, **kwargs): - args[3].fill_(5) - calls.append(kwargs) - return torch.full((4, 2), 11, dtype=torch.float32) - - with ExitStack() as stack: - stack.enter_context( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=SimpleNamespace( - mha_companion_layers=[layer] - ), - ) - ) - stack.enter_context( - patch.object( - radix_attention_module, - "is_in_breakable_cuda_graph", - return_value=breakable, - ) - ) - mocks = { - name: stack.enter_context( - patch.object( - radix_attention_module, - name, - side_effect=( - output_and_lse - if name.endswith("and_lse") - else output_only - ), - ) - ) - for name in op_names.values() - } - result = layer( - query, - key, - value, - forward_batch, - key_value_num_tokens=3, - ) - - selected_name = op_names[(breakable, return_lse)] - for name, mock in mocks.items(): - self.assertEqual(mock.call_count, int(name == selected_name)) - - self.assertEqual( - calls, - [ - { - "use_mha_companion": True, - "key_value_num_tokens": 3, - } - ], - ) - if return_lse: - output, lse = result - self.assertEqual(lse.shape, (4, 2)) - self.assertTrue(torch.all(lse == 11)) - else: - output = result - self.assertEqual(output.shape, query.shape) - self.assertTrue(torch.all(output == 5)) - - def test_sparse_attention_breaks_the_graph_under_breakable(self): - """Breakable replays captured segments with the capture batch's metadata - baked in, so sparse attention must run as an eager break. Called inline, - MiniMax-M3 replayed stale metadata and hit an illegal memory access.""" - layer = self._new_layer() - query = torch.zeros((4, 2, 3)) - idx_q = torch.zeros((4, 1, 3)) - op_names = { - False: "unified_sparse_attention_with_output", - True: "breakable_unified_sparse_attention_with_output", - } - - def fill_outputs(*args, **kwargs): - args[3].fill_(5) - args[4].fill_(7) - - for breakable in (False, True): - with self.subTest(breakable=breakable): - forward_batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) - with ExitStack() as stack: - stack.enter_context( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=SimpleNamespace(), - ) - ) - stack.enter_context( - patch.object( - radix_attention_module, - "is_in_breakable_cuda_graph", - return_value=breakable, - ) - ) - backend = stack.enter_context( - patch.object(radix_attention_module, "get_attn_backend") - ) - backend.return_value.forward.return_value = ( - torch.zeros((4, 3)), - torch.zeros((4, 6)), - ) - mocks = { - name: stack.enter_context( - patch.object( - radix_attention_module, name, side_effect=fill_outputs - ) - ) - for name in op_names.values() - } - idx_out, attn_out = layer( - query, query, query, forward_batch, idx_q=idx_q, idx_k=query - ) - - backend.assert_not_called() - for name, mock in mocks.items(): - self.assertEqual(mock.call_count, int(name == op_names[breakable])) - self.assertTrue(torch.all(attn_out == 5)) - self.assertTrue(torch.all(idx_out == 7)) - - def test_prefill_wrapper_opt_out_preserves_expanded_rows_and_batch(self): - """Expanded attention rows must not be sliced using the runner's token count.""" - layer = RadixAttention( - num_heads=2, - head_dim=3, - scaling=1.0, - num_kv_heads=2, - layer_id=0, - use_prefill_attention_wrapper=False, - ) - runner_batch = self._new_impl_context([layer], num_tokens=2).forward_batch - expanded_batch = SimpleNamespace( - forward_mode=ForwardMode.EXTEND, - global_num_token_non_padded_cpu=None, - out_cache_loc=runner_batch.out_cache_loc, - _attn_output=None, - mha_return_lse=False, - ) - query = torch.arange(48, dtype=torch.float32).view(8, 2, 3) - key = query.flatten(1).clone() - value = -key - backend = _RecordingAttentionBackend(return_lse=False) - - with ( - forward_context(ForwardContext(backend)), - set_tc_piecewise_forward_context( - runner_batch, [layer], None, [], [], num_tokens=8, full_graph=True - ), - ): - output = layer( - query, key, value, expanded_batch, save_kv_cache=False, causal=False - ) - self.assertIs( - get_tc_piecewise_forward_context().forward_batch, runner_batch - ) - - self.assertTrue(torch.equal(output, torch.full_like(query, 3))) - self.assertEqual(len(backend.calls), 1) - call = backend.calls[0] - self.assertIs(call.query, query) - for actual, original in ((call.key, key), (call.value, value)): - self.assertEqual(actual.shape, (8, 2, 3)) - self.assertEqual(actual.data_ptr(), original.data_ptr()) - self.assertTrue(torch.equal(actual.flatten(), original.flatten())) - self.assertIs(call.attention_layer, layer) - self.assertIs(call.forward_batch, expanded_batch) - self.assertIs(expanded_batch.out_cache_loc, runner_batch.out_cache_loc) - self.assertFalse(call.save_kv_cache) - self.assertEqual(call.kwargs, {"causal": False}) - self.assertEqual(runner_batch.global_num_token_non_padded_cpu, 2) - - def test_deferred_norm_rope_operands_follow_real_tokens_on_each_call(self): - layer = self._new_layer() - query = torch.zeros((4, 2, 3)) - positions = torch.arange(4) - temp_scale = torch.arange(4, dtype=torch.float32).reshape(4, 1) - norm_weight = torch.ones(3) - operands = { - "mxfp8_norm_rope_positions": positions, - "mxfp8_norm_rope_temp_scale": temp_scale, - "norm_weight": norm_weight, - } - for breakable in (False, True): - with self.subTest(breakable=breakable): - context = self._new_impl_context([layer]) - forward_batch = context.forward_batch - forward_batch.forward_mode = ForwardMode.EXTEND - original_cache_loc = forward_batch.out_cache_loc - backend = _RecordingAttentionBackend(return_lse=False) - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, - "get_attn_backend", - return_value=backend, - ), - patch.object( - radix_attention_module, - "is_in_breakable_cuda_graph", - return_value=breakable, - ), - patch.object( - radix_attention_module, - "breakable_attention_with_output_extra_kwargs", - side_effect=radix_attention_module.attention_with_output_extra_kwargs, - ) as graph_break, - ): - for real_tokens in (2, 4): - forward_batch.global_num_token_non_padded_cpu = real_tokens - positions.add_(10) - result = layer(query, query, query, forward_batch, **operands) - call = backend.calls[-1] - self.assertEqual(call.query.shape[0], real_tokens) - self.assertTrue( - torch.equal( - call.kwargs["mxfp8_norm_rope_positions"], - positions[:real_tokens], - ) - ) - self.assertTrue( - torch.equal( - call.kwargs["mxfp8_norm_rope_temp_scale"], - temp_scale[:real_tokens], - ) - ) - self.assertIs(call.kwargs["norm_weight"], norm_weight) - self.assertIs(forward_batch.out_cache_loc, original_cache_loc) - self.assertTrue(torch.all(result[:real_tokens] == 3)) - self.assertEqual(positions.shape[0], 4) - self.assertEqual(temp_scale.shape[0], 4) - self.assertEqual(graph_break.call_count, 2 if breakable else 0) - - def test_impl_preserves_attention_identity_and_lse(self): - mqa = SimpleNamespace() - mha = SimpleNamespace() - context = self._new_impl_context([mqa], mha_companion_layers=[mha]) - forward_batch = context.forward_batch - original_out_cache_loc = forward_batch.out_cache_loc - backend = _RecordingAttentionBackend() - query = torch.zeros((4, 2, 3)) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - for use_mha_companion, expected_layer in ((False, mqa), (True, mha)): - with self.subTest(use_mha_companion=use_mha_companion): - output = torch.empty_like(query) - lse = radix_attention_module._unified_attention_with_output_impl( - query, - query, - query, - output, - False, - 0, - use_mha_companion, - True, - ) - - call_record = backend.calls[-1] - self.assertIs(call_record.attention_layer, expected_layer) - self.assertEqual(call_record.query.shape, (2, 2, 3)) - self.assertEqual(call_record.key.shape, (2, 2, 3)) - self.assertEqual(call_record.value.shape, (2, 2, 3)) - self.assertEqual(call_record.output.shape, (2, 2, 3)) - self.assertEqual(call_record.out_cache_loc.tolist(), [0, 1]) - self.assertFalse(call_record.save_kv_cache) - self.assertTrue(torch.all(output[:2] == 3)) - self.assertEqual(lse.shape, (4, 2)) - self.assertTrue(torch.all(lse[:2] == 7)) - self.assertTrue(torch.all(lse[2:] == 0)) - self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) - - def test_extra_kwargs_path_returns_bucket_shaped_lse(self): - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer]) - backend = _RecordingAttentionBackend() - query = torch.zeros((4, 2, 3)) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - lse = radix_attention_module.attention_with_output_extra_kwargs( - query, - query, - query, - torch.empty_like(query), - False, - 0, - {"return_lse": True}, - ) - - self.assertEqual(lse.tolist(), [[7, 7], [7, 7], [0, 0], [0, 0]]) - - def test_impl_uses_independent_query_and_key_value_extents(self): - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer]) - forward_batch = context.forward_batch - original_out_cache_loc = forward_batch.out_cache_loc - backend = _RecordingAttentionBackend() - query = torch.zeros((4, 2, 3)) - key = torch.zeros((6, 2, 3)) - value = torch.zeros((6, 2, 3)) - k_rope = torch.zeros((6, 2, 1)) - output = torch.empty_like(query) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - lse = radix_attention_module._unified_attention_with_output_impl( - query, - key, - value, - output, - False, - 0, - False, - True, - key_value_num_tokens=5, - k_rope=k_rope, - ) - - call_record = backend.calls[-1] - self.assertEqual(call_record.query.shape, (2, 2, 3)) - self.assertEqual(call_record.key.shape, (5, 2, 3)) - self.assertEqual(call_record.value.shape, (5, 2, 3)) - self.assertEqual(call_record.kwargs["k_rope"].shape, (5, 2, 1)) - self.assertEqual(call_record.output.shape, (2, 2, 3)) - self.assertEqual(lse.shape, (4, 2)) - self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) - - def test_impl_preserves_output_only_contract(self): - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer]) - forward_batch = context.forward_batch - original_out_cache_loc = forward_batch.out_cache_loc - backend = _RecordingAttentionBackend(return_lse=False) - query = torch.zeros((4, 2, 3)) - output = torch.empty_like(query) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - lse = radix_attention_module._unified_attention_with_output_impl( - query, - query, - query, - output, - False, - 0, - False, - False, - ) - - self.assertIsNone(lse) - self.assertIs(backend.calls[-1].attention_layer, attention_layer) - self.assertTrue(torch.all(output[:2] == 3)) - self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc) - - def test_impl_zero_real_tokens_returns_zeroed_lse(self): - # Regression: an idle DP rank whose fabricated EXTEND batch is masked to - # 0 real tokens skips attention entirely. The skip must still honor the - # LSE return mode -- unified_attention_with_output_and_lse asserts a - # tensor comes back, so returning a bare None raised AssertionError as - # soon as any 0-real-token call needed LSE (chunked-prefix MHA merge). - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer], real_num_tokens=0) - backend = _RecordingAttentionBackend() - query = torch.zeros((4, 2, 3)) - output = torch.full_like(query, float("nan")) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - lse = radix_attention_module._unified_attention_with_output_impl( - query, - query, - query, - output, - False, - 0, - False, - True, - ) - - self.assertEqual(backend.calls, []) - self.assertTrue(torch.all(output == 0)) - # Same shape/dtype the registered fake impl declares, so - # unified_attention_with_output_and_lse's `assert lse is not None` holds. - self.assertEqual(lse.shape, (4, 2)) - self.assertEqual(lse.dtype, torch.float32) - self.assertTrue(torch.all(lse == 0)) - - def test_impl_zero_real_tokens_output_only_returns_none(self): - # The 0-token skip must not start returning a tensor on the non-LSE - # path: unified_attention_with_output is registered with an inplace - # (None-returning) schema. - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer], real_num_tokens=0) - backend = _RecordingAttentionBackend(return_lse=False) - query = torch.zeros((4, 2, 3)) - output = torch.full_like(query, float("nan")) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - lse = radix_attention_module._unified_attention_with_output_impl( - query, - query, - query, - output, - False, - 0, - False, - False, + result = layer(q, q, q, batch, key_value_num_tokens=3) + output, lse = result if return_lse else (result, None) + torch.testing.assert_close(output[:tokens], torch.full_like(output[:tokens], 3)) + assert torch.count_nonzero(output[tokens:]) == 0 + assert len(calls) == bool(tokens) + assert batch.out_cache_loc is original[0] and batch.positions is original[1] + assert batch._attn_output is None + if return_lse: + assert lse.shape == (4, 2) + assert torch.all(lse[:tokens] == 7) + assert torch.count_nonzero(lse[tokens:]) == 0 + + +def test_extra_kwargs_and_exception_restore(): + layer = RadixAttention(2, 3, 1.0, 2, 0) + batch = make_batch() + sentinel = torch.empty(1) + batch._attn_output = sentinel + original = batch.out_cache_loc, batch.positions + modifier = lambda score: score + q = torch.zeros(4, 2, 3) + aux = [torch.arange(4)] + + def attention(q, k, v, actual_layer, live_batch, save_kv_cache, **kwargs): + assert kwargs["score_mod"] is modifier + assert kwargs["aux_tensors"][0].shape == (2,) + assert kwargs["q_rope"].shape[0] == 2 + assert kwargs["k_rope"].shape[0] == 3 + assert kwargs["q_descale"].shape[0] == 2 + raise RuntimeError("backend failed") + + with attention_scope( + batch, SimpleNamespace(forward=attention), full_graph=True, raw_num_tokens=2 + ): + with pytest.raises(RuntimeError, match="backend failed"): + layer( + q, + q, + q, + batch, + key_value_num_tokens=3, + score_mod=modifier, + aux_tensors=aux, + q_rope=q, + k_rope=q, + q_descale=q, ) + assert batch.out_cache_loc is original[0] and batch.positions is original[1] + assert batch._attn_output is sentinel + assert aux[0].shape == (4,) + + +@pytest.mark.parametrize("tokens", [0, 2]) +@pytest.mark.parametrize("has_index_value", [False, True]) +def test_sparse_full_graph_two_outputs(tokens, has_index_value): + layer = RadixAttention(2, 3, 1.0, 2, 0) + batch = make_batch(tokens=tokens) + q = torch.zeros(4, 2, 3) + idx = torch.zeros(4, 1, 2) + + def attention(q, k, v, layer, batch, save_kv_cache, **kwargs): + assert kwargs["idx_q"].shape == (tokens, 1, 2) + assert kwargs["idx_k"].shape == (tokens, 1, 2) + return ( + torch.full((tokens, 1, 2), 5.0) if has_index_value else None, + torch.ones_like(q), + ) - self.assertIsNone(lse) - self.assertEqual(backend.calls, []) - self.assertTrue(torch.all(output == 0)) - - def test_extra_kwargs_zero_real_tokens_zeroes_output(self): - # Regression: attention_with_output_extra_kwargs (Inkling score_mod / - # aux_tensors) narrowed to query[:0] and copied output[:0], so with 0 - # real tokens the preallocated torch.empty output was never written and - # its garbage (NaN/Inf) flowed into residuals and MoE routing. Only ROCm - # zeroed the padded tail, so every other platform leaked it. - attention_layer = SimpleNamespace() - context = self._new_impl_context([attention_layer], real_num_tokens=0) - backend = _RecordingAttentionBackend(return_lse=False) - query = torch.zeros((4, 2, 3)) - output = torch.full_like(query, float("nan")) - - with ( - patch.object( - radix_attention_module, - "get_tc_piecewise_forward_context", - return_value=context, - ), - patch.object( - radix_attention_module, "get_attn_backend", return_value=backend - ), - ): - radix_attention_module.attention_with_output_extra_kwargs( - query, - query, - query, - output, - False, - 0, - {}, - ) + with attention_scope( + batch, + SimpleNamespace(forward=attention), + full_graph=True, + raw_num_tokens=tokens, + ): + index_output, output = layer(q, q, q, batch, idx_q=idx, idx_k=idx) + assert output.shape == (4, 6) and index_output.shape == (4, 2) + assert torch.all(output[:tokens] == 1) + assert torch.count_nonzero(output[tokens:]) == 0 + if has_index_value: + assert torch.all(index_output[:tokens] == 5) + assert torch.count_nonzero(index_output[tokens:]) == 0 + + +@pytest.mark.parametrize("sparse", [False, True]) +def test_direct_backend_dispatch(sparse): + layer = RadixAttention(2, 3, 1.0, 2, 0) + batch = make_batch() + q = torch.zeros(4, 2, 3) + result = object() + backend = SimpleNamespace(forward=lambda *args, **kwargs: result) + with ( + attention_scope(batch, backend), + patch( + "sglang.srt.layers.radix_attention.is_in_breakable_cuda_graph", + return_value=False, + ), + ): + kwargs = {"idx_q": q, "idx_k": q} if sparse else {} + assert layer(q, q, q, batch, **kwargs) is result - self.assertEqual(backend.calls, []) - self.assertTrue(torch.all(output == 0)) - def test_lse_fake_impl_declares_shape_and_dtype(self): - query = torch.empty((5, 3, 7), dtype=torch.float16) - output = torch.empty_like(query) +def test_output_dtype_follows_values(): + layer = RadixAttention(2, 3, 1.0, 2, 0) + batch = make_batch() + q = torch.zeros(4, 2, 3, dtype=torch.float32) + v = q.to(torch.bfloat16) + backend = SimpleNamespace( + forward=lambda q, k, v, *args, **kwargs: torch.ones_like(v) + ) + with attention_scope(batch, backend, full_graph=True, raw_num_tokens=2): + assert layer(q, q, v, batch).dtype == torch.bfloat16 - lse = radix_attention_module._unified_attention_with_output_and_lse_fake( - query, - None, - None, - output, - False, - 0, - ) - self.assertEqual(lse.shape, (5, 3)) - self.assertEqual(lse.dtype, torch.float32) - self.assertEqual(lse.device, query.device) +@pytest.mark.parametrize( + "return_lse,sparse", [(False, False), (True, False), (False, True)] +) +def test_attention_replay_uses_live_batch_and_static_output_buffers(return_lse, sparse): + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + breakable_cuda_graph as bcg, + ) + + layer = RadixAttention(2, 3, 1.0, 2, 0) + q = torch.zeros(4, 2, 3) + captured_batch = make_batch(return_lse=return_lse) + live_batch = make_batch(tokens=1, return_lse=return_lse) + seen = [] + + def attention(query, key, value, actual_layer, batch, save_kv_cache, **kwargs): + seen.append(batch) + out = torch.full_like(query, float(len(seen))) + lse = torch.full((query.shape[0], 2), float(len(seen))) + if sparse: + return torch.full_like(query, float(len(seen) + 2)), out + return (out, lse) if return_lse else out + + graph = bcg.BreakableCUDAGraph() + capture = SimpleNamespace( + cuda_graph=graph, + _barrier_fn=None, + _end_current_segment=lambda: None, + _begin_new_segment=lambda: None, + ) + backend = SimpleNamespace(forward=attention) + with ( + attention_scope(captured_batch, backend, raw_num_tokens=2), + patch( + "sglang.srt.layers.radix_attention.is_in_breakable_cuda_graph", + return_value=True, + ), + ): + token = bcg._current_capture_var.set(capture) + try: + kwargs = {"idx_q": q, "idx_k": q} if sparse else {} + result = layer(q, q, q, captured_batch, **kwargs) + finally: + bcg._current_capture_var.reset(token) + index_output = None + if sparse: + index_output, output = result + lse = None + else: + output, lse = result if return_lse else (result, None) + assert len(graph._break_fns) == 1 + pointer = output.data_ptr() + index_pointer = index_output.data_ptr() if sparse else None + with attention_scope(live_batch, backend, raw_num_tokens=1): + graph._break_fns[0]() + assert seen == [captured_batch, live_batch] + assert output.data_ptr() == pointer + assert torch.all(output[:1] == 2) and torch.count_nonzero(output[1:]) == 0 + if sparse: + assert index_output.data_ptr() == index_pointer + assert torch.all(index_output[:1] == 4) + assert torch.count_nonzero(index_output[1:]) == 0 + if return_lse: + assert torch.all(lse[:1] == 2) and torch.count_nonzero(lse[1:]) == 0 + + +def test_prefill_graph_policy_restores_after_exception(): + from sglang.srt.model_executor.runner_utils.prefill_graph import ( + get_prefill_raw_num_tokens, + is_in_full_prefill_graph, + ) + + with prefill_graph_scope(full_graph=True, raw_num_tokens=3): + with pytest.raises(RuntimeError): + with prefill_graph_scope(full_graph=False, raw_num_tokens=1): + assert not is_in_full_prefill_graph() + assert get_prefill_raw_num_tokens() == 1 + raise RuntimeError("forward failed") + assert is_in_full_prefill_graph() + assert get_prefill_raw_num_tokens() == 3 + assert not is_in_full_prefill_graph() + assert get_prefill_raw_num_tokens() is None if __name__ == "__main__": - unittest.main() + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/test_radix_attention_padded_extend.py b/test/registered/unit/layers/test_radix_attention_padded_extend.py index 9a62e9bb03c2..848a84fdbfa0 100644 --- a/test/registered/unit/layers/test_radix_attention_padded_extend.py +++ b/test/registered/unit/layers/test_radix_attention_padded_extend.py @@ -4,7 +4,7 @@ import torch -from sglang.srt.layers import radix_attention +from sglang.srt.layers.attention import graph_utils as radix_attention from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase diff --git a/test/registered/unit/layers/test_radix_linear_attention.py b/test/registered/unit/layers/test_radix_linear_attention.py index a11230a8e7a1..ad0384e3fcab 100644 --- a/test/registered/unit/layers/test_radix_linear_attention.py +++ b/test/registered/unit/layers/test_radix_linear_attention.py @@ -66,6 +66,90 @@ def forward(self, *, layer, forward_batch, mixed_qkv, a, b): class TestRadixLinearAttentionPadding(CustomTestCase): + def test_graph_dispatch_distinguishes_prefill_decode_and_verify(self): + layer = radix_linear_attention.RadixLinearAttention( + layer_id=0, + num_q_heads=1, + num_k_heads=1, + num_v_heads=2, + head_q_dim=4, + head_k_dim=4, + head_v_dim=4, + ) + for mode, extend, verify in ( + ("prefill", True, False), + ("decode", False, False), + ("verify", True, True), + ): + for breakable, full in ((False, False), (True, False), (False, True)): + batch = SimpleNamespace( + forward_mode=SimpleNamespace( + is_extend=lambda: extend, + is_target_verify=lambda: verify, + is_extend_without_speculative=lambda: extend and not verify, + ), + global_num_token_non_padded_cpu=3, + out_cache_loc=torch.arange(3), + ) + with ( + self.subTest(mode=mode, breakable=breakable, full=full), + patch.object( + radix_linear_attention, + "is_in_breakable_cuda_graph", + return_value=breakable, + ), + patch.object( + radix_linear_attention, + "is_in_full_prefill_graph", + return_value=full, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FakeAttentionBackend(), + ), + patch.object( + radix_linear_attention, + "_linear_extend_in_graph", + return_value=False, + ), + patch.object(layer, "_eager_linear_attention") as eager, + ): + layer.forward( + batch, torch.zeros(3, 8), torch.zeros(3, 2), torch.zeros(3, 2) + ) + self.assertEqual( + eager.called, extend and (full or (breakable and not verify)) + ) + + def test_linear_extend_in_graph_does_not_create_an_eager_break(self): + layer = radix_linear_attention.RadixLinearAttention(0, 1, 1, 2, 4, 4, 4) + batch = SimpleNamespace( + forward_mode=SimpleNamespace( + is_extend=lambda: True, + is_extend_without_speculative=lambda: True, + ), + global_num_token_non_padded_cpu=3, + out_cache_loc=torch.arange(3), + ) + with ( + patch.object( + radix_linear_attention, "is_in_breakable_cuda_graph", return_value=True + ), + patch.object( + radix_linear_attention, "_linear_extend_in_graph", return_value=True + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FakeAttentionBackend(), + ), + patch.object(layer, "_eager_linear_attention") as eager, + ): + out = layer(batch, torch.zeros(3, 8), torch.zeros(3, 2), torch.zeros(3, 2)) + eager.assert_not_called() + torch.testing.assert_close(out, torch.full((1, 3, 2, 4), 5.0)) + def test_eager_padded_input_is_sliced_and_output_shape_is_restored(self): layer = radix_linear_attention.RadixLinearAttention( layer_id=0, @@ -86,8 +170,8 @@ def test_eager_padded_input_is_sliced_and_output_shape_is_restored(self): with ( patch.object( radix_linear_attention, - "get_tc_piecewise_forward_context", - return_value=None, + "is_in_full_prefill_graph", + return_value=False, ), patch.object( radix_linear_attention, @@ -126,8 +210,8 @@ def test_target_verify_keeps_physical_rows_matching_its_metadata(self): with ( patch.object( radix_linear_attention, - "get_tc_piecewise_forward_context", - return_value=None, + "is_in_full_prefill_graph", + return_value=False, ), patch.object( radix_linear_attention, @@ -165,8 +249,8 @@ def test_eager_backend_failure_restores_out_cache_loc(self): with ( patch.object( radix_linear_attention, - "get_tc_piecewise_forward_context", - return_value=None, + "is_in_full_prefill_graph", + return_value=False, ), patch.object( radix_linear_attention, @@ -201,8 +285,8 @@ def test_padded_output_tail_is_initialized(self): with ( patch.object( radix_linear_attention, - "get_tc_piecewise_forward_context", - return_value=context, + "is_in_full_prefill_graph", + return_value=True, ), patch.object( radix_linear_attention, @@ -210,12 +294,13 @@ def test_padded_output_tail_is_initialized(self): return_value=_FakeAttentionBackend(), ), ): - radix_linear_attention._unified_linear_attention_with_output_impl( + radix_linear_attention._linear_attention_with_output_impl( mixed_qkv=torch.zeros((padded_num_tokens, 8)), a=torch.zeros((padded_num_tokens, 2)), b=torch.zeros((padded_num_tokens, 2)), output=output, - layer_id=0, + attention_layer=context.attention_layers[0], + forward_batch=forward_batch, ) torch.testing.assert_close(output[:, :3], torch.full((1, 3, 2, 4), 5.0)) diff --git a/test/registered/unit/models/test_inkling_graph_fusion_gates.py b/test/registered/unit/models/test_inkling_graph_fusion_gates.py new file mode 100644 index 000000000000..d9d7e781a5f8 --- /dev/null +++ b/test/registered/unit/models/test_inkling_graph_fusion_gates.py @@ -0,0 +1,71 @@ +"""BCG prefill must not change decode or verification fusion eligibility.""" + +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope +from sglang.srt.models.inkling_common.kernels import comm +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +@pytest.mark.parametrize("scope", ["eager", "breakable", "full"]) +@pytest.mark.parametrize("mode", ["extend", "decode", "verify"]) +@pytest.mark.parametrize("scattered", [False, True]) +def test_fusion_gate_preserves_non_prefill_behavior(scope, mode, scattered): + batch = SimpleNamespace( + forward_mode=SimpleNamespace( + is_draft_extend_v2=lambda: False, + is_extend=lambda: mode != "decode", + is_decode=lambda: mode == "decode", + is_target_verify=lambda: mode == "verify", + is_extend_without_speculative=lambda: mode == "extend", + ) + ) + group = SimpleNamespace( + world_size=4, + torch_symm_mem_comm=SimpleNamespace(disabled=False, dtype=torch.bfloat16), + ) + gate = ( + comm.scattered_ar_sconv_fusable + if scattered + else comm.fullwidth_ar_sconv_fusable + ) + with ( + prefill_graph_scope(full_graph=scope == "full"), + patch.object(comm, "is_cuda", return_value=True), + patch.object( + comm, "is_in_breakable_cuda_graph", return_value=scope == "breakable" + ), + patch.object( + comm, + "get_exec", + return_value=SimpleNamespace( + comm=SimpleNamespace(enable_scattered_sconv=scattered) + ), + ), + patch.object( + comm.envs.SGLANG_OPT_USE_INKLING_CUSTOM_AR, "get", return_value=True + ), + patch.object( + comm.envs.SGLANG_OPT_USE_INKLING_FUSED_AR_SCONV, "get", return_value=True + ), + patch.object( + comm, + "_get_inkling_ar_resources", + return_value=SimpleNamespace(ssconv_out=4096 * 128), + ), + ): + expected = (scattered or mode == "extend") and not ( + scope == "breakable" and mode == "extend" + ) + assert gate(group, batch, 4096, 128, torch.bfloat16) == expected + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) From f1692d1a515d8c384fa4064475b7d4aaeb17b61a Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 15:48:46 -0700 Subject: [PATCH 2/7] fix: restrict attention eager dispatch to prefill runner scope --- .../srt/layers/attention/hpc_ops_backend.py | 15 +++++--- python/sglang/srt/layers/radix_attention.py | 2 + .../srt/layers/radix_linear_attention.py | 15 +++++--- .../runner_utils/prefill_graph.py | 5 +++ .../attention_forward_methods/forward_mla.py | 5 ++- python/sglang/srt/models/inkling.py | 9 ++++- .../srt/models/inkling_common/kernels/comm.py | 11 +++++- .../unit/layers/test_radix_attention.py | 27 ++++++++++++++ .../layers/test_radix_linear_attention.py | 37 +++++++++++++++++++ .../models/test_inkling_graph_fusion_gates.py | 13 +++++-- 10 files changed, 121 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/attention/hpc_ops_backend.py b/python/sglang/srt/layers/attention/hpc_ops_backend.py index 5c882e66f393..e9a3a8483f7f 100644 --- a/python/sglang/srt/layers/attention/hpc_ops_backend.py +++ b/python/sglang/srt/layers/attention/hpc_ops_backend.py @@ -53,6 +53,7 @@ from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch from sglang.srt.model_executor.runner_utils.prefill_graph import ( is_in_full_prefill_graph, + is_in_prefill_graph_scope, ) from sglang.srt.runtime_context import get_parallel @@ -411,11 +412,15 @@ def fused_qk_rope_store_kv_fp8( device=qkv.device, ) - if is_extend and ( - is_in_full_prefill_graph() - or ( - is_in_breakable_cuda_graph() - and forward_batch.forward_mode.is_extend_without_speculative() + if ( + is_in_prefill_graph_scope() + and is_extend + and ( + is_in_full_prefill_graph() + or ( + is_in_breakable_cuda_graph() + and forward_batch.forward_mode.is_extend_without_speculative() + ) ) ): self._eager_rope_store_kv( diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index ef5b75d4611d..9d0b67ae02d8 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -39,6 +39,7 @@ from sglang.srt.model_executor.runner_utils.prefill_graph import ( get_prefill_raw_num_tokens, is_in_full_prefill_graph, + is_in_prefill_graph_scope, ) if TYPE_CHECKING: @@ -131,6 +132,7 @@ def forward( full_cg = is_in_full_prefill_graph() use_eager_attention = ( self.use_prefill_attention_wrapper + and is_in_prefill_graph_scope() and forward_batch.forward_mode.is_extend() and (breakable_cg or full_cg) ) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 2ecafa55a18d..00d44400f80e 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -30,6 +30,7 @@ from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch from sglang.srt.model_executor.runner_utils.prefill_graph import ( is_in_full_prefill_graph, + is_in_prefill_graph_scope, ) if TYPE_CHECKING: @@ -86,11 +87,15 @@ def forward( b: torch.Tensor, ) -> torch.Tensor: is_extend = forward_batch.forward_mode.is_extend() - if is_extend and ( - is_in_full_prefill_graph() - or ( - is_in_breakable_cuda_graph() - and forward_batch.forward_mode.is_extend_without_speculative() + if ( + is_in_prefill_graph_scope() + and is_extend + and ( + is_in_full_prefill_graph() + or ( + is_in_breakable_cuda_graph() + and forward_batch.forward_mode.is_extend_without_speculative() + ) ) ): # Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim) diff --git a/python/sglang/srt/model_executor/runner_utils/prefill_graph.py b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py index dfad7e50ad24..1c95be62e135 100644 --- a/python/sglang/srt/model_executor/runner_utils/prefill_graph.py +++ b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py @@ -16,6 +16,11 @@ class _PrefillGraphState: ) +def is_in_prefill_graph_scope() -> bool: + """Whether the prefill runner owns this execution, regardless of graph backend.""" + return _current.get() is not None + + def is_in_full_prefill_graph() -> bool: state = _current.get() return state is not None and state.full_graph diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index e53502706a13..078d653b1721 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -65,6 +65,7 @@ from sglang.srt.model_executor.runner_utils.prefill_graph import ( is_in_full_prefill_graph, + is_in_prefill_graph_scope, ) if TYPE_CHECKING: @@ -159,7 +160,9 @@ def _can_fuse_bmm_into_attention( return False # Like the DSA indexer eager region, this fusion is enabled for # non-speculative CUDA BCG prefill. - if not is_graph_dsa_split_op_surface(forward_batch): + if not is_in_prefill_graph_scope() or not is_graph_dsa_split_op_surface( + forward_batch + ): return False if not self.use_dsa: return False diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py index 0e04d3aa345d..bb27d8ee0be3 100644 --- a/python/sglang/srt/models/inkling.py +++ b/python/sglang/srt/models/inkling.py @@ -42,6 +42,9 @@ is_in_breakable_cuda_graph, ) from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + is_in_prefill_graph_scope, +) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.inkling_common.attn import ( InklingAttention, @@ -465,7 +468,8 @@ def forward( # The eager group reads the live prepared batch from the runner scope. if ( - is_in_breakable_cuda_graph() + is_in_prefill_graph_scope() + and is_in_breakable_cuda_graph() and forward_batch.forward_mode.is_extend_without_speculative() ): # BCG prefill path: the AR fusion is decode-only, so partials never @@ -906,7 +910,8 @@ def forward( # Match the per-layer prefill eager region. scattered = self.layers[-1].scattered_sconv if ( - is_in_breakable_cuda_graph() + is_in_prefill_graph_scope() + and is_in_breakable_cuda_graph() and forward_batch.forward_mode.is_extend_without_speculative() ): # Under scattered sconv the input is the last MoE's [T, H/P] diff --git a/python/sglang/srt/models/inkling_common/kernels/comm.py b/python/sglang/srt/models/inkling_common/kernels/comm.py index 1c094fe6da8e..61a1ca80d5d5 100644 --- a/python/sglang/srt/models/inkling_common/kernels/comm.py +++ b/python/sglang/srt/models/inkling_common/kernels/comm.py @@ -10,6 +10,9 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( is_in_breakable_cuda_graph, ) +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + is_in_prefill_graph_scope, +) from sglang.srt.runtime_context import get_exec from sglang.srt.utils import is_cuda @@ -724,7 +727,11 @@ def scattered_ar_sconv_fusable( # [:raw]), and sentinel request slots have qlen == 0 so the in-kernel # cache update/track skip them. - if is_in_breakable_cuda_graph() and fm.is_extend_without_speculative(): + if ( + is_in_prefill_graph_scope() + and is_in_breakable_cuda_graph() + and fm.is_extend_without_speculative() + ): return False comm = group.torch_symm_mem_comm if ( @@ -1052,7 +1059,7 @@ def fullwidth_ar_sconv_fusable( # pieces can't carry the cross-layer producer contract; the FULL prefill # CUDA-graph backend is supported (capture-safe kernel, in-graph metadata). - if is_in_breakable_cuda_graph(): + if is_in_prefill_graph_scope() and is_in_breakable_cuda_graph(): return False comm = group.torch_symm_mem_comm if ( diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index 1e290a4c09f9..a57c20f0e5b4 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/test/registered/unit/layers/test_radix_attention.py @@ -270,5 +270,32 @@ def test_prefill_graph_policy_restores_after_exception(): assert get_prefill_raw_num_tokens() is None +@pytest.mark.parametrize("mode_name", ["DECODE", "TARGET_VERIFY", "DLLM_EXTEND"]) +@pytest.mark.parametrize("sparse", [False, True]) +def test_decode_bcg_without_prefill_scope_uses_backend(mode_name, sparse): + from sglang.srt.model_executor.forward_batch_info import ForwardMode + from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch + + layer = RadixAttention(1, 2, 1.0, 1, 0) + batch = make_batch(tokens=4) + batch.forward_mode = ForwardMode[mode_name] + q = torch.zeros(4, 1, 2) + result = object() + backend = SimpleNamespace(forward=lambda *args, **kwargs: result) + with ( + forward_context(ForwardContext(backend)), + patch( + "sglang.srt.layers.radix_attention.is_in_breakable_cuda_graph", + return_value=True, + ), + patch.object(layer, "_eager_attention") as eager, + ): + with pytest.raises(RuntimeError, match="No forward batch"): + get_forward_batch() + kwargs = {"idx_q": q, "idx_k": q} if sparse else {} + assert layer(q, q, q, batch, **kwargs) is result + eager.assert_not_called() + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/test_radix_linear_attention.py b/test/registered/unit/layers/test_radix_linear_attention.py index ad0384e3fcab..4b8cede291e2 100644 --- a/test/registered/unit/layers/test_radix_linear_attention.py +++ b/test/registered/unit/layers/test_radix_linear_attention.py @@ -6,6 +6,7 @@ import torch import sglang.srt.layers.radix_linear_attention as radix_linear_attention +from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -93,6 +94,7 @@ def test_graph_dispatch_distinguishes_prefill_decode_and_verify(self): ) with ( self.subTest(mode=mode, breakable=breakable, full=full), + prefill_graph_scope(full_graph=full), patch.object( radix_linear_attention, "is_in_breakable_cuda_graph", @@ -122,6 +124,40 @@ def test_graph_dispatch_distinguishes_prefill_decode_and_verify(self): eager.called, extend and (full or (breakable and not verify)) ) + def test_decode_bcg_without_prefill_scope_uses_backend(self): + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + layer = radix_linear_attention.RadixLinearAttention(0, 1, 1, 2, 4, 4, 4) + for mode in ( + ForwardMode.DECODE, + ForwardMode.TARGET_VERIFY, + ForwardMode.DLLM_EXTEND, + ): + batch = SimpleNamespace( + forward_mode=mode, + global_num_token_non_padded_cpu=3, + out_cache_loc=torch.arange(3), + ) + with ( + self.subTest(mode=mode), + patch.object( + radix_linear_attention, + "is_in_breakable_cuda_graph", + return_value=True, + ), + patch.object( + radix_linear_attention, + "get_attn_backend", + return_value=_FakeAttentionBackend(), + ), + patch.object(layer, "_eager_linear_attention") as eager, + ): + out = layer( + batch, torch.zeros(3, 8), torch.zeros(3, 2), torch.zeros(3, 2) + ) + eager.assert_not_called() + torch.testing.assert_close(out, torch.full((1, 3, 2, 4), 5.0)) + def test_linear_extend_in_graph_does_not_create_an_eager_break(self): layer = radix_linear_attention.RadixLinearAttention(0, 1, 1, 2, 4, 4, 4) batch = SimpleNamespace( @@ -133,6 +169,7 @@ def test_linear_extend_in_graph_does_not_create_an_eager_break(self): out_cache_loc=torch.arange(3), ) with ( + prefill_graph_scope(full_graph=False), patch.object( radix_linear_attention, "is_in_breakable_cuda_graph", return_value=True ), diff --git a/test/registered/unit/models/test_inkling_graph_fusion_gates.py b/test/registered/unit/models/test_inkling_graph_fusion_gates.py index d9d7e781a5f8..9635816a2f47 100644 --- a/test/registered/unit/models/test_inkling_graph_fusion_gates.py +++ b/test/registered/unit/models/test_inkling_graph_fusion_gates.py @@ -1,6 +1,7 @@ """BCG prefill must not change decode or verification fusion eligibility.""" import sys +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import patch @@ -14,7 +15,7 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") -@pytest.mark.parametrize("scope", ["eager", "breakable", "full"]) +@pytest.mark.parametrize("scope", ["eager", "breakable", "full", "decode_breakable"]) @pytest.mark.parametrize("mode", ["extend", "decode", "verify"]) @pytest.mark.parametrize("scattered", [False, True]) def test_fusion_gate_preserves_non_prefill_behavior(scope, mode, scattered): @@ -37,10 +38,16 @@ def test_fusion_gate_preserves_non_prefill_behavior(scope, mode, scattered): else comm.fullwidth_ar_sconv_fusable ) with ( - prefill_graph_scope(full_graph=scope == "full"), + ( + prefill_graph_scope(full_graph=scope == "full") + if scope in ("breakable", "full") + else nullcontext() + ), patch.object(comm, "is_cuda", return_value=True), patch.object( - comm, "is_in_breakable_cuda_graph", return_value=scope == "breakable" + comm, + "is_in_breakable_cuda_graph", + return_value=scope in ("breakable", "decode_breakable"), ), patch.object( comm, From 5bff1604967a7e992482fe6be8ea9ea5d45fb090 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 16:02:42 -0700 Subject: [PATCH 3/7] style: move attention test imports to module scope --- .../breakable/test_breakable_cuda_graph.py | 14 +++++------ .../unit/layers/test_radix_attention.py | 24 ++++++++++--------- .../layers/test_radix_linear_attention.py | 2 +- 3 files changed, 21 insertions(+), 19 deletions(-) diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 7283e09b2977..159feb8aa422 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py @@ -12,6 +12,13 @@ import torch +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.model_executor.runner_utils.forward_batch import ( + set_forward_batch, +) +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + prefill_graph_scope, +) from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.run_eval import run_eval @@ -241,13 +248,6 @@ def scale(src): ) def test_attention_narrows_padded_positions(self): - from sglang.srt.layers.radix_attention import RadixAttention - from sglang.srt.model_executor.runner_utils.forward_batch import ( - set_forward_batch, - ) - from sglang.srt.model_executor.runner_utils.prefill_graph import ( - prefill_graph_scope, - ) num_tokens = 3 padded_num_tokens = 5 diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index a57c20f0e5b4..18f3970fa145 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/test/registered/unit/layers/test_radix_attention.py @@ -9,9 +9,20 @@ import torch from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context -from sglang.srt.model_executor.runner_utils.forward_batch import set_forward_batch -from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + breakable_cuda_graph as bcg, +) +from sglang.srt.model_executor.runner_utils.forward_batch import ( + get_forward_batch, + set_forward_batch, +) +from sglang.srt.model_executor.runner_utils.prefill_graph import ( + get_prefill_raw_num_tokens, + is_in_full_prefill_graph, + prefill_graph_scope, +) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -191,9 +202,6 @@ def test_output_dtype_follows_values(): "return_lse,sparse", [(False, False), (True, False), (False, True)] ) def test_attention_replay_uses_live_batch_and_static_output_buffers(return_lse, sparse): - from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( - breakable_cuda_graph as bcg, - ) layer = RadixAttention(2, 3, 1.0, 2, 0) q = torch.zeros(4, 2, 3) @@ -253,10 +261,6 @@ def attention(query, key, value, actual_layer, batch, save_kv_cache, **kwargs): def test_prefill_graph_policy_restores_after_exception(): - from sglang.srt.model_executor.runner_utils.prefill_graph import ( - get_prefill_raw_num_tokens, - is_in_full_prefill_graph, - ) with prefill_graph_scope(full_graph=True, raw_num_tokens=3): with pytest.raises(RuntimeError): @@ -273,8 +277,6 @@ def test_prefill_graph_policy_restores_after_exception(): @pytest.mark.parametrize("mode_name", ["DECODE", "TARGET_VERIFY", "DLLM_EXTEND"]) @pytest.mark.parametrize("sparse", [False, True]) def test_decode_bcg_without_prefill_scope_uses_backend(mode_name, sparse): - from sglang.srt.model_executor.forward_batch_info import ForwardMode - from sglang.srt.model_executor.runner_utils.forward_batch import get_forward_batch layer = RadixAttention(1, 2, 1.0, 1, 0) batch = make_batch(tokens=4) diff --git a/test/registered/unit/layers/test_radix_linear_attention.py b/test/registered/unit/layers/test_radix_linear_attention.py index 4b8cede291e2..3a47b8c9fb7b 100644 --- a/test/registered/unit/layers/test_radix_linear_attention.py +++ b/test/registered/unit/layers/test_radix_linear_attention.py @@ -6,6 +6,7 @@ import torch import sglang.srt.layers.radix_linear_attention as radix_linear_attention +from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -125,7 +126,6 @@ def test_graph_dispatch_distinguishes_prefill_decode_and_verify(self): ) def test_decode_bcg_without_prefill_scope_uses_backend(self): - from sglang.srt.model_executor.forward_batch_info import ForwardMode layer = radix_linear_attention.RadixLinearAttention(0, 1, 1, 2, 4, 4, 4) for mode in ( From 5dd290e38be3f307d15388a2128d242752f35246 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 16:50:18 -0700 Subject: [PATCH 4/7] refactor: expose bound run_eager entry for graph eager methods --- .../breakable_cuda_graph.py | 26 +++++++- .../attention_forward_methods/forward_mla.py | 3 +- .../model_executor/test_eager_on_graph.py | 64 +++++++++++++++++-- 3 files changed, 85 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index cdd3af540956..eaa2f8e774c6 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py @@ -219,6 +219,26 @@ def _copy_output(dst: Any, src: Any) -> Any: return src +class _EagerOnGraphCallable: + """Bind both graph-aware and explicit eager entry points to the same instance.""" + + def __init__(self, graph_call: Callable, eager_call: Callable): + self._graph_call = graph_call + self.run_eager = eager_call + + def __call__(self, *args, **kwargs): + return self._graph_call(*args, **kwargs) + + def __get__(self, instance, owner=None): + if instance is None: + return self + bound = _EagerOnGraphCallable( + functools.partial(self._graph_call, instance), + functools.partial(self.run_eager, instance), + ) + return bound + + def eager_on_graph( fn: Optional[Callable] = None, capture_stub: Optional[Callable] = None, @@ -230,6 +250,10 @@ def eager_on_graph( Arguments retain their capture-time identity. Tensors must use the static buffers owned by the graph runner; request-specific state must be read inside the eager function at execution time. + + Use ``fn.run_eager(...)`` to call the body inside an existing eager region. + This explicit entry bypasses graph segmentation and capture stubs; methods + bind their instance automatically. """ # Transitional support while model callers migrate to the bare decorator. @@ -286,7 +310,7 @@ def replay_fn(): capture._begin_new_segment() return output - return wrapper + return functools.update_wrapper(_EagerOnGraphCallable(wrapper, inner), inner) return decorator(fn) if fn is not None else decorator diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 078d653b1721..a34aad0dbc02 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -957,8 +957,7 @@ def _eager_bmm_attention( topk_indices: Optional[torch.Tensor] = None, ) -> None: torch.bmm(q_nope_t, w_kc, out=q_nope_out_buf) - self.attn_mqa._eager_attention.__wrapped__( - self.attn_mqa, + self.attn_mqa._eager_attention.run_eager( q_nope_out_view, k_nope, k_nope, diff --git a/test/registered/unit/model_executor/test_eager_on_graph.py b/test/registered/unit/model_executor/test_eager_on_graph.py index aabaf7eca336..5f6dd8bae4ee 100644 --- a/test/registered/unit/model_executor/test_eager_on_graph.py +++ b/test/registered/unit/model_executor/test_eager_on_graph.py @@ -1,7 +1,10 @@ """CPU tests of scoped eager replay without requiring CUDA graph capture.""" +import gc import sys +import weakref from contextlib import contextmanager +from contextvars import Context from types import SimpleNamespace import pytest @@ -104,9 +107,6 @@ def increment(x): def test_replay_does_not_retain_capture_or_serving_batches(): - import gc - import weakref - class Batch: scale = 3 @@ -162,13 +162,67 @@ def test_batch_scope_restores_outer_batch_after_exception(): def test_batch_scope_is_isolated_between_execution_contexts(): - from contextvars import Context - with set_forward_batch(SimpleNamespace(scale=2)): with pytest.raises(RuntimeError, match="No forward batch"): Context().run(get_forward_batch) assert get_forward_batch().scale == 2 +def test_run_eager_binds_instance_and_keeps_one_outer_break(): + calls = [] + + class Layer: + def __init__(self, scale): + self.scale = scale + + def stub(self, x, *, offset=0): + raise AssertionError("run_eager must bypass the capture stub") + + @bcg.eager_on_graph(capture_stub=stub) + def inner(self, x, *, offset=0): + calls.append(self.scale) + return x * self.scale + offset + + @bcg.eager_on_graph + def outer(self, x): + return self.inner.run_eager(x, offset=1) + + first, second = Layer(2), Layer(5) + with recording_capture() as (graph, events): + output = first.outer(torch.ones(1)) + assert events == ["end", "begin"] + assert len(graph._break_fns) == 1 + assert output.item() == 3 + assert second.inner.run_eager(torch.ones(1), offset=1).item() == 6 + graph._break_fns[0]() + assert calls == [2, 5, 2] + assert output.item() == 3 + + +def test_run_eager_plain_function_does_not_record_a_break(): + @bcg.eager_on_graph + def increment(x): + return x + 1 + + with recording_capture() as (graph, events): + assert increment.run_eager(torch.tensor(2)).item() == 3 + assert events == [] and graph._break_fns == [] + + +def test_run_eager_method_torch_compile(): + class Layer(torch.nn.Module): + @bcg.eager_on_graph + def inner(self, x): + return x + 1 + + def forward(self, x): + return self.inner.run_eager(x) + + layer = Layer() + x = torch.ones(2) + compiled = torch.compile(layer, backend="eager", fullgraph=True) + torch.testing.assert_close(compiled(x), layer(x)) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) From 70c0739e43a50b0be51d243d63ae1d8452c20b49 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 16:53:43 -0700 Subject: [PATCH 5/7] refactor: make graph and eager call storage symmetric --- .../breakable_cuda_graph/breakable_cuda_graph.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index eaa2f8e774c6..f33b96018aab 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py @@ -224,17 +224,20 @@ class _EagerOnGraphCallable: def __init__(self, graph_call: Callable, eager_call: Callable): self._graph_call = graph_call - self.run_eager = eager_call + self._eager_call = eager_call def __call__(self, *args, **kwargs): return self._graph_call(*args, **kwargs) + def run_eager(self, *args, **kwargs): + return self._eager_call(*args, **kwargs) + def __get__(self, instance, owner=None): if instance is None: return self bound = _EagerOnGraphCallable( functools.partial(self._graph_call, instance), - functools.partial(self.run_eager, instance), + functools.partial(self._eager_call, instance), ) return bound From 0f423a9551f1670147b442aabf260ef0922f7ab5 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 20:17:24 -0700 Subject: [PATCH 6/7] fix: guard compiled prefill lookups and restore BCG test dispatch --- .../runner_utils/prefill_graph.py | 8 ++++ .../runner_modes/split_op_runner.py | 37 ++++++++++++++++++- .../unit/layers/test_radix_attention.py | 30 ++++++++++++++- 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/runner_utils/prefill_graph.py b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py index 1c95be62e135..34616e51d778 100644 --- a/python/sglang/srt/model_executor/runner_utils/prefill_graph.py +++ b/python/sglang/srt/model_executor/runner_utils/prefill_graph.py @@ -4,6 +4,8 @@ from contextvars import ContextVar from dataclasses import dataclass +import torch + @dataclass(frozen=True) class _PrefillGraphState: @@ -18,10 +20,16 @@ class _PrefillGraphState: def is_in_prefill_graph_scope() -> bool: """Whether the prefill runner owns this execution, regardless of graph backend.""" + # BCG/full prefill runs without torch.compile. Older supported PyTorch + # builds cannot trace ContextVar.get(), including during decode/verify. + if torch.compiler.is_compiling(): + return False return _current.get() is not None def is_in_full_prefill_graph() -> bool: + if torch.compiler.is_compiling(): + return False state = _current.get() return state is not None and state.full_graph diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py index 9542b8c96ff4..47fbc9327e5e 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py @@ -1,8 +1,12 @@ +from contextlib import ExitStack, nullcontext from dataclasses import dataclass, replace from typing import Any, Callable +from unittest.mock import patch import torch +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( enable_breakable_cuda_graph, @@ -13,6 +17,8 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import ( set_tc_piecewise_forward_context as piecewise_forward_context, ) +from sglang.srt.model_executor.runner_utils.forward_batch import set_forward_batch +from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope from ..attention_methods.dense_attention import DEFAULT_DEVICE as DENSE_DEFAULT_DEVICE from ..attention_methods.dense_attention import DEFAULT_DTYPE as DENSE_DEFAULT_DTYPE @@ -304,20 +310,49 @@ def _run_split_op_extend_case( max_context_len=max_context_len, ) + layers = adapter.attention_layers(split_fixture) with ( + ExitStack() as stack, torch.no_grad(), + set_forward_batch(static_batch) if breakable else nullcontext(), + prefill_graph_scope(full_graph=False, raw_num_tokens=raw_num_tokens) + if breakable + else nullcontext(), _split_op_context(breakable=breakable), forward_context(ForwardContext(attn_backend=split_fixture.backend)), piecewise_forward_context( static_batch, - adapter.attention_layers(split_fixture), + layers, None, [], [], ), ): + eager_calls = [] + if breakable: + for layer in layers: + if isinstance(layer, RadixAttention): + eager_calls.append( + stack.enter_context( + patch.object( + layer, "_eager_attention", wraps=layer._eager_attention + ) + ) + ) + elif isinstance(layer, RadixLinearAttention): + eager_calls.append( + stack.enter_context( + patch.object( + layer, + "_eager_linear_attention", + wraps=layer._eager_linear_attention, + ) + ) + ) split_fixture.backend.init_forward_metadata(raw_batch) actual = adapter.run_forward(split_fixture, static_batch, static_inputs) + for eager_call in eager_calls: + eager_call.assert_called() actual = _slice_live_tokens(actual, raw_num_tokens) torch.testing.assert_close(actual, expected, atol=adapter.atol, rtol=adapter.rtol) diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index 18f3970fa145..d5c853a70ba7 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/test/registered/unit/layers/test_radix_attention.py @@ -1,7 +1,7 @@ """CPU contracts shared by the full and breakable prefill attention paths.""" import sys -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import patch @@ -14,6 +14,7 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( breakable_cuda_graph as bcg, ) +from sglang.srt.model_executor.runner_utils import prefill_graph from sglang.srt.model_executor.runner_utils.forward_batch import ( get_forward_batch, set_forward_batch, @@ -299,5 +300,32 @@ def test_decode_bcg_without_prefill_scope_uses_backend(mode_name, sparse): eager.assert_not_called() +@pytest.mark.parametrize("mode", [ForwardMode.DECODE, ForwardMode.TARGET_VERIFY]) +@pytest.mark.parametrize("forbid_context_read", [False, True]) +def test_compiled_forward_without_prefill_scope_does_not_read_context( + mode, forbid_context_read +): + class Backend: + def forward(self, q, k, v, layer, batch, save_kv_cache, **kwargs): + return q + k + v + + class UnreadableContext: + def get(self): + raise AssertionError("compiled dispatch must not read prefill ContextVar") + + layer = RadixAttention(1, 2, 1.0, 1, 0) + batch = make_batch(tokens=4) + batch.forward_mode = mode + q = torch.ones(4, 1, 2) + with ( + forward_context(ForwardContext(Backend())), + patch.object(prefill_graph, "_current", UnreadableContext()) + if forbid_context_read + else nullcontext(), + ): + compiled = torch.compile(layer, backend="eager", fullgraph=True) + torch.testing.assert_close(compiled(q, q, q, batch), q * 3) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) From 0bf044729dd68a0d31a5de0e1657cf6379ecee75 Mon Sep 17 00:00:00 2001 From: Oasis-Git Date: Mon, 12 Oct 2026 20:38:07 -0700 Subject: [PATCH 7/7] Fix runner utility imports causing CPU test circular dependency --- .../srt/model_executor/runner/__init__.py | 12 +++- .../srt/model_executor/runner/eager_runner.py | 2 +- .../runner/prefill_cuda_graph_runner.py | 6 +- .../model_executor/runner_utils/__init__.py | 38 +++---------- .../runner/test_prefill_shared_read_done.py | 2 +- .../runner_utils/test_imports.py | 56 +++++++++++++++++++ 6 files changed, 78 insertions(+), 38 deletions(-) create mode 100644 test/registered/unit/model_executor/runner_utils/test_imports.py diff --git a/python/sglang/srt/model_executor/runner/__init__.py b/python/sglang/srt/model_executor/runner/__init__.py index 9aa65fbf579b..29d15f4471aa 100644 --- a/python/sglang/srt/model_executor/runner/__init__.py +++ b/python/sglang/srt/model_executor/runner/__init__.py @@ -41,16 +41,22 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( # noqa: F401 TCPCG_FAILURE_HINT, ) -from sglang.srt.model_executor.runner_utils import ( # noqa: F401 +from sglang.srt.model_executor.runner_utils.buffers import ( # noqa: F401 DecodeInputBuffers, - DeepEPCudaGraphRunnerAdapter, PrefillInputBuffers, _grouped_foreach_copy_, +) +from sglang.srt.model_executor.runner_utils.capture_mode import ( # noqa: F401 _set_capture_lora_variant, compile_in_capture_mode, get_capture_lora_variant, - get_global_graph_memory_pool, get_is_capture_mode, model_capture_mode, +) +from sglang.srt.model_executor.runner_utils.deepep_adapter import ( # noqa: F401 + DeepEPCudaGraphRunnerAdapter, +) +from sglang.srt.model_executor.runner_utils.pool import ( # noqa: F401 + get_global_graph_memory_pool, set_global_graph_memory_pool, ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index b1f21a26dad0..f976b0403457 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -50,7 +50,7 @@ get_token_to_kv_pool, ) from sglang.srt.model_executor.runner.base_runner import BaseRunner -from sglang.srt.model_executor.runner_utils import ( +from sglang.srt.model_executor.runner_utils.shared_read_event import ( maybe_publish_prefill_shared_read_done, ) from sglang.srt.runtime_context import ( diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 8fbb6ab6cbe9..e17cf2d83891 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -117,9 +117,6 @@ TCPCG_FAILURE_HINT, set_tc_piecewise_forward_context, ) -from sglang.srt.model_executor.runner_utils import ( - maybe_publish_prefill_shared_read_done, -) from sglang.srt.model_executor.runner_utils.buffers import ( PrefillInputBuffers, ) @@ -129,6 +126,9 @@ get_or_create_global_graph_capture_stream, ) from sglang.srt.model_executor.runner_utils.prefill_graph import prefill_graph_scope +from sglang.srt.model_executor.runner_utils.shared_read_event import ( + maybe_publish_prefill_shared_read_done, +) from sglang.srt.model_loader.utils import resolve_language_model from sglang.srt.runtime_context import ( get_exec, diff --git a/python/sglang/srt/model_executor/runner_utils/__init__.py b/python/sglang/srt/model_executor/runner_utils/__init__.py index c2529cb20530..703661faa0ec 100644 --- a/python/sglang/srt/model_executor/runner_utils/__init__.py +++ b/python/sglang/srt/model_executor/runner_utils/__init__.py @@ -1,32 +1,10 @@ -"""Low-level utilities used by the CUDA graph runners. +"""Low-level utilities shared by model layers and CUDA graph runners. -Mirror of cuda_graph_backend_utils/ for runner-side state — buffer -dataclasses, process-global capture flags, the speculative-shared -graph memory pool, and the DeepEP capture/replay adapter. Runners in -cuda_graph_runner/ import from here; nothing here should import -back into cuda_graph_runner/. -""" +Keep this package initializer free of imports: model layers import lightweight +execution state from forward_batch and prefill_graph while their own modules +are still initializing. Loading runner dependencies here (especially DeepEP) +can cycle back through attention, quantization, and disaggregation modules. -from sglang.srt.model_executor.runner_utils.buffers import ( # noqa: F401 - DecodeInputBuffers, - PrefillInputBuffers, - _grouped_foreach_copy_, -) -from sglang.srt.model_executor.runner_utils.capture_mode import ( # noqa: F401 - _set_capture_lora_variant, - compile_in_capture_mode, - get_capture_lora_variant, - get_is_capture_mode, - model_capture_mode, -) -from sglang.srt.model_executor.runner_utils.deepep_adapter import ( # noqa: F401 - DeepEPCudaGraphRunnerAdapter, -) -from sglang.srt.model_executor.runner_utils.pool import ( # noqa: F401 - get_global_graph_memory_pool, - get_or_create_global_graph_capture_stream, - set_global_graph_memory_pool, -) -from sglang.srt.model_executor.runner_utils.shared_read_event import ( # noqa: F401 - maybe_publish_prefill_shared_read_done, -) +Import utilities directly from their defining submodules. The legacy runner +exports remain in model_executor.runner for speculative draft runners. +""" diff --git a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py index 0c96e270aa69..abfc0ddb56b1 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py +++ b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py @@ -9,7 +9,7 @@ SharedReadEnds, ) from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.model_executor.runner_utils import ( +from sglang.srt.model_executor.runner_utils.shared_read_event import ( maybe_publish_prefill_shared_read_done, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm diff --git a/test/registered/unit/model_executor/runner_utils/test_imports.py b/test/registered/unit/model_executor/runner_utils/test_imports.py new file mode 100644 index 000000000000..8d337213d291 --- /dev/null +++ b/test/registered/unit/model_executor/runner_utils/test_imports.py @@ -0,0 +1,56 @@ +"""Exercise import order in fresh interpreters, without a warmed module cache.""" + +import subprocess +import sys +import unittest + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + + +class TestRunnerUtilsImports(unittest.TestCase): + def test_execution_state_does_not_load_runner_dependencies(self): + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys\n" + "from sglang.srt.model_executor.runner_utils.forward_batch " + "import get_forward_batch, set_forward_batch\n" + "from sglang.srt.model_executor.runner_utils.prefill_graph " + "import prefill_graph_scope\n" + "assert 'sglang.srt.model_executor.runner_utils.deepep_adapter' " + "not in sys.modules\n" + "assert 'sglang.srt.model_executor.runner_utils.buffers' " + "not in sys.modules\n" + "batch = object()\n" + "with set_forward_batch(batch), prefill_graph_scope(full_graph=True):\n" + " assert get_forward_batch() is batch\n", + ], + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + + def test_test_utils_imports_before_attention(self): + # CPU suites enter through test_utils -> disaggregation -> quantization + # -> radix_attention. Importing runner_utils must not loop back through + # DeepEP and metrics_collector into partially initialized disaggregation. + result = subprocess.run( + [ + sys.executable, + "-c", + "from sglang.test.test_utils import CustomTestCase\n" + "from sglang.srt.layers.radix_attention import RadixAttention\n", + ], + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + + +if __name__ == "__main__": + unittest.main()