Files
2026-07-13 12:24:33 +08:00

18 KiB
Raw Permalink Blame History

Per-Request Root Span

This document describes the root span design for MPServerTracingSubscriber: how a single "request" OTel span wraps all child operations for one request, and how the span close is deferred correctly when GPU stores are still in flight.

Problem

Before this change, MPServerTracingSubscriber emitted flat, orphaned child spans (mp.store, mp.retrieve, mp.lookup_prefetch) with no parent context. Traces in Tempo/Jaeger showed disconnected spans with no request-level view.

Design

Each request gets one root "request" span that:

  • Opens at MP_REQUEST_START — the first CPU-synchronous touch of a request_id
  • Nests all child spans beneath it via OTel context propagation
  • Closes at MP_REQUEST_END, deferred if async GPU stores are still in flight

New Events

Four new EventType values, all CPU-synchronous:

Event Published from Purpose
MP_REQUEST_START lookup_prefetch_start(), top of method Open root span at true request arrival
MP_STORE_SUBMITTED store(), before publish_on_stream(MP_STORE_START) Register a pending GPU store before it's enqueued
MP_RETRIEVE_SUBMITTED retrieve(), before publish_on_stream(MP_RETRIEVE_START) Register a pending GPU retrieve before it's enqueued
MP_REQUEST_END end_session(), after session_manager.remove() Signal that the session lifecycle is complete

Deferral Protocol

end_session() is CPU-synchronous; GPU store/retrieve callbacks (MP_STORE_END, MP_RETRIEVE_END) fire later via CUDA host callbacks. Without coordination, MP_REQUEST_END can arrive and close the root span before GPU work finishes — producing orphaned child spans.

The fix: MP_STORE_SUBMITTED and MP_RETRIEVE_SUBMITTED are published before the respective GPU work is enqueued, incrementing _pending_store_count and _pending_retrieve_count. When MP_REQUEST_END arrives:

  • If both counters are zero → close root immediately
  • Otherwise → save the REQUEST_END timestamp; the last MP_STORE_END or MP_RETRIEVE_END to decrement its counter to zero (when the other counter is also zero) closes the root using that saved timestamp

Root end-time is always the REQUEST_END timestamp (the logical request end), not the GPU callback timestamp.

Why MP_RETRIEVE_SUBMITTED is needed: vLLM's IPC completion event is recorded on the CUDA stream between MP_RETRIEVE_START and MP_RETRIEVE_END. When vLLM unblocks on that event, it can call end_session() before the GPU callback for MP_RETRIEVE_END fires. EventBus queue becomes: → MP_RETRIEVE_START → MP_REQUEST_END → MP_RETRIEVE_END Without MP_RETRIEVE_SUBMITTED, _on_session_end sees no in-flight work and closes the root span before the retrieve child span ends.

Root Span Attributes

In addition to session_id, the root "request" span carries three hit rate attributes that are set when MP_LOOKUP_PREFETCH_END is processed:

Attribute OTel type Value
hit_tokens int tokens found in L1+L2 (numerator)
requested_tokens int chunk-aligned tokens submitted for lookup (denominator)
hit_rate float hit_tokens / requested_tokens; 0.0 when denominator is zero

hit_rate is stored as a precomputed float because trace UIs (Tempo, Jaeger) cannot derive it from two integer attributes at query time.

Invariant: these attributes are set at MP_LOOKUP_PREFETCH_END time, while the root span is still open. LP_END always precedes MP_REQUEST_END in the event stream, so the root span is guaranteed to be live in the registry when the attributes are written.

Store-only requests (no lookup_prefetch_start() call) never emit MP_LOOKUP_PREFETCH_END, so the root span will not carry these attributes.

CB path — cb.request span

The same three attributes appear on the "cb.request" root span and are set when CB_LOOKUP_END is processed by BlendTracingSubscriber.

CB_LOOKUP_END carries hit_tokens and requested_tokens in its metadata, computed at the emit site in lmcache/v1/multiprocess/modules/blend.py:

Field Value
hit_tokens storage_hits * chunk_size
requested_tokens (num_tokens // chunk_size) * chunk_size (chunk-aligned)

All three CB_LOOKUP_END emit sites (no-fingerprint-match, no-GPU-context, happy path) populate these fields, so hit_rate is always present on the cb.request span.

A fourth attribute is also set on "cb.request" at CB_LOOKUP_END time:

Attribute OTel type Value
prefix_hits int chunks found via the prefix probe (not fingerprint matching)

Prefix probe

cb_lookup_pre_computed has two lookup paths:

  1. Fingerprint pathBlendTokenRangeMatcher.match_sub_sequence finds sub-sequence matches using polynomial rolling hashes. Covers arbitrary (non-prefix) positions in the token sequence.
  2. Prefix probe — a fallback that runs after the fingerprint path and fills in chunks at contiguous prefix positions not already covered by fingerprint results. It calls token_hasher.compute_chunk_hashes(token_ids) to derive the same storage keys used by cb_store_final and cb_store_pre_computed, then creates CBMatchResult(old_st==cur_st) candidates for uncovered slots. These candidates flow through the same prefetch/poll/evict machinery as fingerprint results.

The prefix probe closes the gap between the MP and CB storage paths: chunks written by cb_store_final (which only registers fingerprints when worker_id in [0, None]) and chunks written via the MP store() path (which uses block hashes incompatible with fingerprint matching) are both visible to cb_lookup_pre_computed through the prefix probe.

Lazy registration

When cb_lookup_pre_computed returns results that came entirely from the prefix probe (i.e. fingerprint_results is empty) and the calling worker is rank 0 or the driver (worker_id in [0, None]), the found prefix chunks are registered into BlendTokenRangeMatcher so that future lookups can find them via the faster fingerprint path. Registration is guarded by BlendTokenRangeMatcher.has_chunk(token_hash) to prevent overwriting existing compact-ID assignments when the range matcher already has entries for the same token sequence.

prefix_hits counts the chunks found exclusively through the prefix probe (after deduplication against fingerprint results). When fingerprint_results is non-empty and prefix candidates fill in additional positions, prefix_hits reflects only the prefix-probe portion of the total storage_hits.

Request Scenarios

Scenario 1 — Full Cache Hit

Path: lookup_prefetch → retrieve → store

CPU  ─[REQUEST_START]─[LP_START]─[LP_END]──[RETR_SUBMITTED]──[STORE_SUBMITTED]─[REQUEST_END]─►
GPU  ──────────────────────────────[RETR_START]─[vLLM_IPC]─[RETR_END]──[STORE_START]─[STORE_END]─►

root "request"  [═══════════════════════════════════════════════════════════════════════════════]
  mp.lookup_prefetch    [══════════]
  mp.retrieve                          [══════════════════]
  mp.store                                                        [══════════════════════]

Root closes at REQUEST_END (deferred until both retrieve and store complete).


Scenario 2 — Cache Miss (no retrieve)

Path: lookup_prefetch → store, no retrieve

CPU  ─[REQUEST_START]─[LP_START]─[LP_END]──────────[STORE_SUBMITTED]─[REQUEST_END]─►
GPU  ───────────────────────────────────────────────────────[STORE_START]─[STORE_END]─►

root "request"  [═══════════════════════════════════════════════════════════════════]
  mp.lookup_prefetch    [══════════]
  mp.store                                                    [══════════════════════]

No retrieve occurred, so mp.retrieve is absent.


Scenario 3 — Lookup Only

Path: lookup_prefetch only, no store

CPU  ─[REQUEST_START]─[LP_START]─[LP_END]─[REQUEST_END]─►

root "request"  [════════════════════════════════════════]
  mp.lookup_prefetch    [══════════]

Root closes immediately at REQUEST_END.


Scenario 4 — Store Only (no lookup)

Path: store with no prior lookup_prefetch_start() call

CPU  ─(no REQUEST_START)──────[STORE_SUBMITTED]─[REQUEST_END]─►
GPU  ──────────────────────────────────[STORE_START]─[STORE_END]─►

root "request" (lazy, created at MP_STORE_START)
                                       [═════════════════════════]
  mp.store                             [══════════════]

MP_REQUEST_START is only emitted from lookup_prefetch_start(). If that path was not taken, _get_or_create_request_span() is called lazily on the first child _on_start(). Root start time equals STORE_START timestamp.


Scenario 5 — REQUEST_END Races GPU Store

end_session() called before the GPU store callback fires.

CPU  ─[REQUEST_START]─[LP_START]─[LP_END]─[STORE_SUBMITTED]─[REQUEST_END]────────────────────►
GPU  ──────────────────────────────────────────────[STORE_START]──────────────[STORE_END]─────►
                                                                       ▲
                                              REQUEST_END arrives here─┘ (before STORE_END)

root "request"  [═══════════════════════════════════════════════════════════════════════════]
  mp.lookup_prefetch    [══════════]
  mp.store                                                    [═══════════════════]
                                                                                  ▲
                   STORE_SUBMITTED → count=1                                      │
                   REQUEST_END → count>0 → defer (save ts)                        │
                   STORE_END → count=0 → _close_request_span(deferred_ts) ────────────────┘

Scenario 6 — Multiple Stores, Deferred Close

Two concurrent stores; root stays open until both complete.

CPU  ─[REQUEST_START]─[LP_START]─[LP_END]─[SUBMITTED×2]─[REQUEST_END]──────────────────────────────────►
GPU  ────────────────────────────────────────────────────[S1_START]─[S1_END]─[S2_START]─[S2_END]────────►

root "request"  [═══════════════════════════════════════════════════════════════════════════════════════]
  mp.lookup_prefetch    [══════════]
  mp.store (1)                                                       [══════════]
  mp.store (2)                                                                    [══════════]
                                                                                            ▲
                   count=2 at REQUEST_END → defer                                           │
                   S1_END → count=1 → still open                                            │
                   S2_END → count=0 → _close_request_span(deferred_ts) ─────────────────────────────┘

Summary

Scenario Root opens Root closes
Full hit MP_REQUEST_START last MP_STORE_END / MP_RETRIEVE_END (stamped at REQUEST_END time)
Cache miss MP_REQUEST_START last MP_STORE_END (stamped at REQUEST_END time)
Lookup only MP_REQUEST_START REQUEST_END (immediate)
Store only MP_STORE_START (lazy) REQUEST_END (immediate)
REQUEST_END races store MP_REQUEST_START last MP_STORE_END (stamped at REQUEST_END time)
REQUEST_END races retrieve MP_REQUEST_START last MP_RETRIEVE_END (stamped at REQUEST_END time)
Multiple stores MP_REQUEST_START last MP_STORE_END (stamped at REQUEST_END time)

Implementation

File Change
lmcache/v1/mp_observability/event.py Add MP_REQUEST_START, MP_STORE_SUBMITTED, MP_RETRIEVE_SUBMITTED, MP_REQUEST_END
lmcache/v1/multiprocess/server.py Emit the 4 events at lookup_prefetch_start(), store(), retrieve(), end_session()
lmcache/v1/mp_observability/subscribers/tracing/mp_server.py Root span logic: _pending_store_count, _pending_retrieve_count, _deferred_session_end_ts; handlers _on_request_start, _on_store_submitted, _on_retrieve_submitted, _on_session_end; helpers _get_or_create_request_span, _close_request_span
lmcache/v1/mp_observability/subscribers/tracing/span_registry.py SpanRegistry: shared dict of open spans keyed by (session_id, span_name) for cross-subscriber parent lookup
tests/v1/mp_observability/subscribers/tracing/test_mp_server.py Tests for all scenarios including retrieve deferral
lmcache/v1/multiprocess/modules/blend.py Prefix probe in cb_lookup_pre_computed; lazy registration; has_chunk on BlendTokenRangeMatcher; prefix_hits in CB_LOOKUP_END metadata
lmcache/v1/mp_observability/subscribers/tracing/cb_server.py Stamp prefix_hits on "cb.request" root span from CB_LOOKUP_END
tests/v1/multiprocess/test_blend_server_v2.py has_chunk unit tests
tests/v1/mp_observability/subscribers/tracing/test_cb_server.py prefix_hits attribute tests

Extending the Span Hierarchy

How the registry works

MPServerTracingSubscriber writes every open span into a shared SpanRegistry while it is live:

registry[(session_id, "request")]       → (root_span, root_ctx)       # open: REQUEST_START → REQUEST_END
registry[(session_id, "retrieve")]      → (retrieve_span, ctx)         # open: RETRIEVE_START → RETRIEVE_END
registry[(session_id, "store")]         → (store_span, ctx)            # open: STORE_START → STORE_END
registry[(session_id, "lookup_prefetch")] → (lp_span, ctx)            # open: LP_START → LP_END

Any subscriber that receives the same SpanRegistry instance can call registry.get_context(session_id, "request") (or any other name) to obtain the OTel context needed to nest a new span.


Example 1 — new span at the same level

To add an l1.read span nested directly under the root "request" span, create a new subscriber file and register it with the shared registry. No existing files need to change.

subscribers/tracing/l1.py:

# SPDX-License-Identifier: Apache-2.0
from opentelemetry import trace

from lmcache.v1.mp_observability.event import Event, EventType
from lmcache.v1.mp_observability.event_bus import EventCallback, EventSubscriber
from lmcache.v1.mp_observability.subscribers.tracing.span_registry import SpanRegistry

_tracer = trace.get_tracer("lmcache_mp.l1")

class L1TracingSubscriber(EventSubscriber):
    def __init__(self, registry: SpanRegistry) -> None:
        self._registry = registry
        self._pending: dict[str, object] = {}

    def get_subscriptions(self) -> dict[EventType, EventCallback]:
        return {
            EventType.L1_READ_RESERVED: self._on_start,
            EventType.L1_READ_FINISHED: self._on_end,
        }

    def _on_start(self, event: Event) -> None:
        parent_ctx = self._registry.get_context(event.session_id, "request")
        span = _tracer.start_span(
            "l1.read", context=parent_ctx, start_time=int(event.timestamp * 1e9)
        )
        self._pending[event.session_id] = span

    def _on_end(self, event: Event) -> None:
        span = self._pending.pop(event.session_id, None)
        if span:
            span.end(end_time=int(event.timestamp * 1e9))

config.py (the only change needed):

registry = SpanRegistry()
bus.register_subscriber(MPServerTracingSubscriber(registry))
bus.register_subscriber(L1TracingSubscriber(registry))   # ← add this line

This produces: request → l1.read (alongside mp.retrieve, mp.store, etc.)


Example 2 — sub-span nested under an existing child span

To nest a span inside mp.retrieve (e.g. an L2 disk load that happens during a retrieve), look up "retrieve" as the parent instead of "request". The "retrieve" entry is live in the registry from MP_RETRIEVE_START to MP_RETRIEVE_END.

    def _on_detail_start(self, event: Event) -> None:
        sid = event.session_id
        # Prefer the immediate parent; fall back to root if retrieve has ended.
        parent_ctx = (
            self._registry.get_context(sid, "retrieve")
            or self._registry.get_context(sid, "request")
        )
        span = _tracer.start_span(
            "l2.disk_load", context=parent_ctx, start_time=int(event.timestamp * 1e9)
        )
        self._pending[sid] = span

This produces a three-level trace: request → mp.retrieve → l2.disk_load.