Skip to content

GLiFormer: replay the static part of the forward as CUDA graphs #374

Description

@svonava

Problem

After #365, a knowledgator/gliformer-large-v1 extract call on a short text is bound by launching kernels, not by the GPU. For one typed-decision record (40 to 140 tokens) per call on an L4, float16 weights:

  • the whole extract call takes a median of 51 ms, and the model forward takes about 47 ms of it;
  • the GPU is busy for about 8 ms of the forward;
  • the forward launches about 1,270 CUDA kernels. It also synchronizes the stream about 71 times, 31 of them from .item() in the task heads.

The launch and synchronization counts come from torch.profiler over 10 calls. The DeBERTa backbone, the word-level LSTM, and the task heads each launch small kernels one at a time. At this size, each launch costs more on the CPU than its kernel costs on the GPU.

#367 solved the same problem for GLiClass by recording forwards as CUDA graphs and replaying them.

Proposal

Replay the static part of the GLiFormer forward as CUDA graphs, following #367:

  1. Profile which launches can be graphed. The DeBERTa backbone has fixed shapes for a given padded sequence length. The word-level LSTM (_apply_word_rnn) packs sequences by length, which reads the lengths on the host and gives data-dependent shapes, so it may need a padded variant with the same outputs. The task heads have data-dependent control flow and host synchronization: they select thresholded spans, rank relation candidates, and match record anchors. Measure how many of the ~1,270 launches the graphable part accounts for.
  2. Graph that part with the modes from feat(server): replay GLiClass forwards as CUDA graphs #367: exact shapes, or lengths padded to buckets. The heads stay eager and read the graph's static output buffers.
  3. Carry over feat(server): replay GLiClass forwards as CUDA graphs #367's limits: a per-model memory budget, one recording at a time, no recording under memory pressure, and a way to turn graphs off.
  4. Check numerics against eager execution, like the GLiClass parity checks. Scores should be bit-identical in exact mode, and in bucketed mode they should stay within a stated bound.

Acceptance

  • Median extract latency for short GLiFormer inputs drops by a measured amount on an L4, reported with the eager baseline.
  • Outputs in exact mode are bit-identical to eager execution. The adapter's parity tests pass in every mode.
  • Graph memory stays within the budget. It is released when the model unloads, and under memory pressure once Release GLiClass CUDA graph memory under memory pressure before evicting models #370 lands.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions