Class VerifyCudaGraphSession

java.lang.Object
smile.llm.engine.VerifyCudaGraphSession
All Implemented Interfaces:
AutoCloseable
public final class VerifyCudaGraphSession extends Object implements AutoCloseable
Per-device CUDA graph session for the uniform MTP window-verify forward.

Graphs are bucketed by (batch, windowLen, numPages) — unlike DecodeCudaGraphSession (implicit S == 1), verify's query length is a real bucket dimension. Within a bucket, the flat KV write-index buffer and CSR metadata are updated in place before replay(int).

Stage 4 of the verify-CUDA-graph plan drives this session's shouldCapture(int, int, int, int) bookkeeping against real traffic while VerifyCudaGraph.captureEnabled() is false, so beginCapture(int) is never actually invoked yet — only bucket-key transitions and warmup counting are exercised for real.

  • Field Summary

    Fields
    Modifier and Type
    Field
    Description
    static final AtomicInteger
    CAPTURES
    Diagnostic counter (process-wide): graph captures.
    static final AtomicInteger
    CLOSES
    Diagnostic counter (process-wide): session closes.
    static final AtomicInteger
    RESETS
    Diagnostic counter (process-wide): bucket resets.
    static final AtomicInteger
    WARMUPS
    Diagnostic counter (process-wide): eager warmup steps.

  • Method Summary

    Modifier and Type
    Method
    Description
    boolean
    beginCapture(int deviceIndex)
    Begins CUDA graph capture on deviceIndex (call from TP worker thread).
    boolean
    canReplay(int batch, int windowLen, int numPages)
    Returns whether a graph for (batch, windowLen, numPages) can be replayed.
    int
    capturedBatch()
    Returns the batch size of the captured bucket.
    int
    capturedNumPages()
    Returns the numPages of the captured bucket.
    int
    capturedWindowLen()
    Returns the windowLen of the captured bucket.
    void
    close()
     
    void
    endCapture()
    Ends capture and instantiates the graph.
    long
    lastCaptureMs()
    Returns wall time in milliseconds for the last successful capture.
    void
    logCapture(int tpRank)
    Logs (once per bucket) a successful graph capture.
    void
    replay(int tpRank)
    Replays the captured graph (inputs must already be on device).
    boolean
    shouldCapture(int batch, int windowLen, int numPages, int tpRank)
    Marks one eager warmup step for the current bucket.
    static VerifyCudaGraphSession
    tryCreate()
    Creates a session when verify CUDA graph wiring is enabled.

    Methods inherited from class Object

    clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

  • Field Details

    • RESETS

      public static final AtomicInteger RESETS
      Diagnostic counter (process-wide): bucket resets.

    • WARMUPS

      public static final AtomicInteger WARMUPS
      Diagnostic counter (process-wide): eager warmup steps.

    • CAPTURES

      public static final AtomicInteger CAPTURES
      Diagnostic counter (process-wide): graph captures.

    • CLOSES

      public static final AtomicInteger CLOSES
      Diagnostic counter (process-wide): session closes.

  • Method Details

    • tryCreate

      public static VerifyCudaGraphSession tryCreate()
      Creates a session when verify CUDA graph wiring is enabled.
      Returns:
      new session, or null when unavailable.

    • canReplay

      public boolean canReplay(int batch, int windowLen, int numPages)
      Returns whether a graph for (batch, windowLen, numPages) can be replayed.
      Parameters:
      batch - verify batch size.
      windowLen - query length (numDrafts + 1).
      numPages - KV page count for this verify step.
      Returns:
      true when a graph for this bucket can be replayed.

    • shouldCapture

      public boolean shouldCapture(int batch, int windowLen, int numPages, int tpRank)
      Marks one eager warmup step for the current bucket.
      Parameters:
      batch - verify batch size.
      windowLen - query length (numDrafts + 1).
      numPages - KV page count for this verify step.
      tpRank - tensor-parallel rank for logging (>= 0).
      Returns:
      true when the next forward should capture a new graph.

    • beginCapture

      public boolean beginCapture(int deviceIndex)
      Begins CUDA graph capture on deviceIndex (call from TP worker thread).
      Parameters:
      deviceIndex - CUDA device ordinal for capture.
      Returns:
      true when capture started.

    • lastCaptureMs

      public long lastCaptureMs()
      Returns wall time in milliseconds for the last successful capture.
      Returns:
      capture duration in milliseconds.

    • endCapture

      public void endCapture()
      Ends capture and instantiates the graph. No-op when not capturing.

    • logCapture

      public void logCapture(int tpRank)
      Logs (once per bucket) a successful graph capture.
      Parameters:
      tpRank - tensor-parallel rank for logging.

    • replay

      public void replay(int tpRank)
      Replays the captured graph (inputs must already be on device).
      Parameters:
      tpRank - tensor-parallel rank for one-shot bucket logging.

    • capturedBatch

      public int capturedBatch()
      Returns the batch size of the captured bucket.
      Returns:
      batch size, or -1 when none captured.

    • capturedWindowLen

      public int capturedWindowLen()
      Returns the windowLen of the captured bucket.
      Returns:
      window length, or -1 when none captured.

    • capturedNumPages

      public int capturedNumPages()
      Returns the numPages of the captured bucket.
      Returns:
      page count, or -1 when none captured.

    • close

      public void close()
      Specified by:
      close in interface AutoCloseable