Add embedding KV-cache skip fast path (skip redundant per-layer paged-KV write/gather)

[xtangxtang]

What

Embedding requests are pure prefill, no decode. SGLang's forward_extend still writes K/V into the paged KV pool and gathers it right back per layer — but for an embedding model nothing ever reads it back (there is no decode step). This PR adds an opt-in fast path that, when every safety guard holds, skips both the set_kv_buffer write and the pool gather and runs per-sequence SDPA directly on the freshly-computed local K/V. Bit-exact vs the original path.

How

• Opt-in via SGLANG_SKIP_EMBED_KV_CACHE=1; OFF by default → forward_extend runs the original path byte-for-byte (zero regression).
• The change to torch_native_backend.py is purely additive: a guard block is prepended and only returns early (into _forward_extend_kv_skip) when all guards pass, otherwise control falls through to the unchanged original body.

Guards (all must hold, else fall back to the original path):

1. env switch on (SGLANG_SKIP_EMBED_KV_CACHE == "1")
2. save_kv_cache
3. embedding model (not model_config.is_generation) — no decode reads KV back
4. not cross-attention
5. no cached prefix (extend_prefix_lens == 0 for all sequences)
6. chunked-prefill safe (extend_seq_lens == orig_seq_lens when chunking is possible for the server config)
7. pool stores K/V verbatim (store_dtype == dtype, no fp8/uint8 repacking)
8. sliding-window not enabled for the layer (the local path builds no SWA mask, so any SWA layer conservatively falls back)

Usage

SGLANG_SKIP_EMBED_KV_CACHE=1 \
python3 -m sglang.launch_server \
    --model-path <embedding-model> --is-embedding \
    --attention-backend torch_native \
    --max-total-tokens 16384

Off by default — with SGLANG_SKIP_EMBED_KV_CACHE unset, forward_extend runs the original path byte-for-byte.

Companion config (no code)

The --max-total-tokens 16384 above is an orthogonal, pure-memory optimization (apply it with or without the switch). The KV pool only exists for decode, which embedding never does, so it does not need to be sized for a large concurrent decode population — capping it shrinks the paged KV pool from the default (hundreds of GB reserved) to ~1.8 GB, with no effect on throughput or correctness.

Correctness

• torch-only equivalence proof: skip path vs original gather path is bit-identical (max_abs_diff = 0, torch.equal = True) across single-seq, multi-seq, GQA (16Q/8KV), bf16, MHA (8Q/8KV), and encoder-only (causal=False). Rationale: when extend_prefix_lens == 0, the pool-gathered K/V equals the local K/V (guard fix radix cache match sgl-project/sglang#7 ensures verbatim storage), so both SDPA calls receive identical inputs.
• The prior monkey-patch form of this optimization was validated end-to-end at cos = 1.0 on a real embedding server (including a forced-chunked-prefill fallback case).
• Added a self-contained unit test: test/registered/attention/test_embedding_kv_cache_skip.py.

Measured impact (bf16, Qwen3-Embedding-0.6B, single node, 48 cores)

• bs=30, ~30-token sequences: −16.7% batch latency, lossless.
• The gain tracks the KV write/gather + sync share of the forward: largest for high-concurrency + short sequences (bs=30 token30 measured −17.8%), and shrinks toward a few percent on long sequences (which become compute/GEMM-bound).

---

CI States

Latest PR Test (Base): ❌ Run #29314861004
Latest PR Test (Extra): 🚫 Run #29316328565

[xtangxtang]

/tag-and-rerun-ci extra