Merge pull request #46 from smb209/fix/mxfp8-a-side-offsets-cudagraph

Pre-warm MXFP8 A-side offsets at load so cudagraph capture cannot hit a host sync

Author: RobTand

CHANGELOG.md

@@ -1,5 +1,40 @@
 # Changelog
 
+## Unreleased
+
+- **Fixed: `FULL_DECODE_ONLY` capture aborted the load in the MXFP8 dense
+  lane.** `_SfOffsetCache` computes swizzled-plane offsets on the host and
+  moves them with an unpinned `.to(device)`.
+  `process_weights_after_loading` resolved the **B side** from weight
+  dimensions known at load, but the **A side** is keyed by the runtime row
+  count, so under `cudagraph_mode=FULL_DECODE_ONLY` the first call at each
+  capture size landed *inside* the capture region and raised `Cannot copy
+  between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor
+  is pinned`. It surfaced as `cudaErrorStreamCaptureUnjoined` at
+  `capture_end()`, because DeepSeek-V4 reaches this lane from inside
+  `maybe_execute_in_parallel`, leaving that side stream unjoined. The A-side
+  offsets are now pre-warmed at load for every `cudagraph_capture_sizes` entry
+  — `docs/KERNELS.md` CUDA-graph safety rule 3, mirroring `cb_gemv_v2_prepare`
+  and `cb_moe_persistent_b_prepare`. **No numerics change**: the same offsets,
+  computed earlier. Eager serving is unaffected (the reader returns `()` when
+  no capture sizes are configured, and the pre-warm is then a no-op).
+
+  Measured on one GB10 / DGX Spark (`sm_121`, arm64), vLLM
+  0.26.1rc1.dev515, torch 2.11.0+cu130, a DeepSeek-V4-Flash-0731 NVFP4-CB
+  artifact at TP=1, `--kv-cache-dtype fp8`, `--max-model-len 8192`,
+  `--compilation-config {"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY",
+  "cudagraph_capture_sizes":[1,2,4,8]}`: capture previously failed the load
+  outright; it now succeeds, and single-stream decode improves **12.05 ->
+  14.08 tok/s (+16.8%)**, 83.0 -> 71.0 ms/token.
+
+  Per-token logprobs under capture are **bit-equal** to `--enforce-eager` on
+  7 of the 8 probe prompts. The 8th is excluded on the basis of an
+  eager-vs-eager control run *before* the captured arm was compared: it is a
+  near-tie that flips between two identical eager runs, so it cannot serve as
+  a reference. Protocol: prefix caching off, serial single requests — under
+  the serving default (prefix caching on) and concurrency, eager is not
+  run-to-run reproducible either.
+
 ## 0.8.10 — 2026-08-18
 
 - **0.8.9 could not load a per-expert split-format (mixed) expert bank at

gridbook/mxfp8_dense_lane.py

