From cead77a08eeb1ad1e6a958c476b4f7effe73d150 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sun, 24 May 2026 06:49:45 +0000 Subject: [PATCH] feat(yoco): implement YOCO fast-prefill branch in Gemma4TextModel (PR-B/3) When --kv-sharing-fast-prefill is on AND the served Gemma-4 checkpoint has num_kv_shared_layers > 0 (currently Gemma-4 E4B-it and E2B-it, plus any future Gemma-4 variant that uses KV-sharing), the EXTEND-mode layer loop now splits into a self-decoder + a cross-decoder, where the cross- decoder runs only on the per-request last-extend-token rows. The cross- decoder output is scattered back into a full-shape hidden_states tensor so the LM head, frozen-KV MTP worker, and any other downstream consumer sees the existing [T, H] contract. Note: gemma-4-26b-a4b-it and gemma-4-31B-it both have num_kv_shared_layers = 0 in their HF text_config, so YOCO does NOT fire for those checkpoints (the predicate returns False). The target benchmark model for this stack is google/gemma-4-E4B-it (num_kv_shared_layers=18, num_hidden_layers=42, PLE enabled hidden_size_per_layer_input=256). PR-C runs the bench there. Mechanism mirrors vllm/model_executor/models/gemma4.py:1190-1273 (fast_prefill_forward), reimplemented in SGLang's own forward + attention metadata abstractions: * Gemma4TextModel.__init__: caches first_kv_shared_layer_idx and reads the flag from get_global_server_args(). * Gemma4TextModel._can_run_yoco(forward_batch): predicate that gates the branch. Returns False unless flag-on AND model has KV-shared layers AND forward_mode is EXTEND-not-TARGET_VERIFY AND extend_seq_lens populated AND no input-token logprobs AND no Eagle3 aux-hidden capture AND the back-half boundary actually falls inside the local PP rank's layer range. PLE-enabled variants (E4B/E2B) are supported via the per-layer inputs gather below. * Gemma4TextModel._build_cross_decoder_last_token_index(forward_batch): per-request last-extend-token row index. Matches LogitsProcessor._get_pruned_states cumsum-1 (or the padded-static-len variant for piecewise CUDA graph). * Gemma4TextModel._run_cross_decoder_with_yoco(...): temporarily mutates forward_batch.forward_mode to DECODE so the triton attention backend rebuilds its metadata as qo_indptr=[0..B], max_extend_len=1, kv_indptr=cumsum(seq_lens), kv_indices over the full prefix+extend KV span per request (from the donor layer's KV pool). Runs the KV-shared layers on the gathered Q rows, then restores mode and metadata in a try/finally so the caller sees no externally visible mutation. Also gathers per_layer_inputs to the same rows when PLE is enabled. * Gemma4TextModel.forward: when _can_run_yoco fires, runs layers [start_layer, first_kv_shared_layer_idx) on the full extend-token batch, calls _run_cross_decoder_with_yoco, then scatters via hidden_states.clone().index_copy_(0, last_token_index, cross_hidden) so the final norm sees a full-shape tensor. Tests: test/registered/unit/models/test_gemma4_yoco.py adds 16 unit tests covering: * TestCanRunYoco: 11 cases (eligible, flag-off, no-KV-shared, decode, target-verify, no-extend-lens, PLE-enabled-accepted, input-logprobs, eagle3-capture, PP-no-back-half, PP-no-front-half) * TestBuildCrossDecoderLastTokenIndex: 2 cases (non-padded cumsum-1, padded-static-len) * TestRunCrossDecoderWithYoco: 3 cases (mode+metadata restore, per_layer_inputs gather, exception-still-restores) All 16 tests pass: Ran 16 tests in 0.004s OK Diff in gemma4_causal.py is +188/-5 (modest reflow noise from auto-format). Stack base: pyc/yoco-fast-prefill-config @ 0911a9627 (PR-A: flag plumbing) Next: PR-C/3 (E4B benchmark + parity test + tuning). Plan: .humanize/yoco-gemma4/refined-plan.md Co-authored-by: Claude --- python/sglang/srt/models/gemma4_causal.py | 202 ++++++++++- .../unit/models/test_gemma4_yoco.py | 336 ++++++++++++++++++ 2 files changed, 533 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/models/test_gemma4_yoco.py diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index a943730cc893..9885399271f4 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -50,7 +50,12 @@ from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + ForwardMode, + PPProxyTensors, +) +from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, @@ -843,6 +848,35 @@ def __init__( else: self.norm = PPMissingLayer() self.layers_to_capture = [] + + # KV-sharing fast-prefill (YOCO) configuration. When the user opts + # into ``--kv-sharing-fast-prefill`` and the served model declares + # ``num_kv_shared_layers > 0``, EXTEND-mode forwards split the layer + # loop into a self-decoder (layers ``[0, first_kv_shared_layer_idx)``) + # that runs on the full extend-token batch, and a cross-decoder + # (layers ``[first_kv_shared_layer_idx, num_hidden_layers)``) that + # runs only on the per-request last-extend-token rows. The cross- + # decoder output is scattered back into a full-shape hidden_states + # tensor so downstream consumers (LogitsProcessor, frozen-KV MTP + # worker) see the same ``[T, H]`` contract. Modeled after vLLM's + # ``Gemma4Model.fast_prefill_forward`` + # (vllm/model_executor/models/gemma4.py:1190-1273), reimplemented in + # SGLang's own forward + attention-metadata abstractions. + self._num_kv_shared_layers = int( + getattr(config, "num_kv_shared_layers", 0) or 0 + ) + self._first_kv_shared_layer_idx = ( + config.num_hidden_layers - self._num_kv_shared_layers + ) + try: + _server_args = get_global_server_args() + except Exception: + _server_args = None + self._kv_sharing_fast_prefill_enabled = ( + bool(getattr(_server_args, "kv_sharing_fast_prefill", False)) + and self._num_kv_shared_layers > 0 + ) + self.post_init() def get_input_embeddings(self) -> nn.Embedding: @@ -919,6 +953,128 @@ def project_per_layer_inputs( # Combine: (projection + per_layer_inputs) * scale return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale + # -- KV-sharing fast-prefill (YOCO) helpers ---------------------------- + + def _can_run_yoco(self, forward_batch: ForwardBatch) -> bool: + """Predicate for the YOCO fast-prefill branch. + + Returns True only when (a) the flag is on, (b) the model declares + KV-shared layers, (c) the batch is in EXTEND mode (and not + TARGET_VERIFY), (d) ``extend_seq_lens`` is populated, (e) the model + does not use PLE (gated out for v0 because the per-layer-input slice + would need its own gather), and (f) no request in the batch asks + for input-token logprobs (the LM-head gather then differs from the + per-request last-extend-token gather YOCO produces). + """ + if not self._kv_sharing_fast_prefill_enabled: + return False + if forward_batch.extend_seq_lens is None: + return False + mode = forward_batch.forward_mode + if not mode.is_extend() or mode.is_target_verify(): + return False + # Gate out requests that asked for input-token logprobs: the + # LM-head pruning then uses a different index set than the + # per-request last-extend-token positions YOCO computes. + start_lens = forward_batch.extend_logprob_start_lens_cpu + seq_lens_cpu = forward_batch.extend_seq_lens_cpu + if start_lens and seq_lens_cpu: + if any(s < e for s, e in zip(start_lens, seq_lens_cpu)): + return False + # Gate out Eagle3 aux-hidden-state captures that fall inside the + # back half: under YOCO the back-half ``hidden_states`` are gathered + # ``[B, H]`` and the captured tensor would be the wrong shape. + # Front-half captures (layers in ``[0, first_kv_shared_layer_idx)``) + # are still safe and we could allow them, but for v0 we simply + # bypass YOCO whenever any capture is requested. + if self.layers_to_capture: + return False + # YOCO touches only the back half; nothing to gain when the model + # has only PP first/middle ranks (no back-half layers on this rank) + # or when the back-half boundary falls outside the local layer range. + if self.end_layer <= self._first_kv_shared_layer_idx: + return False + if self.start_layer >= self._first_kv_shared_layer_idx: + return False + return True + + def _build_cross_decoder_last_token_index( + self, forward_batch: ForwardBatch + ) -> torch.Tensor: + """Per-request last-extend-token row index in the flat token buffer. + + Mirrors ``LogitsProcessor._get_pruned_states`` (see + ``layers/logits_processor.py``). Returns an int64 ``[B]`` tensor + suitable for fancy-indexing ``hidden_states`` and ``positions``. + """ + padded = forward_batch.padded_static_len + if padded is None or padded < 0: + return torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1 + idx = torch.arange( + forward_batch.extend_seq_lens.shape[0], + device=forward_batch.extend_seq_lens.device, + ) + return idx * padded + forward_batch.extend_seq_lens - 1 + + def _run_cross_decoder_with_yoco( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + last_token_index: torch.Tensor, + per_layer_inputs: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Run the cross-decoder (KV-shared) layers on the gathered last- + token rows. + + Swaps ``forward_batch.forward_mode`` to ``DECODE`` for the duration + of the back-half so the triton attention backend rebuilds its + metadata with ``qo_indptr=[0..B]`` (Q has one token per request) and + ``kv_indptr=cumsum(seq_lens)`` (K/V covers the full prefix+extend + from the donor layer's KV pool). The original mode and metadata + are restored before returning so the caller sees no externally + visible mutation. + + When ``per_layer_inputs`` is provided (PLE-enabled variants like + Gemma-4 E4B/E2B), the per-layer-input tensor is also gathered to + the same rows so the back-half PLE add operates on the correct + per-request positions. + """ + attn_backend = get_attn_backend() + saved_mode = forward_batch.forward_mode + saved_metadata = getattr(attn_backend, "forward_metadata", None) + + gathered_hidden = hidden_states[last_token_index] + gathered_positions = positions[last_token_index] + gathered_per_layer_inputs = ( + per_layer_inputs[last_token_index] if per_layer_inputs is not None else None + ) + + try: + forward_batch.forward_mode = ForwardMode.DECODE + attn_backend.init_forward_metadata(forward_batch) + cross_hidden = gathered_hidden + for layer_idx in range(self._first_kv_shared_layer_idx, self.end_layer): + if gathered_per_layer_inputs is not None: + per_layer_input = gathered_per_layer_inputs[:, layer_idx, :] + else: + per_layer_input = None + layer_out = self.layers[layer_idx]( + positions=gathered_positions, + hidden_states=cross_hidden, + per_layer_input=per_layer_input, + forward_batch=forward_batch, + **kwargs, + ) + cross_hidden = layer_out[0] + finally: + forward_batch.forward_mode = saved_mode + if saved_metadata is not None: + attn_backend.forward_metadata = saved_metadata + + return cross_hidden + def forward( self, input_ids: torch.Tensor, @@ -943,9 +1099,9 @@ def forward( ) hidden_states = input_embeds else: - assert ( - pp_proxy_tensors is not None - ), "pp_proxy_tensors is required on non-first PP ranks" + assert pp_proxy_tensors is not None, ( + "pp_proxy_tensors is required on non-first PP ranks" + ) hidden_states = pp_proxy_tensors["hidden_states"] # PLE inputs were computed on rank 0 and forwarded along the # pipeline; non-PLE models simply omit the key. @@ -954,7 +1110,17 @@ def forward( aux_hidden_states = [] num_layers = self.config.num_hidden_layers - for layer_idx in range(self.start_layer, self.end_layer): + # KV-sharing fast-prefill (YOCO) branch — when applicable, run the + # back-half (KV-shared) layers only on the per-request last-extend- + # token rows and scatter the result back into the full-shape + # ``hidden_states`` tensor before the final norm. See + # ``_can_run_yoco`` for the predicate and the audit doc at + # ``runs/20260523_gemma4_31b_it_h100_sota_humanize/yoco/draft.md`` + # for the design rationale. + use_yoco = self.pp_group.is_last_rank and self._can_run_yoco(forward_batch) + loop_end = self._first_kv_shared_layer_idx if use_yoco else self.end_layer + + for layer_idx in range(self.start_layer, loop_end): if layer_idx in self.layers_to_capture: aux_hidden_states.append(hidden_states) @@ -974,6 +1140,32 @@ def forward( # Gemma4DecoderLayer.forward always returns (hidden_states, None); # the residual is fused inside the layer, so nothing to thread. + if use_yoco: + # Gather per-request last-extend-token rows, run the KV-shared + # back half on the reduced batch, then scatter the cross-decoder + # output back into a full-shape ``hidden_states`` tensor so the + # downstream final norm + LM head + frozen-KV MTP worker see + # the existing ``[T, H]`` contract. Only the gathered rows + # carry post-cross-decoder values; non-gathered rows retain + # their pre-cross-decoder (self-decoder output) values, which + # is harmless because the sampler reads only the gathered rows + # (see LogitsProcessor._get_pruned_states). + last_token_index = self._build_cross_decoder_last_token_index(forward_batch) + cross_hidden = self._run_cross_decoder_with_yoco( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + last_token_index=last_token_index, + per_layer_inputs=per_layer_inputs, + **kwargs, + ) + # Clone the self-decoder output before scattering so we never + # alias a tensor that downstream consumers (CUDA-graph output + # buffers) may have weakly referenced. + full_hidden = hidden_states.clone() + full_hidden.index_copy_(0, last_token_index, cross_hidden) + hidden_states = full_hidden + if not self.pp_group.is_last_rank: # cuda_graph_runner allocates a fixed PP-proxy schema of # {hidden_states, residual} and KeyErrors if a model omits a key. diff --git a/test/registered/unit/models/test_gemma4_yoco.py b/test/registered/unit/models/test_gemma4_yoco.py new file mode 100644 index 000000000000..b8b326f950c8 --- /dev/null +++ b/test/registered/unit/models/test_gemma4_yoco.py @@ -0,0 +1,336 @@ +"""Unit tests for the YOCO fast-prefill helpers on ``Gemma4TextModel``. + +These tests cover three things: + +* ``_can_run_yoco`` — the predicate that decides whether the YOCO branch + should fire. Covers the eligible case plus every individual reject + condition. +* ``_build_cross_decoder_last_token_index`` — the gather index, matching + ``LogitsProcessor._get_pruned_states`` (cumsum-1 in the non-padded + case). +* ``_run_cross_decoder_with_yoco`` — verifies the temporary mutation of + ``forward_batch.forward_mode`` and the restoration of the attention + backend's ``forward_metadata`` after the back-half runs, with all + external collaborators stubbed. + +The actual end-to-end correctness (tokens match between YOCO-on and +YOCO-off greedy sampling) is exercised by the live integration test that +runs against a Gemma-4 31B-IT server during PR-C benchmarking, not here. +""" + +import types +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.models.gemma4_causal import Gemma4TextModel + + +def _stub_model( + *, + num_kv_shared_layers=5, + num_hidden_layers=60, + hidden_size_per_layer_input=0, + kv_sharing_fast_prefill_enabled=True, + start_layer=0, + end_layer=None, + layers_to_capture=None, +): + """Build a bare ``Gemma4TextModel`` shell with just the fields the YOCO + helpers read. We deliberately avoid ``Gemma4TextModel.__init__`` (which + would try to download weights, allocate embeddings, etc.) and inject + only the attributes ``_can_run_yoco`` and ``_build_cross_decoder_*`` + consult. + """ + if end_layer is None: + end_layer = num_hidden_layers + model = Gemma4TextModel.__new__(Gemma4TextModel) + model._num_kv_shared_layers = num_kv_shared_layers + model._first_kv_shared_layer_idx = num_hidden_layers - num_kv_shared_layers + model._kv_sharing_fast_prefill_enabled = kv_sharing_fast_prefill_enabled + model.hidden_size_per_layer_input = hidden_size_per_layer_input + model.start_layer = start_layer + model.end_layer = end_layer + model.layers_to_capture = layers_to_capture or [] + return model + + +def _stub_forward_batch( + *, + mode=None, + extend_seq_lens=None, + extend_seq_lens_cpu=None, + extend_logprob_start_lens_cpu=None, + padded_static_len=-1, +): + """Lightweight namespace standing in for ``ForwardBatch``.""" + if mode is None: + mode = ForwardMode.EXTEND + return types.SimpleNamespace( + forward_mode=mode, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + extend_logprob_start_lens_cpu=extend_logprob_start_lens_cpu, + padded_static_len=padded_static_len, + ) + + +class TestCanRunYoco(unittest.TestCase): + def setUp(self): + self.model = _stub_model() + self.batch = _stub_forward_batch( + mode=ForwardMode.EXTEND, + extend_seq_lens=torch.tensor([4, 7, 3], dtype=torch.int32), + extend_seq_lens_cpu=[4, 7, 3], + extend_logprob_start_lens_cpu=[4, 7, 3], # no input logprobs + ) + + def test_eligible_batch_returns_true(self): + self.assertTrue(self.model._can_run_yoco(self.batch)) + + def test_flag_off_returns_false(self): + self.model._kv_sharing_fast_prefill_enabled = False + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_no_kv_shared_layers_returns_false(self): + self.model = _stub_model(num_kv_shared_layers=0) + # When no shared layers, ``end_layer <= first_kv_shared_layer_idx`` + # because first_kv_shared_layer_idx == num_hidden_layers. + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_decode_mode_returns_false(self): + self.batch.forward_mode = ForwardMode.DECODE + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_target_verify_returns_false(self): + self.batch.forward_mode = ForwardMode.TARGET_VERIFY + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_missing_extend_seq_lens_returns_false(self): + self.batch.extend_seq_lens = None + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_ple_enabled_still_returns_true(self): + # PLE-enabled variants (Gemma-4 E4B/E2B) are supported: the YOCO + # branch gathers per_layer_inputs alongside hidden_states. + self.model.hidden_size_per_layer_input = 256 + self.assertTrue(self.model._can_run_yoco(self.batch)) + + def test_input_logprobs_returns_false(self): + # Any request with start < extend_len => input logprobs requested. + self.batch.extend_logprob_start_lens_cpu = [0, 7, 3] + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_eagle3_aux_capture_returns_false(self): + self.model.layers_to_capture = [55] # inside back half + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_pp_no_back_half_layers_returns_false(self): + # Front-only PP rank: no layers in the back half on this rank. + self.model = _stub_model(start_layer=0, end_layer=30) + self.assertFalse(self.model._can_run_yoco(self.batch)) + + def test_pp_no_front_half_layers_returns_false(self): + # Back-only PP rank. + self.model = _stub_model(start_layer=55, end_layer=60) + self.assertFalse(self.model._can_run_yoco(self.batch)) + + +class TestBuildCrossDecoderLastTokenIndex(unittest.TestCase): + def test_non_padded_matches_cumsum_minus_one(self): + model = _stub_model() + extend_lens = torch.tensor([4, 7, 3, 1], dtype=torch.int32) + batch = _stub_forward_batch( + extend_seq_lens=extend_lens, + extend_seq_lens_cpu=[4, 7, 3, 1], + extend_logprob_start_lens_cpu=[4, 7, 3, 1], + padded_static_len=-1, + ) + last_idx = model._build_cross_decoder_last_token_index(batch) + # Expected: cumsum([4,7,3,1]) - 1 = [3, 10, 13, 14] + torch.testing.assert_close( + last_idx, + torch.tensor([3, 10, 13, 14], dtype=torch.int32), + check_dtype=False, + ) + + def test_padded_static_len(self): + model = _stub_model() + extend_lens = torch.tensor([3, 5, 2], dtype=torch.int32) + batch = _stub_forward_batch( + extend_seq_lens=extend_lens, + extend_seq_lens_cpu=[3, 5, 2], + extend_logprob_start_lens_cpu=[3, 5, 2], + padded_static_len=8, + ) + last_idx = model._build_cross_decoder_last_token_index(batch) + # Padded layout: each req gets 8 slots; last valid token is at + # i*8 + extend_len[i] - 1 => [0+3-1, 8+5-1, 16+2-1] = [2, 12, 17]. + torch.testing.assert_close( + last_idx, + torch.tensor([2, 12, 17], dtype=torch.int64), + check_dtype=False, + ) + + +class TestRunCrossDecoderWithYoco(unittest.TestCase): + """Verify the mode-mutation + metadata-restore contract.""" + + def _make_model_with_back_half_layers(self, num_back_half=2, hidden=8): + model = _stub_model(num_kv_shared_layers=num_back_half, num_hidden_layers=4) + # The back-half layers must respond to ``__call__`` like Gemma4 + # decoder layers do (return (hidden_states, None)). We replace them + # with simple Linear stubs that increment by a known constant per + # layer so the test can verify which layers ran. + model.layers = [None] * 4 + for i in range(model._first_kv_shared_layer_idx, 4): + layer = MagicMock() + layer.return_value = ( + # Output: hidden_states + (i+1) — distinct per layer + torch.zeros(0), # placeholder; replaced in side_effect below + None, + ) + layer.side_effect = ( + lambda positions, hidden_states, per_layer_input, forward_batch, _i=i: ( + hidden_states + (_i + 1), + None, + ) + ) + model.layers[i] = layer + return model + + def test_mode_and_metadata_are_restored(self): + model = self._make_model_with_back_half_layers(num_back_half=2) + original_mode = ForwardMode.EXTEND + original_metadata = object() # sentinel + + fb = _stub_forward_batch( + mode=original_mode, + extend_seq_lens=torch.tensor([3, 2], dtype=torch.int32), + extend_seq_lens_cpu=[3, 2], + extend_logprob_start_lens_cpu=[3, 2], + ) + + mock_backend = MagicMock() + mock_backend.forward_metadata = original_metadata + # init_forward_metadata replaces forward_metadata with a NEW sentinel + # whenever called inside the YOCO scope. We assert the swap happens + # AND the original is restored on exit. + new_metadata = object() + + def init_side_effect(_fb): + mock_backend.forward_metadata = new_metadata + + mock_backend.init_forward_metadata.side_effect = init_side_effect + + with patch( + "sglang.srt.models.gemma4_causal.get_attn_backend", + return_value=mock_backend, + ): + positions = torch.arange(5) + hidden = torch.zeros(5, 4) # [T=5, H=4] + last_idx = torch.tensor([2, 4]) # cumsum([3,2])-1 + out = model._run_cross_decoder_with_yoco( + positions=positions, + hidden_states=hidden, + forward_batch=fb, + last_token_index=last_idx, + ) + + # The cross-decoder ran on the gathered rows (shape [2, 4]). + self.assertEqual(out.shape, (2, 4)) + # Both back-half layers fired (each added i+1 = 3 then 4 => sum 7). + self.assertTrue(torch.equal(out, torch.full((2, 4), 7.0))) + # Mode is restored. + self.assertEqual(fb.forward_mode, original_mode) + # Metadata is restored. + self.assertIs(mock_backend.forward_metadata, original_metadata) + # init_forward_metadata was called exactly once (to build the + # decode-shaped metadata for the back half). + self.assertEqual(mock_backend.init_forward_metadata.call_count, 1) + + def test_per_layer_inputs_are_gathered(self): + """For PLE-enabled variants, per_layer_inputs must be gathered to + the same rows as hidden_states before the back-half runs.""" + model = self._make_model_with_back_half_layers(num_back_half=2) + captured_per_layer_inputs = [] + + def layer_side_effect(positions, hidden_states, per_layer_input, forward_batch): + captured_per_layer_inputs.append( + None if per_layer_input is None else tuple(per_layer_input.shape) + ) + return (hidden_states + 1, None) + + for layer in model.layers[model._first_kv_shared_layer_idx :]: + layer.side_effect = layer_side_effect + + fb = _stub_forward_batch( + mode=ForwardMode.EXTEND, + extend_seq_lens=torch.tensor([3, 2], dtype=torch.int32), + extend_seq_lens_cpu=[3, 2], + extend_logprob_start_lens_cpu=[3, 2], + ) + mock_backend = MagicMock() + mock_backend.forward_metadata = object() + + # Provide a per_layer_inputs tensor of shape [T=5, num_layers=4, ple_dim=16]. + ple = torch.randn(5, 4, 16) + + with patch( + "sglang.srt.models.gemma4_causal.get_attn_backend", + return_value=mock_backend, + ): + model._run_cross_decoder_with_yoco( + positions=torch.arange(5), + hidden_states=torch.zeros(5, 4), + forward_batch=fb, + last_token_index=torch.tensor([2, 4]), + per_layer_inputs=ple, + ) + + # Each back-half layer should have received a per_layer_input of + # shape [B=2, ple_dim=16] (gathered to last-token rows + sliced by layer). + self.assertEqual(len(captured_per_layer_inputs), 2) + self.assertEqual(captured_per_layer_inputs, [(2, 16), (2, 16)]) + + def test_exception_in_back_half_still_restores(self): + """If a back-half layer raises, mode and metadata must still revert.""" + model = self._make_model_with_back_half_layers(num_back_half=1) + # Make the single back-half layer raise. + model.layers[3].side_effect = RuntimeError("intentional") + + original_mode = ForwardMode.EXTEND + original_metadata = object() + fb = _stub_forward_batch( + mode=original_mode, + extend_seq_lens=torch.tensor([2], dtype=torch.int32), + extend_seq_lens_cpu=[2], + extend_logprob_start_lens_cpu=[2], + ) + mock_backend = MagicMock() + mock_backend.forward_metadata = original_metadata + mock_backend.init_forward_metadata.side_effect = lambda _fb: setattr( + mock_backend, "forward_metadata", object() + ) + + with patch( + "sglang.srt.models.gemma4_causal.get_attn_backend", + return_value=mock_backend, + ): + with self.assertRaises(RuntimeError): + model._run_cross_decoder_with_yoco( + positions=torch.arange(2), + hidden_states=torch.zeros(2, 4), + forward_batch=fb, + last_token_index=torch.tensor([1]), + ) + + # Even on exception, mode and metadata are restored. + self.assertEqual(fb.forward_mode, original_mode) + self.assertIs(mock_backend.forward_metadata, original_metadata) + + +if __name__ == "__main__": + unittest.main()