feat(yoco): implement YOCO fast-prefill branch in Gemma4TextModel (PR-B/3)

[pyc96]

Summary

Stacked on PR-A (#11). Adds the model-side YOCO fast-prefill branch to `Gemma4TextModel.forward` so that, when `--kv-sharing-fast-prefill` is on AND the served checkpoint has `num_kv_shared_layers > 0`, the last N decoder layers run only on the per-request last-extend-token rows during EXTEND-mode forwards.

Mirrors vLLM's `fast_prefill_forward` (`vllm/model_executor/models/gemma4.py:1190-1273`), reimplemented in SGLang's own forward + attention-metadata abstractions.

Target model

The Gemma-4 family checkpoints have:

Checkpoint	num_hidden_layers	num_kv_shared_layers	PLE	YOCO applies?
google/gemma-4-26b-a4b-it	30	0	no	No (no shared layers)
google/gemma-4-31B-it	60	0	no	No (no shared layers)
google/gemma-4-E4B-it	42	18	yes (256)	Yes
google/gemma-4-E2B-it	35	20	yes (256)	Yes
google/gemma-4-26b-a4b-it-assistant	4	4	no	No (DECODE-only path)
google/gemma-4-31B-it-assistant	4	4	no	No (DECODE-only path)

PR-C benches against google/gemma-4-E4B-it (4.5B params, ~9 GB bf16, 42 layers split 24 self + 18 cross).

What's in PR-B

File	Change
`python/sglang/srt/models/gemma4_causal.py`	`Gemma4TextModel` gains `_can_run_yoco` (predicate), `_build_cross_decoder_last_token_index` (gather index, mirrors `LogitsProcessor._get_pruned_states`), and `_run_cross_decoder_with_yoco` (back-half runner with mode-swap + metadata-restore). `forward()` branches on `_can_run_yoco`: runs front half on full batch, then gathers, runs cross-decoder on `[B, H]` (and PLE `[B, num_layers, ple_dim]` slice per layer), scatters back into a cloned `[T, H]` tensor.
`test/registered/unit/models/test_gemma4_yoco.py`	16 new unit tests covering predicate (11 cases including PLE-accepted), gather index (2 cases: non-padded, padded-static-len), back-half runner (3 cases: mode+metadata restore, per_layer_inputs gather, exception-restores).

Mechanism

1. Front half: layers `[start_layer, first_kv_shared_layer_idx)` run on full `[T, H]` hidden_states. These layers write their K/V to the pool with `save_kv_cache=True` exactly as today.
2. Build gather index: `last_token_index = cumsum(extend_seq_lens) - 1` (or the padded variant for piecewise CUDA graph). Identical to the index `LogitsProcessor` uses to feed the LM head.
3. `_run_cross_decoder_with_yoco`:
   • Gather `hidden_states[last_token_index]` → `[B, H]`, `positions[last_token_index]` → `[B]`, and (if PLE) `per_layer_inputs[last_token_index]` → `[B, num_layers, ple_dim]`.
   • Temporarily set `forward_batch.forward_mode = ForwardMode.DECODE` and call `attn_backend.init_forward_metadata(forward_batch)` so the triton backend rebuilds `qo_indptr=[0..B]`, `max_extend_len=1`, `kv_indptr=cumsum(seq_lens)`, with `kv_indices` over the full prefix+extend KV span per request.
   • Run layers `[first_kv_shared_layer_idx, end_layer)` on the gathered tensors. KV reads use the donor layer's pool slot (existing `is_kv_shared_layer` + `save_kv_cache=False` mechanism); KV writes don't happen.
   • In a `try/finally`, restore `forward_mode` and the original `forward_metadata` so the caller sees no externally visible mutation.
4. Scatter back: `full_hidden = hidden_states.clone(); full_hidden.index_copy_(0, last_token_index, cross_hidden)`. Non-gathered rows retain self-decoder output (harmless because the sampler reads only the gathered rows).
5. Final norm + LM head proceed unchanged on the full-shape tensor.

Tests

```
$ python test/registered/unit/models/test_gemma4_yoco.py
................
Ran 16 tests in 0.004s
OK
```

Test class	Cases	Covers
TestCanRunYoco	11	flag-on eligible, flag-off, no-KV-shared, decode, target-verify, missing-extend-lens, PLE-accepted, input-logprobs, eagle3-aux-capture, PP-no-back-half, PP-no-front-half
TestBuildCrossDecoderLastTokenIndex	2	non-padded cumsum-1, padded-static-len i*pad+len-1
TestRunCrossDecoderWithYoco	3	mode+metadata restore on success, per_layer_inputs gather, exception-still-restores via try/finally

Predicate (AC-2)

`_can_run_yoco` returns True only when all of:

• Flag on (`get_global_server_args().kv_sharing_fast_prefill`).
• Model has `num_kv_shared_layers > 0`.
• `forward_mode.is_extend()` AND not `is_target_verify()`.
• `extend_seq_lens` populated.
• No request asked for input-token logprobs (per-prompt-token logits would need a different gather than last-token-per-req).
• No Eagle3 aux-hidden capture is wanted in the back half (those captures would see the wrong tensor shape under YOCO; gated for v0).
• The back-half boundary falls inside the local PP rank's layer range.

KV cache integrity (AC-5)

• Front-half layers: `is_kv_shared_layer == False` → `save_kv_cache=True` → write to pool as today.
• Cross-decoder layers: `is_kv_shared_layer == True` → `save_kv_cache=False` → never write (existing semantics, unchanged).
• No mutation of `req_to_token_pool` or `token_to_kv_pool_allocator`.

Benchmark / MMLU

No runtime change for any currently-benchmarked target (26b-a4b-it and 31B-it have num_kv_shared_layers=0 so the predicate returns False and the existing layer loop runs verbatim). E4B-it bench results land in PR-C.

Stack

• PR-A (feat(yoco): add --kv-sharing-fast-prefill server arg + validators (PR-A/3) #11): flag plumbing
• PR-B (this PR): model-side YOCO branch + 16 unit tests
• PR-C: E4B benchmark + per-prompt parity test + tuning results

Stack base: `pyc/yoco-fast-prefill-config` @ `0911a9627`

Plan: `.humanize/yoco-gemma4/refined-plan.md`

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.