sgl-project/sglang · error · RuntimeError

missing captured KV on {attr}

Error message

missing captured KV on {attr}

What it means

collect_captured_kv_from_blocks reads the cached (K, V) tensors each attention block stored under _cached_kv_pre or _cached_kv_post (depending on mode). If any block's cache is empty, capture was never enabled or the attention forward did not run/populate it, and the collector raises this RuntimeError.

Source

Thrown at python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming_refiner.py:77

    if mode == "pre_rope":
        attr, clear_attr = "_kv_cache_capture", "_cached_kv_pre"
    elif mode == "post_rope":
        attr, clear_attr = "_tf_capture_kv", "_cached_kv_post"
    else:
        raise ValueError(f"unsupported capture mode: {mode}")
    for block in transformer.transformer_blocks:
        setattr(block.attn1, attr, bool(enable))
        if enable and hasattr(block.attn1, clear_attr):
            setattr(block.attn1, clear_attr, None)


def collect_captured_kv_from_blocks(transformer: nn.Module, mode: str):
    attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post"
    out = []
    for block in transformer.transformer_blocks:
        cached = getattr(block.attn1, attr, None)
        if cached is None:
            raise RuntimeError(f"missing captured KV on {attr}")
        out.append(cached)
        setattr(block.attn1, attr, None)
    return out


# --------------------------------------------------------------------------- #
# Absolute-position RoPE (port of refiner.py:721-750)
# --------------------------------------------------------------------------- #
def build_rotary_emb_for_absolute_positions(
    *, transformer, batch_size, frame_positions, height, width, device, fps
):
    rope = transformer.rope
    patch_size_t = int(rope.patch_size_t)
    patch_size = int(rope.patch_size)
    f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device)
    if patch_size_t > 1:
        f_positions = f_positions[::patch_size_t]
    grid_h = torch.arange(0, height, patch_size, dtype=torch.float32, device=device)

View on GitHub (pinned to 0132848349)

Solutions

1. Always pair set_capture_flag_on_blocks(..., enable=True, mode=M) with a forward, then collect_captured_kv_from_blocks(..., mode=M) with the same M
2. Do not collect twice without re-running the forward — the first collect clears the caches
3. Ensure the attention module honors the capture attrs (custom kernels may bypass them)

Example fix

# before
set_capture_flag_on_blocks(tf, enable=True, mode="pre_rope")
kvs = collect_captured_kv_from_blocks(tf, mode="post_rope")  # mode mismatch -> None
# after
set_capture_flag_on_blocks(tf, enable=True, mode="pre_rope")
_ = tf(...)
kvs = collect_captured_kv_from_blocks(tf, mode="pre_rope")

set_capture_flag_on_blocks(tf, enable=True, mode=mode)
_ = tf(*model_inputs)  # forward must run to populate caches
attrs_ok = all(getattr(b.attn1, "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post", None) is not None for b in tf.transformer_blocks)

def kv_captured(tf, mode: str) -> bool:
    attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post"
    return all(getattr(b.attn1, attr, None) is not None for b in tf.transformer_blocks)

try:
    kvs = collect_captured_kv_from_blocks(tf, mode)
except RuntimeError:
    set_capture_flag_on_blocks(tf, enable=True, mode=mode)
    _ = tf(*model_inputs)
    kvs = collect_captured_kv_from_blocks(tf, mode)

When it happens

Trigger: Calling collect_captured_kv_from_blocks before running a forward pass with capture enabled (set_capture_flag_on_blocks not called, or called with a different mode than the collect call); a custom attention implementation ignoring the capture flag.

Common situations: Mismatched mode strings between enable and collect calls; collecting twice (first collect clears the attrs via setattr None); a forward that short-circuited before attention executed.

Related errors

• unsupported capture mode: {mode}
• SANA-WM refiner text encoder must return per-layer hidden_st
• SANA-WM height/width must be divisible by the LTX-2 spatial
• SANA-WM plucker_embedder is not initialized.
• plucker_emb token count {plucker_emb.shape[1]} != latent tok

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/8a702f012ebc64e4. Report an issue: GitHub.