@@ -96,6 +96,54 @@ def get(self, ext, rows: int, k: int, *, is_b: bool,
 _OFFSETS = _SfOffsetCache()
 
 
+def _cudagraph_capture_sizes() -> tuple[int, ...]:
+    """Decode row counts vLLM will capture graphs at, or ``()`` if unknown.
+
+    Read the same way ``ops.warn_if_capture_sizes_exceed_the_decode_gates``
+    reads it: vLLM imported inside the function so ``import gridbook`` stays
+    vLLM-free, and every failure degrading to silence rather than breaking a
+    load, because the shape of the compilation config has moved across
+    releases. Returning ``()`` costs only the pre-warm below; it cannot make
+    serving wrong.
+    """
+    try:
+        from vllm.config import get_current_vllm_config
+
+        compilation = get_current_vllm_config().compilation_config
+        sizes = getattr(compilation, "cudagraph_capture_sizes", None) or ()
+        return tuple(sorted({int(s) for s in sizes if int(s) > 0}))
+    except Exception:  # noqa: BLE001 — advisory only; never break a load
+        return ()
+
+
+def _prewarm_activation_offsets(ext, k: int, device) -> None:
+    """Populate the A-side offset cache for every graph capture size.
+
+    WHY THIS EXISTS. ``_SfOffsetCache.get`` computes offsets on the host and
+    moves them with an unpinned ``.to(device)``. That copy is illegal inside a
+    CUDA graph capture ("Cannot copy between CPU and CUDA tensors during CUDA
+    graph capture unless the CPU tensor is pinned"), so any key first seen
+    *during* capture is a hard failure rather than a slow path.
+
+    The B side never had this problem: ``process_weights_after_loading``
+    resolves it from weight dimensions known at load. The A side is keyed by
+    the runtime row count, which under ``FULL_DECODE_ONLY`` is first seen
+    inside the capture region — one first-time miss per capture size.
+
+    Doing it here is ``docs/KERNELS.md`` CUDA-graph safety rule 3 ("all
+    device-side constants and per-device kernel setup happen once, at model
+    load"), and mirrors ``cb_gemv_v2_prepare`` / ``cb_moe_persistent_b_prepare``
+    / ``cb_fused_fp4v2_prepare``, which are called from this same hook for the
+    same reason.
+
+    Not wrapped in try/except: at load these are the identical calls the
+    forward will make, so a failure here is a real defect and should surface
+    at load rather than mid-capture.
+    """
+    for rows in _cudagraph_capture_sizes():
+        _OFFSETS.get(ext, rows, k, is_b=False, device=device)
+
+
 def _quantize_activations(ext, x: torch.Tensor) -> tuple[torch.Tensor,
                                                          torch.Tensor]:
     """bf16 ``[M, K]`` -> (e4m3 ``[M, K]``, swizzled SF plane)."""
@@ -159,6 +207,12 @@ def process_weights_after_loading(self, layer) -> None:
             ext = _require_lane_ext(device)
             w = layer.weight.data
             n, k = int(w.shape[0]), int(w.shape[1])
+            # A-side (activation) offsets, for graph capture. Both call sites
+            # that need them -- ``_quantize_activations`` and the BMM branch of
+            # ``apply`` -- key on the runtime row count at this same ``k``
+            # (``k == layer.weight.shape[1]`` for both), so one pre-warm here
+            # covers both paths.
+            _prewarm_activation_offsets(ext, k, device)
             sf_rm = layer.weight_scale.data
             del layer.weight_scale
             is_bmm = bool(getattr(layer, "is_bmm", False))

tests/test_mxfp8_dense_lane.py

@@ -110,3 +110,103 @@ def test_audited_and_broken_sets_stay_disjoint_for_fp8_source_entries():
     for wire in (WIRE_FP8_BLOCK128, WIRE_MXFP8_G32):
         fmt = sp.FORMATS[wire]
         assert not (fmt.audited_backends & set(fmt.known_broken_backends))
+
+
+# --- A-side offset pre-warm: the CUDA-graph capture prerequisite ------------
+#
+# ``_SfOffsetCache.get`` computes offsets on the host and moves them with an
+# unpinned ``.to(device)``. Inside a CUDA graph capture that copy is a hard
+# error, so every key must already exist before capture starts. The B side was
+# always resolved at load from weight dimensions; the A side is keyed by the
+# runtime row count, which under FULL_DECODE_ONLY is first seen INSIDE the
+# capture region. These cover the reader and the pre-warm on CPU; the served
+# proof is a graph capture on the GPU image.
+
+class _FakeExt:
+    """Records the offset requests the lane makes."""
+
+    def __init__(self):
+        self.calls = []
+
+    def mxfp8_sf_offsets(self, rows, k, is_b):
+        self.calls.append((int(rows), int(k), bool(is_b)))
+        return torch.zeros(4, dtype=torch.int32)
+
+
+def test_capture_sizes_reader_is_silent_without_vllm():
+    """No vLLM in the CPU tier: the reader must degrade to (), not raise --
+    a load must never fail for a compilation-config schema reason."""
+    from gridbook.mxfp8_dense_lane import _cudagraph_capture_sizes
+    assert _cudagraph_capture_sizes() == ()
+
+
+def test_capture_sizes_reader_dedups_sorts_and_drops_nonpositive(monkeypatch):
+    """Exercise the REAL reader against a stub config.
+
+    ``monkeypatch.setitem`` on ``sys.modules`` is deliberate: it restores the
+    entries afterwards, so this does not leak stub ``vllm`` modules into later
+    files the way CONTRIBUTING.md warns about.
+    """
+    import sys
+    import types
+
+    import gridbook.mxfp8_dense_lane as lane
+
+    config = types.SimpleNamespace(
+        compilation_config=types.SimpleNamespace(
+            cudagraph_capture_sizes=[8, 2, 2, 0, 4, 1, -3]))
+    stub = types.ModuleType("vllm.config")
+    stub.get_current_vllm_config = lambda: config
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(sys.modules, "vllm.config", stub)
+
+    assert lane._cudagraph_capture_sizes() == (1, 2, 4, 8)
+
+
+def test_capture_sizes_reader_degrades_to_silence_on_a_hostile_config(
+        monkeypatch):
+    """Schema drift must not fail a load -- the reason ops.py warns instead of
+    raising. A config whose accessor explodes yields (), not an exception."""
+    import sys
+    import types
+
+    import gridbook.mxfp8_dense_lane as lane
+
+    def boom():
+        raise RuntimeError("compilation_config moved again")
+
+    stub = types.ModuleType("vllm.config")
+    stub.get_current_vllm_config = boom
+    monkeypatch.setitem(sys.modules, "vllm", types.ModuleType("vllm"))
+    monkeypatch.setitem(sys.modules, "vllm.config", stub)
+
+    assert lane._cudagraph_capture_sizes() == ()
+
+
+def test_prewarm_populates_a_side_for_every_capture_size(monkeypatch):
+    """The fix proper: one A-side entry per capture size, at the layer's K."""
+    import gridbook.mxfp8_dense_lane as lane
+
+    monkeypatch.setattr(lane, "_OFFSETS", lane._SfOffsetCache())
+    monkeypatch.setattr(lane, "_cudagraph_capture_sizes", lambda: (1, 2, 4, 8))
+    ext = _FakeExt()
+    lane._prewarm_activation_offsets(ext, 4096, torch.device("cpu"))
+
+    assert ext.calls == [(1, 4096, False), (2, 4096, False),
+                         (4, 4096, False), (8, 4096, False)]
+    # and a subsequent forward-time lookup is a cache HIT, i.e. no host copy
+    # would be issued inside a capture
+    before = len(ext.calls)
+    lane._OFFSETS.get(ext, 4, 4096, is_b=False, device=torch.device("cpu"))
+    assert len(ext.calls) == before
+
+
+def test_prewarm_is_a_noop_when_capture_sizes_are_unknown(monkeypatch):
+    """Eager serving (or an unreadable config) must not pay for this."""
+    import gridbook.mxfp8_dense_lane as lane
+
+    monkeypatch.setattr(lane, "_OFFSETS", lane._SfOffsetCache())
+    monkeypatch.setattr(lane, "_cudagraph_capture_sizes", tuple)
+    ext = _FakeExt()
+    lane._prewarm_activation_offsets(ext, 4096, torch.device("cpu"))
+    assert ext.calls == []