Class DecodeCudaGraphSession

java.lang.Object
smile.llm.engine.DecodeCudaGraphSession
All Implemented Interfaces:
AutoCloseable
public final class DecodeCudaGraphSession extends Object implements AutoCloseable
Per-device CUDA graph session for uniform decode (batch 1 or power-of-two B).

Graphs are bucketed by (batch, numPages). Within a bucket, last_page_len and per-row KV indices are updated in place before replay(int).

  • 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 numPages)
    Returns whether a graph for (batch, numPages) can be replayed.
    int
    capturedBatch()
    Returns the batch size of the captured bucket.
    int
    capturedNumPages()
    Returns the numPages 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
    replay(int tpRank)
    Replays the captured graph (inputs must already be on device).
    boolean
    shouldCapture(int batch, int numPages, int tpRank)
    Marks one eager warmup step for the current bucket.
    boolean
    shouldCapture(int batch, int numPages, int tpRank, boolean prefetch)
    Marks one eager warmup step for the current bucket.
    static DecodeCudaGraphSession
    tryCreate()
    Creates a session when decode CUDA graphs are enabled.

    Methods inherited from class Object

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

  • Method Details

    • tryCreate

      public static DecodeCudaGraphSession tryCreate()
      Creates a session when decode CUDA graphs are enabled.
      Returns:
      new session, or null when unavailable.

    • canReplay

      public boolean canReplay(int batch, int numPages)
      Returns whether a graph for (batch, numPages) can be replayed.
      Parameters:
      batch - decode batch size.
      numPages - KV page count for this decode step.
      Returns:
      true when a graph for (batch, numPages) can be replayed.

    • shouldCapture

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

    • shouldCapture

      public boolean shouldCapture(int batch, int numPages, int tpRank, boolean prefetch)
      Marks one eager warmup step for the current bucket.
      Parameters:
      batch - decode batch size.
      numPages - KV page count for this decode step.
      tpRank - tensor-parallel rank for logging (>= 0).
      prefetch - true for next-bucket prefetch capture logging.
      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.

    • 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.

    • 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